@devilsdev/rag-pipeline-utils 2.4.5 → 2.6.0

package/src/adapters/ai-sdk.js

@@ -0,0 +1,215 @@
+'use strict';
+
+/**
+ * Adapter for the Vercel AI SDK (`ai` package): use an AI SDK embedding
+ * model and language model as this library's `embedder` and `llm`.
+ *
+ * The AI SDK is not a dependency. You pass in the functions you already
+ * import from it, so any AI SDK version from 5 on works, from ESM or
+ * CommonJS, and this package keeps supporting Node 18 (AI SDK 7 needs Node
+ * 22 and is ESM-only).
+ */
+
+const { buildRagPrompt } = require('../connectors/rag-prompt');
+const { componentCacheKey } = require('../cache/component-keys');
+
+/** Identity of an AI SDK model object (provider and model id). */
+const modelId = (model) =>
+  `${model.provider || 'unknown'}:${model.modelId || model.modelName || 'unknown'}`;
+
+/**
+ * @typedef {Object} AiSdkFunctions
+ * @property {Function} [embed] - `embed` from 'ai' (needed for an embedder)
+ * @property {Function} [embedMany] - `embedMany` from 'ai' (batch embedding)
+ * @property {Function} [generateText] - `generateText` from 'ai' (needed for an llm)
+ * @property {Function} [streamText] - `streamText` from 'ai' (streaming)
+ */
+
+function requireFunction(ai, name, purpose) {
+  if (typeof ai[name] !== 'function') {
+    throw new Error(
+      `createAiSdkAdapter: ${purpose} needs ai.${name}. Pass the AI SDK functions, e.g. import { ${name} } from "ai".`,
+    );
+  }
+}
+
+/**
+ * Options for every generateText / streamText call. The RAG system prompt
+ * (retrieved text is untrusted reference data, not instructions) is passed
+ * as both `instructions` (AI SDK 7) and `system` (AI SDK 5 and 6; 7 reads
+ * `instructions ?? system`), so it reaches the model on every version.
+ */
+function textCallOptions(model, query, docs, settings, signal, systemPrompt) {
+  const { system, user } = buildRagPrompt(query, docs, { systemPrompt });
+  return {
+    ...settings,
+    model,
+    ...(system && { system, instructions: system }),
+    prompt: user,
+    ...(signal && { abortSignal: signal }),
+  };
+}
+
+/** Per-call settings the pipeline may pass through to the model. */
+const CALL_SETTINGS = [
+  'temperature',
+  'maxOutputTokens',
+  'topP',
+  'topK',
+  'presencePenalty',
+  'frequencyPenalty',
+  'seed',
+  'stopSequences',
+  'maxRetries',
+];
+
+function pickSettings(options) {
+  const out = {};
+  for (const key of CALL_SETTINGS) {
+    if (options[key] !== undefined) out[key] = options[key];
+  }
+  return out;
+}
+
+/**
+ * Wrap AI SDK models as an embedder and an llm for createRagPipeline().
+ *
+ * @example
+ * import { embed, embedMany, generateText, streamText } from "ai";
+ * import { openai } from "@ai-sdk/openai";
+ * import { createAiSdkAdapter, createRagPipeline } from "@devilsdev/rag-pipeline-utils";
+ *
+ * const { embedder, llm } = createAiSdkAdapter({
+ *   ai: { embed, embedMany, generateText, streamText },
+ *   embeddingModel: openai.embedding("text-embedding-3-small"),
+ *   languageModel: openai("gpt-4o-mini"),
+ * });
+ * const pipeline = createRagPipeline({ embedder, retriever, llm });
+ *
+ * @param {Object} options
+ * @param {AiSdkFunctions} options.ai - Functions imported from 'ai'
+ * @param {Object} [options.embeddingModel] - AI SDK embedding model
+ * @param {Object} [options.languageModel] - AI SDK language model
+ * @param {Object} [options.settings] - Defaults for every text call, e.g.
+ *   `{ temperature: 0, maxOutputTokens: 512 }`
+ * @param {string} [options.systemPrompt] - Replaces the built-in RAG system
+ *   prompt
+ * @returns {{ embedder: Object|null, llm: Object|null }}
+ */
+function createAiSdkAdapter(options = {}) {
+  const {
+    ai,
+    embeddingModel,
+    languageModel,
+    settings = {},
+    systemPrompt,
+  } = options;
+  if (!ai || typeof ai !== 'object') {
+    throw new Error(
+      'createAiSdkAdapter: pass `ai`, the AI SDK functions, e.g. { embed, embedMany, generateText, streamText } from "ai".',
+    );
+  }
+  if (!embeddingModel && !languageModel) {
+    throw new Error(
+      'createAiSdkAdapter: pass an embeddingModel, a languageModel, or both.',
+    );
+  }
+
+  let embedder = null;
+  if (embeddingModel) {
+    requireFunction(ai, 'embed', 'an embeddingModel');
+    const embedOne = async (text, { signal } = {}) => {
+      const { embedding } = await ai.embed({
+        model: embeddingModel,
+        value: String(text),
+        ...(signal && { abortSignal: signal }),
+      });
+      return embedding;
+    };
+    embedder = {
+      metadata: { name: 'ai-sdk-embedder', version: '1.0.0', type: 'embedder' },
+      getCacheKey: (role) =>
+        role === 'embedder'
+          ? componentCacheKey('ai-sdk-embed', {
+              model: modelId(embeddingModel),
+            })
+          : null,
+      embed: embedOne,
+      embedQuery: embedOne,
+      /** One vector per text; uses embedMany when it was passed in. */
+      async embedBatch(texts, { signal } = {}) {
+        if (typeof ai.embedMany !== 'function') {
+          const vectors = [];
+          for (const text of texts)
+            vectors.push(await embedOne(text, { signal }));
+          return vectors;
+        }
+        const { embeddings } = await ai.embedMany({
+          model: embeddingModel,
+          values: texts.map(String),
+          ...(signal && { abortSignal: signal }),
+        });
+        return embeddings;
+      },
+    };
+  }
+
+  let llm = null;
+  if (languageModel) {
+    requireFunction(ai, 'generateText', 'a languageModel');
+    llm = {
+      metadata: { name: 'ai-sdk-llm', version: '1.0.0', type: 'llm' },
+      // The AI SDK leaves an unset temperature to the provider, whose default
+      // is usually 1, so unset counts as non-deterministic: generation
+      // caching needs `settings.temperature: 0`.
+      config: { temperature: settings.temperature ?? 1 },
+      getCacheKey: (role) =>
+        role === 'llm'
+          ? componentCacheKey('ai-sdk', {
+              model: modelId(languageModel),
+              settings: JSON.stringify(settings),
+              systemPrompt,
+            })
+          : null,
+      async generate(query, docs = [], callOptions = {}) {
+        const { text } = await ai.generateText(
+          textCallOptions(
+            languageModel,
+            query,
+            docs,
+            { ...settings, ...pickSettings(callOptions) },
+            callOptions.signal,
+            systemPrompt,
+          ),
+        );
+        return text;
+      },
+    };
+    if (typeof ai.streamText === 'function') {
+      llm.generateStream = async function* generateStream(
+        query,
+        docs = [],
+        callOptions = {},
+      ) {
+        const result = ai.streamText(
+          textCallOptions(
+            languageModel,
+            query,
+            docs,
+            { ...settings, ...pickSettings(callOptions) },
+            callOptions.signal,
+            systemPrompt,
+          ),
+        );
+        for await (const token of result.textStream) {
+          yield { token, done: false };
+        }
+        yield { token: '', done: true };
+      };
+    }
+  }
+
+  return { embedder, llm };
+}
+
+module.exports = { createAiSdkAdapter };

package/src/agentic/agentic-pipeline.js

@@ -5,6 +5,8 @@ const { DAG } = require('../dag/dag-engine');
 const { QueryPlanner } = require('./query-planner');
 const { IterativeRetriever } = require('./iterative-retriever');
 const { SelfCritiqueChecker } = require('./self-critique');
+const { embedQuery, canEmbed } = require('../utils/embed-query');
+const { createRunAbort } = require('../utils/run-abort');
 
 /**
  * Default configuration for the AgenticPipeline.
@@ -14,6 +16,7 @@ const DEFAULT_CONFIG = {
   maxConcurrency: 3,
   timeout: 60000,
   enablePlanning: true,
+  enableLlmPlanning: false,
   enableCritique: true,
   enableIterativeRetrieval: false,
   maxIterations: 3,
@@ -33,6 +36,9 @@ class AgenticPipeline extends EventEmitter {
    * @param {number} [options.maxConcurrency=3] - Max concurrent DAG node execution
    * @param {number} [options.timeout=60000] - Global timeout in ms
    * @param {boolean} [options.enablePlanning=true] - Whether to decompose queries
+   * @param {boolean} [options.enableLlmPlanning=false] - Ask the llm for the
+   *   sub-queries (strictly validated, heuristic fallback); costs one extra
+   *   llm.generate() call per run
    * @param {boolean} [options.enableCritique=true] - Whether to self-critique answers
    * @param {boolean} [options.enableIterativeRetrieval=false] - Whether to use iterative retrieval
    * @param {number} [options.maxIterations=3] - Max iterations for iterative retrieval
@@ -62,7 +68,16 @@ class AgenticPipeline extends EventEmitter {
    * @param {object} params.retriever - Retriever instance with .retrieve() method
    * @param {object} params.llm - LLM instance with .generate() method
    * @param {Array<number>} [params.queryVector] - Optional pre-computed query vector
+   * @param {object} [params.embedder] - Optional embedder; when given, each
+   *   sub-query is embedded on its own instead of reusing queryVector
    * @param {object} [params.options={}] - Extra options passed to llm.generate()
+   * @param {AbortSignal} [params.options.signal] - Cancels the run. The run
+   *   has its own AbortController (also aborted after `timeout`); its signal
+   *   is passed to retriever.retrieve({ ..., signal }), embedder.embedQuery /
+   *   embed(text, { signal }) and llm.generate(query, docs, { ...options,
+   *   signal }), and no later stage starts once it aborts. The returned
+   *   promise then rejects with Error('Execution timeout') or
+   *   Error('Execution aborted').
    * @returns {Promise<{
    *   success: boolean,
    *   query: string,
@@ -74,7 +89,7 @@ class AgenticPipeline extends EventEmitter {
    * }>}
    * @throws {Error} If retriever or llm is missing
    */
-  async run({ query, retriever, llm, queryVector, options = {} }) {
+  async run({ query, retriever, llm, queryVector, embedder, options = {} }) {
     // 1. Validate inputs
     if (!retriever || typeof retriever.retrieve !== 'function') {
       throw new Error(
@@ -92,10 +107,44 @@ class AgenticPipeline extends EventEmitter {
       );
     }
 
+    // One AbortController per run: aborted on timeout or when the caller's
+    // options.signal aborts. Its signal goes to every component call.
+    const runAbort = createRunAbort({
+      timeout: this.config.timeout,
+      signal: options && options.signal,
+    });
+    try {
+      return await runAbort.race(
+        this._execute(
+          { query, retriever, llm, queryVector, embedder, options },
+          runAbort,
+        ),
+      );
+    } finally {
+      runAbort.cleanup();
+    }
+  }
+
+  /**
+   * @param {object} params - Validated run() parameters
+   * @param {object} runAbort - Per-run cancellation (see utils/run-abort)
+   * @returns {Promise<object>}
+   * @private
+   */
+  async _execute(
+    { query, retriever, llm, queryVector, embedder, options },
+    runAbort,
+  ) {
+    const { signal } = runAbort;
+    const llmOptions = { ...(options || {}), signal };
+
     // Phase 1: Plan
     let subQueries;
     if (this.config.enablePlanning) {
-      const plan = this.planner.plan(query);
+      const plan = this.config.enableLlmPlanning
+        ? await this.planner.plan(query, { llm, signal })
+        : this.planner.plan(query);
+      runAbort.throwIfAborted();
       subQueries = plan.subQueries;
       this.emit('planned', plan);
     } else {
@@ -111,19 +160,31 @@ class AgenticPipeline extends EventEmitter {
       dag.addNode(
         `retrieve-${i}`,
         async () => {
+          runAbort.throwIfAborted();
+          // queryVector belongs to the original query; a decomposed
+          // sub-query needs its own embedding when an embedder is available.
+          const needsOwnVector =
+            canEmbed(embedder) && (sq.subQuery !== query || !queryVector);
+          const subQueryVector = needsOwnVector
+            ? await embedQuery(embedder, sq.subQuery, { signal })
+            : queryVector;
+          runAbort.throwIfAborted();
           if (this.iterativeRetriever) {
             const result = await this.iterativeRetriever.retrieve({
               query: sq.subQuery,
               retriever,
               llm,
-              queryVector,
+              queryVector: subQueryVector,
+              embedder,
+              signal,
             });
             return result.results;
           }
           return await retriever.retrieve({
             query: sq.subQuery,
-            queryVector,
+            queryVector: subQueryVector,
             topK: 10,
+            signal,
           });
         },
         { timeout: this.config.timeout / 2, priority: sq.priority },
@@ -132,6 +193,7 @@ class AgenticPipeline extends EventEmitter {
 
     // Add merge node
     dag.addNode('merge', async (inputs) => {
+      runAbort.throwIfAborted();
       const allResults = Array.isArray(inputs)
         ? inputs.flat()
         : [inputs].flat();
@@ -155,8 +217,10 @@ class AgenticPipeline extends EventEmitter {
     // Add critique + generate nodes
     if (this.config.enableCritique) {
       dag.addNode('pre-generate', async (mergedResults) => {
+        runAbort.throwIfAborted();
         // Generate initial answer for critique
-        const answer = await llm.generate(query, mergedResults, options);
+        const answer = await llm.generate(query, mergedResults, llmOptions);
+        runAbort.throwIfAborted();
         const critique = this.critiqueChecker.check(
           String(answer),
           mergedResults,
@@ -167,6 +231,7 @@ class AgenticPipeline extends EventEmitter {
       dag.connect('merge', 'pre-generate');
 
       dag.addNode('generate', async ({ answer, mergedResults, critique }) => {
+        runAbort.throwIfAborted();
         if (critique.approved) {
           return { answer, results: mergedResults, critique };
         }
@@ -178,7 +243,7 @@ class AgenticPipeline extends EventEmitter {
         const refinedAnswer = await llm.generate(
           refinedPrompt,
           mergedResults,
-          options,
+          llmOptions,
         );
         return {
           answer: String(refinedAnswer),
@@ -189,7 +254,8 @@ class AgenticPipeline extends EventEmitter {
       dag.connect('pre-generate', 'generate');
     } else {
       dag.addNode('generate', async (mergedResults) => {
-        const answer = await llm.generate(query, mergedResults, options);
+        runAbort.throwIfAborted();
+        const answer = await llm.generate(query, mergedResults, llmOptions);
         return { answer: String(answer), results: mergedResults };
       });
       dag.connect('merge', 'generate');
@@ -206,6 +272,8 @@ class AgenticPipeline extends EventEmitter {
       continueOnError: true,
       gracefulDegradation: true,
     });
+    // continueOnError lets nodes degrade; an abort must not look like success.
+    runAbort.throwIfAborted();
 
     // Extract final result from generate node
     const output = dagResult?.answer

package/src/agentic/iterative-retriever.js

@@ -3,6 +3,7 @@
 const { EventEmitter } = require('events');
 const { retry } = require('../utils/retry');
 const { tokenize, computeJaccardSimilarity } = require('../evaluate/scoring');
+const { embedQuery, canEmbed } = require('../utils/embed-query');
 
 /**
  * Default configuration for the IterativeRetriever.
@@ -110,10 +111,26 @@ class IterativeRetriever extends EventEmitter {
    * @param {object} params.retriever - Retriever instance with .retrieve() method
    * @param {object} [params.llm] - Optional LLM for query refinement
    * @param {Array<number>} [params.queryVector] - Optional pre-computed query vector
+   * @param {object} [params.embedder] - Optional embedder; when given, each
+   *   refined query is re-embedded instead of reusing queryVector
+   * @param {AbortSignal} [params.signal] - Cancels the loop: passed to
+   *   retriever.retrieve({ ..., signal }), the embedder and llm.generate(...,
+   *   { signal }); no further iteration, retry or refinement starts once it
+   *   has aborted, and the call rejects with the signal's reason
    * @returns {Promise<{results: Array<object>, iterations: number, converged: boolean, coverage: number}>}
    * @throws {Error} If query or retriever is missing
    */
-  async retrieve({ query, retriever, llm, queryVector }) {
+  async retrieve({ query, retriever, llm, queryVector, embedder, signal }) {
+    const throwIfAborted = () => {
+      if (signal && signal.aborted) {
+        throw signal.reason instanceof Error
+          ? signal.reason
+          : new Error('Execution aborted');
+      }
+    };
+    // Only add the signal when there is one, so components keep the
+    // argument shapes they had before.
+    const withSignal = (obj) => (signal ? { ...obj, signal } : obj);
     // 1. Validate inputs
     if (!query || typeof query !== 'string' || !query.trim()) {
       throw new Error(
@@ -135,13 +152,31 @@ class IterativeRetriever extends EventEmitter {
 
     // 3. Retrieval loop
     while (iterations < this.config.maxIterations) {
+      throwIfAborted();
       // a. Retrieve with retry
       const retrieveQuery = currentQuery;
-      const newResults = await retry(
-        () =>
-          retriever.retrieve({ query: retrieveQuery, queryVector, topK: 10 }),
-        this.config.retryOptions,
-      );
+      // Without an embedder the caller's vector is all we have; it only
+      // matches the original query, so refined queries need re-embedding.
+      const retrieveVector =
+        canEmbed(embedder) && (retrieveQuery !== query || !queryVector)
+          ? await embedQuery(
+              embedder,
+              retrieveQuery,
+              signal ? { signal } : undefined,
+            )
+          : queryVector;
+      const newResults = await retry(() => {
+        // Do not retry (or start) a retrieval after an abort.
+        throwIfAborted();
+        return retriever.retrieve(
+          withSignal({
+            query: retrieveQuery,
+            queryVector: retrieveVector,
+            topK: 10,
+          }),
+        );
+      }, this.config.retryOptions);
+      throwIfAborted();
 
       // b. Merge and deduplicate
       const resultsArray = Array.isArray(newResults) ? newResults : [];
@@ -173,16 +208,18 @@ class IterativeRetriever extends EventEmitter {
         // Try LLM-based refinement first
         if (llm && typeof llm.generate === 'function') {
           try {
-            refined = await llm.generate(
-              `Given the query "${query}" and that current results don't fully cover the topic, generate a refined search query to find missing information.`,
-              accumulatedResults,
-            );
+            const prompt = `Given the query "${query}" and that current results don't fully cover the topic, generate a refined search query to find missing information.`;
+            refined = signal
+              ? await llm.generate(prompt, accumulatedResults, { signal })
+              : await llm.generate(prompt, accumulatedResults);
             if (refined && typeof refined === 'string' && refined.trim()) {
               refined = refined.trim();
             } else {
               refined = null;
             }
           } catch (_err) {
+            // An abort is not an LLM failure: stop instead of falling back.
+            throwIfAborted();
             // LLM failed -- fall through to heuristic
             refined = null;
           }