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    #[cfg(feature = "image")]
45    type ImageGeneration = Nothing;
46
47    #[cfg(feature = "audio")]
48    type AudioGeneration = Nothing;
49}
50
51impl DebugExt for OpenRouterExt {}
52
53impl ProviderBuilder for OpenRouterExtBuilder {
54    type Output = OpenRouterExt;
55    type ApiKey = OpenRouterApiKey;
56
57    const BASE_URL: &'static str = OPENROUTER_API_BASE_URL;
58}
59
60impl ProviderClient for Client {
61    type Input = OpenRouterApiKey;
62
63    /// Create a new openrouter client from the `OPENROUTER_API_KEY` environment variable.
64    /// Panics if the environment variable is not set.
65    fn from_env() -> Self {
66        let api_key = std::env::var("OPENROUTER_API_KEY").expect("OPENROUTER_API_KEY not set");
67
68        Self::new(&api_key).unwrap()
69    }
70
71    fn from_val(input: Self::Input) -> Self {
72        Self::new(input).unwrap()
73    }
74}
75
76#[derive(Debug, Deserialize)]
77pub(crate) struct ApiErrorResponse {
78    pub message: String,
79}
80
81#[derive(Debug, Deserialize)]
82#[serde(untagged)]
83pub(crate) enum ApiResponse<T> {
84    Ok(T),
85    Err(ApiErrorResponse),
86}
87
88#[derive(Clone, Debug, Deserialize, Serialize)]
89pub struct Usage {
90    pub prompt_tokens: usize,
91    pub completion_tokens: usize,
92    pub total_tokens: usize,
93}
94
95impl std::fmt::Display for Usage {
96    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
97        write!(
98            f,
99            "Prompt tokens: {} Total tokens: {}",
100            self.prompt_tokens, self.total_tokens
101        )
102    }
103}
104
105impl GetTokenUsage for Usage {
106    fn token_usage(&self) -> Option<crate::completion::Usage> {
107        let mut usage = crate::completion::Usage::new();
108
109        usage.input_tokens = self.prompt_tokens as u64;
110        usage.output_tokens = self.completion_tokens as u64;
111        usage.total_tokens = self.total_tokens as u64;
112
113        Some(usage)
114    }
115}
rig/providers/openrouter/client.rs

rig/providers/openrouter/
client.rs
