use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct AgentInput {
pub system_prompt: String,
pub user_message: String,
pub max_turns: u32,
}
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
pub enum StopReason {
FinalAnswer,
MaxTurnsReached,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentOutput {
pub final_answer: String,
pub stop_reason: StopReason,
pub turns_used: u32,
pub tool_calls: u32,
}
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum Role {
System,
User,
Assistant,
Tool,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct ToolCall {
pub id: String,
pub name: String,
pub args: serde_json::Value,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct ToolResult {
pub call_id: String,
pub output: serde_json::Value,
pub error: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct Message {
pub role: Role,
#[serde(default)]
pub content: String,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub tool_calls: Vec<ToolCall>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub tool_call_id: Option<String>,
}
impl Message {
pub fn system(content: impl Into<String>) -> Self {
Self {
role: Role::System,
content: content.into(),
tool_calls: vec![],
tool_call_id: None,
}
}
pub fn user(content: impl Into<String>) -> Self {
Self {
role: Role::User,
content: content.into(),
tool_calls: vec![],
tool_call_id: None,
}
}
pub fn assistant_text(content: impl Into<String>) -> Self {
Self {
role: Role::Assistant,
content: content.into(),
tool_calls: vec![],
tool_call_id: None,
}
}
pub fn assistant_with_tools(calls: Vec<ToolCall>) -> Self {
Self {
role: Role::Assistant,
content: String::new(),
tool_calls: calls,
tool_call_id: None,
}
}
pub fn tool_result(result: &ToolResult) -> Self {
let content = match &result.error {
Some(err) => format!("ERROR: {err}"),
None => result.output.to_string(),
};
Self {
role: Role::Tool,
content,
tool_calls: vec![],
tool_call_id: Some(result.call_id.clone()),
}
}
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct AgentState {
pub input: AgentInput,
pub history: Vec<Message>,
pub turn: u32,
pub tool_calls_executed: u32,
#[serde(default)]
pub pending_user_messages: Vec<String>,
}
impl AgentState {
pub fn new(input: AgentInput) -> Self {
let history = vec![
Message::system(&input.system_prompt),
Message::user(&input.user_message),
];
Self {
input,
history,
turn: 0,
tool_calls_executed: 0,
pending_user_messages: vec![],
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum LlmResponse {
Final { answer: String },
UseTools { calls: Vec<ToolCall> },
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LlmChatInput {
pub messages: Vec<Message>,
pub tools: Vec<ToolSchema>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolSchema {
pub name: String,
pub description: String,
pub args_schema: serde_json::Value,
}
pub fn compact(state: &AgentState, keep_recent: usize) -> AgentInput {
let mut summary_lines = Vec::new();
let total = state.history.len();
let drop_until = total.saturating_sub(keep_recent);
for msg in state.history.iter().take(drop_until) {
let line = match msg.role {
Role::System if summary_lines.is_empty() => continue,
Role::User => format!("user: {}", truncate(&msg.content, 200)),
Role::Assistant if !msg.tool_calls.is_empty() => {
let names: Vec<&str> = msg.tool_calls.iter().map(|c| c.name.as_str()).collect();
format!("assistant: called tools [{}]", names.join(", "))
}
Role::Assistant => format!("assistant: {}", truncate(&msg.content, 200)),
Role::Tool => format!("tool: {}", truncate(&msg.content, 120)),
Role::System => continue,
};
summary_lines.push(line);
}
let summary = if summary_lines.is_empty() {
String::new()
} else {
format!(
"\n\n[Prior conversation summary, {} messages dropped]\n{}",
drop_until,
summary_lines.join("\n")
)
};
let recent_user = state
.history
.iter()
.rev()
.find(|m| m.role == Role::User)
.map(|m| m.content.clone())
.unwrap_or_default();
AgentInput {
system_prompt: format!("{}{}", state.input.system_prompt, summary),
user_message: recent_user,
max_turns: state.input.max_turns,
}
}
fn truncate(s: &str, max: usize) -> String {
if s.len() <= max {
return s.to_string();
}
let mut boundary = max;
while boundary > 0 && !s.is_char_boundary(boundary) {
boundary -= 1;
}
format!("{}…", &s[..boundary])
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn agent_state_seeds_system_and_user() {
let s = AgentState::new(AgentInput {
system_prompt: "be helpful".into(),
user_message: "hi".into(),
max_turns: 5,
});
assert_eq!(s.history.len(), 2);
assert_eq!(s.history[0].role, Role::System);
assert_eq!(s.history[1].role, Role::User);
assert_eq!(s.turn, 0);
}
#[test]
fn compact_keeps_system_and_recent() {
let mut state = AgentState::new(AgentInput {
system_prompt: "sys".into(),
user_message: "u0".into(),
max_turns: 50,
});
for i in 1..30 {
state.history.push(Message::user(format!("u{i}")));
state.history.push(Message::assistant_text(format!("a{i}")));
}
let compacted = compact(&state, 10);
assert!(compacted.system_prompt.starts_with("sys"));
assert!(
compacted
.system_prompt
.contains("Prior conversation summary")
);
assert_eq!(compacted.max_turns, state.input.max_turns);
}
#[test]
fn truncate_respects_utf8_char_boundary() {
let t = truncate("héllo world", 2);
assert_eq!(t, "h…");
assert_eq!(truncate("hi", 10), "hi");
assert_eq!(truncate("🦀rust", 2), "…");
}
#[test]
fn message_roundtrips_through_json() {
let m = Message::assistant_with_tools(vec![ToolCall {
id: "c1".into(),
name: "add".into(),
args: serde_json::json!({"a": 1, "b": 2}),
}]);
let s = serde_json::to_string(&m).unwrap();
let back: Message = serde_json::from_str(&s).unwrap();
assert_eq!(m, back);
}
}