@dooor-ai/toolkit/src/rag/helpers.ts

Guards, Evals & Observability for AI applications - works seamlessly with LangChain/LangGraph

/**
 * RAG helper functions for DOOOR AI Toolkit
 */

import { getCortexDBClient } from "../observability/cortexdb-client";
import { RAGContext } from "./context";
import { RAGMetadata, RAGResult } from "./types";

export interface RAGRetrievalResult {
  context: string;
  results: RAGResult[];
  metadata: RAGMetadata;
}

/**
 * Retrieve relevant context using RAG
 *
 * @param query - The user query
 * @param ragContext - RAG configuration (files, documents, embeddings, strategy)
 * @returns Retrieved context and metadata
 *
 * @example
 * ```typescript
 * const ragContext = new RAGContext({
 *   files: [pdfFile],
 *   embeddingProvider: "prod-gemini",
 *   strategy: RAGStrategy.HYDE,
 * });
 *
 * const { context, metadata } = await retrieveContext("How to authenticate?", ragContext);
 *
 * const result = await llm.invoke([
 *   { role: "user", content: `Context:\n${context}\n\nQuestion: How to authenticate?` }
 * ]);
 * ```
 */
export async function retrieveContext(
  query: string,
  ragContext: RAGContext
): Promise<RAGRetrievalResult> {
  const client = getCortexDBClient();

  const response = await client.invokeAI({
    prompt: query,
    usage: "chat",
    ragContext: ragContext,
  });

  if (!response.ragMetadata) {
    throw new Error("RAG metadata not returned from gateway");
  }

  // Map chunks from response (new format)
  const chunks: RAGResult[] = (response as any).ragChunks?.map((chunk: any) => ({
    text: chunk.text,
    score: chunk.score,
    metadata: chunk.metadata || {},
  })) || [];

  // Extract context from response
  return {
    context: response.text,
    results: chunks,
    metadata: response.ragMetadata,
  };
}

/**
 * Build a prompt with RAG context injected
 *
 * @param query - The user query
 * @param ragContext - RAG configuration
 * @returns Prompt with context injected
 *
 * @example
 * ```typescript
 * const prompt = await buildRAGPrompt("How to authenticate?", ragContext);
 * const result = await llm.invoke([{ role: "user", content: prompt }]);
 * ```
 */
export async function buildRAGPrompt(
  query: string,
  ragContext: RAGContext
): Promise<string> {
  const { context } = await retrieveContext(query, ragContext);

  return `Context from documents:
${context}

User question: ${query}

Answer based on the context provided above.`;
}