tatlock/src/responses/streaming.py

"""
Streaming event schemas and coordinator for Responses API.

Handles Server-Sent Events (SSE) streaming with proper event types:
- response.reasoning_summary_text.delta
- response.output_text.delta
- response.function_call_arguments.delta
- response.done
"""

from enum import Enum
from typing import Literal, AsyncGenerator
import json
import time

from src.core.models import CustomBaseModel
from src.responses.schemas import Response
from src.agents.registry import ModelRegistry


# ============================================================================
# Stream Event Types
# ============================================================================

class StreamEventType(str, Enum):
    """Streaming event types for Responses API."""
    REASONING_SUMMARY_DELTA = "response.reasoning_summary_text.delta"
    REASONING_SUMMARY_DONE = "response.reasoning_summary_text.done"
    OUTPUT_TEXT_DELTA = "response.output_text.delta"
    OUTPUT_TEXT_DONE = "response.output_text.done"
    FUNCTION_CALL_DELTA = "response.function_call_arguments.delta"
    FUNCTION_CALL_DONE = "response.function_call_arguments.done"
    RESPONSE_DONE = "response.done"
    ERROR = "error"


# ============================================================================
# Stream Event Schemas
# ============================================================================

class ReasoningSummaryDelta(CustomBaseModel):
    """Reasoning summary text delta event."""
    event: Literal[StreamEventType.REASONING_SUMMARY_DELTA] = StreamEventType.REASONING_SUMMARY_DELTA
    delta: str


class ReasoningSummaryDone(CustomBaseModel):
    """Reasoning summary completion event."""
    event: Literal[StreamEventType.REASONING_SUMMARY_DONE] = StreamEventType.REASONING_SUMMARY_DONE


class OutputTextDelta(CustomBaseModel):
    """Output text delta event."""
    event: Literal[StreamEventType.OUTPUT_TEXT_DELTA] = StreamEventType.OUTPUT_TEXT_DELTA
    delta: str


class OutputTextDone(CustomBaseModel):
    """Output text completion event."""
    event: Literal[StreamEventType.OUTPUT_TEXT_DONE] = StreamEventType.OUTPUT_TEXT_DONE


class FunctionCallDelta(CustomBaseModel):
    """Function call arguments delta event."""
    event: Literal[StreamEventType.FUNCTION_CALL_DELTA] = StreamEventType.FUNCTION_CALL_DELTA
    delta: str
    name: str | None = None  # Only in first chunk


class FunctionCallDone(CustomBaseModel):
    """Function call completion event."""
    event: Literal[StreamEventType.FUNCTION_CALL_DONE] = StreamEventType.FUNCTION_CALL_DONE


class ResponseDone(CustomBaseModel):
    """Response completion event with full response."""
    event: Literal[StreamEventType.RESPONSE_DONE] = StreamEventType.RESPONSE_DONE
    response: Response


class ErrorEvent(CustomBaseModel):
    """Error event."""
    event: Literal[StreamEventType.ERROR] = StreamEventType.ERROR
    error: dict


# Union type for all stream events
StreamEvent = (
    ReasoningSummaryDelta |
    ReasoningSummaryDone |
    OutputTextDelta |
    OutputTextDone |
    FunctionCallDelta |
    FunctionCallDone |
    ResponseDone |
    ErrorEvent
)


# ============================================================================
# Streaming Coordinator
# ============================================================================

class StreamingCoordinator:
    """
    Coordinates streaming from agents to SSE format.

    Real production streaming logic that handles:
    1. Reasoning summary chunks
    2. Function call arguments
    3. Output text chunks
    4. Error handling
    5. Final response event
    """

    async def stream_response_with_steward(
        self,
        request: "ResponseRequest"  # type: ignore # Forward reference
    ) -> AsyncGenerator[StreamEvent, None]:
        """
        Stream response with Steward preprocessing and two-phase Tatlock execution.

        Streams in order:
        1. Steward's analysis as reasoning summary
        2. Think slugs during expert delegation (butler-perspective messages)
        3. Synthesized butler-toned response as output text

        Args:
            request: Response request

        Yields:
            StreamEvent: Stream of SSE events
        """
        from src.responses.service import (
            _calculate_usage,
            generate_id,
            _conversation_history,
            _direct_delegation_with_results,
        )
        from src.core.preprocessing import preprocess_request
        from src.core.tool_tracking import ToolCallTracker
        from src.responses.schemas import MessageOutputItem, ReasoningOutputItem, OutputTextContent
        from src.agents.tatlock import TatlockAgent
        from src.agents.delegation import get_think_message, STREAMING_DELEGATION_WRAPPERS
        import asyncio

        output_items = []

        try:
            # 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 = request.input[:-1] if len(request.input) > 1 else []

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

            # Stream Steward's analysis as reasoning summary
            steward_lines = enriched.steward_reasoning.split('\n')
            for line in steward_lines:
                if line.strip():
                    yield ReasoningSummaryDelta(delta=line + "\n")
                    await asyncio.sleep(0.05)

            yield ReasoningSummaryDone()

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

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

            # Check if direct delegation is recommended
            delegation_agents = {"biographer", "librarian", "housekeeper"}
            delegation_only = all(
                cap in delegation_agents
                for cap in enriched.recommendation.recommended_capabilities
            ) and enriched.recommendation.recommended_capabilities

            tatlock = TatlockAgent()

            if delegation_only:
                # Direct delegation path with streaming think slugs
                orchestration_results = await self._stream_direct_delegation(
                    user_message=user_message,
                    recommendation=enriched.recommendation,
                    tracker=tracker,
                    conversation_id=conversation_id,
                )

                # Stream think slugs that were collected during delegation
                # Each think message is complete, so we signal done after each
                for think_msg in orchestration_results.get("think_messages", []):
                    yield ReasoningSummaryDelta(delta=think_msg)
                    yield ReasoningSummaryDone()
                    await asyncio.sleep(0.05)

            else:
                # Phase 1: Orchestrate tool calls
                orchestration_results = await tatlock.orchestrate_tool_calls(
                    user_message=user_message,
                    steward_note=enriched.steward_note,
                    scoped_tools=enriched.scoped_tools,
                    message_history=conversation_history,
                    tool_tracker=tracker,
                )

            # Phase 2: Synthesize butler-toned response
            tatlock_response = await tatlock.synthesize_from_results(
                user_message=user_message,
                orchestration_results=orchestration_results,
                message_history=conversation_history,
            )

            # Stream the synthesized response
            chunk_size = 50
            for i in range(0, len(tatlock_response), chunk_size):
                yield OutputTextDelta(delta=tatlock_response[i:i + chunk_size])
                await asyncio.sleep(0.02)

            yield OutputTextDone()

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

            # Finalize tool tracking
            await tracker.finalize()

            # Calculate usage and build final response
            usage = _calculate_usage(request.input, output_items)

            final_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, final_response)

            yield ResponseDone(response=final_response)

        except Exception as e:
            # Stream error event
            yield self._create_error_event(e)

    async def _stream_direct_delegation(
        self,
        user_message: str,
        recommendation: "StewardRecommendation",  # type: ignore
        tracker: "ToolCallTracker",  # type: ignore
        conversation_id: str,
    ) -> dict:
        """
        Execute direct delegation with streaming think messages.

        Collects think messages as delegations execute for streaming to client.

        Args:
            user_message: User's request
            recommendation: Steward's recommendation
            tracker: Tool call tracker
            conversation_id: Conversation ID

        Returns:
            dict: Orchestration results with think_messages list
        """
        from src.agents.delegation import (
            get_think_message,
            delegate_to_librarian,
            delegate_to_biographer,
            delegate_to_housekeeper,
        )
        import time as time_module

        expert_results = {}
        tools_called = []
        think_messages = []

        for agent in recommendation.recommended_capabilities:
            # Emit start think message
            start_msg = get_think_message(agent, user_message, "start")
            think_messages.append(start_msg + "\n")

            start_time = time_module.time()
            try:
                # Execute delegation
                if agent == "librarian":
                    result = await delegate_to_librarian(task=user_message)
                elif agent == "biographer":
                    result = await delegate_to_biographer(task=user_message)
                elif agent == "housekeeper":
                    result = await delegate_to_housekeeper(task=user_message)
                else:
                    result = None

                duration = time_module.time() - start_time
                await tracker.track_call(f"delegate_to_{agent}", duration)

                if result and result.success:
                    expert_results[agent] = result.output
                    tools_called.append(f"delegate_to_{agent}")
                    # Emit success think message
                    success_msg = get_think_message(agent, user_message, "success")
                    think_messages.append(success_msg + "\n")
                else:
                    error_msg = result.error if result else "Unknown error"
                    expert_results[agent] = f"Error: {error_msg}"
                    # Emit error think message
                    error_think = get_think_message(agent, user_message, "error")
                    think_messages.append(error_think + "\n")

            except Exception as e:
                expert_results[agent] = f"Error: {e}"
                error_think = get_think_message(agent, user_message, "error")
                think_messages.append(error_think + "\n")

        return {
            "tools_called": tools_called,
            "expert_results": expert_results,
            "tool_outputs": {},
            "raw_output": "",
            "think_messages": think_messages,
        }

    async def stream_response(
        self,
        request: "ResponseRequest"  # type: ignore # Forward reference
    ) -> AsyncGenerator[StreamEvent, None]:
        """
        Coordinate streaming from agent to SSE events.

        Args:
            request: Response request

        Yields:
            StreamEvent: Stream of SSE events

        Example SSE output:
            event: response.reasoning_summary_text.delta
            data: {"delta": "Analyzing..."}

            event: response.output_text.delta
            data: {"delta": "Lorem"}

            event: response.done
            data: {"response": {...}}
        """
        from src.responses.service import _calculate_usage, generate_id
        import asyncio

        output_items = []
        last_message_text = ""  # Track last streamed message text to compute deltas

        try:
            # 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)
            # Stream from agent
            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)

                # Stream based on item type
                if item.type == "reasoning":
                    # Stream reasoning summary
                    for step in item.data.get("summary", []):
                        yield ReasoningSummaryDelta(delta=step + " ")
                        await asyncio.sleep(0.05)  # Simulate typing

                    yield ReasoningSummaryDone()

                elif item.type == "function_call":
                    # Stream function call arguments
                    # First chunk includes name
                    yield FunctionCallDelta(
                        name=item.data["name"],
                        delta=""
                    )

                    # Stream arguments in chunks
                    args = item.data["arguments"]
                    chunk_size = 20
                    for i in range(0, len(args), chunk_size):
                        yield FunctionCallDelta(
                            delta=args[i:i+chunk_size]
                        )
                        await asyncio.sleep(0.03)

                    yield FunctionCallDone()

                elif item.type == "message":
                    # Get current accumulated text from agent
                    current_text = item.data["content"][0]["text"]

                    # Only stream the NEW text (delta) since last update
                    if current_text.startswith(last_message_text):
                        # Extract only the new portion
                        delta_text = current_text[len(last_message_text):]

                        if delta_text:
                            # Stream the delta text in chunks while preserving formatting
                            # (newlines, markdown, code blocks, etc.)
                            chunk_size = 50  # characters per chunk

                            for i in range(0, len(delta_text), chunk_size):
                                chunk = delta_text[i:i+chunk_size]

                                # Check stop sequences on full accumulated text
                                stop_found, text_before_stop = self._check_stop_sequence(
                                    current_text,
                                    request.stop
                                )

                                if stop_found:
                                    # Only emit remaining delta before stop
                                    remaining = text_before_stop[len(last_message_text):]
                                    if remaining:
                                        yield OutputTextDelta(delta=remaining)
                                    yield OutputTextDone()
                                    break

                                # Check max tokens on full text
                                output_tokens = self._count_tokens_approx(current_text)
                                if self._check_max_tokens(output_tokens, request.max_output_tokens):
                                    yield OutputTextDone()
                                    break

                                # Normal streaming of delta chunk (preserves all formatting)
                                yield OutputTextDelta(delta=chunk)
                                await asyncio.sleep(0.02)  # Shorter delay since chunks are larger

                        # Update tracking variable
                        last_message_text = current_text

                    # If this is the final message (status=completed), ensure we send done
                    if item.data.get("status") == "completed":
                        yield OutputTextDone()

            # Final response.done event with complete response
            usage = _calculate_usage(request.input, output_items)

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

            # Track conversation history (import here to avoid circular dependency)
            from src.responses.service import _conversation_history
            conversation_id = await _conversation_history.get_conversation_id(request)
            await _conversation_history.add_response(conversation_id, final_response)

            yield ResponseDone(response=final_response)

        except Exception as e:
            # Stream error event
            yield self._create_error_event(e)

    def _convert_output_items(self, items: list) -> list:
        """Convert agent OutputItem objects to schema OutputItem objects."""
        from src.responses.schemas import (
            MessageOutputItem,
            ReasoningOutputItem,
            FunctionCallOutputItem,
            OutputTextContent,
        )

        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

    def _create_error_event(self, error: Exception) -> ErrorEvent:
        """Create error event from exception."""
        from src.core.exceptions import (
            RateLimitError,
            ContextLengthError,
            AppException,
        )

        if isinstance(error, RateLimitError):
            error_type = "rate_limit_exceeded"
            code = 429
        elif isinstance(error, ContextLengthError):
            error_type = "context_length_exceeded"
            code = 400
        elif isinstance(error, AppException):
            error_type = "app_error"
            code = error.status_code
        else:
            error_type = "internal_error"
            code = 500

        return ErrorEvent(
            error={
                "type": error_type,
                "message": str(error),
                "code": code
            }
        )

    def _check_stop_sequence(
        self,
        accumulated_text: str,
        stop_sequences: list[str] | None
    ) -> tuple[bool, str]:
        """
        Check if any stop sequence is encountered.

        Args:
            accumulated_text: Text accumulated so far
            stop_sequences: List of stop sequences to check

        Returns:
            tuple: (stop_found, text_before_stop)
                - stop_found: True if stop sequence found
                - text_before_stop: Text before the stop sequence (or full text if not found)
        """
        if not stop_sequences:
            return False, accumulated_text

        for stop_seq in stop_sequences:
            if stop_seq in accumulated_text:
                # Found stop sequence - return text before it
                idx = accumulated_text.index(stop_seq)
                return True, accumulated_text[:idx]

        return False, accumulated_text

    def _count_tokens_approx(self, text: str) -> int:
        """
        Approximate token count (4 chars per token).

        Args:
            text: Text to count

        Returns:
            int: Approximate token count
        """
        return len(text) // 4

    def _check_max_tokens(
        self,
        current_tokens: int,
        max_tokens: int | None
    ) -> bool:
        """
        Check if max tokens limit reached.

        Args:
            current_tokens: Current token count
            max_tokens: Maximum allowed tokens (None for unlimited)

        Returns:
            bool: True if limit reached
        """
        if max_tokens is None:
            return False
        return current_tokens >= max_tokens