agent_base/engine/
tool_enforcement.rs

1use std::sync::atomic::{AtomicUsize, Ordering};
2
3use async_trait::async_trait;
4
5use crate::engine::middleware::{Middleware, PostLlmCtx};
6use crate::types::AgentResult;
7
8pub struct ToolEnforcementConfig {
9    pub max_nudges: usize,
10    pub nudge_message: String,
11    pub first_turn_only: bool,
12    pub min_tools_threshold: usize,
13}
14
15impl Default for ToolEnforcementConfig {
16    fn default() -> Self {
17        Self {
18            max_nudges: 3,
19            nudge_message: "CRITICAL: You have tools available but did not call any. \
20                             Call the appropriate tool NOW. \
21                             关键提示：你有可用的工具但没有调用。立即使用工具执行。"
22                .to_string(),
23            first_turn_only: true,
24            min_tools_threshold: 1,
25        }
26    }
27}
28
29pub struct ToolEnforcementMiddleware {
30    config: ToolEnforcementConfig,
31    nudge_count: AtomicUsize,
32}
33
34impl ToolEnforcementMiddleware {
35    pub fn new(config: ToolEnforcementConfig) -> Self {
36        Self {
37            config,
38            nudge_count: AtomicUsize::new(0),
39        }
40    }
41}
42
43#[async_trait]
44impl Middleware for ToolEnforcementMiddleware {
45    async fn on_post_llm(&self, ctx: &mut PostLlmCtx) -> AgentResult<()> {
46        if ctx.available_tools.len() < self.config.min_tools_threshold {
47            return Ok(());
48        }
49        if self.config.first_turn_only && ctx.total_tool_calls > 0 {
50            return Ok(());
51        }
52        if ctx.is_tool_call {
53            return Ok(());
54        }
55        if ctx.full_text.is_empty() {
56            return Ok(());
57        }
58
59        let count = self.nudge_count.fetch_add(1, Ordering::SeqCst);
60        if count >= self.config.max_nudges {
61            return Ok(());
62        }
63
64        tracing::info!(
65            nudge_count = count,
66            full_text_len = ctx.full_text.len(),
67            "ToolEnforcement: suppressing text response, injecting nudge"
68        );
69
70        ctx.skip_push = true;
71        ctx.follow_up_message = Some(self.config.nudge_message.clone());
72
73        Ok(())
74    }
75}
agent_base/engine/tool_enforcement.rs

agent_base/engine/
tool_enforcement.rs
