1import { OpenAiClient, OpenAiEmbeddingModel } from "@effect/ai-openai"
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], object: "embedding" },
29 { index: 0, embedding: [10, 11], object: "embedding" }
30 ],
31 model: "text-embedding-3-small",
32 object: "list",
33 usage: {
34 prompt_tokens: 7,
35 total_tokens: 7
36 }
37 }))
38 })
39 ))
40 )
41
42 const response = yield* Effect.gen(function*() {
43 const model = yield* EmbeddingModel.EmbeddingModel
44 return yield* model.embedMany(["first", "second"])
45 }).pipe(
46 Effect.provide(OpenAiEmbeddingModel.layer({ model: "text-embedding-3-small" })),
47 Effect.provide(clientLayer)
48 )
49
50 assert.deepStrictEqual(response.embeddings.map((embedding) => embedding.vector), [[10, 11], [20, 21]])
51 assert.strictEqual(response.usage.inputTokens, 7)
52
53 assert.isDefined(capturedRequest)
54 if (capturedRequest === undefined) {
55 return
56 }
57
58 const requestBody = yield* getRequestBody(capturedRequest)
59 assert.strictEqual(requestBody.model, "text-embedding-3-small")
60 assert.deepStrictEqual(requestBody.input, ["first", "second"])
61 }))
62
63 it.effect("merges config and applies withConfigOverride precedence", () =>
64 Effect.gen(function*() {
65 let capturedRequest: HttpClientRequest.HttpClientRequest | undefined
66
67 const clientLayer = OpenAiClient.layer({ apiKey: Redacted.make("sk-test-key") }).pipe(
68 Layer.provide(Layer.succeed(
69 HttpClient.HttpClient,
70 makeHttpClient((request) => {
71 capturedRequest = request
72 return Effect.succeed(jsonResponse(request, {
73 data: [{ index: 0, embedding: [1, 2, 3], object: "embedding" }],
74 model: "override-model",
75 object: "list",
76 usage: {
77 prompt_tokens: 3,
78 total_tokens: 3
79 }
80 }))
81 })
82 ))
83 )
84
85 yield* Effect.gen(function*() {
86 const model = yield* EmbeddingModel.EmbeddingModel
87 yield* model.embed("hello")
88 }).pipe(
89 OpenAiEmbeddingModel.withConfigOverride({
90 model: "override-model",
91 dimensions: 1024,
92 user: "request-user"
93 }),
94 Effect.provide(OpenAiEmbeddingModel.layer({
95 model: "base-model",
96 config: {
97 dimensions: 256,
98 user: "provider-user"
99 }
100 })),
101 Effect.provide(clientLayer)
102 )
103
104 assert.isDefined(capturedRequest)
105 if (capturedRequest === undefined) {
106 return
107 }
108
109 const requestBody = yield* getRequestBody(capturedRequest)
110 assert.strictEqual(requestBody.model, "override-model")
111 assert.strictEqual(requestBody.dimensions, 1024)
112 assert.strictEqual(requestBody.user, "request-user")
113 assert.deepStrictEqual(requestBody.input, ["hello"])
114 }))
115
116 it.effect("fails with InvalidOutputError when provider returns duplicate indices", () =>
117 Effect.gen(function*() {
118 const clientLayer = OpenAiClient.layer({ apiKey: Redacted.make("sk-test-key") }).pipe(
119 Layer.provide(Layer.succeed(
120 HttpClient.HttpClient,
121 makeHttpClient((request) =>
122 Effect.succeed(jsonResponse(request, {
123 data: [
124 { index: 0, embedding: [1], object: "embedding" },
125 { index: 0, embedding: [2], object: "embedding" }
126 ],
127 model: "text-embedding-3-small",
128 object: "list",
129 usage: {
130 prompt_tokens: 2,
131 total_tokens: 2
132 }
133 }))
134 )
135 ))
136 )
137
138 const error = yield* Effect.gen(function*() {
139 const model = yield* EmbeddingModel.EmbeddingModel
140 return yield* model.embedMany(["a", "b"]).pipe(Effect.flip)
141 }).pipe(
142 Effect.provide(OpenAiEmbeddingModel.layer({ model: "text-embedding-3-small" })),
143 Effect.provide(clientLayer)
144 )
145
146 assert.strictEqual(error._tag, "AiError")
147 assert.strictEqual(error.reason._tag, "InvalidOutputError")
148 }))
149
150 it.effect("fails with InvalidOutputError when provider returns out-of-range indices", () =>
151 Effect.gen(function*() {
152 const clientLayer = OpenAiClient.layer({ apiKey: Redacted.make("sk-test-key") }).pipe(
153 Layer.provide(Layer.succeed(
154 HttpClient.HttpClient,
155 makeHttpClient((request) =>
156 Effect.succeed(jsonResponse(request, {
157 data: [
158 { index: 0, embedding: [1], object: "embedding" },
159 { index: 2, embedding: [2], object: "embedding" }
160 ],
161 model: "text-embedding-3-small",
162 object: "list",
163 usage: {
164 prompt_tokens: 2,
165 total_tokens: 2
166 }
167 }))
168 )
169 ))
170 )
171
172 const error = yield* Effect.gen(function*() {
173 const model = yield* EmbeddingModel.EmbeddingModel
174 return yield* model.embedMany(["a", "b"]).pipe(Effect.flip)
175 }).pipe(
176 Effect.provide(OpenAiEmbeddingModel.layer({ model: "text-embedding-3-small" })),
177 Effect.provide(clientLayer)
178 )
179
180 assert.strictEqual(error._tag, "AiError")
181 assert.strictEqual(error.reason._tag, "InvalidOutputError")
182 }))
183
184 it.effect("fails with InvalidOutputError when provider returns wrong embedding count", () =>
185 Effect.gen(function*() {
186 const clientLayer = OpenAiClient.layer({ apiKey: Redacted.make("sk-test-key") }).pipe(
187 Layer.provide(Layer.succeed(
188 HttpClient.HttpClient,
189 makeHttpClient((request) =>
190 Effect.succeed(jsonResponse(request, {
191 data: [{ index: 0, embedding: [1], object: "embedding" }],
192 model: "text-embedding-3-small",
193 object: "list",
194 usage: {
195 prompt_tokens: 1,
196 total_tokens: 1
197 }
198 }))
199 )
200 ))
201 )
202
203 const error = yield* Effect.gen(function*() {
204 const model = yield* EmbeddingModel.EmbeddingModel
205 return yield* model.embedMany(["a", "b"]).pipe(Effect.flip)
206 }).pipe(
207 Effect.provide(OpenAiEmbeddingModel.layer({ model: "text-embedding-3-small" })),
208 Effect.provide(clientLayer)
209 )
210
211 assert.strictEqual(error._tag, "AiError")
212 assert.strictEqual(error.reason._tag, "InvalidOutputError")
213 }))
214
215 it.effect("fails with InvalidOutputError when provider returns base64 embeddings", () =>
216 Effect.gen(function*() {
217 const clientLayer = OpenAiClient.layer({ apiKey: Redacted.make("sk-test-key") }).pipe(
218 Layer.provide(Layer.succeed(
219 HttpClient.HttpClient,
220 makeHttpClient((request) =>
221 Effect.succeed(jsonResponse(request, {
222 data: [{ index: 0, embedding: "AQID", object: "embedding" }],
223 model: "text-embedding-3-small",
224 object: "list",
225 usage: {
226 prompt_tokens: 1,
227 total_tokens: 1
228 }
229 }))
230 )
231 ))
232 )
233
234 const error = yield* Effect.gen(function*() {
235 const model = yield* EmbeddingModel.EmbeddingModel
236 return yield* model.embed("a").pipe(Effect.flip)
237 }).pipe(
238 Effect.provide(OpenAiEmbeddingModel.layer({ model: "text-embedding-3-small" })),
239 Effect.provide(clientLayer)
240 )
241
242 assert.strictEqual(error._tag, "AiError")
243 assert.strictEqual(error.reason._tag, "InvalidOutputError")
244 }))
245})
246
247const makeHttpClient = (
248 handler: (
249 request: HttpClientRequest.HttpClientRequest
250 ) => Effect.Effect<HttpClientResponse.HttpClientResponse, HttpClientError.HttpClientError>
251) =>
252 HttpClient.makeWith(
253 Effect.fnUntraced(function*(requestEffect) {
254 const request = yield* requestEffect
255 return yield* handler(request)
256 }),
257 Effect.succeed as HttpClient.HttpClient.Preprocess<HttpClientError.HttpClientError, never>
258 )
259
260const jsonResponse = (
261 request: HttpClientRequest.HttpClientRequest,
262 body: unknown
263): HttpClientResponse.HttpClientResponse =>
264 HttpClientResponse.fromWeb(
265 request,
266 new Response(JSON.stringify(body), {
267 status: 200,
268 headers: {
269 "content-type": "application/json"
270 }
271 })
272 )
273
274const getRequestBody = (request: HttpClientRequest.HttpClientRequest) =>
275 Effect.gen(function*() {
276 const body = request.body
277 if (body._tag === "Uint8Array") {
278 const text = new TextDecoder().decode(body.body)
279 return JSON.parse(text)
280 }
281 return yield* Effect.die(new Error("Expected Uint8Array body"))
282 })
283
284const noopOpenAiClient: OpenAiClient.Service = {
285 client: undefined as unknown as OpenAiClient.Service["client"],
286 createResponse: () => Effect.die(new Error("noop")),
287 createResponseStream: () => Effect.die(new Error("noop")),
288 createEmbedding: () => Effect.die(new Error("noop"))
289}