Samuele E. Locatelli
144 lines
5.2 KiB
Python
144 lines
5.2 KiB
Python
# api/v1/chat.py
|
|
import httpx
|
|
import json
|
|
import asyncio
|
|
from fastapi import APIRouter, HTTPException, Query
|
|
from fastapi.responses import StreamingResponse
|
|
from typing import List, Dict, Any, Optional
|
|
from models.chat import ChatRequest, ChatResponse
|
|
from services import redis_service # now using updated service with session support
|
|
from utils.logging import logger
|
|
from config import settings
|
|
|
|
router = APIRouter()
|
|
|
|
MAX_HISTORY_LENGTH = 50
|
|
|
|
|
|
@router.post("/chat", response_model=ChatResponse)
|
|
async def chat_endpoint(payload: ChatRequest):
|
|
try:
|
|
# Create a new session if session_id not provided
|
|
session_id = payload.session_id
|
|
if not session_id:
|
|
meta = redis_service.create_session(payload.user_id, payload.message)
|
|
session_id = meta["session_id"]
|
|
|
|
# Save user message
|
|
redis_service.save_chat(payload.user_id, session_id, {"role": "user", "content": payload.message})
|
|
history = redis_service.get_chat(payload.user_id, session_id, limit=MAX_HISTORY_LENGTH)
|
|
|
|
async with httpx.AsyncClient(timeout=settings.REQUEST_TIMEOUT) as client:
|
|
resp = await client.post(
|
|
settings.LM_STUDIO_URL,
|
|
json={"model": settings.MODEL_NAME, "messages": history},
|
|
)
|
|
resp.raise_for_status()
|
|
data = resp.json()
|
|
|
|
reply = data["choices"][0]["message"]["content"]
|
|
|
|
# Save assistant message
|
|
redis_service.save_chat(payload.user_id, session_id, {"role": "assistant", "content": reply})
|
|
|
|
# Return normal ChatResponse, but could also include session_id if needed
|
|
return ChatResponse(response=reply, session_id=session_id)
|
|
|
|
except Exception:
|
|
logger.exception("Error in /chat endpoint")
|
|
raise HTTPException(status_code=500, detail="Internal server error")
|
|
|
|
|
|
@router.post("/chat-stream")
|
|
async def chat_stream_endpoint(payload: ChatRequest):
|
|
"""
|
|
Streams model output token-by-token using SSE.
|
|
"""
|
|
session_id = payload.session_id
|
|
if not session_id:
|
|
meta = redis_service.create_session(payload.user_id, payload.message)
|
|
session_id = meta["session_id"]
|
|
|
|
redis_service.save_chat(payload.user_id, session_id, {"role": "user", "content": payload.message})
|
|
history = redis_service.get_chat(payload.user_id, session_id, limit=MAX_HISTORY_LENGTH)
|
|
|
|
async def event_generator():
|
|
assistant_text = ""
|
|
try:
|
|
async with httpx.AsyncClient(timeout=None) as client:
|
|
async with client.stream(
|
|
"POST",
|
|
settings.LM_STUDIO_URL,
|
|
json={
|
|
"model": settings.MODEL_NAME,
|
|
"messages": history,
|
|
"stream": True
|
|
}
|
|
) as r:
|
|
async for raw_line in r.aiter_lines():
|
|
if not raw_line:
|
|
continue
|
|
|
|
line = raw_line if raw_line.startswith("data:") else f"data: {raw_line}"
|
|
payload_str = line[len("data: "):].strip()
|
|
|
|
if payload_str == "[DONE]":
|
|
yield "data: [DONE]\n\n"
|
|
break
|
|
|
|
yield f"data: {payload_str}\n\n"
|
|
|
|
try:
|
|
obj = json.loads(payload_str)
|
|
choice = obj.get("choices", [{}])[0]
|
|
delta = choice.get("delta", {})
|
|
piece = delta.get("content") or choice.get("text")
|
|
if piece:
|
|
assistant_text += piece
|
|
except json.JSONDecodeError:
|
|
pass
|
|
|
|
await asyncio.sleep(0)
|
|
except Exception as e:
|
|
logger.exception("Streaming error in /chat-stream")
|
|
yield f"event: error\ndata: {str(e)}\n\n"
|
|
finally:
|
|
if assistant_text:
|
|
redis_service.save_chat(payload.user_id, session_id, {"role": "assistant", "content": assistant_text})
|
|
|
|
headers = {
|
|
"Cache-Control": "no-cache",
|
|
"Connection": "keep-alive",
|
|
"X-Accel-Buffering": "no"
|
|
}
|
|
return StreamingResponse(event_generator(), media_type="text/event-stream", headers=headers)
|
|
|
|
|
|
@router.get("/history")
|
|
async def get_history(
|
|
user_id: str = Query(..., description="User ID"),
|
|
session_id: str = Query(..., description="Session ID"),
|
|
limit: int = Query(MAX_HISTORY_LENGTH, description="Max number of messages to return")
|
|
) -> List[Dict[str, Any]]:
|
|
"""
|
|
Return all history saved for a given user/session.
|
|
"""
|
|
logger.info(f"[GET /history] user_id={user_id}, session_id={session_id}, limit={limit}")
|
|
history = redis_service.get_chat(user_id, session_id, limit=limit)
|
|
return history or []
|
|
|
|
|
|
@router.delete("/history")
|
|
async def delete_history(
|
|
user_id: str = Query(..., description="User ID"),
|
|
session_id: str = Query(..., description="Session ID")
|
|
):
|
|
"""
|
|
Clears history for a given user/session.
|
|
"""
|
|
logger.info(f"[DELETE /history] user_id={user_id}, session_id={session_id}")
|
|
redis_service.clear_chat(user_id, session_id)
|
|
return {"status": "cleared"}
|
|
|
|
|