plugin-custom-llm 1.3.2 → 1.4.0

package/dist/server/llm-providers/custom-llm.js

@@ -542,51 +542,72 @@ function fixEmptyToolProperties(model) {
   };
   return model;
 }
-function sanitizeGenerateResult(result) {
-  if (!result) return result;
-  for (const gen of (result == null ? void 0 : result.generations) ?? []) {
-    const msg = gen == null ? void 0 : gen.message;
-    if (msg == null ? void 0 : msg.tool_calls) {
-      for (const tc of msg.tool_calls) {
-        tc.id = sanitizeToolCallId(tc.id);
-      }
-    }
-  }
-  return result;
-}
-function sanitizeStreamChunk(chunk) {
-  const msg = chunk == null ? void 0 : chunk.message;
-  if (msg == null ? void 0 : msg.tool_call_chunks) {
-    for (const tc of msg.tool_call_chunks) {
+function sanitizeAIMessageToolCalls(msg) {
+  if (!msg) return;
+  if (msg.tool_calls) {
+    for (const tc of msg.tool_calls) {
       tc.id = sanitizeToolCallId(tc.id);
     }
   }
-  if (msg == null ? void 0 : msg.tool_calls) {
-    for (const tc of msg.tool_calls) {
+  if (msg.tool_call_chunks) {
+    for (const tc of msg.tool_call_chunks) {
       tc.id = sanitizeToolCallId(tc.id);
     }
   }
-  return chunk;
 }
-function createSanitizedChatClass(BaseClass) {
-  return class SanitizedChatModel extends BaseClass {
-    async _generate(messages, options, runManager) {
-      const result = await super._generate(messages, options, runManager);
-      return sanitizeGenerateResult(result);
-    }
-    async *_streamResponseChunks(messages, options, runManager) {
-      for await (const chunk of super._streamResponseChunks(messages, options, runManager)) {
-        yield sanitizeStreamChunk(chunk);
-      }
-    }
-    async *_stream(messages, options, runManager) {
-      if (typeof super._stream === "function") {
-        for await (const chunk of super._stream(messages, options, runManager)) {
-          yield sanitizeStreamChunk(chunk);
-        }
+function patchRunnableForSanitization(runnable) {
+  var _a, _b, _c, _d;
+  if (!runnable || runnable.__toolCallSanitized) return runnable;
+  runnable.__toolCallSanitized = true;
+  const origInvoke = (_a = runnable.invoke) == null ? void 0 : _a.bind(runnable);
+  if (origInvoke) {
+    runnable.invoke = async function(...args) {
+      const result = await origInvoke(...args);
+      sanitizeAIMessageToolCalls(result);
+      return result;
+    };
+  }
+  const origStream = (_b = runnable.stream) == null ? void 0 : _b.bind(runnable);
+  if (origStream) {
+    runnable.stream = async function(...args) {
+      var _a2;
+      const iter = await origStream(...args);
+      const origIterator = (_a2 = iter[Symbol.asyncIterator]) == null ? void 0 : _a2.bind(iter);
+      if (origIterator) {
+        iter[Symbol.asyncIterator] = function() {
+          var _a3, _b2;
+          const it = origIterator();
+          return {
+            async next() {
+              const { value, done } = await it.next();
+              if (!done && value) {
+                sanitizeAIMessageToolCalls(value);
+              }
+              return { value, done };
+            },
+            return: (_a3 = it.return) == null ? void 0 : _a3.bind(it),
+            throw: (_b2 = it.throw) == null ? void 0 : _b2.bind(it)
+          };
+        };
       }
-    }
-  };
+      return iter;
+    };
+  }
+  const origBindTools = (_c = runnable.bindTools) == null ? void 0 : _c.bind(runnable);
+  if (origBindTools) {
+    runnable.bindTools = function(...args) {
+      const bound = origBindTools(...args);
+      return patchRunnableForSanitization(bound);
+    };
+  }
+  const origBind = (_d = runnable.bind) == null ? void 0 : _d.bind(runnable);
+  if (origBind) {
+    runnable.bind = function(...args) {
+      const bound = origBind(...args);
+      return patchRunnableForSanitization(bound);
+    };
+  }
+  return runnable;
 }
 class CustomLLMProvider extends import_plugin_ai.LLMProvider {
   get baseURL() {
@@ -625,8 +646,7 @@ class CustomLLMProvider extends import_plugin_ai.LLMProvider {
     if (reqConfig.extraBody && typeof reqConfig.extraBody === "object") {
       Object.assign(modelKwargs, reqConfig.extraBody);
     }
-    const BaseChatClass = enableReasoning ? createReasoningChatClass() : getChatOpenAI();
-    const ChatClass = createSanitizedChatClass(BaseChatClass);
+    const ChatClass = enableReasoning ? createReasoningChatClass() : getChatOpenAI();
     const config = {
       apiKey,
       ...this.modelOptions,
@@ -653,11 +673,12 @@ class CustomLLMProvider extends import_plugin_ai.LLMProvider {
     let model = new ChatClass(config);
     model = fixEmptyToolProperties(model);
     if (streamKeepAlive && !disableStream) {
-      return wrapWithStreamKeepAlive(model, {
+      model = wrapWithStreamKeepAlive(model, {
         intervalMs: Number(keepAliveIntervalMs) || 5e3,
         keepAliveContent: keepAliveContent || "..."
       });
     }
+    model = patchRunnableForSanitization(model);
     return model;
   }
   parseResponseChunk(chunk) {

package/package.json

@@ -3,7 +3,7 @@
   "displayName": "AI LLM: Custom (OpenAI Compatible)",
   "displayName.zh-CN": "AI LLM：自定义（OpenAI 兼容）",
   "description": "OpenAI-compatible LLM provider with auto response format detection for external LLM services.",
-  "version": "1.3.2",
+  "version": "1.4.0",
   "main": "dist/server/index.js",
   "files": [
     "dist",

package/src/server/llm-providers/custom-llm.ts

@@ -727,66 +727,89 @@ function fixEmptyToolProperties(model: any) {
 }
 
 /**
- * Sanitize all tool call IDs in a ChatResult (used after _generate).
+ * Sanitize tool_calls on an AIMessage (mutates in place).
  */
-function sanitizeGenerateResult(result: any): any {
-  if (!result) return result;
-  for (const gen of result?.generations ?? []) {
-    const msg = gen?.message;
-    if (msg?.tool_calls) {
-      for (const tc of msg.tool_calls) {
-        tc.id = sanitizeToolCallId(tc.id);
-      }
-    }
-  }
-  return result;
-}
-
-/**
- * Sanitize tool call IDs in a streaming chunk.
- */
-function sanitizeStreamChunk(chunk: any): any {
-  const msg = chunk?.message;
-  if (msg?.tool_call_chunks) {
-    for (const tc of msg.tool_call_chunks) {
+function sanitizeAIMessageToolCalls(msg: any): void {
+  if (!msg) return;
+  if (msg.tool_calls) {
+    for (const tc of msg.tool_calls) {
       tc.id = sanitizeToolCallId(tc.id);
     }
   }
-  if (msg?.tool_calls) {
-    for (const tc of msg.tool_calls) {
+  if (msg.tool_call_chunks) {
+    for (const tc of msg.tool_call_chunks) {
       tc.id = sanitizeToolCallId(tc.id);
     }
   }
-  return chunk;
 }
 
 /**
- * Create a subclass of the given ChatModel class that sanitizes tool call IDs
- * in all outputs. Gemini models return IDs like `call_xxx__thought__<long_base64>`
- * which are too long for langgraph. Using class-level overrides (instead of
- * instance patching) ensures the sanitization survives bindTools/RunnableBinding.
+ * Patch a runnable (model or bound model) so that `invoke` and `stream`
+ * sanitize tool call IDs on every AIMessage output.
+ * Also patches `bindTools` and `bind` so that derived runnables inherit
+ * the sanitization — this is critical because langgraph calls
+ * `model.bindTools(tools)` and then uses the BOUND model.
  */
-function createSanitizedChatClass(BaseClass: any) {
-  return class SanitizedChatModel extends BaseClass {
-    async _generate(messages: any[], options: any, runManager?: any) {
-      const result = await super._generate(messages, options, runManager);
-      return sanitizeGenerateResult(result);
-    }
+function patchRunnableForSanitization(runnable: any): any {
+  if (!runnable || runnable.__toolCallSanitized) return runnable;
+  runnable.__toolCallSanitized = true;
+
+  // Patch invoke — covers non-streaming and internal streaming-via-invoke
+  const origInvoke = runnable.invoke?.bind(runnable);
+  if (origInvoke) {
+    runnable.invoke = async function (...args: any[]) {
+      const result = await origInvoke(...args);
+      sanitizeAIMessageToolCalls(result);
+      return result;
+    };
+  }
 
-    async *_streamResponseChunks(messages: any[], options: any, runManager?: any) {
-      for await (const chunk of super._streamResponseChunks(messages, options, runManager)) {
-        yield sanitizeStreamChunk(chunk);
+  // Patch stream — covers streaming path
+  const origStream = runnable.stream?.bind(runnable);
+  if (origStream) {
+    runnable.stream = async function (...args: any[]) {
+      const iter = await origStream(...args);
+      // Wrap the async iterable to sanitize each chunk
+      const origIterator = iter[Symbol.asyncIterator]?.bind(iter);
+      if (origIterator) {
+        iter[Symbol.asyncIterator] = function () {
+          const it = origIterator();
+          return {
+            async next() {
+              const { value, done } = await it.next();
+              if (!done && value) {
+                sanitizeAIMessageToolCalls(value);
+              }
+              return { value, done };
+            },
+            return: it.return?.bind(it),
+            throw: it.throw?.bind(it),
+          };
+        };
       }
-    }
+      return iter;
+    };
+  }
 
-    async *_stream(messages: any[], options: any, runManager?: any) {
-      if (typeof super._stream === 'function') {
-        for await (const chunk of super._stream(messages, options, runManager)) {
-          yield sanitizeStreamChunk(chunk);
-        }
-      }
-    }
-  };
+  // Patch bindTools — the result is a RunnableBinding used by langgraph
+  const origBindTools = runnable.bindTools?.bind(runnable);
+  if (origBindTools) {
+    runnable.bindTools = function (...args: any[]) {
+      const bound = origBindTools(...args);
+      return patchRunnableForSanitization(bound);
+    };
+  }
+
+  // Patch bind — bindTools internally calls bind(), some runnables use it directly
+  const origBind = runnable.bind?.bind(runnable);
+  if (origBind) {
+    runnable.bind = function (...args: any[]) {
+      const bound = origBind(...args);
+      return patchRunnableForSanitization(bound);
+    };
+  }
+
+  return runnable;
 }
 
 export class CustomLLMProvider extends LLMProvider {
@@ -829,11 +852,7 @@ export class CustomLLMProvider extends LLMProvider {
     // Issue #4: Use ReasoningChatOpenAI when enableReasoning is set.
     // This ensures reasoning_content is preserved and patched back into
     // assistant messages during tool call round-trips (required by DeepSeek-R1, etc.)
-    // Wrap with tool call ID sanitizer at the class level — ensures
-    // __thought__<base64> suffixes from Gemini are stripped in all code paths
-    // (invoke, stream, bindTools bindings) via prototype chain.
-    const BaseChatClass = enableReasoning ? createReasoningChatClass() : getChatOpenAI();
-    const ChatClass = createSanitizedChatClass(BaseChatClass);
+    const ChatClass = enableReasoning ? createReasoningChatClass() : getChatOpenAI();
     const config: Record<string, any> = {
       apiKey,
       ...this.modelOptions,
@@ -874,12 +893,18 @@ export class CustomLLMProvider extends LLMProvider {
 
     // Wrap with keepalive proxy if enabled (and streaming is not disabled)
     if (streamKeepAlive && !disableStream) {
-      return wrapWithStreamKeepAlive(model, {
+      model = wrapWithStreamKeepAlive(model, {
         intervalMs: Number(keepAliveIntervalMs) || 5000,
         keepAliveContent: keepAliveContent || '...',
       });
     }
 
+    // Sanitize Gemini's __thought__<base64> suffixes in tool call IDs.
+    // Patches invoke/stream/bindTools/bind at the public API level so that
+    // ALL code paths (including langgraph's internal model calls via
+    // RunnableBinding after bindTools) return clean IDs.
+    model = patchRunnableForSanitization(model);
+
     return model;
   }