@effect/ai-openai 4.0.0-beta.5 → 4.0.0-beta.50

package/src/OpenAiClient.ts

@@ -8,20 +8,28 @@
  */
 import * as Array from "effect/Array"
 import type * as Config from "effect/Config"
+import * as Context from "effect/Context"
 import * as Effect from "effect/Effect"
 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 ServiceMap from "effect/ServiceMap"
+import * as Schema from "effect/Schema"
+import * as Scope from "effect/Scope"
+import * as Semaphore from "effect/Semaphore"
 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 +56,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 +66,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>
     ],
@@ -83,7 +91,7 @@ export interface Service {
  * @since 1.0.0
  * @category service
  */
-export class OpenAiClient extends ServiceMap.Service<OpenAiClient, Service>()(
+export class OpenAiClient extends Context.Service<OpenAiClient, Service>()(
   "@effect/ai-openai/OpenAiClient"
 ) {}
 
@@ -142,32 +150,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 +231,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.contextWith((services) => {
+        const socket = Context.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 +294,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 +304,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 +340,255 @@ 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 Context.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 socketScope = yield* Effect.scope
+  const makeRequest = Effect.orDie(client.client.httpClient.preprocess(HttpClientRequest.post("/responses")))
+  const makeWebSocket = yield* Socket.WebSocketConstructor
+
+  const decoder = new TextDecoder()
+
+  const queueRef: RcRef.RcRef<
+    {
+      readonly send: (message: typeof Generated.CreateResponse.Encoded) => Effect.Effect<void, AiError.AiError>
+      readonly incoming: Queue.Dequeue<ResponseStreamEvent, AiError.AiError>
+    }
+  > = yield* RcRef.make({
+    idleTimeToLive: 60_000,
+    acquire: Effect.gen(function*() {
+      const scope = yield* Effect.scope
+      const request = yield* makeRequest
+      const socket = yield* Socket.makeWebSocket(request.url.replace(/^http/, "ws")).pipe(
+        Effect.provideService(Socket.WebSocketConstructor, (url) =>
+          makeWebSocket(url, {
+            headers: request.headers
+          } as any))
+      )
+      const write = yield* socket.writer
+
+      yield* Scope.addFinalizerExit(scope, () => {
+        tracker.clearUnsafe()
+        return Effect.void
+      })
+
+      const incoming = yield* Queue.unbounded<ResponseStreamEvent, AiError.AiError>()
+      const send = (message: typeof Generated.CreateResponse.Encoded) =>
+        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"
+              })
+            })
+          )
+        )
+
+      yield* socket.runRaw((msg) => {
+        const text = typeof msg === "string" ? msg : decoder.decode(msg)
+        try {
+          const event = decodeEvent(text)
+          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(incoming, event)
+        } catch {}
+      }).pipe(
+        Effect.catchTag("SocketError", (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: error.message
+            })
+          }).asEffect()),
+        Effect.catchCause((cause) => Queue.failCause(incoming, cause)),
+        Effect.ensuring(Effect.forkIn(RcRef.invalidate(queueRef), socketScope, {
+          startImmediately: true
+        })),
+        Effect.forkScoped({ startImmediately: true })
+      )
+
+      return { send, incoming } as const
+    })
+  })
+
+  // Prime the websocket
+  yield* Effect.scoped(RcRef.get(queueRef))
+
+  // Websocket mode only allows one request at a time
+  const semaphore = Semaphore.makeUnsafe(1)
+  const request = yield* makeRequest
+
+  return OpenAiSocket.context({
+    createResponseStream(options) {
+      const stream = Stream.unwrap(Effect.gen(function*() {
+        const scope = yield* Effect.scope
+        yield* Effect.acquireRelease(
+          semaphore.take(1),
+          () => semaphore.release(1),
+          { interruptible: true }
+        )
+        const { send, incoming } = yield* RcRef.get(queueRef)
+        let done = false
+
+        yield* Scope.addFinalizerExit(
+          scope,
+          () => done ? Effect.void : RcRef.invalidate(queueRef)
+        )
+
+        yield* send(options).pipe(
+          Effect.forkScoped({ startImmediately: true })
+        )
+
+        return Stream.fromQueue(incoming).pipe(
+          Stream.takeUntil((e) => {
+            done = e.type === "response.completed" || e.type === "response.incomplete"
+            return done
+          })
+        )
+      }))
+
+      return Effect.succeed([
+        HttpClientResponse.fromWeb(request, new Response()),
+        stream
+      ])
+    }
+  }).pipe(
+    Context.add(ResponseIdTracker.ResponseIdTracker, tracker)
+  )
+})
+
+const ErrorEvent = Schema.Struct({
+  type: Schema.Literal("error"),
+  status: Schema.Number.pipe(
+    Schema.withDecodingDefault(Effect.succeed(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.provideContext(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.effectContext(makeSocket)

package/src/OpenAiConfig.ts

@@ -1,16 +1,16 @@
 /**
  * @since 1.0.0
  */
+import * as Context from "effect/Context"
 import * as Effect from "effect/Effect"
 import { dual } from "effect/Function"
-import * as ServiceMap from "effect/ServiceMap"
 import type { HttpClient } from "effect/unstable/http/HttpClient"
 
 /**
  * @since 1.0.0
  * @category services
  */
-export class OpenAiConfig extends ServiceMap.Service<
+export class OpenAiConfig extends Context.Service<
   OpenAiConfig,
   OpenAiConfig.Service
 >()("@effect/ai-openai/OpenAiConfig") {
@@ -18,7 +18,7 @@ export class OpenAiConfig extends ServiceMap.Service<
    * @since 1.0.0
    */
   static readonly getOrUndefined: Effect.Effect<typeof OpenAiConfig.Service | undefined> = Effect.map(
-    Effect.services<never>(),
+    Effect.context<never>(),
     (context) => context.mapUnsafe.get(OpenAiConfig.key)
   )
 }

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 Context from "effect/Context"
+import * as Effect from "effect/Effect"
+import { dual } from "effect/Function"
+import * as Layer from "effect/Layer"
+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 Context.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.context<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 })
+  })