Spaces:
Sleeping
Sleeping
File size: 8,666 Bytes
62a3701 | 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 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 | """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)
|