rig/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 serde::{Deserialize, Serialize};
10use std::fmt::Debug;
11
12// ================================================================
13// Main openrouter Client
14// ================================================================
15const OPENROUTER_API_BASE_URL: &str = "https://openrouter.ai/api/v1";
16
17#[derive(Debug, Default, Clone, Copy)]
18pub struct OpenRouterExt;
19#[derive(Debug, Default, Clone, Copy)]
20pub struct OpenRouterExtBuilder;
21
22type OpenRouterApiKey = BearerAuth;
23
24pub type Client<H = reqwest::Client> = client::Client<OpenRouterExt, H>;
25pub type ClientBuilder<H = reqwest::Client> =
26    client::ClientBuilder<OpenRouterExtBuilder, OpenRouterApiKey, H>;
27
28impl Provider for OpenRouterExt {
29    type Builder = OpenRouterExtBuilder;
30
31    const VERIFY_PATH: &'static str = "/key";
32
33    fn build<H>(
34        _: &crate::client::ClientBuilder<Self::Builder, OpenRouterApiKey, H>,
35    ) -> http_client::Result<Self> {
36        Ok(Self)
37    }
38}
39
40impl<H> Capabilities<H> for OpenRouterExt {
41    type Completion = Capable<super::CompletionModel<H>>;
42    type Embeddings = Nothing;
43    type Transcription = Nothing;
44    type ModelListing = Nothing;
45    #[cfg(feature = "image")]
46    type ImageGeneration = Nothing;
47
48    #[cfg(feature = "audio")]
49    type AudioGeneration = Nothing;
50}
51
52impl DebugExt for OpenRouterExt {}
53
54impl ProviderBuilder for OpenRouterExtBuilder {
55    type Output = OpenRouterExt;
56    type ApiKey = OpenRouterApiKey;
57
58    const BASE_URL: &'static str = OPENROUTER_API_BASE_URL;
59}
60
61impl ProviderClient for Client {
62    type Input = OpenRouterApiKey;
63
64    /// Create a new openrouter client from the `OPENROUTER_API_KEY` environment variable.
65    /// Panics if the environment variable is not set.
66    fn from_env() -> Self {
67        let api_key = std::env::var("OPENROUTER_API_KEY").expect("OPENROUTER_API_KEY not set");
68
69        Self::new(&api_key).unwrap()
70    }
71
72    fn from_val(input: Self::Input) -> Self {
73        Self::new(input).unwrap()
74    }
75}
76
77#[derive(Debug, Deserialize)]
78pub(crate) struct ApiErrorResponse {
79    pub message: String,
80}
81
82#[derive(Debug, Deserialize)]
83#[serde(untagged)]
84pub(crate) enum ApiResponse<T> {
85    Ok(T),
86    Err(ApiErrorResponse),
87}
88
89#[derive(Clone, Debug, Deserialize, Serialize)]
90pub struct Usage {
91    pub prompt_tokens: usize,
92    pub completion_tokens: usize,
93    pub total_tokens: usize,
94}
95
96impl std::fmt::Display for Usage {
97    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
98        write!(
99            f,
100            "Prompt tokens: {} Total tokens: {}",
101            self.prompt_tokens, self.total_tokens
102        )
103    }
104}
105
106impl GetTokenUsage for Usage {
107    fn token_usage(&self) -> Option<crate::completion::Usage> {
108        let mut usage = crate::completion::Usage::new();
109
110        usage.input_tokens = self.prompt_tokens as u64;
111        usage.output_tokens = self.completion_tokens as u64;
112        usage.total_tokens = self.total_tokens as u64;
113
114        Some(usage)
115    }
116}
117#[cfg(test)]
118mod tests {
119    #[test]
120    fn test_client_initialization() {
121        let _client: crate::providers::openrouter::Client =
122            crate::providers::openrouter::Client::new("dummy-key").expect("Client::new() failed");
123        let _client_from_builder: crate::providers::openrouter::Client =
124            crate::providers::openrouter::Client::builder()
125                .api_key("dummy-key")
126                .build()
127                .expect("Client::builder() failed");
128    }
129}
rig/providers/openrouter/client.rs

rig/providers/openrouter/
client.rs
