Skip to main content

cascade_agent/agent/
state.rs

1//! Conversation state management for the agent loop.
2
3use llm_cascade::{Conversation, Message};
4use serde::Serialize;
5
6use crate::error::Result;
7use crate::memory::state::ConversationStateAccess;
8
9/// Wraps llm-cascade Conversation with agent-specific metadata.
10///
11/// Tracks turn count, task ID, and timestamps for persistence and observability.
12/// Implements [`ConversationStateAccess`] so the memory subsystem can operate on it.
13#[derive(Debug, Serialize)]
14pub struct ConversationState {
15    /// The ordered list of conversation messages (user, assistant, tool results).
16    pub messages: Vec<Message>,
17    /// System prompt prepended to every LLM call.
18    pub system_prompt: String,
19    /// Number of user turns processed so far.
20    pub turn_count: usize,
21    /// Unique identifier for the current task.
22    pub task_id: String,
23    /// Timestamp when this state was created.
24    pub created_at: chrono::DateTime<chrono::Utc>,
25}
26
27impl ConversationState {
28    /// Create a new empty conversation state.
29    pub fn new(system_prompt: String, task_id: String) -> Self {
30        Self {
31            messages: Vec::new(),
32            system_prompt,
33            turn_count: 0,
34            task_id,
35            created_at: chrono::Utc::now(),
36        }
37    }
38
39    /// Append a user message and increment the turn counter.
40    pub fn add_user_message(&mut self, content: String) {
41        self.messages.push(Message::user(content));
42        self.turn_count += 1;
43    }
44
45    /// Append an assistant text response.
46    pub fn add_assistant_text(&mut self, content: String) {
47        self.messages.push(Message::assistant(content));
48    }
49
50    /// Append a tool result message associated with a tool call ID.
51    pub fn add_tool_result(&mut self, tool_call_id: &str, result: &str) {
52        self.messages
53            .push(Message::tool(result.to_string(), tool_call_id.to_string()));
54    }
55
56    /// Build an `llm_cascade::Conversation` from current state.
57    ///
58    /// Prepends the system prompt as the first message and attaches no tools
59    /// (tools are added via `with_tools` by the caller when needed).
60    pub fn to_conversation(&self) -> Conversation {
61        let mut all_msgs = vec![Message::system(&self.system_prompt)];
62        all_msgs.extend(self.messages.clone());
63        Conversation::new(all_msgs)
64    }
65
66    /// Get the last assistant text message (excluding tool-call stubs).
67    pub fn last_assistant_text(&self) -> Option<String> {
68        self.messages
69            .iter()
70            .rev()
71            .find(|m| matches!(m.role, llm_cascade::MessageRole::Assistant))
72            .map(|m| m.content.clone())
73    }
74
75    /// Serialize state to a pretty-printed JSON string.
76    pub fn to_json(&self) -> Result<String> {
77        serde_json::to_string_pretty(self).map_err(Into::into)
78    }
79
80    /// Save state to a JSON file in the given directory.
81    pub fn to_json_file(&self, dir: &str) -> Result<std::path::PathBuf> {
82        std::fs::create_dir_all(dir)?;
83        let filename = format!(
84            "state_{}_{}.json",
85            self.created_at.format("%Y%m%d_%H%M%S"),
86            &self.task_id[..8.min(self.task_id.len())]
87        );
88        let path = std::path::PathBuf::from(dir).join(filename);
89        std::fs::write(&path, self.to_json()?)?;
90        Ok(path)
91    }
92}
93
94impl ConversationStateAccess for ConversationState {
95    fn messages(&self) -> &[Message] {
96        &self.messages
97    }
98
99    fn messages_mut(&mut self) -> &mut Vec<Message> {
100        &mut self.messages
101    }
102
103    fn system_prompt(&self) -> &str {
104        &self.system_prompt
105    }
106}