litellm-rs 0.9.0

//! `DefaultRouter` compatibility methods delegated to the canonical runtime.

use super::*;
use crate::core::completion::{
    convert_from_chat_completion_response, convert_messages_to_chat_messages,
    convert_to_chat_completion_request,
};
use crate::core::router::RuntimeHandle;
use crate::core::types::context::RequestContext;
use crate::core::types::model::ProviderCapability;
use futures::stream::StreamExt;

fn reject_provider_overrides(options: &CompletionOptions) -> Result<()> {
    if options.api_key.is_some()
        || options.api_base.is_some()
        || options.api_version.is_some()
        || options.organization.is_some()
        || options.headers.is_some()
        || options.timeout.is_some()
    {
        return Err(GatewayError::from(
            crate::core::providers::ProviderError::invalid_request(
                "completion",
                "request-level provider overrides require a matching canonical runtime deployment",
            ),
        ));
    }

    Ok(())
}

pub(super) async fn complete_with_runtime_handle(
    handle: &RuntimeHandle,
    model: &str,
    messages: Vec<Message>,
    options: CompletionOptions,
) -> Result<CompletionResponse> {
    reject_provider_overrides(&options)?;
    let chat_messages = convert_messages_to_chat_messages(messages);
    let chat_request = convert_to_chat_completion_request(model, chat_messages, options)?;
    let estimated_tokens = RuntimeHandle::estimated_stream_tokens(&chat_request)?;
    let context = RequestContext::new();
    let execution = handle
        .execute_with_selected_deployment_capability_typed(
            model,
            &ProviderCapability::ChatCompletion,
            estimated_tokens,
            move |deployment| {
                let mut request = chat_request.clone();
                let context = context.clone();
                async move {
                    request.model = deployment.model.clone();
                    let response = deployment
                        .provider
                        .chat_completion(request, context)
                        .await
                        .map_err(|error| (error, None))?;
                    let tokens = response
                        .usage
                        .as_ref()
                        .map(|usage| u64::from(usage.total_tokens));
                    Ok((response, tokens))
                }
            },
        )
        .await
        .map_err(GatewayError::from)?;

    convert_from_chat_completion_response(execution.result)
}

pub(super) async fn complete_stream_with_runtime_handle(
    handle: &RuntimeHandle,
    model: &str,
    messages: Vec<Message>,
    options: CompletionOptions,
) -> Result<CompletionStream> {
    reject_provider_overrides(&options)?;
    let chat_messages = convert_messages_to_chat_messages(messages);
    let mut chat_request = convert_to_chat_completion_request(model, chat_messages, options)?;
    chat_request.stream = true;
    chat_request.stream_options = Some(crate::core::types::chat::StreamOptions {
        include_usage: Some(true),
    });
    let estimated_tokens = RuntimeHandle::estimated_stream_tokens(&chat_request)?;
    let started_at = std::time::Instant::now();
    let context = RequestContext::new();
    let (stream, lease) = handle
        .execute_stream_with_selected_deployment_capability_typed(
            model,
            &ProviderCapability::ChatCompletionStream,
            estimated_tokens,
            move |deployment| {
                let mut request = chat_request.clone();
                let context = context.clone();
                async move {
                    request.model = deployment.model.clone();
                    deployment
                        .provider
                        .chat_completion_stream(request, context)
                        .await
                }
            },
        )
        .await
        .map_err(GatewayError::from)?;

    let mut completion = handle.stream_completion(lease, started_at);
    Ok(Box::pin(async_stream::stream! {
        let mut stream = stream;
        while let Some(chunk) = stream.next().await {
            match chunk {
                Ok(chunk) => {
                    completion.observe_chunk(&chunk);
                    yield Ok(chunk);
                }
                Err(error) => {
                    drop(stream);
                    completion.finish_failure(&error).await;
                    yield Err(GatewayError::from(error));
                    return;
                }
            }
        }
        drop(stream);
        completion.finish_success().await;
    }))
}

#[async_trait]
impl Router for DefaultRouter {
    async fn complete(
        &self,
        model: &str,
        messages: Vec<Message>,
        options: CompletionOptions,
    ) -> Result<CompletionResponse> {
        let handle = match &self.runtime_binding {
            Some(binding) => binding.bind(),
            None => default_runtime().map_err(GatewayError::from)?,
        };
        complete_with_runtime_handle(&handle, model, messages, options).await
    }

    async fn complete_stream(
        &self,
        model: &str,
        messages: Vec<Message>,
        options: CompletionOptions,
    ) -> Result<CompletionStream> {
        let handle = match &self.runtime_binding {
            Some(binding) => binding.bind(),
            None => default_runtime().map_err(GatewayError::from)?,
        };
        complete_stream_with_runtime_handle(&handle, model, messages, options).await
    }
}