"""
Refactored OllamaAgent implementation for OllamaTools.
This module provides the OllamaAgent class which serves as the
primary interface for interacting with Ollama and executing functions.
It has been updated to correctly work with Ollama's function calling API.
"""
import os
import re
import json
import inspect
import logging
from typing import Dict, Any, List, Optional, Callable, Union, Tuple
try:
import ollama
except ImportError:
print("Warning: ollama-python not installed. OllamaAgent will not function properly.")
from ollamatools.core.config import Config
from ollamatools.core.security import PermissionManager
from ollamatools.plugins.tools.system import execute_command, read_file, write_file, find_files
from ollamatools.plugins.tools.web import fetch_url, get_weather, search_web
from ollamatools.plugins.tools.doc import generate_markdown, extract_context, summarize_text
# Set up logging
logger = logging.getLogger(__name__)
class OllamaAgent:
"""
Agent for interacting with Ollama and executing functions.
This class serves as the main interface for OllamaTools, providing methods
to chat with Ollama and automatically execute functions based on the
conversation context.
"""
def __init__(
self,
model: str = "llama3:latest",
host: str = "http://localhost:11434",
permission_prompts: bool = True,
sudo_always_prompt: bool = True,
config: Optional[Config] = None
):
"""
Initialize the OllamaAgent with specified configuration.
Args:
model: Name of the Ollama model to use
host: Ollama server address
permission_prompts: Whether to prompt for function execution permission
sudo_always_prompt: Whether to always prompt for sudo commands
config: Configuration object (overrides individual parameters if provided)
"""
if config:
self.config = config
else:
self.config = Config(
model=model,
host=host,
permission_prompts=permission_prompts,
sudo_always_prompt=sudo_always_prompt
)
# Initialize the permission manager
self.permission_manager = PermissionManager(
prompt_enabled=self.config.permission_prompts,
sudo_always_prompt=self.config.sudo_always_prompt
)
# Set up the Ollama client
self.client = ollama.Client(host=self.config.host)
logger.info(f"Initialized Ollama client with host: {self.config.host}")
# Initialize registered functions dictionary and collect available tools
self.registered_functions = {}
self.collect_tools()
logger.info(f"OllamaAgent initialized with model={self.config.model}, host={self.config.host}")
def chat(
self,
message: str,
system_prompt: Optional[str] = None,
context: Optional[List[Dict[str, str]]] = None
) -> str:
"""
Send a message to Ollama, automatically detecting and executing functions.
Args:
message: User message
system_prompt: Optional system prompt
context: Optional conversation history
Returns:
str: Response from Ollama, potentially after executing functions
"""
if context is None:
context = []
logger.info(f"Sending message to Ollama: {message[:50]}...")
try:
# Prepare the chat request
messages = context + [{"role": "user", "content": message}]
options = {}
# Add system prompt if provided
if system_prompt:
options["system"] = system_prompt
# Generate tool definitions
tools = self.generate_tool_definitions()
# Send the request to Ollama with tools
response = self.client.chat(
model=self.config.model,
messages=messages,
options=options,
tools=tools, # Pass tools with each request
stream=False
)
# Extract the assistant's response
assistant_content = ""
tool_calls = []
# Handle different response formats based on Ollama version/client
if hasattr(response, 'message'):
# Direct object access (newer Ollama client)
assistant_content = response.message.content
tool_calls = getattr(response.message, 'tool_calls', []) or []
else:
# Dictionary access (older Ollama client)
assistant_message = response.get("message", {})
assistant_content = assistant_message.get("content", "")
tool_calls = assistant_message.get("tool_calls", [])
# Check if we have tool calls to process
if tool_calls and len(tool_calls) > 0:
logger.info(f"Detected {len(tool_calls)} tool calls")
# Process the function calls
tool_results = []
for tool_call in tool_calls:
# Extract function name and arguments based on response structure
if hasattr(tool_call, 'function'):
# Direct object access
function_name = tool_call.function.name
function_args = tool_call.function.arguments
if isinstance(function_args, str):
function_args = json.loads(function_args)
else:
# Dictionary access
function_data = tool_call.get('function', {})
function_name = function_data.get('name', '')
function_args = function_data.get('arguments', {})
if isinstance(function_args, str):
function_args = json.loads(function_args)
# Check if function exists
if function_name not in self.registered_functions:
logger.error(f"Function '{function_name}' not found in registered functions")
tool_results.append({
"error": f"Function '{function_name}' not found"
})
continue
# Get the function
func = self.registered_functions[function_name]
# Request permission
is_sudo = False
if function_name == "execute_command" and "command" in function_args:
is_sudo = self.permission_manager.detect_sudo(function_args["command"])
permission_granted = self.permission_manager.request_permission(
function_name=function_name,
args=function_args,
is_sudo=is_sudo
)
# Execute function if permission granted
if permission_granted:
try:
result = func(**function_args)
tool_results.append(result)
logger.info(f"Successfully executed {function_name}")
except Exception as e:
logger.error(f"Error executing {function_name}: {str(e)}")
tool_results.append({
"error": f"Error: {str(e)}"
})
else:
logger.warning(f"Permission denied for {function_name}")
tool_results.append({
"error": "Permission denied"
})
# Construct tool messages for the follow-up request
tool_messages = []
for i, result in enumerate(tool_results):
# Get function name from original tool call
if hasattr(tool_calls[i], 'function'):
function_name = tool_calls[i].function.name
else:
function_name = tool_calls[i].get('function', {}).get('name', '')
# Add tool message
tool_messages.append({
"role": "tool",
"name": function_name,
"content": json.dumps(result)
})
# Update conversation with function results
updated_messages = messages + [
{"role": "assistant", "content": assistant_content, "tool_calls": tool_calls}
] + tool_messages
# Send follow-up request with tool results
try:
follow_up_response = self.client.chat(
model=self.config.model,
messages=updated_messages,
options=options,
stream=False
)
# Extract final response
if hasattr(follow_up_response, 'message'):
# Direct object access
final_content = follow_up_response.message.content
else:
# Dictionary access
final_content = follow_up_response.get("message", {}).get("content", "")
logger.info(f"Received final response after tool execution")
return final_content
except Exception as e:
logger.error(f"Error in follow-up request: {str(e)}")
return f"Error processing function results: {str(e)}"
# Return the assistant's response if no tool calls
logger.info(f"Received response from Ollama")
return assistant_content
except Exception as e:
logger.error(f"Error in chat: {str(e)}")
return f"Error communicating with Ollama: {str(e)}"
def collect_tools(self):
"""
Collect all available tools from the tool modules.
This method scans the tool modules and collects function references
to be used when executing tools.
"""
logger.info("Collecting available tools...")
# Collect tool functions
tool_modules = [
("system_tools", [execute_command, read_file, write_file, find_files]),
("web_tools", [fetch_url, get_weather, search_web]),
("doc_tools", [generate_markdown, extract_context, summarize_text])
]
# Register functions
for module_name, functions in tool_modules:
for func in functions:
if func.__name__.startswith('_'):
continue
# Store the function reference
self.registered_functions[func.__name__] = func
logger.info(f"Collected {len(self.registered_functions)} tools")
def generate_tool_definitions(self) -> List[Dict[str, Any]]:
"""
Generate tool definitions for all registered functions.
Returns:
List[Dict[str, Any]]: List of tool definitions in Ollama's format
"""
tools = []
for func_name, func in self.registered_functions.items():
# Get function docstring
func_doc = inspect.getdoc(func) or "No description available"
# Generate tool definition
tool_def = self._generate_tool_definition(func_name, func, func_doc)
tools.append(tool_def)
return tools
def _generate_tool_definition(self, name: str, func: Callable, docstring: str) -> Dict[str, Any]:
"""
Generate a tool definition from a function according to Ollama's format.
Args:
name: Function name
func: Function object
docstring: Function docstring
Returns:
dict: Tool definition compatible with Ollama
"""
# Extract function signature
sig = inspect.signature(func)
# Generate description
description = self._extract_description_from_docstring(docstring)
# Generate parameters
parameters = {
"type": "object",
"properties": {},
"required": []
}
for param_name, param in sig.parameters.items():
# Skip permission_manager parameter
if param_name == "permission_manager":
continue
# Get parameter type
param_type = self._get_param_type(param)
# Get parameter description
param_desc = self._extract_param_description(docstring, param_name)
# Add parameter to properties
parameters["properties"][param_name] = {
"type": param_type,
"description": param_desc
}
# Add to required list if no default value
if param.default is inspect.Parameter.empty:
parameters["required"].append(param_name)
# Create tool definition in Ollama's format
tool_def = {
"type": "function",
"function": {
"name": name,
"description": description,
"parameters": parameters
}
}
return tool_def
def _extract_description_from_docstring(self, docstring: str) -> str:
"""
Extract the main description from a docstring.
Args:
docstring: Function docstring
Returns:
str: Function description
"""
if not docstring:
return "No description available"
# Split the docstring into lines
lines = docstring.strip().split('\n')
# Extract the first paragraph
description = []
for line in lines:
line = line.strip()
if not line or line.startswith('Args:') or line.startswith('Returns:'):
break
description.append(line)
return ' '.join(description) if description else "No description available"
def _extract_param_description(self, docstring: str, param_name: str) -> str:
"""
Extract parameter description from docstring.
Args:
docstring: Function docstring
param_name: Parameter name
Returns:
str: Parameter description
"""
if not docstring:
return "No description available"
# Look for parameters section
pattern = rf"\s+{re.escape(param_name)}:\s*(.*?)(?:\n\s+\w+:|$)"
match = re.search(pattern, docstring, re.DOTALL)
if match:
description = match.group(1).strip()
return description
return "No description available"
def _get_param_type(self, param: inspect.Parameter) -> str:
"""
Get parameter type for Ollama tool definition.
Args:
param: Parameter object
Returns:
str: Parameter type for Ollama
"""
annotation = param.annotation
if annotation is inspect.Parameter.empty:
return "string"
# Convert type annotation to string
annotation_str = str(annotation)
# Determine parameter type
if "str" in annotation_str:
return "string"
elif "int" in annotation_str:
return "integer"
elif "float" in annotation_str:
return "number"
elif "bool" in annotation_str:
return "boolean"
elif "List" in annotation_str or "list" in annotation_str:
return "array"
elif "Dict" in annotation_str or "dict" in annotation_str:
return "object"
elif "Optional" in annotation_str:
# Handle Optional types
if "str" in annotation_str:
return "string"
elif "int" in annotation_str:
return "integer"
elif "float" in annotation_str:
return "number"
elif "bool" in annotation_str:
return "boolean"
elif "List" in annotation_str or "list" in annotation_str:
return "array"
elif "Dict" in annotation_str or "dict" in annotation_str:
return "object"
# Default to string for unknown types
return "string"