tricoredb/
vector.rs

1//! Vectors: fixed-dimension collections, similarity search and metadata
2//! filters.
3
4use std::collections::BTreeMap;
5
6use serde::{Deserialize, Serialize};
7use serde_json::{json, Map, Value};
8
9use crate::client::Client;
10use crate::error::{Error, Result};
11
12/// How a collection measures closeness. Fixed when the collection is created.
13#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
14#[serde(rename_all = "lowercase")]
15pub enum VectorMetric {
16    /// Cosine similarity: higher is closer.
17    #[default]
18    Cosine,
19    /// Dot product: higher is closer.
20    Dot,
21    /// Squared Euclidean distance, **returned negated** so that higher is still
22    /// closer. An l2 score is therefore `<= 0`, and `-0.02` is nearer than
23    /// `-196.0`. See [`VectorHit::score`].
24    L2,
25}
26
27/// How the search index stores vectors in memory.
28///
29/// This is an index-level choice only: the durable records always keep full
30/// `f32` precision, so [`Client::vector_get`] returns byte-identical values
31/// either way.
32#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
33#[serde(rename_all = "lowercase")]
34pub enum VectorQuantization {
35    /// Full `f32` vectors in the index. The default.
36    #[default]
37    None,
38    /// `int8`-quantized vectors in the index, roughly four times smaller, with
39    /// candidates re-ranked against the durable `f32` vectors.
40    Int8,
41}
42
43/// One stored vector and its metadata.
44#[derive(Debug, Clone, PartialEq, Deserialize)]
45pub struct VectorItem {
46    /// The vector's id.
47    pub id: String,
48    /// The vector itself.
49    #[serde(default, deserialize_with = "crate::serde_null::or_default")]
50    pub vector: Vec<f32>,
51    /// Metadata stored alongside it.
52    #[serde(default, deserialize_with = "crate::serde_null::or_default")]
53    pub metadata: Map<String, Value>,
54}
55
56/// One search result.
57///
58/// The score is a similarity, never a distance: **higher is closer under every
59/// metric**, and results come back best first. Sorting these ascending, or
60/// reading the score as a distance, inverts an l2 ranking while a "is the
61/// expected id in the top k" check still passes.
62#[derive(Debug, Clone, PartialEq, Deserialize)]
63pub struct VectorHit {
64    /// The matching vector's id.
65    pub id: String,
66    /// Its similarity to the query vector.
67    #[serde(default)]
68    pub score: f32,
69    /// Its metadata.
70    #[serde(default, deserialize_with = "crate::serde_null::or_default")]
71    pub metadata: Map<String, Value>,
72    /// The vector itself, when the search asked for vectors.
73    #[serde(default)]
74    pub vector: Option<Vec<f32>>,
75    /// Why it matched, when the search asked for an explanation.
76    #[serde(default)]
77    pub reason: Option<String>,
78}
79
80/// A search's hits, and which path answered it.
81#[derive(Debug, Clone, PartialEq, Deserialize)]
82pub struct VectorSearchResult {
83    /// The hits, best first.
84    #[serde(default, deserialize_with = "crate::serde_null::or_default")]
85    pub results: Vec<VectorHit>,
86    /// `"flat"` for the exact scan a small collection gets, `"hnsw"` or
87    /// `"hnsw-int8"` for the approximate index.
88    #[serde(default)]
89    pub index: String,
90}
91
92/// One entry of [`Client::vector_list_collections`].
93#[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
94pub struct VectorCollectionSummary {
95    /// The collection's name.
96    pub name: String,
97    /// The length every vector in it must have.
98    #[serde(default)]
99    pub dimension: usize,
100    /// The metric searches rank by.
101    #[serde(default)]
102    pub metric: VectorMetric,
103}
104
105/// A collection's metadata and its live vector count.
106#[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
107pub struct VectorCollectionInfo {
108    /// The collection's name.
109    pub collection: String,
110    /// The length every vector in it must have.
111    #[serde(default)]
112    pub dimension: usize,
113    /// The metric searches rank by.
114    #[serde(default)]
115    pub metric: VectorMetric,
116    /// How many vectors are stored.
117    #[serde(default)]
118    pub count: i64,
119    /// The index quantization the collection was created with.
120    #[serde(default)]
121    pub quantization: VectorQuantization,
122}
123
124/// One page of [`Client::vector_list_vectors`].
125#[derive(Debug, Clone, PartialEq, Deserialize)]
126pub struct VectorPage {
127    /// The collection's name.
128    pub collection: String,
129    /// How many vectors are on this page.
130    #[serde(default)]
131    pub count: usize,
132    /// This page's vectors.
133    #[serde(default, deserialize_with = "crate::serde_null::or_default")]
134    pub vectors: Vec<VectorItem>,
135    /// Whether more vectors remain past this page.
136    #[serde(default)]
137    pub truncated: bool,
138    /// The collection's full count.
139    #[serde(default)]
140    pub total: i64,
141}
142
143/// The measure a search scores by, when it should differ from the
144/// collection's own metric.
145#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
146#[serde(rename_all = "snake_case")]
147pub enum ScoreMetric {
148    /// Cosine similarity.
149    Cosine,
150    /// Dot product.
151    Dot,
152    /// Squared Euclidean distance, negated.
153    L2,
154    /// Manhattan distance, negated.
155    L1,
156    /// Chebyshev distance, negated.
157    Linf,
158}
159
160/// How a filtered search combines the filter with the index.
161#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
162#[serde(rename_all = "snake_case")]
163pub enum FilterStrategy {
164    /// Let the server choose. The default.
165    #[default]
166    Auto,
167    /// Walk the index and test each candidate.
168    Graph,
169    /// Scan the vectors the filter admits.
170    Exact,
171}
172
173/// Which approximate index a search uses.
174#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
175#[serde(rename_all = "snake_case")]
176pub enum AnnIndex {
177    /// The HNSW graph.
178    Hnsw,
179    /// An inverted file index.
180    Ivf,
181    /// An inverted file index over product-quantized vectors.
182    IvfPq,
183}
184
185/// How [`Client::vector_search_multi`] merges its queries' rankings.
186#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
187#[serde(rename_all = "snake_case")]
188pub enum FusionMethod {
189    /// Reciprocal rank fusion. The default.
190    #[default]
191    Rrf,
192    /// Each vector's best score across the queries.
193    Max,
194    /// Each vector's mean score across the queries.
195    Mean,
196}
197
198/// A sparse vector: parallel lists of term indices and their weights.
199#[derive(Debug, Clone, PartialEq, Default, Serialize, Deserialize)]
200pub struct SparseVector {
201    /// The indices that carry a weight.
202    #[serde(default)]
203    pub indices: Vec<u32>,
204    /// The weight at each index.
205    #[serde(default)]
206    pub values: Vec<f32>,
207}
208
209/// What a batch or example search should apply beyond its query vectors.
210///
211/// Every field left at its default is left off the wire, so the server's own
212/// default applies.
213#[derive(Debug, Clone, Default, PartialEq)]
214pub struct VectorSearchOptions {
215    /// Keep only vectors whose metadata matches this filter.
216    pub filter: Option<Value>,
217    /// Drop hits scoring below this.
218    pub min_score: Option<f32>,
219    /// Return each hit's vector with it.
220    pub include_vectors: bool,
221    /// Skip this many hits from the front.
222    pub offset: Option<usize>,
223    /// Return each hit's metadata. `None` leaves the server's default, which
224    /// is to include it.
225    pub include_metadata: Option<bool>,
226    /// Drop hits farther than this distance.
227    pub max_distance: Option<f32>,
228    /// The size of the candidate list the index walk keeps.
229    pub ef_search: Option<usize>,
230    /// Scan every vector instead of using the index.
231    pub exact: bool,
232    /// Score by this measure instead of the collection's metric.
233    pub score_metric: Option<ScoreMetric>,
234    /// How to combine the filter with the index.
235    pub filter_strategy: Option<FilterStrategy>,
236    /// Search this named vector instead of the collection's own.
237    pub using: Option<String>,
238    /// Restrict the search to this tenant.
239    pub tenant: Option<String>,
240    /// Restrict the search to this namespace.
241    pub namespace: Option<String>,
242    /// A sparse query to combine with the dense one.
243    pub sparse: Option<SparseVector>,
244    /// The dense share of a hybrid score, from 0 to 1.
245    pub hybrid_alpha: Option<f32>,
246    /// Attach to each hit the reason it matched.
247    pub explain: bool,
248    /// An expression that rescores each hit.
249    pub score_expr: Option<String>,
250    /// The query text an explanation should refer to.
251    pub explain_text: Option<String>,
252    /// Which approximate index to use.
253    pub ann: Option<AnnIndex>,
254    /// How many inverted-file cells to probe.
255    pub nprobe: Option<usize>,
256}
257
258impl VectorSearchOptions {
259    fn write(&self, body: &mut Map<String, Value>) {
260        fn put(body: &mut Map<String, Value>, key: &str, value: Option<impl Serialize>) {
261            if let Some(value) = value {
262                body.insert(key.to_string(), json!(value));
263            }
264        }
265        put(body, "filter", self.filter.as_ref());
266        put(body, "min_score", self.min_score);
267        put(
268            body,
269            "include_vectors",
270            self.include_vectors.then_some(true),
271        );
272        put(body, "offset", self.offset);
273        put(body, "include_metadata", self.include_metadata);
274        put(body, "max_distance", self.max_distance);
275        put(body, "ef_search", self.ef_search);
276        put(body, "exact", self.exact.then_some(true));
277        put(body, "score_metric", self.score_metric);
278        put(body, "filter_strategy", self.filter_strategy);
279        put(body, "using", self.using.as_ref());
280        put(body, "tenant", self.tenant.as_ref());
281        put(body, "namespace", self.namespace.as_ref());
282        put(body, "sparse", self.sparse.as_ref());
283        put(body, "hybrid_alpha", self.hybrid_alpha);
284        put(body, "explain", self.explain.then_some(true));
285        put(body, "score_expr", self.score_expr.as_ref());
286        put(body, "explain_text", self.explain_text.as_ref());
287        put(body, "ann", self.ann);
288        put(body, "nprobe", self.nprobe);
289    }
290}
291
292/// One example for [`Client::vector_search_by_examples`]: a stored vector
293/// named by id, or a vector given outright.
294#[derive(Debug, Clone, PartialEq, Serialize)]
295#[serde(untagged)]
296pub enum VectorExample {
297    /// The stored vector with this id.
298    Id(String),
299    /// This vector.
300    Vector(Vec<f32>),
301}
302
303impl VectorExample {
304    /// The stored vector with this id.
305    pub fn id(id: &str) -> Self {
306        VectorExample::Id(id.to_string())
307    }
308}
309
310/// One vector of [`Client::vector_upsert_many`].
311#[derive(Debug, Clone, PartialEq, Serialize)]
312pub struct VectorUpsert {
313    /// The vector's id.
314    pub id: String,
315    /// The vector, at the collection's dimension.
316    pub vector: Vec<f32>,
317    /// Metadata stored alongside it.
318    #[serde(skip_serializing_if = "Option::is_none")]
319    pub metadata: Option<Map<String, Value>>,
320}
321
322impl VectorUpsert {
323    /// A vector with no metadata.
324    pub fn new(id: &str, vector: impl Into<Vec<f32>>) -> Self {
325        Self {
326            id: id.to_string(),
327            vector: vector.into(),
328            metadata: None,
329        }
330    }
331
332    /// Attach metadata.
333    pub fn metadata(mut self, metadata: Map<String, Value>) -> Self {
334        self.metadata = Some(metadata);
335        self
336    }
337}
338
339/// One text of [`Client::vector_upsert_text`].
340#[derive(Debug, Clone, PartialEq, Serialize)]
341pub struct TextUpsert {
342    /// The id the embedded vector is stored under.
343    pub id: String,
344    /// The text to embed. It is kept in the metadata under `text`.
345    pub text: String,
346    /// Metadata stored alongside it.
347    #[serde(skip_serializing_if = "Option::is_none")]
348    pub metadata: Option<Map<String, Value>>,
349}
350
351impl TextUpsert {
352    /// A text with no metadata.
353    pub fn new(id: &str, text: &str) -> Self {
354        Self {
355            id: id.to_string(),
356            text: text.to_string(),
357            metadata: None,
358        }
359    }
360
361    /// Attach metadata.
362    pub fn metadata(mut self, metadata: Map<String, Value>) -> Self {
363        self.metadata = Some(metadata);
364        self
365    }
366}
367
368/// What [`Client::vector_embed`] produced.
369#[derive(Debug, Clone, PartialEq, Deserialize)]
370pub struct Embedding {
371    /// The embedder's name. `hash-v1` is deterministic feature hashing:
372    /// lexical, with no model behind it.
373    #[serde(default)]
374    pub embedder: String,
375    /// The dimension asked for.
376    #[serde(default)]
377    pub dimension: usize,
378    /// One dense vector per text, in the order given.
379    #[serde(default)]
380    pub vectors: Vec<Vec<f32>>,
381    /// One sparse vector per text, in the order given.
382    #[serde(default)]
383    pub sparse: Vec<SparseVector>,
384}
385
386/// One alias and the collection it names.
387#[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
388pub struct VectorAlias {
389    /// The alias.
390    pub alias: String,
391    /// The collection every operation naming the alias acts on.
392    #[serde(default)]
393    pub collection: String,
394}
395
396/// One partition of a collection.
397#[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
398pub struct VectorPartition {
399    /// The partition's value of the field.
400    pub name: String,
401    /// How many vectors carry it.
402    #[serde(default)]
403    pub count: i64,
404}
405
406/// The partitions of a collection, as [`Client::vector_list_partitions`]
407/// reports them.
408#[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
409pub struct VectorPartitions {
410    /// The metadata field the collection was partitioned by.
411    #[serde(default)]
412    pub field: String,
413    /// Each distinct value of the field, with its vector count.
414    #[serde(default)]
415    pub partitions: Vec<VectorPartition>,
416    /// How many vectors carry no value for the field.
417    #[serde(default)]
418    pub unpartitioned: i64,
419}
420
421/// What [`Client::vector_rebuild_index`] did.
422#[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
423pub struct VectorIndexRebuild {
424    /// Whether the index was rebuilt.
425    #[serde(default)]
426    pub rebuilt: bool,
427    /// How many vectors went into it.
428    #[serde(default)]
429    pub vectors: i64,
430    /// How long the build took, in milliseconds.
431    #[serde(default)]
432    pub build_ms: i64,
433    /// Whether the new index was persisted.
434    #[serde(default)]
435    pub persisted: bool,
436    /// How many tombstones the rebuild cleared.
437    #[serde(default)]
438    pub tombstones_cleared: i64,
439}
440
441/// The shape of a collection's search index.
442///
443/// A collection small enough to be scanned has no graph, so everything past
444/// `stored` is `None` for it and [`VectorIndexStats::index`] is `"flat"`.
445#[derive(Debug, Clone, PartialEq, Deserialize)]
446pub struct VectorIndexStats {
447    /// The collection's name.
448    #[serde(default)]
449    pub collection: String,
450    /// `"flat"` for an exact scan, `"hnsw"` for the graph.
451    #[serde(default)]
452    pub index: String,
453    /// How many vectors are stored.
454    #[serde(default)]
455    pub stored: i64,
456    /// Below this many vectors a search scans instead of using a graph.
457    #[serde(default)]
458    pub flat_scan_threshold: Option<i64>,
459    /// How many layers the graph has.
460    #[serde(default)]
461    pub layers: Option<i64>,
462    /// How many nodes the graph holds, tombstones included.
463    #[serde(default)]
464    pub nodes: Option<i64>,
465    /// How many of them are live.
466    #[serde(default)]
467    pub live: Option<i64>,
468    /// How many nodes a delete or re-upsert left behind.
469    #[serde(default)]
470    pub tombstones: Option<i64>,
471    /// How many live nodes the entry point can reach.
472    #[serde(default)]
473    pub reachable: Option<i64>,
474    /// How many it cannot.
475    #[serde(default)]
476    pub unreachable: Option<i64>,
477    /// Whether the server recommends a rebuild.
478    #[serde(default)]
479    pub should_rebuild: Option<bool>,
480    /// Whether the index is persisted.
481    #[serde(default)]
482    pub persisted: Option<bool>,
483    /// The id of the entry point.
484    #[serde(default)]
485    pub entry: Option<String>,
486    /// Live nodes on each layer, the bottom first.
487    #[serde(default)]
488    pub live_per_layer: Option<Vec<i64>>,
489    /// The degree extremes and mean, as `{min, mean, max}`.
490    #[serde(default)]
491    pub degree: Option<Value>,
492    /// The index's memory, as `{vector_bytes, overhead_bytes}`.
493    #[serde(default)]
494    pub memory: Option<Value>,
495    /// The graph's parameters, as `{m, m0, ef_construction, ef_search}`.
496    #[serde(default)]
497    pub parameters: Option<Value>,
498}
499
500/// What [`Client::vector_measure_recall`] should sample. Every `None` leaves
501/// the server's default.
502#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
503pub struct RecallOptions {
504    /// How many stored vectors to search with. The default is 100, capped at
505    /// 1000.
506    pub sample: Option<usize>,
507    /// How long a result list to compare. The default is 10, capped at 100.
508    pub top_k: Option<usize>,
509    /// The candidate list size to measure at.
510    pub ef_search: Option<usize>,
511}
512
513/// The recall [`Client::vector_measure_recall`] measured.
514#[derive(Debug, Clone, PartialEq, Deserialize)]
515pub struct VectorRecall {
516    /// The collection's name.
517    #[serde(default)]
518    pub collection: String,
519    /// `"flat"` or `"hnsw"`.
520    #[serde(default)]
521    pub index: String,
522    /// Whether anything was measured. A collection small enough to be scanned
523    /// has nothing to measure, and its recall is 1 by construction.
524    #[serde(default)]
525    pub measured: bool,
526    /// The mean overlap of the index's results with an exact scan's.
527    #[serde(default)]
528    pub recall: f64,
529    /// The worst single query's overlap.
530    #[serde(default)]
531    pub min_recall: Option<f64>,
532    /// How many queries were run.
533    #[serde(default)]
534    pub queries: Option<i64>,
535    /// The result list length that was compared.
536    #[serde(default)]
537    pub top_k: Option<i64>,
538    /// The candidate list size it was measured at.
539    #[serde(default)]
540    pub ef_search: Option<i64>,
541    /// How many queries fell back to a scan.
542    #[serde(default)]
543    pub flat_fallbacks: Option<i64>,
544    /// Why nothing was measured, when nothing was.
545    #[serde(default)]
546    pub reason: Option<String>,
547}
548
549/// One hit of a fused ranking.
550#[derive(Debug, Clone, PartialEq, Deserialize)]
551pub struct FusedHit {
552    /// The matching vector's id.
553    pub id: String,
554    /// Its fused score. Higher is better, but the scale is the fusion
555    /// method's, not the collection metric's.
556    #[serde(default)]
557    pub score: f32,
558    /// How many of the queries ranked it.
559    #[serde(default)]
560    pub matched: usize,
561    /// Its metadata.
562    #[serde(default, deserialize_with = "crate::serde_null::or_default")]
563    pub metadata: Map<String, Value>,
564}
565
566/// The answer to [`Client::vector_search_multi`].
567#[derive(Debug, Clone, PartialEq, Deserialize)]
568pub struct FusedSearchResult {
569    /// The fusion method that was applied.
570    #[serde(default)]
571    pub fusion: FusionMethod,
572    /// How many query vectors were run.
573    #[serde(default)]
574    pub queries: usize,
575    /// How deep each query ranked before fusing.
576    #[serde(default)]
577    pub depth: usize,
578    /// The fused hits, best first.
579    #[serde(default, deserialize_with = "crate::serde_null::or_default")]
580    pub results: Vec<FusedHit>,
581}
582
583/// What [`Client::vector_set_collection_config`] should change. A `None`
584/// leaves that setting as it is.
585#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
586pub struct VectorConfigChange {
587    /// Store every vector, and rank every query, at unit length. Can only
588    /// change on an empty collection.
589    pub normalize: Option<bool>,
590    /// The HNSW graph's `m`. Takes effect at the next build, which the call
591    /// forces.
592    pub m: Option<usize>,
593    /// The HNSW graph's `ef_construction`.
594    pub ef_construction: Option<usize>,
595    /// Require every vector to carry a string `_tenant`, and every search to
596    /// name one.
597    pub multi_tenant: Option<bool>,
598}
599
600/// One named vector declared on a collection.
601#[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
602pub struct NamedVector {
603    /// Its name, which a search passes as `using`.
604    pub name: String,
605    /// Its dimension.
606    #[serde(default)]
607    pub dimension: usize,
608    /// Its metric.
609    #[serde(default)]
610    pub metric: VectorMetric,
611}
612
613/// A collection's configuration after [`Client::vector_set_collection_config`].
614#[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
615pub struct VectorConfig {
616    /// The collection's name.
617    #[serde(default)]
618    pub collection: String,
619    /// Whether vectors are stored at unit length.
620    #[serde(default)]
621    pub normalize: bool,
622    /// The HNSW graph's `m`.
623    #[serde(default)]
624    pub m: usize,
625    /// The HNSW graph's `ef_construction`.
626    #[serde(default)]
627    pub ef_construction: usize,
628    /// Whether every vector must carry a `_tenant`.
629    #[serde(default)]
630    pub multi_tenant: bool,
631    /// The named vectors declared on the collection.
632    #[serde(default)]
633    pub named: Vec<NamedVector>,
634}
635
636fn op(body: Value) -> Value {
637    json!({"Vector": body})
638}
639
640impl Client {
641    /// Create a collection. The dimension and metric are fixed for its
642    /// lifetime.
643    pub fn vector_create_collection(
644        &mut self,
645        collection: &str,
646        dimension: usize,
647        metric: VectorMetric,
648    ) -> Result<()> {
649        self.vector_create_collection_quantized(
650            collection,
651            dimension,
652            metric,
653            VectorQuantization::None,
654        )
655    }
656
657    /// Create a collection with an explicit index quantization.
658    pub fn vector_create_collection_quantized(
659        &mut self,
660        collection: &str,
661        dimension: usize,
662        metric: VectorMetric,
663        quantization: VectorQuantization,
664    ) -> Result<()> {
665        self.send(op(json!({"CreateCollection": {
666            "collection": collection,
667            "dimension": dimension,
668            "metric": metric,
669            "quantization": quantization,
670        }})))
671        .map(|_| ())
672    }
673
674    /// Drop a collection and every vector in it. Dropping one that does not
675    /// exist is an error.
676    pub fn vector_drop_collection(&mut self, collection: &str) -> Result<()> {
677        self.send(op(json!({"DropCollection": {"collection": collection}})))
678            .map(|_| ())
679    }
680
681    /// Describe every vector collection in the database.
682    pub fn vector_list_collections(&mut self) -> Result<Vec<VectorCollectionSummary>> {
683        #[derive(Deserialize)]
684        struct Details {
685            #[serde(default)]
686            details: Vec<VectorCollectionSummary>,
687        }
688        let response = self.send(op(Value::String("ListCollections".into())))?;
689        let out: Details = response.decode("ListCollections")?;
690        Ok(out.details)
691    }
692
693    /// Read one collection's metadata and live count.
694    pub fn vector_describe_collection(&mut self, collection: &str) -> Result<VectorCollectionInfo> {
695        let response = self.send(op(
696            json!({"DescribeCollection": {"collection": collection}}),
697        ))?;
698        response.decode("DescribeCollection")
699    }
700
701    /// Store a vector under an id, replacing whatever was there.
702    ///
703    /// The vector's length must equal the collection's dimension; a mismatch is
704    /// refused rather than padded or truncated.
705    pub fn vector_upsert(
706        &mut self,
707        collection: &str,
708        id: &str,
709        vector: &[f32],
710        metadata: Option<&Map<String, Value>>,
711    ) -> Result<()> {
712        self.send(op(json!({"Upsert": {
713            "collection": collection,
714            "id": id,
715            "vector": vector,
716            "metadata": metadata,
717        }})))
718        .map(|_| ())
719    }
720
721    /// Fetch one stored vector. `None` when there is no such id.
722    pub fn vector_get(&mut self, collection: &str, id: &str) -> Result<Option<VectorItem>> {
723        let response = self.send(op(json!({"Get": {"collection": collection, "id": id}})))?;
724        let data = response.expect("Json", "Get")?;
725        if data.is_null() {
726            return Ok(None);
727        }
728        response.decode("Get").map(Some)
729    }
730
731    /// Delete one vector. Deleting an absent id is not an error; a missing
732    /// collection is.
733    pub fn vector_delete(&mut self, collection: &str, id: &str) -> Result<()> {
734        self.send(op(json!({"Delete": {"collection": collection, "id": id}})))
735            .map(|_| ())
736    }
737
738    /// The `top_k` nearest vectors to a query vector.
739    ///
740    /// `top_k` must be between 1 and 1000, and the query vector's length must
741    /// equal the collection's dimension.
742    pub fn vector_search(
743        &mut self,
744        collection: &str,
745        vector: &[f32],
746        top_k: usize,
747    ) -> Result<VectorSearchResult> {
748        self.vector_search_filtered(collection, vector, top_k, &BTreeMap::new())
749    }
750
751    /// A search restricted to vectors whose metadata matches every entry of
752    /// `filter`.
753    ///
754    /// Exact equality on top-level fields only: the server implements no ranges
755    /// and no nesting here. An empty filter means no filtering.
756    pub fn vector_search_filtered(
757        &mut self,
758        collection: &str,
759        vector: &[f32],
760        top_k: usize,
761        filter: &BTreeMap<String, Value>,
762    ) -> Result<VectorSearchResult> {
763        let filter = if filter.is_empty() {
764            Value::Null
765        } else {
766            json!(filter)
767        };
768        let response = self.send(op(json!({"Search": {
769            "collection": collection,
770            "vector": vector,
771            "top_k": top_k,
772            "filter": filter,
773        }})))?;
774        response.decode("Search")
775    }
776
777    /// One page of a collection's vectors, ordered by id.
778    ///
779    /// `limit` of `None` asks for the server's default (500, capped at 1000);
780    /// `offset` skips that many.
781    pub fn vector_list_vectors(
782        &mut self,
783        collection: &str,
784        limit: Option<usize>,
785        offset: Option<usize>,
786    ) -> Result<VectorPage> {
787        let mut body = Map::new();
788        body.insert("collection".into(), json!(collection));
789        if let Some(limit) = limit.filter(|l| *l > 0) {
790            body.insert("limit".into(), json!(limit));
791        }
792        if let Some(offset) = offset.filter(|o| *o > 0) {
793            body.insert("offset".into(), json!(offset));
794        }
795        let response = self.send(op(json!({"ListVectors": Value::Object(body)})))?;
796        response.decode("ListVectors")
797    }
798
799    /// Embed texts with the server's built-in embedder and return the
800    /// vectors. Nothing is stored.
801    pub fn vector_embed(&mut self, texts: &[&str], dimension: usize) -> Result<Embedding> {
802        let response = self.send(op(json!({"Embed": {
803            "texts": texts,
804            "dimension": dimension,
805        }})))?;
806        response.decode("Embed")
807    }
808
809    /// Embed each text at the collection's dimension and upsert the vectors as
810    /// one atomic batch, returning how many were written.
811    pub fn vector_upsert_text(&mut self, collection: &str, items: &[TextUpsert]) -> Result<i64> {
812        let response = self.send(op(json!({"UpsertText": {
813            "collection": collection,
814            "items": items,
815        }})))?;
816        Ok(response.decode::<Upserted>("UpsertText")?.upserted)
817    }
818
819    /// Point an alias at a collection; `None` removes it. Every operation
820    /// naming the alias acts on the collection, so a rebuilt collection can
821    /// replace the old one by moving the alias.
822    pub fn vector_set_alias(&mut self, alias: &str, collection: Option<&str>) -> Result<()> {
823        self.send(op(json!({"SetAlias": {
824            "alias": alias,
825            "collection": collection,
826        }})))
827        .map(|_| ())
828    }
829
830    /// Every alias and the collection it names.
831    pub fn vector_list_aliases(&mut self) -> Result<Vec<VectorAlias>> {
832        #[derive(Deserialize)]
833        struct Aliases {
834            #[serde(default)]
835            aliases: Vec<VectorAlias>,
836        }
837        let response = self.send(op(Value::String("ListAliases".into())))?;
838        let out: Aliases = response.decode("ListAliases")?;
839        Ok(out.aliases)
840    }
841
842    /// Declare a named vector on a collection: a second vector per item, with
843    /// its own dimension and metric.
844    pub fn vector_add_named_vector(
845        &mut self,
846        collection: &str,
847        name: &str,
848        dimension: usize,
849        metric: VectorMetric,
850    ) -> Result<()> {
851        self.send(op(json!({"AddNamedVector": {
852            "collection": collection,
853            "name": name,
854            "dimension": dimension,
855            "metric": metric,
856        }})))
857        .map(|_| ())
858    }
859
860    /// Write an item's vectors by name, each with the same metadata. The name
861    /// `default` is the collection's own vector.
862    pub fn vector_upsert_named(
863        &mut self,
864        collection: &str,
865        id: &str,
866        vectors: &BTreeMap<String, Vec<f32>>,
867        metadata: Option<&Map<String, Value>>,
868    ) -> Result<()> {
869        self.send(op(json!({"UpsertNamed": {
870            "collection": collection,
871            "id": id,
872            "vectors": vectors,
873            "metadata": metadata,
874        }})))
875        .map(|_| ())
876    }
877
878    /// Write many vectors as one atomic batch, returning how many were
879    /// written. One bad item refuses the whole batch, and so does a repeated
880    /// id.
881    pub fn vector_upsert_many(&mut self, collection: &str, items: &[VectorUpsert]) -> Result<i64> {
882        let response = self.send(op(json!({"UpsertMany": {
883            "collection": collection,
884            "items": items,
885        }})))?;
886        Ok(response.decode::<Upserted>("UpsertMany")?.upserted)
887    }
888
889    /// Delete many vectors by id as one atomic batch, returning how many of
890    /// the ids were present.
891    pub fn vector_delete_many(&mut self, collection: &str, ids: &[&str]) -> Result<i64> {
892        let response = self.send(op(json!({"DeleteMany": {
893            "collection": collection,
894            "ids": ids,
895        }})))?;
896        Ok(response.decode::<DeletedCount>("DeleteMany")?.deleted)
897    }
898
899    /// Run several query vectors against one collection in one round trip.
900    ///
901    /// One result list per query, in the order asked, each what a single
902    /// search with the same arguments returns. Not a snapshot.
903    pub fn vector_search_batch(
904        &mut self,
905        collection: &str,
906        vectors: &[Vec<f32>],
907        top_k: usize,
908        options: &VectorSearchOptions,
909    ) -> Result<Vec<Vec<VectorHit>>> {
910        #[derive(Deserialize)]
911        struct Batch {
912            #[serde(default)]
913            results: Vec<Vec<VectorHit>>,
914        }
915        let mut body = Map::new();
916        body.insert("collection".into(), json!(collection));
917        body.insert("vectors".into(), json!(vectors));
918        body.insert("top_k".into(), json!(top_k));
919        options.write(&mut body);
920        let response = self.send(op(json!({"SearchBatch": Value::Object(body)})))?;
921        let out: Batch = response.decode("SearchBatch")?;
922        Ok(out.results)
923    }
924
925    /// Search with several query vectors and return one fused ranking.
926    pub fn vector_search_multi(
927        &mut self,
928        collection: &str,
929        vectors: &[Vec<f32>],
930        top_k: usize,
931        fusion: FusionMethod,
932        filter: Option<&Value>,
933    ) -> Result<FusedSearchResult> {
934        let mut body = Map::new();
935        body.insert("collection".into(), json!(collection));
936        body.insert("vectors".into(), json!(vectors));
937        body.insert("top_k".into(), json!(top_k));
938        body.insert("fusion".into(), json!(fusion));
939        if let Some(filter) = filter {
940            body.insert("filter".into(), filter.clone());
941        }
942        let response = self.send(op(json!({"SearchMulti": Value::Object(body)})))?;
943        response.decode("SearchMulti")
944    }
945
946    /// Search with a query built from examples: the sum of the `positive`
947    /// examples minus the sum of the `negative` ones.
948    ///
949    /// With `exclude_examples`, the examples named by id are left out of the
950    /// results — which is the server's default.
951    pub fn vector_search_by_examples(
952        &mut self,
953        collection: &str,
954        positive: &[VectorExample],
955        negative: &[VectorExample],
956        top_k: usize,
957        exclude_examples: bool,
958        options: &VectorSearchOptions,
959    ) -> Result<VectorSearchResult> {
960        let mut body = Map::new();
961        body.insert("collection".into(), json!(collection));
962        body.insert("positive".into(), json!(positive));
963        body.insert("negative".into(), json!(negative));
964        body.insert("top_k".into(), json!(top_k));
965        body.insert("exclude_examples".into(), json!(exclude_examples));
966        options.write(&mut body);
967        let response = self.send(op(json!({"SearchByExamples": Value::Object(body)})))?;
968        response.decode("SearchByExamples")
969    }
970
971    /// Delete every vector whose metadata matches `filter`, as one atomic
972    /// batch, returning how many went. An empty filter is refused by the
973    /// server: emptying a collection is [`Client::vector_drop_collection`]'s
974    /// job.
975    pub fn vector_delete_by_filter(&mut self, collection: &str, filter: &Value) -> Result<i64> {
976        let response = self.send(op(json!({"DeleteByFilter": {
977            "collection": collection,
978            "filter": filter,
979        }})))?;
980        Ok(response.decode::<DeletedCount>("DeleteByFilter")?.deleted)
981    }
982
983    /// Merge `set` into the metadata of every vector matching `filter`, as one
984    /// atomic batch, returning how many changed. Vectors are never touched.
985    pub fn vector_update_by_filter(
986        &mut self,
987        collection: &str,
988        filter: &Value,
989        set: &Map<String, Value>,
990    ) -> Result<i64> {
991        #[derive(Deserialize)]
992        struct Updated {
993            #[serde(default)]
994            updated: i64,
995        }
996        let response = self.send(op(json!({"UpdateByFilter": {
997            "collection": collection,
998            "filter": filter,
999            "set": set,
1000        }})))?;
1001        Ok(response.decode::<Updated>("UpdateByFilter")?.updated)
1002    }
1003
1004    /// How many vectors a collection holds — or, with a filter, how many match
1005    /// it. A count, not an estimate.
1006    pub fn vector_count(&mut self, collection: &str, filter: Option<&Value>) -> Result<i64> {
1007        #[derive(Deserialize)]
1008        struct Count {
1009            #[serde(default)]
1010            count: i64,
1011        }
1012        let mut body = Map::new();
1013        body.insert("collection".into(), json!(collection));
1014        if let Some(filter) = filter {
1015            body.insert("filter".into(), filter.clone());
1016        }
1017        let response = self.send(op(json!({"Count": Value::Object(body)})))?;
1018        Ok(response.decode::<Count>("Count")?.count)
1019    }
1020
1021    /// The partitions of a collection: each distinct value of `field` with its
1022    /// vector count. `field` is `_namespace` or `_tenant`; `None` leaves it to
1023    /// the server, which uses `_namespace`.
1024    pub fn vector_list_partitions(
1025        &mut self,
1026        collection: &str,
1027        field: Option<&str>,
1028    ) -> Result<VectorPartitions> {
1029        let mut body = Map::new();
1030        body.insert("collection".into(), json!(collection));
1031        if let Some(field) = field {
1032            body.insert("field".into(), json!(field));
1033        }
1034        let response = self.send(op(json!({"ListPartitions": Value::Object(body)})))?;
1035        response.decode("ListPartitions")
1036    }
1037
1038    /// Rename a collection, moving every vector with it in one atomic batch,
1039    /// and return how many moved. Refused when `to` already exists.
1040    pub fn vector_rename_collection(&mut self, collection: &str, to: &str) -> Result<i64> {
1041        #[derive(Deserialize)]
1042        struct Moved {
1043            #[serde(default)]
1044            moved: i64,
1045        }
1046        let response = self.send(op(json!({"RenameCollection": {
1047            "collection": collection,
1048            "to": to,
1049        }})))?;
1050        Ok(response.decode::<Moved>("RenameCollection")?.moved)
1051    }
1052
1053    /// Copy a collection — configuration and every vector — to a new name, and
1054    /// return how many vectors were copied. Refused when `to` exists.
1055    pub fn vector_copy_collection(&mut self, collection: &str, to: &str) -> Result<i64> {
1056        #[derive(Deserialize)]
1057        struct Copied {
1058            #[serde(default)]
1059            copied: i64,
1060        }
1061        let response = self.send(op(json!({"CopyCollection": {
1062            "collection": collection,
1063            "to": to,
1064        }})))?;
1065        Ok(response.decode::<Copied>("CopyCollection")?.copied)
1066    }
1067
1068    /// Rebuild the collection's search index from its stored vectors and
1069    /// persist it, clearing every tombstone.
1070    pub fn vector_rebuild_index(&mut self, collection: &str) -> Result<VectorIndexRebuild> {
1071        let response = self.send(op(json!({"RebuildIndex": {"collection": collection}})))?;
1072        response.decode("RebuildIndex")
1073    }
1074
1075    /// The shape of the collection's search index. Loads or builds the index
1076    /// as a search would.
1077    pub fn vector_index_stats(&mut self, collection: &str) -> Result<VectorIndexStats> {
1078        let response = self.send(op(json!({"IndexStats": {"collection": collection}})))?;
1079        response.decode("IndexStats")
1080    }
1081
1082    /// Measure the index's recall: sampled stored vectors are each searched
1083    /// through the index and by an exact scan, and the overlap is reported.
1084    /// Costs one exact scan per sample.
1085    pub fn vector_measure_recall(
1086        &mut self,
1087        collection: &str,
1088        options: &RecallOptions,
1089    ) -> Result<VectorRecall> {
1090        let mut body = Map::new();
1091        body.insert("collection".into(), json!(collection));
1092        if let Some(sample) = options.sample {
1093            body.insert("sample".into(), json!(sample));
1094        }
1095        if let Some(top_k) = options.top_k {
1096            body.insert("top_k".into(), json!(top_k));
1097        }
1098        if let Some(ef_search) = options.ef_search {
1099            body.insert("ef_search".into(), json!(ef_search));
1100        }
1101        let response = self.send(op(json!({"MeasureRecall": Value::Object(body)})))?;
1102        response.decode("MeasureRecall")
1103    }
1104
1105    /// Change a collection's configuration, and return it as it now stands.
1106    pub fn vector_set_collection_config(
1107        &mut self,
1108        collection: &str,
1109        change: &VectorConfigChange,
1110    ) -> Result<VectorConfig> {
1111        if *change == VectorConfigChange::default() {
1112            return Err(Error::invalid(
1113                "a configuration change must set at least one of `normalize`, `m`, `ef_construction` or `multi_tenant`",
1114            ));
1115        }
1116        let mut body = Map::new();
1117        body.insert("collection".into(), json!(collection));
1118        if let Some(normalize) = change.normalize {
1119            body.insert("normalize".into(), json!(normalize));
1120        }
1121        if let Some(m) = change.m {
1122            body.insert("m".into(), json!(m));
1123        }
1124        if let Some(ef_construction) = change.ef_construction {
1125            body.insert("ef_construction".into(), json!(ef_construction));
1126        }
1127        if let Some(multi_tenant) = change.multi_tenant {
1128            body.insert("multi_tenant".into(), json!(multi_tenant));
1129        }
1130        let response = self.send(op(json!({"SetCollectionConfig": Value::Object(body)})))?;
1131        response.decode("SetCollectionConfig")
1132    }
1133}
1134
1135#[derive(Deserialize)]
1136struct Upserted {
1137    #[serde(default)]
1138    upserted: i64,
1139}
1140
1141#[derive(Deserialize)]
1142struct DeletedCount {
1143    #[serde(default)]
1144    deleted: i64,
1145}
1146
1147#[cfg(test)]
1148mod tests {
1149    use super::*;
1150
1151    #[test]
1152    fn metrics_and_quantizations_use_the_servers_spelling() {
1153        assert_eq!(
1154            serde_json::to_value(VectorMetric::Cosine).unwrap(),
1155            "cosine"
1156        );
1157        assert_eq!(serde_json::to_value(VectorMetric::Dot).unwrap(), "dot");
1158        assert_eq!(serde_json::to_value(VectorMetric::L2).unwrap(), "l2");
1159        assert_eq!(
1160            serde_json::to_value(VectorQuantization::None).unwrap(),
1161            "none"
1162        );
1163        assert_eq!(
1164            serde_json::to_value(VectorQuantization::Int8).unwrap(),
1165            "int8"
1166        );
1167    }
1168
1169    #[test]
1170    fn metrics_round_trip_from_the_wire() {
1171        let metric: VectorMetric = serde_json::from_str("\"l2\"").unwrap();
1172        assert_eq!(metric, VectorMetric::L2);
1173        assert_eq!(VectorMetric::default(), VectorMetric::Cosine);
1174        assert_eq!(VectorQuantization::default(), VectorQuantization::None);
1175    }
1176
1177    #[test]
1178    fn a_hit_decodes_with_its_metadata() {
1179        let hit: VectorHit =
1180            serde_json::from_value(json!({"id": "a", "score": -0.02, "metadata": {"kind": "doc"}}))
1181                .unwrap();
1182        assert_eq!(hit.id, "a");
1183        assert_eq!(hit.metadata["kind"], "doc");
1184    }
1185}
tricoredb/vector.rs

tricoredb/
vector.rs
