agents 0.0.0-c3e8618 → 0.0.0-c41ebbc

package/dist/ai-chat-agent.js

@@ -0,0 +1,657 @@
+import {
+  autoTransformMessages
+} from "./chunk-UJVEAURM.js";
+import {
+  Agent
+} from "./chunk-XGMKNUJA.js";
+import "./chunk-ZMMHNOMZ.js";
+import "./chunk-C2OEBJZ2.js";
+import "./chunk-QEVM4BVL.js";
+import "./chunk-BER7KXUJ.js";
+import "./chunk-PR4QN5HX.js";
+
+// src/ai-chat-agent.ts
+import {
+  getToolName,
+  isToolUIPart,
+  parsePartialJson
+} from "ai";
+var decoder = new TextDecoder();
+var AIChatAgent = class extends Agent {
+  constructor(ctx, env) {
+    super(ctx, env);
+    this.sql`create table if not exists cf_ai_chat_agent_messages (
+      id text primary key,
+      message text not null,
+      created_at datetime default current_timestamp
+    )`;
+    const rawMessages = (this.sql`select * from cf_ai_chat_agent_messages` || []).map((row) => {
+      return JSON.parse(row.message);
+    });
+    this.messages = autoTransformMessages(rawMessages);
+    this._chatMessageAbortControllers = /* @__PURE__ */ new Map();
+  }
+  _broadcastChatMessage(message, exclude) {
+    this.broadcast(JSON.stringify(message), exclude);
+  }
+  async onMessage(connection, message) {
+    if (typeof message === "string") {
+      let data;
+      try {
+        data = JSON.parse(message);
+      } catch (_error) {
+        return;
+      }
+      if (data.type === "cf_agent_use_chat_request" /* CF_AGENT_USE_CHAT_REQUEST */ && data.init.method === "POST") {
+        const {
+          // method,
+          // keepalive,
+          // headers,
+          body
+          // we're reading this
+          //
+          // // these might not exist?
+          // dispatcher,
+          // duplex
+        } = data.init;
+        const { messages } = JSON.parse(body);
+        const transformedMessages = autoTransformMessages(messages);
+        this._broadcastChatMessage(
+          {
+            messages: transformedMessages,
+            type: "cf_agent_chat_messages" /* CF_AGENT_CHAT_MESSAGES */
+          },
+          [connection.id]
+        );
+        await this.persistMessages(transformedMessages, [connection.id]);
+        this.observability?.emit(
+          {
+            displayMessage: "Chat message request",
+            id: data.id,
+            payload: {},
+            timestamp: Date.now(),
+            type: "message:request"
+          },
+          this.ctx
+        );
+        const chatMessageId = data.id;
+        const abortSignal = this._getAbortSignal(chatMessageId);
+        return this._tryCatchChat(async () => {
+          const response = await this.onChatMessage(
+            async (_finishResult) => {
+              this._removeAbortController(chatMessageId);
+              this.observability?.emit(
+                {
+                  displayMessage: "Chat message response",
+                  id: data.id,
+                  payload: {},
+                  timestamp: Date.now(),
+                  type: "message:response"
+                },
+                this.ctx
+              );
+            },
+            abortSignal ? { abortSignal } : void 0
+          );
+          if (response) {
+            await this._reply(data.id, response);
+          } else {
+            console.warn(
+              `[AIChatAgent] onChatMessage returned no response for chatMessageId: ${chatMessageId}`
+            );
+            this._broadcastChatMessage(
+              {
+                body: "No response was generated by the agent.",
+                done: true,
+                id: data.id,
+                type: "cf_agent_use_chat_response" /* CF_AGENT_USE_CHAT_RESPONSE */
+              },
+              [connection.id]
+            );
+          }
+        });
+      }
+      if (data.type === "cf_agent_chat_clear" /* CF_AGENT_CHAT_CLEAR */) {
+        this._destroyAbortControllers();
+        this.sql`delete from cf_ai_chat_agent_messages`;
+        this.messages = [];
+        this._broadcastChatMessage(
+          {
+            type: "cf_agent_chat_clear" /* CF_AGENT_CHAT_CLEAR */
+          },
+          [connection.id]
+        );
+      } else if (data.type === "cf_agent_chat_messages" /* CF_AGENT_CHAT_MESSAGES */) {
+        const transformedMessages = autoTransformMessages(data.messages);
+        await this.persistMessages(transformedMessages, [connection.id]);
+      } else if (data.type === "cf_agent_chat_request_cancel" /* CF_AGENT_CHAT_REQUEST_CANCEL */) {
+        this._cancelChatRequest(data.id);
+      }
+    }
+  }
+  async onRequest(request) {
+    return this._tryCatchChat(() => {
+      const url = new URL(request.url);
+      if (url.pathname.endsWith("/get-messages")) {
+        const messages = (this.sql`select * from cf_ai_chat_agent_messages` || []).map((row) => {
+          return JSON.parse(row.message);
+        });
+        return Response.json(messages);
+      }
+      return super.onRequest(request);
+    });
+  }
+  async _tryCatchChat(fn) {
+    try {
+      return await fn();
+    } catch (e) {
+      throw this.onError(e);
+    }
+  }
+  /**
+   * Handle incoming chat messages and generate a response
+   * @param onFinish Callback to be called when the response is finished
+   * @param options.signal A signal to pass to any child requests which can be used to cancel them
+   * @returns Response to send to the client or undefined
+   */
+  async onChatMessage(onFinish, options) {
+    throw new Error(
+      "recieved a chat message, override onChatMessage and return a Response to send to the client"
+    );
+  }
+  /**
+   * Save messages on the server side
+   * @param messages Chat messages to save
+   */
+  async saveMessages(messages) {
+    await this.persistMessages(messages);
+    await this._tryCatchChat(async () => {
+      const response = await this.onChatMessage(() => {
+      });
+      if (response) this._reply(crypto.randomUUID(), response);
+    });
+  }
+  async persistMessages(messages, excludeBroadcastIds = []) {
+    this.sql`delete from cf_ai_chat_agent_messages`;
+    for (const message of messages) {
+      this.sql`insert into cf_ai_chat_agent_messages (id, message) values (${message.id},${JSON.stringify(message)})`;
+    }
+    this.messages = messages;
+    this._broadcastChatMessage(
+      {
+        messages,
+        type: "cf_agent_chat_messages" /* CF_AGENT_CHAT_MESSAGES */
+      },
+      excludeBroadcastIds
+    );
+  }
+  async _reply(id, response) {
+    return this._tryCatchChat(async () => {
+      if (!response.body) {
+        this._broadcastChatMessage({
+          body: "",
+          done: true,
+          id,
+          type: "cf_agent_use_chat_response" /* CF_AGENT_USE_CHAT_RESPONSE */
+        });
+        return;
+      }
+      const reader = response.body.getReader();
+      const message = {
+        id: `assistant_${Date.now()}_${Math.random().toString(36).slice(2, 11)}`,
+        // default
+        role: "assistant",
+        parts: []
+      };
+      let activeTextParts = {};
+      let activeReasoningParts = {};
+      const partialToolCalls = {};
+      function updateDynamicToolPart(options) {
+        const part = message.parts.find(
+          (part2) => part2.type === "dynamic-tool" && part2.toolCallId === options.toolCallId
+        );
+        const anyOptions = options;
+        const anyPart = part;
+        if (part != null) {
+          part.state = options.state;
+          anyPart.toolName = options.toolName;
+          anyPart.input = anyOptions.input;
+          anyPart.output = anyOptions.output;
+          anyPart.errorText = anyOptions.errorText;
+          anyPart.rawInput = anyOptions.rawInput ?? anyPart.rawInput;
+          anyPart.preliminary = anyOptions.preliminary;
+          if (anyOptions.providerMetadata != null && part.state === "input-available") {
+            part.callProviderMetadata = anyOptions.providerMetadata;
+          }
+        } else {
+          message.parts.push({
+            type: "dynamic-tool",
+            toolName: options.toolName,
+            toolCallId: options.toolCallId,
+            state: options.state,
+            input: anyOptions.input,
+            output: anyOptions.output,
+            errorText: anyOptions.errorText,
+            preliminary: anyOptions.preliminary,
+            ...anyOptions.providerMetadata != null ? { callProviderMetadata: anyOptions.providerMetadata } : {}
+          });
+        }
+      }
+      function updateToolPart(options) {
+        const part = message.parts.find(
+          (part2) => isToolUIPart(part2) && part2.toolCallId === options.toolCallId
+        );
+        const anyOptions = options;
+        const anyPart = part;
+        if (part != null) {
+          part.state = options.state;
+          anyPart.input = anyOptions.input;
+          anyPart.output = anyOptions.output;
+          anyPart.errorText = anyOptions.errorText;
+          anyPart.rawInput = anyOptions.rawInput;
+          anyPart.preliminary = anyOptions.preliminary;
+          anyPart.providerExecuted = anyOptions.providerExecuted ?? part.providerExecuted;
+          if (anyOptions.providerMetadata != null && part.state === "input-available") {
+            part.callProviderMetadata = anyOptions.providerMetadata;
+          }
+        } else {
+          message.parts.push({
+            type: `tool-${options.toolName}`,
+            toolCallId: options.toolCallId,
+            state: options.state,
+            input: anyOptions.input,
+            output: anyOptions.output,
+            rawInput: anyOptions.rawInput,
+            errorText: anyOptions.errorText,
+            providerExecuted: anyOptions.providerExecuted,
+            preliminary: anyOptions.preliminary,
+            ...anyOptions.providerMetadata != null ? { callProviderMetadata: anyOptions.providerMetadata } : {}
+          });
+        }
+      }
+      async function updateMessageMetadata(metadata) {
+        if (metadata != null) {
+          const mergedMetadata = message.metadata != null ? { ...message.metadata, ...metadata } : metadata;
+          message.metadata = mergedMetadata;
+        }
+      }
+      try {
+        while (true) {
+          const { done, value } = await reader.read();
+          if (done) {
+            this._broadcastChatMessage({
+              body: "",
+              done: true,
+              id,
+              type: "cf_agent_use_chat_response" /* CF_AGENT_USE_CHAT_RESPONSE */
+            });
+            break;
+          }
+          const chunk = decoder.decode(value);
+          const contentType = response.headers.get("content-type") || "";
+          const isSSE = contentType.includes("text/event-stream");
+          if (isSSE) {
+            const lines = chunk.split("\n");
+            for (const line of lines) {
+              if (line.startsWith("data: ") && line !== "data: [DONE]") {
+                try {
+                  const data = JSON.parse(line.slice(6));
+                  switch (data.type) {
+                    case "text-start": {
+                      const textPart = {
+                        type: "text",
+                        text: "",
+                        providerMetadata: data.providerMetadata,
+                        state: "streaming"
+                      };
+                      activeTextParts[data.id] = textPart;
+                      message.parts.push(textPart);
+                      break;
+                    }
+                    case "text-delta": {
+                      const textPart = activeTextParts[data.id];
+                      textPart.text += data.delta;
+                      textPart.providerMetadata = data.providerMetadata ?? textPart.providerMetadata;
+                      break;
+                    }
+                    case "text-end": {
+                      const textPart = activeTextParts[data.id];
+                      textPart.state = "done";
+                      textPart.providerMetadata = data.providerMetadata ?? textPart.providerMetadata;
+                      delete activeTextParts[data.id];
+                      break;
+                    }
+                    case "reasoning-start": {
+                      const reasoningPart = {
+                        type: "reasoning",
+                        text: "",
+                        providerMetadata: data.providerMetadata,
+                        state: "streaming"
+                      };
+                      activeReasoningParts[data.id] = reasoningPart;
+                      message.parts.push(reasoningPart);
+                      break;
+                    }
+                    case "reasoning-delta": {
+                      const reasoningPart = activeReasoningParts[data.id];
+                      reasoningPart.text += data.delta;
+                      reasoningPart.providerMetadata = data.providerMetadata ?? reasoningPart.providerMetadata;
+                      break;
+                    }
+                    case "reasoning-end": {
+                      const reasoningPart = activeReasoningParts[data.id];
+                      reasoningPart.providerMetadata = data.providerMetadata ?? reasoningPart.providerMetadata;
+                      reasoningPart.state = "done";
+                      delete activeReasoningParts[data.id];
+                      break;
+                    }
+                    case "file": {
+                      message.parts.push({
+                        type: "file",
+                        mediaType: data.mediaType,
+                        url: data.url
+                      });
+                      break;
+                    }
+                    case "source-url": {
+                      message.parts.push({
+                        type: "source-url",
+                        sourceId: data.sourceId,
+                        url: data.url,
+                        title: data.title,
+                        providerMetadata: data.providerMetadata
+                      });
+                      break;
+                    }
+                    case "source-document": {
+                      message.parts.push({
+                        type: "source-document",
+                        sourceId: data.sourceId,
+                        mediaType: data.mediaType,
+                        title: data.title,
+                        filename: data.filename,
+                        providerMetadata: data.providerMetadata
+                      });
+                      break;
+                    }
+                    case "tool-input-start": {
+                      const toolInvocations = message.parts.filter(isToolUIPart);
+                      partialToolCalls[data.toolCallId] = {
+                        text: "",
+                        toolName: data.toolName,
+                        index: toolInvocations.length,
+                        dynamic: data.dynamic
+                      };
+                      if (data.dynamic) {
+                        updateDynamicToolPart({
+                          toolCallId: data.toolCallId,
+                          toolName: data.toolName,
+                          state: "input-streaming",
+                          input: void 0
+                        });
+                      } else {
+                        updateToolPart({
+                          toolCallId: data.toolCallId,
+                          toolName: data.toolName,
+                          state: "input-streaming",
+                          input: void 0
+                        });
+                      }
+                      break;
+                    }
+                    case "tool-input-delta": {
+                      const partialToolCall = partialToolCalls[data.toolCallId];
+                      partialToolCall.text += data.inputTextDelta;
+                      const { value: partialArgs } = await parsePartialJson(
+                        partialToolCall.text
+                      );
+                      if (partialToolCall.dynamic) {
+                        updateDynamicToolPart({
+                          toolCallId: data.toolCallId,
+                          toolName: partialToolCall.toolName,
+                          state: "input-streaming",
+                          input: partialArgs
+                        });
+                      } else {
+                        updateToolPart({
+                          toolCallId: data.toolCallId,
+                          toolName: partialToolCall.toolName,
+                          state: "input-streaming",
+                          input: partialArgs
+                        });
+                      }
+                      break;
+                    }
+                    case "tool-input-available": {
+                      if (data.dynamic) {
+                        updateDynamicToolPart({
+                          toolCallId: data.toolCallId,
+                          toolName: data.toolName,
+                          state: "input-available",
+                          input: data.input,
+                          providerMetadata: data.providerMetadata
+                        });
+                      } else {
+                        updateToolPart({
+                          toolCallId: data.toolCallId,
+                          toolName: data.toolName,
+                          state: "input-available",
+                          input: data.input,
+                          providerExecuted: data.providerExecuted,
+                          providerMetadata: data.providerMetadata
+                        });
+                      }
+                      break;
+                    }
+                    case "tool-input-error": {
+                      if (data.dynamic) {
+                        updateDynamicToolPart({
+                          toolCallId: data.toolCallId,
+                          toolName: data.toolName,
+                          state: "output-error",
+                          input: data.input,
+                          errorText: data.errorText,
+                          providerMetadata: data.providerMetadata
+                        });
+                      } else {
+                        updateToolPart({
+                          toolCallId: data.toolCallId,
+                          toolName: data.toolName,
+                          state: "output-error",
+                          input: void 0,
+                          rawInput: data.input,
+                          errorText: data.errorText,
+                          providerExecuted: data.providerExecuted,
+                          providerMetadata: data.providerMetadata
+                        });
+                      }
+                      break;
+                    }
+                    case "tool-output-available": {
+                      if (data.dynamic) {
+                        const toolInvocations = message.parts.filter(
+                          (part) => part.type === "dynamic-tool"
+                        );
+                        const toolInvocation = toolInvocations.find(
+                          (invocation) => invocation.toolCallId === data.toolCallId
+                        );
+                        if (!toolInvocation)
+                          throw new Error("Tool invocation not found");
+                        updateDynamicToolPart({
+                          toolCallId: data.toolCallId,
+                          toolName: toolInvocation.toolName,
+                          state: "output-available",
+                          input: toolInvocation.input,
+                          output: data.output,
+                          preliminary: data.preliminary
+                        });
+                      } else {
+                        const toolInvocations = message.parts.filter(isToolUIPart);
+                        const toolInvocation = toolInvocations.find(
+                          (invocation) => invocation.toolCallId === data.toolCallId
+                        );
+                        if (!toolInvocation)
+                          throw new Error("Tool invocation not found");
+                        updateToolPart({
+                          toolCallId: data.toolCallId,
+                          toolName: getToolName(toolInvocation),
+                          state: "output-available",
+                          input: toolInvocation.input,
+                          output: data.output,
+                          providerExecuted: data.providerExecuted,
+                          preliminary: data.preliminary
+                        });
+                      }
+                      break;
+                    }
+                    case "tool-output-error": {
+                      if (data.dynamic) {
+                        const toolInvocations = message.parts.filter(
+                          (part) => part.type === "dynamic-tool"
+                        );
+                        const toolInvocation = toolInvocations.find(
+                          (invocation) => invocation.toolCallId === data.toolCallId
+                        );
+                        if (!toolInvocation)
+                          throw new Error("Tool invocation not found");
+                        updateDynamicToolPart({
+                          toolCallId: data.toolCallId,
+                          toolName: toolInvocation.toolName,
+                          state: "output-error",
+                          input: toolInvocation.input,
+                          errorText: data.errorText
+                        });
+                      } else {
+                        const toolInvocations = message.parts.filter(isToolUIPart);
+                        const toolInvocation = toolInvocations.find(
+                          (invocation) => invocation.toolCallId === data.toolCallId
+                        );
+                        if (!toolInvocation)
+                          throw new Error("Tool invocation not found");
+                        updateToolPart({
+                          toolCallId: data.toolCallId,
+                          toolName: getToolName(toolInvocation),
+                          state: "output-error",
+                          input: toolInvocation.input,
+                          rawInput: "rawInput" in toolInvocation ? toolInvocation.rawInput : void 0,
+                          errorText: data.errorText
+                        });
+                      }
+                      break;
+                    }
+                    case "start-step": {
+                      message.parts.push({ type: "step-start" });
+                      break;
+                    }
+                    case "finish-step": {
+                      activeTextParts = {};
+                      activeReasoningParts = {};
+                      break;
+                    }
+                    case "start": {
+                      if (data.messageId != null) {
+                        message.id = data.messageId;
+                      }
+                      await updateMessageMetadata(data.messageMetadata);
+                      break;
+                    }
+                    case "finish": {
+                      await updateMessageMetadata(data.messageMetadata);
+                      break;
+                    }
+                    case "message-metadata": {
+                      await updateMessageMetadata(data.messageMetadata);
+                      break;
+                    }
+                    case "error": {
+                      this._broadcastChatMessage({
+                        error: true,
+                        body: data.errorText ?? JSON.stringify(data),
+                        done: false,
+                        id,
+                        type: "cf_agent_use_chat_response" /* CF_AGENT_USE_CHAT_RESPONSE */
+                      });
+                      break;
+                    }
+                  }
+                  this._broadcastChatMessage({
+                    body: JSON.stringify(data),
+                    done: false,
+                    id,
+                    type: "cf_agent_use_chat_response" /* CF_AGENT_USE_CHAT_RESPONSE */
+                  });
+                } catch (_e) {
+                }
+              }
+            }
+          } else {
+            if (chunk.length > 0) {
+              message.parts.push({ type: "text", text: chunk });
+              this._broadcastChatMessage({
+                body: JSON.stringify({ type: "text-delta", delta: chunk }),
+                done: false,
+                id,
+                type: "cf_agent_use_chat_response" /* CF_AGENT_USE_CHAT_RESPONSE */
+              });
+            }
+          }
+        }
+      } finally {
+        reader.releaseLock();
+      }
+      if (message.parts.length > 0) {
+        await this.persistMessages([...this.messages, message]);
+      }
+    });
+  }
+  /**
+   * For the given message id, look up its associated AbortController
+   * If the AbortController does not exist, create and store one in memory
+   *
+   * returns the AbortSignal associated with the AbortController
+   */
+  _getAbortSignal(id) {
+    if (typeof id !== "string") {
+      return void 0;
+    }
+    if (!this._chatMessageAbortControllers.has(id)) {
+      this._chatMessageAbortControllers.set(id, new AbortController());
+    }
+    return this._chatMessageAbortControllers.get(id)?.signal;
+  }
+  /**
+   * Remove an abort controller from the cache of pending message responses
+   */
+  _removeAbortController(id) {
+    this._chatMessageAbortControllers.delete(id);
+  }
+  /**
+   * Propagate an abort signal for any requests associated with the given message id
+   */
+  _cancelChatRequest(id) {
+    if (this._chatMessageAbortControllers.has(id)) {
+      const abortController = this._chatMessageAbortControllers.get(id);
+      abortController?.abort();
+    }
+  }
+  /**
+   * Abort all pending requests and clear the cache of AbortControllers
+   */
+  _destroyAbortControllers() {
+    for (const controller of this._chatMessageAbortControllers.values()) {
+      controller?.abort();
+    }
+    this._chatMessageAbortControllers.clear();
+  }
+  /**
+   * When the DO is destroyed, cancel all pending requests
+   */
+  async destroy() {
+    this._destroyAbortControllers();
+    await super.destroy();
+  }
+};
+export {
+  AIChatAgent
+};
+//# sourceMappingURL=ai-chat-agent.js.map