tiny_agent_rs/core/
memory.rs

1use super::steps::AgentStep;
2use serde::{Deserialize, Serialize};
3use serde_json::Value;
4use tracing::info;
5
6/// Memory structure that replaces raw `Vec<Value>` messages
7/// Maintains the agent's reasoning steps and converts them to OpenAI format
8#[derive(Debug, Clone, Serialize, Deserialize)]
9pub struct AgentMemory {
10    steps: Vec<AgentStep>,
11    system_prompt: Option<String>,
12}
13
14impl AgentMemory {
15    /// Create a new memory with optional system prompt
16    pub fn new(system_prompt: Option<String>) -> Self {
17        Self {
18            steps: Vec::new(),
19            system_prompt,
20        }
21    }
22
23    /// Create memory with default system prompt
24    pub fn with_default_system() -> Self {
25        Self::new(Some(
26            "You are a helpful assistant with access to tools. Use tools when necessary to provide accurate information. Be concise and helpful. When you are ready to give the final response, you MUST call the `final_answer` tool with an `answer` string instead of replying directly.".to_string()
27        ))
28    }
29
30    /// Add a step to memory
31    pub fn add_step(&mut self, step: AgentStep) {
32        let description = step.describe();
33        info!(target: "tinyagent::steps", "{}", description);
34        self.steps.push(step);
35    }
36
37    /// Get all steps
38    pub fn steps(&self) -> &[AgentStep] {
39        &self.steps
40    }
41
42    /// Get the last step
43    pub fn last_step(&self) -> Option<&AgentStep> {
44        self.steps.last()
45    }
46
47    /// Convert memory to OpenAI message format
48    pub fn as_messages(&self) -> Vec<Value> {
49        let mut messages = Vec::new();
50
51        if let Some(system_prompt) = &self.system_prompt {
52            messages.push(serde_json::json!({
53                "role": "system",
54                "content": system_prompt
55            }));
56        }
57
58        for step in &self.steps {
59            messages.push(step.to_message());
60        }
61
62        messages
63    }
64
65    /// Clear all steps but keep system prompt
66    pub fn clear_steps(&mut self) {
67        self.steps.clear();
68    }
69
70    /// Get number of steps
71    pub fn step_count(&self) -> usize {
72        self.steps.len()
73    }
74
75    /// Check if memory is empty (excluding system prompt)
76    pub fn is_empty(&self) -> bool {
77        self.steps.is_empty()
78    }
79
80    /// Get all steps of a specific type
81    pub fn filter_steps<F>(&self, predicate: F) -> Vec<&AgentStep>
82    where
83        F: Fn(&AgentStep) -> bool,
84    {
85        self.steps.iter().filter(|step| predicate(step)).collect()
86    }
87
88    /// Count steps of a specific variant
89    pub fn count_actions(&self) -> usize {
90        self.steps
91            .iter()
92            .filter(|s| matches!(s, AgentStep::Action { .. }))
93            .count()
94    }
95
96    pub fn count_observations(&self) -> usize {
97        self.steps
98            .iter()
99            .filter(|s| matches!(s, AgentStep::Observation { .. }))
100            .count()
101    }
102
103    pub fn set_final_answer_structured(&mut self, structured: Value) {
104        if let Some(AgentStep::FinalAnswer {
105            structured: slot, ..
106        }) = self
107            .steps
108            .iter_mut()
109            .rev()
110            .find(|step| matches!(step, AgentStep::FinalAnswer { .. }))
111        {
112            *slot = Some(structured);
113        }
114    }
115}
116
117impl Default for AgentMemory {
118    fn default() -> Self {
119        Self::with_default_system()
120    }
121}
122
123impl From<Vec<Value>> for AgentMemory {
124    fn from(messages: Vec<Value>) -> Self {
125        let mut memory = AgentMemory::new(None);
126
127        for msg in messages {
128            if let Some(role) = msg.get("role").and_then(|r| r.as_str()) {
129                match role {
130                    "system" => {
131                        if let Some(content) = msg.get("content").and_then(|c| c.as_str()) {
132                            memory.system_prompt = Some(content.to_string());
133                        }
134                    }
135                    "user" => {
136                        if let Some(content) = msg.get("content").and_then(|c| c.as_str()) {
137                            memory.add_step(AgentStep::Task {
138                                content: content.to_string(),
139                            });
140                        }
141                    }
142                    "assistant" => {
143                        if let Some(tool_calls) = msg.get("tool_calls") {
144                            if let Some(calls_array) = tool_calls.as_array() {
145                                for call in calls_array {
146                                    if let (Some(id), Some(function)) = (
147                                        call.get("id").and_then(|i| i.as_str()),
148                                        call.get("function"),
149                                    ) {
150                                        let name = function
151                                            .get("name")
152                                            .and_then(|n| n.as_str())
153                                            .unwrap_or("unknown");
154                                        let args_str = function
155                                            .get("arguments")
156                                            .and_then(|a| a.as_str())
157                                            .unwrap_or("{}");
158                                        let args: Value =
159                                            serde_json::from_str(args_str).unwrap_or(Value::Null);
160
161                                        memory.add_step(AgentStep::Action {
162                                            tool_name: name.to_string(),
163                                            tool_call_id: id.to_string(),
164                                            arguments: args,
165                                        });
166                                    }
167                                }
168                            }
169                        } else if let Some(content) = msg.get("content").and_then(|c| c.as_str()) {
170                            if !content.is_empty() {
171                                memory.add_step(AgentStep::FinalAnswer {
172                                    answer: content.to_string(),
173                                    structured: None,
174                                });
175                            }
176                        }
177                    }
178                    "tool" => {
179                        if let (Some(id), Some(content)) = (
180                            msg.get("tool_call_id").and_then(|i| i.as_str()),
181                            msg.get("content").and_then(|c| c.as_str()),
182                        ) {
183                            let is_error = detect_tool_error(content);
184                            memory.add_step(AgentStep::Observation {
185                                tool_call_id: id.to_string(),
186                                result: content.to_string(),
187                                is_error,
188                            });
189                        }
190                    }
191                    _ => {}
192                }
193            }
194        }
195
196        memory
197    }
198}
199
200fn detect_tool_error(content: &str) -> bool {
201    match serde_json::from_str::<Value>(content) {
202        Ok(Value::Object(map)) => map.get("error").map(|err| !err.is_null()).unwrap_or(false),
203        _ => false,
204    }
205}
206
207#[cfg(test)]
208mod tests {
209    use super::*;
210
211    #[test]
212    fn test_memory_creation() {
213        let memory = AgentMemory::new(Some("System".to_string()));
214        assert_eq!(memory.step_count(), 0);
215        assert!(memory.is_empty());
216    }
217
218    #[test]
219    fn test_add_steps() {
220        let mut memory = AgentMemory::default();
221        memory.add_step(AgentStep::Task {
222            content: "Test task".to_string(),
223        });
224        assert_eq!(memory.step_count(), 1);
225        assert!(!memory.is_empty());
226    }
227
228    #[test]
229    fn test_as_messages() {
230        let mut memory = AgentMemory::with_default_system();
231        memory.add_step(AgentStep::Task {
232            content: "Hello".to_string(),
233        });
234
235        let messages = memory.as_messages();
236        assert_eq!(messages.len(), 2); // system + task
237        assert_eq!(messages[0]["role"], "system");
238        assert_eq!(messages[1]["role"], "user");
239    }
240
241    #[test]
242    fn test_count_actions() {
243        let mut memory = AgentMemory::default();
244        memory.add_step(AgentStep::Action {
245            tool_name: "test".to_string(),
246            tool_call_id: "1".to_string(),
247            arguments: Value::Null,
248        });
249        memory.add_step(AgentStep::Action {
250            tool_name: "test2".to_string(),
251            tool_call_id: "2".to_string(),
252            arguments: Value::Null,
253        });
254        assert_eq!(memory.count_actions(), 2);
255    }
256}
tiny_agent_rs/core/memory.rs

tiny_agent_rs/core/
memory.rs
