chroma_types/
segment_scope.rs

1use crate::{chroma_proto, ConversionError};
2use chroma_error::{ChromaError, ErrorCodes};
3
4use thiserror::Error;
5
6#[derive(Clone, Debug, PartialEq)]
7pub enum SegmentScope {
8    VECTOR,
9    METADATA,
10    RECORD,
11    SQLITE,
12}
13
14impl From<SegmentScope> for String {
15    fn from(scope: SegmentScope) -> String {
16        match scope {
17            SegmentScope::VECTOR => "VECTOR".to_string(),
18            SegmentScope::METADATA => "METADATA".to_string(),
19            SegmentScope::RECORD => "RECORD".to_string(),
20            SegmentScope::SQLITE => "SQLITE".to_string(),
21        }
22    }
23}
24
25impl TryFrom<&str> for SegmentScope {
26    type Error = SegmentScopeConversionError;
27
28    fn try_from(scope: &str) -> Result<Self, Self::Error> {
29        match scope {
30            "VECTOR" => Ok(SegmentScope::VECTOR),
31            "METADATA" => Ok(SegmentScope::METADATA),
32            "RECORD" => Ok(SegmentScope::RECORD),
33            "SQLITE" => Ok(SegmentScope::SQLITE),
34            _ => Err(SegmentScopeConversionError::InvalidScope),
35        }
36    }
37}
38
39#[derive(Error, Debug)]
40pub enum SegmentScopeConversionError {
41    #[error("Invalid segment scope, valid scopes are: Vector, Metadata")]
42    InvalidScope,
43    #[error(transparent)]
44    DecodeError(#[from] ConversionError),
45}
46
47impl_base_convert_error!(SegmentScopeConversionError, {
48    SegmentScopeConversionError::InvalidScope => ErrorCodes::InvalidArgument,
49});
50
51impl From<chroma_proto::SegmentScope> for SegmentScope {
52    fn from(value: chroma_proto::SegmentScope) -> Self {
53        match value {
54            chroma_proto::SegmentScope::Vector => Self::VECTOR,
55            chroma_proto::SegmentScope::Metadata => Self::METADATA,
56            chroma_proto::SegmentScope::Record => Self::RECORD,
57            chroma_proto::SegmentScope::Sqlite => Self::SQLITE,
58        }
59    }
60}
61
62impl From<SegmentScope> for chroma_proto::SegmentScope {
63    fn from(value: SegmentScope) -> Self {
64        match value {
65            SegmentScope::VECTOR => Self::Vector,
66            SegmentScope::METADATA => Self::Metadata,
67            SegmentScope::RECORD => Self::Record,
68            SegmentScope::SQLITE => Self::Sqlite,
69        }
70    }
71}
72
73impl TryFrom<i32> for SegmentScope {
74    type Error = SegmentScopeConversionError;
75
76    fn try_from(scope: i32) -> Result<Self, Self::Error> {
77        let maybe_scope = chroma_proto::SegmentScope::try_from(scope);
78        match maybe_scope {
79            Ok(scope) => match scope {
80                chroma_proto::SegmentScope::Vector => Ok(SegmentScope::VECTOR),
81                chroma_proto::SegmentScope::Metadata => Ok(SegmentScope::METADATA),
82                chroma_proto::SegmentScope::Record => Ok(SegmentScope::RECORD),
83                chroma_proto::SegmentScope::Sqlite => Ok(SegmentScope::SQLITE),
84            },
85            Err(_) => Err(SegmentScopeConversionError::InvalidScope),
86        }
87    }
88}
89
90#[cfg(test)]
91mod tests {
92    use super::*;
93
94    #[test]
95    fn test_segment_scope_try_from() {
96        let proto_scope = chroma_proto::SegmentScope::Vector;
97        let converted_scope: SegmentScope = proto_scope.into();
98        assert_eq!(converted_scope, SegmentScope::VECTOR);
99
100        let proto_scope = chroma_proto::SegmentScope::Metadata;
101        let converted_scope: SegmentScope = proto_scope.into();
102        assert_eq!(converted_scope, SegmentScope::METADATA);
103
104        let proto_scope = chroma_proto::SegmentScope::Sqlite;
105        let converted_scope: SegmentScope = proto_scope.into();
106        assert_eq!(converted_scope, SegmentScope::SQLITE);
107
108        let proto_scope = chroma_proto::SegmentScope::Record;
109        let converted_scope: SegmentScope = proto_scope.into();
110        assert_eq!(converted_scope, SegmentScope::RECORD);
111    }
112}