76 lines
2.6 KiB
Python
76 lines
2.6 KiB
Python
|
"""SSE streaming endpoint for token-by-token LLM responses."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
import json
|
||
|
|
import uuid
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
from fastapi import APIRouter
|
||
|
|
from fastapi.responses import StreamingResponse
|
||
|
|
|
||
|
|
from fusionagi._logger import logger
|
||
|
|
from fusionagi.api.dependencies import get_orchestrator
|
||
|
|
|
||
|
|
router = APIRouter()
|
||
|
|
|
||
|
|
|
||
|
|
async def _sse_generator(session_id: str, prompt: str) -> Any:
|
||
|
|
"""Generate SSE events for a streaming prompt response."""
|
||
|
|
event_id = str(uuid.uuid4())[:8]
|
||
|
|
|
||
|
|
yield f"event: start\ndata: {json.dumps({'session_id': session_id, 'event_id': event_id})}\n\n"
|
||
|
|
|
||
|
|
orch = get_orchestrator()
|
||
|
|
if orch is None:
|
||
|
|
yield f"event: error\ndata: {json.dumps({'error': 'Orchestrator not available'})}\n\n"
|
||
|
|
return
|
||
|
|
|
||
|
|
try:
|
||
|
|
yield f"event: heads_running\ndata: {json.dumps({'heads': ['logic', 'creativity', 'research', 'safety']})}\n\n"
|
||
|
|
|
||
|
|
from fusionagi.schemas.task import Task
|
||
|
|
task = Task(task_id=f"stream_{event_id}", prompt=prompt)
|
||
|
|
result = orch.run(task)
|
||
|
|
|
||
|
|
if result and hasattr(result, "final_answer"):
|
||
|
|
answer = result.final_answer or ""
|
||
|
|
# Stream token-by-token (simulate chunked response)
|
||
|
|
words = answer.split()
|
||
|
|
for i, word in enumerate(words):
|
||
|
|
chunk = word + (" " if i < len(words) - 1 else "")
|
||
|
|
yield f"event: token\ndata: {json.dumps({'token': chunk, 'index': i})}\n\n"
|
||
|
|
await asyncio.sleep(0.02)
|
||
|
|
|
||
|
|
yield f"event: complete\ndata: {json.dumps({'session_id': session_id, 'full_text': answer})}\n\n"
|
||
|
|
else:
|
||
|
|
yield f"event: complete\ndata: {json.dumps({'session_id': session_id, 'full_text': ''})}\n\n"
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
logger.error("SSE streaming error", extra={"error": str(e), "session_id": session_id})
|
||
|
|
yield f"event: error\ndata: {json.dumps({'error': str(e)})}\n\n"
|
||
|
|
|
||
|
|
|
||
|
|
@router.post("/sessions/{session_id}/stream/sse")
|
||
|
|
async def stream_sse(session_id: str, body: dict[str, Any]) -> StreamingResponse:
|
||
|
|
"""Stream a prompt response as Server-Sent Events.
|
||
|
|
|
||
|
|
Events emitted:
|
||
|
|
- ``start``: Stream began
|
||
|
|
- ``heads_running``: Which heads are processing
|
||
|
|
- ``token``: Individual response token
|
||
|
|
- ``complete``: Final response with full text
|
||
|
|
- ``error``: Error occurred
|
||
|
|
"""
|
||
|
|
prompt = body.get("prompt", "")
|
||
|
|
return StreamingResponse(
|
||
|
|
_sse_generator(session_id, prompt),
|
||
|
|
media_type="text/event-stream",
|
||
|
|
headers={
|
||
|
|
"Cache-Control": "no-cache",
|
||
|
|
"Connection": "keep-alive",
|
||
|
|
"X-Accel-Buffering": "no",
|
||
|
|
},
|
||
|
|
)
|