ai-lib 0.4.0

A unified AI SDK for Rust providing a single interface for multiple AI providers with hybrid architecture
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

//! Request execution module.
//!
//! This module handles single synchronous chat completion requests.
//! It coordinates with:
//! - `ProviderFactory`: Creates provider adapters
//! - `InterceptorPipeline`: Applies interceptors

use super::AiClient;
use crate::rate_limiter::BackpressurePermit;
use crate::types::{AiLibError, ChatCompletionRequest, ChatCompletionResponse};
use tracing::warn;

#[allow(unused_mut)]
pub async fn chat_completion(
    client: &AiClient,
    request: ChatCompletionRequest,
) -> Result<ChatCompletionResponse, AiLibError> {
    // Acquire backpressure permit if configured
    let _bp_permit: Option<BackpressurePermit> = if let Some(ctrl) = &client.backpressure {
        match ctrl.acquire_permit().await {
            Ok(p) => Some(p),
            Err(_) => {
                return Err(AiLibError::RateLimitExceeded(
                    "Backpressure: no permits available".to_string(),
                ))
            }
        }
    } else {
        None
    };

    let provider = client.provider_id();
    let mut processed_request = client.prepare_chat_request(request);

    loop {
        match execute_once(client, &processed_request).await {
            Ok(resp) => return Ok(resp),
            Err(err) => {
                if client.model_resolver.looks_like_invalid_model(&err) {
                    if let Some(resolution) =
                        client.fallback_model_after_invalid(&processed_request.model)
                    {
                        warn!(
                            target = "ai_lib.model",
                            provider = ?provider,
                            failed_model = %processed_request.model,
                            fallback_model = %resolution.model,
                            source = ?resolution.source,
                            "Retrying request with fallback model"
                        );
                        processed_request.model = resolution.model;
                        continue;
                    }
                    let decorated = client.model_resolver.decorate_invalid_model_error(
                        provider,
                        &processed_request.model,
                        err,
                    );
                    return Err(decorated.with_context("AiClient::chat_completion"));
                }

                return Err(err.with_context("AiClient::chat_completion"));
            }
        }
    }
}

async fn execute_once(
    client: &AiClient,
    request: &ChatCompletionRequest,
) -> Result<ChatCompletionResponse, AiLibError> {
    #[cfg(feature = "interceptors")]
    if let Some(p) = &client.interceptor_pipeline {
        let ctx = crate::interceptors::RequestContext {
            provider: client.provider_name().to_lowercase(),
            model: request.model.clone(),
        };
        return p
            .execute(&ctx, request, || client.chat_provider.chat(request.clone()))
            .await;
    }

    client.chat_provider.chat(request.clone()).await
}