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