Tiny RAG

一个轻量级的 RAG（检索增强生成）系统，仅用单个 Python 文件实现完整的知识库问答功能。

项目简介

Tiny RAG 是一个基于硅基流动 API 的最小化 RAG 实现，支持：

本地文档索引（支持 txt 文件）
向量相似度检索
基于知识库的智能问答
流式输出支持

整个系统仅需一个 Python 文件，无需复杂的依赖和配置，即可快速搭建自己的知识库问答系统。

系统设计

架构概览

┌─────────────────┐
│   本地文档目录    │
│   (doc/*.txt)   │
└────────┬────────┘
         │
         ▼
┌─────────────────┐      ┌──────────────────┐
│   文档切片处理    │─────▶│  文本嵌入API(线上) │
│  (chunk_text)   │      │  (硅基流动API)    │
└────────┬────────┘      └──────┬───────────┘
         │                       │
         │                       ▼
         │              ┌──────────────────┐
         │              │  生成向量嵌入     │
         │              │  (embedding)     │
         │              └──────┬───────────┘
         │                     │
         ▼                     ▼
┌─────────────────┐      ┌──────────────────┐
│  SQLite 向量库   │◀─────│  本地向量存储     │
│ (knowledge_base)│      │  (存储向量结果)   │
└────────┬────────┘      └──────────────────┘
         │
         │ 用户提问
         ▼
┌─────────────────┐      ┌──────────────────┐
│   问题向量化     │─────▶│  文本嵌入API(线上) │
│ (fetch_embedding)│      │  (硅基流动API)    │
└────────┬────────┘      └──────┬───────────┘
         │                       │
         │                       ▼
         │              ┌──────────────────┐
         │              │  生成问题向量     │
         │              └──────┬───────────┘
         │                     │
         ▼                     ▼
┌─────────────────┐      ┌──────────────────┐
│   相似度检索     │◀─────│  从SQLite加载向量 │
│(cosine_similarity)│     │  (本地向量库)     │
└────────┬────────┘      └──────────────────┘
         │                       │
         ▼                       ▼
┌─────────────────┐      ┌──────────────┐
│   上下文组装     │      │  Top-K 片段   │
│ (_prepare_prompt)│      │  (retrieve)   │
└────────┬────────┘      └──────────────┘
         │
         ▼
┌─────────────────┐
│   大模型生成     │
│ (chat_completion)│
└─────────────────┘

核心组件

1. 文档处理层

文件扫描 (iter_txt_files): 递归扫描指定目录下的所有 txt 文件
文本切片 (chunk_text): 将长文本按固定大小切片，支持重叠窗口
- 默认切片大小：400 字符
- 默认重叠长度：50 字符

2. 向量存储层

数据库：python3内置的SQLite3（轻量级，无需额外服务）

表结构：

CREATE TABLE documents (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    path TEXT NOT NULL,           -- 源文件路径
    chunk_index INTEGER NOT NULL, -- 片段索引
    content TEXT NOT NULL,        -- 原始文本内容
    embedding TEXT NOT NULL       -- JSON 格式的向量
)

3. 检索层

嵌入模型：BAAI/bge-large-zh-v1.5（中文优化，通过硅基流动 API 调用）
向量生成：调用线上嵌入 API 将文本转换为向量
相似度计算：余弦相似度（Cosine Similarity，本地计算）
检索策略：
- Top-K 召回（默认 3 个片段）
- 相似度阈值过滤（默认 0.3）

4. 生成层

对话模型：Qwen/Qwen3-8B
提示工程：
- System Prompt：定义助手角色和引用规范
- User Prompt：包含检索到的上下文片段和用户问题
- 自动添加 [source=文件路径#chunk编号] 引用标记

工作流程

知识库构建流程

扫描文档目录，收集所有 .txt 文件
对每个文件进行文本切片（支持重叠）
调用嵌入 API 为每个片段生成向量
将片段内容、向量、元数据存入 SQLite

问答流程

问题向量化：调用线上嵌入 API 将用户问题转换为向量
相似度检索：
- 从 SQLite 加载所有文档片段向量（本地）
- 计算问题向量与所有片段的余弦相似度（本地计算）
- 按相似度排序并过滤阈值
- 返回 Top-K 最相关片段
上下文组装：将检索到的片段与问题组合成提示词
模型生成：调用线上大模型 API 生成回答（支持流式输出）

功能特性

单文件实现：所有功能集中在一个 Python 文件中，易于理解和部署
轻量级存储：使用 SQLite，无需额外数据库服务
中文优化：使用中文优化的嵌入模型
流式输出：支持实时流式生成回答
可配置参数：切片大小、重叠长度、Top-K、相似度阈值等均可配置
引用追踪：自动在回答中标注信息来源
交互式 CLI：提供友好的命令行交互界面

快速开始

环境要求

Python 3.7+
硅基流动 API Key

安装依赖

pip install -r requirements.txt

配置 API Key

编辑 tiny_rag.py 文件，修改 API_KEY 常量：

API_KEY = "your_api_key"  # 替换为你的 API Key

或者通过环境变量设置：

export API_KEY="your_api_key"

准备文档

在项目根目录创建 doc 目录，放入需要索引的 .txt 文件：

mkdir doc
# 将你的文档放入 doc 目录

运行程序

python tiny_rag.py

使用说明

交互式命令行

运行程序后，会进入交互式菜单：

=== 硅基流动 RAG 命令行 ===
1. 知识库初始化（索引 doc 目录）
2. 开始对话（RAG 问答）
0. 退出

选项 1：构建知识库

选择操作 1
输入文档目录（默认：doc）
设置切片长度（默认：400）
设置重叠长度（默认：50）

系统会自动：

扫描目录下的所有 txt 文件
对每个文件进行切片和向量化
存储到 knowledge_base.db 数据库

选项 2：开始问答

选择操作 2
输入你的问题
设置召回片段数量（默认：3）
设置相似度阈值（默认：0.3）

系统会：

显示检索到的相关片段及其相似度
调用大模型生成回答（流式输出）
在回答中自动标注信息来源

配置参数

可在 tiny_rag.py 中修改以下常量：

参数	说明	默认值
`API_BASE`	硅基流动 API 基础地址	`https://api.siliconflow.cn/v1`
`CHAT_MODEL`	对话大模型	`Qwen/Qwen3-8B`
`EMBEDDING_MODEL`	嵌入模型	`BAAI/bge-large-zh-v1.5`
`DB_PATH`	知识库数据库路径	`knowledge_base.db`
`DOC_DIR`	文档目录	`doc`
`DEFAULT_TOP_K`	默认召回片段数	`3`
`DEFAULT_THRESHOLD`	默认相似度阈值	`0.3`

技术栈

Python 3.7+: 核心语言
requests: HTTP 请求库
sqlite3: 轻量级数据库（Python 标准库）
硅基流动 API: 提供嵌入和对话模型服务

项目结构

tiny-rag/
├── tiny_rag.py          # 主程序文件（包含所有功能）
├── requirements.txt     # 依赖列表
├── README.md           # 项目文档
├── LICENSE             # 许可证
├── knowledge_base.db   # 知识库数据库（运行后生成）
└── doc/                # 文档目录（需自行创建）
    └── *.txt           # 待索引的文档文件

工作原理详解

文本切片策略

采用滑动窗口方式切片，确保上下文连续性：

原始文本: [========================================]
切片1:    [==========]
切片2:            [==========]  (重叠50字符)
切片3:                      [==========]

向量检索算法

余弦相似度计算：
```
similarity = (A · B) / (||A|| × ||B||)
```
值域：[-1, 1]，越接近 1 表示越相似
检索流程：
- 计算所有片段的相似度分数
- 按分数降序排序
- 过滤低于阈值的片段
- 返回 Top-K 结果

提示词模板

系统使用结构化的提示词模板：

System: 你是一位知识库问答助手。回答必须严格基于提供的片段，
        并在涉及事实的句子末尾使用 [source=文件路径#chunk编号] 进行引用。
        若片段不足以回答，请明确说明。

User: 已知文档片段：
      [source=doc/file1.txt#chunk0] 片段内容1
      [source=doc/file2.txt#chunk5] 片段内容2
      
      问题：用户的问题
      请基于片段回答，确保在相关句子后附上对应的 source 引用。

注意事项

API Key 安全：不要将 API Key 提交到版本控制系统
文档格式：目前仅支持 UTF-8 编码的 .txt 文件
数据库文件：knowledge_base.db 会在首次运行时自动创建
网络连接：需要稳定的网络连接以访问硅基流动 API
成本控制：每次问答会调用嵌入 API 和对话 API，注意 API 调用成本

Name		Name	Last commit message	Last commit date
Latest commit History 2 Commits
doc		doc
LICENSE		LICENSE
README.md		README.md
knowledge_base.db		knowledge_base.db
requirements.txt		requirements.txt
tiny_rag.py		tiny_rag.py

License

peaprincesslab/tiny-rag

Folders and files

Latest commit

History

Repository files navigation