tatlock/src/core/preprocessing.py

"""
Request preprocessing pipeline.

Analyzes requests via the Steward and creates scoped toolsets for Tatlock.
"""
from dataclasses import dataclass
from datetime import datetime
from typing import Any, Optional

from src.agents.steward import analyze_request, format_steward_note
from src.agents.steward.schemas import StewardRecommendation
from src.core.household_registry import get_household_registry
from src.core.logging_config import get_logger
from src.core.tracing import trace_span, SpanType

logger = get_logger(__name__)


def _inject_temporal_context(request: str) -> str:
    """
    Append current time context to user request.

    Provides Tatlock with temporal awareness for time-sensitive queries.

    Args:
        request: Original user request

    Returns:
        Request with appended time context
    """
    now = datetime.now()
    time_str = now.strftime("%Y-%m-%d %H:%M")
    return f"{request}\n\n[Current time: {time_str}]"


@dataclass
class EnrichedRequest:
    """
    Request enriched with Steward's analysis.

    Attributes:
        original_request: The user's original message
        steward_note: Formatted note for Tatlock (includes context analysis)
        scoped_tools: List of tools from recommended capabilities
        recommendation: Full Steward recommendation
        steward_reasoning: Plain text reasoning for streaming to user
    """
    original_request: str
    steward_note: str
    scoped_tools: list[Any]  # PydanticAI tool definitions
    recommendation: StewardRecommendation
    steward_reasoning: str


async def preprocess_request(
    user_request: str,
    conversation_history: list[dict],
    conversation_id: Optional[str] = None,
) -> EnrichedRequest:
    """
    Analyze request via Steward and prepare scoped context for Tatlock.

    This is the main preprocessing pipeline that:
    1. Calls Steward with full conversation history
    2. Gets capability recommendations
    3. Creates scoped toolset from recommended capabilities
    4. Formats a note for Tatlock with context analysis

    Args:
        user_request: Current user message to analyze
        conversation_history: Full conversation history (all previous turns)
        conversation_id: Optional conversation ID for tracking

    Returns:
        EnrichedRequest with scoped tools and Steward analysis

    Example:
        >>> enriched = await preprocess_request(
        ...     "What's sqrt(144)?",
        ...     conversation_history=[],
        ... )
        >>> print(enriched.recommendation.recommended_capabilities)
        ['tatlock_core']
        >>> print(len(enriched.scoped_tools))
        5  # All tatlock_core tools
    """
    # Inject temporal context for time-aware processing
    enriched_request = _inject_temporal_context(user_request)

    logger.info(
        "preprocessing_request",
        request_preview=user_request[:100],
        history_length=len(conversation_history),
        conversation_id=conversation_id,
    )

    # Call Steward with full conversation history (traced)
    async with trace_span(
        "steward_analysis",
        SpanType.STEWARD,
        metadata={
            "request_preview": user_request[:100],
            "history_length": len(conversation_history),
        },
    ) as span:
        recommendation = await analyze_request(
            enriched_request,
            conversation_history=conversation_history,
            conversation_id=conversation_id,
        )

        # Update span with results
        if span:
            span.metadata.update({
                "recommended_capabilities": recommendation.recommended_capabilities,
                "complexity": recommendation.estimated_complexity,
                "has_memory_context": bool(recommendation.memory_context),
                "has_conversation_context": recommendation.conversation_context.has_previous_context,
            })
            span.details["reasoning"] = recommendation.reasoning
            if recommendation.enriched_query:
                span.details["enriched_query"] = recommendation.enriched_query

    # Format note for Tatlock (includes conversation context)
    steward_note = await format_steward_note(recommendation)

    # Get delegation tools from household registry
    # Uses agent-as-tool pattern: expert agents get delegation wrappers,
    # core tools are returned directly
    registry = get_household_registry()
    scoped_tools = registry.get_delegation_tools(
        recommendation.recommended_capabilities
    )

    logger.info(
        "preprocessing_complete",
        recommended_capabilities=recommendation.recommended_capabilities,
        tool_count=len(scoped_tools),
        complexity=recommendation.estimated_complexity,
        has_context=recommendation.conversation_context.has_previous_context,
    )

    return EnrichedRequest(
        original_request=enriched_request,
        steward_note=steward_note,
        scoped_tools=scoped_tools,
        recommendation=recommendation,
        steward_reasoning=recommendation.reasoning,
    )