arche 4.3.1

An opinionated backend foundation for Axum applications, providing batteries-included integrations for cloud services, databases, authentication, middleware, and logging.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96

use crate::config::{resolve_optional, resolve_with_default};
use crate::error::AppError;
use crate::gcp::token::{ServiceAccountKey, TokenSource};
use std::sync::Arc;

#[derive(Debug, Clone, Default)]
pub struct VertexConfig {
    pub api_key: Option<String>,
    pub project_id: Option<String>,
    pub region: Option<String>,
    pub service_account_key: Option<ServiceAccountKey>,
    pub service_account_key_path: Option<String>,
}

impl VertexConfig {
    pub fn with_api_key(mut self, val: impl Into<String>) -> Self {
        self.api_key = Some(val.into());
        self
    }

    pub fn with_project_id(mut self, val: impl Into<String>) -> Self {
        self.project_id = Some(val.into());
        self
    }

    pub fn with_region(mut self, val: impl Into<String>) -> Self {
        self.region = Some(val.into());
        self
    }

    pub fn with_service_account_key(mut self, key: ServiceAccountKey) -> Self {
        self.service_account_key = Some(key);
        self
    }

    pub fn with_service_account_key_path(mut self, val: impl Into<String>) -> Self {
        self.service_account_key_path = Some(val.into());
        self
    }
}

pub(crate) enum ResolvedAuth {
    ApiKey {
        api_key: String,
    },
    ServiceAccount {
        project_id: String,
        region: String,
        token_source: Arc<TokenSource>,
    },
}

pub(crate) async fn resolve_auth(
    config: Option<VertexConfig>,
    http: reqwest::Client,
) -> Result<ResolvedAuth, AppError> {
    let config = config.unwrap_or_default();

    let api_key = resolve_optional(config.api_key, "VERTEX_API_KEY")
        .or_else(|| resolve_optional::<String>(None, "GEMINI_API_KEY"));

    let project_id = resolve_optional(config.project_id, "VERTEX_PROJECT_ID");

    if let Some(api_key) = api_key {
        return Ok(ResolvedAuth::ApiKey { api_key });
    }

    if let Some(project_id) = project_id {
        let region =
            resolve_with_default(config.region, "VERTEX_REGION", "asia-south1".to_string());

        let sa_key = match (config.service_account_key, config.service_account_key_path) {
            (Some(key), _) => key,
            (None, Some(path)) => ServiceAccountKey::from_path(&path).await?,
            (None, None) => {
                return Err(AppError::internal_error(
                    "Vertex AI service-account auth requires service_account_key or service_account_key_path".into(),
                    None,
                ));
            }
        };

        let token_source = Arc::new(TokenSource::new(http, sa_key));

        return Ok(ResolvedAuth::ServiceAccount {
            project_id,
            region,
            token_source,
        });
    }

    Err(AppError::internal_error(
        "Vertex AI requires either VERTEX_API_KEY/GEMINI_API_KEY, or VERTEX_PROJECT_ID + service_account_key/service_account_key_path".into(),
        None,
    ))
}