Skip to main content

rig_core/providers/openrouter/
client.rs

1use 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
13// ================================================================
14// Main openrouter Client
15// ================================================================
16const 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    /// Create a new openrouter client from the `OPENROUTER_API_KEY` environment variable.
74    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    /// OpenAI-compatible prompt-token details, returned by OpenRouter when a
94    /// provider reports cache activity (Anthropic with cache_control, OpenAI
95    /// with server-side automatic caching).
96    #[serde(default, skip_serializing_if = "Option::is_none")]
97    pub prompt_tokens_details: Option<PromptTokensDetails>,
98}
99
100/// Prompt-token breakdown reported by OpenRouter for cached requests.
101// `usize` matches the parent `Usage` struct in this module; the streaming counterpart
102// in `streaming.rs` uses `u32` to match its own parent.
103#[derive(Clone, Debug, Deserialize, Serialize, Default)]
104pub struct PromptTokensDetails {
105    /// Tokens served from cache (cache hit).
106    #[serde(default)]
107    pub cached_tokens: usize,
108    /// Tokens written to cache on this call (cache miss that populated the cache).
109    #[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    /// Attach OpenRouter app-identification headers (`X-OpenRouter-Title` and `HTTP-Referer`)
143    /// to every request made by this client. `title` appears in the dashboard activity feed
144    /// and rankings page; `url` is the primary app identifier required to create an app page
145    /// on OpenRouter. Invalid (non-ASCII) values are silently skipped.
146    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    /// Assign this app to up to two OpenRouter marketplace categories via the
161    /// `X-OpenRouter-Categories` header. Categories must be lowercase and hyphen-separated
162    /// (e.g. `"cli-agent"`, `"ide-extension"`). OpenRouter silently ignores unrecognized
163    /// categories. Extra categories beyond the first two are not sent. Invalid (non-ASCII)
164    /// values are silently skipped.
165    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}