code.py 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377
  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. """Check if all blockedBy dependencies are completed.
  74. Missing dependencies are treated as blocked."""
  75. task = load_task(task_id)
  76. for dep_id in task.blockedBy:
  77. if not _task_path(dep_id).exists():
  78. return False
  79. if load_task(dep_id).status != "completed":
  80. return False
  81. return True
  82. def claim_task(task_id: str, owner: str = "agent") -> str:
  83. task = load_task(task_id)
  84. if task.status != "pending":
  85. return f"任务 {task_id} 当前状态为 {task.status},无法认领"
  86. if not can_start(task_id):
  87. deps = [d for d in task.blockedBy
  88. if not _task_path(d).exists() or load_task(d).status != "completed"]
  89. return f"Blocked by: {deps}"
  90. task.owner = owner
  91. task.status = "in_progress"
  92. save_task(task)
  93. print(f" \033[36m[claim] {task.subject} → in_progress (owner: {owner})\033[0m")
  94. return f"已认领 {task.id} ({task.subject})"
  95. def complete_task(task_id: str) -> str:
  96. task = load_task(task_id)
  97. if task.status != "in_progress":
  98. return f"任务 {task_id} 当前状态为 {task.status},无法完成"
  99. task.status = "completed"
  100. save_task(task)
  101. unblocked = [t.subject for t in list_tasks()
  102. if t.status == "pending" and t.blockedBy and can_start(t.id)]
  103. print(f" \033[32m[complete] {task.subject} ✓\033[0m")
  104. msg = f"已完成 {task.id} ({task.subject})"
  105. if unblocked:
  106. msg += f"\n已解除阻塞:{', '.join(unblocked)}"
  107. print(f" \033[33m[unblocked] {', '.join(unblocked)}\033[0m")
  108. return msg
  109. # ── 提示词组装 (来自 s10,已同步) ──
  110. PROMPT_SECTIONS = {
  111. "identity": "你是一个编码 Agent。直接行动,不要只解释。",
  112. "tools": "可用工具:bash, read_file, write_file, "
  113. "create_task, list_tasks, get_task, claim_task, complete_task.",
  114. "workspace": f"工作目录:{WORKDIR}",
  115. "memory": "有可用的相关记忆时,会在下方注入。",
  116. }
  117. def assemble_system_prompt(context: dict) -> str:
  118. sections = [PROMPT_SECTIONS["identity"],
  119. PROMPT_SECTIONS["tools"],
  120. PROMPT_SECTIONS["workspace"]]
  121. memories = context.get("memories", "")
  122. if memories:
  123. sections.append(f"Relevant memories:\n{memories}")
  124. return "\n\n".join(sections)
  125. _last_context_key, _last_prompt = None, None
  126. def get_system_prompt(context: dict) -> str:
  127. global _last_context_key, _last_prompt
  128. key = json.dumps(context, sort_keys=True, ensure_ascii=False, default=str)
  129. if key == _last_context_key and _last_prompt:
  130. return _last_prompt
  131. _last_context_key = key
  132. _last_prompt = assemble_system_prompt(context)
  133. return _last_prompt
  134. # ── 工具 ──
  135. def safe_path(p: str) -> Path:
  136. path = (WORKDIR / p).resolve()
  137. if not path.is_relative_to(WORKDIR):
  138. raise ValueError(f"路径逃逸出工作区:{p}")
  139. return path
  140. def run_bash(command: str) -> str:
  141. try:
  142. r = subprocess.run(command, shell=True, cwd=WORKDIR,
  143. capture_output=True, text=True, timeout=120)
  144. out = (r.stdout + r.stderr).strip()
  145. return out[:50000] if out else "(无输出)"
  146. except subprocess.TimeoutExpired:
  147. return "错误:执行超时(120 秒)"
  148. def run_read(path: str, limit: int | None = None) -> str:
  149. try:
  150. lines = safe_path(path).read_text().splitlines()
  151. if limit and limit < len(lines):
  152. lines = lines[:limit] + [f"... ({len(lines) - limit} 行更多内容)"]
  153. return "\n".join(lines)
  154. except Exception as e:
  155. return f"错误:{e}"
  156. def run_write(path: str, content: str) -> str:
  157. try:
  158. fp = safe_path(path)
  159. fp.parent.mkdir(parents=True, exist_ok=True)
  160. fp.write_text(content)
  161. return f"已写入 {len(content)} 字节到 {path}"
  162. except Exception as e:
  163. return f"错误:{e}"
  164. # 任务工具
  165. def run_create_task(subject: str, description: str = "",
  166. blockedBy: list[str] | None = None) -> str:
  167. task = create_task(subject, description, blockedBy)
  168. deps = f" (blockedBy: {', '.join(blockedBy)})" if blockedBy else ""
  169. print(f" \033[34m[create] {task.subject}{deps}\033[0m")
  170. return f"已创建 {task.id}: {task.subject}{deps}"
  171. def run_list_tasks() -> str:
  172. 个任务 = list_tasks()
  173. if not 个任务:
  174. return "暂无任务。请使用 create_task 添加任务。"
  175. lines = []
  176. for t in 个任务:
  177. icon = {"pending": "○", "in_progress": "●",
  178. "completed": "✓"}.get(t.status, "?")
  179. deps = f" (blockedBy: {', '.join(t.blockedBy)})" if t.blockedBy else ""
  180. owner = f" [{t.owner}]" if t.owner else ""
  181. lines.append(f" {icon} {t.id}: {t.subject} "
  182. f"[{t.status}]{owner}{deps}")
  183. return "\n".join(lines)
  184. def run_get_task(task_id: str) -> str:
  185. try:
  186. return get_task(task_id)
  187. except FileNotFoundError:
  188. return f"错误:任务 {task_id} 未找到"
  189. def run_claim_task(task_id: str) -> str:
  190. return claim_task(task_id, owner="agent")
  191. def run_complete_task(task_id: str) -> str:
  192. return complete_task(task_id)
  193. TOOLS = [
  194. {"name": "bash", "description": "运行一条 shell 命令。",
  195. "input_schema": {"type": "object",
  196. "properties": {"command": {"type": "string"}},
  197. "required": ["command"]}},
  198. {"name": "read_file", "description": "读取文件内容。",
  199. "input_schema": {"type": "object",
  200. "properties": {"path": {"type": "string"},
  201. "limit": {"type": "integer"}},
  202. "required": ["path"]}},
  203. {"name": "write_file", "description": "向文件写入内容。",
  204. "input_schema": {"type": "object",
  205. "properties": {"path": {"type": "string"},
  206. "content": {"type": "string"}},
  207. "required": ["path", "content"]}},
  208. {"name": "create_task",
  209. "description": "创建一个新任务,可选 blockedBy 依赖。",
  210. "input_schema": {"type": "object",
  211. "properties": {
  212. "subject": {"type": "string"},
  213. "description": {"type": "string"},
  214. "blockedBy": {"type": "array",
  215. "items": {"type": "string"}}},
  216. "required": ["subject"]}},
  217. {"name": "list_tasks",
  218. "description": "列出所有任务及其状态、负责人和依赖。",
  219. "input_schema": {"type": "object", "properties": {},
  220. "required": []}},
  221. {"name": "get_task",
  222. "description": "按 ID 获取指定任务的完整详情。",
  223. "input_schema": {"type": "object",
  224. "properties": {"task_id": {"type": "string"}},
  225. "required": ["task_id"]}},
  226. {"name": "claim_task",
  227. "description": "认领一个待处理任务。设置 owner,并将状态改为 in_progress。",
  228. "input_schema": {"type": "object",
  229. "properties": {"task_id": {"type": "string"}},
  230. "required": ["task_id"]}},
  231. {"name": "complete_task",
  232. "description": "完成一个进行中的任务。报告被解除阻塞的下游任务。",
  233. "input_schema": {"type": "object",
  234. "properties": {"task_id": {"type": "string"}},
  235. "required": ["task_id"]}},
  236. ]
  237. TOOL_HANDLERS = {
  238. "bash": run_bash, "read_file": run_read, "write_file": run_write,
  239. "create_task": run_create_task, "list_tasks": run_list_tasks,
  240. "get_task": run_get_task, "claim_task": run_claim_task,
  241. "complete_task": run_complete_task,
  242. }
  243. # ── 上下文 ──
  244. def update_context(context: dict, messages: list) -> dict:
  245. """Derive context from real state."""
  246. memories = ""
  247. if MEMORY_INDEX.exists():
  248. content = MEMORY_INDEX.read_text().strip()
  249. if content:
  250. memories = content
  251. return {
  252. "enabled_tools": list(TOOL_HANDLERS.keys()),
  253. "workspace": str(WORKDIR),
  254. "memories": memories,
  255. }
  256. # ── Agent 循环(简化版,聚焦任务系统) ──
  257. def agent_loop(messages: list, context: dict):
  258. system = get_system_prompt(context)
  259. while True:
  260. try:
  261. response = client.messages.create(
  262. model=MODEL, system=system, messages=messages,
  263. tools=TOOLS, max_tokens=8000)
  264. except Exception as e:
  265. messages.append({"role": "assistant", "content": [
  266. {"type": "text",
  267. "text": f"[错误] {type(e).__name__}: {e}"}]})
  268. return
  269. messages.append({"role": "assistant", "content": response.content})
  270. if response.stop_reason != "tool_use":
  271. return
  272. results = []
  273. for block in response.content:
  274. if block.type != "tool_use":
  275. continue
  276. print(f"\033[36m> {block.name}\033[0m")
  277. handler = TOOL_HANDLERS.get(block.name)
  278. output = handler(**block.input) if handler else f"未知工具:{block.name}"
  279. print(str(output)[:300])
  280. results.append({"type": "tool_result",
  281. "tool_use_id": block.id, "content": output})
  282. messages.append({"role": "user", "content": results})
  283. context = update_context(context, messages)
  284. system = get_system_prompt(context)
  285. if __name__ == "__main__":
  286. print("s12: 任务系统")
  287. print("输入问题后按回车发送。输入 q 退出。\n")
  288. history = []
  289. context = update_context({}, [])
  290. while True:
  291. try:
  292. query = input("\033[36ms12 >> \033[0m")
  293. except (EOFError, KeyboardInterrupt):
  294. break
  295. if query.strip().lower() in ("q", "exit", ""):
  296. break
  297. history.append({"role": "user", "content": query})
  298. agent_loop(history, context)
  299. context = update_context(context, history)
  300. for block in history[-1]["content"]:
  301. if getattr(block, "type", None) == "text":
  302. print(block.text)
  303. elif isinstance(block, dict) and block.get("type") == "text":
  304. print(block.get("text", ""))
  305. print()