刘金宝
225 lines
8.6 KiB
Python
225 lines
8.6 KiB
Python
from datetime import datetime
|
|
|
|
from sqlalchemy import delete, select
|
|
from sqlalchemy.orm import Session
|
|
|
|
from app.core.config import settings
|
|
from app.models.knowledge_base import (
|
|
KbKnowledgeChunk,
|
|
KbKnowledgeDocument,
|
|
KbKnowledgeIngestionTask,
|
|
KbKnowledgeQueryLog,
|
|
KbKnowledgeSpace,
|
|
)
|
|
|
|
|
|
class KnowledgeBaseRepository:
|
|
"""知识库仓储：集中管理 kb_* 表的查询和写入。"""
|
|
|
|
def __init__(self, db: Session) -> None:
|
|
self.db = db
|
|
|
|
def get_or_create_space(self, institution_id: int, institution_name: str | None) -> KbKnowledgeSpace:
|
|
"""知识空间：按机构和 embedding 版本获取或创建 Milvus collection 映射。"""
|
|
version = "v1"
|
|
space = self.db.scalar(
|
|
select(KbKnowledgeSpace).where(
|
|
KbKnowledgeSpace.institution_id == institution_id,
|
|
KbKnowledgeSpace.embedding_version == version,
|
|
)
|
|
)
|
|
if space:
|
|
return space
|
|
collection_name = self._collection_name(institution_id, version)
|
|
space = KbKnowledgeSpace(
|
|
institution_id=institution_id,
|
|
institution_name=institution_name,
|
|
space_code=f"institution_{institution_id}_{version}",
|
|
collection_name=collection_name,
|
|
embedding_model=settings.embedding_model,
|
|
embedding_dim=settings.embedding_dim,
|
|
embedding_version=version,
|
|
chunk_size=1100,
|
|
chunk_overlap=180,
|
|
top_k_default=settings.rag_top_k,
|
|
score_threshold=settings.rag_score_threshold,
|
|
status="active",
|
|
)
|
|
self.db.add(space)
|
|
self.db.flush()
|
|
return space
|
|
|
|
def get_space(self, institution_id: int) -> KbKnowledgeSpace | None:
|
|
"""知识空间查询：读取机构当前可用的知识库 collection 映射。"""
|
|
return self.db.scalar(
|
|
select(KbKnowledgeSpace)
|
|
.where(KbKnowledgeSpace.institution_id == institution_id, KbKnowledgeSpace.status == "active")
|
|
.order_by(KbKnowledgeSpace.id.desc())
|
|
)
|
|
|
|
def get_document_by_hash(self, institution_id: int, file_sha256: str) -> KbKnowledgeDocument | None:
|
|
"""文档去重：按机构和文件 SHA256 判断是否已上传。"""
|
|
return self.db.scalar(
|
|
select(KbKnowledgeDocument).where(
|
|
KbKnowledgeDocument.institution_id == institution_id,
|
|
KbKnowledgeDocument.file_sha256 == file_sha256,
|
|
)
|
|
)
|
|
|
|
def create_document(
|
|
self,
|
|
*,
|
|
institution_id: int,
|
|
uploaded_by: str,
|
|
file_name: str,
|
|
file_sha256: str,
|
|
file_size: int,
|
|
file_path: str,
|
|
document_title: str | None,
|
|
document_category: str,
|
|
version: str,
|
|
) -> KbKnowledgeDocument:
|
|
"""文档创建：保存内容管理员上传 PDF 的元数据和本地存储路径。"""
|
|
document = KbKnowledgeDocument(
|
|
institution_id=institution_id,
|
|
uploaded_by=uploaded_by,
|
|
file_name=file_name,
|
|
file_sha256=file_sha256,
|
|
file_type="pdf",
|
|
file_size=file_size,
|
|
file_path=file_path,
|
|
document_title=document_title,
|
|
document_category=document_category,
|
|
version=version,
|
|
status="uploaded",
|
|
parse_status="pending",
|
|
embedding_status="pending",
|
|
chunk_count=0,
|
|
)
|
|
self.db.add(document)
|
|
self.db.flush()
|
|
return document
|
|
|
|
def get_document(self, document_id: int, institution_id: int | None = None) -> KbKnowledgeDocument | None:
|
|
"""文档查询：按文档 ID 获取知识库文档，机构参数用于访问隔离。"""
|
|
stmt = select(KbKnowledgeDocument).where(KbKnowledgeDocument.id == document_id)
|
|
if institution_id is not None:
|
|
stmt = stmt.where(KbKnowledgeDocument.institution_id == institution_id)
|
|
return self.db.scalar(stmt)
|
|
|
|
def create_ingestion_task(self, document: KbKnowledgeDocument) -> KbKnowledgeIngestionTask:
|
|
"""入库任务：记录 PDF 解析、分片、向量化和写入 Milvus 的处理进度。"""
|
|
task = KbKnowledgeIngestionTask(
|
|
document_id=document.id,
|
|
institution_id=document.institution_id,
|
|
task_type="document_ingestion",
|
|
status="queued",
|
|
progress=0,
|
|
current_step="queued",
|
|
)
|
|
self.db.add(task)
|
|
self.db.flush()
|
|
return task
|
|
|
|
def get_ingestion_task(self, task_id: int) -> KbKnowledgeIngestionTask | None:
|
|
"""任务查询：按任务 ID 读取知识入库任务。"""
|
|
return self.db.get(KbKnowledgeIngestionTask, task_id)
|
|
|
|
def update_task(
|
|
self,
|
|
task: KbKnowledgeIngestionTask,
|
|
*,
|
|
status: str | None = None,
|
|
progress: int | None = None,
|
|
current_step: str | None = None,
|
|
error_message: str | None = None,
|
|
) -> None:
|
|
"""任务进度：更新入库任务状态，供前端或运维查看。"""
|
|
if status:
|
|
task.status = status
|
|
if status == "running" and task.started_at is None:
|
|
task.started_at = datetime.utcnow()
|
|
if status in {"success", "failed"}:
|
|
task.finished_at = datetime.utcnow()
|
|
if progress is not None:
|
|
task.progress = progress
|
|
if current_step is not None:
|
|
task.current_step = current_step
|
|
if error_message is not None:
|
|
task.error_message = error_message
|
|
|
|
def replace_chunks(self, document: KbKnowledgeDocument, chunks: list[KbKnowledgeChunk]) -> None:
|
|
"""分片替换：重新构建文档时先删除旧分片，再写入新分片。"""
|
|
self.db.execute(delete(KbKnowledgeChunk).where(KbKnowledgeChunk.document_id == document.id))
|
|
for chunk in chunks:
|
|
self.db.add(chunk)
|
|
document.chunk_count = len(chunks)
|
|
self.db.flush()
|
|
|
|
def get_chunks_by_uids(self, institution_id: int, chunk_uids: list[str]) -> list[KbKnowledgeChunk]:
|
|
"""分片查询：根据 Milvus 返回的 chunk_uid 批量读取 MySQL 分片详情。"""
|
|
if not chunk_uids:
|
|
return []
|
|
rows = self.db.scalars(
|
|
select(KbKnowledgeChunk).where(
|
|
KbKnowledgeChunk.institution_id == institution_id,
|
|
KbKnowledgeChunk.chunk_uid.in_(chunk_uids),
|
|
)
|
|
).all()
|
|
order = {chunk_uid: index for index, chunk_uid in enumerate(chunk_uids)}
|
|
return sorted(rows, key=lambda row: order.get(row.chunk_uid, 10_000))
|
|
|
|
def list_documents(self, institution_id: int, limit: int = 20) -> list[KbKnowledgeDocument]:
|
|
"""文档列表：返回机构最近上传的知识库文档。"""
|
|
return list(
|
|
self.db.scalars(
|
|
select(KbKnowledgeDocument)
|
|
.where(KbKnowledgeDocument.institution_id == institution_id)
|
|
.order_by(KbKnowledgeDocument.id.desc())
|
|
.limit(limit)
|
|
)
|
|
)
|
|
|
|
def create_query_log(
|
|
self,
|
|
*,
|
|
user_id: str,
|
|
institution_id: int,
|
|
question: str,
|
|
retrieval_hit: bool,
|
|
retrieved_chunk_ids: list[str],
|
|
answer_summary: str,
|
|
llm_model: str | None,
|
|
top_k: int,
|
|
score_threshold: float,
|
|
embedding_latency_ms: int | None,
|
|
search_latency_ms: int | None,
|
|
llm_latency_ms: int | None,
|
|
total_latency_ms: int | None,
|
|
) -> KbKnowledgeQueryLog:
|
|
"""查询日志：记录 RAG 命中、来源和耗时，支撑后续审计与效果分析。"""
|
|
log = KbKnowledgeQueryLog(
|
|
user_id=user_id,
|
|
institution_id=institution_id,
|
|
question=question,
|
|
retrieval_hit=bool(retrieval_hit),
|
|
retrieved_chunk_ids=retrieved_chunk_ids,
|
|
answer_summary=answer_summary[:1000],
|
|
llm_model=llm_model,
|
|
embedding_model=settings.embedding_model,
|
|
top_k=top_k,
|
|
score_threshold=score_threshold,
|
|
embedding_latency_ms=embedding_latency_ms,
|
|
search_latency_ms=search_latency_ms,
|
|
llm_latency_ms=llm_latency_ms,
|
|
total_latency_ms=total_latency_ms,
|
|
)
|
|
self.db.add(log)
|
|
self.db.flush()
|
|
return log
|
|
|
|
def _collection_name(self, institution_id: int, version: str) -> str:
|
|
"""集合命名：按机构隔离 Milvus collection，名称只使用安全字符。"""
|
|
model_part = "".join(ch if ch.isalnum() else "_" for ch in settings.embedding_model.lower())
|
|
return f"{settings.milvus_collection_prefix}_{institution_id}_{model_part}_{version}"[:120]
|