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

package/src/react/useSmoothText.ts

@@ -2,7 +2,20 @@ import { useEffect, useRef, useState } from "react";
 
 const FPS = 20;
 const MS_PER_FRAME = 1000 / FPS;
-const MAX_TIME_JUMP_MS = 250;
+const INITIAL_CHARS_PER_SEC = 128;
+
+export type SmoothTextOptions = {
+  /**
+   * The number of characters to display per second.
+   */
+  charsPerSec?: number;
+  /**
+   * Whether to initially start streaming.
+   * If this later turns to false, it'll continue streaming.
+   * This will start streaming the first value it sees.
+   */
+  startStreaming?: boolean;
+};
 /**
  * A hook that smoothly displays text as it is streamed.
  *
@@ -16,28 +29,20 @@ const MAX_TIME_JUMP_MS = 250;
 export function useSmoothText(
   text: string,
   {
-    charsPerSec = 256,
+    charsPerSec = INITIAL_CHARS_PER_SEC,
     startStreaming = false,
-  }: {
-    /**
-     * The number of characters to display per second.
-     */
-    charsPerSec?: number;
-    /**
-     * Whether to initially start streaming.
-     * If this later turns to false, it'll continue streaming.
-     * This will start streaming the first value it sees.
-     */
-    startStreaming?: boolean;
-  } = {},
+  }: SmoothTextOptions = {},
 ): [string, { cursor: number; isStreaming: boolean }] {
-  const [visibleText, setVisibleText] = useState(startStreaming ? "" : text);
+  const [visibleText, setVisibleText] = useState(
+    startStreaming ? "" : text || "",
+  );
   const smoothState = useRef({
-    tick: Date.now() + (visibleText.length * 1000) / charsPerSec,
+    tick: Date.now(),
     cursor: visibleText.length,
-    start: Date.now(),
-    initialLength: visibleText.length,
+    lastUpdate: Date.now(),
+    lastUpdateLength: text.length,
     charsPerMs: charsPerSec / 1000,
+    initial: true,
   });
 
   const isStreaming = smoothState.current.cursor < text.length;
@@ -46,36 +51,51 @@ export function useSmoothText(
     if (!isStreaming) {
       return;
     }
-    const latestCharsPerMs =
-      (text.length - smoothState.current.initialLength) /
-      (Date.now() - smoothState.current.start);
-    // Smooth out the charsPerSec by averaging it with the previous value.
-    smoothState.current.charsPerMs = Math.min(
-      (2 * latestCharsPerMs + smoothState.current.charsPerMs) / 3,
-      smoothState.current.charsPerMs * 2,
-    );
+    if (smoothState.current.lastUpdateLength !== text.length) {
+      const timeSinceLastUpdate = Date.now() - smoothState.current.lastUpdate;
+      const latestCharsPerMs =
+        (text.length - smoothState.current.lastUpdateLength) /
+        timeSinceLastUpdate;
+      // Is the rate increasing?
+      const rateError = latestCharsPerMs - smoothState.current.charsPerMs;
+      // Is our visible text falling behind what it could show?
+      const charLag =
+        smoothState.current.lastUpdateLength - smoothState.current.cursor;
+      const lagRate = charLag / timeSinceLastUpdate;
+      const charsPerMs =
+        latestCharsPerMs +
+        (smoothState.current.initial
+          ? 0
+          : Math.max(0, (rateError + lagRate) / 2));
+      smoothState.current.initial = false;
+      // Smooth out the charsPerSec by weighting it with the previous value.
+      smoothState.current.charsPerMs = Math.min(
+        (2 * charsPerMs + smoothState.current.charsPerMs) / 3,
+        smoothState.current.charsPerMs * 2,
+      );
+    }
     smoothState.current.tick = Math.max(
       smoothState.current.tick,
-      Date.now() - 2 * MS_PER_FRAME,
+      Date.now() - MS_PER_FRAME,
     );
+    smoothState.current.lastUpdate = Date.now();
+    smoothState.current.lastUpdateLength = text.length;
 
     function update() {
       if (smoothState.current.cursor >= text.length) {
         return;
       }
       const now = Date.now();
-      const timeSinceLastUpdate = Math.min(
-        MAX_TIME_JUMP_MS,
-        now - smoothState.current.tick,
-      );
-      const chars = Math.floor(
+      const timeSinceLastUpdate = now - smoothState.current.tick;
+      const charsSinceLastUpdate = Math.floor(
         timeSinceLastUpdate * smoothState.current.charsPerMs,
       );
-      smoothState.current.cursor = Math.min(
-        smoothState.current.cursor + chars,
-        text.length,
+      const chars = Math.min(
+        charsSinceLastUpdate,
+        text.length - smoothState.current.cursor,
       );
-      smoothState.current.tick = now;
+      smoothState.current.cursor += chars;
+      smoothState.current.tick += chars / smoothState.current.charsPerMs;
       setVisibleText(text.slice(0, smoothState.current.cursor));
     }
     update();

package/src/react/useStreamingUIMessages.ts

@@ -0,0 +1,143 @@
+"use client";
+import { useMemo, useState, useEffect } from "react";
+import { type UIDataTypes, type UIMessageChunk, type UITools } from "ai";
+import type { StreamQuery, StreamQueryArgs } from "./types.js";
+import { type UIMessage } from "../UIMessages.js";
+import {
+  blankUIMessage,
+  getParts,
+  updateFromUIMessageChunks,
+  deriveUIMessagesFromTextStreamParts,
+} from "../deltas.js";
+import { useDeltaStreams } from "./useDeltaStreams.js";
+
+// Polyfill structuredClone to support readUIMessageStream on ReactNative
+if (!("structuredClone" in globalThis)) {
+  void import("@ungap/structured-clone" as any).then(
+    ({ default: structuredClone }) =>
+      (globalThis.structuredClone = structuredClone),
+  );
+}
+
+/**
+ * A hook that fetches streaming messages from a thread and converts them to UIMessages
+ * using AI SDK's readUIMessageStream.
+ * This ONLY returns streaming UIMessages. To get both full and streaming messages,
+ * use `useUIMessages`.
+ *
+ * @param query The query to use to fetch messages.
+ * It must take as arguments `{ threadId, paginationOpts, streamArgs }` and
+ * return a `streams` object returned from `agent.syncStreams`.
+ * @param args The arguments to pass to the query other than `paginationOpts`
+ * and `streamArgs`. So `{ threadId }` at minimum, plus any other arguments that
+ * you want to pass to the query.
+ * @returns The streaming UIMessages.
+ */
+export function useStreamingUIMessages<
+  METADATA = unknown,
+  DATA_PARTS extends UIDataTypes = UIDataTypes,
+  TOOLS extends UITools = UITools,
+  Query extends StreamQuery<any> = StreamQuery<object>,
+>(
+  query: Query,
+  args: StreamQueryArgs<Query> | "skip",
+  options?: {
+    startOrder?: number;
+    skipStreamIds?: string[];
+  },
+  // TODO: make generic on metadata, etc.
+): UIMessage<METADATA, DATA_PARTS, TOOLS>[] | undefined {
+  const [messageState, setMessageState] = useState<
+    Record<
+      string,
+      {
+        uiMessage: UIMessage<METADATA, DATA_PARTS, TOOLS>;
+        cursor: number;
+      }
+    >
+  >({});
+
+  const streams = useDeltaStreams(query, args, options);
+
+  const threadId = args === "skip" ? undefined : args.threadId;
+
+  useEffect(() => {
+    if (!streams) return;
+    // return if there are no new deltas beyond the cursors
+    let noNewDeltas = true;
+    for (const stream of streams) {
+      const lastDelta = stream.deltas.at(-1);
+      const cursor = messageState[stream.streamMessage.streamId]?.cursor;
+      if (!cursor) {
+        noNewDeltas = false;
+        break;
+      }
+      if (lastDelta && lastDelta.start >= cursor) {
+        noNewDeltas = false;
+        break;
+      }
+    }
+    if (noNewDeltas) {
+      return;
+    }
+    const abortController = new AbortController();
+    void (async () => {
+      const newMessageState: Record<
+        string,
+        {
+          uiMessage: UIMessage<METADATA, DATA_PARTS, TOOLS>;
+          cursor: number;
+        }
+      > = Object.fromEntries(
+        await Promise.all(
+          streams.map(async ({ deltas, streamMessage }) => {
+            const { parts, cursor } = getParts<UIMessageChunk>(deltas, 0);
+            if (streamMessage.format === "UIMessageChunk") {
+              // Unfortunately this can't handle resuming from a UIMessage and
+              // adding more chunks, so we re-create it from scratch each time.
+              const uiMessage = await updateFromUIMessageChunks(
+                blankUIMessage(streamMessage, threadId),
+                parts,
+              );
+              return [
+                streamMessage.streamId,
+                {
+                  uiMessage,
+                  cursor,
+                },
+              ];
+            } else {
+              const [uiMessages] = deriveUIMessagesFromTextStreamParts(
+                threadId,
+                [streamMessage],
+                [],
+                deltas,
+              );
+              return [
+                streamMessage.streamId,
+                {
+                  uiMessage: uiMessages[0],
+                  cursor,
+                },
+              ];
+            }
+          }),
+        ),
+      );
+      if (abortController.signal.aborted) return;
+      setMessageState(newMessageState);
+    })();
+    return () => {
+      abortController.abort();
+    };
+  }, [messageState, streams, threadId]);
+
+  return useMemo(() => {
+    if (!streams) return undefined;
+    return streams
+      .map(
+        ({ streamMessage }) => messageState[streamMessage.streamId]?.uiMessage,
+      )
+      .filter((uiMessage) => uiMessage !== undefined);
+  }, [messageState, streams]);
+}

package/src/react/useThreadMessages.ts

@@ -0,0 +1,262 @@
+"use client";
+import {
+  omit,
+  type BetterOmit,
+  type ErrorMessage,
+  type Expand,
+} from "convex-helpers";
+import { usePaginatedQuery } from "convex-helpers/react";
+import {
+  type PaginatedQueryArgs,
+  type UsePaginatedQueryResult,
+} from "convex/react";
+import type {
+  FunctionArgs,
+  FunctionReference,
+  PaginationOptions,
+  PaginationResult,
+} from "convex/server";
+import { useMemo, useState, useEffect, useRef } from "react";
+import type { SyncStreamsReturnValue } from "../client/types.js";
+import { sorted } from "../shared.js";
+import { fromUIMessages } from "../UIMessages.js";
+import type {
+  Message,
+  MessageDoc,
+  MessageStatus,
+  StreamArgs,
+} from "../validators.js";
+import type { StreamQueryArgs, StreamQuery } from "./types.js";
+import { useStreamingUIMessages } from "./useStreamingUIMessages.js";
+
+export type MessageDocLike = {
+  order: number;
+  stepOrder: number;
+  status: MessageStatus | "streaming";
+  message?: Message;
+};
+
+export type ThreadMessagesQuery<
+  Args = unknown,
+  M extends MessageDocLike = MessageDocLike,
+> = FunctionReference<
+  "query",
+  "public",
+  {
+    threadId: string;
+    paginationOpts: PaginationOptions;
+    /**
+     * 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 }
+>;
+
+export type ThreadMessagesArgs<
+  Query extends ThreadMessagesQuery<unknown, MessageDocLike>,
+> =
+  Query extends ThreadMessagesQuery<unknown, MessageDocLike>
+    ? Expand<BetterOmit<FunctionArgs<Query>, "paginationOpts" | "streamArgs">>
+    : never;
+
+export type ThreadMessagesResult<
+  Query extends ThreadMessagesQuery<unknown, MessageDocLike>,
+> = Query extends ThreadMessagesQuery<unknown, infer M> ? M : never;
+
+/**
+ * A hook that fetches messages from a thread.
+ *
+ * This hook is a wrapper around `usePaginatedQuery` and `useStreamingThreadMessages`.
+ * It will fetch both full messages and streaming messages, and merge them together.
+ *
+ * The query must take as arguments `{ threadId, paginationOpts }` and return a
+ * pagination result of objects that extend `MessageDoc`.
+ *
+ * For streaming, it should look like this:
+ * ```ts
+ * export const listThreadMessages = query({
+ *   args: {
+ *     threadId: v.string(),
+ *     paginationOpts: paginationOptsValidator,
+ *     streamArgs: vStreamArgs,
+ *     ... other arguments you want
+ *   },
+ *   handler: async (ctx, args) => {
+ *     // await authorizeThreadAccess(ctx, threadId);
+ *     // NOTE: listMessages returns MessageDocs, not UIMessages.
+ *     const paginated = await listMessages(ctx, components.agent, args);
+ *     const streams = await syncStreams(ctx, components.agent, args);
+ *     // Here you could filter out / modify the documents & stream deltas.
+ *     return { ...paginated, streams };
+ *   },
+ * });
+ * ```
+ *
+ * Then the hook can be used like this:
+ * ```ts
+ * const { results, status, loadMore } = useThreadMessages(
+ *   api.myModule.listThreadMessages,
+ *   { threadId },
+ *   { initialNumItems: 10, stream: true }
+ * );
+ * ```
+ *
+ * @param query The query to use to fetch messages.
+ * It must take as arguments `{ threadId, paginationOpts }` and return a
+ * pagination result of objects that extend `MessageDoc`.
+ * To support streaming, it must also take in `streamArgs: vStreamArgs` and
+ * return a `streams` object returned from `agent.syncStreams`.
+ * @param args The arguments to pass to the query other than `paginationOpts`
+ * and `streamArgs`. So `{ threadId }` at minimum, plus any other arguments that
+ * you want to pass to the query.
+ * @param options The options for the query. Similar to usePaginatedQuery.
+ * To enable streaming, pass `stream: true`.
+ * @returns The messages. If stream is true, it will return a list of messages
+ *   that includes both full messages and streaming messages.
+ */
+export function useThreadMessages<Query extends ThreadMessagesQuery<any, any>>(
+  query: Query,
+  args: ThreadMessagesArgs<Query> | "skip",
+  options: {
+    initialNumItems: number;
+    stream?: Query extends StreamQuery
+      ? boolean
+      : ErrorMessage<"To enable streaming, your query must take in streamArgs: vStreamArgs and return a streams object returned from syncStreams. See docs.">;
+  },
+): UsePaginatedQueryResult<
+  ThreadMessagesResult<Query> & { streaming: boolean; key: string }
+> {
+  // These are full messages
+  const paginated = usePaginatedQuery(
+    query,
+    args as PaginatedQueryArgs<Query> | "skip",
+    { initialNumItems: options.initialNumItems },
+  );
+
+  let startOrder = paginated.results.at(-1)?.order ?? 0;
+  for (let i = paginated.results.length - 1; i >= 0; i--) {
+    const m = paginated.results[i];
+    if (!m.streaming && m.status === "pending") {
+      // round down to the nearest 10 for some cache benefits
+      startOrder = m.order - (m.order % 10);
+      break;
+    }
+  }
+  // These are streaming messages that will not include full messages.
+  const streamMessages = useStreamingThreadMessages(
+    query as StreamQuery<ThreadMessagesArgs<Query>>,
+    !options.stream ||
+      args === "skip" ||
+      paginated.status === "LoadingFirstPage"
+      ? "skip"
+      : ({ ...args, paginationOpts: { cursor: null, numItems: 0 } } as any),
+    { startOrder },
+  );
+
+  const threadId = args === "skip" ? undefined : args.threadId;
+
+  const merged = useMemo(() => {
+    const streamListMessages =
+      streamMessages?.map((m) => ({
+        ...m,
+        streaming: !m.status || m.status === "pending",
+      })) ?? [];
+    return {
+      ...paginated,
+      results: sorted(
+        paginated.results
+          .map((m) => ({ ...m, streaming: false }))
+          // Note: this is intentionally after paginated results.
+          .concat(streamListMessages) as (MessageDocLike & {
+          streaming: boolean;
+          key: string;
+        })[],
+      ).reduce(
+        (msgs, msg: MessageDocLike & { streaming: boolean; key: string }) => {
+          msg.key = `${threadId}-${msg.order}-${msg.stepOrder}`;
+          const last = msgs.at(-1);
+          if (!last) {
+            return [msg];
+          }
+          if (last.order !== msg.order || last.stepOrder !== msg.stepOrder) {
+            return [...msgs, msg];
+          }
+          if (
+            last.status === "pending" &&
+            (msg.streaming || msg.status !== "pending")
+          ) {
+            // Let's prefer a streaming or finalized message over a pending
+            // one.
+            return [...msgs.slice(0, -1), msg];
+          }
+          // skip the new one if the previous one (listed) was finalized
+          return msgs;
+        },
+        [] as (MessageDocLike & { streaming: boolean; key: string })[],
+      ),
+    };
+  }, [paginated, streamMessages, threadId]);
+
+  return merged as ThreadMessagesResult<Query> & {
+    key: string;
+    streaming: boolean;
+  };
+}
+
+/**
+ * @deprecated FYI `useStreamingUIMessages` is likely better for you.
+ * A hook that fetches streaming messages from a thread.
+ * This ONLY returns streaming messages. To get both, use `useThreadMessages`.
+ *
+ * @param query The query to use to fetch messages.
+ * It must take as arguments `{ threadId, paginationOpts, streamArgs }` and
+ * return a `streams` object returned from `agent.syncStreams`.
+ * @param args The arguments to pass to the query other than `paginationOpts`
+ * and `streamArgs`. So `{ threadId }` at minimum, plus any other arguments that
+ * you want to pass to the query.
+ * @returns The streaming messages.
+ */
+export function useStreamingThreadMessages<Query extends StreamQuery<any>>(
+  query: Query,
+  args:
+    | (StreamQueryArgs<Query> & {
+        /** @deprecated Pass startOrder to the next argument (third argument). */
+        startOrder?: number;
+      })
+    | "skip",
+  options?: {
+    startOrder?: number;
+    skipStreamIds?: string[];
+  },
+): Array<MessageDoc> | undefined {
+  const queryArgs =
+    args === "skip"
+      ? args
+      : (omit(args, ["startOrder"]) as unknown as StreamQueryArgs<Query>);
+  const startOrder =
+    args === "skip" ? undefined : (args.startOrder ?? undefined);
+  const queryOptions = { startOrder, ...options };
+  const uiMessages = useStreamingUIMessages(query, queryArgs, queryOptions);
+  const [messages, setMessages] = useState<Array<MessageDoc> | undefined>();
+  const generationRef = useRef(0);
+
+  useEffect(() => {
+    if (args === "skip" || !uiMessages) {
+      setMessages(undefined);
+      return;
+    }
+    const currentGeneration = ++generationRef.current;
+    (async () => {
+      const nested = await Promise.all(
+        uiMessages.map((m) => fromUIMessages([m], { threadId: args.threadId })),
+      );
+      if (generationRef.current === currentGeneration) {
+        setMessages(nested.flat());
+      }
+    })();
+  }, [uiMessages, args === "skip" ? undefined : args.threadId]);
+
+  return messages;
+}