EFF-726 Add ModelDimensions service to Embeddings module (#1771)

Author: tim-smart

.changeset/eff-726-model-dimensions.md

@@ -0,0 +1,7 @@
+---
+"effect": patch
+"@effect/ai-openai": patch
+"@effect/ai-openai-compat": patch
+---
+
+Add `EmbeddingModel.ModelDimensions` and require dimensions in embedding provider `model` constructors.

packages/ai/openai-compat/src/OpenAiEmbeddingModel.ts

@@ -49,9 +49,25 @@ export class Config extends ServiceMap.Service<
  */
 export const model = (
   model: string,
-  config?: Omit<typeof Config.Service, "model">
-): AiModel.Model<"openai", EmbeddingModel.EmbeddingModel, OpenAiClient> =>
-  AiModel.make("openai", model, layer({ model, config }))
+  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.

packages/ai/openai-compat/test/OpenAiEmbeddingModel.test.ts

@@ -5,6 +5,15 @@ import { EmbeddingModel } from "effect/unstable/ai"
 import { HttpClient, type HttpClientError, type HttpClientRequest, HttpClientResponse } from "effect/unstable/http"
 
 describe("OpenAiEmbeddingModel", () => {
+  it.effect("model provides dimensions service", () =>
+    Effect.gen(function*() {
+      const dimensions = yield* EmbeddingModel.Dimensions
+      assert.strictEqual(dimensions, 1536)
+    }).pipe(
+      Effect.provide(OpenAiEmbeddingModel.model("text-embedding-3-small", { dimensions: 1536 })),
+      Effect.provideService(OpenAiClient.OpenAiClient, noopOpenAiClient)
+    ))
+
   it.effect("reorders embeddings by provider index", () =>
     Effect.gen(function*() {
       let capturedRequest: HttpClientRequest.HttpClientRequest | undefined
@@ -198,3 +207,10 @@ const getRequestBody = (request: HttpClientRequest.HttpClientRequest) =>
     }
     return yield* Effect.die(new Error("Expected Uint8Array body"))
   })
+
+const noopOpenAiClient: OpenAiClient.Service = {
+  client: undefined as unknown as OpenAiClient.Service["client"],
+  createResponse: () => Effect.die(new Error("noop")),
+  createResponseStream: () => Effect.die(new Error("noop")),
+  createEmbedding: () => Effect.die(new Error("noop"))
+}

packages/ai/openai/src/OpenAiEmbeddingModel.ts

@@ -49,9 +49,25 @@ export class Config extends ServiceMap.Service<
  */
 export const model = (
   model: (string & {}) | Model,
-  config?: Omit<typeof Config.Service, "model">
-): AiModel.Model<"openai", EmbeddingModel.EmbeddingModel, OpenAiClient> =>
-  AiModel.make("openai", model, layer({ model, config }))
+  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.

packages/ai/openai/test/OpenAiEmbeddingModel.test.ts

@@ -5,6 +5,15 @@ import { EmbeddingModel } from "effect/unstable/ai"
 import { HttpClient, type HttpClientError, type HttpClientRequest, HttpClientResponse } from "effect/unstable/http"
 
 describe("OpenAiEmbeddingModel", () => {
+  it.effect("model provides dimensions service", () =>
+    Effect.gen(function*() {
+      const dimensions = yield* EmbeddingModel.Dimensions
+      assert.strictEqual(dimensions, 1536)
+    }).pipe(
+      Effect.provide(OpenAiEmbeddingModel.model("text-embedding-3-small", { dimensions: 1536 })),
+      Effect.provideService(OpenAiClient.OpenAiClient, noopOpenAiClient)
+    ))
+
   it.effect("reorders embeddings by provider index", () =>
     Effect.gen(function*() {
       let capturedRequest: HttpClientRequest.HttpClientRequest | undefined
@@ -240,3 +249,10 @@ const getRequestBody = (request: HttpClientRequest.HttpClientRequest) =>
     }
     return yield* Effect.die(new Error("Expected Uint8Array body"))
   })
+
+const noopOpenAiClient: OpenAiClient.Service = {
+  client: undefined as unknown as OpenAiClient.Service["client"],
+  createResponse: () => Effect.die(new Error("noop")),
+  createResponseStream: () => Effect.die(new Error("noop")),
+  createEmbedding: () => Effect.die(new Error("noop"))
+}

packages/effect/src/unstable/ai/EmbeddingModel.ts

@@ -32,6 +32,16 @@ export class EmbeddingModel extends ServiceMap.Service<EmbeddingModel, Service>(
   "effect/unstable/ai/EmbeddingModel"
 ) {}
 
+/**
+ * Service tag that provides the current embedding dimensions.
+ *
+ * @since 4.0.0
+ * @category services
+ */
+export class Dimensions extends ServiceMap.Service<Dimensions, number>()(
+  "effect/unstable/ai/EmbeddingModel/Dimensions"
+) {}
+
 /**
  * Token usage metadata for embedding operations.
  *