"""
Basic OllamaAgent implementation.
This module provides the BasicOllamaAgent class, which is a simple wrapper
around the core OllamaAgent class, providing a clean interface for basic usage
without the additional features of the modular agent.
"""
import logging
from typing import Dict, Any, List, Optional
from ollamatools.core.agent import OllamaAgent
from ollamatools.core.config import Config
from ollamatools.core.utils.logger import get_ollamatools_logger
# Set up logging
logger = get_ollamatools_logger("basic_ollama_agent", level="INFO")
class BasicOllamaAgent:
"""
A simple wrapper around the core OllamaAgent class.
This class provides a clean interface for basic usage of the OllamaAgent
without the additional features of the modular agent.
"""
def __init__(
self,
model: str = "llama3.1:latest",
host: str = "http://localhost:11434",
permission_prompts: bool = True,
sudo_always_prompt: bool = True,
config: Optional[Config] = None
):
"""
Initialize the BasicOllamaAgent 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)
"""
# Initialize the core OllamaAgent
self.agent = OllamaAgent(
model=model,
host=host,
permission_prompts=permission_prompts,
sudo_always_prompt=sudo_always_prompt,
config=config
)
# Initialize conversation context
self.context = []
# Define default system prompt
self.default_system_prompt = """
You are an AI assistant with access to various tools to help the user.
You can execute system commands, fetch data from the web, and generate documentation.
When a user asks for something that requires using one of your tools,
use the appropriate function to complete the task.
Be helpful, concise, and accurate.
"""
logger.info(f"BasicOllamaAgent initialized with model={self.agent.config.model}, host={self.agent.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 (defaults to self.default_system_prompt)
context: Optional conversation history (defaults to self.context)
Returns:
str: Response from Ollama, potentially after executing functions
"""
# Use provided context or instance context
if context is None:
context = self.context
# Use provided system prompt or default
if system_prompt is None:
system_prompt = self.default_system_prompt
logger.info(f"Sending message to Ollama: {message[:50]}...")
try:
# Send the message to the core agent
response = self.agent.chat(
message=message,
system_prompt=system_prompt,
context=context
)
# Update instance context if using it
if context is self.context:
self.context.append({"role": "user", "content": message})
self.context.append({"role": "assistant", "content": response})
# Limit context size to avoid exceeding token limits
if len(self.context) > 10:
# Keep only the last 10 messages
self.context = self.context[-10:]
logger.info(f"Received response from Ollama")
return response
except Exception as e:
logger.error(f"Error in chat: {str(e)}")
return f"Error communicating with Ollama: {str(e)}"
def reset_context(self) -> None:
"""
Reset the conversation context.
"""
self.context = []
logger.info("Conversation context reset")
def get_registered_functions(self) -> List[str]:
"""
Get a list of registered function names.
Returns:
List[str]: List of registered function names
"""
return list(self.agent.registered_functions.keys())
def set_system_prompt(self, prompt: str) -> None:
"""
Set a new default system prompt.
Args:
prompt: New system prompt
"""
self.default_system_prompt = prompt
logger.info("Default system prompt updated")