litellm-rs 0.9.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
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106

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>,
) {
    super::super::execution::completion::observe_usage(super::image_proxy_tokens_used(usage));
    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;
    }

    super::super::spend::capture_ledger_settlement(
        None,
        &provider.provider_name,
        model,
        None,
        Some(cost),
    );
    super::super::spend::settle_api_key_budget_reservation(
        key_budget_reservation,
        cost,
        "image proxy",
    );
    let budget_settlement = async {
        if let Some(reservation) = budget_reservation {
            if let Err(error) =
                crate::server::routes::ai::execution::completion::settle_budget(reservation, cost)
                    .await
            {
                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);
        }
    };
    let usage_record = async {
        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}");
            }
        }
    };
    tokio::join!(budget_settlement, usage_record);
}

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,
            )),
        },
    }
}