litellm-rs 0.4.16

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
107
108
109
110
111
112

//! Image generation endpoint

use crate::core::models::openai::{ImageGenerationRequest, ImageGenerationResponse};
use crate::core::providers::ProviderError;
use crate::core::types::context::RequestContext;
use crate::core::types::image::ImageGenerationRequest as CoreImageRequest;
use crate::core::types::model::ProviderCapability;
use crate::server::state::AppState;
use crate::utils::error::gateway_error::GatewayError;
use actix_web::{HttpRequest, HttpResponse, Result as ActixResult, web};
use tracing::info;

use super::context::handle_ai_request;
use super::provider_selection::select_provider_for_optional_model;

/// Image generation endpoint
///
/// OpenAI-compatible image generation API.
pub async fn image_generations(
    state: web::Data<AppState>,
    req: HttpRequest,
    request: web::Json<ImageGenerationRequest>,
) -> ActixResult<HttpResponse> {
    info!("Image generation request for model: {:?}", request.model);

    handle_ai_request(
        &req,
        request.into_inner(),
        "Image generation",
        |request, context| handle_image_generation_with_state(state.get_ref(), request, context),
    )
    .await
}

/// Handle image generation with app state (UnifiedRouter only)
pub async fn handle_image_generation_with_state(
    state: &AppState,
    request: ImageGenerationRequest,
    context: RequestContext,
) -> Result<ImageGenerationResponse, GatewayError> {
    let unified_router = &state.unified_router;
    handle_image_generation_internal(unified_router, request, context).await
}

async fn handle_image_generation_internal(
    unified_router: &crate::core::router::UnifiedRouter,
    request: ImageGenerationRequest,
    context: RequestContext,
) -> Result<ImageGenerationResponse, GatewayError> {
    let (provider_hint, selected_model) = select_provider_for_optional_model(
        unified_router,
        request.model.as_deref(),
        ProviderCapability::ImageGeneration,
    )?;
    drop(provider_hint);

    let core_request = CoreImageRequest {
        prompt: request.prompt,
        model: Some(selected_model.clone()),
        n: request.n,
        size: request.size,
        response_format: request.response_format,
        user: request.user,
        quality: None,
        style: None,
    };

    let context_for_execution = context.clone();
    let execution = unified_router
        .execute_with_retry(&selected_model, move |deployment_id| {
            let core_request = core_request.clone();
            let context = context_for_execution.clone();
            async move {
                let deployment =
                    unified_router
                        .get_deployment(&deployment_id)
                        .ok_or_else(|| {
                            ProviderError::other("router", "Selected deployment not found")
                        })?;

                let provider = deployment.provider.clone();
                let selected_model = deployment.model.clone();
                drop(deployment);

                let mut request_for_provider = core_request.clone();
                request_for_provider.model = Some(selected_model);
                let response = provider
                    .create_images(request_for_provider, context)
                    .await?;
                Ok((response, 0))
            }
        })
        .await
        .map_err(|(e, _)| GatewayError::Provider(e))?;

    let core_response = execution.0;

    // Convert core response to OpenAI format
    let response = ImageGenerationResponse {
        created: core_response.created,
        data: core_response
            .data
            .into_iter()
            .map(|d| crate::core::models::openai::ImageObject {
                url: d.url,
                b64_json: d.b64_json,
            })
            .collect(),
    };

    Ok(response)
}