Skip to main content

vecgraph_proto/traits/
search_impl.rs

1use crate::{RerankParams, SearchKind, SearchQuery, SearchResult};
2use vecgraph_core::VecGraphError;
3
4impl TryFrom<SearchQuery> for vecgraph_core::SearchQuery {
5    type Error = VecGraphError;
6
7    fn try_from(proto: SearchQuery) -> Result<Self, Self::Error> {
8        let proto_search_kind = SearchKind::try_from(proto.search_kind).map_err(|_| {
9            VecGraphError::Other(format!("invalid search_kind value: {}", proto.search_kind))
10        })?;
11
12        let search_kind: vecgraph_core::SearchKind = proto_search_kind.try_into()?;
13
14        let rerank = proto.rerank.map(|rp| vecgraph_core::RerankParams {
15            vector: rp.vectors,
16            kind: rp.kind,
17            weight: rp.weight,
18        });
19
20        Ok(vecgraph_core::SearchQuery {
21            query_vec: proto.query_vecs,
22            top_k: proto.top_k as usize,
23            search_kind,
24            rerank,
25            exclude_names: proto.exclude_names,
26            kind: proto.kind,
27            namespace: proto.namespace,
28        })
29    }
30}
31
32impl TryFrom<vecgraph_core::SearchQuery> for SearchQuery {
33    type Error = VecGraphError;
34
35    fn try_from(core: vecgraph_core::SearchQuery) -> Result<Self, Self::Error> {
36        let search_kind = SearchKind::from(core.search_kind);
37        let rerank = core.rerank.map(|rp| RerankParams {
38            vectors: rp.vector,
39            kind: rp.kind,
40            weight: rp.weight,
41        });
42
43        Ok(SearchQuery {
44            query_vecs: core.query_vec,
45            top_k: core.top_k as u32,
46            search_kind: search_kind.into(),
47            rerank,
48            exclude_names: core.exclude_names,
49            kind: core.kind,
50            namespace: core.namespace,
51        })
52    }
53}
54
55impl TryFrom<RerankParams> for vecgraph_core::RerankParams {
56    type Error = VecGraphError;
57
58    fn try_from(proto: RerankParams) -> Result<Self, Self::Error> {
59        let vector = proto.vectors;
60
61        Ok(vecgraph_core::RerankParams {
62            vector,
63            kind: proto.kind,
64            weight: proto.weight,
65        })
66    }
67}
68
69impl TryFrom<SearchKind> for vecgraph_core::SearchKind {
70    type Error = VecGraphError;
71
72    fn try_from(proto: SearchKind) -> Result<Self, Self::Error> {
73        match proto {
74            SearchKind::Edge => Ok(vecgraph_core::SearchKind::Edge),
75            SearchKind::Node => Ok(vecgraph_core::SearchKind::Node),
76            SearchKind::AllUnspecified => Ok(vecgraph_core::SearchKind::All),
77        }
78    }
79}
80
81impl TryFrom<SearchResult> for vecgraph_core::SearchResult {
82    type Error = VecGraphError;
83
84    fn try_from(proto: SearchResult) -> Result<Self, Self::Error> {
85        let search_kind = match SearchKind::try_from(proto.hit_kind).ok() {
86            Some(SearchKind::Edge) => vecgraph_core::SearchKind::Edge,
87            Some(SearchKind::Node) => vecgraph_core::SearchKind::Node,
88            Some(SearchKind::AllUnspecified) => vecgraph_core::SearchKind::All,
89            _ => {
90                return Err(VecGraphError::Other(format!(
91                    "Invalid search hit kind: {}",
92                    proto.hit_kind
93                )));
94            }
95        };
96        Ok(vecgraph_core::SearchResult {
97            node_id: proto.node_id.into(),
98            kind: proto.kind,
99            score: proto.score,
100            hit_kind: search_kind,
101        })
102    }
103}
104
105// Reverse
106
107impl From<vecgraph_core::SearchKind> for SearchKind {
108    fn from(core: vecgraph_core::SearchKind) -> Self {
109        match core {
110            vecgraph_core::SearchKind::Edge => SearchKind::Edge,
111            vecgraph_core::SearchKind::Node => SearchKind::Node,
112            vecgraph_core::SearchKind::All => SearchKind::AllUnspecified,
113        }
114    }
115}
116
117impl From<vecgraph_core::SearchResult> for SearchResult {
118    fn from(core: vecgraph_core::SearchResult) -> Self {
119        let hit_kind = SearchKind::from(core.hit_kind);
120        SearchResult {
121            node_id: core.node_id.0,
122            kind: core.kind,
123            score: core.score,
124            hit_kind: hit_kind.into(),
125        }
126    }
127}
128
129impl From<vecgraph_core::RerankParams> for RerankParams {
130    fn from(core: vecgraph_core::RerankParams) -> Self {
131        RerankParams {
132            vectors: core.vector,
133            kind: core.kind,
134            weight: core.weight,
135        }
136    }
137}