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

rig_core/providers/openai/wire/
auth.rs
