@convex-dev/agent 0.2.6-alpha.0 → 0.2.6

package/dist/client/index.js

@@ -1,36 +1,31 @@
-import { embedMany, generateObject, generateText, stepCountIs, streamObject, streamText, } from "ai";
+import { generateObject, generateText, stepCountIs, streamObject, streamText, } from "ai";
 import { assert, omit, pick } from "convex-helpers";
 import { internalActionGeneric, internalMutationGeneric, } from "convex/server";
 import { convexToJson, v } from "convex/values";
-import { validateVectorDimension, } from "../component/vector/tables.js";
+import {} from "../component/vector/tables.js";
 import { deserializeMessage, serializeMessage, serializeNewMessagesInStep, serializeObjectResult, } from "../mapping.js";
-import { extractText, isTool } from "../shared.js";
+import { getModelName, getProviderName } from "../shared.js";
 import { vMessageEmbeddings, vMessageWithMetadata, vSafeObjectArgs, vTextArgs, } from "../validators.js";
-import { wrapTools } from "./createTool.js";
 import { listMessages, saveMessages, } from "./messages.js";
-import { fetchContextMessages, getModelName, getProviderName, } from "./search.js";
+import { embedMany, embedMessages, fetchContextMessages, generateAndSaveEmbeddings, } from "./search.js";
+import { start } from "./start.js";
 import { DeltaStreamer, syncStreams, } from "./streaming.js";
 import { mergeTransforms, serializeTextStreamingPartsV5, } from "./textStreamParts.js";
 import { createThread, getThreadMetadata } from "./threads.js";
-import { inlineMessagesFiles } from "./files.js";
 export { stepCountIs } from "ai";
-export { vMessageDoc, vThreadDoc } from "../component/schema.js";
 export { deserializeMessage, serializeDataOrUrl, serializeMessage, guessMimeType, toUIFilePart, } 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, vSource, vContent, } from "../validators.js";
+export { vAssistantMessage, vContextOptions, vMessage, vMessageDoc, vPaginationResult, vProviderMetadata, vStorageOptions, vStreamArgs, vSystemMessage, vThreadDoc, vToolMessage, vUsage, vUserMessage, vSource, vContent, } from "../validators.js";
 export { definePlaygroundAPI, } from "./definePlaygroundAPI.js";
 export { getFile, storeFile } from "./files.js";
 export { listMessages, saveMessage, saveMessages, } from "./messages.js";
-export { fetchContextMessages, filterOutOrphanedToolMessages, } from "./search.js";
+export { fetchContextMessages, filterOutOrphanedToolMessages, fetchContextWithPrompt, generateAndSaveEmbeddings, embedMessages, embedMany, } from "./search.js";
 export { abortStream, listStreams, syncStreams, vStreamMessagesReturnValue, } from "./streaming.js";
 export { createThread, getThreadMetadata, updateThreadMetadata, searchThreadTitles, } from "./threads.js";
 export { extractText, isTool, sorted } from "../shared.js";
 export { createTool } from "./createTool.js";
 export { mockModel } from "./mockModel.js";
-// 10k characters should be more than enough for most cases, and stays under
-// the 8k token limit for some models.
-const MAX_EMBEDDING_TEXT_LENGTH = 10_000;
 export class Agent {
     component;
     options;
@@ -84,147 +79,17 @@ export class Agent {
      * you pass here.
      */
     args, options) {
-        const { threadId, ...opts } = { ...this.options, ...options };
-        const context = await this._saveMessagesAndFetchContext(ctx, {
-            userId: options?.userId,
-            threadId: options?.threadId,
-            messages: args.messages,
-            prompt: args.prompt,
-            promptMessageId: args.promptMessageId,
-            ...opts,
-        });
-        let pendingMessageId = context.pendingMessageId;
-        const { messages, promptMessageId, order, stepOrder, userId } = context;
-        const savedMessages = context.savedMessages ?? [];
-        const toolCtx = {
-            ...ctx,
-            userId,
-            threadId,
-            promptMessageId,
-            agent: this,
-        };
-        const tools = wrapTools(toolCtx, args.tools ?? this.options.tools);
-        const saveOutput = opts.storageOptions?.saveMessages !== "none";
-        const fail = async (reason) => {
-            if (threadId && promptMessageId) {
-                console.error(`Message failed in thread ${threadId} with promptMessageId ${promptMessageId}: ${reason}`);
-            }
-            if (pendingMessageId) {
-                await ctx.runMutation(this.component.messages.finalizeMessage, {
-                    messageId: pendingMessageId,
-                    result: { status: "failed", error: reason },
-                });
-            }
-        };
-        if (args.abortSignal) {
-            const abortSignal = args.abortSignal;
-            abortSignal.addEventListener("abort", async () => {
-                await fail(abortSignal.reason ?? "Aborted");
-            }, { once: true });
-        }
-        const aiArgs = {
-            ...this.options.callSettings,
-            providerOptions: this.options.providerOptions,
-            ...omit(args, ["messages", "prompt", "promptMessageId"]),
-            model: args.model ?? this.options.languageModel,
+        return start(ctx, this.component, {
+            ...args,
+            tools: (args.tools ?? this.options.tools),
             system: args.system ?? this.options.instructions,
-            messages,
-            stopWhen: args.stopWhen ??
-                this.options.stopWhen ??
-                stepCountIs(this.options.maxSteps ?? 1),
-            tools,
-        };
-        if (pendingMessageId) {
-            if (!aiArgs._internal?.generateId) {
-                aiArgs._internal = {
-                    ...aiArgs._internal,
-                    generateId: () => pendingMessageId ?? crypto.randomUUID(),
-                };
-            }
-        }
-        let activeModel = aiArgs.model;
-        return {
-            args: aiArgs,
-            order: order ?? 0,
-            stepOrder: stepOrder ?? 0,
-            userId,
-            promptMessageId,
-            getSavedMessages: () => savedMessages,
-            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",
-                        });
-                        embeddings?.vectors.push(null);
-                    }
-                    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;
-                            savedMessages.push(...saved.messages);
-                            await fail(lastMessage.error ??
-                                "Aborting - the pending message was marked as failed");
-                        }
-                        else {
-                            pendingMessageId = lastMessage._id;
-                            savedMessages.push(...saved.messages.slice(0, -1));
-                        }
-                    }
-                    else {
-                        pendingMessageId = undefined;
-                        savedMessages.push(...saved.messages);
-                    }
-                }
-                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,
-                    });
-                }
-            },
-        };
+            stopWhen: (args.stopWhen ?? this.options.stopWhen),
+        }, {
+            ...this.options,
+            ...options,
+            agentName: this.options.name,
+            agentForToolCtx: this,
+        });
     }
     /**
      * This behaves like {@link generateText} from the "ai" package except that
@@ -321,9 +186,6 @@ export class Agent {
                 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) {
@@ -343,17 +205,17 @@ export class Agent {
                 return args.onStepFinish?.(step);
             },
         });
+        if ((typeof options?.saveStreamDeltas === "object" &&
+            !options.saveStreamDeltas.returnImmediately) ||
+            options?.saveStreamDeltas === true) {
+            await result.consumeStream();
+        }
         const metadata = {
             promptMessageId,
             order,
             savedMessages: call.getSavedMessages(),
             messageId: promptMessageId,
         };
-        if ((typeof options?.saveStreamDeltas === "object" &&
-            !options.saveStreamDeltas.returnImmediately) ||
-            options?.saveStreamDeltas === true) {
-            await result.consumeStream();
-        }
         return Object.assign(result, metadata);
     }
     /**
@@ -549,7 +411,9 @@ export class Agent {
                 assert("runAction" in ctx);
                 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, {
+                    embedding: (await embedMany(ctx, {
+                        ...this.options,
+                        agentName: this.options.name,
                         userId: args.userId,
                         threadId: args.threadId,
                         values: [text],
@@ -584,40 +448,8 @@ export class Agent {
      * @param messages The messages to get the embeddings for.
      * @returns The embeddings for the messages.
      */
-    async generateEmbeddings(ctx, { userId, threadId, }, messages) {
-        if (!this.options.textEmbeddingModel) {
-            return undefined;
-        }
-        let embeddings;
-        const messageTexts = messages.map((m) => !isTool(m) && extractText(m));
-        // Find the indexes of the messages that have text.
-        const textIndexes = messageTexts
-            .map((t, i) => (t ? i : undefined))
-            .filter((i) => i !== undefined);
-        if (textIndexes.length === 0) {
-            return undefined;
-        }
-        const values = messageTexts
-            .map((t) => t && t.trim().slice(0, MAX_EMBEDDING_TEXT_LENGTH))
-            .filter((t) => !!t);
-        // Then embed those messages.
-        const textEmbeddings = await this.doEmbed(ctx, {
-            userId,
-            threadId,
-            values,
-        });
-        // Then assemble the embeddings into a single array with nulls for the messages without text.
-        const embeddingsOrNull = Array(messages.length).fill(null);
-        textIndexes.forEach((i, j) => {
-            embeddingsOrNull[i] = textEmbeddings.embeddings[j];
-        });
-        if (textEmbeddings.embeddings.length > 0) {
-            const dimension = textEmbeddings.embeddings[0].length;
-            validateVectorDimension(dimension);
-            const model = getModelName(this.options.textEmbeddingModel);
-            embeddings = { vectors: embeddingsOrNull, dimension, model };
-        }
-        return embeddings;
+    async generateEmbeddings(ctx, args, messages) {
+        return embedMessages(ctx, { ...args, ...this.options, agentName: this.options.name }, messages);
     }
     /**
      * Generate embeddings for a set of messages, and save them to the database.
@@ -636,9 +468,6 @@ export class Agent {
                     .filter((id) => !messages.some((m) => m?._id === id))
                     .join(", "));
         }
-        await this._generateAndSaveEmbeddings(ctx, messages);
-    }
-    async _generateAndSaveEmbeddings(ctx, messages) {
         if (messages.some((m) => !m.message)) {
             throw new Error("Some messages don't have a message: " +
                 messages
@@ -646,34 +475,17 @@ export class Agent {
                     .map((m) => m._id)
                     .join(", "));
         }
-        const messagesMissingEmbeddings = messages.filter((m) => !m.embeddingId);
-        if (messagesMissingEmbeddings.length === 0) {
-            return;
+        const { textEmbeddingModel } = this.options;
+        if (!textEmbeddingModel) {
+            throw new Error("No embeddings were generated for the messages. You must pass a textEmbeddingModel to the agent constructor.");
         }
-        const embeddings = await this.generateEmbeddings(ctx, {
-            userId: messagesMissingEmbeddings[0].userId,
-            threadId: messagesMissingEmbeddings[0].threadId,
-        }, messagesMissingEmbeddings.map((m) => deserializeMessage(m.message)));
-        if (!embeddings) {
-            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(", "));
-        }
-        await ctx.runMutation(this.component.vector.index.insertBatch, {
-            vectorDimension: embeddings.dimension,
-            vectors: messagesMissingEmbeddings
-                .map((m, i) => ({
-                messageId: m._id,
-                model: embeddings.model,
-                table: "messages",
-                userId: m.userId,
-                threadId: m.threadId,
-                vector: embeddings.vectors[i],
-            }))
-                .filter((v) => v.vector !== null),
-        });
+        await generateAndSaveEmbeddings(ctx, this.component, {
+            ...this.options,
+            agentName: this.options.name,
+            threadId: messages[0].threadId,
+            userId: messages[0].userId,
+            textEmbeddingModel,
+        }, messages);
     }
     /**
      * Explicitly save a "step" created by the AI SDK.
@@ -848,143 +660,6 @@ export class Agent {
             limit: args.pageSize,
         });
     }
-    async _saveMessagesAndFetchContext(ctx, { userId: argsUserId, threadId, contextOptions, storageOptions, ...args }) {
-        // If only a promptMessageId is provided, this will be empty.
-        const messages = args.messages ?? [];
-        const promptArray = !args.prompt
-            ? []
-            : Array.isArray(args.prompt)
-                ? args.prompt.map((p) => deserializeMessage(p))
-                : [{ role: "user", content: args.prompt }];
-        const userId = argsUserId ??
-            (threadId &&
-                (await ctx.runQuery(this.component.threads.getThread, { threadId }))
-                    ?.userId) ??
-            undefined;
-        // If only a messageId is provided, this will add that message to the end.
-        const contextMessages = await this.fetchContextMessages(ctx, {
-            userId,
-            threadId,
-            upToAndIncludingMessageId: args.promptMessageId,
-            messages,
-            contextOptions,
-        });
-        // If it was a promptMessageId, pop it off context messages
-        // and add to the end of messages.
-        const promptMessageIndex = args.promptMessageId
-            ? contextMessages.findIndex((m) => m._id === args.promptMessageId)
-            : -1;
-        const promptMessage = promptMessageIndex !== -1
-            ? contextMessages.splice(promptMessageIndex, 1)[0]
-            : undefined;
-        let promptMessageId = promptMessage?._id;
-        let order = promptMessage?.order;
-        let stepOrder = promptMessage?.stepOrder;
-        let savedMessages = undefined;
-        let pendingMessageId = undefined;
-        if (threadId && storageOptions?.saveMessages !== "none") {
-            let saved;
-            if (messages.length + promptArray.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,
-                    ...promptArray,
-                ];
-                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" }],
-                    // TODO: sanity check
-                    failPendingSteps: !!args.promptMessageId,
-                });
-                promptMessageId = saved.messages.at(-2)._id;
-            }
-            else {
-                saved = await this.saveMessages(ctx, {
-                    threadId,
-                    userId,
-                    messages: [{ role: "assistant", content: [] }],
-                    metadata: [{ status: "pending" }],
-                    failPendingSteps: !!args.promptMessageId,
-                });
-            }
-            pendingMessageId = saved.messages.at(-1)._id;
-            order = saved.messages.at(-1).order;
-            stepOrder = saved.messages.at(-1).stepOrder;
-            // Don't return the pending message
-            savedMessages = saved.messages.slice(0, -1);
-        }
-        if (promptMessage?.message) {
-            if (!args.prompt) {
-                // If they override the prompt, we skip the existing prompt message.
-                messages.push(promptMessage.message);
-            }
-            // 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.textEmbeddingModel) {
-                await this._generateAndSaveEmbeddings(ctx, [promptMessage]);
-            }
-        }
-        const prePrompt = contextMessages.map((m) => m.message).filter((m) => !!m);
-        let existingResponses = [];
-        if (promptMessageIndex !== -1) {
-            // pull any messages that already responded to the prompt off
-            // and add them after the prompt
-            existingResponses = prePrompt.splice(promptMessageIndex);
-        }
-        let processedMessages = [
-            ...prePrompt,
-            ...messages,
-            ...promptArray,
-            ...existingResponses,
-        ].map((m) => deserializeMessage(m));
-        // 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 inlineMessagesFiles(processedMessages);
-        }
-        return {
-            messages: processedMessages,
-            userId,
-            promptMessageId,
-            pendingMessageId,
-            savedMessages,
-            order,
-            stepOrder,
-        };
-    }
-    async doEmbed(ctx, options) {
-        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,
-            values: options.values,
-            abortSignal: options.abortSignal,
-            headers: options.headers,
-        });
-        if (this.options.usageHandler && result.usage) {
-            await this.options.usageHandler(ctx, {
-                userId: options.userId,
-                threadId: options.threadId,
-                agentName: this.options.name,
-                model: getModelName(embeddingModel),
-                provider: getProviderName(embeddingModel),
-                providerMetadata: undefined,
-                usage: {
-                    inputTokens: result.usage.tokens,
-                    outputTokens: 0,
-                    totalTokens: result.usage.tokens,
-                },
-            });
-        }
-        return { embeddings: result.embeddings };
-    }
     /**
      * WORKFLOW UTILITIES
      */
@@ -1030,22 +705,22 @@ export class Agent {
             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 { userId, threadId, prompt, messages, maxSteps, ...rest } = args;
+                const targetArgs = { userId, threadId };
                 const llmArgs = {
-                    stopWhen: spec?.stopWhen ?? this.options.stopWhen,
+                    stopWhen: spec?.stopWhen,
                     ...overrides,
-                    ...omit(args, ["storageOptions", "contextOptions"]),
-                    messages: args.messages?.map(deserializeMessage),
-                    prompt: Array.isArray(args.prompt)
-                        ? args.prompt.map(deserializeMessage)
-                        : args.prompt,
+                    ...omit(rest, ["storageOptions", "contextOptions", "stream"]),
+                    messages: messages?.map(deserializeMessage),
+                    prompt: Array.isArray(prompt)
+                        ? prompt.map(deserializeMessage)
+                        : prompt,
                     toolChoice: args.toolChoice,
                 };
-                if (args.maxSteps) {
-                    llmArgs.stopWhen = stepCountIs(args.maxSteps);
+                if (maxSteps) {
+                    llmArgs.stopWhen = stepCountIs(maxSteps);
                 }
                 const opts = {
-                    ...this.options,
                     ...pick(spec, ["contextOptions", "storageOptions"]),
                     ...pick(args, ["contextOptions", "storageOptions"]),
                     saveStreamDeltas: stream,
@@ -1062,7 +737,7 @@ export class Agent {
                         promptMessageId: result.promptMessageId,
                         order: result.order,
                         finishReason: await result.finishReason,
-                        warnings: result.warnings,
+                        warnings: await result.warnings,
                         savedMessageIds: result.savedMessages?.map((m) => m._id) ?? [],
                     };
                 }