import json import uuid from typing import Optional from fastapi import APIRouter, Depends, HTTPException, Query from fastapi.responses import StreamingResponse from pydantic import BaseModel, Field from sqlalchemy import text from sqlalchemy.ext.asyncio import create_async_engine from app.core.security import get_current_user, RateLimiter from app.rag.retriever import MixedRetriever, classify_intent from app.rag.reranker import Reranker from app.rag.prompt import build_prompt from app.core.llm_client import llm_client from app.core.config import get_settings settings = get_settings() router = APIRouter(prefix="/chat", tags=["对话"]) retriever = MixedRetriever() reranker = Reranker() class ChatRequest(BaseModel): message: str = Field(..., min_length=1, max_length=2000) conversation_id: str = Field(default_factory=lambda: str(uuid.uuid4())) class ChatResponse(BaseModel): answer: str sources: list[dict] conversation_id: str intent: str class FeedbackRequest(BaseModel): conversation_id: str message_id: int feedback: str # ============================================ # DB helpers # ============================================ async def _ensure_user(openid: str) -> int: engine = create_async_engine(settings.database_url) try: async with engine.begin() as conn: result = await conn.execute( text("SELECT id FROM users WHERE openid = :openid"), {"openid": openid}, ) row = result.fetchone() if row: return row[0] result = await conn.execute( text("INSERT INTO users (openid) VALUES (:openid) RETURNING id"), {"openid": openid}, ) return result.fetchone()[0] finally: await engine.dispose() async def _save_message(conversation_id: str, role: str, content: str, intent: str = None, sources: list = None): engine = create_async_engine(settings.database_url) try: async with engine.begin() as conn: # 确保 conversation 存在（不要求 user_id 外键，因为 dev token 的 user 可能不在库中） await conn.execute( text(""" INSERT INTO conversations (conversation_id, user_id, title) VALUES (:cid, 0, :title) ON CONFLICT (conversation_id) DO NOTHING """), { "cid": conversation_id, "title": content[:50] if role == "user" else "", }, ) await conn.execute( text(""" INSERT INTO messages (conversation_id, role, content, intent, sources) VALUES (:cid, :role, :content, :intent, :sources) """), { "cid": conversation_id, "role": role, "content": content, "intent": intent, "sources": json.dumps(sources) if sources else None, }, ) finally: await engine.dispose() # ============================================ # Chat endpoints # ============================================ @router.post("/ask", response_model=ChatResponse) async def chat_ask(req: ChatRequest, user: dict = Depends(get_current_user)): intent = classify_intent(req.message) documents = await retriever.search(req.message, intent=intent, top_k=20) documents = reranker.rerank(req.message, documents, top_k=5) msgs = build_prompt(req.message, documents, intent=intent) answer = await llm_client.chat(msgs) sources = [ {"name": d.get("drug_name", d.get("source", "")), "section": d.get("section", ""), "source": d.get("source", ""), "score": d.get("score", 0)} for d in documents ] # 保存到 DB await _save_message(req.conversation_id, "user", req.message, intent) await _save_message(req.conversation_id, "assistant", answer, intent, sources) return ChatResponse(answer=answer, sources=sources, conversation_id=req.conversation_id, intent=intent) @router.post("/stream") async def chat_stream(req: ChatRequest, user: dict = Depends(get_current_user)): async def stream_gen(): intent = classify_intent(req.message) yield f"event: intent\ndata: {intent}\n\n" yield "event: status\ndata: 正在检索...\n\n" documents = await retriever.search(req.message, intent=intent, top_k=20) documents = reranker.rerank(req.message, documents, top_k=5) yield f"event: status\ndata: 已匹配 {len(documents)} 条，生成中...\n\n" msgs = build_prompt(req.message, documents, intent=intent) sources = [ {"name": d.get("drug_name", d.get("source", "")), "section": d.get("section", ""), "source": d.get("source", ""), "score": d.get("score", 0)} for d in documents ] yield "event: content\n" full_answer = [] async for token in llm_client.chat_stream(msgs): full_answer.append(token) yield f"data: {token}\n\n" yield "data: [DONE]\n\n" # 元数据追加 import json yield f"event: meta\ndata: {json.dumps({'intent': intent, 'sources': sources, 'cid': req.conversation_id})}\n\n" answer_text = "".join(full_answer) await _save_message(req.conversation_id, "user", req.message, intent) await _save_message(req.conversation_id, "assistant", answer_text, intent, sources) return StreamingResponse(stream_gen(), media_type="text/event-stream") # ============================================ # 对话历史 API # ============================================ @router.get("/history") async def get_history( page: int = Query(1, ge=1), page_size: int = Query(20, ge=1, le=50), user: dict = Depends(get_current_user), ): engine = create_async_engine(settings.database_url) try: async with engine.connect() as conn: offset = (page - 1) * page_size result = await conn.execute( text(""" SELECT c.conversation_id, c.title, c.created_at, COUNT(m.id) as msg_count FROM conversations c LEFT JOIN messages m ON m.conversation_id = c.conversation_id GROUP BY c.id ORDER BY c.created_at DESC LIMIT :limit OFFSET :offset """), {"limit": page_size, "offset": offset}, ) items = [] for row in result.fetchall(): items.append({ "conversation_id": row[0], "title": row[1] or "新的对话", "created_at": row[2].isoformat() if row[2] else "", "message_count": row[3], }) return {"items": items, "page": page, "page_size": page_size} finally: await engine.dispose() @router.get("/history/{conversation_id}") async def get_conversation_detail( conversation_id: str, user: dict = Depends(get_current_user), ): engine = create_async_engine(settings.database_url) try: async with engine.connect() as conn: result = await conn.execute( text(""" SELECT id, role, content, intent, sources, created_at FROM messages WHERE conversation_id = :cid ORDER BY created_at ASC """), {"cid": conversation_id}, ) messages = [] for row in result.fetchall(): messages.append({ "id": row[0], "role": row[1], "content": row[2], "intent": row[3], "sources": row[4] if row[4] else [], "created_at": row[5].isoformat() if row[5] else "", }) return {"conversation_id": conversation_id, "messages": messages} finally: await engine.dispose() @router.post("/feedback") async def submit_feedback(req: FeedbackRequest, user: dict = Depends(get_current_user)): engine = create_async_engine(settings.database_url) try: async with engine.begin() as conn: await conn.execute( text("UPDATE messages SET feedback = :fb WHERE id = :mid"), {"fb": req.feedback, "mid": req.message_id}, ) return {"ok": True, "message_id": req.message_id, "feedback": req.feedback} finally: await engine.dispose() # ============================================ # 管理员：查看全部对话 # ============================================ @router.get("/admin/conversations") async def admin_list_conversations( page: int = Query(1, ge=1), page_size: int = Query(20, ge=1, le=100), keyword: Optional[str] = Query(None, description="搜索用户提问关键词"), user: dict = Depends(get_current_user), ): engine = create_async_engine(settings.database_url) try: async with engine.connect() as conn: offset = (page - 1) * page_size where = "" params = {"limit": page_size, "offset": offset} if keyword: where = "WHERE m.content LIKE :kw" params["kw"] = f"%{keyword}%" result = await conn.execute( text(f""" SELECT DISTINCT ON (c.conversation_id) c.conversation_id, c.title, c.created_at, m.content as last_msg, m.role FROM conversations c JOIN messages m ON m.conversation_id = c.conversation_id {where} ORDER BY c.conversation_id, m.created_at DESC LIMIT :limit OFFSET :offset """), params, ) items = [] for row in result.fetchall(): items.append({ "conversation_id": row[0], "title": row[1] or "新的对话", "created_at": row[2].isoformat() if row[2] else "", "last_message": (row[3] or "")[:200], "last_role": row[4], }) return {"items": items, "page": page, "page_size": page_size} finally: await engine.dispose()