""" RAG Lab endpoints for debugging and experimentation. [AC-ASA-05, AC-ASA-19, AC-ASA-20, AC-ASA-21, AC-ASA-22] RAG experiment with AI output. """ import json import logging import time from typing import Annotated, Any, List from fastapi import APIRouter, Depends, Body from fastapi.responses import JSONResponse, StreamingResponse from pydantic import BaseModel, Field from app.core.config import get_settings from app.core.exceptions import MissingTenantIdException from app.core.tenant import get_tenant_id from app.models import ErrorResponse from app.services.retrieval.vector_retriever import get_vector_retriever from app.services.retrieval.base import RetrievalContext from app.services.llm.factory import get_llm_config_manager logger = logging.getLogger(__name__) router = APIRouter(prefix="/admin/rag", tags=["RAG Lab"]) def get_current_tenant_id() -> str: """Dependency to get current tenant ID or raise exception.""" tenant_id = get_tenant_id() if not tenant_id: raise MissingTenantIdException() return tenant_id class RAGExperimentRequest(BaseModel): query: str = Field(..., description="Query text for retrieval") kb_ids: List[str] | None = Field(default=None, description="Knowledge base IDs to search") top_k: int = Field(default=5, description="Number of results to retrieve") score_threshold: float = Field(default=0.5, description="Minimum similarity score") generate_response: bool = Field(default=True, description="Whether to generate AI response") llm_provider: str | None = Field(default=None, description="Specific LLM provider to use") class AIResponse(BaseModel): content: str prompt_tokens: int = 0 completion_tokens: int = 0 total_tokens: int = 0 latency_ms: float = 0 model: str = "" class RAGExperimentResult(BaseModel): query: str retrieval_results: List[dict] = [] final_prompt: str = "" ai_response: AIResponse | None = None total_latency_ms: float = 0 diagnostics: dict[str, Any] = {} @router.post( "/experiments/run", operation_id="runRagExperiment", summary="Run RAG debugging experiment with AI output", description="[AC-ASA-05, AC-ASA-19, AC-ASA-21, AC-ASA-22] Trigger RAG experiment with retrieval, prompt generation, and AI response.", responses={ 200: {"description": "Experiment results with retrieval, prompt, and AI response"}, 401: {"description": "Unauthorized", "model": ErrorResponse}, 403: {"description": "Forbidden", "model": ErrorResponse}, }, ) async def run_rag_experiment( tenant_id: Annotated[str, Depends(get_current_tenant_id)], request: RAGExperimentRequest = Body(...), ) -> JSONResponse: """ [AC-ASA-05, AC-ASA-19, AC-ASA-21, AC-ASA-22] Run RAG experiment and return retrieval results with AI response. """ start_time = time.time() logger.info( f"[AC-ASA-05] Running RAG experiment: tenant={tenant_id}, " f"query={request.query[:50]}..., kb_ids={request.kb_ids}, " f"generate_response={request.generate_response}" ) settings = get_settings() top_k = request.top_k or settings.rag_top_k threshold = request.score_threshold or settings.rag_score_threshold try: retriever = await get_vector_retriever() retrieval_ctx = RetrievalContext( tenant_id=tenant_id, query=request.query, session_id="rag_experiment", channel_type="admin", metadata={"kb_ids": request.kb_ids}, ) result = await retriever.retrieve(retrieval_ctx) retrieval_results = [ { "content": hit.text, "score": hit.score, "source": hit.source, "metadata": hit.metadata, } for hit in result.hits ] final_prompt = _build_final_prompt(request.query, retrieval_results) logger.info( f"[AC-ASA-05] RAG retrieval complete: hits={len(retrieval_results)}, " f"max_score={result.max_score:.3f}" ) ai_response = None if request.generate_response: ai_response = await _generate_ai_response( final_prompt, provider=request.llm_provider, ) total_latency_ms = (time.time() - start_time) * 1000 return JSONResponse( content={ "query": request.query, "retrieval_results": retrieval_results, "final_prompt": final_prompt, "ai_response": ai_response.model_dump() if ai_response else None, "total_latency_ms": round(total_latency_ms, 2), "diagnostics": result.diagnostics, } ) except Exception as e: logger.error(f"[AC-ASA-05] RAG experiment failed: {e}") fallback_results = _get_fallback_results(request.query) fallback_prompt = _build_final_prompt(request.query, fallback_results) ai_response = None if request.generate_response: ai_response = await _generate_ai_response( fallback_prompt, provider=request.llm_provider, ) total_latency_ms = (time.time() - start_time) * 1000 return JSONResponse( content={ "query": request.query, "retrieval_results": fallback_results, "final_prompt": fallback_prompt, "ai_response": ai_response.model_dump() if ai_response else None, "total_latency_ms": round(total_latency_ms, 2), "diagnostics": { "error": str(e), "fallback": True, }, } ) @router.post( "/experiments/stream", operation_id="runRagExperimentStream", summary="Run RAG experiment with streaming AI output", description="[AC-ASA-20] Trigger RAG experiment with SSE streaming for AI response.", responses={ 200: {"description": "SSE stream with retrieval results and AI response"}, 401: {"description": "Unauthorized", "model": ErrorResponse}, 403: {"description": "Forbidden", "model": ErrorResponse}, }, ) async def run_rag_experiment_stream( tenant_id: Annotated[str, Depends(get_current_tenant_id)], request: RAGExperimentRequest = Body(...), ) -> StreamingResponse: """ [AC-ASA-20] Run RAG experiment with SSE streaming for AI response. """ logger.info( f"[AC-ASA-20] Running RAG experiment stream: tenant={tenant_id}, " f"query={request.query[:50]}..." ) settings = get_settings() top_k = request.top_k or settings.rag_top_k threshold = request.score_threshold or settings.rag_score_threshold async def event_generator(): try: retriever = await get_vector_retriever() retrieval_ctx = RetrievalContext( tenant_id=tenant_id, query=request.query, session_id="rag_experiment_stream", channel_type="admin", metadata={"kb_ids": request.kb_ids}, ) result = await retriever.retrieve(retrieval_ctx) retrieval_results = [ { "content": hit.text, "score": hit.score, "source": hit.source, "metadata": hit.metadata, } for hit in result.hits ] final_prompt = _build_final_prompt(request.query, retrieval_results) yield f"event: retrieval\ndata: {json.dumps({'results': retrieval_results, 'count': len(retrieval_results)})}\n\n" yield f"event: prompt\ndata: {json.dumps({'prompt': final_prompt})}\n\n" if request.generate_response: manager = get_llm_config_manager() client = manager.get_client() full_content = "" async for chunk in client.stream_generate( messages=[{"role": "user", "content": final_prompt}], ): if chunk.delta: full_content += chunk.delta yield f"event: message\ndata: {json.dumps({'delta': chunk.delta})}\n\n" yield f"event: final\ndata: {json.dumps({'content': full_content, 'finish_reason': 'stop'})}\n\n" else: yield f"event: final\ndata: {json.dumps({'content': '', 'finish_reason': 'skipped'})}\n\n" except Exception as e: logger.error(f"[AC-ASA-20] RAG experiment stream failed: {e}") yield f"event: error\ndata: {json.dumps({'error': str(e)})}\n\n" return StreamingResponse( event_generator(), media_type="text/event-stream", headers={ "Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no", }, ) async def _generate_ai_response( prompt: str, provider: str | None = None, ) -> AIResponse | None: """ [AC-ASA-19, AC-ASA-21] Generate AI response from prompt. """ import time try: manager = get_llm_config_manager() client = manager.get_client() start_time = time.time() response = await client.generate( messages=[{"role": "user", "content": prompt}], ) latency_ms = (time.time() - start_time) * 1000 return AIResponse( content=response.content, prompt_tokens=response.usage.get("prompt_tokens", 0), completion_tokens=response.usage.get("completion_tokens", 0), total_tokens=response.usage.get("total_tokens", 0), latency_ms=round(latency_ms, 2), model=response.model, ) except Exception as e: logger.error(f"[AC-ASA-19] AI response generation failed: {e}") return AIResponse( content=f"AI 响应生成失败: {str(e)}", latency_ms=0, ) def _build_final_prompt(query: str, retrieval_results: list[dict]) -> str: """ Build the final prompt from query and retrieval results. """ if not retrieval_results: return f"""用户问题：{query} 未找到相关检索结果，请基于通用知识回答用户问题。""" evidence_text = "\n".join([ f"{i+1}. [Score: {hit['score']:.2f}] {hit['content'][:200]}{'...' if len(hit['content']) > 200 else ''}" for i, hit in enumerate(retrieval_results[:5]) ]) return f"""基于以下检索到的信息，作为一个回答简洁精准的客服，回答用户问题： 用户问题：{query} 检索结果： {evidence_text} 请基于以上信息生成专业、准确的回答，注意输出内容应该格式整齐，不包含json符号等。""" def _get_fallback_results(query: str) -> list[dict]: """ Provide fallback results when retrieval fails. """ return [ { "content": "检索服务暂时不可用，这是模拟结果。", "score": 0.5, "source": "fallback", } ]