@convex-dev/agent
Version:
A agent component for Convex.
141 lines • 6.16 kB
JavaScript
"use client";
import { omit, } from "convex-helpers";
import { usePaginatedQuery } from "convex-helpers/react";
import {} from "convex/react";
import { useMemo } from "react";
import { sorted } from "../shared.js";
import { fromUIMessages } from "../UIMessages.js";
import { useStreamingUIMessages } from "./useStreamingUIMessages.js";
/**
* 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, args, options) {
// These are full messages
const paginated = usePaginatedQuery(query, args, { 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, !options.stream ||
args === "skip" ||
paginated.status === "LoadingFirstPage"
? "skip"
: // eslint-disable-next-line @typescript-eslint/no-explicit-any
{ ...args, paginationOpts: { cursor: null, numItems: 0 } }, { 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)).reduce((msgs, msg) => {
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;
}, []),
};
}, [paginated, streamMessages, threadId]);
return merged;
}
/**
* @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, args, options) {
const queryArgs = args === "skip"
? args
: omit(args, ["startOrder"]);
const startOrder = args === "skip" ? undefined : (args.startOrder ?? undefined);
const queryOptions = { startOrder, ...options };
const uiMessages = useStreamingUIMessages(query, queryArgs, queryOptions);
if (args === "skip") {
return undefined;
}
// TODO: we aren't passing through as much metadata as we could here.
// We could share the stream metadata logic with useStreamingUIMessages.
return uiMessages
?.map((m) => fromUIMessages([m], { threadId: args.threadId }))
.flat();
}
//# sourceMappingURL=useThreadMessages.js.map