Accusys
39a2cbc65b
- get_face_groups_handler: COALESCE(tp.name, tn.label) for name consistency - sync_file_status: compare JSON vs pre_chunks (not chunk table) - face consistency: compare frames.len() not total_faces - cleanup 2 ghost records with NULL file_name/file_path - replace identity_agent with face_dedup in pipeline stages - remove identity_agent_api.rs and all references - update required_processors to match actual processors - update AGENTS.md with team responsibilities - add Studio pipeline changes documentation
791 lines
30 KiB
Python
791 lines
30 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Momentry Tool Calling Module
|
|
============================
|
|
Supports sequential multi-tool execution using Ollama API.
|
|
|
|
Tools:
|
|
- query_postgres: PostgreSQL database queries
|
|
- search_qdrant: Vector similarity search
|
|
- execute_bash: Bash command execution
|
|
- call_api: HTTP API calls
|
|
|
|
Version: 1.1.0
|
|
Updated: 2026-07-26
|
|
Changes:
|
|
- Added embedding server health check
|
|
- Enhanced bash safety patterns
|
|
- Added logging support
|
|
- Environment variable for Qdrant collection
|
|
"""
|
|
|
|
import json
|
|
import logging
|
|
import os
|
|
import re
|
|
import subprocess
|
|
import requests
|
|
from typing import Any, Callable, Dict, List, Optional, Tuple
|
|
from dataclasses import dataclass
|
|
from datetime import datetime, date
|
|
|
|
# Configure logging
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
|
)
|
|
logger = logging.getLogger('tool_caller')
|
|
|
|
|
|
class DateTimeEncoder(json.JSONEncoder):
|
|
"""Custom JSON encoder for datetime objects"""
|
|
def default(self, obj):
|
|
if isinstance(obj, (datetime, date)):
|
|
return obj.isoformat()
|
|
return super().default(obj)
|
|
|
|
|
|
@dataclass
|
|
class ToolResult:
|
|
"""Result from tool execution"""
|
|
success: bool
|
|
data: Any
|
|
error: Optional[str] = None
|
|
execution_time_ms: float = 0
|
|
|
|
|
|
class ToolRegistry:
|
|
"""Registry for available tools"""
|
|
|
|
def __init__(self):
|
|
self._tools: Dict[str, Dict[str, Any]] = {}
|
|
self._executors: Dict[str, Callable] = {}
|
|
|
|
def register(
|
|
self,
|
|
name: str,
|
|
description: str,
|
|
parameters: Dict[str, Any],
|
|
executor: Callable[[Dict[str, Any]], ToolResult]
|
|
):
|
|
"""Register a new tool"""
|
|
self._tools[name] = {
|
|
"type": "function",
|
|
"function": {
|
|
"name": name,
|
|
"description": description,
|
|
"parameters": parameters
|
|
}
|
|
}
|
|
self._executors[name] = executor
|
|
logger.debug(f"Registered tool: {name}")
|
|
|
|
def get_definitions(self) -> List[Dict[str, Any]]:
|
|
"""Get all tool definitions for API calls"""
|
|
return list(self._tools.values())
|
|
|
|
def has_tool(self, name: str) -> bool:
|
|
"""Check if tool exists"""
|
|
return name in self._executors
|
|
|
|
def execute(self, name: str, arguments: Dict[str, Any]) -> ToolResult:
|
|
"""Execute a tool by name"""
|
|
if name not in self._executors:
|
|
logger.error(f"Unknown tool: {name}")
|
|
return ToolResult(
|
|
success=False,
|
|
data=None,
|
|
error=f"Unknown tool: {name}"
|
|
)
|
|
|
|
start_time = datetime.now()
|
|
logger.info(f"Executing tool: {name} with args: {json.dumps(arguments, ensure_ascii=False)[:200]}")
|
|
|
|
try:
|
|
result = self._executors[name](arguments)
|
|
elapsed = (datetime.now() - start_time).total_seconds() * 1000
|
|
result.execution_time_ms = elapsed
|
|
|
|
if result.success:
|
|
logger.info(f"Tool {name} succeeded in {elapsed:.1f}ms")
|
|
else:
|
|
logger.error(f"Tool {name} failed: {result.error}")
|
|
|
|
return result
|
|
except Exception as e:
|
|
elapsed = (datetime.now() - start_time).total_seconds() * 1000
|
|
logger.error(f"Tool {name} exception: {e}")
|
|
return ToolResult(
|
|
success=False,
|
|
data=None,
|
|
error=str(e),
|
|
execution_time_ms=elapsed
|
|
)
|
|
|
|
|
|
class OllamaToolCaller:
|
|
"""Tool caller using Ollama API with improved stability"""
|
|
|
|
def __init__(
|
|
self,
|
|
base_url: str = "http://localhost:11434",
|
|
model: str = "llama3.1:8b",
|
|
max_iterations: int = 10,
|
|
max_tool_calls: int = 5
|
|
):
|
|
self.base_url = base_url
|
|
self.model = model
|
|
self.max_iterations = max_iterations
|
|
self.max_tool_calls = max_tool_calls
|
|
self.registry = ToolRegistry()
|
|
|
|
# System prompt optimized for tool calling
|
|
self.system_prompt = """You are a data assistant with access to tools.
|
|
|
|
CORE RULES:
|
|
1. When you receive a tool result, you MUST provide a final answer - do NOT call more tools
|
|
2. Call exactly ONE tool per response
|
|
3. Answer based on tool results, not assumptions
|
|
4. If tool fails, report the error and stop
|
|
|
|
IMPORTANT: After receiving ANY tool result, respond with a clear answer to the user."""
|
|
|
|
# Track tool calls for loop prevention
|
|
self._tool_call_history: List[str] = []
|
|
|
|
logger.info(f"Initialized OllamaToolCaller with model={model}")
|
|
|
|
def chat(self, messages: List[Dict[str, Any]]) -> Dict[str, Any]:
|
|
"""Send chat request to Ollama"""
|
|
url = f"{self.base_url}/api/chat"
|
|
payload = {
|
|
"model": self.model,
|
|
"messages": messages,
|
|
"tools": self.registry.get_definitions(),
|
|
"stream": False
|
|
}
|
|
|
|
logger.debug(f"Sending chat request to {url}")
|
|
response = requests.post(url, json=payload, timeout=120)
|
|
response.raise_for_status()
|
|
return response.json()
|
|
|
|
def execute_tool_call(self, tool_call: Dict[str, Any]) -> ToolResult:
|
|
"""Execute a single tool call with flexible argument parsing"""
|
|
func = tool_call.get("function", {})
|
|
name = func.get("name", "")
|
|
|
|
# Handle both dict and string arguments
|
|
arguments = func.get("arguments", {})
|
|
if isinstance(arguments, str):
|
|
try:
|
|
arguments = json.loads(arguments)
|
|
except json.JSONDecodeError:
|
|
arguments = {}
|
|
|
|
# Normalize arguments based on tool name
|
|
arguments = self._normalize_arguments(name, arguments)
|
|
|
|
return self.registry.execute(name, arguments)
|
|
|
|
def _normalize_arguments(self, tool_name: str, args: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""Normalize arguments to match expected format"""
|
|
if not args:
|
|
return args
|
|
|
|
# For query_postgres, ensure we have 'query' parameter
|
|
if tool_name == "query_postgres":
|
|
if "query" not in args:
|
|
for key in ["query_text", "sql", "sql_query", "statement"]:
|
|
if key in args:
|
|
args["query"] = args[key]
|
|
break
|
|
|
|
# For search_qdrant, ensure we have 'query_text' parameter
|
|
elif tool_name == "search_qdrant":
|
|
if "query_text" not in args:
|
|
for key in ["query", "search", "text"]:
|
|
if key in args:
|
|
args["query_text"] = args[key]
|
|
break
|
|
|
|
return args
|
|
|
|
def _get_tool_key(self, name: str, arguments: Any) -> str:
|
|
"""Generate unique key for tool call"""
|
|
if isinstance(arguments, str):
|
|
try:
|
|
arguments = json.loads(arguments)
|
|
except:
|
|
arguments = {"raw": arguments}
|
|
return f"{name}:{json.dumps(arguments, sort_keys=True)}"
|
|
|
|
def _is_duplicate_call(self, tool_key: str) -> bool:
|
|
"""Check if this tool call was already made"""
|
|
is_dup = tool_key in self._tool_call_history
|
|
if is_dup:
|
|
logger.warning(f"Duplicate tool call detected: {tool_key[:50]}")
|
|
return is_dup
|
|
|
|
def _add_to_history(self, tool_key: str):
|
|
"""Add tool call to history"""
|
|
self._tool_call_history.append(tool_key)
|
|
|
|
def _extract_tool_from_text(self, text: str) -> Optional[Tuple[str, Dict[str, Any]]]:
|
|
"""Extract tool call from text content (when model outputs JSON instead of tool_calls)"""
|
|
if not text or '{"name":' not in text:
|
|
return None
|
|
|
|
try:
|
|
start_idx = text.find('{"name":')
|
|
if start_idx == -1:
|
|
return None
|
|
|
|
# Find the end of the JSON object
|
|
depth = 0
|
|
end_idx = start_idx
|
|
for i in range(start_idx, len(text)):
|
|
if text[i] == '{':
|
|
depth += 1
|
|
elif text[i] == '}':
|
|
depth -= 1
|
|
if depth == 0:
|
|
end_idx = i + 1
|
|
break
|
|
|
|
if end_idx > start_idx:
|
|
json_str = text[start_idx:end_idx]
|
|
# Fix common JSON issues
|
|
json_str = json_str.replace('\\ ', ' ')
|
|
json_str = json_str.replace('\\n', '\n')
|
|
|
|
tool_data = json.loads(json_str)
|
|
if 'name' in tool_data:
|
|
name = tool_data['name']
|
|
params = tool_data.get('parameters', tool_data.get('arguments', {}))
|
|
if self.registry.has_tool(name):
|
|
logger.debug(f"Extracted tool from text: {name}")
|
|
return (name, params)
|
|
except Exception as e:
|
|
logger.debug(f"Failed to extract tool from text: {e}")
|
|
|
|
return None
|
|
|
|
def _force_answer(self, messages: List[Dict[str, Any]], tool_results: str) -> str:
|
|
"""Force the model to provide a final answer based on tool results"""
|
|
messages.append({
|
|
"role": "user",
|
|
"content": f"Based on the tool results above, please provide your final answer now. Do not call any more tools.\n\nTool results:\n{tool_results}"
|
|
})
|
|
|
|
url = f"{self.base_url}/api/chat"
|
|
payload = {
|
|
"model": self.model,
|
|
"messages": messages,
|
|
"stream": False
|
|
}
|
|
|
|
try:
|
|
logger.info("Forcing final answer from model")
|
|
response = requests.post(url, json=payload, timeout=120)
|
|
response.raise_for_status()
|
|
return response.json().get("message", {}).get("content", "Unable to generate answer.")
|
|
except Exception as e:
|
|
logger.error(f"Failed to force answer: {e}")
|
|
return f"Tool execution completed. Results: {tool_results}"
|
|
|
|
def run(self, user_query: str) -> str:
|
|
"""
|
|
Run tool calling loop with user query.
|
|
|
|
Args:
|
|
user_query: The user's question or request
|
|
|
|
Returns:
|
|
Final text response from the model
|
|
"""
|
|
logger.info(f"Starting tool call loop for query: {user_query[:100]}")
|
|
self._tool_call_history = []
|
|
|
|
messages = [
|
|
{"role": "system", "content": self.system_prompt},
|
|
{"role": "user", "content": user_query}
|
|
]
|
|
|
|
tool_results_collected = []
|
|
|
|
for iteration in range(self.max_iterations):
|
|
logger.info(f"Iteration {iteration + 1}/{self.max_iterations}")
|
|
|
|
# Call LLM
|
|
response = self.chat(messages)
|
|
|
|
message = response.get("message", {})
|
|
content = message.get("content", "")
|
|
tool_calls = message.get("tool_calls", [])
|
|
|
|
# If no tool calls in response, check if content has text tool call
|
|
if not tool_calls:
|
|
extracted = self._extract_tool_from_text(content)
|
|
if extracted:
|
|
name, params = extracted
|
|
tool_key = self._get_tool_key(name, params)
|
|
|
|
if not self._is_duplicate_call(tool_key):
|
|
self._add_to_history(tool_key)
|
|
tool_call = {"function": {"name": name, "arguments": params}}
|
|
tool_result = self.execute_tool_call(tool_call)
|
|
|
|
# Build result string with error handling
|
|
if tool_result.success:
|
|
result_str = json.dumps({
|
|
"success": True,
|
|
"data": tool_result.data
|
|
}, ensure_ascii=False, cls=DateTimeEncoder)
|
|
else:
|
|
result_str = json.dumps({
|
|
"success": False,
|
|
"error": tool_result.error,
|
|
"suggestion": "Try a different tool or rephrase your query."
|
|
}, ensure_ascii=False)
|
|
|
|
messages.append({"role": "assistant", "content": "", "tool_calls": [tool_call]})
|
|
messages.append({"role": "user", "content": f"Tool result: {result_str}"})
|
|
tool_results_collected.append(f"{name}: {result_str[:500]}")
|
|
continue
|
|
|
|
# No tool call - return content
|
|
if content:
|
|
logger.info(f"Returning final answer after {iteration + 1} iterations")
|
|
return content
|
|
else:
|
|
return self._force_answer(messages, "\n".join(tool_results_collected) if tool_results_collected else "No tool was called")
|
|
|
|
# Execute first tool call (sequential)
|
|
tool_call = tool_calls[0]
|
|
func = tool_call.get("function", {})
|
|
tool_name = func.get("name", "")
|
|
arguments = func.get("arguments", {})
|
|
|
|
tool_key = self._get_tool_key(tool_name, arguments)
|
|
|
|
# Check for duplicate
|
|
if self._is_duplicate_call(tool_key):
|
|
return self._force_answer(messages, "\n".join(tool_results_collected) if tool_results_collected else "Tool already executed")
|
|
|
|
# Track and execute
|
|
self._add_to_history(tool_key)
|
|
tool_result = self.execute_tool_call(tool_call)
|
|
|
|
# Check if we've hit the tool call limit
|
|
if len(self._tool_call_history) >= self.max_tool_calls:
|
|
if tool_result.success:
|
|
result_str = json.dumps({
|
|
"success": True,
|
|
"data": tool_result.data
|
|
}, ensure_ascii=False, cls=DateTimeEncoder)
|
|
else:
|
|
result_str = json.dumps({
|
|
"success": False,
|
|
"error": tool_result.error,
|
|
"suggestion": "Try a different tool or rephrase your query."
|
|
}, ensure_ascii=False)
|
|
|
|
tool_results_collected.append(f"{tool_name}: {result_str[:500]}")
|
|
return self._force_answer(messages, "\n".join(tool_results_collected))
|
|
|
|
# Add to messages
|
|
messages.append({"role": "assistant", "content": "", "tool_calls": [tool_call]})
|
|
|
|
if tool_result.success:
|
|
result_str = json.dumps({
|
|
"success": True,
|
|
"data": tool_result.data
|
|
}, ensure_ascii=False, cls=DateTimeEncoder)
|
|
else:
|
|
result_str = json.dumps({
|
|
"success": False,
|
|
"error": tool_result.error,
|
|
"suggestion": "Try a different tool or rephrase your query."
|
|
}, ensure_ascii=False)
|
|
|
|
messages.append({"role": "user", "content": f"Tool result: {result_str}"})
|
|
tool_results_collected.append(f"{tool_name}: {result_str[:500]}")
|
|
|
|
# Max iterations reached
|
|
logger.warning(f"Max iterations ({self.max_iterations}) reached")
|
|
return self._force_answer(messages, "\n".join(tool_results_collected) if tool_results_collected else "Max iterations reached")
|
|
|
|
def register_default_tools(self, db_url: str = None, qdrant_url: str = None):
|
|
"""Register default tools with connection strings"""
|
|
|
|
# Get default collection from environment variable
|
|
default_collection = os.environ.get("QDRANT_DEFAULT_COLLECTION", "momentry_rule1")
|
|
logger.info(f"Using default Qdrant collection: {default_collection}")
|
|
|
|
# PostgreSQL tool
|
|
def query_postgres(args: Dict[str, Any]) -> ToolResult:
|
|
import psycopg2
|
|
import psycopg2.extras
|
|
|
|
query = args.get("query", "")
|
|
if not query:
|
|
return ToolResult(success=False, data=None, error="No query provided")
|
|
|
|
logger.info(f"Executing PostgreSQL query: {query[:100]}")
|
|
|
|
conn = psycopg2.connect(db_url or "postgres://accusys@localhost:5432/momentry")
|
|
try:
|
|
with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
|
|
cur.execute(query)
|
|
if cur.description:
|
|
rows = cur.fetchall()
|
|
logger.info(f"Query returned {len(rows)} rows")
|
|
return ToolResult(
|
|
success=True,
|
|
data={"rows": [dict(r) for r in rows], "row_count": len(rows)}
|
|
)
|
|
else:
|
|
conn.commit()
|
|
logger.info(f"Query affected {cur.rowcount} rows")
|
|
return ToolResult(
|
|
success=True,
|
|
data={"affected_rows": cur.rowcount}
|
|
)
|
|
finally:
|
|
conn.close()
|
|
|
|
self.registry.register(
|
|
name="query_postgres",
|
|
description="Execute SQL query on PostgreSQL database. Use SELECT for queries, INSERT/UPDATE/DELETE for modifications.",
|
|
parameters={
|
|
"type": "object",
|
|
"properties": {
|
|
"query": {
|
|
"type": "string",
|
|
"description": "SQL query to execute"
|
|
}
|
|
},
|
|
"required": ["query"]
|
|
},
|
|
executor=query_postgres
|
|
)
|
|
|
|
# Qdrant tool with health check
|
|
def search_qdrant(args: Dict[str, Any]) -> ToolResult:
|
|
import requests as req
|
|
|
|
collection = args.get("collection", default_collection)
|
|
query_text = args.get("query_text", "")
|
|
limit = args.get("limit", 10)
|
|
|
|
# Get Qdrant API key from args or environment
|
|
api_key = args.get("api_key") or os.environ.get("QDRANT_API_KEY", "Test3200Test3200Test3200")
|
|
|
|
if not query_text:
|
|
return ToolResult(success=False, data=None, error="No query text provided")
|
|
|
|
logger.info(f"Searching Qdrant collection '{collection}' for: {query_text[:50]}")
|
|
|
|
# Check embedding server health
|
|
try:
|
|
embed_health = req.get("http://localhost:11436/health", timeout=5)
|
|
if embed_health.status_code != 200:
|
|
return ToolResult(
|
|
success=False,
|
|
data=None,
|
|
error=f"Embedding server not healthy (status={embed_health.status_code})"
|
|
)
|
|
except req.exceptions.ConnectionError:
|
|
return ToolResult(
|
|
success=False,
|
|
data=None,
|
|
error="Embedding server not available at http://localhost:11436"
|
|
)
|
|
except Exception as e:
|
|
return ToolResult(
|
|
success=False,
|
|
data=None,
|
|
error=f"Embedding server health check failed: {str(e)}"
|
|
)
|
|
|
|
# Get embedding
|
|
try:
|
|
embed_url = "http://localhost:11436/v1/embeddings"
|
|
embed_resp = req.post(embed_url, json={"input": query_text}, timeout=30)
|
|
embed_resp.raise_for_status()
|
|
embed_data = embed_resp.json()
|
|
# Handle both response formats
|
|
if "data" in embed_data:
|
|
embedding = embed_data["data"][0]["embedding"]
|
|
elif "embeddings" in embed_data:
|
|
embedding = embed_data["embeddings"][0]
|
|
else:
|
|
return ToolResult(
|
|
success=False,
|
|
data=None,
|
|
error=f"Unexpected embedding response format: {list(embed_data.keys())}"
|
|
)
|
|
logger.info(f"Generated embedding (dim={len(embedding)})")
|
|
except Exception as e:
|
|
return ToolResult(
|
|
success=False,
|
|
data=None,
|
|
error=f"Embedding generation failed: {str(e)}"
|
|
)
|
|
|
|
# Check Qdrant collection exists
|
|
qdrant_url_base = qdrant_url or "http://localhost:6333"
|
|
headers = {"api-key": api_key}
|
|
try:
|
|
collections_url = f"{qdrant_url_base}/collections"
|
|
coll_resp = req.get(collections_url, headers=headers, timeout=10)
|
|
coll_resp.raise_for_status()
|
|
collections = [c["name"] for c in coll_resp.json().get("result", {}).get("collections", [])]
|
|
logger.info(f"Available Qdrant collections: {collections}")
|
|
|
|
if collection not in collections:
|
|
return ToolResult(
|
|
success=False,
|
|
data=None,
|
|
error=f"Collection '{collection}' not found. Available: {', '.join(collections)}"
|
|
)
|
|
except Exception as e:
|
|
return ToolResult(
|
|
success=False,
|
|
data=None,
|
|
error=f"Qdrant connection failed: {str(e)}"
|
|
)
|
|
|
|
# Execute search
|
|
search_url = f"{qdrant_url_base}/collections/{collection}/points/search"
|
|
search_payload = {"vector": embedding, "limit": limit, "with_payload": True}
|
|
|
|
try:
|
|
search_resp = req.post(search_url, json=search_payload, headers=headers, timeout=30)
|
|
search_resp.raise_for_status()
|
|
results = search_resp.json().get("result", [])
|
|
logger.info(f"Qdrant returned {len(results)} matches")
|
|
except Exception as e:
|
|
return ToolResult(
|
|
success=False,
|
|
data=None,
|
|
error=f"Qdrant search failed: {str(e)}"
|
|
)
|
|
|
|
return ToolResult(
|
|
success=True,
|
|
data={
|
|
"matches": [
|
|
{"id": r.get("id"), "score": r.get("score"), "payload": r.get("payload", {})}
|
|
for r in results
|
|
],
|
|
"match_count": len(results)
|
|
}
|
|
)
|
|
|
|
self.registry.register(
|
|
name="search_qdrant",
|
|
description=f"Search for similar vectors in Qdrant collection. Default collection: {default_collection}",
|
|
parameters={
|
|
"type": "object",
|
|
"properties": {
|
|
"collection": {
|
|
"type": "string",
|
|
"description": f"Qdrant collection name (default: {default_collection})",
|
|
"default": default_collection
|
|
},
|
|
"query_text": {
|
|
"type": "string",
|
|
"description": "Search query text (will be embedded)"
|
|
},
|
|
"limit": {
|
|
"type": "integer",
|
|
"description": "Maximum number of results",
|
|
"default": 10
|
|
}
|
|
},
|
|
"required": ["query_text"]
|
|
},
|
|
executor=search_qdrant
|
|
)
|
|
|
|
# Bash tool with enhanced safety
|
|
def execute_bash(args: Dict[str, Any]) -> ToolResult:
|
|
command = args.get("command", "")
|
|
timeout = args.get("timeout", 30)
|
|
|
|
if not command:
|
|
return ToolResult(success=False, data=None, error="No command provided")
|
|
|
|
# Limit command length
|
|
if len(command) > 2000:
|
|
return ToolResult(success=False, data=None, error="Command too long (max 2000 chars)")
|
|
|
|
logger.info(f"Executing bash command: {command[:100]}")
|
|
|
|
# Enhanced safety patterns
|
|
blocked_patterns = [
|
|
# File system destruction
|
|
"rm -rf /", "rm -rf /*", "mkfs", "dd if=", "> /dev/",
|
|
# Permission escalation
|
|
"sudo ", "su -", "chmod 777", "chown root",
|
|
# Remote execution
|
|
"curl | bash", "curl | sh", "wget | sh", "wget | bash",
|
|
"curl http", "wget http",
|
|
# Fork bomb / DoS
|
|
":(){", "fork", "kill -9 1",
|
|
# Network listeners
|
|
"nc -l", "netcat -l", "socat",
|
|
# Process killing
|
|
"killall", "pkill",
|
|
# Disk operations
|
|
"fdisk", "parted",
|
|
]
|
|
|
|
command_lower = command.lower()
|
|
for pattern in blocked_patterns:
|
|
if pattern in command_lower:
|
|
logger.warning(f"Blocked dangerous command pattern: {pattern}")
|
|
return ToolResult(
|
|
success=False,
|
|
data=None,
|
|
error=f"Blocked dangerous command pattern: {pattern}"
|
|
)
|
|
|
|
try:
|
|
result = subprocess.run(
|
|
command,
|
|
shell=True,
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=timeout
|
|
)
|
|
logger.info(f"Command completed with return code: {result.returncode}")
|
|
return ToolResult(
|
|
success=result.returncode == 0,
|
|
data={
|
|
"stdout": result.stdout[:5000],
|
|
"stderr": result.stderr[:2000],
|
|
"returncode": result.returncode
|
|
}
|
|
)
|
|
except subprocess.TimeoutExpired:
|
|
logger.error(f"Command timed out after {timeout}s")
|
|
return ToolResult(
|
|
success=False,
|
|
data=None,
|
|
error=f"Command timed out after {timeout}s"
|
|
)
|
|
|
|
self.registry.register(
|
|
name="execute_bash",
|
|
description="Execute a bash command on the system. Use for file operations, system checks, etc.",
|
|
parameters={
|
|
"type": "object",
|
|
"properties": {
|
|
"command": {
|
|
"type": "string",
|
|
"description": "Bash command to execute"
|
|
},
|
|
"timeout": {
|
|
"type": "integer",
|
|
"description": "Timeout in seconds",
|
|
"default": 30
|
|
}
|
|
},
|
|
"required": ["command"]
|
|
},
|
|
executor=execute_bash
|
|
)
|
|
|
|
# API tool
|
|
def call_api(args: Dict[str, Any]) -> ToolResult:
|
|
import requests as req
|
|
|
|
url = args.get("url", "")
|
|
method = args.get("method", "GET").upper()
|
|
headers = args.get("headers", {})
|
|
data = args.get("data")
|
|
|
|
if not url:
|
|
return ToolResult(success=False, data=None, error="No URL provided")
|
|
|
|
logger.info(f"Calling API: {method} {url[:100]}")
|
|
|
|
try:
|
|
if method == "GET":
|
|
resp = req.get(url, headers=headers, timeout=30)
|
|
elif method == "POST":
|
|
resp = req.post(url, json=data, headers=headers, timeout=30)
|
|
elif method == "PUT":
|
|
resp = req.put(url, json=data, headers=headers, timeout=30)
|
|
elif method == "DELETE":
|
|
resp = req.delete(url, headers=headers, timeout=30)
|
|
else:
|
|
return ToolResult(
|
|
success=False,
|
|
data=None,
|
|
error=f"Unsupported method: {method}"
|
|
)
|
|
|
|
logger.info(f"API response: {resp.status_code}")
|
|
return ToolResult(
|
|
success=resp.status_code < 400,
|
|
data={
|
|
"status_code": resp.status_code,
|
|
"headers": dict(resp.headers),
|
|
"body": resp.text[:10000]
|
|
}
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"API call failed: {e}")
|
|
return ToolResult(success=False, data=None, error=str(e))
|
|
|
|
self.registry.register(
|
|
name="call_api",
|
|
description="Call an external HTTP API endpoint. Use for REST API interactions.",
|
|
parameters={
|
|
"type": "object",
|
|
"properties": {
|
|
"url": {
|
|
"type": "string",
|
|
"description": "API endpoint URL"
|
|
},
|
|
"method": {
|
|
"type": "string",
|
|
"enum": ["GET", "POST", "PUT", "DELETE"],
|
|
"description": "HTTP method",
|
|
"default": "GET"
|
|
},
|
|
"headers": {
|
|
"type": "object",
|
|
"description": "HTTP headers"
|
|
},
|
|
"data": {
|
|
"type": "object",
|
|
"description": "Request body for POST/PUT"
|
|
}
|
|
},
|
|
"required": ["url"]
|
|
},
|
|
executor=call_api
|
|
)
|
|
|
|
|
|
def main():
|
|
"""Example usage"""
|
|
caller = OllamaToolCaller()
|
|
caller.register_default_tools()
|
|
|
|
query = "How many videos are in the database?"
|
|
print(f"Query: {query}")
|
|
print("-" * 50)
|
|
|
|
result = caller.run(query)
|
|
print(f"Result: {result}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|