feat(ui): add artifact update parts processing (#1016)

Author: PetrBulanek

apps/beeai-ui/src/api/a2a/client.ts

@@ -3,18 +3,17 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import type { TaskStatusUpdateEvent } from '@a2a-js/sdk';
+import type { TaskArtifactUpdateEvent, TaskStatusUpdateEvent } from '@a2a-js/sdk';
 import { A2AClient } from '@a2a-js/sdk/client';
 import { Subject } from 'rxjs';
 import { match } from 'ts-pattern';
 
 import type { UIMessagePart, UIUserMessage } from '#modules/messages/types.ts';
 import type { ContextId, TaskId } from '#modules/tasks/api/types.ts';
 import { getBaseUrl } from '#utils/api/getBaseUrl.ts';
-import { isNotNull } from '#utils/helpers.ts';
 
 import { AGENT_ERROR_MESSAGE } from './constants';
-import { processFilePart, processMessageMetadata, processTextPart } from './part-processors';
+import { processMessageMetadata, processParts } from './part-processors';
 import type { ChatRun } from './types';
 import { createUserMessage, extractTextFromMessage } from './utils';
 
@@ -32,23 +31,17 @@ function handleStatusUpdate(event: TaskStatusUpdateEvent): UIMessagePart[] {
   }
 
   const metadataParts = processMessageMetadata(message);
+  const contentParts = processParts(message.parts);
 
-  const contentParts = message.parts
-    .flatMap((part) => {
-      const processedParts = match(part)
-        .with({ kind: 'text' }, (part) => processTextPart(part))
-        .with({ kind: 'file' }, processFilePart)
-        .otherwise((otherPart) => {
-          console.warn(`Unsupported part - ${otherPart.kind}`);
+  return [...metadataParts, ...contentParts];
+}
 
-          return null;
-        });
+function handleArtifactUpdate(event: TaskArtifactUpdateEvent): UIMessagePart[] {
+  const { artifact } = event;
 
-      return processedParts;
-    })
-    .filter(isNotNull);
+  const contentParts = processParts(artifact.parts);
 
-  return [...metadataParts, ...contentParts];
+  return contentParts;
 }
 
 export const buildA2AClient = (providerId: string) => {
@@ -70,9 +63,16 @@ export const buildA2AClient = (providerId: string) => {
           .with({ kind: 'status-update' }, (event) => {
             taskId = event.taskId;
 
-            const messageParts = handleStatusUpdate(event);
+            const parts = handleStatusUpdate(event);
+
+            messageSubject.next({ parts, taskId });
+          })
+          .with({ kind: 'artifact-update' }, (event) => {
+            taskId = event.taskId;
+
+            const parts = handleArtifactUpdate(event);
 
-            messageSubject.next({ parts: messageParts, taskId });
+            messageSubject.next({ parts, taskId });
           });
       }
       messageSubject.complete();

apps/beeai-ui/src/api/a2a/part-processors.ts

@@ -3,7 +3,8 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import type { FilePart, Message, TextPart } from '@a2a-js/sdk';
+import type { FilePart, Message, Part, TextPart } from '@a2a-js/sdk';
+import { match } from 'ts-pattern';
 import { v4 as uuid } from 'uuid';
 
 import type { UIFilePart, UIMessagePart } from '#modules/messages/types.ts';
@@ -54,3 +55,22 @@ export function processFilePart(part: FilePart): UIMessagePart[] {
 
   return [filePart];
 }
+
+export function processParts(parts: Part[]): UIMessagePart[] {
+  const processedParts = parts
+    .flatMap((part) => {
+      const processedParts = match(part)
+        .with({ kind: 'text' }, (part) => processTextPart(part))
+        .with({ kind: 'file' }, processFilePart)
+        .otherwise((otherPart) => {
+          console.warn(`Unsupported part - ${otherPart.kind}`);
+
+          return null;
+        });
+
+      return processedParts;
+    })
+    .filter(isNotNull);
+
+  return processedParts;
+}