Create tools_system.py

tools_system.py

@@ -0,0 +1,466 @@
+"""
+Helion-V1.5 Tools and Function Calling System
+Enables structured function calls and tool integration
+"""
+
+import json
+import logging
+from typing import List, Dict, Any, Callable, Optional
+from dataclasses import dataclass, asdict
+from enum import Enum
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class ToolType(Enum):
+    """Types of tools available."""
+    FUNCTION = "function"
+    API = "api"
+    RETRIEVAL = "retrieval"
+    CODE_INTERPRETER = "code_interpreter"
+
+
+@dataclass
+class ToolParameter:
+    """Tool parameter specification."""
+    name: str
+    type: str
+    description: str
+    required: bool = True
+    enum: Optional[List[str]] = None
+    default: Optional[Any] = None
+
+
+@dataclass
+class Tool:
+    """Tool/Function definition."""
+    name: str
+    description: str
+    parameters: List[ToolParameter]
+    tool_type: ToolType = ToolType.FUNCTION
+    implementation: Optional[Callable] = None
+    
+    def to_openai_format(self) -> Dict:
+        """Convert to OpenAI function calling format."""
+        properties = {}
+        required = []
+        
+        for param in self.parameters:
+            prop = {
+                "type": param.type,
+                "description": param.description
+            }
+            if param.enum:
+                prop["enum"] = param.enum
+            
+            properties[param.name] = prop
+            
+            if param.required:
+                required.append(param.name)
+        
+        return {
+            "type": "function",
+            "function": {
+                "name": self.name,
+                "description": self.description,
+                "parameters": {
+                    "type": "object",
+                    "properties": properties,
+                    "required": required
+                }
+            }
+        }
+
+
+class ToolRegistry:
+    """Registry for managing available tools."""
+    
+    def __init__(self):
+        self.tools: Dict[str, Tool] = {}
+        self._register_default_tools()
+    
+    def register(self, tool: Tool):
+        """Register a new tool."""
+        self.tools[tool.name] = tool
+        logger.info(f"Registered tool: {tool.name}")
+    
+    def get(self, name: str) -> Optional[Tool]:
+        """Get tool by name."""
+        return self.tools.get(name)
+    
+    def list_tools(self) -> List[str]:
+        """List all registered tools."""
+        return list(self.tools.keys())
+    
+    def get_tools_schema(self) -> List[Dict]:
+        """Get tools in OpenAI schema format."""
+        return [tool.to_openai_format() for tool in self.tools.values()]
+    
+    def _register_default_tools(self):
+        """Register default built-in tools."""
+        
+        # Web search tool
+        search_tool = Tool(
+            name="web_search",
+            description="Search the web for current information",
+            parameters=[
+                ToolParameter(
+                    name="query",
+                    type="string",
+                    description="Search query"
+                ),
+                ToolParameter(
+                    name="num_results",
+                    type="integer",
+                    description="Number of results to return",
+                    required=False,
+                    default=5
+                )
+            ],
+            tool_type=ToolType.API
+        )
+        self.register(search_tool)
+        
+        # Calculator tool
+        calc_tool = Tool(
+            name="calculator",
+            description="Perform mathematical calculations",
+            parameters=[
+                ToolParameter(
+                    name="expression",
+                    type="string",
+                    description="Mathematical expression to evaluate"
+                )
+            ],
+            tool_type=ToolType.FUNCTION,
+            implementation=self._calculator_impl
+        )
+        self.register(calc_tool)
+        
+        # Code execution tool
+        code_tool = Tool(
+            name="execute_python",
+            description="Execute Python code and return the result",
+            parameters=[
+                ToolParameter(
+                    name="code",
+                    type="string",
+                    description="Python code to execute"
+                )
+            ],
+            tool_type=ToolType.CODE_INTERPRETER,
+            implementation=self._python_executor
+        )
+        self.register(code_tool)
+    
+    def _calculator_impl(self, expression: str) -> Dict:
+        """Calculator implementation."""
+        try:
+            # Safe eval for math expressions
+            import ast
+            import operator
+            
+            operators = {
+                ast.Add: operator.add,
+                ast.Sub: operator.sub,
+                ast.Mult: operator.mul,
+                ast.Div: operator.truediv,
+                ast.Pow: operator.pow,
+                ast.USub: operator.neg
+            }
+            
+            def eval_expr(node):
+                if isinstance(node, ast.Num):
+                    return node.n
+                elif isinstance(node, ast.BinOp):
+                    return operators[type(node.op)](
+                        eval_expr(node.left),
+                        eval_expr(node.right)
+                    )
+                elif isinstance(node, ast.UnaryOp):
+                    return operators[type(node.op)](eval_expr(node.operand))
+                else:
+                    raise TypeError(node)
+            
+            result = eval_expr(ast.parse(expression, mode='eval').body)
+            return {"result": result, "success": True}
+        
+        except Exception as e:
+            return {"error": str(e), "success": False}
+    
+    def _python_executor(self, code: str) -> Dict:
+        """Python code executor (sandboxed)."""
+        try:
+            # Basic sandbox - restrict imports and dangerous functions
+            restricted_builtins = {
+                'print': print,
+                'range': range,
+                'len': len,
+                'str': str,
+                'int': int,
+                'float': float,
+                'list': list,
+                'dict': dict,
+                'sum': sum,
+                'max': max,
+                'min': min
+            }
+            
+            output = []
+            
+            def custom_print(*args, **kwargs):
+                output.append(' '.join(str(arg) for arg in args))
+            
+            restricted_builtins['print'] = custom_print
+            
+            exec(code, {"__builtins__": restricted_builtins}, {})
+            
+            return {
+                "output": '\n'.join(output),
+                "success": True
+            }
+        
+        except Exception as e:
+            return {"error": str(e), "success": False}
+
+
+class FunctionCallParser:
+    """Parse function calls from model output."""
+    
+    @staticmethod
+    def extract_function_calls(text: str) -> List[Dict]:
+        """
+        Extract function calls from model output.
+        
+        Args:
+            text: Model output text
+            
+        Returns:
+            List of function call dictionaries
+        """
+        function_calls = []
+        
+        # Look for JSON function call format
+        import re
+        
+        # Pattern: {"function": "name", "parameters": {...}}
+        pattern = r'\{["\']function["\']\s*:\s*["\']([^"\']+)["\']\s*,\s*["\']parameters["\']\s*:\s*(\{[^}]+\})\}'
+        
+        matches = re.finditer(pattern, text)
+        
+        for match in matches:
+            try:
+                func_name = match.group(1)
+                params_str = match.group(2)
+                params = json.loads(params_str)
+                
+                function_calls.append({
+                    "function": func_name,
+                    "parameters": params
+                })
+            except json.JSONDecodeError:
+                continue
+        
+        return function_calls
+    
+    @staticmethod
+    def format_function_result(
+        function_name: str,
+        result: Dict
+    ) -> str:
+        """Format function result for model."""
+        return f"\n[Function {function_name} returned: {json.dumps(result)}]\n"
+
+
+class HelionToolSystem:
+    """
+    Complete tool system for Helion-V1.5.
+    Manages tool registration, execution, and integration.
+    """
+    
+    def __init__(self, model, tokenizer):
+        self.model = model
+        self.tokenizer = tokenizer
+        self.registry = ToolRegistry()
+        self.parser = FunctionCallParser()
+    
+    def add_tool(self, tool: Tool):
+        """Add a custom tool."""
+        self.registry.register(tool)
+    
+    def generate_with_tools(
+        self,
+        messages: List[Dict[str, str]],
+        tools: Optional[List[str]] = None,
+        max_iterations: int = 5,
+        **kwargs
+    ) -> Dict[str, Any]:
+        """
+        Generate response with tool calling capability.
+        
+        Args:
+            messages: Chat messages
+            tools: List of tool names to make available (None = all)
+            max_iterations: Max tool calling iterations
+            **kwargs: Generation parameters
+            
+        Returns:
+            Dict with response and tool execution info
+        """
+        import torch
+        
+        available_tools = tools or self.registry.list_tools()
+        tool_schemas = [
+            self.registry.get(name).to_openai_format()
+            for name in available_tools
+            if self.registry.get(name)
+        ]
+        
+        # Add tools to system message
+        system_msg = {
+            "role": "system",
+            "content": f"""You have access to the following tools:
+
+{json.dumps(tool_schemas, indent=2)}
+
+To use a tool, output JSON in this format:
+{{"function": "tool_name", "parameters": {{"param": "value"}}}}
+
+After receiving tool results, continue your response."""
+        }
+        
+        messages_with_tools = [system_msg] + messages
+        tool_calls = []
+        
+        for iteration in range(max_iterations):
+            # Generate
+            input_ids = self.tokenizer.apply_chat_template(
+                messages_with_tools,
+                add_generation_prompt=True,
+                return_tensors="pt"
+            ).to(self.model.device)
+            
+            with torch.no_grad():
+                output = self.model.generate(
+                    input_ids,
+                    max_new_tokens=kwargs.get('max_new_tokens', 512),
+                    temperature=kwargs.get('temperature', 0.7),
+                    top_p=kwargs.get('top_p', 0.9),
+                    do_sample=kwargs.get('do_sample', True),
+                    pad_token_id=self.tokenizer.pad_token_id,
+                    eos_token_id=self.tokenizer.eos_token_id
+                )
+            
+            response = self.tokenizer.decode(
+                output[0][input_ids.shape[1]:],
+                skip_special_tokens=True
+            )
+            
+            # Check for function calls
+            calls = self.parser.extract_function_calls(response)
+            
+            if not calls:
+                # No more function calls, return final response
+                return {
+                    "response": response.strip(),
+                    "tool_calls": tool_calls,
+                    "iterations": iteration + 1
+                }
+            
+            # Execute function calls
+            for call in calls:
+                func_name = call["function"]
+                params = call["parameters"]
+                
+                tool = self.registry.get(func_name)
+                if not tool or not tool.implementation:
+                    result = {"error": f"Tool {func_name} not found or not executable"}
+                else:
+                    result = tool.implementation(**params)
+                
+                tool_calls.append({
+                    "function": func_name,
+                    "parameters": params,
+                    "result": result
+                })
+                
+                # Add result to conversation
+                result_msg = self.parser.format_function_result(func_name, result)
+                messages_with_tools.append({
+                    "role": "assistant",
+                    "content": response
+                })
+                messages_with_tools.append({
+                    "role": "system",
+                    "content": result_msg
+                })
+        
+        return {
+            "response": "Max iterations reached",
+            "tool_calls": tool_calls,
+            "iterations": max_iterations
+        }
+
+
+# Example custom tool
+def create_weather_tool() -> Tool:
+    """Example: Create a weather lookup tool."""
+    
+    def get_weather(location: str, units: str = "celsius") -> Dict:
+        """Mock weather implementation."""
+        return {
+            "location": location,
+            "temperature": 22 if units == "celsius" else 72,
+            "conditions": "Partly cloudy",
+            "units": units
+        }
+    
+    return Tool(
+        name="get_weather",
+        description="Get current weather for a location",
+        parameters=[
+            ToolParameter(
+                name="location",
+                type="string",
+                description="City name or location"
+            ),
+            ToolParameter(
+                name="units",
+                type="string",
+                description="Temperature units",
+                required=False,
+                enum=["celsius", "fahrenheit"],
+                default="celsius"
+            )
+        ],
+        tool_type=ToolType.API,
+        implementation=get_weather
+    )
+
+
+if __name__ == "__main__":
+    # Demo tool system
+    registry = ToolRegistry()
+    
+    print("Registered Tools:")
+    print("="*60)
+    for tool_name in registry.list_tools():
+        tool = registry.get(tool_name)
+        print(f"\n{tool.name}:")
+        print(f"  Description: {tool.description}")
+        print(f"  Type: {tool.tool_type.value}")
+        print(f"  Parameters: {[p.name for p in tool.parameters]}")
+    
+    # Test calculator
+    print("\n" + "="*60)
+    print("Testing Calculator:")
+    calc = registry.get("calculator")
+    result = calc.implementation(expression="2 + 3 * 4")
+    print(f"  2 + 3 * 4 = {result}")
+    
+    # Test code executor
+    print("\nTesting Code Executor:")
+    executor = registry.get("execute_python")
+    result = executor.implementation(code="print('Hello'); print(sum([1,2,3]))")
+    print(f"  Output: {result}")