Build a RAG System That Actually Works in Production

type DocType = "prose" | "code" | "markdown" | "structured";

function chunkBySemanticUnit(document: string, docType: DocType): string[] {
  switch (docType) {
    case "markdown":
      // Split at headers, keeping header with its content
      const sections = document.split(/(?=^#{1,3} )/m);
      return sections
        .filter(s => s.trim().length > 100)
        .flatMap(section => {
          // Section too large? Split at paragraph boundaries
          if (section.length > 2000) {
            return section
              .split(/\n\n+/)
              .filter(p => p.trim().length > 50);
          }
          return [section];
        });

    case "code":
      // Split at function/class boundaries
      return document
        .split(/(?=\n(?:export\s+)?(?:function|class|const|async\s+function|interface|type)\s)/)
        .filter(s => s.trim().length > 50);

    case "structured":
      // For tables, keep each row with header
      const lines = document.split("\n");
      const headerLine = lines.find(l => l.includes("|"));
      if (!headerLine) return [document];
      return lines
        .filter(l => l.includes("|") && !l.match(/^\|[-|]+\|$/))
        .map(row => `${headerLine}\n${row}`);

    default: // prose
      // Split at paragraph boundaries, merge short ones
      const paragraphs = document
        .split(/\n\n+/)
        .filter(p => p.trim().length > 50);

      const chunks: string[] = [];
      let current = "";

      for (const para of paragraphs) {
        if ((current + para).length > 1500 && current.length > 0) {
          chunks.push(current.trim());
          current = para;
        } else {
          current += "\n\n" + para;
        }
      }
      if (current.trim()) chunks.push(current.trim());
      return chunks;
  }
}

interface DocumentChunk {
  id: string;
  documentId: string;
  content: string;
  metadata: {
    source: string;
    title: string;
    section?: string;
    docType: DocType;
    pageNumber?: number;
    lastUpdated: Date;
    contentHash: string;   // For detecting when chunks need re-indexing
  };
  embedding: number[];
}

interface SearchResult {
  chunk: DocumentChunk;
  semanticScore: number;
  keywordScore: number;
  combinedScore: number;
}

async function hybridSearch(
  query: string,
  chunks: DocumentChunk[],
  options: { topK: number; semanticWeight?: number; filters?: MetadataFilters }
): Promise<SearchResult[]> {
  const { topK, semanticWeight = 0.7, filters } = options;
  const keywordWeight = 1 - semanticWeight;

  // Apply metadata filters first to narrow the search space
  const filteredChunks = filters ? applyFilters(chunks, filters) : chunks;

  // Semantic search via embedding similarity
  const queryEmbedding = await embedText(query);
  const results: SearchResult[] = filteredChunks.map(chunk => ({
    chunk,
    semanticScore: cosineSimilarity(queryEmbedding, chunk.embedding),
    keywordScore: 0,
    combinedScore: 0,
  }));

  // BM25-style keyword scoring
  const queryTerms = query.toLowerCase().split(/\s+/).filter(t => t.length > 2);
  for (const result of results) {
    const content = result.chunk.content.toLowerCase();
    const words = content.split(/\s+/).length;
    let score = 0;
    for (const term of queryTerms) {
      const occurrences = (content.match(new RegExp(`\b${term}\b`, "g")) || []).length;
      if (occurrences > 0) {
        // TF-IDF inspired scoring
        score += (1 + Math.log(occurrences)) / Math.sqrt(words);
      }
    }
    result.keywordScore = score;
  }

  // Normalize scores to [0, 1] range
  const maxSemantic = Math.max(...results.map(r => r.semanticScore), 0.001);
  const maxKeyword = Math.max(...results.map(r => r.keywordScore), 0.001);

  for (const result of results) {
    result.combinedScore =
      semanticWeight * (result.semanticScore / maxSemantic) +
      keywordWeight * (result.keywordScore / maxKeyword);
  }

  return results
    .sort((a, b) => b.combinedScore - a.combinedScore)
    .slice(0, topK);
}

async function expandQuery(originalQuery: string): Promise<string[]> {
  const response = await anthropic.messages.create({
    model: "claude-haiku-4-5",
    max_tokens: 300,
    system: `Generate 3 alternative phrasings of this search query that would match 
different but relevant documents. Focus on:
- Different terminology for the same concept
- More specific versions of the query
- Related questions that would have useful overlapping answers
Return a JSON array of strings only.`,
    messages: [{ role: "user", content: originalQuery }]
  });

  try {
    const text = response.content[0].type === "text" ? response.content[0].text : "[]";
    const alternatives = JSON.parse(text);
    return [originalQuery, ...alternatives.slice(0, 3)];
  } catch {
    return [originalQuery]; // Fall back to original if parsing fails
  }
}

async function retrieveWithExpansion(
  query: string,
  chunks: DocumentChunk[],
  topK: number
): Promise<SearchResult[]> {
  const expandedQueries = await expandQuery(query);

  // Run parallel searches for each query variant
  const allResultArrays = await Promise.all(
    expandedQueries.map(q => hybridSearch(q, chunks, { topK, semanticWeight: 0.7 }))
  );

  // Deduplicate: keep highest score for each chunk seen
  const bestScores = new Map<string, SearchResult>();

  for (const results of allResultArrays) {
    for (const result of results) {
      const existing = bestScores.get(result.chunk.id);
      if (!existing || result.combinedScore > existing.combinedScore) {
        bestScores.set(result.chunk.id, result);
      }
    }
  }

  return Array.from(bestScores.values())
    .sort((a, b) => b.combinedScore - a.combinedScore)
    .slice(0, topK);
}

async function rerankResults(
  query: string,
  candidates: SearchResult[],
  topN: number
): Promise<SearchResult[]> {
  if (candidates.length <= topN) return candidates;

  // Use LLM to score relevance of each candidate
  const scoringPrompt = `Score the relevance of each passage to the query.
Query: "${query}"

${candidates.map((r, i) => `[${i}] ${r.chunk.content.substring(0, 300)}...`).join("\n\n")}

Return a JSON array of numbers from 0-10 representing relevance scores for each passage in order.`;

  const response = await anthropic.messages.create({
    model: "claude-haiku-4-5",
    max_tokens: 200,
    messages: [{ role: "user", content: scoringPrompt }]
  });

  try {
    const text = response.content[0].type === "text" ? response.content[0].text : "[]";
    const scores: number[] = JSON.parse(text);

    return candidates
      .map((result, i) => ({ ...result, rerankedScore: scores[i] || 0 }))
      .sort((a: any, b: any) => b.rerankedScore - a.rerankedScore)
      .slice(0, topN);
  } catch {
    return candidates.slice(0, topN); // Fall back to original ranking
  }
}

function assembleStructuredContext(
  query: string,
  retrievedChunks: SearchResult[]
): string {
  // Group chunks by source document
  const byDocument = new Map<string, SearchResult[]>();
  for (const result of retrievedChunks) {
    const docId = result.chunk.documentId;
    if (!byDocument.has(docId)) byDocument.set(docId, []);
    byDocument.get(docId)!.push(result);
  }

  const contextParts: string[] = [
    `RETRIEVED CONTEXT FOR QUERY: "${query}"`,
    "=".repeat(50),
  ];

  let sourceIndex = 1;
  for (const [, docResults] of byDocument) {
    const meta = docResults[0].chunk.metadata;
    contextParts.push(`\n[Source ${sourceIndex}: ${meta.title} | ${meta.docType} | Updated: ${meta.lastUpdated.toLocaleDateString()}]`);

    // Sort chunks within document by page/section order
    const sorted = docResults.sort(
      (a, b) => (a.chunk.metadata.pageNumber || 0) - (b.chunk.metadata.pageNumber || 0)
    );

    for (const result of sorted) {
      contextParts.push(result.chunk.content);
    }
    sourceIndex++;
  }

  return contextParts.join("\n\n");
}

interface RAGOptions {
  topKRetrieval?: number;    // Initial candidates to retrieve
  topNRerank?: number;       // Final chunks after re-ranking
  semanticWeight?: number;   // Balance semantic vs keyword search
  enableQueryExpansion?: boolean;
  metadataFilters?: MetadataFilters;
}

async function productionRAGQuery(
  userQuery: string,
  knowledgeBase: DocumentChunk[],
  options: RAGOptions = {}
): Promise<{ answer: string; sources: string[] }> {
  const {
    topKRetrieval = 12,
    topNRerank = 5,
    semanticWeight = 0.7,
    enableQueryExpansion = true,
    metadataFilters,
  } = options;

  // Step 1: Retrieve candidates (with or without query expansion)
  const candidates = enableQueryExpansion
    ? await retrieveWithExpansion(userQuery, knowledgeBase, topKRetrieval)
    : await hybridSearch(userQuery, knowledgeBase, { topK: topKRetrieval, semanticWeight, filters: metadataFilters });

  // Step 2: Re-rank to get final context
  const finalChunks = await rerankResults(userQuery, candidates, topNRerank);

  // Step 3: Assemble structured context
  const context = assembleStructuredContext(userQuery, finalChunks);

  // Step 4: Generate grounded response
  const response = await anthropic.messages.create({
    model: "claude-sonnet-4-6",
    max_tokens: 1500,
    system: `You are a helpful assistant. Answer the user's question based ONLY on the provided context.
If the context does not contain enough information to answer fully, say what you know from the context and clearly indicate what is missing.
Always cite sources by referring to the [Source N] labels in the context.`,
    messages: [{
      role: "user",
      content: `${context}\n\n${"-".repeat(50)}\n\nQuestion: ${userQuery}`
    }]
  });

  const answer = response.content[0].type === "text" ? response.content[0].text : "";

  // Extract source citations from the answer
  const sourcesUsed = finalChunks
    .filter((_, i) => answer.includes(`[Source ${i + 1}]`) || answer.includes(`Source ${i + 1}`))
    .map(r => `${r.chunk.metadata.title} (${r.chunk.metadata.source})`);

  return { answer, sources: sourcesUsed };
}

class DocumentIndexer {
  async indexDocument(document: RawDocument): Promise<void> {
    // Detect document type for appropriate chunking
    const docType = this.detectDocType(document);

    // Chunk by semantic unit
    const textChunks = chunkBySemanticUnit(document.content, docType);

    // Generate embeddings in batches to respect rate limits
    const batchSize = 20;
    const chunks: DocumentChunk[] = [];

    for (let i = 0; i < textChunks.length; i += batchSize) {
      const batch = textChunks.slice(i, i + batchSize);
      const embeddings = await this.embedBatch(batch);

      for (let j = 0; j < batch.length; j++) {
        chunks.push({
          id: `${document.id}-chunk-${i + j}`,
          documentId: document.id,
          content: batch[j],
          metadata: {
            source: document.url,
            title: document.title,
            docType,
            lastUpdated: document.updatedAt,
            contentHash: this.hash(batch[j]),
          },
          embedding: embeddings[j],
        });
      }
    }

    // Remove old chunks and insert new ones atomically
    await this.vectorDB.transaction(async (tx) => {
      await tx.deleteWhere({ documentId: document.id });
      await tx.insertBatch(chunks);
    });

    console.log(`Indexed ${chunks.length} chunks for: ${document.title}`);
  }

  async handleDocumentDeletion(documentId: string): Promise<void> {
    await this.vectorDB.deleteWhere({ documentId });
    console.log(`Removed chunks for deleted document: ${documentId}`);
  }
}

interface RAGEvaluation {
  retrievalPrecision: number;  // Fraction of retrieved chunks that are relevant
  retrievalRecall: number;     // Fraction of relevant chunks that were retrieved
  answerFaithfulness: number;  // Fraction of answer claims supported by context
  answerRelevance: number;     // How well the answer addresses the question
}

async function evaluateRAG(
  testCases: Array<{ query: string; relevantDocIds: string[]; expectedAnswer: string }>,
  knowledgeBase: DocumentChunk[]
): Promise<RAGEvaluation> {
  const results = await Promise.all(
    testCases.map(async (tc) => {
      const { answer, sources } = await productionRAGQuery(tc.query, knowledgeBase);
      const retrievedIds = sources.map(s => extractDocId(s));

      const truePositives = retrievedIds.filter(id => tc.relevantDocIds.includes(id)).length;
      const precision = truePositives / Math.max(retrievedIds.length, 1);
      const recall = truePositives / Math.max(tc.relevantDocIds.length, 1);

      return { precision, recall, answer, expectedAnswer: tc.expectedAnswer };
    })
  );

  return {
    retrievalPrecision: average(results.map(r => r.precision)),
    retrievalRecall: average(results.map(r => r.recall)),
    answerFaithfulness: 0, // Requires LLM-as-judge evaluation
    answerRelevance: 0,    // Requires LLM-as-judge evaluation
  };
}