Skip to main content

llm/providers/openrouter/
provider.rs

1use super::types::OpenRouterChatRequest;
2use crate::provider::{error_stream, get_context_window};
3use crate::providers::openai_compatible::{
4    AetherOpenAiConfig, build_chat_request, streaming::create_custom_stream_generic,
5};
6use crate::{
7    Context, LlmError, LlmResponseStream, ProviderAuthMode, ProviderConnectionConfig, ProviderFactory, Result,
8    StreamingModelProvider,
9};
10use async_openai::{Client, config::OpenAIConfig};
11use std::future::ready;
12
13pub struct OpenRouterProvider {
14    client: Client<AetherOpenAiConfig>,
15    model: String,
16}
17
18impl OpenRouterProvider {
19    pub fn new(api_key: String, model: String) -> Result<Self> {
20        let config = openai_config(Some(api_key), ProviderConnectionConfig::default());
21
22        let client = Client::with_config(config);
23        Ok(Self { client, model })
24    }
25
26    pub fn default(model: &str) -> Result<Self> {
27        let api_key = std::env::var("OPENROUTER_API_KEY")
28            .map_err(|_| LlmError::MissingApiKey("OPENROUTER_API_KEY".to_string()))?;
29
30        let config = openai_config(Some(api_key), ProviderConnectionConfig::default());
31
32        let client = Client::with_config(config);
33
34        Ok(Self { client, model: model.to_string() })
35    }
36}
37
38fn openai_config(api_key: Option<String>, connection: ProviderConnectionConfig) -> AetherOpenAiConfig {
39    let api_key = api_key.unwrap_or_default();
40    let api_base = connection.base_url.unwrap_or_else(|| "https://openrouter.ai/api/v1".to_string());
41    let config = OpenAIConfig::new().with_api_key(api_key).with_api_base(api_base);
42    AetherOpenAiConfig::new(config, connection.auth_mode)
43}
44
45impl ProviderFactory for OpenRouterProvider {
46    async fn from_env() -> Result<Self> {
47        Self::from_env_with_connection(ProviderConnectionConfig::default()).await
48    }
49
50    fn from_env_with_connection(connection: ProviderConnectionConfig) -> impl Future<Output = Result<Self>> + Send {
51        ready(provider_from_connection(connection))
52    }
53
54    fn with_model(mut self, model: &str) -> Self {
55        self.model = model.to_string();
56        self
57    }
58}
59
60impl StreamingModelProvider for OpenRouterProvider {
61    fn model(&self) -> Option<crate::LlmModel> {
62        format!("openrouter:{}", self.model).parse().ok()
63    }
64
65    fn context_window(&self) -> Option<u32> {
66        get_context_window("openrouter", &self.model)
67    }
68
69    fn stream_response(&self, context: &Context) -> LlmResponseStream {
70        let mut request = match build_chat_request(&self.model, context, None) {
71            Ok(request) => request,
72            Err(e) => return error_stream(e),
73        };
74        request.prompt_cache_key = context.prompt_cache_key().map(String::from);
75        let request = OpenRouterChatRequest::from_compatible(request, context.session_affinity_key());
76
77        create_custom_stream_generic(&self.client, request)
78    }
79
80    fn display_name(&self) -> String {
81        format!("OpenRouter ({})", self.model)
82    }
83}
84
85fn provider_from_connection(connection: ProviderConnectionConfig) -> Result<OpenRouterProvider> {
86    let api_key = match connection.auth_mode {
87        ProviderAuthMode::Default => Some(
88            std::env::var("OPENROUTER_API_KEY")
89                .map_err(|_| LlmError::MissingApiKey("OPENROUTER_API_KEY".to_string()))?,
90        ),
91        ProviderAuthMode::None => None,
92    };
93    let config = openai_config(api_key, connection);
94    let client = Client::with_config(config);
95
96    Ok(OpenRouterProvider { client, model: String::new() })
97}
98
99#[cfg(test)]
100mod tests {
101    use futures::StreamExt;
102
103    use super::*;
104    use crate::ChatMessage;
105    use crate::providers::test_capture_server::CaptureServer;
106
107    #[tokio::test]
108    async fn stream_response_propagates_prompt_cache_key_and_keeps_cache_control() {
109        let mut server = CaptureServer::start_openrouter().await;
110        let provider = OpenRouterProvider::from_env_with_connection(ProviderConnectionConfig {
111            base_url: Some(server.base_url.clone()),
112            auth_mode: ProviderAuthMode::None,
113            ..Default::default()
114        })
115        .await
116        .unwrap()
117        .with_model("anthropic/claude-haiku-4.5");
118        let mut context = Context::new(vec![ChatMessage::user("Hello")], vec![]);
119        context.set_prompt_cache_key(Some("prefix-abc".to_string()));
120        context.set_session_affinity_key(Some("conversation-abc".to_string()));
121        context.set_reasoning_effort(Some(crate::ReasoningEffort::High));
122
123        let responses = provider.stream_response(&context).collect::<Vec<_>>().await;
124        let captured = server.captured().await;
125
126        assert!(responses.iter().all(Result::is_ok), "{responses:?}");
127        assert!(responses.iter().any(|response| matches!(response, Ok(crate::LlmResponse::Done { .. }))));
128        assert_eq!(captured.path, "/chat/completions");
129        assert_eq!(captured.body["prompt_cache_key"], "prefix-abc");
130        assert_eq!(captured.body["session_id"], "conversation-abc");
131        assert_eq!(captured.body["reasoning_effort"], "high");
132        assert_eq!(captured.body["cache_control"]["type"], "ephemeral");
133        assert_eq!(captured.body["usage"]["include"], true);
134    }
135}