elph-ai 0.0.21

Unified LLM API with provider collections, auth resolution, and streaming
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

use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;

use tokio::sync::Mutex;

use super::types::{Credential, CredentialStore};

/// In-memory credential store with per-provider serialized writes.
pub struct InMemoryCredentialStore {
    credentials: Mutex<HashMap<String, Credential>>,
    chains: Mutex<HashMap<String, Arc<Mutex<()>>>>,
}

impl Default for InMemoryCredentialStore {
    fn default() -> Self {
        Self::new()
    }
}

impl InMemoryCredentialStore {
    pub fn new() -> Self {
        Self {
            credentials: Mutex::new(HashMap::new()),
            chains: Mutex::new(HashMap::new()),
        }
    }

    async fn lock_chain(&self, provider_id: &str) -> Arc<Mutex<()>> {
        let mut chains = self.chains.lock().await;
        chains
            .entry(provider_id.to_string())
            .or_insert_with(|| Arc::new(Mutex::new(())))
            .clone()
    }
}

#[async_trait::async_trait]
impl CredentialStore for InMemoryCredentialStore {
    async fn read(&self, provider_id: &str) -> Option<Credential> {
        self.credentials.lock().await.get(provider_id).cloned()
    }

    async fn modify(
        &self,
        provider_id: &str,
        f: Box<dyn FnOnce(Option<Credential>) -> Pin<Box<dyn Future<Output = Option<Credential>> + Send>> + Send>,
    ) -> Option<Credential> {
        let chain = self.lock_chain(provider_id).await;
        let _guard = chain.lock().await;
        let current = self.credentials.lock().await.get(provider_id).cloned();
        let next = f(current).await;
        if let Some(ref cred) = next {
            self.credentials
                .lock()
                .await
                .insert(provider_id.to_string(), cred.clone());
        }
        if next.is_some() {
            next
        } else {
            self.credentials.lock().await.get(provider_id).cloned()
        }
    }

    async fn delete(&self, provider_id: &str) {
        let chain = self.lock_chain(provider_id).await;
        let _guard = chain.lock().await;
        self.credentials.lock().await.remove(provider_id);
    }
}