@lmnr-ai/lmnr 0.2.0 → 0.2.2

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@lmnr-ai/lmnr",
-  "version": "0.2.0",
+  "version": "0.2.2",
   "description": "TypeScript SDK for Laminar AI",
   "main": "src/index.ts",
   "repository": {
@@ -20,14 +20,17 @@
   },
   "homepage": "https://github.com/lmnr-ai/ts-sdk#readme",
   "dependencies": {
+    "@types/ws": "^8.5.10",
     "browser-or-node": "^3.0.0",
     "ts-node": "^10.9.2",
     "tsup": "^8.1.0",
     "typescript": "^5.5.2",
-    "websocket": "^1.0.35"
+    "uuid": "^10.0.0",
+    "ws": "^8.18.0"
   },
   "devDependencies": {
-    "@types/websocket": "^1.0.10"
+    "@types/uuid": "^10.0.0",
+    "bufferutil": "^4.0.8"
   },
   "scripts": {
     "build": "tsup",

package/src/index.ts

@@ -1,9 +1,9 @@
-import { NodeInput, EndpointRunResponse, EndpointRunRequest, NodeWebSocketMessage, ToolCall } from './types';
+import { NodeInput, EndpointRunResponse, EndpointRunRequest, NodeWebSocketMessage, ToolCall, WebSocketError } from './types';
 import { isBrowser, isNode } from 'browser-or-node';
+import { WebSocket } from 'ws';
 
 export { NodeInput, EndpointRunResponse, EndpointRunRequest, ChatMessage } from './types';
-
-class WebSocketError extends Error {}
+export { RemoteDebugger as LaminarRemoteDebugger } from './remote_debugger';
 
 export class Laminar {
     private readonly projectApiKey: string;
@@ -69,54 +69,42 @@ export class Laminar {
         metadata = {},
         tools = [],
     }: EndpointRunRequest): Promise<EndpointRunResponse> {
-        const WebSocketClient = require('websocket').client;
-        const socket = new WebSocketClient();
+        const socket = new  WebSocket(this.ws_url, [], {headers: {'Authorization': `Bearer ${this.projectApiKey}`}});
         let response: EndpointRunResponse | null = null;
-
-        socket.on('connectFailed', (error: any) => {
-            console.log('Connect Error: ' + error.toString());
-        });
         
-        socket.on('connect', (connection: any) => {
-            connection.send(JSON.stringify({
+        socket.on('open', () => {
+            socket.send(JSON.stringify({
                 inputs,
                 endpoint,
                 env,
                 metadata,
                 tools,
             }));
-            connection.on('message', (data: NodeWebSocketMessage) => {
+        });
+        socket.on('message', (data: NodeWebSocketMessage) => {
+            try {
+                const toolCall = JSON.parse(data.toString()) as ToolCall;
+                const matchingTool = tools.find(tool => tool.name === toolCall.function.name);
+                if (!matchingTool) {
+                    throw new WebSocketError(
+                        `Tool ${toolCall.function.name} not found. Registered tools: ${tools.map(tool => tool.name).join(', ')}`
+                    );
+                }
+                let args = {}
                 try {
-                    const toolCall = JSON.parse(data["utf8Data"]!) as ToolCall;
-                    const matchingTool = tools.find(tool => tool.name === toolCall.function.name);
-                    if (!matchingTool) {
-                        throw new WebSocketError(
-                            `Tool ${toolCall.function.name} not found. Registered tools: ${tools.map(tool => tool.name).join(', ')}`
-                        );
-                    }
-                    let args = {}
-                    try {
-                        args = JSON.parse(toolCall.function.arguments);
-                    } catch (e) {}
-                    const result = matchingTool(args);
-                    connection.send(JSON.stringify(result));
-                } catch (e) {
-                    if (e instanceof WebSocketError) {
-                        throw e;
-                    }
-                    response = JSON.parse(data["utf8Data"]!) as EndpointRunResponse;
-                    connection.close();
+                    args = JSON.parse(toolCall.function.arguments);
+                } catch (e) {}
+                const result = matchingTool(args);
+                socket.send(JSON.stringify(result));
+            } catch (e) {
+                if (e instanceof WebSocketError) {
+                    throw e;
                 }
-            });
+                response = JSON.parse(data.toString()) as EndpointRunResponse;
+                socket.close();
+            }
         });
 
-        socket.connect(
-            this.ws_url,
-            '',  // protocols. If specified, server tries to match one of these, and fails
-            null,  // origin. Makes more sense in a browser context
-            {'Authorization': `Bearer ${this.projectApiKey}`}
-        );
-
         // TODO: Implement a better way to wait for the response
         while (!response) {
             await new Promise((resolve) => setTimeout(resolve, 1000));

package/src/remote_debugger.ts

@@ -0,0 +1,81 @@
+import { NodeWebSocketMessage, RegisterDebuggerRequest, Tool, ToolCall, WebSocketError } from "./types";
+import { v4 as uuidv4 } from 'uuid';
+import { WebSocket } from 'ws';
+
+
+export class RemoteDebugger {
+    private readonly projectApiKey: string;
+    private readonly url: string = 'wss://api.lmnr.ai/v2/endpoint/ws';
+    private readonly tools: Tool[];
+    private sessionId: string | null = null;
+    private socket;
+
+    constructor(projectApiKey: string, tools: Tool[]) {
+        this.projectApiKey = projectApiKey;
+        this.tools = tools;
+        this.socket = new  WebSocket(this.url, [], {headers: {'Authorization': `Bearer ${this.projectApiKey}`}});
+    }
+
+    public async start(): Promise<string> {
+        this.sessionId = this.generateSessionId();
+        console.log(this.formatSessionId());
+
+        this.socket.on('open', () => {
+            this.socket.send(JSON.stringify({
+                debuggerSessionId: this.sessionId
+            } as RegisterDebuggerRequest));
+        });
+
+        this.socket.on('message', (data: NodeWebSocketMessage) => {
+            try {
+                const toolCall = JSON.parse(data.toString()) as ToolCall;
+                const matchingTool = this.tools.find(tool => tool.name === toolCall.function.name);
+                if (!matchingTool) {
+                    console.error(
+                        `Tool ${toolCall.function.name} not found. Registered tools: ${this.tools.map(tool => tool.name).join(', ')}`
+                    );
+                    return;
+                }
+                let args = {}
+                try {
+                    args = JSON.parse(toolCall.function.arguments);
+                } catch (e) {}
+                const result = matchingTool(args);
+                this.socket.send(JSON.stringify(result));
+            } catch (e) {
+                this.socket.send(JSON.stringify({
+                    deregister: true,
+                    debuggerSessionId: this.sessionId
+                } as RegisterDebuggerRequest));
+                this.socket.close();
+            }
+        });
+
+        return this.sessionId;
+    }
+
+    public stop() {
+        this.socket.send(JSON.stringify({
+            deregister: true,
+            debuggerSessionId: this.sessionId
+        } as RegisterDebuggerRequest));
+        this.socket.close();
+    }
+
+    public getSessionId(): string | null {
+        return this.sessionId;
+    }
+
+    private generateSessionId(): string {
+        return uuidv4();
+    }
+
+    private formatSessionId(): string {
+        return `
+========================================
+Debugger Session ID:
+${this.sessionId}
+========================================
+`
+    }
+}

package/src/types.ts

@@ -5,12 +5,14 @@ export type ChatMessage = {
 
 export type NodeInput = ChatMessage[] | string;
 
+export type Tool = (...args: any) => NodeInput | Promise<NodeInput>;
+
 export type EndpointRunRequest = {
     inputs: Record<string, NodeInput>;
     endpoint: string;
     env?: Record<string, string>;
     metadata?: Record<string, string>;
-    tools?: ((params: Record<string, any>) => NodeInput | Promise<NodeInput>)[];
+    tools?: Tool[];
 }
 
 export type EndpointRunResponse = {
@@ -34,4 +36,13 @@ export type NodeWebSocketMessage = {
     utf8Data?: string;
 }
 
-const bar = ({a }: {a: string}): string => a;
+export type RegisterDebuggerRequest = {
+    debuggerSessionId: string;
+}
+
+export type DeregisterDebuggerRequest = {
+    deregister: true;
+    debuggerSessionId: string;
+}
+
+export class WebSocketError extends Error {}