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;
pub const DEFAULT_MODEL: &str = "google/gemma-3n-e4b-it:free";
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, }
pub struct LastData {
pub opts: PromptOpts,
pub messages: Vec<Message>,
}
#[derive(Clone)]
pub struct PromptOpts {
pub prompt: Option<String>,
pub models: Vec<String>,
pub provider: Option<String>,
pub system: Option<String>,
pub priority: Option<Priority>,
pub reasoning: Option<ReasoningConfig>,
pub show_reasoning: Option<bool>,
pub quiet: Option<bool>,
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 {
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() {
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",
)), }
}
}
#[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, Low,
#[default]
Medium,
High,
XHigh, }
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)
}
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 {
Start,
Think(ThinkEvent),
Content(String),
Stats(super::stats::Stats),
Error(String),
#[default]
None,
}
#[derive(Debug, Clone)]
pub enum ThinkEvent {
Start,
Content(String),
Stop,
}