Cohere RAG Pipeline (Core Workflow A)

Overview

End-to-end Retrieval-Augmented Generation using Cohere's three core endpoints: Embed (vectorize), Rerank (sort by relevance), Chat (generate grounded answer with citations).

Prerequisites

• Completed cohere-install-auth setup
• cohere-ai package installed
• Understanding of vector similarity search

Instructions

Step 1: Embed Your Documents

import { CohereClientV2 } from 'cohere-ai';

const cohere = new CohereClientV2();

// Your knowledge base
const documents = [
  { id: 'doc1', text: 'Cohere Command A has 256K context and supports tool use.' },
  { id: 'doc2', text: 'Embed v4 generates 1024-dim vectors with 128K token context.' },
  { id: 'doc3', text: 'Rerank v3.5 scores relevance from 0 to 1 across 100+ languages.' },
  { id: 'doc4', text: 'The Chat API v2 requires model as a mandatory parameter.' },
  { id: 'doc5', text: 'Cohere supports structured JSON output via response_format.' },
];

// Embed documents for storage
const docEmbeddings = await cohere.embed({
  model: 'embed-v4.0',
  texts: documents.map(d => d.text),
  inputType: 'search_document',
  embeddingTypes: ['float'],
});

// Store vectors alongside document text in your vector DB
const vectors = docEmbeddings.embeddings.float;
console.log(`Embedded ${vectors.length} docs, ${vectors[0].length} dimensions each`);

Step 2: Search — Embed the Query

async function searchDocuments(query: string, topK = 10) {
  // Embed the query (note: inputType is 'search_query', not 'search_document')
  const queryEmbedding = await cohere.embed({
    model: 'embed-v4.0',
    texts: [query],
    inputType: 'search_query',
    embeddingTypes: ['float'],
  });

  const queryVector = queryEmbedding.embeddings.float[0];

  // Cosine similarity search (replace with your vector DB query)
  const scores = vectors.map((vec, i) => ({
    index: i,
    score: cosineSimilarity(queryVector, vec),
  }));

  return scores
    .sort((a, b) => b.score - a.score)
    .slice(0, topK)
    .map(s => documents[s.index]);
}

function cosineSimilarity(a: number[], b: number[]): number {
  let dot = 0, magA = 0, magB = 0;
  for (let i = 0; i < a.length; i++) {
    dot += a[i] * b[i];
    magA += a[i] * a[i];
    magB += b[i] * b[i];
  }
  return dot / (Math.sqrt(magA) * Math.sqrt(magB));
}