litellm-rs 0.8.2

A high-performance AI Gateway written in Rust, providing OpenAI-compatible APIs with intelligent routing, load balancing, and enterprise features
Documentation
use std::future::Future;

use crate::core::budget::UnifiedBudgetReservation;
use crate::core::providers::ProviderError;

use super::super::budgeted::{BudgetContext, BudgetReservations, BudgetedCall};

pub(super) async fn reserve_call_settle_media_job<
    T,
    Reserve,
    Call,
    CallFuture,
    Settle,
    SettleFuture,
>(
    budgeted: BudgetedCall,
    reserve: Reserve,
    call: Call,
    settle: Settle,
) -> Result<(T, u64), ProviderError>
where
    Reserve: FnOnce(&BudgetContext) -> Result<Option<UnifiedBudgetReservation>, ProviderError>,
    Call: FnOnce() -> CallFuture,
    CallFuture: Future<Output = Result<T, ProviderError>>,
    Settle: FnOnce(BudgetReservations, BudgetContext) -> SettleFuture,
    SettleFuture: Future<Output = u64>,
{
    let (mut reservations, budget) = budgeted.reserve_for_call(reserve)?;
    match call().await {
        Ok(value) => {
            let tokens_used = settle(reservations, budget).await;
            Ok((value, tokens_used))
        }
        Err(error) if is_accepted_native_task_error(&error) => {
            settle(reservations, budget).await;
            Err(error)
        }
        Err(error) => {
            reservations.cancel();
            Err(error)
        }
    }
}

#[cfg(feature = "providers-extended")]
fn is_accepted_native_task_error(error: &ProviderError) -> bool {
    crate::core::providers::bfl::is_post_submit_error(error)
        || crate::core::providers::stability::is_post_submit_error(error)
}

#[cfg(not(feature = "providers-extended"))]
fn is_accepted_native_task_error(_error: &ProviderError) -> bool {
    false
}