Skip to main content

rig_core/providers/openai/wire/
auth.rs

1//! How an OpenAI-shaped request carries its credential and the caller
2//! identity a gateway requires.
3//!
4//! ```
5//! use rig_core::providers::openai::{OpenAIConfig, wire::Auth};
6//! let config = OpenAIConfig::new("key").with_auth(Auth::OptionalBearer);
7//! assert_eq!(config.auth, Auth::OptionalBearer);
8//! ```
9
10use serde::{Deserialize, Serialize};
11
12use crate::client::env::{self, EnvError};
13use crate::wire::Secret;
14
15use super::{Dialect, OpenAIConfig};
16
17/// Credential header policy, including omission of empty optional tokens.
18#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
19pub enum Auth {
20    /// `Authorization: Bearer <key>`.
21    Bearer,
22    /// `Authorization: Bearer <key>`, omitted entirely when the key is empty.
23    OptionalBearer,
24    /// Azure's `api-key: <key>`.
25    ApiKeyHeader,
26}
27
28/// Alternative credential environment variable and its authentication policy.
29#[derive(Clone, Copy, Debug, PartialEq, Eq)]
30pub struct AuthAlternative {
31    /// The variable holding this credential.
32    pub api_key_env: &'static str,
33    /// How it is sent.
34    pub auth: Auth,
35}
36
37/// The caller identity a gateway requires on every request.
38#[derive(Debug, Clone, Copy, PartialEq, Eq)]
39pub struct Identity {
40    /// The `originator` header's default value.
41    pub originator: &'static str,
42    /// The environment variable overriding `originator`.
43    pub originator_env: &'static str,
44    /// The environment variable overriding `user-agent`.
45    pub user_agent_env: &'static str,
46    /// Whether every request carries a fresh `session_id` header.
47    pub session_ids: bool,
48}
49
50/// The identity a gateway requires on every request, resolved.
51#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
52#[serde(deny_unknown_fields)]
53pub struct CallerIdentity {
54    /// The `originator` header.
55    pub originator: String,
56    /// The `user-agent` header.
57    pub user_agent: String,
58}
59
60/// The user agent a gateway that asks for one is told: the crate, the host,
61/// and who is calling.
62pub(super) fn default_user_agent(originator: &str) -> String {
63    format!(
64        "rig/{} ({} {}; {originator})",
65        env!("CARGO_PKG_VERSION"),
66        std::env::consts::OS,
67        std::env::consts::ARCH,
68    )
69}
70
71impl OpenAIConfig {
72    /// `dialect` with the credential it accepts through its
73    /// [`alternate_auth`](Dialect::alternate_auth) variable, sent with that
74    /// alternative's header.
75    ///
76    /// Azure's account key and its Entra bearer token are both credentials
77    /// for the same account but go out under different headers, so which one
78    /// is held has to be recorded rather than guessed from the value.
79    pub fn with_alternate_key(dialect: &Dialect, api_key: impl Into<Secret>) -> Self {
80        let auth = dialect
81            .alternate_auth
82            .map_or(dialect.quirks.auth, |alternative| alternative.auth);
83        Self {
84            auth,
85            ..Self::with_key(dialect, api_key)
86        }
87    }
88
89    /// The credential `dialect` reads from the environment, and how it is
90    /// sent. Alternative credentials require their own header policy; the
91    /// primary is preferred.
92    pub(crate) fn credential_from_env(dialect: &Dialect) -> Result<(String, Auth), EnvError> {
93        let quirks = &dialect.quirks;
94        Ok(match dialect.alternate_auth {
95            Some(alternative) => match env::optional(dialect.api_key_env)? {
96                Some(api_key) => (api_key, quirks.auth),
97                None => match env::optional(alternative.api_key_env)? {
98                    Some(api_key) => (api_key, alternative.auth),
99                    None => {
100                        return Err(EnvError::Invalid {
101                            name: dialect.api_key_env,
102                            detail: format!(
103                                "either `{}` or `{}` must be set",
104                                dialect.api_key_env, alternative.api_key_env
105                            ),
106                        });
107                    }
108                },
109            },
110            None => (env::required(dialect.api_key_env)?, quirks.auth),
111        })
112    }
113
114    /// Send the credential with `auth`'s header.
115    pub fn with_auth(mut self, auth: Auth) -> Self {
116        self.auth = auth;
117        self
118    }
119
120    /// Name the account the credential belongs to (`ChatGPT-Account-Id`).
121    pub fn with_account_id(mut self, account_id: impl Into<String>) -> Self {
122        self.account_id = Some(account_id.into());
123        self
124    }
125
126    /// Apply the dialect's authentication to a request builder.
127    pub(crate) fn authenticate(&self, builder: http::request::Builder) -> http::request::Builder {
128        match self.auth {
129            Auth::Bearer => {
130                builder.header("Authorization", format!("Bearer {}", self.api_key.expose()))
131            }
132            Auth::OptionalBearer if self.api_key.is_empty() => builder,
133            Auth::OptionalBearer => {
134                builder.header("Authorization", format!("Bearer {}", self.api_key.expose()))
135            }
136            Auth::ApiKeyHeader => builder.header("api-key", self.api_key.expose()),
137        }
138    }
139
140    /// Apply authentication, configured identity, account, and per-request session headers.
141    pub(crate) fn headers(&self, builder: http::request::Builder) -> http::request::Builder {
142        let mut builder = self.authenticate(builder);
143        if let Some(identity) = &self.identity {
144            builder = builder
145                .header("originator", &identity.originator)
146                .header(http::header::USER_AGENT, &identity.user_agent);
147        }
148        if self
149            .dialect
150            .quirks
151            .identity
152            .is_some_and(|identity| identity.session_ids)
153        {
154            // Session identity must be fresh for each request.
155            builder = builder.header("session_id", crate::providers::chatgpt::session_id());
156        }
157        if let Some(account_id) = &self.account_id {
158            builder = builder.header("ChatGPT-Account-Id", account_id);
159        }
160        builder
161    }
162}