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