@xalia/agent 0.5.0 → 0.5.2

package/src/{agentUtils.ts → agent/agentUtils.ts}

@@ -1,7 +1,7 @@
 import { getLogger } from "@xalia/xmcp/sdk";
 import { Agent, AgentProfile, OnMessageCB, OnToolCallCB } from "./agent";
 import { IPlatform } from "./iplatform";
-import { SudoMcpServerManager } from "./sudoMcpServerManager";
+import { SkillManager } from "./sudoMcpServerManager";
 import OpenAI from "openai";
 import { Configuration as SudoMcpConfiguration } from "@xalia/xmcp/sdk";
 import { OpenAILLM } from "./openAILLM";
@@ -13,6 +13,7 @@ import { strict as assert } from "assert";
 const logger = getLogger();
 
 export const DEFAULT_LLM_URL = "http://localhost:5001/v1";
+export const DEFAULT_LLM_MODEL = "gpt-4o";
 
 /**
  * Util function to create an Agent from some config information.
@@ -61,28 +62,28 @@ async function createAgent(
 /**
  * Util function to create and initialize an Agent given an AgentProfile.
  */
-export async function createAgentAndSudoMcpServerManager(
-  url: string,
+export async function createAgentWithSkills(
+  llmUrl: string,
   agentProfile: AgentProfile,
   onMessage: OnMessageCB,
   onToolCall: OnToolCallCB,
   platform: IPlatform,
-  openaiApiKey: string | undefined,
+  llmApiKey: string | undefined,
   sudomcpConfig: SudoMcpConfiguration,
   authorizedUrl: string | undefined,
   conversation: OpenAI.ChatCompletionMessageParam[] | undefined,
   stream: boolean = false
-): Promise<[Agent, SudoMcpServerManager]> {
+): Promise<[Agent, SkillManager]> {
   // Create agent
   logger.debug("[createAgentAndSudoMcpServerManager] creating agent ...");
   const agent = await createAgent(
-    url,
+    llmUrl,
     agentProfile.model,
     agentProfile.system_prompt,
     onMessage,
     onToolCall,
     platform,
-    openaiApiKey,
+    llmApiKey,
     stream
   );
   if (conversation) {
@@ -90,10 +91,8 @@ export async function createAgentAndSudoMcpServerManager(
   }
 
   // Init SudoMcpServerManager
-  logger.debug(
-    "[createAgentAndSudoMcpServerManager] creating SudoMcpServerManager."
-  );
-  const sudoMcpServerManager = await SudoMcpServerManager.initialize(
+  logger.debug("[createAgentWithSkills] creating SudoMcpServerManager.");
+  const sudoMcpServerManager = await SkillManager.initialize(
     agent.getMcpServerManager(),
     platform.openUrl,
     sudomcpConfig.backend_url,
@@ -101,12 +100,12 @@ export async function createAgentAndSudoMcpServerManager(
     authorizedUrl
   );
   logger.debug(
-    "[createAgentAndSudoMcpServerManager] restore mcp settings:" +
+    "[createAgentWithSkills] restore mcp settings:" +
       JSON.stringify(agentProfile.mcp_settings)
   );
   await sudoMcpServerManager.restoreMcpSettings(agentProfile.mcp_settings);
 
-  logger.debug("[createAgentAndSudoMcpServerManager] done");
+  logger.debug("[createAgentWithSkills] done");
   return [agent, sudoMcpServerManager];
 }
 
@@ -135,7 +134,7 @@ export async function createNonInteractiveAgent(
     return false;
   };
 
-  const [agent, _] = await createAgentAndSudoMcpServerManager(
+  const [agent, _] = await createAgentWithSkills(
     url,
     agentProfile,
     onMessage,

package/src/{mcpServerManager.ts → agent/mcpServerManager.ts}

@@ -177,7 +177,8 @@ export class McpServerManager {
     try {
       await client.connect(transport);
     } catch (e) {
-      // TODO: is this catch necessary?
+      // Ensure the socket is closed so the process can exit if there is an
+      // error at connection time.
       await client.close();
       throw e;
     }

package/src/{sudoMcpServerManager.ts → agent/sudoMcpServerManager.ts}

@@ -39,7 +39,7 @@ class SanitizedServerBrief extends McpServerBrief {
  * Manages access to the catalogue of servers hosted by sudomcp.  Supports
  * adding these servers to McpServerManager.
  */
-export class SudoMcpServerManager {
+export class SkillManager {
   private constructor(
     private mcpServerManager: McpServerManager,
     private apiClient: ApiClient,
@@ -69,7 +69,7 @@ export class SudoMcpServerManager {
     sudoMcpUrl?: string,
     sudoMcpApiKey?: string,
     authorized_url?: string
-  ): Promise<SudoMcpServerManager> {
+  ): Promise<SkillManager> {
     // TODO: Keep it on here and pass to `McpServerManager.addMcpServer`
     const apiClient = new ApiClient(
       sudoMcpUrl ?? DEFAULT_SERVER_URL,
@@ -79,7 +79,7 @@ export class SudoMcpServerManager {
     const servers = await apiClient.listServers();
 
     const [mcpServers, mcpServersMap] = buildServersList(servers);
-    return new SudoMcpServerManager(
+    return new SkillManager(
       mcpServerManager,
       apiClient,
       mcpServers,

package/src/chat/apiKeyManager.ts

@@ -0,0 +1,24 @@
+// TODO:
+// - lru-cache
+
+import { Database, UserData } from "./db";
+
+export class ApiKeyManager {
+  constructor(private db: Database) {}
+
+  public async verifyApiKey(apiKey: string): Promise<UserData | undefined> {
+    // TODO: Cache this
+
+    const userInfo = await this.db.getUserDataFromApiKey(apiKey);
+    return userInfo;
+
+    // if (apiKey.startsWith("dummy_key")) {
+    //   return {
+    //     user_uuid: apiKey,
+    //     nickname: apiKey,
+    //   };
+    // }
+
+    // return undefined;
+  }
+}

package/src/chat/asyncQueue.ts

@@ -0,0 +1,51 @@
+export class AsyncQueue<T> {
+  private queue: T[];
+  private process: (queueEntry: T) => Promise<void>;
+  private maxBacklog: number;
+  private running: boolean = false;
+
+  constructor(
+    process: (queueEntry: T) => Promise<void>,
+    maxBacklog: number = 100
+  ) {
+    this.queue = [];
+    this.process = process;
+    this.maxBacklog = maxBacklog;
+  }
+
+  public getLength(): number {
+    return this.queue.length;
+  }
+
+  public getMaxBacklog(): number {
+    return this.maxBacklog;
+  }
+
+  public async enqueueAsync(queueEntry: T): Promise<void> {
+    while (this.maxBacklog > 0 && this.queue.length >= this.maxBacklog) {
+      await new Promise((r) => setTimeout(r, 1));
+    }
+    this.queue.push(queueEntry);
+    this.tryNext();
+  }
+
+  private shift(): T | undefined {
+    return this.queue.shift();
+  }
+
+  private async tryNext(): Promise<void> {
+    if (this.running) {
+      return;
+    }
+
+    const queueEntry = this.shift();
+    if (queueEntry) {
+      this.running = true;
+      await this.process(queueEntry);
+      this.running = false;
+
+      // Check for more tasks on the queue.
+      setTimeout(() => this.tryNext());
+    }
+  }
+}

package/src/chat/client.ts

@@ -0,0 +1,142 @@
+import * as dotenv from "dotenv";
+import { getLogger } from "@xalia/xmcp/sdk";
+import { strict as assert } from "assert";
+import * as websocket from "ws";
+import { ServerToClient, ClientToServer } from "./messages";
+
+dotenv.config();
+
+const logger = getLogger();
+
+type OnMessageCallback = (msg: ServerToClient) => void;
+
+type OnConnectionClosedCallback = () => void;
+
+export class ChatClient {
+  private ws: websocket.WebSocket;
+  private onMessageCB: OnMessageCallback;
+  private onConnectionClosedCB: OnConnectionClosedCallback;
+  private closed: boolean;
+
+  private constructor(
+    ws: websocket.WebSocket,
+    onMessageCB: OnMessageCallback,
+    onConnectionClosedCB: OnConnectionClosedCallback
+  ) {
+    this.ws = ws;
+    this.onMessageCB = onMessageCB;
+    this.onConnectionClosedCB = onConnectionClosedCB;
+    this.closed = false;
+  }
+
+  static async initWithParams(
+    host: string,
+    port: number,
+    token: string,
+    params: Record<string, string>,
+    onMessageCB: OnMessageCallback,
+    onConnectionClosedCB: OnConnectionClosedCallback
+  ): Promise<ChatClient> {
+    return new Promise((r, e) => {
+      const urlParams = new URLSearchParams(params);
+      const url = `ws://${host}:${port}?${urlParams}`;
+
+      const ws = new websocket.WebSocket(url, [token]);
+      logger.info("created ws");
+
+      const client = new ChatClient(ws, onMessageCB, onConnectionClosedCB);
+
+      ws.onopen = async () => {
+        logger.info("opened");
+
+        ws.onmessage = (ev: websocket.MessageEvent) => {
+          try {
+            const msgData = ev.data;
+            if (typeof msgData !== "string") {
+              throw `expected "string" data, got ${typeof msgData}`;
+            }
+            logger.debug(`[client.onmessage]: ${msgData}`);
+            const msg: ServerToClient = JSON.parse(msgData);
+            client.onMessageCB(msg);
+          } catch (e) {
+            client.close();
+            throw e;
+          }
+        };
+        r(client);
+      };
+
+      ws.onclose = (err) => {
+        logger.info("closed");
+        logger.info(
+          `[client] WebSocket connection closed: ${JSON.stringify(err)}`
+        );
+        client.closed = true;
+        onConnectionClosedCB();
+      };
+
+      ws.onerror = (error) => {
+        logger.error("[client] WebSocket error:", JSON.stringify(error));
+        e(error);
+
+        client.closed = true;
+        onConnectionClosedCB();
+      };
+    });
+  }
+
+  // public static async initWithSession(
+  //   host: string,
+  //   port: number,
+  //   token: string,
+  //   sessionId: string,
+  //   onMessageCB: OnMessageCallback,
+  //   onConnectionClosedCB: OnConnectionClosedCallback
+  // ): Promise<ChatClient> {
+  //   return ChatClient.initWithParams(
+  //     host,
+  //     port,
+  //     token,
+  //     { session_id: sessionId },
+  //     onMessageCB,
+  //     onConnectionClosedCB
+  //   );
+  // }
+
+  public static async init(
+    host: string,
+    port: number,
+    token: string,
+    onMessageCB: OnMessageCallback,
+    onConnectionClosedCB: OnConnectionClosedCallback,
+    sessionId: string = "untitled",
+    agentProfileId: string | undefined = undefined
+  ): Promise<ChatClient> {
+    const params: Record<string, string> = { session_id: sessionId };
+    if (agentProfileId) {
+      params["agent_profile_id"] = agentProfileId;
+    }
+    return ChatClient.initWithParams(
+      host,
+      port,
+      token,
+      params,
+      onMessageCB,
+      onConnectionClosedCB
+    );
+  }
+
+  public sendMessage(message: ClientToServer): void {
+    assert(!this.closed);
+    const data = JSON.stringify(message);
+    this.ws.send(data);
+  }
+
+  public close(): void {
+    this.closed = true;
+    this.onConnectionClosedCB();
+    this.ws.close();
+  }
+}
+
+// TODO: remove this

package/src/chat/conversationManager.ts

@@ -0,0 +1,283 @@
+import * as ws from "ws";
+import { v4 as uuidv4 } from "uuid";
+import { strict as assert } from "assert";
+
+import { Configuration, getLogger } from "@xalia/xmcp/sdk";
+
+import {
+  Agent,
+  ChatCompletionMessageParam,
+  ChatCompletionMessageToolCall,
+} from "../agent/agent";
+import { SkillManager } from "../agent/sudoMcpServerManager";
+import { createAgentWithSkills } from "../agent/agentUtils";
+
+import type {
+  ClientToServer,
+  ServerToClient,
+  ClientUserMessage,
+} from "./messages";
+import { AsyncQueue } from "./asyncQueue";
+import { Database, UserData } from "./db";
+import { AsyncLock } from "../utils/asyncLock";
+
+const logger = getLogger();
+
+export class UserAlreadyConnected extends Error {
+  constructor() {
+    super("User already connected to the conversation");
+  }
+}
+
+type QueuedClientMessage<T extends ClientToServer = ClientToServer> = {
+  msg: T;
+  from: string;
+};
+
+/**
+ * Describes a Session (conversation) with connected participants.
+ */
+export class OpenSession {
+  public onEmpty: () => void;
+  public connections: Record<string, ws.WebSocket>;
+  public agent: Agent;
+  public skillManager: SkillManager;
+  public messageQueue: AsyncQueue<QueuedClientMessage>;
+  public curAgentMsgId: string | undefined;
+
+  constructor(
+    agent: Agent,
+    sudoMcpServerManager: SkillManager,
+    onEmpty: () => void
+  ) {
+    // public agent: Agent,
+    this.agent = agent;
+    this.skillManager = sudoMcpServerManager;
+    this.onEmpty = onEmpty;
+    this.connections = {};
+    this.messageQueue = new AsyncQueue<QueuedClientMessage>((m) =>
+      this.onMessage(m)
+    );
+    this.curAgentMsgId = undefined;
+  }
+
+  join(userName: string, ws: ws.WebSocket): void {
+    if (this.connections[userName]) {
+      throw new UserAlreadyConnected();
+    }
+
+    // Inform any other participants, and add the WebSocket to the map.
+
+    this.broadcast({ type: "user_joined", user: userName });
+    this.connections[userName] = ws;
+
+    ws.on("message", async (message: ws.RawData) => {
+      logger.debug(`[convMgr]: got message: (from ${userName}): ${message}`);
+      const msgStr = message.toString();
+      const msg: ClientToServer = JSON.parse(msgStr);
+      await this.messageQueue.enqueueAsync({ msg, from: userName });
+    });
+
+    ws.on("close", () => {
+      logger.debug(`[convMgr]: ${userName} closed`);
+
+      // Remove our connection then inform any other participants
+
+      delete this.connections[userName];
+      this.broadcast({ type: "user_left", user: userName });
+
+      if (Object.keys(this.connections).length == 0) {
+        this.onEmpty();
+      }
+    });
+  }
+
+  /**
+   * Called once for each message.  Messages are queued to avoid overlapping
+   * calls to the LLM.
+   */
+  async onMessage(queuedMessage: QueuedClientMessage): Promise<void> {
+    logger.debug(
+      `[onMessage]: processing (${queuedMessage.from}) ` +
+        `${JSON.stringify(queuedMessage.msg)}`
+    );
+
+    const msg = queuedMessage.msg;
+    switch (msg.type) {
+      case "msg":
+        return this.onChatMessage(
+          queuedMessage as QueuedClientMessage<ClientUserMessage>
+        );
+      default:
+        throw `unknown message: ${JSON.stringify(queuedMessage)}`;
+    }
+  }
+
+  async onAgentMessage(msg: string, end: boolean): Promise<void> {
+    logger.debug(`[onAgentMessage] msg: ${msg}, end: ${end}`);
+    assert(this.curAgentMsgId);
+    this.broadcast({
+      type: "agent_msg_chunk",
+      message_id: this.curAgentMsgId,
+      message: msg,
+      end,
+    });
+  }
+
+  async onChatMessage(
+    queuedMessage: QueuedClientMessage<ClientUserMessage>
+  ): Promise<void> {
+    const msg = queuedMessage.msg;
+    const userToken = queuedMessage.from;
+    const msgId = uuidv4();
+    this.broadcast({
+      type: "user_msg",
+      message_id: msgId,
+      message: msg.message,
+      from: userToken,
+    });
+
+    // Messages will be handled by the Agent.onMessage callback.  We await the
+    // response here before processing further messages.
+
+    this.curAgentMsgId = `${msgId}-resp`;
+    await this.agent.userMessage(msg.message, undefined, userToken);
+  }
+
+  broadcast(msg: ServerToClient) {
+    logger.info(`[broadcast]: broadcast msg: ${JSON.stringify(msg)}`);
+    const msgString = JSON.stringify(msg);
+    for (const ws of Object.values(this.connections)) {
+      ws.send(msgString);
+    }
+  }
+}
+
+/**
+ * Handles forwarding of messages between all participants of a session, as
+ * well as messages to/from and Agent.
+ */
+export class ConversationManager {
+  public db: Database;
+  public llmUrl: string;
+  public xmcpUrl: string;
+  public openSessionsLock: AsyncLock;
+  public openSessions: Record<string, OpenSession>;
+
+  constructor(db: Database, llmUrl: string, xmcpUrl: string) {
+    this.db = db;
+    this.llmUrl = llmUrl;
+    this.xmcpUrl = xmcpUrl;
+    this.openSessionsLock = new AsyncLock();
+    this.openSessions = {};
+  }
+
+  async join(
+    sessionId: string,
+    llmApiKey: string,
+    xmcpApiKey: string,
+    userData: UserData,
+    ws: ws.WebSocket
+  ): Promise<void> {
+    await this.openSessionsLock.lockAndProcess(() => {
+      return this.getOrCreateAndSubscribe(
+        sessionId,
+        llmApiKey,
+        xmcpApiKey,
+        userData.nickname || userData.user_uuid,
+        ws
+      );
+    });
+  }
+  /**
+   * Must be called while holding the openSessionsLock
+   */
+  private async getOrCreateAndSubscribe(
+    sessionId: string,
+    llmApiKey: string,
+    xmcpApiKey: string,
+    userName: string,
+    ws: ws.WebSocket
+  ): Promise<OpenSession> {
+    let openSession = this.openSessions[sessionId];
+    if (openSession) {
+      openSession.join(userName, ws);
+      return openSession;
+    }
+
+    // Create a new session
+
+    // TODO: The owner of llmApiKey and xmcpApiKey may not be the owner of the
+    // session.  Should we create the Agent and SudoMcpServerManager with the
+    // session-owners api key?
+
+    const onEmpty = () => {
+      logger.debug(`session ${sessionId} empty.  removing`);
+      delete this.openSessions[sessionId];
+    };
+
+    const sessionData = await this.db.getSessionById(sessionId);
+    if (!sessionData) {
+      throw `no such session ${sessionId}`;
+    }
+    const conversation =
+      sessionData.conversation as unknown as ChatCompletionMessageParam[];
+
+    const agentProfile = await this.db.getAgentProfileById(
+      sessionData.agent_profile_uuid
+    );
+    if (!agentProfile) {
+      throw `no such agent profile ${sessionData.agent_profile_uuid}`;
+    }
+
+    const platform = {
+      openUrl: (
+        _url: string,
+        _authResultP: Promise<boolean>,
+        _displayName: string
+      ) => {
+        throw "unimpl";
+      },
+      load: (_filename: string): Promise<string> => {
+        throw "unimpl";
+      },
+    };
+
+    // Forward agent messages to the OpenSession (once it is iniailized
+
+    const context: { openSession?: OpenSession } = {};
+    const onMessage = async (msg: string, end: boolean) => {
+      logger.debug(`[onMessage] msg: ${msg}, end: ${end}`);
+      assert(context.openSession);
+      context.openSession.onAgentMessage(msg, end);
+    };
+
+    const onToolCall = async (
+      toolCall: ChatCompletionMessageToolCall
+    ): Promise<boolean> => {
+      logger.debug(`[onToolCall] : ${JSON.stringify(toolCall)}`);
+      return true;
+    };
+
+    const xmcpConfig = Configuration.new(xmcpApiKey, this.xmcpUrl, false);
+
+    const [agent, smsm] = await createAgentWithSkills(
+      this.llmUrl,
+      agentProfile,
+      onMessage,
+      onToolCall,
+      platform,
+      llmApiKey,
+      xmcpConfig,
+      undefined,
+      conversation,
+      true
+    );
+
+    openSession = new OpenSession(agent, smsm, onEmpty);
+    context.openSession = openSession;
+    this.openSessions[sessionId] = openSession;
+    openSession.join(userName, ws);
+    return openSession;
+  }
+}