libsql-search 0.1.6 → 0.2.1

package/dist/index.cjs

@@ -4,7 +4,15 @@ var promises = require('fs/promises');
 var path = require('path');
 var matter = require('gray-matter');
 
-const providerCache = {};
+const DEFAULT_DIMENSIONS = 768;
+const DEFAULT_MAX_LENGTH = 8e3;
+const DEFAULT_TIMEOUT_MS = 3e4;
+const LOCAL_MODEL = "Xenova/all-MiniLM-L6-v2";
+const GEMINI_MODEL = "text-embedding-004";
+const OPENAI_SMALL_MODEL = "text-embedding-3-small";
+const OPENAI_LARGE_MODEL = "text-embedding-3-large";
+const OPENAI_EMBEDDINGS_URL = "https://api.openai.com/v1/embeddings";
+const localModelCacheByModel = /* @__PURE__ */ new Map();
 function getEnvironmentVariable(name) {
   const runtime = globalThis;
   const nodeValue = runtime.process?.env?.[name];
@@ -17,80 +25,421 @@ function getEnvironmentVariable(name) {
     return void 0;
   }
 }
-async function getLocalEmbeddingModel() {
-  if (!providerCache.local) {
-    console.log("Loading local embedding model (Xenova/all-MiniLM-L6-v2)...");
-    const { pipeline } = await import('@xenova/transformers');
-    providerCache.local = await pipeline("feature-extraction", "Xenova/all-MiniLM-L6-v2");
-    console.log("Local model loaded successfully");
+async function getLocalEmbeddingModel(modelName) {
+  const cached = localModelCacheByModel.get(modelName);
+  if (cached?.model) {
+    return cached.model;
   }
-  return providerCache.local;
+  console.log(`Loading local embedding model (${modelName})...`);
+  const { pipeline } = await import('@xenova/transformers');
+  const model = await pipeline("feature-extraction", modelName);
+  localModelCacheByModel.set(modelName, { model });
+  console.log("Local model loaded successfully");
+  return model;
 }
-async function generateEmbedding(text, options = {}) {
-  const {
-    provider = "local",
-    apiKey,
-    dimensions = 768,
-    maxLength = 8e3
-  } = options;
-  const truncated = text.substring(0, maxLength);
-  switch (provider) {
+function getPositiveInteger(value, optionName) {
+  if (!Number.isInteger(value) || value <= 0) {
+    throw new Error(`Invalid ${optionName}: expected a positive integer`);
+  }
+  return value;
+}
+function getTimeoutMs(value) {
+  return getPositiveInteger(value ?? DEFAULT_TIMEOUT_MS, "timeoutMs");
+}
+function resolveProviderName(provider) {
+  switch (provider ?? "local") {
     case "local":
-      return generateLocalEmbedding(truncated, dimensions);
     case "gemini":
-      return generateGeminiEmbedding(truncated, apiKey);
     case "openai":
-      return generateOpenAIEmbedding(truncated, apiKey, dimensions);
+      return provider ?? "local";
+    default:
+      throw new Error(`Unknown embedding provider: ${String(provider)}`);
+  }
+}
+function resolveIntent(intent) {
+  switch (intent ?? "document") {
+    case "document":
+    case "query":
+      return intent ?? "document";
     default:
-      throw new Error(`Unknown embedding provider: ${provider}`);
+      throw new Error(`Unknown embedding intent: ${String(intent)}`);
+  }
+}
+function truncateTexts(texts, maxLength) {
+  return texts.map((text) => text.substring(0, maxLength));
+}
+function getOpenAIModel(dimensions) {
+  return dimensions <= 1536 ? OPENAI_SMALL_MODEL : OPENAI_LARGE_MODEL;
+}
+function hasOwnProperty(value, property) {
+  return Object.prototype.hasOwnProperty.call(value, property);
+}
+function redactErrorText(value, exactSecrets = []) {
+  let raw = value instanceof Error ? value.message : String(value);
+  for (const secret of exactSecrets) {
+    if (secret.length > 0) {
+      raw = raw.split(secret).join("[redacted]");
+    }
   }
+  return raw.replace(/Bearer\s+[A-Za-z0-9._~+/=-]+/gi, "Bearer [redacted]").replace(/([?&](?:api[_-]?key|key|token)=)[^&\s]+/gi, "$1[redacted]").replace(/(authorization\s*[:=]\s*)[^\s,}]+/gi, "$1[redacted]").replace(/https?:\/\/[^\s)]+/gi, (match) => {
+    try {
+      const url = new URL(match);
+      return `${url.origin}${url.pathname}`;
+    } catch {
+      return "[redacted-url]";
+    }
+  }).slice(0, 300);
 }
-async function generateLocalEmbedding(text, targetDimensions) {
-  const model = await getLocalEmbeddingModel();
-  const output = await model(text, {
-    pooling: "mean",
-    normalize: true
+function providerError(provider, message, cause, exactSecrets = []) {
+  const safeCause = cause === void 0 ? "" : `: ${redactErrorText(cause, exactSecrets).trim()}`;
+  return new Error(`${provider} embedding error: ${message}${safeCause}`);
+}
+function getResponseHeader(response, name) {
+  return response.headers.get(name) ?? void 0;
+}
+async function withTimeout(provider, operation, timeoutMs, run, parentSignal, exactSecrets = []) {
+  if (parentSignal?.aborted) {
+    throw providerError(provider, `${operation} was aborted`);
+  }
+  const controller = new AbortController();
+  let abortError;
+  let rejectAbort = () => {
+  };
+  const abortPromise = new Promise((_resolve, reject) => {
+    rejectAbort = reject;
   });
-  const embedding = Array.from(output.data);
-  return padEmbedding(embedding, targetDimensions);
-}
-async function generateGeminiEmbedding(text, apiKey) {
-  const key = apiKey || getEnvironmentVariable("GEMINI_API_KEY");
-  if (!key) {
-    throw new Error("GEMINI_API_KEY is required for Gemini embeddings");
-  }
-  if (!providerCache.gemini) {
-    const { GoogleGenerativeAI } = await import('@google/generative-ai');
-    const genAI = new GoogleGenerativeAI(key);
-    providerCache.gemini = genAI.getGenerativeModel({ model: "text-embedding-004" });
-  }
-  const result = await providerCache.gemini.embedContent(text);
-  return result.embedding.values;
-}
-async function generateOpenAIEmbedding(text, apiKey, dimensions = 1536) {
-  const key = apiKey || getEnvironmentVariable("OPENAI_API_KEY");
-  if (!key) {
-    throw new Error("OPENAI_API_KEY is required for OpenAI embeddings");
-  }
-  const model = dimensions <= 1536 ? "text-embedding-3-small" : "text-embedding-3-large";
-  const response = await fetch("https://api.openai.com/v1/embeddings", {
-    method: "POST",
-    headers: {
-      "Content-Type": "application/json",
-      "Authorization": `Bearer ${key}`
-    },
-    body: JSON.stringify({
-      input: text,
-      model,
-      dimensions
-    })
+  const abort = (error) => {
+    if (abortError) {
+      return;
+    }
+    abortError = error;
+    controller.abort();
+    rejectAbort(error);
+  };
+  const onParentAbort = () => abort(providerError(provider, `${operation} was aborted`));
+  const timeout = setTimeout(() => {
+    abort(providerError(provider, `${operation} timed out after ${timeoutMs}ms`));
+  }, timeoutMs);
+  parentSignal?.addEventListener("abort", onParentAbort, { once: true });
+  const operationPromise = Promise.resolve().then(() => run(controller.signal));
+  operationPromise.catch(() => void 0);
+  try {
+    return await Promise.race([operationPromise, abortPromise]);
+  } catch (error) {
+    if (error === abortError) {
+      throw error;
+    }
+    throw providerError(provider, `${operation} failed`, error, exactSecrets);
+  } finally {
+    clearTimeout(timeout);
+    parentSignal?.removeEventListener("abort", onParentAbort);
+  }
+}
+function validateEmbeddingBatch(items, expectedCount, expectedDimensions, provider) {
+  if (items.length !== expectedCount) {
+    throw providerError(
+      provider,
+      `expected ${expectedCount} embedding result(s), received ${items.length}`
+    );
+  }
+  const hasIndexedItems = items.some((item) => !Array.isArray(item) && hasOwnProperty(item, "index"));
+  const ordered = hasIndexedItems ? reorderIndexedEmbeddings(items, expectedCount, provider) : items.map((item) => Array.isArray(item) ? item : item.embedding);
+  return ordered.map(
+    (embedding, itemIndex) => validateEmbeddingVector(embedding, expectedDimensions, provider, itemIndex)
+  );
+}
+function reorderIndexedEmbeddings(items, expectedCount, provider) {
+  const ordered = new Array(expectedCount);
+  const seen = /* @__PURE__ */ new Set();
+  for (const item of items) {
+    if (Array.isArray(item) || !hasOwnProperty(item, "index")) {
+      throw providerError(provider, "provider returned a partially indexed embedding batch");
+    }
+    if (typeof item.index !== "number" || !Number.isInteger(item.index) || item.index < 0 || item.index >= expectedCount) {
+      throw providerError(provider, `provider returned invalid embedding index ${String(item.index)}`);
+    }
+    if (seen.has(item.index)) {
+      throw providerError(provider, `provider returned duplicate embedding index ${item.index}`);
+    }
+    seen.add(item.index);
+    ordered[item.index] = item.embedding;
+  }
+  if (seen.size !== expectedCount) {
+    throw providerError(provider, "provider returned non-contiguous embedding indices");
+  }
+  return ordered;
+}
+function validateEmbeddingVector(embedding, expectedDimensions, provider, itemIndex) {
+  if (!Array.isArray(embedding)) {
+    throw providerError(provider, `embedding ${itemIndex} is not an array`);
+  }
+  if (embedding.length !== expectedDimensions) {
+    throw providerError(
+      provider,
+      `embedding ${itemIndex} has ${embedding.length} dimensions, expected ${expectedDimensions}`
+    );
+  }
+  for (let i = 0; i < embedding.length; i++) {
+    const value = embedding[i];
+    if (typeof value !== "number" || !Number.isFinite(value)) {
+      throw providerError(provider, `embedding ${itemIndex} contains a non-finite value at dimension ${i}`);
+    }
+  }
+  return embedding;
+}
+function throwIfAborted(provider, operation, signal) {
+  if (signal.aborted) {
+    throw providerError(provider, `${operation} was aborted`);
+  }
+}
+async function embedSequentially(provider, operation, texts, signal, embedOne) {
+  const results = [];
+  for (const text of texts) {
+    throwIfAborted(provider, operation, signal);
+    const embedding = await embedOne(text, signal);
+    throwIfAborted(provider, operation, signal);
+    results.push(embedding);
+  }
+  return results;
+}
+function createProviderMetadata(provider, dimensions) {
+  switch (provider) {
+    case "local":
+      return Object.freeze({
+        name: "local",
+        model: LOCAL_MODEL,
+        dimensions,
+        batch: Object.freeze({ mode: "sequential" })
+      });
+    case "gemini":
+      return Object.freeze({
+        name: "gemini",
+        model: GEMINI_MODEL,
+        dimensions: DEFAULT_DIMENSIONS,
+        batch: Object.freeze({ mode: "sequential" })
+      });
+    case "openai":
+      return Object.freeze({
+        name: "openai",
+        model: getOpenAIModel(dimensions),
+        dimensions,
+        batch: Object.freeze({ mode: "native", maxSize: 2048 })
+      });
+  }
+}
+function getEffectiveDimensions(provider, dimensions) {
+  if (provider === "gemini") {
+    return DEFAULT_DIMENSIONS;
+  }
+  return getPositiveInteger(dimensions ?? DEFAULT_DIMENSIONS, "dimensions");
+}
+function assertBatchSize(metadata, count) {
+  if (metadata.batch.maxSize !== void 0 && count > metadata.batch.maxSize) {
+    throw providerError(
+      metadata.name,
+      `batch size ${count} exceeds maximum ${metadata.batch.maxSize}`
+    );
+  }
+}
+function createEmbeddingBatchResult(metadata, intent, embeddings) {
+  return Object.freeze({
+    embeddings,
+    provider: metadata.name,
+    model: metadata.model,
+    dimensions: metadata.dimensions,
+    intent
   });
-  if (!response.ok) {
-    const error = await response.text();
-    throw new Error(`OpenAI API error: ${error}`);
+}
+class LocalEmbeddingProvider {
+  metadata;
+  #timeoutMs;
+  constructor(metadata, timeoutMs) {
+    this.metadata = metadata;
+    this.#timeoutMs = timeoutMs;
+  }
+  async embed(texts, options = {}) {
+    const intent = resolveIntent(options.intent);
+    if (texts.length === 0) {
+      return createEmbeddingBatchResult(this.metadata, intent, []);
+    }
+    assertBatchSize(this.metadata, texts.length);
+    const vectors = await withTimeout(
+      "local",
+      "model inference",
+      this.#timeoutMs,
+      async (signal) => {
+        const model = await getLocalEmbeddingModel(this.metadata.model);
+        return embedSequentially("local", "model inference", texts, signal, async (text) => {
+          const output = await model(text, {
+            pooling: "mean",
+            normalize: true
+          });
+          const embedding = Array.from(output.data);
+          return padEmbedding(embedding, this.metadata.dimensions);
+        });
+      },
+      options.signal
+    );
+    return createEmbeddingBatchResult(
+      this.metadata,
+      intent,
+      validateEmbeddingBatch(vectors, texts.length, this.metadata.dimensions, "local")
+    );
+  }
+}
+class GeminiEmbeddingProvider {
+  metadata;
+  #apiKey;
+  #timeoutMs;
+  constructor(metadata, apiKey, timeoutMs) {
+    this.metadata = metadata;
+    this.#apiKey = apiKey;
+    this.#timeoutMs = timeoutMs;
+  }
+  async embed(texts, options = {}) {
+    const intent = resolveIntent(options.intent);
+    if (texts.length === 0) {
+      return createEmbeddingBatchResult(this.metadata, intent, []);
+    }
+    assertBatchSize(this.metadata, texts.length);
+    const vectors = await withTimeout(
+      "gemini",
+      "API request",
+      this.#timeoutMs,
+      async (signal) => {
+        const { GoogleGenerativeAI } = await import('@google/generative-ai');
+        const genAI = new GoogleGenerativeAI(this.#apiKey);
+        const model = genAI.getGenerativeModel({ model: this.metadata.model });
+        return embedSequentially("gemini", "API request", texts, signal, async (text, itemSignal) => {
+          const result = await model.embedContent(text, { signal: itemSignal });
+          const values = result.embedding?.values;
+          if (!Array.isArray(values)) {
+            throw new Error("Gemini response did not include embedding values");
+          }
+          return values;
+        });
+      },
+      options.signal,
+      [this.#apiKey]
+    );
+    return createEmbeddingBatchResult(
+      this.metadata,
+      intent,
+      validateEmbeddingBatch(vectors, texts.length, this.metadata.dimensions, "gemini")
+    );
   }
-  const data = await response.json();
-  return data.data[0].embedding;
+}
+class OpenAIEmbeddingProvider {
+  metadata;
+  #apiKey;
+  #timeoutMs;
+  constructor(metadata, apiKey, timeoutMs) {
+    this.metadata = metadata;
+    this.#apiKey = apiKey;
+    this.#timeoutMs = timeoutMs;
+  }
+  async embed(texts, options = {}) {
+    const intent = resolveIntent(options.intent);
+    if (texts.length === 0) {
+      return createEmbeddingBatchResult(this.metadata, intent, []);
+    }
+    assertBatchSize(this.metadata, texts.length);
+    const items = await withTimeout(
+      "openai",
+      "API request",
+      this.#timeoutMs,
+      async (signal) => {
+        const response = await fetch(OPENAI_EMBEDDINGS_URL, {
+          method: "POST",
+          headers: {
+            "Content-Type": "application/json",
+            "Authorization": `Bearer ${this.#apiKey}`
+          },
+          signal,
+          body: JSON.stringify({
+            input: texts,
+            model: this.metadata.model,
+            dimensions: this.metadata.dimensions
+          })
+        });
+        if (!response.ok) {
+          const requestId = getResponseHeader(response, "x-request-id");
+          const context = requestId ? ` status ${response.status}, request ${requestId}` : ` status ${response.status}`;
+          throw new Error(context);
+        }
+        const data = await response.json();
+        if (!Array.isArray(data.data)) {
+          throw new Error("OpenAI response did not include a data array");
+        }
+        return data.data.map((item) => {
+          const typed = item;
+          const embedding = typed.embedding;
+          if (!Array.isArray(embedding)) {
+            throw new Error("OpenAI response item did not include an embedding array");
+          }
+          const parsed = { embedding };
+          if (hasOwnProperty(typed, "index")) {
+            parsed.index = typed.index;
+          }
+          return parsed;
+        });
+      },
+      options.signal,
+      [this.#apiKey]
+    );
+    return createEmbeddingBatchResult(
+      this.metadata,
+      intent,
+      validateEmbeddingBatch(items, texts.length, this.metadata.dimensions, "openai")
+    );
+  }
+}
+function createEmbeddingProvider(options = {}) {
+  const provider = resolveProviderName(options.provider);
+  const metadata = getEmbeddingProviderMetadata(options);
+  const timeoutMs = getTimeoutMs(options.timeoutMs);
+  switch (provider) {
+    case "local":
+      return new LocalEmbeddingProvider(metadata, timeoutMs);
+    case "gemini": {
+      const key = options.apiKey || getEnvironmentVariable("GEMINI_API_KEY");
+      if (!key) {
+        throw new Error("GEMINI_API_KEY is required for Gemini embeddings");
+      }
+      return new GeminiEmbeddingProvider(metadata, key, timeoutMs);
+    }
+    case "openai": {
+      const key = options.apiKey || getEnvironmentVariable("OPENAI_API_KEY");
+      if (!key) {
+        throw new Error("OPENAI_API_KEY is required for OpenAI embeddings");
+      }
+      return new OpenAIEmbeddingProvider(metadata, key, timeoutMs);
+    }
+  }
+}
+function getEmbeddingProviderMetadata(options = {}) {
+  const provider = resolveProviderName(options.provider);
+  return createProviderMetadata(provider, getEffectiveDimensions(provider, options.dimensions));
+}
+async function generateEmbeddings(texts, options = {}) {
+  const maxLength = getPositiveInteger(options.maxLength ?? DEFAULT_MAX_LENGTH, "maxLength");
+  const intent = resolveIntent(options.intent);
+  if (texts.length === 0) {
+    return [];
+  }
+  const provider = createEmbeddingProvider(options);
+  const result = await provider.embed(truncateTexts(texts, maxLength), {
+    intent,
+    ...options.signal ? { signal: options.signal } : {}
+  });
+  return result.embeddings;
+}
+async function generateEmbedding(text, options = {}) {
+  const [embedding] = await generateEmbeddings([text], options);
+  if (!embedding) {
+    throw new Error("Embedding provider returned no embedding");
+  }
+  return embedding;
 }
 function padEmbedding(embedding, targetDimensions) {
   if (embedding.length === targetDimensions) {
@@ -211,7 +560,10 @@ async function processFile(file, embeddingOptions) {
     content: markdown,
     tags
   });
-  const embedding = await generateEmbedding(embeddingText, embeddingOptions);
+  const embedding = await generateEmbedding(embeddingText, {
+    ...embeddingOptions,
+    intent: embeddingOptions.intent ?? "document"
+  });
   return {
     slug,
     title,
@@ -280,7 +632,10 @@ async function search(options) {
   } = options;
   const quotedTableName = quoteSqlIdentifier(tableName, "tableName");
   const resultLimit = normalizeSearchLimit(limit);
-  const queryEmbedding = await generateEmbedding(query, embeddingOptions);
+  const queryEmbedding = await generateEmbedding(query, {
+    ...embeddingOptions,
+    intent: embeddingOptions.intent ?? "query"
+  });
   const results = await client.execute({
     sql: `
       SELECT
@@ -382,13 +737,17 @@ async function getFolders(client, tableName = "articles") {
   return results.rows.map((row) => row.folder);
 }
 
+exports.createEmbeddingProvider = createEmbeddingProvider;
 exports.createTable = createTable;
 exports.generateEmbedding = generateEmbedding;
+exports.generateEmbeddings = generateEmbeddings;
 exports.getAllArticles = getAllArticles;
 exports.getArticleBySlug = getArticleBySlug;
 exports.getArticlesByFolder = getArticlesByFolder;
+exports.getEmbeddingProviderMetadata = getEmbeddingProviderMetadata;
 exports.getFolders = getFolders;
 exports.indexContent = indexContent;
 exports.padEmbedding = padEmbedding;
 exports.prepareTextForEmbedding = prepareTextForEmbedding;
 exports.search = search;
+exports.validateEmbeddingBatch = validateEmbeddingBatch;

package/dist/index.d.ts

@@ -1,14 +1,56 @@
 import { Client } from '@libsql/client';
 
 type EmbeddingProvider = 'local' | 'gemini' | 'openai';
+type EmbeddingIntent = 'document' | 'query';
+type EmbeddingBatchMode = 'native' | 'sequential';
+interface EmbeddingBatchBehavior {
+    mode: EmbeddingBatchMode;
+    maxSize?: number;
+}
+interface EmbeddingProviderMetadata {
+    name: EmbeddingProvider;
+    model: string;
+    dimensions: number;
+    batch: EmbeddingBatchBehavior;
+}
+interface EmbeddingRequestOptions {
+    intent?: EmbeddingIntent;
+    signal?: AbortSignal;
+}
+interface EmbeddingProviderClient {
+    readonly metadata: EmbeddingProviderMetadata;
+    embed(texts: string[], options?: EmbeddingRequestOptions): Promise<EmbeddingBatchResult>;
+}
+interface EmbeddingBatchResult {
+    embeddings: number[][];
+    provider: EmbeddingProvider;
+    model: string;
+    dimensions: number;
+    intent: EmbeddingIntent;
+}
 interface EmbeddingOptions {
     provider?: EmbeddingProvider;
     apiKey?: string;
     dimensions?: number;
     maxLength?: number;
+    intent?: EmbeddingIntent;
+    timeoutMs?: number;
+    signal?: AbortSignal;
 }
+interface EmbeddingBatchItemResult {
+    embedding: number[];
+    index?: unknown;
+}
+type EmbeddingBatchItem = number[] | EmbeddingBatchItemResult;
+declare function validateEmbeddingBatch(items: EmbeddingBatchItem[], expectedCount: number, expectedDimensions: number, provider: EmbeddingProvider): number[][];
+declare function createEmbeddingProvider(options?: EmbeddingOptions): EmbeddingProviderClient;
+declare function getEmbeddingProviderMetadata(options?: EmbeddingOptions): EmbeddingProviderMetadata;
+/**
+ * Generate embeddings using the specified provider.
+ */
+declare function generateEmbeddings(texts: string[], options?: EmbeddingOptions): Promise<number[][]>;
 /**
- * Generate embeddings using the specified provider
+ * Generate an embedding using the specified provider.
  */
 declare function generateEmbedding(text: string, options?: EmbeddingOptions): Promise<number[]>;
 /**
@@ -23,7 +65,7 @@ declare function prepareTextForEmbedding(fields: {
     description?: string;
     content?: string;
     tags?: string[];
-    [key: string]: any;
+    [key: string]: unknown;
 }): string;
 
 /**
@@ -126,5 +168,5 @@ declare function getArticlesByFolder(client: Client, folder: string, tableName?:
  */
 declare function getFolders(client: Client, tableName?: string): Promise<string[]>;
 
-export { createTable, generateEmbedding, getAllArticles, getArticleBySlug, getArticlesByFolder, getFolders, indexContent, padEmbedding, prepareTextForEmbedding, search };
-export type { EmbeddingOptions, EmbeddingProvider, IndexedDocument, IndexerOptions, SearchOptions, SearchResult };
+export { createEmbeddingProvider, createTable, generateEmbedding, generateEmbeddings, getAllArticles, getArticleBySlug, getArticlesByFolder, getEmbeddingProviderMetadata, getFolders, indexContent, padEmbedding, prepareTextForEmbedding, search, validateEmbeddingBatch };
+export type { EmbeddingBatchBehavior, EmbeddingBatchItem, EmbeddingBatchItemResult, EmbeddingBatchMode, EmbeddingBatchResult, EmbeddingIntent, EmbeddingOptions, EmbeddingProvider, EmbeddingProviderClient, EmbeddingProviderMetadata, EmbeddingRequestOptions, IndexedDocument, IndexerOptions, SearchOptions, SearchResult };