jutella 0.6.0

// Copyright (c) 2024 Dmitry Markin
//
// SPDX-License-Identifier: MIT
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

//! OpenAI REST API client.

use crate::chat_client::openai_api::chat_completions::{ChatCompletions, ChatCompletionsBody};
use reqwest::{
    header::{HeaderMap, HeaderName, HeaderValue, InvalidHeaderValue, AUTHORIZATION},
    Client, ClientBuilder, StatusCode,
};
use serde::Deserialize;
use std::{fmt::Display, str::FromStr, time::Duration};

const CHAT_COMPLETIONS_ENDPOINT: &str = "chat/completions";
const REQUEST_TIMEOUT: Duration = Duration::from_secs(300);

/// Authorization header.
///
/// Use `HeaderMap::try_from(auth)` to convert to `reqwest` headers.
#[derive(Debug)]
pub enum Auth {
    /// Auth header `Authorization: Bearer {api_token}`.
    Token(String),
    /// Auth header `api-key: {api_key}`.
    ApiKey(String),
}

impl TryFrom<Auth> for HeaderMap {
    type Error = InvalidHeaderValue;

    fn try_from(auth: Auth) -> Result<Self, InvalidHeaderValue> {
        let headers = match auth {
            Auth::Token(token) => [(
                AUTHORIZATION,
                HeaderValue::from_str(&format!("Bearer {token}"))?,
            )],
            Auth::ApiKey(api_key) => [(
                HeaderName::from_str("api-key").expect("to be valid ASCII"),
                HeaderValue::from_str(&api_key)?,
            )],
        }
        .into_iter()
        .collect();

        Ok(headers)
    }
}

/// OpenAI REST API client.
pub struct OpenAiClient {
    client: Client,
    endpoint: String,
}

impl OpenAiClient {
    /// Create new OpenAI API client.
    pub fn new(auth: Auth, base_url: String, api_version: Option<String>) -> Result<Self, Error> {
        let client = ClientBuilder::new()
            .default_headers(auth.try_into()?)
            .timeout(REQUEST_TIMEOUT)
            .build()?;

        let endpoint = build_url(base_url, api_version);

        Ok(Self { client, endpoint })
    }

    /// Create new OpenAI API client with custom [`reqwest::Client`].
    ///
    /// You are responsible for configuration of authorization headers and the request timeout!
    pub fn new_with_client(client: Client, base_url: String, api_version: Option<String>) -> Self {
        Self {
            client,
            endpoint: build_url(base_url, api_version),
        }
    }

    /// Request chat completion message.
    pub async fn chat_completions(
        &mut self,
        body: ChatCompletionsBody,
    ) -> Result<ChatCompletions, Error> {
        let response = self
            .client
            .post(self.endpoint.clone())
            .json(&body)
            .send()
            .await?;

        if response.status().is_success() {
            Ok(response.json().await?)
        } else {
            let status = response.status();
            let body = response
                .text()
                .await
                .unwrap_or(String::from("<invalid UTF-8>"));

            let description = serde_json::from_str::<ErrorBody>(&body)
                .map(|e| e.error.message)
                .unwrap_or(body);

            Err(ApiError {
                status,
                description,
            }
            .into())
        }
    }
}

fn build_url(base_url: String, api_version: Option<String>) -> String {
    if let Some(version) = api_version {
        format!("{base_url}{CHAT_COMPLETIONS_ENDPOINT}?api-version={version}")
    } else {
        format!("{base_url}{CHAT_COMPLETIONS_ENDPOINT}")
    }
}

/// Errors generated by [`OpenAiClient`].
#[derive(Debug, thiserror::Error)]
pub enum Error {
    /// Invalid API key charcters.
    #[error("Non ASCII / non visible characters in API key")]
    InvalidCharactersInApiKey(#[from] InvalidHeaderValue),

    /// Reqwest error.
    #[error("Request error: {0}")]
    Request(reqwest::Error),

    /// API (HTTP) error.
    #[error("{0}")]
    Api(#[from] ApiError),
}

impl From<reqwest::Error> for Error {
    fn from(error: reqwest::Error) -> Self {
        // Remove potentially sensitive information.
        Self::Request(error.without_url())
    }
}

/// Error in case of HTTP status != 200 OK.
#[derive(Debug, thiserror::Error)]
pub struct ApiError {
    /// HTTP status code.
    pub status: StatusCode,
    /// Error description.
    pub description: String,
}

impl Display for ApiError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}: {}", self.status, self.description)
    }
}

/// Possible error body (might be incomplete type).
#[derive(Debug, Deserialize)]
pub struct ErrorBody {
    /// Internal `error` JSON object.
    error: OpenAiError,
}

/// Possible `error` field (fields other than `message` omitted).
#[derive(Debug, Deserialize)]
pub struct OpenAiError {
    /// Field `message` of `error` JSON object.
    message: String,
}