langchain_rust/agent/chat/
chat_agent.rs

1use std::sync::Arc;
2
3use async_trait::async_trait;
4use serde_json::json;
5
6use crate::{
7    agent::{agent::Agent, chat::prompt::FORMAT_INSTRUCTIONS, AgentError},
8    chain::chain_trait::Chain,
9    language_models::TokenUsage,
10    message_formatter,
11    prompt::{
12        HumanMessagePromptTemplate, MessageFormatterStruct, MessageOrTemplate, PromptArgs,
13        PromptFromatter, DEFAULT_IMAGES_KEY,
14    },
15    prompt_args,
16    schemas::{
17        agent::{AgentAction, AgentEvent},
18        messages::Message,
19    },
20    template_jinja2,
21    tools::Tool,
22};
23
24use super::{output_parser::ChatOutputParser, prompt::TEMPLATE_TOOL_RESPONSE};
25
26pub struct ConversationalAgent {
27    pub(crate) chain: Box<dyn Chain>,
28    pub(crate) tools: Vec<Arc<dyn Tool>>,
29    pub(crate) output_parser: ChatOutputParser,
30}
31
32impl ConversationalAgent {
33    pub fn create_prompt(
34        tools: &[Arc<dyn Tool>],
35        suffix: &str,
36        prefix: &str,
37    ) -> Result<MessageFormatterStruct, AgentError> {
38        let tool_string = tools
39            .iter()
40            .map(|tool| format!("> {}: {}", tool.name(), tool.description()))
41            .collect::<Vec<_>>()
42            .join("\n");
43        let tool_names = tools
44            .iter()
45            .map(|tool| tool.name())
46            .collect::<Vec<_>>()
47            .join(", ");
48
49        let sufix_prompt = template_jinja2!(suffix, "tools", "format_instructions");
50
51        let input_variables_fstring = prompt_args! {
52            "tools" => tool_string,
53            "format_instructions" => FORMAT_INSTRUCTIONS,
54            "tool_names"=>tool_names
55        };
56
57        let sufix_prompt = sufix_prompt.format(input_variables_fstring)?;
58        let formatter = message_formatter![
59            MessageOrTemplate::Message(Message::new_system_message(prefix)),
60            MessageOrTemplate::MessagesPlaceholder("chat_history".to_string()),
61            MessageOrTemplate::Template(
62                HumanMessagePromptTemplate::new(template_jinja2!(
63                    &sufix_prompt.to_string(),
64                    "input"
65                ))
66                .with_images(DEFAULT_IMAGES_KEY)
67                .into()
68            ),
69            MessageOrTemplate::MessagesPlaceholder("agent_scratchpad".to_string()),
70        ];
71        Ok(formatter)
72    }
73
74    fn construct_scratchpad(
75        &self,
76        intermediate_steps: &[(AgentAction, String)],
77    ) -> Result<Vec<Message>, AgentError> {
78        let mut thoughts: Vec<Message> = Vec::new();
79        for (action, observation) in intermediate_steps.iter() {
80            thoughts.push(Message::new_ai_message(&action.log));
81            let tool_response = template_jinja2!(TEMPLATE_TOOL_RESPONSE, "observation")
82                .format(prompt_args!("observation"=>observation))?;
83            thoughts.push(Message::new_human_message(&tool_response));
84        }
85        Ok(thoughts)
86    }
87}
88
89#[async_trait]
90impl Agent for ConversationalAgent {
91    async fn plan(
92        &self,
93        intermediate_steps: &[(AgentAction, String)],
94        inputs: PromptArgs,
95    ) -> Result<AgentEvent, AgentError> {
96        Ok(self.plan_with_usage(intermediate_steps, inputs).await?.0)
97    }
98
99    async fn plan_with_usage(
100        &self,
101        intermediate_steps: &[(AgentAction, String)],
102        inputs: PromptArgs,
103    ) -> Result<(AgentEvent, Option<TokenUsage>), AgentError> {
104        let scratchpad = self.construct_scratchpad(intermediate_steps)?;
105        let mut inputs = inputs.clone();
106        inputs.insert("agent_scratchpad".to_string(), json!(scratchpad));
107        let result = self.chain.call(inputs.clone()).await?;
108        let parsed_output = self.output_parser.parse(&result.generation)?;
109        Ok((parsed_output, result.tokens))
110    }
111
112    fn get_tools(&self) -> Vec<Arc<dyn Tool>> {
113        self.tools.clone()
114    }
115}
116
117#[cfg(test)]
118mod tests {
119    use std::{error::Error, sync::Arc};
120
121    use async_trait::async_trait;
122    use serde_json::Value;
123
124    use crate::{
125        agent::{chat::builder::ConversationalAgentBuilder, executor::AgentExecutor},
126        chain::chain_trait::Chain,
127        llm::openai::{OpenAI, OpenAIModel},
128        memory::SimpleMemory,
129        prompt_args,
130        tools::Tool,
131    };
132
133    struct Calc {}
134
135    #[async_trait]
136    impl Tool for Calc {
137        fn name(&self) -> String {
138            "Calculator".to_string()
139        }
140        fn description(&self) -> String {
141            "Usefull to make calculations".to_string()
142        }
143        async fn run(&self, _input: Value) -> Result<String, Box<dyn Error>> {
144            Ok("25".to_string())
145        }
146    }
147
148    #[tokio::test]
149    #[ignore]
150    async fn test_invoke_agent() {
151        let llm = OpenAI::default().with_model(OpenAIModel::Gpt4.to_string());
152        let memory = SimpleMemory::new();
153        let tool_calc = Calc {};
154        let agent = ConversationalAgentBuilder::new()
155            .tools(&[Arc::new(tool_calc)])
156            .build(llm)
157            .unwrap();
158        let input_variables = prompt_args! {
159            "input" => "hola,Me llamo luis, y tengo 10 anos, y estudio Computer scinence",
160        };
161        let executor = AgentExecutor::from_agent(agent).with_memory(memory.into());
162        match executor.invoke(input_variables).await {
163            Ok(result) => {
164                println!("Result: {:?}", result);
165            }
166            Err(e) => panic!("Error invoking LLMChain: {:?}", e),
167        }
168        let input_variables = prompt_args! {
169            "input" => "cuanta es la edad de luis +10 y que estudia",
170        };
171        match executor.invoke(input_variables).await {
172            Ok(result) => {
173                println!("Result: {:?}", result);
174            }
175            Err(e) => panic!("Error invoking LLMChain: {:?}", e),
176        }
177    }
178}
langchain_rust/agent/chat/chat_agent.rs

langchain_rust/agent/chat/
chat_agent.rs
