RAG Code Search Agent

Python LangGraph ChromaDB Claude API MCP Docker License: MIT

A RAG-based code search agent that indexes repositories and answers natural language queries about codebases. Built with a LangGraph pipeline that rewrites queries, retrieves semantically similar code chunks, re-ranks results, compresses context, and generates precise answers with source citations.

Features

• Natural language code search -- Ask questions like "How does authentication work?" and get answers with file references
• Multi-language support -- Python, JavaScript, TypeScript, Go, Java, Rust, Ruby, C/C++, C#, Swift, Kotlin, Scala, PHP, and more
• Symbol-aware chunking -- Splits code at function/class boundaries to preserve semantic structure
• Incremental indexing -- Only re-indexes changed files using SHA-256 hash tracking
• Dual embedding providers -- Local (sentence-transformers) or OpenAI embeddings
• Cross-encoder re-ranking -- Improves retrieval precision with a second-stage scoring model
• LLM-powered context compression -- Compresses large contexts to fit token limits while preserving signatures
• Heuristic fallbacks -- Gracefully degrades when no API key is configured (abbreviation expansion, similarity-based ranking, rule-based compression)
• MCP server -- Expose search and indexing as tools for AI assistants via stdio or SSE transport
• CLI -- Full command-line interface for indexing, querying, and serving
• Docker-ready -- Containerized deployment with docker-compose

Architecture

The agent is built as a linear LangGraph state machine with five pipeline stages:

┌────────────────┐    ┌───────────┐    ┌──────────┐    ┌────────────────────┐    ┌─────────────────┐
│  Query Rewriter │───>│ Retriever  │───>│ Re-Ranker │───>│ Context Compressor │───>│ Answer Generator │
└────────────────┘    └───────────┘    └──────────┘    └────────────────────┘    └─────────────────┘
       │                   │                │                     │                        │
  Expand abbreviations  Embed query,    Cross-encoder      Compress to fit          Claude Sonnet /
  add synonyms          search ChromaDB  re-score results   token limit              heuristic answer

Each node reads from and writes to a shared AgentState TypedDict, making the pipeline transparent and debuggable.

CLI Commands

Command	Description
rag-search index <repo_path>	Index a repository into the vector store
rag-search index <repo_path> --full	Force full re-index (skip incremental)
rag-search query <question>	Search the indexed codebase with a natural language question
rag-search query <question> --format json	Search and return structured JSON output
rag-search serve	Start the MCP server (stdio transport by default)
rag-search serve --transport sse	Start the MCP server with SSE transport
rag-search serve --transport sse --port 9090	Start SSE server on a custom port
rag-search list	List all indexed repositories

Global option: --db-path overrides the ChromaDB storage path.

MCP Server Tools

Tool	Description
search_codebase	Search the indexed codebase using a natural language query. Returns an answer with source citations and code snippets. Optional repo_path filter.
index_repository	Index a repository into the vector store. Supports incremental indexing to skip unchanged files. Returns indexing statistics.
list_indexed_repos	List all indexed repositories with file counts, chunk counts, and detected languages.

Installation

From Source

git clone https://github.com/your-org/rag-code-search.git
cd rag-code-search
python -m venv .venv
source .venv/bin/activate
pip install -e .

With Dev Dependencies

pip install -e ".[dev]"

Via Docker

docker compose build
docker compose up

The SSE MCP server will be available on http://localhost:8080. The data/ directory is mounted as a volume for persistent storage.

Configuration

Copy .env.example to .env and fill in values:

cp .env.example .env

Environment Variables

Variable	Default	Description
ANTHROPIC_API_KEY	""	API key for Claude (query rewriting, context compression, answer generation)
CHROMA_DB_PATH	data/chroma_db	Path to ChromaDB persistent storage
EMBEDDING_MODEL	sentence-transformers/all-MiniLM-L6-v2	Local embedding model name
EMBEDDING_PROVIDER	local	Embedding provider: local or openai
OPENAI_API_KEY	""	API key for OpenAI embeddings (when provider is openai)
OPENAI_EMBEDDING_MODEL	text-embedding-3-small	OpenAI embedding model name
RERANKER_MODEL	cross-encoder/ms-marco-MiniLM-L-6-v2	Cross-encoder model for re-ranking
CLAUDE_MODEL	claude-sonnet-4-20250514	Claude model for answer generation
CLAUDE_HAIKU_MODEL	claude-haiku-4-20250414	Claude model for query rewriting and context compression
CHUNK_SIZE	500	Maximum tokens per code chunk
CHUNK_OVERLAP	50	Overlap lines between adjacent chunks
RETRIEVAL_TOP_K	20	Number of documents to retrieve from ChromaDB
SIMILARITY_THRESHOLD	0.5	Minimum cosine similarity for retrieval results
RERANK_TOP_K	10	Number of top documents after re-ranking
CONTEXT_TOKEN_LIMIT	4000	Maximum tokens for compressed context
MCP_TRANSPORT	stdio	MCP server transport: stdio or sse
MCP_HOST	127.0.0.1	Host for SSE transport
MCP_PORT	8080	Port for SSE transport

When ANTHROPIC_API_KEY is unset, the agent falls back to heuristic query rewriting (abbreviation expansion), similarity-based re-ranking, rule-based context compression, and structured code snippet answers -- no LLM calls are made.

Usage

Index a Repository

rag-search index /path/to/my-repo

Force a full re-index:

rag-search index /path/to/my-repo --full

Query the Codebase

rag-search query "How does the authentication middleware work?"

JSON output:

rag-search query "Where is the payment processing logic?" --format json

List Indexed Repositories

rag-search list

Start the MCP Server

Stdio transport (for direct process communication):

rag-search serve

SSE transport (for HTTP-based integration):

rag-search serve --transport sse --host 0.0.0.0 --port 8080

How It Works

Chunking

Source files are split into CodeChunk objects using symbol-aware boundary detection. Language-specific regex patterns identify function, class, and type definitions across 13+ languages. Chunks that exceed CHUNK_SIZE tokens are sub-split with overlap (CHUNK_OVERLAP) to preserve context at boundaries.

Embedding

Chunks are embedded using either a local sentence-transformers model or OpenAI's embedding API. Embeddings are L2-normalized for cosine similarity search in ChromaDB.

Retrieval

The rewritten query is embedded and used to search the ChromaDB collection (HNSW index with cosine distance). Results below SIMILARITY_THRESHOLD are filtered out, and the top RETRIEVAL_TOP_K documents are returned.

Re-ranking

Retrieved documents are re-scored using a CrossEncoder model (defaults to cross-encoder/ms-marco-MiniLM-L-6-v2). The cross-encoder evaluates each (query, document) pair and produces a relevance score that is more accurate than embedding similarity alone. If the cross-encoder fails to load, the system falls back to cosine similarity scores.

Context Compression

Ranked documents are concatenated with file path headers. If the total token count exceeds CONTEXT_TOKEN_LIMIT, the compressor uses Claude Haiku to summarize the context while preserving function signatures, class definitions, and import statements. When no API key is available, a heuristic compressor strips boilerplate and truncates long function bodies.

Answer Generation

The compressed context is passed to Claude Sonnet with a system prompt that instructs it to produce a detailed answer with specific file references formatted as `file_path:start_line-end_line (symbol_name)`. Source metadata (file path, line range, language, symbol name, relevance score) is extracted from ranked documents and returned alongside the answer. Without an API key, a heuristic formatter generates a structured code snippet summary.

Project Structure

rag-code-search/
├── .env.example
├── docker-compose.yml
├── Dockerfile
├── pyproject.toml
├── src/
│   └── rag_code_search/
│       ├── __init__.py
│       ├── agent/
│       │   ├── __init__.py
│       │   ├── graph.py          # LangGraph pipeline definition
│       │   ├── nodes.py          # Pipeline node implementations
│       │   └── state.py          # AgentState TypedDict
│       ├── cli/
│       │   ├── __init__.py
│       │   └── main.py           # Click CLI (index, query, serve, list)
│       ├── config/
│       │   ├── __init__.py
│       │   └── settings.py       # Pydantic settings with .env support
│       ├── indexer/
│       │   ├── __init__.py
│       │   ├── chunker.py        # Symbol-aware code chunker
│       │   ├── embedder.py       # Local & OpenAI embedding providers
│       │   └── repository_indexer.py  # Repository walking & incremental indexing
│       ├── mcp_server/
│       │   ├── __init__.py
│       │   └── server.py         # FastMCP server with search/index/list tools
│       └── retrieval/
│           ├── __init__.py
│           ├── context_compressor.py  # LLM & heuristic context compression
│           ├── re_ranker.py      # Cross-encoder & similarity re-ranking
│           └── vector_store.py   # ChromaDB wrapper (add, query, delete)
└── tests/
    ├── __init__.py
    ├── test_chunker.py
    ├── test_context_compressor.py
    ├── test_embedder.py
    ├── test_graph.py
    ├── test_mcp_server.py
    ├── test_re_ranker.py
    └── test_vector_store.py

Testing

pip install -e ".[dev]"
pytest

Tests run with pytest-asyncio in auto mode. All test files are in the tests/ directory as configured in pyproject.toml.

License

MIT