@samanhappy/mcphub 0.12.7 → 0.12.8

package/dist/utils/tokenTruncation.js

@@ -0,0 +1,243 @@
+/**
+ * Token-aware text truncation utilities for embedding generation.
+ *
+ * Provides precise tokenization for known model families and a conservative
+ * heuristic fallback for unknown models.
+ *
+ * Model families and strategies:
+ * - OpenAI / Azure (text-embedding-*):  BPE cl100k_base via gpt-tokenizer (exact)
+ * - BAAI/BGE and HuggingFace models:    AutoTokenizer via @huggingface/transformers (exact)
+ * - Google Gemini (gemini-embedding-*): countTokens API via @google/genai (exact)
+ * - Unknown models:                     heuristic maxTokens * 3 chars (approximate)
+ */
+/**
+ * Per-model token limits.
+ * Order matters: more specific entries must appear before generic ones.
+ * bge-m3 is explicitly listed before the generic 'bge' catch-all because
+ * its real limit (8192 tokens) differs substantially from the conservative
+ * 512 used for other BGE variants.
+ */
+const MODEL_TOKEN_LIMITS = [
+    ['text-embedding-3-small', 8191],
+    ['text-embedding-3-large', 8191],
+    ['text-embedding-ada-002', 8191],
+    ['gemini-embedding-001', 2048],
+    ['bge-m3', 8192],
+];
+/**
+ * Returns the default maximum token limit for a given embedding model name.
+ * Used when no explicit limit is configured via EMBEDDING_MAX_TOKENS or
+ * smartRouting.embeddingMaxTokens.
+ */
+export function getModelDefaultTokenLimit(model) {
+    const lower = model.toLowerCase();
+    for (const [pattern, limit] of MODEL_TOKEN_LIMITS) {
+        if (lower.includes(pattern)) {
+            return limit;
+        }
+    }
+    // For other BGE variants (bge-large-en, bge-small-zh, etc.) use conservative limit
+    if (lower.includes('bge')) {
+        return 512;
+    }
+    // Default conservative limit: safe for entirely unknown models.
+    // Users can raise it with EMBEDDING_MAX_TOKENS if they know their model supports more.
+    return 512;
+}
+// ─────────────────────────────────────────────────────────────────────────────
+// Model family detection helpers
+// ─────────────────────────────────────────────────────────────────────────────
+const OPENAI_MODELS = new Set(['text-embedding-3-small', 'text-embedding-3-large', 'text-embedding-ada-002']);
+function isOpenAIModel(model) {
+    return OPENAI_MODELS.has(model.toLowerCase());
+}
+function isGeminiModel(model) {
+    return model.toLowerCase() === 'gemini-embedding-001';
+}
+function isBgeM3Model(model) {
+    return model.toLowerCase().includes('bge-m3');
+}
+// ─────────────────────────────────────────────────────────────────────────────
+// Branch 1 — OpenAI / Azure: BPE cl100k_base via gpt-tokenizer
+// ─────────────────────────────────────────────────────────────────────────────
+/**
+ * Truncates text using OpenAI's BPE tokenizer (cl100k_base).
+ *
+ * Encodes the input text into tokens using the gpt-tokenizer library,
+ * which implements the exact cl100k_base BPE vocabulary used by OpenAI's
+ * embedding models (text-embedding-3-small, text-embedding-3-large, text-embedding-ada-002).
+ * If the token count exceeds maxTokens, decodes the truncated token sequence back to text.
+ *
+ * @param text       The input text to truncate.
+ * @param maxTokens  The maximum number of tokens allowed.
+ * @returns          The original text if it fits, or a truncated prefix.
+ */
+async function truncateWithGptTokenizer(text, maxTokens) {
+    const { encode, decode } = await import('gpt-tokenizer');
+    const tokens = encode(text);
+    if (tokens.length <= maxTokens) {
+        return text;
+    }
+    return decode(tokens.slice(0, maxTokens));
+}
+// ─────────────────────────────────────────────────────────────────────────────
+// Branch 2 — HuggingFace / BGE: AutoTokenizer (no ONNX, pure JS tokenisation)
+// ─────────────────────────────────────────────────────────────────────────────
+// Tokenizer instances are cached in memory to avoid repeated HF Hub downloads
+const tokenizerCache = new Map();
+/**
+ * Fetches or retrieves a cached HuggingFace tokenizer for a given model.
+ *
+ * The tokenizer is downloaded from HuggingFace Hub (public models like BAAI/bge-m3
+ * do not require authentication). Once downloaded, the tokenizer is cached in memory
+ * to avoid redundant network requests on subsequent calls.
+ *
+ * @param modelId  The fully-qualified HuggingFace Hub model ID (e.g., "BAAI/bge-m3").
+ * @returns        The cached or freshly-downloaded tokenizer instance.
+ */
+async function getHFTokenizer(modelId) {
+    if (!tokenizerCache.has(modelId)) {
+        const { AutoTokenizer } = await import('@huggingface/transformers');
+        const tokenizer = await AutoTokenizer.from_pretrained(modelId);
+        tokenizerCache.set(modelId, tokenizer);
+    }
+    return tokenizerCache.get(modelId);
+}
+/**
+ * Resolves a shorthand model name to a fully-qualified HuggingFace Hub repo ID.
+ * BAAI/bge-m3 is a public model — no HF_TOKEN is required to download its
+ * tokenizer.json file.
+ */
+function getHFModelId(model) {
+    if (model.includes('/')) {
+        // Already fully qualified (e.g. "BAAI/bge-m3", "sentence-transformers/all-MiniLM-L6-v2")
+        return model;
+    }
+    const lower = model.toLowerCase();
+    if (lower.includes('bge-m3')) {
+        return 'BAAI/bge-m3';
+    }
+    if (lower.includes('bge')) {
+        return `BAAI/${model}`;
+    }
+    return model;
+}
+/**
+ * Truncates text using a HuggingFace AutoTokenizer (BAAI/BGE and other transformer models).
+ *
+ * Leverages the @huggingface/transformers library to tokenize input text using
+ * the model's own SentencePiece or WordPiece tokenizer, then decodes a truncated
+ * token sequence back to text. This provides exact tokenization matching the model's
+ * vocabulary and behavior, with tokenizer instances cached to avoid repeated downloads.
+ *
+ * @param text       The input text to truncate.
+ * @param maxTokens  The maximum number of tokens allowed.
+ * @param model      The model identifier (shorthand or fully-qualified HF Hub ID).
+ * @returns          The original text if it fits, or a truncated prefix.
+ */
+async function truncateWithHFTokenizer(text, maxTokens, model) {
+    const modelId = getHFModelId(model);
+    const tokenizer = await getHFTokenizer(modelId);
+    // Tokenize without automatic truncation so we can apply the exact limit
+    const encoded = await tokenizer(text, { padding: false, truncation: false });
+    // input_ids.data is BigInt64Array or Int32Array depending on the model/environment
+    const rawIds = encoded.input_ids.data;
+    const ids = Array.from(rawIds).map(Number);
+    if (ids.length <= maxTokens) {
+        return text;
+    }
+    const truncatedIds = ids.slice(0, maxTokens);
+    return (await tokenizer.decode(truncatedIds, { skip_special_tokens: true }));
+}
+// ─────────────────────────────────────────────────────────────────────────────
+// Branch 3 — Google Gemini: countTokens API with binary-search bisection
+// ─────────────────────────────────────────────────────────────────────────────
+/**
+ * Truncates text using Google Gemini's countTokens API.
+ *
+ * Uses the @google/genai library to query the exact token count from Gemini's
+ * tokenizer (SentencePiece). To minimize API calls, a binary-search algorithm
+ * finds the longest text prefix whose token count does not exceed maxTokens.
+ * If no Google API key is configured, falls back to a conservative 3× char heuristic.
+ *
+ * @param text       The input text to truncate.
+ * @param maxTokens  The maximum number of tokens allowed.
+ * @param model      The Gemini model identifier (e.g., "gemini-embedding-001").
+ * @returns          The original text if it fits, or a truncated prefix.
+ */
+async function truncateWithGeminiAPI(text, maxTokens, model, apiKey) {
+    // Pre-filter: if char count is safely within 2× the token limit the text fits
+    // (each SentencePiece token is at least 1 character), so skip the network call.
+    if (text.length <= maxTokens * 2) {
+        return text;
+    }
+    // Use the apiKey provided from smartRouting config (with priority: env var → settings → default)
+    const finalApiKey = apiKey || '';
+    if (!finalApiKey) {
+        // No Google Gemini API key configured (OPENAI_API_KEY) — fall back to conservative heuristic
+        const maxChars = maxTokens * 3;
+        return text.length <= maxChars ? text : text.substring(0, maxChars);
+    }
+    const { GoogleGenAI } = await import('@google/genai');
+    const ai = new GoogleGenAI({ apiKey: finalApiKey });
+    const countTokens = async (chunk) => {
+        const result = await ai.models.countTokens({ model, contents: chunk });
+        return result.totalTokens ?? 0;
+    };
+    const totalTokens = await countTokens(text);
+    if (totalTokens <= maxTokens) {
+        return text;
+    }
+    // Binary search: find the longest prefix whose token count ≤ maxTokens.
+    // This minimizes the number of countTokens calls (O(log n) on text length).
+    let lo = 0;
+    let hi = text.length;
+    while (lo < hi - 1) {
+        const mid = Math.floor((lo + hi) / 2);
+        const count = await countTokens(text.slice(0, mid));
+        if (count <= maxTokens) {
+            lo = mid;
+        }
+        else {
+            hi = mid;
+        }
+    }
+    return text.slice(0, lo);
+}
+// ─────────────────────────────────────────────────────────────────────────────
+// Main entry point
+// ─────────────────────────────────────────────────────────────────────────────
+/**
+ * Truncates `text` so that its token count does not exceed `maxTokens`,
+ * using the most accurate tokenization strategy available for the given model.
+ *
+ * The function is async because the Gemini branch may require a network call.
+ * All callers in vectorSearchService.ts must use `await`.
+ *
+ * NOTE: When using Google Gemini embeddings (gemini-embedding-*), the Google AI
+ * Studio API key must be provided via the `OPENAI_API_KEY` environment variable,
+ * not `GOOGLE_API_KEY`. This allows centralized API key configuration across all
+ * embedding model families.
+ *
+ * @param text      Input text to potentially truncate.
+ * @param maxTokens Maximum number of tokens allowed.
+ * @param model     Embedding model identifier (selects truncation strategy).
+ * @param apiKey    Optional API key for Gemini models (from smartRouting config).
+ * @returns         The original text if it fits, or a truncated prefix.
+ */
+export async function truncateToTokenLimit(text, maxTokens, model, apiKey) {
+    if (isOpenAIModel(model)) {
+        return truncateWithGptTokenizer(text, maxTokens);
+    }
+    if (isGeminiModel(model)) {
+        return truncateWithGeminiAPI(text, maxTokens, model, apiKey);
+    }
+    if (isBgeM3Model(model)) {
+        return truncateWithHFTokenizer(text, maxTokens, model);
+    }
+    // Fallback heuristic: ~3 chars per token (conservative for CJK/multilingual).
+    // Ratio is safe for English (~4 chars/token) and CJK (~2 chars/token).
+    const maxChars = maxTokens * 3;
+    return text.length <= maxChars ? text : text.substring(0, maxChars);
+}
+//# sourceMappingURL=tokenTruncation.js.map

package/dist/utils/tokenTruncation.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"tokenTruncation.js","sourceRoot":"","sources":["../../src/utils/tokenTruncation.ts"],"names":[],"mappings":"AAAA;;;;;;;;;;;GAWG;AAEH;;;;;;GAMG;AACH,MAAM,kBAAkB,GAA4B;IAClD,CAAC,wBAAwB,EAAE,IAAI,CAAC;IAChC,CAAC,wBAAwB,EAAE,IAAI,CAAC;IAChC,CAAC,wBAAwB,EAAE,IAAI,CAAC;IAChC,CAAC,sBAAsB,EAAE,IAAI,CAAC;IAC9B,CAAC,QAAQ,EAAE,IAAI,CAAC;CACjB,CAAC;AAEF;;;;GAIG;AACH,MAAM,UAAU,yBAAyB,CAAC,KAAa;IACrD,MAAM,KAAK,GAAG,KAAK,CAAC,WAAW,EAAE,CAAC;IAClC,KAAK,MAAM,CAAC,OAAO,EAAE,KAAK,CAAC,IAAI,kBAAkB,EAAE,CAAC;QAClD,IAAI,KAAK,CAAC,QAAQ,CAAC,OAAO,CAAC,EAAE,CAAC;YAC5B,OAAO,KAAK,CAAC;QACf,CAAC;IACH,CAAC;IACD,mFAAmF;IACnF,IAAI,KAAK,CAAC,QAAQ,CAAC,KAAK,CAAC,EAAE,CAAC;QAC1B,OAAO,GAAG,CAAC;IACb,CAAC;IACD,gEAAgE;IAChE,uFAAuF;IACvF,OAAO,GAAG,CAAC;AACb,CAAC;AAED,gFAAgF;AAChF,iCAAiC;AACjC,gFAAgF;AAEhF,MAAM,aAAa,GAAG,IAAI,GAAG,CAAC,CAAC,wBAAwB,EAAE,wBAAwB,EAAE,wBAAwB,CAAC,CAAC,CAAC;AAE9G,SAAS,aAAa,CAAC,KAAa;IAClC,OAAO,aAAa,CAAC,GAAG,CAAC,KAAK,CAAC,WAAW,EAAE,CAAC,CAAC;AAChD,CAAC;AAED,SAAS,aAAa,CAAC,KAAa;IAClC,OAAO,KAAK,CAAC,WAAW,EAAE,KAAK,sBAAsB,CAAC;AACxD,CAAC;AAED,SAAS,YAAY,CAAC,KAAa;IACjC,OAAO,KAAK,CAAC,WAAW,EAAE,CAAC,QAAQ,CAAC,QAAQ,CAAC,CAAC;AAChD,CAAC;AAED,gFAAgF;AAChF,+DAA+D;AAC/D,gFAAgF;AAEhF;;;;;;;;;;;GAWG;AACH,KAAK,UAAU,wBAAwB,CAAC,IAAY,EAAE,SAAiB;IACrE,MAAM,EAAE,MAAM,EAAE,MAAM,EAAE,GAAG,MAAM,MAAM,CAAC,eAAe,CAAC,CAAC;IACzD,MAAM,MAAM,GAAG,MAAM,CAAC,IAAI,CAAC,CAAC;IAC5B,IAAI,MAAM,CAAC,MAAM,IAAI,SAAS,EAAE,CAAC;QAC/B,OAAO,IAAI,CAAC;IACd,CAAC;IACD,OAAO,MAAM,CAAC,MAAM,CAAC,KAAK,CAAC,CAAC,EAAE,SAAS,CAAC,CAAC,CAAC;AAC5C,CAAC;AAED,gFAAgF;AAChF,8EAA8E;AAC9E,gFAAgF;AAEhF,8EAA8E;AAC9E,MAAM,cAAc,GAAG,IAAI,GAAG,EAA2I,CAAC;AAE1K;;;;;;;;;GASG;AAEH,KAAK,UAAU,cAAc,CAAC,OAAe;IAC3C,IAAI,CAAC,cAAc,CAAC,GAAG,CAAC,OAAO,CAAC,EAAE,CAAC;QACjC,MAAM,EAAE,aAAa,EAAE,GAAG,MAAM,MAAM,CAAC,2BAA2B,CAAC,CAAC;QACpE,MAAM,SAAS,GAAG,MAAM,aAAa,CAAC,eAAe,CAAC,OAAO,CAAC,CAAC;QAC/D,cAAc,CAAC,GAAG,CAAC,OAAO,EAAE,SAAS,CAAC,CAAC;IACzC,CAAC;IACD,OAAO,cAAc,CAAC,GAAG,CAAC,OAAO,CAAC,CAAC;AACrC,CAAC;AAED;;;;GAIG;AACH,SAAS,YAAY,CAAC,KAAa;IACjC,IAAI,KAAK,CAAC,QAAQ,CAAC,GAAG,CAAC,EAAE,CAAC;QACxB,yFAAyF;QACzF,OAAO,KAAK,CAAC;IACf,CAAC;IACD,MAAM,KAAK,GAAG,KAAK,CAAC,WAAW,EAAE,CAAC;IAClC,IAAI,KAAK,CAAC,QAAQ,CAAC,QAAQ,CAAC,EAAE,CAAC;QAC7B,OAAO,aAAa,CAAC;IACvB,CAAC;IACD,IAAI,KAAK,CAAC,QAAQ,CAAC,KAAK,CAAC,EAAE,CAAC;QAC1B,OAAO,QAAQ,KAAK,EAAE,CAAC;IACzB,CAAC;IACD,OAAO,KAAK,CAAC;AACf,CAAC;AAED;;;;;;;;;;;;GAYG;AACH,KAAK,UAAU,uBAAuB,CACpC,IAAY,EACZ,SAAiB,EACjB,KAAa;IAEb,MAAM,OAAO,GAAG,YAAY,CAAC,KAAK,CAAC,CAAC;IACpC,MAAM,SAAS,GAAG,MAAM,cAAc,CAAC,OAAO,CAAC,CAAC;IAChD,wEAAwE;IACxE,MAAM,OAAO,GAAG,MAAM,SAAS,CAAC,IAAI,EAAE,EAAE,OAAO,EAAE,KAAK,EAAE,UAAU,EAAE,KAAK,EAAE,CAAC,CAAC;IAC7E,mFAAmF;IACnF,MAAM,MAAM,GAAgC,OAAO,CAAC,SAAkD,CAAC,IAAI,CAAC;IAC5G,MAAM,GAAG,GAAG,KAAK,CAAC,IAAI,CAAC,MAA2B,CAAC,CAAC,GAAG,CAAC,MAAM,CAAC,CAAC;IAChE,IAAI,GAAG,CAAC,MAAM,IAAI,SAAS,EAAE,CAAC;QAC5B,OAAO,IAAI,CAAC;IACd,CAAC;IACD,MAAM,YAAY,GAAG,GAAG,CAAC,KAAK,CAAC,CAAC,EAAE,SAAS,CAAC,CAAC;IAC7C,OAAO,CAAC,MAAM,SAAS,CAAC,MAAM,CAAC,YAAY,EAAE,EAAE,mBAAmB,EAAE,IAAI,EAAE,CAAC,CAAW,CAAC;AACzF,CAAC;AAED,gFAAgF;AAChF,yEAAyE;AACzE,gFAAgF;AAEhF;;;;;;;;;;;;GAYG;AACH,KAAK,UAAU,qBAAqB,CAClC,IAAY,EACZ,SAAiB,EACjB,KAAa,EACb,MAAe;IAEf,8EAA8E;IAC9E,gFAAgF;IAChF,IAAI,IAAI,CAAC,MAAM,IAAI,SAAS,GAAG,CAAC,EAAE,CAAC;QACjC,OAAO,IAAI,CAAC;IACd,CAAC;IAED,iGAAiG;IACjG,MAAM,WAAW,GAAG,MAAM,IAAI,EAAE,CAAC;IACjC,IAAI,CAAC,WAAW,EAAE,CAAC;QACjB,6FAA6F;QAC7F,MAAM,QAAQ,GAAG,SAAS,GAAG,CAAC,CAAC;QAC/B,OAAO,IAAI,CAAC,MAAM,IAAI,QAAQ,CAAC,CAAC,CAAC,IAAI,CAAC,CAAC,CAAC,IAAI,CAAC,SAAS,CAAC,CAAC,EAAE,QAAQ,CAAC,CAAC;IACtE,CAAC;IAED,MAAM,EAAE,WAAW,EAAE,GAAG,MAAM,MAAM,CAAC,eAAe,CAAC,CAAC;IACtD,MAAM,EAAE,GAAG,IAAI,WAAW,CAAC,EAAE,MAAM,EAAE,WAAW,EAAE,CAAC,CAAC;IAEpD,MAAM,WAAW,GAAG,KAAK,EAAE,KAAa,EAAmB,EAAE;QAC3D,MAAM,MAAM,GAAG,MAAM,EAAE,CAAC,MAAM,CAAC,WAAW,CAAC,EAAE,KAAK,EAAE,QAAQ,EAAE,KAAK,EAAE,CAAC,CAAC;QACvE,OAAO,MAAM,CAAC,WAAW,IAAI,CAAC,CAAC;IACjC,CAAC,CAAC;IAEF,MAAM,WAAW,GAAG,MAAM,WAAW,CAAC,IAAI,CAAC,CAAC;IAC5C,IAAI,WAAW,IAAI,SAAS,EAAE,CAAC;QAC7B,OAAO,IAAI,CAAC;IACd,CAAC;IAED,wEAAwE;IACxE,4EAA4E;IAC5E,IAAI,EAAE,GAAG,CAAC,CAAC;IACX,IAAI,EAAE,GAAG,IAAI,CAAC,MAAM,CAAC;IACrB,OAAO,EAAE,GAAG,EAAE,GAAG,CAAC,EAAE,CAAC;QACnB,MAAM,GAAG,GAAG,IAAI,CAAC,KAAK,CAAC,CAAC,EAAE,GAAG,EAAE,CAAC,GAAG,CAAC,CAAC,CAAC;QACtC,MAAM,KAAK,GAAG,MAAM,WAAW,CAAC,IAAI,CAAC,KAAK,CAAC,CAAC,EAAE,GAAG,CAAC,CAAC,CAAC;QACpD,IAAI,KAAK,IAAI,SAAS,EAAE,CAAC;YACvB,EAAE,GAAG,GAAG,CAAC;QACX,CAAC;aAAM,CAAC;YACN,EAAE,GAAG,GAAG,CAAC;QACX,CAAC;IACH,CAAC;IACD,OAAO,IAAI,CAAC,KAAK,CAAC,CAAC,EAAE,EAAE,CAAC,CAAC;AAC3B,CAAC;AAED,gFAAgF;AAChF,mBAAmB;AACnB,gFAAgF;AAEhF;;;;;;;;;;;;;;;;;GAiBG;AACH,MAAM,CAAC,KAAK,UAAU,oBAAoB,CACxC,IAAY,EACZ,SAAiB,EACjB,KAAa,EACb,MAAe;IAEf,IAAI,aAAa,CAAC,KAAK,CAAC,EAAE,CAAC;QACzB,OAAO,wBAAwB,CAAC,IAAI,EAAE,SAAS,CAAC,CAAC;IACnD,CAAC;IACD,IAAI,aAAa,CAAC,KAAK,CAAC,EAAE,CAAC;QACzB,OAAO,qBAAqB,CAAC,IAAI,EAAE,SAAS,EAAE,KAAK,EAAE,MAAM,CAAC,CAAC;IAC/D,CAAC;IACD,IAAI,YAAY,CAAC,KAAK,CAAC,EAAE,CAAC;QACxB,OAAO,uBAAuB,CAAC,IAAI,EAAE,SAAS,EAAE,KAAK,CAAC,CAAC;IACzD,CAAC;IACD,8EAA8E;IAC9E,uEAAuE;IACvE,MAAM,QAAQ,GAAG,SAAS,GAAG,CAAC,CAAC;IAC/B,OAAO,IAAI,CAAC,MAAM,IAAI,QAAQ,CAAC,CAAC,CAAC,IAAI,CAAC,CAAC,CAAC,IAAI,CAAC,SAAS,CAAC,CAAC,EAAE,QAAQ,CAAC,CAAC;AACtE,CAAC"}