asimov_xai_module/
lib.rs

1// This is free and unencumbered software released into the public domain.
2
3#![no_std]
4#![forbid(unsafe_code)]
5
6use asimov_module::{
7    prelude::*,
8    secrecy::{ExposeSecret, SecretString},
9    tracing,
10};
11use core::error::Error;
12use serde_json::{Value, json};
13
14#[derive(Clone, Debug, bon::Builder)]
15#[builder(on(String, into))]
16pub struct Options {
17    #[builder(default = "https://api.x.ai")]
18    pub endpoint: String,
19
20    #[builder(default = "grok-3-mini")]
21    pub model: String,
22
23    pub max_tokens: Option<usize>,
24
25    #[builder(into)]
26    pub api_key: SecretString,
27}
28
29pub fn generate(input: impl AsRef<str>, options: &Options) -> Result<Vec<String>, Box<dyn Error>> {
30    let mut req = json!({
31        "model": options.model,
32        "input": input.as_ref(),
33    });
34
35    if let Some(max_tokens) = options.max_tokens {
36        req["max_output_tokens"] = max_tokens.into();
37    }
38
39    let mut resp = ureq::Agent::config_builder()
40        .http_status_as_error(false)
41        .user_agent("asimov-xai-module")
42        .build()
43        .new_agent()
44        .post(format!("{}/v1/responses", options.endpoint))
45        .header(
46            "Authorization",
47            format!("Bearer {}", options.api_key.expose_secret()),
48        )
49        .header("content-type", "application/json")
50        .send_json(&req)
51        .inspect_err(|e| tracing::error!("HTTP request failed: {e}"))?;
52    tracing::debug!(response = ?resp);
53
54    let status = resp.status();
55    tracing::debug!(status = status.to_string());
56
57    let resp: Value = resp
58        .body_mut()
59        .read_json()
60        .inspect_err(|e| tracing::error!("unable to read HTTP response body: {e}"))?;
61    tracing::debug!(body = resp.to_string());
62
63    if !status.is_success() {
64        tracing::error!("Received an error response: {status}");
65
66        // {
67        //   "code": "Client specified an invalid argument",
68        //   "error": "Incorrect API key provided: fo***ar. You can obtain an API key from https://console.x.ai."
69        // }
70        if let Some(message) = resp["error"].as_str() {
71            return Err(message.into());
72        }
73        if let Some(message) = resp.as_str() {
74            return Err(message.into());
75        }
76    }
77
78    let mut responses = Vec::new();
79
80    // {
81    //   "created_at": 1758188599,
82    //   "id": "...",
83    //   "incomplete_details": null,
84    //   "max_output_tokens": null,
85    //   "metadata": {},
86    //   "model": "grok-3-mini",
87    //   "object": "response",
88    //   "output": [
89    //     {
90    //       "id": "rs_...",
91    //       "status": "completed",
92    //       "summary": [
93    //         {
94    //           "text": "...",
95    //           "type": "summary_text"
96    //         }
97    //       ],
98    //       "type": "reasoning"
99    //     },
100    //     {
101    //       "content": [
102    //         {
103    //           "annotations": [],
104    //           "logprobs": null,
105    //           "text": "...",
106    //           "type": "output_text"
107    //         }
108    //       ],
109    //       "id": "msg_...",
110    //       "role": "assistant",
111    //       "status": "completed",
112    //       "type": "message"
113    //     }
114    //   ],
115    //   "parallel_tool_calls": true,
116    //   "previous_response_id": null,
117    //   "reasoning": {
118    //     "effort": "medium",
119    //     "summary": "detailed"
120    //   },
121    //   "status": "completed",
122    //   "store": true,
123    //   "temperature": null,
124    //   "text": {
125    //     "format": {
126    //       "type": "text"
127    //     }
128    //   },
129    //   "tool_choice": "auto",
130    //   "tools": [],
131    //   "top_p": null,
132    //   "usage": {
133    //     "input_tokens": 8,
134    //     "input_tokens_details": {
135    //       "cached_tokens": 7
136    //     },
137    //     "output_tokens": 199,
138    //     "output_tokens_details": {
139    //       "reasoning_tokens": 188
140    //     },
141    //     "total_tokens": 207
142    //   },
143    //   "user": null
144    // }
145
146    if let Some(chunks) = resp["output"].as_array() {
147        for chunk in chunks {
148            if chunk["type"].as_str().is_none_or(|t| t != "message") {
149                tracing::debug!("skipping non-message chunk in response: {chunk}");
150                continue;
151            }
152            if chunk["role"].as_str().is_none_or(|r| r != "assistant") {
153                tracing::debug!("skipping output chunk not from assistant: {chunk}");
154                continue;
155            }
156
157            if let Some(chunk_contents) = chunk["content"].as_array() {
158                for content in chunk_contents {
159                    if content["type"].as_str().is_none_or(|r| r != "output_text") {
160                        tracing::debug!("skipping non-text message chunk in response: {chunk}");
161                        continue;
162                    }
163
164                    if let Some(text) = content["text"].as_str() {
165                        responses.push(text.to_string());
166                    }
167                }
168            };
169
170            if let Some(status) = chunk["status"].as_str() {
171                tracing::debug!(status);
172            }
173        }
174    }
175
176    Ok(responses)
177}
asimov_xai_module/lib.rs

asimov_xai_module/
lib.rs
