|
|
@@ -0,0 +1,263 @@
|
|
|
+"""向量知识库管理接口 - 直接操作 ChromaDB"""
|
|
|
+from fastapi import APIRouter, HTTPException
|
|
|
+from pydantic import BaseModel
|
|
|
+from typing import Optional, List
|
|
|
+import logging
|
|
|
+import asyncio
|
|
|
+import os
|
|
|
+
|
|
|
+logger = logging.getLogger(__name__)
|
|
|
+
|
|
|
+router = APIRouter(prefix="/api/v1/knowledge-base", tags=["knowledge-base"])
|
|
|
+
|
|
|
+# 后台同步任务引用
|
|
|
+_sync_task = None
|
|
|
+
|
|
|
+
|
|
|
+class DocumentInfo(BaseModel):
|
|
|
+ id: str
|
|
|
+ source: str
|
|
|
+ title: str
|
|
|
+ content: str = ""
|
|
|
+ content_preview: str = ""
|
|
|
+ metadata: dict = {}
|
|
|
+
|
|
|
+
|
|
|
+class ListRequest(BaseModel):
|
|
|
+ source: Optional[str] = None
|
|
|
+ keyword: Optional[str] = None
|
|
|
+ page: int = 1
|
|
|
+ size: int = 20
|
|
|
+
|
|
|
+
|
|
|
+class UpdateDocumentRequest(BaseModel):
|
|
|
+ content: str
|
|
|
+ title: Optional[str] = None
|
|
|
+ metadata_override: Optional[dict] = None
|
|
|
+
|
|
|
+
|
|
|
+@router.post("/stats")
|
|
|
+async def get_stats():
|
|
|
+ """获取向量库统计信息"""
|
|
|
+ try:
|
|
|
+ from app.rag.retriever import RagRetriever
|
|
|
+ from app.rag.embeddings import get_embeddings
|
|
|
+ from app.config import settings
|
|
|
+ from langchain_chroma import Chroma
|
|
|
+
|
|
|
+ retriever = RagRetriever()
|
|
|
+ collection = retriever.vectorstore._collection
|
|
|
+ count = collection.count()
|
|
|
+
|
|
|
+ # 按 source 统计(兼容新旧两种 metadata 字段:source / type)
|
|
|
+ all_data = collection.get(include=["metadatas"])
|
|
|
+ metas = all_data.get("metadatas") or []
|
|
|
+ source_counts = {}
|
|
|
+ type_counts = {}
|
|
|
+ for meta in metas:
|
|
|
+ if not meta:
|
|
|
+ continue
|
|
|
+ src = meta.get("source") or meta.get("type") or "unknown"
|
|
|
+ source_counts[src] = source_counts.get(src, 0) + 1
|
|
|
+ doc_type = meta.get("type") or meta.get("source") or "unknown"
|
|
|
+ type_counts[doc_type] = type_counts.get(doc_type, 0) + 1
|
|
|
+
|
|
|
+ return {
|
|
|
+ "total": count,
|
|
|
+ "by_source": source_counts,
|
|
|
+ "by_type": type_counts,
|
|
|
+ "last_sync": _load_sync_state(),
|
|
|
+ }
|
|
|
+ except Exception as e:
|
|
|
+ logger.error("获取向量库统计失败: %s", e)
|
|
|
+ raise HTTPException(status_code=500, detail=str(e))
|
|
|
+
|
|
|
+
|
|
|
+@router.post("/documents")
|
|
|
+async def list_documents(req: ListRequest):
|
|
|
+ """分页列出向量库文档"""
|
|
|
+ try:
|
|
|
+ from app.rag.retriever import RagRetriever
|
|
|
+ from app.rag.embeddings import get_embeddings
|
|
|
+ from app.config import settings
|
|
|
+ from langchain_chroma import Chroma
|
|
|
+
|
|
|
+ retriever = RagRetriever()
|
|
|
+ collection = retriever.vectorstore._collection
|
|
|
+
|
|
|
+ # 构建过滤条件(兼容新旧两种 metadata 字段:source / type)
|
|
|
+ where_filter = None
|
|
|
+ if req.source:
|
|
|
+ where_filter = {"$or": [{"source": req.source}, {"type": req.source}]}
|
|
|
+
|
|
|
+ # 使用 get() 分页获取,避免加载全部 embedding
|
|
|
+ limit = req.size
|
|
|
+ offset = (req.page - 1) * req.size
|
|
|
+
|
|
|
+ result = collection.get(
|
|
|
+ where=where_filter,
|
|
|
+ include=["documents", "metadatas"],
|
|
|
+ limit=limit,
|
|
|
+ offset=offset,
|
|
|
+ )
|
|
|
+
|
|
|
+ ids = result.get("ids") or []
|
|
|
+ docs = result.get("documents") or []
|
|
|
+ metas = result.get("metadatas") or []
|
|
|
+
|
|
|
+ documents = []
|
|
|
+ for i, doc_id in enumerate(ids):
|
|
|
+ meta = metas[i] if i < len(metas) else {}
|
|
|
+ content = docs[i] if i < len(docs) else ""
|
|
|
+ # 截取内容预览
|
|
|
+ preview = content[:200] + "..." if len(content) > 200 else content
|
|
|
+ documents.append(DocumentInfo(
|
|
|
+ id=doc_id,
|
|
|
+ source=meta.get("source") or meta.get("type") or "unknown",
|
|
|
+ title=meta.get("title", ""),
|
|
|
+ content=content,
|
|
|
+ content_preview=preview,
|
|
|
+ metadata=meta,
|
|
|
+ ))
|
|
|
+
|
|
|
+ # 获取总数
|
|
|
+ count_result = collection.get(
|
|
|
+ where=where_filter,
|
|
|
+ include=[],
|
|
|
+ )
|
|
|
+ total = len(count_result.get("ids", []))
|
|
|
+
|
|
|
+ return {
|
|
|
+ "total": total,
|
|
|
+ "page": req.page,
|
|
|
+ "size": req.size,
|
|
|
+ "documents": documents,
|
|
|
+ }
|
|
|
+ except Exception as e:
|
|
|
+ logger.error("列出文档失败: %s", e)
|
|
|
+ raise HTTPException(status_code=500, detail=str(e))
|
|
|
+
|
|
|
+
|
|
|
+@router.delete("/source/{source_prefix}")
|
|
|
+async def delete_by_source(source_prefix: str):
|
|
|
+ """删除指定 source 前缀的所有文档"""
|
|
|
+ try:
|
|
|
+ from app.rag.retriever import RagRetriever
|
|
|
+ from app.rag.embeddings import get_embeddings
|
|
|
+ from app.config import settings
|
|
|
+ from langchain_chroma import Chroma
|
|
|
+
|
|
|
+ retriever = RagRetriever()
|
|
|
+ collection = retriever.vectorstore._collection
|
|
|
+
|
|
|
+ # 兼容新旧两种 metadata 字段:source / type
|
|
|
+ existing = collection.get(
|
|
|
+ where={"$or": [{"source": source_prefix}, {"type": source_prefix}]},
|
|
|
+ include=[],
|
|
|
+ )
|
|
|
+ ids = existing.get("ids", []) or []
|
|
|
+
|
|
|
+ if not ids:
|
|
|
+ return {"deleted": 0, "message": f"无 {source_prefix} 文档"}
|
|
|
+
|
|
|
+ collection.delete(ids=ids)
|
|
|
+ logger.info("已删除 source=%s 的 %d 个文档", source_prefix, len(ids))
|
|
|
+
|
|
|
+ return {"deleted": len(ids), "source": source_prefix}
|
|
|
+ except Exception as e:
|
|
|
+ logger.error("删除文档失败: %s", e)
|
|
|
+ raise HTTPException(status_code=500, detail=str(e))
|
|
|
+
|
|
|
+
|
|
|
+@router.post("/sync")
|
|
|
+async def trigger_sync():
|
|
|
+ """触发一次完整知识库同步"""
|
|
|
+ global _sync_task
|
|
|
+
|
|
|
+ try:
|
|
|
+ # 如果已有同步任务在运行,返回提示
|
|
|
+ if _sync_task and not _sync_task.done():
|
|
|
+ return {"status": "running", "message": "同步任务已在运行中"}
|
|
|
+
|
|
|
+ _sync_task = asyncio.create_task(_run_sync())
|
|
|
+ return {"status": "started", "message": "已启动知识库同步任务"}
|
|
|
+ except Exception as e:
|
|
|
+ logger.error("触发同步失败: %s", e)
|
|
|
+ raise HTTPException(status_code=500, detail=str(e))
|
|
|
+
|
|
|
+
|
|
|
+async def _run_sync():
|
|
|
+ """执行完整同步"""
|
|
|
+ try:
|
|
|
+ from app.tasks.knowledge_sync import sync_knowledge_base
|
|
|
+ await sync_knowledge_base()
|
|
|
+ logger.info("知识库同步完成")
|
|
|
+ except Exception as e:
|
|
|
+ logger.error("知识库同步失败: %s", e)
|
|
|
+
|
|
|
+
|
|
|
+@router.put("/document/{doc_id}")
|
|
|
+async def update_document(doc_id: str, req: UpdateDocumentRequest):
|
|
|
+ """更新单个文档内容(先获取原 metadata,再删除重建,保持 ID 不变)"""
|
|
|
+ try:
|
|
|
+ from app.rag.retriever import RagRetriever
|
|
|
+ from app.rag.embeddings import get_embeddings
|
|
|
+ from app.config import settings
|
|
|
+ from langchain_chroma import Chroma
|
|
|
+
|
|
|
+ retriever = RagRetriever()
|
|
|
+ collection = retriever.vectorstore._collection
|
|
|
+
|
|
|
+ # 先获取原文档 metadata,避免删除后丢失
|
|
|
+ old_result = collection.get(ids=[doc_id], include=["metadatas", "documents"])
|
|
|
+ old_metas = old_result.get("metadatas") or []
|
|
|
+ old_docs = old_result.get("documents") or []
|
|
|
+
|
|
|
+ if not old_metas:
|
|
|
+ raise HTTPException(status_code=404, detail="文档不存在")
|
|
|
+
|
|
|
+ original_meta = old_metas[0]
|
|
|
+ original_content = old_docs[0] if old_docs else ""
|
|
|
+
|
|
|
+ # 删除旧文档
|
|
|
+ collection.delete(ids=[doc_id])
|
|
|
+ logger.info("已删除旧文档: %s", doc_id)
|
|
|
+
|
|
|
+ # 构建新文档:保留原有 metadata,仅覆盖 content/title
|
|
|
+ meta = req.metadata_override or {}
|
|
|
+ # 保留原有 source/type/title 等字段,除非显式覆盖
|
|
|
+ for k, v in original_meta.items():
|
|
|
+ if k not in meta:
|
|
|
+ meta[k] = v
|
|
|
+
|
|
|
+ if req.title:
|
|
|
+ meta["title"] = req.title
|
|
|
+
|
|
|
+ # 使用原 ID 重新写入(显式传 ids,避免生成新 ID)
|
|
|
+ from langchain_core.documents import Document
|
|
|
+ doc = Document(page_content=req.content, metadata=meta)
|
|
|
+ await retriever.vectorstore.aadd_documents([doc], ids=[doc_id])
|
|
|
+ logger.info("已更新文档: %s", doc_id)
|
|
|
+
|
|
|
+ return {"success": True, "message": "文档已更新", "doc_id": doc_id}
|
|
|
+ except HTTPException:
|
|
|
+ raise
|
|
|
+ except Exception as e:
|
|
|
+ logger.error("更新文档失败: %s", e)
|
|
|
+ raise HTTPException(status_code=500, detail=str(e))
|
|
|
+
|
|
|
+
|
|
|
+def _load_sync_state() -> str:
|
|
|
+ """加载上次同步时间"""
|
|
|
+ try:
|
|
|
+ import json
|
|
|
+ from app.config import settings
|
|
|
+ import os
|
|
|
+ sync_file = os.path.join(settings.chroma_db_path, ".sync_state")
|
|
|
+ if os.path.exists(sync_file):
|
|
|
+ with open(sync_file) as f:
|
|
|
+ state = json.load(f)
|
|
|
+ return state.get("last_sync", "未知")
|
|
|
+ except Exception:
|
|
|
+ pass
|
|
|
+ return "未知"
|