216 lines
6.0 KiB
Python
216 lines
6.0 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
Agent 任务轮询脚本
|
||
每个 Agent 定时运行，查询分配给自己的任务并执行
|
||
"""
|
||
|
||
import os
|
||
import sys
|
||
import time
|
||
import logging
|
||
import requests
|
||
from datetime import datetime
|
||
|
||
# 配置日志
|
||
logging.basicConfig(
|
||
level=logging.INFO,
|
||
format='%(asctime)s - %(levelname)s - %(message)s'
|
||
)
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# 环境变量或配置
|
||
NOTION_TOKEN = os.environ.get("NOTION_TOKEN", "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")
|
||
AGENT_ID = os.environ.get("AGENT_ID", "yunjiang") # 当前 Agent ID
|
||
POLL_INTERVAL = int(os.environ.get("POLL_INTERVAL", "180")) # 轮询间隔（秒），默认3分钟
|
||
|
||
# Database IDs
|
||
TASKS_DB_ID = "32847fe1-da27-8135-af44-eefdbd3b1640"
|
||
AGENTS_DB_ID = "32847fe1-da27-8101-8758-d416db87d4de"
|
||
|
||
# Notion API 基础 URL
|
||
NOTION_API_BASE = "https://api.notion.com/v1"
|
||
|
||
|
||
def notion_request(method, endpoint, **kwargs):
|
||
"""Notion API 请求封装"""
|
||
url = f"{NOTION_API_BASE}{endpoint}"
|
||
headers = {
|
||
"Authorization": f"Bearer {NOTION_TOKEN}",
|
||
"Notion-Version": "2022-06-28",
|
||
"Content-Type": "application/json"
|
||
}
|
||
|
||
if method == "GET":
|
||
response = requests.get(url, headers=headers, **kwargs)
|
||
elif method == "POST":
|
||
response = requests.post(url, headers=headers, **kwargs)
|
||
elif method == "PATCH":
|
||
response = requests.patch(url, headers=headers, **kwargs)
|
||
else:
|
||
raise ValueError(f"Unsupported method: {method}")
|
||
|
||
response.raise_for_status()
|
||
return response.json()
|
||
|
||
|
||
def get_agent_info(agent_id):
|
||
"""获取 Agent 信息"""
|
||
response = notion_request("POST", f"/databases/{AGENTS_DB_ID}/query")
|
||
|
||
for page in response.get("results", []):
|
||
props = page.get("properties", {})
|
||
if "Agent ID" in props:
|
||
title = props["Agent ID"]["title"]
|
||
if title and title[0]["plain_text"] == agent_id:
|
||
return {
|
||
"id": page["id"],
|
||
"name": props["名称"]["rich_text"][0]["plain_text"] if props["名称"]["rich_text"] else agent_id,
|
||
"status": props["状态"]["select"]["name"] if props["状态"].get("select") else "离线"
|
||
}
|
||
return None
|
||
|
||
|
||
def query_todo_tasks(agent_page_id=None):
|
||
"""查询 TODO 任务"""
|
||
logger.info(f"查询 {AGENT_ID} 的 TODO 任务...")
|
||
|
||
# 如果有 agent_page_id，使用 Relation 过滤
|
||
filter_dict = {
|
||
"property": "状态",
|
||
"select": {
|
||
"equals": "TODO"
|
||
}
|
||
}
|
||
|
||
# 这里简化处理：查询所有 TODO 任务
|
||
# 后续应该根据执行者 Relation 过滤
|
||
|
||
try:
|
||
response = notion_request(
|
||
"POST",
|
||
f"/databases/{TASKS_DB_ID}/query",
|
||
json={"filter": filter_dict}
|
||
)
|
||
tasks = response.get("results", [])
|
||
logger.info(f"找到 {len(tasks)} 个 TODO 任务")
|
||
return tasks
|
||
except Exception as e:
|
||
logger.error(f"查询任务失败: {e}")
|
||
return []
|
||
|
||
|
||
def claim_task(task_id):
|
||
"""领取任务：将状态改为进行中"""
|
||
try:
|
||
notion_request(
|
||
"PATCH",
|
||
f"/pages/{task_id}",
|
||
json={
|
||
"properties": {
|
||
"状态": {
|
||
"select": {"name": "进行中"}
|
||
}
|
||
}
|
||
}
|
||
)
|
||
logger.info(f"✓ 领取任务成功: {task_id}")
|
||
return True
|
||
except Exception as e:
|
||
logger.error(f"领取任务失败: {e}")
|
||
return False
|
||
|
||
|
||
def complete_task(task_id, report_link):
|
||
"""完成任务：将状态改为待验收"""
|
||
try:
|
||
notion_request(
|
||
"PATCH",
|
||
f"/pages/{task_id}",
|
||
json={
|
||
"properties": {
|
||
"状态": {
|
||
"select": {"name": "待验收"}
|
||
},
|
||
"报告链接": {
|
||
"url": report_link
|
||
}
|
||
}
|
||
}
|
||
)
|
||
logger.info(f"✓ 完成任务: {task_id}")
|
||
return True
|
||
except Exception as e:
|
||
logger.error(f"完成任务失败: {e}")
|
||
return False
|
||
|
||
|
||
def execute_task(task):
|
||
"""执行任务的逻辑（可自定义）"""
|
||
# 获取任务信息
|
||
props = task.get("properties", {})
|
||
task_name = props.get("任务名", {}).get("title", [{}])[0].get("plain_text", "未命名任务")
|
||
task_id = task["id"]
|
||
|
||
logger.info(f"开始执行任务: {task_name}")
|
||
|
||
# 1. 领取任务
|
||
if not claim_task(task_id):
|
||
return False
|
||
|
||
# 2. 执行任务（这里只是示例，实际应根据任务类型执行不同操作）
|
||
# 模拟执行
|
||
time.sleep(2)
|
||
|
||
# 3. 完成任务（生成报告链接）
|
||
# 这里应该生成实际的 Obsidian 报告
|
||
report_link = f"https://example.com/report/{task_id}"
|
||
|
||
return complete_task(task_id, report_link)
|
||
|
||
|
||
def polling_loop():
|
||
"""轮询主循环"""
|
||
logger.info(f"🚀 Agent {AGENT_ID} 任务轮询启动")
|
||
logger.info(f"轮询间隔: {POLL_INTERVAL} 秒")
|
||
|
||
# 获取 Agent 信息
|
||
agent_info = get_agent_info(AGENT_ID)
|
||
if agent_info:
|
||
logger.info(f"Agent 信息: {agent_info['name']} (状态: {agent_info['status']})")
|
||
else:
|
||
logger.warning(f"未找到 Agent: {AGENT_ID}")
|
||
|
||
while True:
|
||
try:
|
||
# 查询 TODO 任务
|
||
tasks = query_todo_tasks(None)
|
||
|
||
if tasks:
|
||
logger.info(f"发现 {len(tasks)} 个待处理任务")
|
||
for task in tasks:
|
||
execute_task(task)
|
||
else:
|
||
logger.debug("没有待处理任务")
|
||
|
||
except Exception as e:
|
||
logger.error(f"轮询异常: {e}")
|
||
|
||
time.sleep(POLL_INTERVAL)
|
||
|
||
|
||
def main():
|
||
"""主入口"""
|
||
if len(sys.argv) > 1:
|
||
global AGENT_ID
|
||
AGENT_ID = sys.argv[1]
|
||
|
||
if len(sys.argv) > 2:
|
||
global POLL_INTERVAL
|
||
POLL_INTERVAL = int(sys.argv[2])
|
||
|
||
polling_loop()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|