litellm-rs 0.6.0

A high-performance AI Gateway written in Rust, providing OpenAI-compatible APIs with intelligent routing, load balancing, and enterprise features
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

use crate::config::models::gateway::{GatewayPricingConfig, UnpricedModelPolicy};
use crate::core::budget::{BudgetReservation, UnifiedBudgetLimits, UnifiedBudgetReservation};
use crate::core::keys::KeyManager;
use crate::core::pricing_service::{PricingService, PricingUsage};
use crate::utils::error::gateway_error::GatewayError;
use tracing::error;

use super::ImageProxyProvider;

#[allow(clippy::too_many_arguments)]
pub(super) async fn record_image_proxy_spend(
    pricing_config: &GatewayPricingConfig,
    budget_limits: &UnifiedBudgetLimits,
    key_manager: &KeyManager,
    provider: &ImageProxyProvider,
    model: &str,
    usage: &PricingUsage,
    cost: f64,
    unpriced: bool,
    budget_reservation: Option<UnifiedBudgetReservation>,
    api_key_id: Option<uuid::Uuid>,
    key_budget_reservation: Option<BudgetReservation>,
) {
    if unpriced {
        super::super::spend::settle_unpriced_usage(
            pricing_config,
            budget_limits,
            key_manager,
            api_key_id,
            &provider.provider_name,
            model,
            usage,
            budget_reservation,
            key_budget_reservation,
            "image proxy pricing unavailable",
        )
        .await;
        return;
    }

    if let Some(reservation) = budget_reservation {
        if let Err(error) = reservation.settle(cost) {
            error!(
                "failed to settle image proxy budget for provider '{}' model '{}': {error:?}",
                provider.provider_name, model
            );
        }
    } else {
        budget_limits.record_spend(&provider.provider_name, model, cost);
    }
    if let Some(api_key_id) = api_key_id {
        let total_tokens = u64::from(
            usage
                .total_tokens
                .saturating_add(usage.image_tokens.unwrap_or(0)),
        );
        if let Err(error) = key_manager
            .record_usage(api_key_id, total_tokens, cost)
            .await
        {
            error!("failed to record image proxy usage for key {api_key_id}: {error}");
        }
    }
    super::super::spend::settle_api_key_budget_reservation(
        key_budget_reservation,
        cost,
        "image proxy",
    );
}

pub(super) fn image_proxy_cost(
    pricing_service: &PricingService,
    pricing_config: &GatewayPricingConfig,
    provider: &str,
    model: &str,
    usage: &PricingUsage,
) -> Result<(f64, bool), GatewayError> {
    match pricing_service.calculate_loaded_usage_cost_for_provider(provider, model, usage) {
        Ok(breakdown) => Ok((breakdown.total_cost, false)),
        Err(error) => match pricing_config.unpriced_model_policy {
            UnpricedModelPolicy::Reject => Err(GatewayError::Provider(
                super::super::spend::model_not_priced_error(provider, model, error),
            )),
            UnpricedModelPolicy::AllowUnpriced => Ok((
                super::super::spend::fallback_cost_for_usage(pricing_config, usage),
                true,
            )),
        },
    }
}