chore: use message metadata in UI to render citations and trajectory (#995)

Author: tomkis

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

@@ -15,7 +15,7 @@ import { getBaseUrl } from '#utils/api/getBaseUrl.ts';
 import { isNotNull } from '#utils/helpers.ts';
 
 import { AGENT_ERROR_MESSAGE } from './constants';
-import { processFilePart, processTextPart } from './part-processors';
+import { processFilePart, processMessageMetadata, processTextPart } from './part-processors';
 import type { ChatRun } from './types';
 import { createUserMessage, extractTextFromMessage } from './utils';
 
@@ -32,10 +32,12 @@ function handleStatusUpdate(event: TaskStatusUpdateEvent): UIMessagePart[] {
     return [];
   }
 
-  const parts = message.parts
+  const metadataParts = processMessageMetadata(message);
+
+  const contentParts = message.parts
     .flatMap((part) => {
       const processedParts = match(part)
-        .with({ kind: 'text' }, (part) => processTextPart(part, message.messageId))
+        .with({ kind: 'text' }, (part) => processTextPart(part))
         .with({ kind: 'file' }, processFilePart)
         .otherwise((otherPart) => {
           console.warn(`Unsupported part - ${otherPart.kind}`);
@@ -47,7 +49,7 @@ function handleStatusUpdate(event: TaskStatusUpdateEvent): UIMessagePart[] {
     })
     .filter(isNotNull);
 
-  return parts;
+  return [...metadataParts, ...contentParts];
 }
 
 export const buildA2AClient = (providerId: string) => {

apps/beeai-ui/src/api/a2a/extensions/trajectory.ts

@@ -11,10 +11,8 @@ const URI = 'https://a2a-extensions.beeai.dev/ui/trajectory/v1';
 
 const schema = z
   .object({
-    message: z.string(),
-    tool_name: z.string(),
-    tool_input: z.record(z.string(), z.unknown()),
-    tool_outpyut: z.record(z.string(), z.unknown()),
+    title: z.string(),
+    content: z.string(),
   })
   .partial();
 

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

@@ -3,7 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import type { FilePart, TextPart } from '@a2a-js/sdk';
+import type { FilePart, Message, TextPart } from '@a2a-js/sdk';
 import { v4 as uuid } from 'uuid';
 
 import type { UIFilePart, UIMessagePart } from '#modules/messages/types.ts';
@@ -18,25 +18,25 @@ import {
   getFileUri,
 } from './utils';
 
-export function processTextPart(part: TextPart, messageId: string): UIMessagePart[] {
-  const parts: UIMessagePart[] = [];
-  const { metadata, text } = part;
-  const trajectory = extractTrajectory(metadata);
-  const citation = extractCitation(metadata);
+export function processMessageMetadata(message: Message): UIMessagePart[] {
+  const trajectory = extractTrajectory(message.metadata);
+  const citation = extractCitation(message.metadata);
 
   if (trajectory) {
-    parts.push(createTrajectoryPart(trajectory));
+    return [createTrajectoryPart(trajectory)];
   } else if (citation) {
-    const sourcePart = createSourcePart(citation, messageId);
+    const sourcePart = createSourcePart(citation, message.messageId);
 
     if (sourcePart) {
-      parts.push(sourcePart);
+      return [sourcePart];
     }
   }
 
-  parts.push(createTextPart(text));
+  return [];
+}
 
-  return parts;
+export function processTextPart({ text }: TextPart): UIMessagePart[] {
+  return [createTextPart(text)];
 }
 
 export function processFilePart(part: FilePart): UIMessagePart[] {

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

@@ -113,13 +113,13 @@ export function createSourcePart(metadata: CitationMetadata, messageId: string):
 }
 
 export function createTrajectoryPart(metadata: TrajectoryMetadata): UITrajectoryPart {
-  const { message, tool_name } = metadata;
+  const { title, content } = metadata;
 
   const trajectoryPart: UITrajectoryPart = {
     kind: UIMessagePartKind.Trajectory,
     id: uuid(),
-    message: message ?? undefined,
-    toolName: tool_name ?? undefined,
+    title,
+    content,
   };
 
   return trajectoryPart;

apps/beeai-ui/src/modules/messages/types.ts

@@ -53,8 +53,8 @@ export type UISourcePart = {
 export type UITrajectoryPart = {
   kind: UIMessagePartKind.Trajectory;
   id: string;
-  message?: string;
-  toolName?: string;
+  title?: string;
+  content?: string;
 };
 
 export type UITransformPart = {

apps/beeai-ui/src/modules/trajectories/components/TrajectoryItem.tsx

@@ -21,9 +21,9 @@ interface Props {
 export function TrajectoryItem({ trajectory }: Props) {
   const [isOpen, setIsOpen] = useState(false);
 
-  const { toolName, message } = trajectory;
+  const { title, content } = trajectory;
 
-  const isToggleable = isNotNull(message);
+  const isToggleable = isNotNull(content);
 
   return (
     <div className={clsx(classes.root, { [classes.isOpen]: isOpen })}>
@@ -40,17 +40,17 @@ export function TrajectoryItem({ trajectory }: Props) {
           </IconButton>
         )}
 
-        {toolName && (
+        {title && (
           <h3 className={classes.name}>
             {/* <span className={classes.icon}>
             <Icon />
           </span> */}
 
-            <span>{toolName}</span>
+            <span>{title}</span>
           </h3>
         )}
 
-        {message && <div className={classes.message}>{message}</div>}
+        {content && <div className={classes.message}>{content}</div>}
       </header>
 
       {isToggleable && (

apps/beeai-ui/src/modules/trajectories/components/TrajectoryItemContent.tsx

@@ -12,14 +12,12 @@ interface Props {
 }
 
 export function TrajectoryItemContent({ trajectory }: Props) {
-  const { message } = trajectory;
+  const { content } = trajectory;
 
   return (
     <div className={classes.root}>
       <div className={classes.group}>
-        <p className={classes.label}>Thought</p>
-
-        <p className={classes.content}>{message}</p>
+        <p className={classes.content}>{content}</p>
       </div>
     </div>
   );