Skip to main content

ort_openrouter_cli/common/
data.rs

1//! ort: Open Router CLI
2//! https://github.com/grahamking/ort
3//!
4//! MIT License
5//! Copyright (c) 2025 Graham King
6
7#![allow(dead_code)]
8
9use core::str::FromStr;
10
11extern crate alloc;
12use alloc::string::{String, ToString};
13use alloc::vec;
14use alloc::vec::Vec;
15
16use crate::common::base64;
17use crate::utils::filename_read_to_bytes;
18use crate::{ErrorKind, OrtError, OrtResult, ort_error};
19
20const DEFAULT_SHOW_REASONING: bool = false;
21const DEFAULT_QUIET: bool = false;
22const IMAGE_EXT: [&str; 4] = ["jpg", "JPG", "png", "PNG"];
23
24// Keep in sync with src/lib.rs
25pub const DEFAULT_MODEL: &str = "google/gemma-3n-e4b-it:free";
26
27const MIME_TYPES: [(&str, &str); 2] = [("jpg", "image/jpeg"), ("png", "image/png")];
28
29// {
30//  "id":"gen-1756743299-7ytIBcjALWQQShwMQfw9",
31//  "provider":"Meta",
32//  "model":"meta-llama/llama-3.3-8b-instruct:free",
33//  "object":"chat.completion.chunk",
34//  "created":1756743300,
35//  "choices":[
36//      {
37//      "index":0,
38//      "delta":{"role":"assistant","content":""},
39//      "finish_reason":null,
40//      "native_finish_reason":null,
41//      "logprobs":null
42//      }
43//  ],
44//  "usage":{
45//      "prompt_tokens":42,
46//      "completion_tokens":2,
47//      "total_tokens":44,
48//      "cost":0,"
49//      is_byok":false,
50//      "prompt_tokens_details":{"cached_tokens":0,"audio_tokens":0},
51//      "cost_details":{"upstream_inference_cost":null,"upstream_inference_prompt_cost":0,"upstream_inference_completions_cost":0},
52//      "completion_tokens_details":{"reasoning_tokens":0,"image_tokens":0}}
53//  }
54
55pub struct ChatCompletionsResponse {
56    pub provider: Option<String>,
57    pub model: Option<String>,
58    pub choices: Vec<Choice>,
59    pub usage: Option<Usage>,
60}
61
62pub struct Choice {
63    pub delta: Message,
64}
65
66pub struct Usage {
67    pub cost: f32, // In dollars, usually a very small fraction
68}
69
70pub struct LastData {
71    pub opts: PromptOpts,
72    pub messages: Vec<Message>,
73}
74
75#[derive(Clone)]
76pub struct PromptOpts {
77    pub prompt: Option<String>,
78    /// Model IDs, e.g. 'moonshotai/kimi-k2'
79    pub models: Vec<String>,
80    /// Prefered provider slug
81    pub provider: Option<String>,
82    /// System prompt
83    pub system: Option<String>,
84    /// How to choose a provider
85    pub priority: Option<Priority>,
86    /// Reasoning config
87    pub reasoning: Option<ReasoningConfig>,
88    /// Show reasoning output
89    pub show_reasoning: Option<bool>,
90    /// Don't show stats after request
91    pub quiet: Option<bool>,
92    /// Whether to merge in the default settings from config file
93    pub merge_config: bool,
94    /// Images to attach to the request.
95    pub files: Vec<String>,
96}
97
98impl Default for PromptOpts {
99    fn default() -> Self {
100        Self {
101            prompt: None,
102            models: vec![DEFAULT_MODEL.to_string()],
103            provider: None,
104            system: None,
105            priority: None,
106            reasoning: Some(ReasoningConfig::default()),
107            show_reasoning: Some(false),
108            quiet: Some(false),
109            merge_config: true,
110            files: vec![],
111        }
112    }
113}
114
115impl PromptOpts {
116    // Replace any blank or None fields on Self with values from other
117    // or with the defaults.
118    // After this call a PromptOpts is ready to use.
119    pub fn merge(&mut self, o: PromptOpts) {
120        self.prompt.get_or_insert(o.prompt.unwrap_or_default());
121        self.quiet.get_or_insert(o.quiet.unwrap_or(DEFAULT_QUIET));
122        if self.models.is_empty() {
123            // We don't merge the models, otherwise we'd try to query both the
124            // cmd line one, and the config file default.
125            self.models = o.models;
126        }
127        if let Some(provider) = o.provider {
128            self.provider.get_or_insert(provider);
129        }
130        if let Some(system) = o.system {
131            self.system.get_or_insert(system);
132        }
133        if let Some(priority) = o.priority {
134            self.priority.get_or_insert(priority);
135        }
136        self.reasoning
137            .get_or_insert(o.reasoning.unwrap_or_default());
138        self.show_reasoning
139            .get_or_insert(o.show_reasoning.unwrap_or(DEFAULT_SHOW_REASONING));
140        self.files.extend(o.files);
141    }
142}
143
144#[derive(Default, Debug, Clone, Copy)]
145pub enum Priority {
146    Price,
147    #[default]
148    Latency,
149    Throughput,
150}
151
152impl Priority {
153    pub fn as_str(&self) -> &'static str {
154        match self {
155            Priority::Price => "price",
156            Priority::Latency => "latency",
157            Priority::Throughput => "throughput",
158        }
159    }
160}
161
162impl FromStr for Priority {
163    type Err = OrtError;
164
165    fn from_str(s: &str) -> Result<Self, Self::Err> {
166        match s.to_lowercase().as_str() {
167            "price" => Ok(Priority::Price),
168            "latency" => Ok(Priority::Latency),
169            "throughput" => Ok(Priority::Throughput),
170            _ => Err(ort_error(
171                ErrorKind::FormatError,
172                "Priority: Invalid string value",
173            )), // Handle unknown strings
174        }
175    }
176}
177
178#[derive(Default, Debug, Clone)]
179pub struct ReasoningConfig {
180    pub enabled: bool,
181    pub effort: Option<ReasoningEffort>,
182    pub tokens: Option<u32>,
183}
184
185impl ReasoningConfig {
186    pub fn off() -> Self {
187        Self {
188            enabled: false,
189            ..Default::default()
190        }
191    }
192}
193
194#[derive(Default, Debug, Clone, Copy, PartialEq)]
195pub enum ReasoningEffort {
196    None, // GPT 5.x only
197    Low,
198    #[default]
199    Medium,
200    High,
201    XHigh, // GPT 5.x only
202}
203
204impl ReasoningEffort {
205    pub fn as_str(&self) -> &'static str {
206        match self {
207            ReasoningEffort::None => "none",
208            ReasoningEffort::Low => "low",
209            ReasoningEffort::Medium => "medium",
210            ReasoningEffort::High => "high",
211            ReasoningEffort::XHigh => "xhigh",
212        }
213    }
214}
215
216#[derive(Debug, Clone)]
217pub struct Message {
218    pub role: Role,
219    pub content: Vec<Content>,
220    pub reasoning: Option<String>,
221}
222
223impl Message {
224    pub fn new(role: Role, content: Option<String>, reasoning: Option<String>) -> Self {
225        let content = content.map_or_else(Vec::new, |content| vec![Content::Text(content)]);
226        Self::with_content(role, content, reasoning)
227    }
228
229    pub fn with_content(role: Role, content: Vec<Content>, reasoning: Option<String>) -> Self {
230        Message {
231            role,
232            content,
233            reasoning,
234        }
235    }
236    pub fn system(content: String) -> Self {
237        Self::new(Role::System, Some(content), None)
238    }
239    pub fn user(content: String) -> Self {
240        Self::new(Role::User, Some(content), None)
241    }
242    pub fn assistant(content: String) -> Self {
243        Self::new(Role::Assistant, Some(content), None)
244    }
245
246    pub fn with_files(prompt: String, filenames: &[String]) -> OrtResult<Self> {
247        // First message is the user's prompt as Text
248        let mut m = Self::user(prompt);
249        // Then the files as Image
250        for f in filenames {
251            if f.starts_with("http") {
252                m.content.push(Content::ImageUrl(f.clone()));
253            } else {
254                let pf = PromptFile::load(f).map_err(|err| ort_error(ErrorKind::Other, err))?;
255                m.content.push(pf.into_content());
256            }
257        }
258        Ok(m)
259    }
260
261    pub fn text(&self) -> Option<&str> {
262        match self.content.as_slice() {
263            [Content::Text(text)] => Some(text.as_str()),
264            _ => None,
265        }
266    }
267
268    /// Estimate size in bytes
269    pub fn size(&self) -> u32 {
270        let content_len: usize = self.content.iter().map(Content::len).sum();
271        let reasoning_len = self.reasoning.as_ref().map(|c| c.len()).unwrap_or(0);
272        (content_len.max(reasoning_len) + 10) as u32
273    }
274}
275
276#[derive(Debug, Clone)]
277pub enum Content {
278    Text(String),
279    // Just the base64 encoded data
280    Image {
281        mime_type: &'static str,
282        base64: String,
283    },
284    ImageUrl(String),
285    File(PromptFile),
286}
287
288impl Content {
289    pub fn len(&self) -> usize {
290        use Content::*;
291        match self {
292            Text(s) => s.len(),
293            Image { base64, .. } => base64.len(),
294            ImageUrl(s) => s.len(),
295            File(f) => f.len(),
296        }
297    }
298
299    pub fn text(&self) -> Option<&str> {
300        match self {
301            Content::Text(s) => Some(s.as_str()),
302            _ => None,
303        }
304    }
305
306    pub fn content(&self) -> &str {
307        use Content::*;
308        match self {
309            Text(s) => s.as_ref(),
310            Image { base64, .. } => base64.as_ref(),
311            ImageUrl(s) => s.as_ref(),
312            File(f) => f.base64.as_ref(),
313        }
314    }
315}
316
317#[derive(Debug, Copy, Clone)]
318pub enum Role {
319    System,
320    User,
321    Assistant,
322}
323
324impl Role {
325    pub fn as_str(&self) -> &'static str {
326        match self {
327            Role::System => "system",
328            Role::User => "user",
329            Role::Assistant => "assistant",
330        }
331    }
332}
333
334impl FromStr for Role {
335    type Err = &'static str;
336    fn from_str(s: &str) -> Result<Self, Self::Err> {
337        match s.to_lowercase().as_str() {
338            "system" => Ok(Role::System),
339            "user" => Ok(Role::User),
340            "assistant" => Ok(Role::Assistant),
341            _ => Err("Invalid role"),
342        }
343    }
344}
345
346#[derive(Clone, Default)]
347pub enum Response {
348    /// The first time we get anything at all on the SSE stream
349    Start,
350    /// Reasoning events - start, some thoughts, stop
351    Think(ThinkEvent),
352    /// The good stuff
353    Content(String),
354    /// Summary stats at the end of the run
355    Stats(super::stats::Stats),
356    /// Less good things. Often you mistyped the model name.
357    Error(String),
358    /// For default
359    #[default]
360    None,
361}
362
363#[derive(Debug, Clone)]
364pub enum ThinkEvent {
365    Start,
366    Content(String),
367    Stop,
368}
369
370#[derive(Debug, Clone)]
371pub enum PromptFileKind {
372    Image,
373    // Typically a PDF
374    File,
375    //Audio,
376}
377
378#[derive(Debug, Clone)]
379pub struct PromptFile {
380    kind: PromptFileKind,
381    pub filename: String,
382    pub base64: String,
383}
384
385impl PromptFile {
386    /// Load disk file, identify, and base64 encode it
387    pub fn load(filename: &str) -> Result<Self, &'static str> {
388        let kind = if IMAGE_EXT.iter().any(|ext| filename.ends_with(ext)) {
389            PromptFileKind::Image
390        } else {
391            PromptFileKind::File
392        };
393        let data = filename_read_to_bytes(filename)?;
394        Ok(PromptFile {
395            kind,
396            filename: filename.split('/').next_back().unwrap().to_string(),
397            base64: base64::encode(&data),
398        })
399    }
400
401    pub fn len(&self) -> usize {
402        self.base64.len()
403    }
404
405    pub(crate) fn from_parts(kind: PromptFileKind, filename: String, base64: String) -> Self {
406        Self {
407            kind,
408            filename,
409            base64,
410        }
411    }
412
413    pub fn into_content(self) -> Content {
414        match self.kind {
415            PromptFileKind::Image => Content::Image {
416                mime_type: self.mime_type(),
417                base64: self.base64,
418            },
419            PromptFileKind::File => Content::File(self),
420        }
421    }
422
423    pub fn mime_type(&self) -> &'static str {
424        for (ext, mime) in MIME_TYPES {
425            if self.filename.to_lowercase().ends_with(ext) {
426                return mime;
427            }
428        }
429        "application/octet-stream"
430    }
431}