shumengya/mengyanote-rag
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: 初始提交 mengya-rag 知识库项目

Browse Source 
轻量级 Obsidian Markdown RAG 系统，包含：
- Markdown 结构感知分块（标题层级 + 代码块/表格整体保留）
- FastEmbed + BAAI/bge-small-zh-v1.5 本地向量化
- SQLite + sqlite-vec 向量库（无需外部服务）
- BM25 + 向量混合检索，RRF 融合
- 盘点模式（有哪些/列表/目录类问题）
- DeepSeek API 生成回答
- mengya-rag CLI 工具（ask/search/context/read/sync/index/status）
- docs/RAG优化方案.md 待实施优化计划

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
shumengya
2026-06-19 14:12:06 +08:00
commit c410d69970
21 changed files with 4902 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
.env.example Normal file
View File

@@ -0,0 +1,15 @@
DEEPSEEK_API_KEY=
DEEPSEEK_MODEL=deepseek-chat
DEEPSEEK_TEMPERATURE=0.2
NOTES_REMOTE=bigmengya:/shumengya/docker/mengyanote-backend/data/mengyanote/
NOTES_DIR=data/notes
INDEX_DIR=data/index/bm25
TOP_K=6
LIST_TOP_K=20
EMBED_BACKEND=fastembed
EMBED_MODEL=BAAI/bge-small-zh-v1.5
EMBED_BATCH_SIZE=32
EMBED_QUERY_PROMPT_NAME=
VECTOR_WEIGHT=0.65
BM25_WEIGHT=0.35

8
.gitignore vendored Normal file
View File

@@ -0,0 +1,8 @@
.env
.venv/
__pycache__/
*.pyc
.pytest_cache/
data/index/
data/notes/
*.log

346
README.md Normal file
View File

@@ -0,0 +1,346 @@
# 萌芽 RAG 知识库
轻量版 Obsidian 笔记 RAG 项目，基于 LlamaIndex 和 DeepSeek API。
检索使用 Markdown 结构感知分块 + 本地小 embedding + BM25 的混合检索，适合个人 Obsidian 笔记这种规模。
## 项目结构
```text
mengya-rag/
├── data/
│ ├── notes/ # 从 bigmengya 同步来的 Obsidian Markdown 笔记（不入库）
│ └── index/bm25/ # 本地索引目录（不入库）
│ └── rag.sqlite3 # SQLite + sqlite-vec 本地向量库
├── docs/
│ └── RAG优化方案.md # 下一步优化计划
├── scripts/
│ ├── sync_notes.sh # 同步 bigmengya 笔记
│ ├── sync_notes.py
│ ├── build_index.sh # 构建索引
│ ├── build_index.py
│ ├── ask.sh # 命令行问答
│ └── ask.py
├── src/mengya_rag/
│ ├── config.py # 配置读取
│ ├── markdown_chunker.py # Markdown 结构感知分块
│ ├── embeddings.py # Embedding 生成（FastEmbed / SentenceTransformers）
│ ├── vector_store.py # SQLite + sqlite-vec 向量存储
│ ├── indexing.py # 建索引（分块 + 向量写入）
│ ├── retrieval.py # 混合检索、RRF 融合、盘点模式
│ ├── qa.py # DeepSeek 生成回答
│ └── cli.py # CLI 命令入口
├── .env.example
├── pyproject.toml
└── README.md
```
## 技术栈
- RAG 基础框架：LlamaIndex
- 本地 embedding：`BAAI/bge-small-zh-v1.5`
- embedding 运行：FastEmbed + ONNX
- 大模型回答：DeepSeek API
- 当前推荐模型：`deepseek-v4-flash`
- 检索方式：BM25 + 向量检索
- 融合方式：RRF 类融合 + 原始相似度过滤 + 同文件去重
- 数据库：无
- 向量库：SQLite + sqlite-vec
- 服务端：无，当前是命令行工具
## 初始化
```bash
cd /shumengya/project/agent/mengya-rag
cp .env.example .env
```
编辑 `.env`，填入：
```env
DEEPSEEK_API_KEY=你的 DeepSeek API Key
```
安装依赖：
```bash
uv sync
```
## 同步笔记
笔记源路径：
```text
bigmengya:/shumengya/docker/mengyanote-backend/data/mengyanote/
```
同步到本地：
```bash
./scripts/sync_notes.sh
# 或
uv run python scripts/sync_notes.py
```
## 构建索引
```bash
./scripts/build_index.sh
# 或
uv run python scripts/build_index.py
```
## 提问
```bash
./scripts/ask.sh "萌芽笔记是什么"
# 或
uv run python scripts/ask.py "萌芽笔记是什么"
```
也可以直接使用安装后的命令：
```bash
uv run mengya-ask "Docker 常用命令有哪些"
```
## Agent 调用推荐
统一入口是 `mengya-rag`。给大模型 Agent 或脚本调用时，推荐固定使用 `--json`，输出是一行 JSON，便于解析。
```bash
# 查看配置、索引和笔记数量
uv run mengya-rag status --json
# 只检索，不调用大模型，返回命中文档和内容
uv run mengya-rag search "Docker 常用命令" -k 5 --json
# 返回可直接注入上游大模型的 context 字段
uv run mengya-rag context "WireGuard 怎么配置" -k 4 --json
# 按 search/context 返回的 source 读取完整笔记
uv run mengya-rag read "Docker/Docker常用命令总结.md" --json
# 本项目自己调用 DeepSeek 生成回答
uv run mengya-rag ask "Docker 常用命令有哪些" -k 4 --json
# 同步笔记并重建索引
uv run mengya-rag reindex --json
```
Agent 优先使用这几个命令：
- `status --json`：检查索引是否存在、笔记数、节点数、模型配置。
- `search --json`：拿结构化检索结果，包含 `results[].source`、`title`、`heading_path`、`score`、`content`。
- `context --json`：拿拼好的 `context` 字符串，适合直接放进其他大模型提示词。
- `read --json`：按 `source` 读取完整 Markdown 原文，适合检索命中后补充上下文。
- `ask --json`：让本项目完成 RAG + DeepSeek 回答，返回 `answer` 和 `sources`。
- `reindex --json`：一条命令完成同步笔记和重建索引。
常用参数：
```bash
--mode auto|hybrid|inventory # auto 默认；hybrid 普通语义检索；inventory 盘点/目录类检索
-k, --top-k 数字 # 控制返回条数
--env-file /path/.env # 指定配置文件
--notes-dir /path/notes # 覆盖笔记目录
--index-dir /path/index # 覆盖索引目录
```
## 查询模式
项目会自动区分两类问题：
- 普通问答：例如“WireGuard 命令怎么用”，走 Markdown 分块 + 向量/BM25 混合检索。
- 盘点列表：例如“查一下我目前博客文章有哪些”“查一下安卓 Gradle 相关笔记”，走文件级目录检索，优先列出相关 Markdown 文件和内容预览。
## 技术原理
### Markdown 分块
项目不使用固定字符数暴力切割，而是使用结构感知分块：
```text
Markdown 文件
↓
解析 frontmatter
↓
识别 H1-H6 标题
↓
按标题层级组织分块
↓
代码块和表格整体保留
↓
长块使用滑动窗口兜底
```
每个 chunk 前会附加上下文前缀：
```text
文件: AI/控制台Agent工具安装教程.md
标题: 控制台Agent工具安装教程
标题路径: H1 > H2 > H3
```
每个 chunk 会保存这些元数据：
```text
source_file
rel_path
file_name
folder_path
title
heading_path
h1
h2
h3
h4
h5
h6
chunk_index
tags
created_at
```
### 索引构建
执行：
```bash
./scripts/build_index.sh
```
构建流程：
```text
Markdown 文件
↓
结构感知分块 → TextNode
↓
bge-small-zh-v1.5 生成 embedding → 归一化
↓
写入 SQLite + sqlite-vec（rag.sqlite3）
```
当前索引规模参考：
```text
笔记文件数: 495
分块数: ~1800
rag.sqlite3: ~10MB（chunk 表 + sqlite-vec 向量表）
```
### 普通问答检索
普通问答适合“怎么做”“是什么”“原理是什么”这类问题。
```text
用户问题
↓
query 清洗
↓
BM25 关键词检索
↓
本地向量语义检索
↓
RRF 融合
↓
同文件去重
↓
相关性过滤
↓
top-k chunk 注入 DeepSeek
↓
生成回答和来源
```
### 盘点列表检索
盘点列表适合“有哪些”“查一下”“相关笔记”“目录”“清单”这类问题。
```text
用户问题
↓
识别为盘点类问题
↓
构建文件级摘要节点
↓
文件级 BM25 检索
↓
关键词覆盖过滤
↓
输出 Markdown 文件列表、路径和内容预览
```
## 重新更新知识库
笔记变更后按顺序执行：
```bash
./scripts/sync_notes.sh
./scripts/build_index.sh
```
## 设计取舍
- Markdown 笔记按结构感知分块：识别 frontmatter、H1-H6、代码块、表格。
- H1-H6 标题会写入 `h1` 到 `h6` 元数据，同时保存 `heading_path` 面包屑。
- 代码块和表格整体保留，不跨块切割；长块再用滑动窗口兜底。
- 每个 chunk 都带 `source_file`、`heading_path`、`chunk_index`、`tags`、`created_at`。
- 本地 embedding 使用 `BAAI/bge-small-zh-v1.5`，通过 FastEmbed/ONNX 运行，不依赖 PyTorch 和 GPU。
- 检索采用 BM25 + 向量召回，再用 RRF 融合和同文件去重，避免只靠关键词或只靠语义。
- 使用 DeepSeek 只负责生成回答：减少 API 调用，只在提问时调用大模型。
- 本地向量库使用 SQLite + sqlite-vec：无需 Qdrant、Milvus、Postgres 等额外服务。
## 当前效果
目前比较适合处理：
- 某类笔记有哪些
- 某个教程在哪里
- 根据笔记总结某个主题
- 列出博客文章
- 列出 Agent 安装教程
- 查 Gradle 构建相关内容
示例：
```bash
./scripts/ask.sh '查一下我目前博客文章有哪些'
./scripts/ask.sh '查一下我的安卓gradle构建相关笔记'
./scripts/ask.sh '查一下我的Agent安装教程'
```
## 当前限制
- 当前是 SQLite + sqlite-vec 单机向量库，数据量很大后可以换 Qdrant 或 Chroma。
- 没有 reranker，复杂问题可能还会有弱相关内容混入。
- embedding 模型是轻量中文模型，效果够用但不是最强。
- `build_index` 每次是全量重建，还没有增量索引。
- frontmatter 解析是轻量手写版，不是完整 YAML 解析。
- 当前是 CLI 工具，没有 Web 页面和 API 服务。
## 后续优化方向
详见 [docs/RAG优化方案.md](docs/RAG优化方案.md)，主要方向：
- 中文分词换 jieba（BM25 召回率最大瓶颈）
- 放宽文件去重（当前每文件只取 1 chunk，长文章不友好）
- QA Prompt 加强引用约束，降低幻觉
- 注入上下文加 token 预算（节省费用 + 提升精度）
- 可选：换 `BAAI/bge-m3` 嵌入模型（中英混合效果更好）
- 增量索引，只重建变更文件
- 加 reranker，例如 `bge-reranker-v2-m3`
- 数据量变大后接 Qdrant
## 调参
`.env` 中可调整：
```env
TOP_K=6
EMBED_MODEL=BAAI/bge-small-zh-v1.5
EMBED_BATCH_SIZE=32
VECTOR_WEIGHT=0.65
BM25_WEIGHT=0.35
```

262
docs/RAG优化方案.md Normal file
View File

@@ -0,0 +1,262 @@
# mengya-rag 知识库优化方案
> 记录日期：2026-06-19
> 当前状态：待实施
> 目标：低成本 · 低幻觉 · 高命中
---
## 一、现状诊断
### 基础信息
| 项目 | 当前值 |
|------|--------|
| 笔记数量 | 495 篇 Markdown |
| Embedding 模型 | `BAAI/bge-small-zh-v1.5`（90MB，512 dim） |
| 检索方式 | Hybrid（BM25 + 向量），RRF 融合 |
| LLM | DeepSeek（deepseek-v4-flash） |
| 中文分词 | 自制 N-gram 滑窗 |
---
### 问题 1：中文分词太弱（最核心问题）
当前 `chinese_tokenizer` 用的是 N-gram 滑窗，对技术笔记效果极差。
**示例**：搜索"GMP调度系统"
```
N-gram 分词结果：["Go", "Gol", "Gola", "olan", "lang", "中的", "的G", "GMP"]
jieba 分词结果：["Golang", "中", "GMP", "调度", "系统"]
```
N-gram 把"GMP"和"调度"拆散，BM25 根本匹配不上完整关键词，召回率大幅下降。495篇涵盖 Golang、Docker、C++ 等技术笔记，这个问题影响所有技术查询。
---
### 问题 2：文件级强去重损失上下文
`_fuse_results` 里每个文件严格只返回 **1 个 chunk**。
- "字节Agent中台一面凉经.md" 有 1459 行，答案往往跨多个章节
- 只拿最高分那一个 chunk，丢失了大量关联内容
- 对长篇面经、长教程类笔记尤其严重
---
### 问题 3：Prompt 没有引用约束
当前 Prompt 只说"根据以下内容回答"，没有强制引用来源编号。DeepSeek 在上下文不完整时会"补充说明"而不是说"笔记中没有"，导致幻觉。
---
### 问题 4：注入上下文无 token 预算
所有检索到的 chunk 全量拼入 prompt。若 `top_k=6` 且每块 512 token，注入了 3000+ token，低相关内容稀释关键信息，模型注意力分散，精度反而下降，同时增加 API 费用。
---
### 问题 5：嵌入模型偏小
`BAAI/bge-small-zh-v1.5` 对纯中文 ok，但笔记大量混合英文技术词汇（Docker、Golang、nginx……），小模型语义泛化能力有限，向量召回对中英混合查询效果较差。
---
## 二、优化方案
### 方案总览
| # | 改动 | 影响文件 | 预期收益 | 优先级 |
|---|------|----------|----------|--------|
| ① | 中文分词换 jieba | `indexing.py`、`pyproject.toml` | BM25 关键词召回大幅提升 | 🔴 高 |
| ② | 文件去重改为最多 2 chunk/文件 | `retrieval.py` | 长文多节答题完整度提升 | 🟠 中 |
| ③ | QA Prompt 加编号引用 + 严格反幻觉 | `qa.py` | 幻觉率明显下降 | 🟠 中 |
| ④ | 注入上下文加 token 预算（≤ 2000 token） | `qa.py` | 精度提升 + 节省 API 费用 | 🟡 中 |
| ⑤ | 换 `BAAI/bge-m3` 嵌入模型（可选） | `.env` | 中英混合语义理解提升 | 🟢 低 |
---
### ① jieba 中文分词（最高优先级）
**依赖**：`pyproject.toml` 加入 `jieba>=0.42`
**改动位置**：`src/mengya_rag/indexing.py` → `chinese_tokenizer()`
```python
# 旧实现：N-gram 滑窗
def chinese_tokenizer(text: str) -> list[str]:
for part in re.findall(r"[A-Za-z0-9_#+.-]+|[一-鿿]+", text.lower()):
# N-gram 2~4 滑窗...
# 新实现：jieba 分词
import jieba
STOPWORDS = {"请", "根据", "笔记", "简要", "回答", "哪些", "什么", "内容", "一下", "我的", "有哪些", "是什么"}
def chinese_tokenizer(text: str) -> list[str]:
tokens = []
for part in re.findall(r"[A-Za-z0-9_#+.-]+|[一-鿿]+", text.lower()):
if re.fullmatch(r"[A-Za-z0-9_#+.-]+", part):
if part not in {"md", "markdown"} and part not in STOPWORDS:
tokens.append(part)
else:
for word in jieba.cut(part):
if len(word) >= 2 and word not in STOPWORDS:
tokens.append(word)
return tokens
```
**可选增强**：在 `data/` 目录放 `userdict.txt`，补充专有名词：
```
萌芽笔记 5
bigmengya 5
smallmengya 5
GMP调度 5
```
**注意**：改分词后必须重建索引（`mengya-build-index`）。
---
### ② 文件级去重放宽（最多 2 chunk/文件）
**改动位置**：`src/mengya_rag/retrieval.py` → `HybridRetriever._fuse_results()`
```python
# 旧：严格 1 chunk/文件
seen_files: set[str] = set()
...
if rel_path in seen_files:
continue
seen_files.add(rel_path)
# 新：最多 2 chunk/文件
from collections import Counter
seen_files: Counter[str] = Counter()
...
if seen_files[rel_path] >= 2:
continue
seen_files[rel_path] += 1
```
---
### ③ QA Prompt 强化（反幻觉 + 强制引用）
**改动位置**：`src/mengya_rag/qa.py` → `QA_PROMPT` + `format_context()`
```python
QA_PROMPT = """你是萌芽 RAG 知识库助手。请严格根据下方标有编号的笔记片段回答问题。
规则：
1. 回答中必须用 [来源N] 标注依据，例如"根据 [来源1] 可知……"
2. 若所有片段均无法支撑某个细节，必须明确说"笔记中未找到相关记录"，不得凭己见补充
3. 回答要准确、完整，优先使用中文
笔记内容：
---------------------
{context}
---------------------
问题：{question}
回答："""
def format_context(items: list[RetrievedNode]) -> str:
parts: list[str] = []
for index, item in enumerate(items, start=1):
rel_path = item.node.metadata.get("rel_path", "未知来源")
parts.append(
f"[来源{index}]《{rel_path}》\n"
f"{item.node.get_content().strip()}"
)
return "\n\n".join(parts)
```
---
### ④ 上下文 Token 预算（硬限 2000 token）
**改动位置**：`src/mengya_rag/qa.py` → `format_context()` 或 `ask_question()`
```python
CONTEXT_TOKEN_BUDGET = 2000 # 可通过 env 配置
def format_context(items: list[RetrievedNode], budget: int = CONTEXT_TOKEN_BUDGET) -> str:
from .markdown_chunker import estimate_tokens
parts: list[str] = []
used = 0
for index, item in enumerate(items, start=1):
rel_path = item.node.metadata.get("rel_path", "未知来源")
content = item.node.get_content().strip()
block = f"[来源{index}]《{rel_path}》\n{content}"
cost = estimate_tokens(block)
if used + cost > budget and parts:
break
parts.append(block)
used += cost
return "\n\n".join(parts)
```
---
### ⑤ 换更大嵌入模型（可选，需重建索引）
修改 `.env`：
```diff
-EMBED_MODEL=BAAI/bge-small-zh-v1.5
+EMBED_MODEL=BAAI/bge-m3
```
| 对比项 | bge-small-zh | bge-m3 |
|--------|-------------|--------|
| 大小 | ~90 MB | ~570 MB |
| 维度 | 512 | 1024 |
| 语言 | 纯中文 | 中英多语言 |
| 混合技术词召回 | 一般 | 好 |
改完后执行：
```bash
mengya-build-index
```
---
## 三、实施顺序建议
```
第 1 步：换 jieba 分词（改代码 + 重建索引） ← 收益最大
第 2 步：放宽文件去重 + 加 token 预算 ← 改代码，立即生效
第 3 步：加强 Prompt ← 改代码，立即生效
第 4 步（可选）：换 bge-m3 嵌入模型 ← 改配置 + 重建索引
```
---
## 四、改动文件清单
```
pyproject.toml ← 加 jieba 依赖
src/mengya_rag/indexing.py ← chinese_tokenizer 换 jieba
src/mengya_rag/retrieval.py ← _fuse_results 去重策略
src/mengya_rag/qa.py ← QA_PROMPT + format_context token 预算
data/userdict.txt ← （可选）jieba 自定义词典
.env / .env.example ← （可选）EMBED_MODEL 升级
```
---
## 五、重建索引命令
改完代码后执行：
```bash
# 只重建索引（笔记已在本地）
mengya-build-index
# 或先同步笔记再重建
mengya-rag reindex
```

25
pyproject.toml Normal file
View File

@@ -0,0 +1,25 @@
[project]
name = "mengya-rag"
version = "0.1.0"
description = "轻量版萌芽 Obsidian RAG 知识库"
requires-python = ">=3.11"
dependencies = [
"fastembed>=0.7.0",
"llama-index-core>=0.12.0",
"llama-index-llms-deepseek>=0.1.0",
"python-dotenv>=1.0.1",
"sqlite-vec>=0.1.9",
]
[project.scripts]
mengya-rag = "mengya_rag.cli:main"
mengya-sync-notes = "mengya_rag.cli:sync_notes"
mengya-build-index = "mengya_rag.cli:build_index"
mengya-ask = "mengya_rag.cli:ask"
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"
[tool.hatch.build.targets.wheel]
packages = ["src/mengya_rag"]

9
scripts/ask.py Normal file
View File

@@ -0,0 +1,9 @@
import os
os.environ.setdefault("TRANSFORMERS_VERBOSITY", "error")
from mengya_rag.cli import ask
if __name__ == "__main__":
ask()

6
scripts/ask.sh Executable file
View File

@@ -0,0 +1,6 @@
#!/usr/bin/env bash
set -euo pipefail
cd "$(dirname "$0")/.."
export TRANSFORMERS_VERBOSITY=error
uv run mengya-ask "$@"

9
scripts/build_index.py Normal file
View File

@@ -0,0 +1,9 @@
import os
os.environ.setdefault("TRANSFORMERS_VERBOSITY", "error")
from mengya_rag.cli import build_index
if __name__ == "__main__":
build_index()

6
scripts/build_index.sh Executable file
View File

@@ -0,0 +1,6 @@
#!/usr/bin/env bash
set -euo pipefail
cd "$(dirname "$0")/.."
export TRANSFORMERS_VERBOSITY=error
uv run mengya-build-index "$@"

9
scripts/sync_notes.py Normal file
View File

@@ -0,0 +1,9 @@
import os
os.environ.setdefault("TRANSFORMERS_VERBOSITY", "error")
from mengya_rag.cli import sync_notes
if __name__ == "__main__":
sync_notes()

6
scripts/sync_notes.sh Executable file
View File

@@ -0,0 +1,6 @@
#!/usr/bin/env bash
set -euo pipefail
cd "$(dirname "$0")/.."
export TRANSFORMERS_VERBOSITY=error
uv run mengya-sync-notes "$@"

3
src/mengya_rag/__init__.py Normal file
View File

@@ -0,0 +1,3 @@
"""Mengya lightweight RAG."""
__version__ = "0.1.0"

597
src/mengya_rag/cli.py Normal file
View File

@@ -0,0 +1,597 @@
from __future__ import annotations
import argparse
import json
import os
import sys
import time
from pathlib import Path
from typing import Any
os.environ.setdefault("TRANSFORMERS_VERBOSITY", "error")
from .config import Settings, load_settings, override_settings
from .indexing import build_index, load_index_nodes, markdown_files
from .indexing import sync_notes as sync_notes_impl
from .vector_store import db_path
def _positive_int(value: str) -> int:
parsed = int(value)
if parsed <= 0:
raise argparse.ArgumentTypeError("必须是大于 0 的整数")
return parsed
def _existing_file(value: str) -> Path:
path = Path(value)
if not path.exists():
raise argparse.ArgumentTypeError(f"文件不存在: {path}")
if not path.is_file():
raise argparse.ArgumentTypeError(f"不是文件: {path}")
return path
def _add_config_args(parser: argparse.ArgumentParser) -> None:
parser.add_argument(
"--env-file",
type=_existing_file,
default=argparse.SUPPRESS,
help="指定 .env 配置文件",
)
parser.add_argument(
"--notes-dir",
type=Path,
default=argparse.SUPPRESS,
help="覆盖 NOTES_DIR",
)
parser.add_argument(
"--index-dir",
type=Path,
default=argparse.SUPPRESS,
help="覆盖 INDEX_DIR",
)
def _add_output_args(parser: argparse.ArgumentParser) -> None:
parser.add_argument(
"--format",
choices=("text", "json"),
default=argparse.SUPPRESS,
help="输出格式，Agent 调用建议使用 json",
)
parser.add_argument(
"--json",
action="store_true",
default=argparse.SUPPRESS,
help="等同于 --format json",
)
def _add_retrieval_args(parser: argparse.ArgumentParser) -> None:
parser.add_argument(
"--mode",
choices=("auto", "hybrid", "inventory"),
default="auto",
help="检索模式，默认 auto 自动判断",
)
parser.add_argument("-k", "--top-k", type=_positive_int, help="覆盖返回条数")
def _settings_from_args(args: argparse.Namespace) -> Settings:
settings = load_settings(getattr(args, "env_file", None))
return override_settings(
settings,
notes_dir=getattr(args, "notes_dir", None),
index_dir=getattr(args, "index_dir", None),
)
def _json_default(value: object) -> str:
if isinstance(value, Path):
return str(value)
raise TypeError(f"Object of type {type(value).__name__} is not JSON serializable")
def _write_json(payload: dict[str, Any]) -> None:
print(json.dumps(payload, ensure_ascii=False, default=_json_default))
def _wants_json(args: argparse.Namespace) -> bool:
return bool(getattr(args, "json", False) or getattr(args, "format", "text") == "json")
def _emit(args: argparse.Namespace, payload: dict[str, Any], text: str | None = None) -> None:
if _wants_json(args):
_write_json(payload)
return
if text is not None:
print(text)
def _fail(args: argparse.Namespace, message: str, *, code: int = 1) -> None:
if _wants_json(args):
_write_json({"ok": False, "error": message})
raise SystemExit(code)
raise SystemExit(message)
def _question_from_args(args: argparse.Namespace) -> str:
question = " ".join(getattr(args, "question", [])).strip()
if not question and not sys.stdin.isatty():
question = sys.stdin.read().strip()
return question
def _print_sources(sources: list[str]) -> None:
if not sources:
print("\n来源: 无")
return
print("\n来源:")
for source in sources:
print(f"- {source}")
def _preview_text(text: str, limit: int) -> str:
text = " ".join(text.strip().split())
if len(text) <= limit:
return text
return f"{text[:limit].rstrip()}..."
def _resolve_note_path(settings: Settings, source: str) -> Path:
notes_root = settings.notes_dir.resolve()
note_path = (notes_root / source).resolve()
try:
note_path.relative_to(notes_root)
except ValueError as exc:
raise ValueError(f"笔记路径不能越过 NOTES_DIR: {source}") from exc
if not note_path.exists():
raise FileNotFoundError(f"笔记不存在: {source}")
if not note_path.is_file():
raise IsADirectoryError(f"不是笔记文件: {source}")
return note_path
def _source_text(item: Any) -> str:
metadata = item.node.metadata
rel_path = metadata.get("rel_path") or metadata.get("file_path") or "未知来源"
return f"{rel_path} score={item.score:.4f}"
def _node_payload(item: Any, *, include_content: bool = True) -> dict[str, Any]:
metadata = item.node.metadata
payload: dict[str, Any] = {
"source": metadata.get("rel_path") or metadata.get("file_path") or "未知来源",
"title": metadata.get("title", ""),
"heading_path": metadata.get("heading_path", ""),
"chunk_index": metadata.get("chunk_index"),
"score": item.score,
"raw_score": item.raw_score,
"metadata": dict(metadata),
}
if include_content:
payload["content"] = item.node.get_content().strip()
return payload
def _retrieve(args: argparse.Namespace) -> list[Any]:
from .qa import retrieve_nodes
question = _question_from_args(args)
if not question:
_fail(args, "请提供检索内容，例如: mengya-rag search Docker")
settings = _settings_from_args(args)
try:
return retrieve_nodes(
settings,
question,
mode=args.mode,
top_k=args.top_k,
)
except Exception as exc:
_fail(args, f"检索失败: {exc}")
def _sync_notes(args: argparse.Namespace) -> int:
settings = _settings_from_args(args)
started = time.time()
try:
sync_notes_impl(settings)
except Exception as exc:
_fail(args, f"同步失败: {exc}")
payload = {
"ok": True,
"command": "sync",
"notes_remote": settings.notes_remote,
"notes_dir": settings.notes_dir,
"elapsed_seconds": round(time.time() - started, 3),
}
_emit(args, payload, f"笔记已同步到: {settings.notes_dir}")
return 0
def _build_index(args: argparse.Namespace) -> int:
settings = _settings_from_args(args)
started = time.time()
try:
count = build_index(settings)
except Exception as exc:
_fail(args, f"构建索引失败: {exc}")
payload = {
"ok": True,
"command": "index",
"index_dir": settings.index_dir,
"db_path": db_path(settings.index_dir),
"node_count": count,
"elapsed_seconds": round(time.time() - started, 3),
}
text = f"索引已生成: {settings.index_dir}\n索引分块数: {count}"
_emit(args, payload, text)
return 0
def _reindex(args: argparse.Namespace) -> int:
settings = _settings_from_args(args)
started = time.time()
try:
sync_notes_impl(settings)
count = build_index(settings)
except Exception as exc:
_fail(args, f"更新知识库失败: {exc}")
payload = {
"ok": True,
"command": "reindex",
"notes_remote": settings.notes_remote,
"notes_dir": settings.notes_dir,
"index_dir": settings.index_dir,
"db_path": db_path(settings.index_dir),
"node_count": count,
"elapsed_seconds": round(time.time() - started, 3),
}
text = f"知识库已更新: {settings.index_dir}\n索引分块数: {count}"
_emit(args, payload, text)
return 0
def _ask(args: argparse.Namespace) -> int:
from .qa import ask_question
question = _question_from_args(args)
if not question:
_fail(args, "请提供问题，例如: mengya-rag ask Docker 常用命令有哪些")
settings = _settings_from_args(args)
try:
answer, sources = ask_question(
settings,
question,
mode=args.mode,
top_k=args.top_k,
)
except Exception as exc:
_fail(args, f"提问失败: {exc}")
payload = {
"ok": True,
"command": "ask",
"question": question,
"answer": answer,
"sources": sources,
"mode": args.mode,
"top_k": args.top_k,
}
if _wants_json(args):
_write_json(payload)
else:
print(answer)
if not args.no_source:
_print_sources(sources)
return 0
def _search(args: argparse.Namespace) -> int:
question = _question_from_args(args)
if not question:
_fail(args, "请提供检索内容，例如: mengya-rag search Docker")
nodes = _retrieve(args)
include_content = _wants_json(args) or args.show_content
if not nodes:
_emit(
args,
{
"ok": True,
"command": "search",
"question": question,
"results": [],
},
"未检索到相关内容。",
)
return 0
payload = {
"ok": True,
"command": "search",
"question": question,
"mode": args.mode,
"top_k": args.top_k,
"results": [
_node_payload(item, include_content=include_content) for item in nodes
],
}
if _wants_json(args):
_write_json(payload)
return 0
for index, item in enumerate(nodes, start=1):
metadata = item.node.metadata
heading_path = metadata.get("heading_path", "")
print(f"{index}. {_source_text(item)}")
if heading_path:
print(f" 标题路径: {heading_path}")
if args.show_content:
preview = _preview_text(item.node.get_content(), args.preview_chars)
print(f" 预览: {preview}")
return 0
def _context(args: argparse.Namespace) -> int:
question = _question_from_args(args)
if not question:
_fail(args, "请提供问题，例如: mengya-rag context Docker 常用命令")
nodes = _retrieve(args)
context_parts: list[str] = []
sources: list[dict[str, Any]] = []
for index, item in enumerate(nodes, start=1):
metadata = item.node.metadata
source = metadata.get("rel_path") or metadata.get("file_path") or "未知来源"
content = item.node.get_content().strip()
context_parts.append(f"[{index}] 来源: {source}\n{content}")
sources.append(_node_payload(item, include_content=False))
context = "\n\n".join(context_parts)
payload = {
"ok": True,
"command": "context",
"question": question,
"mode": args.mode,
"top_k": args.top_k,
"context": context,
"sources": sources,
}
if _wants_json(args):
_write_json(payload)
else:
print(context or "未检索到相关内容。")
return 0
def _read(args: argparse.Namespace) -> int:
settings = _settings_from_args(args)
source = args.source.strip()
if not source:
_fail(args, "请提供笔记相对路径，例如: mengya-rag read Docker/Docker常用命令总结.md")
try:
path = _resolve_note_path(settings, source)
content = path.read_text(encoding="utf-8")
except Exception as exc:
_fail(args, f"读取笔记失败: {exc}")
original_length = len(content)
if args.max_chars is not None:
content = content[: args.max_chars]
payload = {
"ok": True,
"command": "read",
"source": source,
"path": path,
"content": content,
"truncated": args.max_chars is not None and original_length > len(content),
}
_emit(args, payload, content)
return 0
def _status(args: argparse.Namespace) -> int:
settings = _settings_from_args(args)
db_file = db_path(settings.index_dir)
note_count: int | None = None
note_error: str | None = None
try:
note_count = len(markdown_files(settings.notes_dir))
except Exception as exc:
note_error = str(exc)
node_count: int | None = None
node_error: str | None = None
try:
node_count = len(load_index_nodes(settings))
except Exception as exc:
node_error = str(exc)
payload = {
"ok": True,
"command": "status",
"project_root": Path(__file__).resolve().parents[2],
"notes_dir": settings.notes_dir,
"index_dir": settings.index_dir,
"db_path": db_file,
"db_exists": db_file.exists(),
"deepseek_api_key_configured": bool(settings.deepseek_api_key),
"deepseek_model": settings.deepseek_model,
"embed_backend": settings.embed_backend,
"embed_model": settings.embed_model,
"top_k": settings.top_k,
"list_top_k": settings.list_top_k,
"markdown_note_count": note_count,
"markdown_note_error": note_error,
"index_node_count": node_count,
"index_node_error": node_error,
}
text_lines = [
f"项目目录: {payload['project_root']}",
f"笔记目录: {settings.notes_dir}",
f"索引目录: {settings.index_dir}",
f"向量库: {db_file} ({'存在' if db_file.exists() else '不存在'})",
f"DeepSeek Key: {'已配置' if settings.deepseek_api_key else '未配置'}",
f"DeepSeek 模型: {settings.deepseek_model}",
f"Embedding: {settings.embed_backend} / {settings.embed_model}",
f"默认 top_k: {settings.top_k}",
f"默认 list_top_k: {settings.list_top_k}",
f"Markdown 笔记数: {note_count if note_error is None else f'读取失败 ({note_error})'}",
f"索引节点数: {node_count if node_error is None else f'读取失败 ({node_error})'}",
]
_emit(args, payload, "\n".join(text_lines))
return 0
def _add_question_arg(parser: argparse.ArgumentParser) -> None:
parser.add_argument("question", nargs="*", help="问题；省略时从 stdin 读取")
def build_parser(prog: str = "mengya-rag") -> argparse.ArgumentParser:
config_parent = argparse.ArgumentParser(add_help=False)
_add_config_args(config_parent)
_add_output_args(config_parent)
parser = argparse.ArgumentParser(
prog=prog,
description="萌芽 RAG 知识库命令行工具",
parents=[config_parent],
)
subparsers = parser.add_subparsers(dest="command")
sync_parser = subparsers.add_parser(
"sync",
help="从远端同步 Markdown 笔记",
parents=[config_parent],
)
sync_parser.set_defaults(func=_sync_notes)
index_parser = subparsers.add_parser(
"index",
help="构建本地检索索引",
parents=[config_parent],
)
index_parser.set_defaults(func=_build_index)
reindex_parser = subparsers.add_parser(
"reindex",
help="同步笔记并重建索引",
parents=[config_parent],
)
reindex_parser.set_defaults(func=_reindex)
ask_parser = subparsers.add_parser(
"ask",
help="检索笔记并调用 DeepSeek 回答",
parents=[config_parent],
)
_add_retrieval_args(ask_parser)
ask_parser.add_argument("--no-source", action="store_true", help="不显示来源")
_add_question_arg(ask_parser)
ask_parser.set_defaults(func=_ask)
search_parser = subparsers.add_parser(
"search",
help="只检索来源，不调用大模型",
parents=[config_parent],
)
_add_retrieval_args(search_parser)
search_parser.add_argument(
"--show-content",
action="store_true",
help="显示命中的内容预览",
)
search_parser.add_argument(
"--preview-chars",
type=_positive_int,
default=260,
help="内容预览字符数，默认 260",
)
_add_question_arg(search_parser)
search_parser.set_defaults(func=_search)
context_parser = subparsers.add_parser(
"context",
help="返回可直接注入大模型的检索上下文",
parents=[config_parent],
)
_add_retrieval_args(context_parser)
_add_question_arg(context_parser)
context_parser.set_defaults(func=_context)
read_parser = subparsers.add_parser(
"read",
help="按 search 返回的 source 读取完整笔记",
parents=[config_parent],
)
read_parser.add_argument("source", help="相对 NOTES_DIR 的 Markdown 路径")
read_parser.add_argument(
"--max-chars",
type=_positive_int,
help="最多输出多少字符，默认完整输出",
)
read_parser.set_defaults(func=_read)
status_parser = subparsers.add_parser(
"status",
help="查看配置和索引状态",
parents=[config_parent],
)
status_parser.set_defaults(func=_status)
return parser
def main(argv: list[str] | None = None) -> int:
parser = build_parser()
args = parser.parse_args(argv)
if not hasattr(args, "func"):
parser.print_help()
return 0
return int(args.func(args) or 0)
def sync_notes() -> None:
parser = argparse.ArgumentParser(description="同步萌芽 RAG 笔记")
_add_config_args(parser)
_add_output_args(parser)
args = parser.parse_args()
_sync_notes(args)
def build_index() -> None:
parser = argparse.ArgumentParser(description="构建萌芽 RAG 索引")
_add_config_args(parser)
_add_output_args(parser)
args = parser.parse_args()
_build_index(args)
def ask() -> None:
parser = argparse.ArgumentParser(description="向萌芽 RAG 知识库提问")
_add_config_args(parser)
_add_output_args(parser)
_add_retrieval_args(parser)
parser.add_argument("--no-source", action="store_true", help="不显示来源")
_add_question_arg(parser)
args = parser.parse_args()
_ask(args)
if __name__ == "__main__":
raise SystemExit(main())

71
src/mengya_rag/config.py Normal file
View File

@@ -0,0 +1,71 @@
from __future__ import annotations
import os
from dataclasses import replace
from dataclasses import dataclass
from pathlib import Path
from dotenv import load_dotenv
PROJECT_ROOT = Path(__file__).resolve().parents[2]
def _path_from_env(name: str, default: str) -> Path:
value = os.getenv(name, default)
path = Path(value)
if not path.is_absolute():
path = PROJECT_ROOT / path
return path
@dataclass(frozen=True)
class Settings:
notes_remote: str
notes_dir: Path
index_dir: Path
deepseek_api_key: str
deepseek_model: str
deepseek_temperature: float
top_k: int
list_top_k: int
embed_backend: str
embed_model: str
embed_batch_size: int
embed_query_prompt_name: str
vector_weight: float
bm25_weight: float
def load_settings(env_file: str | Path | None = None) -> Settings:
if env_file is None:
load_dotenv(PROJECT_ROOT / ".env")
else:
load_dotenv(env_file)
return Settings(
notes_remote=os.getenv(
"NOTES_REMOTE",
"bigmengya:/shumengya/docker/mengyanote-backend/data/mengyanote/",
),
notes_dir=_path_from_env("NOTES_DIR", "data/notes"),
index_dir=_path_from_env("INDEX_DIR", "data/index/bm25"),
deepseek_api_key=os.getenv("DEEPSEEK_API_KEY", ""),
deepseek_model=os.getenv("DEEPSEEK_MODEL", "deepseek-chat"),
deepseek_temperature=float(os.getenv("DEEPSEEK_TEMPERATURE", "0.2")),
top_k=int(os.getenv("TOP_K", "6")),
list_top_k=int(os.getenv("LIST_TOP_K", "20")),
embed_backend=os.getenv("EMBED_BACKEND", "fastembed"),
embed_model=os.getenv("EMBED_MODEL", "BAAI/bge-small-zh-v1.5"),
embed_batch_size=int(os.getenv("EMBED_BATCH_SIZE", "32")),
embed_query_prompt_name=os.getenv("EMBED_QUERY_PROMPT_NAME", ""),
vector_weight=float(os.getenv("VECTOR_WEIGHT", "0.65")),
bm25_weight=float(os.getenv("BM25_WEIGHT", "0.35")),
)
def override_settings(settings: Settings, **values: object) -> Settings:
overrides = {key: value for key, value in values.items() if value is not None}
if not overrides:
return settings
return replace(settings, **overrides)

78
src/mengya_rag/embeddings.py Normal file
View File

@@ -0,0 +1,78 @@
from __future__ import annotations
from functools import lru_cache
import numpy as np
from fastembed import TextEmbedding
from .config import Settings
QWEN_QUERY_INSTRUCTION = (
"Instruct: Given a search query in Chinese or English, retrieve relevant "
"Markdown note passages from a personal knowledge base.\nQuery: "
)
@lru_cache(maxsize=2)
def _fastembed_model(model_name: str) -> TextEmbedding:
return TextEmbedding(model_name=model_name)
@lru_cache(maxsize=2)
def _sentence_transformer_model(model_name: str):
from sentence_transformers import SentenceTransformer
return SentenceTransformer(model_name)
def embed_documents(settings: Settings, texts: list[str]) -> np.ndarray:
if settings.embed_backend == "sentence-transformers":
model = _sentence_transformer_model(settings.embed_model)
vectors = model.encode(
texts,
batch_size=settings.embed_batch_size,
normalize_embeddings=True,
show_progress_bar=False,
)
return np.asarray(vectors, dtype=np.float32)
model = _fastembed_model(settings.embed_model)
vectors = np.array(
list(model.embed(texts, batch_size=settings.embed_batch_size)),
dtype=np.float32,
)
return normalize_vectors(vectors)
def embed_query(settings: Settings, query: str) -> np.ndarray:
if settings.embed_backend == "sentence-transformers":
model = _sentence_transformer_model(settings.embed_model)
prompt_name = settings.embed_query_prompt_name or None
if prompt_name:
vector = model.encode(
[query],
prompt_name=prompt_name,
normalize_embeddings=True,
show_progress_bar=False,
)[0]
else:
vector = model.encode(
[query],
normalize_embeddings=True,
show_progress_bar=False,
)[0]
return np.asarray(vector, dtype=np.float32)
model = _fastembed_model(settings.embed_model)
vector = np.array(next(model.query_embed(query)), dtype=np.float32)
return normalize_vector(vector)
def normalize_vectors(vectors: np.ndarray) -> np.ndarray:
norms = np.linalg.norm(vectors, axis=1, keepdims=True)
return vectors / np.clip(norms, 1e-12, None)
def normalize_vector(vector: np.ndarray) -> np.ndarray:
return vector / max(float(np.linalg.norm(vector)), 1e-12)

116
src/mengya_rag/indexing.py Normal file
View File

@@ -0,0 +1,116 @@
from __future__ import annotations
import re
import shutil
import subprocess
from pathlib import Path
import numpy as np
from llama_index.core.schema import TextNode
from .config import Settings
from .embeddings import embed_documents
from .markdown_chunker import chunk_markdown_file
from .vector_store import db_path, load_nodes_from_db, write_vector_db
def sync_notes(settings: Settings) -> None:
settings.notes_dir.mkdir(parents=True, exist_ok=True)
subprocess.run(
[
"rsync",
"-av",
"--delete",
"--exclude=.git/",
"--exclude=.trash/",
"--exclude=.obsidian/workspace*",
settings.notes_remote,
f"{settings.notes_dir}/",
],
check=True,
)
def markdown_files(notes_dir: Path) -> list[Path]:
if not notes_dir.exists():
raise FileNotFoundError(f"笔记目录不存在: {notes_dir}")
return sorted(notes_dir.rglob("*.md"))
def chinese_tokenizer(text: str) -> list[str]:
stopwords = {
"请",
"根据",
"笔记",
"简要",
"回答",
"哪些",
"什么",
"内容",
"一下",
"我的",
"有哪些",
"是什么",
}
tokens: list[str] = []
for part in re.findall(r"[A-Za-z0-9_#+.-]+|[\u4e00-\u9fff]+", text.lower()):
if re.fullmatch(r"[A-Za-z0-9_#+.-]+", part):
if part in {"md", "markdown"}:
continue
if part not in stopwords:
tokens.append(part)
continue
if len(part) == 1:
if part not in stopwords:
tokens.append(part)
continue
max_n = min(4, len(part))
for n in range(2, max_n + 1):
for index in range(0, len(part) - n + 1):
token = part[index : index + n]
if token not in stopwords:
tokens.append(token)
return tokens
def build_index(settings: Settings) -> int:
nodes: list[TextNode] = []
for path in markdown_files(settings.notes_dir):
nodes.extend(chunk_markdown_file(path, settings.notes_dir))
if settings.index_dir.exists():
shutil.rmtree(settings.index_dir)
settings.index_dir.mkdir(parents=True, exist_ok=True)
texts = [node_to_embedding_text(node) for node in nodes]
vectors = embed_documents(settings, texts)
write_vector_db(db_path(settings.index_dir), nodes, vectors)
return len(nodes)
def load_index_nodes(settings: Settings) -> list[TextNode]:
database = db_path(settings.index_dir)
if not database.exists():
raise FileNotFoundError(
f"索引不存在: {database}，请先运行 mengya-build-index"
)
return load_nodes_from_db(database)
def node_to_embedding_text(node: TextNode) -> str:
title = str(node.metadata.get("title", ""))
rel_path = str(node.metadata.get("rel_path", ""))
heading_path = str(node.metadata.get("heading_path", ""))
tags = node.metadata.get("tags", [])
return (
f"文件名: {title}\n"
f"路径: {rel_path}\n"
f"标题路径: {heading_path}\n"
f"标签: {tags}\n"
f"正文:\n{node.get_content()}"
)

285
src/mengya_rag/markdown_chunker.py Normal file
View File

@@ -0,0 +1,285 @@
from __future__ import annotations
import re
from dataclasses import dataclass, field
from pathlib import Path
from llama_index.core.schema import TextNode
TARGET_TOKENS = 512
MAX_TOKENS = 1024
OVERLAP_TOKENS = 96
@dataclass
class MarkdownBlock:
kind: str
text: str
level: int | None = None
title: str = ""
@dataclass
class Section:
headings: dict[int, str]
blocks: list[MarkdownBlock] = field(default_factory=list)
def extract_frontmatter(text: str) -> tuple[dict[str, object], str]:
if not text.startswith("---\n"):
return {}, text
end = text.find("\n---", 4)
if end == -1:
return {}, text
raw = text[4:end].strip()
body = text[end + len("\n---") :].lstrip("\n")
metadata: dict[str, object] = {}
lines = raw.splitlines()
index = 0
while index < len(lines):
line = lines[index]
if ":" not in line:
index += 1
continue
key, value = line.split(":", 1)
key = key.strip()
value = value.strip()
if value:
if value.startswith("[") and value.endswith("]"):
metadata[key] = [
item.strip().strip("'\"")
for item in value[1:-1].split(",")
if item.strip()
]
else:
metadata[key] = value.strip("'\"")
index += 1
continue
items: list[str] = []
index += 1
while index < len(lines) and lines[index].lstrip().startswith("- "):
items.append(lines[index].split("- ", 1)[1].strip().strip("'\""))
index += 1
metadata[key] = items
return metadata, body
def parse_markdown_blocks(text: str) -> list[MarkdownBlock]:
lines = text.splitlines()
blocks: list[MarkdownBlock] = []
paragraph: list[str] = []
index = 0
def flush_paragraph() -> None:
nonlocal paragraph
if paragraph:
blocks.append(MarkdownBlock(kind="text", text="\n".join(paragraph).strip()))
paragraph = []
while index < len(lines):
line = lines[index]
stripped = line.strip()
if stripped.startswith("```") or stripped.startswith("~~~"):
flush_paragraph()
fence = stripped[:3]
code_lines = [line]
index += 1
while index < len(lines):
code_lines.append(lines[index])
if lines[index].strip().startswith(fence):
index += 1
break
index += 1
blocks.append(MarkdownBlock(kind="code", text="\n".join(code_lines)))
continue
heading = re.match(r"^(#{1,6})\s+(.+?)\s*$", line)
if heading:
flush_paragraph()
level = len(heading.group(1))
title = heading.group(2).strip()
blocks.append(
MarkdownBlock(kind="heading", text=line.strip(), level=level, title=title)
)
index += 1
continue
if stripped.startswith("|") and "|" in stripped[1:]:
flush_paragraph()
table_lines = [line]
index += 1
while index < len(lines):
next_line = lines[index]
if not next_line.strip().startswith("|"):
break
table_lines.append(next_line)
index += 1
blocks.append(MarkdownBlock(kind="table", text="\n".join(table_lines)))
continue
if not stripped:
flush_paragraph()
index += 1
continue
paragraph.append(line)
index += 1
flush_paragraph()
return blocks
def build_sections(blocks: list[MarkdownBlock]) -> list[Section]:
sections: list[Section] = []
headings: dict[int, str] = {}
current = Section(headings={})
for block in blocks:
if block.kind != "heading":
current.blocks.append(block)
continue
level = block.level or 1
for old_level in list(headings):
if old_level >= level:
del headings[old_level]
headings[level] = block.title
current_headings = dict(headings)
if level <= 2 or token_count(current.blocks) >= TARGET_TOKENS:
if current.blocks:
sections.append(current)
current = Section(headings=current_headings, blocks=[block])
continue
current.blocks.append(block)
current.headings = current_headings
if current.blocks:
sections.append(current)
return sections
def token_count(blocks: list[MarkdownBlock]) -> int:
text = "\n".join(block.text for block in blocks)
return estimate_tokens(text)
def estimate_tokens(text: str) -> int:
chinese = len(re.findall(r"[\u4e00-\u9fff]", text))
latin = len(re.findall(r"[A-Za-z0-9_#+.-]+", text))
other = max(len(text) - chinese, 0) // 6
return chinese + latin + other
def split_section_blocks(section: Section) -> list[list[MarkdownBlock]]:
chunks: list[list[MarkdownBlock]] = []
current: list[MarkdownBlock] = []
for block in section.blocks:
block_tokens = estimate_tokens(block.text)
current_tokens = token_count(current)
if block.kind in {"code", "table"}:
if current and current_tokens + block_tokens > MAX_TOKENS:
chunks.append(current)
current = []
current.append(block)
continue
if current and current_tokens + block_tokens > MAX_TOKENS:
chunks.append(current)
current = overlap_tail(current)
current.append(block)
if current:
chunks.append(current)
return chunks
def overlap_tail(blocks: list[MarkdownBlock]) -> list[MarkdownBlock]:
kept: list[MarkdownBlock] = []
total = 0
for block in reversed(blocks):
if block.kind in {"code", "table"}:
continue
count = estimate_tokens(block.text)
if kept and total + count > OVERLAP_TOKENS:
break
kept.insert(0, block)
total += count
return kept
def chunk_markdown_file(path: Path, notes_dir: Path) -> list[TextNode]:
text = path.read_text(encoding="utf-8", errors="ignore")
frontmatter, body = extract_frontmatter(text)
blocks = parse_markdown_blocks(body)
sections = build_sections(blocks)
rel_path = path.relative_to(notes_dir)
title = path.stem
mtime = path.stat().st_mtime
tags = frontmatter.get("tags", [])
if isinstance(tags, str):
tags = [tags]
nodes: list[TextNode] = []
chunk_index = 0
for section in sections:
heading_path = heading_path_from_levels(section.headings)
prefix = build_context_prefix(title, rel_path, heading_path)
for chunk_blocks in split_section_blocks(section):
chunk_text = "\n\n".join(block.text for block in chunk_blocks).strip()
if not chunk_text:
continue
node_text = f"{prefix}\n\n{chunk_text}".strip()
nodes.append(
TextNode(
text=node_text,
metadata={
"source_file": str(rel_path),
"rel_path": str(rel_path),
"file_name": path.name,
"folder_path": str(rel_path.parent),
"title": title,
"heading_path": " > ".join(heading_path),
"h1": section.headings.get(1, ""),
"h2": section.headings.get(2, ""),
"h3": section.headings.get(3, ""),
"h4": section.headings.get(4, ""),
"h5": section.headings.get(5, ""),
"h6": section.headings.get(6, ""),
"chunk_index": chunk_index,
"tags": tags,
"created_at": mtime,
},
)
)
chunk_index += 1
return nodes
def heading_path_from_levels(headings: dict[int, str]) -> list[str]:
return [headings[level] for level in sorted(headings)]
def build_context_prefix(title: str, rel_path: Path, heading_path: list[str]) -> str:
lines = [
f"文件: {rel_path}",
f"标题: {title}",
]
if heading_path:
lines.append(f"标题路径: {' > '.join(heading_path)}")
return "\n".join(lines)

86
src/mengya_rag/qa.py Normal file
View File

@@ -0,0 +1,86 @@
from __future__ import annotations
from .config import Settings
from .indexing import load_index_nodes
from .retrieval import HybridRetriever, InventoryRetriever, RetrievedNode, is_inventory_query
QA_PROMPT = """你是萌芽 RAG 知识库助手。请只根据下面检索到的笔记内容回答问题。
如果笔记内容不足以回答，请直接说明没有在笔记中找到足够依据。
回答要准确、完整，优先使用中文。
如果用户在问“有哪些/列表/目录/相关笔记/查一下”，请优先按来源文件列出条目，不要只总结一句话。
笔记内容：
---------------------
{context}
---------------------
问题：{question}
回答："""
def format_source(item: RetrievedNode) -> str:
metadata = item.node.metadata
rel_path = metadata.get("rel_path") or metadata.get("file_path") or "未知来源"
return f"{rel_path} score={item.score:.4f}"
def format_context(items: list[RetrievedNode]) -> str:
parts: list[str] = []
for index, item in enumerate(items, start=1):
rel_path = item.node.metadata.get("rel_path", "未知来源")
parts.append(
f"[{index}] 来源: {rel_path}\n"
f"{item.node.get_content().strip()}"
)
return "\n\n".join(parts)
def retrieve_nodes(
settings: Settings,
question: str,
*,
mode: str = "auto",
top_k: int | None = None,
) -> list[RetrievedNode]:
if mode not in {"auto", "hybrid", "inventory"}:
raise ValueError("mode 只能是 auto、hybrid 或 inventory")
index_nodes = load_index_nodes(settings)
use_inventory = mode == "inventory" or (
mode == "auto" and is_inventory_query(question)
)
limit = top_k or (settings.list_top_k if use_inventory else settings.top_k)
if use_inventory:
retriever = InventoryRetriever(index_nodes)
return retriever.search(question, limit)
retriever = HybridRetriever(settings, index_nodes)
return retriever.search(question, limit)
def ask_question(
settings: Settings,
question: str,
*,
mode: str = "auto",
top_k: int | None = None,
) -> tuple[str, list[str]]:
from llama_index.llms.deepseek import DeepSeek
if not settings.deepseek_api_key:
raise RuntimeError("缺少 DEEPSEEK_API_KEY，请先复制 .env.example 为 .env 并填写。")
nodes = retrieve_nodes(settings, question, mode=mode, top_k=top_k)
llm = DeepSeek(
model=settings.deepseek_model,
api_key=settings.deepseek_api_key,
temperature=settings.deepseek_temperature,
)
prompt = QA_PROMPT.format(context=format_context(nodes), question=question)
response = llm.complete(prompt)
sources = [format_source(node) for node in nodes]
return str(response), sources

359
src/mengya_rag/retrieval.py Normal file
View File

@@ -0,0 +1,359 @@
from __future__ import annotations
import math
import re
from collections import Counter
from dataclasses import dataclass
import numpy as np
from llama_index.core.schema import TextNode
from .config import Settings
from .embeddings import embed_query
from .indexing import chinese_tokenizer
from .vector_store import db_path, search_vectors
@dataclass(frozen=True)
class RetrievedNode:
node: TextNode
score: float
raw_score: float = 0.0
def clean_query(query: str) -> str:
query = query.strip()
for phrase in [
"请根据笔记简要回答",
"请根据笔记回答",
"根据笔记简要回答",
"根据笔记回答",
"简要回答",
"有哪些",
"是什么",
]:
query = query.replace(phrase, " ")
for word in [
"我的",
"我",
"内容",
"请",
"根据",
"简要",
"回答",
"哪些",
"什么",
"一下",
]:
query = query.replace(word, " ")
query = re.sub(r"[，。！？、；：,.!?;:]+", " ", query)
return " ".join(query.split())
def inventory_keywords(query: str) -> list[str]:
stopwords = {
"查一下",
"目前",
"我的",
"我",
"相关",
"笔记",
"教程",
"有哪些",
"有什么",
"列表",
"清单",
"目录",
"安装",
}
words: list[str] = []
cleaned = clean_query(query).lower()
for part in re.findall(r"[a-z0-9_#+.-]+|[\u4e00-\u9fff]+", cleaned):
if part in stopwords:
continue
if re.fullmatch(r"[a-z0-9_#+.-]+", part):
words.append(part)
continue
for stopword in stopwords:
part = part.replace(stopword, "")
if len(part) >= 2:
words.append(part)
return words
class LocalBM25:
def __init__(self, nodes: list[TextNode]) -> None:
self.nodes = nodes
self.doc_tokens = [chinese_tokenizer(self._weighted_text(node)) for node in nodes]
self.doc_freq: Counter[str] = Counter()
self.term_freqs: list[Counter[str]] = []
for tokens in self.doc_tokens:
term_freq = Counter(tokens)
self.term_freqs.append(term_freq)
self.doc_freq.update(term_freq.keys())
self.doc_count = len(nodes)
self.doc_lengths = [len(tokens) for tokens in self.doc_tokens]
self.avg_doc_length = (
sum(self.doc_lengths) / self.doc_count if self.doc_count else 0.0
)
def _weighted_text(self, node: TextNode) -> str:
title = str(node.metadata.get("title", ""))
rel_path = str(node.metadata.get("rel_path", ""))
path_text = rel_path.replace("/", " ").replace(".md", "")
content = node.get_content()
return "\n".join(
[
" ".join([title] * 8),
" ".join([path_text] * 4),
content,
]
)
def search(self, query: str, top_k: int) -> list[RetrievedNode]:
query_terms = chinese_tokenizer(query)
if not query_terms or not self.nodes:
return []
k1 = 1.5
b = 0.75
scored_by_file: dict[str, RetrievedNode] = {}
extra_scores_by_file: Counter[str] = Counter()
for index, term_freq in enumerate(self.term_freqs):
score = 0.0
doc_length = self.doc_lengths[index] or 1
for term in query_terms:
freq = term_freq.get(term, 0)
if freq == 0:
continue
doc_freq = self.doc_freq.get(term, 0)
idf = math.log(1 + (self.doc_count - doc_freq + 0.5) / (doc_freq + 0.5))
denom = freq + k1 * (
1 - b + b * doc_length / (self.avg_doc_length or 1)
)
score += idf * freq * (k1 + 1) / denom
if score <= 0:
continue
node = self.nodes[index]
rel_path = str(node.metadata.get("rel_path", node.node_id))
current = scored_by_file.get(rel_path)
if current is None or score > current.score:
if current is not None:
extra_scores_by_file[rel_path] += current.score * 0.2
scored_by_file[rel_path] = RetrievedNode(
node=node,
score=score,
raw_score=score,
)
else:
extra_scores_by_file[rel_path] += score * 0.2
scored = [
RetrievedNode(
node=item.node,
score=item.score + extra_scores_by_file[rel_path],
raw_score=item.raw_score,
)
for rel_path, item in scored_by_file.items()
]
scored.sort(key=lambda item: item.score, reverse=True)
return scored[:top_k]
class HybridRetriever:
def __init__(self, settings: Settings, nodes: list[TextNode]) -> None:
self.settings = settings
self.nodes = nodes
self.bm25 = LocalBM25(nodes)
def search(self, query: str, top_k: int) -> list[RetrievedNode]:
query = clean_query(query)
candidate_k = min(max(top_k * 8, 30), len(self.nodes))
bm25_results = self.bm25.search(query, candidate_k)
vector_results = self._vector_search(query, candidate_k)
return self._fuse_results(bm25_results, vector_results, top_k)
def _vector_search(self, query: str, top_k: int) -> list[RetrievedNode]:
query_vector = embed_query(self.settings, query)
rows = search_vectors(db_path(self.settings.index_dir), query_vector, top_k)
results: list[RetrievedNode] = []
for rowid, distance in rows:
index = rowid - 1
if index < 0 or index >= len(self.nodes):
continue
score = 1.0 / (1.0 + distance)
results.append(
RetrievedNode(
node=self.nodes[index],
score=score,
raw_score=score,
)
)
return results
def _fuse_results(
self,
bm25_results: list[RetrievedNode],
vector_results: list[RetrievedNode],
top_k: int,
) -> list[RetrievedNode]:
fused: dict[str, tuple[TextNode, float]] = {}
seen_files: set[str] = set()
def add(results: list[RetrievedNode], weight: float) -> None:
if not results:
return
best_raw = max(item.raw_score or item.score for item in results)
if best_raw <= 0:
return
for rank, item in enumerate(results, start=1):
key = item.node.node_id
raw = item.raw_score or item.score
normalized = raw / best_raw
if normalized < 0.18:
continue
score = weight * normalized / (60 + rank)
node, current = fused.get(key, (item.node, 0.0))
fused[key] = (node, current + score)
add(bm25_results, self.settings.bm25_weight)
add(vector_results, self.settings.vector_weight)
ranked = sorted(fused.values(), key=lambda pair: pair[1], reverse=True)
if not ranked:
return []
best_score = ranked[0][1]
final: list[RetrievedNode] = []
for node, score in ranked:
if score < best_score * 0.7 and len(final) >= 1:
continue
rel_path = str(node.metadata.get("rel_path", node.node_id))
if rel_path in seen_files:
continue
seen_files.add(rel_path)
final.append(RetrievedNode(node=node, score=score))
if len(final) >= top_k:
break
return final
def is_inventory_query(query: str) -> bool:
return any(
word in query
for word in ["有哪些", "有什么", "哪些", "列表", "清单", "目录", "查一下", "相关笔记"]
)
def build_file_nodes(nodes: list[TextNode]) -> list[TextNode]:
by_path: dict[str, list[TextNode]] = {}
for node in nodes:
rel_path = str(node.metadata.get("rel_path", node.node_id))
by_path.setdefault(rel_path, []).append(node)
file_nodes: list[TextNode] = []
for rel_path, file_chunks in by_path.items():
first = file_chunks[0]
title = str(first.metadata.get("title", rel_path.rsplit("/", 1)[-1]))
preview = "\n".join(
chunk.get_content().strip() for chunk in file_chunks[:2] if chunk.get_content().strip()
)
text = f"标题: {title}\n路径: {rel_path}\n内容预览:\n{preview[:1200]}"
file_nodes.append(
TextNode(
text=text,
metadata={
"rel_path": rel_path,
"title": title,
"is_file_summary": True,
},
)
)
return file_nodes
class InventoryRetriever:
def __init__(self, nodes: list[TextNode]) -> None:
self.file_nodes = build_file_nodes(nodes)
self.bm25 = LocalBM25(self.file_nodes)
def search(self, query: str, top_k: int) -> list[RetrievedNode]:
query = clean_query(query)
direct = self._direct_directory_results(query, top_k)
if direct:
return direct
results = self.bm25.search(query, top_k)
results = self._filter_by_keywords(query, results)
if not results:
return []
best = results[0].score
kept = [item for item in results if item.score >= best * 0.4]
return kept[:top_k]
def _filter_by_keywords(
self,
query: str,
results: list[RetrievedNode],
) -> list[RetrievedNode]:
keywords = inventory_keywords(query)
if not keywords:
return results
latin_keywords = [
keyword
for keyword in keywords
if re.fullmatch(r"[a-z0-9_#+.-]+", keyword)
]
filtered: list[RetrievedNode] = []
for item in results:
haystack = (
f"{item.node.metadata.get('title', '')} "
f"{item.node.metadata.get('rel_path', '')} "
f"{item.node.get_content()}"
).lower()
if latin_keywords and not all(keyword in haystack for keyword in latin_keywords):
continue
matched = sum(1 for keyword in keywords if keyword in haystack)
if matched == 0:
continue
if len(keywords) >= 2 and matched / len(keywords) < 0.5:
continue
filtered.append(item)
return filtered or results[: min(3, len(results))]
def _direct_directory_results(self, query: str, top_k: int) -> list[RetrievedNode]:
candidates = [
node
for node in self.file_nodes
if str(node.metadata.get("rel_path", "")).startswith(f"{query}/")
]
if not candidates and query.endswith("文章"):
stripped = query[: -len("文章")]
candidates = [
node
for node in self.file_nodes
if str(node.metadata.get("rel_path", "")).startswith(f"{stripped}/")
]
if not candidates:
return []
candidates.sort(key=lambda node: str(node.metadata.get("rel_path", "")))
return [
RetrievedNode(node=node, score=float(top_k - index), raw_score=float(top_k - index))
for index, node in enumerate(candidates[:top_k])
]

117
src/mengya_rag/vector_store.py Normal file
View File

@@ -0,0 +1,117 @@
from __future__ import annotations
import sqlite3
from pathlib import Path
import numpy as np
import sqlite_vec
from llama_index.core.schema import TextNode
DB_FILE = "rag.sqlite3"
def db_path(index_dir: Path) -> Path:
return index_dir / DB_FILE
def connect_vector_db(path: Path) -> sqlite3.Connection:
conn = sqlite3.connect(path)
conn.enable_load_extension(True)
sqlite_vec.load(conn)
conn.enable_load_extension(False)
conn.row_factory = sqlite3.Row
return conn
def init_vector_db(path: Path, dim: int) -> None:
conn = connect_vector_db(path)
try:
conn.execute("drop table if exists chunks")
conn.execute("drop table if exists chunk_vectors")
conn.execute(
"""
create table chunks (
id integer primary key,
node_json text not null,
rel_path text not null,
title text not null,
heading_path text not null,
chunk_index integer not null
)
"""
)
conn.execute(
f"create virtual table chunk_vectors using vec0(embedding float[{dim}])"
)
conn.execute("create index idx_chunks_rel_path on chunks(rel_path)")
conn.commit()
finally:
conn.close()
def write_vector_db(path: Path, nodes: list[TextNode], vectors: np.ndarray) -> None:
if len(nodes) != len(vectors):
raise ValueError("节点数量和向量数量不一致")
init_vector_db(path, int(vectors.shape[1]))
conn = connect_vector_db(path)
try:
chunk_rows = []
vector_rows = []
for rowid, (node, vector) in enumerate(zip(nodes, vectors), start=1):
chunk_rows.append(
(
rowid,
node.model_dump_json(),
str(node.metadata.get("rel_path", "")),
str(node.metadata.get("title", "")),
str(node.metadata.get("heading_path", "")),
int(node.metadata.get("chunk_index", 0)),
)
)
vector_rows.append((rowid, np.asarray(vector, dtype=np.float32).tobytes()))
conn.executemany(
"""
insert into chunks
(id, node_json, rel_path, title, heading_path, chunk_index)
values (?, ?, ?, ?, ?, ?)
""",
chunk_rows,
)
conn.executemany(
"insert into chunk_vectors(rowid, embedding) values (?, ?)",
vector_rows,
)
conn.commit()
finally:
conn.close()
def load_nodes_from_db(path: Path) -> list[TextNode]:
conn = connect_vector_db(path)
try:
rows = conn.execute("select node_json from chunks order by id").fetchall()
return [TextNode.model_validate_json(row["node_json"]) for row in rows]
finally:
conn.close()
def search_vectors(path: Path, query_vector: np.ndarray, top_k: int) -> list[tuple[int, float]]:
conn = connect_vector_db(path)
try:
vector_blob = np.asarray(query_vector, dtype=np.float32).tobytes()
rows = conn.execute(
"""
select rowid, distance
from chunk_vectors
where embedding match ?
order by distance
limit ?
""",
(vector_blob, top_k),
).fetchall()
return [(int(row["rowid"]), float(row["distance"])) for row in rows]
finally:
conn.close()

2489
uv.lock generated Normal file
View File

File diff suppressed because it is too large Load Diff