UNPKG

@lobehub/chat

Version:

Lobe Chat - an open-source, high-performance chatbot framework that supports speech synthesis, multimodal, and extensible Function Call plugin system. Supports one-click free deployment of your private ChatGPT/LLM web application.

835 lines (712 loc) 28.4 kB
/* eslint-disable sort-keys-fix/sort-keys-fix, typescript-sort-keys/interface */ // Disable the auto sort key eslint rule to make the code more logic and readable import { produce } from 'immer'; import { template } from 'lodash-es'; import { StateCreator } from 'zustand/vanilla'; import { LOADING_FLAT, MESSAGE_CANCEL_FLAT } from '@/const/message'; import { TraceEventType, TraceNameMap } from '@/const/trace'; import { isServerMode } from '@/const/version'; import { knowledgeBaseQAPrompts } from '@/prompts/knowledgeBaseQA'; import { chatService } from '@/services/chat'; import { messageService } from '@/services/message'; import { useAgentStore } from '@/store/agent'; import { agentChatConfigSelectors, agentSelectors } from '@/store/agent/selectors'; import { getAgentStoreState } from '@/store/agent/store'; import { aiModelSelectors, aiProviderSelectors } from '@/store/aiInfra'; import { getAiInfraStoreState } from '@/store/aiInfra/store'; import { chatHelpers } from '@/store/chat/helpers'; import { ChatStore } from '@/store/chat/store'; import { messageMapKey } from '@/store/chat/utils/messageMapKey'; import { getFileStoreState } from '@/store/file/store'; import { useSessionStore } from '@/store/session'; import { WebBrowsingManifest } from '@/tools/web-browsing'; import { ChatMessage, CreateMessageParams, SendMessageParams } from '@/types/message'; import { ChatImageItem } from '@/types/message/image'; import { MessageSemanticSearchChunk } from '@/types/rag'; import { setNamespace } from '@/utils/storeDebug'; import { chatSelectors, topicSelectors } from '../../../selectors'; const n = setNamespace('ai'); interface ProcessMessageParams { traceId?: string; isWelcomeQuestion?: boolean; inSearchWorkflow?: boolean; /** * the RAG query content, should be embedding and used in the semantic search */ ragQuery?: string; threadId?: string; inPortalThread?: boolean; } export interface AIGenerateAction { /** * Sends a new message to the AI chat system */ sendMessage: (params: SendMessageParams) => Promise<void>; /** * Regenerates a specific message in the chat */ regenerateMessage: (id: string) => Promise<void>; /** * Deletes an existing message and generates a new one in its place */ delAndRegenerateMessage: (id: string) => Promise<void>; /** * Interrupts the ongoing ai message generation process */ stopGenerateMessage: () => void; // ========= ↓ Internal Method ↓ ========== // // ========================================== // // ========================================== // /** * Executes the core processing logic for AI messages * including preprocessing and postprocessing steps */ internal_coreProcessMessage: ( messages: ChatMessage[], parentId: string, params?: ProcessMessageParams, ) => Promise<void>; /** * Retrieves an AI-generated chat message from the backend service */ internal_fetchAIChatMessage: (input: { messages: ChatMessage[]; messageId: string; params?: ProcessMessageParams; model: string; provider: string; }) => Promise<{ isFunctionCall: boolean; traceId?: string; }>; /** * Resends a specific message, optionally using a trace ID for tracking */ internal_resendMessage: ( id: string, params?: { traceId?: string; messages?: ChatMessage[]; threadId?: string; inPortalThread?: boolean; }, ) => Promise<void>; /** * Toggles the loading state for AI message generation, managing the UI feedback */ internal_toggleChatLoading: ( loading: boolean, id?: string, action?: string, ) => AbortController | undefined; /** * Controls the streaming state of tool calling processes, updating the UI accordingly */ internal_toggleToolCallingStreaming: (id: string, streaming: boolean[] | undefined) => void; /** * Toggles the loading state for AI message reasoning, managing the UI feedback */ internal_toggleChatReasoning: ( loading: boolean, id?: string, action?: string, ) => AbortController | undefined; internal_toggleSearchWorkflow: (loading: boolean, id?: string) => void; } export const generateAIChat: StateCreator< ChatStore, [['zustand/devtools', never]], [], AIGenerateAction > = (set, get) => ({ delAndRegenerateMessage: async (id) => { const traceId = chatSelectors.getTraceIdByMessageId(id)(get()); get().internal_resendMessage(id, { traceId }); get().deleteMessage(id); // trace the delete and regenerate message get().internal_traceMessage(id, { eventType: TraceEventType.DeleteAndRegenerateMessage }); }, regenerateMessage: async (id) => { const traceId = chatSelectors.getTraceIdByMessageId(id)(get()); await get().internal_resendMessage(id, { traceId }); // trace the delete and regenerate message get().internal_traceMessage(id, { eventType: TraceEventType.RegenerateMessage }); }, sendMessage: async ({ message, files, onlyAddUserMessage, isWelcomeQuestion }) => { const { internal_coreProcessMessage, activeTopicId, activeId, activeThreadId } = get(); if (!activeId) return; const fileIdList = files?.map((f) => f.id); const hasFile = !!fileIdList && fileIdList.length > 0; // if message is empty or no files, then stop if (!message && !hasFile) return; set({ isCreatingMessage: true }, false, n('creatingMessage/start')); const newMessage: CreateMessageParams = { content: message, // if message has attached with files, then add files to message and the agent files: fileIdList, role: 'user', sessionId: activeId, // if there is activeTopicId,then add topicId to message topicId: activeTopicId, threadId: activeThreadId, }; const agentConfig = agentChatConfigSelectors.currentChatConfig(getAgentStoreState()); let tempMessageId: string | undefined = undefined; let newTopicId: string | undefined = undefined; // it should be the default topic, then // if autoCreateTopic is enabled, check to whether we need to create a topic if (!onlyAddUserMessage && !activeTopicId && agentConfig.enableAutoCreateTopic) { // check activeTopic and then auto create topic const chats = chatSelectors.activeBaseChats(get()); // we will add two messages (user and assistant), so the finial length should +2 const featureLength = chats.length + 2; // if there is no activeTopicId and the feature length is greater than the threshold // then create a new topic and active it if (!activeTopicId && featureLength >= agentConfig.autoCreateTopicThreshold) { // we need to create a temp message for optimistic update tempMessageId = get().internal_createTmpMessage(newMessage); get().internal_toggleMessageLoading(true, tempMessageId); const topicId = await get().createTopic(); if (topicId) { newTopicId = topicId; newMessage.topicId = topicId; // we need to copy the messages to the new topic or the message will disappear const mapKey = chatSelectors.currentChatKey(get()); const newMaps = { ...get().messagesMap, [messageMapKey(activeId, topicId)]: get().messagesMap[mapKey], }; set({ messagesMap: newMaps }, false, n('moveMessagesToNewTopic')); // make the topic loading get().internal_updateTopicLoading(topicId, true); } } } // update assistant update to make it rerank useSessionStore.getState().triggerSessionUpdate(get().activeId); const id = await get().internal_createMessage(newMessage, { tempMessageId, skipRefresh: !onlyAddUserMessage && newMessage.fileList?.length === 0, }); if (!id) { set({ isCreatingMessage: false }, false, n('creatingMessage/start')); if (!!newTopicId) get().internal_updateTopicLoading(newTopicId, false); return; } if (tempMessageId) get().internal_toggleMessageLoading(false, tempMessageId); // switch to the new topic if create the new topic if (!!newTopicId) { await get().switchTopic(newTopicId, true); await get().internal_fetchMessages(); // delete previous messages // remove the temp message map const newMaps = { ...get().messagesMap, [messageMapKey(activeId, null)]: [] }; set({ messagesMap: newMaps }, false, 'internal_copyMessages'); } // if only add user message, then stop if (onlyAddUserMessage) { set({ isCreatingMessage: false }, false, 'creatingMessage/start'); return; } // Get the current messages to generate AI response const messages = chatSelectors.activeBaseChats(get()); const userFiles = chatSelectors.currentUserFiles(get()).map((f) => f.id); await internal_coreProcessMessage(messages, id, { isWelcomeQuestion, ragQuery: get().internal_shouldUseRAG() ? message : undefined, threadId: activeThreadId, }); set({ isCreatingMessage: false }, false, n('creatingMessage/stop')); const summaryTitle = async () => { // if autoCreateTopic is false, then stop if (!agentConfig.enableAutoCreateTopic) return; // check activeTopic and then auto update topic title if (newTopicId) { const chats = chatSelectors.getBaseChatsByKey(messageMapKey(activeId, newTopicId))(get()); await get().summaryTopicTitle(newTopicId, chats); return; } if (!activeTopicId) return; const topic = topicSelectors.getTopicById(activeTopicId)(get()); if (topic && !topic.title) { const chats = chatSelectors.getBaseChatsByKey(messageMapKey(activeId, topic.id))(get()); await get().summaryTopicTitle(topic.id, chats); } }; // if there is relative files, then add files to agent // only available in server mode const addFilesToAgent = async () => { if (userFiles.length === 0 || !isServerMode) return; await useAgentStore.getState().addFilesToAgent(userFiles, false); }; await Promise.all([summaryTitle(), addFilesToAgent()]); }, stopGenerateMessage: () => { const { chatLoadingIdsAbortController, internal_toggleChatLoading } = get(); if (!chatLoadingIdsAbortController) return; chatLoadingIdsAbortController.abort(MESSAGE_CANCEL_FLAT); internal_toggleChatLoading(false, undefined, n('stopGenerateMessage') as string); }, // the internal process method of the AI message internal_coreProcessMessage: async (originalMessages, userMessageId, params) => { const { internal_fetchAIChatMessage, triggerToolCalls, refreshMessages, activeTopicId } = get(); // create a new array to avoid the original messages array change const messages = [...originalMessages]; const agentStoreState = getAgentStoreState(); const { model, provider, chatConfig } = agentSelectors.currentAgentConfig(agentStoreState); let fileChunks: MessageSemanticSearchChunk[] | undefined; let ragQueryId; // go into RAG flow if there is ragQuery flag if (params?.ragQuery) { // 1. get the relative chunks from semantic search const { chunks, queryId, rewriteQuery } = await get().internal_retrieveChunks( userMessageId, params?.ragQuery, // should skip the last content messages.map((m) => m.content).slice(0, messages.length - 1), ); ragQueryId = queryId; const lastMsg = messages.pop() as ChatMessage; // 2. build the retrieve context messages const knowledgeBaseQAContext = knowledgeBaseQAPrompts({ chunks, userQuery: lastMsg.content, rewriteQuery, knowledge: agentSelectors.currentEnabledKnowledge(agentStoreState), }); // 3. add the retrieve context messages to the messages history messages.push({ ...lastMsg, content: (lastMsg.content + '\n\n' + knowledgeBaseQAContext).trim(), }); fileChunks = chunks.map((c) => ({ id: c.id, similarity: c.similarity })); } // 2. Add an empty message to place the AI response const assistantMessage: CreateMessageParams = { role: 'assistant', content: LOADING_FLAT, fromModel: model, fromProvider: provider, parentId: userMessageId, sessionId: get().activeId, topicId: activeTopicId, // if there is activeTopicId,then add it to topicId threadId: params?.threadId, fileChunks, ragQueryId, }; const assistantId = await get().internal_createMessage(assistantMessage); if (!assistantId) return; // 3. place a search with the search working model if this model is not support tool use const aiInfraStoreState = getAiInfraStoreState(); const isModelSupportToolUse = aiModelSelectors.isModelSupportToolUse( model, provider!, )(aiInfraStoreState); const isProviderHasBuiltinSearch = aiProviderSelectors.isProviderHasBuiltinSearch(provider!)( aiInfraStoreState, ); const isModelHasBuiltinSearch = aiModelSelectors.isModelHasBuiltinSearch( model, provider!, )(aiInfraStoreState); const useModelBuiltinSearch = agentChatConfigSelectors.useModelBuiltinSearch(agentStoreState); const useModelSearch = (isProviderHasBuiltinSearch || isModelHasBuiltinSearch) && useModelBuiltinSearch; const isAgentEnableSearch = agentChatConfigSelectors.isAgentEnableSearch(agentStoreState); if (isAgentEnableSearch && !useModelSearch && !isModelSupportToolUse) { const { model, provider } = agentChatConfigSelectors.searchFCModel(agentStoreState); let isToolsCalling = false; let isError = false; const abortController = get().internal_toggleChatLoading( true, assistantId, n('generateMessage(start)', { messageId: assistantId, messages }) as string, ); get().internal_toggleSearchWorkflow(true, assistantId); await chatService.fetchPresetTaskResult({ params: { messages, model, provider, plugins: [WebBrowsingManifest.identifier] }, onFinish: async (_, { toolCalls, usage }) => { if (toolCalls && toolCalls.length > 0) { get().internal_toggleToolCallingStreaming(assistantId, undefined); // update tools calling await get().internal_updateMessageContent(assistantId, '', { toolCalls, metadata: usage, model, provider, }); } }, trace: { traceId: params?.traceId, sessionId: get().activeId, topicId: get().activeTopicId, traceName: TraceNameMap.SearchIntentRecognition, }, abortController, onMessageHandle: async (chunk) => { if (chunk.type === 'tool_calls') { get().internal_toggleSearchWorkflow(false, assistantId); get().internal_toggleToolCallingStreaming(assistantId, chunk.isAnimationActives); get().internal_dispatchMessage({ id: assistantId, type: 'updateMessage', value: { tools: get().internal_transformToolCalls(chunk.tool_calls) }, }); isToolsCalling = true; } if (chunk.type === 'text') { abortController!.abort('not fc'); } }, onErrorHandle: async (error) => { isError = true; await messageService.updateMessageError(assistantId, error); await refreshMessages(); }, }); get().internal_toggleChatLoading( false, assistantId, n('generateMessage(start)', { messageId: assistantId, messages }) as string, ); get().internal_toggleSearchWorkflow(false, assistantId); // if there is error, then stop if (isError) return; // if it's the function call message, trigger the function method if (isToolsCalling) { await refreshMessages(); await triggerToolCalls(assistantId, { threadId: params?.threadId, inPortalThread: params?.inPortalThread, }); // then story the workflow return; } } // 4. fetch the AI response const { isFunctionCall } = await internal_fetchAIChatMessage({ messages, messageId: assistantId, params, model, provider: provider!, }); // 5. if it's the function call message, trigger the function method if (isFunctionCall) { await refreshMessages(); await triggerToolCalls(assistantId, { threadId: params?.threadId, inPortalThread: params?.inPortalThread, }); } // 6. summary history if context messages is larger than historyCount const historyCount = agentChatConfigSelectors.historyCount(agentStoreState); if ( agentChatConfigSelectors.enableHistoryCount(agentStoreState) && chatConfig.enableCompressHistory && originalMessages.length > historyCount ) { // after generation: [u1,a1,u2,a2,u3,a3] // but the `originalMessages` is still: [u1,a1,u2,a2,u3] // So if historyCount=2, we need to summary [u1,a1,u2,a2] // because user find UI is [u1,a1,u2,a2 | u3,a3] const historyMessages = originalMessages.slice(0, -historyCount + 1); await get().internal_summaryHistory(historyMessages); } }, internal_fetchAIChatMessage: async ({ messages, messageId, params, provider, model }) => { const { internal_toggleChatLoading, refreshMessages, internal_updateMessageContent, internal_dispatchMessage, internal_toggleToolCallingStreaming, internal_toggleChatReasoning, } = get(); const abortController = internal_toggleChatLoading( true, messageId, n('generateMessage(start)', { messageId, messages }) as string, ); const agentConfig = agentSelectors.currentAgentConfig(getAgentStoreState()); const chatConfig = agentChatConfigSelectors.currentChatConfig(getAgentStoreState()); const compiler = template(chatConfig.inputTemplate, { interpolate: /{{\s*(text)\s*}}/g }); // ================================== // // messages uniformly preprocess // // ================================== // // 1. slice messages with config const historyCount = agentChatConfigSelectors.historyCount(getAgentStoreState()); const enableHistoryCount = agentChatConfigSelectors.enableHistoryCount(getAgentStoreState()); let preprocessMsgs = chatHelpers.getSlicedMessages(messages, { includeNewUserMessage: true, enableHistoryCount, historyCount, }); // 2. replace inputMessage template preprocessMsgs = !chatConfig.inputTemplate ? preprocessMsgs : preprocessMsgs.map((m) => { if (m.role === 'user') { try { return { ...m, content: compiler({ text: m.content }) }; } catch (error) { console.error(error); return m; } } return m; }); // 3. add systemRole if (agentConfig.systemRole) { preprocessMsgs.unshift({ content: agentConfig.systemRole, role: 'system' } as ChatMessage); } // 4. handle max_tokens agentConfig.params.max_tokens = chatConfig.enableMaxTokens ? agentConfig.params.max_tokens : undefined; // 5. handle reasoning_effort agentConfig.params.reasoning_effort = chatConfig.enableReasoningEffort ? agentConfig.params.reasoning_effort : undefined; let isFunctionCall = false; let msgTraceId: string | undefined; let output = ''; let thinking = ''; let thinkingStartAt: number; let duration: number; // to upload image const uploadTasks: Map<string, Promise<{ id?: string; url?: string }>> = new Map(); const historySummary = chatConfig.enableCompressHistory ? topicSelectors.currentActiveTopicSummary(get()) : undefined; await chatService.createAssistantMessageStream({ abortController, params: { messages: preprocessMsgs, model, provider, ...agentConfig.params, plugins: agentConfig.plugins, }, historySummary: historySummary?.content, trace: { traceId: params?.traceId, sessionId: get().activeId, topicId: get().activeTopicId, traceName: TraceNameMap.Conversation, }, isWelcomeQuestion: params?.isWelcomeQuestion, onErrorHandle: async (error) => { await messageService.updateMessageError(messageId, error); await refreshMessages(); }, onFinish: async ( content, { traceId, observationId, toolCalls, reasoning, grounding, usage, speed }, ) => { // if there is traceId, update it if (traceId) { msgTraceId = traceId; await messageService.updateMessage(messageId, { traceId, observationId: observationId ?? undefined, }); } // 等待所有图片上传完成 let finalImages: ChatImageItem[] = []; if (uploadTasks.size > 0) { try { // 等待所有上传任务完成 const uploadResults = await Promise.all(uploadTasks.values()); // 使用上传后的 S3 URL 替换原始图像数据 finalImages = uploadResults.filter((i) => !!i.url) as ChatImageItem[]; } catch (error) { console.error('Error waiting for image uploads:', error); } } let parsedToolCalls = toolCalls; if (parsedToolCalls && parsedToolCalls.length > 0) { internal_toggleToolCallingStreaming(messageId, undefined); parsedToolCalls = parsedToolCalls.map((item) => ({ ...item, function: { ...item.function, arguments: !!item.function.arguments ? item.function.arguments : '{}', }, })); isFunctionCall = true; } // update the content after fetch result await internal_updateMessageContent(messageId, content, { toolCalls: parsedToolCalls, reasoning: !!reasoning ? { ...reasoning, duration } : undefined, search: !!grounding?.citations ? grounding : undefined, imageList: finalImages.length > 0 ? finalImages : undefined, metadata: speed ? { ...usage, ...speed } : usage, }); }, onMessageHandle: async (chunk) => { switch (chunk.type) { case 'grounding': { // if there is no citations, then stop if ( !chunk.grounding || !chunk.grounding.citations || chunk.grounding.citations.length <= 0 ) return; internal_dispatchMessage({ id: messageId, type: 'updateMessage', value: { search: { citations: chunk.grounding.citations, searchQueries: chunk.grounding.searchQueries, }, }, }); break; } case 'base64_image': { internal_dispatchMessage({ id: messageId, type: 'updateMessage', value: { imageList: chunk.images.map((i) => ({ id: i.id, url: i.data, alt: i.id })), }, }); const image = chunk.image; const task = getFileStoreState() .uploadBase64FileWithProgress(image.data) .then((value) => ({ id: value?.id, url: value?.url, alt: value?.filename || value?.id, })); uploadTasks.set(image.id, task); break; } case 'text': { output += chunk.text; // if there is no duration, it means the end of reasoning if (!duration) { duration = Date.now() - thinkingStartAt; const isInChatReasoning = chatSelectors.isMessageInChatReasoning(messageId)(get()); if (isInChatReasoning) { internal_toggleChatReasoning( false, messageId, n('toggleChatReasoning/false') as string, ); } } internal_dispatchMessage({ id: messageId, type: 'updateMessage', value: { content: output, reasoning: !!thinking ? { content: thinking, duration } : undefined, }, }); break; } case 'reasoning': { // if there is no thinkingStartAt, it means the start of reasoning if (!thinkingStartAt) { thinkingStartAt = Date.now(); internal_toggleChatReasoning( true, messageId, n('toggleChatReasoning/true') as string, ); } thinking += chunk.text; internal_dispatchMessage({ id: messageId, type: 'updateMessage', value: { reasoning: { content: thinking } }, }); break; } // is this message is just a tool call case 'tool_calls': { internal_toggleToolCallingStreaming(messageId, chunk.isAnimationActives); internal_dispatchMessage({ id: messageId, type: 'updateMessage', value: { tools: get().internal_transformToolCalls(chunk.tool_calls) }, }); isFunctionCall = true; } } }, }); internal_toggleChatLoading(false, messageId, n('generateMessage(end)') as string); return { isFunctionCall, traceId: msgTraceId }; }, internal_resendMessage: async ( messageId, { traceId, messages: outChats, threadId: outThreadId, inPortalThread } = {}, ) => { // 1. 构造所有相关的历史记录 const chats = outChats ?? chatSelectors.mainAIChats(get()); const currentIndex = chats.findIndex((c) => c.id === messageId); if (currentIndex < 0) return; const currentMessage = chats[currentIndex]; let contextMessages: ChatMessage[] = []; switch (currentMessage.role) { case 'tool': case 'user': { contextMessages = chats.slice(0, currentIndex + 1); break; } case 'assistant': { // 消息是 AI 发出的因此需要找到它的 user 消息 const userId = currentMessage.parentId; const userIndex = chats.findIndex((c) => c.id === userId); // 如果消息没有 parentId,那么同 user/function 模式 contextMessages = chats.slice(0, userIndex < 0 ? currentIndex + 1 : userIndex + 1); break; } } if (contextMessages.length <= 0) return; const { internal_coreProcessMessage, activeThreadId } = get(); const latestMsg = contextMessages.findLast((s) => s.role === 'user'); if (!latestMsg) return; const threadId = outThreadId ?? activeThreadId; await internal_coreProcessMessage(contextMessages, latestMsg.id, { traceId, ragQuery: get().internal_shouldUseRAG() ? latestMsg.content : undefined, threadId, inPortalThread, }); }, // ----- Loading ------- // internal_toggleChatLoading: (loading, id, action) => { return get().internal_toggleLoadingArrays('chatLoadingIds', loading, id, action); }, internal_toggleChatReasoning: (loading, id, action) => { return get().internal_toggleLoadingArrays('reasoningLoadingIds', loading, id, action); }, internal_toggleToolCallingStreaming: (id, streaming) => { set( { toolCallingStreamIds: produce(get().toolCallingStreamIds, (draft) => { if (!!streaming) { draft[id] = streaming; } else { delete draft[id]; } }), }, false, `toggleToolCallingStreaming/${!!streaming ? 'start' : 'end'}`, ); }, internal_toggleSearchWorkflow: (loading, id) => { return get().internal_toggleLoadingArrays('searchWorkflowLoadingIds', loading, id); }, });