Cookbook: RAG Pipeline

import { AllternitClient } from '@allternit/sdk';

const client = new AllternitClient({
  apiKey: process.env.ALLTERNIT_API_KEY!,
});

const store = await client.memory.createStore({
  name: 'product-docs',
  embeddingModel: 'text-embedding-3-small',
});

import { readFileSync } from 'fs';
import { globSync } from 'glob';

const docs = globSync('./docs/**/*.md');

for (const path of docs) {
  const content = readFileSync(path, 'utf-8');
  const chunks = chunkText(content, { maxTokens: 512, overlap: 50 });

  for (const chunk of chunks) {
    await client.memory.addEntry(store.id, {
      content: chunk.text,
      metadata: {
        source: path,
        section: chunk.section,
        line_start: chunk.lineStart,
        line_end: chunk.lineEnd,
      },
    });
  }
}

console.log(`Ingested ${docs.length} documents`);

async function retrieve(query: string, topK = 5) {
  const results = await client.memory.search(store.id, {
    query,
    limit: topK,
  });

  return results.entries.map((entry) => ({
    content: entry.content,
    source: entry.metadata.source,
    score: entry.score,
  }));
}

async function askWithRAG(question: string) {
  // Retrieve
  const context = await retrieve(question);

  // Build prompt
  const contextBlock = context
    .map((c, i) => `[Source ${i + 1}: ${c.source}]\n${c.content}`)
    .join('\n\n---\n\n');

  const messages = [
    {
      role: 'system',
      content: `You are a helpful assistant. Answer the user's question using
        the provided context. Always cite sources as [Source N]. If the context
        doesn't contain the answer, say so.`,
    },
    {
      role: 'user',
      content: `Context:\n\n${contextBlock}\n\nQuestion: ${question}`,
    },
  ];

  // Generate
  const response = await client.chat.completions.create({
    model: 'gizzi-pro',
    messages,
    temperature: 0.1,
  });

  return {
    answer: response.choices[0].message.content,
    sources: context.map((c) => c.source),
  };
}

const result = await askWithRAG('How do I configure permission profiles?');
console.log(result.answer);
// "To configure permission profiles, use `gizzi config profile add`..."
// Citations: [Source 1], [Source 2]
console.log('Sources:', result.sources);

function chunkText(
  text: string,
  { maxTokens = 512, overlap = 50 }: { maxTokens: number; overlap: number }
) {
  const words = text.split(/\s+/);
  const chunks = [];
  let start = 0;

  while (start < words.length) {
    const end = Math.min(start + maxTokens, words.length);
    chunks.push({
      text: words.slice(start, end).join(' '),
      lineStart: start,
      lineEnd: end,
    });
    start = end - overlap;
  }

  return chunks;
}