RAG Code Search Agent
by pratishank05
README.md
# RAG Code Search Agent
[](https://www.python.org/)
[](https://github.com/langchain-ai/langgraph)
[](https://www.trychroma.com/)
[](https://docs.anthropic.com/)
[](https://modelcontextprotocol.io/)
[](https://www.docker.com/)
[](https://opensource.org/licenses/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
```bash
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
```bash
pip install -e ".[dev]"
```
### Via Docker
```bash
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:
```bash
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
```bash
rag-search index /path/to/my-repo
```
Force a full re-index:
```bash
rag-search index /path/to/my-repo --full
```
### Query the Codebase
```bash
rag-search query "How does the authentication middleware work?"
```
JSON output:
```bash
rag-search query "Where is the payment processing logic?" --format json
```
### List Indexed Repositories
```bash
rag-search list
```
### Start the MCP Server
Stdio transport (for direct process communication):
```bash
rag-search serve
```
SSE transport (for HTTP-based integration):
```bash
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
```bash
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](https://opensource.org/licenses/MIT)
This server cannot be deployed
Maintenance
ActivityInactive
ResponsivenessNo issues