watsonx_rs/orchestrate/
client.rs

1//! WatsonX Orchestrate client implementation
2//!
3//! This module provides the main client for interacting with WatsonX Orchestrate services,
4//! including custom assistants, document collections, and chat functionality.
5
6use crate::error::{Error, Result};
7use super::types::*;
8use super::config::OrchestrateConfig;
9use reqwest::{Client, ClientBuilder};
10use serde_json::Value;
11use std::collections::HashMap;
12use std::time::Duration;
13
14/// WatsonX Orchestrate client for managing custom assistants and document collections
15pub struct OrchestrateClient {
16    pub(crate) config: OrchestrateConfig,
17    pub(crate) access_token: Option<String>,
18    pub(crate) client: Client,
19}
20
21impl OrchestrateClient {
22    /// Create a new Orchestrate client (matches wxo-client-main pattern)
23    pub fn new(config: OrchestrateConfig) -> Self {
24        let client = ClientBuilder::new()
25            .timeout(Duration::from_secs(300)) // 5 minute timeout for streaming
26            .tcp_keepalive(Duration::from_secs(60))
27            .http1_title_case_headers()
28            .build()
29            .unwrap_or_else(|_| Client::new());
30
31        Self {
32            config,
33            access_token: None,
34            client,
35        }
36    }
37
38    /// Set the access token for authentication
39    pub fn with_token(mut self, token: String) -> Self {
40        self.access_token = Some(token);
41        self
42    }
43
44    /// Set the access token for authentication (mutable)
45    pub fn set_token(&mut self, token: String) {
46        self.access_token = Some(token);
47    }
48
49    /// Get the current configuration
50    pub fn config(&self) -> &OrchestrateConfig {
51        &self.config
52    }
53
54    /// Check if authenticated
55    pub fn is_authenticated(&self) -> bool {
56        self.access_token.is_some()
57    }
58
59    /// Generate IAM Access Token from Watson Orchestrate API key
60    /// This is required for Watson Orchestrate SaaS authentication
61    pub async fn generate_jwt_token(api_key: &str) -> Result<String> {
62        let client = reqwest::Client::new();
63        let body = format!(
64            "grant_type=urn:ibm:params:oauth:grant-type:apikey&apikey={}",
65            api_key
66        );
67
68        let response = client
69            .post("https://iam.cloud.ibm.com/identity/token")
70            .header("Content-Type", "application/x-www-form-urlencoded")
71            .body(body)
72            .send()
73            .await
74            .map_err(|e| Error::Network(format!("Failed to generate IAM token: {}", e)))?;
75
76        if !response.status().is_success() {
77            let status = response.status();
78            let error_text = response
79                .text()
80                .await
81                .unwrap_or_else(|_| "Unknown error".to_string());
82            return Err(Error::Api(format!(
83                "Failed to generate IAM token: {} - {}",
84                status, error_text
85            )));
86        }
87
88        #[derive(serde::Deserialize)]
89        struct TokenResponse {
90            access_token: String,
91        }
92
93        let token_response: TokenResponse = response
94            .json()
95            .await
96            .map_err(|e| Error::Serialization(format!("Failed to parse IAM token response: {}", e)))?;
97
98        Ok(token_response.access_token)
99    }
100
101    // ============================================================================
102    // Custom Assistant Management
103    // ============================================================================
104
105    /// List all custom assistants
106    pub async fn list_assistants(&self) -> Result<Vec<CustomAssistant>> {
107        let access_token = self.access_token.as_ref().ok_or_else(|| {
108            Error::Authentication("Not authenticated. Set access token first.".to_string())
109        })?;
110
111        let url = format!(
112            "{}/v1/assistants",
113            self.config.get_base_url()
114        );
115
116        let response = self
117            .client
118            .get(&url)
119            .header("Authorization", format!("Bearer {}", access_token))
120            .header("Content-Type", "application/json")
121            .send()
122            .await
123            .map_err(|e| Error::Network(e.to_string()))?;
124
125        if !response.status().is_success() {
126            let status = response.status();
127            let error_text = response
128                .text()
129                .await
130                .unwrap_or_else(|_| "Unknown error".to_string());
131            return Err(Error::Api(format!(
132                "Failed to list assistants: {} - {}",
133                status, error_text
134            )));
135        }
136
137        let text = response
138            .text()
139            .await
140            .map_err(|e| Error::Serialization(e.to_string()))?;
141
142        // Try to parse as direct array first
143        if let Ok(assistants) = serde_json::from_str::<Vec<CustomAssistant>>(&text) {
144            return Ok(assistants);
145        }
146
147        // Try to extract from wrapper object
148        if let Ok(obj) = serde_json::from_str::<serde_json::Value>(&text) {
149            if let Some(assistants_array) = obj.get("assistants").and_then(|a| a.as_array()) {
150                let assistants: Result<Vec<CustomAssistant>> = assistants_array
151                    .iter()
152                    .map(|item| {
153                        serde_json::from_value::<CustomAssistant>(item.clone())
154                            .map_err(|e| Error::Serialization(e.to_string()))
155                    })
156                    .collect();
157                return assistants;
158            }
159        }
160
161        Ok(Vec::new())
162    }
163
164    /// Send multiple messages in a batch
165    pub async fn send_batch_messages(&self, request: BatchMessageRequest) -> Result<BatchMessageResponse> {
166        let api_key = self.access_token.as_ref().ok_or_else(|| {
167            Error::Authentication("Not authenticated. Set access token (API key) first.".to_string())
168        })?;
169
170        let base_url = self.config.get_base_url();
171        let url = format!("{}/batch/messages", base_url);
172
173        let response = self
174            .client
175            .post(&url)
176            .header("Authorization", format!("Bearer {}", api_key))
177            .header("Content-Type", "application/json")
178            .json(&request)
179            .send()
180            .await
181            .map_err(|e| Error::Network(e.to_string()))?;
182
183        if !response.status().is_success() {
184            let status = response.status();
185            let error_text = response
186                .text()
187                .await
188                .unwrap_or_else(|_| "Unknown error".to_string());
189            return Err(Error::Api(format!(
190                "Failed to send batch messages: {} - {}",
191                status, error_text
192            )));
193        }
194
195        let batch_response: BatchMessageResponse = response
196            .json()
197            .await
198            .map_err(|e| Error::Serialization(e.to_string()))?;
199
200        Ok(batch_response)
201    }
202
203    // ============================================================================
204    // Skills Management
205    // ============================================================================
206
207    /// List all skills
208    pub async fn list_skills(&self) -> Result<Vec<Skill>> {
209        let api_key = self.access_token.as_ref().ok_or_else(|| {
210            Error::Authentication("Not authenticated. Set access token (API key) first.".to_string())
211        })?;
212
213        let base_url = self.config.get_base_url();
214        let url = format!("{}/skills", base_url);
215
216        let response = self
217            .client
218            .get(&url)
219            .header("Authorization", format!("Bearer {}", api_key))
220            .header("Content-Type", "application/json")
221            .send()
222            .await
223            .map_err(|e| Error::Network(e.to_string()))?;
224
225        if !response.status().is_success() {
226            let status = response.status();
227            let error_text = response
228                .text()
229                .await
230                .unwrap_or_else(|_| "Unknown error".to_string());
231            return Err(Error::Api(format!(
232                "Failed to list skills: {} - {}",
233                status, error_text
234            )));
235        }
236
237        let text = response
238            .text()
239            .await
240            .map_err(|e| Error::Serialization(e.to_string()))?;
241
242        if let Ok(skills) = serde_json::from_str::<Vec<Skill>>(&text) {
243            return Ok(skills);
244        }
245
246        if let Ok(obj) = serde_json::from_str::<serde_json::Value>(&text) {
247            if let Some(skills_array) = obj.get("skills").and_then(|s| s.as_array()) {
248                let skills: Result<Vec<Skill>> = skills_array
249                    .iter()
250                    .map(|skill| {
251                        serde_json::from_value::<Skill>(skill.clone())
252                            .map_err(|e| Error::Serialization(e.to_string()))
253                    })
254                    .collect();
255                return skills;
256            }
257        }
258
259        Ok(Vec::new())
260    }
261
262    /// Get a specific skill by ID
263    pub async fn get_skill(&self, skill_id: &str) -> Result<Skill> {
264        let api_key = self.access_token.as_ref().ok_or_else(|| {
265            Error::Authentication("Not authenticated. Set access token (API key) first.".to_string())
266        })?;
267
268        let base_url = self.config.get_base_url();
269        let url = format!("{}/skills/{}", base_url, skill_id);
270
271        let response = self
272            .client
273            .get(&url)
274            .header("Authorization", format!("Bearer {}", api_key))
275            .header("Content-Type", "application/json")
276            .send()
277            .await
278            .map_err(|e| Error::Network(e.to_string()))?;
279
280        if !response.status().is_success() {
281            let status = response.status();
282            let error_text = response
283                .text()
284                .await
285                .unwrap_or_else(|_| "Unknown error".to_string());
286            return Err(Error::Api(format!(
287                "Failed to get skill {}: {} - {}",
288                skill_id, status, error_text
289            )));
290        }
291
292        let skill: Skill = response
293            .json()
294            .await
295            .map_err(|e| Error::Serialization(e.to_string()))?;
296
297        Ok(skill)
298    }
299}
300
301// Helper structs for SSE parsing
302#[derive(serde::Deserialize)]
303struct ChatChunk {
304    content: Option<String>,
305    metadata: Option<HashMap<String, Value>>,
306}
307
308#[derive(serde::Deserialize)]
309struct EventData {
310    event: String,
311    data: Value,
312}
watsonx_rs/orchestrate/client.rs

watsonx_rs/orchestrate/
client.rs
