File size: 8,080 Bytes
027123c 6bff5d9 027123c 6bff5d9 027123c 6bff5d9 027123c 6bff5d9 027123c 6bff5d9 027123c 6bff5d9 027123c 6bff5d9 027123c 61c746f 027123c 61c746f 027123c 61c746f 027123c 61c746f 027123c 6bff5d9 027123c 6bff5d9 027123c 61c746f 027123c 61c746f 027123c 61c746f 027123c 6bff5d9 61c746f 6bff5d9 027123c 6bff5d9 027123c 6bff5d9 027123c 61c746f 027123c 6bff5d9 61c746f 6bff5d9 027123c | 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 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 | """Chat endpoint with streaming support."""
import uuid
import json
from typing import List, Dict, Any, Optional
from fastapi import APIRouter, Depends, HTTPException
from langchain_core.messages import HumanMessage, AIMessage
from pydantic import BaseModel
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from sse_starlette.sse import EventSourceResponse
from src.agents.chat_handler import ChatHandler
from src.config.settings import settings
from src.db.postgres.connection import get_db
from src.db.postgres.models import ChatMessage, MessageSource
from src.db.redis.connection import get_redis
from src.middlewares.logging import get_logger, log_execution
logger = get_logger("chat_api")
router = APIRouter(prefix="/api/v1", tags=["Chat"])
_GREETINGS = frozenset(["hi", "hello", "hey", "halo", "hai", "hei"])
_GOODBYES = frozenset(["bye", "goodbye", "thanks", "thank you", "terima kasih", "sampai jumpa"])
def _fast_intent(message: str) -> Optional[str]:
"""Return a direct response for obvious greetings/farewells, else None."""
lower = message.lower().strip().rstrip("!.,?")
if lower in _GREETINGS:
return "Hello! How can I assist you today?"
if lower in _GOODBYES:
return "Goodbye! Have a great day!"
return None
class ChatRequest(BaseModel):
user_id: str
room_id: str
message: str
async def get_cached_response(redis, cache_key: str) -> Optional[dict]:
cached = await redis.get(cache_key)
if cached:
data = json.loads(cached)
if isinstance(data, dict) and "response" in data:
return data
# legacy: plain string cached before this change
return {"response": data, "sources": []}
return None
async def cache_response(redis, cache_key: str, response: str, sources: list):
await redis.setex(cache_key, 86400, json.dumps({"response": response, "sources": sources}))
async def load_history(db: AsyncSession, room_id: str, limit: int = 10) -> list:
"""Load recent chat messages for a room as LangChain message objects (oldest-first)."""
result = await db.execute(
select(ChatMessage)
.where(ChatMessage.room_id == room_id)
.order_by(ChatMessage.created_at.asc())
.limit(limit)
)
rows = result.scalars().all()
return [
HumanMessage(content=row.content) if row.role == "user" else AIMessage(content=row.content)
for row in rows
]
async def save_messages(
db: AsyncSession,
room_id: str,
user_content: str,
assistant_content: str,
sources: Optional[List[Dict[str, Any]]] = None,
):
"""Persist user and assistant messages, and attach sources to the assistant message."""
db.add(ChatMessage(id=str(uuid.uuid4()), room_id=room_id, role="user", content=user_content))
assistant_id = str(uuid.uuid4())
db.add(ChatMessage(id=assistant_id, room_id=room_id, role="assistant", content=assistant_content))
for src in (sources or []):
page = src.get("page_label")
db.add(MessageSource(
id=str(uuid.uuid4()),
message_id=assistant_id,
document_id=src.get("document_id"),
filename=src.get("filename"),
page_label=str(page) if page is not None else None,
))
await db.commit()
@router.delete("/chat/cache")
async def clear_chat_cache(room_id: str, message: str):
"""Delete the Redis cache entry for a specific room + message pair."""
redis = await get_redis()
cache_key = f"{settings.redis_prefix}chat:{room_id}:{message}"
deleted = await redis.delete(cache_key)
return {"deleted": deleted > 0, "cache_key": cache_key}
@router.delete("/chat/cache/room/{room_id}")
async def clear_room_cache(room_id: str):
"""Delete all Redis cache entries for a room."""
redis = await get_redis()
pattern = f"{settings.redis_prefix}chat:{room_id}:*"
keys = await redis.keys(pattern)
if keys:
await redis.delete(*keys)
return {"deleted_count": len(keys), "room_id": room_id}
@router.delete("/retrieval/cache/{user_id}")
async def clear_retrieval_cache(user_id: str):
"""Delete all cached retrieval results for a user. Call this after uploading/processing new documents."""
from src.retrieval.router import retrieval_router
deleted = await retrieval_router.invalidate_cache(user_id)
return {"deleted_count": deleted, "user_id": user_id}
@router.post("/chat/stream")
@log_execution(logger)
async def chat_stream(request: ChatRequest, db: AsyncSession = Depends(get_db)):
"""Chat endpoint with streaming response.
SSE event sequence:
1. sources — JSON array of source refs from ChatHandler (table for
structured; deduped document_id/page_label for unstructured)
2. chunk — text fragments of the answer
3. done — signals end of stream
"""
redis = await get_redis()
cache_key = f"{settings.redis_prefix}chat:{request.room_id}:{request.message}"
# Redis cache hit
cached = await get_cached_response(redis, cache_key)
logger.info("cache check", cache_key=cache_key, cache_hit=cached is not None)
if cached:
logger.info("Returning cached response")
cached_text = cached["response"]
cached_sources = cached["sources"]
await save_messages(db, request.room_id, request.message, cached_text, sources=cached_sources)
async def stream_cached():
yield {"event": "sources", "data": json.dumps(cached_sources)}
for i in range(0, len(cached_text), 50):
yield {"event": "chunk", "data": cached_text[i:i + 50]}
yield {"event": "done", "data": ""}
return EventSourceResponse(stream_cached())
try:
# Fast intent: greetings/farewells bypass LLM entirely
direct = _fast_intent(request.message)
if direct:
await cache_response(redis, cache_key, direct, sources=[])
await save_messages(db, request.room_id, request.message, direct, sources=[])
async def stream_direct():
yield {"event": "sources", "data": json.dumps([])}
yield {"event": "chunk", "data": direct}
yield {"event": "done", "data": ""}
return EventSourceResponse(stream_direct())
history = await load_history(db, request.room_id, limit=10)
handler = ChatHandler()
async def stream_response():
logger.info("stream_response started", room_id=request.room_id, user_id=request.user_id)
full_response = ""
sources: List[Dict[str, Any]] = []
async for event in handler.handle(request.message, request.user_id, history):
if event["event"] == "sources":
try:
sources = json.loads(event["data"]) or []
except (TypeError, ValueError):
sources = []
yield event
elif event["event"] == "chunk":
full_response += event["data"]
yield event
elif event["event"] == "done":
await cache_response(redis, cache_key, full_response, sources=sources)
logger.info("saving messages", sources_count=len(sources), sources=sources)
try:
await save_messages(db, request.room_id, request.message, full_response, sources=sources)
except Exception as e:
logger.error("save_messages failed", room_id=request.room_id, error=str(e))
yield event
elif event["event"] == "error":
yield event
return
# "intent" event: consumed internally, not forwarded to frontend
return EventSourceResponse(stream_response())
except Exception as e:
logger.error("Chat failed", error=str(e))
raise HTTPException(status_code=500, detail=f"Chat failed: {str(e)}")
|