@convex-dev/agent
Version:
A agent component for Convex.
208 lines • 6.98 kB
JavaScript
import { smoothStream, } from "ai";
import { omit } from "convex-helpers";
import { serializeTextStreamingPartsV5 } from "../parts.js";
/**
* A function that handles fetching stream deltas, used with the React hooks
* `useThreadMessages` or `useStreamingThreadMessages`.
* @param ctx A ctx object from a query, mutation, or action.
* @param component The agent component, usually `components.agent`.
* @param args.threadId The thread to sync streams for.
* @param args.streamArgs The stream arguments with per-stream cursors.
* @returns The deltas for each stream from their existing cursor.
*/
export async function syncStreams(ctx, component, args) {
if (!args.streamArgs)
return undefined;
if (args.streamArgs.kind === "list") {
return {
kind: "list",
messages: await listStreams(ctx, component, {
threadId: args.threadId,
startOrder: args.streamArgs.startOrder,
includeStatuses: args.includeStatuses,
}),
};
}
else {
return {
kind: "deltas",
deltas: await ctx.runQuery(component.streams.listDeltas, {
threadId: args.threadId,
cursors: args.streamArgs.cursors,
}),
};
}
}
export async function abortStream(ctx, component, args) {
if ("streamId" in args) {
return await ctx.runMutation(component.streams.abort, {
reason: args.reason,
streamId: args.streamId,
});
}
else {
return await ctx.runMutation(component.streams.abortByOrder, {
reason: args.reason,
threadId: args.threadId,
order: args.order,
});
}
}
/**
* List the streaming messages for a thread.
* @param ctx A ctx object from a query, mutation, or action.
* @param component The agent component, usually `components.agent`.
* @param args.threadId The thread to list streams for.
* @param args.startOrder The order of the messages in the thread to start listing from.
* @param args.includeStatuses The statuses to include in the list.
* @returns The streams for the thread.
*/
export async function listStreams(ctx, component, { threadId, startOrder, includeStatuses, }) {
return ctx.runQuery(component.streams.list, {
threadId,
startOrder,
statuses: includeStatuses,
});
}
export const DEFAULT_STREAMING_OPTIONS = {
// This chunks by sentences / clauses. Punctuation followed by whitespace.
chunking: /[\p{P}\s]/u,
throttleMs: 250,
returnImmediately: false,
};
export function mergeTransforms(options, existing) {
if (!options) {
return existing;
}
const chunking = typeof options === "boolean"
? DEFAULT_STREAMING_OPTIONS.chunking
: options.chunking;
const transforms = Array.isArray(existing)
? existing
: existing
? [existing]
: [];
transforms.push(smoothStream({ delayInMs: null, chunking }));
return transforms;
}
export class DeltaStreamer {
component;
ctx;
metadata;
streamId;
options;
#nextParts = [];
#latestWrite = 0;
#ongoingWrite;
#cursor = 0;
abortController;
constructor(component, ctx, options, metadata) {
this.component = component;
this.ctx = ctx;
this.metadata = metadata;
this.options =
typeof options === "boolean"
? DEFAULT_STREAMING_OPTIONS
: { ...DEFAULT_STREAMING_OPTIONS, ...options };
this.#nextParts = [];
this.abortController = new AbortController();
if (metadata.abortSignal) {
metadata.abortSignal.addEventListener("abort", async () => {
if (this.streamId) {
this.abortController.abort();
const finalDelta = this.#createDelta();
await this.#ongoingWrite;
await this.ctx.runMutation(this.component.streams.abort, {
streamId: this.streamId,
reason: "abortSignal",
finalDelta,
});
}
});
}
}
async addParts(parts) {
if (this.abortController.signal.aborted) {
return;
}
if (!this.streamId) {
this.streamId = await this.ctx.runMutation(this.component.streams.create, omit(this.metadata, ["abortSignal"]));
}
this.#nextParts.push(...parts);
if (!this.#ongoingWrite &&
Date.now() - this.#latestWrite >= this.options.throttleMs) {
this.#ongoingWrite = this.#sendDelta();
}
}
async #sendDelta() {
if (this.abortController.signal.aborted) {
return;
}
const delta = this.#createDelta();
if (!delta) {
return;
}
this.#latestWrite = Date.now();
try {
const success = await this.ctx.runMutation(this.component.streams.addDelta, delta);
if (!success) {
this.abortController.abort();
}
}
catch (e) {
this.abortController.abort();
throw e;
}
// Now that we've sent the delta, check if we need to send another one.
if (this.#nextParts.length > 0 &&
Date.now() - this.#latestWrite >= this.options.throttleMs) {
// We send again immediately with the accumulated deltas.
this.#ongoingWrite = this.#sendDelta();
}
else {
this.#ongoingWrite = undefined;
}
}
#createDelta() {
if (this.#nextParts.length === 0) {
return undefined;
}
const start = this.#cursor;
const end = start + this.#nextParts.length;
this.#cursor = end;
const parts = serializeTextStreamingPartsV5(this.#nextParts);
this.#nextParts = [];
if (!this.streamId) {
throw new Error("Creating a delta before the stream is created");
}
return { streamId: this.streamId, start, end, parts };
}
async finish() {
if (!this.streamId) {
return;
}
const finalDelta = this.#createDelta();
await this.#ongoingWrite;
await this.ctx.runMutation(this.component.streams.finish, {
streamId: this.streamId,
finalDelta,
});
}
async fail(reason) {
if (this.abortController.signal.aborted) {
return;
}
this.abortController.abort();
if (!this.streamId) {
return;
}
const finalDelta = this.#createDelta();
await this.#ongoingWrite;
await this.ctx.runMutation(this.component.streams.abort, {
streamId: this.streamId,
reason,
finalDelta,
});
}
}
//# sourceMappingURL=streaming.js.map