use crate::{ModelOutput, Signal};
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
#[derive(Debug, Clone, Serialize, Deserialize)]
#[non_exhaustive]
pub enum Block {
Text(String),
FileRef {
path: String,
hash: Option<String>,
excerpt: Option<String>,
},
Skill { name: String, body: String },
ToolCall {
call_id: String,
name: String,
args: serde_json::Value,
},
ToolResult {
call_id: String,
content: serde_json::Value,
},
Feedback(Vec<Signal>),
Reasoning(String),
Image { media_type: String, base64: String },
Audio { media_type: String, base64: String },
}
impl Block {
pub fn image_bytes(media_type: impl Into<String>, bytes: &[u8]) -> Self {
Block::Image {
media_type: media_type.into(),
base64: base64_encode(bytes),
}
}
pub fn audio_bytes(media_type: impl Into<String>, bytes: &[u8]) -> Self {
Block::Audio {
media_type: media_type.into(),
base64: base64_encode(bytes),
}
}
}
fn base64_encode(input: &[u8]) -> String {
const T: &[u8; 64] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
let mut out = String::with_capacity(input.len().div_ceil(3) * 4);
for chunk in input.chunks(3) {
let b0 = chunk[0] as u32;
let b1 = *chunk.get(1).unwrap_or(&0) as u32;
let b2 = *chunk.get(2).unwrap_or(&0) as u32;
let n = (b0 << 16) | (b1 << 8) | b2;
out.push(T[(n >> 18) as usize & 63] as char);
out.push(T[(n >> 12) as usize & 63] as char);
out.push(if chunk.len() > 1 {
T[(n >> 6) as usize & 63] as char
} else {
'='
});
out.push(if chunk.len() > 2 {
T[n as usize & 63] as char
} else {
'='
});
}
out
}
#[cfg(test)]
mod base64_tests {
use super::base64_encode;
#[test]
fn matches_known_vectors() {
assert_eq!(base64_encode(b""), "");
assert_eq!(base64_encode(b"f"), "Zg==");
assert_eq!(base64_encode(b"fo"), "Zm8=");
assert_eq!(base64_encode(b"foo"), "Zm9v");
assert_eq!(base64_encode(b"foob"), "Zm9vYg==");
assert_eq!(base64_encode(b"fooba"), "Zm9vYmE=");
assert_eq!(base64_encode(b"foobar"), "Zm9vYmFy");
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Turn {
pub role: TurnRole,
pub blocks: Vec<Block>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
#[non_exhaustive]
pub enum TurnRole {
User,
Assistant,
System,
Tool,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Task {
pub description: String,
pub source: Option<String>, pub deadline: Option<i64>,
}
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub struct Policy {
pub max_iters: u32,
pub max_input_tokens: u32,
pub max_output_tokens: u32,
pub self_correct_rounds: u32,
}
impl Default for Policy {
fn default() -> Self {
Self {
max_iters: 50,
max_input_tokens: 150_000,
max_output_tokens: 8_000,
self_correct_rounds: 3,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
#[serde(tag = "type", rename_all = "snake_case")]
#[non_exhaustive]
pub enum ResponseFormat {
#[default]
Free,
JsonObject,
JsonSchema {
name: String,
schema: serde_json::Value,
},
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Context {
pub system: Vec<Block>,
pub guides: Vec<Block>,
pub history: Vec<Turn>,
pub task: Task,
pub policy: Policy,
pub metadata: BTreeMap<String, serde_json::Value>,
pub tools: Vec<crate::ToolSchema>,
#[serde(default, skip_serializing_if = "response_format_is_default")]
pub response_format: ResponseFormat,
}
fn response_format_is_default(f: &ResponseFormat) -> bool {
matches!(f, ResponseFormat::Free)
}
impl Context {
pub fn new(task: Task) -> Self {
Self {
system: Vec::new(),
guides: Vec::new(),
history: Vec::new(),
task,
policy: Policy::default(),
metadata: BTreeMap::new(),
tools: Vec::new(),
response_format: ResponseFormat::Free,
}
}
pub fn push_model_output(&mut self, out: &ModelOutput) {
let mut blocks = Vec::new();
if let Some(r) = &out.reasoning
&& !r.is_empty()
{
blocks.push(Block::Reasoning(r.clone()));
}
if let Some(t) = &out.text
&& !t.is_empty()
{
blocks.push(Block::Text(t.clone()));
}
for c in &out.tool_calls {
blocks.push(Block::ToolCall {
call_id: c.id.clone(),
name: c.name.clone(),
args: c.args.clone(),
});
}
self.history.push(Turn {
role: TurnRole::Assistant,
blocks,
});
}
pub fn push_feedback(&mut self, signals: Vec<Signal>) {
self.history.push(Turn {
role: TurnRole::Tool,
blocks: vec![Block::Feedback(signals)],
});
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Action {
pub tool: String,
pub call_id: String,
pub args: serde_json::Value,
}