# BaseRAG 检索增强生成系统 ## 概述 BaseRAG 是一个灵活的检索增强生成(RAG)框架,支持多种嵌入模型和重排策略,专注于本地部署和HuggingFace生态系统。 ## 🎯 核心特性 ### 1. 多种嵌入模型支持 - **本地HuggingFace模型**: 支持模型名称和本地路径两种方式 - **本地API接口**: 兼容OpenAI API格式的本地嵌入服务 - **自动回退机制**: API不可用时自动切换到本地模型 - **模型缓存**: 智能缓存机制,多实例共享模型 ### 2. 文档重排功能 - **相似度重排**: 基于余弦相似度,无额外依赖 - **CrossEncoder重排**: 专业重排模型,效果优秀 - **BGE重排**: 中文支持良好的重排模型 ### 3. 向量存储与文件管理 - **Chroma数据库**: 自动持久化,支持多集合管理 - **文件处理**: 支持txt、md、doc/docx等多种格式 - **状态追踪**: SQLite数据库管理文件处理状态 - **智能去重**: 自动检测和跳过重复文件 ### 4. 简洁易用的API - **抽象基类设计**: 易于扩展和自定义 - **配置驱动**: 通过配置文件灵活调整模型和参数 - **错误处理**: 完善的错误处理和状态报告 ## 🔧 快速开始 ### 安装依赖 ```bash pip install -r requirements.txt ``` ### 基本使用示例 ```python import asyncio from base_rag.core import BaseRAG, FileStatus class MyRAG(BaseRAG): async def ingest(self, file_paths): """批量导入文档""" results = [] for file_path in file_paths: result = await self.process_file_to_vector_store(file_path) results.append(result) return results async def query(self, question): """问答查询""" docs = await self.similarity_search_with_rerank(question, k=3) # 处理文档并生成答案 context = "\n".join([doc.page_content for doc in docs]) return f"基于检索结果: {context[:200]}..." async def main(): # 初始化RAG系统 rag = MyRAG( vector_store_name="my_knowledge", embedding_config={ "type": "local", "model_name": "BAAI/bge-small-zh-v1.5" }, rerank_config={ "enabled": True, "type": "local", "model": "BAAI/bge-reranker-base" } ) # 处理文档 await rag.ingest(["document1.txt", "document2.txt"]) # 查询 answer = await rag.query("什么是Python?") print(answer) ## 📋 配置选项 ### 嵌入模型配置 #### 本地HuggingFace模型 ```python embedding_config = { "type": "local", "model_name": "BAAI/bge-small-zh-v1.5" } # 或使用本地路径 embedding_config = { "type": "local", "model_path": "/path/to/your/model" } ``` #### 本地API接口 ```python embedding_config = { "type": "api", "api_url": "http://localhost:8000/embeddings", "model": "your-model", "api_key": "your-api-key" } ``` ### 重排配置 #### CrossEncoder重排 ```python rerank_config = { "enabled": True, "type": "local", "model": "BAAI/bge-reranker-base", "top_k": 3 } ``` #### 相似度重排 ```python rerank_config = { "enabled": True, "method": "similarity", "top_k": 3 } ``` ### 完整配置示例 ```python rag = MyRAG( vector_store_name="knowledge_base", retriever_top_k=5, persist_directory="./chroma_db", storage_directory="./documents", status_db_path="./file_status.db", embedding_config={ "type": "local", "model_name": "BAAI/bge-small-zh-v1.5" }, rerank_config={ "enabled": True, "type": "local", "model": "BAAI/bge-reranker-base", "top_k": 3 } ) ``` from pydantic import BaseModel app = FastAPI() rag_instance = MyAsyncRAG() class QueryRequest(BaseModel): question: str @app.post("/query") async def query_endpoint(request: QueryRequest): answer = await rag_instance.query(request.question) ## 🚀 使用示例 ### 1. 文件处理 ```python # 处理单个文件 result = await rag.process_file_to_vector_store("document.txt") print(result) # 批量处理文件 file_paths = ["doc1.txt", "doc2.md", "doc3.docx"] results = await rag.ingest(file_paths) # 查看处理状态 status = await rag.get_file_processing_status() completed_files = await rag.list_files_by_status(FileStatus.COMPLETED) ``` ### 2. 文档检索 ```python # 基本相似性搜索 docs = await rag.similarity_search("Python编程", k=5) # 带重排的搜索 docs = await rag.similarity_search_with_rerank("Python编程", k=3) # 问答查询 answer = await rag.query("什么是Python?") ``` ### 3. 状态管理 ```python from base_rag.core import FileStatus # 查看所有文件状态 all_files = await rag.get_file_processing_status() # 查看特定状态的文件 completed = await rag.list_files_by_status(FileStatus.COMPLETED) failed = await rag.list_files_by_status(FileStatus.ERROR) print(f"已完成: {len(completed)} 个文件") print(f"处理失败: {len(failed)} 个文件") ``` ## 📁 项目结构 ``` base_rag/ ├── src/ │ └── base_rag/ │ ├── __init__.py # 包入口 │ └── core.py # 核心BaseRAG类 ├── examples/ │ └── simple_test.py # 基础使用示例 ├── requirements.txt # 依赖列表 ├── pyproject.toml # 包配置 ├── FILE_PROCESSING_GUIDE.md # 文件处理功能说明 ├── RERANK_GUIDE.md # 重排功能详细说明 └── README.md # 项目说明 ``` ``` ## 🚀 快速开始 1. **安装依赖** ```bash pip install -r requirements.txt ``` 2. **运行示例** ```bash # 异步功能演示 python examples/async_example.py # 性能测试 python examples/performance_test.py # 基础功能演示 python examples/quick_start.py # 重排功能演示 python examples/rerank_demo.py # 本地API配置演示 python examples/local_api_demo.py # FastAPI服务示例 pip install fastapi uvicorn uvicorn examples.async_example:app --reload ``` ## 🚀 异步特性详解 ### 主要异步方法 所有BaseRAG的核心方法都已异步化: ```python # 文件处理 await rag.process_file_to_vector_store("document.txt") # 相似性搜索 docs = await rag.similarity_search("query", k=5) # 带重排的搜索 docs = await rag.similarity_search_with_rerank("query", k=3) # 文件状态管理 status = await rag.get_file_processing_status() files = await rag.list_files_by_status(FileStatus.COMPLETED) # 向量库操作 await rag.add_documents_to_vector_store(documents) retriever = await rag.build_retriever() qa_chain = await rag.build_qa_chain() ``` ### 并发处理示例 ```python async def concurrent_file_processing(rag, file_paths, max_concurrent=3): """并发处理多个文件""" semaphore = asyncio.Semaphore(max_concurrent) async def process_single_file(file_path): async with semaphore: return await rag.process_file_to_vector_store(file_path) tasks = [process_single_file(fp) for fp in file_paths] results = await asyncio.gather(*tasks, return_exceptions=True) return results async def concurrent_queries(rag, queries): """并发处理多个查询""" tasks = [rag.similarity_search_with_rerank(q, k=3) for q in queries] results = await asyncio.gather(*tasks) return results ``` ### 性能优势 **并发查询性能测试结果:** ## 🔍 运行示例 ```bash # 1. 安装依赖 pip install -r requirements.txt # 2. 运行基础示例 python examples/simple_test.py ``` ## 📦 依赖要求 ### 核心依赖 ```txt langchain>=0.3.0 langchain-community>=0.3.0 langchain-chroma>=0.1.0 langchain-huggingface>=0.1.0 chromadb>=0.4.0 sentence-transformers>=2.2.0 numpy>=1.21.0 aiofiles>=23.0.0 aiosqlite>=0.19.0 aiohttp>=3.8.0 ``` ### 文档处理依赖 ```txt unstructured>=0.10.0 python-docx>=0.8.11 ``` ### 可选依赖 ```bash # 本地API接口支持 pip install langchain-openai # BGE重排支持 pip install FlagEmbedding ``` ## 🔍 API参考 ### 核心方法 - `similarity_search(query, k)`: 基础相似性搜索 - `similarity_search_with_rerank(query, k)`: 带重排的搜索 - `process_file_to_vector_store(file_path)`: 处理文件到向量库 - `get_file_processing_status()`: 获取文件处理状态 - `list_files_by_status(status)`: 按状态列出文件 ### 抽象方法(需实现) - `ingest(*args, **kwargs)`: 文档导入逻辑 - `query(question)`: 问答逻辑 ## 💡 最佳实践 1. **开发阶段**: 使用本地模型,快速原型验证 2. **生产环境**: - 小规模: 本地模型 + 相似度重排 - 大规模: 本地API + CrossEncoder/BGE重排 3. **模型选择**: 根据语言和领域需求选择合适的嵌入模型 4. **重排策略**: 在实际数据上测试不同重排方法的效果 ## 🛠️ 技术特点 - **并发安全**: 支持并发访问和模型缓存 - **错误处理**: 完善的异常处理和回退机制 - **灵活配置**: 支持多种配置方式和自定义参数 - **易于扩展**: 抽象设计,便于子类实现特定业务逻辑 ## 📋 注意事项 1. **模型下载**: 首次运行会下载模型,需要网络连接 2. **内存管理**: 模型会被缓存,注意内存使用 3. **文件格式**: 确保文档格式受支持(txt、md、doc、docx) 4. **错误处理**: 注意处理文件加载和模型推理的异常 ## 🔄 版本信息 - **当前版本**: 1.0.0 - **Python要求**: >= 3.8 - **主要特性**: 多模型支持,智能重排,文件管理 ## 📚 文档指南 更多详细信息请参考: - **[文件处理功能说明](FILE_PROCESSING_GUIDE.md)** - 文件处理详细介绍 - **[重排功能详细说明](RERANK_GUIDE.md)** - 重排功能配置和使用 - **[示例代码](examples/)** - 使用示例 - **[配置文件](pyproject.toml)** - 项目配置 --- 🎯 **BaseRAG** - 灵活强大的RAG框架! 1. 首次运行会下载模型,需要网络连接 2. 重排功能会增加查询延迟,但提高结果质量 3. 不同模型对硬件要求不同,请根据实际情况选择 4. 建议在生产环境前进行充分测试 ## 🔄 版本信息 - **当前版本**: 0.1.0 - **Python要求**: >= 3.8 - **主要依赖**: langchain, chromadb, sentence-transformers --- 更多详细信息请参考: - [重排功能详细说明](RERANK_GUIDE.md) - [示例代码](examples/) - [配置文件](pyproject.toml)