1import { OpenAiClient, OpenAiEmbeddingModel } from "@effect/ai-openai-compat"
2import { assert, describe, it } from "@effect/vitest"
3import { Effect, Layer, Redacted } from "effect"
4import { EmbeddingModel } from "effect/ai"
5import { HttpClient, type HttpClientError, type HttpClientRequest, HttpClientResponse } from "effect/http"
6
7describe("OpenAiEmbeddingModel", () => {
8 it.effect("model provides dimensions service", () =>
9 Effect.gen(function*() {
10 const dimensions = yield* EmbeddingModel.Dimensions
11 assert.strictEqual(dimensions, 1536)
12 }).pipe(
13 Effect.provide(OpenAiEmbeddingModel.model("text-embedding-3-small", { dimensions: 1536 })),
14 Effect.provideService(OpenAiClient.OpenAiClient, noopOpenAiClient)
15 ))
16
17 it.effect("reorders embeddings by provider index", () =>
18 Effect.gen(function*() {
19 let capturedRequest: HttpClientRequest.HttpClientRequest | undefined
20
21 const clientLayer = OpenAiClient.layer({ apiKey: Redacted.make("sk-test-key") }).pipe(
22 Layer.provide(Layer.succeed(
23 HttpClient.HttpClient,
24 makeHttpClient((request) => {
25 capturedRequest = request
26 return Effect.succeed(jsonResponse(request, {
27 data: [
28 { index: 1, embedding: [20, 21] },
29 { index: 0, embedding: [10, 11] }
30 ],
31 model: "text-embedding-3-small",
32 usage: {
33 prompt_tokens: 7,
34 total_tokens: 7
35 }
36 }))
37 })
38 ))
39 )
40
41 const response = yield* Effect.gen(function*() {
42 const model = yield* EmbeddingModel.EmbeddingModel
43 return yield* model.embedMany(["first", "second"])
44 }).pipe(
45 Effect.provide(OpenAiEmbeddingModel.layer({ model: "text-embedding-3-small" })),
46 Effect.provide(clientLayer)
47 )
48
49 assert.deepStrictEqual(response.embeddings.map((embedding) => embedding.vector), [[10, 11], [20, 21]])
50 assert.strictEqual(response.usage.inputTokens, 7)
51
52 assert.isDefined(capturedRequest)
53 if (capturedRequest === undefined) {
54 return
55 }
56
57 const requestBody = yield* getRequestBody(capturedRequest)
58 assert.strictEqual(requestBody.model, "text-embedding-3-small")
59 assert.deepStrictEqual(requestBody.input, ["first", "second"])
60 }))
61
62 it.effect("merges config and applies withConfigOverride precedence", () =>
63 Effect.gen(function*() {
64 let capturedRequest: HttpClientRequest.HttpClientRequest | undefined
65
66 const clientLayer = OpenAiClient.layer({ apiKey: Redacted.make("sk-test-key") }).pipe(
67 Layer.provide(Layer.succeed(
68 HttpClient.HttpClient,
69 makeHttpClient((request) => {
70 capturedRequest = request
71 return Effect.succeed(jsonResponse(request, {
72 data: [{ index: 0, embedding: [1, 2, 3] }],
73 model: "override-model"
74 }))
75 })
76 ))
77 )
78
79 yield* Effect.gen(function*() {
80 const model = yield* EmbeddingModel.EmbeddingModel
81 yield* model.embed("hello")
82 }).pipe(
83 OpenAiEmbeddingModel.withConfigOverride({
84 model: "override-model",
85 dimensions: 1024,
86 user: "request-user"
87 }),
88 Effect.provide(OpenAiEmbeddingModel.layer({
89 model: "base-model",
90 config: {
91 dimensions: 256,
92 user: "provider-user"
93 }
94 })),
95 Effect.provide(clientLayer)
96 )
97
98 assert.isDefined(capturedRequest)
99 if (capturedRequest === undefined) {
100 return
101 }
102
103 const requestBody = yield* getRequestBody(capturedRequest)
104 assert.strictEqual(requestBody.model, "override-model")
105 assert.strictEqual(requestBody.dimensions, 1024)
106 assert.strictEqual(requestBody.user, "request-user")
107 assert.deepStrictEqual(requestBody.input, ["hello"])
108 }))
109
110 it.effect("fails with InvalidOutputError when provider returns duplicate indices", () =>
111 Effect.gen(function*() {
112 const clientLayer = OpenAiClient.layer({ apiKey: Redacted.make("sk-test-key") }).pipe(
113 Layer.provide(Layer.succeed(
114 HttpClient.HttpClient,
115 makeHttpClient((request) =>
116 Effect.succeed(jsonResponse(request, {
117 data: [
118 { index: 0, embedding: [1] },
119 { index: 0, embedding: [2] }
120 ],
121 model: "text-embedding-3-small",
122 usage: {
123 prompt_tokens: 2,
124 total_tokens: 2
125 }
126 }))
127 )
128 ))
129 )
130
131 const error = yield* Effect.gen(function*() {
132 const model = yield* EmbeddingModel.EmbeddingModel
133 return yield* model.embedMany(["a", "b"]).pipe(Effect.flip)
134 }).pipe(
135 Effect.provide(OpenAiEmbeddingModel.layer({ model: "text-embedding-3-small" })),
136 Effect.provide(clientLayer)
137 )
138
139 assert.strictEqual(error._tag, "AiError")
140 assert.strictEqual(error.reason._tag, "InvalidOutputError")
141 }))
142
143 it.effect("fails with InvalidOutputError when provider returns non-vector embedding payload", () =>
144 Effect.gen(function*() {
145 const clientLayer = OpenAiClient.layer({ apiKey: Redacted.make("sk-test-key") }).pipe(
146 Layer.provide(Layer.succeed(
147 HttpClient.HttpClient,
148 makeHttpClient((request) =>
149 Effect.succeed(jsonResponse(request, {
150 data: [{ index: 0, embedding: "AQID" }],
151 model: "text-embedding-3-small",
152 usage: {
153 prompt_tokens: 1,
154 total_tokens: 1
155 }
156 }))
157 )
158 ))
159 )
160
161 const error = yield* Effect.gen(function*() {
162 const model = yield* EmbeddingModel.EmbeddingModel
163 return yield* model.embed("a").pipe(Effect.flip)
164 }).pipe(
165 Effect.provide(OpenAiEmbeddingModel.layer({ model: "text-embedding-3-small" })),
166 Effect.provide(clientLayer)
167 )
168
169 assert.strictEqual(error._tag, "AiError")
170 assert.strictEqual(error.reason._tag, "InvalidOutputError")
171 }))
172})
173
174const makeHttpClient = (
175 handler: (
176 request: HttpClientRequest.HttpClientRequest
177 ) => Effect.Effect<HttpClientResponse.HttpClientResponse, HttpClientError.HttpClientError>
178) =>
179 HttpClient.makeWith(
180 Effect.fnUntraced(function*(requestEffect) {
181 const request = yield* requestEffect
182 return yield* handler(request)
183 }),
184 Effect.succeed as HttpClient.HttpClient.Preprocess<HttpClientError.HttpClientError, never>
185 )
186
187const jsonResponse = (
188 request: HttpClientRequest.HttpClientRequest,
189 body: unknown
190): HttpClientResponse.HttpClientResponse =>
191 HttpClientResponse.fromWeb(
192 request,
193 new Response(JSON.stringify(body), {
194 status: 200,
195 headers: {
196 "content-type": "application/json"
197 }
198 })
199 )
200
201const getRequestBody = (request: HttpClientRequest.HttpClientRequest) =>
202 Effect.gen(function*() {
203 const body = request.body
204 if (body._tag === "Uint8Array") {
205 const text = new TextDecoder().decode(body.body)
206 return JSON.parse(text)
207 }
208 return yield* Effect.die(new Error("Expected Uint8Array body"))
209 })
210
211const noopOpenAiClient: OpenAiClient.Service = {
212 client: undefined as unknown as OpenAiClient.Service["client"],
213 createResponse: () => Effect.die(new Error("noop")),
214 createResponseStream: () => Effect.die(new Error("noop")),
215 createEmbedding: () => Effect.die(new Error("noop"))
216}