stream_utils.py

"""
SSE (Server-Sent Events) Utilities
Format streaming responses for real-time chat
"""
import json
from typing import Dict, Any, AsyncGenerator
import asyncio


def format_sse(event: str, data: Any) -> str:
    """
    Format data as SSE message
    
    Args:
        event: Event type (token, status, done, error)
        data: Data payload (string or dict)
    
    Returns:
        Formatted SSE string
        
    Example:
        format_sse("token", "Hello")
        # "event: token\ndata: Hello\n\n"
    """
    if isinstance(data, dict):
        data_str = json.dumps(data, ensure_ascii=False)
    else:
        data_str = str(data)
    
    return f"event: {event}\ndata: {data_str}\n\n"


async def simulate_typing(
    text: str,
    chars_per_chunk: int = 3,
    delay_ms: float = 20
) -> AsyncGenerator[str, None]:
    """
    Simulate typing effect by yielding text in chunks
    
    Args:
        text: Full text to stream
        chars_per_chunk: Characters per chunk
        delay_ms: Milliseconds delay between chunks
        
    Yields:
        Text chunks
        
    Example:
        async for chunk in simulate_typing("Hello world", chars_per_chunk=2):
            yield format_sse("token", chunk)
    """
    for i in range(0, len(text), chars_per_chunk):
        chunk = text[i:i + chars_per_chunk]
        yield chunk
        await asyncio.sleep(delay_ms / 1000)


async def stream_text_slowly(
    text: str,
    event_type: str = "token",
    chars_per_chunk: int = 3,
    delay_ms: float = 20
) -> AsyncGenerator[str, None]:
    """
    Stream text with typing effect in SSE format
    
    Args:
        text: Text to stream
        event_type: SSE event type
        chars_per_chunk: Characters per chunk
        delay_ms: Delay between chunks
        
    Yields:
        SSE formatted chunks
    """
    async for chunk in simulate_typing(text, chars_per_chunk, delay_ms):
        yield format_sse(event_type, chunk)


# Event type constants
EVENT_STATUS = "status"
EVENT_TOKEN = "token"
EVENT_DONE = "done"
EVENT_ERROR = "error"
EVENT_METADATA = "metadata"