v0.2.0: Add purpose_agent/tools.py

purpose_agent/tools.py

@@ -0,0 +1,459 @@
+"""
+Tool Framework — Structured tool definitions with schema, validation, and retry.
+
+Tools are the agent's interface to the world. This module provides:
+  - Tool base class with JSON schema, input validation, retry logic
+  - Built-in tools (search, calculator, python_exec, file_read)
+  - Tool registry with semantic search (for Tool RAG with SLMs)
+  - Automatic tool documentation generation for prompts
+
+SLM optimization: Tools are formatted with minimal token overhead.
+For small context windows, use ToolRegistry.get_relevant_tools()
+to retrieve only the k most relevant tools (TinyAgent pattern).
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import math
+import re
+import time
+import traceback
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Any, Callable
+
+logger = logging.getLogger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# Tool Base Class
+# ---------------------------------------------------------------------------
+
+@dataclass
+class ToolResult:
+    """Result of a tool execution."""
+    output: str
+    success: bool = True
+    error: str | None = None
+    execution_time_s: float = 0.0
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+
+class Tool(ABC):
+    """
+    Abstract tool that an Agent can invoke.
+    
+    Every tool has:
+    - name: Unique identifier
+    - description: What it does (used in agent prompts)
+    - parameters: JSON schema of expected inputs
+    - execute(): The actual implementation
+    
+    Built-in retry logic, input validation, and timeout handling.
+    
+    Example:
+        class SearchTool(Tool):
+            name = "web_search"
+            description = "Search the web for information"
+            parameters = {
+                "type": "object",
+                "properties": {
+                    "query": {"type": "string", "description": "Search query"}
+                },
+                "required": ["query"]
+            }
+            
+            def execute(self, query: str) -> str:
+                return do_search(query)
+    """
+
+    name: str = "unnamed_tool"
+    description: str = "No description"
+    parameters: dict[str, Any] = {}
+    max_retries: int = 2
+    timeout_seconds: float = 30.0
+
+    @abstractmethod
+    def execute(self, **kwargs) -> str:
+        """Execute the tool with the given parameters. Return a string result."""
+        ...
+
+    def run(self, **kwargs) -> ToolResult:
+        """
+        Run the tool with validation, retry, and error handling.
+        
+        This is the method the Orchestrator calls — it wraps execute()
+        with production safeguards.
+        """
+        # Validate inputs against schema
+        validation_error = self._validate_inputs(kwargs)
+        if validation_error:
+            return ToolResult(
+                output="", success=False,
+                error=f"Input validation failed: {validation_error}",
+            )
+
+        # Retry loop
+        last_error = None
+        for attempt in range(self.max_retries + 1):
+            start = time.time()
+            try:
+                result = self.execute(**kwargs)
+                return ToolResult(
+                    output=str(result),
+                    success=True,
+                    execution_time_s=time.time() - start,
+                    metadata={"attempt": attempt + 1},
+                )
+            except Exception as e:
+                last_error = e
+                elapsed = time.time() - start
+                logger.warning(
+                    f"Tool '{self.name}' failed (attempt {attempt + 1}/{self.max_retries + 1}): {e}"
+                )
+                if attempt < self.max_retries:
+                    time.sleep(0.5 * (attempt + 1))  # Exponential backoff
+
+        return ToolResult(
+            output="",
+            success=False,
+            error=f"Tool '{self.name}' failed after {self.max_retries + 1} attempts: {last_error}",
+            execution_time_s=time.time() - start,
+        )
+
+    def _validate_inputs(self, kwargs: dict) -> str | None:
+        """Validate inputs against the JSON schema. Returns error string or None."""
+        if not self.parameters:
+            return None
+
+        required = self.parameters.get("required", [])
+        properties = self.parameters.get("properties", {})
+
+        for req in required:
+            if req not in kwargs:
+                return f"Missing required parameter: '{req}'"
+
+        for key, value in kwargs.items():
+            if key in properties:
+                expected_type = properties[key].get("type")
+                if expected_type == "string" and not isinstance(value, str):
+                    return f"Parameter '{key}' should be string, got {type(value).__name__}"
+                elif expected_type == "integer" and not isinstance(value, int):
+                    return f"Parameter '{key}' should be integer, got {type(value).__name__}"
+                elif expected_type == "number" and not isinstance(value, (int, float)):
+                    return f"Parameter '{key}' should be number, got {type(value).__name__}"
+
+        return None
+
+    def to_schema(self) -> dict[str, Any]:
+        """Return OpenAI-compatible tool schema."""
+        return {
+            "type": "function",
+            "function": {
+                "name": self.name,
+                "description": self.description,
+                "parameters": self.parameters or {"type": "object", "properties": {}},
+            },
+        }
+
+    def to_prompt(self, compact: bool = False) -> str:
+        """
+        Format tool for inclusion in agent prompts.
+        
+        compact=True: Minimal format for SLMs (fewer tokens)
+        compact=False: Full description with parameter details
+        """
+        if compact:
+            params = ", ".join(
+                f"{k}: {v.get('type', 'any')}"
+                for k, v in self.parameters.get("properties", {}).items()
+            )
+            return f"- {self.name}({params}) — {self.description}"
+
+        lines = [f"### {self.name}"]
+        lines.append(f"  {self.description}")
+        if self.parameters.get("properties"):
+            lines.append("  Parameters:")
+            for pname, pinfo in self.parameters["properties"].items():
+                req = "REQUIRED" if pname in self.parameters.get("required", []) else "optional"
+                lines.append(f"    - {pname} ({pinfo.get('type', 'any')}, {req}): {pinfo.get('description', '')}")
+        return "\n".join(lines)
+
+
+# ---------------------------------------------------------------------------
+# Function Tool — Create tools from plain functions
+# ---------------------------------------------------------------------------
+
+class FunctionTool(Tool):
+    """
+    Create a Tool from a plain Python function.
+    
+    Usage:
+        def search(query: str) -> str:
+            '''Search the web for information.'''
+            return requests.get(f"https://api.search.com?q={query}").text
+        
+        tool = FunctionTool.from_function(search)
+        # or
+        tool = FunctionTool(
+            name="search",
+            description="Search the web",
+            parameters={"type": "object", "properties": {"query": {"type": "string"}}, "required": ["query"]},
+            fn=search,
+        )
+    """
+
+    def __init__(
+        self,
+        name: str,
+        description: str,
+        parameters: dict[str, Any],
+        fn: Callable[..., str],
+        **kwargs,
+    ):
+        self.name = name
+        self.description = description
+        self.parameters = parameters
+        self._fn = fn
+        super().__init__(**kwargs)
+
+    def execute(self, **kwargs) -> str:
+        return str(self._fn(**kwargs))
+
+    @classmethod
+    def from_function(cls, fn: Callable, name: str | None = None, description: str | None = None) -> "FunctionTool":
+        """Auto-create a Tool from a function's signature and docstring."""
+        import inspect
+
+        func_name = name or fn.__name__
+        func_desc = description or fn.__doc__ or f"Execute {func_name}"
+        func_desc = func_desc.strip().split("\n")[0]  # First line only
+
+        sig = inspect.signature(fn)
+        properties = {}
+        required = []
+
+        type_map = {str: "string", int: "integer", float: "number", bool: "boolean"}
+
+        for pname, param in sig.parameters.items():
+            ptype = "string"
+            if param.annotation != inspect.Parameter.empty:
+                ptype = type_map.get(param.annotation, "string")
+
+            properties[pname] = {"type": ptype, "description": f"Parameter {pname}"}
+
+            if param.default == inspect.Parameter.empty:
+                required.append(pname)
+
+        return cls(
+            name=func_name,
+            description=func_desc,
+            parameters={"type": "object", "properties": properties, "required": required},
+            fn=fn,
+        )
+
+
+# ---------------------------------------------------------------------------
+# Built-in Tools
+# ---------------------------------------------------------------------------
+
+class CalculatorTool(Tool):
+    """Safe math expression evaluator."""
+    name = "calculator"
+    description = "Evaluate a mathematical expression. Supports +, -, *, /, **, sqrt, sin, cos, abs."
+    parameters = {
+        "type": "object",
+        "properties": {
+            "expression": {"type": "string", "description": "Math expression to evaluate (e.g. '2 + 3 * 4')"}
+        },
+        "required": ["expression"],
+    }
+
+    def execute(self, expression: str) -> str:
+        # Safe eval with math functions only
+        allowed = {
+            "abs": abs, "round": round, "min": min, "max": max,
+            "sqrt": math.sqrt, "sin": math.sin, "cos": math.cos,
+            "tan": math.tan, "log": math.log, "pi": math.pi, "e": math.e,
+        }
+        # Sanitize: only allow digits, operators, parentheses, dots, and allowed function names
+        clean = re.sub(r'[^0-9+\-*/().,%\s]', '', expression.replace("^", "**"))
+        try:
+            result = eval(expression, {"__builtins__": {}}, allowed)
+            return str(result)
+        except Exception as e:
+            return f"Error evaluating '{expression}': {e}"
+
+
+class PythonExecTool(Tool):
+    """Execute Python code in a sandboxed environment."""
+    name = "python_exec"
+    description = "Execute Python code and return the output. Use print() to output results."
+    parameters = {
+        "type": "object",
+        "properties": {
+            "code": {"type": "string", "description": "Python code to execute"}
+        },
+        "required": ["code"],
+    }
+
+    def execute(self, code: str) -> str:
+        import io
+        import contextlib
+
+        output = io.StringIO()
+        try:
+            with contextlib.redirect_stdout(output):
+                exec(code, {"__builtins__": __builtins__}, {})
+            return output.getvalue() or "(no output)"
+        except Exception as e:
+            return f"Error: {e}\n{traceback.format_exc()}"
+
+
+class ReadFileTool(Tool):
+    """Read a local file."""
+    name = "read_file"
+    description = "Read the contents of a file at the given path."
+    parameters = {
+        "type": "object",
+        "properties": {
+            "path": {"type": "string", "description": "File path to read"}
+        },
+        "required": ["path"],
+    }
+
+    def execute(self, path: str) -> str:
+        try:
+            with open(path, "r") as f:
+                content = f.read()
+            if len(content) > 10000:
+                return content[:10000] + f"\n...[truncated, {len(content)} chars total]"
+            return content
+        except Exception as e:
+            return f"Error reading '{path}': {e}"
+
+
+class WriteFileTool(Tool):
+    """Write content to a local file."""
+    name = "write_file"
+    description = "Write content to a file. Creates the file if it doesn't exist."
+    parameters = {
+        "type": "object",
+        "properties": {
+            "path": {"type": "string", "description": "File path to write"},
+            "content": {"type": "string", "description": "Content to write"},
+        },
+        "required": ["path", "content"],
+    }
+
+    def execute(self, path: str, content: str) -> str:
+        try:
+            with open(path, "w") as f:
+                f.write(content)
+            return f"Written {len(content)} chars to {path}"
+        except Exception as e:
+            return f"Error writing '{path}': {e}"
+
+
+# ---------------------------------------------------------------------------
+# Tool Registry — with semantic retrieval for SLMs (Tool RAG)
+# ---------------------------------------------------------------------------
+
+class ToolRegistry:
+    """
+    Registry of available tools with semantic retrieval.
+    
+    For SLMs with small context windows, you can't list all tools in the prompt.
+    Instead, use get_relevant_tools() to retrieve only the k most relevant
+    tools for the current task (TinyAgent pattern, arxiv:2409.00608).
+    
+    Usage:
+        registry = ToolRegistry()
+        registry.register(CalculatorTool())
+        registry.register(SearchTool())
+        registry.register(PythonExecTool())
+        
+        # Get all tools (for LLMs with large context)
+        all_tools = registry.get_all()
+        
+        # Get top-k relevant tools (for SLMs)
+        relevant = registry.get_relevant_tools("calculate 2+2", top_k=3)
+    """
+
+    def __init__(self):
+        self._tools: dict[str, Tool] = {}
+        self._embeddings: dict[str, list[float]] = {}
+
+    def register(self, tool: Tool) -> "ToolRegistry":
+        """Register a tool."""
+        self._tools[tool.name] = tool
+        # Compute embedding for Tool RAG
+        text = f"{tool.name} {tool.description}"
+        self._embeddings[tool.name] = self._embed(text)
+        return self
+
+    def get(self, name: str) -> Tool | None:
+        return self._tools.get(name)
+
+    def get_all(self) -> list[Tool]:
+        return list(self._tools.values())
+
+    def execute(self, name: str, **kwargs) -> ToolResult:
+        """Execute a tool by name."""
+        tool = self._tools.get(name)
+        if not tool:
+            return ToolResult(output="", success=False, error=f"Unknown tool: '{name}'")
+        return tool.run(**kwargs)
+
+    def get_relevant_tools(self, query: str, top_k: int = 5) -> list[Tool]:
+        """
+        Retrieve the k most relevant tools for a query.
+        
+        Uses lightweight trigram embedding + cosine similarity (same as
+        ExperienceReplay). For production, swap in sentence-transformers.
+        """
+        if len(self._tools) <= top_k:
+            return list(self._tools.values())
+
+        query_emb = self._embed(query)
+        scored = []
+        for name, emb in self._embeddings.items():
+            sim = self._cosine_sim(query_emb, emb)
+            scored.append((sim, name))
+        scored.sort(key=lambda x: -x[0])
+
+        return [self._tools[name] for _, name in scored[:top_k]]
+
+    def format_for_prompt(self, tools: list[Tool] | None = None, compact: bool = False) -> str:
+        """Format tools for inclusion in agent prompts."""
+        tools = tools or list(self._tools.values())
+        return "\n".join(t.to_prompt(compact=compact) for t in tools)
+
+    def to_schemas(self, tools: list[Tool] | None = None) -> list[dict]:
+        """Get OpenAI-compatible tool schemas."""
+        tools = tools or list(self._tools.values())
+        return [t.to_schema() for t in tools]
+
+    @staticmethod
+    def _embed(text: str) -> list[float]:
+        """Lightweight embedding (same as ExperienceReplay)."""
+        dim = 64
+        vec = [0.0] * dim
+        text_lower = text.lower()
+        for i in range(len(text_lower) - 2):
+            trigram = text_lower[i:i + 3]
+            h = hash(trigram) % dim
+            vec[h] += 1.0
+        magnitude = math.sqrt(sum(x * x for x in vec))
+        if magnitude > 0:
+            vec = [x / magnitude for x in vec]
+        return vec
+
+    @staticmethod
+    def _cosine_sim(a: list[float], b: list[float]) -> float:
+        if not a or not b or len(a) != len(b):
+            return 0.0
+        dot = sum(x * y for x, y in zip(a, b))
+        return dot  # Vectors are already normalized