@xalia/agent 0.6.7 → 0.6.9

package/src/chat/server/chatContextManager.ts

@@ -2,14 +2,14 @@ import { strict as assert } from "assert";
 
 import { getLogger } from "@xalia/xmcp/sdk";
 
-import { createUserMessage } from "../../agent/agent";
 import {
   AssistantMessageParam,
+  ILLM,
   MessageParam,
   ToolMessageParam,
   UserMessageParam,
 } from "../../agent/llm";
-import { IContextManager } from "../../agent/context";
+import { IContextTransaction } from "../../agent/context";
 import {
   CompressingContextManager,
   createCheckpointMessage,
@@ -39,6 +39,10 @@ import {
   SessionFileDescriptor,
   SessionFileEntry,
 } from "../data/dbSessionFileModels";
+// eslint-disable-next-line max-len
+import { ContextTransactionWithWorkspace } from "../../agent/contextWithWorkspace";
+import { getErrorString } from "./errorUtils";
+import { createUserMessage } from "../../agent/agent";
 
 const logger = getLogger();
 
@@ -57,6 +61,143 @@ export interface ICheckpointWriter {
   writeCheckpoint(checkpoint: SessionCheckpoint): Promise<void>;
 }
 
+export class ChatContextTransaction implements IContextTransaction {
+  private readonly baseTx: ContextTransactionWithWorkspace;
+  private readonly sessionUUID: string;
+  /// Index of final message in the committed context
+  private readonly baseMsgIdx: number | undefined;
+  private readonly startingLLMContextLength: number;
+  private readonly pendingMessages: ConversationMessage[];
+  private curAgentMsgIdx: number;
+
+  constructor(
+    baseTx: ContextTransactionWithWorkspace,
+    sessionUUID: string,
+    baseMsgIdx: number | undefined,
+    pendingUserMessages: ServerUserMessage[],
+    curAgentMsgIdx: number
+  ) {
+    assert(typeof curAgentMsgIdx !== "undefined");
+
+    this.sessionUUID = sessionUUID;
+    this.baseTx = baseTx;
+    this.baseMsgIdx = baseMsgIdx;
+    this.startingLLMContextLength = baseTx.getLLMContextLength();
+    this.pendingMessages = pendingUserMessages;
+    this.curAgentMsgIdx = curAgentMsgIdx;
+  }
+
+  // IContextTransaction.addMessages
+  addMessages(messages: MessageParam[]): number {
+    return this.baseTx.addMessages(messages);
+  }
+
+  // IContextTransaction.addMessage
+  addMessage(message: MessageParam): number {
+    return this.baseTx.addMessage(message);
+  }
+
+  // IContextTransaction.getMessage
+  getMessage(handle: number): MessageParam {
+    return this.baseTx.getMessage(handle);
+  }
+
+  // IContextTransaction.getLLMContext
+  getLLMContext(): MessageParam[] {
+    return this.baseTx.getLLMContext();
+  }
+
+  // IContextTransaction.getLLMContextLength
+  getLLMContextLength(): number {
+    return this.baseTx.getLLMContextLength();
+  }
+
+  getPending(): ConversationMessage[] {
+    return this.pendingMessages;
+  }
+
+  baseMessageIdx(): number | undefined {
+    return this.baseMsgIdx;
+  }
+
+  getBaseTx(): ContextTransactionWithWorkspace {
+    return this.baseTx;
+  }
+
+  /**
+   * Process a FULL Agent message (not chunks from stream).  No message is
+   * required for broadcast as the calling code is expected to broadcast this
+   * as chunks.
+   */
+  processAgentResponse(result: AssistantMessageParam) {
+    // Insert this (full) agent response into the list of agent messages
+    const msg: ServerAgentMessage = {
+      type: "agent_msg",
+      session_id: this.sessionUUID,
+      message_idx: this.getNextMessageSubIdx(),
+      message: result,
+    };
+    this.pendingMessages.push(msg);
+  }
+
+  processAgentMessageChunk(msg: string, end: boolean): ServerAgentMessageChunk {
+    const message: ServerAgentMessageChunk = {
+      type: "agent_msg_chunk",
+      session_id: this.sessionUUID,
+      message_idx: this.getCurrentAgentMessageIdx(),
+      message: msg,
+      end,
+    };
+    return message;
+  }
+
+  processToolCallResult(result: ToolMessageParam): ServerToolCallResult {
+    // Allocate the sub-index for this tool call result.  It should not
+    // have been used already.
+
+    const message_idx = this.getNextMessageSubIdx();
+    const numPending = this.pendingMessages.length;
+    assert(numPending > 0);
+    assert(this.pendingMessages[numPending - 1].message_idx < message_idx);
+
+    const msg: ServerToolCallResult = {
+      type: "tool_call_result",
+      session_id: this.sessionUUID,
+      message_idx,
+      result,
+    };
+    this.pendingMessages.push(msg);
+    return msg;
+  }
+
+  revertAgentResponse(errMsg: string): void {
+    logger.warn(`[ChatContextManager.revertAgentResponse] error: ${errMsg}`);
+
+    // Remove all messages since the user messages were placed on.
+
+    while (this.baseTx.getLLMContextLength() > this.startingLLMContextLength) {
+      const last = this.baseTx.popMessage();
+      assert(last);
+    }
+  }
+
+  newMessages(): MessageParam[] {
+    return this.baseTx.newMessages();
+  }
+
+  private getNextMessageSubIdx(): number {
+    const idx = this.curAgentMsgIdx;
+    this.curAgentMsgIdx += MESSAGE_INDEX_SUB_INCREMENT;
+    return idx;
+  }
+
+  /// Get the current index to use for streaming Agent chunks
+  private getCurrentAgentMessageIdx(): number {
+    assert(typeof this.curAgentMsgIdx !== "undefined");
+    return this.curAgentMsgIdx;
+  }
+}
+
 /**
  * A context manager for Agents interacting with the (potentially multi-user)
  * chat conversations.
@@ -68,27 +209,21 @@ export interface ICheckpointWriter {
  * - Maintain pending user messages, agent loop messages and messages to be
  *   committed to the DB
  */
-export class ChatContextManager
-  implements IContextManager, ISessionFileManagerEventHandler
-{
+export class ChatContextManager implements ISessionFileManagerEventHandler {
   // Including any pending user messages
-  private sessionUUID: string;
-  private conversationMessages: ConversationMessage[];
-  private pendingUserMessages: ServerUserMessage[];
-  private llmContext: CompressingContextManager;
-  private nextMessageIdx: number;
+  private readonly sessionUUID: string;
+  private readonly conversationMessages: ConversationMessage[];
+  private readonly llmContext: CompressingContextManager;
 
-  // State while processing agent messages
-  private startingLLMContextLength: number | undefined;
-  private curAgentMsgIdx: number | undefined; // active agent message
-  private pendingMessages: ConversationMessage[] | undefined;
+  private nextMessageIdx: number;
+  private pendingUserMessages: ServerUserMessage[];
 
   // Compression state
-  private checkpointWriter: ICheckpointWriter;
+  private readonly checkpointWriter: ICheckpointWriter;
   private pendingCompression: boolean;
 
   // FileManager
-  private fileManager: ISessionFileManager;
+  private readonly fileManager: ISessionFileManager;
   private fileManagerDescriptionsDirty: boolean;
 
   constructor(
@@ -97,11 +232,9 @@ export class ChatContextManager
     sessionUUID: string,
     defaultUserName: string,
     checkpoint: SessionCheckpoint | undefined = undefined,
-    compressionAgentUrl: string,
-    compressionAgentModel: string,
-    compressionAgentApiKey: string,
     checkpointWriter: ICheckpointWriter,
-    fileManager: ISessionFileManager
+    fileManager: ISessionFileManager,
+    llm: ILLM
   ) {
     const nextMessageIdx = sessionMessagesToNextIndex(sessionMessages);
     const { messages: llmMessages } = resolveConversationWithCheckpoint(
@@ -113,7 +246,7 @@ export class ChatContextManager
       `[ChatContextManager]: llm messages: ${JSON.stringify(llmMessages)}`
     );
 
-    // Insert a system message placeholder into the context
+    const getLLM = () => Promise.resolve(llm);
 
     this.sessionUUID = sessionUUID;
     this.conversationMessages = sessionMessagesToConversationMessages(
@@ -125,14 +258,9 @@ export class ChatContextManager
     this.llmContext = new CompressingContextManager(
       systemPrompt,
       llmMessages,
-      compressionAgentUrl,
-      compressionAgentModel,
-      compressionAgentApiKey
+      getLLM
     );
     this.nextMessageIdx = nextMessageIdx;
-    this.startingLLMContextLength = undefined;
-    this.curAgentMsgIdx = undefined;
-    this.pendingMessages = undefined;
     this.pendingCompression = false;
     this.checkpointWriter = checkpointWriter;
     this.fileManager = fileManager;
@@ -140,16 +268,6 @@ export class ChatContextManager
     this.fileManagerDescriptionsDirty = true;
   }
 
-  // IContextManager.addMessages
-  addMessages(messages: MessageParam[]): void {
-    this.llmContext.addMessages(messages);
-  }
-
-  // IContextManager.addMessage
-  addMessage(message: MessageParam): void {
-    this.llmContext.addMessage(message);
-  }
-
   // IContextManager.getLLMContext
   getLLMContext(): MessageParam[] {
     if (this.fileManagerDescriptionsDirty) {
@@ -166,6 +284,7 @@ export class ChatContextManager
   getAgentPrompt(): string {
     return this.llmContext.getAgentPrompt();
   }
+
   // IContextManager.setAgentPrompt
   setAgentPrompt(prompt: string): void {
     this.llmContext.setAgentPrompt(prompt);
@@ -247,24 +366,21 @@ export class ChatContextManager
   // class manage the interaction with the agent and ensure this process only
   // happens one-at-a-time.
 
-  startAgentResponse(msgs: ServerUserMessage[]): {
-    llmUserMessages: UserMessageParam[];
+  async startAgentResponse(msgs: ServerUserMessage[]): Promise<{
+    contextTx: ChatContextTransaction;
     agentFirstChunk: ServerAgentMessageChunk;
-  } {
-    // Sanity check the state
-
-    assert(
-      typeof this.startingLLMContextLength === "undefined",
-      "already processing"
-    );
-    assert(typeof this.pendingMessages === "undefined", "already processing");
-    assert(typeof this.curAgentMsgIdx === "undefined", "already processing");
-
-    // Sanity check the messages
+  }> {
+    // Sanity check the state - the incoming user messages should match the
+    // pending user messages.
 
     const numMessages = this.pendingUserMessages.length;
     assert(numMessages > 0);
-    assert(msgs.length === this.pendingUserMessages.length);
+    if (msgs.length !== this.pendingUserMessages.length) {
+      throw new Error(
+        `length mismatch: msgs: ${JSON.stringify(msgs)}, ` +
+          `this.pendingUserMsgs: ${JSON.stringify(this.pendingUserMessages)}`
+      );
+    }
     assert(msgs[0].message_idx === this.pendingUserMessages[0].message_idx);
     assert(
       msgs[numMessages - 1].message_idx ===
@@ -275,10 +391,9 @@ export class ChatContextManager
     // agent messages and tool calls.
 
     const pendingUserMessages = this.pendingUserMessages;
+    const baseMsgIdx = this.lastMessageIdx();
+    const curAgentMsgIdx = this.getNextMessageIdx();
     this.pendingUserMessages = [];
-    this.startingLLMContextLength = this.llmContext.getCommittedLength();
-    this.curAgentMsgIdx = this.getNextMessageIdx();
-    this.pendingMessages = pendingUserMessages as ConversationMessage[];
 
     // Compute the new llm messages
 
@@ -294,47 +409,50 @@ export class ChatContextManager
       }
     }
 
-    return {
-      llmUserMessages,
-      agentFirstChunk: {
-        type: "agent_msg_chunk",
-        session_id: this.sessionUUID,
-        message_idx: this.curAgentMsgIdx,
-        message: "",
-        end: false,
-      },
+    // Return the context tx and first ServerAgentMessageChunk
+
+    const agentFirstChunk: ServerAgentMessageChunk = {
+      type: "agent_msg_chunk",
+      session_id: this.sessionUUID,
+      message_idx: curAgentMsgIdx,
+      message: "",
+      end: false,
     };
-  }
 
-  endAgentResponse(): SessionMessage[] {
-    assert(
-      typeof this.startingLLMContextLength !== "undefined",
-      "agent response not started (startingLLMContextLength)"
-    );
-    assert(
-      typeof this.pendingMessages !== "undefined",
-      "agent response not started (pendingMessages)"
-    );
-    assert(
-      typeof this.curAgentMsgIdx !== "undefined",
-      "agent response not started (curAgentMsgIdx)"
+    const baseTx = await this.llmContext.startTx(llmUserMessages);
+    const contextTx = new ChatContextTransaction(
+      baseTx,
+      this.sessionUUID,
+      baseMsgIdx,
+      pendingUserMessages,
+      curAgentMsgIdx
     );
+    return { agentFirstChunk, contextTx };
+  }
 
-    const numPending = this.pendingMessages.length;
+  async endAgentResponse(tx: IContextTransaction): Promise<SessionMessage[]> {
+    assert(tx instanceof ChatContextTransaction);
+    if (tx.baseMessageIdx() !== this.lastMessageIdx()) {
+      throw new Error(
+        `Tx stale? tx.baseMessageIdx=${String(tx.baseMessageIdx())}, ` +
+          `this.conv: ${JSON.stringify(this.conversationMessages)}`
+      );
+    }
+
+    const pending = tx.getPending();
+    const numPending = pending.length;
     assert(numPending > 0, "no pending"); // at least 1 user message
 
     // Compute DB messages
 
-    const newSessionMessages = chatMessagesToSessionMessages(
-      this.pendingMessages
-    );
-    const newLLMMessages = this.llmContext.getPending();
+    const newSessionMessages = chatMessagesToSessionMessages(pending);
+    const newLLMMessages = tx.newMessages();
 
     const messageListError = (error: string) => {
       throw new Error(
         `${error}:` +
           `\n  newSessionMessages: ${JSON.stringify(newSessionMessages)}` +
-          `\n  this.pendingMessages: ${JSON.stringify(this.pendingMessages)}` +
+          `\n  pending: ${JSON.stringify(pending)}` +
           `\n  newLLMMessages: ${JSON.stringify(newLLMMessages)}`
       );
     };
@@ -353,7 +471,7 @@ export class ChatContextManager
 
     for (let i = 0; i < numPending; ++i) {
       const sMsg = newSessionMessages[i];
-      const pMsg = this.pendingMessages[i];
+      const pMsg = pending[i];
       const lMsg = newLLMMessages[i];
 
       if (sMsg.content.role !== lMsg.role) {
@@ -379,11 +497,8 @@ export class ChatContextManager
     // Update our internal state and return the SessionMessages to write to
     // the DB
 
-    this.llmContext.commit();
-    this.conversationMessages.push(...this.pendingMessages);
-    this.startingLLMContextLength = undefined;
-    this.pendingMessages = undefined;
-    this.curAgentMsgIdx = undefined;
+    await this.llmContext.commit(tx.getBaseTx());
+    this.conversationMessages.push(...pending);
 
     // Kick off a compression?
     this.checkCompression();
@@ -391,98 +506,12 @@ export class ChatContextManager
     return newSessionMessages;
   }
 
-  /**
-   * End the Agent message session with an error.  Caller should not call
-   * `endAgentResponse` after calling this function.
-   *
-   * This function checks that nothing has been entered into the LLM context,
-   * and drops any new user messages or responses before the error.
-   */
-  revertAgentResponse(errMsg: string): void {
-    logger.warn(`[ChatContextManager.revertAgentResponse] error: ${errMsg}`);
-
-    assert(typeof this.startingLLMContextLength !== "undefined");
-    assert(typeof this.pendingMessages !== "undefined");
-    assert(typeof this.curAgentMsgIdx !== "undefined");
-
-    // Sanity check that no new messages were put into the context (The Agent
-    // is expected to only call `addMessage(s)` at the end of the Agent loop.
-    // (Note, we don't check for equality here, just in case the context was
-    // compressed while the Agent was executing).
-
-    const contextLength = this.llmContext.getCommittedLength();
-    if (contextLength > this.startingLLMContextLength) {
-      logger.error(
-        "[ChatContextManager.revertAgentResponse] llmContext has grown " +
-          `despite Agent error (${String(contextLength)}, ` +
-          `${String(this.startingLLMContextLength)})`
-      );
+  private lastMessageIdx(): number | undefined {
+    const numMsgs = this.conversationMessages.length;
+    if (numMsgs > 0) {
+      return this.conversationMessages[numMsgs - 1].message_idx;
     }
-
-    // We simply reset the state, dropping any pending messages.
-
-    this.startingLLMContextLength = undefined;
-    this.pendingMessages = undefined;
-    this.curAgentMsgIdx = undefined;
-  }
-
-  processAgentMessage(msg: string, end: boolean): ServerAgentMessageChunk {
-    assert(typeof this.startingLLMContextLength !== "undefined");
-    assert(typeof this.pendingMessages !== "undefined");
-    assert(typeof this.curAgentMsgIdx !== "undefined");
-
-    const message: ServerAgentMessageChunk = {
-      type: "agent_msg_chunk",
-      session_id: this.sessionUUID,
-      message_idx: this.getCurrentAgentMessageIdx(),
-      message: msg,
-      end,
-    };
-    return message;
-  }
-
-  /**
-   * Process a FULL Agent message (not chunks from stream).  No message is
-   * required for broadcast as the calling code is expected to broadcast this
-   * as chunks.
-   */
-  processAgentResponse(result: AssistantMessageParam) {
-    assert(typeof this.startingLLMContextLength !== "undefined");
-    assert(typeof this.pendingMessages !== "undefined");
-    assert(typeof this.curAgentMsgIdx !== "undefined");
-
-    // Insert this (full) agent response into the list of agent messages
-
-    const msg: ServerAgentMessage = {
-      type: "agent_msg",
-      session_id: this.sessionUUID,
-      message_idx: this.getNextMessageSubIdx(),
-      message: result,
-    };
-    this.pendingMessages.push(msg);
-  }
-
-  processToolCallResult(result: ToolMessageParam): ServerToolCallResult {
-    assert(typeof this.startingLLMContextLength !== "undefined");
-    assert(typeof this.pendingMessages !== "undefined");
-    assert(typeof this.curAgentMsgIdx !== "undefined");
-
-    // Allocate the sub-index for this tool call result.  It should not
-    // have been used already.
-
-    const message_idx = this.getNextMessageSubIdx();
-    const numPending = this.pendingMessages.length;
-    assert(numPending > 0);
-    assert(this.pendingMessages[numPending - 1].message_idx < message_idx);
-
-    const msg: ServerToolCallResult = {
-      type: "tool_call_result",
-      session_id: this.sessionUUID,
-      message_idx,
-      result,
-    };
-    this.pendingMessages.push(msg);
-    return msg;
+    return undefined;
   }
 
   private getNextMessageIdx(): number {
@@ -503,22 +532,6 @@ export class ChatContextManager
     this.nextMessageIdx = messageIdx;
   }
 
-  /// Get the current index to use for streaming Agent chunks
-  private getCurrentAgentMessageIdx(): number {
-    assert(typeof this.pendingMessages !== "undefined");
-    assert(typeof this.curAgentMsgIdx !== "undefined");
-    return this.curAgentMsgIdx;
-  }
-
-  private getNextMessageSubIdx(): number {
-    assert(typeof this.pendingMessages !== "undefined");
-    assert(typeof this.curAgentMsgIdx !== "undefined");
-
-    const idx = this.curAgentMsgIdx;
-    this.curAgentMsgIdx += MESSAGE_INDEX_SUB_INCREMENT;
-    return idx;
-  }
-
   private checkCompression(): void {
     if (this.pendingCompression) {
       return;
@@ -526,7 +539,7 @@ export class ChatContextManager
 
     // TODO: track tokens and use that to trigger compression
 
-    const numCommitted = this.llmContext.getCommittedLength();
+    const numCommitted = this.llmContext.numMessages();
     if (numCommitted < COMPRESSION_TRIGGER_NUM_MESSAGES) {
       return;
     }
@@ -551,7 +564,7 @@ export class ChatContextManager
       await this.checkpointWriter.writeCheckpoint(checkpoint);
     } catch (err: unknown) {
       logger.warn(
-        `[runCompression] error during compression: ${JSON.stringify(err)}`
+        `[runCompression] error during compression: ${getErrorString(err)}`
       );
     } finally {
       this.pendingCompression = false;

package/src/chat/server/connectionManager.ts

@@ -25,7 +25,7 @@ export interface IUserConnectionManager<ServerMsgT> {
    * Send message to all active connections of specific users.
    * Handles user-to-connection routing internally.
    */
-  sendToUsers(userIds: Set<string>, message: ServerMsgT): void;
+  sendToUsers(userIds: Set<string> | string[], message: ServerMsgT): void;
 
   /**
    * Send message to a specific connection.

package/src/chat/server/imageGeneratorTools.ts

@@ -1,25 +1,40 @@
 import { getLogger } from "@xalia/xmcp/sdk";
 
-import { Agent, IAgentToolProvider, ToolCallResult } from "../../agent/agent";
+import { AgentEx, IAgentToolProvider, ToolCallResult } from "../../agent/agent";
 import { ToolDescriptor } from "../../agent/llm";
 import { ImageGenerator } from "../../agent/imageGenerator";
-
 import {
   ChatSessionFileManager,
   ToolCallResultWithFileRef,
 } from "./sessionFileManager";
 import { makeParseArgsFn } from "./tools";
 import { getSessionFileMimeTypeFromDataUrl } from "../data/dbSessionFileModels";
+import { getOpenAIClientParams } from "./openAIRouterLLM";
+import { DEFAULT_IMAGE_GEN_MODEL, ImageGenLLM } from "../../agent/imageGenLLM";
+import { createSpecializedLLM } from "../../agent/agentUtils";
+import { IPlatform } from "../../agent/iplatform";
 
 const logger = getLogger();
 
-function createImageGenerator(
-  llmUrl: string,
-  llmApiKey: string
+async function createImageGenerator(
+  platform: IPlatform
 ): Promise<ImageGenerator> {
-  const imageGenModel = process.env["GEN_IMAGE_MODEL"];
-  logger.debug(`[genImageFileTool] model: ${imageGenModel || "(default)"}`);
-  return ImageGenerator.init(llmUrl, llmApiKey, imageGenModel);
+  const imageGenModel =
+    process.env["GEN_IMAGE_MODEL"] || DEFAULT_IMAGE_GEN_MODEL;
+
+  // Allow the image generator to use a "specialized" or "debug" LLM if
+  // requested.  Otherwise query the router maps for the url and api key and
+  // instantiate an ImageGenLLM.
+
+  let llm = await createSpecializedLLM(imageGenModel, platform);
+  if (!llm) {
+    const { apiKey, baseURL } = getOpenAIClientParams(imageGenModel);
+    logger.debug(
+      `[genImageFileTool] model: ${imageGenModel}, url: ${baseURL}"}`
+    );
+    llm = new ImageGenLLM(apiKey, baseURL, imageGenModel);
+  }
+  return ImageGenerator.init(llm);
 }
 
 // gen_image
@@ -28,8 +43,7 @@ function createImageGenerator(
 // decide how/whether to use the image.
 
 export async function genImageTool(
-  llmUrl: string,
-  llmApiKey: string
+  platform: IPlatform
 ): Promise<IAgentToolProvider> {
   const GEN_IMAGE_DESC: ToolDescriptor = {
     type: "function",
@@ -52,12 +66,12 @@ export async function genImageTool(
       },
     },
   } as const;
-  const imageGenerator = await createImageGenerator(llmUrl, llmApiKey);
+  const imageGenerator = await createImageGenerator(platform);
   const getPromptInputImg = makeParseArgsFn(
     ["prompt"] as const,
     ["input_img"] as const
   );
-  const toolFn = async (_: Agent, args: unknown): Promise<ToolCallResult> => {
+  const toolFn = async (_: AgentEx, args: unknown): Promise<ToolCallResult> => {
     const { prompt, input_img } = getPromptInputImg(args);
     const image = await imageGenerator.generate(prompt, input_img);
     return {
@@ -71,7 +85,7 @@ export async function genImageTool(
 
   return {
     // eslint-disable-next-line @typescript-eslint/require-await
-    setup: async (agent: Agent) => {
+    setup: async (agent: AgentEx) => {
       agent.addAgentTool(GEN_IMAGE_DESC, toolFn);
     },
   };
@@ -80,8 +94,7 @@ export async function genImageTool(
 // gen_image_file
 
 export async function genImageFileTool(
-  llmUrl: string,
-  llmApiKey: string,
+  platform: IPlatform,
   fileManager: ChatSessionFileManager
 ): Promise<IAgentToolProvider> {
   const GEN_IMAGE_FILE_DESC: ToolDescriptor = {
@@ -114,13 +127,13 @@ export async function genImageFileTool(
     },
   } as const;
 
-  const imageGenerator = await createImageGenerator(llmUrl, llmApiKey);
+  const imageGenerator = await createImageGenerator(platform);
   const getPromptNameSummary = makeParseArgsFn(
     ["prompt", "name", "summary"] as const,
     ["input_name"] as const
   );
   const toolFn = async (
-    _: Agent,
+    _: AgentEx,
     args: unknown
   ): Promise<ToolCallResultWithFileRef> => {
     const { prompt, name, summary, input_name } = getPromptNameSummary(args);
@@ -140,7 +153,7 @@ export async function genImageFileTool(
 
   return {
     // eslint-disable-next-line @typescript-eslint/require-await
-    setup: async (agent: Agent) => {
+    setup: async (agent: AgentEx) => {
       agent.addAgentTool(GEN_IMAGE_FILE_DESC, toolFn);
     },
   };