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

lc_agents/function_calling/
agent.rs
