Fix: restore token-by-token streaming for chains using createTextOnlyOutputParser (#6086)

fix: restore token-by-token streaming for chains using createTextOnlyOutputParser

Author: HenryHengZJ

packages/components/src/utils.ts

@@ -11,7 +11,7 @@ import { ICommonObject, IDatabaseEntity, IFileUpload, IMessage, INodeData, IVari
 import { BaseChatModel } from '@langchain/core/language_models/chat_models'
 import { AES, enc } from 'crypto-js'
 import { AIMessage, AIMessageChunk, HumanMessage, BaseMessage } from '@langchain/core/messages'
-import { RunnableLambda } from '@langchain/core/runnables'
+import { Runnable, type RunnableConfig } from '@langchain/core/runnables'
 import { Document } from '@langchain/core/documents'
 import { getFileFromStorage } from './storageUtils'
 import { GetSecretValueCommand, SecretsManagerClient, SecretsManagerClientConfig } from '@aws-sdk/client-secrets-manager'
@@ -303,24 +303,54 @@ export const transformBracesWithColon = (input: string): string => {
 }
 
 /**
- * Creates a RunnableLambda that extracts text content from a chat model response,
- * filtering out reasoning/thinking content blocks that reasoning models may return.
+ * Extracts text content from an AIMessageChunk, filtering out reasoning/thinking
+ * content blocks that reasoning models may return.
+ */
+const extractTextFromChunk = (response: AIMessageChunk): string => {
+    if (typeof response.content === 'string') {
+        return response.content
+    }
+    if (Array.isArray(response.content)) {
+        return response.content
+            .filter((block: any) => block.type === 'text' || block.type === 'text_delta')
+            .map((block: any) => block.text ?? '')
+            .join('')
+    }
+    return ''
+}
+
+/**
+ * Creates a streaming-compatible output parser that extracts text content from
+ * chat model responses, filtering out reasoning/thinking content blocks.
+ * https://github.com/FlowiseAI/Flowise/pull/5893#issuecomment-4045466531
  */
 export const createTextOnlyOutputParser = () => {
-    return new RunnableLambda({
-        func: (response: AIMessageChunk) => {
-            if (typeof response.content === 'string') {
-                return response.content
-            }
-            if (Array.isArray(response.content)) {
-                return response.content
-                    .filter((block: any) => block.type === 'text' || block.type === 'text_delta')
-                    .map((block: any) => block.text ?? '')
-                    .join('')
+    return new TextOnlyOutputParser()
+}
+
+class TextOnlyOutputParser extends Runnable<AIMessageChunk, string> {
+    static lc_name() {
+        return 'TextOnlyOutputParser'
+    }
+
+    lc_namespace = ['flowise', 'output_parsers']
+
+    async invoke(input: AIMessageChunk, _options?: Partial<RunnableConfig>): Promise<string> {
+        return extractTextFromChunk(input)
+    }
+
+    async *_transform(generator: AsyncGenerator<AIMessageChunk>): AsyncGenerator<string> {
+        for await (const chunk of generator) {
+            const text = extractTextFromChunk(chunk)
+            if (text) {
+                yield text
             }
-            return ''
         }
-    })
+    }
+
+    async *transform(generator: AsyncGenerator<AIMessageChunk>, options?: Partial<RunnableConfig>): AsyncGenerator<string> {
+        yield* this._transformStreamWithConfig(generator, this._transform.bind(this), options)
+    }
 }
 
 /**