tool-templates.py 7.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271
  1. """
  2. Tool Templates - Copy and customize these for your agent.
  3. Each tool needs:
  4. 1. Definition (JSON schema for the model)
  5. 2. Implementation (Python function)
  6. """
  7. from pathlib import Path
  8. import subprocess
  9. WORKDIR = Path.cwd()
  10. # =============================================================================
  11. # TOOL DEFINITIONS (for TOOLS list)
  12. # =============================================================================
  13. BASH_TOOL = {
  14. "name": "bash",
  15. "description": "Run a shell command. Use for: ls, find, grep, git, npm, python, etc.",
  16. "input_schema": {
  17. "type": "object",
  18. "properties": {
  19. "command": {
  20. "type": "string",
  21. "description": "The shell command to execute"
  22. }
  23. },
  24. "required": ["command"],
  25. },
  26. }
  27. READ_FILE_TOOL = {
  28. "name": "read_file",
  29. "description": "Read file contents. Returns UTF-8 text.",
  30. "input_schema": {
  31. "type": "object",
  32. "properties": {
  33. "path": {
  34. "type": "string",
  35. "description": "Relative path to the file"
  36. },
  37. "limit": {
  38. "type": "integer",
  39. "description": "Max lines to read (default: all)"
  40. },
  41. },
  42. "required": ["path"],
  43. },
  44. }
  45. WRITE_FILE_TOOL = {
  46. "name": "write_file",
  47. "description": "Write content to a file. Creates parent directories if needed.",
  48. "input_schema": {
  49. "type": "object",
  50. "properties": {
  51. "path": {
  52. "type": "string",
  53. "description": "Relative path for the file"
  54. },
  55. "content": {
  56. "type": "string",
  57. "description": "Content to write"
  58. },
  59. },
  60. "required": ["path", "content"],
  61. },
  62. }
  63. EDIT_FILE_TOOL = {
  64. "name": "edit_file",
  65. "description": "Replace exact text in a file. Use for surgical edits.",
  66. "input_schema": {
  67. "type": "object",
  68. "properties": {
  69. "path": {
  70. "type": "string",
  71. "description": "Relative path to the file"
  72. },
  73. "old_text": {
  74. "type": "string",
  75. "description": "Exact text to find (must match precisely)"
  76. },
  77. "new_text": {
  78. "type": "string",
  79. "description": "Replacement text"
  80. },
  81. },
  82. "required": ["path", "old_text", "new_text"],
  83. },
  84. }
  85. TODO_WRITE_TOOL = {
  86. "name": "TodoWrite",
  87. "description": "Update the task list. Use to plan and track progress.",
  88. "input_schema": {
  89. "type": "object",
  90. "properties": {
  91. "items": {
  92. "type": "array",
  93. "description": "Complete list of tasks",
  94. "items": {
  95. "type": "object",
  96. "properties": {
  97. "content": {"type": "string", "description": "Task description"},
  98. "status": {"type": "string", "enum": ["pending", "in_progress", "completed"]},
  99. "activeForm": {"type": "string", "description": "Present tense, e.g. 'Reading files'"},
  100. },
  101. "required": ["content", "status", "activeForm"],
  102. },
  103. }
  104. },
  105. "required": ["items"],
  106. },
  107. }
  108. TASK_TOOL_TEMPLATE = """
  109. # Generate dynamically with agent types
  110. TASK_TOOL = {
  111. "name": "Task",
  112. "description": f"Spawn a subagent for a focused subtask.\\n\\nAgent types:\\n{get_agent_descriptions()}",
  113. "input_schema": {
  114. "type": "object",
  115. "properties": {
  116. "description": {"type": "string", "description": "Short task name (3-5 words)"},
  117. "prompt": {"type": "string", "description": "Detailed instructions"},
  118. "agent_type": {"type": "string", "enum": list(AGENT_TYPES.keys())},
  119. },
  120. "required": ["description", "prompt", "agent_type"],
  121. },
  122. }
  123. """
  124. # =============================================================================
  125. # TOOL IMPLEMENTATIONS
  126. # =============================================================================
  127. def safe_path(p: str) -> Path:
  128. """
  129. Security: Ensure path stays within workspace.
  130. Prevents ../../../etc/passwd attacks.
  131. """
  132. path = (WORKDIR / p).resolve()
  133. if not path.is_relative_to(WORKDIR):
  134. raise ValueError(f"Path escapes workspace: {p}")
  135. return path
  136. def run_bash(command: str) -> str:
  137. """
  138. Execute shell command with safety checks.
  139. Safety features:
  140. - Blocks obviously dangerous commands
  141. - 60 second timeout
  142. - Output truncated to 50KB
  143. """
  144. dangerous = ["rm -rf /", "sudo", "shutdown", "reboot", "> /dev/"]
  145. if any(d in command for d in dangerous):
  146. return "Error: Dangerous command blocked"
  147. try:
  148. result = subprocess.run(
  149. command,
  150. shell=True,
  151. cwd=WORKDIR,
  152. capture_output=True,
  153. text=True,
  154. timeout=60
  155. )
  156. output = (result.stdout + result.stderr).strip()
  157. return output[:50000] if output else "(no output)"
  158. except subprocess.TimeoutExpired:
  159. return "Error: Command timed out (60s)"
  160. except Exception as e:
  161. return f"Error: {e}"
  162. def run_read_file(path: str, limit: int = None) -> str:
  163. """
  164. Read file contents with optional line limit.
  165. Features:
  166. - Safe path resolution
  167. - Optional line limit for large files
  168. - Output truncated to 50KB
  169. """
  170. try:
  171. text = safe_path(path).read_text()
  172. lines = text.splitlines()
  173. if limit and limit < len(lines):
  174. lines = lines[:limit]
  175. lines.append(f"... ({len(text.splitlines()) - limit} more lines)")
  176. return "\n".join(lines)[:50000]
  177. except Exception as e:
  178. return f"Error: {e}"
  179. def run_write_file(path: str, content: str) -> str:
  180. """
  181. Write content to file, creating parent directories if needed.
  182. Features:
  183. - Safe path resolution
  184. - Auto-creates parent directories
  185. - Returns byte count for confirmation
  186. """
  187. try:
  188. fp = safe_path(path)
  189. fp.parent.mkdir(parents=True, exist_ok=True)
  190. fp.write_text(content)
  191. return f"Wrote {len(content)} bytes to {path}"
  192. except Exception as e:
  193. return f"Error: {e}"
  194. def run_edit_file(path: str, old_text: str, new_text: str) -> str:
  195. """
  196. Replace exact text in a file (surgical edit).
  197. Features:
  198. - Exact string matching (not regex)
  199. - Only replaces first occurrence (safety)
  200. - Clear error if text not found
  201. """
  202. try:
  203. fp = safe_path(path)
  204. content = fp.read_text()
  205. if old_text not in content:
  206. return f"Error: Text not found in {path}"
  207. new_content = content.replace(old_text, new_text, 1)
  208. fp.write_text(new_content)
  209. return f"Edited {path}"
  210. except Exception as e:
  211. return f"Error: {e}"
  212. # =============================================================================
  213. # DISPATCHER PATTERN
  214. # =============================================================================
  215. def execute_tool(name: str, args: dict) -> str:
  216. """
  217. Dispatch tool call to implementation.
  218. This pattern makes it easy to add new tools:
  219. 1. Add definition to TOOLS list
  220. 2. Add implementation function
  221. 3. Add case to this dispatcher
  222. """
  223. if name == "bash":
  224. return run_bash(args["command"])
  225. if name == "read_file":
  226. return run_read_file(args["path"], args.get("limit"))
  227. if name == "write_file":
  228. return run_write_file(args["path"], args["content"])
  229. if name == "edit_file":
  230. return run_edit_file(args["path"], args["old_text"], args["new_text"])
  231. # Add more tools here...
  232. return f"Unknown tool: {name}"