Spaces:
Sleeping
Sleeping
| """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) | |