init_agent.py 9.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279
  1. #!/usr/bin/env python3
  2. """
  3. Agent Scaffold Script - Create a new agent project with best practices.
  4. Usage:
  5. python init_agent.py <agent-name> [--level 0-4] [--path <output-dir>]
  6. Examples:
  7. python init_agent.py my-agent # Level 1 (4 tools)
  8. python init_agent.py my-agent --level 0 # Minimal (bash only)
  9. python init_agent.py my-agent --level 2 # With TodoWrite
  10. python init_agent.py my-agent --path ./bots # Custom output directory
  11. """
  12. import argparse
  13. import sys
  14. from pathlib import Path
  15. # Agent templates for each level
  16. TEMPLATES = {
  17. 0: '''#!/usr/bin/env python3
  18. """
  19. Level 0 Agent - Bash is All You Need (~50 lines)
  20. Core insight: One tool (bash) can do everything.
  21. Subagents via self-recursion: python {name}.py "subtask"
  22. """
  23. from anthropic import Anthropic
  24. from dotenv import load_dotenv
  25. import subprocess
  26. import os
  27. load_dotenv()
  28. client = Anthropic(
  29. api_key=os.getenv("ANTHROPIC_API_KEY"),
  30. base_url=os.getenv("ANTHROPIC_BASE_URL")
  31. )
  32. MODEL = os.getenv("MODEL_NAME", "claude-sonnet-4-20250514")
  33. SYSTEM = """You are a coding agent. Use bash for everything:
  34. - Read: cat, grep, find, ls
  35. - Write: echo 'content' > file
  36. - Subagent: python {name}.py "subtask"
  37. """
  38. TOOL = [{{
  39. "name": "bash",
  40. "description": "Execute shell command",
  41. "input_schema": {{"type": "object", "properties": {{"command": {{"type": "string"}}}}, "required": ["command"]}}
  42. }}]
  43. def run(prompt, history=[]):
  44. history.append({{"role": "user", "content": prompt}})
  45. while True:
  46. r = client.messages.create(model=MODEL, system=SYSTEM, messages=history, tools=TOOL, max_tokens=8000)
  47. history.append({{"role": "assistant", "content": r.content}})
  48. if r.stop_reason != "tool_use":
  49. return "".join(b.text for b in r.content if hasattr(b, "text"))
  50. results = []
  51. for b in r.content:
  52. if b.type == "tool_use":
  53. print(f"> {{b.input['command']}}")
  54. try:
  55. out = subprocess.run(b.input["command"], shell=True, capture_output=True, text=True, timeout=60)
  56. output = (out.stdout + out.stderr).strip() or "(empty)"
  57. except Exception as e:
  58. output = f"Error: {{e}}"
  59. results.append({{"type": "tool_result", "tool_use_id": b.id, "content": output[:50000]}})
  60. history.append({{"role": "user", "content": results}})
  61. if __name__ == "__main__":
  62. h = []
  63. print("{name} - Level 0 Agent\\nType 'q' to quit.\\n")
  64. while (q := input(">> ").strip()) not in ("q", "quit", ""):
  65. print(run(q, h), "\\n")
  66. ''',
  67. 1: '''#!/usr/bin/env python3
  68. """
  69. Level 1 Agent - Model as Agent (~200 lines)
  70. Core insight: 4 tools cover 90% of coding tasks.
  71. The model IS the agent. Code just runs the loop.
  72. """
  73. from anthropic import Anthropic
  74. from dotenv import load_dotenv
  75. from pathlib import Path
  76. import subprocess
  77. import os
  78. load_dotenv()
  79. client = Anthropic(
  80. api_key=os.getenv("ANTHROPIC_API_KEY"),
  81. base_url=os.getenv("ANTHROPIC_BASE_URL")
  82. )
  83. MODEL = os.getenv("MODEL_NAME", "claude-sonnet-4-20250514")
  84. WORKDIR = Path.cwd()
  85. SYSTEM = f"""You are a coding agent at {{WORKDIR}}.
  86. Rules:
  87. - Prefer tools over prose. Act, don't just explain.
  88. - Never invent file paths. Use ls/find first if unsure.
  89. - Make minimal changes. Don't over-engineer.
  90. - After finishing, summarize what changed."""
  91. TOOLS = [
  92. {{"name": "bash", "description": "Run shell command",
  93. "input_schema": {{"type": "object", "properties": {{"command": {{"type": "string"}}}}, "required": ["command"]}}}},
  94. {{"name": "read_file", "description": "Read file contents",
  95. "input_schema": {{"type": "object", "properties": {{"path": {{"type": "string"}}}}, "required": ["path"]}}}},
  96. {{"name": "write_file", "description": "Write content to file",
  97. "input_schema": {{"type": "object", "properties": {{"path": {{"type": "string"}}, "content": {{"type": "string"}}}}, "required": ["path", "content"]}}}},
  98. {{"name": "edit_file", "description": "Replace exact text in file",
  99. "input_schema": {{"type": "object", "properties": {{"path": {{"type": "string"}}, "old_text": {{"type": "string"}}, "new_text": {{"type": "string"}}}}, "required": ["path", "old_text", "new_text"]}}}},
  100. ]
  101. def safe_path(p: str) -> Path:
  102. """Prevent path escape attacks."""
  103. path = (WORKDIR / p).resolve()
  104. if not path.is_relative_to(WORKDIR):
  105. raise ValueError(f"Path escapes workspace: {{p}}")
  106. return path
  107. def execute(name: str, args: dict) -> str:
  108. """Execute a tool and return result."""
  109. if name == "bash":
  110. dangerous = ["rm -rf /", "sudo", "shutdown", "> /dev/"]
  111. if any(d in args["command"] for d in dangerous):
  112. return "Error: Dangerous command blocked"
  113. try:
  114. r = subprocess.run(args["command"], shell=True, cwd=WORKDIR, capture_output=True, text=True, timeout=60)
  115. return (r.stdout + r.stderr).strip()[:50000] or "(empty)"
  116. except subprocess.TimeoutExpired:
  117. return "Error: Timeout (60s)"
  118. except Exception as e:
  119. return f"Error: {{e}}"
  120. if name == "read_file":
  121. try:
  122. return safe_path(args["path"]).read_text()[:50000]
  123. except Exception as e:
  124. return f"Error: {{e}}"
  125. if name == "write_file":
  126. try:
  127. p = safe_path(args["path"])
  128. p.parent.mkdir(parents=True, exist_ok=True)
  129. p.write_text(args["content"])
  130. return f"Wrote {{len(args['content'])}} bytes to {{args['path']}}"
  131. except Exception as e:
  132. return f"Error: {{e}}"
  133. if name == "edit_file":
  134. try:
  135. p = safe_path(args["path"])
  136. content = p.read_text()
  137. if args["old_text"] not in content:
  138. return f"Error: Text not found in {{args['path']}}"
  139. p.write_text(content.replace(args["old_text"], args["new_text"], 1))
  140. return f"Edited {{args['path']}}"
  141. except Exception as e:
  142. return f"Error: {{e}}"
  143. return f"Unknown tool: {{name}}"
  144. def agent(prompt: str, history: list = None) -> str:
  145. """Run the agent loop."""
  146. if history is None:
  147. history = []
  148. history.append({{"role": "user", "content": prompt}})
  149. while True:
  150. response = client.messages.create(
  151. model=MODEL, system=SYSTEM, messages=history, tools=TOOLS, max_tokens=8000
  152. )
  153. history.append({{"role": "assistant", "content": response.content}})
  154. if response.stop_reason != "tool_use":
  155. return "".join(b.text for b in response.content if hasattr(b, "text"))
  156. results = []
  157. for block in response.content:
  158. if block.type == "tool_use":
  159. print(f"> {{block.name}}: {{str(block.input)[:100]}}")
  160. output = execute(block.name, block.input)
  161. print(f" {{output[:100]}}...")
  162. results.append({{"type": "tool_result", "tool_use_id": block.id, "content": output}})
  163. history.append({{"role": "user", "content": results}})
  164. if __name__ == "__main__":
  165. print(f"{name} - Level 1 Agent at {{WORKDIR}}")
  166. print("Type 'q' to quit.\\n")
  167. h = []
  168. while True:
  169. try:
  170. query = input(">> ").strip()
  171. except (EOFError, KeyboardInterrupt):
  172. break
  173. if query in ("q", "quit", "exit", ""):
  174. break
  175. print(agent(query, h), "\\n")
  176. ''',
  177. }
  178. ENV_TEMPLATE = '''# API Configuration
  179. ANTHROPIC_API_KEY=sk-xxx
  180. ANTHROPIC_BASE_URL=https://api.anthropic.com
  181. MODEL_NAME=claude-sonnet-4-20250514
  182. '''
  183. def create_agent(name: str, level: int, output_dir: Path):
  184. """Create a new agent project."""
  185. # Validate level
  186. if level not in TEMPLATES and level not in (2, 3, 4):
  187. print(f"Error: Level {level} not yet implemented in scaffold.")
  188. print("Available levels: 0 (minimal), 1 (4 tools)")
  189. print("For levels 2-4, copy from mini-claude-code repository.")
  190. sys.exit(1)
  191. # Create output directory
  192. agent_dir = output_dir / name
  193. agent_dir.mkdir(parents=True, exist_ok=True)
  194. # Write agent file
  195. agent_file = agent_dir / f"{name}.py"
  196. template = TEMPLATES.get(level, TEMPLATES[1])
  197. agent_file.write_text(template.format(name=name))
  198. print(f"Created: {agent_file}")
  199. # Write .env.example
  200. env_file = agent_dir / ".env.example"
  201. env_file.write_text(ENV_TEMPLATE)
  202. print(f"Created: {env_file}")
  203. # Write .gitignore
  204. gitignore = agent_dir / ".gitignore"
  205. gitignore.write_text(".env\n__pycache__/\n*.pyc\n")
  206. print(f"Created: {gitignore}")
  207. print(f"\nAgent '{name}' created at {agent_dir}")
  208. print(f"\nNext steps:")
  209. print(f" 1. cd {agent_dir}")
  210. print(f" 2. cp .env.example .env")
  211. print(f" 3. Edit .env with your API key")
  212. print(f" 4. pip install anthropic python-dotenv")
  213. print(f" 5. python {name}.py")
  214. def main():
  215. parser = argparse.ArgumentParser(
  216. description="Scaffold a new AI coding agent project",
  217. formatter_class=argparse.RawDescriptionHelpFormatter,
  218. epilog="""
  219. Levels:
  220. 0 Minimal (~50 lines) - Single bash tool, self-recursion for subagents
  221. 1 Basic (~200 lines) - 4 core tools: bash, read, write, edit
  222. 2 Todo (~300 lines) - + TodoWrite for structured planning
  223. 3 Subagent (~450) - + Task tool for context isolation
  224. 4 Skills (~550) - + Skill tool for domain expertise
  225. """
  226. )
  227. parser.add_argument("name", help="Name of the agent to create")
  228. parser.add_argument("--level", type=int, default=1, choices=[0, 1, 2, 3, 4],
  229. help="Complexity level (default: 1)")
  230. parser.add_argument("--path", type=Path, default=Path.cwd(),
  231. help="Output directory (default: current directory)")
  232. args = parser.parse_args()
  233. create_agent(args.name, args.level, args.path)
  234. if __name__ == "__main__":
  235. main()