刘金宝
225 lines
9.8 KiB
Python
225 lines
9.8 KiB
Python
import json
|
|
import time
|
|
from collections.abc import AsyncIterator
|
|
from dataclasses import dataclass
|
|
|
|
from sqlalchemy.orm import Session
|
|
|
|
from app.agents.learning_assistant_agent import LearningAssistantAgent
|
|
from app.core.config import settings
|
|
from app.core.context import UserContext
|
|
from app.core.exceptions import AppError
|
|
from app.repositories.knowledge_base_repository import KnowledgeBaseRepository
|
|
from app.schemas.learning_assistant import LearningAssistantChatRequest, LearningAssistantChatResponse, LearningAssistantSource
|
|
from app.services.knowledge_space_service import KnowledgeSpaceService
|
|
from app.services.vector_search_service import RetrievedChunk, VectorSearchService
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class LearningAssistantRetrieval:
|
|
"""学习助手检索结果：封装知识库命中、耗时和降级原因。"""
|
|
|
|
institution_id: int | None
|
|
score_threshold: float
|
|
sources: list[LearningAssistantSource]
|
|
embedding_latency_ms: int | None = None
|
|
search_latency_ms: int | None = None
|
|
retrieval_error: str | None = None
|
|
|
|
|
|
class LearningAssistantService:
|
|
"""AI 学习助手服务：优先 RAG 检索，知识库不可用时降级为通用流式问答。"""
|
|
|
|
def __init__(
|
|
self,
|
|
db: Session,
|
|
*,
|
|
vector_search_service: VectorSearchService | None = None,
|
|
agent: LearningAssistantAgent | None = None,
|
|
) -> None:
|
|
self.db = db
|
|
self.repo = KnowledgeBaseRepository(db)
|
|
self.space_service = KnowledgeSpaceService(self.repo)
|
|
self.vector_search = vector_search_service or VectorSearchService(db)
|
|
self.agent = agent or LearningAssistantAgent()
|
|
|
|
async def chat(self, ctx: UserContext, payload: LearningAssistantChatRequest) -> LearningAssistantChatResponse:
|
|
"""知识问答调试：检索失败不阻断回答，返回完整文本和检索降级信息。"""
|
|
start = time.perf_counter()
|
|
retrieval = await self._retrieve_sources(ctx, payload)
|
|
llm_started = time.perf_counter()
|
|
response = await self.agent.answer(payload.question, retrieval.sources)
|
|
total_latency_ms = int((time.perf_counter() - start) * 1000)
|
|
llm_latency_ms = response.latency_ms or int((time.perf_counter() - llm_started) * 1000)
|
|
self._write_query_log(
|
|
ctx=ctx,
|
|
payload=payload,
|
|
retrieval=retrieval,
|
|
answer=response.content,
|
|
model=response.model,
|
|
llm_latency_ms=llm_latency_ms,
|
|
total_latency_ms=total_latency_ms,
|
|
)
|
|
return LearningAssistantChatResponse(
|
|
answer=response.content,
|
|
retrieval_hit=bool(retrieval.sources),
|
|
sources=retrieval.sources,
|
|
retrieval_error=retrieval.retrieval_error,
|
|
model=response.model,
|
|
embedding_latency_ms=retrieval.embedding_latency_ms,
|
|
search_latency_ms=retrieval.search_latency_ms,
|
|
llm_latency_ms=llm_latency_ms,
|
|
total_latency_ms=total_latency_ms,
|
|
)
|
|
|
|
async def stream_chat(self, ctx: UserContext, payload: LearningAssistantChatRequest) -> AsyncIterator[str]:
|
|
"""流式知识问答：先返回检索状态，再流式输出 LLM 回答。"""
|
|
start = time.perf_counter()
|
|
retrieval = await self._retrieve_sources(ctx, payload)
|
|
yield self._sse(
|
|
"retrieval_done",
|
|
{
|
|
"retrieval_hit": bool(retrieval.sources),
|
|
"sources": [source.model_dump() for source in retrieval.sources],
|
|
"retrieval_error": retrieval.retrieval_error,
|
|
"embedding_latency_ms": retrieval.embedding_latency_ms,
|
|
"search_latency_ms": retrieval.search_latency_ms,
|
|
},
|
|
)
|
|
|
|
answer_parts: list[str] = []
|
|
llm_latency_ms: int | None = None
|
|
model: str | None = None
|
|
try:
|
|
async for chunk in self.agent.stream_answer(payload.question, retrieval.sources):
|
|
if chunk.done:
|
|
llm_latency_ms = chunk.total_latency_ms
|
|
model = chunk.model
|
|
break
|
|
if chunk.delta:
|
|
answer_parts.append(chunk.delta)
|
|
yield self._sse("answer_delta", {"delta": chunk.delta})
|
|
except AppError as exc:
|
|
yield self._sse("error", {"code": exc.code, "message": exc.message})
|
|
return
|
|
except Exception:
|
|
yield self._sse("error", {"code": "LEARNING_ASSISTANT_LLM_FAILED", "message": "AI 学习助手回答生成失败，请稍后重试"})
|
|
return
|
|
|
|
answer = "".join(answer_parts)
|
|
total_latency_ms = int((time.perf_counter() - start) * 1000)
|
|
self._write_query_log(
|
|
ctx=ctx,
|
|
payload=payload,
|
|
retrieval=retrieval,
|
|
answer=answer,
|
|
model=model,
|
|
llm_latency_ms=llm_latency_ms,
|
|
total_latency_ms=total_latency_ms,
|
|
commit=True,
|
|
)
|
|
yield self._sse("answer_done", {"model": model, "total_latency_ms": total_latency_ms})
|
|
|
|
async def _retrieve_sources(self, ctx: UserContext, payload: LearningAssistantChatRequest) -> LearningAssistantRetrieval:
|
|
"""知识检索：按机构读取知识空间；无空间、Milvus 或 embedding 异常时降级为空来源。"""
|
|
score_threshold = payload.score_threshold if payload.score_threshold is not None else settings.rag_score_threshold
|
|
try:
|
|
institution_id = self.space_service.require_institution_id(ctx)
|
|
except AppError:
|
|
return LearningAssistantRetrieval(
|
|
institution_id=None,
|
|
score_threshold=score_threshold,
|
|
sources=[],
|
|
retrieval_error="当前用户缺少机构信息，已转为大模型通用学习回答。",
|
|
)
|
|
|
|
try:
|
|
space = self.space_service.get_active_space(ctx)
|
|
retrieval = await self.vector_search.search(
|
|
institution_id=space.institution_id,
|
|
collection_name=space.collection_name,
|
|
question=payload.question,
|
|
top_k=payload.top_k,
|
|
score_threshold=payload.score_threshold,
|
|
)
|
|
return LearningAssistantRetrieval(
|
|
institution_id=space.institution_id,
|
|
score_threshold=payload.score_threshold if payload.score_threshold is not None else space.score_threshold,
|
|
sources=self._build_sources(retrieval.chunks),
|
|
embedding_latency_ms=retrieval.embedding_latency_ms,
|
|
search_latency_ms=retrieval.search_latency_ms,
|
|
)
|
|
except AppError as exc:
|
|
if exc.code in {"KNOWLEDGE_SPACE_NOT_FOUND", "MILVUS_COLLECTION_NOT_FOUND", "EMBEDDING_CALL_FAILED"}:
|
|
return LearningAssistantRetrieval(
|
|
institution_id=institution_id,
|
|
score_threshold=score_threshold,
|
|
sources=[],
|
|
retrieval_error="当前机构知识库暂未初始化或检索不可用，已转为大模型通用学习回答。",
|
|
)
|
|
raise
|
|
except Exception:
|
|
return LearningAssistantRetrieval(
|
|
institution_id=institution_id,
|
|
score_threshold=score_threshold,
|
|
sources=[],
|
|
retrieval_error="当前机构知识库检索暂不可用，已转为大模型通用学习回答。",
|
|
)
|
|
|
|
def _write_query_log(
|
|
self,
|
|
*,
|
|
ctx: UserContext,
|
|
payload: LearningAssistantChatRequest,
|
|
retrieval: LearningAssistantRetrieval,
|
|
answer: str,
|
|
model: str | None,
|
|
llm_latency_ms: int | None,
|
|
total_latency_ms: int | None,
|
|
commit: bool = False,
|
|
) -> None:
|
|
"""查询日志：仅在存在机构 ID 时记录 RAG 命中、来源和耗时。"""
|
|
if retrieval.institution_id is None:
|
|
return
|
|
self.repo.create_query_log(
|
|
user_id=ctx.user_id,
|
|
institution_id=retrieval.institution_id,
|
|
question=payload.question,
|
|
retrieval_hit=bool(retrieval.sources),
|
|
retrieved_chunk_ids=[source.chunk_uid for source in retrieval.sources],
|
|
answer_summary=answer,
|
|
llm_model=model,
|
|
top_k=payload.top_k or len(retrieval.sources) or settings.rag_top_k,
|
|
score_threshold=retrieval.score_threshold,
|
|
embedding_latency_ms=retrieval.embedding_latency_ms,
|
|
search_latency_ms=retrieval.search_latency_ms,
|
|
llm_latency_ms=llm_latency_ms,
|
|
total_latency_ms=total_latency_ms,
|
|
)
|
|
if commit:
|
|
self.db.commit()
|
|
|
|
def _build_sources(self, chunks: list[RetrievedChunk]) -> list[LearningAssistantSource]:
|
|
"""来源构建：把检索分片转换为前端可展示的 PDF 来源结构。"""
|
|
sources: list[LearningAssistantSource] = []
|
|
for item in chunks:
|
|
document = self.repo.get_document(item.chunk.document_id, item.chunk.institution_id)
|
|
quote = item.chunk.chunk_text[:500]
|
|
sources.append(
|
|
LearningAssistantSource(
|
|
document_id=item.chunk.document_id,
|
|
document_title=document.document_title if document else None,
|
|
file_name=document.file_name if document else "",
|
|
page_start=item.chunk.page_start,
|
|
page_end=item.chunk.page_end,
|
|
chunk_uid=item.chunk.chunk_uid,
|
|
score=round(item.score, 4),
|
|
quote=quote,
|
|
)
|
|
)
|
|
return sources
|
|
|
|
def _sse(self, event: str, data: dict) -> str:
|
|
"""SSE 封装：统一输出 event + data 格式。"""
|
|
return f"event: {event}\ndata: {json.dumps(data, ensure_ascii=False)}\n\n"
|