412 lines
10 KiB
Markdown
412 lines
10 KiB
Markdown
# 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)
|