khroma/
models.rs

1#![allow(non_snake_case)]
2use serde::{Deserialize, Serialize};
3use std::collections::HashMap;
4use uuid::Uuid;
5
6pub type CollectionUuid = Uuid;
7pub type Metadata = HashMap<String, serde_json::Value>;
8
9#[derive(Serialize, Deserialize, Debug, Clone)]
10pub struct ErrorResponse {
11    pub error: String,
12    pub message: String,
13}
14
15#[derive(Serialize, Deserialize, Debug, Clone)]
16pub struct GetUserIdentityResponse {
17    pub user_id: String,
18    pub tenant: String,
19    pub databases: Vec<String>,
20}
21
22#[derive(Serialize, Deserialize, Debug, Clone)]
23pub struct HeartbeatResponse {
24    #[serde(rename = "nanosecond heartbeat")]
25    pub nanosecond_heartbeat: u64,
26}
27
28#[derive(Serialize, Deserialize, Debug, Clone)]
29pub struct ChecklistResponse {
30    pub max_batch_size: i32,
31    pub supports_base64_encoding: bool,
32}
33
34#[derive(Serialize, Deserialize, Debug, Clone)]
35pub struct CreateTenantPayload {
36    pub name: String,
37}
38
39#[derive(Serialize, Deserialize, Debug, Clone)]
40pub struct CreateTenantResponse {}
41
42#[derive(Serialize, Deserialize, Debug, Clone)]
43pub struct GetTenantResponse {
44    pub name: String,
45}
46
47#[derive(Serialize, Deserialize, Debug, Clone)]
48pub struct CreateDatabasePayload {
49    pub name: String,
50}
51
52#[derive(Serialize, Deserialize, Debug, Clone)]
53pub struct CreateDatabaseResponse {}
54
55#[derive(Serialize, Deserialize, Debug, Clone)]
56pub struct Database {
57    pub id: Uuid,
58    pub name: String,
59    pub tenant: String,
60}
61
62#[derive(Serialize, Deserialize, Debug, Clone)]
63pub struct DeleteDatabaseResponse {}
64
65#[derive(Serialize, Deserialize, Debug, Clone)]
66#[serde(untagged)]
67pub enum EmbeddingFunctionConfiguration {
68    Legacy {
69        r#type: String, // "legacy"
70    },
71    Known {
72        r#type: String, // "known"
73        #[serde(flatten)]
74        config: EmbeddingFunctionNewConfiguration,
75    },
76}
77
78#[derive(Serialize, Deserialize, Debug, Clone)]
79pub struct EmbeddingFunctionNewConfiguration {
80    pub name: String,
81    pub config: serde_json::Value,
82}
83
84#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
85#[serde(rename_all = "lowercase")]
86pub enum HnswSpace {
87    L2,
88    Cosine,
89    Ip,
90}
91
92#[derive(Serialize, Deserialize, Debug, Clone, Default)]
93#[serde(deny_unknown_fields)]
94pub struct HnswConfiguration {
95    #[serde(skip_serializing_if = "Option::is_none")]
96    pub ef_construction: Option<u32>,
97    #[serde(skip_serializing_if = "Option::is_none")]
98    pub ef_search: Option<u32>,
99    #[serde(skip_serializing_if = "Option::is_none")]
100    pub max_neighbors: Option<u32>,
101    #[serde(skip_serializing_if = "Option::is_none")]
102    pub resize_factor: Option<f64>,
103    #[serde(skip_serializing_if = "Option::is_none")]
104    pub space: Option<HnswSpace>,
105    #[serde(skip_serializing_if = "Option::is_none")]
106    pub sync_threshold: Option<u32>,
107}
108
109
110#[derive(Serialize, Deserialize, Debug, Clone, Default)]
111#[serde(deny_unknown_fields)]
112pub struct SpannConfiguration {
113    #[serde(skip_serializing_if = "Option::is_none")]
114    pub ef_construction: Option<u32>,
115    #[serde(skip_serializing_if = "Option::is_none")]
116    pub ef_search: Option<u32>,
117    #[serde(skip_serializing_if = "Option::is_none")]
118    pub max_neighbors: Option<u32>,
119    #[serde(skip_serializing_if = "Option::is_none")]
120    pub merge_threshold: Option<i32>,
121    #[serde(skip_serializing_if = "Option::is_none")]
122    pub reassign_neighbor_count: Option<i32>,
123    #[serde(skip_serializing_if = "Option::is_none")]
124    pub search_nprobe: Option<i32>,
125    #[serde(skip_serializing_if = "Option::is_none")]
126    pub space: Option<HnswSpace>,
127    #[serde(skip_serializing_if = "Option::is_none")]
128    pub split_threshold: Option<i32>,
129    #[serde(skip_serializing_if = "Option::is_none")]
130    pub write_nprobe: Option<i32>,
131}
132
133#[derive(Serialize, Deserialize, Debug, Clone, Default)]
134pub struct CollectionConfiguration {
135    #[serde(skip_serializing_if = "Option::is_none")]
136    pub embedding_function: Option<EmbeddingFunctionConfiguration>,
137    #[serde(skip_serializing_if = "Option::is_none")]
138    pub hnsw: Option<HnswConfiguration>,
139    #[serde(skip_serializing_if = "Option::is_none")]
140    pub spann: Option<SpannConfiguration>,
141}
142
143#[derive(Serialize, Deserialize, Debug, Clone)]
144pub struct Collection {
145    pub id: CollectionUuid,
146    pub name: String,
147    #[serde(skip_serializing_if = "Option::is_none")]
148    pub metadata: Option<Metadata>,
149    pub configuration_json: CollectionConfiguration,
150    pub tenant: String,
151    pub database: String,
152    pub log_position: i64,
153    pub version: i32,
154    #[serde(skip_serializing_if = "Option::is_none")]
155    pub dimension: Option<i32>,
156}
157
158#[derive(Serialize, Deserialize, Debug, Clone, Default)]
159pub struct CreateCollectionPayload {
160    pub name: String,
161    #[serde(skip_serializing_if = "Option::is_none")]
162    pub metadata: Option<Metadata>,
163    #[serde(skip_serializing_if = "Option::is_none")]
164    pub configuration: Option<CollectionConfiguration>,
165    #[serde(skip_serializing_if = "Option::is_none")]
166    pub get_or_create: Option<bool>,
167}
168
169#[derive(Serialize, Deserialize, Debug, Clone)]
170pub struct UpdateCollectionResponse {}
171
172#[derive(Serialize, Deserialize, Debug, Clone, Default)]
173pub struct UpdateCollectionPayload {
174    #[serde(skip_serializing_if = "Option::is_none")]
175    pub new_name: Option<String>,
176    #[serde(skip_serializing_if = "Option::is_none")]
177    pub new_metadata: Option<Metadata>,
178    #[serde(skip_serializing_if = "Option::is_none")]
179    pub new_configuration: Option<UpdateCollectionConfiguration>,
180}
181
182#[derive(Serialize, Deserialize, Debug, Clone, Default)]
183pub struct UpdateCollectionConfiguration {
184    #[serde(skip_serializing_if = "Option::is_none")]
185    pub embedding_function: Option<EmbeddingFunctionConfiguration>,
186    #[serde(skip_serializing_if = "Option::is_none")]
187    pub hnsw: Option<UpdateHnswConfiguration>,
188    #[serde(skip_serializing_if = "Option::is_none")]
189    pub spann: Option<SpannConfiguration>,
190}
191
192#[derive(Serialize, Deserialize, Debug, Clone, Default)]
193#[serde(deny_unknown_fields)]
194pub struct UpdateHnswConfiguration {
195    #[serde(skip_serializing_if = "Option::is_none")]
196    pub batch_size: Option<u32>,
197    #[serde(skip_serializing_if = "Option::is_none")]
198    pub ef_search: Option<u32>,
199    #[serde(skip_serializing_if = "Option::is_none")]
200    pub max_neighbors: Option<u32>,
201    #[serde(skip_serializing_if = "Option::is_none")]
202    pub num_threads: Option<u32>,
203    #[serde(skip_serializing_if = "Option::is_none")]
204    pub resize_factor: Option<f64>,
205    #[serde(skip_serializing_if = "Option::is_none")]
206    pub sync_threshold: Option<u32>,
207}
208
209#[derive(Serialize, Deserialize, Debug, Clone)]
210#[serde(untagged)]
211pub enum EmbeddingsPayload {
212    Float(Vec<Vec<f32>>),
213    String(Vec<String>),
214}
215
216#[derive(Serialize, Deserialize, Debug, Clone)]
217pub struct AddCollectionRecordsPayload {
218    pub ids: Vec<String>,
219    #[serde(skip_serializing_if = "Option::is_none")]
220    pub embeddings: Option<EmbeddingsPayload>,
221    #[serde(skip_serializing_if = "Option::is_none")]
222    pub metadatas: Option<Vec<Option<Metadata>>>,
223    #[serde(skip_serializing_if = "Option::is_none")]
224    pub documents: Option<Vec<Option<String>>>,
225    #[serde(skip_serializing_if = "Option::is_none")]
226    pub uris: Option<Vec<Option<String>>>,
227}
228
229#[derive(Serialize, Deserialize, Debug, Clone)]
230pub struct AddCollectionRecordsResponse {}
231
232#[derive(Serialize, Deserialize, Debug, Clone, Default)]
233pub struct RawWhereFields {
234    #[serde(skip_serializing_if = "Option::is_none")]
235    pub r#where: Option<serde_json::Value>,
236    #[serde(skip_serializing_if = "Option::is_none")]
237    pub where_document: Option<serde_json::Value>,
238}
239
240#[derive(Serialize, Deserialize, Debug, Clone, Default)]
241pub struct DeleteCollectionRecordsPayload {
242    #[serde(flatten)]
243    pub where_fields: RawWhereFields,
244    #[serde(skip_serializing_if = "Option::is_none")]
245    pub ids: Option<Vec<String>>,
246}
247
248#[derive(Serialize, Deserialize, Debug, Clone)]
249pub struct DeleteCollectionRecordsResponse {}
250
251#[derive(Serialize, Deserialize, Debug, Clone)]
252pub struct ForkCollectionPayload {
253    pub new_name: String,
254}
255
256#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)]
257#[serde(rename_all = "lowercase")]
258pub enum Include {
259    Distances,
260    Documents,
261    Embeddings,
262    Metadatas,
263    Uris,
264}
265
266pub type IncludeList = Vec<Include>;
267
268#[derive(Serialize, Deserialize, Debug, Clone, Default)]
269pub struct GetRequestPayload {
270    #[serde(flatten)]
271    pub where_fields: RawWhereFields,
272    #[serde(skip_serializing_if = "Option::is_none")]
273    pub ids: Option<Vec<String>>,
274    #[serde(skip_serializing_if = "Option::is_none")]
275    pub include: Option<IncludeList>,
276    #[serde(skip_serializing_if = "Option::is_none")]
277    pub limit: Option<i32>,
278    #[serde(skip_serializing_if = "Option::is_none")]
279    pub offset: Option<i32>,
280}
281
282#[derive(Serialize, Deserialize, Debug, Clone)]
283pub struct GetResponse {
284    pub ids: Vec<String>,
285    pub include: Vec<Include>,
286    #[serde(skip_serializing_if = "Option::is_none")]
287    pub metadatas: Option<Vec<Option<Metadata>>>,
288    #[serde(skip_serializing_if = "Option::is_none")]
289    pub documents: Option<Vec<Option<String>>>,
290    #[serde(skip_serializing_if = "Option::is_none")]
291    pub uris: Option<Vec<Option<String>>>,
292    #[serde(skip_serializing_if = "Option::is_none")]
293    pub embeddings: Option<Vec<Vec<f32>>>,
294}
295
296#[derive(Serialize, Deserialize, Debug, Clone)]
297pub struct QueryRequestPayload {
298    #[serde(flatten)]
299    pub where_fields: RawWhereFields,
300    pub query_embeddings: Vec<Vec<f32>>,
301    #[serde(skip_serializing_if = "Option::is_none")]
302    pub ids: Option<Vec<String>>,
303    #[serde(skip_serializing_if = "Option::is_none")]
304    pub include: Option<IncludeList>,
305    #[serde(skip_serializing_if = "Option::is_none")]
306    pub n_results: Option<i32>,
307}
308
309#[derive(Serialize, Deserialize, Debug, Clone)]
310pub struct QueryResponse {
311    pub ids: Vec<Vec<String>>,
312    pub include: Vec<Include>,
313    #[serde(skip_serializing_if = "Option::is_none")]
314    pub distances: Option<Vec<Vec<Option<f32>>>>,
315    #[serde(skip_serializing_if = "Option::is_none")]
316    pub metadatas: Option<Vec<Vec<Option<Metadata>>>>,
317    #[serde(skip_serializing_if = "Option::is_none")]
318    pub documents: Option<Vec<Vec<Option<String>>>>,
319    #[serde(skip_serializing_if = "Option::is_none")]
320    pub uris: Option<Vec<Vec<Option<String>>>>,
321    #[serde(skip_serializing_if = "Option::is_none")]
322    pub embeddings: Option<Vec<Vec<Vec<Option<f32>>>>>,
323}
324
325#[derive(Serialize, Deserialize, Debug, Clone)]
326#[serde(untagged)]
327pub enum UpdateEmbeddingsPayload {
328    Float(Vec<Option<Vec<f32>>>),
329    String(Vec<Option<String>>),
330}
331
332#[derive(Serialize, Deserialize, Debug, Clone)]
333pub struct UpdateCollectionRecordsPayload {
334    pub ids: Vec<String>,
335    #[serde(skip_serializing_if = "Option::is_none")]
336    pub embeddings: Option<UpdateEmbeddingsPayload>,
337    #[serde(skip_serializing_if = "Option::is_none")]
338    pub metadatas: Option<Vec<Option<Metadata>>>,
339    #[serde(skip_serializing_if = "Option::is_none")]
340    pub documents: Option<Vec<Option<String>>>,
341    #[serde(skip_serializing_if = "Option::is_none")]
342    pub uris: Option<Vec<Option<String>>>,
343}
344
345#[derive(Serialize, Deserialize, Debug, Clone)]
346pub struct UpdateCollectionRecordsResponse {}
347
348#[derive(Serialize, Deserialize, Debug, Clone)]
349pub struct UpsertCollectionRecordsPayload {
350    pub ids: Vec<String>,
351    #[serde(skip_serializing_if = "Option::is_none")]
352    pub embeddings: Option<EmbeddingsPayload>,
353    #[serde(skip_serializing_if = "Option::is_none")]
354    pub metadatas: Option<Vec<Option<Metadata>>>,
355    #[serde(skip_serializing_if = "Option::is_none")]
356    pub documents: Option<Vec<Option<String>>>,
357    #[serde(skip_serializing_if = "Option::is_none")]
358    pub uris: Option<Vec<Option<String>>>,
359}
360
361#[derive(Serialize, Deserialize, Debug, Clone)]
362pub struct UpsertCollectionRecordsResponse {}
khroma/models.rs

khroma/
models.rs
