agent_base/engine/
tool_enforcement.rs

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

agent_base/engine/
tool_enforcement.rs
