langchain_rust/agent/open_ai_tools/
agent.rs

1use std::collections::HashMap;
2use std::sync::Arc;
3
4use async_trait::async_trait;
5use serde_json::json;
6
7use crate::{
8    agent::{Agent, AgentError},
9    chain::Chain,
10    fmt_message, fmt_placeholder, fmt_template, message_formatter,
11    prompt::{HumanMessagePromptTemplate, MessageFormatterStruct, PromptArgs},
12    schemas::{
13        agent::{AgentAction, AgentEvent, AgentFinish, LogTools},
14        messages::Message,
15        ToolCall,
16    },
17    template_jinja2,
18    tools::Tool,
19};
20
21pub struct OpenAiToolAgent {
22    pub(crate) chain: Box<dyn Chain>,
23    pub(crate) tools: Vec<Arc<dyn Tool>>,
24}
25
26impl OpenAiToolAgent {
27    pub fn create_prompt(prefix: &str) -> Result<MessageFormatterStruct, AgentError> {
28        let prompt = message_formatter![
29            fmt_message!(Message::new_system_message(prefix)),
30            fmt_placeholder!("chat_history"),
31            fmt_template!(HumanMessagePromptTemplate::new(template_jinja2!(
32                "{{input}}",
33                "input"
34            ))),
35            fmt_placeholder!("agent_scratchpad")
36        ];
37
38        Ok(prompt)
39    }
40
41    fn construct_scratchpad(
42        &self,
43        intermediate_steps: &[(AgentAction, String)],
44    ) -> Result<Vec<Message>, AgentError> {
45        let mut thoughts: Vec<Message> = Vec::new();
46
47        let mut tools_ai_message_seen: HashMap<String, ()> = HashMap::default();
48        for (action, observation) in intermediate_steps {
49            // Deserialize directly and embed in method calls to streamline code.
50            // Extract the tool ID and tool calls from the log.
51            let LogTools { tool_id, tools } = serde_json::from_str(&action.log)?;
52            let tools_vec: Vec<ToolCall> = serde_json::from_str(&tools)?;
53
54            // one action can trigger multiple observations.  But make sure we add it to
55            // the scratchpad before the related observations.  There can also be multiple
56            // different actions in the same thought chain.
57            if tools_ai_message_seen.insert(tools, ()).is_none() {
58                thoughts.push(Message::new_ai_message_with_tool_calls("", &tools_vec));
59            }
60
61            // Add a tool message for each observation. Observation is the ouput of the tool call.
62            // tool_id is the id of the tool.
63            thoughts.push(Message::new_tool_message(observation, tool_id));
64        }
65
66        Ok(thoughts)
67    }
68}
69
70#[async_trait]
71impl Agent for OpenAiToolAgent {
72    async fn plan(
73        &self,
74        intermediate_steps: &[(AgentAction, String)],
75        inputs: PromptArgs,
76    ) -> Result<AgentEvent, AgentError> {
77        let mut inputs = inputs.clone();
78        let scratchpad = self.construct_scratchpad(intermediate_steps)?;
79        inputs.insert("agent_scratchpad".to_string(), json!(scratchpad));
80        let result = self.chain.call(inputs).await?;
81        // Providers fill `tool_calls`; LLMs that only encode calls as JSON text in
82        // `generation` (custom impls, older code) still work through the fallback.
83        let calls = if result.has_tool_calls() {
84            Some(result.tool_calls)
85        } else {
86            ToolCall::parse_list(&result.generation).filter(|c| !c.is_empty())
87        };
88        let Some(calls) = calls else {
89            return Ok(AgentEvent::Finish(AgentFinish {
90                output: result.generation,
91            }));
92        };
93
94        // Every action carries the full list: the scratchpad replays it as one AI message.
95        let tools = ToolCall::list_to_string(&calls);
96        let mut actions: Vec<AgentAction> = Vec::new();
97        for call in calls {
98            let log = LogTools {
99                tool_id: call.id,
100                tools: tools.clone(),
101            };
102            actions.push(AgentAction {
103                tool: call.name,
104                tool_input: call.arguments,
105                log: serde_json::to_string(&log)?, //We send this as string to minimise changes
106            });
107        }
108        Ok(AgentEvent::Action(actions))
109    }
110
111    fn get_tools(&self) -> Vec<Arc<dyn Tool>> {
112        self.tools.clone()
113    }
114}
langchain_rust/agent/open_ai_tools/agent.rs

langchain_rust/agent/open_ai_tools/
agent.rs
