lc_agents/function_calling/
agent.rs

1// src/agents/function_calling/agent.rs
2//! Function Calling Agent 实现
3//!
4//! 使用 LLM 原生 Function Calling 的 Agent，不依赖文本解析。
5//! 支持任何实现了 `BaseChatModel` 的 LLM Provider。
6
7use crate::{AgentAction, AgentError, AgentFinish, AgentOutput, AgentStep, BaseAgent, ToolInput};
8use async_trait::async_trait;
9use lc_core::language_models::{BaseChatModel, LLMResult};
10use lc_core::tools::{to_tool_definition, BaseTool, ToolCall, ToolDefinition};
11use lc_providers::ProviderError;
12use lc_schema::Message;
13use std::collections::HashMap;
14use std::sync::Arc;
15
16/// Function Calling Agent
17///
18/// 使用 LLM 原生 Function Calling 的 Agent。
19/// 不依赖文本解析，直接处理 tool_calls。
20/// 支持任何实现了 `BaseChatModel` 的 LLM Provider。
21pub struct FunctionCallingAgent {
22    /// LLM 客户端（已绑定工具）
23    llm: Arc<dyn BaseChatModel<Error = ProviderError> + Send + Sync>,
24
25    /// 可用工具列表
26    tools: Vec<Arc<dyn BaseTool>>,
27
28    /// 自定义系统提示词
29    system_prompt: Option<String>,
30}
31
32impl FunctionCallingAgent {
33    /// 创建新的 Function Calling Agent
34    ///
35    /// # 参数
36    /// * `llm` - LLM 客户端（任何实现了 `BaseChatModel` 的类型）
37    /// * `tools` - 可用工具列表
38    /// * `system_prompt` - 自定义系统提示词（可选）
39    ///
40    /// # 向后兼容
41    /// 旧代码 `FunctionCallingAgent::new(openai_chat, tools, None)` 仍然可用，
42    /// 因为 `OpenAIChat: BaseChatModel` 且 `OpenAIError: Into<Error>`。
43    pub fn new<L>(llm: L, tools: Vec<Arc<dyn BaseTool>>, system_prompt: Option<String>) -> Self
44    where
45        L: BaseChatModel + Send + Sync + 'static,
46        L::Error: Into<ProviderError>,
47    {
48        // 先包装 LLM，将错误类型统一为 ProviderError
49        let wrapped = lc_providers::ChatModelWrapper::new(llm);
50
51        let tool_definitions: Vec<ToolDefinition> = tools
52            .iter()
53            .map(|t| to_tool_definition(t.as_ref()))
54            .collect();
55
56        // 优先用 trait bind_tools（返回 Box<dyn BaseChatModel<Error = ProviderError>>）
57        // Provider 不支持则直接用包装后的 LLM
58        let llm_with_tools: Arc<dyn BaseChatModel<Error = ProviderError> + Send + Sync> = wrapped
59            .bind_tools(tool_definitions)
60            .map(|boxed| {
61                Arc::from(boxed) as Arc<dyn BaseChatModel<Error = ProviderError> + Send + Sync>
62            })
63            .unwrap_or_else(|| Arc::new(wrapped));
64
65        Self {
66            llm: llm_with_tools,
67            tools,
68            system_prompt,
69        }
70    }
71
72    /// 从已包装的 `Arc<dyn BaseChatModel>` 创建 Agent
73    ///
74    /// 适用于已通过 `wrap_chat_model()` 或 `LLMClient` 创建的 LLM 实例。
75    pub fn from_arc(
76        llm: Arc<dyn BaseChatModel<Error = ProviderError> + Send + Sync>,
77        tools: Vec<Arc<dyn BaseTool>>,
78        system_prompt: Option<String>,
79    ) -> Self {
80        let tool_definitions: Vec<ToolDefinition> = tools
81            .iter()
82            .map(|t| to_tool_definition(t.as_ref()))
83            .collect();
84
85        let llm_with_tools = llm
86            .bind_tools(tool_definitions)
87            .map(|boxed| {
88                Arc::from(boxed) as Arc<dyn BaseChatModel<Error = ProviderError> + Send + Sync>
89            })
90            .unwrap_or(llm);
91
92        Self {
93            llm: llm_with_tools,
94            tools,
95            system_prompt,
96        }
97    }
98
99    /// 获取工具数量
100    pub fn tools_count(&self) -> usize {
101        self.tools.len()
102    }
103
104    /// 获取系统提示词
105    pub fn system_prompt(&self) -> Option<&str> {
106        self.system_prompt.as_deref()
107    }
108
109    /// 构建消息
110    fn build_messages(
111        &self,
112        inputs: &HashMap<String, String>,
113        intermediate_steps: &[AgentStep],
114    ) -> Vec<Message> {
115        let mut messages = Vec::new();
116
117        let system_content = self
118            .system_prompt
119            .clone()
120            .unwrap_or_else(|| "你是一个助手，可以使用工具回答问题。".to_string());
121        messages.push(Message::system(&system_content));
122
123        let default_input = String::new();
124        let input = inputs.get("input").unwrap_or(&default_input);
125        messages.push(Message::human(input));
126
127        for step in intermediate_steps {
128            let tool_call = ToolCall::new(
129                &step.action.log,
130                &step.action.tool,
131                match &step.action.tool_input {
132                    ToolInput::String { value: s } => s.clone(),
133                    ToolInput::Object { value: v } => {
134                        serde_json::to_string(v).unwrap_or_else(|_| v.to_string())
135                    }
136                },
137            );
138            messages.push(Message::ai_with_tool_calls("", vec![tool_call]));
139            messages.push(Message::tool(&step.action.log, &step.observation));
140        }
141
142        messages
143    }
144}
145
146#[async_trait]
147impl BaseAgent for FunctionCallingAgent {
148    async fn plan(
149        &self,
150        intermediate_steps: &[AgentStep],
151        inputs: &HashMap<String, String>,
152    ) -> Result<AgentOutput, AgentError> {
153        let messages = self.build_messages(inputs, intermediate_steps);
154
155        let result: LLMResult = self
156            .llm
157            .chat(messages, None)
158            .await
159            .map_err(|e| AgentError::Other(format!("LLM 调用失败: {}", e)))?;
160
161        if let Some(tool_calls) = &result.tool_calls {
162            if !tool_calls.is_empty() {
163                let actions: Vec<AgentAction> = tool_calls
164                    .iter()
165                    .map(|call| {
166                        let tool_input = match serde_json::from_str::<serde_json::Value>(
167                            &call.function.arguments,
168                        ) {
169                            Ok(v) => ToolInput::Object { value: v },
170                            Err(_) => ToolInput::String {
171                                value: call.function.arguments.clone(),
172                            },
173                        };
174
175                        AgentAction {
176                            tool: call.function.name.clone(),
177                            tool_input,
178                            log: call.id.clone(),
179                        }
180                    })
181                    .collect();
182
183                if actions.len() == 1 {
184                    return Ok(AgentOutput::Action(actions.into_iter().next().unwrap()));
185                } else {
186                    return Ok(AgentOutput::Actions(actions));
187                }
188            }
189        }
190
191        Ok(AgentOutput::Finish(AgentFinish::new(
192            result.content.clone(),
193            String::new(),
194        )))
195    }
196
197    fn get_allowed_tools(&self) -> Option<Vec<&str>> {
198        Some(self.tools.iter().map(|t| t.name()).collect())
199    }
200}
201
202impl std::fmt::Debug for FunctionCallingAgent {
203    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
204        f.debug_struct("FunctionCallingAgent")
205            .field("tools_count", &self.tools.len())
206            .field("system_prompt", &self.system_prompt)
207            .finish()
208    }
209}
210
211#[cfg(test)]
212mod tests {
213    use super::*;
214    use lc_providers::{OpenAIChat, OpenAIConfig};
215    use lc_tools::Calculator;
216
217    fn create_test_config() -> OpenAIConfig {
218        OpenAIConfig::new("test_key").with_base_url("http://localhost:8080/v1")
219    }
220
221    #[test]
222    fn test_function_calling_agent_creation() {
223        let config = create_test_config();
224        let llm = OpenAIChat::new(config);
225        let tools: Vec<Arc<dyn BaseTool>> = vec![Arc::new(Calculator::new())];
226
227        let agent = FunctionCallingAgent::new(llm, tools, None);
228        assert_eq!(agent.tools.len(), 1);
229    }
230
231    #[test]
232    fn test_get_allowed_tools() {
233        let config = create_test_config();
234        let llm = OpenAIChat::new(config);
235        let tools: Vec<Arc<dyn BaseTool>> = vec![Arc::new(Calculator::new())];
236
237        let agent = FunctionCallingAgent::new(llm, tools, None);
238
239        assert_eq!(agent.tools.len(), 1);
240        assert!(agent.system_prompt.is_none());
241    }
242
243    #[test]
244    fn test_new_with_system_prompt() {
245        let config = create_test_config();
246        let llm = OpenAIChat::new(config);
247        let tools: Vec<Arc<dyn BaseTool>> = vec![Arc::new(Calculator::new())];
248
249        let agent = FunctionCallingAgent::new(llm, tools, Some("你是一个数学助手".to_string()));
250
251        assert_eq!(agent.system_prompt, Some("你是一个数学助手".to_string()));
252    }
253
254    #[test]
255    fn test_build_messages_empty() {
256        let config = create_test_config();
257        let llm = OpenAIChat::new(config);
258        let tools: Vec<Arc<dyn BaseTool>> = vec![Arc::new(Calculator::new())];
259
260        let agent = FunctionCallingAgent::new(llm, tools, None);
261
262        let mut inputs = HashMap::new();
263        inputs.insert("input".to_string(), "计算 2 + 3".to_string());
264
265        let messages = agent.build_messages(&inputs, &[]);
266
267        assert_eq!(messages.len(), 2);
268        assert_eq!(messages[0].content, "你是一个助手，可以使用工具回答问题。");
269        assert_eq!(messages[1].content, "计算 2 + 3");
270    }
271
272    #[test]
273    fn test_build_messages_with_history() {
274        let config = create_test_config();
275        let llm = OpenAIChat::new(config);
276        let tools: Vec<Arc<dyn BaseTool>> = vec![Arc::new(Calculator::new())];
277
278        let agent = FunctionCallingAgent::new(llm, tools, None);
279
280        let mut inputs = HashMap::new();
281        inputs.insert("input".to_string(), "继续计算".to_string());
282
283        let steps = vec![AgentStep::new(
284            AgentAction {
285                tool: "calculator".to_string(),
286                tool_input: ToolInput::String {
287                    value: "2 + 3".to_string(),
288                },
289                log: "call_123".to_string(),
290            },
291            "5".to_string(),
292        )];
293
294        let messages = agent.build_messages(&inputs, &steps);
295
296        assert_eq!(messages.len(), 4);
297        assert!(messages[2].has_tool_calls());
298    }
299
300    #[test]
301    fn test_from_arc_creation() {
302        let config = create_test_config();
303        let llm = OpenAIChat::new(config);
304        let llm_arc: Arc<dyn BaseChatModel<Error = ProviderError> + Send + Sync> =
305            lc_providers::wrap_chat_model(llm);
306        let tools: Vec<Arc<dyn BaseTool>> = vec![Arc::new(Calculator::new())];
307
308        let agent = FunctionCallingAgent::from_arc(llm_arc, tools, Some("test".into()));
309        assert_eq!(agent.tools.len(), 1);
310        assert_eq!(agent.system_prompt, Some("test".to_string()));
311    }
312}
lc_agents/function_calling/agent.rs

lc_agents/function_calling/
agent.rs
