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, 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
106pub const DEFAULT_STREAM_BUFFER: usize = 1024;
107
108pub fn wrap_call_as_streaming(
109    call_future: BoxFut<'static, Result<AssistantMessage, RuntimeError>>,
110) -> Observable<AssistantMessage> {
111    let (tx, events) = broadcast::channel(DEFAULT_STREAM_BUFFER);
112    let cancel = CancellationToken::new();
113    let cancel_for_task = cancel.clone();
114    let output: BoxFut<'static, Result<AssistantMessage, RuntimeError>> = Box::pin(async move {
115        tokio::select! {
116            biased;
117            _ = cancel_for_task.cancelled() => {
118                let _ = tx.send(NodeEvent::LlmDone { total_tokens: 0 });
119                Err(RuntimeError::Cancelled("call cancelled".into()))
120            }
121            result = call_future => {
122                match &result {
123                    Ok(am) => {
124                        let text = am.text_concat();
125                        if !text.is_empty() {
126                            let _ = tx.send(NodeEvent::LlmChunk {
127                                text: text.clone(),
128                                cumulative_tokens: estimate_tokens(&text),
129                            });
130                        }
131                        let _ = tx.send(NodeEvent::LlmDone { total_tokens: am.token_usage.output });
132                    }
133                    Err(_) => {
134                        let _ = tx.send(NodeEvent::LlmDone { total_tokens: 0 });
135                    }
136                }
137                result
138            }
139        }
140    });
141    Observable {
142        output,
143        events,
144        cancel,
145    }
146}
147
148pub fn estimate_tokens(text: &str) -> u64 {
149    ((text.len() as f64) / 3.5).ceil() as u64
150}
151
152pub fn assistant_message_to_value(am: &AssistantMessage) -> Value {
153    let has_structural_part = am
154        .message
155        .parts
156        .iter()
157        .any(|p| !matches!(p, MessagePart::Text { .. }));
158    if has_structural_part {
159        return Value::Message(am.message.clone());
160    }
161    let text = am.text_concat();
162    if text.is_empty() {
163        return Value::Message(am.message.clone());
164    }
165    match serde_json::from_str::<serde_json::Value>(&text) {
166        Ok(json) => Value::from_json(json),
167        Err(_) => Value::Str(text),
168    }
169}
170
171pub fn user_text_message(text: impl Into<String>) -> Message {
172    Message {
173        role: MessageRole::User,
174        parts: vec![MessagePart::Text { text: text.into() }],
175        turn_id: crate::event::TurnId::now(),
176    }
177}
178
179#[derive(Default, Clone)]
180pub struct ProviderRegistry {
181    providers: HashMap<String, Arc<dyn Provider>>,
182    default: Option<String>,
183}
184
185impl ProviderRegistry {
186    pub fn new() -> Self {
187        Self::default()
188    }
189
190    pub fn register(&mut self, provider: Arc<dyn Provider>) {
191        let name = provider.name().to_string();
192        if self.default.is_none() {
193            self.default = Some(name.clone());
194        }
195        self.providers.insert(name, provider);
196    }
197
198    pub fn set_default(&mut self, name: &str) {
199        if self.providers.contains_key(name) {
200            self.default = Some(name.to_string());
201        }
202    }
203
204    pub fn resolve(&self, model: &str) -> Option<Arc<dyn Provider>> {
205        if let Some(p) = self.providers.get(model) {
206            return Some(p.clone());
207        }
208        if let Some((prefix, _)) = model.split_once('/')
209            && let Some(p) = self.providers.get(prefix)
210        {
211            return Some(p.clone());
212        }
213        if let Some(entry) = crate::model_registry::model_entry(model) {
214            let provider_name = format!("config:{}", entry.model);
215            if let Some(p) = self.providers.get(&provider_name) {
216                return Some(p.clone());
217            }
218            let provider_name = format!("config:{model}");
219            if let Some(p) = self.providers.get(&provider_name) {
220                return Some(p.clone());
221            }
222        }
223        self.default
224            .as_ref()
225            .and_then(|n| self.providers.get(n).cloned())
226    }
227
228    pub fn get(&self, name: &str) -> Option<Arc<dyn Provider>> {
229        self.providers.get(name).cloned()
230    }
231}
atman_runtime/provider.rs

atman_runtime/
provider.rs
