trustgraph/ts/packages/flow/src/model/text-completion/common.ts

import {
  TooManyRequestsError,
  errorMessage,
  type LlmChunk,
} from "@trustgraph/base";
import { Config, Effect } from "effect";
import * as O from "effect/Option";
import * as Predicate from "effect/Predicate";
import * as S from "effect/Schema";

export class TextCompletionConfigError extends S.TaggedErrorClass<TextCompletionConfigError>()(
  "TextCompletionConfigError",
  {
    message: S.String,
    provider: S.String,
    key: S.String,
  },
) {}

export class TextCompletionProviderError extends S.TaggedErrorClass<TextCompletionProviderError>()(
  "TextCompletionProviderError",
  {
    message: S.String,
    provider: S.String,
  },
) {}

export type TextCompletionRuntimeError =
  | TextCompletionProviderError
  | TooManyRequestsError;

export const optionalStringConfig = Effect.fn("TextCompletion.optionalStringConfig")(function*(
  provider: string,
  name: string,
) {
  const value = yield* Config.string(name).pipe(
    Config.option,
    Effect.mapError((cause) =>
      TextCompletionConfigError.make({
        provider,
        key: name,
        message: errorMessage(cause),
      })
    ),
  );
  return O.getOrUndefined(value);
});

export const requiredString = (
  value: string | undefined,
  provider: string,
  key: string,
  message: string,
) =>
  value !== undefined && value.length > 0
    ? Effect.succeed(value)
    : Effect.fail(TextCompletionConfigError.make({ provider, key, message }));

export const providerRuntimeError = (
  provider: string,
  error: unknown,
): TextCompletionRuntimeError =>
  TextCompletionProviderError.make({
    provider,
    message: errorMessage(error),
  });

export const providerStatusError = (
  provider: string,
  error: unknown,
): TextCompletionRuntimeError => {
  const status = Predicate.isObject(error) && Predicate.hasProperty(error, "status")
    ? error.status
    : undefined;
  const statusCode = Predicate.isObject(error) && Predicate.hasProperty(error, "statusCode")
    ? error.statusCode
    : undefined;
  return status === 429 || statusCode === 429
    ? TooManyRequestsError.make({ message: "Rate limit exceeded" })
    : providerRuntimeError(provider, error);
};

export const toAsyncGenerator = (
  iterable: AsyncIterable<LlmChunk>,
  mapError: (error: unknown) => TextCompletionRuntimeError,
): AsyncGenerator<LlmChunk> => {
  const iterator = iterable[Symbol.asyncIterator]();
  let generator: AsyncGenerator<LlmChunk>;
  generator = {
    next: (value?: unknown) => iterator.next(value),
    return: (value?: unknown) =>
      iterator.return === undefined
        ? Promise.resolve({ done: true, value })
        : iterator.return(value),
    throw: (error?: unknown) =>
      iterator.throw === undefined
        ? Promise.reject(mapError(error))
        : iterator.throw(error),
    [Symbol.asyncIterator]: () => generator,
  };
  return generator;
};