Skip to main content

elph_ai/auth/
types.rs

1use std::future::Future;
2use std::pin::Pin;
3use std::sync::Arc;
4
5use crate::types::{ImagesModel, Model, ProviderEnv, ProviderHeaders};
6
7#[derive(Debug, Clone)]
8pub struct ModelAuth {
9    pub api_key: Option<String>,
10    pub headers: Option<ProviderHeaders>,
11    pub base_url: Option<String>,
12}
13
14#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
15pub struct ApiKeyCredential {
16    #[serde(rename = "type")]
17    pub kind: String,
18    pub key: Option<String>,
19    #[serde(default, skip_serializing_if = "Option::is_none")]
20    pub env: Option<ProviderEnv>,
21}
22
23impl ApiKeyCredential {
24    pub fn new(key: impl Into<String>) -> Self {
25        Self {
26            kind: "api_key".to_string(),
27            key: Some(key.into()),
28            env: None,
29        }
30    }
31}
32
33#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
34pub struct OAuthCredential {
35    #[serde(rename = "type")]
36    pub kind: String,
37    pub access: String,
38    pub refresh: String,
39    pub expires: i64,
40    #[serde(default, skip_serializing_if = "Option::is_none")]
41    pub account_id: Option<String>,
42    #[serde(default, skip_serializing_if = "Option::is_none")]
43    pub enterprise_url: Option<String>,
44    #[serde(default, skip_serializing_if = "Option::is_none")]
45    pub available_model_ids: Option<Vec<String>>,
46}
47
48#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
49#[serde(tag = "type")]
50pub enum Credential {
51    #[serde(rename = "api_key")]
52    ApiKey(ApiKeyCredential),
53    #[serde(rename = "oauth")]
54    OAuth(OAuthCredential),
55}
56
57#[async_trait::async_trait]
58pub trait CredentialStore: Send + Sync {
59    async fn read(&self, provider_id: &str) -> Option<Credential>;
60    async fn modify(
61        &self,
62        provider_id: &str,
63        f: Box<dyn FnOnce(Option<Credential>) -> Pin<Box<dyn Future<Output = Option<Credential>> + Send>> + Send>,
64    ) -> Option<Credential>;
65    async fn delete(&self, provider_id: &str);
66}
67
68#[async_trait::async_trait]
69pub trait AuthContext: Send + Sync {
70    async fn env(&self, name: &str) -> Option<String>;
71    async fn file_exists(&self, path: &str) -> bool;
72}
73
74#[derive(Debug, Clone)]
75pub struct AuthResult {
76    pub auth: ModelAuth,
77    pub env: Option<ProviderEnv>,
78    pub source: Option<String>,
79}
80
81#[derive(Debug, Clone)]
82pub enum AuthPrompt {
83    Text {
84        message: String,
85        placeholder: Option<String>,
86    },
87    Secret {
88        message: String,
89        placeholder: Option<String>,
90    },
91    Select {
92        message: String,
93        options: Vec<AuthSelectOption>,
94    },
95    ManualCode {
96        message: String,
97        placeholder: Option<String>,
98    },
99}
100
101#[derive(Debug, Clone)]
102pub struct AuthSelectOption {
103    pub id: String,
104    pub label: String,
105    pub description: Option<String>,
106}
107
108#[derive(Debug, Clone)]
109pub enum AuthEvent {
110    AuthUrl {
111        url: String,
112        instructions: Option<String>,
113    },
114    DeviceCode {
115        user_code: String,
116        verification_uri: String,
117        interval_seconds: Option<u32>,
118        expires_in_seconds: Option<u32>,
119    },
120    Progress {
121        message: String,
122    },
123}
124
125#[async_trait::async_trait]
126pub trait AuthLoginCallbacks: Send + Sync {
127    async fn prompt(&self, prompt: AuthPrompt) -> anyhow::Result<String>;
128    fn notify(&self, event: AuthEvent);
129}
130
131pub type ApiKeyResolveFn =
132    Arc<dyn Fn(AuthResolveInput) -> Pin<Box<dyn Future<Output = Option<AuthResult>> + Send>> + Send + Sync>;
133pub type ApiKeyLoginFn = Arc<
134    dyn Fn(Arc<dyn AuthLoginCallbacks>) -> Pin<Box<dyn Future<Output = anyhow::Result<ApiKeyCredential>> + Send>>
135        + Send
136        + Sync,
137>;
138
139pub struct AuthResolveInput {
140    pub model: AuthModel,
141    pub ctx: Arc<dyn AuthContext>,
142    pub credential: Option<ApiKeyCredential>,
143}
144
145#[derive(Clone)]
146pub enum AuthModel {
147    Chat(Model),
148    Images(ImagesModel),
149}
150
151#[derive(Clone)]
152pub struct ApiKeyAuth {
153    pub name: String,
154    pub resolve: ApiKeyResolveFn,
155    pub login: Option<ApiKeyLoginFn>,
156}
157
158pub type OAuthLoginFn = Arc<
159    dyn Fn(Arc<dyn AuthLoginCallbacks>) -> Pin<Box<dyn Future<Output = anyhow::Result<OAuthCredential>> + Send>>
160        + Send
161        + Sync,
162>;
163pub type OAuthRefreshFn =
164    Arc<dyn Fn(OAuthCredential) -> Pin<Box<dyn Future<Output = anyhow::Result<OAuthCredential>> + Send>> + Send + Sync>;
165pub type OAuthToAuthFn =
166    Arc<dyn Fn(OAuthCredential) -> Pin<Box<dyn Future<Output = anyhow::Result<ModelAuth>> + Send>> + Send + Sync>;
167
168#[derive(Clone)]
169pub struct OAuthAuth {
170    pub name: String,
171    pub login: OAuthLoginFn,
172    pub refresh: OAuthRefreshFn,
173    pub to_auth: OAuthToAuthFn,
174}
175
176#[derive(Clone, Default)]
177pub struct ProviderAuth {
178    pub api_key: Option<ApiKeyAuth>,
179    pub oauth: Option<OAuthAuth>,
180}