@devilsdev/rag-pipeline-utils 2.3.1 → 2.4.0

package/src/connectors/local-embedder.js

@@ -0,0 +1,171 @@
+'use strict';
+
+const { BaseConnector } = require('./base-connector');
+const { tokenize } = require('../evaluate/scoring');
+
+/**
+ * Default configuration for the local TF-IDF embedder.
+ * @type {Object}
+ */
+const DEFAULT_CONFIG = {
+  maxVocabularySize: 10000,
+  minDocFrequency: 1,
+};
+
+/**
+ * TF-IDF based embedder for fully offline use.
+ * No external API required -- builds a vocabulary from training documents
+ * and produces sparse TF-IDF vectors.
+ * @extends BaseConnector
+ */
+class LocalEmbedder extends BaseConnector {
+  /**
+   * @param {Object} options
+   * @param {number} [options.maxVocabularySize=10000] - Maximum vocabulary size
+   * @param {number} [options.minDocFrequency=1] - Minimum document frequency for a token to be included
+   */
+  constructor(options = {}) {
+    super({ name: 'local-embedder', ...options });
+    this.config = { ...DEFAULT_CONFIG, ...options };
+    /** @type {Map<string, number>} token -> vocabulary index */
+    this.vocabulary = new Map();
+    /** @type {Map<string, number>} token -> IDF weight */
+    this.idf = new Map();
+    /** @type {number} */
+    this.docCount = 0;
+    this.connected = true; // always available
+  }
+
+  /**
+   * Build vocabulary and IDF weights from a corpus of documents.
+   * @param {Array<string|{ content?: string, text?: string }>} documents - Training documents
+   */
+  train(documents) {
+    const docs = documents.map((doc) => {
+      if (typeof doc === 'string') return doc;
+      return doc.content || doc.text || String(doc);
+    });
+
+    this.docCount = docs.length;
+
+    // Count document frequency for each token
+    const docFreq = new Map();
+    for (const doc of docs) {
+      const tokens = new Set(tokenize(doc));
+      for (const token of tokens) {
+        docFreq.set(token, (docFreq.get(token) || 0) + 1);
+      }
+    }
+
+    // Compute IDF for tokens that meet the minimum document frequency
+    const idfEntries = [];
+    for (const [token, freq] of docFreq) {
+      if (freq < this.config.minDocFrequency) continue;
+      const idfValue = Math.log(this.docCount / (freq + 1));
+      idfEntries.push([token, idfValue]);
+    }
+
+    // Sort by IDF descending and take top maxVocabularySize tokens
+    idfEntries.sort((a, b) => b[1] - a[1]);
+    const topEntries = idfEntries.slice(0, this.config.maxVocabularySize);
+
+    // Build vocabulary map and IDF map
+    this.vocabulary.clear();
+    this.idf.clear();
+    for (let i = 0; i < topEntries.length; i++) {
+      const [token, idfValue] = topEntries[i];
+      this.vocabulary.set(token, i);
+      this.idf.set(token, idfValue);
+    }
+
+    this.emit('trained', {
+      vocabSize: this.vocabulary.size,
+      docCount: this.docCount,
+    });
+  }
+
+  /**
+   * Compute a TF-IDF vector for the given text.
+   * The vector length equals the vocabulary size.
+   * @param {string} text - Text to embed
+   * @returns {number[]} L2-normalized TF-IDF vector
+   */
+  embed(text) {
+    const vocabSize = this.vocabulary.size;
+
+    if (vocabSize === 0) {
+      this.emit('warning', {
+        message: 'LocalEmbedder has not been trained yet',
+      });
+      return new Array(1).fill(0);
+    }
+
+    const tokens = tokenize(text);
+
+    // Count term frequency
+    const tf = new Map();
+    for (const token of tokens) {
+      tf.set(token, (tf.get(token) || 0) + 1);
+    }
+
+    // Build TF-IDF vector
+    const vector = new Array(vocabSize).fill(0);
+    for (const [token, freq] of tf) {
+      const idx = this.vocabulary.get(token);
+      if (idx !== undefined) {
+        vector[idx] = freq * (this.idf.get(token) || 0);
+      }
+    }
+
+    // L2 normalize
+    let magnitude = 0;
+    for (let i = 0; i < vector.length; i++) {
+      magnitude += vector[i] * vector[i];
+    }
+    magnitude = Math.sqrt(magnitude);
+
+    if (magnitude > 0) {
+      for (let i = 0; i < vector.length; i++) {
+        vector[i] /= magnitude;
+      }
+    }
+
+    return vector;
+  }
+
+  /**
+   * Compute a TF-IDF vector for a query string.
+   * Alias for {@link embed}.
+   * @param {string} query - Query text to embed
+   * @returns {number[]} L2-normalized TF-IDF vector
+   */
+  embedQuery(query) {
+    return this.embed(query);
+  }
+
+  /**
+   * Connect the embedder (always succeeds since it is local).
+   * @returns {Promise<void>}
+   */
+  async connect() {
+    this.connected = true;
+  }
+
+  /**
+   * Check whether the embedder is healthy (vocabulary has been trained).
+   * @returns {Promise<boolean>}
+   */
+  async isHealthy() {
+    return this.vocabulary.size > 0;
+  }
+
+  /**
+   * Get the current vocabulary size.
+   * @returns {number}
+   */
+  getVocabularySize() {
+    return this.vocabulary.size;
+  }
+}
+
+module.exports = { LocalEmbedder, DEFAULT_CONFIG };

package/src/connectors/memory-retriever.js

@@ -0,0 +1,150 @@
+'use strict';
+
+const { BaseConnector } = require('./base-connector');
+
+/**
+ * Default configuration for the in-memory retriever.
+ * @type {Object}
+ */
+const DEFAULT_CONFIG = {
+  maxDocuments: 100000,
+};
+
+/**
+ * In-memory vector store with cosine similarity retrieval.
+ * Stores document vectors and retrieves the most similar ones to a query vector.
+ * @extends BaseConnector
+ */
+class MemoryRetriever extends BaseConnector {
+  /**
+   * @param {Object} options
+   * @param {number} [options.maxDocuments=100000] - Maximum number of documents to store
+   */
+  constructor(options = {}) {
+    super({ name: 'memory-retriever', ...options });
+    this.config = { ...DEFAULT_CONFIG, ...options };
+    /** @type {Map<string, { vector: number[], metadata?: Object, content?: string }>} */
+    this.documents = new Map();
+    this.connected = true;
+  }
+
+  /**
+   * Store one or more vector documents.
+   * @param {Array<{ id: string, vector: number[], metadata?: Object, content?: string }>} vectors
+   * @returns {Promise<void>}
+   */
+  async store(vectors) {
+    for (const doc of vectors) {
+      this.documents.set(doc.id, {
+        vector: doc.vector,
+        metadata: doc.metadata || {},
+        content: doc.content || '',
+      });
+    }
+
+    if (this.documents.size > this.config.maxDocuments) {
+      this.emit('warning', {
+        message: `Document count (${this.documents.size}) exceeds maxDocuments (${this.config.maxDocuments})`,
+      });
+    }
+
+    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} [k=10] - Number of results to return (used when queryOrObj is an array)
+   * @returns {Promise<Array<{ id: string, score: number, metadata: Object, content: string }>>}
+   */
+  async retrieve(queryOrObj, k = 10) {
+    let queryVector;
+    let topK = k;
+
+    if (Array.isArray(queryOrObj)) {
+      queryVector = queryOrObj;
+    } else {
+      queryVector = queryOrObj.queryVector;
+      topK = queryOrObj.topK || k;
+    }
+
+    const results = [];
+
+    for (const [id, doc] of this.documents) {
+      const score = this._cosineSimilarity(queryVector, doc.vector);
+      results.push({
+        id,
+        score,
+        metadata: doc.metadata,
+        content: doc.content,
+      });
+    }
+
+    results.sort((a, b) => b.score - a.score);
+
+    return results.slice(0, topK);
+  }
+
+  /**
+   * Compute cosine similarity between two vectors.
+   * @param {number[]} a - First vector
+   * @param {number[]} b - Second vector
+   * @returns {number} Cosine similarity in the range [-1, 1], or 0 for zero vectors
+   * @private
+   */
+  _cosineSimilarity(a, b) {
+    const len = Math.min(a.length, b.length);
+    let dot = 0;
+    let magA = 0;
+    let magB = 0;
+
+    for (let i = 0; i < len; i++) {
+      dot += a[i] * b[i];
+      magA += a[i] * a[i];
+      magB += b[i] * b[i];
+    }
+
+    const magnitude = Math.sqrt(magA) * Math.sqrt(magB);
+    if (magnitude === 0) return 0;
+
+    return dot / magnitude;
+  }
+
+  /**
+   * Remove all stored documents.
+   * @returns {Promise<void>}
+   */
+  async clear() {
+    this.documents.clear();
+  }
+
+  /**
+   * Get statistics about the store.
+   * @returns {{ documentCount: number, maxDocuments: number }}
+   */
+  getStats() {
+    return {
+      documentCount: this.documents.size,
+      maxDocuments: this.config.maxDocuments,
+    };
+  }
+
+  /**
+   * Connect the retriever (always succeeds since it is in-memory).
+   * @returns {Promise<void>}
+   */
+  async connect() {
+    this.connected = true;
+  }
+
+  /**
+   * Check whether the retriever is healthy (always true).
+   * @returns {Promise<boolean>}
+   */
+  async isHealthy() {
+    return true;
+  }
+}
+
+module.exports = { MemoryRetriever, DEFAULT_CONFIG };

package/src/connectors/ollama-connector.js

@@ -0,0 +1,205 @@
+'use strict';
+
+const axios = require('axios');
+const { BaseConnector } = require('./base-connector');
+
+/**
+ * Default configuration for the Ollama connector.
+ * @type {Object}
+ */
+const DEFAULT_CONFIG = {
+  baseURL: 'http://localhost:11434',
+  model: 'llama3',
+  embeddingModel: 'nomic-embed-text',
+  timeout: 60000,
+  temperature: 0.7,
+};
+
+/**
+ * Connector for Ollama local LLM server.
+ * Supports text generation (sync and streaming) and embeddings.
+ * @extends BaseConnector
+ */
+class OllamaConnector extends BaseConnector {
+  /**
+   * @param {Object} options
+   * @param {string} [options.baseURL='http://localhost:11434'] - Ollama server URL
+   * @param {string} [options.model='llama3'] - Default generation model
+   * @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
+   */
+  constructor(options = {}) {
+    super({ name: 'ollama', ...options });
+    this.config = { ...DEFAULT_CONFIG, ...options };
+    this.client = axios.create({
+      baseURL: this.config.baseURL,
+      timeout: this.config.timeout,
+    });
+  }
+
+  /**
+   * Connect to the Ollama server by verifying it is reachable.
+   * @returns {Promise<void>}
+   * @throws {Error} If the Ollama server is not reachable
+   */
+  async connect() {
+    try {
+      await this.client.get('/api/tags');
+      await super.connect();
+    } catch (err) {
+      throw new Error(
+        `Cannot connect to Ollama at ${this.config.baseURL}. Is Ollama running? (ollama serve)`,
+      );
+    }
+  }
+
+  /**
+   * Check whether the Ollama server is reachable.
+   * @returns {Promise<boolean>}
+   */
+  async isHealthy() {
+    try {
+      await this.client.get('/api/tags');
+      return true;
+    } catch {
+      return false;
+    }
+  }
+
+  /**
+   * Generate a response for the given query with optional context documents.
+   * @param {string} query - The user query
+   * @param {Array<Object|string>} [context=[]] - Context documents for RAG
+   * @param {Object} [options={}] - Override options
+   * @param {string} [options.model] - Override the generation model
+   * @param {number} [options.temperature] - Override temperature
+   * @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;
+  }
+
+  /**
+   * Stream a response token-by-token for the given query.
+   * @param {string} query - The user query
+   * @param {Array<Object|string>} [context=[]] - Context documents for RAG
+   * @param {Object} [options={}] - Override options
+   * @param {string} [options.model] - Override the generation model
+   * @param {number} [options.temperature] - Override temperature
+   * @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);
+
+    let response;
+    try {
+      response = await this.client.post(
+        '/api/generate',
+        {
+          model,
+          prompt,
+          stream: true,
+          options: { temperature },
+        },
+        { responseType: 'stream' },
+      );
+    } catch (err) {
+      throw new Error(`Ollama streaming request failed: ${err.message}`);
+    }
+
+    let buffer = '';
+    for await (const chunk of response.data) {
+      buffer += chunk.toString();
+      const lines = buffer.split('\n');
+      // Keep the last potentially incomplete line in the buffer
+      buffer = lines.pop() || '';
+
+      for (const line of lines) {
+        const trimmed = line.trim();
+        if (!trimmed) continue;
+        try {
+          const parsed = JSON.parse(trimmed);
+          yield { token: parsed.response || '', done: !!parsed.done };
+        } catch {
+          // Skip malformed lines
+        }
+      }
+    }
+
+    // Process any remaining data in the buffer
+    if (buffer.trim()) {
+      try {
+        const parsed = JSON.parse(buffer.trim());
+        yield { token: parsed.response || '', done: !!parsed.done };
+      } catch {
+        // Skip malformed trailing data
+      }
+    }
+  }
+
+  /**
+   * Generate an embedding vector for the given text.
+   * @param {string} text - Text to embed
+   * @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;
+  }
+
+  /**
+   * Generate an embedding vector for a query string.
+   * Alias for {@link embed}.
+   * @param {string} query - Query text to embed
+   * @returns {Promise<number[]>} Embedding vector
+   */
+  async embedQuery(query) {
+    return this.embed(query);
+  }
+
+  /**
+   * Build a prompt string incorporating context documents.
+   * @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:`;
+  }
+}
+
+module.exports = { OllamaConnector, DEFAULT_CONFIG };

package/src/connectors/openai-connector.js

@@ -0,0 +1,173 @@
+'use strict';
+
+const { BaseConnector } = require('./base-connector');
+
+/**
+ * Default configuration for the OpenAI connector.
+ * @type {Object}
+ */
+const DEFAULT_CONFIG = {
+  model: 'gpt-4o-mini',
+  embeddingModel: 'text-embedding-3-small',
+  temperature: 0.7,
+  maxTokens: 4096,
+  apiKey: process.env.OPENAI_API_KEY || '',
+};
+
+/**
+ * Full LLM + embedder connector for the OpenAI API.
+ * Supports chat completions (with streaming) and text embeddings.
+ *
+ * @extends BaseConnector
+ */
+class OpenAIConnector extends BaseConnector {
+  /**
+   * @param {Object} [options]
+   * @param {string} [options.model='gpt-4o-mini'] - Chat model to use
+   * @param {string} [options.embeddingModel='text-embedding-3-small'] - Embedding model
+   * @param {number} [options.temperature=0.7] - Sampling temperature
+   * @param {number} [options.maxTokens=4096] - Max tokens in completion
+   * @param {string} [options.apiKey] - OpenAI API key
+   */
+  constructor(options = {}) {
+    super({ name: 'openai', ...options });
+    this.config = { ...DEFAULT_CONFIG, ...options };
+    this._client = null; // lazy-loaded
+  }
+
+  /**
+   * Lazily initialise and return the OpenAI client.
+   * @returns {Object} OpenAI client instance
+   * @private
+   */
+  _getClient() {
+    if (!this._client) {
+      const OpenAI = require('openai');
+      this._client = new OpenAI({ apiKey: this.config.apiKey });
+    }
+    return this._client;
+  }
+
+  /**
+   * Connect to OpenAI. Verifies the API key is present.
+   * @returns {Promise<void>}
+   * @throws {Error} If no API key is configured
+   */
+  async connect() {
+    if (!this.config.apiKey) {
+      throw new Error(
+        'OpenAI API key required. Set OPENAI_API_KEY or pass apiKey option.',
+      );
+    }
+    await super.connect();
+  }
+
+  /**
+   * Generate a chat completion using the provided query and optional RAG context.
+   *
+   * @param {string} query - The user query / prompt
+   * @param {Array<Object|string>} [context=[]] - Retrieved documents for RAG
+   * @param {Object} [options={}] - Override model, temperature, maxTokens per-call
+   * @returns {Promise<string>} The generated text
+   */
+  async generate(query, context = [], options = {}) {
+    return this.withRetry(async () => {
+      const client = this._getClient();
+      const messages = this._buildMessages(query, context);
+
+      const response = await client.chat.completions.create({
+        model: options.model || this.config.model,
+        messages,
+        temperature: options.temperature ?? this.config.temperature,
+        max_tokens: options.maxTokens ?? this.config.maxTokens,
+      });
+
+      return response.choices[0].message.content;
+    });
+  }
+
+  /**
+   * Stream a chat completion token-by-token.
+   *
+   * @param {string} query - The user query / prompt
+   * @param {Array<Object|string>} [context=[]] - Retrieved documents for RAG
+   * @param {Object} [options={}] - Override model, temperature, maxTokens per-call
+   * @yields {{ token: string, done: boolean }}
+   */
+  async *generateStream(query, context = [], options = {}) {
+    const client = this._getClient();
+    const messages = this._buildMessages(query, context);
+
+    const stream = await client.chat.completions.create({
+      model: options.model || this.config.model,
+      messages,
+      temperature: options.temperature ?? this.config.temperature,
+      max_tokens: options.maxTokens ?? this.config.maxTokens,
+      stream: true,
+    });
+
+    for await (const chunk of stream) {
+      const token = chunk.choices[0]?.delta?.content || '';
+      yield { token, done: false };
+    }
+
+    yield { token: '', done: true };
+  }
+
+  /**
+   * Generate an embedding vector for the given text.
+   *
+   * @param {string} text - Input text to embed
+   * @returns {Promise<number[]>} Embedding vector
+   */
+  async embed(text) {
+    return this.withRetry(async () => {
+      const client = this._getClient();
+
+      const response = await client.embeddings.create({
+        model: this.config.embeddingModel,
+        input: text,
+      });
+
+      return response.data[0].embedding;
+    });
+  }
+
+  /**
+   * Generate an embedding vector for a search query.
+   * Alias for {@link OpenAIConnector#embed}.
+   *
+   * @param {string} query - Query text to embed
+   * @returns {Promise<number[]>} Embedding vector
+   */
+  async embedQuery(query) {
+    return this.embed(query);
+  }
+
+  /**
+   * Build the messages array for a chat completion request.
+   *
+   * @param {string} query - User query
+   * @param {Array<Object|string>} context - RAG context documents
+   * @returns {Array<{ role: string, content: string }>}
+   * @private
+   */
+  _buildMessages(query, context) {
+    const messages = [];
+
+    if (context && context.length > 0) {
+      const contextText = context
+        .map((doc) => doc.content || doc.text || String(doc))
+        .join('\n\n');
+      messages.push({
+        role: 'system',
+        content: `Use the following context to answer the question:\n\n${contextText}`,
+      });
+    }
+
+    messages.push({ role: 'user', content: query });
+    return messages;
+  }
+}
+
+module.exports = { OpenAIConnector, DEFAULT_CONFIG };