tatlock/src/responses/service.py

"""
Response service for creating responses.

Handles both streaming and non-streaming response generation.
Tracks conversation history for analytics and future vector memory.
Integrates with Steward preprocessing for Phase 2 two-tier architecture.
"""

import time
import secrets
from typing import AsyncGenerator

from src.agents.registry import ModelRegistry
from src.responses.schemas import (
    Response,
    ResponseRequest,
    ResponseUsage,
    MessageOutputItem,
    ReasoningOutputItem,
    FunctionCallOutputItem,
    OutputTextContent,
)
from src.responses.streaming import StreamingCoordinator
from src.responses.history import ConversationHistory
from src.responses.context import ContextWindow
from src.core.preprocessing import preprocess_request
from src.core.tool_tracking import ToolCallTracker
from src.core.logging_config import get_logger

logger = get_logger(__name__)

# Global conversation history tracker
# In production, this would be backed by a database or Redis
_conversation_history = ConversationHistory(max_turns=20)

# Global context window manager (4096 token default)
_context_window = ContextWindow(max_tokens=4096)


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


def _calculate_usage(input_messages: list[dict], output_items: list) -> ResponseUsage:
    """
    Calculate token usage for the response.

    For now: approximate token counting.
    Future: Use tiktoken or similar for accurate counting.

    Args:
        input_messages: Input messages
        output_items: Output items generated

    Returns:
        ResponseUsage: Token usage statistics
    """
    # Approximate input tokens (chars / 4)
    input_text = " ".join(str(msg) for msg in input_messages)
    input_tokens = len(input_text) // 4

    # Approximate output tokens
    output_tokens = 0
    reasoning_tokens = 0

    for item in output_items:
        # Check if it's a schema object (has summary/content attributes directly)
        if isinstance(item, ReasoningOutputItem):
            reasoning_text = " ".join(item.summary)
            reasoning_tokens += len(reasoning_text) // 4
        elif isinstance(item, MessageOutputItem):
            message_text = item.content[0].text
            output_tokens += len(message_text) // 4
        elif isinstance(item, FunctionCallOutputItem):
            func_text = item.arguments
            output_tokens += len(func_text) // 4
        elif hasattr(item, 'type'):
            # Agent OutputItem objects (backward compatibility)
            if item.type == "reasoning":
                reasoning_text = " ".join(item.data.get("summary", []))
                reasoning_tokens += len(reasoning_text) // 4
            elif item.type == "message":
                message_text = item.data["content"][0]["text"]
                output_tokens += len(message_text) // 4
            elif item.type == "function_call":
                func_text = item.data["arguments"]
                output_tokens += len(func_text) // 4

    total_tokens = input_tokens + output_tokens + reasoning_tokens

    return ResponseUsage(
        input_tokens=input_tokens,
        output_tokens=output_tokens,
        reasoning_tokens=reasoning_tokens,
        total_tokens=total_tokens
    )


async def create_response(request: ResponseRequest) -> Response:
    """
    Create non-streaming response.

    Tracks conversation history if conversation_id is provided in metadata.

    Args:
        request: Response request

    Returns:
        Response: Complete response object

    Example:
        request = ResponseRequest(
            model="lorem-tester",
            input=[{"role": "user", "content": "Hello"}],
            reasoning={"effort": "medium", "summary": "auto"},
            metadata={"conversation_id": "conv_abc123"}  # Optional
        )
        response = await create_response(request)
    """
    # Get or generate conversation ID
    conversation_id = await _conversation_history.get_conversation_id(request)

    # Strip pipeline prefix if present (e.g., "pipeline.model" -> "model")
    model_id = request.model
    if "." in model_id:
        model_id = model_id.split(".", 1)[1]

    # Get agent for model
    agent = ModelRegistry.get_agent(model_id)

    # Collect all output items from agent
    output_items = []
    async for item in agent.generate_response(
        messages=request.input,
        reasoning=request.reasoning,
        tools=request.tools,
        temperature=request.temperature,
        max_tokens=request.max_output_tokens,
        stop=request.stop,
    ):
        output_items.append(item)

    # Convert agent OutputItems to schema OutputItems
    converted_items = _convert_output_items(output_items)

    # Calculate token usage
    usage = _calculate_usage(request.input, output_items)

    response = Response(
        id=f"resp_{generate_id()}",
        created_at=int(time.time()),
        model=request.model,
        status="completed",
        output=converted_items,
        usage=usage
    )

    # Track conversation history (for analytics and future vector memory)
    await _conversation_history.add_response(conversation_id, response)

    return response


async def create_response_with_steward(request: ResponseRequest) -> Response:
    """
    Create response using Steward preprocessing (Phase 2 flow).

    This is the two-tier architecture where:
    1. Steward analyzes the request and recommends capabilities
    2. Tatlock runs with scoped tools based on recommendations
    3. Tool usage is tracked for benchmarking

    Args:
        request: Response request

    Returns:
        Response: Complete response object with Steward analysis included

    Example:
        request = ResponseRequest(
            model="tatlock",
            input=[{"role": "user", "content": "What's sqrt(144)?"}],
            metadata={"conversation_id": "conv_abc123"}
        )
        response = await create_response_with_steward(request)
    """
    # Get or generate conversation ID
    conversation_id = await _conversation_history.get_conversation_id(request)

    # Extract user message and conversation history
    user_message = ""
    for msg in reversed(request.input):
        if msg.get("role") == "user":
            user_message = msg.get("content", "")
            break

    # Conversation history is all messages except the current one
    conversation_history = request.input[:-1] if len(request.input) > 1 else []

    logger.info(
        "creating_response_with_steward",
        user_message_preview=user_message[:100],
        history_length=len(conversation_history),
        conversation_id=conversation_id,
    )

    # Phase 1: Steward preprocessing
    enriched = await preprocess_request(
        user_message,
        conversation_history=conversation_history,
        conversation_id=conversation_id,
    )

    # Phase 2: Initialize tool tracker
    tracker = ToolCallTracker(
        recommended_capabilities=enriched.recommendation.recommended_capabilities,
        conversation_id=conversation_id,
    )

    # Phase 3: Run Tatlock with scoped tools
    from src.agents.tatlock import TatlockAgent
    tatlock = TatlockAgent()

    tatlock_response = await tatlock.run_with_scoped_tools(
        user_message=user_message,
        steward_note=enriched.steward_note,
        scoped_tools=enriched.scoped_tools,
        message_history=conversation_history,
        tool_tracker=tracker,
    )

    # Phase 4: Finalize tool tracking
    await tracker.finalize()

    # Build response output items
    output_items = []

    # Add Steward reasoning as a reasoning output item
    output_items.append(ReasoningOutputItem(
        id=f"reasoning_{generate_id()}",
        summary=[
            "🎩 Steward's Analysis:",
            enriched.steward_reasoning,
        ],
        status="completed"
    ))

    # Add Tatlock's message
    output_items.append(MessageOutputItem(
        id=f"msg_{generate_id()}",
        role="assistant",
        content=[OutputTextContent(
            type="output_text",
            text=tatlock_response,
            annotations=[]
        )],
        status="completed"
    ))

    # Calculate usage (approximate)
    usage = _calculate_usage(request.input, output_items)

    response = Response(
        id=f"resp_{generate_id()}",
        created_at=int(time.time()),
        model=request.model,
        status="completed",
        output=output_items,
        usage=usage
    )

    # Track conversation history
    await _conversation_history.add_response(conversation_id, response)

    logger.info(
        "response_with_steward_complete",
        response_id=response.id,
        recommended_capabilities=enriched.recommendation.recommended_capabilities,
        tool_summary=tracker.get_summary(),
    )

    return response


async def create_response_stream(
    request: ResponseRequest
) -> AsyncGenerator[dict, None]:
    """
    Create streaming response.

    Yields SSE-formatted events for streaming to client.

    Args:
        request: Response request

    Yields:
        dict: SSE event dict with 'event' and 'data' keys

    Example:
        async for event in create_response_stream(request):
            # event = {"event": "response.output_text.delta", "data": "..."}
            yield event
    """
    coordinator = StreamingCoordinator()

    async for event in coordinator.stream_response(request):
        yield {
            "event": event.event,
            "data": event.model_dump_json()
        }


async def get_conversation_history(conversation_id: str) -> list[Response]:
    """
    Get conversation history for a given conversation ID.

    Args:
        conversation_id: Conversation identifier

    Returns:
        list: List of Response objects
    """
    return await _conversation_history.get_history(conversation_id)


async def get_conversation_stats() -> dict:
    """
    Get conversation history statistics.

    Returns:
        dict: Statistics including total conversations tracked
    """
    return {
        "total_conversations": await _conversation_history.get_conversation_count(),
        "max_turns_per_conversation": _conversation_history._max_turns
    }


async def clear_conversation(conversation_id: str) -> bool:
    """
    Clear a specific conversation history.

    Args:
        conversation_id: Conversation identifier

    Returns:
        bool: True if conversation was cleared
    """
    return await _conversation_history.clear_conversation(conversation_id)


def get_context_window() -> ContextWindow:
    """
    Get the global context window manager.

    Returns:
        ContextWindow: Context window manager instance
    """
    return _context_window


def _convert_output_items(items: list) -> list:
    """
    Convert agent OutputItem objects to schema OutputItem objects.

    Args:
        items: List of agent OutputItem objects

    Returns:
        list: List of schema OutputItem objects
    """
    converted = []

    for item in items:
        if item.type == "message":
            converted.append(MessageOutputItem(
                id=item.id,
                content=[OutputTextContent(**c) for c in item.data["content"]],
                status=item.data.get("status", "completed")
            ))
        elif item.type == "reasoning":
            converted.append(ReasoningOutputItem(
                id=item.id,
                summary=item.data["summary"],
                status=item.data.get("status", "completed")
            ))
        elif item.type == "function_call":
            converted.append(FunctionCallOutputItem(
                id=item.id,
                name=item.data["name"],
                arguments=item.data["arguments"],
                status=item.data.get("status", "completed")
            ))

    return converted