Real-time RAG at Scale: Multi-Stage Retrieval & Edge Caching for Low-Latency AI

Introduction & Industry Context

The Core Problem & Business/Technical Impact

• High Latency:
  • Vector Search Latency: Querying large vector databases can be slow, especially with increasing data volume and complex indexing.
  • LLM Inference Latency: Passing a large context window to powerful LLMs (e.g., Claude 3 Opus, GPT-4) incurs significant processing time, often measured in seconds.
  • Sequential Operations: Each step (retrieval, concatenation, LLM call) adds to the overall response time, making real-time interactions challenging.
• Exorbitant Costs:
  • Token Consumption: Large context windows, especially when many retrieved documents are only tangentially relevant, lead to high token consumption and increased LLM API costs.
  • Inefficient Retrieval: Retrieving and processing irrelevant documents wastes computational resources and LLM tokens.
• Suboptimal Accuracy & Hallucinations:
  • Context Dilution: If the initial retrieval pulls too many low-relevance documents, the LLM's attention may be diluted, leading to less accurate responses or even hallucinations based on weaker context.
  • Lack of Precision: A single-stage vector search may miss crucial nuances required for complex queries.

Architectural Concept & Solution Blueprint

1. Multi-Stage Retrieval: Instead of a single, broad retrieval, we introduce a two-phase retrieval process:
   • Initial Broad Retrieval: A fast, efficient vector search retrieves a larger set of potentially relevant documents (e.g., top 50-100 chunks). This prioritizes recall over precision.
   • Re-ranking with a Cross-Encoder: A smaller, specialized re-ranking model (often a transformer-based cross-encoder like sentence-transformers or a dedicated re-ranker API like Cohere's) takes the initial query and the retrieved documents and re-scores their relevance. This significantly improves precision, allowing us to select only the top 5-10 *most* relevant documents to pass to the final LLM.
2. Edge Caching with Cloudflare Workers: Many RAG queries are repetitive or common. Caching the responses at the edge (close to the user) can dramatically reduce latency and backend load. Cloudflare Workers, with their low-latency global network and built-in KV storage or Cache API, are ideal for this.

Client Request
     |
     V
Cloudflare Worker (Edge Cache)
     | (Cache Hit: Respond Directly)
     | (Cache Miss: Forward to Origin)
     V
API Gateway (e.g., Next.js 15 API Route / Node.js Express Service)
     |
     V
RAG Service Backend (Node.js/TypeScript)
     |---> 1. Initial Retrieval (Vector DB e.g., Qdrant, Pinecone)
     |---> 2. Re-ranking (Cross-Encoder / Re-ranker API)
     |---> 3. LLM Orchestration (e.g., Claude, OpenAI)
     |
     V
Response (Cached by Cloudflare Worker for future identical requests)

Step-by-Step Implementation

// src/types.ts
export interface Document {
  id: string;
  content: string;
  metadata?: Record<string, any>;
}

export interface Embedding {
  vector: number[];
  documentId: string;
}

export interface RetrievalResult {
  document: Document;
  score: number;
}

export interface VectorDBClient {
  upsert(embeddings: Embedding[], documents: Document[]): Promise<void>;
  query(vector: number[], topK: number): Promise<RetrievalResult[]>;
}

export interface EmbeddingClient {
  embed(text: string): Promise<number[]>;
}

export interface ReRankerClient {
  reRank(query: string, documents: Document[]): Promise<RetrievalResult[]>;
}

export interface LLMClient {
  generate(prompt: string, context: string): Promise<string>;
}

// src/ragService.ts
import { Document, EmbeddingClient, VectorDBClient, ReRankerClient, LLMClient, RetrievalResult } from './types';

export class RAGService {
  constructor(
    private embeddingClient: EmbeddingClient,
    private vectorDBClient: VectorDBClient,
    private reRankerClient: ReRankerClient,
    private llmClient: LLMClient
  ) {}

  async ingestDocument(doc: Document): Promise<void> {
    const embedding = await this.embeddingClient.embed(doc.content);
    await this.vectorDBClient.upsert([{ vector: embedding, documentId: doc.id }], [doc]);
    console.log(`Document "${doc.id}" ingested.`);
  }

  async getResponse(query: string): Promise<string> {
    console.log(`Processing query: "${query}"`);

    // Stage 1: Initial Broad Retrieval (Vector Search)
    const queryEmbedding = await this.embeddingClient.embed(query);
    const initialRetrievals = await this.vectorDBClient.query(queryEmbedding, 50); // Get top 50 docs
    console.log(`Initial retrieval found ${initialRetrievals.length} documents.`);

    const retrievedDocuments = initialRetrievals.map(r => r.document);

    // Stage 2: Re-ranking for Precision
    const reRankedResults = await this.reRankerClient.reRank(query, retrievedDocuments);
    const finalContextDocs = reRankedResults.slice(0, 5); // Take top 5 after re-ranking
    console.log(`Re-ranking selected ${finalContextDocs.length} most relevant documents.`);

    if (finalContextDocs.length === 0) {
      return "I couldn't find enough relevant information to answer your question.";
    }

    const contextString = finalContextDocs.map(r => r.document.content).join('\n---\n');

    // Stage 3: LLM Generation
    const prompt = `Based on the following context, answer the query accurately and concisely:\n\nContext:\n${contextString}\n\nQuery: ${query}\n\nAnswer:`;
    const llmResponse = await this.llmClient.generate(prompt, contextString);
    console.log('LLM generated response.');

    return llmResponse;
  }
}

// --- Mock Implementations for Demonstration ---
// In a real application, these would be calls to actual APIs/services (e.g., OpenAI, Qdrant, Cohere)

class MockEmbeddingClient implements EmbeddingClient {
  async embed(text: string): Promise<number[]> {
    // Simulate embedding with a simple hash or fixed vector for demo
    return Array.from({ length: 1536 }, () => Math.random()); // Example vector size
  }
}

class MockVectorDBClient implements VectorDBClient {
  private store: Map<string, Document> = new Map();
  private embeddings: Map<string, number[]> = new Map();

  async upsert(embeddings: Embedding[], documents: Document[]): Promise<void> {
    documents.forEach(doc => this.store.set(doc.id, doc));
    embeddings.forEach(emb => this.embeddings.set(emb.documentId, emb.vector));
  }

  async query(vector: number[], topK: number): Promise<RetrievalResult[]> {
    // Simulate finding topK closest documents (very basic simulation)
    const results: RetrievalResult[] = Array.from(this.store.values()).map(doc => ({
      document: doc,
      score: Math.random() // Simulate varying relevance
    }));
    results.sort((a, b) => b.score - a.score);
    return results.slice(0, topK);
  }
}

class MockReRankerClient implements ReRankerClient {
  async reRank(query: string, documents: Document[]): Promise<RetrievalResult[]> {
    // Simulate re-ranking based on a keyword match or simply random re-score
    const reRanked = documents.map(doc => ({
      document: doc,
      score: doc.content.includes(query.split(' ')[0]) ? Math.random() * 2 : Math.random() * 0.5 // Higher score for keyword match
    }));
    reRanked.sort((a, b) => b.score - a.score);
    return reRanked;
  }
}

class MockLLMClient implements LLMClient {
  async generate(prompt: string, context: string): Promise<string> {
    // Simulate LLM response based on context
    if (context.includes("Next.js 15") && prompt.includes("new features")) {
      return "Next.js 15 introduces enhanced React Server Components, asset optimization, and improved caching mechanisms for dynamic content. These features enable faster loading times and better developer experience.";
    }
    return "This is a simulated LLM response based on the provided context. I have processed your query.";
  }
}

// Example Usage (for Node.js backend)
async function runDemo() {
  const ragService = new RAGService(
    new MockEmbeddingClient(),
    new MockVectorDBClient(),
    new MockReRankerClient(),
    new MockLLMClient()
  );

  // Ingest some example documents
  await ragService.ingestDocument({
    id: 'doc1',
    content: 'Next.js 15 introduces a stable App Router, server actions, and improved asset optimization.',
    metadata: { source: 'blog' }
  });
  await ragService.ingestDocument({
    id: 'doc2',
    content: 'Cloudflare Workers provide a serverless execution environment at the edge, ideal for caching and routing.',
    metadata: { source: 'docs' }
  });
  await ragService.ingestDocument({
    id: 'doc3',
    content: 'React 19 brings new features like React Compiler, use hook, and enhanced server component support.',
    metadata: { source: 'docs' }
  });

  const response = await ragService.getResponse('What are the new features in Next.js 15?');
  console.log('\nFinal RAG Response:');
  console.log(response);
}

runDemo();

// worker.ts (Cloudflare Worker)

interface Env {
  RAG_API_URL: string; // Your RAG service backend URL
  // Could also add a KV namespace for more persistent caching
}

export default {
  async fetch(
    request: Request,
    env: Env,
    ctx: ExecutionContext
  ): Promise<Response> {
    const url = new URL(request.url);

    // Construct a cache key based on the request URL and relevant query parameters
    // For RAG, we might cache based on the 'query' parameter
    const cacheKey = new Request(url.toString(), request);
    const cache = caches.default; // Default Workers cache

    // Check if the response is already in cache
    let response = await cache.match(cacheKey);

    if (response) {
      console.log('Cache Hit for: ' + url.pathname + url.search);
      return response;
    }

    console.log('Cache Miss for: ' + url.pathname + url.search);

    // If not in cache, fetch from the origin RAG API
    // IMPORTANT: Make sure your backend can handle the original query parameter
    const originUrl = new URL(env.RAG_API_URL);
    originUrl.pathname = url.pathname;
    originUrl.search = url.search;

    // Clone the request to modify URL for origin fetch
    const originRequest = new Request(originUrl.toString(), {
      method: request.method,
      headers: request.headers,
      body: request.body,
      redirect: request.redirect,
      signal: request.signal,
    });

    response = await fetch(originRequest);

    // Ensure response is cacheable (e.g., status 200, no Set-Cookie headers)
    // Cache for 1 hour (3600 seconds)
    const cacheableResponse = new Response(response.body, response);
    cacheableResponse.headers.append('Cache-Control', 's-maxage=3600');

    // Store the response in cache
    // Await caching only if you want to ensure it happens before responding
    // ctx.waitUntil(cache.put(cacheKey, cacheableResponse.clone()));
    // Or simply put it and let it run in background
    cache.put(cacheKey, cacheableResponse.clone());

    return cacheableResponse;
  },
};

Performance Optimization & Best Practices

• Vector Database Optimization:
  • Indexing: Utilize Hierarchical Navigable Small Worlds (HNSW) or other efficient indexing algorithms for faster approximate nearest neighbor (ANN) searches.
  • Filtering: Apply metadata filters during vector search to narrow down the search space before similarity calculation, improving both speed and relevance.
  • Sharding & Replication: For extremely large datasets, shard your vector database and replicate it across regions to reduce latency and increase availability.
• Re-ranker Selection:
  • Choose a re-ranking model that balances performance and cost. Smaller, specialized cross-encoders (e.g., MiniLM-L6, BGE-Reranker) are often more cost-effective and faster than using a large LLM for re-ranking.
  • Consider managed re-ranking services (like Cohere Rerank) for ease of deployment and scalability.
• LLM Prompt Engineering & Model Selection:
  • Craft concise, clear prompts to guide the LLM effectively and minimize token usage.
  • Utilize smaller, fine-tuned LLMs for specific tasks when possible, saving costs and latency over general-purpose large models.
  • Experiment with different LLMs (Claude, GPT, Llama, Mixtral) to find the best balance of quality, speed, and cost for your specific use case.
• Caching Strategy (Cloudflare Workers):
  • Cache Key Design: Ensure cache keys are deterministic and capture all relevant aspects of the request (e.g., query parameters, headers). For RAG, the user's query is paramount.
  • TTL Management: Set appropriate Time-To-Live (TTL) values for cached responses. For rapidly changing data, a shorter TTL is necessary, while static knowledge might allow for longer caching.
  • Cache Invalidation: Implement mechanisms to purge or invalidate cache entries when the underlying knowledge base changes. This can be done via Cloudflare's API or by generating unique versions for cache keys.
  • Stale-While-Revalidate: Use stale-while-revalidate HTTP headers to serve cached content immediately while asynchronously updating it in the background for improved perceived performance.
• Monitoring & Observability:
  • Track key metrics: end-to-end latency, cache hit ratio, LLM token usage, vector database query times, and re-ranker latency.
  • Implement distributed tracing (e.g., OpenTelemetry) to pinpoint bottlenecks across the RAG pipeline.
  • Use Cloudflare Workers analytics for insights into cache performance and worker execution.

Business ROI & Future Outlook

• Significant Cost Reduction: By precisely selecting relevant context and aggressively caching, you can reduce LLM token consumption by 30-50%, leading to substantial savings on API costs.
• Dramatic Performance Improvements: Edge caching can slash latency for repeat queries from hundreds of milliseconds or even seconds down to tens of milliseconds. Multi-stage retrieval reduces the context window size, leading to faster LLM inference for cache misses.
• Enhanced User Experience: Faster, more accurate AI responses lead to higher user satisfaction, increased engagement, and improved conversion rates for AI-powered features.
• Scalability: Offloading traffic to the edge and optimizing backend processing allows your RAG system to handle significantly higher loads without requiring costly infrastructure upgrades.

Conclusion