litellm_rs/core/completion/default_router/
router_impl.rs

1//! `DefaultRouter` compatibility methods delegated to the canonical runtime.
2
3use super::*;
4use crate::core::completion::{
5    convert_from_chat_completion_response, convert_messages_to_chat_messages,
6    convert_to_chat_completion_request,
7};
8use crate::core::router::RuntimeHandle;
9use crate::core::types::context::RequestContext;
10use crate::core::types::model::ProviderCapability;
11use futures::stream::StreamExt;
12
13fn reject_provider_overrides(options: &CompletionOptions) -> Result<()> {
14    if options.api_key.is_some()
15        || options.api_base.is_some()
16        || options.api_version.is_some()
17        || options.organization.is_some()
18        || options.headers.is_some()
19        || options.timeout.is_some()
20    {
21        return Err(GatewayError::from(
22            crate::core::providers::ProviderError::invalid_request(
23                "completion",
24                "request-level provider overrides require a matching canonical runtime deployment",
25            ),
26        ));
27    }
28
29    Ok(())
30}
31
32pub(super) async fn complete_with_runtime_handle(
33    handle: &RuntimeHandle,
34    model: &str,
35    messages: Vec<Message>,
36    options: CompletionOptions,
37) -> Result<CompletionResponse> {
38    reject_provider_overrides(&options)?;
39    let chat_messages = convert_messages_to_chat_messages(messages);
40    let chat_request = convert_to_chat_completion_request(model, chat_messages, options)?;
41    let estimated_tokens = RuntimeHandle::estimated_stream_tokens(&chat_request)?;
42    let context = RequestContext::new();
43    let execution = handle
44        .execute_with_selected_deployment_capability_typed(
45            model,
46            &ProviderCapability::ChatCompletion,
47            estimated_tokens,
48            move |deployment| {
49                let mut request = chat_request.clone();
50                let context = context.clone();
51                async move {
52                    request.model = deployment.model.clone();
53                    let response = deployment
54                        .provider
55                        .chat_completion(request, context)
56                        .await
57                        .map_err(|error| (error, None))?;
58                    let tokens = response
59                        .usage
60                        .as_ref()
61                        .map(|usage| u64::from(usage.total_tokens));
62                    Ok((response, tokens))
63                }
64            },
65        )
66        .await
67        .map_err(GatewayError::from)?;
68
69    convert_from_chat_completion_response(execution.result)
70}
71
72pub(super) async fn complete_stream_with_runtime_handle(
73    handle: &RuntimeHandle,
74    model: &str,
75    messages: Vec<Message>,
76    options: CompletionOptions,
77) -> Result<CompletionStream> {
78    reject_provider_overrides(&options)?;
79    let chat_messages = convert_messages_to_chat_messages(messages);
80    let mut chat_request = convert_to_chat_completion_request(model, chat_messages, options)?;
81    chat_request.stream = true;
82    chat_request.stream_options = Some(crate::core::types::chat::StreamOptions {
83        include_usage: Some(true),
84    });
85    let estimated_tokens = RuntimeHandle::estimated_stream_tokens(&chat_request)?;
86    let started_at = std::time::Instant::now();
87    let context = RequestContext::new();
88    let (stream, lease) = handle
89        .execute_stream_with_selected_deployment_capability_typed(
90            model,
91            &ProviderCapability::ChatCompletionStream,
92            estimated_tokens,
93            move |deployment| {
94                let mut request = chat_request.clone();
95                let context = context.clone();
96                async move {
97                    request.model = deployment.model.clone();
98                    deployment
99                        .provider
100                        .chat_completion_stream(request, context)
101                        .await
102                }
103            },
104        )
105        .await
106        .map_err(GatewayError::from)?;
107
108    let mut completion = handle.stream_completion(lease, started_at);
109    Ok(Box::pin(async_stream::stream! {
110        let mut stream = stream;
111        while let Some(chunk) = stream.next().await {
112            match chunk {
113                Ok(chunk) => {
114                    completion.observe_chunk(&chunk);
115                    yield Ok(chunk);
116                }
117                Err(error) => {
118                    drop(stream);
119                    completion.finish_failure(&error).await;
120                    yield Err(GatewayError::from(error));
121                    return;
122                }
123            }
124        }
125        drop(stream);
126        completion.finish_success().await;
127    }))
128}
129
130#[async_trait]
131impl Router for DefaultRouter {
132    async fn complete(
133        &self,
134        model: &str,
135        messages: Vec<Message>,
136        options: CompletionOptions,
137    ) -> Result<CompletionResponse> {
138        let handle = match &self.runtime_binding {
139            Some(binding) => binding.bind(),
140            None => default_runtime().map_err(GatewayError::from)?,
141        };
142        complete_with_runtime_handle(&handle, model, messages, options).await
143    }
144
145    async fn complete_stream(
146        &self,
147        model: &str,
148        messages: Vec<Message>,
149        options: CompletionOptions,
150    ) -> Result<CompletionStream> {
151        let handle = match &self.runtime_binding {
152            Some(binding) => binding.bind(),
153            None => default_runtime().map_err(GatewayError::from)?,
154        };
155        complete_stream_with_runtime_handle(&handle, model, messages, options).await
156    }
157}
litellm_rs/core/completion/default_router/router_impl.rs

litellm_rs/core/completion/default_router/
router_impl.rs
