1use std::collections::BTreeMap;
9
10use crate::{
11 Dispatch, MethodSig, PackageDeclaration, Param, PropertySig, Span, Type, TypeKind, TypeSymbol,
12 ValueKind, ValueSymbol,
13};
14
15use super::MODULE_NAME;
16
17const EMBEDDINGS: &str = "Embeddings";
18const EMBEDDING_MODEL: &str = "EmbeddingModel";
19
20pub fn package_declaration() -> PackageDeclaration {
21 let mut defs = PackageDeclaration::with_package(MODULE_NAME);
22 insert_embeddings_interface(&mut defs);
23 insert_model_interface(&mut defs);
24 insert_fn(
25 &mut defs,
26 "embed",
27 vec![
28 Param::new("model", Type::String),
29 Param::new("texts", Type::Array(Box::new(Type::String))),
30 Param::new("purpose", purpose_type()),
31 ],
32 interface_type(EMBEDDINGS),
33 "/**\n * Embed `texts` with the alias `model` and return one sealed `Embeddings`: \
34 `result.vector(i)` is the vector of `texts[i]`, in input order.\n *\n * `purpose` \
35 is `\"query\"` for text you will search with and `\"document\"` for text you will \
36 search over. Some providers embed the two differently; where the model has no \
37 notion of purpose it has no effect. Results from one alias share an `identity` \
38 whichever purpose they used.\n *\n * The batch either fully succeeds or throws: \
39 there are no per-text outcomes. An empty `texts` throws a `RangeError`. Every \
40 vector is unit length.\n *\n * Nothing is truncated. A text longer than the \
41 alias's `maxInputBytes` (see `models()`) throws a `RangeError` naming its index \
42 (numbered from 0) before anything is sent and before any budget is charged, and \
43 so does a call over the per-call limits of 128 texts and 2 MiB of text. Throws a \
44 `QuotaExceededError` when the execution or server embedding budget cannot cover \
45 the call, a `TypeError` when the provider returns an unusable response, and a \
46 catchable `Error` with a fixed reason such as `rate-limited` or `transport` for \
47 other provider failures. An unknown alias throws an error naming the aliases this \
48 caller may use. No error ever quotes the input text or a vector.\n * @param model \
49 Alias name. Call `models()` for the ones this runtime serves.\n * @param texts \
50 The texts to embed, each at most the alias's `maxInputBytes` in UTF-8 bytes.\n * \
51 @param purpose `\"query\"` or `\"document\"`.\n * @capability embedding.embed { \
52 model: $model, input_count: $texts.length }\n */",
53 );
54 insert_fn(
55 &mut defs,
56 "models",
57 Vec::new(),
58 Type::Array(Box::new(interface_type(EMBEDDING_MODEL))),
59 "/**\n * The embedding aliases this runtime serves and this caller may use.\n *\n * \
60 Each candidate is filtered under `embedding.embed` by the same `model` filter \
61 that gates `embed`, so a listing never offers an alias the caller would be \
62 denied at `embed` time. The list can therefore come back short, or empty, and \
63 nothing in it reveals how many candidates were filtered out.\n *\n * Size inputs \
64 by `maxInputBytes`: an input within it is never refused for length before \
65 sending. Compare `identity` before mixing vectors from two results; vectors \
66 from different identities are not comparable even when `dimensions` match.\n * \
67 @capability embedding.embed { model: $model, input_count: 0 }\n */",
68 );
69 defs
70}
71
72fn purpose_type() -> Type {
75 Type::union(vec![
76 Type::StringLiteral("query".to_string()),
77 Type::StringLiteral("document".to_string()),
78 ])
79}
80
81fn interface_type(name: &str) -> Type {
82 Type::InterfaceRef {
83 mangled: crate::mangle::package_symbol(MODULE_NAME, name),
84 package: crate::Package(MODULE_NAME.to_string()),
85 name: name.to_string(),
86 args: Vec::new(),
87 }
88}
89
90fn optional_number() -> Type {
91 Type::union(vec![Type::Number, Type::Undefined])
92}
93
94fn optional_string() -> Type {
95 Type::union(vec![Type::String, Type::Undefined])
96}
97
98fn insert_embeddings_interface(defs: &mut PackageDeclaration) {
99 let mut properties = BTreeMap::new();
100 insert_property(
101 &mut properties,
102 "count",
103 Type::Number,
104 "/** The number of vectors, equal to the number of texts embedded. */",
105 );
106 insert_property(
107 &mut properties,
108 "dimensions",
109 Type::Number,
110 "/** The length of every vector. */",
111 );
112 insert_property(
113 &mut properties,
114 "identity",
115 Type::String,
116 "/** The embedding-space identity of every vector here. Opaque: compare the whole string for equality and never parse it. Vectors with different identities are not comparable. It covers what the blueprint configures for the alias (provider, model id, dimensions, normalization, purpose handling); it cannot detect a vendor changing a model behind an unchanged name. */",
117 );
118 insert_property(
119 &mut properties,
120 "model",
121 Type::String,
122 "/** The alias these vectors were embedded with. */",
123 );
124 insert_property(
125 &mut properties,
126 "inputTokens",
127 optional_number(),
128 "/** Input tokens the provider reported for the whole call, or `undefined` when it reported none. `undefined` means unknown, not free: the call may still have been billed. */",
129 );
130
131 let mut methods = BTreeMap::new();
132 methods.insert(
133 "vector".to_string(),
134 MethodSig {
135 optional: false,
136 generics: Vec::new(),
137 params: vec![Param::new("index", Type::Number)],
138 ret: Type::Array(Box::new(Type::Number)),
139 predicate: None,
140 doc: crate::doc(
141 crate::FileId::EMBEDDING,
142 "/**\n * The vector of `texts[index]` as a `number[]` of `dimensions` numbers. Built on demand, so read the ones you need rather than every vector: a plain array costs far more run memory than the sealed result does.\n * @param index Zero-based position; throws a `RangeError` unless it is an integer below `count`.\n */",
143 ),
144 },
145 );
146 methods.insert(
147 "bytes".to_string(),
148 MethodSig {
149 optional: false,
150 generics: Vec::new(),
151 params: vec![Param::new("index", Type::Number)],
152 ret: Type::Uint8Array,
153 predicate: None,
154 doc: crate::doc(
155 crate::FileId::EMBEDDING,
156 "/**\n * The vector of `texts[index]` as little-endian 32-bit floats: `dimensions * 4` bytes, the compact form for storing a vector.\n * @param index Zero-based position; throws a `RangeError` unless it is an integer below `count`.\n */",
157 ),
158 },
159 );
160 insert_interface(
161 defs,
162 EMBEDDINGS,
163 properties,
164 methods,
165 "/** The sealed result of `embed`: vectors held by the runtime at 4 bytes per number, in input order, labeled with the embedding-space `identity`. Constructed only by `embed`. Memory is released when the program drops the result. */",
166 );
167}
168
169fn insert_model_interface(defs: &mut PackageDeclaration) {
170 let mut properties = BTreeMap::new();
171 insert_property(
172 &mut properties,
173 "name",
174 Type::String,
175 "/** The alias name, exactly as `embed` expects it. */",
176 );
177 insert_property(
178 &mut properties,
179 "description",
180 optional_string(),
181 "/** Operator-authored deployment intent, or `undefined` when none was declared. Advice, not fact: it is free text. Always a single line, within a fixed length bound. */",
182 );
183 insert_property(
184 &mut properties,
185 "dimensions",
186 Type::Number,
187 "/** The length of every vector this alias returns. */",
188 );
189 insert_property(
190 &mut properties,
191 "maxInputTokens",
192 optional_number(),
193 "/** The alias's input limit in tokens, or `undefined` when none is known. */",
194 );
195 insert_property(
196 &mut properties,
197 "maxInputBytes",
198 Type::Number,
199 "/** The UTF-8 byte length above which one text is refused before sending. Size chunks by this, not by tokens: a text within it is never refused for length. */",
200 );
201 insert_property(
202 &mut properties,
203 "identity",
204 Type::String,
205 "/** The embedding-space identity results from this alias carry. Opaque: compare for equality only. */",
206 );
207 insert_interface(
208 defs,
209 EMBEDDING_MODEL,
210 properties,
211 BTreeMap::new(),
212 "/** One alias this caller may embed with. Constructed only by `models()`, and the list carries nothing derived from the candidates the policy filtered out. */",
213 );
214}
215
216fn insert_interface(
217 defs: &mut PackageDeclaration,
218 name: &str,
219 properties: BTreeMap<String, PropertySig>,
220 methods: BTreeMap<String, MethodSig>,
221 doc: &str,
222) {
223 defs.types.insert(
224 name.to_string(),
225 TypeSymbol {
226 name: name.to_string(),
227 mangled_name: crate::mangle::package_symbol(MODULE_NAME, name),
228 declaration_span: Span::at(crate::FileId::EMBEDDING),
229 kind: TypeKind::Interface {
230 index: None,
231 generics: Vec::new(),
232 methods,
233 properties,
234 dispatch: Dispatch::Direct,
235 doc: crate::doc(crate::FileId::EMBEDDING, doc),
236 },
237 },
238 );
239}
240
241fn insert_property(
242 properties: &mut BTreeMap<String, PropertySig>,
243 name: &str,
244 ty: Type,
245 doc: &str,
246) {
247 properties.insert(
248 name.to_string(),
249 PropertySig {
250 ty,
251 readonly: true,
252 optional: false,
253 intrinsic: false,
254 doc: crate::doc(crate::FileId::EMBEDDING, doc),
255 },
256 );
257}
258
259fn insert_fn(defs: &mut PackageDeclaration, name: &str, params: Vec<Param>, ret: Type, doc: &str) {
260 defs.values.insert(
261 name.to_string(),
262 ValueSymbol {
263 name: name.to_string(),
264 mangled_name: crate::mangle::package_symbol(MODULE_NAME, name),
265 declaration_span: Span::at(crate::FileId::EMBEDDING),
266 kind: ValueKind::Function {
267 generics: Vec::new(),
268 params,
269 ret,
270 type_predicate: None,
271 doc: crate::doc(crate::FileId::EMBEDDING, doc),
272 },
273 },
274 );
275}