diff --git a/packages/server-shared/src/types/websocket/events.ts b/packages/server-shared/src/types/websocket/events.ts index a7ac76e0f..bcf69405c 100644 --- a/packages/server-shared/src/types/websocket/events.ts +++ b/packages/server-shared/src/types/websocket/events.ts @@ -25,7 +25,11 @@ interface InputSource { } interface OutputSource { - 'gen-ai:chat': string + 'gen-ai:chat': { + input: UserMessage + contexts: Record + composedMessage: Array + } } export enum ContextUpdateStrategy { @@ -118,11 +122,14 @@ export interface WebSocketEvents { message: AssistantMessage } & Partial> & Partial> 'output:gen-ai:chat:complete': { - input: UserMessage - contexts: Record - composedMessage: Array message: AssistantMessage toolCalls: ToolMessage[] + usage: { + promptTokens: number + completionTokens: number + totalTokens: number + source: 'provider-based' | 'estimate-based' + } } & Partial> & Partial> 'context:update': ContextUpdate diff --git a/packages/stage-ui/src/stores/chat.ts b/packages/stage-ui/src/stores/chat.ts index b03aaef5c..4115b34ec 100644 --- a/packages/stage-ui/src/stores/chat.ts +++ b/packages/stage-ui/src/stores/chat.ts @@ -2,7 +2,7 @@ import type { ChatProvider } from '@xsai-ext/shared-providers' import type { CommonContentPart, Message, SystemMessage, ToolMessage } from '@xsai/shared-chat' import type { StreamEvent, StreamOptions } from '../stores/llm' -import type { ChatAssistantMessage, ChatHistoryItem, ChatSlices, ContextMessage, StreamingAssistantMessage } from '../types/chat' +import type { ChatAssistantMessage, ChatHistoryItem, ChatSlices, ChatStreamEventContext, ContextMessage, StreamingAssistantMessage } from '../types/chat' import { ContextUpdateStrategy } from '@proj-airi/server-sdk' import { useLocalStorage } from '@vueuse/core' @@ -90,77 +90,63 @@ export const useChatStore = defineStore('chat', () => { }) // ----- Hooks (UI callbacks) ----- - const onBeforeMessageComposedHooks = ref Promise>>([]) - const onAfterMessageComposedHooks = ref Promise>>([]) - const onBeforeSendHooks = ref Promise>>([]) - const onAfterSendHooks = ref Promise>>([]) - const onTokenLiteralHooks = ref Promise>>([]) - const onTokenSpecialHooks = ref Promise>>([]) - const onStreamEndHooks = ref Promise>>([]) - const onAssistantResponseEndHooks = ref Promise>>([]) - const onAssistantMessageHooks = ref Promise>>([]) - const onChatTurnCompleteHooks = ref - composedMessage: Message[] - output: StreamingAssistantMessage - outputText: string - toolCalls: ToolMessage[] - }) => Promise>>([]) + const onBeforeMessageComposedHooks = ref) => Promise>>([]) + const onAfterMessageComposedHooks = ref Promise>>([]) + const onBeforeSendHooks = ref Promise>>([]) + const onAfterSendHooks = ref Promise>>([]) + const onTokenLiteralHooks = ref Promise>>([]) + const onTokenSpecialHooks = ref Promise>>([]) + const onStreamEndHooks = ref Promise>>([]) + const onAssistantResponseEndHooks = ref Promise>>([]) + const onAssistantMessageHooks = ref Promise>>([]) + const onChatTurnCompleteHooks = ref Promise>>([]) - function onBeforeMessageComposed(cb: (message: string) => Promise) { + function onBeforeMessageComposed(cb: (message: string, context: Omit) => Promise) { onBeforeMessageComposedHooks.value.push(cb) return () => onBeforeMessageComposedHooks.value = onBeforeMessageComposedHooks.value.filter(hook => hook !== cb) // return remove listener callback } - function onAfterMessageComposed(cb: (message: string) => Promise) { + function onAfterMessageComposed(cb: (message: string, context: ChatStreamEventContext) => Promise) { onAfterMessageComposedHooks.value.push(cb) return () => onAfterMessageComposedHooks.value = onAfterMessageComposedHooks.value.filter(hook => hook !== cb) // return remove listener callback } - function onBeforeSend(cb: (message: string) => Promise) { + function onBeforeSend(cb: (message: string, context: ChatStreamEventContext) => Promise) { onBeforeSendHooks.value.push(cb) return () => onBeforeSendHooks.value = onBeforeSendHooks.value.filter(hook => hook !== cb) // return remove listener callback } - function onAfterSend(cb: (message: string) => Promise) { + function onAfterSend(cb: (message: string, context: ChatStreamEventContext) => Promise) { onAfterSendHooks.value.push(cb) return () => onAfterSendHooks.value = onAfterSendHooks.value.filter(hook => hook !== cb) // return remove listener callback } - function onTokenLiteral(cb: (literal: string) => Promise) { + function onTokenLiteral(cb: (literal: string, context: ChatStreamEventContext) => Promise) { onTokenLiteralHooks.value.push(cb) return () => onTokenLiteralHooks.value = onTokenLiteralHooks.value.filter(hook => hook !== cb) // return remove listener callback } - function onTokenSpecial(cb: (special: string) => Promise) { + function onTokenSpecial(cb: (special: string, context: ChatStreamEventContext) => Promise) { onTokenSpecialHooks.value.push(cb) return () => onTokenSpecialHooks.value = onTokenSpecialHooks.value.filter(hook => hook !== cb) // return remove listener callback } - function onStreamEnd(cb: () => Promise) { + function onStreamEnd(cb: (context: ChatStreamEventContext) => Promise) { onStreamEndHooks.value.push(cb) return () => onStreamEndHooks.value = onStreamEndHooks.value.filter(hook => hook !== cb) // return remove listener callback } - function onAssistantResponseEnd(cb: (message: string) => Promise) { + function onAssistantResponseEnd(cb: (message: string, context: ChatStreamEventContext) => Promise) { onAssistantResponseEndHooks.value.push(cb) return () => onAssistantResponseEndHooks.value = onAssistantResponseEndHooks.value.filter(hook => hook !== cb) // return remove listener callback } - function onAssistantMessage(cb: (message: StreamingAssistantMessage, messageText?: string) => Promise) { + function onAssistantMessage(cb: (message: StreamingAssistantMessage, messageText: string, context: ChatStreamEventContext) => Promise) { onAssistantMessageHooks.value.push(cb) return () => onAssistantMessageHooks.value = onAssistantMessageHooks.value.filter(hook => hook !== cb) // return remove listener callback } - function onChatTurnComplete(cb: (chat: { - input: ChatHistoryItem - contexts: Record - composedMessage: Message[] - output: StreamingAssistantMessage - outputText: string - toolCalls: ToolMessage[] - }) => Promise) { + function onChatTurnComplete(cb: (chat: { output: StreamingAssistantMessage, outputText: string, toolCalls: ToolMessage[] }, context: ChatStreamEventContext) => Promise) { onChatTurnCompleteHooks.value.push(cb) return () => onChatTurnCompleteHooks.value = onChatTurnCompleteHooks.value.filter(hook => hook !== cb) // return remove listener callback } @@ -178,61 +164,54 @@ export const useChatStore = defineStore('chat', () => { onChatTurnCompleteHooks.value = [] } - async function emitBeforeMessageComposedHooks(message: string) { + async function emitBeforeMessageComposedHooks(message: string, context: Omit) { for (const hook of onBeforeMessageComposedHooks.value) - await hook(message) + await hook(message, context) } - async function emitAfterMessageComposedHooks(message: string) { + async function emitAfterMessageComposedHooks(message: string, context: ChatStreamEventContext) { for (const hook of onAfterMessageComposedHooks.value) - await hook(message) + await hook(message, context) } - async function emitBeforeSendHooks(message: string) { + async function emitBeforeSendHooks(message: string, context: ChatStreamEventContext) { for (const hook of onBeforeSendHooks.value) - await hook(message) + await hook(message, context) } - async function emitAfterSendHooks(message: string) { + async function emitAfterSendHooks(message: string, context: ChatStreamEventContext) { for (const hook of onAfterSendHooks.value) - await hook(message) + await hook(message, context) } - async function emitTokenLiteralHooks(literal: string) { + async function emitTokenLiteralHooks(literal: string, context: ChatStreamEventContext) { for (const hook of onTokenLiteralHooks.value) - await hook(literal) + await hook(literal, context) } - async function emitTokenSpecialHooks(special: string) { + async function emitTokenSpecialHooks(special: string, context: ChatStreamEventContext) { for (const hook of onTokenSpecialHooks.value) - await hook(special) + await hook(special, context) } - async function emitStreamEndHooks() { + async function emitStreamEndHooks(context: ChatStreamEventContext) { for (const hook of onStreamEndHooks.value) - await hook() + await hook(context) } - async function emitAssistantResponseEndHooks(message: string) { + async function emitAssistantResponseEndHooks(message: string, context: ChatStreamEventContext) { for (const hook of onAssistantResponseEndHooks.value) - await hook(message) + await hook(message, context) } - async function emitAssistantMessageHooks(message: StreamingAssistantMessage, messageText?: string) { + async function emitAssistantMessageHooks(message: StreamingAssistantMessage, messageText: string, context: ChatStreamEventContext) { for (const hook of onAssistantMessageHooks.value) - await hook(message, messageText) + await hook(message, messageText, context) } - async function emitChatTurnCompleteHooks(chat: { - input: ChatHistoryItem - contexts: Record - composedMessage: Message[] - output: StreamingAssistantMessage - outputText: string - toolCalls: ToolMessage[] - }) { + async function emitChatTurnCompleteHooks(chat: { output: StreamingAssistantMessage, outputText: string, toolCalls: ToolMessage[] }, context: ChatStreamEventContext) { for (const hook of onChatTurnCompleteHooks.value) - await hook(chat) + await hook(chat, context) } // ----- Session state helpers ----- @@ -369,6 +348,12 @@ export const useChatStore = defineStore('chat', () => { ensureSession(sessionId) const sendingCreatedAt = Date.now() + const streamingMessageContext: ChatStreamEventContext = { + input: { role: 'user', content: sendingMessage, createdAt: sendingCreatedAt }, + contexts: { ...activeContexts.value }, + composedMessage: [], + } + const isStaleGeneration = () => getSessionGeneration(sessionId) !== generation const shouldAbort = () => isStaleGeneration() if (shouldAbort()) @@ -380,7 +365,7 @@ export const useChatStore = defineStore('chat', () => { trackFirstMessage() try { - await emitBeforeMessageComposedHooks(sendingMessage) + await emitBeforeMessageComposedHooks(sendingMessage, streamingMessageContext) const contentParts: CommonContentPart[] = [{ type: 'text', text: sendingMessage }] @@ -398,6 +383,7 @@ export const useChatStore = defineStore('chat', () => { } const finalContent = contentParts.length > 1 ? contentParts : sendingMessage + streamingMessageContext.input.content = finalContent if (shouldAbort()) return @@ -409,7 +395,8 @@ export const useChatStore = defineStore('chat', () => { onLiteral: async (literal) => { if (shouldAbort()) return - await emitTokenLiteralHooks(literal) + + await emitTokenLiteralHooks(literal, streamingMessageContext) streamingMessage.value.content += literal @@ -429,7 +416,7 @@ export const useChatStore = defineStore('chat', () => { if (shouldAbort()) return - await emitTokenSpecialHooks(special) + await emitTokenSpecialHooks(special, streamingMessageContext) }, minLiteralEmitLength: 24, // Avoid emitting literals too fast. This is a magic number and can be changed later. }) @@ -488,8 +475,10 @@ export const useChatStore = defineStore('chat', () => { ] } - await emitAfterMessageComposedHooks(sendingMessage) - await emitBeforeSendHooks(sendingMessage) + streamingMessageContext.composedMessage = newMessages as Message[] + + await emitAfterMessageComposedHooks(sendingMessage, streamingMessageContext) + await emitBeforeSendHooks(sendingMessage, streamingMessageContext) let fullText = '' const headers = (options.providerConfig?.headers || {}) as Record @@ -540,24 +529,21 @@ export const useChatStore = defineStore('chat', () => { // Instruct the TTS pipeline to flush by calling hooks directly const flushSignal = `${TTS_FLUSH_INSTRUCTION}${TTS_FLUSH_INSTRUCTION}` - await emitTokenLiteralHooks(flushSignal) + await emitTokenLiteralHooks(flushSignal, streamingMessageContext) // Call the end-of-stream hooks - await emitStreamEndHooks() + await emitStreamEndHooks(streamingMessageContext) // Call the end-of-response hooks with the full text - await emitAssistantResponseEndHooks(fullText) + await emitAssistantResponseEndHooks(fullText, streamingMessageContext) - await emitAfterSendHooks(sendingMessage) - await emitAssistantMessageHooks({ ...streamingMessage.value }, fullText) + await emitAfterSendHooks(sendingMessage, streamingMessageContext) + await emitAssistantMessageHooks({ ...streamingMessage.value }, fullText, streamingMessageContext) await emitChatTurnCompleteHooks({ - input: { role: 'user', content: sendingMessage, createdAt: sendingCreatedAt }, - contexts: { ...activeContexts.value }, - composedMessage: newMessages as Message[], output: { ...streamingMessage.value }, outputText: fullText, toolCalls: sessionMessagesForSend.filter(msg => msg.role === 'tool') as ToolMessage[], - }) + }, streamingMessageContext) // Reset the streaming message for the next turn streamingMessage.value = { role: 'assistant', content: '', slices: [], tool_results: [] } diff --git a/packages/stage-ui/src/stores/mods/api/context-bridge.ts b/packages/stage-ui/src/stores/mods/api/context-bridge.ts index 80180f402..e2c7f2a05 100644 --- a/packages/stage-ui/src/stores/mods/api/context-bridge.ts +++ b/packages/stage-ui/src/stores/mods/api/context-bridge.ts @@ -40,117 +40,130 @@ export const useContextBridgeStore = defineStore('mods:api:context-bridge', () = })) disposeHookFns.value.push( - chatStore.onBeforeMessageComposed(async (message) => { + chatStore.onBeforeMessageComposed(async (message, context) => { if (isProcessingRemoteStream) return - broadcastStreamEvent({ type: 'before-compose', message, sessionId: chatStore.activeSessionId }) + broadcastStreamEvent({ type: 'before-compose', message, sessionId: chatStore.activeSessionId, context }) }), - chatStore.onAfterMessageComposed(async (message) => { + chatStore.onAfterMessageComposed(async (message, context) => { if (isProcessingRemoteStream) return - broadcastStreamEvent({ type: 'after-compose', message, sessionId: chatStore.activeSessionId }) + broadcastStreamEvent({ type: 'after-compose', message, sessionId: chatStore.activeSessionId, context }) }), - chatStore.onBeforeSend(async (message) => { + chatStore.onBeforeSend(async (message, context) => { if (isProcessingRemoteStream) return - broadcastStreamEvent({ type: 'before-send', message, sessionId: chatStore.activeSessionId }) + broadcastStreamEvent({ type: 'before-send', message, sessionId: chatStore.activeSessionId, context }) }), - chatStore.onAfterSend(async (message) => { + chatStore.onAfterSend(async (message, context) => { if (isProcessingRemoteStream) return - broadcastStreamEvent({ type: 'after-send', message, sessionId: chatStore.activeSessionId }) + broadcastStreamEvent({ type: 'after-send', message, sessionId: chatStore.activeSessionId, context }) }), - chatStore.onTokenLiteral(async (literal) => { + chatStore.onTokenLiteral(async (literal, context) => { if (isProcessingRemoteStream) return - broadcastStreamEvent({ type: 'token-literal', literal, sessionId: chatStore.activeSessionId }) + broadcastStreamEvent({ type: 'token-literal', literal, sessionId: chatStore.activeSessionId, context }) }), - chatStore.onTokenSpecial(async (special) => { + chatStore.onTokenSpecial(async (special, context) => { if (isProcessingRemoteStream) return - broadcastStreamEvent({ type: 'token-special', special, sessionId: chatStore.activeSessionId }) + broadcastStreamEvent({ type: 'token-special', special, sessionId: chatStore.activeSessionId, context }) }), - chatStore.onStreamEnd(async () => { + chatStore.onStreamEnd(async (context) => { if (isProcessingRemoteStream) return - broadcastStreamEvent({ type: 'stream-end', sessionId: chatStore.activeSessionId }) + broadcastStreamEvent({ type: 'stream-end', sessionId: chatStore.activeSessionId, context }) }), - chatStore.onAssistantResponseEnd(async (message) => { + chatStore.onAssistantResponseEnd(async (message, context) => { if (isProcessingRemoteStream) return - broadcastStreamEvent({ type: 'assistant-end', message, sessionId: chatStore.activeSessionId }) + broadcastStreamEvent({ type: 'assistant-end', message, sessionId: chatStore.activeSessionId, context }) + }), + + chatStore.onAssistantMessage(async (message, messageText, context) => { + serverChannelStore.send({ + type: 'output:gen-ai:chat:message', + data: { + message, + 'stage-web': isStageWeb(), + 'stage-tamagotchi': isStageTamagotchi(), + 'gen-ai:chat': { + input: context.input as UserMessage, + composedMessage: context.composedMessage, + contexts: context.contexts, + }, + }, + }) + }), + + chatStore.onChatTurnComplete(async (chat, context) => { + serverChannelStore.send({ + type: 'output:gen-ai:chat:complete', + data: { + 'message': chat.output, + 'toolCalls': [], + 'stage-web': isStageWeb(), + 'stage-tamagotchi': isStageTamagotchi(), + // TODO: Properly calculate usage data + 'usage': { + promptTokens: 0, + completionTokens: 0, + totalTokens: 0, + source: 'estimate-based', + }, + 'gen-ai:chat': { + input: context.input as UserMessage, + composedMessage: context.composedMessage, + contexts: context.contexts, + }, + }, + }) }), ) - disposeHookFns.value.push(chatStore.onAssistantMessage(async (message, messageText) => { - serverChannelStore.send({ - type: 'output:gen-ai:chat:message', - data: { - message, - 'stage-web': isStageWeb(), - 'stage-tamagotchi': isStageTamagotchi(), - 'gen-ai:chat': messageText || '', - }, - }) - })) - - disposeHookFns.value.push(chatStore.onChatTurnComplete(async (chat) => { - serverChannelStore.send({ - type: 'output:gen-ai:chat:complete', - data: { - 'input': chat.input as UserMessage, - 'composedMessage': chat.composedMessage, - 'contexts': chat.contexts, - 'message': chat.output, - 'toolCalls': [], - 'stage-web': isStageWeb(), - 'stage-tamagotchi': isStageTamagotchi(), - 'gen-ai:chat': chat.outputText, - }, - }) - })) - const { stop: stopIncomingStreamWatch } = watch(incomingStreamEvent, async (event) => { if (!event) return isProcessingRemoteStream = true + try { if (event.sessionId && chatStore.activeSessionId !== event.sessionId) chatStore.setActiveSession(event.sessionId) switch (event.type) { case 'before-compose': - await chatStore.emitBeforeMessageComposedHooks(event.message) + await chatStore.emitBeforeMessageComposedHooks(event.message, event.context) break case 'after-compose': - await chatStore.emitAfterMessageComposedHooks(event.message) + await chatStore.emitAfterMessageComposedHooks(event.message, event.context) break case 'before-send': - await chatStore.emitBeforeSendHooks(event.message) + await chatStore.emitBeforeSendHooks(event.message, event.context) break case 'after-send': - await chatStore.emitAfterSendHooks(event.message) + await chatStore.emitAfterSendHooks(event.message, event.context) break case 'token-literal': - await chatStore.emitTokenLiteralHooks(event.literal) + await chatStore.emitTokenLiteralHooks(event.literal, event.context) break case 'token-special': - await chatStore.emitTokenSpecialHooks(event.special) + await chatStore.emitTokenSpecialHooks(event.special, event.context) break case 'stream-end': - await chatStore.emitStreamEndHooks() + await chatStore.emitStreamEndHooks(event.context) break case 'assistant-end': - await chatStore.emitAssistantResponseEndHooks(event.message) + await chatStore.emitAssistantResponseEndHooks(event.message, event.context) break } } diff --git a/packages/stage-ui/src/types/chat.ts b/packages/stage-ui/src/types/chat.ts index 310952b1e..2b65f94b1 100644 --- a/packages/stage-ui/src/types/chat.ts +++ b/packages/stage-ui/src/types/chat.ts @@ -1,5 +1,5 @@ import type { ContextUpdate, WebSocketEventSource } from '@proj-airi/server-sdk' -import type { AssistantMessage, CommonContentPart, CompletionToolCall, SystemMessage, ToolMessage, UserMessage } from '@xsai/shared-chat' +import type { AssistantMessage, CommonContentPart, CompletionToolCall, Message, SystemMessage, ToolMessage, UserMessage } from '@xsai/shared-chat' export interface ChatSlicesText { type: 'text' @@ -41,14 +41,21 @@ export interface ContextMessage extends ContextUpdate { export type ChatHistoryItem = (ChatMessage | ErrorMessage) & { context?: ContextMessage } & { createdAt?: number } +export interface ChatStreamEventContext { + input: ChatHistoryItem + contexts: Record + composedMessage: Message[] +} + export type ChatStreamEvent - = | { type: 'before-compose', message: string, sessionId: string } - | { type: 'after-compose', message: string, sessionId: string } - | { type: 'before-send', message: string, sessionId: string } - | { type: 'after-send', message: string, sessionId: string } - | { type: 'token-literal', literal: string, sessionId: string } - | { type: 'token-special', special: string, sessionId: string } - | { type: 'stream-end', sessionId: string } - | { type: 'assistant-end', message: string, sessionId: string } + = | { type: 'before-compose', message: string, sessionId: string, context: Omit } + | { type: 'after-compose', message: string, sessionId: string, context: ChatStreamEventContext } + | { type: 'before-send', message: string, sessionId: string, context: ChatStreamEventContext } + | { type: 'after-send', message: string, sessionId: string, context: ChatStreamEventContext } + | { type: 'token-literal', literal: string, sessionId: string, context: ChatStreamEventContext } + | { type: 'token-special', special: string, sessionId: string, context: ChatStreamEventContext } + | { type: 'stream-end', sessionId: string, context: ChatStreamEventContext } + | { type: 'assistant-end', message: string, sessionId: string, context: ChatStreamEventContext } + | { type: 'assistant-message', message: ChatAssistantMessage, sessionId: string, messageText: string, context: ChatStreamEventContext } export type StreamingAssistantMessage = ChatAssistantMessage & { context?: ContextMessage } & { createdAt?: number }