use serde::{Deserialize, Serialize};
use std::time::Duration;
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum Role {
System,
User,
Assistant,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Message {
pub role: Role,
pub content: String,
}
impl Message {
pub fn system(content: impl Into<String>) -> Self {
Self {
role: Role::System,
content: content.into(),
}
}
pub fn user(content: impl Into<String>) -> Self {
Self {
role: Role::User,
content: content.into(),
}
}
pub fn assistant(content: impl Into<String>) -> Self {
Self {
role: Role::Assistant,
content: content.into(),
}
}
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct UsageInfo {
pub prompt_tokens: u32,
pub completion_tokens: u32,
pub total_tokens: u32,
}
impl UsageInfo {
pub fn accumulate(&mut self, other: &UsageInfo) {
self.prompt_tokens += other.prompt_tokens;
self.completion_tokens += other.completion_tokens;
self.total_tokens += other.total_tokens;
}
}
#[derive(Debug, Clone)]
pub struct CompletionResponse {
pub content: String,
pub model: String,
pub usage: UsageInfo,
}
#[derive(Debug, Clone)]
pub struct ReplResult {
pub stdout: String,
pub stderr: String,
pub final_answer: Option<FinalAnswer>,
pub execution_time: Duration,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum FinalAnswer {
Direct(String),
Var(String),
}
#[derive(Debug, Clone)]
pub struct CodeBlock {
pub code: String,
pub result: Option<ReplResult>,
}
#[derive(Debug, Clone)]
pub struct RlmCompletion {
pub response: String,
pub iterations: u32,
pub total_usage: UsageInfo,
}