Go to file
李如威 83b7f45174 feat: 调整代码结构 2025-08-09 13:43:12 +08:00
documents feat: 调整代码结构 2025-08-09 13:43:12 +08:00
examples feat: 调整代码结构 2025-08-09 13:43:12 +08:00
scripts feat: build 2025-07-28 11:13:41 +08:00
src/base_rag feat: 切分检测 2025-08-09 10:46:55 +08:00
storage feat: 调整代码结构 2025-08-09 13:43:12 +08:00
test_files feat: 删除多余文件 2025-08-09 10:33:44 +08:00
.gitignore feat: init project 2025-07-28 10:44:56 +08:00
LICENSE feat: init project 2025-07-28 10:44:56 +08:00
README.md feat: 转换 async 2025-08-08 14:48:41 +08:00
file_status.db feat: 调整代码结构 2025-08-09 13:43:12 +08:00
pyproject.toml feat: 删除多余文件 2025-07-28 23:32:33 +08:00
requirements.txt feat: 转换 async 2025-08-08 14:48:41 +08:00

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
  1. 运行示例
# 异步功能演示
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): 问答逻辑

💡 最佳实践

  1. 开发阶段: 使用本地模型,快速原型验证
  2. 生产环境:
    • 小规模: 本地模型 + 相似度重排
    • 大规模: 本地API + CrossEncoder/BGE重排
  3. 模型选择: 根据语言和领域需求选择合适的嵌入模型
  4. 重排策略: 在实际数据上测试不同重排方法的效果

🛠️ 技术特点

  • 并发安全: 支持并发访问和模型缓存
  • 错误处理: 完善的异常处理和回退机制
  • 灵活配置: 支持多种配置方式和自定义参数
  • 易于扩展: 抽象设计,便于子类实现特定业务逻辑

📋 注意事项

  1. 模型下载: 首次运行会下载模型,需要网络连接
  2. 内存管理: 模型会被缓存,注意内存使用
  3. 文件格式: 确保文档格式受支持txt、md、doc、docx
  4. 错误处理: 注意处理文件加载和模型推理的异常

🔄 版本信息

  • 当前版本: 1.0.0
  • Python要求: >= 3.8
  • 主要特性: 多模型支持,智能重排,文件管理

📚 文档指南

更多详细信息请参考:


🎯 BaseRAG - 灵活强大的RAG框架

  1. 首次运行会下载模型,需要网络连接
  2. 重排功能会增加查询延迟,但提高结果质量
  3. 不同模型对硬件要求不同,请根据实际情况选择
  4. 建议在生产环境前进行充分测试

🔄 版本信息

  • 当前版本: 0.1.0
  • Python要求: >= 3.8
  • 主要依赖: langchain, chromadb, sentence-transformers

更多详细信息请参考: