AI-RAG-embed

Preview

Local RAG: CLI index build & retrieval, and retrieval output examples.

Overview

Drop your .txt/.md/.typ files into documents/, embed them locally to build the index, then ask questions and get answers from a remote LLM based on your content.

Quick Start

1. Install dependencies

uv sync

2. Configure

Copy config_example.json to config.json and fill in your API key.

# bash
cp config_example.json config.json

# pwsh
Copy-Item config_example.json config.json

3. Add documents

Put your .txt / .md / .typ files in documents/. Embedding models are downloaded to local cache on first run.

4. Build and ask

python rag_qa.py --build    # build index (incremental, skips if no changes)
python rag_qa.py            # interactive mode (multi-turn)

Note: The code uses hf-mirror.com as the default HuggingFace endpoint. Do not override HF_ENDPOINT with huggingface.co -- it will cause connection timeouts.

CLI Quick Reference

python rag_qa.py --build              # build index (incremental)
python rag_qa.py --rebuild            # force rebuild index
python rag_qa.py --search "question"  # retrieve only, no answer
python rag_qa.py "question"           # single question
python rag_qa.py                      # interactive mode
python rag_qa.py --help               # show all commands and options

MCP Server

Exposes the RAG system as MCP tools for agent integration (Claude Code, etc.).

python servers/rag_server.py          # stdio transport (for agent connection)

Three tools are registered:

Tool	Description	Parameters
`rag_search`	Retrieve relevant chunks, no LLM generation	`question` (required), `enhance`, `k`
`rag_ask`	Retrieve + generate answer via LLM	`question` (required), `enhance`, `k`
`rag_get_info`	Get RAG system config, indexed documents, and paths	(none)

enhance enables query enhancement (same as --enhance CLI flag). k overrides the retrieval count (default from config). To connect, add the server to your agent's MCP config:

{
  "mcpServers": {
    "rag-qa": {
      "type": "stdio",
      "command": "mcp",
      "args": ["run", "D:\\path\\to\\AI-RAG-embed\\servers\\rag_server.py", "-t", "stdio"],
      "env": {}
    }
  }
}

Web UI

A Flask wrapper that runs rag_qa.py per request and streams the answer to the browser.

uv run python web.py   # open http://localhost:5000

Usage

Build Index

python rag_qa.py --build    # incremental: skip if no files changed
python rag_qa.py --rebuild  # force full re-embed

--build detects file content changes and embedding model changes. Switching models in config.json triggers an automatic full rebuild -- no manual --rebuild needed.

Interactive Mode

Multi-turn conversation with history.

python rag_qa.py

Ask a question. Be specific. /quit or /q to quit.

>>> What is exponential smoothing?
Exponential smoothing is a forecasting method...

>>> How many types are there?
The system understands "types" refers to exponential smoothing types from our previous discussion.

>>> /quit

Single-Question Mode

One-shot, saves result to output/.

python rag_qa.py "What is exponential smoothing?"

# With query enhancement (retrieval-optimized rewriting before searching)
python rag_qa.py --enhance "What is exponential smoothing?"

Use --retrieval_k, --retrieval_distance_threshold, --strict_context, --debug to override config settings temporarily. --debug prints query params, query path, question rewriting, per-chunk scores, and source previews.

Search-Only Mode

Retrieve relevant document chunks without generating an answer.

python rag_qa.py --search "What is exponential smoothing?"

# With query enhancement (retrieval-optimized rewriting before searching)
python rag_qa.py --search --enhance "What is exponential smoothing?"

Returns the top retrieval_k chunks (default: 3) directly to stdout. With --enhance, the query is processed through the enhancer before retrieval, same as cmd_ask. See Query Enhancement for mode differences.

Query Tips

The enhancer rewrites queries for better vector similarity, but it cannot retrieve what the documents do not contain.

Stay within the knowledge base's domain and use terms that appear in the documents.
Provide enough context to disambiguate. "Prediction intervals are too narrow" is ambiguous; "ARIMA prediction intervals are too narrow" is not.
Local mode only translates, so how you phrase the question matters more than in LLM mode.

Configuration

Edit config.json to configure the system. Relative paths (./) are resolved against the project root. "Default" below refers to the code's hardcoded fallback when a key is omitted.

Model Slots

The system has 3 local model slots and 1 remote LLM slot. Each local slot is gated by its own config flag — disabled slots never load.

Slot	Config	Type	Purpose
Embedding	`vector_enabled`	Local	Vectorize documents and queries for similarity search
Enhancer	`query_enhance_enabled`	Local or Remote	Rewrite or translate queries before retrieval
Reranker	`reranker_enabled`	Local	Re-score retrieved chunks with a cross-encoder
LLM	always on	Remote (API)	Generate the final answer

Enhancer in local mode loads a MarianMT model; in LLM mode it calls the remote API (shared with the LLM slot). All three local slots can be active simultaneously. A BM25-only setup loads zero local models.

Document & Indexing

Key	Description
`docs_dir`	Folder containing your document files (including subfolders). Use `.doc_loader_ignore` to exclude files (`.gitignore` syntax).
`docs_lang`	Language of your documents (`"en"`, `"zh"`, etc.). Controls which embedding model is selected and which query prefix is used.
`chunking`	Chunking config object. Mode-specific keys are listed below.
`embedding_model_name`	Embedding model config. Can be a string (single model) or an object mapping `docs_lang` values to model IDs. See below.
`chroma_persist_dir`	Folder where the vector database is saved.

embedding_model_name accepts two formats:

// Simple: one model for all languages
"embedding_model_name": "BAAI/bge-small-zh-v1.5"

// Per-language: auto-selected based on docs_lang
"embedding_model_name": {
    "zh": "BAAI/bge-small-zh-v1.5",
    "en": "mixedbread-ai/mxbai-embed-large-v1"
}

When using the per-language format, changing docs_lang automatically switches the model and query prefix. Model dimension changes are handled automatically -- --build detects the switch and rebuilds the index from scratch. Models are downloaded to local cache on first run: bge-small-zh-v1.5 ~92MB, mxbai-embed-large-v1 ~641MB.

Chunking (`chunking`)

Controls how documents are split into chunks for embedding.

Key	Description
`mode`	`"auto"` (default) = smart splitting. `"fixed"` = fixed-length with separator fallback.

Auto mode (chunking.auto): Heading-aware for .md and .typ files (headings, tables, code blocks are detected; preamble and comments are skipped for .typ). Paragraph-first for .txt files. Code blocks and tables are never split.

Key	Description
`target_chars`	Target chunk size. Atomic units (code blocks, tables) may exceed this. Default: `700`.
`split_at_level`	Heading level to split at (1-6). `2` = split at `#` and `##`. Default: `3`.
`min_chars`	Minimum characters per chunk. Shorter sections are dropped. Default: `100`.
`include_heading`	Prepend section heading to each chunk (as `> heading`). Default: `false`.

Fixed mode (chunking.fixed): Two sub-modes controlled by split_by:

"char" (default): splits at max_chars (hard ceiling), falling back to the best separator by priority (\n\n > \n > punctuation > space). overlap_chars controls adjacent chunk overlap.
"line": splits by line count (max_lines), preserving complete line boundaries. Prefers blank lines (paragraph breaks) as split points. overlap_lines controls adjacent chunk overlap.

Key	Description
`split_by`	`"char"` (default) = split by character count. `"line"` = split by line count.
`char.max_chars`	Hard limit per chunk (char mode). Default: `700`.
`char.overlap_chars`	Overlap between adjacent chunks (char mode). Default: `70`.
`line.max_lines`	Max lines per chunk (line mode). Default: `20`.
`line.overlap_lines`	Overlap in lines between adjacent chunks (line mode). Keep below `max_lines // 2` for best break-point quality. Default: `3`.

Query Enhancement

Rewrites your question before searching to improve retrieval quality. The enhanced output is used only for retrieval -- the answer LLM always receives the original question.

Key	Description
`query_enhance_enabled`	Enable query enhancement. Default: `false`.

Enhancement is configured under retrieval.enhancer. See the Retrieval section below.

Retrieval & Reranking

Key	Description
`retrieval.mode`	Retrieval strategy: `"llm"` = LLM-generated retrieval paragraph. `"local"` = MarianMT translation. `null` = direct retrieval (no enhancement).
`retrieval.k`	Number of chunks to retrieve per query. Default: `3`.
`retrieval.distance_threshold`	Cosine distance threshold for vector retrieval. Chunks with distance above this value are filtered out. Only effective when `vector_enabled` is `true`. Set `null` to disable filtering. Default: `0.3`.
`vector_enabled`	Enable vector (embedding) retrieval. When `false`, the embedding model is not loaded — faster startup, lower memory. Default: `true`.
`bm25_enabled`	Enable BM25 keyword retrieval. Can be used alone or combined with vector search; combined results are merged via Reciprocal Rank Fusion (RRF). Default: `false`.

LLM mode (retrieval.enhancer.llm): Generates a dense retrieval paragraph covering key terms, related concepts, and likely document content. For follow-up questions, rewrites using conversation history first.

Key	Description
`api_base_url`, `api_key`, `model`	LLM API connection settings.
`temperature`	Default: `0.0`.
`thinking_mode`	Default: `false`.

Local mode (retrieval.enhancer.local): Translates the question to docs_lang via MarianMT. No term replacement, no context rewrite.

Key	Description
`query_lang`	Language you ask questions in (e.g., `"zh"`, `"en"`).
`model_name`	HuggingFace model ID (optional). Auto-selects `Helsinki-NLP/opus-mt-{query_lang}-{docs_lang}` if omitted (zh-en ~599MB, en-zh ~301MB).

At least one of vector_enabled or bm25_enabled must be true. The four combinations:

`vector_enabled`	`bm25_enabled`	Behavior
`true`	`true`	Hybrid: vector + BM25, merged via RRF
`true`	`false`	Vector-only retrieval
`false`	`true`	BM25-only retrieval (no embedding model loaded)
`false`	`false`	Error: at least one must be enabled

Reranker (optional): Two-stage retrieval -- coarse retrieval of retrieval_k * 4 candidates (vector or BM25), then a cross-encoder jointly scores each candidate against the query and keeps the top retrieval_k. Improves precision when term matching alone is insufficient.

Key	Description
`reranker_enabled`	Enable cross-encoder reranker. Default: `false`.
`reranker.model_name`	Cross-encoder model ID. Default: `"BAAI/bge-reranker-v2-m3"`.
`reranker.top_k`	Number of chunks after reranking. Default: `null` (uses `retrieval_k`).

First load: Downloads the model (~2.2GB) to the local HuggingFace cache on startup. Per-query cost: Cross-encoder scoring adds ~1-2s on CPU (~1.8s for 80 candidates); scales linearly with candidate count.

Answer Generation (`llm`)

The LLM that generates the final answer using the retrieved chunks. Fields: api_base_url, api_key, model, temperature (Default: 0.3), thinking_mode.

Chat Behavior

Key	Description
`max_history_rounds`	Recent conversation rounds to keep for context. Default: `10`.
`strict_context`	`true` = answer only from retrieved chunks. `false` = LLM may supplement with its own knowledge. Default: `false`.

System Prompt Customization: To customize the system prompt, edit lib/prompt_templates.py directly.

Document Filtering

Place a .doc_loader_ignore file in any subdirectory under documents/. Uses .gitignore syntax. Patterns apply to the containing directory and all subdirectories.

# documents/.doc_loader_ignore
r4ds_textbook/
fpp3_textbook/

# documents/fpp3_textbook/.doc_loader_ignore
README.md
*.log
_draft/

Output

Conversation logs are saved to output/<session>/ as numbered Markdown files.

output/
└── What_is_exponential_smoothing_20260519_173334/
    ├── 01_round.md
    └── 02_round.md

Each round file contains the question, answer, processed question (labeled "Enhanced Question" in LLM mode or "Translated Question" in local mode), and the retrieved chunks used.

Project Structure

rag_qa.py               # entry point (--build | --rebuild | --search | question | interactive)
web.py                  # Flask web UI wrapper around rag_qa.py
config.json             # your configuration (gitignored)
config_example.json     # configuration template
templates/
└── index.html          # web UI template (used by web.py)
documents/              # put your .txt / .md / .typ files here
chroma_db/              # persisted vector database (generated)
output/                 # conversation exports (generated)
docs/
└── ARCHITECTURE.md     # design philosophy, workflows, module internals
tests/
├── test_chunking.py    # pytest: chunking design-property tests
└── thinking_probe.py   # networked API probe (reads OPENAI_API_KEY env)
servers/
└── rag_server.py       # MCP server entry point (stdio transport)
tools/
├── __init__.py         # package marker (all logging via lib/log.py to stderr)
├── shared_store.py     # shared store cache for MCP tools
├── rag_search.py       # MCP tool: retrieve chunks without LLM
├── rag_ask.py          # MCP tool: retrieve + LLM answer
└── rag_get_info.py     # MCP tool: system config and indexed documents
lib/
├── __init__.py         # package marker
├── log.py              # unified logging (all progress/debug/error to stderr)
├── engine.py           # shared init/config logic (CLI + MCP)
├── doc_loader.py       # file I/O + text chunking + ignore patterns
├── embed_engine.py     # embedding model wrapper (sentence-transformers)
├── vector_db.py        # Chroma vector store + hybrid retrieval
├── bm25_retriever.py   # BM25 keyword retriever (jieba + rank-bm25)
├── llm_api.py          # remote LLM API client (OpenAI-compatible)
├── query_enhancer.py   # query enhancement (retrieval-optimized rewriting)
├── local_translator.py # MarianMT local translation backend
├── reranker.py         # cross-encoder reranker for precision re-ranking
└── prompt_templates.py # system prompts + message assembly

Requirements

Dependencies

Python 3.11+ (pinned 3.12.10 via .python-version)
sentence-transformers, chromadb, openai, pathspec (core)
mcp[cli] (MCP server), flask (web UI)
jieba, rank-bm25 (BM25 hybrid retrieval)
transformers, sentencepiece, sacremoses (only needed for mode: "local")

GPU Setup (Optional)

Check if you have an NVIDIA GPU and driver: nvidia-smi

Install CUDA-enabled PyTorch (default pip install torch is CPU-only):

pip install torch --index-url https://download.pytorch.org/whl/cu121

Verify: python -c "import torch; print(torch.cuda.is_available())" -- should print True

Without GPU, everything still works, but much slower.

Model Cache

All local models are downloaded to ~/.cache/huggingface/hub/ on first run, then read from cache. To list cached models and their sizes:

# bash
du -sh ~/.cache/huggingface/hub/models--* | awk '{sub(/.*models--/, "", $2); sub(/--/, "/", $2); print $2": "$1}'

# pwsh
Get-ChildItem "$env:USERPROFILE\.cache\huggingface\hub\models--*" -Directory | ForEach-Object { $name = $_.Name -replace 'models--', '' -replace '--', '/'; $size = (Get-ChildItem $_.FullName -Recurse -File | Measure-Object -Property Length -Sum).Sum; "$name`: $([math]::Round($size/1MB)) MB" }

See docs/ARCHITECTURE.md for build/query workflow, module internals, and environment setup (sentence-transformers version).

Developers/maintainers: see AGENTS.md for working rules, test commands, and the doc map.

License

Released under the MIT License.

Contributing

Personal project; questions and suggestions welcome via Issues.

Maintainer docs map → AGENTS.md.

Name		Name	Last commit message	Last commit date
.agents/notes		.agents/notes
.github/workflows		.github/workflows
docs		docs
documents		documents
lib		lib
screenshots		screenshots
servers		servers
templates		templates
tests		tests
tools		tools
.gitattributes		.gitattributes
.gitignore		.gitignore
.pre-commit-config.yaml		.pre-commit-config.yaml
.python-version		.python-version
AGENTS.md		AGENTS.md
LICENSE		LICENSE
README.md		README.md
README_zh.md		README_zh.md
config_example.json		config_example.json
pyproject.toml		pyproject.toml
rag_qa.py		rag_qa.py
uv.lock		uv.lock
web.py		web.py

Latest commit

History

Folders and files

Repository files navigation

AI-RAG-embed

Preview

Overview

Quick Start

1. Install dependencies

2. Configure

3. Add documents

4. Build and ask

CLI Quick Reference

MCP Server

Web UI

Usage

Build Index

Interactive Mode

Single-Question Mode

Search-Only Mode

Query Tips

Configuration

Model Slots

Document & Indexing

Chunking (chunking)

Query Enhancement

Retrieval & Reranking

Answer Generation (llm)

Chat Behavior

Document Filtering

Output

Project Structure

Requirements

Dependencies

GPU Setup (Optional)

Model Cache

License

Contributing

About

Topics

Resources

Code of conduct

Contributing

Security policy

Stars

Watchers

Forks

Used by

Contributors

Languages

Chunking (`chunking`)

Answer Generation (`llm`)