init: RAG 知识库服务初始提交
- 后端 API（Flask + Gunicorn） - RAG 引擎（混合检索 + 云端 Reranker + 引用溯源） - 文档解析（MinerU + 多格式支持） - Docker 生产部署配置 - 排除前端项目、敏感配置、模型文件
This commit is contained in:
lacerate551
266
knowledge/document.py
Normal file
266
knowledge/document.py
Normal file
@@ -0,0 +1,266 @@
|
||||
"""
|
||||
知识库管理器 - 文档管理 Mixin
|
||||
|
||||
包含文档级别的管理方法
|
||||
"""
|
||||
|
||||
import os
|
||||
import logging
|
||||
from datetime import datetime
|
||||
from typing import List, Dict, Optional
|
||||
|
||||
from .base import _get_doc_type
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class DocumentMixin:
|
||||
"""文档管理方法"""
|
||||
|
||||
def get_document_count(self, kb_name: str) -> int:
|
||||
"""获取向量库中的文档数量"""
|
||||
collection = self.get_collection(kb_name)
|
||||
return collection.count() if collection else 0
|
||||
|
||||
def list_documents(self, kb_name: str) -> List[dict]:
|
||||
"""列出向量库中的文档"""
|
||||
collection = self.get_collection(kb_name)
|
||||
if not collection:
|
||||
return []
|
||||
|
||||
result = collection.get()
|
||||
|
||||
from collections import Counter
|
||||
file_chunks = Counter()
|
||||
|
||||
for meta in result.get('metadatas', []):
|
||||
source = meta.get('source', 'unknown')
|
||||
file_chunks[source] += 1
|
||||
|
||||
return [
|
||||
{"source": source, "chunks": count}
|
||||
for source, count in file_chunks.items()
|
||||
]
|
||||
|
||||
def delete_document(self, kb_name: str, filename: str) -> int:
|
||||
"""从向量库删除文档"""
|
||||
collection = self.get_collection(kb_name)
|
||||
if not collection:
|
||||
return 0
|
||||
|
||||
result = collection.get(where={"source": filename})
|
||||
|
||||
if not result['ids']:
|
||||
return 0
|
||||
|
||||
collection.delete(ids=result['ids'])
|
||||
deleted = len(result['ids'])
|
||||
|
||||
logger.info(f"从 {kb_name} 删除文档: {filename}, 片段数: {deleted}")
|
||||
return deleted
|
||||
|
||||
def deprecate_document(
|
||||
self,
|
||||
kb_name: str,
|
||||
filename: str,
|
||||
reason: str = "制度废止",
|
||||
deprecated_by: str = ""
|
||||
) -> Dict:
|
||||
"""软删除文档 - 将chunks状态标记为deprecated"""
|
||||
collection = self.get_collection(kb_name)
|
||||
if not collection:
|
||||
return {"success": False, "error": "向量库不存在"}
|
||||
|
||||
result = collection.get(where={"source": filename})
|
||||
|
||||
if not result['ids']:
|
||||
return {"success": False, "error": "文档不存在"}
|
||||
|
||||
deprecated_date = datetime.now().isoformat()
|
||||
updated_metadatas = [
|
||||
{
|
||||
**m,
|
||||
"status": "deprecated",
|
||||
"deprecated_date": deprecated_date,
|
||||
"deprecated_reason": reason,
|
||||
"deprecated_by": deprecated_by
|
||||
}
|
||||
for m in result['metadatas']
|
||||
]
|
||||
|
||||
collection.update(
|
||||
ids=result['ids'],
|
||||
metadatas=updated_metadatas
|
||||
)
|
||||
|
||||
self.rebuild_bm25_index(kb_name)
|
||||
|
||||
logger.info(f"软删除文档: {kb_name}/{filename}, chunks: {len(result['ids'])}, 原因: {reason}")
|
||||
|
||||
return {
|
||||
"success": True,
|
||||
"deprecated_chunks": len(result['ids']),
|
||||
"document_id": filename,
|
||||
"collection": kb_name,
|
||||
"deprecated_date": deprecated_date
|
||||
}
|
||||
|
||||
def restore_document(self, kb_name: str, filename: str) -> Dict:
|
||||
"""恢复已废止的文档"""
|
||||
collection = self.get_collection(kb_name)
|
||||
if not collection:
|
||||
return {"success": False, "error": "向量库不存在"}
|
||||
|
||||
result = collection.get(
|
||||
where={
|
||||
"$and": [
|
||||
{"source": filename},
|
||||
{"status": "deprecated"}
|
||||
]
|
||||
}
|
||||
)
|
||||
|
||||
if not result['ids']:
|
||||
return {"success": False, "error": "未找到已废止的文档"}
|
||||
|
||||
updated_metadatas = [
|
||||
{
|
||||
**m,
|
||||
"status": "active",
|
||||
"deprecated_date": None,
|
||||
"deprecated_reason": None
|
||||
}
|
||||
for m in result['metadatas']
|
||||
]
|
||||
|
||||
collection.update(
|
||||
ids=result['ids'],
|
||||
metadatas=updated_metadatas
|
||||
)
|
||||
|
||||
self.rebuild_bm25_index(kb_name)
|
||||
|
||||
logger.info(f"恢复文档: {kb_name}/{filename}, chunks: {len(result['ids'])}")
|
||||
|
||||
return {
|
||||
"success": True,
|
||||
"restored_chunks": len(result['ids']),
|
||||
"document_id": filename,
|
||||
"collection": kb_name
|
||||
}
|
||||
|
||||
def get_document_chunks(
|
||||
self,
|
||||
kb_name: str,
|
||||
filename: str,
|
||||
status: str = None
|
||||
) -> List[Dict]:
|
||||
"""获取文档的chunks列表"""
|
||||
collection = self.get_collection(kb_name)
|
||||
if not collection:
|
||||
return []
|
||||
|
||||
where_filter = {"source": filename}
|
||||
if status:
|
||||
where_filter["status"] = status
|
||||
|
||||
result = collection.get(where=where_filter)
|
||||
|
||||
return [
|
||||
{
|
||||
"id": id,
|
||||
"document": doc,
|
||||
"metadata": meta,
|
||||
"status": meta.get("status", "active"),
|
||||
"version": meta.get("version", "v1")
|
||||
}
|
||||
for id, doc, meta in zip(
|
||||
result['ids'],
|
||||
result['documents'],
|
||||
result['metadatas']
|
||||
)
|
||||
]
|
||||
|
||||
def get_document_info(self, kb_name: str, filename: str) -> Optional[Dict]:
|
||||
"""获取文档基本信息"""
|
||||
collection = self.get_collection(kb_name)
|
||||
if not collection:
|
||||
return None
|
||||
|
||||
result = collection.get(where={"source": filename})
|
||||
|
||||
if not result['ids']:
|
||||
return None
|
||||
|
||||
status_counts = {}
|
||||
for meta in result['metadatas']:
|
||||
status = meta.get("status", "active")
|
||||
status_counts[status] = status_counts.get(status, 0) + 1
|
||||
|
||||
main_status = "active"
|
||||
if status_counts.get("deprecated", 0) > status_counts.get("active", 0):
|
||||
main_status = "deprecated"
|
||||
elif status_counts.get("superseded", 0) > 0:
|
||||
main_status = "superseded"
|
||||
|
||||
first_meta = result['metadatas'][0] if result['metadatas'] else {}
|
||||
|
||||
return {
|
||||
"document_id": filename,
|
||||
"collection": kb_name,
|
||||
"total_chunks": len(result['ids']),
|
||||
"status": main_status,
|
||||
"status_counts": status_counts,
|
||||
"version": first_meta.get("version", "v1"),
|
||||
"effective_date": first_meta.get("effective_date"),
|
||||
"deprecated_date": first_meta.get("deprecated_date"),
|
||||
"deprecated_reason": first_meta.get("deprecated_reason"),
|
||||
"security_level": first_meta.get("security_level", "public")
|
||||
}
|
||||
|
||||
def list_documents_by_status(
|
||||
self,
|
||||
kb_name: str,
|
||||
status: str = None
|
||||
) -> List[Dict]:
|
||||
"""按状态列出文档"""
|
||||
collection = self.get_collection(kb_name)
|
||||
if not collection:
|
||||
return []
|
||||
|
||||
result = collection.get()
|
||||
|
||||
doc_info = {}
|
||||
for meta in result.get('metadatas', []):
|
||||
source = meta.get('source', 'unknown')
|
||||
chunk_status = meta.get('status', 'active')
|
||||
|
||||
if source not in doc_info:
|
||||
doc_info[source] = {
|
||||
"source": source,
|
||||
"chunks": 0,
|
||||
"status_counts": {},
|
||||
"collection": kb_name
|
||||
}
|
||||
|
||||
doc_info[source]["chunks"] += 1
|
||||
doc_info[source]["status_counts"][chunk_status] = \
|
||||
doc_info[source]["status_counts"].get(chunk_status, 0) + 1
|
||||
|
||||
result_list = []
|
||||
for doc in doc_info.values():
|
||||
counts = doc["status_counts"]
|
||||
if counts.get("deprecated", 0) > counts.get("active", 0):
|
||||
doc["status"] = "deprecated"
|
||||
elif counts.get("superseded", 0) > 0:
|
||||
doc["status"] = "superseded"
|
||||
else:
|
||||
doc["status"] = "active"
|
||||
|
||||
if status is None or doc["status"] == status:
|
||||
result_list.append(doc)
|
||||
|
||||
return result_list
|
||||
|
||||
# add_file_to_kb 方法较长，暂时保留在 manager.py 中
|
||||
# 后续可以拆分到单独的 document_processing.py
|
||||
Reference in New Issue
Block a user