Advanced RAG Optimization: Chunking, Re-ranking & Hybrid Retrieval

typescript
function fixedSizeChunk(text: string, chunkSize: number, overlap: number): string[] {
  const words = text.split(/\s+/);
  const chunks: string[] = [];
  let start = 0;

  while (start < words.length) {
    const end = Math.min(start + chunkSize, words.length);
    chunks.push(words.slice(start, end).join(" "));
    start += chunkSize - overlap;
  }

  return chunks;
}

// Usage: 500-word chunks with 100-word overlap
const chunks = fixedSizeChunk(documentText, 500, 100);

typescript
function sentenceChunk(text: string, maxTokens: number): string[] {
  // Split on sentence boundaries
  const sentences = text.match(/[^.!?]+[.!?]+/g) || [text];
  const chunks: string[] = [];
  let currentChunk: string[] = [];
  let currentTokens = 0;

  for (const sentence of sentences) {
    const sentenceTokens = estimateTokens(sentence);

    if (currentTokens + sentenceTokens > maxTokens && currentChunk.length > 0) {
      chunks.push(currentChunk.join(" "));
      currentChunk = [];
      currentTokens = 0;
    }

    currentChunk.push(sentence.trim());
    currentTokens += sentenceTokens;
  }

  if (currentChunk.length > 0) {
    chunks.push(currentChunk.join(" "));
  }

  return chunks;
}

python
import numpy as np
from openai import OpenAI

client = OpenAI()

def semantic_chunk(text: str, threshold: float = 0.75, max_tokens: int = 800) -> list[str]:
    """Split text at semantic boundaries using embedding similarity."""
    sentences = split_sentences(text)

    # Embed each sentence
    response = client.embeddings.create(
        input=sentences,
        model="text-embedding-3-small"
    )
    embeddings = [e.embedding for e in response.data]

    # Compute cosine similarity between consecutive sentences
    chunks = []
    current_chunk = [sentences[0]]
    current_tokens = estimate_tokens(sentences[0])

    for i in range(1, len(sentences)):
        sim = cosine_similarity(embeddings[i - 1], embeddings[i])
        sent_tokens = estimate_tokens(sentences[i])

        # Break if similarity drops below threshold OR token budget exceeded
        if sim < threshold or current_tokens + sent_tokens > max_tokens:
            chunks.append(" ".join(current_chunk))
            current_chunk = [sentences[i]]
            current_tokens = sent_tokens
        else:
            current_chunk.append(sentences[i])
            current_tokens += sent_tokens

    if current_chunk:
        chunks.append(" ".join(current_chunk))

    return chunks

typescript
interface DocumentSection {
  heading: string;
  level: number; // h1=1, h2=2, etc.
  content: string;
  codeBlocks: string[];
}

function documentAwareChunk(
  markdown: string,
  maxTokens: number = 800
): Array<{ content: string; metadata: { heading: string; level: number } }> {
  const sections = parseMarkdownSections(markdown);
  const chunks: Array<{ content: string; metadata: { heading: string; level: number } }> = [];

  for (const section of sections) {
    const sectionText = `## ${section.heading}\n\n${section.content}`;
    const tokens = estimateTokens(sectionText);

    if (tokens <= maxTokens) {
      // Section fits in one chunk --- keep it whole
      chunks.push({
        content: sectionText,
        metadata: { heading: section.heading, level: section.level },
      });
    } else {
      // Section too large --- split by paragraphs, keeping heading as prefix
      const paragraphs = section.content.split(/\n\n+/);
      let currentContent = `## ${section.heading}\n\n`;
      let currentTokens = estimateTokens(currentContent);

      for (const para of paragraphs) {
        const paraTokens = estimateTokens(para);

        if (currentTokens + paraTokens > maxTokens && currentTokens > estimateTokens(`## ${section.heading}\n\n`)) {
          chunks.push({
            content: currentContent.trim(),
            metadata: { heading: section.heading, level: section.level },
          });
          currentContent = `## ${section.heading} (continued)\n\n`;
          currentTokens = estimateTokens(currentContent);
        }

        currentContent += para + "\n\n";
        currentTokens += paraTokens;
      }

      if (currentTokens > estimateTokens(`## ${section.heading}\n\n`)) {
        chunks.push({
          content: currentContent.trim(),
          metadata: { heading: section.heading, level: section.level },
        });
      }
    }
  }

  return chunks;
}

typescript
interface ParentChildChunk {
  parentId: string;
  parentContent: string;  // ~1500 tokens
  children: Array<{
    childId: string;
    content: string;       // ~300 tokens
    embedding: number[];
  }>;
}

async function parentChildRetrieval(
  query: string,
  topK: number = 5,
  expandToParent: boolean = true
): Promise<string[]> {
  const queryEmbedding = await embed(query);

  // Search against child chunks for precision
  const childResults = await vectorDb.search({
    embedding: queryEmbedding,
    table: "child_chunks",
    limit: topK,
  });

  if (!expandToParent) {
    return childResults.map((r) => r.content);
  }

  // Expand to parent chunks, deduplicating
  const parentIds = [...new Set(childResults.map((r) => r.parentId))];
  const parents = await db.query(
    `SELECT content FROM parent_chunks WHERE id = ANY($1)`,
    [parentIds]
  );

  return parents.map((p) => p.content);
}

Bi-encoder (embedding search):
  Query  → Encoder → query_vec  ─┐
                                  ├→ cosine_sim → score
  Doc    → Encoder → doc_vec   ─┘

Cross-encoder (re-ranking):
  [Query + Doc] → Encoder → relevance_score

typescript
import { CohereClient } from "cohere-ai";

const cohere = new CohereClient({ token: process.env.COHERE_API_KEY });

async function retrieveAndRerank(
  query: string,
  topK: number = 5
): Promise<Array<{ content: string; relevanceScore: number }>> {
  // Stage 1: Retrieve top-50 candidates via vector search
  const candidates = await vectorSearch(query, 50);

  // Stage 2: Re-rank with Cohere
  const reranked = await cohere.rerank({
    model: "rerank-english-v3.0",
    query,
    documents: candidates.map((c) => ({ text: c.content })),
    topN: topK,
  });

  return reranked.results.map((r) => ({
    content: candidates[r.index].content,
    relevanceScore: r.relevanceScore,
  }));
}

python
# ColBERT-style scoring (simplified)
import torch

def colbert_score(query_embeddings: torch.Tensor, doc_embeddings: torch.Tensor) -> float:
    """
    query_embeddings: (num_query_tokens, dim)
    doc_embeddings: (num_doc_tokens, dim)
    """
    # For each query token, find max similarity with any doc token
    similarity_matrix = query_embeddings @ doc_embeddings.T  # (Q, D)
    max_sim_per_query_token = similarity_matrix.max(dim=1).values  # (Q,)
    return max_sim_per_query_token.sum().item()

typescript
async function hydeRetrieval(query: string, topK: number = 5): Promise<string[]> {
  // Step 1: Generate a hypothetical answer
  const hypothetical = await llm.chat({
    model: "gpt-4o-mini",
    messages: [
      {
        role: "system",
        content: `Given a question, write a short paragraph that would be a good
answer to it. Do not say "I don't know." Write as if you are a knowledgeable
support article, even if you have to guess.`,
      },
      { role: "user", content: query },
    ],
  });

  // Step 2: Embed the hypothetical answer (not the original query)
  const hydeEmbedding = await embed(hypothetical.content);

  // Step 3: Search with the hypothetical embedding
  return vectorSearch(hydeEmbedding, topK);
}

typescript
async function multiQueryRetrieval(
  query: string,
  topK: number = 5
): Promise<string[]> {
  // Generate 3 reformulations
  const expansions = await llm.chat({
    model: "gpt-4o-mini",
    messages: [
      {
        role: "system",
        content: `Generate 3 alternative phrasings of the following question.
Return as a JSON array of strings. Each rephrasing should approach the question
from a different angle.`,
      },
      { role: "user", content: query },
    ],
    response_format: { type: "json_object" },
  });

  const queries = [query, ...JSON.parse(expansions.content).queries];

  // Retrieve with each query
  const allResults = await Promise.all(
    queries.map((q) => vectorSearch(q, topK * 2))
  );

  // Merge with Reciprocal Rank Fusion
  return reciprocalRankFusion(allResults, topK);
}

typescript
// Original: "Why does my webhook return 403 when using OAuth?"
// Step-back: "How does webhook authentication work?"

async function stepBackRetrieval(query: string, topK: number = 5): Promise<string[]> {
  const stepBack = await llm.chat({
    model: "gpt-4o-mini",
    messages: [
      {
        role: "system",
        content: `Given a specific question, generate a more general question that
would retrieve useful background context. Return only the question.`,
      },
      { role: "user", content: query },
    ],
  });

  const [specificResults, broadResults] = await Promise.all([
    vectorSearch(query, topK),
    vectorSearch(stepBack.content, topK),
  ]);

  // Combine: prioritize specific results, supplement with broad context
  return deduplicateAndMerge(specificResults, broadResults, topK);
}

typescript
interface RetrievalResult {
  chunkId: string;
  content: string;
  score: number;
  source: "bm25" | "vector" | "metadata";
}

async function hybridRetrieval(
  query: string,
  filters?: { category?: string; updatedAfter?: string },
  topK: number = 10
): Promise<RetrievalResult[]> {
  // Run all three retrieval methods in parallel
  const [bm25Results, vectorResults, metadataResults] = await Promise.all([
    // BM25 keyword search (handles exact matches, error codes, product names)
    bm25Search(query, topK * 3),
    // Dense vector search (handles semantic similarity)
    vectorSearch(query, topK * 3),
    // Metadata-filtered search (narrows by category, date, etc.)
    filters ? metadataFilteredSearch(query, filters, topK * 2) : Promise.resolve([]),
  ]);

  // Merge with Reciprocal Rank Fusion
  return reciprocalRankFusion(
    [bm25Results, vectorResults, metadataResults],
    topK
  );
}

function reciprocalRankFusion(
  resultSets: RetrievalResult[][],
  topK: number,
  k: number = 60 // RRF constant
): RetrievalResult[] {
  const scores = new Map<string, { score: number; result: RetrievalResult }>();

  for (const results of resultSets) {
    for (let rank = 0; rank < results.length; rank++) {
      const result = results[rank];
      const rrfScore = 1 / (k + rank + 1);
      const existing = scores.get(result.chunkId);

      if (existing) {
        existing.score += rrfScore;
      } else {
        scores.set(result.chunkId, { score: rrfScore, result });
      }
    }
  }

  return [...scores.values()]
    .sort((a, b) => b.score - a.score)
    .slice(0, topK)
    .map((s) => ({ ...s.result, score: s.score }));
}

typescript
async function advancedRAGPipeline(
  query: string,
  conversationHistory: Message[]
): Promise<string> {
  // 1. Query transformation: expand into multiple queries
  const expandedQueries = await expandQuery(query);

  // 2. Hybrid retrieval for each expanded query
  const allResults = await Promise.all(
    expandedQueries.map((q) => hybridRetrieval(q, undefined, 20))
  );

  // 3. Fuse results across all queries
  const fused = reciprocalRankFusion(allResults, 20);

  // 4. Re-rank with cross-encoder
  const reranked = await cohereRerank(query, fused, 5);

  // 5. Assemble context and generate
  const context = reranked.map((r) => r.content).join("\n\n---\n\n");

  const response = await llm.chat({
    model: "gpt-4o",
    messages: [
      {
        role: "system",
        content: `Answer the user's question using ONLY the provided context.
If the context does not contain enough information, say so.
Do not make up information.

Context:
${context}`,
      },
      ...conversationHistory,
      { role: "user", content: query },
    ],
  });

  return response.content;
}

python
def mean_reciprocal_rank(queries: list[dict]) -> float:
    """
    queries: [{ "query": str, "retrieved": [str], "relevant": set[str] }]
    """
    reciprocal_ranks = []
    for q in queries:
        for rank, doc_id in enumerate(q["retrieved"], 1):
            if doc_id in q["relevant"]:
                reciprocal_ranks.append(1.0 / rank)
                break
        else:
            reciprocal_ranks.append(0.0)
    return sum(reciprocal_ranks) / len(reciprocal_ranks)

python
import numpy as np

def ndcg_at_k(retrieved: list[str], relevant: dict[str, int], k: int) -> float:
    """
    relevant: { doc_id: relevance_grade } where grade is 0, 1, 2, 3
    """
    dcg = 0.0
    for i, doc_id in enumerate(retrieved[:k]):
        rel = relevant.get(doc_id, 0)
        dcg += (2**rel - 1) / np.log2(i + 2)  # i+2 because log2(1) = 0

    # Ideal DCG: sort by relevance grade descending
    ideal_rels = sorted(relevant.values(), reverse=True)[:k]
    idcg = sum((2**rel - 1) / np.log2(i + 2) for i, rel in enumerate(ideal_rels))

    return dcg / idcg if idcg > 0 else 0.0

python
def recall_at_k(retrieved: list[str], relevant: set[str], k: int) -> float:
    retrieved_set = set(retrieved[:k])
    return len(retrieved_set & relevant) / len(relevant) if relevant else 0.0