Skip to main content

submilli_engine/stdlib/embedding/
declaration.rs

1//! Type surface of `submilli:embedding`.
2//!
3//! `Embeddings` is the sealed result: its scalar facts are properties, and its
4//! vectors are reachable only through the `vector` and `bytes` methods, because
5//! a plain `number[][]` costs about 60 bytes per number of run memory while the
6//! host holds the same vectors at 4.
7
8use 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
72/// `\"query\" | \"document\"`. The host re-validates at run time, so a cast that
73/// bypasses the compile-time check still gets a `RangeError`, not a guess.
74fn 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}