rig_core/providers/openrouter/
client.rs1#[cfg(feature = "image")]
2use crate::client::Nothing;
3use crate::{
4 client::{
5 self, BearerAuth, Capabilities, Capable, DebugExt, Provider, ProviderBuilder,
6 ProviderClient,
7 },
8 completion::GetTokenUsage,
9 http_client,
10};
11use http::HeaderValue;
12use serde::{Deserialize, Serialize};
13use std::fmt::Debug;
14
15const OPENROUTER_API_BASE_URL: &str = "https://openrouter.ai/api/v1";
19
20#[derive(Debug, Default, Clone, Copy)]
21pub struct OpenRouterExt;
22#[derive(Debug, Default, Clone, Copy)]
23pub struct OpenRouterExtBuilder;
24
25type OpenRouterApiKey = BearerAuth;
26
27pub type Client<H = reqwest::Client> = client::Client<OpenRouterExt, H>;
28pub type ClientBuilder<H = crate::markers::Missing> =
29 client::ClientBuilder<OpenRouterExtBuilder, OpenRouterApiKey, H>;
30
31impl Provider for OpenRouterExt {
32 type Builder = OpenRouterExtBuilder;
33
34 const VERIFY_PATH: &'static str = "/key";
35}
36
37impl<H> Capabilities<H> for OpenRouterExt {
38 type Completion = Capable<super::CompletionModel<H>>;
39 type Embeddings = Capable<super::EmbeddingModel<H>>;
40 type Transcription = Capable<super::transcription::TranscriptionModel<H>>;
41 type ModelListing = Capable<super::OpenRouterModelLister<H>>;
42 #[cfg(feature = "image")]
43 type ImageGeneration = Nothing;
44
45 #[cfg(feature = "audio")]
46 type AudioGeneration = Capable<super::audio_generation::AudioGenerationModel<H>>;
47}
48
49impl DebugExt for OpenRouterExt {}
50
51impl ProviderBuilder for OpenRouterExtBuilder {
52 type Extension<H>
53 = OpenRouterExt
54 where
55 H: http_client::HttpClientExt;
56 type ApiKey = OpenRouterApiKey;
57
58 const BASE_URL: &'static str = OPENROUTER_API_BASE_URL;
59
60 fn build<H>(
61 _builder: &crate::client::ClientBuilder<Self, Self::ApiKey, H>,
62 ) -> http_client::Result<Self::Extension<H>>
63 where
64 H: http_client::HttpClientExt,
65 {
66 Ok(OpenRouterExt)
67 }
68}
69
70impl ProviderClient for Client {
71 type Input = OpenRouterApiKey;
72 type Error = crate::client::ProviderClientError;
73
74 fn from_env() -> Result<Self, Self::Error> {
76 let api_key = crate::client::required_env_var("OPENROUTER_API_KEY")?;
77
78 Self::new(&api_key).map_err(Into::into)
79 }
80
81 fn from_val(input: Self::Input) -> Result<Self, Self::Error> {
82 Self::new(input).map_err(Into::into)
83 }
84}
85
86#[derive(Debug, Deserialize)]
87pub(crate) struct ApiErrorResponse {
88 pub message: String,
89}
90
91#[derive(Debug, Deserialize)]
92#[serde(untagged)]
93pub(crate) enum ApiResponse<T> {
94 Ok(T),
95 Err(ApiErrorResponse),
96}
97
98#[derive(Clone, Debug, Deserialize, Serialize)]
99pub struct Usage {
100 pub prompt_tokens: usize,
101 #[serde(default)]
102 pub completion_tokens: usize,
103 pub total_tokens: usize,
104 #[serde(default)]
105 pub cost: f64,
106 #[serde(default, skip_serializing_if = "Option::is_none")]
110 pub prompt_tokens_details: Option<PromptTokensDetails>,
111}
112
113#[derive(Clone, Debug, Deserialize, Serialize, Default)]
117pub struct PromptTokensDetails {
118 #[serde(default)]
120 pub cached_tokens: usize,
121 #[serde(default)]
123 pub cache_write_tokens: usize,
124}
125
126impl std::fmt::Display for Usage {
127 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
128 write!(
129 f,
130 "Prompt tokens: {} Total tokens: {}",
131 self.prompt_tokens, self.total_tokens
132 )
133 }
134}
135
136impl GetTokenUsage for Usage {
137 fn token_usage(&self) -> Option<crate::completion::Usage> {
138 let (cached_input, cache_creation) = self
139 .prompt_tokens_details
140 .as_ref()
141 .map(|d| (d.cached_tokens as u64, d.cache_write_tokens as u64))
142 .unwrap_or((0, 0));
143 Some(crate::completion::Usage {
144 input_tokens: self.prompt_tokens as u64,
145 output_tokens: self.completion_tokens as u64,
146 total_tokens: self.total_tokens as u64,
147 cached_input_tokens: cached_input,
148 cache_creation_input_tokens: cache_creation,
149 tool_use_prompt_tokens: 0,
150 reasoning_tokens: 0,
151 })
152 }
153}
154impl<ApiKey, H> client::ClientBuilder<OpenRouterExtBuilder, ApiKey, H> {
155 pub fn with_app_identity(mut self, title: impl AsRef<str>, url: impl AsRef<str>) -> Self {
160 if let Ok(val) = HeaderValue::from_str(title.as_ref()) {
161 self.headers_mut().insert(
162 http::header::HeaderName::from_static("x-openrouter-title"),
163 val,
164 );
165 }
166 if let Ok(val) = HeaderValue::from_str(url.as_ref()) {
167 self.headers_mut()
168 .insert(http::header::HeaderName::from_static("http-referer"), val);
169 }
170 self
171 }
172
173 pub fn with_app_categories<S>(mut self, categories: &[S]) -> Self
179 where
180 S: AsRef<str>,
181 {
182 let joined = categories
183 .iter()
184 .take(2)
185 .map(|c| c.as_ref())
186 .collect::<Vec<_>>()
187 .join(",");
188 if !joined.is_empty()
189 && let Ok(val) = HeaderValue::from_str(&joined)
190 {
191 self.headers_mut().insert(
192 http::header::HeaderName::from_static("x-openrouter-categories"),
193 val,
194 );
195 }
196 self
197 }
198}
199
200#[cfg(test)]
201mod tests {
202 #[test]
203 fn test_client_initialization() {
204 let _client =
205 crate::providers::openrouter::Client::new("dummy-key").expect("Client::new() failed");
206 let _client_from_builder = crate::providers::openrouter::Client::builder()
207 .api_key("dummy-key")
208 .build()
209 .expect("Client::builder() failed");
210 }
211
212 #[test]
213 fn test_with_app_identity_sets_headers() {
214 let client = crate::providers::openrouter::Client::builder()
215 .with_app_identity("My App", "https://myapp.example.com")
216 .api_key("dummy-key")
217 .build()
218 .expect("Client::builder() failed");
219
220 let headers = client.headers();
221 assert_eq!(
222 headers
223 .get("x-openrouter-title")
224 .and_then(|v| v.to_str().ok()),
225 Some("My App"),
226 );
227 assert_eq!(
228 headers.get("http-referer").and_then(|v| v.to_str().ok()),
229 Some("https://myapp.example.com"),
230 );
231 }
232
233 #[test]
234 fn test_without_app_identity_no_extra_headers() {
235 let client = crate::providers::openrouter::Client::builder()
236 .api_key("dummy-key")
237 .build()
238 .expect("Client::builder() failed");
239
240 let headers = client.headers();
241 assert!(headers.get("x-openrouter-title").is_none());
242 assert!(headers.get("http-referer").is_none());
243 }
244
245 #[test]
246 fn test_with_app_categories_sets_header() {
247 let client = crate::providers::openrouter::Client::builder()
248 .with_app_categories(&["cli-agent", "ide-extension"])
249 .api_key("dummy-key")
250 .build()
251 .expect("Client::builder() failed");
252
253 assert_eq!(
254 client
255 .headers()
256 .get("x-openrouter-categories")
257 .and_then(|v| v.to_str().ok()),
258 Some("cli-agent,ide-extension"),
259 );
260 }
261
262 #[test]
263 fn test_with_app_categories_sends_at_most_two_categories() {
264 let client = crate::providers::openrouter::Client::builder()
265 .with_app_categories(&["cli-agent", "ide-extension", "chat"])
266 .api_key("dummy-key")
267 .build()
268 .expect("Client::builder() failed");
269
270 assert_eq!(
271 client
272 .headers()
273 .get("x-openrouter-categories")
274 .and_then(|v| v.to_str().ok()),
275 Some("cli-agent,ide-extension"),
276 );
277 }
278
279 #[test]
280 fn test_with_app_categories_empty_list_no_header() {
281 let empty: [&str; 0] = [];
282 let client = crate::providers::openrouter::Client::builder()
283 .with_app_categories(&empty)
284 .api_key("dummy-key")
285 .build()
286 .expect("Client::builder() failed");
287
288 assert!(client.headers().get("x-openrouter-categories").is_none());
289 }
290
291 #[test]
292 fn test_without_app_categories_no_header() {
293 let client = crate::providers::openrouter::Client::builder()
294 .api_key("dummy-key")
295 .build()
296 .expect("Client::builder() failed");
297
298 assert!(client.headers().get("x-openrouter-categories").is_none());
299 }
300}