@hautechai/sdk 0.1.4 → 0.2.0

package/dist/sdk/index.js

@@ -10,6 +10,9 @@ import pipelines from './pipelines';
 import stacks from './stacks';
 import storage from './storage';
 import utils from './utils';
+import { OperationsListener } from './listeners';
+import { getAPI, getBaseUrl } from './api';
+import { OperationsApi } from '../autogenerated';
 export const createSDK = (options) => {
     let token = undefined;
     const authToken = async () => {
@@ -23,6 +26,17 @@ export const createSDK = (options) => {
         return token;
     };
     const optionsWithTokenRefresher = { ...options, authToken };
+    const operationsListener = new OperationsListener({
+        ws: options.useWebsocket
+            ? {
+                endpoint: getBaseUrl(options),
+                token: authToken,
+            }
+            : null,
+        // TODO: Refactor the API initialization
+        operations: () => getAPI(OperationsApi, optionsWithTokenRefresher),
+    });
+    operationsListener.subscribe();
     return {
         access: access(optionsWithTokenRefresher),
         accounts: accounts(optionsWithTokenRefresher),
@@ -30,10 +44,13 @@ export const createSDK = (options) => {
         collections: collections(optionsWithTokenRefresher),
         groups: groups(optionsWithTokenRefresher),
         images: images(optionsWithTokenRefresher),
-        operations: operations(optionsWithTokenRefresher),
+        operations: operations(optionsWithTokenRefresher, operationsListener),
         pipelines: pipelines(optionsWithTokenRefresher),
         stacks: stacks(optionsWithTokenRefresher),
         storage: storage(optionsWithTokenRefresher),
         utils: utils(optionsWithTokenRefresher),
+        close: () => {
+            operationsListener.close();
+        },
     };
 };

package/dist/sdk/listeners/index.d.ts

@@ -0,0 +1,25 @@
+import { OperationEntity, OperationsApi } from '../../autogenerated';
+import { WebSocket } from 'ws';
+export declare class OperationsListener {
+    useWebsocket: {
+        endpoint: string;
+        token: () => string | Promise<string>;
+    } | null;
+    ws: WebSocket | null;
+    operations: () => Promise<OperationsApi>;
+    constructor({ ws, operations, }: {
+        ws: {
+            endpoint: string;
+            token: () => string | Promise<string>;
+        } | null;
+        operations: () => Promise<OperationsApi>;
+    });
+    operationsStore: Record<string, OperationEntity>;
+    getOperation(id: string): Promise<OperationEntity | null>;
+    private updateOperation;
+    subscribe(): Promise<void>;
+    onOpen(): void;
+    onMessage(msg: string): void;
+    onClose(number: number, reason: Buffer): void;
+    close(): void;
+}

package/dist/sdk/listeners/index.js

@@ -0,0 +1,81 @@
+import { WebSocket } from 'ws';
+import { HautechException } from '../exceptions';
+// This is pretty much a dirty solution until we need to rework this part.
+export class OperationsListener {
+    constructor({ ws, operations, }) {
+        this.useWebsocket = null;
+        this.ws = null;
+        this.operationsStore = {};
+        if (ws) {
+            this.useWebsocket = {
+                endpoint: ws?.endpoint,
+                token: ws?.token,
+            };
+        }
+        this.operations = operations;
+    }
+    async getOperation(id) {
+        const isWsReady = this.ws?.readyState == WebSocket.OPEN;
+        if (!this.operationsStore[id] || !isWsReady) {
+            const api = await this.operations();
+            const operation = await api.operationsControllerGetOperationV1(id);
+            if (operation.status == 200)
+                this.updateOperation(id, operation.data);
+        }
+        return this.operationsStore[id] || null;
+    }
+    updateOperation(id, operation) {
+        const stored = this.operationsStore[id];
+        if (!stored || stored.updatedAt < operation.updatedAt) {
+            this.operationsStore[id] = operation;
+        }
+    }
+    async subscribe() {
+        if (!this.useWebsocket)
+            return;
+        try {
+            this.ws = new WebSocket(this.useWebsocket.endpoint, ['1', await this.useWebsocket.token()]);
+            this.ws.on('open', () => this.onOpen());
+            this.ws.on('message', (msg) => this.onMessage(msg));
+            this.ws.on('close', (number, reason) => this.onClose(number, reason));
+        }
+        catch (err) {
+            throw new HautechException(`SDK failed to open websocket: ${err}`);
+        }
+    }
+    onOpen() {
+        if (!this.ws)
+            throw new HautechException('Semantics error: this is a bug.');
+        this.ws.send(JSON.stringify({
+            event: 'subscribe',
+            data: {
+                channel: 'own_resources',
+            },
+        }));
+    }
+    onMessage(msg) {
+        if (!this.ws)
+            throw new HautechException('Semantics error: this is a bug.');
+        const { event, data } = JSON.parse(msg);
+        switch (event) {
+            case 'subscription:created':
+                break;
+            case 'operation:created':
+            case 'operation:updated':
+                this.updateOperation(data.id, data);
+                break;
+        }
+    }
+    onClose(number, reason) {
+        if (!this.ws)
+            throw new HautechException('Semantics error: this is a bug.');
+        // Reset dirty state.
+        this.operationsStore = {};
+        this.ws = null;
+    }
+    close() {
+        if (!this.ws)
+            return;
+        this.ws.close();
+    }
+}

package/dist/sdk/operations/index.d.ts

@@ -1,5 +1,6 @@
 import { CompositeV1Input, CompositeV1Response, CutV1Input, CutV1Response, GiseleVtonV1Input, GPTV1Input, GptV1Response, HauteLindaV1Response, HauteNaomiV1Response, ImagineKateV1Response, KateImagineV1Input, KateInpaintV1Input, LindaHauteV1Input, NaomiHauteV1Input, NegateImageV1Input, NegateImageV1Response, ObjectDetectionV1Input, ObjectDetectionV1Response, OperationEntity, OperationEntityStatusEnum, PoseEstimationV1Input, PoseEstimationV1Response, SegmentAnythingEmbeddingsV1Input, SegmentAnythingEmbeddingsV1Response, SegmentAnythingMaskV1Input, SegmentAnythingMaskV1Response, UpscaleV1Input, UpscaleV1Response, VtonGiseleV1Response } from '../../autogenerated';
 import { ListProps, ListResponse, SDKOptions } from '../../types';
+import { OperationsListener } from '../listeners';
 type Waited<T extends OperationEntity> = T & ({
     status: typeof OperationEntityStatusEnum.Failed;
     output: null;
@@ -10,7 +11,7 @@ type Waited<T extends OperationEntity> = T & ({
     status: typeof OperationEntityStatusEnum.Finished;
     output: NonNullable<T['output']>;
 });
-declare const operations: (options: SDKOptions) => {
+declare const operations: (options: SDKOptions, relevantOperations: OperationsListener) => {
     run: {
         haute: {
             linda: {
@@ -120,6 +121,6 @@ declare const operations: (options: SDKOptions) => {
     }) => Promise<void>;
     wait: <T extends OperationEntity | {
         id: string;
-    }>(props: T) => Promise<Waited<T extends OperationEntity ? T : OperationEntity>>;
+    }>(props: T, timeoutMs?: number) => Promise<Waited<T extends OperationEntity ? T : OperationEntity> | null>;
 };
 export default operations;

package/dist/sdk/operations/index.js

@@ -1,7 +1,8 @@
 import { OperationsApi, } from '../../autogenerated';
 import { useAutogeneratedAPI } from '../api';
 import { transformToListResponse } from '../transformers';
-const operations = (options) => {
+const sleep = (ms) => new Promise((resolve) => setTimeout(resolve, ms));
+const operations = (options, relevantOperations) => {
     const api = useAutogeneratedAPI({ API: OperationsApi, options });
     const createOperation = (callMethod) => (props) => api.call({ run: (methods) => callMethod(methods, props) });
     return {
@@ -72,20 +73,23 @@ const operations = (options) => {
         updateMetadata: async (props) => api.call({
             run: (methods) => methods.operationsControllerUpdateMetadataV1(props.id, { overwrite: props.metadata }),
         }),
-        wait: async (props) => new Promise((resolve, reject) => {
-            const initialDelay = 5000;
-            const delay = 2000;
-            let timeoutId;
-            const poll = async () => {
-                const operation = await api.call({
-                    run: (methods) => methods.operationsControllerGetOperationV1(props.id),
-                });
-                if (operation.status !== 'pending')
-                    return resolve(operation);
-                timeoutId = setTimeout(poll, delay);
+        wait: async (props, timeoutMs) => {
+            const deadline = timeoutMs ? Date.now() + timeoutMs : undefined;
+            const delay = 1000;
+            const poll = async (id) => {
+                const operation = await relevantOperations.getOperation(id);
+                if (operation?.status !== 'pending')
+                    return operation;
+                return null;
             };
-            timeoutId = setTimeout(poll, initialDelay);
-        }),
+            while (!deadline || Date.now() < deadline) {
+                const polled = await poll(props.id);
+                if (polled)
+                    return polled;
+                await sleep(delay);
+            }
+            return null;
+        },
     };
 };
 export default operations;