lance_context_api/
lib.rs

1use base64::{engine::general_purpose::STANDARD as BASE64, Engine};
2use chrono::{DateTime, Utc};
3use serde::{Deserialize, Serialize};
4use serde_json::Value;
5use std::future::Future;
6
7// ---------------------------------------------------------------------------
8// Unified error
9// ---------------------------------------------------------------------------
10
11#[derive(Debug, thiserror::Error)]
12pub enum ContextError {
13    #[error("{0}")]
14    NotFound(String),
15    #[error("{0}")]
16    AlreadyExists(String),
17    #[error("{0}")]
18    InvalidRequest(String),
19    #[error("{0}")]
20    Internal(String),
21    #[error("Compaction already in progress")]
22    CompactionInProgress,
23}
24
25pub type ContextResult<T> = Result<T, ContextError>;
26
27// ---------------------------------------------------------------------------
28// Unified trait
29// ---------------------------------------------------------------------------
30
31pub trait ContextStoreApi {
32    fn add(
33        &mut self,
34        records: &[AddRecordRequest],
35    ) -> impl Future<Output = ContextResult<AddRecordsResponse>> + Send;
36
37    fn get(&self, id: &str) -> impl Future<Output = ContextResult<Option<RecordDto>>> + Send;
38
39    fn list(
40        &self,
41        limit: Option<usize>,
42        offset: Option<usize>,
43    ) -> impl Future<Output = ContextResult<Vec<RecordDto>>> + Send;
44
45    fn search(
46        &self,
47        query: &[f32],
48        limit: Option<usize>,
49    ) -> impl Future<Output = ContextResult<Vec<SearchResultDto>>> + Send;
50
51    fn version(&self) -> u64;
52
53    fn checkout(&mut self, version: u64) -> impl Future<Output = ContextResult<()>> + Send;
54
55    fn compact(
56        &mut self,
57        options: Option<CompactRequest>,
58    ) -> impl Future<Output = ContextResult<CompactResponse>> + Send;
59
60    fn compaction_stats(&self) -> impl Future<Output = ContextResult<CompactStatsResponse>> + Send;
61}
62
63// ---------------------------------------------------------------------------
64// Context lifecycle
65// ---------------------------------------------------------------------------
66
67#[derive(Debug, Serialize, Deserialize)]
68pub struct CreateContextRequest {
69    pub name: String,
70    #[serde(default)]
71    pub storage_options: Option<std::collections::HashMap<String, String>>,
72    #[serde(default)]
73    pub id_index_type: Option<String>,
74    #[serde(default)]
75    pub blob_columns: Option<Vec<String>>,
76}
77
78#[derive(Debug, Serialize, Deserialize)]
79pub struct ContextInfo {
80    pub name: String,
81    pub uri: String,
82    pub version: u64,
83}
84
85#[derive(Debug, Serialize, Deserialize)]
86pub struct ListContextsResponse {
87    pub contexts: Vec<ContextInfo>,
88}
89
90// ---------------------------------------------------------------------------
91// Records
92// ---------------------------------------------------------------------------
93
94#[derive(Debug, Clone, Serialize, Deserialize)]
95pub struct StateMetadataDto {
96    #[serde(default, skip_serializing_if = "Option::is_none")]
97    pub step: Option<i32>,
98    #[serde(default, skip_serializing_if = "Option::is_none")]
99    pub active_plan_id: Option<String>,
100    #[serde(default, skip_serializing_if = "Option::is_none")]
101    pub tokens_used: Option<i32>,
102    #[serde(default, skip_serializing_if = "Option::is_none")]
103    pub custom: Option<String>,
104}
105
106#[derive(Debug, Clone, Default, Serialize, Deserialize)]
107pub struct AddRecordRequest {
108    #[serde(default = "default_role")]
109    pub role: String,
110    #[serde(default = "default_content_type")]
111    pub content_type: String,
112    #[serde(default, skip_serializing_if = "Option::is_none")]
113    pub text_payload: Option<String>,
114    #[serde(
115        default,
116        skip_serializing_if = "Option::is_none",
117        serialize_with = "serialize_base64_opt",
118        deserialize_with = "deserialize_base64_opt"
119    )]
120    pub binary_payload: Option<Vec<u8>>,
121    #[serde(default, skip_serializing_if = "Option::is_none")]
122    pub embedding: Option<Vec<f32>>,
123    #[serde(default, skip_serializing_if = "Option::is_none")]
124    pub bot_id: Option<String>,
125    #[serde(default, skip_serializing_if = "Option::is_none")]
126    pub session_id: Option<String>,
127    #[serde(default, skip_serializing_if = "Option::is_none")]
128    pub external_id: Option<String>,
129    #[serde(default, skip_serializing_if = "Option::is_none")]
130    pub state_metadata: Option<StateMetadataDto>,
131    #[serde(default, skip_serializing_if = "Option::is_none")]
132    pub metadata: Option<Value>,
133    #[serde(default, skip_serializing_if = "Option::is_none")]
134    pub expires_at: Option<DateTime<Utc>>,
135    #[serde(default, skip_serializing_if = "Option::is_none")]
136    pub retention_policy: Option<String>,
137    #[serde(default, skip_serializing_if = "Option::is_none")]
138    pub supersedes_id: Option<String>,
139}
140
141#[derive(Debug, Serialize, Deserialize)]
142pub struct AddRecordsRequest {
143    pub records: Vec<AddRecordRequest>,
144}
145
146#[derive(Debug, Serialize, Deserialize)]
147pub struct AddRecordsResponse {
148    pub version: u64,
149    pub ids: Vec<String>,
150    pub count: usize,
151}
152
153#[derive(Debug, Clone, Serialize, Deserialize)]
154pub struct RecordDto {
155    pub id: String,
156    #[serde(default, skip_serializing_if = "Option::is_none")]
157    pub external_id: Option<String>,
158    pub run_id: String,
159    #[serde(default, skip_serializing_if = "Option::is_none")]
160    pub bot_id: Option<String>,
161    #[serde(default, skip_serializing_if = "Option::is_none")]
162    pub session_id: Option<String>,
163    pub created_at: DateTime<Utc>,
164    pub role: String,
165    pub content_type: String,
166    #[serde(default, skip_serializing_if = "Option::is_none")]
167    pub text_payload: Option<String>,
168    #[serde(
169        default,
170        skip_serializing_if = "Option::is_none",
171        serialize_with = "serialize_base64_opt",
172        deserialize_with = "deserialize_base64_opt"
173    )]
174    pub binary_payload: Option<Vec<u8>>,
175    #[serde(default, skip_serializing_if = "Option::is_none")]
176    pub embedding: Option<Vec<f32>>,
177    #[serde(default, skip_serializing_if = "Option::is_none")]
178    pub state_metadata: Option<StateMetadataDto>,
179    #[serde(default, skip_serializing_if = "Option::is_none")]
180    pub metadata: Option<Value>,
181    #[serde(default, skip_serializing_if = "Option::is_none")]
182    pub expires_at: Option<DateTime<Utc>>,
183    #[serde(default, skip_serializing_if = "Option::is_none")]
184    pub retention_policy: Option<String>,
185    pub lifecycle_status: String,
186    #[serde(default, skip_serializing_if = "Option::is_none")]
187    pub retired_at: Option<DateTime<Utc>>,
188    #[serde(default, skip_serializing_if = "Option::is_none")]
189    pub retired_reason: Option<String>,
190    #[serde(default, skip_serializing_if = "Option::is_none")]
191    pub supersedes_id: Option<String>,
192    #[serde(default, skip_serializing_if = "Option::is_none")]
193    pub superseded_by_id: Option<String>,
194}
195
196#[derive(Debug, Serialize, Deserialize)]
197pub struct ListRecordsResponse {
198    pub records: Vec<RecordDto>,
199}
200
201// ---------------------------------------------------------------------------
202// Single record lookup
203// ---------------------------------------------------------------------------
204
205#[derive(Debug, Serialize, Deserialize)]
206pub struct GetRecordResponse {
207    pub record: Option<RecordDto>,
208}
209
210// ---------------------------------------------------------------------------
211// Search
212// ---------------------------------------------------------------------------
213
214#[derive(Debug, Serialize, Deserialize)]
215pub struct SearchRequest {
216    pub query: Vec<f32>,
217    #[serde(default = "default_search_limit")]
218    pub limit: usize,
219}
220
221#[derive(Debug, Serialize, Deserialize)]
222pub struct SearchResultDto {
223    pub record: RecordDto,
224    pub distance: f32,
225}
226
227#[derive(Debug, Serialize, Deserialize)]
228pub struct SearchResponse {
229    pub results: Vec<SearchResultDto>,
230}
231
232// ---------------------------------------------------------------------------
233// Versioning
234// ---------------------------------------------------------------------------
235
236#[derive(Debug, Serialize, Deserialize)]
237pub struct VersionResponse {
238    pub version: u64,
239}
240
241#[derive(Debug, Serialize, Deserialize)]
242pub struct CheckoutRequest {
243    pub version: u64,
244}
245
246// ---------------------------------------------------------------------------
247// Compaction
248// ---------------------------------------------------------------------------
249
250#[derive(Debug, Default, Serialize, Deserialize)]
251pub struct CompactRequest {
252    #[serde(default, skip_serializing_if = "Option::is_none")]
253    pub target_rows_per_fragment: Option<usize>,
254    #[serde(default, skip_serializing_if = "Option::is_none")]
255    pub materialize_deletions: Option<bool>,
256}
257
258#[derive(Debug, Serialize, Deserialize)]
259pub struct CompactResponse {
260    pub fragments_removed: usize,
261    pub fragments_added: usize,
262    pub files_removed: usize,
263    pub files_added: usize,
264}
265
266#[derive(Debug, Serialize, Deserialize)]
267pub struct CompactStatsResponse {
268    pub total_fragments: usize,
269    pub is_compacting: bool,
270    #[serde(default, skip_serializing_if = "Option::is_none")]
271    pub last_compaction: Option<DateTime<Utc>>,
272    #[serde(default, skip_serializing_if = "Option::is_none")]
273    pub last_error: Option<String>,
274    pub total_compactions: u64,
275}
276
277// ---------------------------------------------------------------------------
278// Error
279// ---------------------------------------------------------------------------
280
281#[derive(Debug, Serialize, Deserialize)]
282pub struct ErrorBody {
283    pub code: String,
284    pub message: String,
285}
286
287#[derive(Debug, Serialize, Deserialize)]
288pub struct ErrorResponse {
289    pub error: ErrorBody,
290}
291
292// ---------------------------------------------------------------------------
293// Helpers
294// ---------------------------------------------------------------------------
295
296fn default_content_type() -> String {
297    "text/plain".to_string()
298}
299
300fn default_role() -> String {
301    "user".to_string()
302}
303
304fn default_search_limit() -> usize {
305    10
306}
307
308fn serialize_base64_opt<S>(data: &Option<Vec<u8>>, serializer: S) -> Result<S::Ok, S::Error>
309where
310    S: serde::Serializer,
311{
312    match data {
313        Some(bytes) => serializer.serialize_some(&BASE64.encode(bytes)),
314        None => serializer.serialize_none(),
315    }
316}
317
318fn deserialize_base64_opt<'de, D>(deserializer: D) -> Result<Option<Vec<u8>>, D::Error>
319where
320    D: serde::Deserializer<'de>,
321{
322    let opt: Option<String> = Option::deserialize(deserializer)?;
323    match opt {
324        Some(s) => BASE64
325            .decode(&s)
326            .map(Some)
327            .map_err(serde::de::Error::custom),
328        None => Ok(None),
329    }
330}
lance_context_api/lib.rs

lance_context_api/
lib.rs
