mio-ai/src/platforms/openai/features/text-completion.ts

import { TextCompletionModel } from '../models';
import {
  DataEventId,
  FinishableChoiceBase,
  CreatedResource,
} from '../common';
import {
  CompletionUsage,
  UsageMetadata,
} from '../usage';
import { ConsumeStream, DoFetch } from '../../../packages/request';
import { PlatformApiError } from '../../../common';

export enum DataEventObjectType {
  TEXT_COMPLETION = 'text_completion',
}

export interface CreateTextCompletionParams {
  model: TextCompletionModel;
  prompt: string;
  temperature?: number;
  topP?: number;
  n?: number;
  stop?: string | string[];
  maxTokens?: number;
  logprobs?: number;
  echo?: boolean;
  bestOf?: number;
  logitBias?: Record<string, number>;
  user?: string;
  presencePenalty?: number;
  frequencyPenalty?: number;
}

export interface TextCompletionChoice extends FinishableChoiceBase {
  text: string;
  logprobs?: number;
}

export interface CreateTextCompletionDataEvent<
  C extends Partial<FinishableChoiceBase>
> extends CreatedResource {
  id: DataEventId;
  object: DataEventObjectType;
  model: TextCompletionModel;
  choices: C[];
}

export interface TextCompletion
  extends CreateTextCompletionDataEvent<Partial<TextCompletionChoice>>,
    UsageMetadata<CompletionUsage> {}

export type TextCompletionChunkDataEvent = CreateTextCompletionDataEvent<TextCompletionChoice>;

export function createTextCompletion(
  this: NodeJS.EventEmitter,
  doFetch: DoFetch,
  consumeStream: ConsumeStream,
  params: CreateTextCompletionParams,
) {
  doFetch('POST', '/completions', {
    model: params.model ?? TextCompletionModel.TEXT_DAVINCI_003,
    prompt: params.prompt,
    temperature: params.temperature ?? 1,
    top_p: params.topP ?? 1,
    n: params.n ?? 1,
    stop: params.stop ?? null,
    stream: true,
    max_tokens: params.maxTokens,
    logprobs: params.logprobs,
    echo: params.echo,
    best_of: params.bestOf,
    logit_bias: params.logitBias,
    user: params.user,
    presence_penalty: params.presencePenalty,
    frequency_penalty: params.frequencyPenalty,
  } as Record<string, unknown>)
    .then(async (response) => {
      if (!response.ok) {
        this.emit('error', new PlatformApiError(
          // eslint-disable-next-line @typescript-eslint/restrict-template-expressions
          `Create text completion returned with status: ${response.status}`,
          response,
        ));
        this.emit('end');
        return;
      }

      await consumeStream(response);
      this.emit('end');
    })
    .catch((err) => {
      this.emit('error', err as Error);
      this.emit('end');
    });
}