toast_api/
api.rs

1use crate::config::BASE_URL;
2use anyhow::{anyhow, Result};
3use chrono::Local;
4use rquest::header::{HeaderMap, HeaderValue};
5use rquest::Client;
6use rquest_util::Emulation;
7use serde_json::json;
8use std::time::Duration;
9use uuid::Uuid;
10
11/// Session data for Claude API
12#[derive(Clone)]
13pub struct Session {
14    pub cookie: String,
15    pub user_agent: String,
16    pub organization_id: String,
17}
18/// Representation of a file attachment for send_message
19#[derive(Clone, Debug)]
20pub struct Attachment {
21    /// File name (basename)
22    pub file_name: String,
23    /// File size in bytes
24    pub size: u64,
25    /// File content as text
26    pub content: String,
27}
28
29/// Claude API client
30pub struct Claude {
31    http: Client,
32    session: Session,
33    model: &'static str,
34}
35
36impl Claude {
37    /// Create a new Claude client
38    pub fn new(session: Session, model: &'static str) -> Result<Self> {
39        let http = Client::builder()
40            .emulation(Emulation::Firefox136)
41            .timeout(Duration::from_secs(240))
42            // Add retry settings to handle transient errors
43            .connect_timeout(Duration::from_secs(30))
44            .build()?;
45        Ok(Self {
46            http,
47            session,
48            model,
49        })
50    }
51
52    fn default_headers(&self) -> HeaderMap {
53        let mut headers = HeaderMap::new();
54        headers.insert("Host", HeaderValue::from_static("claude.ai"));
55        headers.insert(
56            "User-Agent",
57            HeaderValue::from_str(&self.session.user_agent).unwrap(),
58        );
59        headers
60    }
61
62    /// Create a new chat conversation
63    pub async fn create_chat(&self) -> Result<String> {
64        let uuid = Uuid::new_v4().to_string();
65        let url = format!(
66            "{}/api/organizations/{}/chat_conversations",
67            BASE_URL, self.session.organization_id
68        );
69        let body = json!({ "name": "", "uuid": uuid }).to_string();
70
71        // Add retry logic
72        let mut attempts = 0;
73        let max_attempts = 3;
74
75        while attempts < max_attempts {
76            attempts += 1;
77
78            match self
79                .http
80                .post(&url)
81                .headers(self.default_headers())
82                .header("Content-Type", "application/json")
83                .header("Cookie", &self.session.cookie)
84                .body(body.clone())
85                .send()
86                .await
87            {
88                Ok(res) => {
89                    let status = res.status();
90
91                    if status == 201 {
92                        return Ok(uuid);
93                    } else {
94                        match res.text().await {
95                            Ok(text) => {
96                                if attempts == max_attempts {
97                                    return Err(anyhow!("create_chat failed: {:?}", text));
98                                }
99                            }
100                            Err(e) => {
101                                if attempts == max_attempts {
102                                    return Err(anyhow!(
103                                        "create_chat failed: unable to read response: {}",
104                                        e
105                                    ));
106                                }
107                            }
108                        }
109                    }
110                }
111                Err(e) => {
112                    if attempts == max_attempts {
113                        return Err(anyhow!("create_chat failed: network error: {}", e));
114                    }
115                }
116            }
117
118            // Wait before retrying
119            tokio::time::sleep(Duration::from_secs(1)).await;
120        }
121
122        Err(anyhow!(
123            "create_chat failed after {} attempts",
124            max_attempts
125        ))
126    }
127
128    /// Delete an existing chat conversation
129    pub async fn delete_chat(&self, chat_id: &str) -> Result<()> {
130        let url = format!(
131            "{}/api/organizations/{}/chat_conversations/{}",
132            BASE_URL, self.session.organization_id, chat_id
133        );
134
135        let body = format!("\"{}\"", chat_id);
136
137        // Add retry logic
138        let mut attempts = 0;
139        let max_attempts = 3;
140
141        while attempts < max_attempts {
142            attempts += 1;
143
144            match self
145                .http
146                .delete(&url)
147                .headers(self.default_headers())
148                .header("Content-Type", "application/json")
149                .header("Cookie", &self.session.cookie)
150                .body(body.clone())
151                .send()
152                .await
153            {
154                Ok(res) => {
155                    let status = res.status();
156
157                    if status == 204 {
158                        return Ok(());
159                    } else {
160                        // Only fail on final attempt
161                        if attempts == max_attempts {
162                            match res.text().await {
163                                Ok(text) => {
164                                    return Err(anyhow!("delete_chat failed: {:?}", text));
165                                }
166                                Err(e) => {
167                                    return Err(anyhow!(
168                                        "delete_chat failed: unable to read response: {}",
169                                        e
170                                    ));
171                                }
172                            }
173                        }
174                    }
175                }
176                Err(e) => {
177                    if attempts == max_attempts {
178                        return Err(anyhow!("delete_chat failed: network error: {}", e));
179                    }
180                }
181            }
182
183            // Wait before retrying
184            tokio::time::sleep(Duration::from_secs(1)).await;
185        }
186
187        Err(anyhow!(
188            "delete_chat failed after {} attempts",
189            max_attempts
190        ))
191    }
192
193    /// Send a message to the chat and return the completion
194    /// Send a message to the chat with optional attachments and return the completion
195    pub async fn send_message(
196        &self,
197        chat_id: &str,
198        prompt: &str,
199        attachments: &[Attachment],
200    ) -> Result<String> {
201        let url = format!(
202            "{}/api/organizations/{}/chat_conversations/{}/completion",
203            BASE_URL, self.session.organization_id, chat_id
204        );
205        let mut payload = json!({
206            "attachments": [],
207            "files": [],
208            "prompt": prompt,
209            "model": self.model,
210            "timezone": Local::now().offset().to_string()
211        });
212        if !attachments.is_empty() {
213            let atts: Vec<_> = attachments
214                .iter()
215                .map(|a| {
216                    json!({
217                        "extracted_content": a.content,
218                        "file_name": a.file_name,
219                        "file_size": a.size.to_string(),
220                        "file_type": "text/plain"
221                    })
222                })
223                .collect();
224            payload["attachments"] = serde_json::Value::Array(atts);
225        }
226
227        // Add retry logic for network errors
228        let mut attempts = 0;
229        let max_attempts = 3;
230        let mut last_error = None;
231
232        while attempts < max_attempts {
233            attempts += 1;
234
235            match self
236                .http
237                .post(&url)
238                .headers(self.default_headers())
239                .header("Content-Type", "application/json")
240                .header("Accept", "text/event-stream, text/event-stream")
241                .header("Cookie", &self.session.cookie)
242                .body(payload.to_string())
243                .send()
244                .await
245            {
246                Ok(res) => {
247                    let status = res.status();
248
249                    // Check status before proceeding
250                    if !status.is_success() {
251                        // Try to get more details about the error
252                        let error_text = match res.text().await {
253                            Ok(text) => text,
254                            Err(_) => "Unable to read error response".to_string(),
255                        };
256                        return Err(anyhow!(
257                            "API request failed with status: {} - {}",
258                            status,
259                            error_text
260                        ));
261                    }
262
263                    // Process the body
264                    match super::utils::decode_body(res).await {
265                        Ok(bytes) => {
266                            let _preview = std::str::from_utf8(&bytes[..bytes.len().min(500)])
267                                .unwrap_or("Invalid UTF-8")
268                                .replace('\n', "\\n");
269                            // eprintln!("RESPONSE BODY: {} (first 500 bytes)", preview);
270
271                            let result = super::utils::parse_stream(&bytes)?;
272                            return Ok(result);
273                        }
274                        Err(e) => {
275                            eprintln!("RESPONSE DECODE ERROR: {}", e);
276                            last_error = Some(e);
277                            if attempts < max_attempts {
278                                // Wait before retrying
279                                tokio::time::sleep(Duration::from_secs(1)).await;
280                                continue;
281                            }
282                        }
283                    }
284                }
285                Err(e) => {
286                    eprintln!("NETWORK ERROR: {}", e);
287                    last_error = Some(anyhow!("Network error: {}", e));
288                    if attempts < max_attempts {
289                        // Wait before retrying
290                        tokio::time::sleep(Duration::from_secs(1)).await;
291                        continue;
292                    }
293                }
294            }
295        }
296
297        // All attempts failed
298        Err(last_error
299            .unwrap_or_else(|| anyhow!("Failed to connect to API after {} attempts", max_attempts)))
300    }
301}
toast_api/api.rs

toast_api/
api.rs
