# custom_rag

**Repository Path**: asd12580/custom_rag

## Basic Information

- **Project Name**: custom_rag
- **Description**: 当前用AI生成的小模型+RAG解决方案，目前只有雏形，代码全AI生成
- **Primary Language**: Unknown
- **License**: Not specified
- **Default Branch**: master
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 0
- **Forks**: 0
- **Created**: 2026-06-11
- **Last Updated**: 2026-06-15

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# Custom RAG 系统

一个轻量级、模块化的 RAG（检索增强生成）系统，基于 **LlamaIndex + Chroma** 构建，支持 Ollama 本地模型与 OpenAI 风格 API。

---

## 项目结构

```
custom_rag/
├── config.py                    # 全局配置
├── api.py                       # FastAPI Web 服务入口
├── main.py                      # 命令行入口示例（添加文档 + 交互问答）
├── requirements.txt             # Python 依赖
├── core/
│   ├── embedding.py             # 文本嵌入模型封装（Ollama / sentence-transformers）
│   ├── llama_index_engine.py    # RAG 核心引擎（LlamaIndex + Chroma）
│   ├── llm.py                   # LLM 封装（API / Ollama / 本地 多模式）
│   ├── rag_engine.py            # 旧版 RAG 引擎（已迁移到 llama_index_engine.py）
│   └── vector_store.py          # 旧版内存向量库（已迁移到 Chroma）
└── README.md                    # 本文件
```

> 注意：向量存储已从自定义 NumPy+JSON 方案迁移到 **Chroma**，检索框架已迁移到 **LlamaIndex**。旧版 `vector_store.json` 会在服务首次启动时自动导入到 `chroma_db/` 并备份为 `vector_store.json.bak`。

---

## 核心模块说明

| 模块 | 职责 | 说明 |
|------|------|------|
| `config.py` | 全局配置 | 模型路径、API Key、分块参数、RAG 参数等 |
| `core/embedding.py` | 文本向量化 | 支持 `sentence-transformers` 与 Ollama `nomic-embed-text` |
| `core/llama_index_engine.py` | RAG 核心引擎 | 基于 LlamaIndex `VectorStoreIndex` + Chroma 持久化存储 |
| `core/llm.py` | 大模型封装 | 抽象基类设计，支持 `API` / `Ollama` / `HuggingFace` 本地模型 |

---

## 快速开始

### 1. 安装依赖

```bash
pip install -r requirements.txt
```

> 依赖版本已严格固定，避免 pip 解析冲突。如果之前安装过不同版本，建议先 `pip install --upgrade pip`，再重新安装。

### 2. 配置模型

编辑 `config.py`，根据你的情况选择 **API 模式**、**Ollama 模式** 或 **本地模式**。

#### 2.1 嵌入模型配置

```python
# 方式一：Ollama 嵌入（推荐，无需本地下载大模型）
EMBEDDING_MODEL = "nomic-embed-text"

# 方式二：HuggingFace 在线模型（首次自动下载）
EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"

# 方式三：本地模型（填写模型文件夹绝对路径）
EMBEDDING_MODEL = "D:/models/all-MiniLM-L6-v2"

EMBEDDING_DEVICE = "cpu"  # 有 NVIDIA 显卡可改为 "cuda"
```

#### 2.2 LLM 配置

**API 模式（需要联网）**

```python
LLM_TYPE = "api"
```

不同服务商示例：

| 服务商 | `LLM_BASE_URL` | `LLM_MODEL` |
|--------|----------------|-------------|
| OpenAI | `https://api.openai.com/v1` | `gpt-3.5-turbo` |
| DeepSeek | `https://api.deepseek.com/v1` | `deepseek-chat` |
| 智谱 AI | `https://open.bigmodel.cn/api/paas/v4` | `glm-4` |
| 阿里云 | `https://dashscope.aliyuncs.com/compatible-mode/v1` | `qwen-turbo` |
| 本地部署 | `http://localhost:8000/v1` | 看你部署的模型名 |

```python
# 示例：使用 DeepSeek
LLM_TYPE = "api"
LLM_API_KEY = "sk-your-deepseek-key"
LLM_BASE_URL = "https://api.deepseek.com/v1"
LLM_MODEL = "deepseek-chat"
```

**Ollama 模式（本地运行，支持多模态）**

```python
LLM_TYPE = "ollama"
OLLAMA_BASE_URL = "http://localhost:11434"
OLLAMA_MODEL = "gemma4:12b"
```

需要先启动 Ollama 服务：
```bash
ollama run gemma4:12b
```

**本地模式（无需联网）**

```python
LLM_TYPE = "local"

# 方式一：HuggingFace 模型名（首次自动下载）
LOCAL_MODEL_PATH = "microsoft/Phi-3-mini-4k-instruct"

# 方式二：本地模型文件夹绝对路径
LOCAL_MODEL_PATH = "D:/models/Qwen-1_8B-Chat"

LOCAL_MODEL_DEVICE = "cpu"  # 有 GPU 可改为 "cuda"
```

本地模型文件夹内应包含：`config.json`、`tokenizer_config.json`、以及 `.bin` 或 `.safetensors` 权重文件。

**模型缓存位置**

如果之前用过 HuggingFace 下载模型，缓存默认在：

| 系统 | 默认缓存路径 |
|------|-------------|
| Windows | `C:\Users\<用户名>\.cache\huggingface\hub\` |
| Linux / Mac | `~/.cache/huggingface/hub/` |

可以直接复制里面的模型文件夹路径填到配置里。

### 3. 运行

#### 方式一：Web 服务（推荐）

```bash
venv\Scripts\uvicorn.exe api:app --port 8000
```

浏览器打开：http://127.0.0.1:8000

#### 方式二：命令行演示

```bash
venv\Scripts\python.exe main.py
```

程序会自动：
1. 加载嵌入模型
2. 将 `main.py` 中的示例文档添加到 Chroma 向量库
3. 进入交互式问答模式

---

## 多模态文档支持

系统支持图片和 PDF 的多模态理解，通过 **Ollama + 多模态模型** 实现：

```python
# 配置 Ollama 多模态模型
LLM_TYPE = "ollama"
OLLAMA_BASE_URL = "http://localhost:11434"
OLLAMA_MODEL = "gemma4:12b"
```

**添加多模态文档：**

```python
# 支持 .png .jpg .jpeg .gif .bmp .webp .pdf
engine.add_multimodal_documents([
    "./document.pdf",      # PDF（自动提取文本+图片描述）
    "./screenshot.png",    # 图片（自动生成描述）
])
engine.save_store()
```

**处理流程：**

```
图片/PDF → 多模态模型生成描述 → 文本向量化 → 存入 Chroma → 检索回答
```

**特点：**
- PDF 会同时提取**文本**和**图片**，分别处理
- 图片由多模态模型生成详细描述，再入向量库
- 查询时与普通文本 RAG 完全一致

---

## 核心流程

```
用户输入 query
    ↓
[Embedding 模型] 编码 query → 向量
    ↓
[LlamaIndex + Chroma] 向量检索 + BM25 混合检索 → top-k 文档片段
    ↓
[可选：查询改写 / 重排序 / 多跳检索 / 联网搜索]
    ↓
[Prompt 构建] 将 query + 检索到的上下文拼接成 prompt
    ↓
[LLM] 生成最终答案
```

---

## 扩展方向

| 方向 | 建议 |
|------|------|
| **文档加载** | 当前 `main.py` 是硬编码文本，可扩展支持 PDF、Word、网页等解析 |
| **分块策略** | 当前是简单滑动窗口，可升级为按段落、按句子、递归分块等 |
| **Prompt 优化** | 支持多轮对话历史、引用溯源、拒答策略、RAG 模板等 |
| **评估监控** | 添加检索准确率、生成质量评估、日志记录等 |
| **接口化** | 已提供 FastAPI 服务，可继续扩展鉴权、会话管理、文件上传等 |
| **多模态** | 已支持图片、PDF，可扩展音频、视频等多模态文档的 RAG |

---

## 注意事项

- **首次运行** 若使用 HuggingFace 模型会自动下载（约几十MB），请确保网络畅通
- 本地模型模式需要较大内存和磁盘空间，建议使用 4GB 以下的小模型（如 Phi-3-mini、Qwen-1.5B）
- 向量库已迁移到 **Chroma**，数据持久化在 `./chroma_db/` 目录
- 如果存在旧版 `vector_store.json`，服务首次启动会自动导入并备份为 `vector_store.json.bak`
- 启动 Web 服务前请确保 Ollama 或 API 服务可正常访问