Function makeDataLoader

packages/ai/ai/src/EmbeddingModel.ts:276–318

(options: {
  /**
   * A method which processes a batch of text inputs and returns embedding
   * results.
   */
  readonly embedMany: (input: ReadonlyArray<string>) => Effect.Effect<Array<Result>, AiError.AiError>
  /**
   * The duration between batch requests during which requests are collected and
   * added to the current batch.
   */
  readonly window: Duration.DurationInput
  /**
   * Optional maximum number of requests to add to the batch before a batch
   * request must be sent.
   */
  readonly maxBatchSize?: number
})

 * @category Constructors
 */
export const makeDataLoader = (options: {
  /**
   * A method which processes a batch of text inputs and returns embedding
   * results.
   */
  readonly embedMany: (input: ReadonlyArray<string>) => Effect.Effect<Array<Result>, AiError.AiError>
  /**
   * The duration between batch requests during which requests are collected and
   * added to the current batch.
   */
  readonly window: Duration.DurationInput
  /**
   * Optional maximum number of requests to add to the batch before a batch
   * request must be sent.
   */
  readonly maxBatchSize?: number
}) =>
  Effect.gen(function*() {
    const resolver = makeBatchedResolver(options.embedMany)
    const resolverDelayed = yield* dataLoader(resolver, {
      window: options.window,
      maxBatchSize: options.maxBatchSize
    })

    function embed(input: string) {
      return Effect.request(new EmbeddingRequest({ input }), resolverDelayed).pipe(
        Effect.withSpan("EmbeddingModel.embed", { captureStackTrace: false })
      )
    }

    function embedMany(inputs: ReadonlyArray<string>, options?: {
      readonly concurrency?: Types.Concurrency | undefined
    }) {
      return Effect.forEach(inputs, embed, { batching: true, concurrency: options?.concurrency }).pipe(
        Effect.withSpan("EmbeddingModel.embedMany", { captureStackTrace: false })
      )
    }

    return EmbeddingModel.of({
      embed,
      embedMany
    })
  })