diff --git a/querybook/webapp/lib/ai-assistant/ai-assistant-socketio.ts b/querybook/webapp/lib/ai-assistant/ai-assistant-socketio.ts
index fc5449886..0b402e69f 100644
--- a/querybook/webapp/lib/ai-assistant/ai-assistant-socketio.ts
+++ b/querybook/webapp/lib/ai-assistant/ai-assistant-socketio.ts
@@ -13,7 +13,7 @@ class AIAssistantSocket {
     private socket: Socket = null;
     private socketPromise: Promise<Socket> = null;
 
-    constructor() {
+    public constructor() {
         this.setupSocket();
     }
 
@@ -21,6 +21,18 @@ class AIAssistantSocket {
         socket.emit('subscribe');
     }
 
+    public requestAIAssistant = (command: AICommandType, payload: object) => {
+        this.socket.emit(AI_ASSISTANT_REQUEST_EVENT, command, payload);
+    };
+
+    public addAIListener = (listener) => {
+        this.socket.on(AI_ASSISTANT_RESPONSE_EVENT, listener);
+    };
+
+    public removeAIListener = (listener) => {
+        this.socket.off(AI_ASSISTANT_RESPONSE_EVENT, listener);
+    };
+
     private setupSocket = async () => {
         if (this.socket) {
             return this.socket;
@@ -44,18 +56,6 @@ class AIAssistantSocket {
 
         return this.socket;
     };
-
-    public requestAIAssistant = (command: AICommandType, payload: object) => {
-        this.socket.emit(AI_ASSISTANT_REQUEST_EVENT, command, payload);
-    };
-
-    public addAIListener = (listener) => {
-        this.socket.on(AI_ASSISTANT_RESPONSE_EVENT, listener);
-    };
-
-    public removeAIListener = (listener) => {
-        this.socket.off(AI_ASSISTANT_RESPONSE_EVENT, listener);
-    };
 }
 
 const defaultSocket = new AIAssistantSocket();