chroma_types/
api_types.rs

1use crate::collection_configuration::InternalCollectionConfiguration;
2use crate::collection_configuration::InternalUpdateCollectionConfiguration;
3use crate::error::QueryConversionError;
4use crate::operator::GetResult;
5use crate::operator::Key;
6use crate::operator::KnnBatchResult;
7use crate::operator::KnnProjectionRecord;
8use crate::operator::ProjectionRecord;
9use crate::operator::SearchResult;
10use crate::plan::PlanToProtoError;
11use crate::plan::SearchPayload;
12use crate::validators::{
13    validate_metadata_vec, validate_name, validate_non_empty_collection_update_metadata,
14    validate_optional_metadata, validate_schema, validate_update_metadata_vec,
15};
16use crate::Collection;
17use crate::CollectionConfigurationToInternalConfigurationError;
18use crate::CollectionConversionError;
19use crate::CollectionUuid;
20use crate::DistributedSpannParametersFromSegmentError;
21use crate::HnswParametersFromSegmentError;
22use crate::InternalSchema;
23use crate::Metadata;
24use crate::SchemaError;
25use crate::SegmentConversionError;
26use crate::SegmentScopeConversionError;
27use crate::UpdateMetadata;
28use crate::Where;
29use chroma_error::ChromaValidationError;
30use chroma_error::{ChromaError, ErrorCodes};
31use serde::Deserialize;
32use serde::Serialize;
33use std::time::SystemTimeError;
34use thiserror::Error;
35use tonic::Status;
36use uuid::Uuid;
37use validator::Validate;
38use validator::ValidationError;
39
40#[cfg(feature = "pyo3")]
41use pyo3::types::PyAnyMethods;
42
43#[derive(Debug, Error)]
44pub enum GetSegmentsError {
45    #[error("Could not parse segment")]
46    SegmentConversion(#[from] SegmentConversionError),
47    #[error("Unknown segment scope")]
48    UnknownScope(#[from] SegmentScopeConversionError),
49    #[error(transparent)]
50    Internal(#[from] Box<dyn ChromaError>),
51}
52
53impl ChromaError for GetSegmentsError {
54    fn code(&self) -> ErrorCodes {
55        match self {
56            GetSegmentsError::SegmentConversion(_) => ErrorCodes::Internal,
57            GetSegmentsError::UnknownScope(_) => ErrorCodes::Internal,
58            GetSegmentsError::Internal(err) => err.code(),
59        }
60    }
61}
62
63#[derive(Debug, Error)]
64pub enum GetCollectionWithSegmentsError {
65    #[error("Failed to convert proto collection")]
66    CollectionConversionError(#[from] CollectionConversionError),
67    #[error("Duplicate segment")]
68    DuplicateSegment,
69    #[error("Missing field: [{0}]")]
70    Field(String),
71    #[error("Failed to convert proto segment")]
72    SegmentConversionError(#[from] SegmentConversionError),
73    #[error("Failed to get segments")]
74    GetSegmentsError(#[from] GetSegmentsError),
75    #[error("Grpc error: {0}")]
76    Grpc(#[from] Status),
77    #[error("Collection [{0}] does not exist.")]
78    NotFound(String),
79    #[error(transparent)]
80    Internal(#[from] Box<dyn ChromaError>),
81}
82
83impl ChromaError for GetCollectionWithSegmentsError {
84    fn code(&self) -> ErrorCodes {
85        match self {
86            GetCollectionWithSegmentsError::CollectionConversionError(
87                collection_conversion_error,
88            ) => collection_conversion_error.code(),
89            GetCollectionWithSegmentsError::DuplicateSegment => ErrorCodes::Internal,
90            GetCollectionWithSegmentsError::Field(_) => ErrorCodes::FailedPrecondition,
91            GetCollectionWithSegmentsError::SegmentConversionError(segment_conversion_error) => {
92                segment_conversion_error.code()
93            }
94            GetCollectionWithSegmentsError::Grpc(status) => status.code().into(),
95            GetCollectionWithSegmentsError::GetSegmentsError(get_segments_error) => {
96                get_segments_error.code()
97            }
98            GetCollectionWithSegmentsError::NotFound(_) => ErrorCodes::NotFound,
99            GetCollectionWithSegmentsError::Internal(err) => err.code(),
100        }
101    }
102
103    fn should_trace_error(&self) -> bool {
104        if let Self::Grpc(status) = self {
105            status.code() != ErrorCodes::NotFound.into()
106        } else {
107            true
108        }
109    }
110}
111
112#[derive(Debug, Error)]
113pub enum BatchGetCollectionVersionFilePathsError {
114    #[error("Grpc error: {0}")]
115    Grpc(#[from] Status),
116    #[error("Could not parse UUID from string {1}: {0}")]
117    Uuid(uuid::Error, String),
118}
119
120impl ChromaError for BatchGetCollectionVersionFilePathsError {
121    fn code(&self) -> ErrorCodes {
122        match self {
123            BatchGetCollectionVersionFilePathsError::Grpc(status) => status.code().into(),
124            BatchGetCollectionVersionFilePathsError::Uuid(_, _) => ErrorCodes::InvalidArgument,
125        }
126    }
127}
128
129#[derive(Debug, Error)]
130pub enum BatchGetCollectionSoftDeleteStatusError {
131    #[error("Grpc error: {0}")]
132    Grpc(#[from] Status),
133    #[error("Could not parse UUID from string {1}: {0}")]
134    Uuid(uuid::Error, String),
135}
136
137impl ChromaError for BatchGetCollectionSoftDeleteStatusError {
138    fn code(&self) -> ErrorCodes {
139        match self {
140            BatchGetCollectionSoftDeleteStatusError::Grpc(status) => status.code().into(),
141            BatchGetCollectionSoftDeleteStatusError::Uuid(_, _) => ErrorCodes::InvalidArgument,
142        }
143    }
144}
145
146#[derive(Serialize)]
147#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
148pub struct ResetResponse {}
149
150#[derive(Debug, Error)]
151pub enum ResetError {
152    #[error(transparent)]
153    Cache(Box<dyn ChromaError>),
154    #[error(transparent)]
155    Internal(#[from] Box<dyn ChromaError>),
156    #[error("Reset is disabled by config")]
157    NotAllowed,
158}
159
160impl ChromaError for ResetError {
161    fn code(&self) -> ErrorCodes {
162        match self {
163            ResetError::Cache(err) => err.code(),
164            ResetError::Internal(err) => err.code(),
165            ResetError::NotAllowed => ErrorCodes::PermissionDenied,
166        }
167    }
168}
169
170#[derive(Serialize)]
171#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
172pub struct ChecklistResponse {
173    pub max_batch_size: u32,
174    pub supports_base64_encoding: bool,
175}
176
177#[derive(Debug, Error)]
178#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
179pub enum HeartbeatError {
180    #[error("system time error: {0}")]
181    CouldNotGetTime(String),
182}
183
184impl From<SystemTimeError> for HeartbeatError {
185    fn from(err: SystemTimeError) -> Self {
186        HeartbeatError::CouldNotGetTime(err.to_string())
187    }
188}
189
190impl ChromaError for HeartbeatError {
191    fn code(&self) -> ErrorCodes {
192        ErrorCodes::Internal
193    }
194}
195
196#[non_exhaustive]
197#[derive(Serialize, Validate, Deserialize)]
198#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
199pub struct CreateTenantRequest {
200    #[validate(length(min = 3))]
201    pub name: String,
202}
203
204impl CreateTenantRequest {
205    pub fn try_new(name: String) -> Result<Self, ChromaValidationError> {
206        let request = Self { name };
207        request.validate().map_err(ChromaValidationError::from)?;
208        Ok(request)
209    }
210}
211
212#[derive(Serialize, Deserialize)]
213#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
214pub struct CreateTenantResponse {}
215
216#[derive(Debug, Error)]
217pub enum CreateTenantError {
218    #[error("Tenant [{0}] already exists")]
219    AlreadyExists(String),
220    #[error(transparent)]
221    Internal(#[from] Box<dyn ChromaError>),
222}
223
224impl ChromaError for CreateTenantError {
225    fn code(&self) -> ErrorCodes {
226        match self {
227            CreateTenantError::AlreadyExists(_) => ErrorCodes::AlreadyExists,
228            CreateTenantError::Internal(err) => err.code(),
229        }
230    }
231}
232
233#[non_exhaustive]
234#[derive(Validate, Serialize)]
235#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
236pub struct GetTenantRequest {
237    pub name: String,
238}
239
240impl GetTenantRequest {
241    pub fn try_new(name: String) -> Result<Self, ChromaValidationError> {
242        let request = Self { name };
243        request.validate().map_err(ChromaValidationError::from)?;
244        Ok(request)
245    }
246}
247
248#[derive(Serialize)]
249#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
250#[cfg_attr(feature = "pyo3", pyo3::pyclass)]
251pub struct GetTenantResponse {
252    pub name: String,
253    pub resource_name: Option<String>,
254}
255
256#[cfg(feature = "pyo3")]
257#[pyo3::pymethods]
258impl GetTenantResponse {
259    #[getter]
260    pub fn name(&self) -> &String {
261        &self.name
262    }
263
264    #[getter]
265    pub fn resource_name(&self) -> Option<String> {
266        self.resource_name.clone()
267    }
268}
269
270#[derive(Debug, Error)]
271pub enum GetTenantError {
272    #[error(transparent)]
273    Internal(#[from] Box<dyn ChromaError>),
274    #[error("Tenant [{0}] not found")]
275    NotFound(String),
276}
277
278impl ChromaError for GetTenantError {
279    fn code(&self) -> ErrorCodes {
280        match self {
281            GetTenantError::Internal(err) => err.code(),
282            GetTenantError::NotFound(_) => ErrorCodes::NotFound,
283        }
284    }
285}
286
287#[non_exhaustive]
288#[derive(Validate, Serialize)]
289#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
290pub struct UpdateTenantRequest {
291    pub tenant_id: String,
292    pub resource_name: String,
293}
294
295impl UpdateTenantRequest {
296    pub fn try_new(
297        tenant_id: String,
298        resource_name: String,
299    ) -> Result<Self, ChromaValidationError> {
300        let request = Self {
301            tenant_id,
302            resource_name,
303        };
304        request.validate().map_err(ChromaValidationError::from)?;
305        Ok(request)
306    }
307}
308
309#[derive(Serialize)]
310#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
311#[cfg_attr(feature = "pyo3", pyo3::pyclass)]
312pub struct UpdateTenantResponse {}
313
314#[cfg(feature = "pyo3")]
315#[pyo3::pymethods]
316impl UpdateTenantResponse {}
317
318#[derive(Error, Debug)]
319pub enum UpdateTenantError {
320    #[error("Failed to set resource name")]
321    FailedToSetResourceName(#[from] tonic::Status),
322    #[error(transparent)]
323    Internal(#[from] Box<dyn ChromaError>),
324    #[error("Tenant [{0}] not found")]
325    NotFound(String),
326}
327
328impl ChromaError for UpdateTenantError {
329    fn code(&self) -> ErrorCodes {
330        match self {
331            UpdateTenantError::FailedToSetResourceName(_) => ErrorCodes::AlreadyExists,
332            UpdateTenantError::Internal(err) => err.code(),
333            UpdateTenantError::NotFound(_) => ErrorCodes::NotFound,
334        }
335    }
336}
337
338#[non_exhaustive]
339#[derive(Validate, Serialize)]
340#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
341pub struct CreateDatabaseRequest {
342    pub database_id: Uuid,
343    pub tenant_id: String,
344    #[validate(length(min = 3))]
345    pub database_name: String,
346}
347
348impl CreateDatabaseRequest {
349    pub fn try_new(
350        tenant_id: String,
351        database_name: String,
352    ) -> Result<Self, ChromaValidationError> {
353        let database_id = Uuid::new_v4();
354        let request = Self {
355            database_id,
356            tenant_id,
357            database_name,
358        };
359        request.validate().map_err(ChromaValidationError::from)?;
360        Ok(request)
361    }
362}
363
364#[derive(Serialize)]
365#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
366pub struct CreateDatabaseResponse {}
367
368#[derive(Error, Debug)]
369pub enum CreateDatabaseError {
370    #[error("Database [{0}] already exists")]
371    AlreadyExists(String),
372    #[error(transparent)]
373    Internal(#[from] Box<dyn ChromaError>),
374}
375
376impl ChromaError for CreateDatabaseError {
377    fn code(&self) -> ErrorCodes {
378        match self {
379            CreateDatabaseError::AlreadyExists(_) => ErrorCodes::AlreadyExists,
380            CreateDatabaseError::Internal(status) => status.code(),
381        }
382    }
383}
384
385#[derive(Serialize, Deserialize, Debug, Clone, Default)]
386#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
387#[cfg_attr(feature = "pyo3", pyo3::pyclass)]
388pub struct Database {
389    pub id: Uuid,
390    pub name: String,
391    pub tenant: String,
392}
393
394#[cfg(feature = "pyo3")]
395#[pyo3::pymethods]
396impl Database {
397    #[getter]
398    fn id<'py>(&self, py: pyo3::Python<'py>) -> pyo3::PyResult<pyo3::Bound<'py, pyo3::PyAny>> {
399        let res = pyo3::prelude::PyModule::import(py, "uuid")?
400            .getattr("UUID")?
401            .call1((self.id.to_string(),))?;
402        Ok(res)
403    }
404
405    #[getter]
406    pub fn name(&self) -> &str {
407        &self.name
408    }
409
410    #[getter]
411    pub fn tenant(&self) -> &str {
412        &self.tenant
413    }
414}
415
416#[non_exhaustive]
417#[derive(Validate, Serialize)]
418#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
419pub struct ListDatabasesRequest {
420    pub tenant_id: String,
421    pub limit: Option<u32>,
422    pub offset: u32,
423}
424
425impl ListDatabasesRequest {
426    pub fn try_new(
427        tenant_id: String,
428        limit: Option<u32>,
429        offset: u32,
430    ) -> Result<Self, ChromaValidationError> {
431        let request = Self {
432            tenant_id,
433            limit,
434            offset,
435        };
436        request.validate().map_err(ChromaValidationError::from)?;
437        Ok(request)
438    }
439}
440
441pub type ListDatabasesResponse = Vec<Database>;
442
443#[derive(Debug, Error)]
444pub enum ListDatabasesError {
445    #[error(transparent)]
446    Internal(#[from] Box<dyn ChromaError>),
447    #[error("Invalid database id [{0}]")]
448    InvalidID(String),
449}
450
451impl ChromaError for ListDatabasesError {
452    fn code(&self) -> ErrorCodes {
453        match self {
454            ListDatabasesError::Internal(status) => status.code(),
455            ListDatabasesError::InvalidID(_) => ErrorCodes::InvalidArgument,
456        }
457    }
458}
459
460#[non_exhaustive]
461#[derive(Validate, Serialize)]
462#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
463pub struct GetDatabaseRequest {
464    pub tenant_id: String,
465    pub database_name: String,
466}
467
468impl GetDatabaseRequest {
469    pub fn try_new(
470        tenant_id: String,
471        database_name: String,
472    ) -> Result<Self, ChromaValidationError> {
473        let request = Self {
474            tenant_id,
475            database_name,
476        };
477        request.validate().map_err(ChromaValidationError::from)?;
478        Ok(request)
479    }
480}
481
482pub type GetDatabaseResponse = Database;
483
484#[derive(Error, Debug)]
485pub enum GetDatabaseError {
486    #[error(transparent)]
487    Internal(#[from] Box<dyn ChromaError>),
488    #[error("Invalid database id [{0}]")]
489    InvalidID(String),
490    #[error("Database [{0}] not found. Are you sure it exists?")]
491    NotFound(String),
492}
493
494impl ChromaError for GetDatabaseError {
495    fn code(&self) -> ErrorCodes {
496        match self {
497            GetDatabaseError::Internal(err) => err.code(),
498            GetDatabaseError::InvalidID(_) => ErrorCodes::InvalidArgument,
499            GetDatabaseError::NotFound(_) => ErrorCodes::NotFound,
500        }
501    }
502}
503
504#[non_exhaustive]
505#[derive(Validate, Serialize)]
506#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
507pub struct DeleteDatabaseRequest {
508    pub tenant_id: String,
509    pub database_name: String,
510}
511
512impl DeleteDatabaseRequest {
513    pub fn try_new(
514        tenant_id: String,
515        database_name: String,
516    ) -> Result<Self, ChromaValidationError> {
517        let request = Self {
518            tenant_id,
519            database_name,
520        };
521        request.validate().map_err(ChromaValidationError::from)?;
522        Ok(request)
523    }
524}
525
526#[derive(Serialize)]
527#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
528pub struct DeleteDatabaseResponse {}
529
530#[derive(Debug, Error)]
531pub enum DeleteDatabaseError {
532    #[error(transparent)]
533    Internal(#[from] Box<dyn ChromaError>),
534    #[error("Invalid database id [{0}]")]
535    InvalidID(String),
536    #[error("Database [{0}] not found")]
537    NotFound(String),
538}
539
540impl ChromaError for DeleteDatabaseError {
541    fn code(&self) -> ErrorCodes {
542        match self {
543            DeleteDatabaseError::Internal(err) => err.code(),
544            DeleteDatabaseError::InvalidID(_) => ErrorCodes::InvalidArgument,
545            DeleteDatabaseError::NotFound(_) => ErrorCodes::NotFound,
546        }
547    }
548}
549
550#[derive(Debug, Error)]
551pub enum FinishDatabaseDeletionError {
552    #[error(transparent)]
553    Internal(#[from] Box<dyn ChromaError>),
554}
555
556impl ChromaError for FinishDatabaseDeletionError {
557    fn code(&self) -> ErrorCodes {
558        match self {
559            FinishDatabaseDeletionError::Internal(err) => err.code(),
560        }
561    }
562}
563
564#[non_exhaustive]
565#[derive(Validate, Debug, Serialize)]
566#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
567pub struct ListCollectionsRequest {
568    pub tenant_id: String,
569    pub database_name: String,
570    pub limit: Option<u32>,
571    pub offset: u32,
572}
573
574impl ListCollectionsRequest {
575    pub fn try_new(
576        tenant_id: String,
577        database_name: String,
578        limit: Option<u32>,
579        offset: u32,
580    ) -> Result<Self, ChromaValidationError> {
581        let request = Self {
582            tenant_id,
583            database_name,
584            limit,
585            offset,
586        };
587        request.validate().map_err(ChromaValidationError::from)?;
588        Ok(request)
589    }
590}
591
592pub type ListCollectionsResponse = Vec<Collection>;
593
594#[non_exhaustive]
595#[derive(Validate, Serialize)]
596#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
597pub struct CountCollectionsRequest {
598    pub tenant_id: String,
599    pub database_name: String,
600}
601
602impl CountCollectionsRequest {
603    pub fn try_new(
604        tenant_id: String,
605        database_name: String,
606    ) -> Result<Self, ChromaValidationError> {
607        let request = Self {
608            tenant_id,
609            database_name,
610        };
611        request.validate().map_err(ChromaValidationError::from)?;
612        Ok(request)
613    }
614}
615
616pub type CountCollectionsResponse = u32;
617
618#[non_exhaustive]
619#[derive(Validate, Clone, Serialize)]
620#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
621pub struct GetCollectionRequest {
622    pub tenant_id: String,
623    pub database_name: String,
624    pub collection_name: String,
625}
626
627impl GetCollectionRequest {
628    pub fn try_new(
629        tenant_id: String,
630        database_name: String,
631        collection_name: String,
632    ) -> Result<Self, ChromaValidationError> {
633        let request = Self {
634            tenant_id,
635            database_name,
636            collection_name,
637        };
638        request.validate().map_err(ChromaValidationError::from)?;
639        Ok(request)
640    }
641}
642
643pub type GetCollectionResponse = Collection;
644
645#[derive(Debug, Error)]
646pub enum GetCollectionError {
647    #[error("Failed to reconcile schema: {0}")]
648    InvalidSchema(#[from] SchemaError),
649    #[error(transparent)]
650    Internal(#[from] Box<dyn ChromaError>),
651    #[error("Collection [{0}] does not exist")]
652    NotFound(String),
653}
654
655impl ChromaError for GetCollectionError {
656    fn code(&self) -> ErrorCodes {
657        match self {
658            GetCollectionError::InvalidSchema(e) => e.code(),
659            GetCollectionError::Internal(err) => err.code(),
660            GetCollectionError::NotFound(_) => ErrorCodes::NotFound,
661        }
662    }
663}
664
665#[non_exhaustive]
666#[derive(Clone, Debug, Validate, Serialize)]
667#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
668pub struct CreateCollectionRequest {
669    pub tenant_id: String,
670    pub database_name: String,
671    #[validate(custom(function = "validate_name"))]
672    pub name: String,
673    #[validate(custom(function = "validate_optional_metadata"))]
674    pub metadata: Option<Metadata>,
675    pub configuration: Option<InternalCollectionConfiguration>,
676    #[validate(custom(function = "validate_schema"))]
677    pub schema: Option<InternalSchema>,
678    pub get_or_create: bool,
679}
680
681impl CreateCollectionRequest {
682    pub fn try_new(
683        tenant_id: String,
684        database_name: String,
685        name: String,
686        metadata: Option<Metadata>,
687        configuration: Option<InternalCollectionConfiguration>,
688        schema: Option<InternalSchema>,
689        get_or_create: bool,
690    ) -> Result<Self, ChromaValidationError> {
691        let request = Self {
692            tenant_id,
693            database_name,
694            name,
695            metadata,
696            configuration,
697            schema,
698            get_or_create,
699        };
700        request.validate().map_err(ChromaValidationError::from)?;
701        Ok(request)
702    }
703}
704
705pub type CreateCollectionResponse = Collection;
706
707#[derive(Debug, Error)]
708pub enum CreateCollectionError {
709    #[error("Invalid HNSW parameters: {0}")]
710    InvalidHnswParameters(#[from] HnswParametersFromSegmentError),
711    #[error("Could not parse config: {0}")]
712    InvalidConfig(#[from] CollectionConfigurationToInternalConfigurationError),
713    #[error("Invalid Spann parameters: {0}")]
714    InvalidSpannParameters(#[from] DistributedSpannParametersFromSegmentError),
715    #[error("Collection [{0}] already exists")]
716    AlreadyExists(String),
717    #[error("Database [{0}] does not exist")]
718    DatabaseNotFound(String),
719    #[error("Could not fetch collections: {0}")]
720    Get(#[from] GetCollectionsError),
721    #[error("Could not deserialize configuration: {0}")]
722    Configuration(serde_json::Error),
723    #[error("Could not serialize schema: {0}")]
724    Schema(#[from] SchemaError),
725    #[error(transparent)]
726    Internal(#[from] Box<dyn ChromaError>),
727    #[error("The operation was aborted, {0}")]
728    Aborted(String),
729    #[error("SPANN is still in development. Not allowed to created spann indexes")]
730    SpannNotImplemented,
731    #[error("HNSW is not supported on this platform")]
732    HnswNotSupported,
733    #[error("Failed to parse db id")]
734    DatabaseIdParseError,
735    #[error("Failed to reconcile schema: {0}")]
736    InvalidSchema(String),
737}
738
739impl ChromaError for CreateCollectionError {
740    fn code(&self) -> ErrorCodes {
741        match self {
742            CreateCollectionError::InvalidHnswParameters(_) => ErrorCodes::InvalidArgument,
743            CreateCollectionError::InvalidConfig(_) => ErrorCodes::InvalidArgument,
744            CreateCollectionError::InvalidSpannParameters(_) => ErrorCodes::InvalidArgument,
745            CreateCollectionError::AlreadyExists(_) => ErrorCodes::AlreadyExists,
746            CreateCollectionError::DatabaseNotFound(_) => ErrorCodes::InvalidArgument,
747            CreateCollectionError::Get(err) => err.code(),
748            CreateCollectionError::Configuration(_) => ErrorCodes::Internal,
749            CreateCollectionError::Internal(err) => err.code(),
750            CreateCollectionError::Aborted(_) => ErrorCodes::Aborted,
751            CreateCollectionError::SpannNotImplemented => ErrorCodes::InvalidArgument,
752            CreateCollectionError::HnswNotSupported => ErrorCodes::InvalidArgument,
753            CreateCollectionError::DatabaseIdParseError => ErrorCodes::Internal,
754            CreateCollectionError::InvalidSchema(_) => ErrorCodes::InvalidArgument,
755            CreateCollectionError::Schema(e) => e.code(),
756        }
757    }
758}
759
760#[derive(Debug, Error)]
761pub enum CountCollectionsError {
762    #[error("Internal error in getting count")]
763    Internal,
764}
765
766impl ChromaError for CountCollectionsError {
767    fn code(&self) -> ErrorCodes {
768        match self {
769            CountCollectionsError::Internal => ErrorCodes::Internal,
770        }
771    }
772}
773
774#[derive(Debug, Error)]
775pub enum GetCollectionsError {
776    #[error(transparent)]
777    Internal(#[from] Box<dyn ChromaError>),
778    #[error("Could not deserialize configuration")]
779    Configuration(#[from] serde_json::Error),
780    #[error("Could not deserialize collection ID")]
781    CollectionId(#[from] uuid::Error),
782    #[error("Could not deserialize database ID")]
783    DatabaseId,
784}
785
786impl ChromaError for GetCollectionsError {
787    fn code(&self) -> ErrorCodes {
788        match self {
789            GetCollectionsError::Internal(err) => err.code(),
790            GetCollectionsError::Configuration(_) => ErrorCodes::Internal,
791            GetCollectionsError::CollectionId(_) => ErrorCodes::Internal,
792            GetCollectionsError::DatabaseId => ErrorCodes::Internal,
793        }
794    }
795}
796
797#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
798#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
799pub struct ChromaResourceName {
800    pub tenant_resource_name: String,
801    pub database_name: String,
802    pub collection_name: String,
803}
804#[non_exhaustive]
805#[derive(Clone, Serialize)]
806#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
807pub struct GetCollectionByCrnRequest {
808    pub parsed_crn: ChromaResourceName,
809}
810
811impl GetCollectionByCrnRequest {
812    pub fn try_new(crn: String) -> Result<Self, ChromaValidationError> {
813        let parsed_crn = parse_and_validate_crn(&crn)?;
814        Ok(Self { parsed_crn })
815    }
816}
817
818fn parse_and_validate_crn(crn: &str) -> Result<ChromaResourceName, ChromaValidationError> {
819    let mut parts = crn.splitn(4, ':');
820    if let (Some(p1), Some(p2), Some(p3), None) =
821        (parts.next(), parts.next(), parts.next(), parts.next())
822    {
823        if !p1.is_empty() && !p2.is_empty() && !p3.is_empty() {
824            return Ok(ChromaResourceName {
825                tenant_resource_name: p1.to_string(),
826                database_name: p2.to_string(),
827                collection_name: p3.to_string(),
828            });
829        }
830    }
831    let mut err = ValidationError::new("invalid_crn_format");
832    err.message = Some(
833        "CRN must be in the format <tenant_resource_name>:<database_name>:<collection_name> with non-empty parts"
834            .into(),
835    );
836    Err(ChromaValidationError::from(("crn", err)))
837}
838
839pub type GetCollectionByCrnResponse = Collection;
840
841#[derive(Debug, Error)]
842pub enum GetCollectionByCrnError {
843    #[error(transparent)]
844    Internal(#[from] Box<dyn ChromaError>),
845    #[error("Collection [{0}] does not exist")]
846    NotFound(String),
847}
848
849impl ChromaError for GetCollectionByCrnError {
850    fn code(&self) -> ErrorCodes {
851        match self {
852            GetCollectionByCrnError::Internal(err) => err.code(),
853            GetCollectionByCrnError::NotFound(_) => ErrorCodes::NotFound,
854        }
855    }
856}
857
858#[derive(Clone, Deserialize, Serialize, Debug)]
859#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
860pub enum CollectionMetadataUpdate {
861    ResetMetadata,
862    UpdateMetadata(UpdateMetadata),
863}
864
865#[non_exhaustive]
866#[derive(Clone, Validate, Debug, Serialize)]
867#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
868pub struct UpdateCollectionRequest {
869    pub collection_id: CollectionUuid,
870    #[validate(custom(function = "validate_name"))]
871    pub new_name: Option<String>,
872    #[validate(custom(function = "validate_non_empty_collection_update_metadata"))]
873    pub new_metadata: Option<CollectionMetadataUpdate>,
874    pub new_configuration: Option<InternalUpdateCollectionConfiguration>,
875}
876
877impl UpdateCollectionRequest {
878    pub fn try_new(
879        collection_id: CollectionUuid,
880        new_name: Option<String>,
881        new_metadata: Option<CollectionMetadataUpdate>,
882        new_configuration: Option<InternalUpdateCollectionConfiguration>,
883    ) -> Result<Self, ChromaValidationError> {
884        let request = Self {
885            collection_id,
886            new_name,
887            new_metadata,
888            new_configuration,
889        };
890        request.validate().map_err(ChromaValidationError::from)?;
891        Ok(request)
892    }
893}
894
895#[derive(Serialize)]
896#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
897pub struct UpdateCollectionResponse {}
898
899#[derive(Error, Debug)]
900pub enum UpdateCollectionError {
901    #[error("Collection [{0}] does not exist")]
902    NotFound(String),
903    #[error("Metadata reset unsupported")]
904    MetadataResetUnsupported,
905    #[error("Could not serialize configuration")]
906    Configuration(#[from] serde_json::Error),
907    #[error(transparent)]
908    Internal(#[from] Box<dyn ChromaError>),
909    #[error("Could not parse config: {0}")]
910    InvalidConfig(#[from] CollectionConfigurationToInternalConfigurationError),
911    #[error("SPANN is still in development. Not allowed to created spann indexes")]
912    SpannNotImplemented,
913}
914
915impl ChromaError for UpdateCollectionError {
916    fn code(&self) -> ErrorCodes {
917        match self {
918            UpdateCollectionError::NotFound(_) => ErrorCodes::NotFound,
919            UpdateCollectionError::MetadataResetUnsupported => ErrorCodes::InvalidArgument,
920            UpdateCollectionError::Configuration(_) => ErrorCodes::Internal,
921            UpdateCollectionError::Internal(err) => err.code(),
922            UpdateCollectionError::InvalidConfig(_) => ErrorCodes::InvalidArgument,
923            UpdateCollectionError::SpannNotImplemented => ErrorCodes::InvalidArgument,
924        }
925    }
926}
927
928#[non_exhaustive]
929#[derive(Clone, Validate, Serialize)]
930#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
931pub struct DeleteCollectionRequest {
932    pub tenant_id: String,
933    pub database_name: String,
934    pub collection_name: String,
935}
936
937impl DeleteCollectionRequest {
938    pub fn try_new(
939        tenant_id: String,
940        database_name: String,
941        collection_name: String,
942    ) -> Result<Self, ChromaValidationError> {
943        let request = Self {
944            tenant_id,
945            database_name,
946            collection_name,
947        };
948        request.validate().map_err(ChromaValidationError::from)?;
949        Ok(request)
950    }
951}
952
953#[derive(Serialize)]
954#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
955pub struct DeleteCollectionResponse {}
956
957#[derive(Error, Debug)]
958pub enum DeleteCollectionError {
959    #[error("Collection [{0}] does not exist")]
960    NotFound(String),
961    #[error(transparent)]
962    Validation(#[from] ChromaValidationError),
963    #[error(transparent)]
964    Get(#[from] GetCollectionError),
965    #[error(transparent)]
966    Internal(#[from] Box<dyn ChromaError>),
967}
968
969impl ChromaError for DeleteCollectionError {
970    fn code(&self) -> ErrorCodes {
971        match self {
972            DeleteCollectionError::Validation(err) => err.code(),
973            DeleteCollectionError::NotFound(_) => ErrorCodes::NotFound,
974            DeleteCollectionError::Get(err) => err.code(),
975            DeleteCollectionError::Internal(err) => err.code(),
976        }
977    }
978}
979
980#[non_exhaustive]
981#[derive(Clone, Validate, Serialize)]
982#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
983pub struct ForkCollectionRequest {
984    pub tenant_id: String,
985    pub database_name: String,
986    pub source_collection_id: CollectionUuid,
987    pub target_collection_name: String,
988}
989
990impl ForkCollectionRequest {
991    pub fn try_new(
992        tenant_id: String,
993        database_name: String,
994        source_collection_id: CollectionUuid,
995        target_collection_name: String,
996    ) -> Result<Self, ChromaValidationError> {
997        let request = Self {
998            tenant_id,
999            database_name,
1000            source_collection_id,
1001            target_collection_name,
1002        };
1003        request.validate().map_err(ChromaValidationError::from)?;
1004        Ok(request)
1005    }
1006}
1007
1008pub type ForkCollectionResponse = Collection;
1009
1010#[derive(Clone, Debug)]
1011pub struct ForkLogsResponse {
1012    pub compaction_offset: u64,
1013    pub enumeration_offset: u64,
1014}
1015
1016#[derive(Error, Debug)]
1017pub enum ForkCollectionError {
1018    #[error("Collection [{0}] already exists")]
1019    AlreadyExists(String),
1020    #[error("Failed to convert proto collection")]
1021    CollectionConversionError(#[from] CollectionConversionError),
1022    #[error("Duplicate segment")]
1023    DuplicateSegment,
1024    #[error("Missing field: [{0}]")]
1025    Field(String),
1026    #[error("Collection forking is unsupported for local chroma")]
1027    Local,
1028    #[error(transparent)]
1029    Internal(#[from] Box<dyn ChromaError>),
1030    #[error("Collection [{0}] does not exist")]
1031    NotFound(String),
1032    #[error("Failed to convert proto segment")]
1033    SegmentConversionError(#[from] SegmentConversionError),
1034    #[error("Failed to reconcile schema: {0}")]
1035    InvalidSchema(#[from] SchemaError),
1036}
1037
1038impl ChromaError for ForkCollectionError {
1039    fn code(&self) -> ErrorCodes {
1040        match self {
1041            ForkCollectionError::NotFound(_) => ErrorCodes::NotFound,
1042            ForkCollectionError::AlreadyExists(_) => ErrorCodes::AlreadyExists,
1043            ForkCollectionError::CollectionConversionError(e) => e.code(),
1044            ForkCollectionError::DuplicateSegment => ErrorCodes::Internal,
1045            ForkCollectionError::Field(_) => ErrorCodes::FailedPrecondition,
1046            ForkCollectionError::Local => ErrorCodes::Unimplemented,
1047            ForkCollectionError::Internal(e) => e.code(),
1048            ForkCollectionError::SegmentConversionError(e) => e.code(),
1049            ForkCollectionError::InvalidSchema(e) => e.code(),
1050        }
1051    }
1052}
1053
1054#[derive(Debug, Error)]
1055pub enum CountForksError {
1056    #[error("Collection [{0}] does not exist")]
1057    NotFound(String),
1058    #[error(transparent)]
1059    Internal(#[from] Box<dyn ChromaError>),
1060    #[error("Count forks is unsupported for local chroma")]
1061    Local,
1062}
1063
1064impl ChromaError for CountForksError {
1065    fn code(&self) -> ErrorCodes {
1066        match self {
1067            CountForksError::NotFound(_) => ErrorCodes::NotFound,
1068            CountForksError::Internal(chroma_error) => chroma_error.code(),
1069            CountForksError::Local => ErrorCodes::Unimplemented,
1070        }
1071    }
1072}
1073
1074#[derive(Debug, Error)]
1075pub enum GetCollectionSizeError {
1076    #[error(transparent)]
1077    Internal(#[from] Box<dyn ChromaError>),
1078    #[error("Collection [{0}] does not exist")]
1079    NotFound(String),
1080}
1081
1082impl ChromaError for GetCollectionSizeError {
1083    fn code(&self) -> ErrorCodes {
1084        match self {
1085            GetCollectionSizeError::Internal(err) => err.code(),
1086            GetCollectionSizeError::NotFound(_) => ErrorCodes::NotFound,
1087        }
1088    }
1089}
1090
1091#[derive(Error, Debug)]
1092pub enum ListCollectionVersionsError {
1093    #[error(transparent)]
1094    Internal(#[from] Box<dyn ChromaError>),
1095    #[error("Collection [{0}] does not exist")]
1096    NotFound(String),
1097}
1098
1099impl ChromaError for ListCollectionVersionsError {
1100    fn code(&self) -> ErrorCodes {
1101        match self {
1102            ListCollectionVersionsError::Internal(err) => err.code(),
1103            ListCollectionVersionsError::NotFound(_) => ErrorCodes::NotFound,
1104        }
1105    }
1106}
1107
1108////////////////////////// Metadata Key Constants //////////////////////////
1109
1110pub const CHROMA_KEY: &str = "chroma:";
1111pub const CHROMA_DOCUMENT_KEY: &str = "chroma:document";
1112pub const CHROMA_URI_KEY: &str = "chroma:uri";
1113
1114////////////////////////// AddCollectionRecords //////////////////////////
1115
1116#[non_exhaustive]
1117#[derive(Debug, Clone, Validate, Serialize)]
1118#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
1119pub struct AddCollectionRecordsRequest {
1120    pub tenant_id: String,
1121    pub database_name: String,
1122    pub collection_id: CollectionUuid,
1123    pub ids: Vec<String>,
1124    #[validate(custom(function = "validate_embeddings"))]
1125    pub embeddings: Vec<Vec<f32>>,
1126    pub documents: Option<Vec<Option<String>>>,
1127    pub uris: Option<Vec<Option<String>>>,
1128    #[validate(custom(function = "validate_metadata_vec"))]
1129    pub metadatas: Option<Vec<Option<Metadata>>>,
1130}
1131
1132impl AddCollectionRecordsRequest {
1133    #[allow(clippy::too_many_arguments)]
1134    pub fn try_new(
1135        tenant_id: String,
1136        database_name: String,
1137        collection_id: CollectionUuid,
1138        ids: Vec<String>,
1139        embeddings: Vec<Vec<f32>>,
1140        documents: Option<Vec<Option<String>>>,
1141        uris: Option<Vec<Option<String>>>,
1142        metadatas: Option<Vec<Option<Metadata>>>,
1143    ) -> Result<Self, ChromaValidationError> {
1144        let request = Self {
1145            tenant_id,
1146            database_name,
1147            collection_id,
1148            ids,
1149            embeddings,
1150            documents,
1151            uris,
1152            metadatas,
1153        };
1154        request.validate().map_err(ChromaValidationError::from)?;
1155        Ok(request)
1156    }
1157}
1158
1159fn validate_embeddings(embeddings: &[Vec<f32>]) -> Result<(), ValidationError> {
1160    if embeddings.iter().any(|e| e.is_empty()) {
1161        return Err(ValidationError::new("embedding_minimum_dimensions")
1162            .with_message("Each embedding must have at least 1 dimension".into()));
1163    }
1164    Ok(())
1165}
1166
1167#[derive(Serialize, Default, Deserialize)]
1168#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
1169pub struct AddCollectionRecordsResponse {}
1170
1171#[derive(Error, Debug)]
1172pub enum AddCollectionRecordsError {
1173    #[error("Failed to get collection: {0}")]
1174    Collection(#[from] GetCollectionError),
1175    #[error("Backoff and retry")]
1176    Backoff,
1177    #[error(transparent)]
1178    Other(#[from] Box<dyn ChromaError>),
1179}
1180
1181impl ChromaError for AddCollectionRecordsError {
1182    fn code(&self) -> ErrorCodes {
1183        match self {
1184            AddCollectionRecordsError::Collection(err) => err.code(),
1185            AddCollectionRecordsError::Backoff => ErrorCodes::ResourceExhausted,
1186            AddCollectionRecordsError::Other(err) => err.code(),
1187        }
1188    }
1189}
1190
1191////////////////////////// UpdateCollectionRecords //////////////////////////
1192
1193#[non_exhaustive]
1194#[derive(Debug, Clone, Validate, Serialize)]
1195#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
1196pub struct UpdateCollectionRecordsRequest {
1197    pub tenant_id: String,
1198    pub database_name: String,
1199    pub collection_id: CollectionUuid,
1200    pub ids: Vec<String>,
1201    pub embeddings: Option<Vec<Option<Vec<f32>>>>,
1202    pub documents: Option<Vec<Option<String>>>,
1203    pub uris: Option<Vec<Option<String>>>,
1204    #[validate(custom(function = "validate_update_metadata_vec"))]
1205    pub metadatas: Option<Vec<Option<UpdateMetadata>>>,
1206}
1207
1208impl UpdateCollectionRecordsRequest {
1209    #[allow(clippy::too_many_arguments)]
1210    pub fn try_new(
1211        tenant_id: String,
1212        database_name: String,
1213        collection_id: CollectionUuid,
1214        ids: Vec<String>,
1215        embeddings: Option<Vec<Option<Vec<f32>>>>,
1216        documents: Option<Vec<Option<String>>>,
1217        uris: Option<Vec<Option<String>>>,
1218        metadatas: Option<Vec<Option<UpdateMetadata>>>,
1219    ) -> Result<Self, ChromaValidationError> {
1220        let request = Self {
1221            tenant_id,
1222            database_name,
1223            collection_id,
1224            ids,
1225            embeddings,
1226            documents,
1227            uris,
1228            metadatas,
1229        };
1230        request.validate().map_err(ChromaValidationError::from)?;
1231        Ok(request)
1232    }
1233}
1234
1235#[derive(Serialize, Deserialize)]
1236#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
1237pub struct UpdateCollectionRecordsResponse {}
1238
1239#[derive(Error, Debug)]
1240pub enum UpdateCollectionRecordsError {
1241    #[error("Backoff and retry")]
1242    Backoff,
1243    #[error(transparent)]
1244    Other(#[from] Box<dyn ChromaError>),
1245}
1246
1247impl ChromaError for UpdateCollectionRecordsError {
1248    fn code(&self) -> ErrorCodes {
1249        match self {
1250            UpdateCollectionRecordsError::Backoff => ErrorCodes::ResourceExhausted,
1251            UpdateCollectionRecordsError::Other(err) => err.code(),
1252        }
1253    }
1254}
1255
1256////////////////////////// UpsertCollectionRecords //////////////////////////
1257
1258#[non_exhaustive]
1259#[derive(Debug, Clone, Validate, Serialize)]
1260#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
1261pub struct UpsertCollectionRecordsRequest {
1262    pub tenant_id: String,
1263    pub database_name: String,
1264    pub collection_id: CollectionUuid,
1265    pub ids: Vec<String>,
1266    #[validate(custom(function = "validate_embeddings"))]
1267    pub embeddings: Vec<Vec<f32>>,
1268    pub documents: Option<Vec<Option<String>>>,
1269    pub uris: Option<Vec<Option<String>>>,
1270    #[validate(custom(function = "validate_update_metadata_vec"))]
1271    pub metadatas: Option<Vec<Option<UpdateMetadata>>>,
1272}
1273
1274impl UpsertCollectionRecordsRequest {
1275    #[allow(clippy::too_many_arguments)]
1276    pub fn try_new(
1277        tenant_id: String,
1278        database_name: String,
1279        collection_id: CollectionUuid,
1280        ids: Vec<String>,
1281        embeddings: Vec<Vec<f32>>,
1282        documents: Option<Vec<Option<String>>>,
1283        uris: Option<Vec<Option<String>>>,
1284        metadatas: Option<Vec<Option<UpdateMetadata>>>,
1285    ) -> Result<Self, ChromaValidationError> {
1286        let request = Self {
1287            tenant_id,
1288            database_name,
1289            collection_id,
1290            ids,
1291            embeddings,
1292            documents,
1293            uris,
1294            metadatas,
1295        };
1296        request.validate().map_err(ChromaValidationError::from)?;
1297        Ok(request)
1298    }
1299}
1300
1301#[derive(Serialize, Deserialize)]
1302#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
1303pub struct UpsertCollectionRecordsResponse {}
1304
1305#[derive(Error, Debug)]
1306pub enum UpsertCollectionRecordsError {
1307    #[error("Backoff and retry")]
1308    Backoff,
1309    #[error(transparent)]
1310    Other(#[from] Box<dyn ChromaError>),
1311}
1312
1313impl ChromaError for UpsertCollectionRecordsError {
1314    fn code(&self) -> ErrorCodes {
1315        match self {
1316            UpsertCollectionRecordsError::Backoff => ErrorCodes::ResourceExhausted,
1317            UpsertCollectionRecordsError::Other(err) => err.code(),
1318        }
1319    }
1320}
1321
1322////////////////////////// DeleteCollectionRecords //////////////////////////
1323
1324#[non_exhaustive]
1325#[derive(Debug, Clone, Validate, Serialize)]
1326#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
1327pub struct DeleteCollectionRecordsRequest {
1328    pub tenant_id: String,
1329    pub database_name: String,
1330    pub collection_id: CollectionUuid,
1331    pub ids: Option<Vec<String>>,
1332    pub r#where: Option<Where>,
1333}
1334
1335impl DeleteCollectionRecordsRequest {
1336    pub fn try_new(
1337        tenant_id: String,
1338        database_name: String,
1339        collection_id: CollectionUuid,
1340        ids: Option<Vec<String>>,
1341        r#where: Option<Where>,
1342    ) -> Result<Self, ChromaValidationError> {
1343        if ids.as_ref().map(|ids| ids.is_empty()).unwrap_or(false) && r#where.is_none() {
1344            return Err(ChromaValidationError::from((
1345                ("ids, where"),
1346                ValidationError::new("filter")
1347                    .with_message("Either ids or where must be specified".into()),
1348            )));
1349        }
1350
1351        let request = Self {
1352            tenant_id,
1353            database_name,
1354            collection_id,
1355            ids,
1356            r#where,
1357        };
1358        request.validate().map_err(ChromaValidationError::from)?;
1359        Ok(request)
1360    }
1361}
1362
1363#[derive(Serialize, Deserialize)]
1364#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
1365pub struct DeleteCollectionRecordsResponse {}
1366
1367#[derive(Error, Debug)]
1368pub enum DeleteCollectionRecordsError {
1369    #[error("Failed to resolve records for deletion: {0}")]
1370    Get(#[from] ExecutorError),
1371    #[error("Backoff and retry")]
1372    Backoff,
1373    #[error(transparent)]
1374    Internal(#[from] Box<dyn ChromaError>),
1375}
1376
1377impl ChromaError for DeleteCollectionRecordsError {
1378    fn code(&self) -> ErrorCodes {
1379        match self {
1380            DeleteCollectionRecordsError::Get(err) => err.code(),
1381            DeleteCollectionRecordsError::Backoff => ErrorCodes::ResourceExhausted,
1382            DeleteCollectionRecordsError::Internal(err) => err.code(),
1383        }
1384    }
1385}
1386
1387////////////////////////// Include //////////////////////////
1388
1389#[derive(Error, Debug)]
1390#[error("Invalid include value: {0}")]
1391pub struct IncludeParsingError(String);
1392
1393impl ChromaError for IncludeParsingError {
1394    fn code(&self) -> ErrorCodes {
1395        ErrorCodes::InvalidArgument
1396    }
1397}
1398
1399#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
1400#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
1401pub enum Include {
1402    #[serde(rename = "distances")]
1403    Distance,
1404    #[serde(rename = "documents")]
1405    Document,
1406    #[serde(rename = "embeddings")]
1407    Embedding,
1408    #[serde(rename = "metadatas")]
1409    Metadata,
1410    #[serde(rename = "uris")]
1411    Uri,
1412}
1413
1414impl TryFrom<&str> for Include {
1415    type Error = IncludeParsingError;
1416
1417    fn try_from(value: &str) -> Result<Self, Self::Error> {
1418        match value {
1419            "distances" => Ok(Include::Distance),
1420            "documents" => Ok(Include::Document),
1421            "embeddings" => Ok(Include::Embedding),
1422            "metadatas" => Ok(Include::Metadata),
1423            "uris" => Ok(Include::Uri),
1424            _ => Err(IncludeParsingError(value.to_string())),
1425        }
1426    }
1427}
1428
1429#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
1430#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
1431#[cfg_attr(feature = "pyo3", pyo3::pyclass)]
1432pub struct IncludeList(pub Vec<Include>);
1433
1434impl IncludeList {
1435    pub fn empty() -> Self {
1436        Self(Vec::new())
1437    }
1438
1439    pub fn default_query() -> Self {
1440        Self(vec![
1441            Include::Document,
1442            Include::Metadata,
1443            Include::Distance,
1444        ])
1445    }
1446    pub fn default_get() -> Self {
1447        Self(vec![Include::Document, Include::Metadata])
1448    }
1449    pub fn all() -> Self {
1450        Self(vec![
1451            Include::Document,
1452            Include::Metadata,
1453            Include::Distance,
1454            Include::Embedding,
1455            Include::Uri,
1456        ])
1457    }
1458}
1459
1460impl TryFrom<Vec<String>> for IncludeList {
1461    type Error = IncludeParsingError;
1462
1463    fn try_from(value: Vec<String>) -> Result<Self, Self::Error> {
1464        let mut includes = Vec::new();
1465        for v in value {
1466            // "data" is only used by single node Chroma
1467            if v == "data" {
1468                includes.push(Include::Metadata);
1469                continue;
1470            }
1471
1472            includes.push(Include::try_from(v.as_str())?);
1473        }
1474        Ok(IncludeList(includes))
1475    }
1476}
1477
1478////////////////////////// Count //////////////////////////
1479
1480#[non_exhaustive]
1481#[derive(Clone, Deserialize, Serialize, Validate)]
1482pub struct CountRequest {
1483    pub tenant_id: String,
1484    pub database_name: String,
1485    pub collection_id: CollectionUuid,
1486}
1487
1488impl CountRequest {
1489    pub fn try_new(
1490        tenant_id: String,
1491        database_name: String,
1492        collection_id: CollectionUuid,
1493    ) -> Result<Self, ChromaValidationError> {
1494        let request = Self {
1495            tenant_id,
1496            database_name,
1497            collection_id,
1498        };
1499        request.validate().map_err(ChromaValidationError::from)?;
1500        Ok(request)
1501    }
1502}
1503
1504pub type CountResponse = u32;
1505
1506////////////////////////// Get //////////////////////////
1507
1508#[non_exhaustive]
1509#[derive(Debug, Clone, Validate, Serialize)]
1510#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
1511pub struct GetRequest {
1512    pub tenant_id: String,
1513    pub database_name: String,
1514    pub collection_id: CollectionUuid,
1515    pub ids: Option<Vec<String>>,
1516    pub r#where: Option<Where>,
1517    pub limit: Option<u32>,
1518    pub offset: u32,
1519    pub include: IncludeList,
1520}
1521
1522impl GetRequest {
1523    #[allow(clippy::too_many_arguments)]
1524    pub fn try_new(
1525        tenant_id: String,
1526        database_name: String,
1527        collection_id: CollectionUuid,
1528        ids: Option<Vec<String>>,
1529        r#where: Option<Where>,
1530        limit: Option<u32>,
1531        offset: u32,
1532        include: IncludeList,
1533    ) -> Result<Self, ChromaValidationError> {
1534        let request = Self {
1535            tenant_id,
1536            database_name,
1537            collection_id,
1538            ids,
1539            r#where,
1540            limit,
1541            offset,
1542            include,
1543        };
1544        request.validate().map_err(ChromaValidationError::from)?;
1545        Ok(request)
1546    }
1547}
1548
1549#[derive(Clone, Deserialize, Serialize, Debug, Default)]
1550#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
1551#[cfg_attr(feature = "pyo3", pyo3::pyclass)]
1552pub struct GetResponse {
1553    pub ids: Vec<String>,
1554    pub embeddings: Option<Vec<Vec<f32>>>,
1555    pub documents: Option<Vec<Option<String>>>,
1556    pub uris: Option<Vec<Option<String>>>,
1557    // TODO(hammadb): Add metadata & include to the response
1558    pub metadatas: Option<Vec<Option<Metadata>>>,
1559    pub include: Vec<Include>,
1560}
1561
1562impl GetResponse {
1563    pub fn sort_by_ids(&mut self) {
1564        let mut indices: Vec<usize> = (0..self.ids.len()).collect();
1565        indices.sort_by(|&a, &b| self.ids[a].cmp(&self.ids[b]));
1566
1567        let sorted_ids = indices.iter().map(|&i| self.ids[i].clone()).collect();
1568        self.ids = sorted_ids;
1569
1570        if let Some(ref mut embeddings) = self.embeddings {
1571            let sorted_embeddings = indices.iter().map(|&i| embeddings[i].clone()).collect();
1572            *embeddings = sorted_embeddings;
1573        }
1574
1575        if let Some(ref mut documents) = self.documents {
1576            let sorted_docs = indices.iter().map(|&i| documents[i].clone()).collect();
1577            *documents = sorted_docs;
1578        }
1579
1580        if let Some(ref mut uris) = self.uris {
1581            let sorted_uris = indices.iter().map(|&i| uris[i].clone()).collect();
1582            *uris = sorted_uris;
1583        }
1584
1585        if let Some(ref mut metadatas) = self.metadatas {
1586            let sorted_metas = indices.iter().map(|&i| metadatas[i].clone()).collect();
1587            *metadatas = sorted_metas;
1588        }
1589    }
1590}
1591
1592#[cfg(feature = "pyo3")]
1593#[pyo3::pymethods]
1594impl GetResponse {
1595    #[getter]
1596    pub fn ids(&self) -> &Vec<String> {
1597        &self.ids
1598    }
1599
1600    #[getter]
1601    pub fn embeddings(&self) -> Option<Vec<Vec<f32>>> {
1602        self.embeddings.clone()
1603    }
1604
1605    #[getter]
1606    pub fn documents(&self) -> Option<Vec<Option<String>>> {
1607        self.documents.clone()
1608    }
1609
1610    #[getter]
1611    pub fn uris(&self) -> Option<Vec<Option<String>>> {
1612        self.uris.clone()
1613    }
1614
1615    #[getter]
1616    pub fn metadatas(&self) -> Option<Vec<Option<Metadata>>> {
1617        self.metadatas.clone()
1618    }
1619}
1620
1621impl From<(GetResult, IncludeList)> for GetResponse {
1622    fn from((result, IncludeList(include_vec)): (GetResult, IncludeList)) -> Self {
1623        let mut res = Self {
1624            ids: Vec::new(),
1625            embeddings: include_vec
1626                .contains(&Include::Embedding)
1627                .then_some(Vec::new()),
1628            documents: include_vec
1629                .contains(&Include::Document)
1630                .then_some(Vec::new()),
1631            uris: include_vec.contains(&Include::Uri).then_some(Vec::new()),
1632            metadatas: include_vec
1633                .contains(&Include::Metadata)
1634                .then_some(Vec::new()),
1635            include: include_vec,
1636        };
1637        for ProjectionRecord {
1638            id,
1639            document,
1640            embedding,
1641            mut metadata,
1642        } in result.result.records
1643        {
1644            res.ids.push(id);
1645            if let (Some(emb), Some(embeddings)) = (embedding, res.embeddings.as_mut()) {
1646                embeddings.push(emb);
1647            }
1648            if let Some(documents) = res.documents.as_mut() {
1649                documents.push(document);
1650            }
1651            let uri = metadata.as_mut().and_then(|meta| {
1652                meta.remove(CHROMA_URI_KEY).and_then(|v| {
1653                    if let crate::MetadataValue::Str(uri) = v {
1654                        Some(uri)
1655                    } else {
1656                        None
1657                    }
1658                })
1659            });
1660            if let Some(uris) = res.uris.as_mut() {
1661                uris.push(uri);
1662            }
1663
1664            let metadata = metadata.map(|m| {
1665                m.into_iter()
1666                    .filter(|(k, _)| !k.starts_with(CHROMA_KEY))
1667                    .collect()
1668            });
1669            if let Some(metadatas) = res.metadatas.as_mut() {
1670                metadatas.push(metadata);
1671            }
1672        }
1673        res
1674    }
1675}
1676
1677////////////////////////// Query //////////////////////////
1678
1679#[non_exhaustive]
1680#[derive(Debug, Clone, Validate, Serialize)]
1681#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
1682pub struct QueryRequest {
1683    pub tenant_id: String,
1684    pub database_name: String,
1685    pub collection_id: CollectionUuid,
1686    pub ids: Option<Vec<String>>,
1687    pub r#where: Option<Where>,
1688    pub embeddings: Vec<Vec<f32>>,
1689    pub n_results: u32,
1690    pub include: IncludeList,
1691}
1692
1693impl QueryRequest {
1694    #[allow(clippy::too_many_arguments)]
1695    pub fn try_new(
1696        tenant_id: String,
1697        database_name: String,
1698        collection_id: CollectionUuid,
1699        ids: Option<Vec<String>>,
1700        r#where: Option<Where>,
1701        embeddings: Vec<Vec<f32>>,
1702        n_results: u32,
1703        include: IncludeList,
1704    ) -> Result<Self, ChromaValidationError> {
1705        let request = Self {
1706            tenant_id,
1707            database_name,
1708            collection_id,
1709            ids,
1710            r#where,
1711            embeddings,
1712            n_results,
1713            include,
1714        };
1715        request.validate().map_err(ChromaValidationError::from)?;
1716        Ok(request)
1717    }
1718}
1719
1720#[derive(Clone, Deserialize, Serialize, Debug)]
1721#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
1722#[cfg_attr(feature = "pyo3", pyo3::pyclass)]
1723pub struct QueryResponse {
1724    pub ids: Vec<Vec<String>>,
1725    pub embeddings: Option<Vec<Vec<Option<Vec<f32>>>>>,
1726    pub documents: Option<Vec<Vec<Option<String>>>>,
1727    pub uris: Option<Vec<Vec<Option<String>>>>,
1728    pub metadatas: Option<Vec<Vec<Option<Metadata>>>>,
1729    pub distances: Option<Vec<Vec<Option<f32>>>>,
1730    pub include: Vec<Include>,
1731}
1732
1733impl QueryResponse {
1734    pub fn sort_by_ids(&mut self) {
1735        fn reorder<T: Clone>(v: &mut [T], indices: &[usize]) {
1736            let old = v.to_owned();
1737            for (new_pos, &i) in indices.iter().enumerate() {
1738                v[new_pos] = old[i].clone();
1739            }
1740        }
1741
1742        for i in 0..self.ids.len() {
1743            let mut indices: Vec<usize> = (0..self.ids[i].len()).collect();
1744
1745            indices.sort_unstable_by(|&a, &b| self.ids[i][a].cmp(&self.ids[i][b]));
1746
1747            reorder(&mut self.ids[i], &indices);
1748
1749            if let Some(embeddings) = &mut self.embeddings {
1750                reorder(&mut embeddings[i], &indices);
1751            }
1752
1753            if let Some(documents) = &mut self.documents {
1754                reorder(&mut documents[i], &indices);
1755            }
1756
1757            if let Some(uris) = &mut self.uris {
1758                reorder(&mut uris[i], &indices);
1759            }
1760
1761            if let Some(metadatas) = &mut self.metadatas {
1762                reorder(&mut metadatas[i], &indices);
1763            }
1764
1765            if let Some(distances) = &mut self.distances {
1766                reorder(&mut distances[i], &indices);
1767            }
1768        }
1769    }
1770}
1771
1772#[cfg(feature = "pyo3")]
1773#[pyo3::pymethods]
1774impl QueryResponse {
1775    #[getter]
1776    pub fn ids(&self) -> &Vec<Vec<String>> {
1777        &self.ids
1778    }
1779
1780    #[getter]
1781    pub fn embeddings(&self) -> Option<Vec<Vec<Option<Vec<f32>>>>> {
1782        self.embeddings.clone()
1783    }
1784
1785    #[getter]
1786    pub fn documents(&self) -> Option<Vec<Vec<Option<String>>>> {
1787        self.documents.clone()
1788    }
1789
1790    #[getter]
1791    pub fn uris(&self) -> Option<Vec<Vec<Option<String>>>> {
1792        self.uris.clone()
1793    }
1794
1795    #[getter]
1796    pub fn metadatas(&self) -> Option<Vec<Vec<Option<Metadata>>>> {
1797        self.metadatas.clone()
1798    }
1799
1800    #[getter]
1801    pub fn distances(&self) -> Option<Vec<Vec<Option<f32>>>> {
1802        self.distances.clone()
1803    }
1804}
1805
1806impl From<(KnnBatchResult, IncludeList)> for QueryResponse {
1807    fn from((result, IncludeList(include_vec)): (KnnBatchResult, IncludeList)) -> Self {
1808        let mut res = Self {
1809            ids: Vec::new(),
1810            embeddings: include_vec
1811                .contains(&Include::Embedding)
1812                .then_some(Vec::new()),
1813            documents: include_vec
1814                .contains(&Include::Document)
1815                .then_some(Vec::new()),
1816            uris: include_vec.contains(&Include::Uri).then_some(Vec::new()),
1817            metadatas: include_vec
1818                .contains(&Include::Metadata)
1819                .then_some(Vec::new()),
1820            distances: include_vec
1821                .contains(&Include::Distance)
1822                .then_some(Vec::new()),
1823            include: include_vec,
1824        };
1825        for query_result in result.results {
1826            let mut ids = Vec::new();
1827            let mut embeddings = Vec::new();
1828            let mut documents = Vec::new();
1829            let mut uris = Vec::new();
1830            let mut metadatas = Vec::new();
1831            let mut distances = Vec::new();
1832            for KnnProjectionRecord {
1833                record:
1834                    ProjectionRecord {
1835                        id,
1836                        document,
1837                        embedding,
1838                        mut metadata,
1839                    },
1840                distance,
1841            } in query_result.records
1842            {
1843                ids.push(id);
1844                embeddings.push(embedding);
1845                documents.push(document);
1846
1847                let uri = metadata.as_mut().and_then(|meta| {
1848                    meta.remove(CHROMA_URI_KEY).and_then(|v| {
1849                        if let crate::MetadataValue::Str(uri) = v {
1850                            Some(uri)
1851                        } else {
1852                            None
1853                        }
1854                    })
1855                });
1856                uris.push(uri);
1857
1858                let metadata = metadata.map(|m| {
1859                    m.into_iter()
1860                        .filter(|(k, _)| !k.starts_with(CHROMA_KEY))
1861                        .collect()
1862                });
1863                metadatas.push(metadata);
1864
1865                distances.push(distance);
1866            }
1867            res.ids.push(ids);
1868
1869            if let Some(res_embs) = res.embeddings.as_mut() {
1870                res_embs.push(embeddings);
1871            }
1872            if let Some(res_docs) = res.documents.as_mut() {
1873                res_docs.push(documents);
1874            }
1875            if let Some(res_uri) = res.uris.as_mut() {
1876                res_uri.push(uris);
1877            }
1878            if let Some(res_metas) = res.metadatas.as_mut() {
1879                res_metas.push(metadatas);
1880            }
1881            if let Some(res_dists) = res.distances.as_mut() {
1882                res_dists.push(distances);
1883            }
1884        }
1885        res
1886    }
1887}
1888
1889#[non_exhaustive]
1890#[derive(Clone, Debug, Serialize, Validate)]
1891#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
1892pub struct SearchRequest {
1893    pub tenant_id: String,
1894    pub database_name: String,
1895    pub collection_id: CollectionUuid,
1896    pub searches: Vec<SearchPayload>,
1897}
1898
1899impl SearchRequest {
1900    pub fn try_new(
1901        tenant_id: String,
1902        database_name: String,
1903        collection_id: CollectionUuid,
1904        searches: Vec<SearchPayload>,
1905    ) -> Result<Self, ChromaValidationError> {
1906        let request = Self {
1907            tenant_id,
1908            database_name,
1909            collection_id,
1910            searches,
1911        };
1912        request.validate().map_err(ChromaValidationError::from)?;
1913        Ok(request)
1914    }
1915}
1916
1917#[derive(Clone, Deserialize, Serialize, Debug)]
1918#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
1919pub struct SearchResponse {
1920    pub ids: Vec<Vec<String>>,
1921    pub documents: Vec<Option<Vec<Option<String>>>>,
1922    pub embeddings: Vec<Option<Vec<Option<Vec<f32>>>>>,
1923    pub metadatas: Vec<Option<Vec<Option<Metadata>>>>,
1924    pub scores: Vec<Option<Vec<Option<f32>>>>,
1925    pub select: Vec<Vec<Key>>,
1926}
1927
1928impl From<(SearchResult, Vec<SearchPayload>)> for SearchResponse {
1929    fn from((result, payloads): (SearchResult, Vec<SearchPayload>)) -> Self {
1930        let num_payloads = payloads.len();
1931        let mut res = Self {
1932            ids: Vec::with_capacity(num_payloads),
1933            documents: Vec::with_capacity(num_payloads),
1934            embeddings: Vec::with_capacity(num_payloads),
1935            metadatas: Vec::with_capacity(num_payloads),
1936            scores: Vec::with_capacity(num_payloads),
1937            select: Vec::with_capacity(num_payloads),
1938        };
1939
1940        for (payload_result, payload) in result.results.into_iter().zip(payloads) {
1941            // Get the sorted keys for this payload
1942            let mut payload_select = Vec::from_iter(payload.select.keys.iter().cloned());
1943            payload_select.sort();
1944
1945            let num_records = payload_result.records.len();
1946            let mut ids = Vec::with_capacity(num_records);
1947            let mut documents = Vec::with_capacity(num_records);
1948            let mut embeddings = Vec::with_capacity(num_records);
1949            let mut metadatas = Vec::with_capacity(num_records);
1950            let mut scores = Vec::with_capacity(num_records);
1951
1952            for record in payload_result.records {
1953                ids.push(record.id);
1954                documents.push(record.document);
1955                embeddings.push(record.embedding);
1956                metadatas.push(record.metadata);
1957                scores.push(record.score);
1958            }
1959
1960            res.ids.push(ids);
1961            res.select.push(payload_select.clone());
1962
1963            // Push documents if requested by this payload, otherwise None
1964            res.documents.push(
1965                payload_select
1966                    .binary_search(&Key::Document)
1967                    .is_ok()
1968                    .then_some(documents),
1969            );
1970
1971            // Push embeddings if requested by this payload, otherwise None
1972            res.embeddings.push(
1973                payload_select
1974                    .binary_search(&Key::Embedding)
1975                    .is_ok()
1976                    .then_some(embeddings),
1977            );
1978
1979            // Push metadatas if requested by this payload, otherwise None
1980            // Include if either Key::Metadata is present or any Key::MetadataField(_)
1981            let has_metadata = payload_select.binary_search(&Key::Metadata).is_ok()
1982                || payload_select
1983                    .last()
1984                    .is_some_and(|field| matches!(field, Key::MetadataField(_)));
1985            res.metadatas.push(has_metadata.then_some(metadatas));
1986
1987            // Push scores if requested by this payload, otherwise None
1988            res.scores.push(
1989                payload_select
1990                    .binary_search(&Key::Score)
1991                    .is_ok()
1992                    .then_some(scores),
1993            );
1994        }
1995
1996        res
1997    }
1998}
1999
2000#[derive(Error, Debug)]
2001pub enum QueryError {
2002    #[error("Error executing plan: {0}")]
2003    Executor(#[from] ExecutorError),
2004    #[error(transparent)]
2005    Other(#[from] Box<dyn ChromaError>),
2006}
2007
2008impl ChromaError for QueryError {
2009    fn code(&self) -> ErrorCodes {
2010        match self {
2011            QueryError::Executor(e) => e.code(),
2012            QueryError::Other(err) => err.code(),
2013        }
2014    }
2015}
2016
2017#[derive(Serialize)]
2018#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
2019pub struct HealthCheckResponse {
2020    pub is_executor_ready: bool,
2021    pub is_log_client_ready: bool,
2022}
2023
2024impl HealthCheckResponse {
2025    pub fn get_status_code(&self) -> tonic::Code {
2026        if self.is_executor_ready && self.is_log_client_ready {
2027            tonic::Code::Ok
2028        } else {
2029            tonic::Code::Unavailable
2030        }
2031    }
2032}
2033
2034#[derive(Debug, Error)]
2035pub enum ExecutorError {
2036    #[error("Error converting: {0}")]
2037    Conversion(#[from] QueryConversionError),
2038    #[error("Error converting plan to proto: {0}")]
2039    PlanToProto(#[from] PlanToProtoError),
2040    #[error(transparent)]
2041    Grpc(#[from] Status),
2042    #[error("Inconsistent data")]
2043    InconsistentData,
2044    #[error("Collection is missing HNSW configuration")]
2045    CollectionMissingHnswConfiguration,
2046    #[error("Internal error: {0}")]
2047    Internal(Box<dyn ChromaError>),
2048    #[error("Error sending backfill request to compactor: {0}")]
2049    BackfillError(Box<dyn ChromaError>),
2050    #[error("Not implemented: {0}")]
2051    NotImplemented(String),
2052}
2053
2054impl ChromaError for ExecutorError {
2055    fn code(&self) -> ErrorCodes {
2056        match self {
2057            ExecutorError::Conversion(_) => ErrorCodes::InvalidArgument,
2058            ExecutorError::PlanToProto(_) => ErrorCodes::Internal,
2059            ExecutorError::Grpc(e) => e.code().into(),
2060            ExecutorError::InconsistentData => ErrorCodes::Internal,
2061            ExecutorError::CollectionMissingHnswConfiguration => ErrorCodes::Internal,
2062            ExecutorError::Internal(e) => e.code(),
2063            ExecutorError::BackfillError(e) => e.code(),
2064            ExecutorError::NotImplemented(_) => ErrorCodes::Unimplemented,
2065        }
2066    }
2067}
2068
2069////////////////////////// Task Operations //////////////////////////
2070
2071#[non_exhaustive]
2072#[derive(Clone, Debug, Deserialize, Serialize, Validate)]
2073#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
2074pub struct CreateTaskRequest {
2075    #[validate(length(min = 1))]
2076    pub task_name: String,
2077    pub operator_name: String,
2078    pub output_collection_name: String,
2079    #[serde(default = "default_empty_json_object")]
2080    pub params: serde_json::Value,
2081}
2082
2083fn default_empty_json_object() -> serde_json::Value {
2084    serde_json::json!({})
2085}
2086
2087impl CreateTaskRequest {
2088    pub fn try_new(
2089        task_name: String,
2090        operator_name: String,
2091        output_collection_name: String,
2092        params: serde_json::Value,
2093    ) -> Result<Self, ChromaValidationError> {
2094        let request = Self {
2095            task_name,
2096            operator_name,
2097            output_collection_name,
2098            params,
2099        };
2100        request.validate().map_err(ChromaValidationError::from)?;
2101        Ok(request)
2102    }
2103}
2104
2105#[derive(Clone, Debug, Serialize)]
2106#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
2107pub struct CreateTaskResponse {
2108    pub success: bool,
2109    pub task_id: String,
2110}
2111
2112#[derive(Error, Debug)]
2113pub enum AddTaskError {
2114    #[error("Task with name [{0}] already exists")]
2115    AlreadyExists(String),
2116    #[error("Input collection [{0}] does not exist")]
2117    InputCollectionNotFound(String),
2118    #[error("Output collection [{0}] already exists")]
2119    OutputCollectionExists(String),
2120    #[error(transparent)]
2121    Validation(#[from] ChromaValidationError),
2122    #[error(transparent)]
2123    Internal(#[from] Box<dyn ChromaError>),
2124}
2125
2126impl ChromaError for AddTaskError {
2127    fn code(&self) -> ErrorCodes {
2128        match self {
2129            AddTaskError::AlreadyExists(_) => ErrorCodes::AlreadyExists,
2130            AddTaskError::InputCollectionNotFound(_) => ErrorCodes::NotFound,
2131            AddTaskError::OutputCollectionExists(_) => ErrorCodes::AlreadyExists,
2132            AddTaskError::Validation(err) => err.code(),
2133            AddTaskError::Internal(err) => err.code(),
2134        }
2135    }
2136}
2137
2138#[non_exhaustive]
2139#[derive(Clone, Debug, Deserialize, Validate, Serialize)]
2140#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
2141pub struct RemoveTaskRequest {
2142    #[validate(length(min = 1))]
2143    pub task_name: String,
2144    /// Whether to delete the output collection as well
2145    #[serde(default)]
2146    pub delete_output: bool,
2147}
2148
2149impl RemoveTaskRequest {
2150    pub fn try_new(task_name: String, delete_output: bool) -> Result<Self, ChromaValidationError> {
2151        let request = Self {
2152            task_name,
2153            delete_output,
2154        };
2155        request.validate().map_err(ChromaValidationError::from)?;
2156        Ok(request)
2157    }
2158}
2159
2160#[derive(Clone, Debug, Serialize)]
2161#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
2162pub struct RemoveTaskResponse {
2163    pub success: bool,
2164}
2165
2166#[derive(Error, Debug)]
2167pub enum RemoveTaskError {
2168    #[error("Task with name [{0}] does not exist")]
2169    NotFound(String),
2170    #[error(transparent)]
2171    Validation(#[from] ChromaValidationError),
2172    #[error(transparent)]
2173    Internal(#[from] Box<dyn ChromaError>),
2174}
2175
2176impl ChromaError for RemoveTaskError {
2177    fn code(&self) -> ErrorCodes {
2178        match self {
2179            RemoveTaskError::NotFound(_) => ErrorCodes::NotFound,
2180            RemoveTaskError::Validation(err) => err.code(),
2181            RemoveTaskError::Internal(err) => err.code(),
2182        }
2183    }
2184}
2185
2186#[cfg(test)]
2187mod test {
2188    use super::*;
2189    use crate::{MetadataValue, SparseVector, UpdateMetadataValue};
2190    use std::collections::HashMap;
2191
2192    #[test]
2193    fn test_create_database_min_length() {
2194        let request = CreateDatabaseRequest::try_new("default_tenant".to_string(), "a".to_string());
2195        assert!(request.is_err());
2196    }
2197
2198    #[test]
2199    fn test_create_tenant_min_length() {
2200        let request = CreateTenantRequest::try_new("a".to_string());
2201        assert!(request.is_err());
2202    }
2203
2204    #[test]
2205    fn test_add_request_validates_sparse_vectors() {
2206        let mut metadata = HashMap::new();
2207        // Add unsorted sparse vector - should fail validation
2208        metadata.insert(
2209            "sparse".to_string(),
2210            MetadataValue::SparseVector(SparseVector::new(vec![3, 1, 2], vec![0.3, 0.1, 0.2])),
2211        );
2212
2213        let result = AddCollectionRecordsRequest::try_new(
2214            "tenant".to_string(),
2215            "database".to_string(),
2216            CollectionUuid(uuid::Uuid::new_v4()),
2217            vec!["id1".to_string()],
2218            vec![vec![0.1, 0.2]],
2219            None,
2220            None,
2221            Some(vec![Some(metadata)]),
2222        );
2223
2224        // Should fail because sparse vector is not sorted
2225        assert!(result.is_err());
2226    }
2227
2228    #[test]
2229    fn test_update_request_validates_sparse_vectors() {
2230        let mut metadata = HashMap::new();
2231        // Add unsorted sparse vector - should fail validation
2232        metadata.insert(
2233            "sparse".to_string(),
2234            UpdateMetadataValue::SparseVector(SparseVector::new(
2235                vec![3, 1, 2],
2236                vec![0.3, 0.1, 0.2],
2237            )),
2238        );
2239
2240        let result = UpdateCollectionRecordsRequest::try_new(
2241            "tenant".to_string(),
2242            "database".to_string(),
2243            CollectionUuid(uuid::Uuid::new_v4()),
2244            vec!["id1".to_string()],
2245            None,
2246            None,
2247            None,
2248            Some(vec![Some(metadata)]),
2249        );
2250
2251        // Should fail because sparse vector is not sorted
2252        assert!(result.is_err());
2253    }
2254
2255    #[test]
2256    fn test_upsert_request_validates_sparse_vectors() {
2257        let mut metadata = HashMap::new();
2258        // Add unsorted sparse vector - should fail validation
2259        metadata.insert(
2260            "sparse".to_string(),
2261            UpdateMetadataValue::SparseVector(SparseVector::new(
2262                vec![3, 1, 2],
2263                vec![0.3, 0.1, 0.2],
2264            )),
2265        );
2266
2267        let result = UpsertCollectionRecordsRequest::try_new(
2268            "tenant".to_string(),
2269            "database".to_string(),
2270            CollectionUuid(uuid::Uuid::new_v4()),
2271            vec!["id1".to_string()],
2272            vec![vec![0.1, 0.2]],
2273            None,
2274            None,
2275            Some(vec![Some(metadata)]),
2276        );
2277
2278        // Should fail because sparse vector is not sorted
2279        assert!(result.is_err());
2280    }
2281}
chroma_types/api_types.rs

chroma_types/
api_types.rs
