atman_runtime/
provider.rs

1use std::collections::HashMap;
2use std::sync::Arc;
3
4use tokio::sync::broadcast;
5use tokio_util::sync::CancellationToken;
6
7use crate::error::RuntimeError;
8use crate::event::{NodeEvent, Observable};
9use crate::message::{Message, MessageOrigin, MessagePart, MessageRole};
10use crate::tool::BoxFut;
11use crate::value::Value;
12
13#[derive(Debug, Clone)]
14pub struct LlmRequest {
15    pub model: String,
16    pub messages: Vec<Message>,
17    pub system: Option<String>,
18    pub input: Value,
19    pub schema: Option<String>,
20    pub cache_prompt: bool,
21    pub tools: Vec<crate::tool::ToolSpec>,
22    pub thinking_enabled: bool,
23    /// Seconds without a streaming chunk before the call is cancelled and
24    /// retried.  Default 120 s.  0 disables stall detection.
25    pub stall_timeout_secs: u64,
26}
27
28#[derive(Debug, Clone, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
29pub struct TokenUsage {
30    pub input: u64,
31    pub cached_input: u64,
32    pub output: u64,
33    pub cache_write: u64,
34    pub reasoning_tokens: u64,
35}
36
37impl TokenUsage {
38    pub fn total(&self) -> u64 {
39        self.input
40            .saturating_add(self.cached_input)
41            .saturating_add(self.output)
42            .saturating_add(self.cache_write)
43    }
44}
45
46#[derive(Debug, Clone, Default, PartialEq, Eq)]
47pub struct CallTiming {
48    pub total_ms: u64,
49    pub ttft_ms: Option<u64>,
50}
51
52impl CallTiming {
53    pub fn tokens_per_second(&self, output_tokens: u64) -> Option<f64> {
54        let ttft = self.ttft_ms? as f64;
55        let total = self.total_ms as f64;
56        let gen_ms = total - ttft;
57        if gen_ms <= 0.0 || output_tokens == 0 {
58            return None;
59        }
60        Some(output_tokens as f64 / (gen_ms / 1000.0))
61    }
62}
63
64#[derive(Debug, Clone, PartialEq, Eq)]
65pub enum StopReason {
66    End,
67    ToolUse,
68    Length,
69    Cancelled,
70}
71
72#[derive(Debug, Clone)]
73pub struct AssistantMessage {
74    pub message: Message,
75    pub stop_reason: StopReason,
76    pub token_usage: TokenUsage,
77    #[allow(dead_code)]
78    pub timing: CallTiming,
79    pub model: String,
80    pub response_id: Option<String>,
81}
82
83impl AssistantMessage {
84    pub fn text_only(msg: Message) -> Self {
85        Self {
86            message: msg,
87            stop_reason: StopReason::End,
88            token_usage: TokenUsage::default(),
89            timing: CallTiming::default(),
90            model: String::new(),
91            response_id: None,
92        }
93    }
94
95    pub fn text_concat(&self) -> String {
96        self.message.text_concat()
97    }
98}
99
100pub trait Provider: Send + Sync {
101    fn name(&self) -> &str;
102    fn call<'a>(&'a self, req: LlmRequest) -> BoxFut<'a, Result<AssistantMessage, RuntimeError>>;
103    fn call_streaming(&self, req: LlmRequest) -> Observable<AssistantMessage>;
104
105    fn discover_models(&self) -> BoxFut<'static, Vec<DiscoveredModel>> {
106        Box::pin(async { vec![] })
107    }
108
109    fn test_connection(&self) -> BoxFut<'_, Result<String, String>> {
110        Box::pin(async { Err("test_connection not implemented".into()) })
111    }
112}
113
114#[derive(Debug, Clone)]
115pub struct DiscoveredModel {
116    pub slug: String,
117    pub context_budget: Option<u64>,
118    pub thinking: bool,
119}
120
121pub const DEFAULT_STREAM_BUFFER: usize = 1024;
122
123pub fn wrap_call_as_streaming(
124    call_future: BoxFut<'static, Result<AssistantMessage, RuntimeError>>,
125) -> Observable<AssistantMessage> {
126    let (tx, events) = broadcast::channel(DEFAULT_STREAM_BUFFER);
127    let cancel = CancellationToken::new();
128    let cancel_for_task = cancel.clone();
129    let output: BoxFut<'static, Result<AssistantMessage, RuntimeError>> = Box::pin(async move {
130        tokio::select! {
131            biased;
132            _ = cancel_for_task.cancelled() => {
133                let _ = tx.send(NodeEvent::LlmDone { total_tokens: 0 });
134                Err(RuntimeError::Cancelled("call cancelled".into()))
135            }
136            result = call_future => {
137                match &result {
138                    Ok(am) => {
139                        let text = am.text_concat();
140                        if !text.is_empty() {
141                            let _ = tx.send(NodeEvent::LlmChunk {
142                                text: text.clone(),
143                                cumulative_tokens: estimate_tokens(&text),
144                            });
145                        }
146                        let _ = tx.send(NodeEvent::LlmDone { total_tokens: am.token_usage.output });
147                    }
148                    Err(_) => {
149                        let _ = tx.send(NodeEvent::LlmDone { total_tokens: 0 });
150                    }
151                }
152                result
153            }
154        }
155    });
156    Observable {
157        output,
158        events,
159        cancel,
160    }
161}
162
163pub fn estimate_tokens(text: &str) -> u64 {
164    ((text.len() as f64) / 3.5).ceil() as u64
165}
166
167pub fn assistant_message_to_value(am: &AssistantMessage) -> Value {
168    let has_structural_part = am
169        .message
170        .parts
171        .iter()
172        .any(|p| !matches!(p, MessagePart::Text { .. }));
173    if has_structural_part {
174        return Value::Message(am.message.clone());
175    }
176    let text = am.text_concat();
177    if text.is_empty() {
178        return Value::Message(am.message.clone());
179    }
180    match serde_json::from_str::<serde_json::Value>(&text) {
181        Ok(json) => Value::from_json(json),
182        Err(_) => Value::Str(text),
183    }
184}
185
186pub fn user_text_message(text: impl Into<String>) -> Message {
187    Message {
188        role: MessageRole::User,
189        parts: vec![MessagePart::Text { text: text.into() }],
190        turn_id: crate::event::TurnId::now(),
191        origin: MessageOrigin::User,
192    }
193}
194
195#[derive(Default, Clone)]
196pub struct ProviderRegistry {
197    providers: std::sync::Arc<std::sync::RwLock<HashMap<String, Arc<dyn Provider>>>>,
198    default: std::sync::Arc<std::sync::RwLock<Option<String>>>,
199}
200
201impl ProviderRegistry {
202    pub fn new() -> Self {
203        Self::default()
204    }
205
206    pub fn register(&self, provider: Arc<dyn Provider>) {
207        let name = provider.name().to_string();
208        let mut defaults = self.default.write().unwrap();
209        if defaults.is_none() {
210            *defaults = Some(name.clone());
211        }
212        drop(defaults);
213        self.providers.write().unwrap().insert(name, provider);
214    }
215
216    pub fn set_default(&self, name: &str) {
217        if self.providers.read().unwrap().contains_key(name) {
218            *self.default.write().unwrap() = Some(name.to_string());
219        }
220    }
221
222    pub fn resolve(&self, model: &str) -> Option<Arc<dyn Provider>> {
223        let providers = self.providers.read().unwrap();
224        if let Some(p) = providers.get(model) {
225            return Some(p.clone());
226        }
227        if let Some((prefix, _)) = model.split_once('/')
228            && let Some(p) = providers.get(prefix)
229        {
230            return Some(p.clone());
231        }
232        if let Some(entry) = crate::model_registry::model_entry(model)
233            && let Some(ref provider_name) = entry.provider
234        {
235            if !crate::model_registry::is_provider_enabled(provider_name) {
236                return None;
237            }
238            let config_key = format!("config:{provider_name}");
239            if let Some(p) = providers
240                .get(&config_key)
241                .or_else(|| providers.get(provider_name))
242            {
243                return Some(p.clone());
244            }
245        }
246        None
247    }
248
249    pub fn get(&self, name: &str) -> Option<Arc<dyn Provider>> {
250        self.providers.read().unwrap().get(name).cloned()
251    }
252}
253
254#[cfg(test)]
255mod tests {
256    use super::*;
257    use crate::providers::mock::MockProvider;
258
259    /// Helper: build a registry with a "codex" provider and an "openai" default.
260    fn fixture_registry() -> ProviderRegistry {
261        let reg = ProviderRegistry::new();
262        let codex = Arc::new(MockProvider::new("codex"));
263        reg.register(codex);
264        let openai = Arc::new(MockProvider::new("openai"));
265        reg.register(openai);
266        reg
267    }
268
269    #[test]
270    fn resolve_prefix_match_codex_slash_model() {
271        // "codex/gpt-5.6-terra" → split '/' → prefix "codex" → found
272        let reg = fixture_registry();
273        let p = reg.resolve("codex/gpt-5.6-terra").expect("should resolve");
274        assert_eq!(p.name(), "codex");
275    }
276
277    #[test]
278    fn resolve_returns_none_for_unknown() {
279        let reg = fixture_registry();
280        assert!(reg.resolve("some-unknown-model").is_none());
281    }
282
283    #[test]
284    fn resolve_model_registry_provider_field_takes_priority() {
285        // Simulate the Codex bootstrap: register model entry with provider="codex",
286        // resolve by model name that has no '/' separator.
287        crate::model_registry::register_model_entries(vec![(
288            "codex-auto-review".into(),
289            crate::model_registry::ModelEntry {
290                model: "codex-auto-review".into(),
291                provider: Some("codex".into()),
292                ..Default::default()
293            },
294        )]);
295
296        let reg = fixture_registry();
297        let p = reg
298            .resolve("codex-auto-review")
299            .expect("should resolve via model registry provider field");
300        assert_eq!(p.name(), "codex");
301    }
302}
atman_runtime/provider.rs

atman_runtime/
provider.rs
