# Context Retrieval System Powered by RAG and MCP Server

A hybrid context-retrieval system built on the Model Context Protocol (MCP) and Retrieval-Augmented Generation (RAG). It exposes multimodal PDF parsing and a fully local vector-search pipeline as MCP tools, so AI agents (Claude Desktop, Cursor, or any MCP-compatible client) can fetch real-time document context on demand instead of relying on pre-loaded, static knowledge.

---

## Highlights


- **MCP-native** — 7 JSON-in/JSON-out tools served over stdio (local) or streamable-HTTP (network), built on the official `mcp` Python SDK.
- **Multimodal parsing** — PyMuPDF extracts structured text and base64-encoded images; Camelot (lattice/stream, with a PyMuPDF fallback) extracts complex tables straight into LLM context windows.
- **Local RAG pipeline** — LangChain `RecursiveCharacterTextSplitter` chunking, on-device fastembed embeddings (BGE-small, ONNX), and a FAISS vector store. No external embeddings API, no per-query cost, no data leaving the machine.
- **Sub-linear search** — cosine similarity over L2-normalised vectors; above a configurable threshold the store switches to an IVF (inverted-file) index, keeping query latency approximately O(log N) as the corpus grows.
- **Persistent & re-entrant** — the FAISS index, raw vectors and chunk metadata persist to disk; re-indexing a changed document replaces its chunks (content-hash doc IDs).
- **Safe by default** — every tool path is validated (null bytes, existence, extension, optional allow-list roots) and returns structured JSON errors instead of crashing the server.
- **Containerized + CI** — Dockerfile (with the embedding model pre-baked) and GitHub Actions running lint, the full test suite including end-to-end MCP protocol tests, and a containerized smoke test.

---

## Setup

```bash
git clone https://github.com/ESPChong/context-retrieval-system-RAG-MCP.git
cd context-retrieval-system-RAG-MCP

python3.12 -m venv .venv
.venv/bin/pip install -r requirements.txt
.venv/bin/pip install -e . --no-deps      # optional: `python -m context_retrieval` from anywhere

# stdio mode is what Claude Desktop / Cursor launch:
.venv/bin/python -m context_retrieval
```

The embedding model (~35 MB) downloads from Hugging Face once and is cached locally; after that the system runs fully offline.

Try the pipeline without any MCP client:

```bash
.venv/bin/python scripts/demo.py
```

---

## Tool reference

| Tool | Purpose | Key parameters |
|---|---|---|
| `extract_document_text` | Per-page structured text + PDF metadata | `pdf_path`, `pages="1-3"` |
| `extract_document_images` | Embedded images as base64 PNG + geometry | `pdf_path`, `include_data=false` to inventory only, `max_images`, `pages` |
| `extract_document_tables` | Tables via Camelot lattice/stream, PyMuPDF fallback; `auto` tries all | `pdf_path`, `pages`, `flavor` |
| `index_document` | Chunk → embed (local) → FAISS, persists the store | `pdf_path`, `chunk_size`, `chunk_overlap`, `replace` |
| `search_context` | Top-k semantic search with scores + source/page provenance | `query`, `k`, `source_filter` |
| `index_info` | Store status: docs, chunks, dimension, index type | — |
| `reset_index` | Drop all indexed documents | — |

---

## Connect an MCP client

<details>
<summary><b>Claude Desktop</b> — Settings → Developer → Edit Config</summary>

```json
{
  "mcpServers": {
    "context-retrieval": {
      "command": "/absolute/path/to/context-retrieval-system-RAG-MCP/.venv/bin/python",
      "args": ["-m", "context_retrieval"],
      "env": {
        "CONTEXT_RETRIEVAL_DATA_DIR": "/absolute/path/to/context-retrieval-system-RAG-MCP/data",
        "CONTEXT_RETRIEVAL_ALLOWED_ROOTS": "/absolute/path/to/your/pdfs"
      }
    }
  }
}
```
(ready-to-edit copy in [`docs/claude-desktop-config.json`](docs/claude-desktop-config.json))
</details>

<details>
<summary><b>Cursor</b> — Settings → MCP → Add new global server</summary>

Same `mcpServers` object; copy in [`docs/cursor-mcp.json`](docs/cursor-mcp.json).
</details>

<details>
<summary><b>Docker</b></summary>

```bash
# stdio mode (pipe JSON-RPC, e.g. from an MCP host):
docker build -t context-retrieval .
docker run -i --rm -v "$PWD/docs:/docs" context-retrieval

# network mode (streamable-http on http://localhost:8000/mcp):
docker compose up --build
```
</details>

Once connected, just ask: *"Index /docs/report.pdf and tell me what it says about retention policy."* — the client will chain `index_document` → `search_context` (or the extraction tools) automatically.

---

## Configuration (env vars)

| Variable | Default | Meaning |
|---|---|---|
| `CONTEXT_RETRIEVAL_DATA_DIR` | `~/.context_retrieval` | Where `index.faiss` / `vectors.npy` / `meta.json` persist |
| `CONTEXT_RETRIEVAL_EMBED_MODEL` | `BAAI/bge-small-en-v1.5` | Local fastembed model (384-dim) |
| `CONTEXT_RETRIEVAL_CHUNK_SIZE` / `_OVERLAP` | `1000` / `150` | LangChain splitter parameters |
| `CONTEXT_RETRIEVAL_IVF_THRESHOLD` | `256` | Vectors above this switch Flat → IVF index |
| `CONTEXT_RETRIEVAL_IVF_NPROBE` | `16` | IVF clusters probed per query |
| `CONTEXT_RETRIEVAL_ALLOWED_ROOTS` | *(unrestricted)* | Comma-separated directory allow-list for tool paths |
| `CONTEXT_RETRIEVAL_MAX_*` | images 10 / 400 KB / tables 30 / rows 200 / pages 500 | Context-window guard rails |

---

## Testing

```bash
.venv/bin/python -m pytest -v
```

19 tests, including **end-to-end MCP protocol tests** that spawn the real server over stdio and exercise the full handshake → tool-call → RAG-roundtrip path. A deterministic 3-page sample PDF (prose / ruled table / embedded image) is generated by the suite itself.

```bash
.venv/bin/python scripts/make_sample_pdf.py my-sample.pdf   # inspect it yourself
```

---

## Design notes

- **Why local embeddings?** Removing the external embeddings API eliminates per-query cost and network latency, and keeps document content on the machine. fastembed runs quantized-ready ONNX models — a 384-dim BGE-small that embeds ~2,600 chunks/sec on a laptop CPU.
- **What "≈ O(log N)" means here.** Flat exact search scans all N vectors. The IVF index partitions the space into `nlist ≈ 4·√N` clusters and compares the query against only `nprobe` of them plus their members — sub-linear, logarithmic-style growth in practice, at the cost of a small recall trade-off tunable via `nprobe`.
- **Chunk provenance.** Chunking runs per page, so every retrieved fragment carries `page` + `source`, letting the host model cite where an answer came from — the antidote to hallucinated citations.
- **Graceful degradation.** Camelot's lattice mode needs Ghostscript; when it is missing (thin environments), `auto` silently falls back to stream, then to PyMuPDF's detector, and table extraction still succeeds.