RAG Pipeline

What is RAG?Copy Anchor Link

Retrieval-Augmented Generation grounds LLM responses in your actual data. Instead of relying solely on the model's training data, a RAG pipeline retrieves relevant documents at query time and includes them in the prompt context.

@cogitator-ai/rag provides a complete pipeline: load documents from any source, chunk them into embeddable segments, store embeddings in a vector database, retrieve relevant chunks at query time, and optionally rerank results for precision.

InstallationCopy Anchor Link

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

The RAG package uses @cogitator-ai/memory for embedding storage and search. You'll also need an embedding provider — see Embedding Providers for setup.

Optional peer dependencies based on your document sources:

pnpm add cheerio        # HTMLLoader, WebLoader
pnpm add papaparse      # CSVLoader
pnpm add pdf-parse      # PDFLoader

Quick StartCopy Anchor Link

import { RAGPipelineBuilder, MarkdownLoader } from '@cogitator-ai/rag';
import { InMemoryEmbeddingAdapter, OpenAIEmbeddingService } from '@cogitator-ai/memory';

const embeddings = new OpenAIEmbeddingService({
  apiKey: process.env.OPENAI_API_KEY!,
});

const store = new InMemoryEmbeddingAdapter();

const pipeline = new RAGPipelineBuilder()
  .withLoader(new MarkdownLoader({ stripFrontmatter: true }))
  .withEmbeddingService(embeddings)
  .withEmbeddingAdapter(store)
  .withConfig({
    chunking: { strategy: 'recursive', chunkSize: 512, chunkOverlap: 50 },
  })
  .build();

await pipeline.ingest('./docs/');

const results = await pipeline.query('How do agents use tools?');
for (const result of results) {
  console.log(`[${result.score.toFixed(3)}] ${result.content.slice(0, 100)}...`);
}

Pipeline ArchitectureCopy Anchor Link

┌──────────┐    ┌──────────┐    ┌────────────┐    ┌─────────┐
│  Loader  │───▶│ Chunker  │───▶│ Embeddings │───▶│  Store  │
└──────────┘    └──────────┘    └────────────┘    └─────────┘
                                                       │
┌──────────┐    ┌──────────┐                           │
│ Reranker │◀───│Retriever │◀──────────────────────────┘
└──────────┘    └──────────┘

Ingest flow: Loader reads documents from a source (files, URLs, etc.) → Chunker splits them into smaller segments → Embedding service converts chunks to vectors → Embedding adapter stores vectors for later search.

Query flow: Retriever embeds the query and searches the store for similar chunks → Reranker (optional) re-scores results using an LLM or Cohere for higher precision.

For persistent storage use PostgresAdapter (pgvector) as the embedding adapter — call setVectorDimensions(embeddings.dimensions) before connect(); see Memory Adapters.

RAGPipelineBuilderCopy Anchor Link

The builder is the primary way to construct a pipeline. All components are pluggable.

const pipeline = new RAGPipelineBuilder()
  .withLoader(loader) // required
  .withChunker(chunker) // optional — defaults based on config
  .withEmbeddingService(svc) // required
  .withEmbeddingAdapter(store) // required
  .withRetriever(retriever) // optional — defaults to the config's retrieval.strategy
  .withHybridSearch(search, weights) // needed for retrieval.strategy 'hybrid'
  .withQueryExpander(expand) // needed for retrieval.strategy 'multi-query'
  .withReranker(reranker) // optional
  .withConfig(config) // required (chunking config is required even with a custom chunker)
  .build();

If you don't provide a chunker, one is created automatically from your config's chunking.strategy. If you don't provide a retriever, one is created from retrieval.strategy: similarity → SimilarityRetriever, mmr → MMRRetriever (using mmrLambda), hybrid → HybridRetriever (requires withHybridSearch()), multi-query → MultiQueryRetriever over a similarity (or hybrid) base, capped at multiQueryCount variants (requires withQueryExpander()). A missing requirement (or an invalid config) fails at build() with a clear error. The retriever is fixed at build time, so retrieval.strategy cannot be switched per query.

You can also construct new RAGPipeline(config, { loader, chunker, embeddingService, embeddingAdapter, retriever, reranker }) directly when you already have every component; the config is validated the same way.

RAGPipelineCopy Anchor Link

The built pipeline exposes two main methods:

ingest(source)Copy Anchor Link

Loads documents from source, chunks them, embeds each document's chunks with one embedBatch() call, and stores the embeddings (at most 16 concurrent writes). A failed write throws with the adapter's error. Each stored chunk keeps the loader's document metadata (frontmatter, page numbers, CSV columns, page title/URL…) plus documentId, source, sourceType, order, startOffset and endOffset. Whitespace-only chunks are skipped.

const { documents, chunks } = await pipeline.ingest('./data/knowledge-base/');
console.log(`Ingested ${documents} documents (${chunks} chunks)`);

query(text, options?)Copy Anchor Link

Retrieves relevant chunks for the given query. Options (topK, threshold, mmrLambda, multiQueryCount) override the configured retrieval defaults; options left undefined keep the configured values. Each RetrievalResult has chunkId, documentId, content, score, source (the document path or URL) and metadata. If a reranker is configured and reranking.enabled is true in the config, results are reranked (keeping reranking.topN) before returning.

const results = await pipeline.query('What is the agent lifecycle?', {
  topK: 5,
  threshold: 0.7,
});

getStats()Copy Anchor Link

Returns ingestion and query statistics.

const stats = pipeline.getStats();
// { documentsIngested: 12, chunksStored: 347, queriesProcessed: 5 }

ConfigurationCopy Anchor Link

Pipeline config is validated with Zod at build time (RAGPipelineConfigSchema, exported along with ChunkingConfigSchema, RetrievalConfigSchema and RerankingConfigSchema):

import type { RAGPipelineConfigInput } from '@cogitator-ai/rag';

const config: RAGPipelineConfigInput = {
  chunking: {
    strategy: 'recursive', // 'fixed' | 'recursive' | 'semantic'
    chunkSize: 512,
    chunkOverlap: 50,
    separators: ['\n\n', '\n', '. ', ' '], // recursive only
  },
  retrieval: {
    strategy: 'similarity', // 'similarity' | 'mmr' | 'hybrid' | 'multi-query'
    topK: 10,
    threshold: 0.0,
    mmrLambda: 0.7, // required for 'mmr'
    multiQueryCount: 3, // required for 'multi-query' (max generated variants)
  },
  reranking: {
    enabled: true,
    topN: 5,
  },
};

Agent ToolsCopy Anchor Link

ragTools(pipeline) returns rag_search and rag_ingest tools whose Zod schemas plug straight into tool() from @cogitator-ai/core:

import { Agent, tool } from '@cogitator-ai/core';
import { ragTools } from '@cogitator-ai/rag';

const [search, ingest] = ragTools(pipeline, {
  allowedRoots: ['./knowledge-base'],
  allowUrls: false,
});

const agent = new Agent({
  name: 'docs-assistant',
  model: 'google/gemini-3.8-flash',
  instructions: 'Use rag_search before answering and cite the source.',
  tools: [tool(search), tool(ingest)],
});

rag_search takes { query, limit?, threshold? } and falls back to the pipeline's configured topK/threshold when the model omits limit/threshold; it returns { success, query, results, count }. rag_ingest takes { source } and returns { success, source, documents, chunks }. Both return { success: false, error } instead of throwing. createSearchTool(pipeline) and createIngestTool(pipeline, options) build them individually. rag_ingest reads whatever source the model passes, so restrict it with allowedRoots (paths are canonicalized — .. traversal and symlink escapes are rejected) and allowUrls: false when it is exposed to untrusted input.

Next StepsCopy Anchor Link

• Document Loaders — load from text, markdown, JSON, CSV, HTML, PDF, or web URLs
• Chunking Strategies — fixed, recursive, and semantic chunking
• Retrieval Strategies — similarity, MMR, hybrid, and multi-query retrieval
• Reranking — LLM-based and Cohere reranking