Beyond Naive RAG: Advanced Chunking and Embedding Strategies for Superior Retrieval

type TextChunk = {
  content: string;
  metadata?: Record<string, any>;
};
 
/**
 * Chunks text by a fixed character size with a specified overlap.
 * Ensures chunks attempt to end/start on natural breaks (sentences) within overlap.
 * @param text The input document text.
 * @param chunkSize The maximum size of each chunk.
 * @param overlapSize The size of the overlap between chunks.
 * @returns An array of TextChunk objects.
 */
function fixedSizeChunker(text: string, chunkSize: number, overlapSize: number = 0): TextChunk[] {
  if (overlapSize >= chunkSize) {
    throw new Error("Overlap size must be less than chunk size.");
  }
 
  const chunks: TextChunk[] = [];
  let currentPos = 0;
 
  while (currentPos < text.length) {
    let endPos = Math.min(currentPos + chunkSize, text.length);
    let startPos = currentPos;
 
    // Adjust endPos to not cut sentences in half if possible, within the chunk boundary
    if (endPos < text.length) {
      const sentenceEnd = text.lastIndexOf('.', endPos - 1) + 1;
      const paragraphEnd = text.lastIndexOf('\n\n', endPos - 1) + 2;
      const lastRelevantBreak = Math.max(sentenceEnd, paragraphEnd);
 
      if (lastRelevantBreak > startPos && lastRelevantBreak < endPos) {
        endPos = lastRelevantBreak;
      }
    }
 
    const chunkContent = text.substring(startPos, endPos).trim();
    if (chunkContent) {
      chunks.push({ content: chunkContent });
    }
 
    // Move currentPos forward, accounting for overlap
    currentPos = endPos - overlapSize;
    if (currentPos <= startPos) { // Handle cases where overlap pushes currentPos back too far
      currentPos = startPos + chunkSize; // Force progress if chunk size is small or overlap large
    }
  }
  return chunks;
}
 
// Example usage
// const doc = "This is the first sentence. This is the second sentence. And here is the third sentence. A new paragraph starts here. With more content.";
// const fixedChunks = fixedSizeChunker(doc, 100, 20);
// console.log("Fixed Chunks:", fixedChunks);

// Assuming a simple sentence splitter (e.g., from Agno or a custom regex)
// For simplicity, I'll use a basic split for this example, but a robust solution
// needs a proper NLP sentence tokenizer.
function splitIntoSentences(text: string): string[] {
  // A simple regex split, not production-ready for all languages/edge cases
  return text.match(/[^.!?]+[.!?]+/g) || [];
}
 
// Mock embedding function for demonstration
// In a real scenario, this would call an API or a local model.
async function getSentenceEmbedding(sentence: string): Promise<number[]> {
  // Simulate an embedding call
  const encoder = new TextEncoder();
  const data = encoder.encode(sentence);
  let hash = 0;
  for (let i = 0; i < data.length; i++) {
    hash = (hash << 5) - hash + data[i];
    hash |= 0; // Convert to 32bit integer
  }
  return [hash % 1000 / 1000, (hash * 2) % 1000 / 1000, (hash * 3) % 1000 / 1000]; // Dummy 3D embedding
}
 
function cosineSimilarity(vecA: number[], vecB: number[]): number {
  let dotProduct = 0;
  let magnitudeA = 0;
  let magnitudeB = 0;
  for (let i = 0; i < vecA.length; i++) {
    dotProduct += vecA[i] * vecB[i];
    magnitudeA += vecA[i] * vecA[i];
    magnitudeB += vecB[i] * vecB[i];
  }
  magnitudeA = Math.sqrt(magnitudeA);
  magnitudeB = Math.sqrt(magnitudeB);
  if (magnitudeA === 0 || magnitudeB === 0) return 0;
  return dotProduct / (magnitudeA * magnitudeB);
}
 
/**
 * Chunks text semantically by detecting significant topic shifts.
 * @param text The input document text.
 * @param similarityThreshold Threshold below which a similarity score indicates a topic shift (e.g., 0.6-0.7).
 * @param minChunkSize Minimum number of sentences per chunk.
 * @returns An array of TextChunk objects.
 */
async function semanticChunker(
  text: string,
  similarityThreshold: number = 0.7,
  minChunkSize: number = 3
): Promise<TextChunk[]> {
  const sentences = splitIntoSentences(text.trim());
  if (sentences.length === 0) return [];
 
  const sentenceEmbeddings: number[][] = await Promise.all(
    sentences.map(s => getSentenceEmbedding(s))
  );
 
  const chunks: TextChunk[] = [];
  let currentChunkSentences: string[] = [];
  let currentChunkEmbeddings: number[][] = [];
 
  for (let i = 0; i < sentences.length; i++) {
    currentChunkSentences.push(sentences[i]);
    currentChunkEmbeddings.push(sentenceEmbeddings[i]);
 
    // If we have enough sentences, check for a semantic break
    if (currentChunkSentences.length >= minChunkSize && i < sentences.length - 1) {
      const currentAvgEmbedding = currentChunkEmbeddings.reduce(
        (acc, val) => acc.map((v, idx) => v + val[idx]),
        Array(sentenceEmbeddings[0].length).fill(0)
      ).map(v => v / currentChunkEmbeddings.length);
 
      const nextSentenceEmbedding = sentenceEmbeddings[i + 1];
      const similarityToNext = cosineSimilarity(currentAvgEmbedding, nextSentenceEmbedding);
 
      if (similarityToNext < similarityThreshold) {
        // Semantic break detected
        chunks.push({ content: currentChunkSentences.join(' ').trim() });
        currentChunkSentences = [];
        currentChunkEmbeddings = [];
      }
    }
  }
 
  // Add any remaining sentences as the last chunk
  if (currentChunkSentences.length > 0) {
    chunks.push({ content: currentChunkSentences.join(' ').trim() });
  }
 
  return chunks;
}
 
// Example usage (will need an actual embedding model)
// const complexDoc = "Quantum entanglement is a phenomenon where two particles become linked. Their states are interdependent. This has profound implications for quantum computing. On a completely different note, the stock market closed higher today. Technology stocks saw significant gains.";
// (async () => {
//   const semanticChunks = await semanticChunker(complexDoc, 0.7);
//   console.log("Semantic Chunks:", semanticChunks);
// })();

// Assume 'documents' is an array of strings, 'queries' is an array of strings,
// and 'groundTruth' is a Map<string, Set<string>> where key=query, value=set of original relevant chunk texts.
 
interface Retriever {
  retrieve(query: string, k: number): Promise<TextChunk[]>;
}
 
// Dummy Retriever implementation for demonstration
class DummyRetriever implements Retriever {
  private chunks: TextChunk[]; // All document chunks for a given strategy
  private embeddings: Map<string, number[]>; // Pre-computed chunk embeddings
  private embeddingFunction: (text: string) => Promise<number[]>;
 
  constructor(
    chunks: TextChunk[],
    embeddingFunction: (text: string) => Promise<number[]>
  ) {
    this.chunks = chunks;
    this.embeddingFunction = embeddingFunction;
    this.embeddings = new Map();
  }
 
  async initialize() {
    for (const chunk of this.chunks) {
      this.embeddings.set(chunk.content, await this.embeddingFunction(chunk.content));
    }
  }
 
  async retrieve(query: string, k: number): Promise<TextChunk[]> {
    const queryEmbedding = await this.embeddingFunction(query);
    const similarities: { chunk: TextChunk; score: number }[] = [];
 
    for (const chunk of this.chunks) {
      const chunkEmbedding = this.embeddings.get(chunk.content);
      if (chunkEmbedding) {
        similarities.push({
          chunk: chunk,
          score: cosineSimilarity(queryEmbedding, chunkEmbedding)
        });
      }
    }
 
    // Sort by similarity and return top k
    return similarities
      .sort((a, b) => b.score - a.score)
      .slice(0, k)
      .map(s => s.chunk);
  }
}
 
async function runEvaluation(
  retriever: Retriever,
  queries: string[],
  groundTruth: Map<string, Set<string>>,
  k: number = 5
): Promise<{ hitRate: number; mrr: number }> {
  let totalHit = 0;
  let totalReciprocalRank = 0;
 
  for (const query of queries) {
    const retrievedChunks = await retriever.retrieve(query, k);
    const relevantChunks = groundTruth.get(query) || new Set<string>();
 
    let hit = false;
    let rankOfFirstRelevant = 0;
 
    for (let i = 0; i < retrievedChunks.length; i++) {
      const retrievedContent = retrievedChunks[i].content;
      // Check if any part of the retrieved content overlaps significantly with a ground truth chunk
      // (This requires a more robust 'isRelevant' check than exact string match in real-world)
      const isRelevant = Array.from(relevantChunks).some(gt => retrievedContent.includes(gt) || gt.includes(retrievedContent));
 
      if (isRelevant) {
        hit = true;
        if (rankOfFirstRelevant === 0) {
          rankOfFirstRelevant = i + 1; // 1-based rank
        }
      }
    }
 
    if (hit) {
      totalHit++;
    }
    if (rankOfFirstRelevant > 0) {
      totalReciprocalRank += 1 / rankOfFirstRelevant;
    }
  }
 
  const hitRate = totalHit / queries.length;
  const mrr = totalReciprocalRank / queries.length;
 
  return { hitRate, mrr };
}
 
// Example usage with a real embedding function (e.g., calling OpenAI API)
// class OpenAIBaseEmbedding implements (text: string) => Promise<number[]> {
//   private apiKey: string;
//   constructor(apiKey: string) { this.apiKey = apiKey; }
//   async embed(text: string): Promise<number[]> {
//     const response = await fetch("https://api.openai.com/v1/embeddings", {
//       method: "POST",
//       headers: {
//         "Content-Type": "application/json",
//         "Authorization": `Bearer ${this.apiKey}`,
//       },
//       body: JSON.stringify({
//         input: text,
//         model: "text-embedding-3-small",
//       }),
//     });
//     const data = await response.json();
//     return data.data[0].embedding;
//   }
// }
 
// (async () => {
//   const documentText = `
//     Quantum computing is a rapidly emerging field that harnesses the principles of quantum mechanics to solve complex computational problems. Unlike classical computers which use bits, quantum computers use qubits, which can exist in multiple states simultaneously due to superposition. This allows quantum computers to process vast amounts of information in parallel.
 
//     One of the most promising applications of quantum computing is in drug discovery and materials science. Simulating molecular interactions with classical computers is extremely difficult, but quantum computers could revolutionize this area by accurately modeling complex chemical reactions.
 
//     Another significant application is in cryptography. Shor's algorithm, for example, can efficiently factor large numbers, posing a threat to many current encryption methods like RSA. However, quantum-safe cryptography is also an active area of research.
 
//     The development of quantum hardware is still in its early stages, with challenges in maintaining qubit coherence and scalability. Various technologies, including superconducting circuits, trapped ions, and photonic systems, are being explored to build stable quantum processors.
 
//     Beyond quantum computing, artificial intelligence continues to advance at an unprecedented pace. Large Language Models (LLMs) are transforming how we interact with information, enabling natural language understanding and generation. The integration of LLMs with specialized knowledge bases via RAG is crucial for building robust AI systems. My personal research focuses on combining these with MoE architectures for more efficient and intelligent agents, mimicking human cognitive processes.
//   `;
 
//   const queries = [
//     "What are the applications of quantum computing?",
//     "How do quantum computers differ from classical computers?",
//     "What are the challenges in quantum hardware development?",
//     "What is my research focus?", // A query specifically for the last paragraph
//   ];
 
//   const groundTruth = new Map<string, Set<string>>([
//     ["What are the applications of quantum computing?", new Set(["One of the most promising applications of quantum computing is in drug discovery and materials science.", "Another significant application is in cryptography."])],
//     ["How do quantum computers differ from classical computers?", new Set(["Unlike classical computers which use bits, quantum computers use qubits, which can exist in multiple states simultaneously due to superposition."])],
//     ["What are the challenges in quantum hardware development?", new Set(["The development of quantum hardware is still in its early stages, with challenges in maintaining qubit coherence and scalability."])],
//     ["What is my research focus?", new Set(["My personal research focuses on combining these with MoE architectures for more efficient and intelligent agents, mimicking human cognitive processes."])]
//   ]);
 
//   // === Run with Fixed-Size Chunks ===
//   const fixedChunks = fixedSizeChunker(documentText, 300, 50);
//   const openAIEmbedder = new OpenAIBaseEmbedding("YOUR_OPENAI_API_KEY"); // Replace with actual API key
//   const fixedRetriever = new DummyRetriever(fixedChunks, openAIEmbedder.embed);
//   await fixedRetriever.initialize();
//   const fixedResults = await runEvaluation(fixedRetriever, queries, groundTruth);
//   console.log("Fixed-Size Chunking Results:", fixedResults);
 
//   // === Run with Semantic Chunks ===
//   const semanticChunks = await semanticChunker(documentText, 0.7, 3); // Re-run init with actual embedder
//   const semanticRetriever = new DummyRetriever(semanticChunks, openAIEmbedder.embed);
//   await semanticRetriever.initialize();
//   const semanticResults = await runEvaluation(semanticRetriever, queries, groundTruth);
//   console.log("Semantic Chunking Results:", semanticResults);
// })();