@devilsdev/rag-pipeline-utils 2.4.5 → 2.6.0

package/src/connectors/memory-retriever.js

@@ -1,6 +1,8 @@
 'use strict';
+const { randomBytes } = require('crypto');
 
 const { BaseConnector } = require('./base-connector');
+const { throwIfAborted } = require('./connector-utils');
 
 /**
  * Default configuration for the in-memory retriever.
@@ -8,32 +10,101 @@ const { BaseConnector } = require('./base-connector');
  */
 const DEFAULT_CONFIG = {
   maxDocuments: 100000,
+  // 'error' rejects a store() that would exceed maxDocuments; 'fifo' evicts
+  // the oldest documents (insertion order) to make room.
+  eviction: 'error',
 };
 
 /**
  * In-memory vector store with cosine similarity retrieval.
  * Stores document vectors and retrieves the most similar ones to a query vector.
+ * All vectors in one store must have the same dimension.
  * @extends BaseConnector
  */
 class MemoryRetriever extends BaseConnector {
   /**
    * @param {Object} options
    * @param {number} [options.maxDocuments=100000] - Maximum number of documents to store
+   * @param {'error'|'fifo'} [options.eviction='error'] - What to do when a store()
+   *   would exceed maxDocuments: throw (default) or evict the oldest documents
    */
   constructor(options = {}) {
     super({ name: 'memory-retriever', ...options });
     this.config = { ...DEFAULT_CONFIG, ...options };
+    if (!['error', 'fifo'].includes(this.config.eviction)) {
+      throw new Error(
+        `MemoryRetriever: unknown eviction "${this.config.eviction}" (use "error" or "fifo")`,
+      );
+    }
     /** @type {Map<string, { vector: number[], metadata?: Object, content?: string }>} */
     this.documents = new Map();
+    /** @type {number|null} Vector dimension, fixed by the first stored vector */
+    this.dimensions = null;
     this.connected = true;
+    // Cache identity: unique per instance, bumped on every change, so a
+    // shared cache never returns another store's or an older state's results.
+    this._cacheId = randomBytes(8).toString('hex');
+    this._revision = 0;
+  }
+
+  /**
+   * Cache identity for createRagPipeline({ cache }).
+   * @param {'retriever'} role
+   * @returns {string|null}
+   */
+  getCacheKey(role) {
+    return role === 'retriever'
+      ? `memory|${this._cacheId}|${this._revision}`
+      : null;
   }
 
   /**
-   * Store one or more vector documents.
+   * Store one or more vector documents. The batch is validated before
+   * anything is inserted, so a rejected call leaves the store unchanged.
+   *
    * @param {Array<{ id: string, vector: number[], metadata?: Object, content?: string }>} vectors
    * @returns {Promise<void>}
+   * @throws {Error} When a vector's dimension differs from the store's, or
+   *   (eviction 'error') when the store would exceed maxDocuments
    */
   async store(vectors) {
+    let dimensions = this.dimensions;
+    const newIds = new Set();
+    for (const doc of vectors) {
+      const length = doc && doc.vector ? doc.vector.length : undefined;
+      if (typeof length !== 'number') {
+        throw new Error(
+          `MemoryRetriever.store(): document "${doc && doc.id}" has no vector`,
+        );
+      }
+      if (dimensions === null) {
+        dimensions = length;
+      } else if (length !== dimensions) {
+        throw new Error(
+          `MemoryRetriever.store(): document "${doc.id}" has ${length} dimensions ` +
+            `but the store holds ${dimensions}-dimension vectors. Use the same embedder for all documents.`,
+        );
+      }
+      if (!this.documents.has(doc.id)) newIds.add(doc.id);
+    }
+
+    const { maxDocuments, eviction } = this.config;
+    const projected = this.documents.size + newIds.size;
+    if (projected > maxDocuments) {
+      if (eviction !== 'fifo') {
+        throw new Error(
+          `MemoryRetriever.store(): storing ${newIds.size} new document(s) would bring the ` +
+            `total to ${projected}, exceeding maxDocuments (${maxDocuments}). Raise ` +
+            'maxDocuments, call clear(), or construct with { eviction: \'fifo\' }.',
+        );
+      }
+      if (newIds.size > maxDocuments) {
+        throw new Error(
+          `MemoryRetriever.store(): batch of ${newIds.size} new documents exceeds maxDocuments (${maxDocuments})`,
+        );
+      }
+    }
+
     for (const doc of vectors) {
       this.documents.set(doc.id, {
         vector: doc.vector,
@@ -41,32 +112,63 @@ class MemoryRetriever extends BaseConnector {
         content: doc.content || '',
       });
     }
+    if (vectors.length > 0) this.dimensions = dimensions;
 
-    if (this.documents.size > this.config.maxDocuments) {
-      this.emit('warning', {
-        message: `Document count (${this.documents.size}) exceeds maxDocuments (${this.config.maxDocuments})`,
-      });
+    let evicted = 0;
+    if (this.documents.size > maxDocuments) {
+      // FIFO eviction: Map iteration follows insertion order. Never evict a
+      // document from the current batch.
+      for (const id of this.documents.keys()) {
+        if (this.documents.size <= maxDocuments) break;
+        if (newIds.has(id)) continue;
+        this.documents.delete(id);
+        evicted++;
+      }
+      this.emit('evicted', { count: evicted, total: this.documents.size });
     }
 
+    this._revision += 1;
     this.emit('stored', { count: vectors.length, total: this.documents.size });
   }
 
   /**
    * Retrieve the top-k most similar documents to a query vector.
-   * @param {number[]|{ query?: string, queryVector: number[], topK?: number }} queryOrObj
-   *   Either a raw vector array or an object with queryVector and optional topK.
+   * @param {number[]|{ query?: string, queryVector: number[], topK?: number, signal?: AbortSignal }} queryOrObj
+   *   Either a raw vector array or an object with queryVector, optional topK
+   *   and optional AbortSignal.
    * @param {number} [k=10] - Number of results to return (used when queryOrObj is an array)
    * @returns {Promise<Array<{ id: string, score: number, metadata: Object, content: string }>>}
+   * @throws {Error} When the query vector's dimension differs from the stored vectors
    */
   async retrieve(queryOrObj, k = 10) {
     let queryVector;
     let topK = k;
+    let signal;
 
     if (Array.isArray(queryOrObj)) {
       queryVector = queryOrObj;
     } else {
-      queryVector = queryOrObj.queryVector;
-      topK = queryOrObj.topK || k;
+      queryVector = queryOrObj?.queryVector;
+      topK = queryOrObj?.topK || k;
+      signal = queryOrObj?.signal;
+    }
+
+    throwIfAborted(signal);
+
+    if (!Array.isArray(queryVector) && !ArrayBuffer.isView(queryVector)) {
+      throw new Error(
+        'MemoryRetriever.retrieve(): queryVector is required. Pass an ' +
+          '`embedder` to createRagPipeline or supply `queryVector` to run().',
+      );
+    }
+
+    if (this.dimensions !== null && this.documents.size > 0) {
+      if (queryVector.length !== this.dimensions) {
+        throw new Error(
+          `Query vector has ${queryVector.length} dimensions but the store holds ` +
+            `${this.dimensions}-dimension vectors. Use the same embedder for queries and documents.`,
+        );
+      }
     }
 
     const results = [];
@@ -83,23 +185,29 @@ class MemoryRetriever extends BaseConnector {
 
     results.sort((a, b) => b.score - a.score);
 
+    throwIfAborted(signal);
     return results.slice(0, topK);
   }
 
   /**
-   * Compute cosine similarity between two vectors.
+   * Compute cosine similarity between two vectors of equal length.
    * @param {number[]} a - First vector
    * @param {number[]} b - Second vector
    * @returns {number} Cosine similarity in the range [-1, 1], or 0 for zero vectors
+   * @throws {Error} When the vectors have different lengths
    * @private
    */
   _cosineSimilarity(a, b) {
-    const len = Math.min(a.length, b.length);
+    if (a.length !== b.length) {
+      throw new Error(
+        `Cannot compare vectors of different dimensions (${a.length} vs ${b.length})`,
+      );
+    }
     let dot = 0;
     let magA = 0;
     let magB = 0;
 
-    for (let i = 0; i < len; i++) {
+    for (let i = 0; i < a.length; i++) {
       dot += a[i] * b[i];
       magA += a[i] * a[i];
       magB += b[i] * b[i];
@@ -112,21 +220,24 @@ class MemoryRetriever extends BaseConnector {
   }
 
   /**
-   * Remove all stored documents.
+   * Remove all stored documents and reset the vector dimension.
    * @returns {Promise<void>}
    */
   async clear() {
     this.documents.clear();
+    this.dimensions = null;
+    this._revision += 1;
   }
 
   /**
    * Get statistics about the store.
-   * @returns {{ documentCount: number, maxDocuments: number }}
+   * @returns {{ documentCount: number, maxDocuments: number, dimensions: number|null }}
    */
   getStats() {
     return {
       documentCount: this.documents.size,
       maxDocuments: this.config.maxDocuments,
+      dimensions: this.dimensions,
     };
   }
 

package/src/connectors/ollama-connector.js

@@ -1,7 +1,16 @@
 'use strict';
 
-const axios = require('axios');
+// axios loads on first use, so importing the package does not pay for it
+// (about 125 ms) unless this connector is constructed.
+const loadAxios = () => require('axios'); // eslint-disable-line global-require
 const { BaseConnector } = require('./base-connector');
+const {
+  raceAbort,
+  signalOptions,
+  throwIfAborted,
+} = require('./connector-utils');
+const { buildRagPrompt } = require('./rag-prompt');
+const { componentCacheKey } = require('../cache/component-keys');
 
 /**
  * Default configuration for the Ollama connector.
@@ -13,14 +22,41 @@ const DEFAULT_CONFIG = {
   embeddingModel: 'nomic-embed-text',
   timeout: 60000,
   temperature: 0.7,
+  maxRetries: 3,
+  retryDelay: 1000,
 };
 
 /**
  * Connector for Ollama local LLM server.
  * Supports text generation (sync and streaming) and embeddings.
+ * Ollama needs no API key.
  * @extends BaseConnector
  */
 class OllamaConnector extends BaseConnector {
+  /**
+   * Cache identity for createRagPipeline({ cache }).
+   * @param {'embedder'|'llm'} role
+   * @returns {string|null}
+   */
+  getCacheKey(role) {
+    const c = this.config;
+    if (role === 'embedder') {
+      return componentCacheKey('ollama-embed', {
+        baseURL: c.baseURL,
+        model: c.embeddingModel,
+      });
+    }
+    if (role === 'llm') {
+      return componentCacheKey('ollama', {
+        baseURL: c.baseURL,
+        model: c.model,
+        temperature: c.temperature,
+        systemPrompt: c.systemPrompt,
+      });
+    }
+    return null;
+  }
+
   /**
    * @param {Object} options
    * @param {string} [options.baseURL='http://localhost:11434'] - Ollama server URL
@@ -28,11 +64,14 @@ class OllamaConnector extends BaseConnector {
    * @param {string} [options.embeddingModel='nomic-embed-text'] - Default embedding model
    * @param {number} [options.timeout=60000] - Request timeout in ms
    * @param {number} [options.temperature=0.7] - Generation temperature
+   * @param {number} [options.maxRetries=3] - Retries for transient failures
+   * @param {number} [options.retryDelay=1000] - Base backoff delay in ms
+   * @param {string} [options.systemPrompt] - System prompt used when RAG context is supplied
    */
   constructor(options = {}) {
     super({ name: 'ollama', ...options });
     this.config = { ...DEFAULT_CONFIG, ...options };
-    this.client = axios.create({
+    this.client = loadAxios().create({
       baseURL: this.config.baseURL,
       timeout: this.config.timeout,
     });
@@ -67,6 +106,46 @@ class OllamaConnector extends BaseConnector {
     }
   }
 
+  /**
+   * POST with the connector's retry policy and cancellation.
+   * @private
+   */
+  _post(path, body, signal) {
+    throwIfAborted(signal);
+    return this.withRetry(
+      async () => {
+        const response = await this.client.post(
+          path,
+          body,
+          signal ? { signal } : undefined,
+        );
+        return response.data;
+      },
+      this.config.maxRetries,
+      this.config.retryDelay,
+      { signal },
+    );
+  }
+
+  /**
+   * Build the /api/generate body. Retrieved context goes into the prompt as
+   * delimited, untrusted reference data; instructions go in `system`.
+   * @private
+   */
+  _buildBody(query, context, options, stream) {
+    const prompt = buildRagPrompt(query, context, {
+      systemPrompt: this.config.systemPrompt,
+    });
+    const body = {
+      model: options.model || this.config.model,
+      prompt: prompt.user,
+      stream,
+      options: { temperature: options.temperature ?? this.config.temperature },
+    };
+    if (prompt.system) body.system = prompt.system;
+    return body;
+  }
+
   /**
    * Generate a response for the given query with optional context documents.
    * @param {string} query - The user query
@@ -74,24 +153,13 @@ class OllamaConnector extends BaseConnector {
    * @param {Object} [options={}] - Override options
    * @param {string} [options.model] - Override the generation model
    * @param {number} [options.temperature] - Override temperature
+   * @param {AbortSignal} [options.signal] - Cancels the request
    * @returns {Promise<string>} Generated response text
    */
   async generate(query, context = [], options = {}) {
-    const model = options.model || this.config.model;
-    const temperature = options.temperature ?? this.config.temperature;
-    const prompt = this._buildPrompt(query, context);
-
-    const result = await this.withRetry(async () => {
-      const response = await this.client.post('/api/generate', {
-        model,
-        prompt,
-        stream: false,
-        options: { temperature },
-      });
-      return response.data.response;
-    });
-
-    return result;
+    const body = this._buildBody(query, context, options, false);
+    const data = await this._post('/api/generate', body, options.signal);
+    return data.response;
   }
 
   /**
@@ -101,31 +169,31 @@ class OllamaConnector extends BaseConnector {
    * @param {Object} [options={}] - Override options
    * @param {string} [options.model] - Override the generation model
    * @param {number} [options.temperature] - Override temperature
+   * @param {AbortSignal} [options.signal] - Cancels the stream
    * @yields {{ token: string, done: boolean }}
    */
   async *generateStream(query, context = [], options = {}) {
-    const model = options.model || this.config.model;
-    const temperature = options.temperature ?? this.config.temperature;
-    const prompt = this._buildPrompt(query, context);
+    const { signal } = options;
+    throwIfAborted(signal);
+    const body = this._buildBody(query, context, options, true);
 
     let response;
     try {
-      response = await this.client.post(
-        '/api/generate',
-        {
-          model,
-          prompt,
-          stream: true,
-          options: { temperature },
-        },
-        { responseType: 'stream' },
+      response = await raceAbort(
+        this.client.post('/api/generate', body, {
+          responseType: 'stream',
+          ...signalOptions(signal),
+        }),
+        signal,
       );
     } catch (err) {
+      if (err && err.name === 'AbortError') throw err;
       throw new Error(`Ollama streaming request failed: ${err.message}`);
     }
 
     let buffer = '';
     for await (const chunk of response.data) {
+      throwIfAborted(signal);
       buffer += chunk.toString();
       const lines = buffer.split('\n');
       // Keep the last potentially incomplete line in the buffer
@@ -157,48 +225,41 @@ class OllamaConnector extends BaseConnector {
   /**
    * Generate an embedding vector for the given text.
    * @param {string} text - Text to embed
+   * @param {Object} [options={}]
+   * @param {AbortSignal} [options.signal] - Cancels the request
    * @returns {Promise<number[]>} Embedding vector
    */
-  async embed(text) {
-    const result = await this.withRetry(async () => {
-      const response = await this.client.post('/api/embeddings', {
-        model: this.config.embeddingModel,
-        prompt: text,
-      });
-      return response.data.embedding;
-    });
-
-    return result;
+  async embed(text, options = {}) {
+    const data = await this._post(
+      '/api/embeddings',
+      { model: this.config.embeddingModel, prompt: text },
+      options.signal,
+    );
+    return data.embedding;
   }
 
   /**
    * Generate an embedding vector for a query string.
    * Alias for {@link embed}.
    * @param {string} query - Query text to embed
+   * @param {Object} [options={}] - Same as embed()
    * @returns {Promise<number[]>} Embedding vector
    */
-  async embedQuery(query) {
-    return this.embed(query);
+  async embedQuery(query, options = {}) {
+    return this.embed(query, options);
   }
 
   /**
    * Build a prompt string incorporating context documents.
+   * Kept for backward compatibility; context is delimited and labelled as
+   * untrusted reference data.
    * @param {string} query - The user query
    * @param {Array<Object|string>} context - Context documents
    * @returns {string} Formatted prompt
    * @private
    */
   _buildPrompt(query, context) {
-    if (!context || context.length === 0) return query;
-
-    const contextText = context
-      .map((doc) => {
-        const text = doc.content || doc.text || String(doc);
-        return text;
-      })
-      .join('\n\n');
-
-    return `Context:\n${contextText}\n\nQuestion: ${query}\n\nAnswer:`;
+    return buildRagPrompt(query, context).user;
   }
 }