@effect/ai-openai 4.0.0-beta.4 → 4.0.0-beta.41

package/src/OpenAiClient.ts

@@ -7,21 +7,32 @@
  * @since 1.0.0
  */
 import * as Array from "effect/Array"
+import * as Cause from "effect/Cause"
 import type * as Config from "effect/Config"
 import * as Effect from "effect/Effect"
+import * as Exit from "effect/Exit"
 import { identity } from "effect/Function"
+import * as Function from "effect/Function"
 import * as Layer from "effect/Layer"
 import * as Predicate from "effect/Predicate"
+import * as Queue from "effect/Queue"
+import * as RcRef from "effect/RcRef"
 import * as Redacted from "effect/Redacted"
+import * as Schedule from "effect/Schedule"
+import * as Schema from "effect/Schema"
+import * as Scope from "effect/Scope"
+import * as Semaphore from "effect/Semaphore"
 import * as ServiceMap from "effect/ServiceMap"
 import * as Stream from "effect/Stream"
-import type * as AiError from "effect/unstable/ai/AiError"
+import * as AiError from "effect/unstable/ai/AiError"
+import * as ResponseIdTracker from "effect/unstable/ai/ResponseIdTracker"
 import * as Sse from "effect/unstable/encoding/Sse"
 import * as Headers from "effect/unstable/http/Headers"
 import * as HttpBody from "effect/unstable/http/HttpBody"
 import * as HttpClient from "effect/unstable/http/HttpClient"
 import * as HttpClientRequest from "effect/unstable/http/HttpClientRequest"
-import type * as HttpClientResponse from "effect/unstable/http/HttpClientResponse"
+import * as HttpClientResponse from "effect/unstable/http/HttpClientResponse"
+import * as Socket from "effect/unstable/socket/Socket"
 import * as Generated from "./Generated.ts"
 import * as Errors from "./internal/errors.ts"
 import { OpenAiConfig } from "./OpenAiConfig.ts"
@@ -48,7 +59,7 @@ export interface Service {
   readonly createResponse: (
     options: typeof Generated.CreateResponse.Encoded
   ) => Effect.Effect<
-    [body: typeof Generated.Response.Type, response: HttpClientResponse.HttpClientResponse],
+    readonly [body: typeof Generated.Response.Type, response: HttpClientResponse.HttpClientResponse],
     AiError.AiError
   >
 
@@ -58,7 +69,7 @@ export interface Service {
   readonly createResponseStream: (
     options: Omit<typeof Generated.CreateResponse.Encoded, "stream">
   ) => Effect.Effect<
-    [
+    readonly [
       response: HttpClientResponse.HttpClientResponse,
       stream: Stream.Stream<typeof Generated.ResponseStreamEvent.Type, AiError.AiError>
     ],
@@ -142,32 +153,33 @@ const RedactedOpenAiHeaders = {
  * @category constructors
  */
 export const make = Effect.fnUntraced(
-  function*(options: Options): Effect.fn.Return<Service, never, HttpClient.HttpClient> {
+  function*(
+    options: Options
+  ): Effect.fn.Return<Service, never, HttpClient.HttpClient> {
     const baseClient = yield* HttpClient.HttpClient
+    const apiUrl = options.apiUrl ?? "https://api.openai.com/v1"
 
     const httpClient = baseClient.pipe(
-      HttpClient.mapRequest((request) =>
-        request.pipe(
-          HttpClientRequest.prependUrl(options.apiUrl ?? "https://api.openai.com/v1"),
-          Predicate.isNotUndefined(options.apiKey)
-            ? HttpClientRequest.bearerToken(Redacted.value(options.apiKey))
-            : identity,
-          Predicate.isNotUndefined(options.organizationId)
-            ? HttpClientRequest.setHeader(
-              RedactedOpenAiHeaders.OpenAiOrganization,
-              Redacted.value(options.organizationId)
-            )
-            : identity,
-          Predicate.isNotUndefined(options.projectId)
-            ? HttpClientRequest.setHeader(
-              RedactedOpenAiHeaders.OpenAiProject,
-              Redacted.value(options.projectId)
-            )
-            : identity,
-          HttpClientRequest.acceptJson
-        )
-      ),
-      Predicate.isNotUndefined(options.transformClient)
+      HttpClient.mapRequest(Function.flow(
+        HttpClientRequest.prependUrl(apiUrl),
+        options.apiKey
+          ? HttpClientRequest.bearerToken(Redacted.value(options.apiKey))
+          : identity,
+        options.organizationId
+          ? HttpClientRequest.setHeader(
+            RedactedOpenAiHeaders.OpenAiOrganization,
+            Redacted.value(options.organizationId)
+          )
+          : identity,
+        options.projectId
+          ? HttpClientRequest.setHeader(
+            RedactedOpenAiHeaders.OpenAiProject,
+            Redacted.value(options.projectId)
+          )
+          : identity,
+        HttpClientRequest.acceptJson
+      )),
+      options.transformClient
         ? options.transformClient
         : identity
     )
@@ -222,17 +234,21 @@ export const make = Effect.fnUntraced(
     }
 
     const createResponseStream: Service["createResponseStream"] = (payload) =>
-      httpClientOk.execute(
-        HttpClientRequest.post("/responses", {
-          body: HttpBody.jsonUnsafe({ ...payload, stream: true })
-        })
-      ).pipe(
-        Effect.map(buildResponseStream),
-        Effect.catchTag(
-          "HttpClientError",
-          (error) => Errors.mapHttpClientError(error, "createResponseStream")
+      Effect.servicesWith((services) => {
+        const socket = ServiceMap.getOrUndefined(services, OpenAiSocket)
+        if (socket) return socket.createResponseStream(payload)
+        return httpClientOk.execute(
+          HttpClientRequest.post("/responses", {
+            body: HttpBody.jsonUnsafe({ ...payload, stream: true })
+          })
+        ).pipe(
+          Effect.map(buildResponseStream),
+          Effect.catchTag(
+            "HttpClientError",
+            (error) => Errors.mapHttpClientError(error, "createResponseStream")
+          )
         )
-      )
+      })
 
     const createEmbedding = (
       payload: typeof Generated.CreateEmbeddingRequest.Encoded
@@ -281,7 +297,7 @@ export const layerConfig = (options?: {
   /**
    * The config value to load for the API key.
    */
-  readonly apiKey?: Config.Config<Redacted.Redacted<string>> | undefined
+  readonly apiKey?: Config.Config<Redacted.Redacted<string> | undefined> | undefined
 
   /**
    * The config value to load for the API URL.
@@ -291,12 +307,12 @@ export const layerConfig = (options?: {
   /**
    * The config value to load for the organization ID.
    */
-  readonly organizationId?: Config.Config<Redacted.Redacted<string>> | undefined
+  readonly organizationId?: Config.Config<Redacted.Redacted<string> | undefined> | undefined
 
   /**
    * The config value to load for the project ID.
    */
-  readonly projectId?: Config.Config<Redacted.Redacted<string>> | undefined
+  readonly projectId?: Config.Config<Redacted.Redacted<string> | undefined> | undefined
 
   /**
    * Optional transformer for the HTTP client.
@@ -327,3 +343,270 @@ export const layerConfig = (options?: {
       })
     })
   )
+
+// =============================================================================
+// Websocket mode
+// =============================================================================
+
+/**
+ * @since 1.0.0
+ * @category Events
+ */
+export type ResponseStreamEvent = typeof Generated.ResponseStreamEvent.Type
+
+/**
+ * @since 1.0.0
+ * @category Websocket mode
+ */
+export class OpenAiSocket extends ServiceMap.Service<OpenAiSocket, {
+  /**
+   * Create a streaming response using the OpenAI responses endpoint.
+   */
+  readonly createResponseStream: (
+    options: Omit<typeof Generated.CreateResponse.Encoded, "stream">
+  ) => Effect.Effect<
+    readonly [
+      response: HttpClientResponse.HttpClientResponse,
+      stream: Stream.Stream<ResponseStreamEvent, AiError.AiError>
+    ],
+    AiError.AiError
+  >
+}>()("@effect/ai-openai/OpenAiClient/OpenAiSocket") {}
+
+const makeSocket = Effect.gen(function*() {
+  const client = yield* OpenAiClient
+  const tracker = yield* ResponseIdTracker.make
+  const request = yield* Effect.orDie(client.client.httpClient.preprocess(HttpClientRequest.post("/responses")))
+
+  const socket = yield* Socket.makeWebSocket(request.url.replace(/^http/, "ws")).pipe(
+    Effect.updateService(Socket.WebSocketConstructor, (f) => (url) =>
+      f(url, {
+        headers: request.headers
+      } as any))
+  )
+
+  const queueRef: RcRef.RcRef<
+    {
+      readonly send: (
+        queue: Queue.Enqueue<ResponseStreamEvent, AiError.AiError | Cause.Done>,
+        message: typeof Generated.CreateResponse.Encoded
+      ) => Effect.Effect<void, AiError.AiError>
+      readonly reset: () => void
+    }
+  > = yield* RcRef.make({
+    idleTimeToLive: 60_000,
+    acquire: Effect.gen(function*() {
+      const write = yield* socket.writer
+
+      let currentQueue: Queue.Enqueue<ResponseStreamEvent, AiError.AiError | Cause.Done> | null = null
+      const send = (
+        queue: Queue.Enqueue<ResponseStreamEvent, AiError.AiError | Cause.Done>,
+        message: typeof Generated.CreateResponse.Encoded
+      ) =>
+        Effect.suspend(() => {
+          currentQueue = queue
+          return write(JSON.stringify({
+            type: "response.create",
+            ...message
+          }))
+        }).pipe(
+          Effect.mapError((_error) =>
+            AiError.make({
+              module: "OpenAiClient",
+              method: "createResponseStream",
+              reason: new AiError.NetworkError({
+                reason: "TransportError",
+                request: {
+                  method: "POST",
+                  url: request.url,
+                  urlParams: [],
+                  hash: undefined,
+                  headers: request.headers
+                },
+                description: "Failed to send message over WebSocket"
+              })
+            })
+          )
+        )
+      const reset = () => {
+        currentQueue = null
+      }
+
+      const decoder = new TextDecoder()
+      yield* socket.runRaw((msg) => {
+        if (!currentQueue) return
+        const text = typeof msg === "string" ? msg : decoder.decode(msg)
+        try {
+          const event = decodeEvent(text)
+          if (event.type === "error") {
+            tracker.clearUnsafe()
+          }
+          if (event.type === "error" && "status" in event) {
+            const json = JSON.stringify(event.error)
+            return Effect.fail(
+              AiError.make({
+                module: "OpenAiClient",
+                method: "createResponseStream",
+                reason: AiError.reasonFromHttpStatus({
+                  description: json,
+                  status: event.status,
+                  metadata: event.error,
+                  http: {
+                    body: json,
+                    request: {
+                      method: "POST",
+                      url: request.url,
+                      urlParams: [],
+                      hash: undefined,
+                      headers: request.headers
+                    }
+                  }
+                })
+              })
+            )
+          }
+          Queue.offerUnsafe(currentQueue, event)
+        } catch {}
+      }).pipe(
+        Effect.catchCause((cause) => {
+          tracker.clearUnsafe()
+          return currentQueue ?
+            Queue.fail(
+              currentQueue,
+              AiError.make({
+                module: "OpenAiClient",
+                method: "createResponseStream",
+                reason: new AiError.NetworkError({
+                  reason: "TransportError",
+                  request: {
+                    method: "POST",
+                    url: request.url,
+                    urlParams: [],
+                    hash: undefined,
+                    headers: request.headers
+                  },
+                  description: Cause.pretty(cause)
+                })
+              })
+            ) :
+            Effect.void
+        }),
+        Effect.repeat(
+          Schedule.exponential(100, 1.5).pipe(
+            Schedule.either(Schedule.spaced({ seconds: 5 })),
+            Schedule.jittered
+          )
+        ),
+        Effect.forkScoped
+      )
+
+      return { send, reset } as const
+    })
+  })
+
+  // Websocket mode only allows one request at a time
+  const semaphore = Semaphore.makeUnsafe(1)
+
+  return OpenAiSocket.serviceMap({
+    createResponseStream(options) {
+      const stream = Effect.gen(function*() {
+        yield* Effect.acquireRelease(
+          semaphore.take(1),
+          () => semaphore.release(1),
+          { interruptible: true }
+        )
+        const { send, reset } = yield* RcRef.get(queueRef)
+        const incoming = yield* Queue.unbounded<ResponseStreamEvent, AiError.AiError | Cause.Done>()
+        let done = false
+
+        yield* Effect.acquireRelease(
+          send(incoming, options),
+          (_, exit) => {
+            reset()
+            if (Exit.isFailure(exit) && !Exit.hasInterrupts(exit)) return Effect.void
+            else if (done) return Effect.void
+            return RcRef.invalidate(queueRef)
+          },
+          { interruptible: true }
+        ).pipe(
+          Effect.forkScoped({ startImmediately: true })
+        )
+        return Stream.fromQueue(incoming).pipe(
+          Stream.takeUntil((e) => {
+            done = e.type === "response.completed" || e.type === "response.incomplete"
+            return done
+          })
+        )
+      }).pipe(Stream.unwrap)
+
+      return Effect.succeed([
+        HttpClientResponse.fromWeb(request, new Response()),
+        stream
+      ])
+    }
+  }).pipe(
+    ServiceMap.add(ResponseIdTracker.ResponseIdTracker, tracker)
+  )
+})
+
+const ErrorEvent = Schema.Struct({
+  type: Schema.Literal("error"),
+  status: Schema.Number.pipe(
+    Schema.withDecodingDefault(() => 500)
+  ),
+  error: Schema.Struct({
+    type: Schema.String,
+    message: Schema.String
+  })
+})
+
+const AllEvents = Schema.Union([ErrorEvent, Generated.ResponseStreamEvent])
+const decodeEvent = Schema.decodeUnknownSync(Schema.fromJsonString(AllEvents))
+
+/**
+ * Uses OpenAI's websocket mode for all responses within the provided effect.
+ *
+ * Note: This only works with the following WebSocket constructor layers:
+ *
+ * - `NodeSocket.layerWebSocketConstructorWS`
+ * - `BunSocket.layerWebSocketConstructor`
+ *
+ * This is because it needs to use non-standard options for setting the
+ * Authorization header.
+ *
+ * @since 1.0.0
+ * @category Websocket mode
+ */
+export const withWebSocketMode = <A, E, R>(
+  effect: Effect.Effect<A, E, R>
+): Effect.Effect<
+  A,
+  E,
+  Exclude<R, OpenAiSocket | ResponseIdTracker.ResponseIdTracker> | OpenAiClient | Socket.WebSocketConstructor
+> =>
+  Effect.scopedWith((scope) =>
+    Effect.flatMap(
+      Scope.provide(makeSocket, scope),
+      (services) => Effect.provideServices(effect, services)
+    )
+  )
+
+/**
+ * Uses OpenAI's websocket mode for all responses that use the Layer.
+ *
+ * Note: This only works with the following WebSocket constructor layers:
+ *
+ * - `NodeSocket.layerWebSocketConstructorWS`
+ * - `BunSocket.layerWebSocketConstructor`
+ *
+ * This is because it needs to use non-standard options for setting the
+ * Authorization header.
+ *
+ * @since 1.0.0
+ * @category Websocket mode
+ */
+export const layerWebSocketMode: Layer.Layer<
+  OpenAiSocket | ResponseIdTracker.ResponseIdTracker,
+  never,
+  OpenAiClient | Socket.WebSocketConstructor
+> = Layer.effectServices(makeSocket)

package/src/OpenAiEmbeddingModel.ts

@@ -0,0 +1,200 @@
+/**
+ * OpenAI Embedding Model implementation.
+ *
+ * Provides an EmbeddingModel implementation for OpenAI's embeddings API.
+ *
+ * @since 1.0.0
+ */
+import * as Effect from "effect/Effect"
+import { dual } from "effect/Function"
+import * as Layer from "effect/Layer"
+import * as ServiceMap from "effect/ServiceMap"
+import type { Simplify } from "effect/Types"
+import * as AiError from "effect/unstable/ai/AiError"
+import * as EmbeddingModel from "effect/unstable/ai/EmbeddingModel"
+import * as AiModel from "effect/unstable/ai/Model"
+import type * as Generated from "./Generated.ts"
+import { OpenAiClient } from "./OpenAiClient.ts"
+
+/**
+ * @since 1.0.0
+ * @category models
+ */
+export type Model = "text-embedding-ada-002" | "text-embedding-3-small" | "text-embedding-3-large"
+
+/**
+ * Service definition for OpenAI embedding model configuration.
+ *
+ * @since 1.0.0
+ * @category services
+ */
+export class Config extends ServiceMap.Service<
+  Config,
+  Simplify<
+    & Partial<
+      Omit<
+        typeof Generated.CreateEmbeddingRequest.Encoded,
+        "input"
+      >
+    >
+    & {
+      readonly [x: string]: unknown
+    }
+  >
+>()("@effect/ai-openai/OpenAiEmbeddingModel/Config") {}
+
+/**
+ * @since 1.0.0
+ * @category constructors
+ */
+export const model = (
+  model: (string & {}) | Model,
+  options: {
+    readonly dimensions: number
+    readonly config?: Omit<typeof Config.Service, "model" | "dimensions">
+  }
+): AiModel.Model<"openai", EmbeddingModel.EmbeddingModel | EmbeddingModel.Dimensions, OpenAiClient> =>
+  AiModel.make(
+    "openai",
+    model,
+    Layer.merge(
+      layer({
+        model,
+        config: {
+          ...options.config,
+          dimensions: options.dimensions
+        }
+      }),
+      Layer.succeed(EmbeddingModel.Dimensions, options.dimensions)
+    )
+  )
+
+/**
+ * Creates an OpenAI embedding model service.
+ *
+ * @since 1.0.0
+ * @category constructors
+ */
+export const make = Effect.fnUntraced(function*({ model, config: providerConfig }: {
+  readonly model: (string & {}) | Model
+  readonly config?: Omit<typeof Config.Service, "model"> | undefined
+}): Effect.fn.Return<EmbeddingModel.Service, never, OpenAiClient> {
+  const client = yield* OpenAiClient
+
+  const makeConfig = Effect.gen(function*() {
+    const services = yield* Effect.services<never>()
+    return { model, ...providerConfig, ...services.mapUnsafe.get(Config.key) }
+  })
+
+  return yield* EmbeddingModel.make({
+    embedMany: Effect.fnUntraced(function*({ inputs }) {
+      const config = yield* makeConfig
+      const response = yield* client.createEmbedding({ ...config, input: inputs })
+      return yield* mapProviderResponse(inputs.length, response)
+    })
+  })
+})
+
+/**
+ * Creates a layer for the OpenAI embedding model.
+ *
+ * @since 1.0.0
+ * @category layers
+ */
+export const layer = (options: {
+  readonly model: (string & {}) | Model
+  readonly config?: Omit<typeof Config.Service, "model"> | undefined
+}): Layer.Layer<EmbeddingModel.EmbeddingModel, never, OpenAiClient> =>
+  Layer.effect(EmbeddingModel.EmbeddingModel, make(options))
+
+/**
+ * Provides config overrides for OpenAI embedding model operations.
+ *
+ * @since 1.0.0
+ * @category configuration
+ */
+export const withConfigOverride: {
+  /**
+   * Provides config overrides for OpenAI embedding model operations.
+   *
+   * @since 1.0.0
+   * @category configuration
+   */
+  (overrides: typeof Config.Service): <A, E, R>(self: Effect.Effect<A, E, R>) => Effect.Effect<A, E, Exclude<R, Config>>
+  /**
+   * Provides config overrides for OpenAI embedding model operations.
+   *
+   * @since 1.0.0
+   * @category configuration
+   */
+  <A, E, R>(self: Effect.Effect<A, E, R>, overrides: typeof Config.Service): Effect.Effect<A, E, Exclude<R, Config>>
+} = dual<
+  /**
+   * Provides config overrides for OpenAI embedding model operations.
+   *
+   * @since 1.0.0
+   * @category configuration
+   */
+  (overrides: typeof Config.Service) => <A, E, R>(self: Effect.Effect<A, E, R>) => Effect.Effect<A, E, Exclude<R, Config>>,
+  /**
+   * Provides config overrides for OpenAI embedding model operations.
+   *
+   * @since 1.0.0
+   * @category configuration
+   */
+  <A, E, R>(self: Effect.Effect<A, E, R>, overrides: typeof Config.Service) => Effect.Effect<A, E, Exclude<R, Config>>
+>(2, (self, overrides) =>
+  Effect.flatMap(
+    Effect.serviceOption(Config),
+    (config) =>
+      Effect.provideService(self, Config, {
+        ...(config._tag === "Some" ? config.value : {}),
+        ...overrides
+      })
+  ))
+
+const mapProviderResponse = (
+  inputLength: number,
+  response: typeof Generated.CreateEmbeddingResponse.Type
+): Effect.Effect<EmbeddingModel.ProviderResponse, AiError.AiError> => {
+  if (response.data.length !== inputLength) {
+    return Effect.fail(
+      invalidOutput("Provider returned " + response.data.length + " embeddings but expected " + inputLength)
+    )
+  }
+
+  const results = new Array<Array<number>>(inputLength)
+  const seen = new Set<number>()
+
+  for (const entry of response.data) {
+    if (!Number.isInteger(entry.index) || entry.index < 0 || entry.index >= inputLength) {
+      return Effect.fail(invalidOutput("Provider returned invalid embedding index: " + entry.index))
+    }
+    if (seen.has(entry.index)) {
+      return Effect.fail(invalidOutput("Provider returned duplicate embedding index: " + entry.index))
+    }
+
+    seen.add(entry.index)
+    results[entry.index] = [...entry.embedding]
+  }
+
+  if (seen.size !== inputLength) {
+    return Effect.fail(
+      invalidOutput("Provider returned embeddings for " + seen.size + " inputs but expected " + inputLength)
+    )
+  }
+
+  return Effect.succeed({
+    results,
+    usage: {
+      inputTokens: response.usage?.prompt_tokens
+    }
+  })
+}
+
+const invalidOutput = (description: string): AiError.AiError =>
+  AiError.make({
+    module: "OpenAiEmbeddingModel",
+    method: "embedMany",
+    reason: new AiError.InvalidOutputError({ description })
+  })