RAG Pipeline

medium10 minrag-pipeline.ts

You have a folder of documents and want to ask questions against it.

What it shows

Assemble loader, embeddings and store with RAGPipelineBuilder
Ingest a directory and query by similarity

The code

rag-pipeline.ts

import { GoogleEmbeddingService, InMemoryEmbeddingAdapter } from '@cogitator-ai/memory';
import { RAGPipelineBuilder, TextLoader } from '@cogitator-ai/rag';
import { mkdir, writeFile } from 'node:fs/promises';
import { join } from 'node:path';

const apiKey = process.env.GOOGLE_API_KEY;
if (!apiKey) throw new Error('Set GOOGLE_API_KEY');

const docsDir = join(process.cwd(), 'knowledge');
await mkdir(docsDir, { recursive: true });
await writeFile(
  join(docsDir, 'typescript.txt'),
  'TypeScript adds static type checking to JavaScript at compile time. It supports interfaces, generics and type inference.'
);
await writeFile(
  join(docsDir, 'space.txt'),
  'The James Webb Space Telescope orbits the Sun at the L2 point, about 1.5 million kilometers from Earth.'
);

const pipeline = new RAGPipelineBuilder()
  .withLoader(new TextLoader())
  .withEmbeddingService(new GoogleEmbeddingService({ apiKey }))
  .withEmbeddingAdapter(new InMemoryEmbeddingAdapter())
  .withConfig({
    chunking: { strategy: 'recursive', chunkSize: 300, chunkOverlap: 50 },
    retrieval: { strategy: 'similarity', topK: 3, threshold: 0.3 },
  })
  .build();

const ingested = await pipeline.ingest(docsDir);
console.log(`Ingested ${ingested.documents} documents -> ${ingested.chunks} chunks`);

for (const question of ['How does TypeScript improve JavaScript?', 'How far is JWST from Earth?']) {
  const results = await pipeline.query(question);
  console.log(question);
  for (const result of results) {
    console.log(`  [${result.score.toFixed(3)}] ${result.content.slice(0, 80)}`);
  }
}

Run it

install · bash

pnpm add @cogitator-ai/memory @cogitator-ai/rag

EnvironmentGOOGLE_API_KEY

run · bash

GOOGLE_API_KEY=your-key npx tsx rag-pipeline.ts

Or run the full example from a clone of the repository (keys in .env):

repository · bash

npx tsx examples/rag/01-basic-retrieval.ts

RAG Pipeline

What it shows

The code

Run it

Learn more