@convex-dev/agent 0.1.4-alpha.0 → 0.1.4

package/src/react/deltas.ts

@@ -0,0 +1,340 @@
+import type { TextPart, ToolCallPart, ToolResultPart } from "ai";
+import type { MessageDoc } from "../client";
+import type {
+  Message,
+  StreamDelta,
+  StreamMessage,
+  TextStreamPart,
+} from "../validators";
+import type { UIMessage } from "./toUIMessages";
+import { toUIMessages } from "./toUIMessages";
+
+export { toUIMessages, type UIMessage };
+
+export function mergeDeltas(
+  threadId: string,
+  streamMessages: StreamMessage[],
+  existingStreams: Array<{
+    streamId: string;
+    cursor: number;
+    messages: MessageDoc[];
+  }>,
+  allDeltas: StreamDelta[]
+): [
+  MessageDoc[],
+  Array<{ streamId: string; cursor: number; messages: MessageDoc[] }>,
+  boolean,
+] {
+  const newStreams: Array<{
+    streamId: string;
+    cursor: number;
+    messages: MessageDoc[];
+  }> = [];
+  // Seed the existing chunks
+  let changed = false;
+  for (const streamMessage of streamMessages) {
+    const deltas = allDeltas.filter(
+      (d) => d.streamId === streamMessage.streamId
+    );
+    const existing = existingStreams.find(
+      (s) => s.streamId === streamMessage.streamId
+    );
+    const [newStream, messageChanged] = applyDeltasToStreamMessage(
+      threadId,
+      streamMessage,
+      existing,
+      deltas
+    );
+    newStreams.push(newStream);
+    if (messageChanged) changed = true;
+  }
+  for (const { streamId } of existingStreams) {
+    if (!newStreams.find((s) => s.streamId === streamId)) {
+      // There's a stream that's no longer active.
+      changed = true;
+    }
+  }
+  const messages = newStreams
+    .map((s) => s.messages)
+    .flat()
+    .sort((a, b) => a.order - b.order || a.stepOrder - b.stepOrder);
+  return [messages, newStreams, changed];
+}
+
+// exported for testing
+export function applyDeltasToStreamMessage(
+  threadId: string,
+  streamMessage: StreamMessage,
+  existing:
+    | { streamId: string; cursor: number; messages: MessageDoc[] }
+    | undefined,
+  deltas: StreamDelta[]
+): [{ streamId: string; cursor: number; messages: MessageDoc[] }, boolean] {
+  let changed = false;
+  let cursor = existing?.cursor ?? 0;
+  let parts: TextStreamPart[] = [];
+  for (const delta of deltas.sort((a, b) => a.start - b.start)) {
+    if (delta.parts.length === 0) {
+      console.warn(`Got delta with no parts: ${JSON.stringify(delta)}`);
+      continue;
+    }
+    if (cursor !== delta.start) {
+      if (cursor >= delta.end) {
+        console.debug(
+          `Got duplicate delta for stream ${delta.streamId} at ${delta.start}`
+        );
+        continue;
+      } else if (cursor < delta.start) {
+        console.warn(
+          `Got delta for stream ${delta.streamId} that has a gap ${cursor} -> ${delta.start}`
+        );
+        continue;
+      } else {
+        throw new Error(
+          `Got unexpected delta for stream ${delta.streamId}: delta: ${delta.start} -> ${delta.end} existing cursor: ${cursor}`
+        );
+      }
+    }
+    changed = true;
+    cursor = delta.end;
+    parts.push(...delta.parts);
+  }
+  if (!changed) {
+    return [
+      existing ?? { streamId: streamMessage.streamId, cursor, messages: [] },
+      false,
+    ];
+  }
+
+  const existingMessages = existing?.messages ?? [];
+
+  let currentMessage: MessageDoc;
+  if (existingMessages.length > 0) {
+    // replace the last message with a new one
+    const lastMessage = existingMessages.at(-1)!;
+    currentMessage = {
+      ...lastMessage,
+      message: cloneMessageAndContent(lastMessage.message),
+    };
+  } else {
+    const newMessage = createStreamingMessage(
+      threadId,
+      streamMessage,
+      parts[0]!,
+      existingMessages.length
+    );
+    parts = parts.slice(1);
+    currentMessage = newMessage;
+  }
+  const newStream = {
+    streamId: streamMessage.streamId,
+    cursor,
+    messages: [...existingMessages.slice(0, -1), currentMessage],
+  };
+  let lastContent = getLastContent(currentMessage);
+  for (const part of parts) {
+    let contentToAdd:
+      | TextPart
+      | ToolCallPart
+      | { type: "reasoning"; text: string }
+      | ToolResultPart
+      | undefined;
+    const isToolRole = part.type === "source" || part.type === "tool-result";
+    if (isToolRole !== (currentMessage.message!.role === "tool")) {
+      currentMessage = createStreamingMessage(
+        threadId,
+        streamMessage,
+        part,
+        newStream.messages.length
+      );
+      lastContent = getLastContent(currentMessage);
+      newStream.messages.push(currentMessage);
+      continue;
+    }
+    switch (part.type) {
+      case "text-delta":
+        currentMessage.text += part.textDelta;
+        if (lastContent?.type === "text") {
+          lastContent.text += part.textDelta;
+        } else {
+          contentToAdd = {
+            type: "text",
+            text: part.textDelta,
+          };
+        }
+        break;
+      case "tool-call-streaming-start":
+        currentMessage.tool = true;
+        contentToAdd = {
+          type: "tool-call",
+          toolCallId: part.toolCallId,
+          toolName: part.toolName,
+          args: "",
+        };
+        break;
+      case "tool-call-delta":
+        {
+          currentMessage.tool = true;
+          if (lastContent?.type !== "tool-call") {
+            throw new Error("Expected last content to be a tool call");
+          }
+          if (typeof lastContent.args !== "string") {
+            throw new Error("Expected args to be a string");
+          }
+          lastContent.args += part.argsTextDelta;
+        }
+        break;
+      case "tool-call":
+        currentMessage.tool = true;
+        contentToAdd = part;
+        break;
+      case "reasoning":
+        if (lastContent?.type === "reasoning") {
+          lastContent.text += part.textDelta;
+        } else {
+          contentToAdd = {
+            type: "reasoning",
+            text: part.textDelta,
+          };
+        }
+        break;
+      case "source":
+        if (!currentMessage.sources) {
+          currentMessage.sources = [];
+        }
+        currentMessage.sources.push(part.source);
+        break;
+      case "tool-result":
+        contentToAdd = part;
+        break;
+      default:
+        console.warn(`Received unexpected part: ${JSON.stringify(part)}`);
+        break;
+    }
+    if (contentToAdd) {
+      if (!currentMessage.message!.content) {
+        currentMessage.message!.content = [];
+      }
+      if (!Array.isArray(currentMessage.message?.content)) {
+        throw new Error("Expected message content to be an array");
+      }
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      currentMessage.message.content.push(contentToAdd as any);
+      lastContent = contentToAdd;
+    }
+  }
+  return [newStream, true];
+}
+
+function cloneMessageAndContent(
+  message: Message | undefined
+): Message | undefined {
+  return (
+    message &&
+    ({
+      ...message,
+      content: Array.isArray(message.content)
+        ? [...message.content]
+        : message.content,
+    } as typeof message)
+  );
+}
+
+function getLastContent(message: MessageDoc) {
+  if (Array.isArray(message.message?.content)) {
+    return message.message.content.at(-1);
+  }
+  return undefined;
+}
+
+export function createStreamingMessage(
+  threadId: string,
+  message: StreamMessage,
+  part: TextStreamPart,
+  index: number
+): MessageDoc {
+  const { streamId, ...rest } = message;
+  const metadata: MessageDoc = {
+    _id: `${streamId}-${index}`,
+    _creationTime: Date.now(),
+    status: "pending",
+    threadId,
+    tool: false,
+    ...rest,
+  };
+  switch (part.type) {
+    case "text-delta":
+      return {
+        ...metadata,
+        message: {
+          role: "assistant",
+          content: [{ type: "text", text: part.textDelta }],
+        },
+        text: part.textDelta,
+      };
+    case "tool-call-streaming-start":
+      return {
+        ...metadata,
+        tool: true,
+        message: {
+          role: "assistant",
+          content: [
+            {
+              type: "tool-call",
+              toolName: part.toolName,
+              toolCallId: part.toolCallId,
+              args: "", // when it's a string, it's a partial call
+            },
+          ],
+        },
+      };
+    case "reasoning":
+      return {
+        ...metadata,
+        message: {
+          role: "assistant",
+          content: [{ type: "reasoning", text: part.textDelta }],
+        },
+        reasoning: part.textDelta,
+      };
+    case "source":
+      console.warn("Received source part first??");
+      return {
+        ...metadata,
+        tool: true,
+        message: { role: "tool", content: [] },
+        sources: [part.source],
+      };
+    case "tool-call":
+      return {
+        ...metadata,
+        tool: true,
+        message: { role: "assistant", content: [part] },
+      };
+    case "tool-call-delta":
+      console.warn("Received tool call delta part first??");
+      return {
+        ...metadata,
+        tool: true,
+        message: {
+          role: "assistant",
+          content: [
+            {
+              type: "tool-call",
+              toolCallId: part.toolCallId,
+              toolName: part.toolName,
+              args: part.argsTextDelta,
+            },
+          ],
+        },
+      };
+    case "tool-result":
+      return {
+        ...metadata,
+        tool: true,
+        message: { role: "tool", content: [part] },
+      };
+    default:
+      throw new Error(`Unexpected part type: ${JSON.stringify(part)}`);
+  }
+}