mio-ai/src/platforms/openai/features/image.ts

import * as FormDataUtils from '../../../packages/form-data';
import {
  CreatedResource,
} from '../common';
import { DoFetch } from '../../../packages/request';
import { PlatformApiError } from '../../../common';

export enum ImageSize {
  SQUARE_256 = '256x256',
  SQUARE_512 = '512x512',
  SQUARE_1024 = '1024x1024',
}

export enum ImageResponseFormat {
  URL = 'url',
  BASE64_JSON = 'b64_json',
}

export interface CreateImageParams {
  prompt: string;
  n? : number;
  size?: ImageSize;
  user?: string;
}

export interface ImageData {
  b64_json: string;
}

export interface ImageDataEvent extends CreatedResource {
  data: Buffer[];
}

export function createImage(
  this: NodeJS.EventEmitter,
  doFetch: DoFetch,
  params: CreateImageParams,
) {
  doFetch('POST', '/images/generations', {
    prompt: params.prompt,
    n: params.n ?? 1,
    size: params.size ?? ImageSize.SQUARE_1024,
    user: params.user,
    response_format: ImageResponseFormat.BASE64_JSON,
  } as Record<string, unknown>)
    .then(async (response) => {
      if (!response.ok) {
        this.emit('error', new PlatformApiError(
          // eslint-disable-next-line @typescript-eslint/restrict-template-expressions
          `Request from platform returned with status: ${response.status}`,
          response,
        ));
        this.emit('end');
        return;
      }

      const responseData = await response.json() as Record<string, unknown>;
      const data = responseData.data as ImageData[];
      this.emit('data', {
        ...responseData,
        data: data.map((item) => Buffer.from(item.b64_json, 'base64')),
      });
      this.emit('end');
    })
    .catch((err) => {
      this.emit('error', err as Error);
      this.emit('end');
    });

  return this;
}

export interface CreateImageEditParams {
  image: Buffer;
  mask?: Buffer;
  prompt: string;
  n?: number;
  size?: ImageSize;
  user?: string;
}

export function createImageEdit(
  this: NodeJS.EventEmitter,
  doFetch: DoFetch,
  params: CreateImageEditParams,
) {
  doFetch('POST', '/images/edits', FormDataUtils.fromJson({
    prompt: params.prompt,
    image: params.image,
    mask: params.mask,
    n: params.n ?? 1,
    size: params.size ?? ImageSize.SQUARE_1024,
    response_format: ImageResponseFormat.BASE64_JSON,
  }))
    .then(async (response) => {
      if (!response.ok) {
        this.emit('error', new PlatformApiError(
          // eslint-disable-next-line @typescript-eslint/restrict-template-expressions
          `Request from platform returned with status: ${response.status}`,
          response,
        ));
        this.emit('end');
        return;
      }

      const responseData = await response.json() as Record<string, unknown>;
      const data = responseData.data as ImageData[];
      const emittedData = {
        ...responseData,
        data: data.map((item) => Buffer.from(item.b64_json, 'base64')),
      } as ImageDataEvent;
      this.emit('data', emittedData);
      this.emit('end');
    })
    .catch((err) => {
      this.emit('error', err as Error);
      this.emit('end');
    });

  return this;
}

export interface CreateImageVariationParams {
  image: Buffer;
  n?: number;
  size?: ImageSize;
  user?: string;
}

export function createImageVariation(
  this: NodeJS.EventEmitter,
  doFetch: DoFetch,
  params: CreateImageVariationParams,
) {
  doFetch('POST', '/images/variations', FormDataUtils.fromJson({
    image: params.image,
    n: params.n ?? 1,
    size: params.size ?? ImageSize.SQUARE_1024,
    response_format: ImageResponseFormat.BASE64_JSON,
  }))
    .then(async (response) => {
      if (!response.ok) {
        this.emit('error', new PlatformApiError(
          // eslint-disable-next-line @typescript-eslint/restrict-template-expressions
          `Request from platform returned with status: ${response.status}`,
          response,
        ));
        this.emit('end');
        return;
      }

      const responseData = await response.json() as Record<string, unknown>;
      const data = responseData.data as ImageData[];
      this.emit('data', {
        ...responseData,
        data: data.map((item) => Buffer.from(item.b64_json, 'base64')),
      });
      this.emit('end');
    })
    .catch((err) => {
      this.emit('error', err as Error);
      this.emit('end');
    });

  return this;
}