@xalia/agent 0.6.10 → 0.6.11

package/src/agent/agent.ts

@@ -1,699 +0,0 @@
-import { strict as assert } from "assert";
-
-export { AgentProfile } from "@xalia/xmcp/sdk";
-import { AgentProfile, getLogger } from "@xalia/xmcp/sdk";
-
-import * as openai from "./openAI";
-import { McpServerManager } from "./mcpServerManager";
-import {
-  ILLM,
-  ToolDescriptor,
-  MessageParam,
-  AssistantMessageParam,
-  UserMessageParam,
-  MessageToolCall,
-  Completion,
-  Message,
-  ToolMessageParam,
-} from "./llm";
-import { IAgentEventHandler } from "./iAgentEventHandler";
-import { IContextManager, IContextTransaction } from "./context";
-import { ChatCompletionContentPartImage } from "openai/resources";
-
-import { MAX_TOOL_CALL_RESPONSE_LENGTH } from "./toolSettings";
-
-export const DEFAULT_LLM_URL = "http://localhost:5001/v1";
-
-/**
- * The message to append to the agent output if the agent is interrupted by a
- * signal from the user.
- */
-export const USER_STOP_MESSAGE = " AGENT INTERRUPTED";
-
-/**
- * An agent's response, with optional extra image data.
- * `ChatCompletionMessageParam` may one day be updated to allow image data (as
- * it does for audio data), but for now image data is not included and so we
- * keep it separate.
- */
-export type AssistantResponse = {
-  message: MessageParam;
-  images?: ChatCompletionContentPartImage[];
-};
-
-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: AgentEx): Promise<void>;
-}
-
-export type ToolCallResult<
-  Meta extends Record<string, string> = Record<string, string>,
-> = {
-  /**
-   * Response to pass to the LLM
-   */
-  response: string;
-
-  /**
-   * An (optional) structured response.
-   */
-  structuredContent?: unknown;
-
-  /**
-   * (Optional) Application-specific meta data about the tool call result.
-   */
-  _meta?: Meta;
-
-  /**
-   * 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: AgentEx,
-  args: unknown
-) => Promise<ToolCallResult>;
-
-const logger = getLogger();
-
-export interface IConversation {
-  userMessage(msg?: string, imageB64?: string): void;
-  getConversation(): MessageParam[];
-
-  getAgentProfile(): AgentProfile;
-
-  setSystemPrompt(systemPrompt: string): void;
-  setModel(model: string): void;
-
-  shutdown(): Promise<void>;
-}
-
-type RegisteredTools = {
-  handler: ToolHandler;
-};
-
-/**
- * An agent attached to an ILLM which updates a context via an
- * IContextTransaction interface (where IContextTransaction is like a DB tx or
- * DB writer, for staging changes and reading back state as-if those changes
- * were applied).
- */
-export class AgentEx {
-  mcpServerManager: McpServerManager;
-  llm: ILLM;
-
-  /// Flag to stop the Agent loop.
-  stopFlag: boolean;
-  /// Function to stop the LLM (only present while it is active)
-  stopFn: ((msg: string) => void) | undefined;
-
-  /// The full list of tools, ready to pass to the LLM
-  private tools: ToolDescriptor[] = [];
-
-  /// Handlers for "agent" (or "built-in") tools.  These do not require
-  /// approval from the user.
-  private agentTools = new Map<string, RegisteredTools>();
-
-  constructor(mcpServerManager: McpServerManager, llm: ILLM) {
-    this.mcpServerManager = mcpServerManager;
-    this.llm = llm;
-    this.stopFlag = false;
-    this.stopFn = undefined;
-  }
-
-  public async shutdown(): Promise<void> {
-    this.stop("shutting down");
-    return this.mcpServerManager.shutdown();
-  }
-
-  public stop(msg?: string) {
-    this.stopFlag = true;
-    if (this.stopFn) {
-      this.stopFn(msg || USER_STOP_MESSAGE);
-    }
-  }
-
-  public getMcpServerManager(): McpServerManager {
-    return this.mcpServerManager;
-  }
-
-  // TODO: rename
-  public async userMessagesRaw(
-    contextTx: IContextTransaction,
-    eventHandler: IAgentEventHandler
-  ): Promise<AssistantResponse | undefined> {
-    this.stopFlag = false;
-
-    // New user messages have already been added to the `contextTx`.
-
-    // Image and audio handling
-    //
-    // `ChatCompletions` (responses from the LLM) can contain `audio` and
-    // `images` tags.  However, the `ChatCompletionMessageParam` type does not
-    // allow for "assistant" messages with images / audio.
-    //
-    // As such, our current approach is to extract all assistant-generated
-    // media and return it separately.
-
-    const images: ChatCompletionContentPartImage[] = [];
-
-    // We convert the `ChatCompletionsMessage` into a
-    // `ChatCompletionAssistantMessageParam` and extract image data.
-
-    let completion = await this.chatCompletion(
-      contextTx.getLLMContext(),
-      eventHandler
-    );
-    let message = this.processCompletion(completion, images, eventHandler);
-    contextTx.addMessage(message);
-
-    // While there are tool calls to make, invoke them and loop
-
-    while (message.tool_calls && message.tool_calls.length > 0) {
-      // Signal the event handler of the assistant message with tool calls
-      // BEFORE processing tool results. This ensures the order of messages
-      // in pendingMessages matches the order in the LLM context:
-      // [user, assistant(tool_calls), tool_result, assistant(final)]
-      eventHandler.onCompletion(message);
-
-      // TODO: Execute all tool calls in parallel
-
-      // [indexInContext, ToolCallResult][]
-      const toolCallResults: [number, ToolCallResult][] = [];
-      for (const toolCall of message.tool_calls ?? []) {
-        // 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, eventHandler);
-        const toolResult: ToolMessageParam = {
-          role: "tool",
-          tool_call_id: toolCall.id,
-          content: result.response,
-          ...(result._meta ? { _meta: result._meta } : {}),
-          ...(result.structuredContent
-            ? { structuredContent: result.structuredContent }
-            : {}),
-        };
-
-        const toolResultHandle = contextTx.addMessage(toolResult);
-        toolCallResults.push([toolResultHandle, result]);
-
-        // Immediately broadcast the tool result to the frontend for UI
-        // feedback. This ensures the frontend knows the tool executed
-        // successfully without waiting for the next LLM completion to
-        // finish streaming
-        eventHandler.onToolCallResult(toolResult);
-
-        // 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)}`);
-        }
-      }
-
-      // 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(
-          contextTx.getLLMContext(),
-          eventHandler
-        );
-        message = this.processCompletion(completion, images, eventHandler);
-        contextTx.addMessage(message);
-      } finally {
-        // Now that the tool call results have been passed to the LLM, perform
-        // any updates on them if overwriteResponse was requested. If so, send
-        // the updated tool result to the frontend to replace the original.
-
-        toolCallResults.forEach(([handle, tcr]) => {
-          if (tcr.overwriteResponse) {
-            const ctxMsg = contextTx.getMessage(handle);
-            ctxMsg.content = tcr.overwriteResponse;
-            assert(ctxMsg.role === "tool");
-            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.
-    eventHandler.onCompletion(message);
-
-    return { message, images: images.length === 0 ? undefined : images };
-  }
-
-  async chatCompletion(
-    context: MessageParam[],
-    eventHandler: IAgentEventHandler
-  ): Promise<Completion> {
-    if (this.stopFlag) {
-      return {
-        id: "user_stopped",
-        choices: [
-          {
-            finish_reason: "stop",
-            index: 0,
-            message: {
-              content: USER_STOP_MESSAGE,
-              role: "assistant",
-              refusal: null,
-            },
-            logprobs: null,
-          },
-        ],
-        created: Date.now(),
-        model: this.llm.getModel(),
-        object: "chat.completion",
-      };
-    }
-
-    // Compute the full list of available tools
-
-    let tools: ToolDescriptor[] | undefined;
-    const mcpTools = this.mcpServerManager.getOpenAITools();
-    logger.debug(`[chatCompletion] mcpTools: ${JSON.stringify(mcpTools)}`);
-    const enabledTools = this.tools.concat(mcpTools);
-    if (enabledTools.length > 0) {
-      tools = enabledTools;
-    }
-    logger.debug(`[chatCompletion] tools: ${JSON.stringify(tools)}`);
-
-    // Log system prompt length
-    if (context.length > 0 && context[0].role === "system") {
-      const systemPrompt = context[0].content as string;
-      logger.info(
-        `[chatCompletion] System prompt length: ${String(systemPrompt.length)}`
-      );
-    }
-
-    const { stop, completion: completionP } =
-      await this.llm.getConversationResponse(
-        context,
-        tools,
-        eventHandler.onAgentMessage.bind(eventHandler),
-        eventHandler.onReasoning.bind(eventHandler)
-      );
-
-    this.stopFn = stop;
-    const completion = await completionP;
-    this.stopFn = undefined;
-
-    logger.debug(`Received chat completion ${JSON.stringify(completion)}`);
-    return completion;
-  }
-
-  public addAgentToolProvider(toolProvider: IAgentToolProvider): Promise<void> {
-    return toolProvider.setup(this);
-  }
-
-  public addAgentTool(tool: ToolDescriptor, handler: ToolHandler) {
-    const name = tool.function.name;
-    if (this.agentTools.has(name)) {
-      throw new Error(`tool ${name} already added`);
-    }
-
-    logger.debug(`Adding tool ${name}`);
-
-    this.tools.push(tool);
-    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);
-  }
-
-  /**
-   * Handle the details of getting approval (if required), invoking the tool
-   * handler, informing the IAgentEventHandler of the result, and returns the
-   * ChatCompletionToolMessageParam to be used in the conversation.
-   */
-  private async doToolCall(
-    toolCall: MessageToolCall,
-    eventHandler: IAgentEventHandler
-  ): Promise<ToolCallResult> {
-    if (this.stopFlag) {
-      return { response: USER_STOP_MESSAGE };
-    }
-
-    // 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;
-
-      if (isAgentTool) {
-        // Internal (agent) tool
-        if (!(await eventHandler.onToolCall(toolCall, true))) {
-          result = { response: "User denied tool request." };
-        } else {
-          const args: unknown = JSON.parse(toolCall.function.arguments || "{}");
-          result = await agentTool.handler(this, args);
-        }
-      } else {
-        // McpServer tool call (agentTool === undefined).  Sanity check the
-        // tool call data, get approval, and then invoke.
-
-        const args: unknown = JSON.parse(toolCall.function.arguments || "{}");
-        const tc = this.mcpServerManager.verifyToolCall(toolName, args);
-        if (!(await eventHandler.onToolCall(toolCall, false))) {
-          result = { response: "User denied tool request." };
-        } else {
-          result = await this.mcpServerManager.invoke(tc);
-        }
-        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 result;
-  }
-
-  private processCompletion(
-    completion: Completion,
-    images: ChatCompletionContentPartImage[],
-    eventHandler: IAgentEventHandler
-  ): AssistantMessageParam {
-    // Add any images into the list, and call the event handler
-
-    const compMessage = completion.choices[0].message;
-    if (compMessage.images) {
-      for (const image of compMessage.images) {
-        eventHandler.onImage(image);
-        images.push(image);
-      }
-    }
-
-    return completionToAssistantMessageParam(compMessage);
-  }
-}
-
-/**
- * Higher-level abstraction over AgentEx, which abstracts out the context
- * transactions.  A single agent is associated with an IContextManager and
- * internally creates and commits transactions during each call to
- * `userMessage*`.
- */
-export class Agent implements IConversation {
-  private eventHandler: IAgentEventHandler;
-  private contextManager: IContextManager;
-  private agentEx: AgentEx;
-
-  private constructor(
-    eventHandler: IAgentEventHandler,
-    mcpServerManager: McpServerManager,
-    llm: ILLM,
-    contextManager: IContextManager
-  ) {
-    this.eventHandler = eventHandler;
-    this.contextManager = contextManager;
-    this.agentEx = new AgentEx(mcpServerManager, llm);
-  }
-
-  public static initializeWithLLM(
-    eventHandler: IAgentEventHandler,
-    llm: ILLM,
-    contextManager: IContextManager,
-    mcpServerManager?: McpServerManager
-  ): Agent {
-    return new Agent(
-      eventHandler,
-      mcpServerManager ?? new McpServerManager(),
-      llm,
-      contextManager
-    );
-  }
-
-  public async shutdown(): Promise<void> {
-    return this.agentEx.shutdown();
-  }
-
-  public getAgentProfile(): AgentProfile {
-    return new AgentProfile(
-      this.agentEx.llm.getModel(),
-      this.getSystemPrompt(),
-      this.agentEx.mcpServerManager.getMcpServerSettings()
-    );
-  }
-
-  public getConversation(): MessageParam[] {
-    const llmMessages = this.contextManager.getLLMContext();
-    assert(
-      llmMessages[0].role === "system",
-      "first message must have system role"
-    );
-    return [...llmMessages.slice(1)];
-  }
-
-  public getMcpServerManager(): McpServerManager {
-    return this.agentEx.mcpServerManager;
-  }
-
-  /**
-   * Like `userMessage`, but can be awaited, and accepts the user name.
-   */
-  public async userMessageEx(
-    msg?: string,
-    imageB64?: string,
-    name?: string
-  ): Promise<AssistantResponse | undefined> {
-    const userMessage = createUserMessage(msg, imageB64, name);
-    if (!userMessage) {
-      return undefined;
-    }
-    return this.userMessageRaw(userMessage);
-  }
-
-  public async userMessageRaw(
-    userMessage: UserMessageParam
-  ): Promise<AssistantResponse | undefined> {
-    return this.userMessagesRaw([userMessage]);
-  }
-
-  public async userMessagesRaw(
-    userMessages: UserMessageParam[]
-  ): Promise<AssistantResponse | undefined> {
-    const tx = await this.contextManager.startTx(userMessages);
-    const result = await this.agentEx.userMessagesRaw(tx, this.eventHandler);
-    await this.contextManager.commit(tx);
-    return result;
-  }
-
-  public userMessage(msg?: string, imageB64?: string): void {
-    void this.userMessageEx(msg, imageB64);
-  }
-
-  public getModel(): string {
-    return this.agentEx.llm.getModel();
-  }
-
-  public setModel(model: string) {
-    logger.debug(`Set model ${model}`);
-    this.agentEx.llm.setModel(model);
-  }
-
-  public getSystemPrompt(): string {
-    return this.contextManager.getAgentPrompt();
-  }
-
-  /**
-   * Set the system prompt
-   */
-  public setSystemPrompt(systemMsg: string) {
-    this.contextManager.setAgentPrompt(systemMsg);
-  }
-
-  public addAgentToolProvider(toolProvider: IAgentToolProvider): Promise<void> {
-    return this.agentEx.addAgentToolProvider(toolProvider);
-  }
-
-  public addAgentTool(tool: ToolDescriptor, handler: ToolHandler) {
-    this.agentEx.addAgentTool(tool, handler);
-  }
-}
-
-/**
- * Returns the ChatCompletionMessageParam constructed from (optional) text and
- * (optional) image.  If neither is given (null message), then undefined is
- * returned.
- **/
-export function createUserMessage(
-  msg?: string,
-  imageB64?: string,
-  name?: string
-): UserMessageParam | undefined {
-  const content = (() => {
-    if (!imageB64) {
-      if (!msg) {
-        return undefined;
-      }
-      return msg;
-    }
-
-    const content: openai.ChatCompletionContentPart[] = [];
-    if (msg) {
-      content.push({
-        type: "text",
-        text: msg,
-      });
-    }
-    if (imageB64) {
-      content.push({
-        type: "image_url",
-        image_url: {
-          url: imageB64,
-        },
-      });
-    }
-    return content;
-  })();
-
-  if (!content) {
-    return undefined;
-  }
-
-  return {
-    role: "user",
-    content,
-    name,
-  };
-}
-
-export function createUserMessageEnsure(
-  msg?: string,
-  imageB64?: string,
-  name?: string
-): UserMessageParam {
-  const userMsg = createUserMessage(msg, imageB64, name);
-  assert(userMsg, "createUserMessageEnsure");
-  return userMsg;
-}
-
-export function completionToAssistantMessageParam(
-  compMessage: Message
-): AssistantMessageParam {
-  // Strip down the `ChatCompletionMessage` to a
-  // `ChatCompletionAssistantMessageParam`, only including the non-null
-  // elements.  For reference:
-  //
-  // Response from the LLM:
-  //
-  //   export interface ChatCompletionMessage {
-  //     role: 'assistant';
-  //     audio?: ChatCompletionAudio | null;
-  //     content: string | null;
-  //     refusal: string | null;
-  //     tool_calls?: Array<ChatCompletionMessageToolCall>;
-  //
-  //     annotations?: Array<ChatCompletionMessage.Annotation>;
-  //     // openrouter
-  //     images?: Array<ChatCompletionContentPartImage>
-  //   }
-  //
-  // Input to the LLM
-  //
-  //   export interface ChatCompletionAssistantMessageParam {
-  //     role: "assistant";
-  //     audio?: ChatCompletionAssistantMessageParam.Audio | null;
-  //     content?:
-  //       | string
-  //       | Array<ChatCompletionContentPartText |
-  //               ChatCompletionContentPartRefusal>
-  //       | null;
-  //     refusal?: string | null;
-  //     tool_calls?: Array<ChatCompletionMessageToolCall>;
-  //
-  //     name?: string;
-  //   }
-
-  const message: AssistantMessageParam = {
-    role: "assistant",
-  };
-  if (compMessage.audio) {
-    message.audio = compMessage.audio;
-  }
-  if (compMessage.content) {
-    message.content = compMessage.content;
-  }
-  if (compMessage.refusal) {
-    message.refusal = compMessage.refusal;
-  }
-  if (compMessage.tool_calls) {
-    message.tool_calls = compMessage.tool_calls;
-  }
-  return message;
-}