tatlock/src/agents/lorem_tester.py

"""
Lorem Tester agent - Mock agent for testing Responses API features.

This agent implements all Responses API features using mock Lorem Ipsum
content. It's designed to test the plumbing (streaming, reasoning display,
tool calling, error handling) before connecting real LLM integration.

The ONLY mock part is the PydanticAI interface - all surrounding
infrastructure is real production code.
"""

import asyncio
import random
import secrets
from typing import AsyncGenerator, Any

from src.agents.base import AgentInterface, OutputItem
from src.core.exceptions import (
    RateLimitError,
    ContextLengthError,
    APIError,
)


# Mock lorem ipsum content
LOREM_PARAGRAPHS = [
    "Lorem ipsum dolor sit amet, consectetur adipiscing elit. Sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.",
    "Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat.",
    "Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur.",
    "Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum.",
]

# Mock reasoning steps
REASONING_STEPS = [
    "Analyzing the user's request and understanding the context...",
    "Considering the available information and identifying key requirements...",
    "Evaluating different approaches and their potential outcomes...",
    "Formulating a comprehensive response strategy...",
    "Selecting appropriate content and structuring the answer...",
]

# Mock tool definitions
MOCK_TOOLS = [
    {
        "name": "search_knowledge",
        "description": "Search the knowledge base for relevant information",
        "parameters": {
            "type": "object",
            "properties": {
                "query": {"type": "string", "description": "Search query"}
            },
            "required": ["query"]
        }
    },
    {
        "name": "calculate",
        "description": "Perform mathematical calculations",
        "parameters": {
            "type": "object",
            "properties": {
                "expression": {"type": "string", "description": "Math expression"}
            },
            "required": ["expression"]
        }
    },
    {
        "name": "get_weather",
        "description": "Get current weather for a location",
        "parameters": {
            "type": "object",
            "properties": {
                "location": {"type": "string", "description": "City name"}
            },
            "required": ["location"]
        }
    },
]


def generate_id() -> str:
    """Generate unique ID for output items."""
    return secrets.token_hex(16)


class LoremTesterAgent(AgentInterface):
    """
    Mock agent that implements all Responses API features.

    Features:
    - Reasoning summaries (thinking steps)
    - Function calling (mock tool execution)
    - Streaming responses
    - Error scenarios (trigger keywords)
    - Multi-turn conversations

    Error Triggers:
    - "trigger_rate_limit" → 429 rate limit error
    - "trigger_timeout" → timeout after delay
    - "trigger_context_overflow" → context length exceeded
    - "trigger_partial_failure" → error mid-stream
    - "trigger_invalid_tool" → invalid tool call
    """

    async def generate_response(
        self,
        messages: list[dict],
        reasoning: dict | None = None,
        tools: list[dict] | None = None,
        temperature: float = 1.0,
        max_tokens: int | None = None,
        stop: list[str] | None = None,
        **kwargs: Any
    ) -> AsyncGenerator[OutputItem, None]:
        """
        Generate mock response with reasoning, tools, and content.

        This is where PydanticAI would be called in the real implementation.
        """

        # Check for error triggers in last message
        await self._check_error_triggers(messages)

        # 1. Yield reasoning item if requested
        if reasoning and reasoning.get("summary") == "auto":
            yield await self._create_reasoning_item(
                messages,
                effort=reasoning.get("effort", "medium")
            )

        # 2. Randomly yield function calls if tools available (30% chance)
        if tools and random.random() < 0.3:
            async for tool_item in self._create_tool_calls(tools):
                yield tool_item

        # 3. Yield final message item
        yield await self._create_message_item(messages, temperature)

    async def supports_tools(self) -> bool:
        """Lorem Tester supports tool calling."""
        return True

    async def supports_reasoning(self) -> bool:
        """Lorem Tester supports reasoning summaries."""
        return True

    async def get_capabilities(self) -> dict:
        """Return full capabilities."""
        return {
            "streaming": True,
            "reasoning": True,
            "tools": True,
            "vision": False,  # Not yet
            "audio": False,   # Not yet
        }

    # Private helper methods

    async def _check_error_triggers(self, messages: list[dict]) -> None:
        """Check for error trigger keywords and raise appropriate errors."""
        if not messages:
            return

        last_message = str(messages[-1]).lower()

        if "trigger_rate_limit" in last_message:
            raise RateLimitError("Rate limit exceeded (mock trigger)")

        if "trigger_timeout" in last_message:
            # Simulate long delay
            await asyncio.sleep(5)  # Shortened for testing
            raise TimeoutError("Request timeout (mock trigger)")

        if "trigger_context_overflow" in last_message:
            raise ContextLengthError(
                "Context length exceeded: 5000 tokens > 4096 max (mock trigger)"
            )

        if "trigger_invalid_tool" in last_message:
            raise APIError("Invalid tool call: tool 'nonexistent' not found (mock trigger)")

    async def _create_reasoning_item(
        self,
        messages: list[dict],
        effort: str = "medium"
    ) -> OutputItem:
        """Create a reasoning output item with mock thinking steps."""

        # Adjust number of steps based on effort
        effort_steps = {
            "none": 0,
            "minimal": 1,
            "low": 2,
            "medium": 3,
            "high": 4,
            "xhigh": 5,
        }
        num_steps = effort_steps.get(effort, 3)

        # Select random reasoning steps
        steps = random.sample(REASONING_STEPS, min(num_steps, len(REASONING_STEPS)))

        return OutputItem(
            type="reasoning",
            id=f"rs_{generate_id()}",
            summary=steps,
            status="completed"
        )

    async def _create_tool_calls(
        self,
        tools: list[dict]
    ) -> AsyncGenerator[OutputItem, None]:
        """Create mock function call output items."""

        # Randomly select 1-2 tools to "call"
        num_calls = random.randint(1, 2)
        selected_tools = random.sample(
            MOCK_TOOLS[:min(len(MOCK_TOOLS), len(tools))],
            min(num_calls, len(MOCK_TOOLS), len(tools))
        )

        for tool in selected_tools:
            # Generate mock arguments
            args = self._generate_mock_args(tool)

            yield OutputItem(
                type="function_call",
                id=f"fc_{generate_id()}",
                name=tool["name"],
                arguments=args,
                status="completed"
            )

    def _generate_mock_args(self, tool: dict) -> str:
        """Generate mock arguments for a tool call."""
        import json

        name = tool["name"]

        # Generate contextual mock arguments
        if name == "search_knowledge":
            queries = ["lorem ipsum", "dolor sit amet", "consectetur adipiscing"]
            return json.dumps({"query": random.choice(queries)})

        elif name == "calculate":
            expressions = ["2 + 2", "10 * 5", "100 / 4"]
            return json.dumps({"expression": random.choice(expressions)})

        elif name == "get_weather":
            cities = ["New York", "London", "Tokyo", "Paris"]
            return json.dumps({"location": random.choice(cities)})

        else:
            # Generic mock arguments
            return json.dumps({"input": "mock_value"})

    async def _create_message_item(
        self,
        messages: list[dict],
        temperature: float
    ) -> OutputItem:
        """Create final message output item with lorem ipsum content."""

        # Select random lorem ipsum paragraphs
        # Temperature affects variety: higher temp = more paragraphs
        num_paragraphs = 1 if temperature < 0.5 else random.randint(1, 2)
        content = " ".join(random.sample(LOREM_PARAGRAPHS, num_paragraphs))

        return OutputItem(
            type="message",
            id=f"msg_{generate_id()}",
            role="assistant",
            content=[{
                "type": "output_text",
                "text": content,
                "annotations": []
            }],
            status="completed"
        )