1/**
2 * Defines the provider-neutral service for text embeddings.
3 *
4 * An `EmbeddingModel` turns text into numeric vectors. It supports single-input
5 * embedding and ordered batch embedding, and represents provider failures as
6 * `AiError` values. This module also includes the embedding dimensions service,
7 * request and response models, usage metadata, provider contracts, and a
8 * constructor that adapts a provider batch implementation into the service.
9 * Single `embed` calls can be batched together internally.
10 *
11 * @stability unstable
12 * @since 4.0.0
13 */
14import * as Context from "../Context.ts"
15import * as Effect from "../Effect.ts"
16import * as Exit from "../Exit.ts"
17import * as Request from "../Request.ts"
18import * as RequestResolver from "../RequestResolver.ts"
19import * as Schema from "../Schema.ts"
20import * as AiError from "./AiError.ts"
21
22/**
23 * Service key for embedding text into vectors.
24 *
25 * @see {@link make} for constructing an embedding model service from a provider
26 * @see {@link Dimensions} for the current embedding vector size service
27 *
28 * @stability unstable
29 * @category services
30 * @since 4.0.0
31 */
32export const EmbeddingModel: Context.Service<EmbeddingModel, EmbeddingModel> = Context.Service(
33 "effect/ai/EmbeddingModel"
34)
35
36/**
37 * Brand type for `EmbeddingModel`.
38 *
39 * @stability unstable
40 * @category type IDs
41 * @since 4.0.0
42 */
43export type TypeId = "~effect/ai/EmbeddingModel"
44
45/**
46 * Brand for `EmbeddingModel` implementations.
47 *
48 * @stability unstable
49 * @category type IDs
50 * @since 4.0.0
51 */
52export const TypeId: TypeId = "~effect/ai/EmbeddingModel"
53
54/**
55 * Service tag that provides the current embedding dimensions.
56 *
57 * **When to use**
58 *
59 * Use to retrieve or provide the configured embedding vector size through
60 * context.
61 *
62 * @see {@link EmbeddingModel} for the embedding service that uses these dimensions
63 *
64 * @stability unstable
65 * @category services
66 * @since 4.0.0
67 */
68export class Dimensions extends Context.Service<Dimensions, number>()(
69 "effect/ai/EmbeddingModel/Dimensions"
70) {}
71
72/**
73 * Represents token usage metadata for embedding operations.
74 *
75 * **Details**
76 *
77 * Contains optional provider-reported `inputTokens`. The value may be
78 * `undefined` when the provider does not report usage or when `embedMany([])`
79 * bypasses the provider.
80 *
81 * @stability unstable
82 * @category models
83 * @since 4.0.0
84 */
85export class EmbeddingUsage extends Schema.Class<EmbeddingUsage>(
86 "effect/ai/EmbeddingModel/EmbeddingUsage"
87)({
88 inputTokens: Schema.optional(Schema.Finite)
89}) {}
90
91/**
92 * Response for a single embedding request.
93 *
94 * @stability unstable
95 * @category models
96 * @since 4.0.0
97 */
98export class EmbedResponse extends Schema.Class<EmbedResponse>(
99 "effect/ai/EmbeddingModel/EmbedResponse"
100)({
101 vector: Schema.Array(Schema.Finite)
102}) {}
103
104/**
105 * Response for batch embedding requests containing per-input embeddings and usage
106 * metadata.
107 *
108 * **Details**
109 *
110 * `embeddings` preserves batch order, and `usage` carries token metadata for
111 * the operation.
112 *
113 * @see {@link EmbedResponse} for individual embedding responses
114 * @see {@link EmbeddingUsage} for token usage metadata
115 *
116 * @stability unstable
117 * @category models
118 * @since 4.0.0
119 */
120export class EmbedManyResponse extends Schema.Class<EmbedManyResponse>(
121 "effect/ai/EmbeddingModel/EmbedManyResponse"
122)({
123 embeddings: Schema.Array(EmbedResponse),
124 usage: EmbeddingUsage
125}) {}
126
127/**
128 * Provider input options for embedding requests.
129 *
130 * @stability unstable
131 * @category options
132 * @since 4.0.0
133 */
134export interface ProviderOptions {
135 readonly inputs: ReadonlyArray<string>
136}
137
138/**
139 * Provider response for batch embedding requests.
140 *
141 * @stability unstable
142 * @category models
143 * @since 4.0.0
144 */
145export interface ProviderResponse {
146 readonly results: Array<Array<number>>
147 readonly usage: {
148 readonly inputTokens: number | undefined
149 }
150}
151
152/**
153 * Represents a tagged request used by request resolvers for embedding operations.
154 *
155 * **When to use**
156 *
157 * Use when you need a typed request for one embedding input while building or
158 * calling a low-level embedding request resolver.
159 *
160 * @see {@link EmbeddingModel} for the resolver-bearing service contract
161 * @see {@link make} for constructing the request resolver from a provider implementation
162 * @see {@link EmbedResponse} for the response produced by this request
163 *
164 * @stability unstable
165 * @category constructors
166 * @since 4.0.0
167 */
168export class EmbeddingRequest extends Request.TaggedClass("EmbeddingRequest")<
169 { readonly input: string },
170 EmbedResponse,
171 AiError.AiError
172> {}
173
174/**
175 * Single-input and batch embedding operations.
176 *
177 * @stability unstable
178 * @category models
179 * @since 4.0.0
180 */
181export interface EmbeddingModel {
182 readonly [TypeId]: TypeId
183 readonly resolver: RequestResolver.RequestResolver<EmbeddingRequest>
184 readonly embed: (input: string) => Effect.Effect<EmbedResponse, AiError.AiError>
185 readonly embedMany: (input: ReadonlyArray<string>) => Effect.Effect<EmbedManyResponse, AiError.AiError>
186}
187
188const invalidProviderResponse = (description: string): AiError.AiError =>
189 AiError.make({
190 module: "EmbeddingModel",
191 method: "embedMany",
192 reason: new AiError.InvalidOutputError({ description })
193 })
194
195/**
196 * Creates an EmbeddingModel service from a provider embedMany implementation.
197 *
198 * **When to use**
199 *
200 * Use to adapt a provider's batch embedding implementation into an
201 * `EmbeddingModel` that offers single-input and batch embedding
202 * operations.
203 *
204 * **Details**
205 *
206 * The returned service builds single-input `embed` calls through a request
207 * resolver, so concurrent `embed` requests can be batched into one provider
208 * `embedMany` call. Direct `embedMany` calls pass the input array to the
209 * provider, while `embedMany([])` returns an empty response without calling the
210 * provider.
211 *
212 * **Gotchas**
213 *
214 * Provider responses are interpreted positionally and must contain exactly one
215 * result for each requested input. If the provider returns a different number
216 * of results, `embed` and `embedMany` fail with `AiError.InvalidOutputError`.
217 *
218 * @see {@link EmbeddingModel} for the service shape returned by this constructor
219 * @see {@link ProviderOptions} for the input passed to the provider implementation
220 * @see {@link ProviderResponse} for the provider response contract consumed by this constructor
221 *
222 * @stability unstable
223 * @category constructors
224 * @since 4.0.0
225 */
226export const make: (params: {
227 readonly embedMany: (options: ProviderOptions) => Effect.Effect<ProviderResponse, AiError.AiError>
228}) => Effect.Effect<EmbeddingModel> = Effect.fnUntraced(function*(params) {
229 const resolver = RequestResolver.make<EmbeddingRequest>((entries) =>
230 Effect.flatMap(
231 params.embedMany({
232 inputs: entries.map((entry) => entry.request.input)
233 }),
234 (response) =>
235 Effect.map(mapProviderResults(entries.length, response.results), (embeddings) => {
236 for (let i = 0; i < entries.length; i++) {
237 entries[i].completeUnsafe(Exit.succeed(embeddings[i]))
238 }
239 })
240 )
241 ).pipe(
242 RequestResolver.withSpan("EmbeddingModel.resolver")
243 )
244
245 return EmbeddingModel.of({
246 [TypeId]: TypeId,
247 resolver,
248 embed: (input) =>
249 Effect.request(new EmbeddingRequest({ input }), resolver).pipe(
250 Effect.withSpan("EmbeddingModel.embed")
251 ),
252 embedMany: (input) =>
253 (input.length === 0
254 ? Effect.succeed(
255 new EmbedManyResponse({
256 embeddings: [],
257 usage: new EmbeddingUsage({ inputTokens: undefined })
258 })
259 )
260 : params.embedMany({ inputs: input }).pipe(
261 Effect.flatMap((response) =>
262 mapProviderResults(input.length, response.results).pipe(
263 Effect.map((embeddings) =>
264 new EmbedManyResponse({
265 embeddings,
266 usage: new EmbeddingUsage({
267 inputTokens: response.usage.inputTokens
268 })
269 })
270 )
271 )
272 )
273 )).pipe(Effect.withSpan("EmbeddingModel.embedMany"))
274 })
275})
276
277const mapProviderResults = (
278 inputLength: number,
279 results: Array<Array<number>>
280): Effect.Effect<Array<EmbedResponse>, AiError.AiError> => {
281 const embeddings = new Array<EmbedResponse>(inputLength)
282 if (results.length !== inputLength) {
283 return Effect.fail(
284 invalidProviderResponse(
285 `Provider returned ${results.length} embeddings but expected ${inputLength}`
286 )
287 )
288 }
289 for (let i = 0; i < results.length; i++) {
290 const vector = results[i]
291 embeddings[i] = new EmbedResponse({ vector })
292 }
293 return Effect.succeed(embeddings)
294}