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: Default::default(),
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: Default::default(),
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`. Returned 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. Returned 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: Default::default(),
232 dispatch: if methods.is_empty() {
233 Dispatch::VTable
234 } else {
235 Dispatch::Direct
236 },
237 methods,
238 properties,
239 doc: crate::doc(crate::FileId::EMBEDDING, doc),
240 },
241 },
242 );
243}
244
245fn insert_property(
246 properties: &mut BTreeMap<String, PropertySig>,
247 name: &str,
248 ty: Type,
249 doc: &str,
250) {
251 properties.insert(
252 name.to_string(),
253 PropertySig {
254 ty,
255 readonly: true,
256 optional: false,
257 intrinsic: false,
258 doc: crate::doc(crate::FileId::EMBEDDING, doc),
259 },
260 );
261}
262
263fn insert_fn(defs: &mut PackageDeclaration, name: &str, params: Vec<Param>, ret: Type, doc: &str) {
264 defs.values.insert(
265 name.to_string(),
266 ValueSymbol {
267 name: name.to_string(),
268 mangled_name: crate::mangle::package_symbol(MODULE_NAME, name),
269 declaration_span: Span::at(crate::FileId::EMBEDDING),
270 kind: ValueKind::Function {
271 generics: Default::default(),
272 params,
273 ret,
274 type_predicate: None,
275 doc: crate::doc(crate::FileId::EMBEDDING, doc),
276 },
277 },
278 );
279}