python -m venv venv
source venv/bin/activate
pip install openai pydantic httpx tenacity
export OPENAI_API_KEY="sk-your-key"

# registry.py
from __future__ import annotations
import json
import inspect
import logging
from typing import Any, Callable, get_type_hints
from pydantic import BaseModel, create_model
from pydantic.fields import FieldInfo

logger = logging.getLogger(__name__)


class ToolSchema(BaseModel):
    """Schema for a registered tool."""
    name: str
    description: str
    parameters: dict[str, Any]
    func: Callable[..., Any]
    is_async: bool = False
    timeout: float = 30.0
    max_retries: int = 3
    tags: list[str] = []

    class Config:
        arbitrary_types_allowed = True


class ToolRegistry:
    """
    Central registry for tool management.
    
    Supports decorator-based registration, auto-schema generation
    from type hints, and OpenAI-compatible format export.
    
    Usage:
        registry = ToolRegistry()
        
        @registry.register(name="search", description="Search the web")
        async def web_search(query: str, num_results: int = 5) -> str:
            ...
    """
    
    def __init__(self) -> None:
        self._tools: dict[str, ToolSchema] = {}
        logger.debug("Initialized empty ToolRegistry")

    def register(
        self,
        name: str | None = None,
        description: str | None = None,
        timeout: float = 30.0,
        max_retries: int = 3,
        tags: list[str] | None = None,
    ) -> Callable:
        """
        Decorator to register a function as a tool.
        
        Args:
            name: Tool name (defaults to function name)
            description: Tool description (defaults to docstring)
            timeout: Maximum execution time in seconds
            max_retries: Number of retry attempts on failure
            tags: Optional tags for categorization
            
        Returns:
            Decorated function
        """
        def decorator(func: Callable) -> Callable:
            tool_name = name or func.__name__
            tool_desc = description or func.__doc__ or f"Execute {tool_name}"
            is_coro = inspect.iscoroutinefunction(func)
            schema = self._generate_schema(func)
            
            self._tools[tool_name] = ToolSchema(
                name=tool_name,
                description=tool_desc,
                parameters=schema,
                func=func,
                is_async=is_coro,
                timeout=timeout,
                max_retries=max_retries,
                tags=tags or [],
            )
            logger.info(
                f"Registered tool: {tool_name} "
                f"(async={is_coro}, timeout={timeout}s)"
            )
            return func
        return decorator

    def register_function(
        self,
        func: Callable,
        name: str | None = None,
        description: str | None = None,
        timeout: float = 30.0,
        max_retries: int = 3,
        tags: list[str] | None = None,
    ) -> None:
        """Register a function programmatically (non-decorator)."""
        tool_name = name or func.__name__
        tool_desc = description or func.__doc__ or f"Execute {tool_name}"
        is_coro = inspect.iscoroutinefunction(func)
        schema = self._generate_schema(func)
        
        self._tools[tool_name] = ToolSchema(
            name=tool_name,
            description=tool_desc,
            parameters=schema,
            func=func,
            is_async=is_coro,
            timeout=timeout,
            max_retries=max_retries,
            tags=tags or [],
        )
        logger.info(f"Registered tool via function: {tool_name}")

    def _generate_schema(self, func: Callable) -> dict[str, Any]:
        """
        Auto-generate JSON Schema from function signature and type hints.
        
        Uses Pydantic to create a model from the function parameters,
        then converts to JSON Schema format.
        """
        sig = inspect.signature(func)
        hints = get_type_hints(func)
        
        fields = {}
        for param_name, param in sig.parameters.items():
            if param_name == "self":
                continue
            
            param_type = hints.get(param_name, str)
            
            # Handle default values
            if param.default is not inspect.Parameter.empty:
                fields[param_name] = (param_type, FieldInfo(default=param.default))
            else:
                fields[param_name] = (param_type, FieldInfo(...))
        
        # Create Pydantic model
        model = create_model(
            f"{func.__name__}Args",
            **fields,
        )
        
        return model.model_json_schema()

    def _python_type_to_json(self, python_type) -> str:
        """Convert Python type to JSON Schema type string."""
        type_map = {
            str: "string",
            int: "integer",
            float: "number",
            bool: "boolean",
            list: "array",
            dict: "object",
        }
        return type_map.get(python_type, "string")

    def get_openai_tools(self) -> list[dict[str, Any]]:
        """Export tools in OpenAI function calling format."""
        tools = []
        for tool in self._tools.values():
            tools.append({
                "type": "function",
                "function": {
                    "name": tool.name,
                    "description": tool.description,
                    "parameters": tool.parameters,
                },
            })
        return tools

    def get_tool(self, name: str) -> ToolSchema | None:
        """Retrieve a tool by name."""
        return self._tools.get(name)

    def list_tools(self) -> list[str]:
        """List all registered tool names."""
        return list(self._tools.keys())

    def get_tools_by_tag(self, tag: str) -> list[ToolSchema]:
        """Get all tools matching a specific tag."""
        return [t for t in self._tools.values() if tag in t.tags]

# executor.py
from __future__ import annotations
import asyncio
import json
import time
import logging
from typing import Any
from dataclasses import dataclass, field
from registry import ToolRegistry, ToolSchema
from tenacity import (
    retry,
    stop_after_attempt,
    wait_exponential,
    retry_if_exception_type,
)

logger = logging.getLogger(__name__)


@dataclass
class ToolExecutionResult:
    """Result of a single tool execution."""
    tool_name: str
    arguments: dict[str, Any]
    result: str
    success: bool
    duration_ms: float
    error: str | None = None
    retries_used: int = 0

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary for serialization."""
        return {
            "tool_name": self.tool_name,
            "arguments": self.arguments,
            "result": self.result,
            "success": self.success,
            "duration_ms": self.duration_ms,
            "error": self.error,
            "retries_used": self.retries_used,
        }


@dataclass
class ExecutionStats:
    """Statistics for tool execution tracking."""
    total: int = 0
    successes: int = 0
    failures: int = 0
    avg_duration_ms: float = 0.0
    tool_counts: dict[str, int] = field(default_factory=dict)
    tool_latencies: dict[str, list[float]] = field(default_factory=dict)

    def record(self, result: ToolExecutionResult) -> None:
        """Record an execution result."""
        self.total += 1
        if result.success:
            self.successes += 1
        else:
            self.failures += 1
        
        # Update averages
        self.avg_duration_ms = (
            (self.avg_duration_ms * (self.total - 1) + result.duration_ms)
            / self.total
        )
        
        # Track per-tool stats
        self.tool_counts[result.tool_name] = (
            self.tool_counts.get(result.tool_name, 0) + 1
        )
        if result.tool_name not in self.tool_latencies:
            self.tool_latencies[result.tool_name] = []
        self.tool_latencies[result.tool_name].append(result.duration_ms)


class ToolExecutor:
    """
    Executes tools with retry logic, timeout handling, and logging.
    
    Features:
    - Parallel execution of independent tool calls
    - Exponential backoff for retries
    - Timeout protection
    - Comprehensive execution logging
    - Per-tool performance statistics
    """
    
    def __init__(self, registry: ToolRegistry) -> None:
        self.registry = registry
        self.execution_log: list[ToolExecutionResult] = []
        self.stats = ExecutionStats()
        logger.info("ToolExecutor initialized")

    async def execute(
        self,
        tool_name: str,
        arguments: dict[str, Any],
        dry_run: bool = False,
    ) -> ToolExecutionResult:
        """
        Execute a single tool call.
        
        Args:
            tool_name: Name of the tool to execute
            arguments: Arguments to pass to the tool
            dry_run: If True, validate but don't execute
            
        Returns:
            ToolExecutionResult with success status and output
        """
        tool = self.registry.get_tool(tool_name)
        
        if not tool:
            logger.warning(f"Unknown tool requested: {tool_name}")
            return ToolExecutionResult(
                tool_name=tool_name,
                arguments=arguments,
                result=f"Error: Unknown tool '{tool_name}'",
                success=False,
                duration_ms=0,
                error=f"Tool '{tool_name}' not found in registry",
            )
        
        if dry_run:
            return ToolExecutionResult(
                tool_name=tool_name,
                arguments=arguments,
                result="[DRY RUN] Tool validated successfully",
                success=True,
                duration_ms=0,
            )

        start_time = time.monotonic()
        retries_used = 0
        
        try:
            if tool.is_async:
                result = await self._execute_async(tool, arguments)
            else:
                loop = asyncio.get_event_loop()
                result = await asyncio.wait_for(
                    loop.run_in_executor(None, lambda: tool.func(**arguments)),
                    timeout=tool.timeout,
                )
            
            duration = (time.monotonic() - start_time) * 1000
            execution = ToolExecutionResult(
                tool_name=tool_name,
                arguments=arguments,
                result=str(result),
                success=True,
                duration_ms=duration,
                retries_used=retries_used,
            )
            
        except asyncio.TimeoutError:
            duration = (time.monotonic() - start_time) * 1000
            error_msg = f"Tool '{tool_name}' timed out after {tool.timeout}s"
            logger.error(error_msg)
            execution = ToolExecutionResult(
                tool_name=tool_name,
                arguments=arguments,
                result=f"Error: {error_msg}",
                success=False,
                duration_ms=duration,
                error="Timeout",
                retries_used=retries_used,
            )
            
        except Exception as e:
            duration = (time.monotonic() - start_time) * 1000
            error_msg = f"Tool '{tool_name}' failed: {str(e)}"
            logger.error(error_msg, exc_info=True)
            execution = ToolExecutionResult(
                tool_name=tool_name,
                arguments=arguments,
                result=f"Error: {str(e)}",
                success=False,
                duration_ms=duration,
                error=str(e),
                retries_used=retries_used,
            )

        self.execution_log.append(execution)
        self.stats.record(execution)
        return execution

    async def _execute_async(
        self, tool: ToolSchema, arguments: dict[str, Any]
    ) -> str:
        """Execute async tool with retry logic."""
        retries_used = 0
        
        @retry(
            stop=stop_after_attempt(tool.max_retries),
            wait=wait_exponential(multiplier=1, min=0.5, max=10),
            retry=retry_if_exception_type((ConnectionError, TimeoutError)),
        )
        async def _inner():
            nonlocal retries_used
            retries_used += 1
            return await tool.func(**arguments)
        
        return await _inner()

    async def execute_parallel(
        self, calls: list[dict[str, Any]]
    ) -> list[ToolExecutionResult]:
        """
        Execute multiple tool calls in parallel.
        
        Args:
            calls: List of {"tool": name, "arguments": args} dicts
            
        Returns:
            List of ToolExecutionResults in same order as calls
        """
        tasks = [
            self.execute(call["tool"], call["arguments"])
            for call in calls
        ]
        results = await asyncio.gather(*tasks)
        
        logger.info(
            f"Parallel execution completed: {len(results)} tools, "
            f"{sum(1 for r in results if r.success)}/{len(results)} successful"
        )
        return results

    def get_stats(self) -> dict[str, Any]:
        """Get execution statistics."""
        return {
            "total": self.stats.total,
            "successes": self.stats.successes,
            "failures": self.stats.failures,
            "success_rate": (
                self.stats.successes / self.stats.total * 100
                if self.stats.total > 0 else 0
            ),
            "avg_duration_ms": round(self.stats.avg_duration_ms, 2),
            "tool_counts": self.stats.tool_counts,
        }

# agent.py
from __future__ import annotations
import json
import logging
from openai import AsyncOpenAI
from registry import ToolRegistry
from executor import ToolExecutor

logger = logging.getLogger(__name__)

TOOL_AGENT_SYSTEM = """You are a helpful assistant with access to tools.

When you need to use a tool, call the appropriate function.
You can call multiple tools in parallel when needed.
Always explain your reasoning before and after tool calls.
If a tool fails, try an alternative approach.
Be concise in your responses after tool execution."""

MAX_TOOL_ITERATIONS = 10  # Prevent infinite loops


class ToolUsingAgent:
    """
    Production-grade tool-using agent with OpenAI function calling.
    
    Features:
    - Automatic tool registration via decorator
    - Parallel tool execution
    - Error recovery and retry logic
    - Conversation history management
    - Token usage tracking
    """
    
    def __init__(
        self,
        model: str = "gpt-4-turbo-preview",
        temperature: float = 0.0,
        max_iterations: int = MAX_TOOL_ITERATIONS,
    ):
        self.client = AsyncOpenAI()
        self.model = model
        self.temperature = temperature
        self.max_iterations = max_iterations
        self.registry = ToolRegistry()
        self.executor = ToolExecutor(self.registry)
        self.messages: list[dict] = []
        self.total_tokens_used = 0
        
        logger.info(f"Initialized ToolUsingAgent with model={model}")

    def register(
        self, name: str = None, description: str = None, **kwargs
    ):
        """Decorator to register a tool with this agent."""
        return self.registry.register(
            name=name, description=description, **kwargs
        )

    async def chat(self, user_input: str) -> str:
        """
        Process a user message and return response.
        
        May involve multiple tool call iterations.
        
        Args:
            user_input: User's message
            
        Returns:
            Assistant's final response text
        """
        self.messages.append({"role": "user", "content": user_input})
        
        for iteration in range(self.max_iterations):
            tools = self.registry.get_openai_tools()
            
            response = await self.client.chat.completions.create(
                model=self.model,
                messages=self.messages,
                tools=tools if tools else None,
                temperature=self.temperature,
            )
            
            # Track token usage
            if response.usage:
                self.total_tokens_used += response.usage.total_tokens
            
            choice = response.choices[0]
            message = choice.message
            
            # Handle tool calls
            if message.tool_calls:
                logger.info(
                    f"Tool calls requested: "
                    f"{[tc.function.name for tc in message.tool_calls]}"
                )
                
                # Add assistant message with tool calls
                self.messages.append({
                    "role": "assistant",
                    "content": message.content,
                    "tool_calls": [
                        {
                            "id": tc.id,
                            "type": "function",
                            "function": {
                                "name": tc.function.name,
                                "arguments": tc.function.arguments,
                            },
                        }
                        for tc in message.tool_calls
                    ],
                })
                
                # Parse and execute tool calls
                calls = []
                for tc in message.tool_calls:
                    try:
                        args = json.loads(tc.function.arguments)
                    except json.JSONDecodeError as e:
                        logger.warning(
                            f"Failed to parse args for {tc.function.name}: {e}"
                        )
                        args = {}
                    calls.append({
                        "tool": tc.function.name,
                        "arguments": args,
                    })
                
                # Execute all tools in parallel
                results = await self.executor.execute_parallel(calls)
                
                # Add tool results to messages
                for tc, result in zip(message.tool_calls, results):
                    self.messages.append({
                        "role": "tool",
                        "tool_call_id": tc.id,
                        "content": result.result,
                    })
                    
            elif message.content:
                # Final text response
                self.messages.append({
                    "role": "assistant",
                    "content": message.content,
                })
                return message.content
        
        # Safety fallback
        logger.warning("Max iterations reached")
        return "I apologize, but I'm having trouble completing this request. Could you please rephrase?"

    def get_stats(self) -> dict:
        """Get agent performance statistics."""
        return {
            **self.executor.get_stats(),
            "total_tokens_used": self.total_tokens_used,
            "messages_in_history": len(self.messages),
        }

    def clear_history(self) -> None:
        """Clear conversation history."""
        self.messages = []
        logger.info("Conversation history cleared")

# tools/web.py
import httpx
import logging
from registry import ToolRegistry

logger = logging.getLogger(__name__)

registry = ToolRegistry()


@registry.register(
    name="web_search",
    description="Search the web for current information on any topic",
    timeout=15.0,
    max_retries=2,
    tags=["search", "web"],
)
async def web_search(query: str, num_results: int = 5) -> str:
    """
    Search the web for information.
    
    Args:
        query: Search query string
        num_results: Number of results to return (default: 5)
        
    Returns:
        Formatted search results
    """
    async with httpx.AsyncClient() as client:
        try:
            response = await client.get(
                "https://api.searchapi.io/api/v1/search",
                params={"q": query, "engine": "google"},
                timeout=10.0,
            )
            response.raise_for_status()
            data = response.json()
            
            results = data.get("organic_results", [])[:num_results]
            if not results:
                return f"No results found for: {query}"
            
            formatted = []
            for i, r in enumerate(results, 1):
                title = r.get("title", "No title")
                snippet = r.get("snippet", "No snippet")
                link = r.get("link", "")
                formatted.append(f"{i}. {title}\n   {snippet}\n   {link}")
            
            return "\n\n".join(formatted)
            
        except httpx.HTTPStatusError as e:
            logger.error(f"HTTP error during web search: {e}")
            return f"Search failed: HTTP {e.response.status_code}"
        except httpx.TimeoutException:
            return "Search timed out. Please try a simpler query."
        except Exception as e:
            logger.error(f"Unexpected error during web search: {e}")
            return f"Search failed: {str(e)}"


# tools/calculator.py
import math
import logging
from registry import ToolRegistry

logger = logging.getLogger(__name__)

calc_registry = ToolRegistry()

SAFE_FUNCTIONS = {
    "sqrt": math.sqrt,
    "log": math.log,
    "log10": math.log10,
    "sin": math.sin,
    "cos": math.cos,
    "tan": math.tan,
    "pi": math.pi,
    "e": math.e,
    "abs": abs,
    "round": round,
    "floor": math.floor,
    "ceil": math.ceil,
    "pow": pow,
}


@calc_registry.register(
    name="calculate",
    description="Evaluate mathematical expressions safely. Supports basic arithmetic, trigonometry, logarithms, and more.",
    timeout=5.0,
    max_retries=1,
    tags=["math", "calculation"],
)
def calculate(expression: str) -> str:
    """
    Safely evaluate a mathematical expression.
    
    Args:
        expression: Mathematical expression to evaluate
            Examples: "2 + 2", "sqrt(144)", "sin(pi/2)"
            
    Returns:
        Result of the expression as a string
    """
    try:
        # Security: Only allow safe functions
        result = eval(expression, {"__builtins__": {}}, SAFE_FUNCTIONS)
        return str(result)
    except ZeroDivisionError:
        return "Error: Division by zero"
    except ValueError as e:
        return f"Error: Invalid value - {str(e)}"
    except SyntaxError:
        return f"Error: Invalid expression syntax"
    except Exception as e:
        logger.error(f"Calculation error: {e}")
        return f"Error: Could not evaluate expression"

# tests/test_tools.py
import pytest
import asyncio
from registry import ToolRegistry
from executor import ToolExecutor, ToolExecutionResult


@pytest.fixture
def registry():
    """Create a test registry with sample tools."""
    reg = ToolRegistry()

    @reg.register(name="add", description="Add two numbers")
    def add(a: int, b: int) -> int:
        return a + b

    @reg.register(name="multiply", description="Multiply two numbers")
    def multiply(a: int, b: int) -> int:
        return a * b

    @reg.register(name="failing_tool", description="Always fails")
    def failing_tool() -> str:
        raise ValueError("Intentional failure")

    return reg


@pytest.fixture
def executor(registry):
    """Create a test executor."""
    return ToolExecutor(registry)


def test_tool_registration(registry):
    """Test that tools are registered correctly."""
    assert "add" in registry.list_tools()
    tools = registry.get_openai_tools()
    assert len(tools) == 3
    assert tools[0]["function"]["name"] == "add"


def test_schema_generation(registry):
    """Test auto-generated schema matches expected format."""
    tools = registry.get_openai_tools()
    schema = tools[0]["function"]["parameters"]
    assert schema["type"] == "object"
    assert "properties" in schema
    assert "required" in schema


@pytest.mark.asyncio
async def test_tool_execution(executor):
    """Test successful tool execution."""
    result = await executor.execute("add", {"a": 2, "b": 3})
    assert result.success
    assert result.result == "5"
    assert result.duration_ms > 0


@pytest.mark.asyncio
async def test_unknown_tool(executor):
    """Test handling of unknown tool names."""
    result = await executor.execute("nonexistent", {})
    assert not result.success
    assert "not found" in result.error


@pytest.mark.asyncio
async def test_tool_failure(executor):
    """Test handling of tool execution failures."""
    result = await executor.execute("failing_tool", {})
    assert not result.success
    assert "Intentional failure" in result.error


@pytest.mark.asyncio
async def test_parallel_execution(executor):
    """Test parallel execution of multiple tools."""
    calls = [
        {"tool": "add", "arguments": {"a": 1, "b": 2}},
        {"tool": "multiply", "arguments": {"a": 3, "b": 4}},
    ]
    results = await executor.execute_parallel(calls)
    assert len(results) == 2
    assert all(r.success for r in results)


def test_stats(executor):
    """Test statistics tracking."""
    stats = executor.get_stats()
    assert stats["total"] == 0
    assert stats["success_rate"] == 0

agent = ToolUsingAgent(model="gpt-4-turbo-preview")

@agent.register(name="search_arxiv", description="Search arXiv papers")
async def search_arxiv(query: str, max_results: int = 5) -> str:
    # Implementation...
    pass

@agent.register(name="calculate_stats", description="Calculate statistics")
def calculate_stats(numbers: list[float]) -> str:
    import statistics
    return f"Mean: {statistics.mean(numbers)}, Std: {statistics.stdev(numbers)}"

response = await agent.chat(
    "Find 5 recent papers on transformer efficiency and calculate "
    "the average citation count"
)

@agent.register(name="query_db", description="Execute SQL query")
async def query_db(query: str, database: str = "production") -> str:
    # Execute query with timeout and logging
    pass

@agent.register(name="generate_chart", description="Create visualization")
def generate_chart(data: str, chart_type: str = "bar") -> str:
    # Generate chart from data
    pass

# Example: Human-in-the-loop for destructive actions
@registry.register(
    name="delete_file",
    description="Delete a file from the system",
    requires_approval=True,  # Flag for human approval
)
def delete_file(path: str) -> str:
    # Check if approval is granted
    if not check_approval("delete_file", {"path": path}):
        return "Error: Human approval required for file deletion"
    # Proceed with deletion
    os.remove(path)
    return f"Deleted: {path}"

# Function calling approach
response = await client.chat.completions.create(
    model="gpt-4-turbo-preview",
    messages=messages,
    tools=[{
        "type": "function",
        "function": {
            "name": "search",
            "parameters": {"type": "object", "properties": {"q": {"type": "string"}}}
        }
    }]
)
# Returns structured tool_calls

# Prompt-based approach
messages.append({"role": "system", "content": "Use Action: search(q=...)"})
# Returns text that may or may not match expected format

# Parallel execution example
if message.tool_calls:
    # Execute all tools in parallel
    calls = [{"tool": tc.function.name, "arguments": json.loads(tc.function.arguments)} 
             for tc in message.tool_calls]
    results = await executor.execute_parallel(calls)
    
    # Return all results at once
    for tc, result in zip(message.tool_calls, results):
        messages.append({"role": "tool", "tool_call_id": tc.id, "content": result.result})

# Example: Graceful error handling with fallback
try:
    result = await primary_search(query)
except SearchError:
    # Fallback to alternative search
    result = await backup_search(query)
    logger.warning(f"Primary search failed, used backup for: {query}")