Skip to main content

kcode_openai_api/
api.rs

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/// Cloneable asynchronous OpenAI client for transcription and image generation.
55#[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    /// Opens a client and retains the supplied OpenAI API key in memory.
74    ///
75    /// Clones share the key and HTTP connection pool. Dropping the last clone
76    /// discards the retained key with best-effort zeroization.
77    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    /// Transcribes one in-memory recording with `gpt-4o-transcribe`.
104    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    /// Generates one image with the latest GPT Image model, `gpt-image-2`.
137    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}