fix: Use extractReasoningMiddleware to properly parse thinking blocks

- Import extractReasoningMiddleware and wrapLanguageModel from AI SDK
- Wrap reasoning-capable models with extractReasoningMiddleware
- Use 'think' as tagName to parse <think> XML blocks from models
- Enable startWithReasoning option to handle models missing reasoning tags
- Extract reasoning from result.reasoning after stream completes
- Remove manual reasoning event handling and provider-specific config

This fixes the issue where thinking blocks appeared as raw XML tags
instead of parsed reasoning content. The middleware automatically
extracts reasoning from <think> tags and exposes it via result.reasoning.
This commit is contained in:
Glen Maddern 2025-06-24 16:47:04 +10:00
commit 25857bc354

View file

@ -1,5 +1,5 @@
import { useState, useRef } from 'react' import { useState, useRef } from 'react'
import { streamText, tool, jsonSchema } from 'ai' import { streamText, tool, jsonSchema, wrapLanguageModel, extractReasoningMiddleware } from 'ai'
import { createGroq } from '@ai-sdk/groq' import { createGroq } from '@ai-sdk/groq'
import { createAnthropic } from '@ai-sdk/anthropic' import { createAnthropic } from '@ai-sdk/anthropic'
import { type Message, type Conversation, type UserMessage, type AssistantMessage, type SystemMessage } from '../types' import { type Message, type Conversation, type UserMessage, type AssistantMessage, type SystemMessage } from '../types'
@ -94,10 +94,13 @@ export const useStreamResponse = ({
} }
const getModelInstance = (model: Model, apiKey: string) => { const getModelInstance = (model: Model, apiKey: string) => {
let baseModel
switch (model.provider.id) { switch (model.provider.id) {
case 'groq': { case 'groq': {
const groqProvider = createGroq({ apiKey }) const groqProvider = createGroq({ apiKey })
return groqProvider(model.modelId) baseModel = groqProvider(model.modelId)
break
} }
case 'anthropic': { case 'anthropic': {
const anthropicProvider = createAnthropic({ const anthropicProvider = createAnthropic({
@ -106,16 +109,30 @@ export const useStreamResponse = ({
'anthropic-dangerous-direct-browser-access': 'true', 'anthropic-dangerous-direct-browser-access': 'true',
}, },
}) })
return anthropicProvider(model.modelId) baseModel = anthropicProvider(model.modelId)
break
} }
default: default:
throw new Error(`Unsupported provider: ${model.provider.id}`) throw new Error(`Unsupported provider: ${model.provider.id}`)
} }
// Wrap reasoning-capable models with extractReasoningMiddleware
if (supportsReasoning(model)) {
debugLog(`[useStreamResponse] Wrapping model with extractReasoningMiddleware for ${model.modelId}`)
return wrapLanguageModel({
model: baseModel,
middleware: extractReasoningMiddleware({
tagName: 'think',
startWithReasoning: true // This ensures reasoning tags are prepended if missing
}),
})
}
return baseModel
} }
const streamResponse = async (messages: Message[]) => { const streamResponse = async (messages: Message[]) => {
let aiResponse = '' let aiResponse = ''
let reasoning = ''
let assistantMessageIndex = -1 // Track the index of our assistant message let assistantMessageIndex = -1 // Track the index of our assistant message
let assistantMessageCreated = false // Track if we've created the assistant message yet let assistantMessageCreated = false // Track if we've created the assistant message yet
@ -169,27 +186,6 @@ export const useStreamResponse = ({
abortSignal: controller.signal, abortSignal: controller.signal,
} }
// Add reasoning parameters for supported models
if (supportsReasoning(selectedModel)) {
debugLog(`[useStreamResponse] Enabling reasoning for model: ${selectedModel.modelId}`)
if (selectedModel.provider.id === 'groq') {
// For Groq models, use reasoning_format parameter
streamOptions.experimental_providerOptions = {
groq: {
reasoning_format: 'parsed'
}
}
} else if (selectedModel.provider.id === 'anthropic') {
// For Anthropic models, enable thinking
streamOptions.experimental_providerOptions = {
anthropic: {
thinking: { type: 'enabled', budgetTokens: 12000 }
}
}
}
}
const result = await streamText(streamOptions) const result = await streamText(streamOptions)
debugLog(`[useStreamResponse] streamText result obtained, starting to process fullStream`) debugLog(`[useStreamResponse] streamText result obtained, starting to process fullStream`)
@ -201,10 +197,7 @@ export const useStreamResponse = ({
try { try {
debugLog(`[useStreamResponse] Full stream event:`, event.type, event) debugLog(`[useStreamResponse] Full stream event:`, event.type, event)
if (event.type === 'reasoning') { if (event.type === 'tool-call') {
debugLog(`[useStreamResponse] Reasoning event:`, event)
reasoning += (event as any).textDelta || ''
} else if (event.type === 'tool-call') {
const eventKey = `tool-call-${event.toolCallId}` const eventKey = `tool-call-${event.toolCallId}`
if (processedEvents.has(eventKey)) { if (processedEvents.has(eventKey)) {
debugLog(`[useStreamResponse] Skipping duplicate tool call:`, eventKey) debugLog(`[useStreamResponse] Skipping duplicate tool call:`, eventKey)
@ -247,7 +240,7 @@ export const useStreamResponse = ({
const updated = [...prev] const updated = [...prev]
const conv = updated.find((c) => c.id === conversationId) const conv = updated.find((c) => c.id === conversationId)
if (conv) { if (conv) {
conv.messages.push({ role: 'assistant', content: '', reasoning: reasoning || undefined }) conv.messages.push({ role: 'assistant', content: '' })
assistantMessageIndex = conv.messages.length - 1 assistantMessageIndex = conv.messages.length - 1
debugLog(`[useStreamResponse] Created assistant message at index ${assistantMessageIndex}`) debugLog(`[useStreamResponse] Created assistant message at index ${assistantMessageIndex}`)
} }
@ -284,10 +277,6 @@ export const useStreamResponse = ({
const assistantMessage = conv.messages[assistantMessageIndex] const assistantMessage = conv.messages[assistantMessageIndex]
if (assistantMessage && hasContent(assistantMessage) && assistantMessage.role === 'assistant') { if (assistantMessage && hasContent(assistantMessage) && assistantMessage.role === 'assistant') {
assistantMessage.content = aiResponse assistantMessage.content = aiResponse
// Update reasoning if available
if (reasoning) {
assistantMessage.reasoning = reasoning
}
debugLog(`[useStreamResponse] Updated assistant message at index ${assistantMessageIndex}`) debugLog(`[useStreamResponse] Updated assistant message at index ${assistantMessageIndex}`)
} else { } else {
debugLog(`[useStreamResponse] Could not find assistant message at index ${assistantMessageIndex}`, assistantMessage) debugLog(`[useStreamResponse] Could not find assistant message at index ${assistantMessageIndex}`, assistantMessage)
@ -306,6 +295,24 @@ export const useStreamResponse = ({
} }
debugLog(`[useStreamResponse] Finished processing full stream. Final response length: ${aiResponse.length}`) debugLog(`[useStreamResponse] Finished processing full stream. Final response length: ${aiResponse.length}`)
// Extract reasoning from the result after streaming completes
const finalReasoning = await result.reasoning
if (finalReasoning && assistantMessageCreated) {
debugLog(`[useStreamResponse] Extracted reasoning:`, finalReasoning)
setConversations((prev) => {
const updated = [...prev]
const conv = updated.find((c) => c.id === conversationId)
if (conv && assistantMessageIndex >= 0) {
const assistantMessage = conv.messages[assistantMessageIndex]
if (assistantMessage && hasContent(assistantMessage) && assistantMessage.role === 'assistant') {
assistantMessage.reasoning = finalReasoning
debugLog(`[useStreamResponse] Added reasoning to assistant message at index ${assistantMessageIndex}`)
}
}
return updated
})
}
} catch (error: unknown) { } catch (error: unknown) {
if (controller.signal.aborted) { if (controller.signal.aborted) {
debugLog('[useStreamResponse] Stream aborted') debugLog('[useStreamResponse] Stream aborted')