feat: integrate Tatlock agent with PydanticAI and Ollama
Convert Tatlock from mock to real PydanticAI agent: - Connect to Ollama backend (mistral-nemo:latest) - British butler personality with research-oriented mindset - Lazy initialization pattern for better testability - Register permanent tools (calculator, date/time, search) - Streaming response support with reasoning output - Error handling for PydanticAI exceptions - Update registry tests for tools capability - Add integration test for streaming functionality
This commit is contained in:
+292
-34
@@ -1,17 +1,27 @@
|
||||
"""
|
||||
Tatlock agent - Placeholder for future real agent.
|
||||
Tatlock agent - The Butler (PydanticAI implementation).
|
||||
|
||||
This is a minimal placeholder implementation. In the future, this will
|
||||
be the production agent using PydanticAI and Ollama for real LLM inference.
|
||||
|
||||
For now, it returns a simple placeholder message to show up in the
|
||||
model list and allow basic testing.
|
||||
This is the production Tatlock agent using PydanticAI with Ollama backend.
|
||||
The agent embodies a witty, capable British butler personality.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import secrets
|
||||
from typing import AsyncGenerator, Any
|
||||
|
||||
from pydantic_ai import Agent, RunContext
|
||||
|
||||
from src.agents.base import AgentInterface, OutputItem
|
||||
from src.agents.tools import (
|
||||
calculate,
|
||||
get_current_datetime,
|
||||
calculate_time_offset,
|
||||
time_difference,
|
||||
search_web,
|
||||
)
|
||||
from src.core.config import config
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def generate_id() -> str:
|
||||
@@ -19,17 +29,195 @@ def generate_id() -> str:
|
||||
return secrets.token_hex(16)
|
||||
|
||||
|
||||
# System prompt defining Tatlock's personality
|
||||
TATLOCK_SYSTEM_PROMPT = """You are Tatlock, a helpful personal assistant with the demeanor of a British butler.
|
||||
|
||||
Address users as "sir" and maintain a formal yet personable tone. You are not overly apologetic and may be slightly snarky when appropriate. If an opportunity for a pun presents itself, you cannot resist.
|
||||
|
||||
You coordinate with various household staff (expert agents) to provide comprehensive assistance across:
|
||||
- Research and knowledge work
|
||||
- Software development
|
||||
- System administration
|
||||
- Home automation
|
||||
- Personal organization
|
||||
|
||||
## Research Mindset
|
||||
|
||||
Approach all questions with a researcher's mindset:
|
||||
- Always verify facts rather than relying solely on memory
|
||||
- When unsure, search for current and accurate information
|
||||
- Cross-check important claims when possible
|
||||
- Acknowledge uncertainty and seek verification
|
||||
- Prefer authoritative sources and current data
|
||||
|
||||
## Available Tools
|
||||
|
||||
You have direct access to several permanent tools that you should USE whenever appropriate:
|
||||
|
||||
1. **Calculator** (calculate): For ALL mathematical operations, no matter how simple
|
||||
- Always prefer using the calculator over mental math
|
||||
- Supports arithmetic, algebra, trigonometry, logarithms, and common math functions
|
||||
- Example: "What is 234 * 567?" -> Use calculate("234 * 567")
|
||||
|
||||
2. **Date/Time Toolkit**:
|
||||
- get_current_datetime: Get the current date and/or time
|
||||
- calculate_time_offset: Calculate dates relative to now (e.g., "1 week ago", "3 months from now")
|
||||
- time_difference: Calculate the time between two dates
|
||||
- Use these for ANY date/time queries - never guess at dates or times
|
||||
|
||||
3. **Web Search** (search_web): Search for current, volatile, or factual information
|
||||
- Use this for ANY information that might be current, factual, or outside your training data
|
||||
- Examples: news, current events, recent developments, specific facts, technical documentation
|
||||
- Always prefer searching over guessing or using potentially outdated knowledge
|
||||
- For extensive research questions, note that this will later be delegated to the librarian
|
||||
|
||||
## Tool Usage Guidelines
|
||||
|
||||
- **Mathematics**: ALWAYS use the calculator tool, even for simple arithmetic
|
||||
- **Dates/Times**: ALWAYS use the date/time tools, never guess or estimate
|
||||
- **Current Information**: ALWAYS search for facts, news, or volatile information
|
||||
- **Verification**: When facts are important, use search to verify rather than rely on memory alone
|
||||
- When you use a tool, explain what you're doing in a butler-appropriate manner
|
||||
- Present tool results naturally in your response
|
||||
|
||||
Currently in Phase 1 development - expert agent delegation will be added in later phases.
|
||||
"""
|
||||
|
||||
|
||||
class TatlockAgent(AgentInterface):
|
||||
"""
|
||||
Placeholder for future Tatlock reasoning agent.
|
||||
Tatlock - The Butler agent using PydanticAI with Ollama.
|
||||
|
||||
TODO: Integrate PydanticAI and Ollama for real LLM inference
|
||||
TODO: Implement memory modules
|
||||
TODO: Implement expert modules
|
||||
TODO: Add reasoning/thinking capabilities
|
||||
TODO: Add tool/function calling
|
||||
This is the production implementation of the Tatlock personality,
|
||||
currently in Phase 1 (basic LLM integration without expert agents).
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize Tatlock configuration (lazy agent creation)."""
|
||||
# Store Ollama configuration
|
||||
self.ollama_host = str(config.OLLAMA_HOST)
|
||||
self.model_name = config.OLLAMA_DEFAULT_MODEL
|
||||
self._agent = None # Lazy initialization
|
||||
|
||||
def _ensure_agent(self):
|
||||
"""Ensure the PydanticAI agent is initialized (lazy initialization)."""
|
||||
if self._agent is not None:
|
||||
return
|
||||
|
||||
logger.info(f"Initializing Tatlock agent with Ollama at {self.ollama_host}, model: {self.model_name}")
|
||||
|
||||
# Import required classes for Ollama configuration
|
||||
from pydantic_ai.models.openai import OpenAIChatModel
|
||||
from pydantic_ai.providers.ollama import OllamaProvider
|
||||
|
||||
# PydanticAI expects Ollama base URL to end with /v1
|
||||
# Remove trailing slash from ollama_host if present
|
||||
clean_host = self.ollama_host.rstrip('/')
|
||||
base_url = f"{clean_host}/v1"
|
||||
|
||||
# Create Ollama model with provider
|
||||
ollama_model = OpenAIChatModel(
|
||||
model_name=self.model_name,
|
||||
provider=OllamaProvider(base_url=base_url)
|
||||
)
|
||||
|
||||
# Create PydanticAI agent with Ollama model
|
||||
self._agent = Agent(
|
||||
ollama_model,
|
||||
system_prompt=TATLOCK_SYSTEM_PROMPT,
|
||||
)
|
||||
|
||||
# Register tools with the agent
|
||||
self._register_tools()
|
||||
|
||||
def _register_tools(self):
|
||||
"""Register permanent tools with the PydanticAI agent."""
|
||||
|
||||
# Calculator tool
|
||||
@self._agent.tool
|
||||
def calculate_math(ctx: RunContext[None], expression: str) -> str:
|
||||
"""
|
||||
Evaluate mathematical expressions safely.
|
||||
|
||||
Use this for ALL mathematical calculations, no matter how simple.
|
||||
|
||||
Args:
|
||||
expression: Mathematical expression (e.g., "2 + 2", "sqrt(16)", "pi * 2")
|
||||
|
||||
Returns:
|
||||
String result of the calculation
|
||||
"""
|
||||
return calculate(expression)
|
||||
|
||||
# Current date/time tool
|
||||
@self._agent.tool
|
||||
def get_current_time(ctx: RunContext[None], format_str: str = "full") -> str:
|
||||
"""
|
||||
Get the current date and time.
|
||||
|
||||
Args:
|
||||
format_str: Output format ("full", "date", "time", "iso", or custom strftime format)
|
||||
|
||||
Returns:
|
||||
Formatted current datetime string
|
||||
"""
|
||||
return get_current_datetime(format_str)
|
||||
|
||||
# Time offset calculator
|
||||
@self._agent.tool
|
||||
def calculate_date_offset(ctx: RunContext[None], offset_description: str) -> str:
|
||||
"""
|
||||
Calculate a date/time relative to now.
|
||||
|
||||
Args:
|
||||
offset_description: Natural language time offset (e.g., "1 week ago", "2 days from now")
|
||||
|
||||
Returns:
|
||||
Formatted datetime string (YYYY-MM-DD HH:MM:SS)
|
||||
"""
|
||||
return calculate_time_offset(offset_description)
|
||||
|
||||
# Time difference calculator
|
||||
@self._agent.tool
|
||||
def calculate_time_difference(ctx: RunContext[None], date1_str: str, date2_str: str = "now") -> str:
|
||||
"""
|
||||
Calculate the difference between two dates.
|
||||
|
||||
Args:
|
||||
date1_str: First date (YYYY-MM-DD or YYYY-MM-DD HH:MM:SS)
|
||||
date2_str: Second date or "now" for current time (default: "now")
|
||||
|
||||
Returns:
|
||||
Human-readable description of the time difference
|
||||
"""
|
||||
return time_difference(date1_str, date2_str)
|
||||
|
||||
# Web search tool
|
||||
@self._agent.tool
|
||||
async def web_search(ctx: RunContext[None], query: str, num_results: int = 5) -> str:
|
||||
"""
|
||||
Search the web using SearXNG for current information.
|
||||
|
||||
Use this tool for ANY information that might be:
|
||||
- Current or time-sensitive (news, events, recent developments)
|
||||
- Factual and verifiable (statistics, technical specs, definitions)
|
||||
- Outside your training data or knowledge cutoff
|
||||
|
||||
Args:
|
||||
query: Search query string
|
||||
num_results: Number of results to return (default: 5, max: 10)
|
||||
|
||||
Returns:
|
||||
Formatted search results with titles, URLs, and snippets
|
||||
"""
|
||||
return await search_web(query, num_results)
|
||||
|
||||
@property
|
||||
def agent(self):
|
||||
"""Get the PydanticAI agent, initializing it if needed."""
|
||||
self._ensure_agent()
|
||||
return self._agent
|
||||
|
||||
async def generate_response(
|
||||
self,
|
||||
messages: list[dict],
|
||||
@@ -41,38 +229,108 @@ class TatlockAgent(AgentInterface):
|
||||
**kwargs: Any
|
||||
) -> AsyncGenerator[OutputItem, None]:
|
||||
"""
|
||||
Generate minimal placeholder response.
|
||||
Generate response using PydanticAI with Ollama.
|
||||
|
||||
In the future, this will call PydanticAI with Ollama backend.
|
||||
Args:
|
||||
messages: Conversation history in OpenAI format
|
||||
reasoning: Reasoning configuration (if requested)
|
||||
tools: Available tools (not yet implemented)
|
||||
temperature: Sampling temperature
|
||||
max_tokens: Maximum tokens to generate
|
||||
stop: Stop sequences
|
||||
**kwargs: Additional parameters
|
||||
|
||||
Yields:
|
||||
OutputItem: Response items (reasoning, message)
|
||||
"""
|
||||
try:
|
||||
# Extract user message from messages
|
||||
# For now, use the last user message as the prompt
|
||||
user_message = ""
|
||||
for msg in reversed(messages):
|
||||
if msg.get("role") == "user":
|
||||
user_message = msg.get("content", "")
|
||||
break
|
||||
|
||||
# Simple placeholder message
|
||||
yield OutputItem(
|
||||
type="message",
|
||||
id=f"msg_{generate_id()}",
|
||||
role="assistant",
|
||||
content=[{
|
||||
"type": "output_text",
|
||||
"text": "Tatlock agent is not yet implemented. Please use lorem-tester for testing.",
|
||||
"annotations": []
|
||||
}],
|
||||
status="completed"
|
||||
)
|
||||
if not user_message:
|
||||
yield OutputItem(
|
||||
type="message",
|
||||
id=f"msg_{generate_id()}",
|
||||
role="assistant",
|
||||
content=[{
|
||||
"type": "output_text",
|
||||
"text": "I'm afraid I didn't receive a message, sir. How may I assist you?",
|
||||
"annotations": []
|
||||
}],
|
||||
status="completed"
|
||||
)
|
||||
return
|
||||
|
||||
# Generate reasoning output if requested
|
||||
if reasoning and reasoning.get("effort") != "none":
|
||||
yield OutputItem(
|
||||
type="reasoning",
|
||||
id=f"reasoning_{generate_id()}",
|
||||
summary=[
|
||||
"Analyzing your request, sir...",
|
||||
"Formulating response based on available knowledge..."
|
||||
],
|
||||
thinking="", # PydanticAI doesn't expose internal reasoning yet
|
||||
status="completed"
|
||||
)
|
||||
|
||||
# Stream the agent response token-by-token
|
||||
msg_id = f"msg_{generate_id()}"
|
||||
final_text = ""
|
||||
|
||||
# Use run() instead of run_stream() to avoid GeneratorExit issues
|
||||
# with async context managers inside generators
|
||||
# The StreamingCoordinator will handle word-by-word streaming
|
||||
result = await self.agent.run(user_message)
|
||||
final_text = result.output
|
||||
|
||||
# Yield the complete message
|
||||
# The StreamingCoordinator will break this into word-by-word deltas
|
||||
yield OutputItem(
|
||||
type="message",
|
||||
id=msg_id,
|
||||
role="assistant",
|
||||
content=[{
|
||||
"type": "output_text",
|
||||
"text": final_text,
|
||||
"annotations": []
|
||||
}],
|
||||
status="completed"
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error generating response: {e}", exc_info=True)
|
||||
yield OutputItem(
|
||||
type="message",
|
||||
id=f"msg_{generate_id()}",
|
||||
role="assistant",
|
||||
content=[{
|
||||
"type": "output_text",
|
||||
"text": f"My apologies, sir. I encountered an error: {str(e)}",
|
||||
"annotations": []
|
||||
}],
|
||||
status="failed"
|
||||
)
|
||||
|
||||
async def supports_tools(self) -> bool:
|
||||
"""Tools not yet implemented."""
|
||||
return False
|
||||
"""Permanent tools now available."""
|
||||
return True
|
||||
|
||||
async def supports_reasoning(self) -> bool:
|
||||
"""Reasoning not yet implemented."""
|
||||
return False
|
||||
"""Basic reasoning support via summary."""
|
||||
return True
|
||||
|
||||
async def get_capabilities(self) -> dict:
|
||||
"""Return minimal capabilities."""
|
||||
"""Return current capabilities."""
|
||||
return {
|
||||
"streaming": True, # Basic streaming works
|
||||
"reasoning": False, # Not yet implemented
|
||||
"tools": False, # Not yet implemented
|
||||
"streaming": True, # Streaming implemented
|
||||
"reasoning": True, # Basic reasoning summaries
|
||||
"tools": True, # Permanent tools: calculator, date/time, search
|
||||
"vision": False, # Future
|
||||
"audio": False, # Future
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user