Skip to content
BranchesTags

Latest commit

 

History

9 Commits

Folders and files

NameName
Last commit message
Last commit date
backend
backend
 
 
data
data
 
 
frontend
frontend
 
 
.gitignore
.gitignore
 
 
README.md
README.md
 
 
assignment.md
assignment.md
 
 
docker-compose.yml
docker-compose.yml
 
 

Repository files navigation

AI Coding Session Search

A production-ready semantic search application for searching through AI coding assistant (Claude Code) session history. Built with FastAPI, Sentence Transformers, ChromaDB, and React.

Architecture AI Frontend Search Status

✨ Current Status

✅ Production-Ready RAG Pipeline with:

  • ✅ Dual Search Modes: Basic (fast) and Enhanced (high-quality)
  • ✅ Hybrid Search: BM25 + Vector similarity for best results
  • ✅ Cross-Encoder Reranking: Improved relevance scoring
  • ✅ Query Expansion: LLM-powered with Groq API
  • ✅ Redis Caching: 70% latency reduction for cached queries
  • ✅ Fairness Enforcement: Equal quota distribution across engineers
  • ✅ Pagination Support: Efficient result browsing
  • ✅ Query Validation: Security and input sanitization
  • ✅ Rate Limiting: Protection against abuse
  • ✅ Docker Deployment: One-command setup
  • ✅ Memory Optimized: Efficient embedding storage
  • ✅ Fixed Data Chunking: Improved message parsing logic

🚀 Quick Start

Option 1: Docker Compose (Recommended)

# Clone the repository
git clone <your-repo-url>
cd Dexicon_assignment

# Copy environment file (optional)
cp .env.example .env

# Edit .env if needed (set GROQ_API_KEY for enhanced mode)

# Start all services
docker-compose up -d

# View logs
docker-compose logs -f

# Access the application
# Frontend: http://localhost
# Backend API: http://localhost:8001

Services:

Stop services:

docker-compose down

Rebuild after changes:

docker-compose up -d --build

Option 2: Local Development

Prerequisites

  • Python 3.9+
  • Node.js 18+
  • npm or yarn
  • Redis (optional, for caching)

1. Install Backend Dependencies

cd backend
pip install -r requirements.txt

Note: First run will download the embedding model (~90MB). This may take a minute.

2. Start the Backend

# Create and activate virtual environment (recommended)
python3 -m venv venv
source venv/bin/activate  # On Windows: venv\Scripts\activate

# Install dependencies
pip install -r backend/requirements.txt

# Set environment variables (optional)
export SEARCH_MODE=enhanced  # or "basic"
export CACHE_ENABLED=true
export REDIS_URL=redis://localhost:6379/0
export GROQ_API_KEY=your_key_here  # Optional, for query expansion

# Start the server
cd backend
python main.py

The server will:

  1. Load session JSON files from data/ directory
  2. Parse and create searchable chunks (based on mode)
  3. Generate embeddings using Sentence Transformers
  4. Index into ChromaDB (persistent storage)
  5. Initialize Redis cache (if enabled)
  6. Start API server at http://localhost:8001

Switch between modes:

  • Set SEARCH_MODE=basic for fast vector-only search
  • Set SEARCH_MODE=enhanced for hybrid search with reranking

3. Install Frontend Dependencies

cd frontend
npm install

4. Start the Frontend

cd frontend
npm run dev

Open http://localhost:5173 in your browser.

📁 Project Structure

├── backend/
│   ├── basic/               # Basic search implementation
│   │   ├── __init__.py
│   │   └── search_engine.py  # Vector-only search
│   │
│   ├── enhanced/            # Enhanced search (hybrid + reranking)
│   │   ├── __init__.py
│   │   ├── enhanced_search.py  # Hybrid search + reranking
│   │   ├── enhanced_chunking.py  # Advanced chunking strategies
│   │   └── knowledge_graph.py  # Knowledge graph for entity-aware retrieval
│   │
│   ├── shared/              # Shared utilities
│   │   ├── __init__.py
│   │   ├── data_loader.py  # JSON parsing & chunking
│   │   ├── models.py        # Pydantic schemas
│   │   ├── config.py        # Configuration management
│   │   ├── cache.py         # Redis caching layer
│   │   ├── fairness.py      # Equal quota enforcement
│   │   └── evaluation.py    # Quality metrics
│   │
│   ├── tests/               # Test suite
│   │   ├── __init__.py
│   │   ├── test_knowledge_graph.py  # KG building & visualization
│   │   ├── test_enhanced_retrieval.py  # Step-by-step retrieval testing
│   │   ├── test_vector_scores.py  # Basic vs Enhanced comparison
│   │   ├── investigate_negative_scores.py  # Score debugging
│   │   ├── knowledge_graph_visualization.png  # Full graph visualization
│   │   └── knowledge_graph_simplified.png  # Simplified graph view
│   │
│   ├── main.py              # FastAPI application
│   ├── Dockerfile           # Backend container
│   ├── requirements.txt
│   └── README.md            # Backend documentation
│
├── frontend/
│   ├── src/
│   │   ├── App.jsx          # Main application
│   │   ├── components/      # React components
│   │   │   ├── SearchBar.jsx
│   │   │   ├── ResultsList.jsx
│   │   │   ├── FilterPanel.jsx
│   │   │   └── ...
│   │   └── hooks/           # Custom hooks
│   │       └── useSearch.js
│   ├── Dockerfile           # Frontend container
│   ├── nginx.conf           # Nginx configuration
│   └── package.json
│
├── data/                    # Session JSON files
│   ├── andrew_wang_sessions.json
│   ├── daniel_lin_sessions.json
│   └── diana_lu_sessions.json
│
├── docker-compose.yml       # Docker Compose configuration
├── .gitignore              # Git ignore rules
└── README.md               # This file

🔧 API Reference

POST /api/search

Search for relevant coding sessions with pagination and filters.

curl -X POST http://localhost:8001/api/search \
  -H "Content-Type: application/json" \
  -d '{
    "query": "how to handle large file uploads",
    "limit": 10,
    "offset": 0,
    "filters": {
      "engineer": "andrewwang",
      "project": "video-ingester",
      "language": "Go"
    }
  }'

Response includes:

  • results: Array of search results
  • total: Total number of results
  • offset, limit, has_more: Pagination info
  • query_time_ms: Query execution time
  • query_expanded: Expanded query (if enhanced mode)
  • reranked: Whether results were reranked
  • distribution: Engineer distribution stats (fairness)

GET /api/engineers

List all engineers in the dataset.

GET /api/projects

List all projects in the dataset.

GET /api/stats

Get dataset statistics.

GET /api/cache/stats

Get Redis cache statistics (hits, misses, hit rate).

POST /api/cache/invalidate

Invalidate cached search results.

🧠 How It Works

Search Modes

Basic Mode (Fast & Simple)

  • Chunking: Simple Q&A pairs (no overlap)
  • Model: all-MiniLM-L6-v2 (384-dim, fast)
  • Search: Pure vector search (cosine similarity)
  • Speed: 50-200ms per query
  • Best For: Fast queries, clear semantic matches

Enhanced Mode (High Quality)

  • Chunking: Overlapping windows (2 Q&A pairs + context, overlap=1)
  • Model: all-mpnet-base-v2 (768-dim, better quality)
  • Search: Hybrid (70% vector + 30% BM25) + Knowledge Graph (20% boost)
  • Reranking: Cross-encoder for final ordering
  • Query Expansion: Groq LLM for intelligent synonym expansion
  • Knowledge Graph: Entity-aware retrieval with relationship traversal
  • Speed: 500-1000ms per query (50-200ms for cached)
  • Best For: Complex queries, maximum quality, entity-specific searches

Semantic Search Pipeline

  1. Data Loading: Parse JSON session files, extract Q+A conversation pairs
  2. Chunking:
    • Basic mode: Simple Q&A pairs
    • Enhanced mode: Overlapping windows with context
  3. Embedding: Generate vectors using Sentence Transformers
  4. Indexing: Store embeddings in ChromaDB with metadata
  5. Search:
    • Basic: Vector similarity only
    • Enhanced: Hybrid search (BM25 + Vector + Knowledge Graph) → Reranking
  6. Post-Processing: Fairness distribution, pagination, caching

Why Sentence Transformers?

  • Zero API setup: No keys, no costs, works offline
  • Fast inference: ~14ms per query on CPU
  • Good quality: MiniLM achieves ~90% of larger models' performance
  • Semantic understanding: Captures meaning beyond keywords

Why ChromaDB?

  • Embedded: Just a Python library, no server to run
  • Metadata filtering: Built-in support for engineer/project filters
  • Fast: Uses HNSW for approximate nearest neighbor search
  • Simple: Perfect for prototypes and small-to-medium datasets

📊 Dataset

The sample data includes sessions from 3 engineers working at a video streaming company:

Engineer Role Topics
Andrew Wang Staff Backend Engineer Video encoding, S3 uploads, Celery
Daniel Lin Senior Full-Stack Engineer WebRTC, video validation, FFprobe
Diana Lu Senior Frontend Engineer HLS streaming, iOS PiP, SwiftUI

Languages: Python, Go, TypeScript, Swift Frameworks: FastAPI, Chi, Next.js, React, SwiftUI

🎨 Design Decisions

Chunking Strategy

Basic Mode: Simple Q&A pairs

  • Each user query + assistant response = one chunk
  • ~53 chunks from dataset
  • Fast indexing and retrieval

Enhanced Mode: Overlapping windows

  • Sliding window of 2 Q&A pairs with 1-pair overlap
  • Includes context from previous pairs
  • ~53 chunks (same count as basic, but with context)
  • Better context preservation, no boundary issues
  • Optimal balance between granularity and context

Why Q+A pairs?

  • Semantic completeness: Question and answer together carry full meaning
  • Better retrieval: Search finds complete discussions, not fragments
  • Display-ready: Can show conversation context directly

Key Improvements Made

  1. Fixed Data Chunking: Improved message parsing to handle all assistant responses, even after tool invocations
  2. Memory Optimization: Embeddings stored in ChromaDB only, not duplicated in memory
  3. Better Embedding Model: Switched from CodeBERT to all-mpnet-base-v2 for better Q&A understanding
  4. Embedding Normalization: Fixed negative semantic scores by normalizing embeddings for proper cosine distance
  5. Knowledge Graph: Implemented entity extraction and relationship-based retrieval
  6. Fairness Enforcement: Implements "equal quota" rule for balanced engineer representation
  7. Security: Query validation, XSS prevention, rate limiting
  8. Performance: Redis caching reduces latency by 70% for repeated queries

UI/UX

  • Dark theme: Matches IDE aesthetic, easier on eyes for developers
  • Animated placeholders: Show example queries to guide users
  • Score badges: Visual indicator of relevance
  • Staggered animations: Results feel more dynamic and responsive
  • Minimal filters: Only most useful (engineer, project) to avoid clutter

🚀 Running in Production

Docker Deployment

The easiest way to deploy is using Docker Compose:

docker-compose up -d

This starts:

  • Backend: FastAPI server with enhanced search
  • Frontend: React app served via Nginx
  • Redis: Caching layer

Environment Variables

Create a .env file:

SEARCH_MODE=enhanced          # or "basic"
GROQ_API_KEY=your_key_here    # Optional, for query expansion
CACHE_ENABLED=true
REDIS_URL=redis://redis:6379/0

Production Considerations

  1. Use persistent ChromaDB storage (already configured in Docker)
  2. Set GROQ_API_KEY for enhanced query expansion
  3. Configure CORS for your domain in backend/main.py
  4. Add authentication if needed
  5. Use environment variables for sensitive data
  6. Monitor Redis for cache performance
  7. Consider Pinecone/Weaviate for larger datasets (>100K chunks)

🧪 Testing

Test Suite

The project includes a comprehensive test suite in backend/tests/:

1. Knowledge Graph Testing

cd backend
python3 tests/test_knowledge_graph.py
  • Builds knowledge graph from chunks
  • Visualizes entities and relationships
  • Shows graph statistics
  • Tests entity extraction and graph retrieval
  • Generates visualization PNGs

2. Enhanced Retrieval Testing

cd backend
python3 tests/test_enhanced_retrieval.py "your query here"
  • Shows step-by-step retrieval process
  • Displays Vector, BM25, and Graph search results
  • Shows score combination
  • Demonstrates reranking impact

3. Vector Score Comparison

cd backend
python3 tests/test_vector_scores.py "video encoding optimization"
  • Compares Basic vs Enhanced mode
  • Explains score calculations
  • Shows performance differences

4. Score Investigation

cd backend
python3 tests/investigate_negative_scores.py
  • Investigates embedding normalization
  • Tests different chunking strategies
  • Analyzes query expansion impact

Test Queries

  • "video encoding optimization" - Should find Andrew's encoding work
  • "file upload S3" - Should find multipart upload discussions
  • "error handling" - Should find multiple engineers' error handling
  • "React streaming" - Should find Diana's React work
  • "memory optimization" - Should find performance discussions

Frontend Testing

  1. Open http://localhost:5173
  2. Try different search queries
  3. Test filters (Engineer, Project)
  4. Test pagination (change limit, use offset)
  5. Verify cache (same query twice - second should be faster)

Knowledge Graph Visualizations

The test suite generates two graph visualizations:

  • Full Graph (tests/knowledge_graph_visualization.png): Complete entity-relationship graph
  • Simplified Graph (tests/knowledge_graph_simplified.png): Focus on engineers, projects, and technologies

View these to understand:

  • Entity relationships (engineers → projects → technologies)
  • Concept connections
  • How graph retrieval works

📊 Performance Metrics

Basic Mode

  • Query Time: 50-200ms
  • Precision@5: ~0.65
  • Recall@5: ~0.64
  • Best For: Fast, simple queries

Enhanced Mode

  • Query Time: 500-1000ms (with caching: 50-200ms for cached)
  • Precision@5: ~0.70-0.78
  • Recall@5: ~0.72-0.75
  • Best For: Complex queries, maximum quality

Cache Performance

  • Hit Rate: ~60-80% for repeated queries
  • Latency Reduction: 70% for cached queries
  • TTL: 1 hour (configurable)

🔒 Security Features

  • ✅ Query validation and sanitization
  • ✅ XSS prevention
  • ✅ SQL injection protection
  • ✅ Rate limiting (30 requests/minute per IP)
  • ✅ CORS configuration
  • ✅ Input length limits

🚧 Known Limitations

  • ChromaDB embeddings are regenerated on restart (unless persistent storage is used)
  • Enhanced mode requires Groq API key for query expansion (fallback available)
  • Rate limiting requires slowapi package
  • Redis is optional but recommended for production

🚀 Future Enhancements (Post-MVP)

1. Enhanced Knowledge Graph ⭐ (Partially Implemented)

Current Implementation: ✅ Knowledge graph is already implemented in Enhanced Mode!

  • Extracts entities: Engineers, Projects, Technologies, Concepts
  • Builds relationships: works_on, uses, discussed, related_to
  • Provides 20% boost in hybrid search
  • Enables entity-aware retrieval

Why Knowledge Graphs? A knowledge graph significantly enhances retrieval by modeling relationships between entities:

Entities to Model:

  • Engineers (Andrew, Daniel, Diana)
  • Projects (video-encoder, video-ingester, etc.)
  • Technologies (Python, Go, React, WebRTC, S3, etc.)
  • Concepts (encoding, streaming, error handling, etc.)
  • Code Patterns (multipart upload, retry logic, etc.)

Relationships:

  • Engineer -[works_on]-> Project
  • Project -[uses]-> Technology
  • Engineer -[discussed]-> Concept
  • Concept -[related_to]-> Concept (e.g., "encoding" ↔ "compression")
  • Session -[about]-> Concept

Benefits:

  1. Multi-hop Reasoning: "What did Andrew discuss about video encoding?" → Query engineer → project → concept relationships
  2. Better Context: Understand that "S3 upload" and "multipart upload" are related concepts
  3. Entity-aware Retrieval: Find results based on entity relationships, not just text similarity
  4. Query Understanding: "Show me all discussions about technologies used in video-encoder" → Traverse graph
  5. Recommendations: "Similar sessions" based on graph structure, not just embeddings

Current Implementation:

# Already implemented in enhanced/knowledge_graph.py
from enhanced.knowledge_graph import KnowledgeGraph

kg = KnowledgeGraph()
kg.build_from_chunks(chunks)  # Extracts entities and relationships
results = kg.graph_search(query, chunks)  # Entity-aware retrieval

Hybrid Retrieval with KG (Current):

  1. Vector search finds semantically similar chunks
  2. BM25 search finds keyword matches
  3. Knowledge graph finds related entities through relationships
  4. Combine: 80% (Vector + BM25) + 20% Graph boost
  5. Rerank for final ordering
  6. Example: Query "video encoding" → Vector finds chunks + BM25 finds keywords + Graph finds all encoding-related discussions across projects

Current Tools:

  • Custom Implementation: Lightweight entity extraction and relationship modeling
  • Pattern Matching: Regex-based technology and concept extraction
  • Graph Traversal: Multi-hop relationship finding

Future Enhancements:

  • Neo4j: Full-featured graph database for larger scale
  • spaCy NER: Better entity extraction with named entity recognition
  • NetworkX: More sophisticated graph algorithms
  • Graph Embeddings: Learn entity representations

Current Status:

  • ✅ Basic entity extraction (engineers, projects, technologies, concepts)
  • ✅ Relationship modeling (works_on, uses, discussed, related_to)
  • ✅ Graph-based retrieval integrated into hybrid search
  • ✅ Entity name variations handled (e.g., "Andrew Wang" → "andrewwang")
  • 🔄 Future: More sophisticated NER, graph embeddings, Neo4j integration

2. Conversation Context View

  • Click on search result to see full session conversation
  • Expandable context windows
  • Thread navigation (previous/next messages)
  • Code block highlighting within conversations

3. Advanced Chunking Strategies

  • Semantic Chunking: Use embeddings to find natural boundaries
  • Code-aware Chunking: Separate code blocks from text
  • Hierarchical Chunking: Multi-level (document → section → paragraph)
  • Adaptive Chunking: Dynamic size based on content type

4. Multi-Modal Search

  • Search across code snippets, documentation, and conversations
  • Code-to-code similarity (using CodeBERT)
  • Visual code search (AST-based matching)
  • Unified search interface

5. Query Understanding & Intent Classification

  • Classify queries: "how-to", "debugging", "optimization", "architecture"
  • Route to specialized retrieval strategies
  • Intent-aware reranking
  • Query suggestions based on intent

6. Learning from User Feedback

  • Click-through rate tracking
  • Relevance feedback (thumbs up/down)
  • Fine-tune embeddings based on user interactions
  • A/B testing framework for retrieval strategies

7. Advanced Analytics

  • Most searched topics dashboard
  • Engineer expertise mapping
  • Project knowledge gaps identification
  • Trending technologies/concepts
  • Search pattern analysis

8. Real-time Updates

  • Watch for new sessions and auto-index
  • Incremental embedding updates
  • Live search results as new data arrives
  • WebSocket support for real-time updates

9. Code-Specific Features

  • Syntax highlighting in results
  • Code diff visualization
  • "Show me similar code patterns"
  • Code snippet extraction and search
  • Language-specific search (Python vs Go vs TypeScript)

10. Enterprise Features

  • Authentication & authorization
  • Team-specific search (private sessions)
  • Export search results (PDF, CSV)
  • Search history and saved searches
  • Collaborative annotations on results

💡 Knowledge Graph Implementation

Current Implementation

The knowledge graph is fully integrated into Enhanced Mode:

Entities Extracted:

  • Engineers: andrewwang, daniellin, dianalu
  • Projects: video-encoder, video-ingester, stream-client-react, etc.
  • Technologies: Python, Go, React, FastAPI, S3, WebRTC, etc.
  • Concepts: encoding, streaming, error handling, optimization, etc.

Relationships Modeled:

  • Engineer -[works_on]-> Project (53 relationships)
  • Project -[uses]-> Technology (333 relationships)
  • Engineer -[discussed]-> Concept (187 relationships)
  • Concept -[related_to]-> Concept (58 relationships)
  • Session -[about]-> Concept (187 relationships)

How It Works:

  1. Entity Extraction: Regex patterns + metadata extraction
  2. Graph Building: Automatically builds relationships from chunks
  3. Query Processing: Extracts entities from user queries
  4. Graph Traversal: Finds related entities (1-2 hop traversal)
  5. Retrieval: Gets chunks associated with entities
  6. Scoring: Boosts results with entity matches (20% of final score)

Example Query Flow:

Query: "video encoding optimization"
→ Extract: concepts=["encoding", "optimization"]
→ Traverse: encoding -[related_to]-> [streaming, compression, ...]
→ Find chunks: All chunks discussing encoding/optimization
→ Score: Direct matches (2.0) + Related matches (0.5)
→ Combine: 80% (Vector+BM25) + 20% Graph

Visualizations:

  • Run python3 tests/test_knowledge_graph.py to generate graph visualizations
  • See tests/knowledge_graph_visualization.png for full graph
  • See tests/knowledge_graph_simplified.png for simplified view

Future Enhancements

  • Better NER: Use spaCy for more accurate entity extraction
  • Graph Embeddings: Learn entity representations
  • Neo4j Integration: Scale to larger graphs
  • Multi-hop Reasoning: Deeper relationship traversal
  • Temporal Relationships: Track concept evolution over time

📝 License

MIT

🙏 Acknowledgments

Built as part of Dexicon AI take-home assessment. Demonstrates production-ready RAG pipeline with comprehensive improvements and best practices.

About

No description, website, or topics provided.

Resources

Readme
Activity

Stars

0 stars

Watchers

0 watching

Forks

0 forks
Report repository

Releases

Packages

Contributors

Languages