litellm_rust/
client.rs

1use crate::config::{Config, ProviderConfig, ProviderKind};
2use crate::error::{LiteLLMError, Result};
3use crate::providers::{anthropic, gemini, openai_compat};
4use crate::registry::Registry;
5use crate::router::{resolve_model, ResolvedModel};
6use crate::stream::ChatStream;
7use crate::types::{
8    ChatRequest, ChatResponse, EmbeddingRequest, EmbeddingResponse, ImageEditRequest, ImageRequest,
9    ImageResponse, VideoRequest, VideoResponse,
10};
11use reqwest::Client;
12
13#[derive(Debug, Clone)]
14pub struct LiteLLM {
15    config: Config,
16    client: Client,
17    registry: Registry,
18}
19
20impl LiteLLM {
21    pub fn new() -> Result<Self> {
22        let registry = Registry::load_embedded()?;
23        Ok(Self {
24            config: Config::default(),
25            client: Client::builder()
26                .timeout(std::time::Duration::from_secs(60))
27                .build()
28                .map_err(LiteLLMError::from)?,
29            registry,
30        })
31    }
32
33    pub fn with_default_provider(mut self, provider: impl Into<String>) -> Self {
34        self.config.default_provider = Some(provider.into());
35        self
36    }
37
38    pub fn with_provider(mut self, name: impl Into<String>, config: ProviderConfig) -> Self {
39        self.config.providers.insert(name.into(), config);
40        self
41    }
42
43    pub fn with_client(mut self, client: Client) -> Self {
44        self.client = client;
45        self
46    }
47
48    pub fn registry(&self) -> &Registry {
49        &self.registry
50    }
51
52    pub async fn completion(&self, mut req: ChatRequest) -> Result<ChatResponse> {
53        let resolved = resolve_model(&req.model, &self.config)?;
54        req.model = resolved.model.clone();
55        dispatch_chat(&self.client, resolved, req).await
56    }
57
58    pub async fn stream_completion(&self, mut req: ChatRequest) -> Result<ChatStream> {
59        let resolved = resolve_model(&req.model, &self.config)?;
60        req.model = resolved.model.clone();
61        match resolved.config.kind {
62            ProviderKind::OpenAICompatible => {
63                openai_compat::chat_stream(&self.client, &resolved.config, req).await
64            }
65            ProviderKind::OpenAIResponses => {
66                openai_compat::responses_chat_stream(&self.client, &resolved.config, req).await
67            }
68            ProviderKind::Anthropic => {
69                anthropic::chat_stream(&self.client, &resolved.config, req).await
70            }
71            ProviderKind::Gemini => gemini::chat_stream(&self.client, &resolved.config, req).await,
72        }
73    }
74
75    pub async fn embedding(&self, mut req: EmbeddingRequest) -> Result<EmbeddingResponse> {
76        let resolved = resolve_model(&req.model, &self.config)?;
77        req.model = resolved.model.clone();
78        match resolved.config.kind {
79            ProviderKind::OpenAICompatible => {
80                openai_compat::embeddings(&self.client, &resolved.config, req).await
81            }
82            ProviderKind::OpenAIResponses => Err(LiteLLMError::Unsupported(
83                "embeddings are not part of the Responses transport".into(),
84            )),
85            ProviderKind::Gemini => gemini::embeddings(&self.client, &resolved.config, req).await,
86            ProviderKind::Anthropic => Err(LiteLLMError::Unsupported(
87                "embeddings not supported for provider".into(),
88            )),
89        }
90    }
91
92    pub async fn image_generation(&self, mut req: ImageRequest) -> Result<ImageResponse> {
93        let resolved = resolve_model(&req.model, &self.config)?;
94        req.model = resolved.model.clone();
95        match resolved.config.kind {
96            ProviderKind::OpenAICompatible => {
97                openai_compat::image_generation(&self.client, &resolved.config, req).await
98            }
99            ProviderKind::OpenAIResponses => Err(LiteLLMError::Unsupported(
100                "image generation is not part of the Responses transport".into(),
101            )),
102            ProviderKind::Gemini => {
103                gemini::image_generation(&self.client, &resolved.config, req).await
104            }
105            _ => Err(LiteLLMError::Unsupported(
106                "image generation not supported for provider".into(),
107            )),
108        }
109    }
110
111    pub async fn image_editing(&self, mut req: ImageEditRequest) -> Result<ImageResponse> {
112        let resolved = resolve_model(&req.model, &self.config)?;
113        req.model = resolved.model.clone();
114        match resolved.config.kind {
115            ProviderKind::OpenAICompatible | ProviderKind::OpenAIResponses => Err(
116                LiteLLMError::Unsupported("image editing not supported for provider".into()),
117            ),
118            ProviderKind::Gemini => {
119                gemini::image_editing(&self.client, &resolved.config, req).await
120            }
121            _ => Err(LiteLLMError::Unsupported(
122                "image editing not supported for provider".into(),
123            )),
124        }
125    }
126
127    pub async fn video_generation(&self, mut req: VideoRequest) -> Result<VideoResponse> {
128        let resolved = resolve_model(&req.model, &self.config)?;
129        req.model = resolved.model.clone();
130        match resolved.config.kind {
131            ProviderKind::OpenAICompatible => {
132                openai_compat::video_generation(&self.client, &resolved.config, req).await
133            }
134            ProviderKind::OpenAIResponses => Err(LiteLLMError::Unsupported(
135                "video generation is not part of the Responses transport".into(),
136            )),
137            ProviderKind::Gemini => {
138                gemini::video_generation(&self.client, &resolved.config, req).await
139            }
140            _ => Err(LiteLLMError::Unsupported(
141                "video generation not supported for provider".into(),
142            )),
143        }
144    }
145
146    pub fn estimate_cost(&self, model: &str, input_tokens: u32, output_tokens: u32) -> Option<f64> {
147        self.registry
148            .estimate_cost(model, input_tokens, output_tokens)
149    }
150
151    /// See [`Registry::estimate_cost_with_cache`] for the token-accounting
152    /// contract: `input_tokens` must exclude `cache_read_tokens` and
153    /// `cache_creation_tokens`. Anthropic's `input_tokens` already does this.
154    /// OpenAI-compatible and Gemini prompt counts include cached tokens, so
155    /// subtract the cached count first.
156    pub fn estimate_cost_with_cache(
157        &self,
158        model: &str,
159        input_tokens: u32,
160        output_tokens: u32,
161        cache_read_tokens: u32,
162        cache_creation_tokens: u32,
163    ) -> Option<f64> {
164        self.registry.estimate_cost_with_cache(
165            model,
166            input_tokens,
167            output_tokens,
168            cache_read_tokens,
169            cache_creation_tokens,
170        )
171    }
172}
173
174async fn dispatch_chat(
175    client: &Client,
176    resolved: ResolvedModel,
177    req: ChatRequest,
178) -> Result<ChatResponse> {
179    match resolved.config.kind {
180        ProviderKind::OpenAICompatible => openai_compat::chat(client, &resolved.config, req).await,
181        ProviderKind::OpenAIResponses => {
182            openai_compat::responses_chat(client, &resolved.config, req).await
183        }
184        ProviderKind::Anthropic => anthropic::chat(client, &resolved.config, req).await,
185        ProviderKind::Gemini => gemini::chat(client, &resolved.config, req).await,
186    }
187}
litellm_rust/client.rs

litellm_rust/
client.rs
