|
||
---|---|---|
documents | ||
examples | ||
scripts | ||
src/base_rag | ||
storage | ||
test_files | ||
.gitignore | ||
LICENSE | ||
README.md | ||
file_status.db | ||
pyproject.toml | ||
requirements.txt |
README.md
BaseRAG 检索增强生成系统
概述
BaseRAG 是一个灵活的检索增强生成(RAG)框架,支持多种嵌入模型和重排策略,专注于本地部署和HuggingFace生态系统。
🎯 核心特性
1. 多种嵌入模型支持
- 本地HuggingFace模型: 支持模型名称和本地路径两种方式
- 本地API接口: 兼容OpenAI API格式的本地嵌入服务
- 自动回退机制: API不可用时自动切换到本地模型
- 模型缓存: 智能缓存机制,多实例共享模型
2. 文档重排功能
- 相似度重排: 基于余弦相似度,无额外依赖
- CrossEncoder重排: 专业重排模型,效果优秀
- BGE重排: 中文支持良好的重排模型
3. 向量存储与文件管理
- Chroma数据库: 自动持久化,支持多集合管理
- 文件处理: 支持txt、md、doc/docx等多种格式
- 状态追踪: SQLite数据库管理文件处理状态
- 智能去重: 自动检测和跳过重复文件
4. 简洁易用的API
- 抽象基类设计: 易于扩展和自定义
- 配置驱动: 通过配置文件灵活调整模型和参数
- 错误处理: 完善的错误处理和状态报告
🔧 快速开始
安装依赖
pip install -r requirements.txt
基本使用示例
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接口
embedding_config = {
"type": "api",
"api_url": "http://localhost:8000/embeddings",
"model": "your-model",
"api_key": "your-api-key"
}
重排配置
CrossEncoder重排
rerank_config = {
"enabled": True,
"type": "local",
"model": "BAAI/bge-reranker-base",
"top_k": 3
}
相似度重排
rerank_config = {
"enabled": True,
"method": "similarity",
"top_k": 3
}
完整配置示例
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. 文件处理
# 处理单个文件
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. 文档检索
# 基本相似性搜索
docs = await rag.similarity_search("Python编程", k=5)
# 带重排的搜索
docs = await rag.similarity_search_with_rerank("Python编程", k=3)
# 问答查询
answer = await rag.query("什么是Python?")
3. 状态管理
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
- 运行示例
# 异步功能演示
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的核心方法都已异步化:
# 文件处理
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()
并发处理示例
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
性能优势
并发查询性能测试结果:
🔍 运行示例
# 1. 安装依赖
pip install -r requirements.txt
# 2. 运行基础示例
python examples/simple_test.py
📦 依赖要求
核心依赖
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
文档处理依赖
unstructured>=0.10.0
python-docx>=0.8.11
可选依赖
# 本地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)
: 问答逻辑
💡 最佳实践
- 开发阶段: 使用本地模型,快速原型验证
- 生产环境:
- 小规模: 本地模型 + 相似度重排
- 大规模: 本地API + CrossEncoder/BGE重排
- 模型选择: 根据语言和领域需求选择合适的嵌入模型
- 重排策略: 在实际数据上测试不同重排方法的效果
🛠️ 技术特点
- 并发安全: 支持并发访问和模型缓存
- 错误处理: 完善的异常处理和回退机制
- 灵活配置: 支持多种配置方式和自定义参数
- 易于扩展: 抽象设计,便于子类实现特定业务逻辑
📋 注意事项
- 模型下载: 首次运行会下载模型,需要网络连接
- 内存管理: 模型会被缓存,注意内存使用
- 文件格式: 确保文档格式受支持(txt、md、doc、docx)
- 错误处理: 注意处理文件加载和模型推理的异常
🔄 版本信息
- 当前版本: 1.0.0
- Python要求: >= 3.8
- 主要特性: 多模型支持,智能重排,文件管理
📚 文档指南
更多详细信息请参考:
🎯 BaseRAG - 灵活强大的RAG框架!
- 首次运行会下载模型,需要网络连接
- 重排功能会增加查询延迟,但提高结果质量
- 不同模型对硬件要求不同,请根据实际情况选择
- 建议在生产环境前进行充分测试
🔄 版本信息
- 当前版本: 0.1.0
- Python要求: >= 3.8
- 主要依赖: langchain, chromadb, sentence-transformers
更多详细信息请参考: