mio-ai/test/platforms/openai/api.test.ts

import { config } from 'dotenv';
import {
  beforeAll,
  beforeEach,
  describe,
  expect,
  it,
} from 'vitest';
import {
  createAiClient,
  PlatformEventEmitter,
  OpenAi,
} from '../../../src';
import { getPromptTokens } from '../../../src/platforms/openai/usage';
import { ChatCompletionModel } from '../../../src/platforms/openai';

describe('OpenAI', () => {
  beforeAll(() => {
    config();
  });

  describe.skip('API', () => {
    let aiClient: PlatformEventEmitter;

    beforeEach(() => {
      aiClient = createAiClient({
        platform: OpenAi.PLATFORM_ID,
        platformConfiguration: {
          apiKey: process.env.OPENAI_API_KEY as string,
          organizationId: process.env.OPENAI_ORGANIZATION_ID as string,
          apiVersion: OpenAi.ApiVersion.V1,
        },
      });
    });

    describe('createChatCompletion', () => {
      let result: Partial<OpenAi.ChatCompletion> | undefined;

      let prompt: string;

      beforeEach(() => {
        result = undefined;

        aiClient.on<OpenAi.ChatCompletionChunkDataEvent>('data', (d) => {
          d.choices.forEach((c) => {
            if (!result) {
              const promptTokens = getPromptTokens(
                prompt,
                ChatCompletionModel.GPT_3_5_TURBO,
              )
                .length;

              result = {
                id: d.id,
                object: OpenAi.ChatCompletionDataEventObjectType.CHAT_COMPLETION,
                created: d.created,
                model: d.model,
                usage: {
                  prompt_tokens: promptTokens,
                  completion_tokens: 0,
                  total_tokens: promptTokens,
                },
              };
            }

            if (!Array.isArray(result?.choices)) {
              result.choices = [];
            }

            if (!result.choices[c.index]) {
              result.choices[c.index] = {
                message: { content: '' },
                index: c.index,
                finish_reason: c.finish_reason,
              };
            }

            if (result.choices[c.index].message) {
              if (c.delta.role) {
                (result.choices[c.index].message as Record<string, unknown>).role = c.delta.role;
              }

              if (c.delta.content) {
                (result.choices[c.index].message as Record<string, unknown>)
                  .content += c.delta.content;

                result.usage!.completion_tokens += 1;
                result.usage!.total_tokens = (
                  result.usage!.prompt_tokens + result.usage!.completion_tokens
                );
              }
            }

            if (c.finish_reason) {
              result.choices[c.index].finish_reason = c.finish_reason;
            }
          });
        });
      });

      it('works', () => new Promise<void>((resolve, reject) => {
        aiClient.on('end', () => {
          expect(result).toHaveProperty('id', expect.any(String));
          expect(result).toHaveProperty('object', OpenAi.ChatCompletionDataEventObjectType.CHAT_COMPLETION);
          expect(result).toHaveProperty('model', expect.any(String));
          expect(result).toHaveProperty('created', expect.any(Number));
          expect(result).toHaveProperty('choices', expect.any(Array));
          resolve();
        });

        aiClient.on('error', (error: Error) => {
          reject(error);
        });

        prompt = 'Count from 1 to 20 in increments of a random number from 1 to 10.';
        aiClient.createChatCompletion({
          messages: prompt,
          model: OpenAi.ChatCompletionModel.GPT_3_5_TURBO,
          n: 2,
        });
      }), { timeout: 10000 });
    });

    describe('createImage', () => {
      it('works', () => new Promise<void>((resolve, reject) => {
        aiClient.on<OpenAi.CreateImageDataEvent>('data', (r) => {
          expect(r).toHaveProperty('created', expect.any(Number));
          expect(r).toHaveProperty('data', expect.any(Array));
          expect(r.data.every((d) => d instanceof Buffer)).toBe(true);
        });

        aiClient.on('end', () => {
          resolve();
        });

        aiClient.on('error', (error: Error) => {
          reject(error);
        });

        aiClient.createImage({
          prompt: 'A photo of a cat',
          size: OpenAi.CreateImageSize.SQUARE_256,
        });
      }), { timeout: 10000 });
    });

    describe('createCompletion', () => {
      let result: Partial<OpenAi.TextCompletion> | undefined;

      beforeEach(() => {
        result = undefined;

        aiClient.on<OpenAi.TextCompletionChunkDataEvent>('data', (d) => {
          d.choices.forEach((c) => {
            if (!result) {
              result = {
                id: d.id,
                object: OpenAi.TextCompletionDataEventObjectType.TEXT_COMPLETION,
                created: d.created,
                model: d.model,
              };
            }

            if (!Array.isArray(result?.choices)) {
              result.choices = [];
            }

            if (!result.choices[c.index]) {
              result.choices[c.index] = {
                text: '',
                index: c.index,
                finish_reason: c.finish_reason,
                logprobs: c.logprobs, // TODO dunno how to use this?
              };
            }

            if (c.text) {
              result.choices[c.index].text += c.text;
            }

            if (c.finish_reason) {
              result.choices[c.index].finish_reason = c.finish_reason;
            }
          });
        });
      });

      it('works', () => new Promise<void>((resolve, reject) => {
        aiClient.on('end', () => {
          expect(result).toHaveProperty('id', expect.any(String));
          expect(result).toHaveProperty('object', OpenAi.TextCompletionDataEventObjectType.TEXT_COMPLETION);
          expect(result).toHaveProperty('model', expect.any(String));
          expect(result).toHaveProperty('created', expect.any(Number));
          expect(result).toHaveProperty('choices', expect.any(Array));
          resolve();
        });

        aiClient.on('error', (error: Error) => {
          reject(error);
        });

        aiClient.createCompletion({
          prompt: 'Say this is a test',
          model: OpenAi.TextCompletionModel.TEXT_DAVINCI_003,
          maxTokens: 7,
          temperature: 0,
        });
      }), { timeout: 10000 });
    });

    describe('createEdit', () => {
      it('works', () => new Promise<void>((resolve, reject) => {
        aiClient.on<OpenAi.CreateEditDataEvent>('data', (r) => {
          expect(r).toHaveProperty('object', OpenAi.EditDataEventObjectType.EDIT);
          expect(r).toHaveProperty('created', expect.any(Number));
          expect(r).toHaveProperty('choices', expect.any(Array));
        });

        aiClient.on('end', () => {
          resolve();
        });

        aiClient.on('error', (error: Error) => {
          reject(error);
        });

        aiClient.createEdit({
          model: OpenAi.EditModel.TEXT_DAVINCI_EDIT_001,
          input: 'What day of the wek is it?',
          instruction: 'Fix the spelling mistakes',
        });
      }), { timeout: 10000 });
    });
  });
});