1use std::time::{Duration, Instant};
2
3pub use context69_contracts::{
4 AuthMeResponse, BatchGetDocumentsRequest, BatchGetDocumentsResponse,
5 CreateMetadataIndexRequest, DeleteBatchRequest, DocumentChunkResponse, DocumentKey,
6 DocumentResponse, EnsureScopeResponse, FileBatchItem, FileBatchRequest, GenericTaskRequest,
7 GroupKind, GroupResponse, HealthResponse, ImportLibraryFileFromUrlRequest as UrlBatchItem,
8 LibraryFileUploadMetadata as FileMetadata, LibraryTextContentFormat as TextContentFormat,
9 MetadataDataType, MetadataFilter, MetadataFilterOperator, MetadataValueKind,
10 ScopeMetadataIndex, ScopeSpec, SearchRequest, TaskItemResponse, TaskItemStatus,
11 TaskItemsResponse, TaskKind, TaskListQuery, TaskPageResponse, TaskProgress, TaskRef,
12 TaskResponse, TaskRetryResponse, TaskStatus, TextBatchRequest, TranslationDirective,
13 TranslationStatus, UpsertLibraryTextRequest as TextBatchItem, UrlBatchRequest, Visibility,
14};
15use reqwest::Method;
16use uuid::Uuid;
17
18use super::{Context69Client, transport::group_path};
19use crate::Error;
20
21#[derive(Debug, Clone)]
22pub struct WaitOptions {
23 pub timeout: Duration,
24 pub initial_backoff: Duration,
25 pub max_backoff: Duration,
26}
27
28impl Default for WaitOptions {
29 fn default() -> Self {
30 Self {
31 timeout: Duration::from_secs(30 * 60),
32 initial_backoff: Duration::from_millis(100),
33 max_backoff: Duration::from_secs(2),
34 }
35 }
36}
37
38#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
39pub struct CompactSearchHit {
40 pub document_id: i64,
41 pub external_id: String,
42 pub title: String,
43 pub summary: Option<String>,
44 pub source_uri: String,
45 pub published_at: Option<chrono::DateTime<chrono::Utc>>,
46 pub score: f32,
47 pub snippet: String,
48}
49
50#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
51pub struct CompactSearchResponse {
52 pub query: String,
53 pub hits: Vec<CompactSearchHit>,
54}
55
56impl Context69Client {
57 pub async fn ensure_scope(&self, spec: &ScopeSpec) -> Result<EnsureScopeResponse, Error> {
58 self.execute_json(
59 self.authorized_request(Method::POST, "/v1/scopes/ensure")
60 .await?
61 .json(spec),
62 )
63 .await
64 }
65
66 pub async fn text_batch(
67 &self,
68 group_path_value: &str,
69 request: &TextBatchRequest,
70 ) -> Result<TaskRef, Error> {
71 self.submit_batch(group_path(group_path_value, "/batch/text"), request)
72 .await
73 }
74
75 pub async fn url_batch(
76 &self,
77 group_path_value: &str,
78 request: &UrlBatchRequest,
79 ) -> Result<TaskRef, Error> {
80 self.submit_batch(group_path(group_path_value, "/batch/url"), request)
81 .await
82 }
83
84 pub async fn file_batch(
85 &self,
86 group_path_value: &str,
87 request: &FileBatchRequest,
88 ) -> Result<TaskRef, Error> {
89 self.submit_batch(group_path(group_path_value, "/batch/file"), request)
90 .await
91 }
92
93 pub async fn delete_batch(
94 &self,
95 group_path_value: &str,
96 request: &DeleteBatchRequest,
97 ) -> Result<TaskRef, Error> {
98 self.submit_batch(group_path(group_path_value, "/batch/delete"), request)
99 .await
100 }
101
102 pub async fn submit_task(&self, request: &GenericTaskRequest) -> Result<TaskRef, Error> {
103 self.submit_batch("/v1/tasks".to_string(), request).await
104 }
105
106 async fn submit_batch<T: serde::Serialize>(
107 &self,
108 path: String,
109 body: &T,
110 ) -> Result<TaskRef, Error> {
111 let key = idempotency_key(&path, body)?;
112 self.execute_json(
113 self.authorized_request(Method::POST, &path)
114 .await?
115 .header("Idempotency-Key", key)
116 .json(body),
117 )
118 .await
119 }
120
121 pub async fn task(&self, task_id: Uuid) -> Result<TaskResponse, Error> {
122 let path = format!("/v1/tasks/{task_id}");
123 self.execute_json(self.authorized_request(Method::GET, &path).await?)
124 .await
125 }
126
127 pub async fn tasks(&self, query: &TaskListQuery) -> Result<TaskPageResponse, Error> {
128 self.execute_json(
129 self.authorized_request(Method::GET, "/v1/tasks")
130 .await?
131 .query(query),
132 )
133 .await
134 }
135
136 pub async fn task_items(
137 &self,
138 task_id: Uuid,
139 cursor: Option<&str>,
140 ) -> Result<TaskItemsResponse, Error> {
141 let path = format!("/v1/tasks/{task_id}/items");
142 let query = [
143 ("limit", "200".to_string()),
144 ("cursor", cursor.unwrap_or("0").to_string()),
145 ];
146 self.execute_json(
147 self.authorized_request(Method::GET, &path)
148 .await?
149 .query(&query),
150 )
151 .await
152 }
153
154 pub async fn wait(&self, task_id: Uuid, timeout: Duration) -> Result<TaskResponse, Error> {
155 self.wait_with_options(
156 task_id,
157 WaitOptions {
158 timeout,
159 ..WaitOptions::default()
160 },
161 )
162 .await
163 }
164
165 pub async fn wait_with_options(
166 &self,
167 task_id: Uuid,
168 options: WaitOptions,
169 ) -> Result<TaskResponse, Error> {
170 if options.timeout.is_zero()
171 || options.initial_backoff.is_zero()
172 || options.max_backoff.is_zero()
173 {
174 return Err(Error::InvalidTimeout(options.timeout));
175 }
176 let started = Instant::now();
177 let mut delay = options.initial_backoff;
178 loop {
179 if started.elapsed() >= options.timeout {
180 return Err(Error::TaskWaitTimeout {
181 task_id,
182 timeout: options.timeout,
183 });
184 }
185 match self.task(task_id).await {
186 Ok(task) => {
187 if matches!(
188 task.status,
189 context69_contracts::TaskStatus::Succeeded
190 | context69_contracts::TaskStatus::Failed
191 | context69_contracts::TaskStatus::Cancelled
192 ) {
193 return Ok(task);
194 }
195 }
196 Err(error) if error.is_retryable() => {}
197 Err(error) => return Err(error),
198 }
199 tokio::time::sleep(delay.min(options.timeout.saturating_sub(started.elapsed()))).await;
200 delay = (delay + delay).min(options.max_backoff);
201 }
202 }
203
204 pub async fn retry_task(&self, task_id: Uuid) -> Result<TaskRetryResponse, Error> {
205 let path = format!("/v1/tasks/{task_id}/retry");
206 self.execute_json(self.authorized_request(Method::POST, &path).await?)
207 .await
208 }
209
210 pub async fn cancel_task(&self, task_id: Uuid) -> Result<(), Error> {
211 let path = format!("/v1/tasks/{task_id}/cancel");
212 self.execute_empty(self.authorized_request(Method::POST, &path).await?)
213 .await
214 }
215
216 pub async fn search_compact(
217 &self,
218 request: &SearchRequest,
219 ) -> Result<CompactSearchResponse, Error> {
220 let response: context69_contracts::SearchResponse = self
221 .execute_json(
222 self.authorized_request(Method::POST, "/v1/search")
223 .await?
224 .json(request),
225 )
226 .await?;
227 Ok(CompactSearchResponse {
228 query: response.query,
229 hits: response
230 .hits
231 .into_iter()
232 .map(|hit| CompactSearchHit {
233 document_id: hit.document_id,
234 external_id: hit.external_id,
235 title: hit.title,
236 summary: hit.summary,
237 source_uri: hit.source_uri,
238 published_at: hit.published_at,
239 score: hit.score,
240 snippet: hit.chunk_text.chars().take(320).collect(),
241 })
242 .collect(),
243 })
244 }
245
246 pub async fn get_document(
247 &self,
248 document_id: i64,
249 locale: Option<&str>,
250 ) -> Result<DocumentResponse, Error> {
251 let path = format!("/v1/documents/{document_id}");
252 self.execute_json(
253 self.authorized_request(Method::GET, &path)
254 .await?
255 .query(&[("locale", locale.unwrap_or("").to_string())]),
256 )
257 .await
258 }
259
260 pub async fn get_document_by_key(
261 &self,
262 group_path_value: &str,
263 key: &DocumentKey,
264 locale: Option<&str>,
265 ) -> Result<DocumentResponse, Error> {
266 let path = group_path(group_path_value, "/documents/by-external-id");
267 self.execute_json(self.authorized_request(Method::GET, &path).await?.query(&[
268 ("source_key", key.source_key.clone()),
269 ("external_id", key.external_id.clone()),
270 ("locale", locale.unwrap_or("").to_string()),
271 ]))
272 .await
273 }
274
275 pub async fn get_documents(
276 &self,
277 group_path_value: &str,
278 request: &BatchGetDocumentsRequest,
279 ) -> Result<BatchGetDocumentsResponse, Error> {
280 let path = group_path(group_path_value, "/documents/batch-get");
281 self.execute_json(
282 self.authorized_request(Method::POST, &path)
283 .await?
284 .json(request),
285 )
286 .await
287 }
288
289 pub async fn me(&self) -> Result<AuthMeResponse, Error> {
290 self.execute_json(self.authorized_request(Method::GET, "/v1/auth/me").await?)
291 .await
292 }
293
294 pub async fn healthz(&self) -> Result<HealthResponse, Error> {
295 let response = self.client.get(self.url("/healthz")?).send().await?;
296 self.read_json_response(response).await
297 }
298}
299
300fn idempotency_key<T: serde::Serialize>(path: &str, body: &T) -> Result<String, Error> {
301 use sha2::{Digest, Sha256};
302
303 let mut hasher = Sha256::new();
304 hasher.update(path.as_bytes());
305 hasher.update([0]);
306 hasher.update(serde_json::to_vec(body)?);
307 Ok(format!(
308 "ctx69-sdk-{}",
309 hasher
310 .finalize()
311 .iter()
312 .map(|byte| format!("{byte:02x}"))
313 .collect::<String>()
314 ))
315}
316
317#[cfg(test)]
318mod tests {
319 use super::idempotency_key;
320
321 #[test]
322 fn batch_idempotency_key_is_stable_for_the_same_request() {
323 let first = idempotency_key(
324 "/v1/groups/by-path/research/batch/text",
325 &serde_json::json!({"items":[{"external_id":"a"}]}),
326 )
327 .expect("key");
328 let second = idempotency_key(
329 "/v1/groups/by-path/research/batch/text",
330 &serde_json::json!({"items":[{"external_id":"a"}]}),
331 )
332 .expect("key");
333 let different = idempotency_key(
334 "/v1/groups/by-path/research/batch/text",
335 &serde_json::json!({"items":[{"external_id":"b"}]}),
336 )
337 .expect("key");
338 assert_eq!(first, second);
339 assert_ne!(first, different);
340 }
341}