src/tools.py

"""
ToolRegistry - JSON-structured function calling for MORPH-AI v6.
Provides tool registration, validation, and execution.
"""

import json
import re
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Optional


@dataclass
class ToolCall:
    tool_name: str
    arguments: Dict[str, Any]
    confidence: float = 1.0
    raw: str = ""

    def to_dict(self) -> dict:
        return {
            "tool": self.tool_name,
            "arguments": self.arguments,
            "confidence": self.confidence,
        }


@dataclass
class Tool:
    name: str
    description: str
    parameters: Dict[str, Any]
    handler: Callable
    required: List[str] = field(default_factory=list)


class ToolRegistry:
    """Registry for tools that can be called by the model."""

    def __init__(self):
        self._tools: Dict[str, Tool] = {}
        self._register_defaults()

    def _register_defaults(self):
        """Register built-in tools."""
        self.register(Tool(
            name="calculator",
            description="Evaluate a mathematical expression",
            parameters={
                "type": "object",
                "properties": {
                    "expression": {"type": "string", "description": "Math expression to evaluate"}
                },
                "required": ["expression"]
            },
            handler=self._calc_handler,
            required=["expression"]
        ))

        self.register(Tool(
            name="search",
            description="Search the web for information",
            parameters={
                "type": "object",
                "properties": {
                    "query": {"type": "string", "description": "Search query"}
                },
                "required": ["query"]
            },
            handler=self._search_handler,
            required=["query"]
        ))

        self.register(Tool(
            name="code_exec",
            description="Execute Python code safely",
            parameters={
                "type": "object",
                "properties": {
                    "code": {"type": "string", "description": "Python code to execute"}
                },
                "required": ["code"]
            },
            handler=self._code_exec_handler,
            required=["code"]
        ))

        self.register(Tool(
            name="current_time",
            description="Get the current date and time",
            parameters={
                "type": "object",
                "properties": {
                    "timezone": {"type": "string", "description": "Timezone (optional)"}
                },
                "required": []
            },
            handler=self._time_handler,
            required=[]
        ))

    def register(self, tool: Tool):
        self._tools[tool.name] = tool

    def get_tool_schema(self) -> str:
        """Return JSON schema of all registered tools for prompt injection."""
        tools = []
        for t in self._tools.values():
            tools.append({
                "name": t.name,
                "description": t.description,
                "parameters": t.parameters
            })
        return json.dumps(tools, indent=2)

    def parse_calls(self, text: str) -> List[ToolCall]:
        """Extract JSON tool calls from model response."""
        calls = []
        pattern = r'```json\s*(\{.*?\})\s*```'
        for m in re.finditer(pattern, text, re.DOTALL):
            try:
                data = json.loads(m.group(1))
                if "tool" in data:
                    calls.append(ToolCall(
                        tool_name=data["tool"],
                        arguments=data.get("arguments", {}),
                        confidence=data.get("confidence", 1.0),
                        raw=m.group(1)
                    ))
            except json.JSONDecodeError:
                continue
        return calls

    def validate(self, call: ToolCall) -> bool:
        tool = self._tools.get(call.tool_name)
        if not tool:
            return False
        for req in tool.required:
            if req not in call.arguments:
                return False
        return True

    def execute(self, call: ToolCall) -> str:
        if not self.validate(call):
            return f"Error: invalid tool call {call.tool_name}"
        tool = self._tools[call.tool_name]
        try:
            result = tool.handler(**call.arguments)
            return str(result)
        except Exception as e:
            return f"Error executing {call.tool_name}: {e}"

    def _calc_handler(self, expression: str) -> Any:
        try:
            result = eval(expression, {"__builtins__": {}}, {})
            return result
        except Exception as e:
            return f"Calculation error: {e}"

    def _search_handler(self, query: str) -> str:
        from search import SearchClient
        client = SearchClient()
        results = client.search(query, num=3)
        return "\n".join(f"- {r.title}: {r.snippet}" for r in results)

    def _code_exec_handler(self, code: str) -> str:
        from architecture import CodeSandbox
        sandbox = CodeSandbox()
        result = sandbox.execute(code)
        if result["success"]:
            return f"Output: {result['output']}"
        return f"Error: {result['error']}"

    def _time_handler(self, timezone: Optional[str] = None) -> str:
        from datetime import datetime, timezone as tz
        import pytz
        if timezone:
            try:
                tz_obj = pytz.timezone(timezone)
                now = datetime.now(tz_obj)
            except Exception:
                now = datetime.now(tz.UTC)
        else:
            now = datetime.now()
        return now.strftime("%Y-%m-%d %H:%M:%S %Z")