File size: 5,790 Bytes
82f262a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
"""
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")