396 lines
15 KiB
TypeScript
396 lines
15 KiB
TypeScript
import { useState, useRef } from 'react'
|
|
import { streamText, tool, jsonSchema } from 'ai'
|
|
import { createGroq } from '@ai-sdk/groq'
|
|
import { createAnthropic } from '@ai-sdk/anthropic'
|
|
import { type Message, type Conversation, type UserMessage, type AssistantMessage, type SystemMessage } from '../types'
|
|
import { type Model } from '../types/models'
|
|
import { getApiKey } from '../utils/apiKeys'
|
|
import { type Tool } from 'use-mcp/react'
|
|
|
|
// Debug logging for message operations
|
|
const debugMessages = (...args: any[]) => {
|
|
if (typeof window !== 'undefined' && localStorage.getItem('USE_MCP_DEBUG') === 'true') {
|
|
console.log('[Messages]', ...args)
|
|
}
|
|
}
|
|
|
|
// Type guard for messages with content
|
|
const hasContent = (message: Message): message is UserMessage | AssistantMessage | SystemMessage => {
|
|
return 'content' in message
|
|
}
|
|
|
|
interface UseStreamResponseProps {
|
|
conversationId?: number
|
|
setConversations: React.Dispatch<React.SetStateAction<Conversation[]>>
|
|
scrollToBottom: (force?: boolean) => void
|
|
selectedModel: Model
|
|
onApiKeyRequired: (model: Model) => Promise<boolean>
|
|
mcpTools: Tool[]
|
|
}
|
|
|
|
export const useStreamResponse = ({
|
|
conversationId,
|
|
setConversations,
|
|
scrollToBottom,
|
|
selectedModel,
|
|
onApiKeyRequired,
|
|
mcpTools,
|
|
}: UseStreamResponseProps) => {
|
|
const [isLoading, setIsLoading] = useState(false)
|
|
const [streamStarted, setStreamStarted] = useState(false)
|
|
const [controller, setController] = useState(new AbortController())
|
|
const aiResponseRef = useRef<string>('')
|
|
|
|
// Convert MCP tools to the format expected by the 'ai' package
|
|
const convertMcpToolsToAiTools = (mcpTools: Tool[]) => {
|
|
const aiTools: Record<string, any> = {}
|
|
|
|
mcpTools.forEach((mcpTool) => {
|
|
const schema = mcpTool.inputSchema || { type: 'object', properties: {} }
|
|
|
|
aiTools[mcpTool.name] = tool({
|
|
description: mcpTool.description,
|
|
parameters: jsonSchema(schema as any),
|
|
execute: async (args: unknown) => {
|
|
try {
|
|
// Call the MCP tool using the callTool method attached to the tool
|
|
if ('callTool' in mcpTool && typeof mcpTool.callTool === 'function') {
|
|
const result = await mcpTool.callTool(args as Record<string, unknown>)
|
|
return result
|
|
} else {
|
|
throw new Error(`Tool ${mcpTool.name} does not have a callable implementation`)
|
|
}
|
|
} catch (error) {
|
|
console.error(`Error calling MCP tool ${mcpTool.name}:`, error)
|
|
throw error
|
|
}
|
|
},
|
|
})
|
|
})
|
|
return aiTools
|
|
}
|
|
|
|
// Check if a model supports reasoning
|
|
const supportsReasoning = (model: Model) => {
|
|
const reasoningModels = [
|
|
'qwen/qwen3-32b',
|
|
'qwen-qwq-32b',
|
|
'deepseek-r1-distill-llama-70b',
|
|
'claude-4-opus-20250514',
|
|
'claude-4-sonnet-20250514',
|
|
'claude-3-7-sonnet-20250219',
|
|
]
|
|
return reasoningModels.includes(model.modelId)
|
|
}
|
|
|
|
const getModelInstance = (model: Model, apiKey: string) => {
|
|
let baseModel
|
|
|
|
switch (model.provider.id) {
|
|
case 'groq': {
|
|
const groqProvider = createGroq({ apiKey })
|
|
baseModel = groqProvider(model.modelId)
|
|
break
|
|
}
|
|
case 'anthropic': {
|
|
const anthropicProvider = createAnthropic({
|
|
apiKey,
|
|
headers: {
|
|
'anthropic-dangerous-direct-browser-access': 'true',
|
|
},
|
|
})
|
|
baseModel = anthropicProvider(model.modelId)
|
|
break
|
|
}
|
|
default:
|
|
throw new Error(`Unsupported provider: ${model.provider.id}`)
|
|
}
|
|
|
|
// For reasoning-capable models, we'll handle reasoning events manually
|
|
// since the middleware expects <think> tags but the API sends reasoning fields directly
|
|
|
|
return baseModel
|
|
}
|
|
|
|
const streamResponse = async (messages: Message[]) => {
|
|
// Check if API key is available
|
|
let apiKey = getApiKey(selectedModel.provider.id)
|
|
if (!apiKey) {
|
|
const keyProvided = await onApiKeyRequired(selectedModel)
|
|
if (!keyProvided) {
|
|
return // User cancelled
|
|
}
|
|
apiKey = getApiKey(selectedModel.provider.id)
|
|
if (!apiKey) {
|
|
throw new Error('No API key provided')
|
|
}
|
|
}
|
|
|
|
try {
|
|
const modelInstance = getModelInstance(selectedModel, apiKey)
|
|
|
|
setStreamStarted(true)
|
|
|
|
// Convert MCP tools to AI package format
|
|
const aiTools = convertMcpToolsToAiTools(mcpTools)
|
|
|
|
// Prepare stream options
|
|
const streamOptions = {
|
|
model: modelInstance,
|
|
messages: messages
|
|
.filter((msg) => ['user', 'assistant', 'system'].includes(msg.role) && hasContent(msg))
|
|
.map((msg) => {
|
|
const contentMsg = msg as UserMessage | AssistantMessage | SystemMessage
|
|
return {
|
|
role: contentMsg.role,
|
|
content: contentMsg.content,
|
|
}
|
|
}),
|
|
tools: Object.keys(aiTools).length > 0 ? aiTools : undefined,
|
|
maxSteps: 5, // Allow up to 5 steps for tool calling
|
|
abortSignal: controller.signal,
|
|
}
|
|
|
|
const result = await streamText(streamOptions)
|
|
|
|
// Use fullStream to get all events including tool calls, results, and text
|
|
for await (const event of result.fullStream) {
|
|
console.log({ 'result.fullStream.event': event })
|
|
setConversations((prev) => {
|
|
const conv = prev.find((c) => c.id === conversationId)
|
|
if (!conv) {
|
|
console.error(`Missing conversation for ID ${conversationId}! Ignoring ${JSON.stringify(event)}`)
|
|
return prev
|
|
}
|
|
const lastMessage = conv.messages?.at(-1)
|
|
console.log({ setConversations: event, conv: JSON.parse(JSON.stringify(conv)) })
|
|
let updatedMessage: Message | undefined
|
|
let newMessage: Message | undefined
|
|
try {
|
|
if (event.type === 'reasoning') {
|
|
if (lastMessage?.role === 'assistant' && lastMessage.type === 'reasoning') {
|
|
// We have an existing reasoning block!
|
|
updatedMessage = {
|
|
...lastMessage,
|
|
content: lastMessage.content + event.textDelta.trim(),
|
|
}
|
|
} else {
|
|
// We need a new reasoning block!
|
|
const currentReasoningStartTime = Date.now()
|
|
|
|
// Create a new assistant message for reasoning (even if we already have one)
|
|
// This handles multiple reasoning sessions in the same conversation turn
|
|
newMessage = {
|
|
role: 'assistant',
|
|
type: 'reasoning',
|
|
content: event.textDelta.trim(),
|
|
isReasoningStreaming: true,
|
|
reasoningStartTime: currentReasoningStartTime!,
|
|
}
|
|
}
|
|
} else {
|
|
// We're not thinking, so maybe close the last message?
|
|
if (lastMessage?.role === 'assistant' && lastMessage.type === 'reasoning') {
|
|
updatedMessage = {
|
|
...lastMessage,
|
|
reasoningEndTime: Date.now(),
|
|
isReasoningStreaming: false,
|
|
}
|
|
}
|
|
|
|
if (event.type === 'tool-call') {
|
|
if (event.toolName) {
|
|
newMessage = {
|
|
role: 'tool-call',
|
|
toolName: event.toolName,
|
|
toolArgs: event.args || {},
|
|
callId: event.toolCallId || 'unknown',
|
|
}
|
|
scrollToBottom(true)
|
|
} else {
|
|
console.warn('Tool call event missing toolName:', event)
|
|
}
|
|
} else if (event.type === 'tool-result') {
|
|
console.info(`TODO TOOL RESULT`, event)
|
|
// if (event.toolName && event.toolCallId) {
|
|
// setConversations((prev) => {
|
|
// const updated = [...prev]
|
|
// const conv = updated.find((c) => c.id === conversationId)
|
|
// if (conv) {
|
|
// conv.messages.push({
|
|
// role: 'tool-result',
|
|
// toolName: event.toolName,
|
|
// toolArgs: event.args || {},
|
|
// toolResult: event.result,
|
|
// callId: event.toolCallId,
|
|
// })
|
|
// debugMessages('Added new tool-result message:', event.toolName, 'callId:', event.toolCallId)
|
|
// debugMessages('Current messages:', JSON.stringify(conv.messages))
|
|
// }
|
|
// return updated
|
|
// })
|
|
// scrollToBottom(true)
|
|
// } else {
|
|
// console.warn('Tool result event missing toolName or toolCallId:', event)
|
|
// }
|
|
} else if (event.type === 'text-delta') {
|
|
console.info(`TODO TEXT DELTA`, event)
|
|
|
|
// // Check if we have an existing reasoning message that needs to be finalized
|
|
// let needsNewAssistantMessage = false
|
|
// if (currentReasoningStartTime && !currentReasoningEndTime) {
|
|
// currentReasoningEndTime = Date.now()
|
|
//
|
|
// // Update reasoning end time and stop streaming
|
|
// 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.isReasoningStreaming = false
|
|
// assistantMessage.reasoningEndTime = currentReasoningEndTime
|
|
// }
|
|
// }
|
|
// return updated
|
|
// })
|
|
//
|
|
// // Reset reasoning state after reasoning ends
|
|
// currentReasoningStartTime = undefined
|
|
// currentReasoningEndTime = undefined
|
|
// // Force creation of new assistant message for content
|
|
// needsNewAssistantMessage = true
|
|
// assistantMessageCreated = false
|
|
// }
|
|
//
|
|
// // Create assistant message on first text delta if not already created OR if we just finished reasoning
|
|
// if (!assistantMessageCreated || needsNewAssistantMessage) {
|
|
// setConversations((prev) => {
|
|
// const updated = [...prev]
|
|
// const conv = updated.find((c) => c.id === conversationId)
|
|
// if (conv) {
|
|
// conv.messages.push({ role: 'assistant', content: '' })
|
|
// assistantMessageIndex = conv.messages.length - 1
|
|
// debugMessages(
|
|
// 'Added new content assistant message at index',
|
|
// assistantMessageIndex,
|
|
// '(after reasoning:',
|
|
// needsNewAssistantMessage,
|
|
// ')'
|
|
// )
|
|
// debugMessages('Current messages:', JSON.stringify(conv.messages))
|
|
// }
|
|
// return updated
|
|
// })
|
|
// assistantMessageCreated = true
|
|
// }
|
|
//
|
|
// aiResponse += event.textDelta
|
|
// aiResponseRef.current = aiResponse
|
|
//
|
|
// //custom extraction of <chat-title> tag
|
|
// const titleMatch = aiResponse.match(/<chat-title>(.*?)<\/chat-title>/)
|
|
// if (titleMatch) {
|
|
// const title = titleMatch[1].trim()
|
|
// setConversations((prev) => {
|
|
// const updated = [...prev]
|
|
// const conv = updated.find((c) => c.id === conversationId)
|
|
// if (conv) {
|
|
// conv.title = title
|
|
// }
|
|
// return updated
|
|
// })
|
|
// aiResponse = aiResponse.replace(/<chat-title>.*?<\/chat-title>/, '').trim()
|
|
// }
|
|
//
|
|
// setConversations((prev) => {
|
|
// const updated = [...prev]
|
|
// const conv = updated.find((c) => c.id === conversationId)
|
|
// if (conv) {
|
|
// // Update the specific assistant message we created, not just the last message
|
|
// const assistantMessage = conv.messages[assistantMessageIndex]
|
|
// if (assistantMessage && hasContent(assistantMessage) && assistantMessage.role === 'assistant') {
|
|
// assistantMessage.content = aiResponse
|
|
// }
|
|
// } else {
|
|
// }
|
|
// return updated
|
|
// })
|
|
//
|
|
// scrollToBottom(true)
|
|
}
|
|
}
|
|
} catch (e) {
|
|
console.error('Error in full stream processing:', e)
|
|
}
|
|
let updated = prev
|
|
if (updatedMessage || newMessage) {
|
|
updated = prev.map((c) => {
|
|
if (c.id !== conversationId) return c
|
|
const messages = c.messages
|
|
|
|
if (updatedMessage) {
|
|
messages.splice(-1, 1, updatedMessage)
|
|
}
|
|
|
|
if (newMessage) {
|
|
messages.push(newMessage)
|
|
}
|
|
|
|
return { ...c, messages }
|
|
})
|
|
}
|
|
console.log({ updatedMessage, newMessage })
|
|
return updated
|
|
})
|
|
}
|
|
|
|
// Show final conversation state
|
|
setConversations((prev) => {
|
|
const conv = prev.find((c) => c.id === conversationId)
|
|
if (conv) {
|
|
debugMessages('Final conversation state:', JSON.stringify(conv.messages))
|
|
}
|
|
return prev
|
|
})
|
|
//
|
|
// // Final cleanup - ensure reasoning streaming is stopped
|
|
// if (assistantMessageCreated && currentReasoningStartTime) {
|
|
// if (!currentReasoningEndTime) {
|
|
// currentReasoningEndTime = Date.now()
|
|
// }
|
|
//
|
|
// 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.isReasoningStreaming = false
|
|
// assistantMessage.reasoningEndTime = currentReasoningEndTime
|
|
// }
|
|
// }
|
|
// return updated
|
|
// })
|
|
// }
|
|
} catch (error: unknown) {
|
|
if (controller.signal.aborted) {
|
|
} else {
|
|
console.error('Error generating response:', error)
|
|
}
|
|
} finally {
|
|
setStreamStarted(false)
|
|
setController(new AbortController())
|
|
debugMessages('Stream response function completed')
|
|
}
|
|
}
|
|
|
|
return {
|
|
isLoading,
|
|
setIsLoading,
|
|
streamStarted,
|
|
controller,
|
|
streamResponse,
|
|
aiResponseRef,
|
|
}
|
|
}
|