RAG Code Search Pipeline

A production-ready implementation of a Retrieval-Augmented Generation (RAG) pipeline for semantic code search and documentation generation using vector databases.

This project demonstrates how to build a semantic code search system that understands meaning, not just keywords. Search for "authentication with JWT" and find relevant code even if functions are named verify_token() or validate_credentials().

• 🔍 Semantic Code Search: Find code by meaning, not just keywords
• 📚 Automated Documentation: Generate contextual explanations using RAG
• ⚡ High Performance: Sub-second query responses for large codebases
• 🎯 Smart Chunking: Function-level extraction with context preservation
• 🔄 Incremental Updates: Git-based change detection for efficient re-indexing
• 🎨 Rich Metadata: Filter by file path, function name, and more
• 🔀 Hybrid Search: Combines semantic similarity with keyword matching

┌─────────────────┐
│  Code Repository │
└────────┬────────┘
         │
         ▼
┌─────────────────┐
│  Code Chunker   │  ← AST-based function extraction
└────────┬────────┘
         │
         ▼
┌─────────────────┐
│ Embedding Model │  ← OpenAI / CodeBERT
└────────┬────────┘
         │
         ▼
┌─────────────────┐
│ Vector Database │  ← Chroma / Pinecone
└────────┬────────┘
         │
         ▼
┌─────────────────┐
│  RAG Pipeline   │  ← Query + Context Assembly
└────────┬────────┘
         │
         ▼
┌─────────────────┐
│   LLM Response  │  ← GPT-4 / Claude
└─────────────────┘

Based on production deployment with 500K+ lines of code:

• Python 3.9 or higher
• OpenAI API key
• Git repository with Python code

# Clone the repository
git clone https://github.com/dinesh-k-elumalai/rag-code-search-pipeline.git
cd rag-code-search-pipeline

# Install dependencies
pip install -r requirements.txt

# Set up environment variables
cp .env.example .env
# Edit .env and add your OPENAI_API_KEY

python scripts/index_codebase.py --repo-path /path/to/your/repo

from src.rag_assistant import RAGCodeAssistant

# Initialize the assistant
assistant = RAGCodeAssistant()

# Search and get explanation
result = assistant.search_and_explain(
    "How do we handle authentication?"
)

print(result)

python scripts/interactive_search.py

rag-code-search-pipeline/
├── src/
│   ├── chunker.py              # Code chunking and AST parsing
│   ├── embedder.py             # Embedding generation
│   ├── vector_store.py         # Vector database operations
│   ├── rag_assistant.py        # RAG pipeline orchestration
├── scripts/
│   ├── index_codebase.py       # Batch indexing script
│   └── interactive_search.py   # CLI interface
├── config/
│   └── config.yaml             # Configuration settings
├── requirements.txt
├── .env.example
├── .gitignore
└── README.md

Edit config/config.yaml to customize:

embedding:
  model: "text-embedding-3-large"
  batch_size: 100

chunking:
  context_lines: 5
  min_function_length: 3

vector_store:
  type: "chroma"  # or "pinecone", "weaviate"
  collection_name: "code_search"
  
search:
  semantic_weight: 0.6
  top_k: 5

llm:
  model: "gpt-4-turbo"
  temperature: 0.3
  max_tokens: 2000

Combine semantic similarity with keyword matching:

results = assistant.hybrid_search(
    query="JWT authentication",
    semantic_weight=0.6,  # 60% semantic, 40% keyword
    top_k=10
)

Filter searches by file path, function name, or custom metadata:

results = assistant.search(
    query="validate input",
    filter={
        "file_path": {"$contains": "auth"},
        "has_docstring": True
    }
)

docker build -t rag-code-search .
docker run -e OPENAI_API_KEY=your_key rag-code-search

1. Batch Processing: Process embeddings in batches of 100
2. Caching: Enable embedding cache for repeated queries
3. Incremental Updates: Use git-based change detection
4. Connection Pooling: Reuse database connections

See OPTIMIZATION.md for detailed tuning guide.

Based on typical usage with OpenAI API:

Monthly estimate for 10-person team: $15-30

Issue: Slow indexing performance

• Solution: Increase batch_size in config.yaml
• Solution: Use incremental updates instead of full re-indexing

Issue: Poor search relevance

• Solution: Adjust context_lines for more/less function context
• Solution: Tune semantic_weight in hybrid search
• Solution: Ensure proper docstrings in your codebase

Issue: Out of memory during indexing

• Solution: Process files in smaller batches
• Solution: Use streaming mode for large repositories

See TROUBLESHOOTING.md for more solutions.

Contributions are welcome! Please see CONTRIBUTING.md for guidelines.

• Support for additional programming languages (Java, Go, TypeScript)
• Web UI for code search
• Integration with popular IDEs (VS Code, IntelliJ)
• Multi-repository indexing
• Custom embedding model fine-tuning
• Real-time code analysis and suggestions

This project is licensed under the MIT License - see the LICENSE file for details.

If you use this project in your research or production systems, please cite:

@software{rag_code_search_2024,
  title = {RAG Code Search Pipeline},
  author = {Dinesh Kumar Elumalai},
  year = {2025},
  url = {https://github.com/dinesh-k-elumalai/rag-code-search-pipeline}
}

• Built with ChromaDB for vector storage
• Powered by OpenAI Embeddings
• Inspired by production needs in enterprise software development

• 📧 Email: dinesh.k.elumalai@ieee.org
• 🐛 Issues: GitHub Issues

Built with ❤️ for developers who deserve better code search