RAG: Answering from Your Own Data

A model's training data is not a reliable source for your private documents or today's information. Retrieval-Augmented Generation (RAG) adds a retrieval step before generation: find relevant chunks, put them in the prompt and ask the model to answer from that context.

What RAG gives you

  • Fresh answers: Re-index documents when they change instead of retraining the model.
  • Private context: Keep company information in your own retrieval layer.
  • Traceability: Return the chunks used to support the answer.
  • A narrower prompt: Give the model the few passages relevant to this question rather than the whole knowledge base.

The pipeline

Building a Simple RAG System

I’ll build a small-scale RAG system in TypeScript. For this example, I’ll store the documents and embeddings in memory. In a real application, you would use a dedicated vector database.

First, I need a way to calculate cosine similarity to find the most relevant documents.

// lib/rag/similarity.ts
export function cosineSimilarity(vecA: number[], vecB: number[]): number {
  if (vecA.length !== vecB.length) {
    throw new Error('Embeddings must have the same dimensions');
  }
  let dotProduct = 0, normA = 0, normB = 0;
  for (let i = 0; i < vecA.length; i++) {
    dotProduct += vecA[i] * vecB[i];
    normA += vecA[i] ** 2;
    normB += vecB[i] ** 2;
  }
  if (normA === 0 || normB === 0) return 0;
  return dotProduct / (Math.sqrt(normA) * Math.sqrt(normB));
}

Now I can create the RAG system. It will manage documents, find relevant context and generate answers.

// lib/rag/simple-rag.ts
import { embed, embedMany, generateText } from 'ai';
import { openai } from '@ai-sdk/openai';
import { cosineSimilarity } from './similarity';
import 'dotenv/config';

interface Document {
  id: string;
  text: string;
  embedding: number[];
}

class SimpleRAG {
  private documents: Map<string, Document> = new Map();

  // Step 1: Ingest and embed documents
  async addDocuments(texts: string[]) {
    const { embeddings } = await embedMany({
      model: openai.embedding('text-embedding-3-small'),
      values: texts,
    });

    for (let i = 0; i < texts.length; i++) {
      const id = `doc-${this.documents.size + 1}`;
      this.documents.set(id, {
        id,
        text: texts[i],
        embedding: embeddings[i],
      });
    }
    console.log(`Indexed ${texts.length} new documents.`);
  }

  // Step 2: Retrieve relevant documents
  private async retrieve(queryEmbedding: number[], topK: number): Promise<Document[]> {
    const allDocs = Array.from(this.documents.values());

    const similarities = allDocs.map(doc => ({
      ...doc,
      similarity: cosineSimilarity(queryEmbedding, doc.embedding),
    }));

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

  // Step 3: Generate the answer
  async query(question: string) {
    // Embed the user's question
    const { embedding: queryEmbedding } = await embed({
      model: openai.embedding('text-embedding-3-small'),
      value: question,
    });

    // Retrieve the top 3 most relevant documents
    const contextDocs = await this.retrieve(queryEmbedding, 3);

    const contextText = contextDocs
      .map((doc, i) => `Source [${i+1}]: ${doc.text}`)
      .join('\n\n');

    // Generate the final response using the context
    const { text } = await generateText({
      model: openai('gpt-4o'),
      system: "You are a helpful assistant. Answer the user's question based on the provided context. Cite your sources using the [Source #] format.",
      prompt: `Context:\n${contextText}\n\nQuestion:\n${question}`,
    });

    console.log('---');
    console.log('Answer:', text);
    console.log('\nSources:');
    contextDocs.forEach((doc, i) => console.log(`- [Source ${i+1}] ${doc.text}`));
    console.log('---');
  }
}

async function main() {
  const rag = new SimpleRAG();

  await rag.addDocuments([
    "The company's sick leave policy allows for 10 paid days per year.",
    'To submit an expense report, use the online portal and attach all receipts.',
    'The dress code is business casual, but jeans are allowed on Fridays.',
    'All employees receive a 15% discount on company products.',
    'Paid time off (PTO) requests must be submitted at least two weeks in advance.',
  ]);

  await rag.query('How many sick days do I get?');
  await rag.query('How do I ask for a day off?');
}

main();

The short version

  • Ingestion: You first need to process your documents (PDFs, text files, etc.), break them into manageable chunks, and generate embeddings for each chunk.
  • Storage: These embeddings are stored in a vector database, which is optimized for fast similarity searches.
  • Retrieval: When a user asks a question, you embed their query and use it to find the most relevant document chunks from your database.
  • Generation: You then pass the user's original question and the retrieved document chunks to an LLM, instructing it to answer the question based only on the provided context.

RAG is a practical pattern for grounding AI responses in your own data. It will not fix every problem, but it is often the right starting point when accuracy, freshness and source context matter.