ort_openrouter_cli/common/

data.rs

//! ort: Open Router CLI
//! https://github.com/grahamking/ort
//!
//! MIT License
//! Copyright (c) 2025 Graham King

use core::str::FromStr;

extern crate alloc;
use alloc::string::{String, ToString};
use alloc::vec;
use alloc::vec::Vec;

use crate::{ErrorKind, OrtError, ort_error};

const DEFAULT_SHOW_REASONING: bool = false;
const DEFAULT_QUIET: bool = false;

// Keep in sync with src/lib.rs
pub const DEFAULT_MODEL: &str = "google/gemma-3n-e4b-it:free";

// {
//  "id":"gen-1756743299-7ytIBcjALWQQShwMQfw9",
//  "provider":"Meta",
//  "model":"meta-llama/llama-3.3-8b-instruct:free",
//  "object":"chat.completion.chunk",
//  "created":1756743300,
//  "choices":[
//      {
//      "index":0,
//      "delta":{"role":"assistant","content":""},
//      "finish_reason":null,
//      "native_finish_reason":null,
//      "logprobs":null
//      }
//  ],
//  "usage":{
//      "prompt_tokens":42,
//      "completion_tokens":2,
//      "total_tokens":44,
//      "cost":0,"
//      is_byok":false,
//      "prompt_tokens_details":{"cached_tokens":0,"audio_tokens":0},
//      "cost_details":{"upstream_inference_cost":null,"upstream_inference_prompt_cost":0,"upstream_inference_completions_cost":0},
//      "completion_tokens_details":{"reasoning_tokens":0,"image_tokens":0}}
//  }

pub struct ChatCompletionsResponse {
    pub provider: Option<String>,
    pub model: Option<String>,
    pub choices: Vec<Choice>,
    pub usage: Option<Usage>,
}

pub struct Choice {
    pub delta: Message,
}

pub struct Usage {
    pub cost: f32, // In dollars, usually a very small fraction
}

pub struct LastData {
    pub opts: PromptOpts,
    pub messages: Vec<Message>,
}

#[derive(Clone)]
pub struct PromptOpts {
    pub prompt: Option<String>,
    /// Model IDs, e.g. 'moonshotai/kimi-k2'
    pub models: Vec<String>,
    /// Prefered provider slug
    pub provider: Option<String>,
    /// System prompt
    pub system: Option<String>,
    /// How to choose a provider
    pub priority: Option<Priority>,
    /// Reasoning config
    pub reasoning: Option<ReasoningConfig>,
    /// Show reasoning output
    pub show_reasoning: Option<bool>,
    /// Don't show stats after request
    pub quiet: Option<bool>,
    /// Whether to merge in the default settings from config file
    pub merge_config: bool,
}

impl Default for PromptOpts {
    fn default() -> Self {
        Self {
            prompt: None,
            models: vec![DEFAULT_MODEL.to_string()],
            provider: None,
            system: None,
            priority: None,
            reasoning: Some(ReasoningConfig::default()),
            show_reasoning: Some(false),
            quiet: Some(false),
            merge_config: true,
        }
    }
}

impl PromptOpts {
    // Replace any blank or None fields on Self with values from other
    // or with the defaults.
    // After this call a PromptOpts is ready to use.
    pub fn merge(&mut self, o: PromptOpts) {
        self.prompt.get_or_insert(o.prompt.unwrap_or_default());
        self.quiet.get_or_insert(o.quiet.unwrap_or(DEFAULT_QUIET));
        if self.models.is_empty() {
            // We don't merge the models, otherwise we'd try to query both the
            // cmd line one, and the config file default.
            self.models = o.models;
        }
        if let Some(provider) = o.provider {
            self.provider.get_or_insert(provider);
        }
        if let Some(system) = o.system {
            self.system.get_or_insert(system);
        }
        if let Some(priority) = o.priority {
            self.priority.get_or_insert(priority);
        }
        self.reasoning
            .get_or_insert(o.reasoning.unwrap_or_default());
        self.show_reasoning
            .get_or_insert(o.show_reasoning.unwrap_or(DEFAULT_SHOW_REASONING));
    }
}

#[derive(Default, Debug, Clone, Copy)]
pub enum Priority {
    Price,
    #[default]
    Latency,
    Throughput,
}

impl Priority {
    pub fn as_str(&self) -> &'static str {
        match self {
            Priority::Price => "price",
            Priority::Latency => "latency",
            Priority::Throughput => "throughput",
        }
    }
}

impl FromStr for Priority {
    type Err = OrtError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s.to_lowercase().as_str() {
            "price" => Ok(Priority::Price),
            "latency" => Ok(Priority::Latency),
            "throughput" => Ok(Priority::Throughput),
            _ => Err(ort_error(
                ErrorKind::FormatError,
                "Priority: Invalid string value",
            )), // Handle unknown strings
        }
    }
}

#[derive(Default, Debug, Clone)]
pub struct ReasoningConfig {
    pub enabled: bool,
    pub effort: Option<ReasoningEffort>,
    pub tokens: Option<u32>,
}

impl ReasoningConfig {
    pub fn off() -> Self {
        Self {
            enabled: false,
            ..Default::default()
        }
    }
}

#[derive(Default, Debug, Clone, Copy, PartialEq)]
pub enum ReasoningEffort {
    None, // GPT 5.x only
    Low,
    #[default]
    Medium,
    High,
    XHigh, // GPT 5.x only
}

impl ReasoningEffort {
    pub fn as_str(&self) -> &'static str {
        match self {
            ReasoningEffort::None => "none",
            ReasoningEffort::Low => "low",
            ReasoningEffort::Medium => "medium",
            ReasoningEffort::High => "high",
            ReasoningEffort::XHigh => "xhigh",
        }
    }
}

#[derive(Debug, Clone)]
pub struct Message {
    pub role: Role,
    pub content: Option<String>,
    pub reasoning: Option<String>,
}

impl Message {
    pub fn new(role: Role, content: Option<String>, reasoning: Option<String>) -> Self {
        Message {
            role,
            content,
            reasoning,
        }
    }
    pub fn system(content: String) -> Self {
        Self::new(Role::System, Some(content), None)
    }
    pub fn user(content: String) -> Self {
        Self::new(Role::User, Some(content), None)
    }
    pub fn assistant(content: String) -> Self {
        Self::new(Role::Assistant, Some(content), None)
    }

    /// Estimate size in bytes
    pub fn size(&self) -> u32 {
        self.content
            .as_ref()
            .or(self.reasoning.as_ref())
            .map(|c| c.len())
            .unwrap_or(0) as u32
            + 10
    }
}

#[derive(Debug, Copy, Clone)]
pub enum Role {
    System,
    User,
    Assistant,
}

impl Role {
    pub fn as_str(&self) -> &'static str {
        match self {
            Role::System => "system",
            Role::User => "user",
            Role::Assistant => "assistant",
        }
    }
}

impl FromStr for Role {
    type Err = &'static str;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s.to_lowercase().as_str() {
            "system" => Ok(Role::System),
            "user" => Ok(Role::User),
            "assistant" => Ok(Role::Assistant),
            _ => Err("Invalid role"),
        }
    }
}

#[derive(Clone, Default)]
pub enum Response {
    /// The first time we get anything at all on the SSE stream
    Start,
    /// Reasoning events - start, some thoughts, stop
    Think(ThinkEvent),
    /// The good stuff
    Content(String),
    /// Summary stats at the end of the run
    Stats(super::stats::Stats),
    /// Less good things. Often you mistyped the model name.
    Error(String),
    /// For default
    #[default]
    None,
}

#[derive(Debug, Clone)]
pub enum ThinkEvent {
    Start,
    Content(String),
    Stop,
}