lc_agents/function_calling/
agent.rs

1// src/agents/function_calling/agent.rs
2//! Function Calling Agent implementation
3//!
4//! An agent that uses the LLM's native Function Calling, without text parsing.
5//! Supports any LLM provider that implements `BaseChatModel`.
6
7use crate::{AgentAction, AgentError, AgentFinish, AgentOutput, AgentStep, BaseAgent, ToolInput};
8use async_trait::async_trait;
9use futures_util::StreamExt;
10use lc_core::language_models::{BaseChatModel, LLMResult, TokenUsage};
11use lc_core::runnables::RunnableConfig;
12use lc_core::tools::{to_tool_definition, BaseTool, ToolCall, ToolDefinition};
13use lc_providers::ProviderError;
14use lc_schema::Message;
15use std::collections::HashMap;
16use std::future::Future;
17use std::pin::Pin;
18use std::sync::Arc;
19
20/// Function Calling Agent
21///
22/// An agent that uses the LLM's native Function Calling.
23/// Does not rely on text parsing; handles `tool_calls` directly.
24/// Supports any LLM provider that implements `BaseChatModel`.
25pub struct FunctionCallingAgent {
26    /// LLM client (with tools bound)
27    llm: Arc<dyn BaseChatModel<Error = ProviderError> + Send + Sync>,
28
29    /// Available tools
30    tools: Vec<Arc<dyn BaseTool>>,
31
32    /// Custom system prompt
33    system_prompt: Option<String>,
34
35    /// Token usage from the most recent `plan()` call (P1-5).
36    last_token_usage: std::sync::Mutex<Option<TokenUsage>>,
37}
38
39impl FunctionCallingAgent {
40    /// Creates a new Function Calling Agent
41    ///
42    /// # Parameters
43    /// * `llm` - LLM client (any type implementing `BaseChatModel`)
44    /// * `tools` - available tools
45    /// * `system_prompt` - custom system prompt (optional)
46    ///
47    /// # Backward compatibility
48    /// Legacy code `FunctionCallingAgent::new(openai_chat, tools, None)` still works,
49    /// because `OpenAIChat: BaseChatModel` and `OpenAIError: Into<Error>`.
50    pub fn new<L>(llm: L, tools: Vec<Arc<dyn BaseTool>>, system_prompt: Option<String>) -> Self
51    where
52        L: BaseChatModel + Send + Sync + 'static,
53        L::Error: Into<ProviderError>,
54    {
55        // Wrap the LLM first, unifying the error type to ProviderError
56        let wrapped = lc_providers::ChatModelWrapper::new(llm);
57
58        let tool_definitions: Vec<ToolDefinition> = tools
59            .iter()
60            .map(|t| to_tool_definition(t.as_ref()))
61            .collect();
62
63        // Prefer the trait bind_tools (returns Box<dyn BaseChatModel<Error = ProviderError>>)
64        // Fall back to the wrapped LLM if the provider does not support it
65        let llm_with_tools: Arc<dyn BaseChatModel<Error = ProviderError> + Send + Sync> = wrapped
66            .bind_tools(tool_definitions)
67            .map(|boxed| {
68                Arc::from(boxed) as Arc<dyn BaseChatModel<Error = ProviderError> + Send + Sync>
69            })
70            .unwrap_or_else(|| Arc::new(wrapped));
71
72        Self {
73            llm: llm_with_tools,
74            tools,
75            system_prompt,
76            last_token_usage: std::sync::Mutex::new(None),
77        }
78    }
79
80    /// Creates an agent from an already-wrapped `Arc<dyn BaseChatModel>`
81    ///
82    /// For LLM instances already created via `wrap_chat_model()` or `LLMClient`.
83    pub fn from_arc(
84        llm: Arc<dyn BaseChatModel<Error = ProviderError> + Send + Sync>,
85        tools: Vec<Arc<dyn BaseTool>>,
86        system_prompt: Option<String>,
87    ) -> Self {
88        let tool_definitions: Vec<ToolDefinition> = tools
89            .iter()
90            .map(|t| to_tool_definition(t.as_ref()))
91            .collect();
92
93        let llm_with_tools = llm
94            .bind_tools(tool_definitions)
95            .map(|boxed| {
96                Arc::from(boxed) as Arc<dyn BaseChatModel<Error = ProviderError> + Send + Sync>
97            })
98            .unwrap_or(llm);
99
100        Self {
101            llm: llm_with_tools,
102            tools,
103            system_prompt,
104            last_token_usage: std::sync::Mutex::new(None),
105        }
106    }
107
108    /// Returns the number of tools
109    pub fn tools_count(&self) -> usize {
110        self.tools.len()
111    }
112
113    /// Returns the system prompt
114    pub fn system_prompt(&self) -> Option<&str> {
115        self.system_prompt.as_deref()
116    }
117
118    /// Builds the messages
119    fn build_messages(
120        &self,
121        inputs: &HashMap<String, String>,
122        intermediate_steps: &[AgentStep],
123    ) -> Vec<Message> {
124        let mut messages = Vec::new();
125
126        let system_content = self
127            .system_prompt
128            .clone()
129            .unwrap_or_else(|| "你是一个助手，可以使用工具回答问题。".to_string());
130        messages.push(Message::system(&system_content));
131
132        let default_input = String::new();
133        let input = inputs.get("input").unwrap_or(&default_input);
134        messages.push(Message::human(input));
135
136        for step in intermediate_steps {
137            let tool_call = ToolCall::builder(&step.action.log)
138                .name(&step.action.tool)
139                .arguments(match &step.action.tool_input {
140                    ToolInput::String { value: s } => s.clone(),
141                    ToolInput::Object { value: v } => {
142                        serde_json::to_string(v).unwrap_or_else(|_| v.to_string())
143                    }
144                })
145                .build();
146            messages.push(Message::ai_with_tool_calls("", vec![tool_call]));
147            messages.push(Message::tool(&step.action.log, &step.observation));
148        }
149
150        messages
151    }
152
153    /// Converts complete [`ToolCall`]s from an LLM response into an [`AgentOutput`].
154    ///
155    /// Shared by the non-streaming [`BaseAgent::plan`] and the streaming
156    /// [`BaseAgent::plan_stream`]: both receive the same native `tool_calls`
157    /// shape, so they build actions identically. 0.20.0 S3.2: `plan_stream`
158    /// now reaches this through the streamed chunks instead of a non-streaming
159    /// fallback.
160    fn output_from_tool_calls(tool_calls: &[ToolCall]) -> AgentOutput {
161        let actions: Vec<AgentAction> = tool_calls
162            .iter()
163            .map(|call| {
164                let tool_input =
165                    match serde_json::from_str::<serde_json::Value>(&call.function.arguments) {
166                        Ok(v) => ToolInput::Object { value: v },
167                        Err(_) => ToolInput::String {
168                            value: call.function.arguments.clone(),
169                        },
170                    };
171
172                AgentAction {
173                    tool: call.function.name.clone(),
174                    tool_input,
175                    log: call.id.clone(),
176                }
177            })
178            .collect();
179
180        if actions.len() == 1 {
181            AgentOutput::Action(actions.into_iter().next().expect("checked len == 1"))
182        } else {
183            AgentOutput::Actions(actions)
184        }
185    }
186}
187
188#[async_trait]
189impl BaseAgent for FunctionCallingAgent {
190    async fn plan(
191        &self,
192        intermediate_steps: &[AgentStep],
193        inputs: &HashMap<String, String>,
194        config: Option<&RunnableConfig>,
195    ) -> Result<AgentOutput, AgentError> {
196        let messages = self.build_messages(inputs, intermediate_steps);
197
198        // A18: forward the executor's config so the provider dispatches the
199        // planning-round `on_llm_*` callbacks.
200        let result: LLMResult = crate::retry::retry_chat(
201            self.llm.as_ref(),
202            messages,
203            config.cloned(),
204            &crate::retry::RetryConfig::default(),
205        )
206        .await
207        .map_err(|e| AgentError::Other(format!("LLM call failed: {}", e)))?;
208
209        // P1-5: record token usage for the executor's metrics.
210        if let Ok(mut guard) = self.last_token_usage.lock() {
211            *guard = result.token_usage.clone();
212        }
213
214        if let Some(tool_calls) = &result.tool_calls {
215            if !tool_calls.is_empty() {
216                return Ok(Self::output_from_tool_calls(tool_calls));
217            }
218        }
219
220        Ok(AgentOutput::Finish(AgentFinish::new(
221            result.content.clone(),
222            String::new(),
223        )))
224    }
225
226    /// Streaming plan (S2 + 0.20.0 S3.2): goes through `stream_chat`, forwarding
227    /// model text token by token and accumulating usage **and tool calls**.
228    ///
229    /// The function-calling agent's **final answer** streams out as text token by
230    /// token (typewriter effect). **Tool-call steps** now stream natively too:
231    /// providers that support streaming `tool_calls` (OpenAI / Azure and their
232    /// delegates) attach the complete tool calls to the terminal `StreamChunk`
233    /// (`StreamChunk.tool_calls`), which `plan_stream` accumulates and converts
234    /// into [`AgentOutput::Action`]/[`AgentOutput::Actions`] — no non-streaming
235    /// fallback needed, so the agent loop does not emit a fake "empty Finish".
236    /// **Mixed steps** (text + tool calls) keep both: the text already streamed
237    /// via `on_token`, the tool calls preserved here.
238    ///
239    /// The non-streaming [`BaseAgent::plan`] fallback remains only as a safety net
240    /// for: `stream_chat` failing immediately, or a provider that yields neither
241    /// text nor `tool_calls` on the stream (e.g. one without streaming tool-call
242    /// support when the model makes a tool call).
243    async fn plan_stream(
244        &self,
245        intermediate_steps: &[AgentStep],
246        inputs: &HashMap<String, String>,
247        on_token: &mut (dyn FnMut(String) -> Pin<Box<dyn Future<Output = ()> + Send>> + Send),
248        config: Option<&RunnableConfig>,
249    ) -> Result<AgentOutput, AgentError> {
250        let messages = self.build_messages(inputs, intermediate_steps);
251
252        // A18: stream planning carries the same config as non-streaming plan.
253        let mut stream = match self.llm.stream_chat(messages, config.cloned()).await {
254            Ok(s) => s,
255            Err(e) => {
256                log::warn!(
257                    "stream_chat unavailable ({}), falling back to non-streaming plan",
258                    e
259                );
260                let output = self.plan(intermediate_steps, inputs, config).await?;
261                if let AgentOutput::Finish(finish) = &output {
262                    on_token(finish.output().unwrap_or("").to_string()).await;
263                }
264                return Ok(output);
265            }
266        };
267
268        // Token by token: forward non-empty text, accumulate the full text + the
269        // last non-None usage, and the last non-empty set of complete tool calls
270        // (providers attach them to the terminal chunk).
271        let mut full = String::new();
272        let mut usage: Option<TokenUsage> = None;
273        let mut tool_calls: Option<Vec<ToolCall>> = None;
274        while let Some(chunk) = stream.next().await {
275            let chunk = chunk.map_err(|e| AgentError::Other(format!("LLM stream error: {}", e)))?;
276            if !chunk.text.is_empty() {
277                on_token(chunk.text.clone()).await;
278            }
279            full.push_str(&chunk.text);
280            if chunk.token_usage.is_some() {
281                usage = chunk.token_usage;
282            }
283            // 0.20.0 S3.2: take the last non-empty tool_calls set so a tool-call
284            // step is returned natively below instead of falling back.
285            if let Some(tc) = &chunk.tool_calls {
286                if !tc.is_empty() {
287                    tool_calls = Some(tc.clone());
288                }
289            }
290        }
291        if let Ok(mut guard) = self.last_token_usage.lock() {
292            *guard = usage;
293        }
294
295        // 0.20.0 S3.2: tool-call steps stream natively — return the accumulated
296        // tool_calls as Action/Actions. Mixed steps keep both: the text already
297        // streamed via on_token, the tool calls preserved here.
298        if let Some(tc) = &tool_calls {
299            return Ok(Self::output_from_tool_calls(tc));
300        }
301
302        // Neither text nor tool_calls on the stream (a provider without streaming
303        // tool-call support on a tool-call step, or an empty reply): fall back to
304        // non-streaming plan() so the agent loop does not end early on an empty
305        // Finish.
306        if full.trim().is_empty() {
307            log::debug!(
308                "streamed plan produced neither text nor tool_calls, \
309                 falling back to non-streaming plan"
310            );
311            let output = self.plan(intermediate_steps, inputs, config).await?;
312            if let AgentOutput::Finish(finish) = &output {
313                on_token(finish.output().unwrap_or("").to_string()).await;
314            }
315            return Ok(output);
316        }
317
318        Ok(AgentOutput::Finish(AgentFinish::new(full, String::new())))
319    }
320
321    fn get_allowed_tools(&self) -> Option<Vec<&str>> {
322        Some(self.tools.iter().map(|t| t.name()).collect())
323    }
324
325    /// Reports the token usage from the most recent `plan()` call (P1-5).
326    fn last_token_usage(&self) -> Option<TokenUsage> {
327        self.last_token_usage.lock().ok().and_then(|g| g.clone())
328    }
329}
330
331impl std::fmt::Debug for FunctionCallingAgent {
332    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
333        f.debug_struct("FunctionCallingAgent")
334            .field("tools_count", &self.tools.len())
335            .field("system_prompt", &self.system_prompt)
336            .field(
337                "has_token_usage",
338                &self.last_token_usage.lock().ok().is_some(),
339            )
340            .finish()
341    }
342}
343
344#[cfg(test)]
345mod tests {
346    use super::*;
347    use futures_util::Stream;
348    use lc_core::language_models::{BaseLanguageModel, StreamChunk};
349    use lc_core::runnables::{Runnable, RunnableConfig};
350    use lc_providers::{AssistantError, OpenAIChat, OpenAIConfig};
351    use lc_tools::Calculator;
352    use std::sync::Mutex;
353
354    fn create_test_config() -> OpenAIConfig {
355        OpenAIConfig::new("test_key").with_base_url("http://localhost:8080/v1")
356    }
357
358    #[test]
359    fn test_function_calling_agent_creation() {
360        let config = create_test_config();
361        let llm = OpenAIChat::new(config);
362        let tools: Vec<Arc<dyn BaseTool>> = vec![Arc::new(Calculator::new())];
363
364        let agent = FunctionCallingAgent::new(llm, tools, None);
365        assert_eq!(agent.tools.len(), 1);
366    }
367
368    #[test]
369    fn test_get_allowed_tools() {
370        let config = create_test_config();
371        let llm = OpenAIChat::new(config);
372        let tools: Vec<Arc<dyn BaseTool>> = vec![Arc::new(Calculator::new())];
373
374        let agent = FunctionCallingAgent::new(llm, tools, None);
375
376        assert_eq!(agent.tools.len(), 1);
377        assert!(agent.system_prompt.is_none());
378    }
379
380    #[test]
381    fn test_new_with_system_prompt() {
382        let config = create_test_config();
383        let llm = OpenAIChat::new(config);
384        let tools: Vec<Arc<dyn BaseTool>> = vec![Arc::new(Calculator::new())];
385
386        let agent = FunctionCallingAgent::new(llm, tools, Some("你是一个数学助手".to_string()));
387
388        assert_eq!(agent.system_prompt, Some("你是一个数学助手".to_string()));
389    }
390
391    #[test]
392    fn test_build_messages_empty() {
393        let config = create_test_config();
394        let llm = OpenAIChat::new(config);
395        let tools: Vec<Arc<dyn BaseTool>> = vec![Arc::new(Calculator::new())];
396
397        let agent = FunctionCallingAgent::new(llm, tools, None);
398
399        let mut inputs = HashMap::new();
400        inputs.insert("input".to_string(), "计算 2 + 3".to_string());
401
402        let messages = agent.build_messages(&inputs, &[]);
403
404        assert_eq!(messages.len(), 2);
405        assert_eq!(messages[0].content, "你是一个助手，可以使用工具回答问题。");
406        assert_eq!(messages[1].content, "计算 2 + 3");
407    }
408
409    #[test]
410    fn test_build_messages_with_history() {
411        let config = create_test_config();
412        let llm = OpenAIChat::new(config);
413        let tools: Vec<Arc<dyn BaseTool>> = vec![Arc::new(Calculator::new())];
414
415        let agent = FunctionCallingAgent::new(llm, tools, None);
416
417        let mut inputs = HashMap::new();
418        inputs.insert("input".to_string(), "继续计算".to_string());
419
420        let steps = vec![AgentStep::new(
421            AgentAction {
422                tool: "calculator".to_string(),
423                tool_input: ToolInput::String {
424                    value: "2 + 3".to_string(),
425                },
426                log: "call_123".to_string(),
427            },
428            "5".to_string(),
429        )];
430
431        let messages = agent.build_messages(&inputs, &steps);
432
433        assert_eq!(messages.len(), 4);
434        assert!(messages[2].has_tool_calls());
435    }
436
437    #[test]
438    fn test_from_arc_creation() {
439        let config = create_test_config();
440        let llm = OpenAIChat::new(config);
441        let llm_arc: Arc<dyn BaseChatModel<Error = ProviderError> + Send + Sync> =
442            lc_providers::wrap_chat_model(llm);
443        let tools: Vec<Arc<dyn BaseTool>> = vec![Arc::new(Calculator::new())];
444
445        let agent = FunctionCallingAgent::from_arc(llm_arc, tools, Some("test".into()));
446        assert_eq!(agent.tools.len(), 1);
447        assert_eq!(agent.system_prompt, Some("test".to_string()));
448    }
449
450    /// S2 streaming mock: configurable `stream_chat` returns (normal chunks /
451    /// immediate failure) and `chat` returns (tool call / final answer), and it
452    /// records the call sequence, to verify FunctionCallingAgent's `plan_stream`
453    /// override: token-by-token forwarding, empty-stream fallback to
454    /// non-streaming plan, and immediate-failure fallback.
455    struct MockFuncLLM {
456        stream_chunks: Option<Vec<StreamChunk>>,
457        chat_result: LLMResult,
458        calls: Arc<Mutex<Vec<String>>>,
459    }
460
461    impl MockFuncLLM {
462        fn new(stream_chunks: Option<Vec<StreamChunk>>, chat_result: LLMResult) -> Self {
463            Self {
464                stream_chunks,
465                chat_result,
466                calls: Arc::new(Mutex::new(Vec::new())),
467            }
468        }
469
470        fn calls(&self) -> Vec<String> {
471            self.calls.lock().unwrap_or_else(|e| e.into_inner()).clone()
472        }
473    }
474
475    #[async_trait]
476    impl Runnable<Vec<Message>, LLMResult> for MockFuncLLM {
477        type Error = ProviderError;
478        async fn invoke(
479            &self,
480            input: Vec<Message>,
481            config: Option<RunnableConfig>,
482        ) -> Result<LLMResult, Self::Error> {
483            self.chat(input, config).await
484        }
485    }
486
487    #[async_trait]
488    impl BaseLanguageModel<Vec<Message>, LLMResult> for MockFuncLLM {
489        fn model_name(&self) -> &str {
490            "mock-func"
491        }
492        fn get_num_tokens(&self, t: &str) -> usize {
493            t.len()
494        }
495        fn with_temperature(self, _: f32) -> Self {
496            self
497        }
498        fn with_max_tokens(self, _: usize) -> Self {
499            self
500        }
501    }
502
503    #[async_trait]
504    impl BaseChatModel for MockFuncLLM {
505        async fn chat(
506            &self,
507            _messages: Vec<Message>,
508            _config: Option<RunnableConfig>,
509        ) -> Result<LLMResult, Self::Error> {
510            self.calls
511                .lock()
512                .unwrap_or_else(|e| e.into_inner())
513                .push("chat".to_string());
514            Ok(self.chat_result.clone())
515        }
516        async fn stream_chat(
517            &self,
518            _messages: Vec<Message>,
519            _config: Option<RunnableConfig>,
520        ) -> Result<Pin<Box<dyn Stream<Item = Result<StreamChunk, Self::Error>> + Send>>, Self::Error>
521        {
522            self.calls
523                .lock()
524                .unwrap_or_else(|e| e.into_inner())
525                .push("stream_chat".to_string());
526            match &self.stream_chunks {
527                Some(chunks) => {
528                    let items: Vec<Result<StreamChunk, ProviderError>> =
529                        chunks.iter().cloned().map(Ok).collect();
530                    Ok(Box::pin(futures_util::stream::iter(items)))
531                }
532                None => Err(ProviderError::Assistant(AssistantError::Api(
533                    "stream_chat unavailable".to_string(),
534                ))),
535            }
536        }
537    }
538
539    fn calculator_call_result() -> LLMResult {
540        let call = ToolCall::builder("call_1")
541            .name("calculator")
542            .arguments(r#"{"expression": "2+3"}"#)
543            .build();
544        LLMResult {
545            content: String::new(),
546            model: "mock-func".to_string(),
547            token_usage: Some(TokenUsage {
548                prompt_tokens: 10,
549                completion_tokens: 5,
550                total_tokens: 15,
551            }),
552            tool_calls: Some(vec![call]),
553            thinking_content: None,
554        }
555    }
556
557    fn text_result(content: &str) -> LLMResult {
558        LLMResult {
559            content: content.to_string(),
560            model: "mock-func".to_string(),
561            token_usage: Some(TokenUsage {
562                prompt_tokens: 8,
563                completion_tokens: 6,
564                total_tokens: 14,
565            }),
566            tool_calls: None,
567            thinking_content: None,
568        }
569    }
570
571    fn streaming_agent(llm: MockFuncLLM) -> (FunctionCallingAgent, Arc<MockFuncLLM>) {
572        let arc: Arc<MockFuncLLM> = Arc::new(llm);
573        let agent = FunctionCallingAgent::from_arc(
574            arc.clone() as Arc<dyn BaseChatModel<Error = ProviderError> + Send + Sync>,
575            vec![],
576            None,
577        );
578        (agent, arc)
579    }
580
581    /// S2: the final answer streams out token by token (Text events), and the streaming usage is recorded for the budget gate.
582    #[tokio::test]
583    async fn test_function_calling_plan_stream_streams_final_answer() {
584        let llm = MockFuncLLM::new(
585            Some(vec![
586                StreamChunk::new("Final "),
587                StreamChunk {
588                    text: "Answer: 42".to_string(),
589                    token_usage: Some(TokenUsage {
590                        prompt_tokens: 10,
591                        completion_tokens: 5,
592                        total_tokens: 15,
593                    }),
594                    tool_calls: None,
595                },
596            ]),
597            text_result("unused"),
598        );
599        let (agent, llm) = streaming_agent(llm);
600
601        let mut inputs = HashMap::new();
602        inputs.insert("input".to_string(), "计算 6 * 7".to_string());
603        let mut received = String::new();
604        let mut on_token = |text: String| {
605            received.push_str(&text);
606            Box::pin(async move {}) as Pin<Box<dyn Future<Output = ()> + Send>>
607        };
608
609        let output = agent
610            .plan_stream(&[], &inputs, &mut on_token, None)
611            .await
612            .expect("plan_stream should succeed");
613
614        assert_eq!(received, "Final Answer: 42");
615        assert!(matches!(
616            output,
617            AgentOutput::Finish(f) if f.output() == Some("Final Answer: 42")
618        ));
619        let usage = agent.last_token_usage().expect("streaming usage recorded");
620        assert_eq!(usage.total_tokens, 15);
621        // Took the streaming path: only stream_chat was called, no chat fallback.
622        assert_eq!(llm.calls(), vec!["stream_chat"]);
623    }
624
625    /// S2 + 0.20.0 S3.2: a provider **without** streaming tool_calls support
626    /// yields only empty text + usage chunks on a tool-call step — `plan_stream`
627    /// falls back to non-streaming `plan()` to get the native tool calls, without
628    /// a fake "empty Finish" stream. Providers that DO stream tool_calls take the
629    /// native path (see `test_..._streams_tool_call_natively`).
630    #[tokio::test]
631    async fn test_function_calling_plan_stream_falls_back_when_no_streaming_tool_calls() {
632        // Simulate a tool-call step: the stream only returns empty text + usage chunks.
633        let llm = MockFuncLLM::new(
634            Some(vec![StreamChunk {
635                text: String::new(),
636                token_usage: Some(TokenUsage {
637                    prompt_tokens: 5,
638                    completion_tokens: 0,
639                    total_tokens: 5,
640                }),
641                tool_calls: None,
642            }]),
643            calculator_call_result(),
644        );
645        let (agent, llm) = streaming_agent(llm);
646
647        let mut inputs = HashMap::new();
648        inputs.insert("input".to_string(), "计算 2 + 3".to_string());
649        let mut emitted: Vec<String> = Vec::new();
650        let mut on_token = |text: String| {
651            emitted.push(text);
652            Box::pin(async move {}) as Pin<Box<dyn Future<Output = ()> + Send>>
653        };
654
655        let output = agent
656            .plan_stream(&[], &inputs, &mut on_token, None)
657            .await
658            .expect("plan_stream should succeed");
659
660        assert!(
661            matches!(&output, AgentOutput::Action(a) if a.tool == "calculator"),
662            "tool-call step must return Action"
663        );
664        assert!(emitted.is_empty(), "tool-call step emits no free text");
665        // The fallback path's usage comes from non-streaming plan() (the budget gate still gets real usage).
666        let usage = agent.last_token_usage().expect("usage via fallback plan");
667        assert_eq!(usage.total_tokens, 15);
668        // Tool-call step: start the stream (empty text) first, then fall back to non-streaming chat for native tool_calls.
669        assert_eq!(llm.calls(), vec!["stream_chat", "chat"]);
670    }
671
672    /// S2: when `stream_chat` fails immediately, `plan_stream` falls back to
673    /// non-streaming `plan()`, forwarding the final answer as a single Text event
674    /// (consistent with the old non-streaming path).
675    #[tokio::test]
676    async fn test_function_calling_plan_stream_falls_back_on_immediate_error() {
677        let llm = MockFuncLLM::new(None, text_result("Final Answer: 42"));
678        let (agent, llm) = streaming_agent(llm);
679
680        let mut inputs = HashMap::new();
681        inputs.insert("input".to_string(), "计算 6 * 7".to_string());
682        let mut received = String::new();
683        let mut on_token = |text: String| {
684            received.push_str(&text);
685            Box::pin(async move {}) as Pin<Box<dyn Future<Output = ()> + Send>>
686        };
687
688        let output = agent
689            .plan_stream(&[], &inputs, &mut on_token, None)
690            .await
691            .expect("fallback plan should succeed");
692
693        assert_eq!(received, "Final Answer: 42");
694        assert!(matches!(output, AgentOutput::Finish(_)));
695        // stream_chat fails immediately → fall back to non-streaming chat, the whole answer forwarded as a single Text event.
696        assert_eq!(llm.calls(), vec!["stream_chat", "chat"]);
697    }
698
699    /// 0.20.0 S3.2: a provider WITH streaming tool_calls support lets a pure
700    /// tool-call step stream natively — the terminal chunk carries the tool calls,
701    /// so `plan_stream` returns Action without a non-streaming fallback (no `chat`
702    /// call at all).
703    #[tokio::test]
704    async fn test_function_calling_plan_stream_streams_tool_call_natively() {
705        let tool_chunk = StreamChunk {
706            text: String::new(),
707            token_usage: Some(TokenUsage {
708                prompt_tokens: 5,
709                completion_tokens: 0,
710                total_tokens: 5,
711            }),
712            tool_calls: Some(vec![ToolCall::builder("call_1")
713                .name("calculator")
714                .arguments(r#"{"expression": "2+3"}"#)
715                .build()]),
716        };
717        let llm = MockFuncLLM::new(Some(vec![tool_chunk]), text_result("unused"));
718        let (agent, llm) = streaming_agent(llm);
719
720        let mut inputs = HashMap::new();
721        inputs.insert("input".to_string(), "计算 2 + 3".to_string());
722        let mut emitted: Vec<String> = Vec::new();
723        let mut on_token = |text: String| {
724            emitted.push(text);
725            Box::pin(async move {}) as Pin<Box<dyn Future<Output = ()> + Send>>
726        };
727
728        let output = agent
729            .plan_stream(&[], &inputs, &mut on_token, None)
730            .await
731            .expect("plan_stream should succeed");
732
733        assert!(
734            matches!(&output, AgentOutput::Action(a) if a.tool == "calculator"),
735            "tool-call step must return Action natively"
736        );
737        assert!(emitted.is_empty(), "no free text on a pure tool-call step");
738        // Native path: only stream_chat ran — no non-streaming fallback.
739        assert_eq!(llm.calls(), vec!["stream_chat"]);
740    }
741
742    /// 0.20.0 S3.2: a mixed step (model emits text AND a tool call in one stream)
743    /// keeps both — the text streams out token by token via on_token, and the
744    /// tool call is returned as an Action (previously the tool call was silently
745    /// dropped).
746    #[tokio::test]
747    async fn test_function_calling_plan_stream_mixed_step_keeps_text_and_tool_call() {
748        let llm = MockFuncLLM::new(
749            Some(vec![
750                StreamChunk::new("Let me compute"),
751                StreamChunk {
752                    text: String::new(),
753                    token_usage: Some(TokenUsage {
754                        prompt_tokens: 5,
755                        completion_tokens: 0,
756                        total_tokens: 5,
757                    }),
758                    tool_calls: Some(vec![ToolCall::builder("call_1")
759                        .name("calculator")
760                        .arguments(r#"{"expression": "2+3"}"#)
761                        .build()]),
762                },
763            ]),
764            text_result("unused"),
765        );
766        let (agent, llm) = streaming_agent(llm);
767
768        let mut inputs = HashMap::new();
769        inputs.insert("input".to_string(), "计算 2 + 3".to_string());
770        let mut emitted: Vec<String> = Vec::new();
771        let mut on_token = |text: String| {
772            emitted.push(text);
773            Box::pin(async move {}) as Pin<Box<dyn Future<Output = ()> + Send>>
774        };
775
776        let output = agent
777            .plan_stream(&[], &inputs, &mut on_token, None)
778            .await
779            .expect("plan_stream should succeed");
780
781        assert_eq!(emitted, vec!["Let me compute"], "preamble streams out");
782        assert!(
783            matches!(&output, AgentOutput::Action(a) if a.tool == "calculator"),
784            "tool call preserved, not dropped"
785        );
786        // Native path throughout — no non-streaming fallback.
787        assert_eq!(llm.calls(), vec!["stream_chat"]);
788    }
789}
lc_agents/function_calling/agent.rs

lc_agents/function_calling/
agent.rs
