xai_rust/models/
collection.rs

1//! Collections API types for document storage and retrieval.
2
3use serde::{Deserialize, Serialize};
4
5/// A collection for storing and searching documents.
6#[derive(Debug, Clone, Deserialize)]
7pub struct Collection {
8    /// Unique identifier for the collection.
9    pub id: String,
10    /// Name of the collection.
11    pub name: String,
12    /// Description of the collection.
13    #[serde(default)]
14    pub description: Option<String>,
15    /// Number of documents in the collection.
16    #[serde(default)]
17    pub document_count: u64,
18    /// Creation timestamp.
19    #[serde(default)]
20    pub created_at: Option<i64>,
21    /// Last updated timestamp.
22    #[serde(default)]
23    pub updated_at: Option<i64>,
24}
25
26/// Response from listing collections.
27#[derive(Debug, Clone, Deserialize)]
28pub struct CollectionListResponse {
29    /// List of collections.
30    pub data: Vec<Collection>,
31    /// Pagination token for next page.
32    #[serde(default)]
33    pub next_token: Option<String>,
34}
35
36/// A document in a collection.
37#[derive(Debug, Clone, Serialize, Deserialize)]
38pub struct Document {
39    /// Unique identifier for the document.
40    #[serde(skip_serializing_if = "Option::is_none")]
41    pub id: Option<String>,
42    /// Content of the document.
43    pub content: String,
44    /// Optional metadata.
45    #[serde(default, skip_serializing_if = "Option::is_none")]
46    pub metadata: Option<serde_json::Value>,
47}
48
49/// Request to update a collection.
50#[derive(Debug, Clone, Serialize)]
51pub struct UpdateCollectionRequest {
52    /// Optional name update.
53    #[serde(skip_serializing_if = "Option::is_none")]
54    pub name: Option<String>,
55    /// Optional description update.
56    #[serde(skip_serializing_if = "Option::is_none")]
57    pub description: Option<String>,
58}
59
60impl UpdateCollectionRequest {
61    /// Create a new empty update request.
62    pub fn new() -> Self {
63        Self {
64            name: None,
65            description: None,
66        }
67    }
68
69    /// Set collection name.
70    pub fn name(mut self, name: impl Into<String>) -> Self {
71        self.name = Some(name.into());
72        self
73    }
74
75    /// Set collection description.
76    pub fn description(mut self, description: impl Into<String>) -> Self {
77        self.description = Some(description.into());
78        self
79    }
80}
81
82impl Default for UpdateCollectionRequest {
83    fn default() -> Self {
84        Self::new()
85    }
86}
87
88/// Request for fetching documents in a batch.
89#[derive(Debug, Clone, Serialize)]
90pub struct BatchGetDocumentsRequest {
91    /// Target document identifiers.
92    pub ids: Vec<String>,
93}
94
95impl BatchGetDocumentsRequest {
96    /// Create a new batch get request.
97    pub fn new(ids: Vec<String>) -> Self {
98        Self { ids }
99    }
100
101    /// Create from a slice of IDs.
102    pub fn from_ids(ids: &[&str]) -> Self {
103        Self {
104            ids: ids.iter().map(|id| (*id).to_string()).collect(),
105        }
106    }
107}
108
109impl Document {
110    /// Create a new document with content.
111    pub fn new(content: impl Into<String>) -> Self {
112        Self {
113            id: None,
114            content: content.into(),
115            metadata: None,
116        }
117    }
118
119    /// Create a document with a custom ID.
120    pub fn with_id(id: impl Into<String>, content: impl Into<String>) -> Self {
121        Self {
122            id: Some(id.into()),
123            content: content.into(),
124            metadata: None,
125        }
126    }
127
128    /// Set metadata for the document.
129    pub fn metadata(mut self, metadata: serde_json::Value) -> Self {
130        self.metadata = Some(metadata);
131        self
132    }
133}
134
135/// Response from listing documents in a collection.
136#[derive(Debug, Clone, Deserialize)]
137pub struct DocumentListResponse {
138    /// List of documents.
139    pub data: Vec<Document>,
140    /// Pagination token for next page.
141    #[serde(default)]
142    pub next_token: Option<String>,
143}
144
145/// Request to create a collection.
146#[derive(Debug, Clone, Serialize)]
147pub struct CreateCollectionRequest {
148    /// Name of the collection.
149    pub name: String,
150    /// Optional description.
151    #[serde(skip_serializing_if = "Option::is_none")]
152    pub description: Option<String>,
153}
154
155impl CreateCollectionRequest {
156    /// Create a new request with a name.
157    pub fn new(name: impl Into<String>) -> Self {
158        Self {
159            name: name.into(),
160            description: None,
161        }
162    }
163
164    /// Set the description.
165    pub fn description(mut self, description: impl Into<String>) -> Self {
166        self.description = Some(description.into());
167        self
168    }
169}
170
171/// Request to add documents to a collection.
172#[derive(Debug, Clone, Serialize)]
173pub struct AddDocumentsRequest {
174    /// Documents to add.
175    pub documents: Vec<Document>,
176}
177
178/// Response from adding documents.
179#[derive(Debug, Clone, Deserialize)]
180pub struct AddDocumentsResponse {
181    /// IDs of the added documents.
182    pub ids: Vec<String>,
183}
184
185/// Search request for querying a collection.
186#[derive(Debug, Clone, Serialize)]
187pub struct SearchRequest {
188    /// The search query.
189    pub query: String,
190    /// Maximum number of results to return.
191    #[serde(skip_serializing_if = "Option::is_none")]
192    pub limit: Option<u32>,
193    /// Minimum similarity score threshold.
194    #[serde(skip_serializing_if = "Option::is_none")]
195    pub score_threshold: Option<f32>,
196}
197
198impl SearchRequest {
199    /// Create a new search request.
200    pub fn new(query: impl Into<String>) -> Self {
201        Self {
202            query: query.into(),
203            limit: None,
204            score_threshold: None,
205        }
206    }
207
208    /// Set the maximum number of results.
209    pub fn limit(mut self, limit: u32) -> Self {
210        self.limit = Some(limit);
211        self
212    }
213
214    /// Set the minimum similarity score threshold.
215    pub fn score_threshold(mut self, threshold: f32) -> Self {
216        self.score_threshold = Some(threshold);
217        self
218    }
219}
220
221/// A search result from querying a collection.
222#[derive(Debug, Clone, Deserialize)]
223pub struct SearchResult {
224    /// The matched document.
225    pub document: Document,
226    /// Similarity score.
227    pub score: f32,
228}
229
230/// Response from searching a collection.
231#[derive(Debug, Clone, Deserialize)]
232pub struct SearchResponse {
233    /// Search results.
234    pub results: Vec<SearchResult>,
235}
236
237#[cfg(test)]
238mod tests {
239    use super::*;
240    use serde_json::json;
241
242    #[test]
243    fn document_builder_sets_optional_fields() {
244        let document = Document::new("simple content");
245        assert_eq!(document.id, None);
246        assert_eq!(document.content, "simple content");
247        assert!(document.metadata.is_none());
248
249        let with_metadata = Document::with_id("doc-1", "metadata content").metadata(json!({
250            "author": "tests",
251            "version": 1
252        }));
253        assert_eq!(with_metadata.id.as_deref(), Some("doc-1"));
254        assert_eq!(with_metadata.content, "metadata content");
255        assert_eq!(
256            with_metadata.metadata.as_ref().unwrap()["author"].as_str(),
257            Some("tests")
258        );
259    }
260
261    #[test]
262    fn collection_request_builder_populates_description() {
263        let request = CreateCollectionRequest::new("research").description("notes corpus");
264        assert_eq!(request.name, "research");
265        assert_eq!(request.description.as_deref(), Some("notes corpus"));
266    }
267
268    #[test]
269    fn search_request_builder_sets_options() {
270        let request = SearchRequest::new("query term")
271            .limit(10)
272            .score_threshold(0.85);
273        assert_eq!(request.query, "query term");
274        assert_eq!(request.limit, Some(10));
275        assert_eq!(request.score_threshold, Some(0.85));
276    }
277
278    #[test]
279    fn collection_update_request_builder() {
280        let request = UpdateCollectionRequest::new()
281            .name("research")
282            .description("notes");
283        assert_eq!(request.name.as_deref(), Some("research"));
284        assert_eq!(request.description.as_deref(), Some("notes"));
285    }
286
287    #[test]
288    fn batch_get_documents_request_from_ids() {
289        let request = BatchGetDocumentsRequest::from_ids(&["d1", "d2"]);
290        assert_eq!(request.ids, vec!["d1".to_string(), "d2".to_string()]);
291    }
292}
xai_rust/models/collection.rs

xai_rust/models/
collection.rs
