diff --git a/apps/beeai-ui/src/api/a2a/extensions/citation.ts b/apps/beeai-ui/src/api/a2a/extensions/citation.ts index 9b71741f4..bc354acd4 100644 --- a/apps/beeai-ui/src/api/a2a/extensions/citation.ts +++ b/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; +export type Citation = z.infer; + export const citationExtension: A2AExtension = { getSchema: () => z.object({ [URI]: schema }).partial(), getUri: () => URI, diff --git a/apps/beeai-ui/src/api/a2a/part-processors.ts b/apps/beeai-ui/src/api/a2a/part-processors.ts index 8d612e900..fc104da3a 100644 --- a/apps/beeai-ui/src/api/a2a/part-processors.ts +++ b/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 []; diff --git a/apps/beeai-ui/src/api/a2a/utils.ts b/apps/beeai-ui/src/api/a2a/utils.ts index 69779d42d..237da733e 100644 --- a/apps/beeai-ui/src/api/a2a/utils.ts +++ b/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;