rig_core/providers/openrouter/
client.rs

1#[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 serde::{Deserialize, Serialize};
12use std::fmt::Debug;
13
14// ================================================================
15// Main openrouter Client
16// ================================================================
17const OPENROUTER_API_BASE_URL: &str = "https://openrouter.ai/api/v1";
18
19#[derive(Debug, Default, Clone, Copy)]
20pub struct OpenRouterExt;
21#[derive(Debug, Default, Clone, Copy)]
22pub struct OpenRouterExtBuilder;
23
24type OpenRouterApiKey = BearerAuth;
25
26pub type Client<H = reqwest::Client> = client::Client<OpenRouterExt, H>;
27pub type ClientBuilder<H = crate::markers::Missing> =
28    client::ClientBuilder<OpenRouterExtBuilder, OpenRouterApiKey, H>;
29
30impl Provider for OpenRouterExt {
31    type Builder = OpenRouterExtBuilder;
32
33    const VERIFY_PATH: &'static str = "/key";
34}
35
36impl<H> Capabilities<H> for OpenRouterExt {
37    type Completion = Capable<super::CompletionModel<H>>;
38    type Embeddings = Capable<super::EmbeddingModel<H>>;
39    type Transcription = Capable<super::transcription::TranscriptionModel<H>>;
40    type ModelListing = Capable<super::OpenRouterModelLister<H>>;
41    #[cfg(feature = "image")]
42    type ImageGeneration = Nothing;
43
44    #[cfg(feature = "audio")]
45    type AudioGeneration = Capable<super::audio_generation::AudioGenerationModel<H>>;
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(Debug, Deserialize)]
86pub(crate) struct ApiErrorResponse {
87    pub message: String,
88}
89
90#[derive(Debug, Deserialize)]
91#[serde(untagged)]
92pub(crate) enum ApiResponse<T> {
93    Ok(T),
94    Err(ApiErrorResponse),
95}
96
97#[derive(Clone, Debug, Deserialize, Serialize)]
98pub struct Usage {
99    pub prompt_tokens: usize,
100    #[serde(default)]
101    pub completion_tokens: usize,
102    pub total_tokens: usize,
103    #[serde(default)]
104    pub cost: f64,
105}
106
107impl std::fmt::Display for Usage {
108    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
109        write!(
110            f,
111            "Prompt tokens: {} Total tokens: {}",
112            self.prompt_tokens, self.total_tokens
113        )
114    }
115}
116
117impl GetTokenUsage for Usage {
118    fn token_usage(&self) -> Option<crate::completion::Usage> {
119        Some(crate::providers::internal::completion_usage(
120            self.prompt_tokens as u64,
121            self.completion_tokens as u64,
122            self.total_tokens as u64,
123            0,
124        ))
125    }
126}
127#[cfg(test)]
128mod tests {
129    #[test]
130    fn test_client_initialization() {
131        let _client =
132            crate::providers::openrouter::Client::new("dummy-key").expect("Client::new() failed");
133        let _client_from_builder = crate::providers::openrouter::Client::builder()
134            .api_key("dummy-key")
135            .build()
136            .expect("Client::builder() failed");
137    }
138}
rig_core/providers/openrouter/client.rs

rig_core/providers/openrouter/
client.rs
