langchain_rust/agent/
executor.rs

1use std::{collections::HashMap, sync::Arc};
2
3use async_trait::async_trait;
4use serde_json::json;
5use tokio::sync::Mutex;
6
7use super::{agent::Agent, AgentError};
8use crate::schemas::{LogTools, Message};
9use crate::{
10    chain::{chain_trait::Chain, ChainError},
11    language_models::GenerateResult,
12    prompt::PromptArgs,
13    schemas::{
14        agent::{AgentAction, AgentEvent},
15        memory::BaseMemory,
16    },
17    tools::Tool,
18};
19
20pub struct AgentExecutor<A>
21where
22    A: Agent,
23{
24    agent: A,
25    max_iterations: Option<i32>,
26    break_if_error: bool,
27    pub memory: Option<Arc<Mutex<dyn BaseMemory>>>,
28}
29
30impl<A> AgentExecutor<A>
31where
32    A: Agent,
33{
34    pub fn from_agent(agent: A) -> Self {
35        Self {
36            agent,
37            max_iterations: Some(10),
38            break_if_error: false,
39            memory: None,
40        }
41    }
42
43    pub fn with_max_iterations(mut self, max_iterations: i32) -> Self {
44        self.max_iterations = Some(max_iterations);
45        self
46    }
47
48    pub fn with_memory(mut self, memory: Arc<Mutex<dyn BaseMemory>>) -> Self {
49        self.memory = Some(memory);
50        self
51    }
52
53    pub fn with_break_if_error(mut self, break_if_error: bool) -> Self {
54        self.break_if_error = break_if_error;
55        self
56    }
57
58    fn get_name_to_tools(&self) -> HashMap<String, Arc<dyn Tool>> {
59        let mut name_to_tool = HashMap::new();
60        for tool in self.agent.get_tools().iter() {
61            log::debug!("Loading Tool:{}", tool.name());
62            name_to_tool.insert(tool.name().trim().replace(" ", "_"), tool.clone());
63        }
64        name_to_tool
65    }
66}
67
68#[async_trait]
69impl<A> Chain for AgentExecutor<A>
70where
71    A: Agent + Send + Sync,
72{
73    async fn call(&self, input_variables: PromptArgs) -> Result<GenerateResult, ChainError> {
74        let mut input_variables = input_variables.clone();
75        let name_to_tools = self.get_name_to_tools();
76        let mut steps: Vec<(AgentAction, String)> = Vec::new();
77        log::debug!("steps: {:?}", steps);
78        if let Some(memory) = &self.memory {
79            let messages = memory.lock().await.messages().await?;
80            input_variables.insert("chat_history".to_string(), json!(messages));
81        } else {
82            input_variables.insert("chat_history".to_string(), json!(Vec::<Message>::new()));
83        }
84
85        loop {
86            let agent_event = self
87                .agent
88                .plan(&steps, input_variables.clone())
89                .await
90                .map_err(|e| ChainError::AgentError(format!("Error in agent planning: {}", e)))?;
91            match agent_event {
92                AgentEvent::Action(actions) => {
93                    for action in actions {
94                        log::debug!("Action: {:?}", action.tool_input);
95                        let tool = name_to_tools
96                            .get(&action.tool.trim().replace(" ", "_"))
97                            .ok_or_else(|| {
98                                AgentError::ToolError(format!("Tool {} not found", action.tool))
99                            })
100                            .map_err(|e| ChainError::AgentError(e.to_string()))?;
101
102                        let observation_result = tool.call(&action.tool_input).await;
103
104                        let observation = match observation_result {
105                            Ok(result) => result,
106                            Err(err) => {
107                                log::info!(
108                                    "The tool return the following error: {}",
109                                    err.to_string()
110                                );
111                                if self.break_if_error {
112                                    return Err(ChainError::AgentError(
113                                        AgentError::ToolError(err.to_string()).to_string(),
114                                    ));
115                                } else {
116                                    format!("The tool return the following error: {}", err)
117                                }
118                            }
119                        };
120
121                        steps.push((action, observation));
122                    }
123                }
124                AgentEvent::Finish(finish) => {
125                    if let Some(memory) = &self.memory {
126                        let mut memory = memory.lock().await;
127
128                        memory
129                            .add_user_message(match &input_variables["input"] {
130                                // This avoids adding extra quotes to the user input in the history.
131                                serde_json::Value::String(s) => s,
132                                x => x, // this the json encoded value.
133                            })
134                            .await?;
135
136                        let mut tools_ai_message_seen: HashMap<String, ()> = HashMap::default();
137                        for (action, observation) in steps {
138                            // Only OpenAI-tools style agents encode `LogTools` in the log;
139                            // other agents (e.g. ConversationalAgent) store raw LLM text.
140                            let Ok(LogTools { tool_id, tools }) =
141                                serde_json::from_str::<LogTools>(&action.log)
142                            else {
143                                continue;
144                            };
145                            let tools_value: serde_json::Value = serde_json::from_str(&tools)?;
146                            if tools_ai_message_seen.insert(tools, ()).is_none() {
147                                memory
148                                    .add_message(
149                                        Message::new_ai_message("").with_tool_calls(tools_value),
150                                    )
151                                    .await?;
152                            }
153                            memory
154                                .add_message(Message::new_tool_message(observation, tool_id))
155                                .await?;
156                        }
157
158                        memory.add_ai_message(&finish.output).await?;
159                    }
160                    return Ok(GenerateResult {
161                        generation: finish.output,
162                        ..Default::default()
163                    });
164                }
165            }
166
167            if let Some(max_iterations) = self.max_iterations {
168                if steps.len() >= max_iterations as usize {
169                    return Ok(GenerateResult {
170                        generation: "Max iterations reached".to_string(),
171                        ..Default::default()
172                    });
173                }
174            }
175        }
176    }
177
178    async fn invoke(&self, input_variables: PromptArgs) -> Result<String, ChainError> {
179        let result = self.call(input_variables).await?;
180        Ok(result.generation)
181    }
182}
langchain_rust/agent/executor.rs

langchain_rust/agent/
executor.rs
