law-kb/app/routers/rag.py

"""RAG 问答 API — /api/rag(SSE 流式)"""
import json
from fastapi import APIRouter, Request, HTTPException
from fastapi.responses import StreamingResponse
from typing import Optional

from app.services.retriever import Retriever
from app.services.embedding import embed_text
from app.services.llm import stream_chat
from app.services.prompt_config import get_template_text, get_system_prompt
from app.services.history import add_chat_history, update_chat_history
from app.config import RAG_TOP_K, RAG_SIMILARITY_THRESHOLD, MAX_QUERY_LENGTH
from app.models import RagRequest

router = APIRouter()


def _build_context(results: list) -> str:
    """将检索结果拼装为 context 文本"""
    lines = []
    for i, r in enumerate(results, 1):
        meta = f"《{r['law_name']}》"
        if r.get("clause_no"):
            meta += f" {r['clause_no']}"
        if r.get("chapter"):
            meta += f"({r['chapter']})"
        if r.get("province"):
            meta += f" [{r['province']}]"
        lines.append(f"{i}. {meta}\n{r['content']}")
    return "\n\n".join(lines)


@router.post("/rag")
async def rag(req: RagRequest, request: Request):
    """RAG 问答(SSE 流式)

    流程:向量化问题 → 检索 Top-10 → 检查相似度阈值 → 拼 prompt → 流式生成
    """
    trace_id = getattr(request.state, "trace_id", "")

    if len(req.question) > MAX_QUERY_LENGTH:
        raise HTTPException(status_code=400, detail="问题过长")

    # 向量化问题
    try:
        query_vec = await embed_text(req.question)
    except Exception as e:
        raise HTTPException(
            status_code=503,
            detail=f"embedding 服务不可用: {str(e)}",
        )

    # 检索
    retriever = Retriever.get_instance()
    if not retriever.is_ready():
        raise HTTPException(status_code=503, detail="索引未加载")

    results = retriever.search(
        query_vec=query_vec,
        top_k=RAG_TOP_K,
        category=req.category,
        province=req.province,
    )

    # 相似度阈值检查
    if not results or results[0]["score"] < RAG_SIMILARITY_THRESHOLD:
        async def no_result_stream():
            yield f"data: {json.dumps({'type': 'no_result', 'message': '未找到相关法规,无法回答'}, ensure_ascii=False)}\n\n"
        return StreamingResponse(no_result_stream(), media_type="text/event-stream")

    # 拼 prompt(优先级:请求自定义 > 数据库配置 > 默认)
    if req.template_text and "{context}" in req.template_text and "{question}" in req.template_text:
        template = req.template_text
    else:
        template = get_template_text(req.template)  # 从数据库读取
    context = _build_context(results)
    prompt = template.format(context=context, question=req.question)

    # 系统提示词(优先级:请求自定义 > 数据库配置 > 默认)
    system_prompt = req.system_prompt if req.system_prompt else get_system_prompt()

    # LLM 参数(从数据库读取)
    from app.services.prompt_config import get_llm_params
    llm_params = get_llm_params()

    # 引用条文(先发送)
    citations = [
        {
            "law_id": r["law_id"],
            "law_name": r["law_name"],
            "category": r["category"],
            "chapter": r.get("chapter"),
            "clause_no": r.get("clause_no"),
            "content": r["content"],
            "score": r["score"],
            "file_path": r["file_path"],
            "province": r.get("province"),
        }
        for r in results
    ]

    async def event_stream():
        # 先创建对话历史记录(status=streaming)
        chat_id = add_chat_history(
            question=req.question,
            template=req.template,
            category=req.category,
            province=req.province,
            thinking_enabled=llm_params["thinking_enabled"],
            status="streaming",
        )
        # 发送 chat_id 给前端(用于关联历史记录)
        yield f"data: {json.dumps({'type': 'chat_id', 'chat_id': chat_id}, ensure_ascii=False)}\n\n"
        # 发送引用条文
        yield f"data: {json.dumps({'type': 'citations', 'clauses': citations}, ensure_ascii=False)}\n\n"
        # 发送 LLM 参数(前端用于显示 thinking 状态等)
        yield f"data: {json.dumps({'type': 'params', 'thinking_enabled': llm_params['thinking_enabled']}, ensure_ascii=False)}\n\n"
        # 流式生成答案(区分 thinking 和 answer)
        thinking_acc = ""
        answer_acc = ""
        try:
            async for chunk_type, chunk_content in stream_chat(
                prompt,
                system=system_prompt,
                temperature=llm_params["temperature"],
                max_tokens=llm_params["max_tokens"],
                thinking_enabled=llm_params["thinking_enabled"],
                thinking_budget=llm_params["thinking_budget"],
            ):
                if chunk_type == "thinking":
                    thinking_acc += chunk_content
                elif chunk_type == "answer":
                    answer_acc += chunk_content
                yield f"data: {json.dumps({'type': chunk_type, 'content': chunk_content}, ensure_ascii=False)}\n\n"
            yield f"data: {json.dumps({'type': 'done'})}\n\n"
            # 流式完成,更新历史记录
            update_chat_history(
                chat_id=chat_id,
                thinking=thinking_acc,
                answer=answer_acc,
                citations=json.dumps(citations, ensure_ascii=False),
                status="done",
            )
        except Exception as e:
            yield f"data: {json.dumps({'type': 'error', 'message': str(e)}, ensure_ascii=False)}\n\n"
            # 记录错误状态
            update_chat_history(
                chat_id=chat_id,
                thinking=thinking_acc,
                answer=answer_acc,
                citations=json.dumps(citations, ensure_ascii=False),
                status="error",
            )

    return StreamingResponse(
        event_stream(),
        media_type="text/event-stream",
        headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
    )