@convex-dev/agent 0.5.0-alpha.1 → 0.6.0-alpha.0

package/src/UIMessages.ts

@@ -0,0 +1,739 @@
+import {
+  convertToModelMessages,
+  type UIMessage as AIUIMessage,
+  type DeepPartial,
+  type DynamicToolUIPart,
+  type ReasoningUIPart,
+  type SourceDocumentUIPart,
+  type SourceUrlUIPart,
+  type StepStartUIPart,
+  type TextUIPart,
+  type ToolResultPart,
+  type ToolUIPart,
+  type UIDataTypes,
+  type UITools,
+} from "ai";
+import type { Infer } from "convex/values";
+import { toModelMessage, fromModelMessage, toUIFilePart } from "./mapping.js";
+import {
+  extractReasoning,
+  extractText,
+  isTool,
+  joinText,
+  sorted,
+} from "./shared.js";
+import type {
+  MessageDoc,
+  MessageStatus,
+  ProviderOptions,
+  SourcePart,
+  vSource,
+} from "./validators.js";
+import { omit, pick } from "convex-helpers";
+
+export type UIStatus = "streaming" | MessageStatus;
+
+export type UIMessage<
+  METADATA = unknown,
+  DATA_PARTS extends UIDataTypes = UIDataTypes,
+  TOOLS extends UITools = UITools,
+> = AIUIMessage<METADATA, DATA_PARTS, TOOLS> & {
+  key: string;
+  order: number;
+  stepOrder: number;
+  status: UIStatus;
+  agentName?: string;
+  userId?: string;
+  text: string;
+  _creationTime: number;
+};
+
+/**
+ * Converts a list of UIMessages to MessageDocs, along with extra metadata that
+ * may be available to associate with the MessageDocs.
+ * @param messages - The UIMessages to convert to MessageDocs.
+ * @param meta - The metadata to add to the MessageDocs.
+ * @returns
+ */
+export async function fromUIMessages<METADATA = unknown>(
+  messages: UIMessage<METADATA>[],
+  meta: {
+    threadId: string;
+    userId?: string;
+    model?: string;
+    provider?: string;
+    providerOptions?: ProviderOptions;
+    metadata?: METADATA;
+  },
+): Promise<(MessageDoc & { streaming: boolean; metadata?: METADATA })[]> {
+  const nested = await Promise.all(
+    messages.map(async (uiMessage) => {
+      const stepOrder = uiMessage.stepOrder;
+      const commonFields = {
+        ...pick(meta, [
+          "threadId",
+          "userId",
+          "model",
+          "provider",
+          "providerOptions",
+          "metadata",
+        ]),
+        ...omit(uiMessage, ["parts", "role", "key", "text", "userId"]),
+        userId: uiMessage.userId ?? meta.userId,
+        status: uiMessage.status === "streaming" ? "pending" : "success",
+        streaming: uiMessage.status === "streaming",
+        // to override
+        _id: uiMessage.id,
+        tool: false,
+      } satisfies MessageDoc & { streaming: boolean; metadata?: METADATA };
+      const modelMessages = await convertToModelMessages([uiMessage]);
+      return modelMessages
+        .map((modelMessage, i) => {
+          if (modelMessage.content.length === 0) {
+            return undefined;
+          }
+          const message = fromModelMessage(modelMessage);
+          const tool = isTool(message);
+          const doc: MessageDoc & { streaming: boolean; metadata?: METADATA } =
+            {
+              ...commonFields,
+              _id: uiMessage.id + `-${i}`,
+              stepOrder: stepOrder + i,
+              message,
+              tool,
+              text: extractText(message),
+              reasoning: extractReasoning(message),
+              finishReason: tool ? "tool-calls" : "stop",
+              sources: fromSourceParts(uiMessage.parts),
+            };
+          if (Array.isArray(modelMessage.content)) {
+            // Find a content part with providerOptions (type assertion needed for SDK compatibility)
+            const partWithProviderOptions = modelMessage.content.find(
+              (c): c is typeof c & { providerOptions: unknown } =>
+                "providerOptions" in c && c.providerOptions !== undefined,
+            );
+            if (partWithProviderOptions?.providerOptions) {
+              // convertToModelMessages changes providerMetadata to providerOptions
+              const providerOptions =
+                partWithProviderOptions.providerOptions as
+                  | Record<string, Record<string, unknown>>
+                  | undefined;
+              if (providerOptions) {
+                doc.providerMetadata = providerOptions;
+                doc.providerOptions ??= providerOptions;
+              }
+            }
+          }
+          return doc;
+        })
+        .filter((d) => d !== undefined);
+    }),
+  );
+  return nested.flat();
+}
+
+function fromSourceParts(parts: UIMessage["parts"]): Infer<typeof vSource>[] {
+  return parts
+    .map((part) => {
+      if (part.type === "source-url") {
+        return {
+          type: "source",
+          sourceType: "url",
+          url: part.url,
+          id: part.sourceId,
+          providerMetadata: part.providerMetadata,
+          title: part.title,
+        } satisfies Infer<typeof vSource>;
+      }
+      if (part.type === "source-document") {
+        return {
+          type: "source",
+          sourceType: "document",
+          mediaType: part.mediaType,
+          id: part.sourceId,
+          providerMetadata: part.providerMetadata,
+          title: part.title,
+        } satisfies Infer<typeof vSource>;
+      }
+      return undefined;
+    })
+    .filter((p) => p !== undefined);
+}
+
+type ExtraFields<METADATA = unknown> = {
+  streaming?: boolean;
+  metadata?: METADATA;
+};
+
+/**
+ * Converts a list of MessageDocs to UIMessages.
+ * This is somewhat lossy, as many fields are not supported by UIMessages, e.g.
+ * the model, provider, userId, etc.
+ * The UIMessage type is the augmented type that includes more fields such as
+ * key, order, stepOrder, status, agentName, text, etc.
+ */
+export function toUIMessages<
+  METADATA = unknown,
+  DATA_PARTS extends UIDataTypes = UIDataTypes,
+  TOOLS extends UITools = UITools,
+>(
+  messages: (MessageDoc & ExtraFields<METADATA>)[],
+): UIMessage<METADATA, DATA_PARTS, TOOLS>[] {
+  // Group assistant and tool messages together
+  const assistantGroups = groupAssistantMessages(sorted(messages));
+
+  const uiMessages: UIMessage<METADATA, DATA_PARTS, TOOLS>[] = [];
+  for (const group of assistantGroups) {
+    if (group.role === "system") {
+      uiMessages.push(createSystemUIMessage(group.message));
+    } else if (group.role === "user") {
+      uiMessages.push(createUserUIMessage(group.message));
+    } else {
+      // Assistant/tool group
+      uiMessages.push(createAssistantUIMessage(group.messages));
+    }
+  }
+
+  return uiMessages;
+}
+
+type Group<METADATA = unknown> =
+  | {
+      role: "user";
+      message: MessageDoc & ExtraFields<METADATA>;
+    }
+  | {
+      role: "system";
+      message: MessageDoc & ExtraFields<METADATA>;
+    }
+  | {
+      role: "assistant";
+      messages: (MessageDoc & ExtraFields<METADATA>)[];
+    };
+
+function groupAssistantMessages<METADATA = unknown>(
+  messages: (MessageDoc & ExtraFields<METADATA>)[],
+): Group<METADATA>[] {
+  const groups: Group<METADATA>[] = [];
+
+  let currentAssistantGroup: (MessageDoc & ExtraFields<METADATA>)[] = [];
+  let currentOrder: number | undefined;
+
+  for (const message of messages) {
+    const coreMessage = message.message && toModelMessage(message.message);
+    if (!coreMessage) continue;
+
+    if (coreMessage.role === "user" || coreMessage.role === "system") {
+      // Finish any current assistant group
+      if (currentAssistantGroup.length > 0) {
+        groups.push({
+          role: "assistant",
+          messages: currentAssistantGroup,
+        });
+        currentAssistantGroup = [];
+        currentOrder = undefined;
+      }
+      // Add singleton group
+      groups.push({
+        role: coreMessage.role,
+        message,
+      });
+    } else {
+      // Assistant or tool message
+
+      // Start new group if order changes or this is the first assistant/tool message
+      if (currentOrder !== undefined && message.order !== currentOrder) {
+        if (currentAssistantGroup.length > 0) {
+          groups.push({
+            role: "assistant",
+            messages: currentAssistantGroup,
+          });
+          currentAssistantGroup = [];
+        }
+      }
+
+      currentOrder = message.order;
+      currentAssistantGroup.push(message);
+
+      // End group if this is an assistant message without tool calls
+      if (coreMessage.role === "assistant" && !message.tool) {
+        groups.push({
+          role: "assistant",
+          messages: currentAssistantGroup,
+        });
+        currentAssistantGroup = [];
+        currentOrder = undefined;
+      }
+    }
+  }
+
+  // Add any remaining assistant group
+  if (currentAssistantGroup.length > 0) {
+    groups.push({
+      role: "assistant",
+      messages: currentAssistantGroup,
+    });
+  }
+
+  return groups;
+}
+
+function createSystemUIMessage<
+  METADATA = unknown,
+  DATA_PARTS extends UIDataTypes = UIDataTypes,
+  TOOLS extends UITools = UITools,
+>(
+  message: MessageDoc & ExtraFields<METADATA>,
+): UIMessage<METADATA, DATA_PARTS, TOOLS> {
+  const text = extractTextFromMessageDoc(message);
+  const partCommon = {
+    state: message.streaming ? ("streaming" as const) : ("done" as const),
+    ...(message.providerMetadata
+      ? { providerMetadata: message.providerMetadata }
+      : {}),
+  };
+
+  return {
+    id: message._id,
+    _creationTime: message._creationTime,
+    order: message.order,
+    stepOrder: message.stepOrder,
+    status: message.streaming ? ("streaming" as const) : message.status,
+    key: `${message.threadId}-${message.order}-${message.stepOrder}`,
+    text,
+    role: "system",
+    agentName: message.agentName,
+    userId: message.userId,
+    parts: [{ type: "text", text, ...partCommon } satisfies TextUIPart],
+    metadata: message.metadata,
+  };
+}
+
+function extractTextFromMessageDoc(message: MessageDoc): string {
+  return (
+    (message.message && extractText(message.message)) || message.text || ""
+  );
+}
+
+function createUserUIMessage<
+  METADATA = unknown,
+  DATA_PARTS extends UIDataTypes = UIDataTypes,
+  TOOLS extends UITools = UITools,
+>(
+  message: MessageDoc & ExtraFields<METADATA>,
+): UIMessage<METADATA, DATA_PARTS, TOOLS> {
+  const text = extractTextFromMessageDoc(message);
+  const coreMessage = toModelMessage(message.message!);
+  const content = coreMessage.content;
+  const nonStringContent =
+    content && typeof content !== "string" ? content : [];
+
+  const partCommon = {
+    state: message.streaming ? ("streaming" as const) : ("done" as const),
+    ...(message.providerMetadata
+      ? { providerMetadata: message.providerMetadata }
+      : {}),
+  };
+
+  const parts: UIMessage<METADATA, DATA_PARTS, TOOLS>["parts"] = [];
+  if (text && !nonStringContent.length) {
+    parts.push({ type: "text", text });
+  }
+  for (const contentPart of nonStringContent) {
+    switch (contentPart.type) {
+      case "text":
+        parts.push({ type: "text", text: contentPart.text, ...partCommon });
+        break;
+      case "file":
+      case "image":
+        parts.push(toUIFilePart(contentPart));
+        break;
+      default:
+        console.warn("Unknown content part type for user", contentPart);
+        break;
+    }
+  }
+
+  return {
+    id: message._id,
+    _creationTime: message._creationTime,
+    order: message.order,
+    stepOrder: message.stepOrder,
+    status: message.streaming ? ("streaming" as const) : message.status,
+    key: `${message.threadId}-${message.order}-${message.stepOrder}`,
+    text,
+    role: "user",
+    userId: message.userId,
+    parts,
+    metadata: message.metadata,
+  };
+}
+
+function createAssistantUIMessage<
+  METADATA = unknown,
+  DATA_PARTS extends UIDataTypes = UIDataTypes,
+  TOOLS extends UITools = UITools,
+>(
+  groupUnordered: (MessageDoc & ExtraFields<METADATA>)[],
+): UIMessage<METADATA, DATA_PARTS, TOOLS> {
+  const group = sorted(groupUnordered);
+  const firstMessage = group[0];
+
+  // Use first message for special fields
+  const common = {
+    id: firstMessage._id,
+    _creationTime: firstMessage._creationTime,
+    order: firstMessage.order,
+    stepOrder: firstMessage.stepOrder,
+    key: `${firstMessage.threadId}-${firstMessage.order}-${firstMessage.stepOrder}`,
+    agentName: firstMessage.agentName,
+    userId: firstMessage.userId,
+  };
+
+  // Get status from last message
+  const lastMessage = group[group.length - 1];
+  const status = lastMessage.streaming
+    ? ("streaming" as const)
+    : lastMessage.status;
+
+  // Collect all parts from all messages
+  const allParts: UIMessage<METADATA, DATA_PARTS, TOOLS>["parts"] = [];
+
+  for (const message of group) {
+    const coreMessage = message.message && toModelMessage(message.message);
+    if (!coreMessage) continue;
+
+    const content = coreMessage.content;
+    const nonStringContent =
+      content && typeof content !== "string" ? content : [];
+    const text = extractTextFromMessageDoc(message);
+
+    const partCommon = {
+      state: message.streaming ? ("streaming" as const) : ("done" as const),
+      ...(message.providerMetadata
+        ? { providerMetadata: message.providerMetadata }
+        : {}),
+    };
+
+    // Add reasoning parts
+    if (
+      message.reasoning &&
+      !nonStringContent.some((c) => c.type === "reasoning")
+    ) {
+      allParts.push({
+        type: "reasoning",
+        text: message.reasoning,
+        ...partCommon,
+      } satisfies ReasoningUIPart);
+    }
+
+    // Add text parts if no structured content
+    if (text && !nonStringContent.length) {
+      allParts.push({
+        type: "text",
+        text: text,
+        ...partCommon,
+      } satisfies TextUIPart);
+    }
+
+    // Add all structured content parts
+    for (const contentPart of nonStringContent) {
+      switch (contentPart.type) {
+        case "text":
+          allParts.push({
+            ...partCommon,
+            ...contentPart,
+          } satisfies TextUIPart);
+          break;
+        case "reasoning":
+          allParts.push({
+            ...partCommon,
+            ...contentPart,
+          } satisfies ReasoningUIPart);
+          break;
+        case "file":
+        case "image":
+          allParts.push(toUIFilePart(contentPart));
+          break;
+        case "tool-call": {
+          allParts.push({
+            type: "step-start",
+          } satisfies StepStartUIPart);
+          const toolPart: ToolUIPart<TOOLS> = {
+            type: `tool-${contentPart.toolName as keyof TOOLS & string}`,
+            toolCallId: contentPart.toolCallId,
+            input: contentPart.input as DeepPartial<
+              TOOLS[keyof TOOLS & string]["input"]
+            >,
+            providerExecuted: contentPart.providerExecuted,
+            ...(message.streaming
+              ? { state: "input-streaming" }
+              : {
+                  state: "input-available",
+                  callProviderMetadata: message.providerMetadata,
+                }),
+          };
+          allParts.push(toolPart);
+          break;
+        }
+        case "tool-result": {
+          const typedPart = contentPart as unknown as ToolResultPart & {
+            output: { type: string; value?: unknown; reason?: string };
+          };
+
+          // Check if this is an execution-denied result
+          if (typedPart.output?.type === "execution-denied") {
+            const call = allParts.find(
+              (part) =>
+                part.type === `tool-${contentPart.toolName}` &&
+                "toolCallId" in part &&
+                part.toolCallId === contentPart.toolCallId,
+            ) as ToolUIPart | undefined;
+
+            if (call) {
+              call.state = "output-denied";
+              if (!("approval" in call) || !call.approval) {
+                (call as ToolUIPart & { approval?: object }).approval = {
+                  id: "",
+                  approved: false,
+                  reason: typedPart.output.reason,
+                };
+              } else {
+                const approval = (
+                  call as ToolUIPart & {
+                    approval: { approved?: boolean; reason?: string };
+                  }
+                ).approval;
+                approval.approved = false;
+                approval.reason = typedPart.output.reason;
+              }
+            }
+            break;
+          }
+
+          const output =
+            typeof typedPart.output?.type === "string"
+              ? typedPart.output.value
+              : typedPart.output;
+          // Check for error at both the content part level (isError) and message level
+          // isError may exist on stored tool results but isn't in ToolResultPart type
+          const hasError =
+            (contentPart as { isError?: boolean }).isError || message.error;
+          const errorText =
+            message.error || (hasError ? String(output) : undefined);
+          const call = allParts.find(
+            (part) =>
+              part.type === `tool-${contentPart.toolName}` &&
+              "toolCallId" in part &&
+              part.toolCallId === contentPart.toolCallId,
+          ) as ToolUIPart | undefined;
+          if (call) {
+            if (hasError) {
+              call.state = "output-error";
+              call.errorText = errorText ?? "Unknown error";
+              call.output = output;
+            } else {
+              call.state = "output-available";
+              call.output = output;
+            }
+          } else {
+            // Tool call is on a previous page - create standalone tool part
+            if (hasError) {
+              allParts.push({
+                type: `tool-${contentPart.toolName}`,
+                toolCallId: contentPart.toolCallId,
+                state: "output-error",
+                input: undefined,
+                errorText: errorText ?? "Unknown error",
+                callProviderMetadata: message.providerMetadata,
+              } satisfies ToolUIPart<TOOLS>);
+            } else {
+              allParts.push({
+                type: `tool-${contentPart.toolName}`,
+                toolCallId: contentPart.toolCallId,
+                state: "output-available",
+                input: undefined,
+                output,
+                callProviderMetadata: message.providerMetadata,
+              } satisfies ToolUIPart<TOOLS>);
+            }
+          }
+          break;
+        }
+        case "tool-approval-request": {
+          // Find the matching tool call
+          const typedPart = contentPart as {
+            toolCallId: string;
+            approvalId: string;
+          };
+          const toolCallPart = allParts.find(
+            (part) =>
+              "toolCallId" in part && part.toolCallId === typedPart.toolCallId,
+          ) as ToolUIPart | undefined;
+
+          if (toolCallPart) {
+            toolCallPart.state = "approval-requested";
+            (toolCallPart as ToolUIPart & { approval?: object }).approval = {
+              id: typedPart.approvalId,
+            };
+          } else {
+            console.warn(
+              "Tool approval request without preceding tool call",
+              contentPart,
+            );
+          }
+          break;
+        }
+        case "tool-approval-response": {
+          // Find the tool call that has this approval by matching approval.id
+          const typedPart = contentPart as {
+            approvalId: string;
+            approved: boolean;
+            reason?: string;
+          };
+          const toolCallPart = allParts.find(
+            (part) =>
+              "approval" in part &&
+              (part as ToolUIPart & { approval?: { id: string } }).approval
+                ?.id === typedPart.approvalId,
+          ) as ToolUIPart | undefined;
+
+          if (toolCallPart) {
+            if (typedPart.approved) {
+              toolCallPart.state = "approval-responded";
+              (toolCallPart as ToolUIPart & { approval?: object }).approval = {
+                id: typedPart.approvalId,
+                approved: true,
+                reason: typedPart.reason,
+              };
+            } else {
+              toolCallPart.state = "output-denied";
+              (toolCallPart as ToolUIPart & { approval?: object }).approval = {
+                id: typedPart.approvalId,
+                approved: false,
+                reason: typedPart.reason,
+              };
+            }
+          } else {
+            console.warn(
+              "Tool approval response without matching approval request",
+              contentPart,
+            );
+          }
+          break;
+        }
+        default: {
+          const maybeSource = contentPart as unknown as SourcePart;
+          if (maybeSource.type === "source") {
+            allParts.push(toSourcePart(maybeSource));
+          } else {
+            console.warn(
+              "Unknown content part type for assistant",
+              contentPart,
+            );
+          }
+        }
+      }
+    }
+
+    // Add source parts
+    for (const source of message.sources ?? []) {
+      allParts.push(toSourcePart(source));
+    }
+  }
+
+  return {
+    ...common,
+    role: "assistant",
+    text: joinText(allParts),
+    status,
+    parts: allParts,
+    metadata: group.find((m) => m.metadata)?.metadata,
+  };
+}
+
+function toSourcePart(
+  part: SourcePart | Infer<typeof vSource>,
+): SourceUrlUIPart | SourceDocumentUIPart {
+  if (part.sourceType === "url") {
+    return {
+      type: "source-url",
+      url: part.url,
+      sourceId: part.id,
+      providerMetadata: part.providerMetadata,
+      title: part.title,
+    } satisfies SourceUrlUIPart;
+  }
+  return {
+    type: "source-document",
+    mediaType: part.mediaType,
+    sourceId: part.id,
+    title: part.title,
+    filename: part.filename,
+    providerMetadata: part.providerMetadata,
+  } satisfies SourceDocumentUIPart;
+}
+
+export function combineUIMessages(messages: UIMessage[]): UIMessage[] {
+  const combined = messages.reduce((acc, message) => {
+    if (!acc.length) {
+      return [message];
+    }
+    const previous = acc.at(-1)!;
+    if (
+      message.order !== previous.order ||
+      previous.role !== message.role ||
+      message.role !== "assistant"
+    ) {
+      acc.push(message);
+      return acc;
+    }
+    // We will replace it with a combined message
+    acc.pop();
+    const newParts = [...previous.parts];
+    for (const part of message.parts) {
+      const toolCallId = getToolCallId(part);
+      if (!toolCallId) {
+        newParts.push(part);
+        continue;
+      }
+      const previousPartIndex = newParts.findIndex(
+        (p) => getToolCallId(p) === toolCallId,
+      );
+      if (previousPartIndex === -1) {
+        // Tool call not found in previous parts, add it as new
+        newParts.push(part);
+        continue;
+      }
+      const previousPart = newParts.splice(previousPartIndex, 1)[0];
+      newParts.push(mergeParts(previousPart, part));
+    }
+    acc.push({
+      ...previous,
+      ...pick(message, ["status", "metadata", "agentName"]),
+      parts: newParts,
+      text: joinText(newParts),
+    });
+    return acc;
+  }, [] as UIMessage[]);
+  return combined;
+}
+
+function getToolCallId(
+  part: UIMessage["parts"][number] & { toolCallId?: string },
+) {
+  return part.toolCallId;
+}
+
+function mergeParts(
+  previousPart: UIMessage["parts"][number],
+  part: UIMessage["parts"][number],
+): UIMessage["parts"][number] {
+  const merged: Record<string, unknown> = { ...previousPart };
+  for (const [key, value] of Object.entries(part)) {
+    if (value !== undefined) {
+      merged[key] = value;
+    }
+  }
+  return merged as ToolUIPart | DynamicToolUIPart;
+}