5ba0d782e2
- 图片选择新增 section_path 字段，支持子章节级过滤 - _filter_images_by_answer 增加 primary_sections 参数和叶子节点匹配 - 检索发散检测：primary_leaf_names > 3 时全局阈值+1 - _FIGURE_ANSWER_KEYWORDS 移除单字"图""表"，正则图号兜底防误触发 - lazy_enhance 后台增强流程优化 - 意图分析与 LLM 工具层改进 评测：图片选择 F1 从 52.4% 提升至 66.3%（+13.9pp），Precision +16.5pp
345 lines
11 KiB
Python
345 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}")
|
||
|
||
# 清理不再被引用的图片和 VLM 缓存文件
|
||
try:
|
||
from knowledge.image_cleanup import cleanup_image_orphans
|
||
cleanup_result = cleanup_image_orphans(self, collections=[kb_name])
|
||
if cleanup_result['deleted_images'] or cleanup_result['deleted_caches']:
|
||
logger.info(
|
||
f"清理孤儿文件: {cleanup_result['deleted_images']} 图片 + "
|
||
f"{cleanup_result['deleted_caches']} VLM缓存"
|
||
)
|
||
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
|