chronicle_proxy/
config.rs

1use std::{collections::BTreeMap, time::Duration};
2
3use serde::{Deserialize, Serialize};
4
5use crate::providers::custom::{CustomProvider, ProviderRequestFormat};
6
7/// Configuration for the proxy
8#[derive(Serialize, Deserialize, Debug, Clone, Default)]
9pub struct ProxyConfig {
10    /// Model providers that the proxy should use
11    #[serde(default)]
12    pub providers: Vec<CustomProviderConfig>,
13    /// Aliases that map to providers and models
14    #[serde(default)]
15    pub aliases: Vec<AliasConfig>,
16    /// API keys that the proxy should use
17    #[serde(default)]
18    pub api_keys: Vec<ApiKeyConfig>,
19    /// The default timeout for requests
20    pub default_timeout: Option<Duration>,
21    /// Whether to log to the database or not.
22    pub log_to_database: Option<bool>,
23    /// The user agent to use when making requests
24    pub user_agent: Option<String>,
25}
26
27/// An alias configuration mape a single name to a list of provider-model pairs
28#[derive(Serialize, Deserialize, Debug, Clone, sqlx::FromRow)]
29pub struct AliasConfig {
30    /// A name for this alias
31    pub name: String,
32    /// If true, start from a random provider.
33    /// If false, always start with the first provider, and only use later providers on retry.
34    #[serde(default)]
35    pub random_order: bool,
36    /// The providers and models that this alias represents.
37    pub models: Vec<AliasConfigProvider>,
38}
39
40/// A provider and model to use in an alias
41#[derive(Serialize, Deserialize, Debug, Clone, sqlx::FromRow)]
42pub struct AliasConfigProvider {
43    /// The model to use
44    pub model: String,
45    /// The provider to use
46    pub provider: String,
47    /// An API key configuration to use
48    pub api_key_name: Option<String>,
49}
50
51sqlx_transparent_json_decode::sqlx_json_decode!(AliasConfigProvider);
52
53/// An API key, or where to find one
54#[derive(Serialize, Deserialize, Clone, sqlx::FromRow)]
55pub struct ApiKeyConfig {
56    /// A name for this key
57    pub name: String,
58    /// If "env", the key is an environment variable name to read, rather than the key itself.
59    /// Eventually this will support other pluggable sources.
60    pub source: String,
61    /// The key itself, or if `source` is "env", the name of the environment variable to read.
62    pub value: String,
63}
64
65impl std::fmt::Debug for ApiKeyConfig {
66    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
67        f.debug_struct("ApiKeyConfig")
68            .field("name", &self.name)
69            .field("source", &self.source)
70            .field(
71                "value",
72                if self.source == "env" {
73                    &self.value
74                } else {
75                    &"***"
76                },
77            )
78            .finish_non_exhaustive()
79    }
80}
81
82/// A declarative definition of a model provider
83#[derive(Serialize, Deserialize, Debug, Clone)]
84pub struct CustomProviderConfig {
85    /// The name of the provider, as referenced in proxy requests
86    pub name: String,
87    /// A human-readable name for the provider
88    pub label: Option<String>,
89    /// The url to use
90    pub url: String,
91    /// The API token to pass along
92    pub api_key: Option<String>,
93    /// Where to retrieve the value for `api_key`.
94    /// If `api_key_source` is "env" then `api_key` is an environment variable.
95    /// If it is empty, then `api_key` is assumed to be the token itself, if provided.
96    /// In the future the key sources will be pluggable, to support external secret sources.
97    pub api_key_source: Option<String>,
98    /// What kind of request format this provider uses. Defaults to OpenAI-compatible
99    #[serde(default)]
100    pub format: ProviderRequestFormat,
101    /// Extra headers to pass with the request
102    #[serde(default)]
103    pub headers: BTreeMap<String, String>,
104    /// Models starting with this prefix will use this provider by default.
105    pub prefix: Option<String>,
106}
107
108impl CustomProviderConfig {
109    /// Generate a [CustomProvider] object from the configuration
110    pub fn into_provider(mut self, client: reqwest::Client) -> CustomProvider {
111        if self.api_key_source.as_deref().unwrap_or_default() == "env" {
112            if let Some(token) = self
113                .api_key
114                .as_deref()
115                .and_then(|var| std::env::var(&var).ok())
116            {
117                self.api_key = Some(token);
118            }
119        }
120
121        CustomProvider::new(self, client)
122    }
123
124    /// Add an API token to the [CustomProviderConfig], or if one is not provided, then configure
125    /// it to read from the given environment variable.
126    pub fn with_token_or_env(mut self, token: Option<String>, env: &str) -> Self {
127        match token {
128            Some(token) => {
129                self.api_key = Some(token);
130                self.api_key_source = None;
131            }
132            None => {
133                self.api_key = Some(env.to_string());
134                self.api_key_source = Some("env".to_string());
135            }
136        }
137
138        self
139    }
140}
chronicle_proxy/config.rs

chronicle_proxy/
config.rs
