reflex/vectordb/
model.rs

1use half::f16;
2use qdrant_client::qdrant::ScoredPoint;
3use qdrant_client::qdrant::point_id::PointIdOptions;
4
5use super::VectorDbError;
6
7#[derive(Debug, Clone)]
8/// Point payload used for upserts.
9pub struct VectorPoint {
10    /// Point id.
11    pub id: u64,
12    /// Vector values.
13    pub vector: Vec<f32>,
14    /// Tenant identifier.
15    pub tenant_id: u64,
16    /// Context hash.
17    pub context_hash: u64,
18    /// Unix timestamp.
19    pub timestamp: i64,
20    /// Optional storage key for loading the full entry.
21    pub storage_key: Option<String>,
22}
23
24impl VectorPoint {
25    /// Creates a point with `timestamp=0` and no storage key.
26    pub fn new(id: u64, vector: Vec<f32>, tenant_id: u64, context_hash: u64) -> Self {
27        Self {
28            id,
29            vector,
30            tenant_id,
31            context_hash,
32            timestamp: 0,
33            storage_key: None,
34        }
35    }
36
37    /// Builds a point from f16 embedding bytes.
38    pub fn from_embedding_bytes(
39        id: u64,
40        embedding_bytes: &[u8],
41        tenant_id: u64,
42        context_hash: u64,
43    ) -> Result<Self, VectorDbError> {
44        let vector = embedding_bytes_to_f32(embedding_bytes)?;
45        Ok(Self::new(id, vector, tenant_id, context_hash))
46    }
47
48    /// Sets the timestamp.
49    pub fn with_timestamp(mut self, timestamp: i64) -> Self {
50        self.timestamp = timestamp;
51        self
52    }
53
54    /// Sets the storage key.
55    pub fn with_storage_key(mut self, key: String) -> Self {
56        self.storage_key = Some(key);
57        self
58    }
59}
60
61#[derive(Debug, Clone)]
62/// Result returned by a vector search.
63pub struct SearchResult {
64    /// Point id.
65    pub id: u64,
66    /// Similarity score (higher is better).
67    pub score: f32,
68    /// Tenant identifier.
69    pub tenant_id: u64,
70    /// Context hash.
71    pub context_hash: u64,
72    /// Unix timestamp.
73    pub timestamp: i64,
74    /// Optional storage key for loading the full entry.
75    pub storage_key: Option<String>,
76}
77
78impl SearchResult {
79    /// Converts a Qdrant `ScoredPoint` into a typed result (returns `None` if unsupported id).
80    pub fn from_scored_point(point: ScoredPoint) -> Option<Self> {
81        let id = match point.id.and_then(|pid| pid.point_id_options) {
82            Some(PointIdOptions::Num(n)) => n,
83            _ => return None,
84        };
85
86        let payload = point.payload;
87
88        let tenant_id = payload
89            .get("tenant_id")
90            .and_then(|v| v.as_integer())
91            .map(|i| i as u64)
92            .unwrap_or(0);
93
94        let context_hash = payload
95            .get("context_hash")
96            .and_then(|v| v.as_integer())
97            .map(|i| i as u64)
98            .unwrap_or(0);
99
100        let timestamp = payload
101            .get("timestamp")
102            .and_then(|v| v.as_integer())
103            .unwrap_or(0);
104
105        let storage_key = payload
106            .get("storage_key")
107            .and_then(|v| v.as_str())
108            .map(|s| s.to_string());
109
110        Some(SearchResult {
111            id,
112            score: point.score,
113            tenant_id,
114            context_hash,
115            timestamp,
116            storage_key,
117        })
118    }
119}
120
121/// Convert little-endian f16 bytes to f32 values.
122pub fn embedding_bytes_to_f32(bytes: &[u8]) -> Result<Vec<f32>, VectorDbError> {
123    if bytes.len() != crate::constants::EMBEDDING_F16_BYTES {
124        return Err(VectorDbError::InvalidEmbeddingBytesLength {
125            expected: crate::constants::EMBEDDING_F16_BYTES,
126            actual: bytes.len(),
127        });
128    }
129
130    if !bytes.len().is_multiple_of(2) {
131        return Err(VectorDbError::InvalidEmbeddingBytesLength {
132            expected: crate::constants::EMBEDDING_F16_BYTES,
133            actual: bytes.len(),
134        });
135    }
136
137    Ok(bytes
138        .chunks_exact(2)
139        .map(|chunk| {
140            let bits = u16::from_le_bytes([chunk[0], chunk[1]]);
141            f16::from_bits(bits).to_f32()
142        })
143        .collect())
144}
145
146/// Convert f32 values to little-endian f16 bytes.
147pub fn f32_to_embedding_bytes(vector: &[f32]) -> Vec<u8> {
148    vector
149        .iter()
150        .flat_map(|&v| f16::from_f32(v).to_le_bytes())
151        .collect()
152}
153
154/// Generate a point id from tenant and context.
155pub fn generate_point_id(tenant_id: u64, context_hash: u64) -> u64 {
156    tenant_id
157        .wrapping_mul(0x517cc1b727220a95)
158        .wrapping_add(context_hash)
159}