1use std::{fmt, sync::Arc, time::Duration};
2
3use base64::{Engine as _, engine::general_purpose::STANDARD};
4use reqwest::{
5 Client, RequestBuilder, Response, StatusCode, Url,
6 header::{AUTHORIZATION, HeaderValue},
7 multipart::{Form, Part},
8 redirect::Policy,
9};
10use serde_json::Value;
11use zeroize::Zeroizing;
12
13use crate::{
14 Error, GPT_4O_TRANSCRIBE, GeneratedImage, ImageFormat, ImageGeneration, ImageGenerationRequest,
15 ImageQuality, ImageTokenDetails, ImageUsage, Result, Transcription, TranscriptionRequest,
16 TranscriptionTokenDetails, TranscriptionTokenUsage, TranscriptionUsage,
17 error::{clean_message, transport},
18};
19
20const API_BASE: &str = "https://api.openai.com/v1/";
21const DEFAULT_TIMEOUT: Duration = Duration::from_secs(5 * 60);
22const MAX_RESPONSE_BYTES: usize = 128 * 1024 * 1024;
23
24struct ApiKey(Zeroizing<String>);
25
26impl ApiKey {
27 fn new(value: impl Into<String>) -> Result<Self> {
28 let supplied = Zeroizing::new(value.into());
29 let value = supplied.trim();
30 if value.is_empty() {
31 return Err(Error::InvalidApiKey);
32 }
33 let authorization = Zeroizing::new(format!("Bearer {value}"));
34 if HeaderValue::from_str(&authorization).is_err() {
35 return Err(Error::InvalidApiKey);
36 }
37 Ok(Self(Zeroizing::new(value.to_owned())))
38 }
39
40 fn sensitive_authorization(&self) -> Result<HeaderValue> {
41 let authorization = Zeroizing::new(format!("Bearer {}", self.0.as_str()));
42 let mut value = HeaderValue::from_str(&authorization).map_err(|_| Error::InvalidApiKey)?;
43 value.set_sensitive(true);
44 Ok(value)
45 }
46}
47
48impl fmt::Debug for ApiKey {
49 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
50 f.write_str("ApiKey([REDACTED])")
51 }
52}
53
54#[derive(Clone)]
56pub struct OpenAi {
57 api_key: Arc<ApiKey>,
58 client: Client,
59 transcription_endpoint: Url,
60 image_generation_endpoint: Url,
61}
62
63impl fmt::Debug for OpenAi {
64 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
65 f.debug_struct("OpenAi")
66 .field("api_key", &self.api_key)
67 .field("api_base", &API_BASE)
68 .finish_non_exhaustive()
69 }
70}
71
72impl OpenAi {
73 pub fn open(api_key: impl Into<String>) -> Result<Self> {
78 let api_key = ApiKey::new(api_key)?;
79 let base = Url::parse(API_BASE)
80 .map_err(|_| Error::Protocol("compiled API base URL is invalid".into()))?;
81 let client = Client::builder()
82 .timeout(DEFAULT_TIMEOUT)
83 .redirect(Policy::none())
84 .retry(reqwest::retry::never())
85 .referer(false)
86 .no_proxy()
87 .https_only(true)
88 .user_agent(concat!("kcode-openai-api/", env!("CARGO_PKG_VERSION")))
89 .build()
90 .map_err(transport)?;
91 Ok(Self {
92 api_key: Arc::new(api_key),
93 client,
94 transcription_endpoint: base
95 .join("audio/transcriptions")
96 .map_err(|_| Error::Protocol("compiled transcription URL is invalid".into()))?,
97 image_generation_endpoint: base
98 .join("images/generations")
99 .map_err(|_| Error::Protocol("compiled image generation URL is invalid".into()))?,
100 })
101 }
102
103 pub async fn transcribe(&self, request: TranscriptionRequest) -> Result<Transcription> {
105 request.validate()?;
106 let TranscriptionRequest {
107 audio,
108 prompt,
109 language,
110 } = request;
111 let (file_name, mime_type, data) = audio.into_parts();
112 let part = Part::bytes(data)
113 .file_name(file_name)
114 .mime_str(&mime_type)
115 .map_err(|_| Error::InvalidInput("audio MIME type is invalid".into()))?;
116 let mut form = Form::new()
117 .part("file", part)
118 .text("model", GPT_4O_TRANSCRIBE)
119 .text("response_format", "json");
120 if let Some(prompt) = prompt {
121 form = form.text("prompt", prompt);
122 }
123 if let Some(language) = language {
124 form = form.text("language", language);
125 }
126 let (payload, request_id) = self
127 .execute(
128 self.client
129 .post(self.transcription_endpoint.clone())
130 .multipart(form),
131 )
132 .await?;
133 parse_transcription(&payload, request_id)
134 }
135
136 pub async fn generate_image(&self, request: ImageGenerationRequest) -> Result<ImageGeneration> {
138 request.validate()?;
139 let requested_format = request.output_format;
140 let (payload, request_id) = self
141 .execute(
142 self.client
143 .post(self.image_generation_endpoint.clone())
144 .json(&request.payload()),
145 )
146 .await?;
147 parse_image_generation(&payload, requested_format, request_id)
148 }
149
150 async fn execute(&self, request: RequestBuilder) -> Result<(Value, Option<String>)> {
151 let response = request
152 .header(AUTHORIZATION, self.api_key.sensitive_authorization()?)
153 .send()
154 .await
155 .map_err(transport)?;
156 let status = response.status();
157 let request_id = response
158 .headers()
159 .get("x-request-id")
160 .and_then(|value| value.to_str().ok())
161 .map(|value| clean_message(value, 200));
162 let body = bounded_body(response).await?;
163 if !status.is_success() {
164 return Err(provider_error(status, &body, request_id));
165 }
166 let payload = serde_json::from_slice(&body)
167 .map_err(|_| Error::Protocol("response was not valid JSON".into()))?;
168 Ok((payload, request_id))
169 }
170}
171
172async fn bounded_body(mut response: Response) -> Result<Vec<u8>> {
173 if response
174 .content_length()
175 .is_some_and(|value| value > MAX_RESPONSE_BYTES as u64)
176 {
177 return Err(Error::Protocol("response exceeded 128 MiB".into()));
178 }
179 let initial_capacity = response
180 .content_length()
181 .and_then(|value| usize::try_from(value).ok())
182 .unwrap_or(0)
183 .min(MAX_RESPONSE_BYTES);
184 let mut body = Vec::with_capacity(initial_capacity);
185 while let Some(chunk) = response.chunk().await.map_err(transport)? {
186 let length = body
187 .len()
188 .checked_add(chunk.len())
189 .ok_or_else(|| Error::Protocol("response exceeded 128 MiB".into()))?;
190 if length > MAX_RESPONSE_BYTES {
191 return Err(Error::Protocol("response exceeded 128 MiB".into()));
192 }
193 body.extend_from_slice(&chunk);
194 }
195 Ok(body)
196}
197
198fn parse_transcription(payload: &Value, request_id: Option<String>) -> Result<Transcription> {
199 let text = payload
200 .get("text")
201 .and_then(Value::as_str)
202 .map(str::trim)
203 .filter(|value| !value.is_empty())
204 .ok_or_else(|| Error::Protocol("transcription response omitted non-empty text".into()))?
205 .to_owned();
206 let usage = payload
207 .get("usage")
208 .filter(|value| !value.is_null())
209 .map(parse_transcription_usage)
210 .transpose()?;
211 Ok(Transcription {
212 text,
213 usage,
214 request_id,
215 })
216}
217
218fn parse_transcription_usage(value: &Value) -> Result<TranscriptionUsage> {
219 let usage_type = value.get("type").and_then(Value::as_str);
220 if usage_type == Some("duration") {
221 let seconds = value
222 .get("seconds")
223 .and_then(Value::as_f64)
224 .ok_or_else(|| {
225 Error::Protocol("duration transcription usage omitted seconds".into())
226 })?;
227 if !seconds.is_finite() || seconds < 0.0 {
228 return Err(Error::Protocol(
229 "duration transcription usage contained invalid seconds".into(),
230 ));
231 }
232 return Ok(TranscriptionUsage::DurationSeconds(seconds));
233 }
234 if !matches!(usage_type, None | Some("tokens")) {
235 return Err(Error::Protocol(
236 "transcription usage returned an unsupported type".into(),
237 ));
238 }
239 let input_tokens = required_u64(value, "input_tokens", "transcription usage")?;
240 let output_tokens = required_u64(value, "output_tokens", "transcription usage")?;
241 let total_tokens = required_u64(value, "total_tokens", "transcription usage")?;
242 let input_details = value
243 .get("input_token_details")
244 .filter(|details| !details.is_null())
245 .map(|details| {
246 Ok(TranscriptionTokenDetails {
247 audio_tokens: optional_u64(details, "audio_tokens", "transcription usage")?,
248 text_tokens: optional_u64(details, "text_tokens", "transcription usage")?,
249 })
250 })
251 .transpose()?;
252 Ok(TranscriptionUsage::Tokens(TranscriptionTokenUsage {
253 input_tokens,
254 output_tokens,
255 total_tokens,
256 input_details,
257 }))
258}
259
260fn parse_image_generation(
261 payload: &Value,
262 requested_format: ImageFormat,
263 request_id: Option<String>,
264) -> Result<ImageGeneration> {
265 let created = required_u64(payload, "created", "image generation response")?;
266 let data = payload
267 .get("data")
268 .and_then(Value::as_array)
269 .ok_or_else(|| Error::Protocol("image generation response omitted image data".into()))?;
270 if data.len() != 1 {
271 return Err(Error::Protocol(
272 "single-image request did not return exactly one image".into(),
273 ));
274 }
275 let encoded = data[0]
276 .get("b64_json")
277 .and_then(Value::as_str)
278 .ok_or_else(|| Error::Protocol("generated image omitted base64 data".into()))?;
279 let decoded = STANDARD
280 .decode(encoded)
281 .map_err(|_| Error::Protocol("generated image contained invalid base64".into()))?;
282 if decoded.is_empty() {
283 return Err(Error::Protocol("generated image was empty".into()));
284 }
285 let format = match payload.get("output_format").and_then(Value::as_str) {
286 Some(value) => ImageFormat::parse(value)
287 .ok_or_else(|| Error::Protocol("generated image used an unknown format".into()))?,
288 None => requested_format,
289 };
290 let quality = payload
291 .get("quality")
292 .and_then(Value::as_str)
293 .and_then(ImageQuality::parse);
294 let size = payload
295 .get("size")
296 .and_then(Value::as_str)
297 .map(|value| clean_message(value, 40));
298 let usage = payload
299 .get("usage")
300 .filter(|value| !value.is_null())
301 .map(parse_image_usage)
302 .transpose()?;
303 Ok(ImageGeneration {
304 created,
305 image: GeneratedImage {
306 data: decoded,
307 format,
308 },
309 size,
310 quality,
311 usage,
312 request_id,
313 })
314}
315
316fn parse_image_usage(value: &Value) -> Result<ImageUsage> {
317 Ok(ImageUsage {
318 input_tokens: required_u64(value, "input_tokens", "image usage")?,
319 output_tokens: required_u64(value, "output_tokens", "image usage")?,
320 total_tokens: required_u64(value, "total_tokens", "image usage")?,
321 input_details: parse_image_token_details(
322 value
323 .get("input_tokens_details")
324 .ok_or_else(|| Error::Protocol("image usage omitted input token details".into()))?,
325 )?,
326 output_details: value
327 .get("output_tokens_details")
328 .filter(|details| !details.is_null())
329 .map(parse_image_token_details)
330 .transpose()?,
331 })
332}
333
334fn parse_image_token_details(value: &Value) -> Result<ImageTokenDetails> {
335 Ok(ImageTokenDetails {
336 text_tokens: required_u64(value, "text_tokens", "image token details")?,
337 image_tokens: required_u64(value, "image_tokens", "image token details")?,
338 })
339}
340
341fn required_u64(value: &Value, field: &str, context: &str) -> Result<u64> {
342 value
343 .get(field)
344 .and_then(Value::as_u64)
345 .ok_or_else(|| Error::Protocol(format!("{context} omitted {field}")))
346}
347
348fn optional_u64(value: &Value, field: &str, context: &str) -> Result<Option<u64>> {
349 match value.get(field) {
350 None | Some(Value::Null) => Ok(None),
351 Some(value) => value
352 .as_u64()
353 .map(Some)
354 .ok_or_else(|| Error::Protocol(format!("{context} returned invalid {field}"))),
355 }
356}
357
358fn provider_error(status: StatusCode, body: &[u8], request_id: Option<String>) -> Error {
359 let payload = serde_json::from_slice::<Value>(body).ok();
360 let code = payload
361 .as_ref()
362 .and_then(|value| {
363 value
364 .pointer("/error/code")
365 .and_then(Value::as_str)
366 .or_else(|| value.pointer("/error/type").and_then(Value::as_str))
367 })
368 .map(|value| clean_message(value, 100));
369 let message = payload
370 .as_ref()
371 .and_then(|value| value.pointer("/error/message"))
372 .and_then(Value::as_str)
373 .map(|value| clean_message(value, 400))
374 .unwrap_or_else(|| format!("provider request failed with HTTP {status}"));
375 Error::Provider {
376 status: status.as_u16(),
377 code,
378 message,
379 request_id,
380 }
381}
382
383#[cfg(test)]
384mod tests {
385 use serde_json::json;
386
387 use super::*;
388 use crate::{ImageGenerationRequest, ImageSize};
389
390 #[test]
391 fn debug_and_authorization_header_redact_api_key() {
392 let client = OpenAi::open("secret-api-key").unwrap();
393 let debug = format!("{client:?}");
394 assert!(debug.contains("[REDACTED]"));
395 assert!(!debug.contains("secret-api-key"));
396
397 let header = client.api_key.sensitive_authorization().unwrap();
398 assert!(header.is_sensitive());
399 let request = client
400 .client
401 .post(client.transcription_endpoint.clone())
402 .header(AUTHORIZATION, header);
403 assert!(!format!("{request:?}").contains("secret-api-key"));
404 }
405
406 #[test]
407 fn transcription_response_normalizes_token_usage() {
408 let payload = json!({
409 "text": " hello world ",
410 "usage": {
411 "type": "tokens",
412 "input_tokens": 12,
413 "output_tokens": 3,
414 "total_tokens": 15,
415 "input_token_details": {"audio_tokens": 10, "text_tokens": 2}
416 }
417 });
418 let parsed = parse_transcription(&payload, Some("req_123".into())).unwrap();
419 assert_eq!(parsed.text, "hello world");
420 assert_eq!(parsed.request_id.as_deref(), Some("req_123"));
421 assert_eq!(
422 parsed.usage,
423 Some(TranscriptionUsage::Tokens(TranscriptionTokenUsage {
424 input_tokens: 12,
425 output_tokens: 3,
426 total_tokens: 15,
427 input_details: Some(TranscriptionTokenDetails {
428 audio_tokens: Some(10),
429 text_tokens: Some(2),
430 }),
431 }))
432 );
433 }
434
435 #[test]
436 fn image_response_decodes_bytes_and_usage() {
437 let payload = json!({
438 "created": 1_721_000_000_u64,
439 "background": "opaque",
440 "output_format": "png",
441 "quality": "high",
442 "size": "2048x2048",
443 "data": [{"b64_json": "AQID"}],
444 "usage": {
445 "input_tokens": 10,
446 "output_tokens": 20,
447 "total_tokens": 30,
448 "input_tokens_details": {"text_tokens": 10, "image_tokens": 0},
449 "output_tokens_details": {"text_tokens": 0, "image_tokens": 20}
450 }
451 });
452 let parsed = parse_image_generation(&payload, ImageFormat::Png, None).unwrap();
453 assert_eq!(parsed.image.data, vec![1, 2, 3]);
454 assert_eq!(parsed.image.format, ImageFormat::Png);
455 assert_eq!(parsed.size.as_deref(), Some("2048x2048"));
456 assert_eq!(parsed.quality, Some(ImageQuality::High));
457 assert_eq!(
458 parsed.usage.unwrap().output_details.unwrap().image_tokens,
459 20
460 );
461 }
462
463 #[test]
464 fn gpt_image_payload_supports_flexible_dimensions() {
465 let mut request = ImageGenerationRequest::new("draw a quiet library");
466 request.size = ImageSize::dimensions(1536, 864).unwrap();
467 let payload = request.payload();
468 assert_eq!(payload["model"], "gpt-image-2");
469 assert_eq!(payload["size"], "1536x864");
470 }
471
472 #[test]
473 fn provider_errors_are_sanitized_and_keep_request_ids() {
474 let error = provider_error(
475 StatusCode::BAD_REQUEST,
476 br#"{"error":{"code":"moderation_blocked","message":"bad\nrequest"}}"#,
477 Some("req_456".into()),
478 );
479 assert!(matches!(
480 error,
481 Error::Provider {
482 status: 400,
483 code: Some(code),
484 message,
485 request_id: Some(request_id),
486 } if code == "moderation_blocked" && message == "bad request" && request_id == "req_456"
487 ));
488 }
489}