@convex-dev/agent
Version:
A agent component for Convex.
145 lines • 6.49 kB
JavaScript
import { streamText as streamTextAi } from "ai";
import { compressUIMessageChunks, DeltaStreamer, mergeTransforms, } from "./streaming.js";
import { startGeneration } from "./start.js";
import { getModelName, getProviderName } from "../shared.js";
import { errorToString, willContinue } from "./utils.js";
/**
* This behaves like {@link streamText} from the "ai" package except that
* it add context based on the userId and threadId and saves the input and
* resulting messages to the thread, if specified.
* Use {@link continueThread} to get a version of this function already scoped
* to a thread (and optionally userId).
*/
export async function streamText(ctx, component,
/**
* The arguments to the streamText function, similar to the ai sdk's
* {@link streamText} function, along with Agent prompt options.
*/
streamTextArgs,
/**
* The {@link ContextOptions} and {@link StorageOptions}
* options to use for fetching contextual messages and saving input/output messages.
*/
options) {
const { threadId } = options ?? {};
const { args, userId, order, stepOrder, promptMessageId, ...call } = await startGeneration(ctx, component, streamTextArgs, options);
const steps = [];
// Track the final step for atomic save with stream finish (issue #181).
// Only used when streamText awaits stream consumption itself; the
// `returnImmediately` path saves inline instead (see onStepFinish below).
let pendingFinalStep;
// Whether streamText will await stream consumption before returning.
// When false (saveStreamDeltas.returnImmediately === true), we cannot
// defer the final-step save to a post-await block — the function has
// already returned by the time onStepFinish fires. See issue #265.
const willAwaitStream = Boolean(threadId) &&
(options.saveStreamDeltas === true ||
(typeof options.saveStreamDeltas === "object" &&
!options.saveStreamDeltas.returnImmediately));
const streamer = threadId && options.saveStreamDeltas
? new DeltaStreamer(component, ctx, {
throttleMs: typeof options.saveStreamDeltas === "object"
? options.saveStreamDeltas.throttleMs
: undefined,
onAsyncAbort: call.fail,
compress: compressUIMessageChunks,
abortSignal: args.abortSignal,
}, {
threadId,
userId,
agentName: options?.agentName,
model: getModelName(args.model),
provider: getProviderName(args.model),
providerOptions: args.providerOptions,
format: "UIMessageChunk",
order,
stepOrder,
})
: undefined;
const result = streamTextAi({
...args,
abortSignal: streamer?.abortController.signal ?? args.abortSignal,
experimental_transform: mergeTransforms(options?.saveStreamDeltas, streamTextArgs.experimental_transform),
onError: async (error) => {
console.error("onError", error);
await call.fail(errorToString(error.error));
await streamer?.fail(errorToString(error.error));
return streamTextArgs.onError?.(error);
},
prepareStep: async (options) => {
const result = await streamTextArgs.prepareStep?.(options);
if (result) {
const model = result.model ?? options.model;
call.updateModel(model);
// streamer?.updateMetadata({
// model: getModelName(model),
// provider: getProviderName(model),
// providerOptions: options.messages.at(-1)?.providerOptions,
// });
return result;
}
return undefined;
},
onStepFinish: async (step) => {
steps.push(step);
const createPendingMessage = await willContinue(steps, args.stopWhen);
if (!createPendingMessage && streamer) {
// Final step with streaming enabled.
streamer.markFinishedExternally();
if (willAwaitStream) {
// We're about to `await stream` below — defer the save so it
// happens atomically with stream finish (issue #181).
pendingFinalStep = step;
}
else {
// returnImmediately path: streamText is about to return without
// awaiting consumption, so the deferred-save block below won't
// see this step. Save inline now (issue #265).
const finishStreamId = await streamer.getOrCreateStreamId();
await call.save({ step }, false, finishStreamId);
}
}
else {
await call.save({ step }, createPendingMessage);
}
return args.onStepFinish?.(step);
},
});
const stream = streamer?.consumeStream(result.toUIMessageStream());
if (willAwaitStream) {
try {
await stream;
await result.consumeStream();
}
catch (e) {
// If the stream errored (e.g. onStepFinish threw), the DeltaStreamer's
// finish() was never called, leaving the streaming message stuck in
// "streaming" state. Clean it up by marking it as aborted.
await streamer?.fail(e instanceof Error ? e.message : String(e));
// Save the deferred final step if it was already generated but not yet persisted
if (pendingFinalStep) {
try {
await call.save({ step: pendingFinalStep }, false);
}
catch (saveError) {
console.error("Failed to save deferred final step:", saveError);
}
pendingFinalStep = undefined;
}
throw e;
}
}
// If we deferred the final step save, do it now with atomic stream finish.
if (pendingFinalStep && streamer) {
const finishStreamId = await streamer.getOrCreateStreamId();
await call.save({ step: pendingFinalStep }, false, finishStreamId);
}
const metadata = {
promptMessageId,
order,
savedMessages: call.getSavedMessages(),
messageId: promptMessageId,
};
return Object.assign(result, metadata);
}
//# sourceMappingURL=streamText.js.map