@sogni-ai/sogni-client 0.3.3 → 0.4.0-aplha.10

package/src/Account/index.ts

@@ -17,6 +17,14 @@ import { SupernetType } from '../ApiClient/WebSocketClient/types';
 
 /**
  * Account API methods that let you interact with the user's account.
+ * Can be accessed via `client.account`. Look for more samples below.
+ *
+ * @example Retrieve the current account balance
+ * ```typescript
+ * const balance = await client.account.refreshBalance();
+ * console.log(balance);
+ * ```
+ *
  */
 class AccountApi extends ApiGroup {
   readonly currentAccount = new CurrentAccount();
@@ -43,13 +51,27 @@ class AccountApi extends ApiGroup {
     this.currentAccount._clear();
   }
 
-  async getNonce(walletAddress: string): Promise<string> {
+  private async getNonce(walletAddress: string): Promise<string> {
     const res = await this.client.rest.post<ApiReponse<Nonce>>('/v1/account/nonce', {
       walletAddress
     });
     return res.data.nonce;
   }
 
+  /**
+   * Create Ethers.js Wallet instance from username and password.
+   * This method is used internally to create a wallet for the user.
+   * You can use this method to create a wallet if you need to sign transactions.
+   *
+   * @example Create a wallet from username and password
+   * ```typescript
+   * const wallet = client.account.getWallet('username', 'password');
+   * console.log(wallet.address);
+   * ```
+   *
+   * @param username - Sogni account username
+   * @param password - Sogni account password
+   */
   getWallet(username: string, password: string): Wallet {
     const pwd = toUtf8Bytes(username.toLowerCase() + password);
     const salt = toUtf8Bytes('sogni-salt-value');
@@ -57,6 +79,16 @@ class AccountApi extends ApiGroup {
     return new Wallet(pkey, this.provider);
   }
 
+  /**
+   * Create a new account with the given username, email, and password.
+   * @internal
+   *
+   * @param username
+   * @param email
+   * @param password
+   * @param subscribe
+   * @param referralCode
+   */
   async create(
     username: string,
     email: string,
@@ -83,6 +115,31 @@ class AccountApi extends ApiGroup {
     return res.data;
   }
 
+  /**
+   * Restore session with username and access token.
+   *
+   * You can save access token that you get from the login method and restore the session with this method.
+   *
+   * @example Store access token to local storage
+   * ```typescript
+   * const { username, token } = await client.account.login('username', 'password');
+   * localStorage.setItem('sogni-username', username);
+   * localStorage.setItem('sogni-token', token);
+   * ```
+   *
+   * @example Restore session from local storage
+   * ```typescript
+   * const username = localStorage.getItem('sogni-username');
+   * const token = localStorage.getItem('sogni-token');
+   * if (username && token) {
+   *  client.account.setToken(username, token);
+   *  console.log('Session restored');
+   * }
+   * ```
+   *
+   * @param username
+   * @param token
+   */
   setToken(username: string, token: string): void {
     this.client.authenticate(token);
     this.currentAccount._update({
@@ -93,6 +150,13 @@ class AccountApi extends ApiGroup {
 
   /**
    * Login with username and password. WebSocket connection is established after successful login.
+   *
+   * @example Login with username and password
+   * ```typescript
+   * await client.account.login('username', 'password');
+   * console.log('Logged in');
+   * ```
+   *
    * @param username
    * @param password
    */
@@ -113,6 +177,12 @@ class AccountApi extends ApiGroup {
 
   /**
    * Logout the user and close the WebSocket connection.
+   *
+   * @example Logout the user
+   * ```typescript
+   * await client.account.logout();
+   * console.log('Logged out');
+   * ```
    */
   async logout(): Promise<void> {
     this.client.rest.post('/v1/account/logout').catch((e) => {
@@ -124,6 +194,16 @@ class AccountApi extends ApiGroup {
 
   /**
    * Refresh the balance of the current account.
+   *
+   * Usually, you don't need to call this method manually. Balance is updated automatically
+   * through WebSocket events. But you can call this method to force a balance refresh.
+   *
+   * @example Refresh user account balance
+   * ```typescript
+   * const balance = await client.account.refreshBalance();
+   * console.log(balance);
+   * // { net: '100.000000', settled: '100.000000', credit: '0.000000', debit: '0.000000' }
+   * ```
    */
   async refreshBalance(): Promise<BalanceData> {
     const res = await this.client.rest.get<ApiReponse<BalanceData>>('/v1/account/balance');
@@ -132,7 +212,18 @@ class AccountApi extends ApiGroup {
   }
 
   /**
-   * Get the balance of a account wallet.
+   * Get the balance of the wallet address.
+   *
+   * This method is used to get the balance of the wallet address. It returns $SOGNI and ETH balance.
+   *
+   * @example Get the balance of the wallet address
+   * ```typescript
+   * const address = client.account.currentAccount.walletAddress;
+   * const balance = await client.account.walletBalance(address);
+   * console.log(balance);
+   * // { token: '100.000000', ether: '0.000000' }
+   * ```
+   *
    * @param walletAddress
    */
   async walletBalance(walletAddress: string) {
@@ -145,6 +236,11 @@ class AccountApi extends ApiGroup {
     return res.data;
   }
 
+  /**
+   * Validate the username before signup
+   * @internal
+   * @param username
+   */
   async validateUsername(username: string) {
     try {
       return await this.client.rest.post<ApiReponse<undefined>>('/v1/account/username/validate', {
@@ -163,21 +259,44 @@ class AccountApi extends ApiGroup {
 
   /**
    * Switch between fast and relaxed networks.
-   * Note: This method will close the current WebSocket connection and establish a new one.
-   * Do not call this method if you have any active projects.
-   * @param network
+   * This will change default network used to process projects. After switching, you will updated
+   * list of AI models available for on selected network.
+   *
+   * @example Switch to the fast network
+   * ```typescript
+   * await client.account.switchNetwork('fast');
+   * console.log('Switched to the fast network, now lets wait until we get list of models');
+   * await client.projects.waitForModels();
+   * ```
+   * @param network - Network type to switch to
    */
-  async switchNetwork(network: SupernetType) {
+  async switchNetwork(network: SupernetType): Promise<SupernetType> {
     this.currentAccount._update({
-      networkStatus: 'connecting',
+      networkStatus: 'switching',
       network: null
     });
-    this.client.socket.switchNetwork(network);
+    const newNetwork = await this.client.socket.switchNetwork(network);
+    this.currentAccount._update({
+      networkStatus: 'connected',
+      network: newNetwork
+    });
+    return newNetwork;
   }
 
   /**
    * Get the transaction history of the current account.
-   * @param params
+   *
+   * @example Get the transaction history
+   * ```typescript
+   * const { entries, next } = await client.account.transactionHistory({
+   *  status: 'completed',
+   *  limit: 10,
+   *  address: client.account.currentAccount.walletAddress
+   * });
+   * ```
+   *
+   * @param params - Transaction history query parameters
+   * @returns Transaction history entries and next query parameters
    */
   async transactionHistory(
     params: TxHistoryParams
@@ -211,6 +330,10 @@ class AccountApi extends ApiGroup {
     };
   }
 
+  /**
+   * Get the rewards of the current account.
+   * @internal
+   */
   async rewards(): Promise<Reward[]> {
     const r =
       await this.client.rest.get<ApiReponse<{ rewards: RewardRaw[] }>>('/v2/account/rewards');
@@ -233,6 +356,11 @@ class AccountApi extends ApiGroup {
     );
   }
 
+  /**
+   * Claim rewards by reward IDs.
+   * @internal
+   * @param rewardIds
+   */
   async claimRewards(rewardIds: string[]): Promise<void> {
     await this.client.rest.post('/v2/account/reward/claim', {
       claims: rewardIds

package/src/ApiClient/WebSocketClient/events.ts

@@ -12,7 +12,7 @@ export type JobErrorData = {
   imgID?: string;
   isFromWorker: boolean;
   error_message: string;
-  error: number;
+  error: number | string;
 };
 
 export type JobProgressData = {
@@ -63,6 +63,10 @@ export type SocketEventMap = {
    * @event WebSocketClient#balanceUpdate - Received balance update
    */
   balanceUpdate: BalanceData;
+  /**
+   * @event WebSocketClient#changeNetwork - Default network changed
+   */
+  changeNetwork: { network: SupernetType };
   /**
    * @event WebSocketClient#jobError - Job error occurred
    */

package/src/ApiClient/WebSocketClient/index.ts

@@ -19,9 +19,13 @@ class WebSocketClient extends RestClient<SocketEventMap> {
   private _pingInterval: NodeJS.Timeout | null = null;
 
   constructor(baseUrl: string, appId: string, supernetType: SupernetType, logger: Logger) {
-    super(baseUrl, logger);
+    const _baseUrl = new URL(baseUrl);
+    if (_baseUrl.protocol === 'wss:') {
+      _baseUrl.protocol = 'https:';
+    }
+    super(_baseUrl.toString(), logger);
     this.appId = appId;
-    this.baseUrl = baseUrl;
+    this.baseUrl = _baseUrl.toString();
     this._supernetType = supernetType;
   }
 
@@ -56,6 +60,7 @@ class WebSocketClient extends RestClient<SocketEventMap> {
     }
     const userAgent = `Sogni/${LIB_VERSION} (sogni-client)`;
     const url = new URL(this.baseUrl);
+    url.protocol = 'wss:';
     url.searchParams.set('appId', this.appId);
     url.searchParams.set('clientName', userAgent);
     url.searchParams.set('clientType', 'artist');
@@ -108,10 +113,14 @@ class WebSocketClient extends RestClient<SocketEventMap> {
     }
   }
 
-  switchNetwork(supernetType: SupernetType) {
-    this._supernetType = supernetType;
-    this.disconnect();
-    this.connect();
+  switchNetwork(supernetType: SupernetType): Promise<SupernetType> {
+    return new Promise<SupernetType>(async (resolve, reject) => {
+      this.once('changeNetwork', ({ network }) => {
+        this._supernetType = network;
+        resolve(network);
+      });
+      await this.send('changeNetwork', supernetType);
+    });
   }
 
   /**

package/src/ApiClient/WebSocketClient/messages.ts

@@ -1,7 +1,9 @@
 import { JobRequestRaw } from '../../Projects/createJobRequestMessage';
+import { SupernetType } from './types';
 
 export interface SocketMessageMap {
   jobRequest: JobRequestRaw;
+  changeNetwork: SupernetType;
 }
 
 export type MessageType = keyof SocketMessageMap;

package/src/Projects/Job.ts

@@ -1,5 +1,9 @@
 import DataEntity, { EntityEvents } from '../lib/DataEntity';
 import ErrorData from '../types/ErrorData';
+import { RawJob, RawProject } from './types/RawProject';
+import ProjectsApi from './index';
+import { Logger } from '../lib/DefaultLogger';
+import getUUID from '../lib/getUUID';
 
 export type JobStatus =
   | 'pending'
@@ -9,14 +13,27 @@ export type JobStatus =
   | 'failed'
   | 'canceled';
 
+const JOB_STATUS_MAP: Record<RawJob['status'], JobStatus> = {
+  created: 'pending',
+  queued: 'pending',
+  assigned: 'initiating',
+  initiatingModel: 'initiating',
+  jobStarted: 'processing',
+  jobProgress: 'processing',
+  jobCompleted: 'completed',
+  jobError: 'failed'
+};
+
 /**
  * @inline
  */
 export interface JobData {
   id: string;
+  projectId: string;
   status: JobStatus;
   step: number;
   stepCount: number;
+  workerName?: string;
   seed?: number;
   isNSFW?: boolean;
   userCanceled?: boolean;
@@ -31,9 +48,37 @@ export interface JobEventMap extends EntityEvents {
   failed: ErrorData;
 }
 
+export interface JobOptions {
+  api: ProjectsApi;
+  logger: Logger;
+}
+
 class Job extends DataEntity<JobData, JobEventMap> {
-  constructor(data: JobData) {
+  static fromRaw(rawProject: RawProject, rawJob: RawJob, options: JobOptions) {
+    return new Job(
+      {
+        id: rawJob.imgID || getUUID(),
+        projectId: rawProject.id,
+        status: JOB_STATUS_MAP[rawJob.status],
+        step: rawJob.performedSteps,
+        stepCount: rawProject.stepCount,
+        workerName: rawJob.worker.name,
+        seed: rawJob.seedUsed,
+        isNSFW: rawJob.triggeredNSFWFilter
+      },
+      options
+    );
+  }
+
+  private readonly _api: ProjectsApi;
+  private readonly _logger: Logger;
+
+  constructor(data: JobData, options: JobOptions) {
     super(data);
+
+    this._api = options.api;
+    this._logger = options.logger;
+
     this.on('updated', this.handleUpdated.bind(this));
   }
 
@@ -41,6 +86,10 @@ class Job extends DataEntity<JobData, JobEventMap> {
     return this.data.id;
   }
 
+  get projectId() {
+    return this.data.projectId;
+  }
+
   /**
    * Current status of the job.
    */
@@ -48,6 +97,10 @@ class Job extends DataEntity<JobData, JobEventMap> {
     return this.data.status;
   }
 
+  get finished() {
+    return ['completed', 'failed', 'canceled'].includes(this.status);
+  }
+
   /**
    * Progress of the job in percentage (0-100).
    */
@@ -108,6 +161,42 @@ class Job extends DataEntity<JobData, JobEventMap> {
     return !!this.data.isNSFW;
   }
 
+  /**
+   * Name of the worker that is processing this job.
+   */
+  get workerName() {
+    return this.data.workerName;
+  }
+
+  /**
+   * Syncs the job data with the data received from the REST API.
+   * @internal
+   * @param data
+   */
+  async _syncWithRestData(data: RawJob) {
+    const delta: Partial<JobData> = {
+      step: data.performedSteps,
+      workerName: data.worker.name,
+      seed: data.seedUsed,
+      isNSFW: data.triggeredNSFWFilter
+    };
+    if (JOB_STATUS_MAP[data.status]) {
+      delta.status = JOB_STATUS_MAP[data.status];
+    }
+    if (!this.data.resultUrl && delta.status === 'completed' && !data.triggeredNSFWFilter) {
+      try {
+        delta.resultUrl = await this._api.downloadUrl({
+          jobId: this.projectId,
+          imageId: this.id,
+          type: 'complete'
+        });
+      } catch (error) {
+        this._logger.error(error);
+      }
+    }
+    this._update(delta);
+  }
+
   private handleUpdated(keys: string[]) {
     if (keys.includes('step') || keys.includes('stepCount')) {
       this.emit('progress', this.progress);

package/src/Projects/Project.ts

@@ -1,11 +1,34 @@
-import Job, { JobData } from './Job';
+import Job, { JobData, JobStatus } from './Job';
 import DataEntity, { EntityEvents } from '../lib/DataEntity';
 import { ProjectParams } from './types';
 import cloneDeep from 'lodash/cloneDeep';
 import ErrorData from '../types/ErrorData';
 import getUUID from '../lib/getUUID';
+import { RawJob, RawProject } from './types/RawProject';
+import ProjectsApi from './index';
+import { Logger } from '../lib/DefaultLogger';
 
-export type ProjectStatus = 'pending' | 'queued' | 'processing' | 'completed' | 'failed';
+// If project is not finished and had no updates for 1 minute, force refresh
+const PROJECT_TIMEOUT = 60 * 1000;
+const MAX_FAILED_SYNC_ATTEMPTS = 3;
+
+export type ProjectStatus =
+  | 'pending'
+  | 'queued'
+  | 'processing'
+  | 'completed'
+  | 'failed'
+  | 'canceled';
+
+const PROJECT_STATUS_MAP: Record<RawProject['status'], ProjectStatus> = {
+  pending: 'pending',
+  active: 'queued',
+  assigned: 'processing',
+  progress: 'processing',
+  completed: 'completed',
+  errored: 'failed',
+  cancelled: 'canceled'
+};
 
 /**
  * @inline
@@ -31,11 +54,20 @@ export interface ProjectEventMap extends EntityEvents {
   jobFailed: Job;
 }
 
+export interface ProjectOptions {
+  api: ProjectsApi;
+  logger: Logger;
+}
+
 class Project extends DataEntity<ProjectData, ProjectEventMap> {
   private _jobs: Job[] = [];
   private _lastEmitedProgress = -1;
+  private readonly _api: ProjectsApi;
+  private readonly _logger: Logger;
+  private _timeout: NodeJS.Timeout | null = null;
+  private _failedSyncAttempts = 0;
 
-  constructor(data: ProjectParams) {
+  constructor(data: ProjectParams, options: ProjectOptions) {
     super({
       id: getUUID(),
       startedAt: new Date(),
@@ -44,6 +76,11 @@ class Project extends DataEntity<ProjectData, ProjectEventMap> {
       status: 'pending'
     });
 
+    this._api = options.api;
+    this._logger = options.logger;
+
+    this._timeout = setInterval(this._checkForTimeout.bind(this), PROJECT_TIMEOUT);
+
     this.on('updated', this.handleUpdated.bind(this));
   }
 
@@ -59,6 +96,10 @@ class Project extends DataEntity<ProjectData, ProjectEventMap> {
     return this.data.status;
   }
 
+  get finished() {
+    return ['completed', 'failed', 'canceled'].includes(this.status);
+  }
+
   get error() {
     return this.data.error;
   }
@@ -137,6 +178,11 @@ class Project extends DataEntity<ProjectData, ProjectEventMap> {
       this.emit('progress', progress);
       this._lastEmitedProgress = progress;
     }
+    // If project is finished stop watching for timeout
+    if (this._timeout && this.finished) {
+      clearInterval(this._timeout!);
+      this._timeout = null;
+    }
     if (keys.includes('status') || keys.includes('jobs')) {
       const allJobsDone = this.jobs.every((job) =>
         ['completed', 'failed', 'canceled'].includes(job.status)
@@ -155,21 +201,104 @@ class Project extends DataEntity<ProjectData, ProjectEventMap> {
    * @internal
    * @param data
    */
-  _addJob(data: JobData) {
-    const job = new Job(data);
+  _addJob(data: JobData | Job) {
+    const job =
+      data instanceof Job ? data : new Job(data, { api: this._api, logger: this._logger });
     this._jobs.push(job);
     job.on('updated', () => {
+      this.lastUpdated = new Date();
       this.emit('updated', ['jobs']);
     });
     job.on('completed', () => {
       this.emit('jobCompleted', job);
+      this._handleJobFinished(job);
     });
     job.on('failed', () => {
       this.emit('jobFailed', job);
+      this._handleJobFinished(job);
     });
     return job;
   }
 
+  private _handleJobFinished(job: Job) {
+    const finalStatus: JobStatus[] = ['completed', 'failed', 'canceled'];
+    const allJobsDone = this.jobs.every((job) => finalStatus.includes(job.status));
+    // If all jobs are done and project is not already failed or completed, update the project status
+    if (allJobsDone && this.status !== 'failed' && this.status !== 'completed') {
+      const allJobsFailed = this.jobs.every((job) => job.status === 'failed');
+      if (allJobsFailed) {
+        this._update({ status: 'failed' });
+      } else {
+        this._update({ status: 'completed' });
+      }
+    }
+  }
+
+  private _checkForTimeout() {
+    if (this.lastUpdated.getTime() + PROJECT_TIMEOUT < Date.now()) {
+      this._syncToServer().catch((error) => {
+        this._logger.error(error);
+        this._failedSyncAttempts++;
+        if (this._failedSyncAttempts > MAX_FAILED_SYNC_ATTEMPTS) {
+          this._logger.error(
+            `Failed to sync project data after ${MAX_FAILED_SYNC_ATTEMPTS} attempts. Stopping further attempts.`
+          );
+          clearInterval(this._timeout!);
+          this._timeout = null;
+        }
+      });
+    }
+  }
+
+  /**
+   * Sync project data with the data received from the REST API.
+   * @internal
+   */
+  async _syncToServer() {
+    const data = await this._api.get(this.id);
+    const jobData = data.completedWorkerJobs.reduce((acc: Record<string, RawJob>, job) => {
+      const jobId = job.imgID || getUUID();
+      acc[jobId] = job;
+      return acc;
+    }, {});
+    for (const job of this._jobs) {
+      const restJob = jobData[job.id];
+      // This should never happen, but just in case we log a warning
+      if (!restJob) {
+        this._logger.warn(`Job with id ${job.id} not found in the REST project data`);
+        return;
+      }
+      try {
+        await job._syncWithRestData(restJob);
+      } catch (error) {
+        this._logger.error(error);
+        this._logger.error(`Failed to sync job ${job.id}`);
+      }
+      delete jobData[job.id];
+    }
+
+    // If there are any jobs left in jobData, it means they are new jobs that are not in the project yet
+    if (Object.keys(jobData).length) {
+      for (const job of Object.values(jobData)) {
+        const jobInstance = Job.fromRaw(data, job, { api: this._api, logger: this._logger });
+        this._addJob(jobInstance);
+      }
+    }
+
+    const delta: Partial<ProjectData> = {
+      params: {
+        ...this.data.params,
+        numberOfImages: data.imageCount,
+        steps: data.stepCount,
+        numberOfPreviews: data.previewCount
+      }
+    };
+    if (PROJECT_STATUS_MAP[data.status]) {
+      delta.status = PROJECT_STATUS_MAP[data.status];
+    }
+    this._update(delta);
+  }
+
   /**
    * Get full project data snapshot. Can be used to serialize the project and store it in a database.
    */

package/src/Projects/createJobRequestMessage.ts

@@ -67,7 +67,7 @@ function getTemplate() {
 
 function createJobRequestMessage(id: string, params: ProjectParams) {
   const template = getTemplate();
-  return {
+  const jobRequest: Record<string, any> = {
     ...template,
     keyFrames: [
       {
@@ -92,6 +92,10 @@ function createJobRequestMessage(id: string, params: ProjectParams) {
     jobID: id,
     disableSafety: !!params.disableNSFWFilter
   };
+  if (params.network) {
+    jobRequest.network = params.network;
+  }
+  return jobRequest;
 }
 
 export type JobRequestRaw = ReturnType<typeof createJobRequestMessage>;