Skip to content

feat: support streaming reasoning chunks #6607

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
WiiliamC wants to merge 1 commit into
base: main
Choose a base branch
from
+339 −93
Open

feat: support streaming reasoning chunks #6607

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
72 changes: 25 additions & 47 deletions packages/components/nodes/agentflow/Agent/Agent.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,8 @@ import {
convertMultiOptionsToStringArray,
processTemplateVariables,
configureStructuredOutput,
extractResponseContent
extractResponseContent,
extractStreamingChunkContent
} from '../../../src/utils'
import { sanitizeFileName } from '../../../src/validator'
import { getModelConfigByModelName, MODEL_TYPE } from '../../../src/modelLoader'
Expand Down Expand Up @@ -1876,62 +1877,39 @@ class Agent_Agentflow implements INode {
try {
for await (const chunk of await llmNodeInstance.stream(messages, { signal: abortController?.signal })) {
if (sseStreamer && !isStructuredOutput) {
let content = ''

if (chunk.contentBlocks?.length) {
for (const block of chunk.contentBlocks) {
if (isLastNode) {
// As soon as we see the first non-reasoning block, send last thinking event with duration
if (block.type !== 'reasoning' && wasThinking && !sentLastThinkingEvent && thinkingStartTime != null) {
thinkingDuration = Math.round((Date.now() - thinkingStartTime) / 1000)
sseStreamer.streamThinkingEvent(chatId, '', thinkingDuration)
sentLastThinkingEvent = true
}
if (block.type === 'reasoning' && (block as { reasoning?: string }).reasoning) {
if (!thinkingStartTime) {
thinkingStartTime = Date.now()
}
wasThinking = true
const reasoningContent = (block as { reasoning: string }).reasoning
sseStreamer.streamThinkingEvent(chatId, reasoningContent)
reasonContent += reasoningContent
}
const { text: content, reasoning } = extractStreamingChunkContent(chunk, { includeCodeExecution: true })

if (reasoning) {
reasonContent += reasoning
if (isLastNode) {
if (!thinkingStartTime) {
thinkingStartTime = Date.now()
}
wasThinking = true
sseStreamer.streamThinkingEvent(chatId, reasoning)
}
}

if (typeof chunk === 'string') {
content = chunk
} else if (Array.isArray(chunk.content) && chunk.content.length > 0) {
content = chunk.content
.map((item: any) => {
if ((item.text && !item.type) || (item.type === 'text' && item.text)) {
return item.text
} else if (item.type === 'executableCode' && item.executableCode) {
const language = item.executableCode.language?.toLowerCase() || 'python'
return `\n\`\`\`${language}\n${item.executableCode.code}\n\`\`\`\n`
} else if (item.type === 'codeExecutionResult' && item.codeExecutionResult) {
const outcome = item.codeExecutionResult.outcome || 'OUTCOME_OK'
const output = item.codeExecutionResult.output || ''
if (outcome === 'OUTCOME_OK' && output) {
return `**Code Output:**\n\`\`\`\n${output}\n\`\`\`\n`
} else if (outcome !== 'OUTCOME_OK') {
return `**Code Execution Error:**\n\`\`\`\n${output}\n\`\`\`\n`
}
}
return ''
})
.filter((text: string) => text)
.join('')
} else if (chunk.content) {
content = chunk.content.toString()
if (content && isLastNode && wasThinking && !sentLastThinkingEvent && thinkingStartTime != null) {
thinkingDuration = Math.round((Date.now() - thinkingStartTime) / 1000)
sseStreamer.streamThinkingEvent(chatId, '', thinkingDuration)
sentLastThinkingEvent = true
}

if (content) {
sseStreamer.streamTokenEvent(chatId, content)
}
sseStreamer.streamTokenEvent(chatId, content)
}

const messageChunk = typeof chunk === 'string' ? new AIMessageChunk(chunk) : chunk
response = response.concat(messageChunk)
}

if (sseStreamer && !isStructuredOutput && isLastNode && wasThinking && !sentLastThinkingEvent && thinkingStartTime != null) {
thinkingDuration = Math.round((Date.now() - thinkingStartTime) / 1000)
sseStreamer.streamThinkingEvent(chatId, '', thinkingDuration)
sentLastThinkingEvent = true
}
} catch (error) {
console.error('Error during streaming:', error)
throw error
Expand Down
39 changes: 39 additions & 0 deletions packages/components/nodes/agentflow/Agent/AgentStreaming.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,39 @@
import { AIMessageChunk } from '@langchain/core/messages'

const { nodeClass: AgentAgentflow } = require('./Agent')

describe('Agent AgentFlow streaming reasoning chunks', () => {
it('streams reasoning-only chunks as thinking content without empty token events', async () => {
const agentNode = new AgentAgentflow()
const chunk = new AIMessageChunk({
content: '',
additional_kwargs: {
reasoning_content: 'Think through the tool plan.'
}
})
const llmNodeInstance = {
stream: jest.fn(async function* () {
yield chunk
})
}
const sseStreamer = {
streamThinkingEvent: jest.fn(),
streamTokenEvent: jest.fn()
}

const response = await (agentNode as any).handleStreamingResponse(
sseStreamer,
llmNodeInstance,
[],
'chat-id',
new AbortController(),
false,
true
)

expect(sseStreamer.streamThinkingEvent).toHaveBeenCalledWith('chat-id', 'Think through the tool plan.')
expect(sseStreamer.streamThinkingEvent).toHaveBeenCalledWith('chat-id', '', expect.any(Number))
expect(sseStreamer.streamTokenEvent).not.toHaveBeenCalled()
expect(response.additional_kwargs?.reasoning_content).toBe('Think through the tool plan.')
})
})
58 changes: 29 additions & 29 deletions packages/components/nodes/agentflow/LLM/LLM.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,12 @@ import {
replaceInlineDataWithFileReferences,
updateFlowState
} from '../utils'
import { processTemplateVariables, configureStructuredOutput, extractResponseContent } from '../../../src/utils'
import {
processTemplateVariables,
configureStructuredOutput,
extractResponseContent,
extractStreamingChunkContent
} from '../../../src/utils'
import { getModelConfigByModelName, MODEL_TYPE } from '../../../src/modelLoader'
import { flatten } from 'lodash'

Expand Down Expand Up @@ -867,44 +872,39 @@ class LLM_Agentflow implements INode {
try {
for await (const chunk of await llmNodeInstance.stream(messages, { signal: abortController?.signal })) {
if (sseStreamer && !isStructuredOutput) {
let content = ''

if (chunk.contentBlocks?.length) {
for (const block of chunk.contentBlocks) {
if (isLastNode) {
// As soon as we see the first non-reasoning block, send last thinking event with duration (only when isLastNode)
if (block.type !== 'reasoning' && wasThinking && !sentLastThinkingEvent && thinkingStartTime != null) {
thinkingDuration = Math.round((Date.now() - thinkingStartTime) / 1000)
sseStreamer.streamThinkingEvent(chatId, '', thinkingDuration)
sentLastThinkingEvent = true
}
if (block.type === 'reasoning' && (block as { reasoning?: string }).reasoning) {
if (!thinkingStartTime) {
thinkingStartTime = Date.now()
}
wasThinking = true
const reasoningContent = (block as { reasoning: string }).reasoning
sseStreamer.streamThinkingEvent(chatId, reasoningContent)
reasonContent += reasoningContent
}
const { text: content, reasoning } = extractStreamingChunkContent(chunk)

if (reasoning) {
reasonContent += reasoning
if (isLastNode) {
if (!thinkingStartTime) {
thinkingStartTime = Date.now()
}
wasThinking = true
sseStreamer.streamThinkingEvent(chatId, reasoning)
}
}

if (typeof chunk === 'string') {
content = chunk
} else if (Array.isArray(chunk.content) && chunk.content.length > 0) {
const contents = chunk.content as ContentBlock.Text[]
content = contents.map((item) => item.text).join('')
} else if (chunk.content) {
content = chunk.content.toString()
if (content && isLastNode && wasThinking && !sentLastThinkingEvent && thinkingStartTime != null) {
thinkingDuration = Math.round((Date.now() - thinkingStartTime) / 1000)
sseStreamer.streamThinkingEvent(chatId, '', thinkingDuration)
sentLastThinkingEvent = true
}

if (content) {
sseStreamer.streamTokenEvent(chatId, content)
}
sseStreamer.streamTokenEvent(chatId, content)
}

const messageChunk = typeof chunk === 'string' ? new AIMessageChunk(chunk) : chunk
response = response.concat(messageChunk)
}

if (sseStreamer && !isStructuredOutput && isLastNode && wasThinking && !sentLastThinkingEvent && thinkingStartTime != null) {
thinkingDuration = Math.round((Date.now() - thinkingStartTime) / 1000)
sseStreamer.streamThinkingEvent(chatId, '', thinkingDuration)
sentLastThinkingEvent = true
}
} catch (error) {
console.error('Error during streaming:', error)
throw error
Expand Down
39 changes: 39 additions & 0 deletions packages/components/nodes/agentflow/LLM/LLMStreaming.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,39 @@
import { AIMessageChunk } from '@langchain/core/messages'

const { nodeClass: LLMAgentflow } = require('./LLM')

describe('LLM AgentFlow streaming reasoning chunks', () => {
it('streams reasoning-only chunks as thinking content without empty token events', async () => {
const llmNode = new LLMAgentflow()
const chunk = new AIMessageChunk({
content: '',
additional_kwargs: {
reasoning_content: 'Think before answering.'
}
})
const llmNodeInstance = {
stream: jest.fn(async function* () {
yield chunk
})
}
const sseStreamer = {
streamThinkingEvent: jest.fn(),
streamTokenEvent: jest.fn()
}

const response = await (llmNode as any).handleStreamingResponse(
sseStreamer,
llmNodeInstance,
[],
'chat-id',
new AbortController(),
false,
true
)

expect(sseStreamer.streamThinkingEvent).toHaveBeenCalledWith('chat-id', 'Think before answering.')
expect(sseStreamer.streamThinkingEvent).toHaveBeenCalledWith('chat-id', '', expect.any(Number))
expect(sseStreamer.streamTokenEvent).not.toHaveBeenCalled()
expect(response.additional_kwargs?.reasoning_content).toBe('Think before answering.')
})
})
3 changes: 2 additions & 1 deletion packages/components/nodes/chatmodels/ChatOpenAICustom/ChatOpenAICustom.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -140,7 +140,8 @@ class ChatOpenAICustom_ChatModels implements INode {
modelName,
openAIApiKey,
apiKey: openAIApiKey,
streaming: streaming ?? true
streaming: streaming ?? true,
__includeRawResponse: true
}

if (maxTokens) obj.maxTokens = parseInt(maxTokens, 10)
Expand Down
86 changes: 85 additions & 1 deletion packages/components/src/utils.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,10 @@ import {
convertRequireToImport,
COMMONJS_REQUIRE_REGEX,
IMPORT_EXTRACTION_REGEX,
executeJavaScriptCode
executeJavaScriptCode,
extractStreamingChunkContent
} from './utils'
import { AIMessageChunk } from '@langchain/core/messages'

describe('removeInvalidImageMarkdown', () => {
describe('strips non-http/https image markdown', () => {
Expand Down Expand Up @@ -236,6 +238,88 @@ describe('Import extraction regex (utils.ts line 1596 pattern)', () => {
})
})

describe('extractStreamingChunkContent', () => {
it('separates vLLM/OpenAI-compatible reasoning_content from visible text', () => {
const chunk = new AIMessageChunk({
content: '',
additional_kwargs: {
reasoning_content: 'I should inspect the inputs.'
}
})

expect(extractStreamingChunkContent(chunk)).toEqual({
text: '',
reasoning: 'I should inspect the inputs.'
})
})

it('extracts reasoning deltas from preserved OpenAI-compatible raw responses', () => {
const chunk = new AIMessageChunk({
content: '',
additional_kwargs: {
__raw_response: {
choices: [
{
delta: {
reasoning: 'Raw reasoning.',
reasoning_content: ' Raw content.'
}
}
]
}
}
})

expect(extractStreamingChunkContent(chunk)).toEqual({
text: '',
reasoning: 'Raw reasoning. Raw content.'
})
})

it('ignores empty role and usage chunks', () => {
expect(
extractStreamingChunkContent(
new AIMessageChunk({
content: '',
additional_kwargs: {
__raw_response: {
choices: [
{
delta: {
role: 'assistant'
}
}
],
usage: {
prompt_tokens: 1,
completion_tokens: 0
}
}
}
})
)
).toEqual({
text: '',
reasoning: ''
})
})

it('preserves visible text extraction while collecting reasoning content blocks', () => {
const chunk = new AIMessageChunk({
content: [
{ type: 'reasoning', reasoning: 'First think.' },
{ type: 'text_delta', text: 'Final ' },
{ type: 'text', text: 'answer.' }
] as any
})

expect(extractStreamingChunkContent(chunk)).toEqual({
text: 'Final answer.',
reasoning: 'First think.'
})
})
})

// ---------------------------------------------------------------------------
// NodeVM sandbox — availableDependencies allowlist
// ---------------------------------------------------------------------------
Expand Down
Loading