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