chatgpt-api/src/utils.ts

import { customAlphabet, urlAlphabet } from 'nanoid'
import type { ThrottledFunction } from 'p-throttle'

import * as types from './types'

/**
 * Extracts a JSON object string from a given string.
 *
 * @param text - string from which to extract the JSON object
 * @returns extracted JSON object string, or `undefined` if no JSON object is found
 */
export function extractJSONObjectFromString(text: string): string | undefined {
  return text.match(/\{(.|\n)*\}/gm)?.[0] // FIXME: This breaks if there are multiple JSON objects in the string
}

/**
 * Extracts a JSON array string from a given string.
 *
 * @param text - string from which to extract the JSON array
 * @returns extracted JSON array string, or `undefined` if no JSON array is found
 */
export function extractJSONArrayFromString(text: string): string | undefined {
  return text.match(/\[(.|\n)*\]/gm)?.[0] // FIXME: This breaks if there are multiple JSON arrays in the string
}

/**
 * Pauses the execution of a function for a specified time.
 *
 * @param ms - number of milliseconds to pause
 * @returns promise that resolves after the specified number of milliseconds
 */
export function sleep(ms: number) {
  return new Promise((resolve) => setTimeout(resolve, ms))
}

/**
 * A default ID generator function that uses a custom alphabet based on URL safe symbols.
 */
export const defaultIDGeneratorFn: types.IDGeneratorFunction =
  customAlphabet(urlAlphabet)

const taskNameRegex = /^[a-zA-Z_][a-zA-Z0-9_-]{0,63}$/
export function isValidTaskIdentifier(id: string): boolean {
  return !!id && taskNameRegex.test(id)
}

export function extractFunctionIdentifierFromString(
  text: string
): string | undefined {
  text = text?.trim()

  if (!text) {
    return
  }

  if (isValidTaskIdentifier(text)) {
    return text
  }

  const splits = text
    .split(/[^a-zA-Z0-9_-]/)
    .map((s) => {
      s = s.trim()
      return isValidTaskIdentifier(s) ? s : undefined
    })
    .filter(Boolean)

  return splits[splits.length - 1]
}

/**
 * Chunks a string into an array of chunks.
 *
 * @param text - string to chunk
 * @param maxLength - maximum length of each chunk
 * @returns array of chunks
 */
export function chunkString(text: string, maxLength: number): string[] {
  const words = text.split(' ')
  const chunks: string[] = []
  let chunk = ''

  for (const word of words) {
    if (word.length > maxLength) {
      // Truncate the word if it's too long and indicate that it was truncated:
      chunks.push(word.substring(0, maxLength - 3) + '...')
    } else if ((chunk + word + 1).length > maxLength) {
      // +1 accounts for the space between words
      chunks.push(chunk.trim())
      chunk = word
    } else {
      chunk += (chunk ? ' ' : '') + word
    }
  }

  if (chunk) {
    chunks.push(chunk.trim())
  }

  return chunks
}

/**
 * Stringifies a JSON value for use in an LLM prompt.
 *
 * @param json - JSON value to stringify
 * @returns stringified value with all double quotes around object keys removed
 */
export function stringifyForModel(json: types.Jsonifiable): string {
  const UNIQUE_PREFIX = defaultIDGeneratorFn()
  return (
    JSON.stringify(json, replacer)
      // Remove all double quotes around keys:
      .replace(new RegExp('"' + UNIQUE_PREFIX + '(.*?)"', 'g'), '$1')
  )

  /**
   * Replacer function prefixing all keys with a unique identifier.
   */
  function replacer(_: string, value: any) {
    if (value && typeof value === 'object') {
      if (Array.isArray(value)) {
        return value
      }

      const replacement = {}

      for (const k in value) {
        if (Object.hasOwnProperty.call(value, k)) {
          replacement[UNIQUE_PREFIX + k] = value[k]
        }
      }

      return replacement
    }

    return value
  }
}

export function pick<T extends object, U = T>(obj: T, ...keys: string[]): U {
  return Object.fromEntries(
    keys.filter((key) => key in obj).map((key) => [key, obj[key]])
  ) as U
}

export function omit<T extends object, U = T>(obj: T, ...keys: string[]): U {
  return Object.fromEntries<T>(
    Object.entries(obj).filter(([key]) => !keys.includes(key))
  ) as U
}

const noop = () => undefined

/**
 * Throttles HTTP requests made by a ky instance. Very useful for enforcing rate limits.
 */
export function throttleKy(
  ky: types.KyInstance,
  throttleFn: <Argument extends readonly unknown[], ReturnValue>(
    function_: (...args: Argument) => ReturnValue
  ) => ThrottledFunction<Argument, ReturnValue>
) {
  return ky.extend({
    hooks: {
      beforeRequest: [throttleFn(noop)]
    }
  })
}