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
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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 "<!DOCTYPE html>" in full_response or "<html" in full_response, \
|
||||
"Should contain HTML5 boilerplate elements"
|
||||
|
||||
# Verify indentation is preserved (check for multiple spaces in a row)
|
||||
# This indicates that code formatting with indentation is maintained
|
||||
assert " " in full_response, "Should preserve indentation (multiple spaces)"
|
||||
|
||||
# Log the response for debugging if test fails
|
||||
if "```" not in full_response or newline_count < 5:
|
||||
print("\n=== Full Response ===")
|
||||
print(repr(full_response)) # Use repr to see escaped characters
|
||||
print("\n=== Newline count ===")
|
||||
print(f"Found {newline_count} newlines")
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
def test_tatlock_markdown_non_streaming(client: TestClient):
|
||||
"""
|
||||
Integration test: Verify markdown in non-streaming mode.
|
||||
"""
|
||||
request_data = {
|
||||
"model": "Tatlock",
|
||||
"input": [{"role": "user", "content": "Give me a simple Python hello world code"}],
|
||||
"stream": False
|
||||
}
|
||||
|
||||
response = client.post("/v1/responses", json=request_data, timeout=30.0)
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
|
||||
# Get the message content
|
||||
message_item = next((item for item in data["output"] if item["type"] == "message"), None)
|
||||
assert message_item is not None
|
||||
|
||||
text = message_item["content"][0]["text"]
|
||||
|
||||
# Verify markdown code block
|
||||
assert "```" in text, "Should contain code block markers"
|
||||
assert "\n" in text, "Should contain newlines"
|
||||
Reference in New Issue
Block a user