@xalia/agent 0.6.1 → 0.6.2

package/src/chat/server/connectionManager.ts

@@ -10,7 +10,6 @@ import {
   ServerConnectionMessage,
 } from "../protocol/connectionMessages";
 import { ChatErrorMessage, ChatFatalError } from "../protocol/errors";
-import { ApiKeyManager } from "../data/apiKeyManager";
 
 const logger = getLogger();
 
@@ -32,18 +31,15 @@ export interface IUserConnectionManager<ServerMsgT> {
    * Send message to a specific connection.
    */
   sendToConnection(connectionId: string, message: ServerMsgT): void;
-
-  /**
-   * Get cached api key for a user. This key will not be available
-   * after the user disconnects.
-   */
-  getLiveUserApiKey(userId: string): string | undefined;
 }
 
 /**
- * Interface for processing messages from the client.
+ * Interface for processing messages from the client.  Also handles guest
+ * admission.
  */
 export interface IMessageProcessor<ClientMsgT> {
+  authenticate(token: string): Promise<string | undefined>;
+
   processMessage(
     connectionId: string,
     userId: string,
@@ -54,25 +50,18 @@ export interface IMessageProcessor<ClientMsgT> {
    * Handle user disconnect - clean up from all sessions.
    * Called when a connection is closed to ensure proper cleanup.
    */
-  handleUserDisconnect(userId: string): void;
+  handleUserDisconnect(userId: string): Promise<void>;
 }
 
 class Connection {
   public ws: ws.WebSocket;
   public userId: string;
-  public apiKey: string;
   public heartbeat: NodeJS.Timeout;
   public lastMessageTime: number;
 
-  constructor(
-    ws: ws.WebSocket,
-    userId: string,
-    apiKey: string,
-    heartbeat: NodeJS.Timeout
-  ) {
+  constructor(ws: ws.WebSocket, userId: string, heartbeat: NodeJS.Timeout) {
     this.ws = ws;
     this.userId = userId;
-    this.apiKey = apiKey;
     this.heartbeat = heartbeat;
     this.lastMessageTime = Date.now();
   }
@@ -96,18 +85,15 @@ export class ConnectionManager<ClientMsgT, ServerMsgT>
 
   // this is a singleton object and can only be initialized by .init method
   private constructor(
-    private apiKeyManager: ApiKeyManager,
     private sessionMessageProcessor: IMessageProcessor<ClientMsgT>
   ) {}
 
   static init<ClientMsgT, ServerMsgT>(
-    apiKeyManager: ApiKeyManager,
     createMsgProcessor: (
       cm: IUserConnectionManager<ServerMsgT>
     ) => IMessageProcessor<ClientMsgT>
   ): ConnectionManager<ClientMsgT, ServerMsgT> {
     const connManager = new ConnectionManager(
-      apiKeyManager,
       {} as IMessageProcessor<ClientMsgT>
     );
     connManager.sessionMessageProcessor = createMsgProcessor(connManager);
@@ -120,17 +106,15 @@ export class ConnectionManager<ClientMsgT, ServerMsgT>
    */
   async handleConnection(
     webSocket: ws.WebSocket,
-    apiKey: string,
+    token: string,
     req?: { headers: { [key: string]: string | string[] | undefined } }
   ): Promise<void> {
     const connectionId = uuidv4();
     logger.info(`[ConnectionManager] New connection: ${connectionId}`);
 
     try {
-      // Verify API key
-      const userData = await this.apiKeyManager.verifyApiKey(apiKey);
-      if (!userData) {
-        logger.error(`[ConnectionManager] Invalid API key: ${apiKey}`);
+      const userUUID = await this.sessionMessageProcessor.authenticate(token);
+      if (!userUUID) {
         throw new ChatFatalError("invalid api key");
       }
 
@@ -142,14 +126,14 @@ export class ConnectionManager<ClientMsgT, ServerMsgT>
       // Register connection
       this.connections.set(
         connectionId,
-        new Connection(webSocket, userData.uuid, apiKey, heartbeat)
+        new Connection(webSocket, userUUID, heartbeat)
       );
 
       // Add to user connections
-      if (!this.userToConnections.has(userData.uuid)) {
-        this.userToConnections.set(userData.uuid, new Set());
+      if (!this.userToConnections.has(userUUID)) {
+        this.userToConnections.set(userUUID, new Set());
       }
-      const userConnections = this.userToConnections.get(userData.uuid);
+      const userConnections = this.userToConnections.get(userUUID);
       if (userConnections) {
         userConnections.add(connectionId);
       }
@@ -164,7 +148,7 @@ export class ConnectionManager<ClientMsgT, ServerMsgT>
       const response: ServerConnectionReady = {
         t: "ready",
         c_id: connectionId,
-        user_uuid: userData.uuid,
+        user_uuid: userUUID,
       };
 
       // Check if connection still exists before sending
@@ -179,7 +163,7 @@ export class ConnectionManager<ClientMsgT, ServerMsgT>
 
       logger.info(
         `[ConnectionManager] Connection ${connectionId} registered ` +
-          `for user ${userData.uuid}, version: ${clientVersion}`
+          `for user ${userUUID}, version: ${clientVersion}`
       );
     } catch (error) {
       const errorMessage =
@@ -383,7 +367,9 @@ export class ConnectionManager<ClientMsgT, ServerMsgT>
         userConnections.delete(connectionId);
         if (userConnections.size === 0) {
           this.userToConnections.delete(conn.userId);
-          this.sessionMessageProcessor.handleUserDisconnect(conn.userId);
+          const _ = this.sessionMessageProcessor.handleUserDisconnect(
+            conn.userId
+          );
         }
       }
 
@@ -517,25 +503,4 @@ export class ConnectionManager<ClientMsgT, ServerMsgT>
     const connections = this.userToConnections.get(userId);
     return connections ? Array.from(connections) : [];
   }
-
-  /**
-   * Implementation of IUserConnectionManager interface.
-   * Get cached api key for a user. This key will not be available
-   * after the user disconnects.
-   */
-  public getLiveUserApiKey(userId: string): string | undefined {
-    const connectionIds = this.getUserConnections(userId);
-
-    // Return API key from the first active connection
-    // (all connections for a user should have the same API key)
-    for (const connectionId of connectionIds) {
-      const conn = this.connections.get(connectionId);
-      if (conn) {
-        return conn.apiKey;
-      }
-    }
-
-    // No active connections for this user
-    return undefined;
-  }
 }

package/src/chat/server/conversation.ts

@@ -3,11 +3,13 @@
 
 import { strict as assert } from "assert";
 
+import { createUserMessage } from "../../agent/agent";
 import {
+  ChatCompletionAssistantMessageParam,
   ChatCompletionMessageParam,
+  ChatCompletionToolMessageParam,
   ChatCompletionUserMessageParam,
-  createUserMessage,
-} from "../../agent/agent";
+} from "../../agent/llm";
 import { SessionMessage, UserMessageData } from "../data/dataModels";
 import {
   ServerAgentMessage,
@@ -96,45 +98,109 @@ export function sessionMessagesToLLMConversation(
   return { firstIndex, conversation };
 }
 
-export function sessionMessagesToChatMessages(
+/**
+ * Convert the DB data (SessionMessage) for a ChatCompletionUserMessageParam
+ * into a ServerUserMessage message.
+ */
+export function userMessageToConversationMessage(
+  userMessage: ChatCompletionUserMessageParam,
+  message_idx: number,
+  defaultUserUuid: string,
+  session_id: string
+): ServerUserMessage {
+  // The name on the message should be the uuid
+
+  const userMsgData = llmUserMessageToUserMessageData(userMessage);
+  const user_uuid = userMessage.name || defaultUserUuid;
+  const msg: ServerUserMessage = {
+    type: "user_msg",
+    message: userMsgData.message,
+    message_idx,
+    user_uuid,
+    session_id,
+  };
+  if (userMsgData.imageB64) {
+    msg.imageB64 = userMsgData.imageB64;
+  }
+  return msg;
+}
+
+/**
+ * Convert the DB data (SessionMessage) for a
+ * ChatCompletionAssistantMessageParam into a ServerAgentMessage message.
+ */
+export function assistantMessageToConversationMessage(
+  msg: ChatCompletionAssistantMessageParam,
+  message_idx: number,
+  session_id: string
+): ServerAgentMessage {
+  assert(!msg.audio);
+  return {
+    type: "agent_msg",
+    message: msg,
+    message_idx,
+    session_id,
+  };
+}
+
+/**
+ * Convert the DB data (SessionMessage) for a ChatCompletionToolMessageParam
+ * into a ServerToolCallResult message.
+ */
+export function toolResultToConversationMessage(
+  toolCall: ChatCompletionToolMessageParam,
+  message_idx: number,
+  session_id: string
+): ServerToolCallResult {
+  return {
+    type: "tool_call_result",
+    result: toolCall,
+    message_idx,
+    session_id,
+  };
+}
+
+/**
+ * Convert a set of SessionMessages (stored in the DB) into Conversation
+ * messages (protocol which can be sent over the wire).
+ */
+export function sessionMessagesToConversationMessages(
   sessionMessages: SessionMessage[],
   defaultUserUuid: string,
   session_id: string
-): (ServerUserMessage | ServerAgentMessage)[] {
-  const msgs: (ServerUserMessage | ServerAgentMessage)[] = [];
+): ConversationMessage[] {
+  const msgs: ConversationMessage[] = [];
   for (const sm of sessionMessages) {
     const ccmp = sm.content;
     const message_idx = sm.message_idx;
-    switch (ccmp.role) {
-      case "developer":
-        throw new Error("developer messages not handled yet");
+    const msgRole = ccmp.role;
+    switch (msgRole) {
+      case "system":
+        throw new Error("cannot convert system to ConversationMessage");
       case "assistant":
-        assert(!ccmp.audio);
-        if (ccmp.content) {
-          msgs.push({
-            type: "agent_msg",
-            message: ccmp,
-            message_idx,
-            session_id,
-          });
-        }
-        // TODO: do we want to convert tool calls etc?
+        msgs.push(
+          assistantMessageToConversationMessage(ccmp, message_idx, session_id)
+        );
         break;
       case "user":
-        {
-          const msg = userMessageToChatMessage(
+        msgs.push(
+          userMessageToConversationMessage(
             ccmp,
             message_idx,
             defaultUserUuid,
             session_id
-          );
-          if (msg) {
-            msgs.push(msg);
-          }
-        }
+          )
+        );
         break;
-      default:
+      case "tool":
+        msgs.push(
+          toolResultToConversationMessage(ccmp, message_idx, session_id)
+        );
         break;
+      default: {
+        const _: never = ccmp;
+        throw new Error(`unexpected message role: ${msgRole as string}`);
+      }
     }
   }
 
@@ -143,7 +209,7 @@ export function sessionMessagesToChatMessages(
 
 export function llmUserMessageToUserMessageData(
   userMessage: ChatCompletionUserMessageParam
-): UserMessageData | undefined {
+): UserMessageData {
   if (typeof userMessage.content === "string") {
     return {
       message: userMessage.content,
@@ -163,15 +229,15 @@ export function llmUserMessageToUserMessageData(
         break;
       case "input_audio":
         throw new ChatErrorMessage(
-          "userMessageToChatMessage: audio content not supported"
+          "llmUserMessageToUserMessageData: audio content not supported"
         );
       case "file":
         throw new ChatErrorMessage(
-          "userMessageToChatMessage: file content not supported"
+          "llmUserMessageToUserMessageData: file content not supported"
         );
       default:
         throw new ChatErrorMessage(
-          "userMessageToChatMessage: unexpected content.type"
+          "llmUserMessageToUserMessageData: unexpected content.type"
         );
     }
   }
@@ -185,33 +251,6 @@ export function llmUserMessageToUserMessageData(
   return finalMsg;
 }
 
-export function userMessageToChatMessage(
-  userMessage: ChatCompletionUserMessageParam,
-  message_idx: number,
-  defaultUserUuid: string,
-  session_id: string
-): ServerUserMessage | undefined {
-  // The name on the message should be the uuid
-
-  const userMsgData = llmUserMessageToUserMessageData(userMessage);
-  if (!userMsgData) {
-    return undefined;
-  }
-
-  const user_uuid = userMessage.name || defaultUserUuid;
-  const msg: ServerUserMessage = {
-    type: "user_msg",
-    message: userMsgData.message,
-    message_idx,
-    user_uuid,
-    session_id,
-  };
-  if (userMsgData.imageB64) {
-    msg.imageB64 = userMsgData.imageB64;
-  }
-  return msg;
-}
-
 export function chatToolResultMessageToSessionMessage(
   chatMessage: ServerToolCallResult
 ): SessionMessage {

package/src/chat/server/imageGeneratorTools.ts

@@ -0,0 +1,138 @@
+import { getLogger } from "@xalia/xmcp/sdk";
+
+import { Agent, IAgentToolProvider, ToolCallResult } from "../../agent/agent";
+import { ChatCompletionTool } from "../../agent/llm";
+import { ImageGenerator } from "../../agent/imageGenerator";
+
+import { ChatSessionFileManager } from "./sessionFileManager";
+import { makeParseArgsFn } from "./tools";
+import { getMimeTypeFromDataUrl } from "../data/mimeTypes";
+
+const logger = getLogger();
+
+function createImageGenerator(
+  llmUrl: string,
+  llmApiKey: string
+): Promise<ImageGenerator> {
+  const imageGenModel = process.env["GEN_IMAGE_MODEL"];
+  logger.debug(`[genImageFileTool] model: ${imageGenModel || "(default)"}`);
+  return ImageGenerator.init(llmUrl, llmApiKey, imageGenModel);
+}
+
+// gen_image
+//
+// Simple tool to generate image output as tool output.  The user can then
+// decide how/whether to use the image.
+
+export async function genImageTool(
+  llmUrl: string,
+  llmApiKey: string
+): Promise<IAgentToolProvider> {
+  const GEN_IMAGE_DESC: ChatCompletionTool = {
+    type: "function",
+    function: {
+      name: "gen_image",
+      description: "Generate image",
+      parameters: {
+        type: "object",
+        properties: {
+          prompt: {
+            type: "string",
+            description: "Prompt/instructions for image generator",
+          },
+          input_img: {
+            type: "string",
+            description: "(Optional) input image data",
+          },
+        },
+        required: ["prompt"],
+      },
+    },
+  } as const;
+  const imageGenerator = await createImageGenerator(llmUrl, llmApiKey);
+  const getPromptInputImg = makeParseArgsFn(
+    ["prompt"] as const,
+    ["input_img"] as const
+  );
+  const toolFn = async (_: Agent, args: unknown): Promise<ToolCallResult> => {
+    const { prompt, input_img } = getPromptInputImg(args);
+    const image = await imageGenerator.generate(prompt, input_img);
+    return {
+      response: image,
+      overwriteArgs: JSON.stringify({ prompt }),
+      // overwriteResponse: "Image generated",
+      // // For now, let the UI see the image.
+      // passOriginalResponseToEventHandler: true,
+    };
+  };
+
+  return {
+    // eslint-disable-next-line @typescript-eslint/require-await
+    setup: async (agent: Agent) => {
+      agent.addAgentTool(GEN_IMAGE_DESC, toolFn);
+    },
+  };
+}
+
+// gen_image_file
+
+export async function genImageFileTool(
+  llmUrl: string,
+  llmApiKey: string,
+  fileManager: ChatSessionFileManager
+): Promise<IAgentToolProvider> {
+  const GEN_IMAGE_FILE_DESC: ChatCompletionTool = {
+    type: "function",
+    function: {
+      name: "gen_image_file",
+      description: "Generate image (into session file manager)",
+      parameters: {
+        type: "object",
+        properties: {
+          prompt: {
+            type: "string",
+            description: "Prompt/instructions for image generator",
+          },
+          name: {
+            type: "string",
+            description: "Name of file to (over)write",
+          },
+          summary: {
+            type: "string",
+            description: "Summary (for file manager)",
+          },
+          input_name: {
+            type: "string",
+            description: "(Optional) input image filename",
+          },
+        },
+        required: ["prompt", "name", "summary"],
+      },
+    },
+  } as const;
+
+  const imageGenerator = await createImageGenerator(llmUrl, llmApiKey);
+  const getPromptNameSummary = makeParseArgsFn(
+    ["prompt", "name", "summary"] as const,
+    ["input_name"] as const
+  );
+  const toolFn = async (_: Agent, args: unknown): Promise<ToolCallResult> => {
+    const { prompt, name, summary, input_name } = getPromptNameSummary(args);
+    const input_image = input_name
+      ? await fileManager.getFileContent(input_name)
+      : undefined;
+
+    const image = await imageGenerator.generate(prompt, input_image);
+    const mimeType = getMimeTypeFromDataUrl(image);
+    await fileManager.putFileContent(name, summary, image);
+    const url = fileManager.getSessionFileRelativeUrl(name);
+    return { response: url, metadata: { type: mimeType } };
+  };
+
+  return {
+    // eslint-disable-next-line @typescript-eslint/require-await
+    setup: async (agent: Agent) => {
+      agent.addAgentTool(GEN_IMAGE_FILE_DESC, toolFn);
+    },
+  };
+}