lambdagent/agentpaas/engine/stream.py

"""
engine.stream — SSE and WebSocket stream adapters.
"""
from __future__ import annotations
import json
from typing import AsyncIterator
from fastapi import WebSocket
from fastapi.responses import StreamingResponse


async def sse_response(event_iter: AsyncIterator[dict]) -> StreamingResponse:
    """Convert async event iterator to SSE StreamingResponse."""
    async def generate():
        async for event in event_iter:
            evt_type = event.get("event", "message")
            data = json.dumps(event.get("data", {}), ensure_ascii=False)
            yield f"event: {evt_type}\ndata: {data}\n\n"

    return StreamingResponse(generate(), media_type="text/event-stream",
                             headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"})


async def ws_handler(websocket: WebSocket, execute_fn):
    """WebSocket bidirectional handler for multi-turn agent interaction."""
    await websocket.accept()
    try:
        while True:
            data = await websocket.receive_json()
            if data.get("type") == "close":
                break
            input_text = data.get("input", "")
            result = await execute_fn(input_text)
            await websocket.send_json({"type": "response", "content": result.output,
                                       "usage": {"duration_ms": result.duration_ms, "steps": result.steps}})
    except Exception:
        pass
    finally:
        await websocket.close()