Spaces:
Sleeping
Sleeping
File size: 5,888 Bytes
3a7eb07 e86dfae 3a7eb07 e86dfae 3a7eb07 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 | """
Streaming Chat Endpoint (SSE)
Server-Sent Events streaming for real-time AI responses.
Delivers three event types to the client:
1. 'sources' — document citations (emitted first so UI renders immediately)
2. 'token' — streamed LLM response tokens
3. 'done' — signals stream completion with metadata
4. 'error' — on failure
"""
import json
import logging
from datetime import datetime
from typing import List, Optional
from fastapi import APIRouter, Depends, Query
from fastapi.responses import StreamingResponse
from sqlalchemy.orm import Session
from app.api.deps import get_current_user_id
from app.config import settings
from app.db.session import get_db
from app.models.chat import ChatMessage, ChatSession
from app.schemas.chat import ChatRequest
logger = logging.getLogger(__name__)
router = APIRouter()
@router.post("/stream")
async def stream_chat(
request: ChatRequest,
user_id: str = Depends(get_current_user_id),
db: Session = Depends(get_db),
):
"""
Stream AI response using Server-Sent Events.
Event format:
event: sources
data: [{"document_title": ..., "file_name": ..., "page_numbers": [...], ...}]
event: token
data: {"text": "word by word..."}
event: done
data: {"session_id": "...", "sources_count": 3}
Usage with fetch():
const source = new EventSource('/api/v1/chat/stream', {...})
source.addEventListener('token', (e) => appendToken(JSON.parse(e.data).text))
source.addEventListener('sources', (e) => renderSources(JSON.parse(e.data)))
source.addEventListener('done', () => source.close())
"""
from app.services.chat_service import ChatService
from app.services.guardrails import MiningGuardrails
# ── Input Guardrails ────────────────────────────────────────────────────
validated_query = MiningGuardrails.validate_input(request.content)
# Get or create session
if request.session_id:
session = (
db.query(ChatSession)
.filter(
ChatSession.id == request.session_id, ChatSession.user_id == user_id
)
.first()
)
if not session:
from app.core.exceptions import NotFoundError
raise NotFoundError("Chat session", request.session_id)
else:
session = ChatSession(user_id=user_id, title="New Chat")
db.add(session)
db.commit()
db.refresh(session)
# Save user message immediately
user_message = ChatMessage(
session_id=session.id,
role="user",
content=request.content,
)
db.add(user_message)
db.commit()
chat_service = ChatService()
async def event_generator():
full_response = []
sources = []
tokens_used = None
try:
async for event in chat_service.generate_response_stream(
query=validated_query,
user_id=user_id,
document_ids=request.document_ids,
db=db,
):
# Forward each SSE event to client
yield event
# Track sources from the sources event for DB persistence
if event.startswith("event: sources\n"):
data_line = event.split("data: ", 1)[-1].strip()
try:
sources = json.loads(data_line)
except Exception:
pass
# Accumulate tokens for DB persistence
elif event.startswith("event: token\n"):
data_line = event.split("data: ", 1)[-1].strip()
try:
token_data = json.loads(data_line)
full_response.append(token_data.get("text", ""))
except Exception:
pass
elif event.startswith("event: done\n"):
data_line = event.split("data: ", 1)[-1].strip()
try:
done_data = json.loads(data_line)
tokens_used = done_data.get("tokens_used")
except Exception:
pass
except Exception as e:
logger.error(f"Streaming error: {e}", exc_info=True)
yield f"event: error\ndata: {json.dumps({'message': str(e)})}\n\n"
finally:
# Persist assistant message after stream completes
if full_response:
response_text = "".join(full_response)
assistant_message = ChatMessage(
session_id=session.id,
role="assistant",
content=response_text,
sources=sources if request.include_sources else [],
model_used=settings.GEMINI_MODEL,
tokens_used=tokens_used,
)
db.add(assistant_message)
# Auto-title on first message
msg_count = (
db.query(ChatMessage)
.filter(ChatMessage.session_id == session.id)
.count()
)
if msg_count <= 2:
session.title = request.content[:50] + (
"..." if len(request.content) > 50 else ""
)
session.updated_at = datetime.utcnow()
db.commit()
return StreamingResponse(
event_generator(),
media_type="text/event-stream",
headers={
"Cache-Control": "no-cache",
"Connection": "keep-alive",
"X-Accel-Buffering": "no", # Disable nginx buffering
},
)
|