code.py 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382
  1. #!/usr/bin/env python3
  2. """
  3. s12: 任务系统 — 用文件持久化带 blockedBy 依赖的任务图。
  4. 运行: python s12_task_system/code.py
  5. 需要: pip install anthropic python-dotenv + .env 中配置 ANTHROPIC_API_KEY
  6. 相对 s11 的变化:
  7. - 任务 dataclass(id、subject、description、status、owner、blockedBy)
  8. - TASKS_DIR = .tasks/,用于持久化 JSON 存储
  9. - create_task / save_task / load_task / list_tasks / get_task
  10. - can_start:检查 blockedBy 是否全部完成(缺失依赖 = 被阻塞)
  11. - claim_task:设置 owner + pending -> in_progress
  12. - complete_task:设置 completed + 报告下游解锁任务
  13. - 5 个新工具:create_task、list_tasks、get_task、claim_task、complete_task
  14. 说明:教学代码保留一个基础 Agent 循环,以便聚焦任务系统。
  15. S11 的完整错误恢复(RecoveryState、退避、升级、reactive compact、备用模型)被省略;
  16. 在真实 CC 中,tasks.ts 和 withRetry 是可以自然组合的独立层。
  17. """
  18. import os, subprocess, json, time, random
  19. from pathlib import Path
  20. from dataclasses import dataclass, asdict
  21. try:
  22. import readline
  23. readline.parse_and_bind('set bind-tty-special-chars off')
  24. except ImportError:
  25. pass
  26. from anthropic import Anthropic
  27. from dotenv import load_dotenv
  28. load_dotenv(override=True)
  29. if os.getenv("ANTHROPIC_BASE_URL"):
  30. os.environ.pop("ANTHROPIC_AUTH_TOKEN", None)
  31. WORKDIR = Path.cwd()
  32. MEMORY_DIR = WORKDIR / ".memory"
  33. MEMORY_INDEX = MEMORY_DIR / "MEMORY.md"
  34. client = Anthropic(base_url=os.getenv("ANTHROPIC_BASE_URL"))
  35. MODEL = os.environ["MODEL_ID"]
  36. # ── 任务系统 ──
  37. TASKS_DIR = WORKDIR / ".tasks"
  38. TASKS_DIR.mkdir(exist_ok=True)
  39. @dataclass
  40. class Task:
  41. id: str
  42. subject: str
  43. description: str
  44. status: str # pending | in_progress | completed
  45. owner: str | None # Agent 名称(多 Agent 场景)
  46. blockedBy: list[str] # 依赖任务 ID
  47. def _task_path(task_id: str) -> Path:
  48. return TASKS_DIR / f"{task_id}.json"
  49. def create_task(subject: str, description: str = "",
  50. blockedBy: list[str] | None = None) -> Task:
  51. task = Task(
  52. id=f"task_{int(time.time())}_{random.randint(0, 9999):04d}",
  53. subject=subject,
  54. description=description,
  55. status="pending",
  56. owner=None,
  57. blockedBy=blockedBy or [],
  58. )
  59. save_task(task)
  60. return task
  61. def save_task(task: Task):
  62. _task_path(task.id).write_text(json.dumps(asdict(task), indent=2))
  63. def load_task(task_id: str) -> Task:
  64. return Task(**json.loads(_task_path(task_id).read_text()))
  65. def list_tasks() -> list[Task]:
  66. return [Task(**json.loads(p.read_text()))
  67. for p in sorted(TASKS_DIR.glob("task_*.json"))]
  68. def get_task(task_id: str) -> str:
  69. """以 JSON 返回完整任务详情。"""
  70. task = load_task(task_id)
  71. return json.dumps(asdict(task), indent=2)
  72. def can_start(task_id: str) -> bool:
  73. """检查所有 blockedBy 依赖是否已完成;缺失依赖视为阻塞。"""
  74. task = load_task(task_id)
  75. for dep_id in task.blockedBy:
  76. if not _task_path(dep_id).exists():
  77. return False
  78. if load_task(dep_id).status != "completed":
  79. return False
  80. return True
  81. def claim_task(task_id: str, owner: str = "agent") -> str:
  82. task = load_task(task_id)
  83. if task.status != "pending":
  84. status = {"pending": "待处理", "in_progress": "进行中",
  85. "completed": "已完成"}.get(task.status, task.status)
  86. return f"任务 {task_id} 当前状态为 {status},无法认领"
  87. if not can_start(task_id):
  88. deps = [d for d in task.blockedBy
  89. if not _task_path(d).exists() or load_task(d).status != "completed"]
  90. return f"被阻塞于:{deps}"
  91. task.owner = owner
  92. task.status = "in_progress"
  93. save_task(task)
  94. print(f" \033[36m[认领] {task.subject} → in_progress(负责人:{owner})\033[0m")
  95. return f"已认领 {task.id} ({task.subject})"
  96. def complete_task(task_id: str) -> str:
  97. task = load_task(task_id)
  98. if task.status != "in_progress":
  99. status = {"pending": "待处理", "in_progress": "进行中",
  100. "completed": "已完成"}.get(task.status, task.status)
  101. return f"任务 {task_id} 当前状态为 {status},无法完成"
  102. task.status = "completed"
  103. save_task(task)
  104. unblocked = [t.subject for t in list_tasks()
  105. if t.status == "pending" and t.blockedBy and can_start(t.id)]
  106. print(f" \033[32m[完成] {task.subject} ✓\033[0m")
  107. msg = f"已完成 {task.id} ({task.subject})"
  108. if unblocked:
  109. msg += f"\n已解除阻塞:{', '.join(unblocked)}"
  110. print(f" \033[33m[已解除阻塞] {', '.join(unblocked)}\033[0m")
  111. return msg
  112. # ── 提示词组装 (来自 s10,已同步) ──
  113. PROMPT_SECTIONS = {
  114. "identity": "你是一个编码 Agent。直接行动,不要只解释。",
  115. "tools": "可用工具:bash, read_file, write_file, "
  116. "create_task, list_tasks, get_task, claim_task, complete_task.",
  117. "workspace": f"工作目录:{WORKDIR}",
  118. "memory": "有可用的相关记忆时,会在下方注入。",
  119. }
  120. def assemble_system_prompt(context: dict) -> str:
  121. sections = [PROMPT_SECTIONS["identity"],
  122. PROMPT_SECTIONS["tools"],
  123. PROMPT_SECTIONS["workspace"]]
  124. memories = context.get("memories", "")
  125. if memories:
  126. sections.append(f"相关记忆:\n{memories}")
  127. return "\n\n".join(sections)
  128. _last_context_key, _last_prompt = None, None
  129. def get_system_prompt(context: dict) -> str:
  130. global _last_context_key, _last_prompt
  131. key = json.dumps(context, sort_keys=True, ensure_ascii=False, default=str)
  132. if key == _last_context_key and _last_prompt:
  133. return _last_prompt
  134. _last_context_key = key
  135. _last_prompt = assemble_system_prompt(context)
  136. return _last_prompt
  137. # ── 工具 ──
  138. def safe_path(p: str) -> Path:
  139. path = (WORKDIR / p).resolve()
  140. if not path.is_relative_to(WORKDIR):
  141. raise ValueError(f"路径逃逸出工作区:{p}")
  142. return path
  143. def run_bash(command: str) -> str:
  144. try:
  145. r = subprocess.run(command, shell=True, cwd=WORKDIR,
  146. capture_output=True, text=True, timeout=120)
  147. out = (r.stdout + r.stderr).strip()
  148. return out[:50000] if out else "(无输出)"
  149. except subprocess.TimeoutExpired:
  150. return "错误:执行超时(120 秒)"
  151. def run_read(path: str, limit: int | None = None) -> str:
  152. try:
  153. lines = safe_path(path).read_text().splitlines()
  154. if limit and limit < len(lines):
  155. lines = lines[:limit] + [f"... ({len(lines) - limit} 行更多内容)"]
  156. return "\n".join(lines)
  157. except Exception as e:
  158. return f"错误:{e}"
  159. def run_write(path: str, content: str) -> str:
  160. try:
  161. fp = safe_path(path)
  162. fp.parent.mkdir(parents=True, exist_ok=True)
  163. fp.write_text(content)
  164. return f"已写入 {len(content)} 字节到 {path}"
  165. except Exception as e:
  166. return f"错误:{e}"
  167. # 任务工具
  168. def run_create_task(subject: str, description: str = "",
  169. blockedBy: list[str] | None = None) -> str:
  170. task = create_task(subject, description, blockedBy)
  171. deps = f"(依赖:{', '.join(blockedBy)})" if blockedBy else ""
  172. print(f" \033[34m[创建] {task.subject}{deps}\033[0m")
  173. return f"已创建 {task.id}: {task.subject}{deps}"
  174. def run_list_tasks() -> str:
  175. 个任务 = list_tasks()
  176. if not 个任务:
  177. return "暂无任务。请使用 create_task 添加任务。"
  178. lines = []
  179. for t in 个任务:
  180. icon = {"pending": "○", "in_progress": "●",
  181. "completed": "✓"}.get(t.status, "?")
  182. deps = f"(依赖:{', '.join(t.blockedBy)})" if t.blockedBy else ""
  183. owner = f" [负责人:{t.owner}]" if t.owner else ""
  184. status = {"pending": "待处理", "in_progress": "进行中",
  185. "completed": "已完成"}.get(t.status, t.status)
  186. lines.append(f" {icon} {t.id}: {t.subject} "
  187. f"[{status}]{owner}{deps}")
  188. return "\n".join(lines)
  189. def run_get_task(task_id: str) -> str:
  190. try:
  191. return get_task(task_id)
  192. except FileNotFoundError:
  193. return f"错误:任务 {task_id} 未找到"
  194. def run_claim_task(task_id: str) -> str:
  195. return claim_task(task_id, owner="agent")
  196. def run_complete_task(task_id: str) -> str:
  197. return complete_task(task_id)
  198. TOOLS = [
  199. {"name": "bash", "description": "运行一条 shell 命令。",
  200. "input_schema": {"type": "object",
  201. "properties": {"command": {"type": "string"}},
  202. "required": ["command"]}},
  203. {"name": "read_file", "description": "读取文件内容。",
  204. "input_schema": {"type": "object",
  205. "properties": {"path": {"type": "string"},
  206. "limit": {"type": "integer"}},
  207. "required": ["path"]}},
  208. {"name": "write_file", "description": "向文件写入内容。",
  209. "input_schema": {"type": "object",
  210. "properties": {"path": {"type": "string"},
  211. "content": {"type": "string"}},
  212. "required": ["path", "content"]}},
  213. {"name": "create_task",
  214. "description": "创建一个新任务,可选 blockedBy 依赖。",
  215. "input_schema": {"type": "object",
  216. "properties": {
  217. "subject": {"type": "string"},
  218. "description": {"type": "string"},
  219. "blockedBy": {"type": "array",
  220. "items": {"type": "string"}}},
  221. "required": ["subject"]}},
  222. {"name": "list_tasks",
  223. "description": "列出所有任务及其状态、负责人和依赖。",
  224. "input_schema": {"type": "object", "properties": {},
  225. "required": []}},
  226. {"name": "get_task",
  227. "description": "按 ID 获取指定任务的完整详情。",
  228. "input_schema": {"type": "object",
  229. "properties": {"task_id": {"type": "string"}},
  230. "required": ["task_id"]}},
  231. {"name": "claim_task",
  232. "description": "认领一个待处理任务。设置 owner,并将状态改为 in_progress。",
  233. "input_schema": {"type": "object",
  234. "properties": {"task_id": {"type": "string"}},
  235. "required": ["task_id"]}},
  236. {"name": "complete_task",
  237. "description": "完成一个进行中的任务。报告被解除阻塞的下游任务。",
  238. "input_schema": {"type": "object",
  239. "properties": {"task_id": {"type": "string"}},
  240. "required": ["task_id"]}},
  241. ]
  242. TOOL_HANDLERS = {
  243. "bash": run_bash, "read_file": run_read, "write_file": run_write,
  244. "create_task": run_create_task, "list_tasks": run_list_tasks,
  245. "get_task": run_get_task, "claim_task": run_claim_task,
  246. "complete_task": run_complete_task,
  247. }
  248. # ── 上下文 ──
  249. def update_context(context: dict, messages: list) -> dict:
  250. """从真实状态推导上下文。"""
  251. memories = ""
  252. if MEMORY_INDEX.exists():
  253. content = MEMORY_INDEX.read_text().strip()
  254. if content:
  255. memories = content
  256. return {
  257. "enabled_tools": list(TOOL_HANDLERS.keys()),
  258. "workspace": str(WORKDIR),
  259. "memories": memories,
  260. }
  261. # ── Agent 循环(简化版,聚焦任务系统) ──
  262. def agent_loop(messages: list, context: dict):
  263. system = get_system_prompt(context)
  264. while True:
  265. try:
  266. response = client.messages.create(
  267. model=MODEL, system=system, messages=messages,
  268. tools=TOOLS, max_tokens=8000)
  269. except Exception as e:
  270. messages.append({"role": "assistant", "content": [
  271. {"type": "text",
  272. "text": f"[错误] {type(e).__name__}: {e}"}]})
  273. return
  274. messages.append({"role": "assistant", "content": response.content})
  275. if response.stop_reason != "tool_use":
  276. return
  277. results = []
  278. for block in response.content:
  279. if block.type != "tool_use":
  280. continue
  281. print(f"\033[36m> {block.name}\033[0m")
  282. handler = TOOL_HANDLERS.get(block.name)
  283. output = handler(**block.input) if handler else f"未知工具:{block.name}"
  284. print(str(output)[:300])
  285. results.append({"type": "tool_result",
  286. "tool_use_id": block.id, "content": output})
  287. messages.append({"role": "user", "content": results})
  288. context = update_context(context, messages)
  289. system = get_system_prompt(context)
  290. if __name__ == "__main__":
  291. print("s12: 任务系统")
  292. print("输入问题后按回车发送。输入 q 退出。\n")
  293. history = []
  294. context = update_context({}, [])
  295. while True:
  296. try:
  297. query = input("\033[36ms12 >> \033[0m")
  298. except (EOFError, KeyboardInterrupt):
  299. break
  300. if query.strip().lower() in ("q", "exit", ""):
  301. break
  302. history.append({"role": "user", "content": query})
  303. agent_loop(history, context)
  304. context = update_context(context, history)
  305. for block in history[-1]["content"]:
  306. if getattr(block, "type", None) == "text":
  307. print(block.text)
  308. elif isinstance(block, dict) and block.get("type") == "text":
  309. print(block.get("text", ""))
  310. print()