Tool-Using Agents

"""Tool Definition — The foundation of tool-using agents."""
from dataclasses import dataclass, field
from typing import Callable, Any
import json

@dataclass
class ToolParameter:
    name: str
    type: str  # "string", "integer", "boolean", "array"
    description: str
    required: bool = True
    enum: list = field(default_factory=list)

@dataclass
class Tool:
    name: str
    description: str
    parameters: list[ToolParameter]
    function: Callable[..., str]
    requires_confirmation: bool = False  # Human-in-the-loop
    timeout_seconds: int = 30
    
    def to_openai_schema(self) -> dict:
        """Convert to OpenAI function calling format."""
        properties = {}
        required = []
        for param in self.parameters:
            properties[param.name] = {
                "type": param.type,
                "description": param.description,
            }
            if param.enum:
                properties[param.name]["enum"] = param.enum
            if param.required:
                required.append(param.name)
        
        return {
            "type": "function",
            "function": {
                "name": self.name,
                "description": self.description,
                "parameters": {
                    "type": "object",
                    "properties": properties,
                    "required": required,
                },
            },
        }

# Example: Define a web search tool
search_tool = Tool(
    name="web_search",
    description="Search the web for current information. Use for facts, news, or data you don't know.",
    parameters=[
        ToolParameter("query", "string", "The search query"),
        ToolParameter("num_results", "integer", "Number of results (1-10)", required=False),
    ],
    function=lambda query, num_results=5: tavily_search(query, num_results),
    timeout_seconds=15,
)

"""Complete Tool Registry with validation, execution, and error handling."""
import time
import traceback
from concurrent.futures import ThreadPoolExecutor, TimeoutError

class ToolRegistry:
    """Central registry for all available tools."""
    
    def __init__(self):
        self.tools: dict[str, Tool] = {}
        self.execution_log: list[dict] = []
    
    def register(self, tool: Tool):
        """Register a tool, checking for name conflicts."""
        if tool.name in self.tools:
            raise ValueError(f"Tool '{tool.name}' already registered")
        self.tools[tool.name] = tool
    
    def get_schemas(self) -> list[dict]:
        """Get all tool schemas for the LLM."""
        return [tool.to_openai_schema() for tool in self.tools.values()]
    
    def get_tool_descriptions(self) -> str:
        """Human-readable tool list for system prompts."""
        lines = []
        for tool in self.tools.values():
            params = ", ".join(f"{p.name}: {p.type}" for p in tool.parameters)
            lines.append(f"- {tool.name}({params}): {tool.description}")
        return "\n".join(lines)
    
    def validate_args(self, tool_name: str, args: dict) -> tuple[bool, str]:
        """Validate arguments against tool schema."""
        tool = self.tools.get(tool_name)
        if not tool:
            return False, f"Unknown tool: {tool_name}"
        
        for param in tool.parameters:
            if param.required and param.name not in args:
                return False, f"Missing required parameter: {param.name}"
            if param.name in args and param.enum:
                if args[param.name] not in param.enum:
                    return False, f"Invalid value for {param.name}. Must be one of: {param.enum}"
        
        return True, "Valid"
    
    def execute(self, tool_name: str, args: dict, 
                confirm_fn: Callable = None) -> dict:
        """Execute a tool with timeout, validation, and logging."""
        start_time = time.time()
        
        # Validate
        valid, error = self.validate_args(tool_name, args)
        if not valid:
            return {"success": False, "error": error, "duration": 0}
        
        tool = self.tools[tool_name]
        
        # Human-in-the-loop confirmation
        if tool.requires_confirmation and confirm_fn:
            approved = confirm_fn(tool_name, args)
            if not approved:
                return {"success": False, "error": "User rejected execution", "duration": 0}
        
        # Execute with timeout
        try:
            with ThreadPoolExecutor(max_workers=1) as executor:
                future = executor.submit(tool.function, **args)
                result = future.result(timeout=tool.timeout_seconds)
            
            duration = time.time() - start_time
            log_entry = {
                "tool": tool_name, "args": args, "result": result,
                "success": True, "duration": duration,
            }
            self.execution_log.append(log_entry)
            return {"success": True, "result": result, "duration": duration}
        
        except TimeoutError:
            return {"success": False, "error": f"Tool timed out after {tool.timeout_seconds}s", 
                    "duration": tool.timeout_seconds}
        except Exception as e:
            return {"success": False, "error": f"{type(e).__name__}: {str(e)}", 
                    "duration": time.time() - start_time}

"""Full agent loop using OpenAI function calling + tool registry."""
import openai

client = openai.OpenAI()

def run_tool_agent(query: str, registry: ToolRegistry, 
                   max_iterations: int = 10, max_retries: int = 2) -> str:
    """Agent loop with automatic error recovery."""
    
    messages = [
        {"role": "system", "content": f"""You are a helpful research assistant with access to tools.
Use tools to find information, then synthesize a comprehensive answer.
If a tool fails, try a different approach or rephrase your query.

Available tools:
{registry.get_tool_descriptions()}"""},
        {"role": "user", "content": query},
    ]
    
    tool_schemas = registry.get_schemas()
    
    for iteration in range(max_iterations):
        response = client.chat.completions.create(
            model="gpt-4o",
            messages=messages,
            tools=tool_schemas,
            tool_choice="auto",  # Let model decide
        )
        
        msg = response.choices[0].message
        messages.append(msg)
        
        # If no tool calls, we have our final answer
        if not msg.tool_calls:
            return msg.content
        
        # Process each tool call
        for tool_call in msg.tool_calls:
            fn_name = tool_call.function.name
            fn_args = json.loads(tool_call.function.arguments)
            
            # Execute with retries
            result = None
            for attempt in range(max_retries + 1):
                result = registry.execute(fn_name, fn_args)
                if result["success"]:
                    break
                # On failure, wait briefly before retry
                time.sleep(1)
            
            # Format result for the LLM
            if result["success"]:
                tool_response = result["result"]
            else:
                tool_response = f"ERROR: {result['error']}. Please try a different approach."
            
            messages.append({
                "role": "tool",
                "tool_call_id": tool_call.id,
                "content": str(tool_response),
            })
    
    return "Agent reached maximum iterations. Partial results may be in the conversation."

"""Tools that chain: search → read → extract → synthesize."""

# Tool: Read a webpage and extract content
read_page_tool = Tool(
    name="read_webpage",
    description="Read the full text content of a webpage URL. Use after searching to get details.",
    parameters=[
        ToolParameter("url", "string", "The full URL to read"),
    ],
    function=lambda url: fetch_and_extract(url),
    timeout_seconds=20,
)

# Tool: Summarize long text to save context space
summarize_tool = Tool(
    name="summarize_text",
    description="Summarize a long text into key points. Use when text is too long to include directly.",
    parameters=[
        ToolParameter("text", "string", "The text to summarize"),
        ToolParameter("focus", "string", "What aspect to focus the summary on", required=False),
    ],
    function=lambda text, focus="": llm_summarize(text, focus),
    timeout_seconds=30,
)

# Tool: Write structured report
write_report_tool = Tool(
    name="write_report",
    description="Write a structured report from gathered research. Use as your final step.",
    parameters=[
        ToolParameter("topic", "string", "The report topic"),
        ToolParameter("findings", "string", "Key findings to include (as bullet points)"),
        ToolParameter("format", "string", "Report format", enum=["brief", "detailed", "executive"]),
    ],
    function=lambda topic, findings, format="detailed": generate_report(topic, findings, format),
    requires_confirmation=True,  # Confirm before generating long content
)

"""Sandboxing patterns for tool execution."""

class SandboxedToolRegistry(ToolRegistry):
    """Registry with security constraints."""
    
    def __init__(self, allowed_domains: list[str] = None,
                 blocked_tools: list[str] = None,
                 rate_limits: dict[str, int] = None):
        super().__init__()
        self.allowed_domains = allowed_domains or []
        self.blocked_tools = blocked_tools or []
        self.rate_limits = rate_limits or {}  # tool_name -> max calls per minute
        self._call_counts: dict[str, list[float]] = {}
    
    def execute(self, tool_name: str, args: dict, **kwargs) -> dict:
        # Check if tool is blocked
        if tool_name in self.blocked_tools:
            return {"success": False, "error": "Tool is disabled by security policy"}
        
        # Check rate limits
        if tool_name in self.rate_limits:
            now = time.time()
            calls = self._call_counts.get(tool_name, [])
            calls = [t for t in calls if now - t < 60]  # Last minute
            if len(calls) >= self.rate_limits[tool_name]:
                return {"success": False, "error": "Rate limit exceeded"}
            calls.append(now)
            self._call_counts[tool_name] = calls
        
        # URL allowlist check
        if "url" in args and self.allowed_domains:
            from urllib.parse import urlparse
            domain = urlparse(args["url"]).netloc
            if not any(domain.endswith(d) for d in self.allowed_domains):
                return {"success": False, "error": f"Domain '{domain}' not in allowlist"}
        
        return super().execute(tool_name, args, **kwargs)

# Usage
registry = SandboxedToolRegistry(
    allowed_domains=["wikipedia.org", "arxiv.org", "github.com"],
    rate_limits={"web_search": 10, "read_webpage": 20},
)