1import { assert, describe, it } from "@effect/vitest"
2import { Effect, Layer, Schema } from "effect"
3import { EmbeddingModel } from "effect/ai"
4import * as AiError from "effect/ai/AiError"
5
6const makeLayer = (
7 embedMany: (
8 options: EmbeddingModel.ProviderOptions
9 ) => Effect.Effect<EmbeddingModel.ProviderResponse, AiError.AiError>
10) =>
11 Layer.effect(
12 EmbeddingModel.EmbeddingModel,
13 EmbeddingModel.make({
14 embedMany
15 })
16 )
17
18describe("EmbeddingModel", () => {
19 it.effect("round trips usage with undefined input tokens through JSON", () =>
20 Effect.gen(function*() {
21 const usage = new EmbeddingModel.EmbeddingUsage({ inputTokens: undefined })
22
23 const encoded = yield* Schema.encodeEffect(EmbeddingModel.EmbeddingUsage)(usage)
24 const json = JSON.parse(JSON.stringify(encoded))
25 const decoded = yield* Schema.decodeUnknownEffect(EmbeddingModel.EmbeddingUsage)(json)
26
27 assert.deepStrictEqual(json, {}, "encoded JSON")
28 assert.deepStrictEqual(decoded, new EmbeddingModel.EmbeddingUsage({}), "decoded usage")
29 }))
30
31 it.effect("embed returns a vector", () => {
32 const calls: Array<ReadonlyArray<string>> = []
33
34 return Effect.gen(function*() {
35 const model = yield* EmbeddingModel.EmbeddingModel
36 const response = yield* model.embed("hello")
37
38 assert.deepStrictEqual(response.vector, [1, 2, 3])
39 assert.deepStrictEqual(calls, [["hello"]])
40 }).pipe(
41 Effect.provide(
42 makeLayer(({ inputs }) =>
43 Effect.sync(() => {
44 calls.push(inputs)
45 return {
46 results: [[1, 2, 3]],
47 usage: { inputTokens: 5 }
48 }
49 })
50 )
51 )
52 )
53 })
54
55 it.effect("embedMany returns ordered vectors with usage", () => {
56 const calls: Array<ReadonlyArray<string>> = []
57
58 return Effect.gen(function*() {
59 const model = yield* EmbeddingModel.EmbeddingModel
60 const response = yield* model.embedMany(["hello", "world"])
61
62 assert.deepStrictEqual(
63 response.embeddings.map((embedding) => embedding.vector),
64 [[1, 2], [3, 4]]
65 )
66 assert.strictEqual(response.usage.inputTokens, 9)
67 assert.deepStrictEqual(calls, [["hello", "world"]])
68 }).pipe(
69 Effect.provide(
70 makeLayer(({ inputs }) => {
71 calls.push(inputs)
72 return Effect.succeed({
73 results: [[1, 2], [3, 4]],
74 usage: { inputTokens: 9 }
75 })
76 })
77 )
78 )
79 })
80
81 it.effect("concurrent embed calls are batched into one provider embedMany call", () => {
82 let calls = 0
83 const batches: Array<ReadonlyArray<string>> = []
84
85 return Effect.gen(function*() {
86 const model = yield* EmbeddingModel.EmbeddingModel
87 const responses = yield* Effect.all([
88 model.embed("a"),
89 model.embed("b"),
90 model.embed("c")
91 ], { concurrency: "unbounded" })
92
93 assert.strictEqual(calls, 1)
94 assert.deepStrictEqual(batches, [["a", "b", "c"]])
95 assert.deepStrictEqual(
96 responses.map((response) => response.vector),
97 [[97], [98], [99]]
98 )
99 }).pipe(
100 Effect.provide(
101 makeLayer(({ inputs }) => {
102 calls++
103 batches.push(inputs)
104 return Effect.succeed({
105 results: inputs.map((input) => [input.charCodeAt(0)]),
106 usage: { inputTokens: undefined }
107 })
108 })
109 )
110 )
111 })
112
113 it.effect("provider AiError propagates through embed", () => {
114 const error = AiError.make({
115 module: "EmbeddingModelTest",
116 method: "embedMany",
117 reason: new AiError.UnknownError({ description: "boom" })
118 })
119
120 return Effect.gen(function*() {
121 const model = yield* EmbeddingModel.EmbeddingModel
122 const result = yield* model.embed("hello").pipe(
123 Effect.match({ onFailure: (error) => error, onSuccess: () => undefined })
124 )
125
126 assert.strictEqual(result, error)
127 }).pipe(
128 Effect.provide(
129 makeLayer(() => Effect.fail(error))
130 )
131 )
132 })
133
134 it.effect("embed fails with InvalidOutputError when provider misses results", () =>
135 Effect.gen(function*() {
136 const model = yield* EmbeddingModel.EmbeddingModel
137 const error = yield* model.embed("hello").pipe(
138 Effect.match({ onFailure: (error) => error, onSuccess: () => undefined })
139 )
140
141 assert.notStrictEqual(error, undefined)
142 if (error !== undefined) {
143 assert.strictEqual(error.reason._tag, "InvalidOutputError")
144 }
145 }).pipe(
146 Effect.provide(
147 makeLayer(() =>
148 Effect.succeed({
149 results: [],
150 usage: { inputTokens: 0 }
151 })
152 )
153 )
154 ))
155
156 it.effect("embedMany fails with InvalidOutputError when provider misses results", () =>
157 Effect.gen(function*() {
158 const model = yield* EmbeddingModel.EmbeddingModel
159 const error = yield* model.embedMany(["hello", "world"]).pipe(
160 Effect.match({ onFailure: (error) => error, onSuccess: () => undefined })
161 )
162
163 assert.notStrictEqual(error, undefined)
164 if (error !== undefined) {
165 assert.strictEqual(error.reason._tag, "InvalidOutputError")
166 }
167 }).pipe(
168 Effect.provide(
169 makeLayer(() =>
170 Effect.succeed({
171 results: [[1, 2, 3]],
172 usage: { inputTokens: 3 }
173 })
174 )
175 )
176 ))
177
178 it.effect("embedMany([]) bypasses provider", () => {
179 let calls = 0
180
181 return Effect.gen(function*() {
182 const model = yield* EmbeddingModel.EmbeddingModel
183 const response = yield* model.embedMany([])
184
185 assert.strictEqual(calls, 0)
186 assert.deepStrictEqual(response.embeddings, [])
187 assert.strictEqual(response.usage.inputTokens, undefined)
188 }).pipe(
189 Effect.provide(
190 makeLayer(() => {
191 calls++
192 return Effect.die("provider should not be called")
193 })
194 )
195 )
196 })
197})