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}
47
48/// The identity a gateway requires on every request, resolved.
49#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
50#[serde(deny_unknown_fields)]
51pub struct CallerIdentity {
52    /// The `originator` header.
53    pub originator: String,
54    /// The `user-agent` header.
55    pub user_agent: String,
56}
57
58/// The user agent a gateway that asks for one is told: the crate, the host,
59/// and who is calling.
60pub(super) fn default_user_agent(originator: &str) -> String {
61    format!(
62        "rig/{} ({} {}; {originator})",
63        env!("CARGO_PKG_VERSION"),
64        std::env::consts::OS,
65        std::env::consts::ARCH,
66    )
67}
68
69impl OpenAIConfig {
70    /// `dialect` with the credential it accepts through its
71    /// [`alternate_auth`](Dialect::alternate_auth) variable, sent with that
72    /// alternative's header.
73    ///
74    /// Azure's account key and its Entra bearer token are both credentials
75    /// for the same account but go out under different headers, so which one
76    /// is held has to be recorded rather than guessed from the value.
77    pub fn with_alternate_key(dialect: &Dialect, api_key: impl Into<Secret>) -> Self {
78        let auth = dialect
79            .alternate_auth
80            .map_or(dialect.quirks.auth, |alternative| alternative.auth);
81        Self {
82            auth,
83            ..Self::with_key(dialect, api_key)
84        }
85    }
86
87    /// The credential `dialect` reads from the environment, and how it is
88    /// sent. Alternative credentials require their own header policy; the
89    /// primary is preferred.
90    pub(crate) fn credential_from_env(dialect: &Dialect) -> Result<(String, Auth), EnvError> {
91        let quirks = &dialect.quirks;
92        Ok(match dialect.alternate_auth {
93            Some(alternative) => match env::optional(dialect.api_key_env)? {
94                Some(api_key) => (api_key, quirks.auth),
95                None => match env::optional(alternative.api_key_env)? {
96                    Some(api_key) => (api_key, alternative.auth),
97                    None => {
98                        return Err(EnvError::Invalid {
99                            name: dialect.api_key_env,
100                            detail: format!(
101                                "either `{}` or `{}` must be set",
102                                dialect.api_key_env, alternative.api_key_env
103                            ),
104                        });
105                    }
106                },
107            },
108            None => (env::required(dialect.api_key_env)?, quirks.auth),
109        })
110    }
111
112    /// Send the credential with `auth`'s header.
113    pub fn with_auth(mut self, auth: Auth) -> Self {
114        self.auth = auth;
115        self
116    }
117
118    /// Name the account the credential belongs to (`ChatGPT-Account-Id`).
119    pub fn with_account_id(mut self, account_id: impl Into<String>) -> Self {
120        self.account_id = Some(account_id.into());
121        self
122    }
123
124    /// Apply the dialect's authentication to a request builder.
125    pub(crate) fn authenticate(&self, builder: http::request::Builder) -> http::request::Builder {
126        match self.auth {
127            Auth::Bearer => {
128                builder.header("Authorization", format!("Bearer {}", self.api_key.expose()))
129            }
130            Auth::OptionalBearer if self.api_key.is_empty() => builder,
131            Auth::OptionalBearer => {
132                builder.header("Authorization", format!("Bearer {}", self.api_key.expose()))
133            }
134            Auth::ApiKeyHeader => builder.header("api-key", self.api_key.expose()),
135        }
136    }
137
138    /// Apply authentication, configured identity, account, and the
139    /// dialect's [`session_header`](super::Quirks::session_header) where it
140    /// names one: `session`, such as a completion's cache key, so requests
141    /// of one conversation share the backend's routing and prompt cache, or
142    /// a fresh id without one.
143    pub(crate) fn headers_in_session(
144        &self,
145        builder: http::request::Builder,
146        session: Option<&str>,
147    ) -> http::request::Builder {
148        let mut builder = self.authenticate(builder);
149        if let Some(identity) = &self.identity {
150            builder = builder
151                .header("originator", &identity.originator)
152                .header(http::header::USER_AGENT, &identity.user_agent);
153        }
154        if let Some(name) = self.dialect.quirks.session_header {
155            builder = builder.header(name, crate::providers::internal::session_id(session));
156        }
157        if let Some(account_id) = &self.account_id {
158            builder = builder.header("ChatGPT-Account-Id", account_id);
159        }
160        builder
161    }
162}