Skip to content

Latest commit

 

History

History
589 lines (441 loc) · 26.5 KB

File metadata and controls

589 lines (441 loc) · 26.5 KB

🧠 Relay — Adaptive RAG Architecture

A full deep-dive into the Retrieval-Augmented Generation pipeline powering Relay's AI workspace assistant. Built from scratch in Node.js — no LangChain, no LlamaIndex, no black boxes.


Table of Contents

  1. What is RAG and Why We Built It From Scratch
  2. Architecture Overview
  3. Knowledge Base — MongoDB Atlas Vector Search
  4. Encoder & Embeddings — Gemini embedding-001
  5. Stage 1 — Adaptive Retrieval Classifier
  6. Stage 2 — Parallel Vector Search
  7. Stage 3 — Relevance Threshold Filtering
  8. Stage 4 — Adaptive Temperature
  9. Stage 5 — Prompt Engineering & Grounded Generation
  10. LLM Orchestration & Failover
  11. Hallucination Prevention — Multi-Layered Strategy
  12. Search Architecture — BM25 vs Vector Search
  13. File Reference Map
  14. Key Parameters at a Glance
  15. Future Improvements

1. What is RAG and Why We Built It From Scratch

RAG (Retrieval-Augmented Generation) is a technique that augments a Large Language Model (LLM) with external knowledge retrieved at query time, rather than relying entirely on what the model learned during training.

Without RAG:

User: "How does my binary search function work?"
LLM:  [Makes up a generic binary search explanation — has no idea what YOUR code looks like]

With RAG:

User: "How does my binary search function work?"
System: [Fetches your actual binary search code from your personal vault]
LLM:  "Your implementation uses an iterative while-loop with a mid-point calculation..."
         [Grounded in your real code — no hallucination]

Why No LangChain or LlamaIndex?

Framework abstraction hides the mechanics that matter. By building the pipeline manually in Node.js, we:

  • Have complete control over every parameter (threshold, temperature, numCandidates)
  • Can implement custom optimisations (adaptive retrieval, adaptive temperature) that generic frameworks don't support out of the box
  • Have zero abstraction overhead — each step is a direct API call, no magic
  • Can explain every single line of the pipeline to an interviewer

2. Architecture Overview

User Query
    │
    ▼
┌─────────────────────────────┐
│   Adaptive RAG Classifier   │  ← needsRetrieval(query) — regex intent detection
│   [vectorSearch.js]         │
└─────────────┬───────────────┘
              │
    ┌─────────┴──────────┐
    │                    │
    ▼                    ▼
GENERAL QUERY        PERSONAL QUERY
(skip retrieval)     (run retrieval)
    │                    │
    │                    ▼
    │          ┌──────────────────────┐
    │          │  Gemini Embedding    │  ← gemini-embedding-001, 768-dim
    │          │  [geminiEmbed.js]    │
    │          └──────────┬───────────┘
    │                     │
    │                     ▼
    │          ┌─────────────────────────┐
    │          │  Parallel $vectorSearch │  ← Notes + Snippets simultaneously
    │          │  [vectorSearch.js]      │    numCandidates:100, limit:5 each
    │          └──────────┬──────────────┘
    │                     │
    │                     ▼
    │          ┌──────────────────────┐
    │          │  Threshold Filter    │  ← cosine similarity ≥ 0.60
    │          │  [vectorSearch.js]   │    discard low-relevance docs
    │          └──────────┬───────────┘
    │                     │
    └─────────┬───────────┘
              │
              ▼
    ┌─────────────────────────────┐
    │   Adaptive Temperature      │  ← 0.2 (RAG/code) or 0.7 (general)
    │   getTemperature()          │
    │   [aiRoutes.js]             │
    └─────────────┬───────────────┘
                  │
                  ▼
    ┌─────────────────────────────┐
    │   Prompt Engineering        │  ← XML context injection + citation instructions
    │   [aiRoutes.js]             │
    └─────────────┬───────────────┘
                  │
                  ▼
    ┌─────────────────────────────┐
    │   LLM Generation            │  ← Gemini 2.5-flash → 2.0-flash (failover)
    │   [aiRoutes.js]             │
    └─────────────┬───────────────┘
                  │
                  ▼
            Final Answer
       (cited, grounded, accurate)

3. Knowledge Base — MongoDB Atlas Vector Search

What is stored?

The knowledge base is the user's personal workspace, split across two MongoDB collections:

Collection Content Embedded Fields Retrieved
notes Title + stripped HTML description _id, title, description, score
codesnippets Title + Language + Tags + raw code _id, title, language, code, score

Why MongoDB Atlas Vector Search?

MongoDB Atlas Vector Search uses an HNSW (Hierarchical Navigable Small World) graph index — the same approximate nearest-neighbour algorithm used by Pinecone, Weaviate, and Qdrant. Running it inside MongoDB means:

  • No separate vector database to manage or pay for
  • Vector and metadata live in the same document — one query fetches both
  • The filter parameter on $vectorSearch applies a user-scoped filter (only fetch THIS user's documents) before traversing the HNSW graph — preventing data leakage between users

Embedding Storage on Document

// Note.js
embedding: {
  type: [Number],  // Array of 768 floats
  default: [],
}

The embedding is stored directly on the document alongside the text. No join operation is needed — the vector is always co-located with its source document.

Automatic Re-Embedding on Save

The embedding is never stale. A Mongoose pre('save') hook regenerates it automatically whenever the content changes:

// Note.js
NoteSchema.pre("save", async function (next) {
  if (this.isModified("title") || this.isModified("description")) {
    const strippedDesc = this.description.replace(/<[^>]*>?/gm, '');
    const textToEmbed = `Title: ${this.title}\nContent:\n${strippedDesc}`;
    const embeddingValues = await generateEmbedding(textToEmbed);
    if (embeddingValues && embeddingValues.length > 0) {
      this.embedding = embeddingValues;
    }
  }
  next();
});

Key detail: HTML tags are stripped from TipTap's rich-text output before embedding. Raw HTML like <p><strong>Binary Search</strong></p> would produce a noisier, less semantically accurate vector than the clean text Binary Search.


4. Encoder & Embeddings — Gemini embedding-001

The Model

Property Value
Model gemini-embedding-001 (successor to text-embedding-004)
Provider Google AI (via @google/generative-ai SDK)
Native Dimensions 3072
Stored Dimensions 768 (truncated for index compatibility)
Task Type RETRIEVAL_DOCUMENT

Why 768 Dimensions?

The MongoDB Atlas Vector Search index was originally created with 768 dimensions. Dropping and rebuilding an index on a live collection is expensive and causes downtime. By passing outputDimensionality: 768 to the API, we request a truncated vector that is still semantically meaningful but matches our existing index — zero migration required.

// geminiEmbed.js
const result = await model.embedContent({
  content: { parts: [{ text: contentToEmbed }], role: "user" },
  taskType: "RETRIEVAL_DOCUMENT",
  outputDimensionality: 768, // match existing Atlas Vector Search index dimensions
});

Why RETRIEVAL_DOCUMENT task type?

Gemini's embedding model is optimised for different tasks. RETRIEVAL_DOCUMENT tells the model that this embedding will be stored in a database and later searched by a query embedding. This produces vectors optimised for asymmetric similarity — where a short query (e.g. "binary search") can correctly match a long document (e.g. a full implementation with 50 lines of code).

What Gets Embedded?

The text fed to the encoder is carefully composed to include all searchable metadata:

For Notes:

Title: <note title>
Content:
<description with HTML stripped>

For Code Snippets:

Title: <snippet title>
Language: <language>
Tags: <tag1>, <tag2>, <tag3>
Code:
<raw code>

Including the language and tags in the snippet's embedded text means a query like "python sorting algorithm" will correctly surface a Python snippet tagged #sorting, even if those exact words don't appear in the code body itself.


5. Stage 1 — Adaptive Retrieval Classifier

File: backend/utils/vectorSearch.jsneedsRetrieval(query)

This is the most important optimisation in the pipeline. Standard RAG systems blindly embed every query and run a vector search, even for questions that have nothing to do with the user's personal data. This wastes one API call and 200–350ms per query.

How It Works

The classifier runs regex patterns against the raw query string before any network call is made:

function needsRetrieval(query) {
  // Personal workspace patterns → ALWAYS retrieve
  const personalPatterns = [
    /\bmy\s+(code|function|snippet|note|implementation|...)\b/i,
    /\bi\s+(wrote|saved|created|added|stored|implemented)\b/i,
    /\bshow me my\b/i,
    /\bfind.*\b(snippet|note|code)\b/i,
    // ...
  ];

  // General knowledge patterns → SKIP retrieval
  const generalPatterns = [
    /^what\s+is\s+(a\s+|an\s+)?[a-z\s]+\??\s*$/i,   // "what is a linked list?"
    /^(explain|define)\s+[a-z\s]+\??\s*$/i,           // "explain recursion"
    /^how\s+does\s+[a-z\s]+\s+work\??\s*$/i,          // "how does quicksort work?"
    /^difference\s+between\b/i,                        // "difference between X and Y"
    /^write\s+(a\s+)?...to\b/i,                        // "write a function to..."
    // ...
  ];

  if (personalPatterns.some(p => p.test(query))) return true;   // force retrieve
  if (generalPatterns.some(p => p.test(query))) return false;   // skip
  return true;  // default: retrieve (conservative — never silently drop context)
}

Decision Priority

1. Personal patterns checked FIRST — "my code" always retrieves
2. General patterns checked SECOND — only if personal patterns don't match
3. Default: retrieve — err on the side of providing context

Cost / Latency Saved Per General Query

Operation Skipped Estimated Savings
gemini-embedding-001 API call ~150–300ms + API credit
MongoDB $vectorSearch (×2, parallel) ~20–50ms
Total saved per general query ~200–350ms + API cost

6. Stage 2 — Parallel Vector Search

File: backend/utils/vectorSearch.jsperformVectorSearch()

If retrieval is needed, both collections are searched simultaneously using Promise.all:

const [snippetResults, noteResults] = await Promise.all([
  CodeSnippet.aggregate([{
    $vectorSearch: {
      index: "vector_index",
      path: "embedding",
      queryVector: queryVector,
      numCandidates: 100,
      limit: 5,
      filter: { user: new mongoose.Types.ObjectId(userId) }
    }
  }, {
    $project: { _id:1, title:1, language:1, code:1, score: { $meta: "vectorSearchScore" } }
  }]),

  Note.aggregate([{ $vectorSearch: { ...same structure... } }])
]);

Key Parameters

Parameter Value Reasoning
numCandidates 100 HNSW explores 100 nearest neighbours before selecting limit results. 10× the limit is the standard MongoDB Atlas recommendation for high recall.
limit 5 Retrieve up to 5 from each collection before threshold filtering. Provides enough candidates to discard low-quality results while keeping context concise.
filter { user: userId } Pre-filter by owner ID inside the HNSW traversal. No other user's documents ever enter the result set. Critical for data privacy.

Why Parallel?

  • Serial: Search notes (~40ms) then snippets (~40ms) = ~80ms total
  • Parallel: Promise.all runs both simultaneously = ~40ms total

2× lower latency at zero cost.


7. Stage 3 — Relevance Threshold Filtering

File: backend/utils/vectorSearch.js

const RELEVANCE_THRESHOLD = 0.60;

const filteredSnippets = snippetResults.filter(s => s.score >= RELEVANCE_THRESHOLD);
const filteredNotes    = noteResults.filter(n    => n.score >= RELEVANCE_THRESHOLD);

What is Cosine Similarity?

Cosine similarity measures the angular distance between two vectors in high-dimensional space:

  • 1.0 = identical meaning (parallel vectors)
  • 0.0 = completely unrelated (perpendicular vectors)

In practice for semantic embeddings: relevant matches typically score 0.70–0.95, loosely related content scores 0.40–0.60.

Why 0.60?

Empirically tuned. Below 0.60, results are tangentially related at best. Example:

  • Query: "React hooks lifecycle"
  • Database contains: Python algorithms, Java data structures, no React
  • Best match: Python quicksort at score 0.45
  • Without threshold: Python quicksort is injected → LLM hallucinates a React-flavoured Python answer
  • With threshold: Score 0.45 < 0.60 → discarded → LLM answers from general knowledge correctly

Over-Fetch Strategy

We request limit: 5 from HNSW, knowing the threshold will further reduce this:

  • All 5 below threshold → inject nothing → LLM answers from general knowledge ✅
  • 3 of 5 above threshold → inject those 3 → LLM answers with partial context ✅
  • All 5 above threshold → inject all 5 → LLM has rich context ✅

8. Stage 4 — Adaptive Temperature

File: backend/routes/aiRoutes.jsgetTemperature()

Temperature controls LLM output randomness. Low = deterministic and focused. High = creative and diverse.

function getTemperature(skippedRetrieval) {
  return skippedRetrieval ? 0.7 : 0.2;
}

Why Adaptive?

Scenario Query Example Retrieval Temperature Reasoning
Personal RAG "explain my quicksort" ✅ Ran 0.2 Must stay strictly grounded in retrieved code. High temp risks hallucinating variable names not in the actual code.
General Knowledge "what is a binary tree?" ❌ Skipped 0.7 No personal context to be faithful to. A richer, more expressive educational explanation is better served with slightly higher creativity.
Code Actions explain/refactor/convert N/A 0.2 (default) Operating on code the user pasted inline. Accuracy and consistency are paramount.

Zero-Cost Implementation

The skippedRetrieval flag is already computed as a by-product of performVectorSearch():

// vectorSearch.js already returns this — no extra work:
return { snippets, notes, skippedRetrieval: true/false };

// aiRoutes.js reuses it to set temperature:
const context = await performVectorSearch(code, req.user._id);
temperature = getTemperature(context.skippedRetrieval);

// Passed into generationConfig for both primary AND fallback:
const generationConfig = { temperature };
const model = genAI.getGenerativeModel({ model: PRIMARY_MODEL, generationConfig });

Zero extra API calls. Zero extra latency. One line of logic.


9. Stage 5 — Prompt Engineering & Grounded Generation

File: backend/routes/aiRoutes.js

Prompt Structure

[System Persona]
"You are Relay's intelligent workspace assistant helping a developer with their personal code vault and notes."

[Retrieval Status Note]
→ General: "Retrieval was skipped — this appears to be a general knowledge question."
→ Personal: "Retrieved N snippet(s) and M note(s) above the 0.60 relevance threshold."

[Context Block — only injected if documents passed the threshold]
<workspace_context>
--- Snippet 1: QuickSort (javascript) [score: 0.847] ---
function quickSort(arr) { ... }

--- Note 1: Algorithm Study Notes [score: 0.782] ---
Binary search requires a sorted array and runs in O(log n)...
</workspace_context>

[User Question]
User Question: "How does my quicksort handle duplicates?"

[Strict Instructions]
- If workspace context is relevant, answer using it and CITE the exact Snippet or Note title.
- If context is present but irrelevant, ignore it and answer from general knowledge.
- If no context was retrieved, answer from general engineering knowledge.
- Be concise and technically precise.

Why XML-like Tags?

The <workspace_context> tags create a clear structural boundary that modern LLMs recognise as "injected external data, not user instruction." This:

  • Reduces prompt injection risk
  • Makes it clear to the model where retrieved context ends and the user question begins
  • Allows the model to evaluate context relevance independently of the instruction

Why Include Scores in the Injected Context?

[score: 0.847]

Including the cosine similarity score gives the LLM meta-information about relative relevance. If two snippets are injected with scores 0.85 and 0.62, the LLM can implicitly weight the more relevant one more heavily.

Citation Enforcement

"cite the exact Snippet or Note title"

This makes the AI's answer auditable. Users can verify "The AI referenced my QuickSort snippet — let me check if its explanation matches." This is the LLM equivalent of a footnote citation — turning a black-box answer into a traceable, verifiable one.


10. LLM Orchestration & Failover

File: backend/routes/aiRoutes.js

const PRIMARY_MODEL = "gemini-2.5-flash";
const BACKUP_MODEL  = "gemini-2.0-flash";

A dual-model failover strategy ensures zero dead-end errors:

const generationConfig = { temperature };  // same config for both models

try {
  const model = genAI.getGenerativeModel({ model: PRIMARY_MODEL, generationConfig });
  text = (await model.generateContent(prompt)).response.text();

} catch (primaryError) {
  // Primary failed (rate limit, quota, or service outage)
  const fallbackModel = genAI.getGenerativeModel({ model: BACKUP_MODEL, generationConfig });
  text = (await fallbackModel.generateContent(prompt)).response.text();
}

Critical: The same generationConfig (including the adaptive temperature) is passed to both models. If the primary is rate-limited, the fallback answers with identical parameters — same grounded, temperature-appropriate response with zero user-visible degradation.


11. Hallucination Prevention — Multi-Layered Strategy

Hallucination prevention is not a single feature — it is a defence-in-depth strategy with five independent layers:

Layer Mechanism Implementation
1 — Don't retrieve irrelevant context Adaptive classifier skips vector search entirely for general questions — irrelevant documents never enter the pipeline needsRetrieval() in vectorSearch.js
2 — Filter low-quality results Cosine threshold 0.60 discards loosely-related documents before they reach the prompt RELEVANCE_THRESHOLD filter in vectorSearch.js
3 — Low temperature on code tasks Temperature 0.2 makes the model deterministic and discourages creative deviation from retrieved context getTemperature() in aiRoutes.js
4 — Explicit escape hatch in prompt Prompt instructs: "if context is irrelevant, ignore it; if no context, use general knowledge" — the LLM is given a clear path instead of being forced to fill gaps Prompt template in aiRoutes.js
5 — Citation enforcement Requiring citations makes hallucinations immediately detectable — if the LLM cites a snippet, users can verify the answer against their actual code Prompt instructions in aiRoutes.js

Each layer alone reduces hallucinations. Together they make grounded responses the strong default.


12. Search Architecture — BM25 vs Vector Search

Relay uses two different search systems for two different use cases:

Standard Search Bar → BM25 via Apache Lucene

File: backend/routes/searchRoutes.js Technology: MongoDB Atlas Search ($search aggregation)

MongoDB Atlas Search is built on Apache Lucene — the same engine powering Elasticsearch. Lucene uses BM25 (Okapi BM25) for relevance scoring — a probabilistic TF-IDF model standard across the industry.

Features active in Relay's search:

  • Field boosting: title boosted 3× over description — title matches rank higher
  • Compound queries: must/should boolean clauses
  • User scoping: filter by createdBy/user field
  • Regex fallback: if Atlas Search is unavailable, falls back to $regex matching

AI Workspace (RAG) → Dense Vector Search

File: backend/utils/vectorSearch.js Technology: MongoDB Atlas Vector Search ($vectorSearch aggregation)

Pure dense retrieval using HNSW graph traversal and cosine similarity. Designed for semantic matching — understanding meaning, not keyword overlap.

Side-by-Side Comparison

Property BM25 (Atlas Search) Vector Search ($vectorSearch)
Matching type Keyword / lexical Semantic / conceptual
Handles typos ✅ Fuzzy matching (Levenshtein) ✅ Naturally (similar vectors)
Handles synonyms ❌ Requires explicit config ✅ Naturally
Handles exact variable names ✅ Precise token match ⚠️ Depends on embedding quality
Relevance scoring BM25 (TF-IDF based) Cosine similarity (0.0–1.0)
Use case in Relay Search bar — find specific notes fast RAG — find semantically related context for the LLM

What's Missing: True Hybrid RAG

The ideal RAG pipeline would run $search (BM25) AND $vectorSearch in parallel, then merge ranked lists using RRF (Reciprocal Rank Fusion). This ensures both semantic understanding AND exact keyword matching — critical for variable names like useDebounce or notation like O(n log n). This is the planned next evolution of Relay's RAG pipeline.


13. File Reference Map

File Role in the Pipeline
backend/utils/geminiEmbed.js Encoder — calls Gemini embedding API, returns 768-dim float array
backend/utils/vectorSearch.js Core RAG Logic — classifier + parallel vector search + threshold filter
backend/routes/aiRoutes.js Orchestrator — adaptive temperature + prompt building + LLM generation + failover
backend/models/Note.js Knowledge Base Schemaembedding field + pre('save') auto-embedding hook
backend/models/CodeSnippet.js Knowledge Base Schemaembedding field + pre('save') auto-embedding hook
backend/routes/searchRoutes.js BM25 Text Search (not RAG) — Atlas Search + Regex fallback for the search bar
backend/scripts/backfill_embeddings.js Maintenance Utility — regenerates embeddings for all pre-existing documents

14. Key Parameters at a Glance

Parameter Value File Reasoning
Embedding model gemini-embedding-001 geminiEmbed.js Latest Google embedding model
Embedding dimensions 768 geminiEmbed.js Truncated from 3072 for index compatibility
Task type RETRIEVAL_DOCUMENT geminiEmbed.js Optimised for asymmetric query-document similarity
numCandidates 100 vectorSearch.js 10× multiplier for HNSW recall (MongoDB recommendation)
limit 5 vectorSearch.js Max candidates per collection before threshold filtering
Relevance threshold 0.60 vectorSearch.js Empirically tuned — below this, results are rarely useful
Temperature (RAG) 0.2 aiRoutes.js Deterministic — strictly grounded in retrieved context
Temperature (general) 0.7 aiRoutes.js Expressive — suitable for educational explanations
Temperature (code actions) 0.2 aiRoutes.js Accuracy over creativity for explain/refactor/convert
Primary LLM gemini-2.5-flash aiRoutes.js Latest, fastest Gemini reasoning model
Fallback LLM gemini-2.0-flash aiRoutes.js Automatic failover if primary is rate-limited or down

15. Future Improvements

Improvement Description Benefit
True Hybrid Search (RRF) Run $search (BM25) AND $vectorSearch in parallel, merge with Reciprocal Rank Fusion Catches exact variable name matches that pure vector search can miss
Cross-Encoder Re-ranking After retrieval, pass candidates through a cross-encoder re-ranker before threshold filtering Higher precision — cross-encoders understand full query-document interaction, not just vector proximity
Streaming Responses Use generateContentStream() for token-by-token streaming to the frontend Perceived latency improvement — users see the answer appear in real-time
HyDE (Hypothetical Document Embedding) Rewrite the user query using the LLM before embedding ("imagine what the answer would look like") Closes the query-document vocabulary gap for better recall
Chunk-Level Embeddings Split long notes into overlapping chunks, embed each separately Better granularity — retrieve the relevant paragraph, not an entire 5000-word document
Conversation Memory Maintain a sliding window of the last N turns, embed and retrieve past context Makes the AI assistant stateful — remembers what was discussed earlier in a session

Last updated: July 2026 · Relay by Manas Raj