smolagents_rs/
agents.rs

1
2//! This module contains the agents that can be used to solve tasks.
3//!
4//! Currently, there are two agents:
5//! - The function calling agent. This agent is used for models that have tool calling capabilities.
6//! - The code agent. This agent takes tools and can write simple python code that is executed to solve the task. 
7//! To use this agent you need to enable the `code-agent` feature.
8//!
9//! You can also implement your own agents by implementing the `Agent` trait.
10//!
11//! Planning agent is not implemented yet and will be added in the future.
12//!
13use crate::errors::AgentError;
14use crate::models::model_traits::Model;
15use crate::models::openai::ToolCall;
16use crate::models::types::Message;
17use crate::models::types::MessageRole;
18use crate::prompts::{
19    user_prompt_plan, TOOL_CALLING_SYSTEM_PROMPT, SYSTEM_PROMPT_FACTS, SYSTEM_PROMPT_PLAN,
20};
21use crate::tools::{AnyTool, FinalAnswerTool, ToolGroup, ToolInfo};
22use std::collections::HashMap;
23
24use crate::logger::LOGGER;
25use anyhow::Result;
26use colored::Colorize;
27use log::info;
28
29use serde_json::json;
30#[cfg(feature = "code-agent")]
31use {
32    crate::errors::InterpreterError, crate::local_python_interpreter::LocalPythonInterpreter,
33    crate::models::openai::FunctionCall, crate::prompts::CODE_SYSTEM_PROMPT, regex::Regex,
34};
35
36
37
38const DEFAULT_TOOL_DESCRIPTION_TEMPLATE: &str = r#"
39{{ tool.name }}: {{ tool.description }}
40    Takes inputs: {{tool.inputs}}
41"#;
42
43use std::fmt::Debug;
44
45pub fn get_tool_description_with_args(tool: &ToolInfo) -> String {
46    let mut description = DEFAULT_TOOL_DESCRIPTION_TEMPLATE.to_string();
47    description = description.replace("{{ tool.name }}", tool.function.name);
48    description = description.replace("{{ tool.description }}", tool.function.description);
49    description = description.replace(
50        "{{tool.inputs}}",
51        json!(&tool.function.parameters.schema)["properties"]
52            .to_string()
53            .as_str(),
54    );
55
56    description
57}
58
59pub fn get_tool_descriptions(tools: &[ToolInfo]) -> Vec<String> {
60    tools.iter().map(get_tool_description_with_args).collect()
61}
62pub fn format_prompt_with_tools(tools: Vec<ToolInfo>, prompt_template: &str) -> String {
63    let tool_descriptions = get_tool_descriptions(&tools);
64    let mut prompt = prompt_template.to_string();
65    prompt = prompt.replace("{{tool_descriptions}}", &tool_descriptions.join("\n"));
66    if prompt.contains("{{tool_names}}") {
67        let tool_names: Vec<String> = tools
68            .iter()
69            .map(|tool| tool.function.name.to_string())
70            .collect();
71        prompt = prompt.replace("{{tool_names}}", &tool_names.join(", "));
72    }
73    prompt
74}
75
76pub fn show_agents_description(managed_agents: &HashMap<String, Box<dyn Agent>>) -> String {
77    let mut managed_agent_description = r#"You can also give requests to team members.
78Calling a team member works the same as for calling a tool: simply, the only argument you can give in the call is 'request', a long string explaining your request.
79Given that this team member is a real human, you should be very verbose in your request.
80Here is a list of the team members that you can call:"#.to_string();
81
82    for (name, agent) in managed_agents.iter() {
83        managed_agent_description.push_str(&format!("{}: {:?}\n", name, agent.description()));
84    }
85
86    managed_agent_description
87}
88
89pub fn format_prompt_with_managed_agent_description(
90    prompt_template: String,
91    managed_agents: &HashMap<String, Box<dyn Agent>>,
92    agent_descriptions_placeholder: Option<&str>,
93) -> Result<String> {
94    let agent_descriptions_placeholder =
95        agent_descriptions_placeholder.unwrap_or("{{managed_agents_descriptions}}");
96
97    if managed_agents.keys().len() > 0 {
98        Ok(prompt_template.replace(
99            agent_descriptions_placeholder,
100            &show_agents_description(managed_agents),
101        ))
102    } else {
103        Ok(prompt_template.replace(agent_descriptions_placeholder, ""))
104    }
105}
106
107pub trait Agent {
108    fn name(&self) -> &'static str;
109    fn get_max_steps(&self) -> usize;
110    fn get_step_number(&self) -> usize;
111    fn increment_step_number(&mut self);
112    fn get_logs_mut(&mut self) -> &mut Vec<Step>;
113    fn set_task(&mut self, task: &str);
114    fn get_system_prompt(&self) -> &str;
115    fn description(&self) -> String {
116        "".to_string()
117    }
118    fn model(&self) -> &dyn Model;
119    fn step(&mut self, log_entry: &mut Step) -> Result<Option<String>>;
120    fn direct_run(&mut self, _task: &str) -> Result<String> {
121        let mut final_answer: Option<String> = None;
122        while final_answer.is_none() && self.get_step_number() < self.get_max_steps() {
123            let mut step_log = Step::ActionStep(AgentStep {
124                agent_memory: None,
125                llm_output: None,
126                tool_call: None,
127                error: None,
128                observations: None,
129                _step: self.get_step_number(),
130            });
131
132            final_answer = self.step(&mut step_log)?;
133            self.get_logs_mut().push(step_log);
134            self.increment_step_number();
135        }
136
137        if final_answer.is_none() && self.get_step_number() >= self.get_max_steps() {
138            final_answer = self.provide_final_answer(_task)?;
139        }
140        info!(
141            "Final answer: {}",
142            final_answer
143                .clone()
144                .unwrap_or("Could not find answer".to_string())
145        );
146        Ok(final_answer.unwrap_or_else(|| "Max steps reached without final answer".to_string()))
147    }
148    fn stream_run(&mut self, _task: &str) -> Result<String> {
149        todo!()
150    }
151    fn run(&mut self, task: &str, stream: bool, reset: bool) -> Result<String> {
152        // self.task = task.to_string();
153        self.set_task(task);
154
155        let system_prompt_step = Step::SystemPromptStep(self.get_system_prompt().to_string());
156        if reset {
157            self.get_logs_mut().clear();
158            self.get_logs_mut().push(system_prompt_step);
159        } else if self.get_logs_mut().is_empty() {
160            self.get_logs_mut().push(system_prompt_step);
161        } else {
162            self.get_logs_mut()[0] = system_prompt_step;
163        }
164        self.get_logs_mut().push(Step::TaskStep(task.to_string()));
165        match stream {
166            true => self.stream_run(task),
167            false => self.direct_run(task),
168        }
169    }
170    fn provide_final_answer(&mut self, task: &str) -> Result<Option<String>> {
171        let mut input_messages = vec![Message {
172            role: MessageRole::System,
173            content: "An agent tried to answer a user query but it got stuck and failed to do so. You are tasked with providing an answer instead. Here is the agent's memory:".to_string(),
174        }];
175
176        input_messages.extend(self.write_inner_memory_from_logs(Some(true))?[1..].to_vec());
177        input_messages.push(Message {
178            role: MessageRole::User,
179            content: format!("Based on the above, please provide an answer to the following user request: \n```\n{}", task),
180        });
181        let response = self
182            .model()
183            .run(input_messages, vec![], None, None)?
184            .get_response()?;
185        Ok(Some(response))
186    }
187
188    fn write_inner_memory_from_logs(&mut self, summary_mode: Option<bool>) -> Result<Vec<Message>> {
189        let mut memory = Vec::new();
190        let summary_mode = summary_mode.unwrap_or(false);
191        for log in self.get_logs_mut() {
192            match log {
193                Step::ToolCall(_) => {}
194                Step::PlanningStep(plan, facts) => {
195                    memory.push(Message {
196                        role: MessageRole::Assistant,
197                        content: "[PLAN]:\n".to_owned() + plan.as_str(),
198                    });
199
200                    if !summary_mode {
201                        memory.push(Message {
202                            role: MessageRole::Assistant,
203                            content: "[FACTS]:\n".to_owned() + facts.as_str(),
204                        });
205                    }
206                }
207                Step::TaskStep(task) => {
208                    memory.push(Message {
209                        role: MessageRole::User,
210                        content: "New Task: ".to_owned() + task.as_str(),
211                    });
212                }
213                Step::SystemPromptStep(prompt) => {
214                    memory.push(Message {
215                        role: MessageRole::System,
216                        content: prompt.to_string(),
217                    });
218                }
219                Step::ActionStep(step_log) => {
220                    if step_log.llm_output.is_some() && !summary_mode {
221                        memory.push(Message {
222                            role: MessageRole::Assistant,
223                            content: step_log.llm_output.clone().unwrap_or_default(),
224                        });
225                    }
226                    if step_log.tool_call.is_some() {
227                        let tool_call_message = Message {
228                            role: MessageRole::Assistant,
229                            content: serde_json::to_string(&step_log.tool_call.clone().unwrap())?,
230                        };
231                        memory.push(tool_call_message);
232                    }
233                    if step_log.tool_call.is_none() && step_log.error.is_some() {
234                        let message_content = "Error: ".to_owned() + step_log.error.clone().unwrap().message()+"\nNow let's retry: take care not to repeat previous errors! If you have retried several times, try a completely different approach.\n";
235                        memory.push(Message {
236                            role: MessageRole::Assistant,
237                            content: message_content,
238                        });
239                    }
240                    if step_log.tool_call.is_some()
241                        && (step_log.error.is_some() || step_log.observations.is_some())
242                    {
243                        let mut message_content = "".to_string();
244                        if step_log.error.is_some() {
245                            message_content = "Error: ".to_owned() + step_log.error.as_ref().unwrap().message()+"\nNow let's retry: take care not to repeat previous errors! If you have retried several times, try a completely different approach.\n";
246                        } else if step_log.observations.is_some() {
247                            message_content = "Observations: ".to_owned()
248                                + step_log.observations.as_ref().unwrap().as_str();
249                        }
250                        let tool_response_message = {
251                            Message {
252                                role: MessageRole::User,
253                                content: format!(
254                                    "Call id: {}\n{}",
255                                    step_log
256                                        .tool_call
257                                        .as_ref()
258                                        .unwrap()
259                                        .id
260                                        .clone()
261                                        .unwrap_or_default(),
262                                    message_content
263                                ),
264                            }
265                        };
266                        memory.push(tool_response_message);
267                    }
268
269                    if step_log.observations.is_some() || step_log.error.is_some() {
270                        let mut message_content = "".to_string();
271                        if step_log.error.is_some() {
272                            message_content = "Error: ".to_owned() + step_log.error.as_ref().unwrap().message()+"\nNow let's retry: take care not to repeat previous errors! If you have retried several times, try a completely different approach.\n";
273                        } else if step_log.observations.is_some() {
274                            message_content = "Observations: ".to_owned()
275                                + step_log.observations.as_ref().unwrap().as_str();
276                        }
277                        memory.push(Message {
278                            role: MessageRole::Assistant,
279                            content: message_content,
280                        });
281                    }
282                }
283            }
284        }
285        Ok(memory)
286    }
287}
288
289#[derive(Debug)]
290pub enum Step {
291    PlanningStep(String, String),
292    TaskStep(String),
293    SystemPromptStep(String),
294    ActionStep(AgentStep),
295    ToolCall(ToolCall),
296}
297
298#[derive(Debug, Clone)]
299pub struct AgentStep {
300    agent_memory: Option<Vec<Message>>,
301    llm_output: Option<String>,
302    tool_call: Option<ToolCall>,
303    error: Option<AgentError>,
304    observations: Option<String>,
305    _step: usize,
306}
307
308// Define a trait for the parent functionality
309
310pub struct MultiStepAgent<M: Model> {
311    pub model: M,
312    pub tools: Vec<Box<dyn AnyTool>>,
313    pub system_prompt_template: String,
314    pub name: &'static str,
315    pub managed_agents: Option<HashMap<String, Box<dyn Agent>>>,
316    pub description: String,
317    pub max_steps: usize,
318    pub step_number: usize,
319    pub task: String,
320    pub input_messages: Option<Vec<Message>>,
321    pub logs: Vec<Step>,
322}
323
324impl<M: Model + Debug> Agent for MultiStepAgent<M> {
325    fn name(&self) -> &'static str {
326        self.name
327    }
328    fn get_max_steps(&self) -> usize {
329        self.max_steps
330    }
331    fn get_step_number(&self) -> usize {
332        self.step_number
333    }
334    fn set_task(&mut self, task: &str) {
335        self.task = task.to_string();
336    }
337    fn get_system_prompt(&self) -> &str {
338        &self.system_prompt_template
339    }
340    fn increment_step_number(&mut self) {
341        self.step_number += 1;
342    }
343    fn get_logs_mut(&mut self) -> &mut Vec<Step> {
344        &mut self.logs
345    }
346    fn description(&self) -> String {
347        self.description.clone()
348    }
349    fn model(&self) -> &dyn Model {
350        &self.model
351    }
352
353    /// Perform one step in the ReAct framework: the agent thinks, acts, and observes the result.
354    ///
355    /// Returns None if the step is not final.
356    fn step(&mut self, _: &mut Step) -> Result<Option<String>> {
357        todo!()
358    }
359}
360
361impl<M: Model> MultiStepAgent<M> {
362    pub fn new(
363        model: M,
364        mut tools: Vec<Box<dyn AnyTool>>,
365        system_prompt: Option<&str>,
366        managed_agents: Option<HashMap<String, Box<dyn Agent>>>,
367        description: Option<&str>,
368        max_steps: Option<usize>,
369    ) -> Result<Self> {
370        // Initialize logger
371        log::set_logger(&LOGGER).unwrap();
372        log::set_max_level(log::LevelFilter::Info);
373
374        let name = "MultiStepAgent";
375
376        let system_prompt_template = match system_prompt {
377            Some(prompt) => prompt.to_string(),
378            None => TOOL_CALLING_SYSTEM_PROMPT.to_string(),
379        };
380        let description = match description {
381            Some(desc) => desc.to_string(),
382            None => "A multi-step agent that can solve tasks using a series of tools".to_string(),
383        };
384
385        let final_answer_tool = FinalAnswerTool::new();
386        tools.push(Box::new(final_answer_tool));
387
388        let mut agent = MultiStepAgent {
389            model,
390            tools,
391            system_prompt_template,
392            name,
393            managed_agents,
394            description,
395            max_steps: max_steps.unwrap_or(10),
396            step_number: 0,
397            task: "".to_string(),
398            logs: Vec::new(),
399            input_messages: None,
400        };
401
402        agent.initialize_system_prompt()?;
403        Ok(agent)
404    }
405
406    fn initialize_system_prompt(&mut self) -> Result<String> {
407        let tools = self.tools.tool_info();
408        self.system_prompt_template = format_prompt_with_tools(tools, &self.system_prompt_template);
409        match &self.managed_agents {
410            Some(managed_agents) => {
411                self.system_prompt_template = format_prompt_with_managed_agent_description(
412                    self.system_prompt_template.clone(),
413                    managed_agents,
414                    None,
415                )?;
416            }
417            None => {
418                self.system_prompt_template = format_prompt_with_managed_agent_description(
419                    self.system_prompt_template.clone(),
420                    &HashMap::new(),
421                    None,
422                )?;
423            }
424        }
425        self.system_prompt_template = self
426            .system_prompt_template
427            .replace("{{current_time}}", &chrono::Local::now().to_string());
428        Ok(self.system_prompt_template.clone())
429    }
430
431    pub fn planning_step(&mut self, task: &str, is_first_step: bool, _step: usize) {
432        if is_first_step {
433            let message_prompt_facts = Message {
434                role: MessageRole::System,
435                content: SYSTEM_PROMPT_FACTS.to_string(),
436            };
437            let message_prompt_task = Message {
438                role: MessageRole::User,
439                content: format!(
440                    "Here is the task: ```
441                    {}
442                    ```
443                    Now Begin!
444                    ",
445                    task
446                ),
447            };
448
449            let answer_facts = self
450                .model
451                .run(
452                    vec![message_prompt_facts, message_prompt_task],
453                    vec![],
454                    None,
455                    None,
456                )
457                .unwrap()
458                .get_response()
459                .unwrap_or("".to_string());
460            let message_system_prompt_plan = Message {
461                role: MessageRole::System,
462                content: SYSTEM_PROMPT_PLAN.to_string(),
463            };
464            let tool_descriptions = serde_json::to_string(
465                &self
466                    .tools
467                    .iter()
468                    .map(|tool| tool.tool_info())
469                    .collect::<Vec<_>>(),
470            )
471            .unwrap();
472            let message_user_prompt_plan = Message {
473                role: MessageRole::User,
474                content: user_prompt_plan(
475                    task,
476                    &tool_descriptions,
477                    &show_agents_description(
478                        self.managed_agents.as_ref().unwrap_or(&HashMap::new()),
479                    ),
480                    &answer_facts,
481                ),
482            };
483            let answer_plan = self
484                .model
485                .run(
486                    vec![message_system_prompt_plan, message_user_prompt_plan],
487                    vec![],
488                    None,
489                    Some(HashMap::from([(
490                        "stop".to_string(),
491                        vec!["Observation:".to_string()],
492                    )])),
493                )
494                .unwrap()
495                .get_response()
496                .unwrap();
497            let final_plan_redaction = format!(
498                "Here is the plan of action that I will follow for the task: \n{}",
499                answer_plan
500            );
501            let final_facts_redaction =
502                format!("Here are the facts that I know so far: \n{}", answer_facts);
503            self.logs.push(Step::PlanningStep(
504                final_plan_redaction.clone(),
505                final_facts_redaction,
506            ));
507            info!("Plan: {}", final_plan_redaction.blue().bold());
508        }
509    }
510}
511
512pub struct FunctionCallingAgent<M: Model> {
513    base_agent: MultiStepAgent<M>,
514}
515
516impl<M: Model + Debug> FunctionCallingAgent<M> {
517    pub fn new(
518        model: M,
519        tools: Vec<Box<dyn AnyTool>>,
520        system_prompt: Option<&str>,
521        managed_agents: Option<HashMap<String, Box<dyn Agent>>>,
522        description: Option<&str>,
523        max_steps: Option<usize>,
524    ) -> Result<Self> {
525        let system_prompt = system_prompt.unwrap_or(TOOL_CALLING_SYSTEM_PROMPT);
526        let base_agent = MultiStepAgent::new(
527            model,
528            tools,
529            Some(system_prompt),
530            managed_agents,
531            description,
532            max_steps,
533        )?;
534        Ok(Self { base_agent })
535    }
536}
537
538impl<M: Model + Debug> Agent for FunctionCallingAgent<M> {
539    fn name(&self) -> &'static str {
540        self.base_agent.name()
541    }
542    fn set_task(&mut self, task: &str) {
543        self.base_agent.set_task(task);
544    }
545    fn get_system_prompt(&self) -> &str {
546        self.base_agent.get_system_prompt()
547    }
548    fn get_max_steps(&self) -> usize {
549        self.base_agent.get_max_steps()
550    }
551    fn get_step_number(&self) -> usize {
552        self.base_agent.get_step_number()
553    }
554    fn increment_step_number(&mut self) {
555        self.base_agent.increment_step_number();
556    }
557    fn get_logs_mut(&mut self) -> &mut Vec<Step> {
558        self.base_agent.get_logs_mut()
559    }
560    fn model(&self) -> &dyn Model {
561        self.base_agent.model()
562    }
563
564    /// Perform one step in the ReAct framework: the agent thinks, acts, and observes the result.
565    ///
566    /// Returns None if the step is not final.
567    fn step(&mut self, log_entry: &mut Step) -> Result<Option<String>> {
568        match log_entry {
569            Step::ActionStep(step_log) => {
570                let agent_memory = self.base_agent.write_inner_memory_from_logs(None)?;
571                self.base_agent.input_messages = Some(agent_memory.clone());
572                step_log.agent_memory = Some(agent_memory.clone());
573                let tools = self
574                    .base_agent
575                    .tools
576                    .iter()
577                    .map(|tool| tool.tool_info())
578                    .collect::<Vec<_>>();
579                let model_message = self
580                    .base_agent
581                    .model
582                    .run(
583                        self.base_agent.input_messages.as_ref().unwrap().clone(),
584                        tools,
585                        None,
586                        Some(HashMap::from([(
587                            "stop".to_string(),
588                            vec!["Observation:".to_string()],
589                        )])),
590                    )
591                    .unwrap();
592
593                let mut observations = Vec::new();
594
595                if let Ok(response) = model_message.get_response() {
596                    if !response.trim().is_empty() {
597                        observations.push(response);
598                    }
599                }
600
601                let tools = model_message.get_tools_used()?;
602
603                for tool in tools {
604                    let function_name = tool.clone().function.name;
605
606                    match function_name.as_str() {
607                        "final_answer" => {
608                            info!("Executing tool call: {}", function_name);
609                            let answer = self.base_agent.tools.call(&tool.function)?;
610                            return Ok(Some(answer));
611                        }
612                        _ => {
613                            step_log.tool_call = Some(tool.clone());
614
615                            info!(
616                                "Executing tool call: {} with arguments: {:?}",
617                                function_name, tool.function.arguments
618                            );
619                            let observation = self.base_agent.tools.call(&tool.function);
620                            match observation {
621                                Ok(observation) => {
622                                    observations.push(format!(
623                                        "Observation from {}: {}",
624                                        function_name,
625                                        observation.chars().take(3000).collect::<String>()
626                                    ));
627                                }
628                                Err(e) => {
629                                    step_log.error = Some(AgentError::Execution(e.to_string()));
630                                    info!("Error: {}", e);
631                                }
632                            }
633                            step_log.observations = Some(observations.join("\n"));
634                            info!(
635                                "Observation: {}",
636                                step_log.observations.clone().unwrap_or_default().trim()
637                            );
638                        }
639                    }
640                }
641                Ok(None)
642            }
643            _ => {
644                todo!()
645            }
646        }
647    }
648}
649
650#[cfg(feature = "code-agent")]
651pub struct CodeAgent<M: Model> {
652    base_agent: MultiStepAgent<M>,
653    local_python_interpreter: LocalPythonInterpreter,
654}
655
656#[cfg(feature = "code-agent")]
657impl<M: Model> CodeAgent<M> {
658    pub fn new(
659        model: M,
660        tools: Vec<Box<dyn AnyTool>>,
661        system_prompt: Option<&str>,
662        managed_agents: Option<HashMap<String, Box<dyn Agent>>>,
663        description: Option<&str>,
664        max_steps: Option<usize>,
665    ) -> Result<Self> {
666        let system_prompt = system_prompt.unwrap_or(CODE_SYSTEM_PROMPT);
667
668        let base_agent = MultiStepAgent::new(
669            model,
670            tools,
671            Some(system_prompt),
672            managed_agents,
673            description,
674            max_steps,
675        )?;
676        let local_python_interpreter = LocalPythonInterpreter::new(
677            base_agent
678                .tools
679                .iter()
680                .map(|tool| tool.clone_box())
681                .collect(),
682        );
683
684        Ok(Self {
685            base_agent,
686            local_python_interpreter,
687        })
688    }
689}
690
691#[cfg(feature = "code-agent")]
692impl<M: Model + Debug> Agent for CodeAgent<M> {
693    fn name(&self) -> &'static str {
694        self.base_agent.name()
695    }
696    fn get_max_steps(&self) -> usize {
697        self.base_agent.get_max_steps()
698    }
699    fn get_step_number(&self) -> usize {
700        self.base_agent.get_step_number()
701    }
702    fn increment_step_number(&mut self) {
703        self.base_agent.increment_step_number()
704    }
705    fn get_logs_mut(&mut self) -> &mut Vec<Step> {
706        self.base_agent.get_logs_mut()
707    }
708    fn set_task(&mut self, task: &str) {
709        self.base_agent.set_task(task);
710    }
711    fn get_system_prompt(&self) -> &str {
712        self.base_agent.get_system_prompt()
713    }
714    fn model(&self) -> &dyn Model {
715        self.base_agent.model()
716    }
717    fn step(&mut self, log_entry: &mut Step) -> Result<Option<String>> {
718        match log_entry {
719            Step::ActionStep(step_log) => {
720                let agent_memory = self.base_agent.write_inner_memory_from_logs(None)?;
721                self.base_agent.input_messages = Some(agent_memory.clone());
722                step_log.agent_memory = Some(agent_memory);
723
724                let llm_output = self.base_agent.model.run(
725                    self.base_agent.input_messages.as_ref().unwrap().clone(),
726                    vec![],
727                    None,
728                    Some(HashMap::from([(
729                        "stop".to_string(),
730                        vec!["Observation:".to_string(), "<end_code>".to_string()],
731                    )])),
732                )?;
733
734                let response = llm_output.get_response()?;
735
736                let code = match parse_code_blobs(&response) {
737                    Ok(code) => code,
738                    Err(e) => {
739                        step_log.error = Some(e);
740                        return Ok(None);
741                    }
742                };
743
744                info!("Code: {}", code);
745                step_log.tool_call = Some(ToolCall {
746                    id: None,
747                    call_type: Some("function".to_string()),
748                    function: FunctionCall {
749                        name: "python_interpreter".to_string(),
750                        arguments: serde_json::json!({ "code": code }),
751                    },
752                });
753                let result = self.local_python_interpreter.forward(&code, &mut None);
754                match result {
755                    Ok(result) => {
756                        let (result, execution_logs) = result;
757                        let mut observation = if !execution_logs.is_empty() {
758                            format!(
759                                "Observation: {}\nExecution logs: {}",
760                                result, execution_logs
761                            )
762                        } else {
763                            format!("Observation: {}", result)
764                        };
765                        if observation.len() > 20000 {
766                            observation = observation.chars().take(20000).collect::<String>();
767                            observation = format!("{} \n....This content has been truncated due to the 20000 character limit.....", observation);
768                        }
769                        info!("{}", observation);
770
771                        step_log.observations =
772                            Some(observation.chars().take(20000).collect::<String>());
773                    }
774                    Err(e) => match e {
775                        InterpreterError::FinalAnswer(answer) => {
776                            return Ok(Some(answer));
777                        }
778                        _ => {
779                            step_log.error = Some(AgentError::Execution(e.to_string()));
780                            info!("Error: {}", e);
781                        }
782                    },
783                }
784            }
785            _ => {
786                todo!()
787            }
788        }
789
790        Ok(None)
791    }
792}
793
794#[cfg(feature = "code-agent")]
795pub fn parse_code_blobs(code_blob: &str) -> Result<String, AgentError> {
796    let pattern = r"```(?:py|python)?\n([\s\S]*?)\n```";
797    let re = Regex::new(pattern).map_err(|e| AgentError::Execution(e.to_string()))?;
798
799    let matches: Vec<String> = re
800        .captures_iter(code_blob)
801        .map(|cap| cap[1].trim().to_string())
802        .collect();
803
804    if matches.is_empty() {
805        // Check if it's a direct code blob or final answer
806        if code_blob.contains("final") && code_blob.contains("answer") {
807            return Err(AgentError::Parsing(
808                "The code blob is invalid. It seems like you're trying to return the final answer. Use:\n\
809                Code:\n\
810                ```py\n\
811                final_answer(\"YOUR FINAL ANSWER HERE\")\n\
812                ```".to_string(),
813            ));
814        }
815
816        return Err(AgentError::Parsing(
817            "The code blob is invalid. Make sure to include code with the correct pattern, for instance:\n\
818            Thoughts: Your thoughts\n\
819            Code:\n\
820            ```py\n\
821            # Your python code here\n\
822            ```".to_string(),
823        ));
824    }
825
826    Ok(matches.join("\n\n"))
827}
smolagents_rs/agents.rs

smolagents_rs/
agents.rs
