watsonx_rs/orchestrate/
collection.rs

1//! Document collection management operations
2
3use crate::error::{Error, Result};
4use super::types::{DocumentCollection, Document, SearchRequest, SearchResponse};
5use super::OrchestrateClient;
6
7impl OrchestrateClient {
8    /// List all document collections
9    pub async fn list_collections(&self) -> Result<Vec<DocumentCollection>> {
10        let api_key = self.access_token.as_ref().ok_or_else(|| {
11            Error::Authentication("Not authenticated. Set access token (API key) first.".to_string())
12        })?;
13
14        let base_url = self.config.get_base_url();
15        let url = format!("{}/collections", base_url);
16
17        let response = self
18            .client
19            .get(&url)
20            .header("Authorization", format!("Bearer {}", api_key))
21            .header("Content-Type", "application/json")
22            .send()
23            .await
24            .map_err(|e| Error::Network(e.to_string()))?;
25
26        if !response.status().is_success() {
27            let status = response.status();
28            let error_text = response
29                .text()
30                .await
31                .unwrap_or_else(|_| "Unknown error".to_string());
32            return Err(Error::Api(format!(
33                "Failed to list collections: {} - {}",
34                status, error_text
35            )));
36        }
37
38        let text = response
39            .text()
40            .await
41            .map_err(|e| Error::Serialization(e.to_string()))?;
42
43        if let Ok(collections) = serde_json::from_str::<Vec<DocumentCollection>>(&text) {
44            return Ok(collections);
45        }
46
47        if let Ok(obj) = serde_json::from_str::<serde_json::Value>(&text) {
48            if let Some(collections_array) = obj.get("collections").and_then(|c| c.as_array()) {
49                let collections: Result<Vec<DocumentCollection>> = collections_array
50                    .iter()
51                    .map(|coll| {
52                        serde_json::from_value::<DocumentCollection>(coll.clone())
53                            .map_err(|e| Error::Serialization(e.to_string()))
54                    })
55                    .collect();
56                return collections;
57            }
58        }
59
60        Ok(Vec::new())
61    }
62
63    /// Get a specific document collection
64    pub async fn get_collection(&self, collection_id: &str) -> Result<DocumentCollection> {
65        let api_key = self.access_token.as_ref().ok_or_else(|| {
66            Error::Authentication("Not authenticated. Set access token (API key) first.".to_string())
67        })?;
68
69        let base_url = self.config.get_base_url();
70        let url = format!("{}/collections/{}", base_url, collection_id);
71
72        let response = self
73            .client
74            .get(&url)
75            .header("Authorization", format!("Bearer {}", api_key))
76            .header("Content-Type", "application/json")
77            .send()
78            .await
79            .map_err(|e| Error::Network(e.to_string()))?;
80
81        if !response.status().is_success() {
82            let status = response.status();
83            let error_text = response
84                .text()
85                .await
86                .unwrap_or_else(|_| "Unknown error".to_string());
87            return Err(Error::Api(format!(
88                "Failed to get collection {}: {} - {}",
89                collection_id, status, error_text
90            )));
91        }
92
93        let collection: DocumentCollection = response
94            .json()
95            .await
96            .map_err(|e| Error::Serialization(e.to_string()))?;
97
98        Ok(collection)
99    }
100
101    /// Get a specific document from a collection
102    pub async fn get_document(&self, collection_id: &str, document_id: &str) -> Result<Document> {
103        let api_key = self.access_token.as_ref().ok_or_else(|| {
104            Error::Authentication("Not authenticated. Set access token (API key) first.".to_string())
105        })?;
106
107        let base_url = self.config.get_base_url();
108        let url = format!("{}/collections/{}/documents/{}", base_url, collection_id, document_id);
109
110        let response = self
111            .client
112            .get(&url)
113            .header("Authorization", format!("Bearer {}", api_key))
114            .header("Content-Type", "application/json")
115            .send()
116            .await
117            .map_err(|e| Error::Network(e.to_string()))?;
118
119        if !response.status().is_success() {
120            let status = response.status();
121            let error_text = response
122                .text()
123                .await
124                .unwrap_or_else(|_| "Unknown error".to_string());
125            return Err(Error::Api(format!(
126                "Failed to get document {}: {} - {}",
127                document_id, status, error_text
128            )));
129        }
130
131        let document: Document = response
132            .json()
133            .await
134            .map_err(|e| Error::Serialization(e.to_string()))?;
135
136        Ok(document)
137    }
138
139    /// Delete a document from a collection
140    pub async fn delete_document(&self, collection_id: &str, document_id: &str) -> Result<()> {
141        let api_key = self.access_token.as_ref().ok_or_else(|| {
142            Error::Authentication("Not authenticated. Set access token (API key) first.".to_string())
143        })?;
144
145        let base_url = self.config.get_base_url();
146        let url = format!("{}/collections/{}/documents/{}", base_url, collection_id, document_id);
147
148        let response = self
149            .client
150            .delete(&url)
151            .header("Authorization", format!("Bearer {}", api_key))
152            .send()
153            .await
154            .map_err(|e| Error::Network(e.to_string()))?;
155
156        if !response.status().is_success() {
157            let status = response.status();
158            let error_text = response
159                .text()
160                .await
161                .unwrap_or_else(|_| "Unknown error".to_string());
162            return Err(Error::Api(format!(
163                "Failed to delete document {}: {} - {}",
164                document_id, status, error_text
165            )));
166        }
167
168        Ok(())
169    }
170
171    /// Search documents in a collection
172    pub async fn search_documents(&self, collection_id: &str, request: SearchRequest) -> Result<SearchResponse> {
173        let api_key = self.access_token.as_ref().ok_or_else(|| {
174            Error::Authentication("Not authenticated. Set access token (API key) first.".to_string())
175        })?;
176
177        let base_url = self.config.get_base_url();
178        let url = format!("{}/collections/{}/search", base_url, collection_id);
179
180        let response = self
181            .client
182            .post(&url)
183            .header("Authorization", format!("Bearer {}", api_key))
184            .header("Content-Type", "application/json")
185            .json(&request)
186            .send()
187            .await
188            .map_err(|e| Error::Network(e.to_string()))?;
189
190        if !response.status().is_success() {
191            let status = response.status();
192            let error_text = response
193                .text()
194                .await
195                .unwrap_or_else(|_| "Unknown error".to_string());
196            return Err(Error::Api(format!(
197                "Failed to search documents: {} - {}",
198                status, error_text
199            )));
200        }
201
202        let search_response: SearchResponse = response
203            .json()
204            .await
205            .map_err(|e| Error::Serialization(e.to_string()))?;
206
207        Ok(search_response)
208    }
209}
watsonx_rs/orchestrate/collection.rs

watsonx_rs/orchestrate/
collection.rs
