llmshim 0.12.1

Blazing fast LLM API translation layer in pure Rust
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

use super::ratelimit::{penalty_duration, Backpressure, RateKey, RateLimiter};
use crate::policy::{
    AttemptEvent, AttemptIdentity, AttemptOutcome, AttemptPolicy, AttemptPolicyError,
    AttemptPolicyFuture, AttemptPolicyRefusal, AttemptPolicyRefusalKind, DispatchPolicyContext,
    PreparedAttempt,
};
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use tokio::sync::OwnedSemaphorePermit;

struct ProxyAttemptPolicy {
    limiter: Arc<dyn RateLimiter>,
    backpressure: Backpressure,
    active_permits: Mutex<HashMap<uuid::Uuid, OwnedSemaphorePermit>>,
}

pub(crate) fn context(
    limiter: Arc<dyn RateLimiter>,
    backpressure: Backpressure,
    logical_deadline: Option<tokio::time::Instant>,
) -> DispatchPolicyContext {
    let context = DispatchPolicyContext::new(Arc::new(ProxyAttemptPolicy {
        limiter,
        backpressure,
        active_permits: Mutex::new(HashMap::new()),
    }));
    match logical_deadline {
        Some(deadline) => context.with_logical_deadline(deadline),
        None => context,
    }
}

impl AttemptPolicy for ProxyAttemptPolicy {
    fn acquire<'a>(
        &'a self,
        attempt: &'a PreparedAttempt<'a>,
    ) -> AttemptPolicyFuture<'a, Result<(), AttemptPolicyRefusal>> {
        Box::pin(async move {
            let permit = self.backpressure.acquire().await.map_err(|_| {
                AttemptPolicyRefusal::new(AttemptPolicyRefusalKind::CoordinatorUnavailable, None)
            })?;
            let key = RateKey::provider(attempt.identity().provider_name());
            self.limiter
                .acquire(&key, super::ratelimit::estimate_attempt_tokens(attempt))
                .await
                .map_err(|retry| {
                    AttemptPolicyRefusal::new(
                        AttemptPolicyRefusalKind::ProviderLimit,
                        Some(retry.0),
                    )
                })?;
            self.active_permits
                .lock()
                .unwrap()
                .insert(attempt.identity().id(), permit);
            Ok(())
        })
    }

    fn observe<'a>(
        &'a self,
        attempt: &'a AttemptIdentity,
        event: AttemptEvent<'a>,
    ) -> AttemptPolicyFuture<'a, Result<(), AttemptPolicyError>> {
        Box::pin(async move {
            match event {
                AttemptEvent::ResponseHeaders { status: 429 } => {
                    self.limiter
                        .penalize(
                            &RateKey::provider(attempt.provider_name()),
                            penalty_duration(),
                        )
                        .await;
                    Ok(())
                }
                AttemptEvent::Finished(_) => {
                    self.active_permits.lock().unwrap().remove(&attempt.id());
                    Ok(())
                }
                AttemptEvent::ResponseHeaders { .. } | AttemptEvent::Usage { .. } => Ok(()),
            }
        })
    }

    fn observe_abandoned(
        &self,
        attempt: &AttemptIdentity,
        _outcome: AttemptOutcome,
    ) -> Result<(), AttemptPolicyError> {
        self.active_permits.lock().unwrap().remove(&attempt.id());
        Ok(())
    }
}