@convex-dev/agent 0.1.3 → 0.1.4

package/src/react/index.ts

@@ -1,31 +1,27 @@
-import type { TextPart, ToolCallPart, ToolResultPart } from "ai";
-import type { BetterOmit, ErrorMessage, Expand } from "convex-helpers";
+import type { ErrorMessage } from "convex-helpers";
 import {
-  insertAtTop,
   type PaginatedQueryArgs,
   usePaginatedQuery,
   type UsePaginatedQueryResult,
   useQuery,
 } from "convex/react";
-import type {
-  FunctionArgs,
-  FunctionReference,
-  PaginationOptions,
-  PaginationResult,
-} from "convex/server";
-import { useEffect, useMemo, useRef, useState } from "react";
+import type { FunctionArgs } from "convex/server";
+import { useMemo, useState } from "react";
 import type { MessageDoc } from "../client";
 import type { SyncStreamsReturnValue } from "../client/types";
-import type {
-  StreamArgs,
-  StreamDelta,
-  StreamMessage,
-  TextStreamPart,
-} from "../validators";
+import type { StreamArgs } from "../validators";
 import type { UIMessage } from "./toUIMessages";
 import { toUIMessages } from "./toUIMessages";
-import { OptimisticLocalStore } from "convex/browser";
+import { mergeDeltas } from "./deltas";
+import type {
+  ThreadQuery,
+  ThreadStreamQuery,
+  ThreadMessagesArgs,
+  ThreadMessagesResult,
+} from "./types";
 
+export { optimisticallySendMessage } from "./optimisticallySendMessage";
+export { useSmoothText } from "./useSmoothText";
 export { toUIMessages, type UIMessage };
 
 /**
@@ -102,8 +98,10 @@ export function useThreadMessages<
 
   // These are streaming messages that will not include full messages.
   const streamMessages = useStreamingThreadMessages(
-    // eslint-disable-next-line @typescript-eslint/no-explicit-any
-    query as ThreadStreamQuery<any, any>,
+    query as ThreadStreamQuery<
+      ThreadMessagesArgs<Query>,
+      ThreadMessagesResult<Query>
+    >,
     !options.stream ? "skip" : args
   );
 
@@ -157,7 +155,7 @@ export function useStreamingThreadMessages<
   // cursor. There can be multiple messages in the same stream, e.g. for tool
   // calls.
   const [streams, setStreams] = useState<
-    Array<{ stream: StreamMessage; cursor: number; messages: MessageDoc[] }>
+    Array<{ streamId: string; cursor: number; messages: MessageDoc[] }>
   >([]);
   // Get all the active streams
   const streamList = useQuery(
@@ -179,8 +177,7 @@ export function useStreamingThreadMessages<
       throw new Error("Expected list streams");
     }
     return streamList.streams.messages.map(({ streamId }) => {
-      const stream = streams.find((s) => s.stream.streamId === streamId);
-      // Because of the invariant, we can just take the last chunk's end.
+      const stream = streams.find((s) => s.streamId === streamId);
       const cursor = stream?.cursor ?? 0;
       return { streamId, cursor };
     });
@@ -199,21 +196,22 @@ export function useStreamingThreadMessages<
     | { streams: Extract<SyncStreamsReturnValue, { kind: "deltas" }> }
     | undefined;
   // Merge any deltas into the streamChunks, keeping it unmodified if unchanged.
+  const threadId = args === "skip" ? undefined : args.threadId;
   const [messages, newStreams, changed] = useMemo(() => {
-    if (args === "skip") return [undefined, streams, false];
-    if (!streamList) return [undefined, streams, false];
+    if (!threadId) return [undefined, [], false];
+    if (!streamList) return [undefined, [], false];
     if (cursorQuery && cursorQuery.streams?.kind !== "deltas") {
       throw new Error("Expected deltas streams");
     }
     return mergeDeltas(
-      args.threadId,
+      threadId,
       streamList.streams.messages,
       streams,
       cursorQuery?.streams?.deltas ?? []
     );
-  }, [cursorQuery, streams, streamList]);
+  }, [threadId, cursorQuery, streams, streamList]);
   // Now assemble the chunks into messages
-  if (args === "skip") {
+  if (!threadId) {
     return undefined;
   }
   if (changed) {
@@ -222,567 +220,6 @@ export function useStreamingThreadMessages<
   return messages as ThreadMessagesResult<Query>[] | undefined;
 }
 
-/**
- * A hook that smoothly displays text as it is streamed.
- *
- * @param text The text to display. Pass in the full text each time.
- * @param charsPerSec The number of characters to display per second.
- * @returns A tuple of the visible text and the state of the smooth text,
- * including the current cursor position and whether it's still streaming.
- * This allows you to decide if it's too far behind and you want to adjust
- * the charsPerSec or just prefer the full text.
- */
-export function useSmoothText(
-  text: string,
-  {
-    charsPerSec = 512,
-  }: {
-    /**
-     * The number of characters to display per second.
-     */
-    charsPerSec?: number;
-  } = {}
-): [string, { cursor: number; isStreaming: boolean }] {
-  const [visibleText, setVisibleText] = useState(text);
-  const smoothState = useRef({ lastUpdated: Date.now(), cursor: text.length });
-
-  const isStreaming = smoothState.current.cursor < text.length;
-
-  useEffect(() => {
-    if (!isStreaming) {
-      return;
-    }
-    function update() {
-      if (smoothState.current.cursor >= text.length) {
-        return;
-      }
-      const now = Date.now();
-      const timeSinceLastUpdate = now - smoothState.current.lastUpdated;
-      const chars = Math.floor((timeSinceLastUpdate * charsPerSec) / 1000);
-      smoothState.current.cursor = Math.min(
-        smoothState.current.cursor + chars,
-        text.length
-      );
-      smoothState.current.lastUpdated = now;
-      setVisibleText(text.slice(0, smoothState.current.cursor));
-    }
-    update();
-    const interval = setInterval(() => {
-      update();
-    }, 1000 / 60);
-    return () => clearInterval(interval);
-  }, [text, isStreaming, charsPerSec]);
-
-  return [visibleText, { cursor: smoothState.current.cursor, isStreaming }];
-}
-
-export function optimisticallySendMessage(
-  query: ThreadQuery<unknown, MessageDoc>
-): (
-  store: OptimisticLocalStore,
-  args: { threadId: string; prompt: string }
-) => void {
-  return (store, args) => {
-    const queries = store.getAllQueries(query);
-    let maxOrder = 0;
-    let maxStepOrder = 0;
-    for (const q of queries) {
-      if (q.args?.threadId !== args.threadId) continue;
-      if (q.args.streamArgs) continue;
-      for (const m of q.value?.page ?? []) {
-        maxOrder = Math.max(maxOrder, m.order);
-        maxStepOrder = Math.max(maxStepOrder, m.stepOrder);
-      }
-    }
-    const order = maxOrder + 1;
-    const stepOrder = 0;
-    insertAtTop({
-      paginatedQuery: query,
-      argsToMatch: { threadId: args.threadId, streamArgs: undefined },
-      item: {
-        _creationTime: Date.now(),
-        _id: crypto.randomUUID(),
-        order,
-        stepOrder,
-        status: "pending",
-        threadId: args.threadId,
-        tool: false,
-        message: {
-          role: "user",
-          content: args.prompt,
-        },
-        text: args.prompt,
-      },
-      localQueryStore: store,
-    });
-  };
-}
-
-function mergeDeltas(
-  threadId: string,
-  streamMessages: StreamMessage[],
-  existingStreams: Array<{
-    stream: StreamMessage;
-    cursor: number;
-    messages: MessageDoc[];
-  }>,
-  allDeltas: StreamDelta[]
-): [
-  MessageDoc[],
-  Array<{ stream: StreamMessage; cursor: number; messages: MessageDoc[] }>,
-  boolean,
-] {
-  const newStreams: Array<{
-    stream: StreamMessage;
-    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)
-      .sort((a, b) => a.start - b.start);
-    const existing = existingStreams.find(
-      (s) => s.stream.streamId === streamMessage.streamId
-    );
-    const [newStream, messageChanged] = applyDeltasToStreamMessages(
-      threadId,
-      streamMessage,
-      existing,
-      deltas
-    );
-    newStreams.push(newStream);
-    if (messageChanged) changed = true;
-  }
-  for (const { stream } of existingStreams) {
-    if (!newStreams.find((s) => s.stream.streamId === stream.streamId)) {
-      // There's a stream that's no longer active.
-      changed = true;
-    }
-  }
-  const messages = newStreams
-    .sort(
-      (a, b) =>
-        a.stream.order - b.stream.order ||
-        a.stream.stepOrder - b.stream.stepOrder
-    )
-    .map((s) => s.messages)
-    .flat();
-  return [messages, newStreams, changed];
-}
-
-function applyDeltasToStreamMessages(
-  threadId: string,
-  streamMessage: StreamMessage,
-  existing:
-    | { stream: StreamMessage; cursor: number; messages: MessageDoc[] }
-    | undefined,
-  deltas: StreamDelta[]
-): [
-  { stream: StreamMessage; cursor: number; messages: MessageDoc[] },
-  boolean,
-] {
-  let changed = false;
-  const newStream = {
-    stream: streamMessage,
-    cursor: existing?.cursor ?? 0,
-    messages: existing?.messages ?? [],
-  };
-  let parts: TextStreamPart[] = [];
-  for (const delta of deltas) {
-    if (delta.parts.length === 0) {
-      console.warn(`Got delta for stream ${delta.streamId} with no parts`);
-      continue;
-    }
-    if (newStream.cursor !== delta.start) {
-      if (newStream.cursor >= delta.end) {
-        console.debug(
-          `Got duplicate delta for stream ${delta.streamId} at ${delta.start}`
-        );
-        continue;
-      } else if (newStream.cursor < delta.start) {
-        console.warn(
-          `Got delta for stream ${delta.streamId} that has a gap ${newStream.cursor} -> ${delta.start}`
-        );
-        continue;
-      } else {
-        throw new Error(
-          `Got unexpected delta for stream ${delta.streamId}: delta: ${delta.start} -> ${delta.end} existing cursor: ${newStream.cursor}`
-        );
-      }
-    }
-    changed = true;
-    newStream.cursor = delta.end;
-    parts.push(...delta.parts);
-  }
-  if (!changed) {
-    return [existing ?? newStream, false];
-  }
-
-  if (!newStream.messages.at(-1)) {
-    newStream.messages.push(
-      createStreamingMessage(
-        threadId,
-        streamMessage,
-        parts[0]!,
-        newStream.messages.length
-      )
-    );
-    parts = parts.slice(1);
-  }
-  let currentMessage = newStream.messages.at(-1)!;
-  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);
-    }
-    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 getLastContent(message: MessageDoc) {
-  if (Array.isArray(message.message?.content)) {
-    return message.message.content.at(-1);
-  }
-  return undefined;
-}
-
-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)}`);
-  }
-}
-
-type ThreadQuery<
-  Args = unknown,
-  M extends MessageDoc = MessageDoc,
-> = FunctionReference<
-  "query",
-  "public",
-  {
-    threadId: string;
-    paginationOpts: PaginationOptions;
-    // TODO: will this allow passing a function that doesn't have this param?
-    /**
-     * If { stream: true } is passed, it will also query for stream deltas.
-     * In order for this to work, the query must take as an argument streamArgs.
-     */
-    streamArgs?: StreamArgs;
-  } & Args,
-  PaginationResult<M> & { streams?: SyncStreamsReturnValue }
->;
-
-type ThreadStreamQuery<
-  Args = Record<string, unknown>,
-  M extends MessageDoc = MessageDoc,
-> = FunctionReference<
-  "query",
-  "public",
-  {
-    threadId: string;
-    paginationOpts: PaginationOptions;
-    streamArgs?: StreamArgs; // required for stream query
-  } & Args,
-  PaginationResult<M> & { streams: SyncStreamsReturnValue }
->;
-
-type ThreadMessagesArgs<Query extends ThreadQuery<unknown, MessageDoc>> =
-  Query extends ThreadQuery<unknown, MessageDoc>
-    ? Expand<BetterOmit<FunctionArgs<Query>, "paginationOpts" | "streamArgs">>
-    : never;
-
-type ThreadMessagesResult<Query extends ThreadQuery<unknown, MessageDoc>> =
-  Query extends ThreadQuery<unknown, infer M> ? M : never;
-
-// TODO: pass in the messages we need to watch? that way it can be consistent..
-
-// export function streamMessagesToUIMessages(
-//   messages: StreamsChunk[],
-//   keyOrder: string[]
-// ): UIMessageOrdered[] {
-//   const uiMessagesByMessageId: Record<string, UIMessageOrdered[]> = {};
-//   for (const message of messages) {
-//     if (message.deltas.length === 0) {
-//       continue;
-//     }
-//     if (!uiMessagesByMessageId[message.key]) {
-//       uiMessagesByMessageId[message.key] = [];
-//     }
-//     if (uiMessagesByMessageId[message.key].length === 0) {
-//       uiMessagesByMessageId[message.key] = [
-//         createUIMessageFromPart(message.deltas[0], {
-//           id: message.key,
-//           createdAt: new Date(),
-//         }),
-//       ];
-//     }
-//     const currentMessage = uiMessagesByMessageId[message.key].at(-1)!;
-//     const lastPart = currentMessage.parts.at(-1);
-//     for (const delta of message.deltas) {
-//       switch (delta.type) {
-//         case "text-delta":
-//           currentMessage.content += delta.textDelta;
-//           if (lastPart?.type === "text") {
-//             lastPart.text += delta.textDelta;
-//           } else {
-//             currentMessage.parts.push({
-//               type: "text",
-//               text: delta.textDelta,
-//             });
-//           }
-//           break;
-//         case "tool-call-delta": {
-//           let lastToolInvocation: ToolInvocationUIPart | undefined;
-//           for (let i = currentMessage.parts.length - 1; i >= 0; i--) {
-//             const part = currentMessage.parts[i];
-//             if (
-//               part.type === "tool-invocation" &&
-//               part.toolInvocation.state === "partial-call"
-//             ) {
-//               lastToolInvocation = part;
-//               break;
-//             }
-//           }
-//           if (lastToolInvocation) {
-//             lastToolInvocation.toolInvocation = {
-//               ...lastToolInvocation.toolInvocation,
-//               state: "call",
-//             };
-//             break;
-//           }
-//           console.error(
-//             `Received a tool call delta without a previous tool invocation: ${JSON.stringify(currentMessage.parts)}, creating one anyways...`
-//           );
-//         }
-//         // fallthrough
-//         case "tool-call-streaming-start":
-//           currentMessage.parts.push({
-//             type: "tool-invocation",
-//             toolInvocation: {
-//               toolCallId: delta.toolCallId,
-//               toolName: delta.toolName,
-//               args: {},
-//               state: "partial-call",
-//               step:
-//                 currentMessage.parts.filter(
-//                   (part) => part.type === "tool-invocation"
-//                 ).length + 1,
-//             },
-//           });
-//           break;
-//         case "tool-call":
-//           currentMessage.parts.push({
-//             type: "tool-invocation",
-//             toolInvocation: {
-//               toolCallId: delta.toolCallId,
-//               toolName: delta.toolName,
-//               args: delta.args,
-//               state: "call",
-//               step:
-//                 currentMessage.parts.filter(
-//                   (part) => part.type === "tool-invocation"
-//                 ).length + 1,
-//             },
-//           });
-//           break;
-//         case "reasoning":
-//           if (lastPart?.type === "reasoning") {
-//             lastPart.reasoning += delta.textDelta;
-//           } else {
-//             currentMessage.parts.push({
-//               type: "reasoning",
-//               reasoning: delta.textDelta,
-//               details: [],
-//             });
-//           }
-//           break;
-//         case "source":
-//           currentMessage.parts.push({
-//             type: "source",
-//             source: delta.source,
-//           });
-//           break;
-//         default:
-//           console.warn(`Received unexpected part: ${JSON.stringify(delta)}`);
-//           break;
-//       }
-//     }
-//   }
-//   return keyOrder.map((key) => uiMessagesByMessageId[key]).flat();
-// }
-
 if (typeof window === "undefined") {
   throw new Error("this is frontend code, but it's running somewhere else!");
 }
@@ -849,7 +286,7 @@ export function useStreamingText(
         setLoading(false);
       }
     },
-    [threadId, token]
+    [threadId, token, url]
   );
   return [{ text, loading, error }, readStream] as const;
 }