Merge pull request #18961 from Budibase/fix/channel-drift

fix: RAG drift within Teams / Slack channels vs preview chat

Author: Dakuan

packages/server/src/api/controllers/ai/chatConversations.ts

@@ -23,9 +23,10 @@ import {
 } from "@budibase/types"
 import {
   consumeStream,
+  readUIMessageStream,
   type LanguageModelUsage,
-  type StreamTextResult,
-  type ToolSet,
+  type UIMessage,
+  type UIMessageChunk,
 } from "ai"
 import sdk from "../../../sdk"
 import {
@@ -270,10 +271,35 @@ const resolveChatStreamRequest = async (
   }
 }
 
-export type WebhookAssistantStream = StreamTextResult<
-  ToolSet,
-  never
->["fullStream"]
+export type WebhookAssistantStream = AsyncIterable<string>
+
+const getAssistantMessageText = (assistantMessage?: UIMessage) =>
+  assistantMessage?.parts
+    ?.flatMap(part => (part.type === "text" ? [part.text] : []))
+    .join("") || ""
+
+const createAssistantTextStream = async function* (
+  stream: ReadableStream<UIMessageChunk>
+): AsyncGenerator<string, void, void> {
+  let previousText = ""
+  for await (const assistantMessage of readUIMessageStream({ stream })) {
+    const currentText = getAssistantMessageText(assistantMessage)
+    if (!currentText || currentText === previousText) {
+      continue
+    }
+
+    if (currentText.startsWith(previousText)) {
+      const delta = currentText.slice(previousText.length)
+      if (delta) {
+        yield delta
+      }
+    } else {
+      yield currentText
+    }
+
+    previousText = currentText
+  }
+}
 
 export async function webhookChat({
   chat,
@@ -320,15 +346,37 @@ export async function webhookChat({
 
   const result = await run.stream()
 
-  const streamTask = onAssistantStream
-    ? onAssistantStream(result.fullStream)
-    : Promise.resolve()
+  const uiMessageStream = result.toUIMessageStream({
+    generateMessageId: v4,
+    sendReasoning: true,
+  })
+  let assistantStreamForCapture: ReadableStream<UIMessageChunk> =
+    uiMessageStream
+  let streamTask: Promise<void> = Promise.resolve()
+  if (onAssistantStream) {
+    const [deliveryStream, captureStream] = uiMessageStream.tee()
+    assistantStreamForCapture = captureStream
+    streamTask = onAssistantStream(createAssistantTextStream(deliveryStream))
+  }
+
+  const assistantMessageTask = (async () => {
+    let assistantMessage: ChatConversation["messages"][number] | undefined
+    for await (const message of readUIMessageStream<
+      ChatConversation["messages"][number]
+    >({
+      stream: assistantStreamForCapture,
+    })) {
+      assistantMessage = message
+    }
+    return assistantMessage
+  })()
 
-  const [textResult, responseResult, streamOutcome] = await Promise.allSettled([
-    result.text,
-    result.response,
-    streamTask,
-  ])
+  const [assistantMessageResult, responseResult, streamOutcome] =
+    await Promise.allSettled([
+      assistantMessageTask,
+      result.response,
+      streamTask,
+    ])
 
   if (streamOutcome.status === "rejected") {
     console.error("Chat webhook stream delivery failed", streamOutcome.reason)
@@ -346,19 +394,19 @@ export async function webhookChat({
   run.sessionLogIndexer.addRequestId(requestId)
   await run.sessionLogIndexer.index()
 
-  if (textResult.status === "rejected") {
+  if (assistantMessageResult.status === "rejected") {
     console.error("Agent streaming error", {
       agentId,
       chatAppId,
       sessionId,
-      error: textResult.reason,
+      error: assistantMessageResult.reason,
     })
     events.action.aiAgentFailed({
       agentId,
       reason: ActionFailureReason.ERROR,
-      errorMessage: getErrorMessage(textResult.reason),
+      errorMessage: getErrorMessage(assistantMessageResult.reason),
     })
-    throw textResult.reason
+    throw assistantMessageResult.reason
   }
   if (responseResult.status === "rejected") {
     console.error("Agent response metadata error", {
@@ -376,12 +424,18 @@ export async function webhookChat({
   }
 
   events.action.aiAgentExecuted({ agentId })
-
-  const assistantText = textResult.value
+  const ragSources = run.getUsedKnowledgeSourcesMetadata()
+
+  const finalAssistantMessage =
+    assistantMessageResult.value ||
+    ({
+      id: v4(),
+      role: "assistant",
+      parts: [{ type: "text", text: "" }],
+    } satisfies ChatConversation["messages"][number])
+  const assistantText = getAssistantMessageText(finalAssistantMessage)
   const assistantMessage: ChatConversation["messages"][number] = {
-    id: v4(),
-    role: "assistant",
-    parts: [{ type: "text", text: assistantText || "" }],
+    ...finalAssistantMessage,
   }
 
   return {
@@ -391,6 +445,7 @@ export async function webhookChat({
     allowKnowledgeSourceDownload:
       run.selectedOperation?.allowKnowledgeSourceDownload,
     title,
+    ...(ragSources?.length ? { ragSources } : {}),
   }
 }