mio-ai/src/platforms/openai/events.ts

import { PassThrough } from 'stream';
import { EventEmitter } from 'events';
import fetchPonyfill from 'fetch-ponyfill';
import { DoFetchBody, processRequest } from '../../packages/request';
import * as AllPlatformsCommon from '../../common';
import { Configuration, DEFAULT_BASE_URL } from './common';
import { createTextCompletion, CreateTextCompletionParams } from './features/text-completion';
import { CreateChatCompletionParams, createChatCompletion } from './features/chat-completion';
import {
  CreateImageParams,
  createImage,
  CreateImageEditParams,
  createImageEdit,
  CreateImageVariationParams,
  createImageVariation,
} from './features/image';
import { CreateEditParams, createEdit } from './features/edit';
import { listModels, ModelId, retrieveModel } from './features/model';
import { createEmbedding, CreateEmbeddingParams } from './features/embedding';

export interface PlatformEventEmitter extends AllPlatformsCommon.PlatformEventEmitter {
  listModels(): void;
  retrieveModel(modelId: ModelId): void;
  createCompletion(params: CreateTextCompletionParams): void;
  createChatCompletion(params: CreateChatCompletionParams): void;
  createEdit(params: CreateEditParams): void;
  createImage(params: CreateImageParams): void;
  createImageEdit(params: CreateImageEditParams): void;
  createImageVariation(params: CreateImageVariationParams): void;
  createEmbedding(params: CreateEmbeddingParams): void;
}

export class PlatformEventEmitterImpl extends EventEmitter implements PlatformEventEmitter {
  readonly listModels: PlatformEventEmitter['listModels'];

  readonly retrieveModel: PlatformEventEmitter['retrieveModel'];

  readonly createCompletion: PlatformEventEmitter['createCompletion'];

  readonly createChatCompletion: PlatformEventEmitter['createChatCompletion'];

  readonly createEdit: PlatformEventEmitter['createEdit'];

  readonly createImage: PlatformEventEmitter['createImage'];

  readonly createImageEdit: PlatformEventEmitter['createImageEdit'];

  readonly createImageVariation: PlatformEventEmitter['createImageVariation'];

  readonly createEmbedding: PlatformEventEmitter['createEmbedding'];

  constructor(configParams: Configuration) {
    super();
    const platformHeaders: Record<string, string> = {
      Authorization: `Bearer ${configParams.apiKey}`,
    };

    if (configParams.organizationId) {
      platformHeaders['OpenAI-Organization'] = configParams.organizationId;
    }

    const { fetch: fetchInstance } = fetchPonyfill();
    const doFetch = (method: string, path: string, body?: DoFetchBody) => {
      let finalBody: BodyInit | undefined;
      let finalHeaders = {
        ...platformHeaders,
      };
      if (body) {
        const finalRequest = processRequest(body, finalHeaders);
        finalBody = finalRequest.body;
        if (finalRequest.headers) {
          finalHeaders = finalRequest.headers;
        }
      }

      const theFetchParams: Record<string, unknown> = {
        method,
        headers: finalHeaders,
      };

      if (finalBody) {
        theFetchParams.body = finalBody;
      }

      const url = new URL(
        `/${configParams.apiVersion}${path}`,
        configParams.baseUrl ?? DEFAULT_BASE_URL,
      ).toString();

      this.emit('start', {
        ...theFetchParams,
        url,
      });

      return fetchInstance(url, theFetchParams);
    };

    const consumeStream = async (response: Response) => {
      // eslint-disable-next-line no-restricted-syntax
      for await (const chunk of response.body as unknown as PassThrough) {
        const chunkStringMaybeMultiple = chunk.toString();
        const chunkStrings = chunkStringMaybeMultiple
          .trim()
          .split('\n')
          .filter((chunkString: string) => chunkString.length > 0);
        chunkStrings.forEach((chunkString: string) => {
          const dataRaw = chunkString.split('data: ').at(1);
          if (!dataRaw) {
            return;
          }
          if (dataRaw === '[DONE]') {
            return;
          }
          const data = JSON.parse(dataRaw);
          this.emit('data', data);
        });
      }
    };
    this.listModels = listModels.bind(this, doFetch);
    this.retrieveModel = retrieveModel.bind(this, doFetch);
    this.createCompletion = createTextCompletion.bind(this, doFetch, consumeStream);
    this.createChatCompletion = createChatCompletion.bind(this, doFetch, consumeStream);
    this.createEdit = createEdit.bind(this, doFetch);
    this.createImage = createImage.bind(this, doFetch);
    this.createImageEdit = createImageEdit.bind(this, doFetch);
    this.createImageVariation = createImageVariation.bind(this, doFetch);
    this.createEmbedding = createEmbedding.bind(this, doFetch);
  }
}