""" 文档服务 - 文档管理和处理服务 """ from typing import List, Dict, Any, Optional from pathlib import Path from loguru import logger from langchain.schema import Document from ..core.document_processor import document_processor from ..core.vector_store import vector_store_manager from ..utils.cache import cache_manager class DocumentService: """文档服务""" def __init__(self): self.default_store_type = "chroma" self.default_collection = "documents" def process_and_store_document(self, file_path: str, store_type: str = None, collection_name: str = None, chunk_size: int = 1000, chunk_overlap: int = 200, additional_metadata: Dict[str, Any] = None) -> Dict[str, Any]: """处理并存储文档""" try: # 处理文档 documents = document_processor.process_document_pipeline( file_path, chunk_size, chunk_overlap, additional_metadata ) # 存储到向量数据库 vector_store_manager.add_documents( documents, store_type or self.default_store_type, collection_name or self.default_collection ) # 获取统计信息 stats = document_processor.get_document_stats(documents) collection_stats = vector_store_manager.get_collection_stats( store_type or self.default_store_type, collection_name or self.default_collection ) result = { "file_path": file_path, "processed_documents": len(documents), "document_stats": stats, "collection_stats": collection_stats, "store_type": store_type or self.default_store_type, "collection_name": collection_name or self.default_collection } logger.info(f"文档处理完成: {file_path}") return result except Exception as e: logger.error(f"文档处理失败: {file_path}, 错误: {e}") raise def process_and_store_directory(self, directory_path: str, store_type: str = None, collection_name: str = None, recursive: bool = True, chunk_size: int = 1000, chunk_overlap: int = 200, additional_metadata: Dict[str, Any] = None) -> Dict[str, Any]: """处理并存储目录中的所有文档""" try: # 加载目录中的所有文档 documents = document_processor.load_directory(directory_path, recursive) if not documents: logger.warning(f"目录中没有找到可处理的文档: {directory_path}") return { "directory_path": directory_path, "processed_documents": 0, "message": "没有找到可处理的文档" } # 分割文档 split_docs = document_processor.split_documents( documents, chunk_size, chunk_overlap ) # 添加元数据 if additional_metadata: split_docs = document_processor.add_metadata(split_docs, additional_metadata) # 存储到向量数据库 vector_store_manager.add_documents( split_docs, store_type or self.default_store_type, collection_name or self.default_collection ) # 获取统计信息 stats = document_processor.get_document_stats(split_docs) collection_stats = vector_store_manager.get_collection_stats( store_type or self.default_store_type, collection_name or self.default_collection ) result = { "directory_path": directory_path, "processed_documents": len(split_docs), "original_documents": len(documents), "document_stats": stats, "collection_stats": collection_stats, "store_type": store_type or self.default_store_type, "collection_name": collection_name or self.default_collection } logger.info(f"目录处理完成: {directory_path}") return result except Exception as e: logger.error(f"目录处理失败: {directory_path}, 错误: {e}") raise def search_documents(self, query: str, k: int = 4, store_type: str = None, collection_name: str = None) -> List[Dict[str, Any]]: """搜索文档""" documents = vector_store_manager.similarity_search( query, k, store_type or self.default_store_type, collection_name or self.default_collection ) results = [] for doc in documents: result = { "content": doc.page_content, "metadata": doc.metadata, "source": doc.metadata.get("source", "unknown") } results.append(result) return results def search_documents_with_scores(self, query: str, k: int = 4, store_type: str = None, collection_name: str = None) -> List[Dict[str, Any]]: """带分数的文档搜索""" results_with_scores = vector_store_manager.similarity_search_with_score( query, k, store_type or self.default_store_type, collection_name or self.default_collection ) results = [] for doc, score in results_with_scores: result = { "content": doc.page_content, "metadata": doc.metadata, "source": doc.metadata.get("source", "unknown"), "similarity_score": score } results.append(result) return results def get_document_collections(self, store_type: str = None) -> List[str]: """获取文档集合列表""" return vector_store_manager.list_collections( store_type or self.default_store_type ) def get_collection_info(self, store_type: str = None, collection_name: str = None) -> Dict[str, Any]: """获取集合信息""" return vector_store_manager.get_collection_stats( store_type or self.default_store_type, collection_name or self.default_collection ) def delete_collection(self, store_type: str = None, collection_name: str = None) -> None: """删除文档集合""" vector_store_manager.delete_collection( store_type or self.default_store_type, collection_name or self.default_collection ) logger.info(f"已删除文档集合: {collection_name or self.default_collection}") def export_documents(self, store_type: str = None, collection_name: str = None, format: str = "json") -> str: """导出文档集合""" # 获取所有文档（这里简化处理，实际可能需要分批获取） documents = vector_store_manager.similarity_search( "", 1000, store_type or self.default_store_type, collection_name or self.default_collection ) if format == "json": import json export_data = { "collection_name": collection_name or self.default_collection, "store_type": store_type or self.default_store_type, "document_count": len(documents), "documents": [ { "content": doc.page_content, "metadata": doc.metadata } for doc in documents ] } return json.dumps(export_data, ensure_ascii=False, indent=2) else: raise ValueError(f"不支持的导出格式: {format}") def get_supported_formats(self) -> List[str]: """获取支持的文件格式""" return list(document_processor.supported_extensions.keys()) def validate_document(self, file_path: str) -> Dict[str, Any]: """验证文档""" file_path = Path(file_path) if not file_path.exists(): return { "valid": False, "error": "文件不存在" } extension = file_path.suffix.lower() if extension not in document_processor.supported_extensions: return { "valid": False, "error": f"不支持的文件格式: {extension}", "supported_formats": self.get_supported_formats() } try: # 尝试加载文档 documents = document_processor.load_document(str(file_path)) stats = document_processor.get_document_stats(documents) return { "valid": True, "file_path": str(file_path), "file_size": file_path.stat().st_size, "extension": extension, "document_count": len(documents), "stats": stats } except Exception as e: return { "valid": False, "error": str(e) } # 全局文档服务实例 document_service = DocumentService()