From 67481515cca95a71bf9301d1a0d57a8a5fc66ca1 Mon Sep 17 00:00:00 2001 From: Jeroen Schweitzer Date: Sun, 7 Dec 2025 00:12:39 +0100 Subject: [PATCH] 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 --- src/agents/tatlock.py | 326 ++++++++++++++-- tests/agents/test_registry.py | 18 +- tests/integration/test_tatlock_streaming.py | 410 ++++++++++++++++++++ 3 files changed, 711 insertions(+), 43 deletions(-) create mode 100644 tests/integration/test_tatlock_streaming.py diff --git a/src/agents/tatlock.py b/src/agents/tatlock.py index c9f164e..85ab6c8 100644 --- a/src/agents/tatlock.py +++ b/src/agents/tatlock.py @@ -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 } diff --git a/tests/agents/test_registry.py b/tests/agents/test_registry.py index 8af4209..cb44a1a 100644 --- a/tests/agents/test_registry.py +++ b/tests/agents/test_registry.py @@ -22,7 +22,7 @@ async def test_list_models(): # Check model IDs model_ids = [m["id"] for m in models] assert "lorem-tester" in model_ids - assert "tatlock" in model_ids + assert "Tatlock" in model_ids # Check structure for model in models: @@ -57,16 +57,16 @@ async def test_lorem_tester_capabilities(): async def test_tatlock_capabilities(): """Test tatlock model capabilities.""" models = await ModelRegistry.list_models() - tatlock_model = next(m for m in models if m["id"] == "tatlock") + tatlock_model = next(m for m in models if m["id"] == "Tatlock") capabilities = tatlock_model["capabilities"] - # Tatlock is placeholder - minimal capabilities + # Tatlock Phase 1 - basic streaming, reasoning, and permanent tools assert capabilities["streaming"] is True - assert capabilities["reasoning"] is False # Not yet - assert capabilities["tools"] is False # Not yet - assert capabilities["vision"] is False - assert capabilities["audio"] is False + assert capabilities["reasoning"] is True # Basic reasoning summaries + assert capabilities["tools"] is True # Permanent tools: calculator, date/time, search + assert capabilities["vision"] is False # Future + assert capabilities["audio"] is False # Future @pytest.mark.unit @@ -80,7 +80,7 @@ def test_get_agent_lorem_tester(): @pytest.mark.unit def test_get_agent_tatlock(): """Test getting tatlock agent instance.""" - agent = ModelRegistry.get_agent("tatlock") + agent = ModelRegistry.get_agent("Tatlock") assert isinstance(agent, TatlockAgent) @@ -98,7 +98,7 @@ def test_get_agent_not_found(): def test_model_exists(): """Test checking if model exists.""" assert ModelRegistry.model_exists("lorem-tester") is True - assert ModelRegistry.model_exists("tatlock") is True + assert ModelRegistry.model_exists("Tatlock") is True assert ModelRegistry.model_exists("nonexistent") is False diff --git a/tests/integration/test_tatlock_streaming.py b/tests/integration/test_tatlock_streaming.py new file mode 100644 index 0000000..1d97231 --- /dev/null +++ b/tests/integration/test_tatlock_streaming.py @@ -0,0 +1,410 @@ +""" +Integration tests for Tatlock agent streaming through full API stack. + +These tests verify the complete streaming flow from API endpoint through +StreamingCoordinator to TatlockAgent, ensuring no text duplication and +proper delta calculation. +""" +import json +import pytest +from httpx import AsyncClient +from fastapi.testclient import TestClient + + +@pytest.mark.integration +@pytest.mark.asyncio +async def test_tatlock_streaming_no_duplication(async_client: AsyncClient): + """ + Integration test: Verify Tatlock streaming produces no text duplication. + + This test catches the bug where accumulated text from PydanticAI was + being re-streamed multiple times by the StreamingCoordinator. + """ + request_data = { + "model": "Tatlock", + "input": [{"role": "user", "content": "Say hello"}], + "stream": True + } + + collected_deltas = [] + + async with async_client.stream( + "POST", + "/v1/responses", + json=request_data, + timeout=30.0, # Give enough time for Ollama response + ) as response: + assert response.status_code == 200 + assert response.headers["content-type"] == "text/event-stream; charset=utf-8" + + async for line in response.aiter_lines(): + if not line.strip(): + continue + + if line.startswith("event: "): + event_type = line[7:].strip() + elif line.startswith("data: "): + data_str = line[6:].strip() + if data_str != "[DONE]": + try: + chunk = json.loads(data_str) + + # Collect output text deltas + if chunk.get("event") == "response.output_text.delta": + collected_deltas.append(chunk["delta"]) + + except json.JSONDecodeError: + pass + + # Reconstruct full text from deltas + full_text = "".join(collected_deltas) + + # Verify we got some response + assert len(full_text) > 0, "Should have received some text" + + # Verify no obvious duplication patterns + # Check that common words don't appear excessively repeated + words = full_text.lower().split() + if len(words) > 0: + # Check for consecutive duplicate words (sign of duplication bug) + consecutive_dupes = sum( + 1 for i in range(len(words) - 1) + if words[i] == words[i + 1] and len(words[i]) > 3 + ) + # Allow a few duplicates (natural language), but not excessive + assert consecutive_dupes < len(words) * 0.1, \ + f"Too many consecutive duplicate words: {consecutive_dupes}/{len(words)}" + + +@pytest.mark.integration +@pytest.mark.asyncio +async def test_tatlock_chat_streaming_no_duplication(async_client: AsyncClient): + """ + Integration test: Verify Tatlock streaming through Chat Completions API. + + Tests the full stack through the chat completions wrapper to ensure + streaming works correctly without duplication. + """ + request_data = { + "model": "Tatlock", + "messages": [{"role": "user", "content": "Hello"}], + "stream": True + } + + collected_content = [] + + async with async_client.stream( + "POST", + "/v1/chat/completions", + json=request_data, + timeout=30.0, + ) as response: + assert response.status_code == 200 + + async for line in response.aiter_lines(): + if not line.strip(): + continue + + if line.startswith("data: "): + data_str = line[6:].strip() + if data_str == "[DONE]": + break + + try: + chunk = json.loads(data_str) + + # Collect content deltas from choices + if "choices" in chunk and len(chunk["choices"]) > 0: + delta = chunk["choices"][0].get("delta", {}) + if "content" in delta and delta["content"]: + collected_content.append(delta["content"]) + + except json.JSONDecodeError: + pass + + # Reconstruct full response + full_response = "".join(collected_content) + + # Verify we got a response + assert len(full_response) > 0, "Should have received response content" + + # Check for duplication patterns + words = full_response.lower().split() + if len(words) > 0: + consecutive_dupes = sum( + 1 for i in range(len(words) - 1) + if words[i] == words[i + 1] and len(words[i]) > 3 + ) + assert consecutive_dupes < len(words) * 0.1, \ + f"Too many consecutive duplicate words in chat response: {consecutive_dupes}/{len(words)}" + + +@pytest.mark.integration +def test_tatlock_non_streaming_responses_api(client: TestClient): + """ + Integration test: Verify Tatlock non-streaming through Responses API. + """ + request_data = { + "model": "Tatlock", + "input": [{"role": "user", "content": "Say hello"}], + "stream": False + } + + response = client.post("/v1/responses", json=request_data, timeout=30.0) + + assert response.status_code == 200 + data = response.json() + + # Verify response structure + assert data["status"] == "completed" + assert "output" in data + assert len(data["output"]) > 0 + + # Get the message content + message_item = next((item for item in data["output"] if item["type"] == "message"), None) + assert message_item is not None, "Should have a message output item" + assert len(message_item["content"]) > 0 + + text = message_item["content"][0]["text"] + assert len(text) > 0, "Should have response text" + + +@pytest.mark.integration +def test_tatlock_non_streaming_chat_api(client: TestClient): + """ + Integration test: Verify Tatlock non-streaming through Chat Completions API. + """ + request_data = { + "model": "Tatlock", + "messages": [{"role": "user", "content": "Hello"}], + "stream": False + } + + response = client.post("/v1/chat/completions", json=request_data, timeout=30.0) + + assert response.status_code == 200 + data = response.json() + + # Verify OpenAI-compatible structure + assert "id" in data + assert data["object"] == "chat.completion" + assert "choices" in data + assert len(data["choices"]) > 0 + + # Verify content + choice = data["choices"][0] + assert choice["message"]["role"] == "assistant" + assert len(choice["message"]["content"]) > 0 + + +@pytest.mark.integration +@pytest.mark.asyncio +async def test_tatlock_streaming_delta_accumulation(async_client: AsyncClient): + """ + Integration test: Verify deltas accumulate correctly without duplication. + + This test explicitly checks that when we accumulate all deltas, + we get a coherent response without repeated text. + """ + request_data = { + "model": "Tatlock", + "input": [{"role": "user", "content": "Count to three"}], + "stream": True + } + + collected_deltas = [] + previous_full_text = "" + + async with async_client.stream( + "POST", + "/v1/responses", + json=request_data, + timeout=30.0, + ) as response: + assert response.status_code == 200 + + async for line in response.aiter_lines(): + if not line.strip(): + continue + + if line.startswith("data: "): + data_str = line[6:].strip() + if data_str != "[DONE]": + try: + chunk = json.loads(data_str) + + if chunk.get("event") == "response.output_text.delta": + delta = chunk["delta"] + collected_deltas.append(delta) + + # Verify each delta is new content + current_full = "".join(collected_deltas) + assert current_full.startswith(previous_full_text), \ + "Deltas should accumulate progressively" + previous_full_text = current_full + + except json.JSONDecodeError: + pass + + full_text = "".join(collected_deltas) + assert len(full_text) > 0 + + +@pytest.mark.integration +@pytest.mark.asyncio +async def test_tatlock_with_reasoning(async_client: AsyncClient): + """ + Integration test: Verify Tatlock with reasoning enabled. + """ + request_data = { + "model": "Tatlock", + "input": [{"role": "user", "content": "Hello"}], + "reasoning": {"effort": "medium", "summary": "auto"}, + "stream": True + } + + has_reasoning = False + has_output = False + + async with async_client.stream( + "POST", + "/v1/responses", + json=request_data, + timeout=30.0, + ) as response: + assert response.status_code == 200 + + async for line in response.aiter_lines(): + if not line.strip(): + continue + + if line.startswith("data: "): + data_str = line[6:].strip() + if data_str != "[DONE]": + try: + chunk = json.loads(data_str) + + if chunk.get("event") == "response.reasoning_summary_text.delta": + has_reasoning = True + elif chunk.get("event") == "response.output_text.delta": + has_output = True + + except json.JSONDecodeError: + pass + + assert has_reasoning, "Should have reasoning summary" + assert has_output, "Should have output text" + + +@pytest.mark.integration +@pytest.mark.asyncio +async def test_tatlock_markdown_formatting_preserved(async_client: AsyncClient): + """ + Integration test: Verify markdown formatting is preserved in responses. + + Tests that code blocks, newlines, and other markdown formatting + are properly preserved through the streaming pipeline. + """ + request_data = { + "model": "Tatlock", + "input": [{"role": "user", "content": "Can you give me an HTML5 boilerplate template?"}], + "stream": True + } + + collected_deltas = [] + + async with async_client.stream( + "POST", + "/v1/responses", + json=request_data, + timeout=45.0, # Give extra time for code generation + ) as response: + assert response.status_code == 200 + + async for line in response.aiter_lines(): + if not line.strip(): + continue + + if line.startswith("data: "): + data_str = line[6:].strip() + if data_str != "[DONE]": + try: + chunk = json.loads(data_str) + + if chunk.get("event") == "response.output_text.delta": + collected_deltas.append(chunk["delta"]) + + except json.JSONDecodeError: + pass + + # Reconstruct full response + full_response = "".join(collected_deltas) + + # Always print the response for debugging + print("\n" + "="*80) + print("FULL RESPONSE (repr):") + print("="*80) + print(repr(full_response)) + print("\n" + "="*80) + print("FULL RESPONSE (formatted):") + print("="*80) + print(full_response) + print("="*80 + "\n") + + # Verify we got a response + assert len(full_response) > 100, "Should have a substantial response" + + # Verify markdown code block is present + assert "```" in full_response, "Response should contain markdown code blocks" + + # Verify newlines are preserved (not all collapsed to spaces) + newline_count = full_response.count('\n') + assert newline_count > 5, f"Should have multiple newlines preserved, got {newline_count}" + + # Verify code block markers are complete + code_block_starts = full_response.count("```") + # Should have at least opening and closing markers (even count) + assert code_block_starts % 2 == 0, "Code blocks should have matching opening/closing markers" + assert code_block_starts >= 2, "Should have at least one complete code block" + + # Verify HTML tags are present (indicates code block content is preserved) + assert "" in full_response or "