RAG Code Search Agent

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

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

MCP Server Tools

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

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