grok_client/
types.rs

1pub mod api {
2    // ####################
3    // AI API TYPES
4    // ####################
5    use crate::ai;
6    use serde::{Deserialize, Serialize};
7
8    #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
9    pub struct ApiKey {
10        pub redacted_api_key: String,
11        pub user_id: String,
12        pub name: String,
13        pub create_time: String,
14        pub modify_time: String,
15        pub modified_by: String,
16        pub team_id: String,
17        pub acls: Vec<String>,
18        pub api_key_id: String,
19        pub team_blocked: bool,
20        pub api_key_blocked: bool,
21        pub api_key_disabled: bool,
22    }
23
24    #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
25    pub struct Model {
26        pub id: String,
27        pub created: i64,
28        pub object: String,
29        pub owned_by: String,
30    }
31
32    #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
33    pub struct LanguageModel {
34        pub id: String,
35        pub fingerprint: String,
36        pub created: i64,
37        pub object: String,
38        pub owned_by: String,
39        pub version: String,
40        pub cached_prompt_text_token_price: i64,
41        pub completion_text_token_price: i64,
42        pub prompt_image_token_price: i64,
43        pub prompt_text_token_price: i64,
44        pub input_modalities: Vec<String>,
45        pub output_modalities: Vec<String>,
46        pub aliases: Vec<String>,
47    }
48
49    #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
50    pub struct LanguageModels {
51        pub models: Vec<LanguageModel>,
52    }
53
54    #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
55    pub struct ImageModel {
56        pub id: String,
57        pub fingerprint: String,
58        pub max_prompt_length: i64,
59        pub created: i64,
60        pub object: String,
61        pub owned_by: String,
62        pub version: String,
63        pub image_price: i64,
64        pub input_modalities: Vec<String>,
65        pub output_modalities: Vec<String>,
66        pub aliases: Vec<String>,
67    }
68
69    #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
70    pub struct ImageModels {
71        pub models: Vec<ImageModel>,
72    }
73
74    #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
75    pub struct TokenizeRequest {
76        pub model: String,
77        pub text: String,
78    }
79
80    impl TokenizeRequest {
81        pub fn init(model: ai::LanguageModel, text: String) -> Self {
82            Self {
83                model: model.to_string(),
84                text,
85            }
86        }
87    }
88
89    #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
90    pub struct Token {
91        pub token_id: i64,
92        pub string_token: String,
93        pub token_bytes: Vec<i64>,
94    }
95
96    #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
97    pub struct TokenizeResponse {
98        pub token_ids: Vec<Token>,
99    }
100}
101
102pub mod chat {
103    // ####################
104    // AI CHAT API TYPES
105    // ####################
106    use crate::ai;
107    use serde::{Deserialize, Serialize};
108
109    pub mod stream {
110        use crate::types::chat::Usage;
111        use serde::{Deserialize, Serialize};
112
113        #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
114        pub struct Options {
115            pub include_usage: Option<bool>,
116        }
117
118        impl Options {
119            pub fn default() -> Self {
120                Self {
121                    include_usage: Some(true),
122                }
123            }
124        }
125
126        #[derive(Debug, Serialize, Deserialize)]
127        pub struct ChatCompletionChunk {
128            pub id: String,
129            pub object: String,
130            pub created: i64,
131            pub model: String,
132            pub choices: Vec<Choice>,
133            pub usage: Option<Usage>,
134            pub citations: Option<Vec<String>>,
135            pub system_fingerprint: String,
136        }
137
138        #[derive(Debug, Serialize, Deserialize)]
139        pub struct Choice {
140            pub index: i32,
141            pub delta: Delta,
142        }
143
144        #[derive(Debug, Serialize, Deserialize)]
145        pub struct Delta {
146            pub content: Option<String>,
147            pub reasoning_content: Option<String>,
148            pub role: Option<String>,
149        }
150    }
151
152    pub mod search {
153        use serde::{Deserialize, Serialize};
154
155        #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
156        pub struct Parameters {
157            #[serde(default, skip_serializing_if = "Option::is_none")]
158            pub from_date: Option<String>,
159            #[serde(
160                default = "default_max_search_results",
161                skip_serializing_if = "Option::is_none"
162            )]
163            pub max_search_results: Option<i32>,
164            #[serde(default = "default_mode", skip_serializing_if = "Option::is_none")]
165            pub mode: Option<String>,
166            #[serde(
167                default = "default_return_citations",
168                skip_serializing_if = "Option::is_none"
169            )]
170            pub return_citations: Option<bool>,
171            #[serde(default, skip_serializing_if = "Option::is_none")]
172            pub sources: Option<Vec<Source>>,
173            #[serde(default, skip_serializing_if = "Option::is_none")]
174            pub to_date: Option<String>,
175        }
176
177        impl Parameters {
178            pub fn default() -> Self {
179                Self {
180                    from_date: None,
181                    max_search_results: default_max_search_results(),
182                    mode: default_mode(),
183                    return_citations: default_return_citations(),
184                    sources: None,
185                    to_date: None,
186                }
187            }
188        }
189
190        #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
191        pub struct Source {
192            #[serde(default, skip_serializing_if = "Option::is_none")]
193            pub excluded_x_handles: Option<Vec<String>>,
194            #[serde(default, skip_serializing_if = "Option::is_none")]
195            pub included_x_handles: Option<Vec<String>>,
196            #[serde(default, skip_serializing_if = "Option::is_none")]
197            pub x_handles: Option<Vec<String>>,
198            #[serde(default, skip_serializing_if = "Option::is_none")]
199            pub post_favorite_count: Option<i32>,
200            #[serde(default, skip_serializing_if = "Option::is_none")]
201            pub post_view_count: Option<i32>,
202            #[serde(default, skip_serializing_if = "Option::is_none")]
203            pub allowed_websites: Option<Vec<String>>,
204            #[serde(default, skip_serializing_if = "Option::is_none")]
205            pub country: Option<String>,
206            #[serde(default, skip_serializing_if = "Option::is_none")]
207            pub excluded_websites: Option<Vec<String>>,
208            #[serde(
209                default = "default_safe_search",
210                skip_serializing_if = "Option::is_none"
211            )]
212            pub safe_search: Option<bool>,
213            #[serde(default, skip_serializing_if = "Option::is_none")]
214            pub links: Option<Vec<String>>,
215            pub r#type: String,
216        }
217
218        fn default_max_search_results() -> Option<i32> {
219            Some(15)
220        }
221
222        fn default_mode() -> Option<String> {
223            Some(String::from("auto"))
224        }
225
226        fn default_return_citations() -> Option<bool> {
227            Some(true)
228        }
229
230        fn default_safe_search() -> Option<bool> {
231            Some(false)
232        }
233    }
234
235    pub mod tool {
236        use serde::{Deserialize, Serialize};
237
238        // "none" "auto" "required" or a single tool::Tool
239        #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
240        pub enum Choice {
241            None,
242            Auto,
243            Required,
244            Instance(Tool),
245        }
246
247        #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
248        pub struct Tool {
249            pub r#type: String, // Using `r#type` to match the schema's `type` field
250            pub function: Function,
251
252            // CALL
253            #[serde(default, skip_serializing_if = "Option::is_none")]
254            pub id: Option<String>,
255            #[serde(default, skip_serializing_if = "Option::is_none")]
256            pub index: Option<i32>,
257        }
258
259        #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
260        pub struct Function {
261            pub name: String,
262
263            // DEFINITION
264            #[serde(default, skip_serializing_if = "Option::is_none")]
265            pub parameters: Option<serde_json::Value>, // Using Value to handle arbitrary JSON schema
266            #[serde(default, skip_serializing_if = "Option::is_none")]
267            pub description: Option<String>,
268
269            // CALL
270            // #[serde(
271            //     default,
272            //     skip_serializing_if = "Option::is_none",
273            //     deserialize_with = "deserialize_to_object"
274            // )]
275            #[serde(default, skip_serializing_if = "Option::is_none")]
276            pub arguments: Option<serde_json::Value>, // Using Value to handle arbitrary JSON schema
277        }
278    }
279
280    // Struct for chat completion request
281    #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
282    pub struct ChatCompletionRequest {
283        pub model: String,
284        pub messages: Vec<Message>,
285        // STREAM
286        #[serde(default, skip_serializing_if = "Option::is_none")]
287        pub stream: Option<bool>,
288        #[serde(default, skip_serializing_if = "Option::is_none")]
289        pub stream_options: Option<stream::Options>,
290        // DEFER
291        #[serde(default, skip_serializing_if = "Option::is_none")]
292        pub deferred: Option<bool>,
293        // SEARCH
294        #[serde(default, skip_serializing_if = "Option::is_none")]
295        pub search_parameters: Option<search::Parameters>,
296        // TOOLS
297        #[serde(default, skip_serializing_if = "Option::is_none")]
298        pub parallel_tool_calls: Option<bool>,
299        #[serde(default, skip_serializing_if = "Option::is_none")]
300        pub tools: Option<Vec<tool::Tool>>,
301        // #[serde(default, skip_serializing_if = "Option::is_none")]
302        // pub tool_choice: Option<tool::Choice>, // "none" "auto" "required" or a single tool::Tool
303        // #[serde(default, skip_serializing_if = "Option::is_none")]
304        // pub n: Option<i32>,
305    }
306
307    impl ChatCompletionRequest {
308        pub fn init(model: ai::LanguageModel, messages: Vec<Message>) -> Self {
309            Self {
310                model: model.to_string(),
311                messages,
312                deferred: Some(false),
313                stream: Some(false),
314                stream_options: None,
315                search_parameters: None,
316                parallel_tool_calls: None,
317                tools: None,
318                // tool_choice: None,
319            }
320        }
321    }
322
323    // Message object (used in both request and response)
324    #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
325    pub struct Message {
326        pub role: String,
327        pub content: String, // Could also be and Vec<Struct>
328        #[serde(default, skip_serializing_if = "Option::is_none")]
329        pub reasoning_content: Option<String>,
330        #[serde(default, skip_serializing_if = "Option::is_none")]
331        pub refusal: Option<String>,
332        #[serde(default, skip_serializing_if = "Option::is_none")]
333        pub tool_calls: Option<Vec<tool::Tool>>,
334        #[serde(default, skip_serializing_if = "Option::is_none")]
335        pub tool_call_id: Option<String>,
336    }
337
338    impl Message {
339        pub fn init(role: ai::Role, content: String) -> Self {
340            Self {
341                role: role.to_string(),
342                content,
343                reasoning_content: None,
344                refusal: None,
345                tool_calls: None,
346                tool_call_id: None,
347            }
348        }
349        pub fn tool(tool_call_id: String, content: String) -> Self {
350            Self {
351                role: "tool".to_string(),
352                tool_call_id: Some(tool_call_id),
353                content,
354                reasoning_content: None,
355                refusal: None,
356                tool_calls: None,
357            }
358        }
359    }
360
361    #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
362    pub struct DeferredChatCompletionResponse {
363        pub request_id: String,
364    }
365
366    #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
367    pub struct ChatCompletionResponse {
368        pub id: String,
369        pub object: String, // Always "chat.response"
370        pub created: i64,   // Unix timestamp
371        pub model: String,
372        pub choices: Vec<Choice>, // Length corresponds to 'n' in request
373        #[serde(default, skip_serializing_if = "Option::is_none")]
374        pub usage: Option<Usage>,
375        #[serde(default, skip_serializing_if = "Option::is_none")]
376        pub citations: Option<Vec<String>>,
377        #[serde(default, skip_serializing_if = "Option::is_none")]
378        pub system_fingerprint: Option<String>,
379    }
380
381    impl ChatCompletionResponse {
382        pub fn new(i: i64) -> Self {
383            ChatCompletionResponse {
384                id: format!("chatcmpl-{}", i),
385                object: "chat.response".to_string(),
386                created: 1697056000 + i as i64,
387                model: format!("grok-beta-{}", i),
388                choices: vec![],
389                usage: None,
390                citations: None,
391                system_fingerprint: None,
392            }
393        }
394
395        pub fn mock(count: i64) -> Vec<Self> {
396            (0..count).map(|i| ChatCompletionResponse::new(i)).collect()
397        }
398    }
399
400    // Choice object within the response
401    #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
402    pub struct Choice {
403        pub index: i32,
404        pub message: Message,
405        pub finish_reason: String, // "stop" "tool_calls" "length" "end_turn"
406    }
407
408    // Usage object
409    #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
410    pub struct Usage {
411        pub prompt_tokens: i32,
412        pub prompt_tokens_details: PromptTokensDetails,
413        pub completion_tokens: i32,
414        pub completion_tokens_details: CompletionTokensDetails,
415        pub total_tokens: i32,
416        #[serde(default, skip_serializing_if = "Option::is_none")]
417        pub num_sources_used: Option<i32>, // Optional for backward compatibility
418    }
419
420    // Details for prompt tokens
421    #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
422    pub struct PromptTokensDetails {
423        pub text_tokens: i32,
424        pub audio_tokens: i32,
425        pub image_tokens: i32,
426        pub cached_tokens: i32,
427    }
428
429    // Details for completion tokens
430    #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
431    pub struct CompletionTokensDetails {
432        pub reasoning_tokens: i32,
433        pub audio_tokens: i32,
434        pub accepted_prediction_tokens: i32,
435        pub rejected_prediction_tokens: i32,
436    }
437}
438
439pub mod image {
440    // ####################
441    // AI IMAGE API TYPES
442    // ####################
443    use serde::{Deserialize, Serialize};
444    use strum::{Display, EnumIter, EnumString};
445
446    #[derive(
447        Clone, Copy, Debug, Display, PartialEq, EnumIter, EnumString, Serialize, Deserialize,
448    )]
449    pub enum ImageResponseFormat {
450        #[strum(serialize = "url")]
451        URL,
452        #[strum(serialize = "b64_json")]
453        B64JSON,
454    }
455
456    impl ImageResponseFormat {
457        pub fn err_ivalid_response_format(format: &String) -> String {
458            format!("Invalid response format '{format}'")
459        }
460    }
461
462    #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
463    pub struct ImageRequest {
464        prompt: String,
465        model: String,
466        response_format: String,
467        n: u32,
468    }
469
470    impl ImageRequest {
471        pub fn init(
472            prompt: String,
473            model: String,
474            response_format: ImageResponseFormat,
475            n: u32,
476        ) -> Self {
477            Self {
478                prompt,
479                model,
480                response_format: response_format.to_string(),
481                n,
482            }
483        }
484    }
485
486    #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
487    pub struct ImageData {
488        pub revised_prompt: String,
489        #[serde(default, skip_serializing_if = "Option::is_none")]
490        pub url: Option<String>,
491        #[serde(default, skip_serializing_if = "Option::is_none")]
492        pub b64_json: Option<String>,
493    }
494
495    #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
496    pub struct ImageResponse {
497        pub data: Vec<ImageData>,
498    }
499}
grok_client/types.rs

grok_client/
types.rs
