@xalia/agent 0.5.8 → 0.6.1

package/src/agent/agent.ts

@@ -1,13 +1,51 @@
 import * as dotenv from "dotenv";
 import { OpenAI } from "openai";
 import { McpServerManager } from "./mcpServerManager";
-import { ChatCompletionContentPart } from "openai/resources.mjs";
 import { strict as assert } from "assert";
 import { ILLM } from "./llm";
 import { AgentProfile, getLogger } from "@xalia/xmcp/sdk";
+import { IAgentEventHandler } from "./iAgentEventHandler";
+import { IContextManager } from "./context";
 export { AgentProfile } from "@xalia/xmcp/sdk";
 
-export type ToolHandler = (args: unknown) => string;
+const MAX_TOOL_CALL_RESPONSE_LENGTH = 4000;
+
+export interface IAgentToolProvider {
+  /**
+   * Any initial setup to be performed by the tool (loading data, etc).  This
+   * function is responsible for registering any tools that this provider
+   * exposes.
+   */
+  setup(agent: Agent): Promise<void>;
+}
+
+export type ToolCallResult = {
+  /**
+   * Response to pass to the LLM
+   */
+  response: string;
+
+  /**
+   * If set, `response` is used in the next round of the Agent loop (if any),
+   * but `overwriteResponse` is passed to the ContextManager to be stored, and
+   * to be used for future LLM invocations.
+   */
+  overwriteResponse?: string;
+
+  /**
+   * If set, the arguments to this tool all, as stored in the context and used
+   * for future LLM invocations, will be overwritten by this value.  This is
+   * intended for the case where the LLM generates a large amount of data to
+   * be saved which we do not want to appear in the context at every
+   * iteration.
+   */
+  overwriteArgs?: string;
+};
+
+export type ToolHandler = (
+  agent: Agent,
+  args: unknown
+) => Promise<ToolCallResult>;
 
 export type McpServerUrls = (name: string) => string;
 
@@ -22,69 +60,64 @@ export type ChatCompletionAssistantMessageParam =
 export type ChatCompletionUserMessageParam =
   OpenAI.ChatCompletionUserMessageParam;
 
-// Role: If content, give it to UI
-export type OnMessageCB = {
-  (msg: string, msgEnd: boolean): Promise<void>;
-};
+export type ChatCompletionToolMessageParam =
+  OpenAI.ChatCompletionToolMessageParam;
 
-// Role: If tool calls, prompt for permission to handle them
-export type OnToolCallCB = {
-  (msg: ChatCompletionMessageToolCall): Promise<boolean>;
-};
+dotenv.config();
+const logger = getLogger();
 
 export interface IConversation {
   userMessage(msg?: string, imageB64?: string): void;
   getConversation(): ChatCompletionMessageParam[];
 
-  // TODO: remove?
-  resetConversation(): void;
-
   getAgentProfile(): AgentProfile;
-  getSystemPrompt(): string;
+
   setSystemPrompt(systemPrompt: string): void;
-  getModel(): string;
   setModel(model: string): void;
 
   shutdown(): Promise<void>;
 }
 
-dotenv.config();
-const logger = getLogger();
+type RegisteredTools = {
+  handler: ToolHandler;
+};
 
 export class Agent implements IConversation {
-  private toolHandlers: { [toolName: string]: ToolHandler } = {};
+  private eventHandler: IAgentEventHandler;
+  private mcpServerManager: McpServerManager;
+  private llm: ILLM;
+  private contextManager: IContextManager;
+
+  /// The full list of tools, ready to pass to the LLM
+  private tools: OpenAI.ChatCompletionTool[] = [];
+
+  /// Handlers for "agent" (or "built-in") tools.  These do not require
+  /// approval from the user.
+  private agentTools = new Map<string, RegisteredTools>();
 
   private constructor(
-    public onMessage: OnMessageCB,
-    public onToolCall: OnToolCallCB,
-    private messages: ChatCompletionMessageParam[],
-    private mcpServerManager: McpServerManager,
-    private tools: OpenAI.ChatCompletionTool[],
-    private llm: ILLM
-  ) {}
-
-  public static async initializeWithLLM(
-    onMessage: OnMessageCB,
-    onToolCall: OnToolCallCB,
-    systemPrompt: string | undefined,
+    eventHandler: IAgentEventHandler,
+    mcpServerManager: McpServerManager,
     llm: ILLM,
-    mcpServerManager?: McpServerManager
-  ): Promise<Agent> {
-    // Initialize messages with system prompt
-    const messages = [
-      {
-        role: "system",
-        content: systemPrompt ?? "You are a helpful assistant",
-      } as OpenAI.ChatCompletionMessageParam,
-    ];
+    contextManager: IContextManager
+  ) {
+    this.eventHandler = eventHandler;
+    this.mcpServerManager = mcpServerManager;
+    this.llm = llm;
+    this.contextManager = contextManager;
+  }
 
+  public static initializeWithLLM(
+    eventHandler: IAgentEventHandler,
+    llm: ILLM,
+    contextManager: IContextManager,
+    mcpServerManager?: McpServerManager
+  ): Agent {
     return new Agent(
-      onMessage,
-      onToolCall,
-      messages,
+      eventHandler,
       mcpServerManager ?? new McpServerManager(),
-      [],
-      llm
+      llm,
+      contextManager
     );
   }
 
@@ -101,24 +134,12 @@ export class Agent implements IConversation {
   }
 
   public getConversation(): ChatCompletionMessageParam[] {
+    const llmMessages = this.contextManager.getLLMContext();
     assert(
-      this.messages[0].role == "system",
+      llmMessages[0].role === "system",
       "first message must have system role"
     );
-    // Return a copy so future modifications to `this.messages` don't impact
-    // the callers copy.
-    return structuredClone(this.messages.slice(1));
-  }
-
-  public setConversation(messages: ChatCompletionMessageParam[]) {
-    assert(this.messages[0].role == "system");
-    assert(
-      messages.length === 0 || messages[0].role != "system",
-      "conversation contains system msg"
-    );
-
-    const newMessages: ChatCompletionMessageParam[] = [this.messages[0]];
-    this.messages = newMessages.concat(structuredClone(messages));
+    return [...llmMessages.slice(1)];
   }
 
   public getMcpServerManager(): McpServerManager {
@@ -143,49 +164,103 @@ export class Agent implements IConversation {
   public async userMessageRaw(
     userMessage: ChatCompletionUserMessageParam
   ): Promise<ChatCompletionMessageParam | undefined> {
-    this.messages.push(userMessage);
-    let completion = await this.chatCompletion();
+    return this.userMessagesRaw([userMessage]);
+  }
+
+  public async userMessagesRaw(
+    userMessages: ChatCompletionUserMessageParam[]
+  ): Promise<ChatCompletionMessageParam | undefined> {
+    // Note: `getLLMContext` returns a copy to we can mutate this array
+    const context = this.contextManager.getLLMContext();
+    const newMessagesIdx = context.length;
+
+    // Add the new user messages
+    context.push(...userMessages);
 
+    let completion = await this.chatCompletion(context);
     let message = completion.choices[0].message;
-    this.messages.push(message);
+    context.push(message);
 
-    // While there are tool calls to make, make them and loop
+    // While there are tool calls to make, invoke them and loop
 
     while (message.tool_calls && message.tool_calls.length > 0) {
+      // TODO: Execute all tool calls in parallel
+
+      // [indexInContext, ToolCallResult][]
+      const toolCallResults: [number, ToolCallResult][] = [];
       for (const toolCall of message.tool_calls ?? []) {
-        const approval = await this.onToolCall(toolCall);
-        if (approval) {
-          try {
-            const result = await this.doToolCall(toolCall);
-            logger.debug(`tool call result ${JSON.stringify(result)}`);
-            this.messages.push(result);
-          } catch (e) {
-            logger.error(`tool call error: ${e}`);
-            this.messages.push({
-              role: "tool",
-              tool_call_id: toolCall.id,
-              content: "Tool call failed.",
-            });
-          }
-        } else {
-          this.messages.push({
-            role: "tool",
-            tool_call_id: toolCall.id,
-            content: "User denied tool use request.",
-          });
+        // Execute the tool call, add the result to the context as an LLM
+        // mesage, and record the index of the message alongside the result in
+        // `toolCallResults`.
+
+        const result = await this.doToolCall(toolCall);
+        toolCallResults.push([context.length, result]);
+        context.push({
+          role: "tool",
+          tool_call_id: toolCall.id,
+          content: result.response,
+        });
+
+        // If the tool call requested that its args be redacted, this can be
+        // done now - before the next LLM invocation.
+
+        if (result.overwriteArgs) {
+          logger.debug(
+            `updating args for toolcall ${toolCall.id}: ${result.overwriteArgs}`
+          );
+          toolCall.function.arguments = result.overwriteArgs;
+          logger.debug(`agent message after update ${JSON.stringify(message)}`);
         }
       }
 
-      completion = await this.chatCompletion();
-      message = completion.choices[0].message;
-      this.messages.push(message);
+      // Now that any args have been overwritten, signal the event handler of
+      // the prevoius completion.
+
+      this.eventHandler.onCompletion(message);
+
+      // Get a new completion using the untouched tool call results.  Note
+      // that, since we are deferring the `onToolCallResult` calls (so they
+      // can be redacted), we must take care that the errors in
+      // `chatCompletion` do not disrupt this, so the caller has a consistent
+      // view of the conversation state.
+
+      try {
+        completion = await this.chatCompletion(context); // CAN THROW
+        message = completion.choices[0].message;
+        context.push(message);
+      } finally {
+        // Now that the tool call results have been passed to the LLM, perform
+        // any updates on them.  Pass the (updated) tool-call-result LLM
+        // messages to the event handler - note, we want to do this even if
+        // the an error occured, so that the caller has an up-to-date picture
+        // of the context state when the error occured.
+
+        toolCallResults.forEach(([indexInContext, tcr]) => {
+          const ctxMsg = context[indexInContext];
+          if (tcr.overwriteResponse) {
+            ctxMsg.content = tcr.overwriteResponse;
+          }
+
+          assert(ctxMsg.role === "tool");
+          this.eventHandler.onToolCallResult(ctxMsg);
+        });
+
+        // Note, if an error DID occur, the ContextManager does not see any of
+        // the new context.
+      }
     }
 
+    // Signal the event handler of the final completion.
+    this.eventHandler.onCompletion(message);
+
+    // Add all new new messages to the context
+    this.contextManager.addMessages(context.slice(newMessagesIdx));
+
     return completion.choices[0].message;
   }
 
   public userMessage(msg?: string, imageB64?: string): void {
-    this.userMessageEx(msg, imageB64);
+    void this.userMessageEx(msg, imageB64);
   }
 
   public getModel(): string {
@@ -197,85 +272,135 @@ export class Agent implements IConversation {
     this.llm.setModel(model);
   }
 
-  /**
-   * Clear the conversation.
-   */
-  public resetConversation() {
-    assert(this.messages.length > 0);
-    // Keep only the system message
-    this.messages.splice(1);
-  }
-
   public getSystemPrompt(): string {
-    assert(this.messages[0].role === "system");
-    return this.messages[0].content as string;
+    return this.contextManager.getAgentPrompt();
   }
 
   /**
    * Set the system prompt
    */
   public setSystemPrompt(systemMsg: string) {
-    assert(this.messages[0].role === "system");
-    this.messages[0].content = systemMsg;
+    this.contextManager.setAgentPrompt(systemMsg);
   }
 
-  async chatCompletion(): Promise<OpenAI.Chat.Completions.ChatCompletion> {
+  async chatCompletion(
+    context: ChatCompletionMessageParam[]
+  ): Promise<OpenAI.Chat.Completions.ChatCompletion> {
+    // Compute the full list of available tools
+
     let tools: OpenAI.ChatCompletionTool[] | undefined;
-    const enabledTools = this.tools.concat(
-      this.mcpServerManager.getOpenAITools()
-    );
+    const mcpTools = this.mcpServerManager.getOpenAITools();
+    const enabledTools = this.tools.concat(mcpTools);
     if (enabledTools.length > 0) {
       tools = enabledTools;
     }
-    // logger.debug(
-    //   `chatCompletion: tools: ${JSON.stringify(tools, undefined, 2)}`
-    // );
     const completion = await this.llm.getConversationResponse(
-      this.messages,
+      context,
       tools,
-      this.onMessage
+      this.eventHandler.onAgentMessage.bind(this.eventHandler)
     );
     logger.debug(`Received chat completion ${JSON.stringify(completion)}`);
     return completion;
   }
 
-  public toolNames(): string[] {
-    return this.mcpServerManager
-      .getOpenAITools()
-      .map((tool) => tool.function.name);
+  public addAgentToolProvider(toolProvider: IAgentToolProvider): Promise<void> {
+    return toolProvider.setup(this);
   }
 
-  public addTool(tool: OpenAI.ChatCompletionTool, handler: ToolHandler) {
+  public addAgentTool(tool: OpenAI.ChatCompletionTool, handler: ToolHandler) {
     const name = tool.function.name;
-    if (this.toolHandlers[name]) {
-      throw `tool ${name} already added`;
+    if (this.agentTools.has(name)) {
+      throw new Error(`tool ${name} already added`);
     }
 
     logger.debug(`Adding tool ${name}`);
 
     this.tools.push(tool);
-    this.toolHandlers[name] = handler;
+    this.agentTools.set(name, { handler });
+  }
+
+  public removeAgentTool(name: string) {
+    if (!this.agentTools.has(name)) {
+      logger.warn(`[removeTool] tool ${name} not present`);
+    }
+
+    // Find idx of the tool in the list
+    const idx = (() => {
+      let idx = 0;
+      while (idx < this.tools.length) {
+        if (this.tools[idx].function.name === name) {
+          return idx;
+        }
+        idx++;
+      }
+      return -1;
+    })();
+    assert(idx > -1);
+
+    // Remove entries
+    this.tools.splice(idx, 1);
+    this.agentTools.delete(name);
   }
 
-  async doToolCall(
+  /**
+   * Handle the details of getting approval (if required), invoking the tool
+   * handler, informing the IAgentEventHandler of the result, and returns the
+   * OpenAI.ChatCompletionToolMessageParam to be used in the conversation.
+   */
+  private async doToolCall(
     toolCall: ChatCompletionMessageToolCall
-  ): Promise<OpenAI.ChatCompletionToolMessageParam> {
-    const name = toolCall.function.name;
-    const args = JSON.parse(toolCall.function.arguments);
-
-    let result: string | undefined = undefined;
-    const handler = this.toolHandlers[name];
-    if (handler) {
-      logger.debug(`  found agent tool ${name} ...`);
-      result = handler(args);
-    } else {
-      result = await this.mcpServerManager.invoke(name, args);
+  ): Promise<ToolCallResult> {
+    // If the tool is and "agent" (internal) tool, we can just execute it.
+    // Otherwise, call the event handler to get permission and invoke the
+    // external tool handler.
+
+    let result: ToolCallResult;
+    try {
+      const toolName = toolCall.function.name;
+      const agentTool = this.agentTools.get(toolName);
+      const isAgentTool = !!agentTool;
+      const approve = await this.eventHandler.onToolCall(toolCall, isAgentTool);
+      if (!approve) {
+        result = { response: "User denied tool request." };
+      } else if (isAgentTool) {
+        // Internal (agent) tool
+        const args: unknown = JSON.parse(toolCall.function.arguments);
+        result = await agentTool.handler(this, args);
+      } else {
+        // McpServer tool call (agentTool === undefined)
+        const args: unknown = JSON.parse(toolCall.function.arguments);
+        result = {
+          response: await this.mcpServerManager.invoke(toolName, args),
+        };
+        logger.debug(`tool call result ${JSON.stringify(result)}`);
+      }
+    } catch (e) {
+      let msg: string;
+      if (e instanceof Error) {
+        msg = e.message;
+      } else if (typeof e === "string") {
+        msg = e;
+      } else {
+        msg = String(e);
+      }
+      logger.error(`tool call error: ${msg}`);
+      result = {
+        response: `tool call error: ${msg}`,
+      };
+    }
+
+    // Final sanity check on the tool call response length.
+    if (result.response.length > MAX_TOOL_CALL_RESPONSE_LENGTH) {
+      logger.warn(
+        "[Agent.doToolCall]: truncating tool call result.response for call:\n" +
+          JSON.stringify(toolCall)
+      );
+      result.response =
+        result.response.slice(0, MAX_TOOL_CALL_RESPONSE_LENGTH) +
+        " ..truncated";
     }
-    return {
-      role: "tool",
-      tool_call_id: toolCall.id,
-      content: result.toString(),
-    };
+
+    return result;
   }
 }
 
@@ -297,7 +422,7 @@ export function createUserMessage(
       return msg;
     }
 
-    const content: ChatCompletionContentPart[] = [];
+    const content: OpenAI.ChatCompletionContentPart[] = [];
     if (msg) {
       content.push({
         type: "text",
@@ -325,3 +450,13 @@ export function createUserMessage(
     name,
   };
 }
+
+export function createUserMessageEnsure(
+  msg?: string,
+  imageB64?: string,
+  name?: string
+): ChatCompletionUserMessageParam {
+  const userMsg = createUserMessage(msg, imageB64, name);
+  assert(userMsg);
+  return userMsg;
+}