@convex-dev/agent 0.2.0-alpha.2 → 0.2.0

package/dist/client/index.js

@@ -1,29 +1,30 @@
 import { embedMany, generateObject, generateText, stepCountIs, streamObject, streamText, } from "ai";
 import { assert, omit, pick } from "convex-helpers";
 import { internalActionGeneric, internalMutationGeneric, } from "convex/server";
-import { v } from "convex/values";
+import { convexToJson, v } from "convex/values";
 import { validateVectorDimension, } from "../component/vector/tables.js";
 import { deserializeMessage, serializeMessage, serializeNewMessagesInStep, serializeObjectResult, } from "../mapping.js";
 import { extractText, isTool } from "../shared.js";
 import { vMessageEmbeddings, vMessageWithMetadata, vSafeObjectArgs, vTextArgs, } from "../validators.js";
 import { createTool, wrapTools } from "./createTool.js";
 import { listMessages, saveMessages, } from "./messages.js";
-import { createThread, getThreadMetadata } from "./threads.js";
 import { fetchContextMessages, getModelName, getProviderName, } from "./search.js";
 import { DeltaStreamer, mergeTransforms, syncStreams, } from "./streaming.js";
+import { createThread, getThreadMetadata } from "./threads.js";
+import { inlineMessagesFiles } from "./files.js";
 export { stepCountIs } from "ai";
 export { vMessageDoc, vThreadDoc } from "../component/schema.js";
-export { serializeDataOrUrl, deserializeMessage, serializeMessage, } from "../mapping.js";
+export { deserializeMessage, serializeDataOrUrl, serializeMessage, } from "../mapping.js";
 // NOTE: these are also exported via @convex-dev/agent/validators
 // a future version may put them all here or move these over there
 export { vAssistantMessage, vContextOptions, vMessage, vPaginationResult, vProviderMetadata, vStorageOptions, vStreamArgs, vSystemMessage, vToolMessage, vUsage, vUserMessage, } from "../validators.js";
+export { definePlaygroundAPI, } from "./definePlaygroundAPI.js";
 export { getFile, storeFile } from "./files.js";
-export { filterOutOrphanedToolMessages, fetchContextMessages, } from "./search.js";
-export { abortStream, listStreams, syncStreams } from "./streaming.js";
 export { listMessages, saveMessage, saveMessages, } from "./messages.js";
+export { fetchContextMessages, filterOutOrphanedToolMessages, } from "./search.js";
+export { abortStream, listStreams, syncStreams } from "./streaming.js";
 export { createThread, getThreadMetadata } from "./threads.js";
 export { createTool, extractText, isTool };
-export { definePlaygroundAPI, } from "./definePlaygroundAPI.js";
 export class Agent {
     component;
     options;
@@ -39,8 +40,6 @@ export class Agent {
         const { thread } = await this.continueThread(ctx, {
             threadId,
             userId: args?.userId,
-            usageHandler: args?.usageHandler,
-            tools: args?.tools,
         });
         return { threadId, thread };
     }
@@ -70,90 +69,188 @@ export class Agent {
             },
         };
     }
+    async start(ctx, 
     /**
-     * This behaves like {@link generateText} 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).
-     * @param ctx The context passed from the action function calling this.
-     * @param { userId, threadId }: The user and thread to associate the message with
-     * @param args The arguments to the generateText function, along with extra controls
-     * for the {@link ContextOptions} and {@link StorageOptions}.
-     * @returns The result of the generateText function.
+     * These are the arguments you'll pass to the LLM call such as
+     * `generateText` or `streamText`. This function will look up the context
+     * and provide functions to save the steps, abort the generation, and more.
+     * The type of the arguments returned infers from the type of the arguments
+     * you pass here.
      */
-    async generateText(ctx, { userId: argsUserId, threadId, tools: threadTools, ...usageHandler }, args, options) {
-        const opts = { ...this.options, ...options, ...usageHandler };
+    args, options) {
+        const { threadId, ...opts } = { ...this.options, ...options };
         const context = await this._saveMessagesAndFetchContext(ctx, args, {
-            userId: argsUserId ?? undefined,
-            threadId,
+            userId: options?.userId,
+            threadId: options?.threadId,
             ...opts,
         });
-        const { args: aiArgs, messageId, order, userId } = context;
+        let pendingMessageId = context.pendingMessageId;
+        // TODO: extract pending message if one exists
+        const { args: aiArgs, promptMessageId, order, stepOrder, userId } = context;
         const messages = context.savedMessages ?? [];
+        if (pendingMessageId) {
+            if (!aiArgs._internal?.generateId) {
+                aiArgs._internal = {
+                    ...aiArgs._internal,
+                    generateId: () => pendingMessageId ?? crypto.randomUUID(),
+                };
+            }
+        }
         const toolCtx = {
             ...ctx,
             userId,
             threadId,
-            messageId,
+            promptMessageId,
             agent: this,
         };
-        const tools = wrapTools(toolCtx, args.tools ?? threadTools ?? this.options.tools);
+        const tools = wrapTools(toolCtx, args.tools ?? this.options.tools);
         const saveOutput = opts.storageOptions?.saveMessages !== "none";
-        try {
-            const result = (await generateText({
-                // Can be overridden
-                stopWhen: this.options.stopWhen,
+        const fail = async (reason) => {
+            if (threadId && promptMessageId) {
+                console.error("RollbackMessage", promptMessageId, reason);
+            }
+            if (pendingMessageId) {
+                await ctx.runMutation(this.component.messages.finalizeMessage, {
+                    messageId: pendingMessageId,
+                    result: { status: "failed", error: reason },
+                });
+            }
+        };
+        let activeModel = aiArgs.model;
+        if (aiArgs.abortSignal) {
+            const abortSignal = aiArgs.abortSignal;
+            aiArgs.abortSignal.addEventListener("abort", async () => {
+                await fail(abortSignal.reason ?? "Aborted");
+            }, { once: true });
+        }
+        return {
+            args: {
+                // eslint-disable-next-line @typescript-eslint/no-explicit-any
+                stopWhen: args.stopWhen ?? this.options.stopWhen,
                 ...aiArgs,
                 tools,
-                onStepFinish: async (step) => {
-                    if (threadId && messageId && saveOutput) {
-                        const saved = await this.saveStep(ctx, {
-                            userId,
-                            threadId,
-                            promptMessageId: messageId,
-                            model: aiArgs.model.modelId,
-                            provider: aiArgs.model.provider,
-                            step,
+                // abortSignal: abortController.signal,
+            },
+            order: order ?? 0,
+            stepOrder: stepOrder ?? 0,
+            userId,
+            promptMessageId,
+            getSavedMessages: () => messages,
+            updateModel: (model) => {
+                if (model) {
+                    activeModel = model;
+                }
+            },
+            fail,
+            save: async (toSave, createPendingMessage) => {
+                if (threadId && promptMessageId && saveOutput) {
+                    const metadata = {
+                        // TODO: get up to date one when user selects mid-generation
+                        model: getModelName(activeModel),
+                        provider: getProviderName(activeModel),
+                    };
+                    const serialized = "object" in toSave
+                        ? await serializeObjectResult(ctx, this.component, toSave.object, metadata)
+                        : await serializeNewMessagesInStep(ctx, this.component, toSave.step, metadata);
+                    const embeddings = await this.generateEmbeddings(ctx, { userId, threadId }, serialized.messages.map((m) => m.message));
+                    if (createPendingMessage) {
+                        serialized.messages.push({
+                            message: { role: "assistant", content: [] },
+                            status: "pending",
                         });
-                        messages.push(...saved.messages);
+                        embeddings?.vectors.push(null);
                     }
-                    if (this.options.rawRequestResponseHandler) {
-                        await this.options.rawRequestResponseHandler(ctx, {
-                            userId,
-                            threadId,
-                            agentName: this.options.name,
-                            request: step.request,
-                            response: step.response,
-                        });
+                    const saved = await ctx.runMutation(this.component.messages.addMessages, {
+                        userId,
+                        threadId,
+                        agentName: this.options.name,
+                        promptMessageId,
+                        pendingMessageId,
+                        messages: serialized.messages,
+                        embeddings,
+                        failPendingSteps: false,
+                    });
+                    const lastMessage = saved.messages.at(-1);
+                    if (createPendingMessage) {
+                        if (lastMessage.status === "failed") {
+                            pendingMessageId = undefined;
+                            messages.push(...saved.messages);
+                            await fail(lastMessage.error ??
+                                "Aborting - the pending message was marked as failed");
+                        }
+                        else {
+                            pendingMessageId = lastMessage._id;
+                            messages.push(...saved.messages.slice(0, -1));
+                        }
                     }
-                    if (opts.usageHandler && step.usage) {
-                        await opts.usageHandler(ctx, {
-                            userId,
-                            threadId,
-                            agentName: this.options.name,
-                            model: aiArgs.model.modelId,
-                            provider: aiArgs.model.provider,
-                            usage: step.usage,
-                            providerMetadata: step.providerMetadata,
-                        });
+                    else {
+                        pendingMessageId = undefined;
+                        messages.push(...saved.messages);
                     }
-                    return args.onStepFinish?.(step);
+                }
+                const output = "object" in toSave ? toSave.object : toSave.step;
+                if (this.options.rawRequestResponseHandler) {
+                    await this.options.rawRequestResponseHandler(ctx, {
+                        userId,
+                        threadId,
+                        agentName: this.options.name,
+                        request: output.request,
+                        response: output.response,
+                    });
+                }
+                if (opts.usageHandler && output.usage) {
+                    await opts.usageHandler(ctx, {
+                        userId,
+                        threadId,
+                        agentName: this.options.name,
+                        model: getModelName(activeModel),
+                        provider: getProviderName(activeModel),
+                        usage: output.usage,
+                        providerMetadata: output.providerMetadata,
+                    });
+                }
+            },
+        };
+    }
+    /**
+     * This behaves like {@link generateText} 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).
+     * @param ctx The context passed from the action function calling this.
+     * @param { userId, threadId }: The user and thread to associate the message with
+     * @param generateTextArgs The arguments to the generateText function, along with extra controls
+     * for the {@link ContextOptions} and {@link StorageOptions}.
+     * @returns The result of the generateText function.
+     */
+    async generateText(ctx, threadOpts, generateTextArgs, options) {
+        const { args, promptMessageId, order, ...call } = await this.start(ctx, generateTextArgs, { ...threadOpts, ...options });
+        const steps = [];
+        try {
+            const result = (await generateText({
+                ...args,
+                prepareStep: async (options) => {
+                    const result = await generateTextArgs.prepareStep?.(options);
+                    call.updateModel(result?.model ?? options.model);
+                    return result;
+                },
+                onStepFinish: async (step) => {
+                    steps.push(step);
+                    await call.save({ step }, await willContinue(steps, args.stopWhen));
+                    return generateTextArgs.onStepFinish?.(step);
                 },
             }));
-            result.messageId = messageId;
-            result.order = order;
-            result.messages = messages;
-            return result;
+            const metadata = {
+                promptMessageId,
+                order,
+                savedMessages: call.getSavedMessages(),
+                messageId: promptMessageId,
+            };
+            return Object.assign(result, metadata);
         }
         catch (error) {
-            if (threadId && messageId) {
-                console.error("RollbackMessage", messageId);
-                await ctx.runMutation(this.component.messages.rollbackMessage, {
-                    messageId,
-                    error: error.message,
-                });
-            }
+            await call.fail(errorToString(error));
             throw error;
         }
     }
@@ -164,112 +261,78 @@ export class Agent {
      * Use {@link continueThread} to get a version of this function already scoped
      * to a thread (and optionally userId).
      */
-    async streamText(ctx, { userId: argsUserId, threadId, 
-    /** Note: to get better type inference, pass tools in the next arg */
-    tools: threadTools, ...usageHandler }, 
+    async streamText(ctx, threadOpts, 
     /**
      * The arguments to the streamText function, similar to the ai `streamText` function.
      */
-    args, 
+    streamTextArgs, 
     /**
      * The {@link ContextOptions} and {@link StorageOptions}
      * options to use for fetching contextual messages and saving input/output messages.
      */
     options) {
-        const opts = { ...this.options, ...options, ...usageHandler };
-        const context = await this._saveMessagesAndFetchContext(ctx, args, {
-            userId: argsUserId ?? undefined,
-            threadId,
-            ...opts,
-        });
-        const { args: aiArgs, messageId, order, stepOrder, userId } = context;
-        const messages = context.savedMessages ?? [];
-        const toolCtx = {
-            ...ctx,
-            userId,
-            threadId,
-            messageId,
-            agent: this,
-        };
-        const tools = wrapTools(toolCtx, args.tools ?? threadTools ?? this.options.tools);
-        const saveOutput = opts.storageOptions?.saveMessages !== "none";
+        const { threadId } = threadOpts;
+        const { args, userId, order, stepOrder, promptMessageId, ...call } = await this.start(ctx, streamTextArgs, { ...threadOpts, ...options });
+        const steps = [];
+        const opts = { ...this.options, ...options };
         const streamer = threadId && opts.saveStreamDeltas
             ? new DeltaStreamer(this.component, ctx, opts.saveStreamDeltas, {
                 threadId,
                 userId,
                 agentName: this.options.name,
-                model: aiArgs.model.modelId,
-                provider: aiArgs.model.provider,
-                providerOptions: aiArgs.providerOptions,
+                model: getModelName(args.model),
+                provider: getProviderName(args.model),
+                providerOptions: args.providerOptions,
                 order,
                 stepOrder,
-                abortSignal: aiArgs.abortSignal,
+                abortSignal: args.abortSignal,
             })
             : undefined;
         const result = streamText({
-            // eslint-disable-next-line @typescript-eslint/no-explicit-any
-            stopWhen: this.options.stopWhen, // Can be overridden
-            ...aiArgs,
-            tools,
-            abortSignal: streamer?.abortController.signal ?? aiArgs.abortSignal,
-            experimental_transform: mergeTransforms(options?.saveStreamDeltas, args.experimental_transform),
+            ...args,
+            abortSignal: streamer?.abortController.signal ?? args.abortSignal,
+            // TODO: this is probably why reasoning isn't streaming
+            experimental_transform: mergeTransforms(options?.saveStreamDeltas, streamTextArgs.experimental_transform),
             onChunk: async (event) => {
                 await streamer?.addParts([event.chunk]);
                 // console.log("onChunk", chunk);
-                return args.onChunk?.(event);
+                return streamTextArgs.onChunk?.(event);
             },
             onError: async (error) => {
                 console.error("onError", error);
-                if (threadId && messageId && saveOutput) {
-                    await ctx.runMutation(this.component.messages.rollbackMessage, {
-                        messageId,
-                        error: error.error.message,
-                    });
+                await call.fail(errorToString(error.error));
+                await streamer?.fail(errorToString(error.error));
+                return streamTextArgs.onError?.(error);
+            },
+            // onFinish: async (event) => {
+            //   return streamTextArgs.onFinish?.(event);
+            // },
+            prepareStep: async (options) => {
+                const result = await streamTextArgs.prepareStep?.(options);
+                if (result) {
+                    const model = result.model ?? options.model;
+                    call.updateModel(model);
+                    return result;
                 }
-                // TODO: update the streamer to error state
-                return args.onError?.(error);
+                return undefined;
             },
             onStepFinish: async (step) => {
-                // console.log("onStepFinish", step);
-                if (threadId && messageId && saveOutput) {
-                    const saved = await this.saveStep(ctx, {
-                        userId,
-                        threadId,
-                        model: aiArgs.model.modelId,
-                        provider: aiArgs.model.provider,
-                        promptMessageId: messageId,
-                        step,
-                    });
-                    await streamer?.finish(saved.messages);
-                    messages.push(...saved.messages);
-                }
-                if (this.options.rawRequestResponseHandler) {
-                    await this.options.rawRequestResponseHandler(ctx, {
-                        userId,
-                        threadId,
-                        agentName: this.options.name,
-                        request: step.request,
-                        response: step.response,
-                    });
-                }
-                if (opts.usageHandler && step.usage) {
-                    await opts.usageHandler(ctx, {
-                        userId,
-                        threadId,
-                        agentName: this.options.name,
-                        model: aiArgs.model.modelId,
-                        provider: aiArgs.model.provider,
-                        usage: step.usage,
-                        providerMetadata: step.providerMetadata,
-                    });
+                steps.push(step);
+                const createPendingMessage = await willContinue(steps, args.stopWhen);
+                await call.save({ step }, createPendingMessage);
+                if (!createPendingMessage) {
+                    await streamer?.finish();
                 }
                 return args.onStepFinish?.(step);
             },
         });
-        result.messageId = messageId;
-        result.order = order;
-        result.messages = messages;
-        return result;
+        const metadata = {
+            promptMessageId,
+            order,
+            savedMessages: call.getSavedMessages(),
+            messageId: promptMessageId,
+        };
+        return Object.assign(result, metadata);
     }
     /**
      * This behaves like {@link generateObject} from the "ai" package except that
@@ -278,72 +341,30 @@ export class Agent {
      * Use {@link continueThread} to get a version of this function already scoped
      * to a thread (and optionally userId).
      */
-    async generateObject(ctx, { userId: argsUserId, threadId, ...usageHandler }, 
+    async generateObject(ctx, threadOpts, 
     /**
      * The arguments to the generateObject function, similar to the ai.generateObject function.
      */
-    args, 
+    generateObjectArgs, 
     /**
      * The {@link ContextOptions} and {@link StorageOptions}
      * options to use for fetching contextual messages and saving input/output messages.
      */
     options) {
-        const opts = { ...this.options, ...options, ...usageHandler };
-        const context = await this._saveMessagesAndFetchContext(ctx, args, {
-            userId: argsUserId ?? undefined,
-            threadId,
-            ...opts,
-        });
-        const { args: aiArgs, messageId, order, userId } = context;
-        const messages = context.savedMessages ?? [];
-        const saveOutput = opts.storageOptions?.saveMessages !== "none";
+        const { args, promptMessageId, order, fail, save, getSavedMessages } = await this.start(ctx, generateObjectArgs, { ...threadOpts, ...options });
         try {
-            const result = (await generateObject(
-            // eslint-disable-next-line @typescript-eslint/no-explicit-any
-            aiArgs));
-            if (threadId && messageId && saveOutput) {
-                const saved = await this.saveObject(ctx, {
-                    threadId,
-                    promptMessageId: messageId,
-                    result,
-                    userId,
-                    model: aiArgs.model.modelId,
-                    provider: aiArgs.model.provider,
-                });
-                messages.push(...saved.messages);
-            }
-            result.messageId = messageId;
-            result.order = order;
-            result.messages = messages;
-            if (this.options.rawRequestResponseHandler) {
-                await this.options.rawRequestResponseHandler(ctx, {
-                    userId,
-                    threadId,
-                    agentName: this.options.name,
-                    request: result.request,
-                    response: result.response,
-                });
-            }
-            if (opts.usageHandler && result.usage) {
-                await opts.usageHandler(ctx, {
-                    userId,
-                    threadId,
-                    agentName: this.options.name,
-                    model: aiArgs.model.modelId,
-                    provider: aiArgs.model.provider,
-                    usage: result.usage,
-                    providerMetadata: result.providerMetadata,
-                });
-            }
-            return result;
+            const result = (await generateObject(args));
+            await save({ object: result });
+            const metadata = {
+                promptMessageId,
+                order,
+                savedMessages: getSavedMessages(),
+                messageId: promptMessageId,
+            };
+            return Object.assign(result, metadata);
         }
         catch (error) {
-            if (threadId && messageId) {
-                await ctx.runMutation(this.component.messages.rollbackMessage, {
-                    messageId,
-                    error: error.message,
-                });
-            }
+            await fail(errorToString(error));
             throw error;
         }
     }
@@ -354,83 +375,50 @@ export class Agent {
      * Use {@link continueThread} to get a version of this function already scoped
      * to a thread (and optionally userId).
      */
-    async streamObject(ctx, { userId: argsUserId, threadId, ...usageHandler }, 
+    async streamObject(ctx, threadOpts, 
     /**
      * The arguments to the streamObject function, similar to the ai `streamObject` function.
      */
-    args, 
+    streamObjectArgs, 
     /**
      * The {@link ContextOptions} and {@link StorageOptions}
      * options to use for fetching contextual messages and saving input/output messages.
      */
     options) {
-        // TODO: unify all this shared code between all the generate* and stream* functions
-        const opts = { ...this.options, ...options, ...usageHandler };
-        const context = await this._saveMessagesAndFetchContext(ctx, args, {
-            userId: argsUserId ?? undefined,
-            threadId,
-            ...opts,
-        });
-        const { args: aiArgs, messageId, order, userId } = context;
-        const messages = context.savedMessages ?? [];
-        const saveOutput = opts.storageOptions?.saveMessages !== "none";
-        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        const { args, promptMessageId, order, fail, save, getSavedMessages } = await this.start(ctx, streamObjectArgs, { ...threadOpts, ...options });
         const stream = streamObject({
             // eslint-disable-next-line @typescript-eslint/no-explicit-any
-            ...aiArgs,
+            ...args,
             onError: async (error) => {
-                console.error("onError", error);
+                console.error(" streamObject onError", error);
+                // TODO: content that we have so far
+                // content: stream.fullStream.
+                await fail(errorToString(error.error));
                 return args.onError?.(error);
             },
             onFinish: async (result) => {
-                if (threadId && messageId && saveOutput) {
-                    const saved = await this.saveObject(ctx, {
-                        userId,
-                        threadId,
-                        promptMessageId: messageId,
-                        result: {
-                            object: result.object,
-                            finishReason: "stop",
-                            usage: result.usage,
-                            warnings: result.warnings,
-                            request: await stream.request,
-                            response: result.response,
-                            providerMetadata: result.providerMetadata,
-                            toJsonResponse: stream.toTextStreamResponse,
-                        },
-                        model: aiArgs.model.modelId,
-                        provider: aiArgs.model.provider,
-                    });
-                    messages.push(...saved.messages);
-                }
-                if (opts.usageHandler && result.usage) {
-                    await opts.usageHandler(ctx, {
-                        userId,
-                        threadId,
-                        agentName: this.options.name,
-                        model: aiArgs.model.modelId,
-                        provider: aiArgs.model.provider,
+                await save({
+                    object: {
+                        object: result.object,
+                        finishReason: result.error ? "error" : "stop",
                         usage: result.usage,
-                        providerMetadata: result.providerMetadata,
-                    });
-                }
-                if (this.options.rawRequestResponseHandler) {
-                    await this.options.rawRequestResponseHandler(ctx, {
-                        userId,
-                        threadId,
-                        agentName: this.options.name,
+                        warnings: result.warnings,
                         request: await stream.request,
                         response: result.response,
-                    });
-                }
-                // eslint-disable-next-line @typescript-eslint/no-explicit-any
+                        providerMetadata: result.providerMetadata,
+                        toJsonResponse: stream.toTextStreamResponse,
+                    },
+                });
                 return args.onFinish?.(result);
             },
         });
-        stream.messageId = messageId;
-        stream.order = order;
-        stream.messages = messages;
-        return stream;
+        const metadata = {
+            promptMessageId,
+            order,
+            savedMessages: getSavedMessages(),
+            messageId: promptMessageId,
+        };
+        return Object.assign(stream, metadata);
     }
     /**
      * Save a message to the thread.
@@ -451,6 +439,7 @@ export class Agent {
                 : [args.message],
             metadata: args.metadata ? [args.metadata] : undefined,
             skipEmbeddings: args.skipEmbeddings,
+            pendingMessageId: args.pendingMessageId,
         });
         const message = messages.at(-1);
         return { messageId: message._id, message };
@@ -469,7 +458,7 @@ export class Agent {
         if (args.embeddings) {
             embeddings = args.embeddings;
         }
-        else if (!skipEmbeddings && this.options.textEmbedding) {
+        else if (!skipEmbeddings && this.options.textEmbeddingModel) {
             if (!("runAction" in ctx)) {
                 console.warn("You're trying to save messages and generate embeddings, but you're in a mutation. " +
                     "Pass `skipEmbeddings: true` to skip generating embeddings in the mutation and skip this warning. " +
@@ -535,14 +524,14 @@ export class Agent {
             contextOptions,
             getEmbedding: async (text) => {
                 assert("runAction" in ctx);
-                assert(this.options.textEmbedding, "A textEmbedding model is required to be set on the Agent that you're doing vector search with");
+                assert(this.options.textEmbeddingModel, "A textEmbeddingModel is required to be set on the Agent that you're doing vector search with");
                 return {
                     embedding: (await this.doEmbed(ctx, {
                         userId: args.userId,
                         threadId: args.threadId,
                         values: [text],
                     })).embeddings[0],
-                    embeddingModel: this.options.textEmbedding,
+                    textEmbeddingModel: this.options.textEmbeddingModel,
                 };
             },
         });
@@ -573,7 +562,7 @@ export class Agent {
      * @returns The embeddings for the messages.
      */
     async generateEmbeddings(ctx, { userId, threadId, }, messages) {
-        if (!this.options.textEmbedding) {
+        if (!this.options.textEmbeddingModel) {
             return undefined;
         }
         let embeddings;
@@ -585,11 +574,12 @@ export class Agent {
         if (textIndexes.length === 0) {
             return undefined;
         }
+        const values = messageTexts.filter((t) => !!t);
         // Then embed those messages.
         const textEmbeddings = await this.doEmbed(ctx, {
             userId,
             threadId,
-            values: messageTexts,
+            values,
         });
         // Then assemble the embeddings into a single array with nulls for the messages without text.
         const embeddingsOrNull = Array(messages.length).fill(null);
@@ -599,7 +589,7 @@ export class Agent {
         if (textEmbeddings.embeddings.length > 0) {
             const dimension = textEmbeddings.embeddings[0].length;
             validateVectorDimension(dimension);
-            const model = getModelName(this.options.textEmbedding);
+            const model = getModelName(this.options.textEmbeddingModel);
             embeddings = { vectors: embeddingsOrNull, dimension, model };
         }
         return embeddings;
@@ -640,8 +630,8 @@ export class Agent {
             threadId: messagesMissingEmbeddings[0].threadId,
         }, messagesMissingEmbeddings.map((m) => deserializeMessage(m.message)));
         if (!embeddings) {
-            if (!this.options.textEmbedding) {
-                throw new Error("No embeddings were generated for the messages. You must pass a textEmbedding model to the agent constructor.");
+            if (!this.options.textEmbeddingModel) {
+                throw new Error("No embeddings were generated for the messages. You must pass a textEmbeddingModel to the agent constructor.");
             }
             throw new Error("No embeddings were generated for these messages: " +
                 messagesMissingEmbeddings.map((m) => m._id).join(", "));
@@ -666,9 +656,9 @@ export class Agent {
      * @param args The Step generated by the AI SDK.
      */
     async saveStep(ctx, args) {
-        const messages = await serializeNewMessagesInStep(ctx, this.component, args.step, {
-            provider: args.provider ?? getProviderName(this.options.chat),
-            model: args.model ?? getModelName(this.options.chat),
+        const { messages } = await serializeNewMessagesInStep(ctx, this.component, args.step, {
+            provider: args.provider ?? getProviderName(this.options.languageModel),
+            model: args.model ?? getModelName(this.options.languageModel),
         });
         const embeddings = await this.generateEmbeddings(ctx, { userId: args.userId, threadId: args.threadId }, messages.map((m) => m.message));
         return ctx.runMutation(this.component.messages.addMessages, {
@@ -690,10 +680,14 @@ export class Agent {
      */
     async saveObject(ctx, args) {
         const { messages } = await serializeObjectResult(ctx, this.component, args.result, {
-            model: args.model ?? getModelName(this.options.chat),
-            provider: args.provider ?? getProviderName(this.options.chat),
+            model: args.model ??
+                args.metadata?.model ??
+                getModelName(this.options.languageModel),
+            provider: args.provider ??
+                args.metadata?.provider ??
+                getProviderName(this.options.languageModel),
         });
-        const embeddings = await this.generateEmbeddings(ctx, { userId: args.userId, threadId: args.threadId }, messages.map((m) => deserializeMessage(m.message)));
+        const embeddings = await this.generateEmbeddings(ctx, { userId: args.userId, threadId: args.threadId }, messages.map((m) => m.message));
         return ctx.runMutation(this.component.messages.addMessages, {
             userId: args.userId,
             threadId: args.threadId,
@@ -712,19 +706,11 @@ export class Agent {
      * @param args What message to save. Generally the parent message sent into
      *   the generateText call.
      */
-    async completeMessage(ctx, args) {
-        const result = args.result;
-        if (result.kind === "success") {
-            await ctx.runMutation(this.component.messages.commitMessage, {
-                messageId: args.messageId,
-            });
-        }
-        else {
-            await ctx.runMutation(this.component.messages.rollbackMessage, {
-                messageId: args.messageId,
-                error: result.error,
-            });
-        }
+    async finalizeMessage(ctx, args) {
+        await ctx.runMutation(this.component.messages.finalizeMessage, {
+            messageId: args.messageId,
+            result: args.result,
+        });
     }
     /**
      * Update a message by its id.
@@ -848,7 +834,8 @@ export class Agent {
         const userId = argsUserId ??
             (threadId &&
                 (await ctx.runQuery(this.component.threads.getThread, { threadId }))
-                    ?.userId);
+                    ?.userId) ??
+            undefined;
         // If only a messageId is provided, this will add that message to the end.
         const contextMessages = await this.fetchContextMessages(ctx, {
             userId,
@@ -865,30 +852,45 @@ export class Agent {
         const promptMessage = promptMessageIndex !== -1
             ? contextMessages.splice(promptMessageIndex, 1)[0]
             : undefined;
-        let messageId = promptMessage?._id;
+        let promptMessageId = promptMessage?._id;
         let order = promptMessage?.order;
         let stepOrder = promptMessage?.stepOrder;
         let savedMessages = undefined;
-        if (threadId &&
-            messages.length + prompt.length &&
-            storageOptions?.saveMessages !== "none" &&
-            // If it was a promptMessageId, we don't want to save it again.
-            (!args.promptMessageId || storageOptions?.saveMessages === "all")) {
-            const saveAll = storageOptions?.saveMessages === "all";
-            const coreMessages = [...messages, ...prompt];
-            const toSave = saveAll ? coreMessages : coreMessages.slice(-1);
-            const metadata = Array.from({ length: toSave.length }, () => ({}));
-            const saved = await this.saveMessages(ctx, {
-                threadId,
-                userId,
-                messages: toSave,
-                metadata,
-                failPendingSteps: true,
-            });
-            messageId = saved.messages.at(-1)._id;
+        let pendingMessageId = undefined;
+        if (threadId && storageOptions?.saveMessages !== "none") {
+            let saved;
+            if (messages.length + prompt.length &&
+                // If it was a promptMessageId, we don't want to save it again.
+                (!args.promptMessageId || storageOptions?.saveMessages === "all")) {
+                const saveAll = storageOptions?.saveMessages === "all";
+                const coreMessages = [...messages, ...prompt];
+                const toSave = saveAll ? coreMessages : coreMessages.slice(-1);
+                const metadata = Array.from({ length: toSave.length }, () => ({}));
+                saved = await this.saveMessages(ctx, {
+                    threadId,
+                    userId,
+                    messages: [...toSave, { role: "assistant", content: [] }],
+                    metadata: [...metadata, { status: "pending" }],
+                    failPendingSteps: true,
+                    pendingMessageId: args.pendingMessageId,
+                });
+                promptMessageId = saved.messages.at(-2)._id;
+            }
+            else {
+                saved = await this.saveMessages(ctx, {
+                    threadId,
+                    userId,
+                    messages: [{ role: "assistant", content: [] }],
+                    metadata: [{ status: "pending" }],
+                    failPendingSteps: true,
+                    pendingMessageId: args.pendingMessageId,
+                });
+            }
+            pendingMessageId = saved.messages.at(-1)._id;
             order = saved.messages.at(-1).order;
             stepOrder = saved.messages.at(-1).stepOrder;
-            savedMessages = saved.messages;
+            // Don't return the pending message
+            savedMessages = saved.messages.slice(0, -1);
         }
         if (promptMessage?.message) {
             if (!args.prompt) {
@@ -898,7 +900,7 @@ export class Agent {
             // Lazily generate embeddings for the prompt message, if it doesn't have
             // embeddings yet. This can happen if the message was saved in a mutation
             // where the LLM is not available.
-            if (!promptMessage.embeddingId && this.options.textEmbedding) {
+            if (!promptMessage.embeddingId && this.options.textEmbeddingModel) {
                 await this._generateAndSaveEmbeddings(ctx, [promptMessage]);
             }
         }
@@ -924,27 +926,28 @@ export class Agent {
         }
         // Process messages to inline localhost files (if not, file urls pointing to localhost will be sent to LLM providers)
         if (process.env.CONVEX_CLOUD_URL?.startsWith("http://127.0.0.1")) {
-            processedMessages = await this._inlineMessagesFiles(processedMessages);
+            processedMessages = await inlineMessagesFiles(processedMessages);
         }
         const { prompt: _, model, ...rest } = args;
         return {
             args: {
                 ...this.options.callSettings,
                 ...rest,
-                model: model ?? this.options.chat,
+                model: model ?? this.options.languageModel,
                 system: args.system ?? this.options.instructions,
                 messages: processedMessages,
             },
             userId,
-            messageId,
+            promptMessageId,
+            pendingMessageId,
             savedMessages,
             order,
             stepOrder,
         };
     }
     async doEmbed(ctx, options) {
-        const embeddingModel = this.options.textEmbedding;
-        assert(embeddingModel, "a textEmbedding model is required to be set on the Agent that you're doing vector search with");
+        const embeddingModel = this.options.textEmbeddingModel;
+        assert(embeddingModel, "a textEmbeddingModel is required to be set on the Agent that you're doing vector search with");
         const result = await embedMany({
             ...this.options.callSettings,
             model: embeddingModel,
@@ -969,64 +972,6 @@ export class Agent {
         }
         return { embeddings: result.embeddings };
     }
-    /**
-     * Process messages to inline file and image URLs that point to localhost
-     * by converting them to base64. This solves the problem of LLMs not being
-     * able to access localhost URLs.
-     */
-    async _inlineMessagesFiles(messages) {
-        // Process each message to convert localhost URLs to base64
-        return Promise.all(messages.map(async (message) => {
-            if ((message.role !== "user" && message.role !== "assistant") ||
-                typeof message.content === "string" ||
-                !Array.isArray(message.content)) {
-                return message;
-            }
-            const processedContent = await Promise.all(message.content.map(async (part) => {
-                if (part.type === "image" && part.image instanceof URL) {
-                    assert(message.role === "user", "Images can only be in user messages");
-                    if (this._isLocalhostUrl(part.image)) {
-                        const imageData = await this._downloadFile(part.image);
-                        return { ...part, image: imageData };
-                    }
-                }
-                // Handle file parts
-                if (part.type === "file" && part.data instanceof URL) {
-                    if (this._isLocalhostUrl(part.data)) {
-                        const fileData = await this._downloadFile(part.data);
-                        return { ...part, data: fileData };
-                    }
-                }
-                return part;
-            }));
-            if (message.role === "user") {
-                return { ...message, content: processedContent };
-            }
-            else {
-                return { ...message, content: processedContent };
-            }
-        }));
-    }
-    /**
-     * Check if a URL points to localhost
-     */
-    _isLocalhostUrl(url) {
-        return (url.hostname === "localhost" ||
-            url.hostname === "127.0.0.1" ||
-            url.hostname === "::1" ||
-            url.hostname === "0.0.0.0");
-    }
-    /**
-     * Download a file from a URL
-     */
-    async _downloadFile(url) {
-        // Fetch the file
-        const response = await fetch(url);
-        if (!response.ok) {
-            throw new Error(`Failed to fetch ${url}: ${response.statusText}`);
-        }
-        return await response.arrayBuffer();
-    }
     /**
      * WORKFLOW UTILITIES
      */
@@ -1067,15 +1012,15 @@ export class Agent {
      * @param spec Configuration for the agent acting as an action, including
      *   {@link ContextOptions}, {@link StorageOptions}, and {@link stopWhen}.
      */
-    asTextAction(spec) {
-        const stopWhen = spec?.stopWhen ?? this.options.stopWhen;
+    asTextAction(spec, overrides) {
         return internalActionGeneric({
             args: vTextArgs,
             handler: async (ctx_, args) => {
                 const stream = args.stream === true ? spec?.stream || true : (spec?.stream ?? false);
                 const targetArgs = { userId: args.userId, threadId: args.threadId };
                 const llmArgs = {
-                    stopWhen,
+                    stopWhen: spec?.stopWhen ?? this.options.stopWhen,
+                    ...overrides,
                     ...omit(args, ["storageOptions", "contextOptions"]),
                     messages: args.messages?.map(deserializeMessage),
                     prompt: Array.isArray(args.prompt)
@@ -1088,7 +1033,7 @@ export class Agent {
                 }
                 const opts = {
                     ...this.options,
-                    ...(spec && pick(spec, ["contextOptions", "storageOptions"])),
+                    ...pick(spec, ["contextOptions", "storageOptions"]),
                     ...pick(args, ["contextOptions", "storageOptions"]),
                     saveStreamDeltas: stream,
                 };
@@ -1101,10 +1046,11 @@ export class Agent {
                     await result.consumeStream();
                     return {
                         text: await result.text,
-                        messageId: result.messageId,
+                        promptMessageId: result.promptMessageId,
                         order: result.order,
                         finishReason: await result.finishReason,
                         warnings: result.warnings,
+                        savedMessageIds: result.savedMessages?.map((m) => m._id) ?? [],
                     };
                 }
                 else {
@@ -1112,10 +1058,11 @@ export class Agent {
                     const res = await this.generateText(ctx, targetArgs, llmArgs, opts);
                     return {
                         text: res.text,
-                        messageId: res.messageId,
+                        promptMessageId: res.promptMessageId,
                         order: res.order,
                         finishReason: res.finishReason,
                         warnings: res.warnings,
+                        savedMessageIds: res.savedMessages?.map((m) => m._id) ?? [],
                     };
                 }
             },
@@ -1128,22 +1075,37 @@ export class Agent {
      * the normal parameters to {@link generateObject}, plus {@link ContextOptions}
      * and stopWhen.
      */
-    asObjectAction(spec, options) {
+    asObjectAction(objectArgs, options) {
         return internalActionGeneric({
             args: vSafeObjectArgs,
-            handler: async (ctx, args) => {
-                const { userId, threadId, ...rest } = args;
+            handler: async (ctx_, args) => {
+                const { userId, threadId, callSettings, ...rest } = args;
                 const overrides = pick(rest, ["contextOptions", "storageOptions"]);
-                const value = await this.generateObject(ctx, { userId, threadId }, {
-                    ...spec,
-                    ...omit(rest, ["contextOptions", "storageOptions"]),
-                }, { ...this.options, ...options, ...overrides });
+                const targetArgs = { userId, threadId };
+                const llmArgs = {
+                    ...objectArgs,
+                    ...callSettings,
+                    ...omit(rest, ["storageOptions", "contextOptions"]),
+                    messages: args.messages?.map(deserializeMessage),
+                    prompt: Array.isArray(args.prompt)
+                        ? args.prompt.map(deserializeMessage)
+                        : args.prompt,
+                };
+                const ctx = (options?.customCtx
+                    ? { ...ctx_, ...options.customCtx(ctx_, targetArgs, llmArgs) }
+                    : ctx_);
+                const value = await this.generateObject(ctx, targetArgs, llmArgs, {
+                    ...this.options,
+                    ...options,
+                    ...overrides,
+                });
                 return {
-                    object: value.object,
-                    messageId: value.messageId,
+                    object: convexToJson(value.object),
+                    promptMessageId: value.promptMessageId,
                     order: value.order,
                     finishReason: value.finishReason,
                     warnings: value.warnings,
+                    savedMessageIds: value.savedMessages?.map((m) => m._id) ?? [],
                 };
             },
         });
@@ -1193,4 +1155,28 @@ export class Agent {
         });
     }
 }
+async function willContinue(
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+steps, 
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+stopWhen) {
+    const step = steps.at(-1);
+    // we aren't doing another round after a tool result
+    // TODO: whether to handle continuing after too much context used..
+    if (step.finishReason !== "tool-calls")
+        return false;
+    // we don't have a tool result, so we'll wait for more
+    if (step.toolCalls.length > step.toolResults.length)
+        return false;
+    if (Array.isArray(stopWhen)) {
+        return (await Promise.all(stopWhen.map(async (s) => s({ steps })))).every((stop) => !stop);
+    }
+    return !(await stopWhen?.({ steps }));
+}
+function errorToString(error) {
+    if (error instanceof Error) {
+        return error.message;
+    }
+    return String(error);
+}
 //# sourceMappingURL=index.js.map