feat(ui): allow multiple citations in CitationMetadata (#1008)

Signed-off-by: Petr Bulánek <bulanek.petr@gmail.com>

Author: PetrBulanek

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

@@ -9,7 +9,7 @@ import type { A2AExtension } from './types';
 
 const URI = 'https://a2a-extensions.beeai.dev/ui/citation/v1';
 
-const schema = z
+const citationSchema = z
   .object({
     url: z.string(),
     start_index: z.number(),
@@ -19,8 +19,14 @@ const schema = z
   })
   .partial();
 
+const schema = z.object({
+  citations: z.array(citationSchema),
+});
+
 export type CitationMetadata = z.infer<typeof schema>;
 
+export type Citation = z.infer<typeof citationSchema>;
+
 export const citationExtension: A2AExtension<typeof URI, CitationMetadata> = {
   getSchema: () => z.object({ [URI]: schema }).partial(),
   getUri: () => URI,

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

@@ -8,6 +8,7 @@ import { v4 as uuid } from 'uuid';
 
 import type { UIFilePart, UIMessagePart } from '#modules/messages/types.ts';
 import { UIMessagePartKind } from '#modules/messages/types.ts';
+import { isNotNull } from '#utils/helpers.ts';
 
 import {
   createSourcePart,
@@ -20,16 +21,14 @@ import {
 
 export function processMessageMetadata(message: Message): UIMessagePart[] {
   const trajectory = extractTrajectory(message.metadata);
-  const citation = extractCitation(message.metadata);
+  const citations = extractCitation(message.metadata)?.citations;
 
   if (trajectory) {
     return [createTrajectoryPart(trajectory)];
-  } else if (citation) {
-    const sourcePart = createSourcePart(citation, message.taskId);
+  } else if (citations) {
+    const sourceParts = citations.map((citation) => createSourcePart(citation, message.taskId)).filter(isNotNull);
 
-    if (sourcePart) {
-      return [sourcePart];
-    }
+    return [...sourceParts];
   }
 
   return [];

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

@@ -18,7 +18,7 @@ import { UIMessagePartKind } from '#modules/messages/types.ts';
 import type { ContextId, TaskId } from '#modules/tasks/api/types.ts';
 import { isNotNull } from '#utils/helpers.ts';
 
-import type { CitationMetadata } from './extensions/citation';
+import type { Citation } from './extensions/citation';
 import { citationExtension } from './extensions/citation';
 import type { TrajectoryMetadata } from './extensions/trajectory';
 import { trajectoryExtension } from './extensions/trajectory';
@@ -101,8 +101,8 @@ export function getFileUri(file: FilePart['file']): string {
   return `data:${mimeType};base64,${bytes}`;
 }
 
-export function createSourcePart(metadata: CitationMetadata, messageId: string | undefined): UISourcePart | null {
-  const { url, start_index, end_index, title, description } = metadata;
+export function createSourcePart(citation: Citation, messageId: string | undefined): UISourcePart | null {
+  const { url, start_index, end_index, title, description } = citation;
 
   if (!url || !messageId) {
     return null;