Visual Flow
QUERY: "authentication code"
STAGE 1: Fast Retrieval
┌──────────────────────────────┐
│ Qwen3-Embedding-0.6B │
│ (bi-encoder) │
│ • 1024 dims │ ┌──────────────────┐
│ • <100ms │─────▶│ ChromaDB │
└──────────────────────────────┘ │ Vector Search │
│ • Cosine │
│ → Top 20 │
└────────┬─────────┘
│
STAGE 2: Precise Reranking ▼
┌────────────────────────────────────────────────────┐
│ ms-marco-MiniLM-L-12-v2 (cross-encoder) │
│ • Full query-document attention │
│ • Processes all 20 candidates │
│ • Returns top N ranked by relevance │
└────────────────────┬───────────────────────────────┘
│
▼
┌──────────────────┐
│ Ranked Results │
│ • File paths │
│ • Code chunks │
│ • Rerank scores │
└──────────────────┘
Overview
MCP server for semantic code search. Two-stage retrieval: Qwen3-Embedding-0.6B gets top-20 candidates via ChromaDB vector search, then cross-encoder/ms-marco-MiniLM-L-12-v2 reranks with full query-document attention for final results.
Bi-encoders are fast but encode query and document independently - no interaction between them. Cross-encoders process both together with transformer attention - more accurate but O(n) complexity makes them too slow for initial retrieval. Split the work: fast model filters to ~20 candidates, slow model ranks those 20. Gets cross-encoder accuracy at bi-encoder speed.
Pre-built index approach. build-index.py reads project paths from projects.txt, recursively scans files while excluding node_modules/venv/.git, chunks into 800-char segments with 100-char overlap, generates embeddings, stores in ChromaDB. Run once. rag-mcp.py loads the database and serves search - no re-indexing on startup. 6 projects, 304 files, 19,600 chunks. Hash IDs include project name to handle duplicate paths across projects.
Additional considerations for production use
Reasoning for model selection, auth, rate limiting, monitoring, incremental re-indexing, middleware.
Requirements | freeze after installing
chromadb
sentence-transformers
fastmcp
torch
rag-mcp.py
import hashlib
import os
import socket
import sys
from pathlib import Path
import chromadb
from chromadb.utils.embedding_functions import SentenceTransformerEmbeddingFunction
from sentence_transformers import CrossEncoder
from fastmcp import FastMCP
chroma_path = os.environ.get("CHROMA_PATH", "./chroma")
client = chromadb.PersistentClient(path=chroma_path)
try:
import torch
device = "cuda" if torch.cuda.is_available() else "cpu"
except Exception:
device = "cpu"
embed = SentenceTransformerEmbeddingFunction(
model_name="Qwen/Qwen3-Embedding-0.6B",
device=device,
)
rerank_model = os.environ.get("RAG_RERANK_MODEL", "cross-encoder/ms-marco-MiniLM-L-12-v2")
reranker = CrossEncoder(rerank_model, device=device)
root = Path(os.environ.get("RAG_ROOT", ".")).resolve()
max_chars = int(os.environ.get("RAG_CHUNK_SIZE", "800"))
overlap = int(os.environ.get("RAG_CHUNK_OVERLAP", "100"))
rerank_candidates = int(os.environ.get("RAG_RERANK_CANDIDATES", "20"))
documents = []
ids = []
metadatas = []
def _chunk_text(text: str, max_chars: int, overlap: int):
if max_chars <= 0:
return
start = 0
text_len = len(text)
while start < text_len:
end = min(start + max_chars, text_len)
yield text[start:end], start, end
if end == text_len:
break
start = max(0, end - overlap)
collection = client.get_or_create_collection(
name="codebase",
embedding_function=embed,
)
existing_count = collection.count()
if existing_count > 0:
print(f"Using existing index with {existing_count} chunks", file=sys.stderr)
else:
print("No index found. Run build-index.py first!", file=sys.stderr)
print(f" RAG_ROOT={root} python build-index.py", file=sys.stderr)
mcp = FastMCP("RAG")
@mcp.tool()
def search(query: str, n: int = 1) -> dict:
"""Search the semantic index and return results."""
query = (query or "").strip()
if not query:
return {"error": "Missing query"}
fetch_k = max(n, rerank_candidates)
result = collection.query(query_texts=[query], n_results=fetch_k)
documents = (result.get("documents") or [[]])[0]
metadatas = (result.get("metadatas") or [[]])[0]
ids = (result.get("ids") or [[]])[0]
if documents:
pairs = [(query, doc) for doc in documents]
scores = reranker.predict(pairs)
ranked = sorted(range(len(documents)), key=lambda i: scores[i], reverse=True)
ranked = ranked[:n]
result = {
"documents": [[documents[i] for i in ranked]],
"metadatas": [[metadatas[i] for i in ranked]],
"ids": [[ids[i] for i in ranked]],
"rerank_scores": [[float(scores[i]) for i in ranked]],
}
return {
"server": {
"host": socket.gethostname(),
"root": str(root),
"collection": "codebase",
"model": embed.model_name,
"reranker": rerank_model,
},
"results": result,
}
if __name__ == "__main__":
mcp.run()
.mcp.json
{
"mcpServers": {
"rag-codebase": {
"type": "stdio",
"command": "/{codebasepath}/mcp-rag-rerank/.venv/bin/python",
"args": [
"/{codebasepath}/mcp-rag-rerank/rag-mcp.py"
],
"env": {
"CHROMA_PATH": "/{codebasepath}/mcp-rag-rerank/chroma"
}
}
}
}
projects.txt
project folders to include (codebase)
path
path
path
Build ChromdaDB chunking from codebase
#!/usr/bin/env python3
"""Build ChromaDB index from codebase."""
import hashlib
import os
import sys
from pathlib import Path
import chromadb
from chromadb.utils.embedding_functions import SentenceTransformerEmbeddingFunction
try:
import torch
device = "cuda" if torch.cuda.is_available() else "cpu"
except Exception:
device = "cpu"
# Configuration
chroma_path = os.environ.get("CHROMA_PATH", "./chroma")
embed_model = os.environ.get("EMBED_MODEL", "Qwen/Qwen3-Embedding-0.6B")
max_chars = int(os.environ.get("RAG_CHUNK_SIZE", "800"))
overlap = int(os.environ.get("RAG_CHUNK_OVERLAP", "100"))
projects_file = Path(__file__).parent / "projects.txt"
# Load project directories from projects.txt or use RAG_ROOT
project_roots = []
if projects_file.exists():
print(f"Loading projects from: {projects_file}")
with open(projects_file) as f:
for line in f:
line = line.strip()
if line and not line.startswith("#"):
project_path = Path(line).expanduser().resolve()
if project_path.exists():
project_roots.append(project_path)
print(f" ✓ {project_path}")
else:
print(f" ✗ Not found: {project_path}")
else:
# Fallback to RAG_ROOT if projects.txt doesn't exist
root = Path(os.environ.get("RAG_ROOT", ".")).resolve()
project_roots = [root]
print(f"Building index for: {root}")
if not project_roots:
print("No valid project directories found!")
sys.exit(1)
print(f"Using device: {device}")
print(f"Chunk size: {max_chars}, overlap: {overlap}")
print(f"Chroma DB: {chroma_path}")
print()
# Initialize embedding function
embed = SentenceTransformerEmbeddingFunction(
model_name=embed_model,
device=device,
)
# Initialize ChromaDB
client = chromadb.PersistentClient(path=chroma_path)
collection = client.get_or_create_collection(
name="codebase",
embedding_function=embed,
)
def _chunk_text(text: str, max_chars: int, overlap: int):
"""Chunk text into overlapping segments."""
if max_chars <= 0:
return
start = 0
text_len = len(text)
while start < text_len:
end = min(start + max_chars, text_len)
yield text[start:end], start, end
if end == text_len:
break
start = max(0, end - overlap)
# Scan and index files
documents = []
ids = []
metadatas = []
file_count = 0
skipped_count = 0
# Directories to exclude (blacklist)
EXCLUDE_DIRS = {
'node_modules', 'venv', 'env', '.venv', '__pycache__',
'dist', 'build', 'target', 'out', '.next', '.nuxt',
'vendor', '.gradle', '.idea', '.vscode',
'coverage', '.pytest_cache', '.mypy_cache', '.tox',
'eggs', '.eggs', 'lib', 'lib64', 'parts', 'sdist',
'var', 'wheels', '*.egg-info', '.installed.cfg', '*.egg',
}
print("Scanning files...")
for root in project_roots:
print(f"\nScanning: {root}")
for file_path in root.rglob("*"):
if not file_path.is_file():
continue
# Skip excluded directories and hidden files (blacklist)
rel_parts = file_path.relative_to(root).parts
if any(part in EXCLUDE_DIRS or (part.startswith(".") and part != ".") for part in rel_parts):
if skipped_count < 10: # Debug: show first 10 skipped
print(f" Skipping: {file_path.relative_to(root)}")
skipped_count += 1
continue
# Skip binary files
try:
with file_path.open("rb") as handle:
if b"\x00" in handle.read(2048):
skipped_count += 1
continue
except Exception:
skipped_count += 1
continue
# Read file content
try:
content = file_path.read_text(encoding="utf-8", errors="ignore")
except Exception:
skipped_count += 1
continue
if not content.strip():
skipped_count += 1
continue
# Chunk and add to index
rel_path = str(file_path.relative_to(root))
project_name = root.name
file_count += 1
if file_count <= 5: # Debug: show first 5 files
print(f" Indexing: {rel_path}")
for i, (chunk, start, end) in enumerate(_chunk_text(content, max_chars, overlap)):
documents.append(chunk)
# Include project name in ID to avoid duplicates
ids.append(hashlib.sha1(f"{project_name}:{rel_path}:{i}".encode("utf-8")).hexdigest())
metadatas.append({"path": rel_path, "project": project_name, "chunk": i, "start": start, "end": end})
if file_count % 100 == 0:
print(f"Processed {file_count} files, {len(documents)} chunks...")
# Safety limit to prevent runaway indexing
if len(documents) > 100000:
print(f"\n⚠ Safety limit reached: {len(documents)} chunks")
print("Consider refining RAG_ROOT or checking for large directories")
break
print(f"\nScanned: {file_count} files")
print(f"Skipped: {skipped_count} files")
print(f"Total chunks: {len(documents)}")
if documents:
print("\nAdding to ChromaDB...")
# Add in batches to avoid max batch size limit
batch_size = 5000
for i in range(0, len(documents), batch_size):
end = min(i + batch_size, len(documents))
print(f" Adding batch {i//batch_size + 1}: chunks {i} to {end}")
collection.add(
documents=documents[i:end],
ids=ids[i:end],
metadatas=metadatas[i:end]
)
print("✓ Index built successfully!")
else:
print("⚠ No documents to index")
sys.exit(1)