cb75b9b274
多库检索与存储修复: - RRF 融合去重改用 (collection, chunk_id) 复合键，修复同名文件结果被吞 - DocStore 存储路径加 collection 前缀，修复跨库同名切片数据覆盖 - search_multiple 去重改用复合键 - chunk_id 解析改用 rsplit 兼容下划线文件名 上传与版本管理修复: - 同名文件上传改为覆盖模式，自动清理旧切片 - 修复首次上传不创建版本记录 - 修复覆盖上传版本号回退到 v1 - sync ADDED 分支改用动态版本号生成 - _generate_version_id 改为基于全部版本递增 - 废止/恢复操作同步 SQLite 版本记录 - mark_document_as_superseded 改为仅更新 SQLite 删除清理修复: - 删除文档时同步清理 SQLite 版本记录和变更日志 - 删除向量库时同步清理该库所有版本记录 - cleanup 改为清理 SQLite 记录而非 ChromaDB 测试: - test_version_management.py: 27 条版本管理单元测试 - test_edge_cases.py: 28 条边界用例测试 - test_upload_dedup.py: 5 条上传去重测试 - e2e_risk_test.py: 27 条端到端风险测试 文档: - 新增风险边界问题修复注意事项.md（面向后端的对接文档） - 新增向量库边界风险分析.md - 更新多篇现有文档
333 lines
11 KiB
Python
333 lines
11 KiB
Python
"""
|
||
知识库管理器 - 文档管理 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:
|
||
"""从向量库删除文档，并清理 SQLite 版本记录"""
|
||
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'])
|
||
|
||
# 清理 SQLite 版本记录和变更日志
|
||
try:
|
||
from data.db import get_connection
|
||
with get_connection("knowledge") as conn:
|
||
conn.execute(
|
||
"DELETE FROM document_versions WHERE collection = ? AND document_id = ?",
|
||
(kb_name, filename)
|
||
)
|
||
conn.execute(
|
||
"DELETE FROM version_change_logs WHERE collection = ? AND document_id = ?",
|
||
(kb_name, filename)
|
||
)
|
||
conn.commit()
|
||
except Exception as e:
|
||
logger.warning(f"清理版本记录失败: {e}")
|
||
|
||
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，并同步 SQLite 版本记录"""
|
||
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}")
|
||
|
||
# 同步 SQLite 版本记录
|
||
try:
|
||
from knowledge.document_versions import get_version_query
|
||
from data.db import get_connection
|
||
vq = get_version_query()
|
||
active = vq.get_active_version(kb_name, filename)
|
||
if active:
|
||
with get_connection("knowledge") as conn:
|
||
conn.execute("""
|
||
UPDATE document_versions
|
||
SET status='deprecated', deprecated_date=?, deprecated_reason=?
|
||
WHERE collection=? AND document_id=? AND version=?
|
||
""", (deprecated_date, reason, kb_name, filename, active.version))
|
||
conn.commit()
|
||
vq.log_version_change(
|
||
kb_name, filename,
|
||
change_type="deprecate",
|
||
old_version=active.version,
|
||
old_status="active", new_status="deprecated",
|
||
reason=reason, changed_by=deprecated_by
|
||
)
|
||
except Exception as e:
|
||
logger.warning(f"同步版本记录失败: {e}")
|
||
|
||
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:
|
||
"""恢复已废止的文档，并同步 SQLite 版本记录"""
|
||
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'])}")
|
||
|
||
# 同步 SQLite 版本记录
|
||
try:
|
||
from knowledge.document_versions import get_version_query
|
||
from data.db import get_connection
|
||
vq = get_version_query()
|
||
history = vq.get_document_history(kb_name, filename)
|
||
deprecated_ver = next(
|
||
(v for v in history if v.status.value == 'deprecated'), None
|
||
)
|
||
if deprecated_ver:
|
||
with get_connection("knowledge") as conn:
|
||
conn.execute("""
|
||
UPDATE document_versions
|
||
SET status='active', deprecated_date=NULL, deprecated_reason=NULL
|
||
WHERE collection=? AND document_id=? AND version=?
|
||
""", (kb_name, filename, deprecated_ver.version))
|
||
conn.commit()
|
||
vq.log_version_change(
|
||
kb_name, filename,
|
||
change_type="restore",
|
||
old_version=deprecated_ver.version,
|
||
old_status="deprecated", new_status="active"
|
||
)
|
||
except Exception as e:
|
||
logger.warning(f"同步版本记录失败: {e}")
|
||
|
||
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
|