diff --git a/Dockerfile.mcp b/Dockerfile.mcp index 22524111..c6c85d63 100644 --- a/Dockerfile.mcp +++ b/Dockerfile.mcp @@ -13,6 +13,10 @@ RUN pip install --no-cache-dir --upgrade mcp fastmcp qdrant-client fastembed # Create cache directory with proper permissions RUN mkdir -p /tmp/cache && chmod 755 /tmp/cache +# Ensure rerank volumes are writable when containers run as non-root (uid 1000). +RUN mkdir -p /tmp/rerank_events /tmp/rerank_weights \ + && chmod 777 /tmp/rerank_events /tmp/rerank_weights + # Bake scripts into image so server can run even when /work points elsewhere COPY scripts /app/scripts diff --git a/Dockerfile.mcp-indexer b/Dockerfile.mcp-indexer index 5cfe1286..9e7a6281 100644 --- a/Dockerfile.mcp-indexer +++ b/Dockerfile.mcp-indexer @@ -14,6 +14,12 @@ RUN apt-get update && apt-get install -y --no-install-recommends git ca-certific COPY requirements.txt /tmp/requirements.txt RUN pip install --no-cache-dir --upgrade -r /tmp/requirements.txt +# Ensure rerank volumes are writable when containers run as non-root (uid 1000). +# On first mount of an empty named volume, Docker copies directory contents (incl perms) +# from the image into the volume. +RUN mkdir -p /tmp/rerank_events /tmp/rerank_weights \ + && chmod 777 /tmp/rerank_events /tmp/rerank_weights + # Download reranker model and tokenizer during build # Cross-encoder for reranking (ms-marco-MiniLM) + BGE tokenizer for micro-chunking ARG RERANKER_ONNX_URL=https://huggingface.co/cross-encoder/ms-marco-MiniLM-L-6-v2/resolve/main/onnx/model.onnx diff --git a/docker-compose.yml b/docker-compose.yml index a8b26557..b77fc8d6 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -59,11 +59,18 @@ services: - TOOL_STORE_DESCRIPTION=${TOOL_STORE_DESCRIPTION} - TOOL_FIND_DESCRIPTION=${TOOL_FIND_DESCRIPTION} - FASTMCP_HEALTH_PORT=18000 + # Learning reranker configuration + - RERANK_LEARNING=${RERANK_LEARNING:-1} + - RERANKER_WEIGHTS_DIR=/tmp/rerank_weights + - RERANK_EVENTS_DIR=/tmp/rerank_events + - RERANK_EVENTS_ENABLED=${RERANK_EVENTS_ENABLED:-1} ports: - "18000:18000" - "8000:8000" volumes: - workspace_pvc:/work:ro + - rerank_data:/tmp/rerank_weights:rw + - rerank_events:/tmp/rerank_events:rw networks: - dev-remote-network @@ -120,15 +127,57 @@ services: - INDEX_UPSERT_BATCH=${INDEX_UPSERT_BATCH:-512} - INDEX_UPSERT_RETRIES=${INDEX_UPSERT_RETRIES:-5} - MAX_MICRO_CHUNKS_PER_FILE=${MAX_MICRO_CHUNKS_PER_FILE:-200} + # Learning reranker configuration + - RERANK_LEARNING=${RERANK_LEARNING:-1} + - RERANKER_WEIGHTS_DIR=/tmp/rerank_weights + - RERANK_EVENTS_DIR=/tmp/rerank_events + - RERANK_EVENTS_ENABLED=${RERANK_EVENTS_ENABLED:-1} ports: - "${FASTMCP_INDEXER_PORT:-8001}:8001" - "18001:18001" volumes: - workspace_pvc:/work:rw - codebase_pvc:/work/.codebase:rw + - rerank_data:/tmp/rerank_weights:rw + - rerank_events:/tmp/rerank_events:rw networks: - dev-remote-network + # Learning reranker worker - processes training events in background + learning_worker: + build: + context: . + dockerfile: Dockerfile.mcp-indexer + container_name: learning-worker-dev-remote + user: "1000:1000" + command: ["sh", "-c", "mkdir -p /tmp/huggingface/hub /tmp/huggingface/transformers /tmp/huggingface/fastembed && exec python /app/scripts/learning_reranker_worker.py --daemon"] + depends_on: + - qdrant + - mcp_indexer + env_file: + - .env + environment: + - QDRANT_URL=${QDRANT_URL} + - COLLECTION_NAME=${COLLECTION_NAME} + - HF_HOME=/tmp/huggingface + - HF_HUB_CACHE=/tmp/huggingface/hub + - TRANSFORMERS_CACHE=/tmp/huggingface/transformers + - FASTEMBED_CACHE_PATH=/tmp/huggingface/fastembed + - EMBEDDING_MODEL=${EMBEDDING_MODEL} + - EMBEDDING_PROVIDER=${EMBEDDING_PROVIDER} + - RERANK_EVENTS_DIR=/tmp/rerank_events + - RERANKER_WEIGHTS_DIR=/tmp/rerank_weights + - RERANK_LEARNING_BATCH_SIZE=${RERANK_LEARNING_BATCH_SIZE:-32} + - RERANK_LEARNING_POLL_INTERVAL=${RERANK_LEARNING_POLL_INTERVAL:-30} + - RERANK_LEARNING_RATE=${RERANK_LEARNING_RATE:-0.001} + volumes: + - workspace_pvc:/work:rw + - rerank_data:/tmp/rerank_weights:rw + - rerank_events:/tmp/rerank_events:rw + networks: + - dev-remote-network + restart: unless-stopped + # MCP HTTP search service - same as base compose mcp_http: build: @@ -171,11 +220,18 @@ services: - TOOL_STORE_DESCRIPTION=${TOOL_STORE_DESCRIPTION} - TOOL_FIND_DESCRIPTION=${TOOL_FIND_DESCRIPTION} - FASTMCP_HEALTH_PORT=18000 + # Learning reranker configuration + - RERANK_LEARNING=${RERANK_LEARNING:-1} + - RERANKER_WEIGHTS_DIR=/tmp/rerank_weights + - RERANK_EVENTS_DIR=/tmp/rerank_events + - RERANK_EVENTS_ENABLED=${RERANK_EVENTS_ENABLED:-1} ports: - "${FASTMCP_HTTP_HEALTH_PORT:-18002}:18000" - "${FASTMCP_HTTP_PORT:-8002}:8000" volumes: - workspace_pvc:/work:ro + - rerank_data:/tmp/rerank_weights:rw + - rerank_events:/tmp/rerank_events:rw networks: - dev-remote-network @@ -233,12 +289,19 @@ services: - INDEX_UPSERT_BATCH=${INDEX_UPSERT_BATCH:-512} - INDEX_UPSERT_RETRIES=${INDEX_UPSERT_RETRIES:-5} - MAX_MICRO_CHUNKS_PER_FILE=${MAX_MICRO_CHUNKS_PER_FILE:-200} + # Learning reranker configuration + - RERANK_LEARNING=${RERANK_LEARNING:-1} + - RERANKER_WEIGHTS_DIR=/tmp/rerank_weights + - RERANK_EVENTS_DIR=/tmp/rerank_events + - RERANK_EVENTS_ENABLED=${RERANK_EVENTS_ENABLED:-1} ports: - "${FASTMCP_INDEXER_HTTP_PORT:-8003}:8001" - "${FASTMCP_INDEXER_HTTP_HEALTH_PORT:-18003}:18001" volumes: - workspace_pvc:/work:rw - codebase_pvc:/work/.codebase:rw + - rerank_data:/tmp/rerank_weights:rw + - rerank_events:/tmp/rerank_events:rw networks: - dev-remote-network @@ -494,6 +557,14 @@ volumes: qdrant_storage_dev_remote: driver: local + # Learning reranker weights storage (shared between indexer and worker) + rerank_data: + driver: local + + # Learning reranker events storage (shared between indexer and worker) + rerank_events: + driver: local + # Custom network for service discovery networks: dev-remote-network: diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index cc6edd43..b7549761 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -8,6 +8,7 @@ - [Overview](#overview) - [Core Principles](#core-principles) - [System Architecture](#system-architecture) +- [Learning Reranker System](#5-learning-reranker-system) - [Data Flow](#data-flow) - [ReFRAG Pipeline](#refrag-pipeline) @@ -122,6 +123,110 @@ Context Engine is a production-ready MCP (Model Context Protocol) retrieval stac - **Local LLM Integration**: llama.cpp for offline expansion - **Caching**: Expanded query results cached for reuse +### 5. Learning Reranker System (Optional) + +The Learning Reranker is an **optional** self-improving ranking system that learns from search patterns to provide increasingly relevant results over time. It is enabled by default but can be disabled via `RERANK_LEARNING=0` and `RERANK_EVENTS_ENABLED=0` environment variables. See [Configuration](CONFIGURATION.md#learning-reranker) for all options. + +#### Architecture Overview + +``` +┌─────────────────┐ ┌──────────────────┐ ┌─────────────────┐ +│ Search Query │────►│ Hybrid Search │────►│ TinyScorer │ +│ │ │ (initial rank) │ │ (learned rank) │ +└─────────────────┘ └──────────────────┘ └─────────────────┘ + │ + ┌──────────────────┐ │ + │ Event Logger │◄────────────┘ + │ (NDJSON files) │ + └────────┬─────────┘ + │ + ┌────────▼─────────┐ + │ Learning Worker │ + │ (background) │ + └────────┬─────────┘ + │ + ┌────────▼─────────┐ + │ ONNX Teacher │ + │ (cross-encoder) │ + └────────┬─────────┘ + │ + ┌────────▼─────────┐ + │ Weight Updates │ + │ (.npz files) │ + └──────────────────┘ +``` + +#### Components + +**TinyScorer** (`scripts/rerank_recursive.py`) +- 2-layer MLP neural network (~3MB per collection) +- Scores query-document pairs based on learned patterns +- Hot-reloads weights every 60 seconds from disk +- Per-collection weights (each repo learns independently) + +**Event Logger** (`scripts/rerank_events.py`) +- Logs every search to NDJSON files at `/tmp/rerank_events/` +- Records: query, candidates, initial scores, timestamps +- Hourly file rotation with configurable retention + +**Learning Worker** (`scripts/learning_reranker_worker.py`) +- Background daemon that processes logged events +- Uses ONNX cross-encoder as "teacher" model +- Trains TinyScorer via knowledge distillation +- Saves versioned weight checkpoints atomically + +#### Learning Flow + +1. **Event Capture**: Every search logs query + candidates to NDJSON +2. **Teacher Scoring**: ONNX cross-encoder scores the candidates +3. **Student Training**: TinyScorer learns to match teacher rankings +4. **Weight Update**: New weights saved atomically with versioning +5. **Hot Reload**: Serving path picks up new weights within 60s +6. **Score Integration**: `learning_score` blends with other signals + +#### Configuration + +| Variable | Description | Default | +|----------|-------------|---------| +| `RERANKER_WEIGHTS_DIR` | Directory for weight files | `/tmp/rerank_weights` | +| `RERANKER_WEIGHTS_RELOAD_INTERVAL` | Hot-reload check interval (seconds) | 60 | +| `RERANKER_MAX_CHECKPOINTS` | Number of weight versions to keep | 5 | +| `RERANKER_LR_DECAY_STEPS` | Steps between learning rate decay | 1000 | +| `RERANKER_LR_DECAY_RATE` | Learning rate decay multiplier | 0.95 | +| `RERANKER_MIN_LR` | Minimum learning rate | 0.0001 | +| `RERANK_EVENTS_DIR` | Directory for event logs | `/tmp/rerank_events` | +| `RERANK_EVENTS_RETENTION_DAYS` | Days to keep event files | 7 | +| `RERANK_LEARNING_BATCH_SIZE` | Events per training batch | 32 | +| `RERANK_LEARNING_POLL_INTERVAL` | Worker poll interval (seconds) | 30 | +| `RERANK_LEARNING_RATE` | Initial learning rate | 0.001 | + +#### Observability + +Search results include learning metrics in the `why` field: +```json +{ + "score": 3.2, + "why": ["lexical:1.0", "dense_rrf:0.05", "learning:3", "score:3.2"], + "components": { + "learning_score": 3.2, + "learning_iterations": 3 + } +} +``` + +Worker logs show training progress: +``` +[codebase] Processed 5 events | v12 | lr=0.001 | avg_loss=1.8 | converged=False +``` + +#### Benefits + +- **Zero Manual Training**: Learns automatically from usage +- **Per-Collection Specialization**: Each codebase gets tuned rankings +- **Fast Inference**: TinyScorer adds <1ms to search latency +- **Continuous Improvement**: Rankings improve over time +- **Offline Capable**: Teacher runs locally, no external API calls + #### MCP Router (`scripts/mcp_router.py`) - **Intent Classification**: Determines which MCP tool to call based on query - **Tool Orchestration**: Routes to search, answer, memory, or index tools diff --git a/docs/CONFIGURATION.md b/docs/CONFIGURATION.md index c2e93b7e..b7f90ca4 100644 --- a/docs/CONFIGURATION.md +++ b/docs/CONFIGURATION.md @@ -13,6 +13,7 @@ Complete environment variable reference for Context Engine. - [Query Optimization](#query-optimization) - [Watcher Settings](#watcher-settings) - [Reranker](#reranker) +- [Learning Reranker](#learning-reranker) - [Decoder (llama.cpp / GLM / MiniMax)](#decoder-llamacpp--glm--minimax) - [ReFRAG](#refrag) - [Ports](#ports) @@ -128,6 +129,61 @@ Dynamic HNSW_EF tuning and intelligent query routing for 2x faster simple querie | RERANKER_TOKENIZER_PATH | Tokenizer path for reranker | unset | | RERANKER_ENABLED | Enable reranker by default | 1 (enabled) | +## Learning Reranker + +The learning reranker trains a lightweight neural network (TinyScorer) to improve search rankings over time. See [Architecture](ARCHITECTURE.md#5-learning-reranker-system) for details. + +**This feature is optional and enabled by default.** To disable: + +```bash +# Disable learning scorer in search results +RERANK_LEARNING=0 + +# Disable event logging (no training data collected) +RERANK_EVENTS_ENABLED=0 + +# Or simply don't run the learning_worker container +``` + +### Enable/Disable + +| Name | Description | Default | +|------|-------------|---------| +| RERANK_LEARNING | Enable learning scorer in search results | 1 (enabled) | +| RERANK_EVENTS_ENABLED | Enable event logging for training | 1 (enabled) | +| RERANK_EVENTS_SAMPLE_RATE | Fraction of events to log (0.0-1.0) | 0.33 | + +### Weight Management + +| Name | Description | Default | +|------|-------------|---------| +| RERANKER_WEIGHTS_DIR | Directory for learned weight files | /tmp/rerank_weights | +| RERANKER_WEIGHTS_RELOAD_INTERVAL | How often to check for new weights (seconds) | 60 | +| RERANKER_MAX_CHECKPOINTS | Number of weight versions to retain | 5 | + +### Learning Rate + +| Name | Description | Default | +|------|-------------|---------| +| RERANKER_LR_DECAY_STEPS | Updates between learning rate decay | 1000 | +| RERANKER_LR_DECAY_RATE | Decay multiplier (e.g., 0.95 = 5% reduction) | 0.95 | +| RERANKER_MIN_LR | Minimum learning rate floor | 0.0001 | + +### Event Logging + +| Name | Description | Default | +|------|-------------|---------| +| RERANK_EVENTS_DIR | Directory for search event logs | /tmp/rerank_events | +| RERANK_EVENTS_RETENTION_DAYS | Days to keep event files before cleanup | 7 | + +### Learning Worker + +| Name | Description | Default | +|------|-------------|---------| +| RERANK_LEARNING_BATCH_SIZE | Number of events per training batch | 32 | +| RERANK_LEARNING_POLL_INTERVAL | Seconds between checking for new events | 30 | +| RERANK_LEARNING_RATE | Initial learning rate for TinyScorer | 0.001 | + ## Decoder (llama.cpp / GLM / MiniMax) | Name | Description | Default | diff --git a/scripts/rerank_recursive.py b/scripts/rerank_recursive.py index 8a886e55..84915922 100644 --- a/scripts/rerank_recursive.py +++ b/scripts/rerank_recursive.py @@ -1012,7 +1012,11 @@ def rerank_with_learning( teacher_scores = None try: - from scripts.rerank_events import log_training_event + # Try both import paths for Docker (/app/scripts) and local (scripts/) + try: + from rerank_events import log_training_event + except ImportError: + from scripts.rerank_events import log_training_event log_training_event( query=query, candidates=candidates,