backend/app/core/ai_stream.py

"""SSE Response Streaming — real-time token streaming for AI endpoints."""
import asyncio
import json
from fastapi import APIRouter
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from app.core.model_router import route_task, TaskType

router = APIRouter(prefix="/api/v1/ai", tags=["ai-streaming"])

class StreamRequest(BaseModel):
    prompt: str
    task: str = "fast"  # fast|cheap|complex|bulk|code
    prefer: str = "fast"  # fast|cheap

async def _stream_ollama(prompt: str, model: str):
    """Stream from Ollama."""
    import httpx
    async with httpx.AsyncClient(timeout=120) as c:
        async with c.stream("POST", "http://localhost:11434/api/generate", json={
            "model": model, "prompt": prompt
        }) as r:
            async for line in r.aiter_lines():
                if line:
                    try:
                        chunk = json.loads(line)
                        if chunk.get("done"):
                            yield f"data: {json.dumps({'done': True, 'model': model})}\n\n"
                            break
                        yield f"data: {json.dumps({'token': chunk.get('response', '')})}\n\n"
                    except json.JSONDecodeError:
                        continue

@router.post("/stream")
async def stream_ai(req: StreamRequest):
    """Stream AI response in real-time. Tokens appear as generated."""
    decision = route_task(TaskType(req.task), req.prefer)
    
    if decision.provider == "ollama":
        return StreamingResponse(
            _stream_ollama(req.prompt, decision.model),
            media_type="text/event-stream",
            headers={"X-Model": decision.model, "X-Provider": decision.provider, "X-Latency-Ms": str(decision.estimated_latency_ms)}
        )
    
    # For non-Ollama providers, fall back to blocking + stream result
    async def _blocking_stream():
        from app.core.model_router import smart_route
        result = await smart_route(req.prompt, req.task, req.prefer)
        if result and "response" in result:
            words = result["response"].split()
            for word in words:
                yield f"data: {json.dumps({'token': word + ' '})}\n\n"
                await asyncio.sleep(0.05)
        yield f"data: {json.dumps({'done': True, 'model': decision.model})}\n\n"
    
    return StreamingResponse(
        _blocking_stream(),
        media_type="text/event-stream",
        headers={"X-Model": decision.model, "X-Provider": decision.provider}
    )

@router.post("/route")
async def ai_route(req: StreamRequest):
    """Non-streaming: auto-route to best model, return complete response."""
    from app.core.model_router import smart_route
    result = await smart_route(req.prompt, req.task, req.prefer)
    return result if result else {"error": "All providers failed"}

@router.get("/providers")
async def list_providers():
    """List all available AI providers with capabilities."""
    from app.core.model_router import ROUTING_TABLE
    return {
        "providers": {
            task.value: [{"model": m[0], "provider": m[1], "latency_ms": m[2]*1000, "cost_per_1M_input": m[3]}
                        for m in models]
            for task, models in ROUTING_TABLE.items()
        }
    }