""" 知识库管理器 - 文档管理 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