Streaming payload models and serializers

Author: asimurka

src/models/common/streaming/__init__.py

@@ -0,0 +1,37 @@
+"""Streaming payload models and event type exports."""
+
+from models.common.streaming.stream_payloads import (
+    EndEventData,
+    EndStreamPayload,
+    ErrorEventData,
+    ErrorStreamPayload,
+    InterruptedEventData,
+    InterruptedStreamPayload,
+    StartEventData,
+    StartStreamPayload,
+    StreamEventPayload,
+    StreamPayloadBase,
+    TokenChunkData,
+    TokenStreamPayload,
+    ToolCallStreamPayload,
+    ToolResultStreamPayload,
+    TurnCompleteStreamPayload,
+)
+
+__all__ = [
+    "StreamPayloadBase",
+    "ErrorEventData",
+    "StartEventData",
+    "InterruptedEventData",
+    "EndEventData",
+    "ErrorStreamPayload",
+    "StartStreamPayload",
+    "InterruptedStreamPayload",
+    "EndStreamPayload",
+    "TokenChunkData",
+    "TokenStreamPayload",
+    "TurnCompleteStreamPayload",
+    "ToolCallStreamPayload",
+    "ToolResultStreamPayload",
+    "StreamEventPayload",
+]

src/models/common/streaming/stream_payloads.py

@@ -0,0 +1,120 @@
+"""Typed JSON bodies for SSE streaming events."""
+
+from typing import Annotated, Literal, Optional, TypeAlias
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from models.common import ReferencedDocument, ToolCallSummary, ToolResultSummary
+
+
+class StreamPayloadBase(BaseModel):
+    """Base for streaming SSE JSON payloads."""
+
+    model_config = ConfigDict(extra="forbid")
+
+
+class ErrorEventData(StreamPayloadBase):
+    """Payload for event: "error"."""
+
+    status_code: int
+    response: str
+    cause: str
+
+
+class StartEventData(StreamPayloadBase):
+    """Payload for event: "start"."""
+
+    conversation_id: str
+    request_id: str
+
+
+class InterruptedEventData(StreamPayloadBase):
+    """Payload for event: "interrupted"."""
+
+    request_id: str
+
+
+class EndEventData(StreamPayloadBase):
+    """Nested data for event: "end"."""
+
+    referenced_documents: list[ReferencedDocument]
+    truncated: Optional[bool]
+    input_tokens: int
+    output_tokens: int
+
+
+class ErrorStreamPayload(StreamPayloadBase):
+    """SSE error event body (event + typed data)."""
+
+    event: Literal["error"] = "error"
+    data: ErrorEventData
+
+
+class StartStreamPayload(StreamPayloadBase):
+    """SSE stream start body."""
+
+    event: Literal["start"] = "start"
+    data: StartEventData
+
+
+class InterruptedStreamPayload(StreamPayloadBase):
+    """SSE interrupted stream body."""
+
+    event: Literal["interrupted"] = "interrupted"
+    data: InterruptedEventData
+
+
+class EndStreamPayload(StreamPayloadBase):
+    """SSE end-of-stream body (includes available_quotas beside data)."""
+
+    event: Literal["end"] = "end"
+    data: EndEventData
+    available_quotas: dict[str, int]
+
+
+class TokenChunkData(StreamPayloadBase):
+    """Structured data for token and turn-complete stream lines."""
+
+    id: int
+    token: str
+
+
+class TokenStreamPayload(StreamPayloadBase):
+    """SSE token delta (event: "token")."""
+
+    event: Literal["token"] = "token"
+    data: TokenChunkData
+
+
+class TurnCompleteStreamPayload(StreamPayloadBase):
+    """SSE turn completion (same data shape as token)."""
+
+    event: Literal["turn_complete"] = "turn_complete"
+    data: TokenChunkData
+
+
+class ToolCallStreamPayload(StreamPayloadBase):
+    """SSE tool call summary."""
+
+    event: Literal["tool_call"] = "tool_call"
+    data: ToolCallSummary
+
+
+class ToolResultStreamPayload(StreamPayloadBase):
+    """SSE tool result summary."""
+
+    event: Literal["tool_result"] = "tool_result"
+    data: ToolResultSummary
+
+
+StreamEventPayload: TypeAlias = Annotated[
+    TokenStreamPayload
+    | TurnCompleteStreamPayload
+    | ToolCallStreamPayload
+    | ToolResultStreamPayload
+    | EndStreamPayload
+    | ErrorStreamPayload
+    | InterruptedStreamPayload
+    | StartStreamPayload,
+    Field(discriminator="event"),
+]

src/utils/streaming/event_serializers.py

@@ -0,0 +1,161 @@
+"""Shared streaming event formatting utilities."""
+
+import json
+from functools import singledispatch
+from typing import Optional
+
+from constants import MEDIA_TYPE_JSON
+from log import get_logger
+from models.api.responses.error import (
+    AbstractErrorResponse,
+)
+from models.common.streaming import (
+    EndEventData,
+    EndStreamPayload,
+    ErrorEventData,
+    ErrorStreamPayload,
+    InterruptedEventData,
+    InterruptedStreamPayload,
+    StartEventData,
+    StartStreamPayload,
+    StreamEventPayload,
+    StreamPayloadBase,
+    TokenStreamPayload,
+    ToolCallStreamPayload,
+    ToolResultStreamPayload,
+    TurnCompleteStreamPayload,
+)
+from models.common.turn_summary import ReferencedDocument
+from utils.token_counter import TokenCounter
+
+logger = get_logger(__name__)
+
+
+def format_stream_data(data: StreamPayloadBase) -> str:
+    """Format a Pydantic payload as an SSE ``data:`` line."""
+    return f"data: {json.dumps(data.model_dump(mode='json'))}\n\n"
+
+
+def serialize_http_error_event(
+    error: AbstractErrorResponse,
+    media_type: Optional[str] = None,
+) -> str:
+    """Serialize an API error to an SSE or plain-text client response."""
+    logger.error("Error while obtaining answer for user question")
+    resolved_media_type = media_type or MEDIA_TYPE_JSON
+    payload = ErrorStreamPayload(
+        data=ErrorEventData(
+            status_code=error.status_code,
+            response=error.detail.response,
+            cause=error.detail.cause,
+        ),
+    )
+    return serialize_event(payload, resolved_media_type)
+
+
+def serialize_start_event(
+    conversation_id: str,
+    request_id: str,
+    media_type: str = MEDIA_TYPE_JSON,
+) -> str:
+    """Serialize the stream start payload to an SSE line."""
+    payload = StartStreamPayload(
+        data=StartEventData(
+            conversation_id=conversation_id,
+            request_id=request_id,
+        ),
+    )
+    return serialize_event(payload, media_type)
+
+
+def serialize_interrupted_event(
+    request_id: str, media_type: str = MEDIA_TYPE_JSON
+) -> str:
+    """Serialize an interrupted-stream payload to an SSE line."""
+    payload = InterruptedStreamPayload(
+        data=InterruptedEventData(request_id=request_id),
+    )
+    return serialize_event(payload, media_type)
+
+
+def serialize_end_event(
+    token_usage: TokenCounter,
+    available_quotas: dict[str, int],
+    referenced_documents: list[ReferencedDocument],
+    media_type: Optional[str] = None,
+) -> str:
+    """Serialize the stream end payload for JSON SSE or plain-text clients."""
+    resolved_media_type = media_type or MEDIA_TYPE_JSON
+    payload = EndStreamPayload(
+        data=EndEventData(
+            referenced_documents=referenced_documents,
+            truncated=None,
+            input_tokens=token_usage.input_tokens,
+            output_tokens=token_usage.output_tokens,
+        ),
+        available_quotas=available_quotas,
+    )
+    return serialize_event(payload, resolved_media_type)
+
+
+def serialize_event(
+    payload: StreamEventPayload,
+    media_type: str = MEDIA_TYPE_JSON,
+) -> str:
+    """Serialize an LLM stream payload (token, tool, turn complete) for the client."""
+    if media_type == MEDIA_TYPE_JSON:
+        return format_stream_data(payload)
+    return serialize_event_text(payload)
+
+
+@singledispatch
+def serialize_event_text(_payload: StreamPayloadBase) -> str:
+    """Serialize stream payload to plain text for text media type."""
+    return ""
+
+
+@serialize_event_text.register
+def _(payload: TokenStreamPayload) -> str:
+    """Serialize token stream payload to plain text."""
+    return payload.data.token
+
+
+@serialize_event_text.register
+def _(_payload: TurnCompleteStreamPayload) -> str:
+    """Serialize turn complete stream payload to plain text."""
+    return ""
+
+
+@serialize_event_text.register
+def _(payload: ToolCallStreamPayload) -> str:
+    """Serialize tool call stream payload to plain text."""
+    return f"[Tool Call: {payload.data.name}]\n"
+
+
+@serialize_event_text.register
+def _(_payload: ToolResultStreamPayload) -> str:
+    """Serialize tool result stream payload to plain text."""
+    return "[Tool Result]\n"
+
+
+@serialize_event_text.register
+def _(payload: EndStreamPayload) -> str:
+    """Serialize end stream payload to plain text."""
+    ref_docs_string = "\n".join(
+        f"{doc.doc_title}: {doc.doc_url}"
+        for doc in payload.data.referenced_documents
+        if doc.doc_url and doc.doc_title
+    )
+    return f"\n\n---\n\n{ref_docs_string}" if ref_docs_string else ""
+
+
+@serialize_event_text.register
+def _(payload: ErrorStreamPayload) -> str:
+    """Serialize error stream payload to plain text."""
+    return f"Status: {payload.data.status_code} - {payload.data.response} - {payload.data.cause}"
+
+
+@serialize_event_text.register
+def _(_payload: StartStreamPayload) -> str:
+    """Serialize start stream payload to plain text."""
+    return ""