tatlock/src/agents/tatlock.py

"""
Tatlock agent - The Butler (PydanticAI implementation).

This is the production Tatlock agent using PydanticAI with Ollama backend.
The agent embodies a witty, capable British butler personality.
"""

import secrets
from typing import AsyncGenerator, Any
from dataclasses import dataclass, field

from pydantic_ai import Agent, RunContext

from src.agents.base import AgentInterface, OutputItem
from src.agents.tatlock_core.tools import (
    calculate,
    get_current_datetime,
    calculate_time_offset,
    time_difference,
    search_web,
)
from src.core.config import config
from src.core.logging_config import get_logger

logger = get_logger(__name__)


@dataclass
class ToolCallTracker:
    """Tracks tool calls for reporting to reasoning output."""
    calls: list[str] = field(default_factory=list)

    def log_call(self, message: str):
        """Log a tool call."""
        self.calls.append(message)


def generate_id() -> str:
    """Generate unique ID for output items."""
    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):
    """
    Tatlock - The Butler agent using PydanticAI with Ollama.

    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(
            "tatlock_agent_initializing",
            ollama_host=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[ToolCallTracker], 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
            """
            # Log the calculation to reasoning output
            if ctx.deps:
                ctx.deps.log_call(f"🧮 Calculating: {expression}")
            return calculate(expression)

        # Current date/time tool
        @self._agent.tool
        def get_current_time(ctx: RunContext[ToolCallTracker], 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
            """
            if ctx.deps:
                ctx.deps.log_call(f"🕐 Getting current time (format: {format_str})")
            return get_current_datetime(format_str)

        # Time offset calculator
        @self._agent.tool
        def calculate_date_offset(ctx: RunContext[ToolCallTracker], 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)
            """
            if ctx.deps:
                ctx.deps.log_call(f"🕐 Calculating date offset: {offset_description}")
            return calculate_time_offset(offset_description)

        # Time difference calculator
        @self._agent.tool
        def calculate_time_difference(ctx: RunContext[ToolCallTracker], 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
            """
            if ctx.deps:
                ctx.deps.log_call(f"🕐 Calculating time difference between {date1_str} and {date2_str}")
            return time_difference(date1_str, date2_str)

        # Web search tool
        @self._agent.tool
        async def web_search(ctx: RunContext[ToolCallTracker], 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
            """
            # Log the search query to reasoning output
            if ctx.deps:
                ctx.deps.log_call(f"🔍 Searching for: '{query}'")
            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],
        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 response using PydanticAI with Ollama.

        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:
            # Convert OpenAI-format messages to PydanticAI format
            # PydanticAI uses: {"role": "user"/"assistant", "content": "text"}
            # OpenAI format is the same, so we can use messages directly

            # Extract the latest user message for the prompt
            user_message = ""
            for msg in reversed(messages):
                if msg.get("role") == "user":
                    user_message = msg.get("content", "")
                    break

            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

            # Build message history (all messages except the last user message)
            # PydanticAI expects history as list of ModelRequest/ModelResponse objects
            from pydantic_ai.messages import ModelRequest, ModelResponse, UserPromptPart, TextPart

            message_history = []
            for i, msg in enumerate(messages[:-1]):  # All messages except the last one
                role = msg.get("role")
                content = msg.get("content", "")

                # Skip messages with empty content (can cause Ollama errors)
                if not content or not content.strip():
                    logger.warning(f"Skipping message {i} with empty content: role={role}")
                    continue

                # Debug: Check for problematic content
                if '"' in content or "'" in content:
                    logger.debug(f"Message {i} ({role}) contains quotes. Content preview: {content[:100]}...")

                # Convert to PydanticAI message format
                try:
                    if role == "user":
                        message_history.append(
                            ModelRequest(parts=[UserPromptPart(content=content)])
                        )
                    elif role == "assistant":
                        message_history.append(
                            ModelResponse(parts=[TextPart(content=content)])
                        )
                except Exception as e:
                    logger.error(f"Error creating message history item {i}: {e}")
                    logger.error(f"Problematic content: {repr(content)}")
                    raise

            # Debug: Log the message history summary
            logger.info(f"Built message history with {len(message_history)} messages")
            if message_history:
                for i, hist_msg in enumerate(message_history):
                    msg_type = type(hist_msg).__name__
                    content_preview = str(hist_msg.parts[0].content)[:50] if hist_msg.parts else "no parts"
                    logger.info(f"  History[{i}]: {msg_type} - {content_preview}...")

            # 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"
                )

            # Create a tool call tracker for this request
            tracker = ToolCallTracker()

            # 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
            # Pass message_history to maintain conversation context and tracker for tool logging
            result = await self.agent.run(
                user_message,
                message_history=message_history if message_history else None,
                deps=tracker
            )
            final_text = result.output

            # If tools were called, yield a reasoning item showing what was done
            if tracker.calls:
                yield OutputItem(
                    type="reasoning",
                    id=f"reasoning_tools_{generate_id()}",
                    summary=tracker.calls,
                    thinking="",
                    status="completed"
                )

            # 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:
        """Permanent tools now available."""
        return True

    async def supports_reasoning(self) -> bool:
        """Basic reasoning support via summary."""
        return True

    async def run_with_scoped_tools(
        self,
        user_message: str,
        steward_note: str,
        scoped_tools: list[Any],
        message_history: list[dict],
        tool_tracker: Any = None,
    ) -> str:
        """
        Run Tatlock with scoped tools from Steward preprocessing.

        This is the Phase 2 request flow where the Steward has already
        analyzed the request and provided scoped tools.

        Args:
            user_message: The user's original message
            steward_note: Note from Steward (prepended to request, invisible to user)
            scoped_tools: List of tool definitions from household registry
            message_history: Conversation history in PydanticAI format
            tool_tracker: Optional tool call tracker for benchmarking

        Returns:
            str: Tatlock's response text

        Example:
            >>> response = await tatlock.run_with_scoped_tools(
            ...     "What's sqrt(144)?",
            ...     steward_note="Simple math request...",
            ...     scoped_tools=[calculator_tool, ...],
            ...     message_history=[],
            ...     tool_tracker=tracker,
            ... )
        """
        from pydantic_ai.models.openai import OpenAIChatModel
        from pydantic_ai.providers.ollama import OllamaProvider

        logger.info(
            "tatlock_run_with_scoped_tools",
            user_message_preview=user_message[:100],
            scoped_tool_count=len(scoped_tools),
            history_length=len(message_history),
        )

        # Create a fresh agent instance with scoped tools only
        # This ensures Tatlock can ONLY use tools recommended by the Steward
        clean_host = self.ollama_host.rstrip('/')
        base_url = f"{clean_host}/v1"

        ollama_model = OpenAIChatModel(
            model_name=self.model_name,
            provider=OllamaProvider(base_url=base_url)
        )

        # Create agent with scoped tools
        # Tools from household registry are already PydanticAI Tool objects
        scoped_agent = Agent(
            ollama_model,
            system_prompt=TATLOCK_SYSTEM_PROMPT,
            tools=scoped_tools,  # Pass tools directly to Agent constructor
        )

        # Prepend Steward's note to the request (invisible to user, visible to Tatlock)
        enriched_message = f"{steward_note}\n\n{user_message}"

        # Convert message history to PydanticAI format
        from pydantic_ai.messages import ModelRequest, ModelResponse, UserPromptPart, TextPart

        pydantic_history = []
        for msg in message_history:
            role = msg.get("role")
            content = msg.get("content", "")

            if not content or not content.strip():
                continue

            if role == "user":
                pydantic_history.append(
                    ModelRequest(parts=[UserPromptPart(content=content)])
                )
            elif role == "assistant":
                pydantic_history.append(
                    ModelResponse(parts=[TextPart(content=content)])
                )

        # Run with scoped tools and tracker
        result = await scoped_agent.run(
            enriched_message,
            message_history=pydantic_history if pydantic_history else None,
            deps=tool_tracker
        )

        logger.info(
            "tatlock_response_generated",
            response_preview=result.output[:100],
        )

        return result.output

    async def run_with_scoped_tools_stream(
        self,
        user_message: str,
        steward_note: str,
        scoped_tools: list,
        message_history: list[dict],
        tool_tracker: "ToolCallTracker",
    ):
        """
        Run Tatlock with scoped tools recommended by Steward (streaming version).

        This is the Phase 2 execution flow where Steward has preprocessed
        the request and provided:
        - steward_note: Instructions for Tatlock (invisible to user)
        - scoped_tools: Only the tools Steward recommended

        Args:
            user_message: Original user message
            steward_note: Steward's instructions for Tatlock
            scoped_tools: List of PydanticAI Tool objects to use
            message_history: Previous conversation turns
            tool_tracker: Tracker for tool call analytics

        Yields:
            Text chunks from the streaming response
        """
        from pydantic_ai.models.openai import OpenAIChatModel
        from pydantic_ai.providers.ollama import OllamaProvider

        logger.info(
            "tatlock_run_with_scoped_tools_stream",
            user_message_preview=user_message[:100],
            scoped_tool_count=len(scoped_tools),
            history_length=len(message_history),
        )

        # Create a fresh agent instance with scoped tools only
        clean_host = self.ollama_host.rstrip('/')
        base_url = f"{clean_host}/v1"

        ollama_model = OpenAIChatModel(
            model_name=self.model_name,
            provider=OllamaProvider(base_url=base_url)
        )

        # Create agent with scoped tools
        scoped_agent = Agent(
            ollama_model,
            system_prompt=TATLOCK_SYSTEM_PROMPT,
            tools=scoped_tools,
        )

        # Prepend Steward's note to the request
        enriched_message = f"{steward_note}\n\n{user_message}"

        # Convert message history to PydanticAI format
        from pydantic_ai.messages import ModelRequest, ModelResponse, UserPromptPart, TextPart

        pydantic_history = []
        for msg in message_history:
            role = msg.get("role")
            content = msg.get("content", "")

            if not content or not content.strip():
                continue

            if role == "user":
                pydantic_history.append(
                    ModelRequest(parts=[UserPromptPart(content=content)])
                )
            elif role == "assistant":
                pydantic_history.append(
                    ModelResponse(parts=[TextPart(content=content)])
                )

        # Use run() instead of run_stream() to avoid Ollama 400 bug
        # with streaming + tool calls (PydanticAI issues #1292, #2256)
        # We yield the final response in chunks to maintain streaming interface
        result = await scoped_agent.run(
            enriched_message,
            message_history=pydantic_history if pydantic_history else None,
            deps=tool_tracker
        )

        # Stream the final response in chunks to maintain UX
        response_text = result.output
        chunk_size = 50  # characters per chunk

        for i in range(0, len(response_text), chunk_size):
            yield response_text[i:i + chunk_size]

        logger.info("tatlock_scoped_run_complete")

    async def get_capabilities(self) -> dict:
        """Return current capabilities."""
        return {
            "streaming": True,   # Streaming implemented
            "reasoning": True,   # Basic reasoning summaries
            "tools": True,       # Permanent tools: calculator, date/time, search
            "vision": False,     # Future
            "audio": False,      # Future
        }