"""Model inference runner with tool calling support.""" import json import re from typing import Any, Dict, List, Optional from huggingface_hub import InferenceClient from metrics.performance_tracker import PerformanceTracker, estimate_token_count from tools.tool_executor import ToolExecutor class ModelRunner: """Runs model inference with tool calling support.""" def __init__( self, model_id: str, model_name: str, hf_token: Optional[str] = None, max_iterations: int = 5, ): """Initialize model runner. Args: model_id: Hugging Face model ID model_name: Display name for the model hf_token: Optional Hugging Face API token max_iterations: Maximum number of tool calling iterations """ self.model_id = model_id self.model_name = model_name self.max_iterations = max_iterations self.client = InferenceClient(token=hf_token) self.tool_executor = ToolExecutor() self.tracker = PerformanceTracker(model_name) def run( self, prompt: str, tools: List[Dict[str, Any]], system_message: Optional[str] = None, ) -> Dict[str, Any]: """Run model inference with tool calling. Args: prompt: User prompt tools: List of available tools system_message: Optional system message Returns: Dictionary with 'output', 'metrics', 'tools_used', and 'conversation_history' """ self.tracker.reset() self.tracker.start() self.tool_executor.clear_history() try: # Build system message with tools full_system_message = self._build_system_message(tools, system_message) # Initialize conversation messages = [] if full_system_message: messages.append({"role": "system", "content": full_system_message}) messages.append({"role": "user", "content": prompt}) conversation_history = [] final_output = "" # Tool calling loop for iteration in range(self.max_iterations): # Get model response response = self._get_model_response(messages) if not response: break # Track tokens self.tracker.record_tokens(estimate_token_count(response)) conversation_history.append({"role": "assistant", "content": response}) # Check for tool calls in response tool_calls = self._extract_tool_calls(response) if not tool_calls: # No more tool calls, we're done final_output = response break # Execute tool calls tool_results = [] for tool_call in tool_calls: self.tracker.start_tool_execution() result = self.tool_executor.execute( tool_call["name"], tool_call["arguments"], ) self.tracker.end_tool_execution() tool_results.append(result) # Add tool results to conversation tool_response = self._format_tool_results(tool_calls, tool_results) messages.append({"role": "assistant", "content": response}) messages.append({"role": "user", "content": f"Tool results:\n{tool_response}"}) conversation_history.append({"role": "tool", "content": tool_response}) # If we hit max iterations without final output if not final_output and conversation_history: final_output = conversation_history[-1].get("content", "") self.tracker.end(success=True) return { "output": final_output, "metrics": self.tracker.get_metrics(), "tools_used": [h["tool"] for h in self.tool_executor.get_execution_history()], "conversation_history": conversation_history, } except Exception as e: self.tracker.end(success=False, error_message=str(e)) return { "output": f"Error: {str(e)}", "metrics": self.tracker.get_metrics(), "tools_used": [], "conversation_history": [], } def _build_system_message( self, tools: List[Dict[str, Any]], custom_message: Optional[str] = None, ) -> str: """Build system message with tool descriptions. Args: tools: List of available tools custom_message: Optional custom system message Returns: Complete system message """ base_message = custom_message or "You are a helpful AI assistant with access to tools." if not tools: return base_message tool_descriptions = [] for tool in tools: func = tool["function"] tool_desc = f"- **{func['name']}**: {func['description']}" tool_descriptions.append(tool_desc) tools_section = "\n\nAvailable tools:\n" + "\n".join(tool_descriptions) tools_section += ( "\n\nTo use a tool, respond with: TOOL_CALL: {\"name\": \"tool_name\", " '\"arguments\": {\'arg1\': \'value1\'}}' ) return base_message + tools_section def _get_model_response(self, messages: List[Dict[str, str]]) -> str: """Get response from model. Args: messages: Conversation messages Returns: Model response text """ try: # Format messages for the API formatted_prompt = self._format_messages_for_api(messages) # Call Hugging Face Inference API response = self.client.text_generation( formatted_prompt, model=self.model_id, max_new_tokens=512, temperature=0.7, return_full_text=False, ) return response.strip() if response else "" except Exception as e: raise RuntimeError(f"Model inference failed: {str(e)}") def _format_messages_for_api(self, messages: List[Dict[str, str]]) -> str: """Format messages for the inference API. Args: messages: List of message dictionaries Returns: Formatted prompt string """ # Simple chat template formatting formatted_parts = [] for msg in messages: role = msg["role"] content = msg["content"] if role == "system": formatted_parts.append(f"System: {content}") elif role == "user": formatted_parts.append(f"User: {content}") elif role == "assistant": formatted_parts.append(f"Assistant: {content}") formatted_parts.append("Assistant:") return "\n\n".join(formatted_parts) def _extract_tool_calls(self, response: str) -> List[Dict[str, Any]]: """Extract tool calls from model response. Args: response: Model response text Returns: List of tool call dictionaries """ tool_calls = [] # Look for TOOL_CALL: {json} pattern pattern = r'TOOL_CALL:\s*(\{[^}]+\})' matches = re.findall(pattern, response, re.IGNORECASE) for match in matches: try: tool_call = json.loads(match) if "name" in tool_call and "arguments" in tool_call: tool_calls.append(tool_call) except json.JSONDecodeError: continue return tool_calls def _format_tool_results( self, tool_calls: List[Dict[str, Any]], results: List[Dict[str, Any]], ) -> str: """Format tool results for conversation. Args: tool_calls: List of tool calls results: List of tool execution results Returns: Formatted results string """ formatted = [] for tool_call, result in zip(tool_calls, results): tool_name = tool_call["name"] if result["success"]: formatted.append(f"{tool_name}: {result['result']}") else: formatted.append(f"{tool_name}: Error - {result.get('error', 'Unknown error')}") return "\n".join(formatted)