tokenmiser_providers/
lib.rs

1//! Upstream provider clients, normalized to the OpenAI `chat/completions`
2//! wire shape in both directions. `Provider` is the single seam: proxy code
3//! never branches on provider kind.
4
5use async_trait::async_trait;
6use serde::{Deserialize, Serialize};
7use thiserror::Error;
8use tokenmiser_config::ProviderConfig;
9
10pub mod anthropic;
11pub mod ollama;
12pub mod openai;
13pub mod registry;
14
15pub use registry::ProviderRegistry;
16
17/// Canonical OpenAI-shaped chat completion request.
18#[derive(Debug, Clone, Serialize, Deserialize)]
19pub struct ChatRequest {
20    pub model: String,
21    pub messages: Vec<ChatMessage>,
22    #[serde(skip_serializing_if = "Option::is_none")]
23    pub temperature: Option<f32>,
24    #[serde(skip_serializing_if = "Option::is_none")]
25    pub max_tokens: Option<u32>,
26    #[serde(skip_serializing_if = "Option::is_none")]
27    pub top_p: Option<f32>,
28    #[serde(skip_serializing_if = "Option::is_none")]
29    pub stream: Option<bool>,
30    /// Pass-through for unmodeled fields (`tools`, `response_format`, …).
31    #[serde(flatten)]
32    pub extra: serde_json::Map<String, serde_json::Value>,
33}
34
35#[derive(Debug, Clone, Serialize, Deserialize)]
36pub struct ChatMessage {
37    pub role: String,
38    #[serde(default)]
39    pub content: serde_json::Value,
40    #[serde(flatten)]
41    pub extra: serde_json::Map<String, serde_json::Value>,
42}
43
44/// Canonical OpenAI-shaped chat completion response.
45#[derive(Debug, Clone, Serialize, Deserialize)]
46pub struct ChatResponse {
47    pub id: String,
48    #[serde(default = "default_object")]
49    pub object: String,
50    #[serde(default)]
51    pub created: u64,
52    pub model: String,
53    pub choices: Vec<ChatChoice>,
54    pub usage: Usage,
55    /// Unmodeled fields returned by the upstream.
56    #[serde(flatten, default)]
57    pub extra: serde_json::Map<String, serde_json::Value>,
58}
59
60fn default_object() -> String {
61    "chat.completion".into()
62}
63
64#[derive(Debug, Clone, Serialize, Deserialize)]
65pub struct ChatChoice {
66    pub index: u32,
67    pub message: ChatMessage,
68    #[serde(default)]
69    pub finish_reason: Option<String>,
70    /// Token logprobs, scored by the cascade router for cheap-model
71    /// confidence and passed through to clients opaquely.
72    #[serde(default, skip_serializing_if = "Option::is_none")]
73    pub logprobs: Option<serde_json::Value>,
74}
75
76/// True when the visible `message.content` is empty or whitespace.
77/// Reasoning-mode models can emit every token into a `reasoning` field and
78/// leave content blank, which is a broken response to the caller.
79pub fn response_visible_content_empty(resp: &ChatResponse) -> bool {
80    let Some(choice) = resp.choices.first() else {
81        return true;
82    };
83    match &choice.message.content {
84        serde_json::Value::String(s) => s.trim().is_empty(),
85        serde_json::Value::Null => true,
86        serde_json::Value::Array(arr) => !arr.iter().any(|item| {
87            item.get("text")
88                .and_then(|t| t.as_str())
89                .map(|s| !s.trim().is_empty())
90                .unwrap_or(false)
91        }),
92        _ => false,
93    }
94}
95
96#[derive(Debug, Clone, Default, Serialize, Deserialize)]
97pub struct Usage {
98    #[serde(default)]
99    pub prompt_tokens: u64,
100    #[serde(default)]
101    pub completion_tokens: u64,
102    #[serde(default)]
103    pub total_tokens: u64,
104}
105
106#[derive(Debug, Error)]
107pub enum ProviderError {
108    #[error("provider {name} not registered")]
109    NotFound { name: String },
110    #[error("unknown model `{model}`: no alias, no `provider:model` prefix, no family heuristic match. Known providers: {known_providers}. Hint: use a `provider:model` prefix (e.g. `ollama:llama3.2`) or set `routing.default_provider` in config.")]
111    UnknownModel {
112        model: String,
113        known_providers: String,
114    },
115    #[error("missing api key env var: {0}")]
116    MissingApiKey(String),
117    #[error("upstream http error: {0}")]
118    Http(#[from] reqwest::Error),
119    #[error("upstream returned status {status}: {body}")]
120    Upstream { status: u16, body: String },
121    #[error("response parse error: {0}")]
122    Parse(#[from] serde_json::Error),
123    #[error("provider returned malformed response: {0}")]
124    Malformed(String),
125}
126
127/// A single SSE event from an upstream provider, passed to the client as-is.
128#[derive(Debug, Clone)]
129pub enum StreamChunk {
130    /// Raw SSE bytes, typically `data: {...}\n\n`.
131    Sse(bytes::Bytes),
132    /// Sent after the upstream closes naturally.
133    Done,
134}
135
136#[async_trait]
137pub trait Provider: Send + Sync {
138    fn name(&self) -> &str;
139    fn config(&self) -> &ProviderConfig;
140    async fn complete(&self, req: &ChatRequest) -> Result<ChatResponse, ProviderError>;
141
142    /// Streaming variant. The default falls back to `complete()` and emits
143    /// the whole response as one SSE chunk; natively streaming providers
144    /// override this.
145    async fn stream(
146        &self,
147        req: &ChatRequest,
148    ) -> Result<
149        futures::stream::BoxStream<'static, Result<StreamChunk, ProviderError>>,
150        ProviderError,
151    > {
152        // Wrap the non-streaming response as a single SSE chunk.
153        let resp = self.complete(req).await?;
154        let json = serde_json::to_vec(&resp)?;
155        let mut sse = b"data: ".to_vec();
156        sse.extend_from_slice(&json);
157        sse.extend_from_slice(b"\n\n");
158        let chunk = StreamChunk::Sse(bytes::Bytes::from(sse));
159        let done = StreamChunk::Sse(bytes::Bytes::from_static(b"data: [DONE]\n\n"));
160        use futures::stream::StreamExt;
161        let s = futures::stream::iter(vec![Ok(chunk), Ok(done), Ok(StreamChunk::Done)]);
162        Ok(s.boxed())
163    }
164}
tokenmiser_providers/lib.rs

tokenmiser_providers/
lib.rs
