awaken_ext_deferred_tools/
config.rs

1//! Configuration types for deferred tool loading.
2
3use std::collections::HashMap;
4
5use awaken_contract::PluginConfigKey;
6use awaken_tool_pattern::wildcard_match;
7use schemars::JsonSchema;
8use serde::{Deserialize, Serialize};
9
10/// Loading mode for a tool.
11#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
12#[serde(rename_all = "snake_case")]
13pub enum ToolLoadMode {
14    Eager,
15    #[default]
16    Deferred,
17}
18
19/// A single rule mapping a tool name pattern to a load mode.
20#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
21pub struct DeferralRule {
22    pub tool: String,
23    pub mode: ToolLoadMode,
24}
25
26/// Parameters for the discounted Beta probability model.
27#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
28#[serde(default)]
29pub struct DiscBetaParams {
30    /// Discount factor per turn. Effective memory ≈ 1/(1-ω) turns.
31    pub omega: f64,
32    /// Prior strength in equivalent observations.
33    pub n0: f64,
34    /// Consecutive idle turns before re-defer is considered.
35    pub defer_after: u64,
36    /// Defer threshold multiplier against breakeven frequency.
37    pub thresh_mult: f64,
38    /// Estimated ToolSearch call cost in tokens (for breakeven calculation).
39    pub gamma: f64,
40}
41
42impl Default for DiscBetaParams {
43    fn default() -> Self {
44        Self {
45            omega: 0.95,
46            n0: 5.0,
47            defer_after: 5,
48            thresh_mult: 0.5,
49            gamma: 2000.0,
50        }
51    }
52}
53
54/// Top-level configuration for the deferred tools plugin.
55#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
56#[serde(default)]
57pub struct DeferredToolsConfig {
58    /// `None` = auto-enable when total schema savings > beta overhead.
59    /// `Some(true)` = always enable. `Some(false)` = always disable.
60    pub enabled: Option<bool>,
61    /// Ordered rules — first match wins.
62    pub rules: Vec<DeferralRule>,
63    /// Mode for tools not matching any rule.
64    pub default_mode: ToolLoadMode,
65    /// Per-turn overhead of ToolSearch schema + deferred list prompt (tokens).
66    pub beta_overhead: f64,
67    /// Agent-level prior per-turn presence frequencies per tool (from historical data).
68    #[serde(default, skip_serializing_if = "HashMap::is_empty")]
69    pub agent_priors: HashMap<String, f64>,
70    /// Discounted Beta model parameters.
71    pub disc_beta: DiscBetaParams,
72}
73
74impl Default for DeferredToolsConfig {
75    fn default() -> Self {
76        Self {
77            enabled: None,
78            rules: Vec::new(),
79            default_mode: ToolLoadMode::default(),
80            beta_overhead: 1136.0,
81            agent_priors: HashMap::new(),
82            disc_beta: DiscBetaParams::default(),
83        }
84    }
85}
86
87impl DeferredToolsConfig {
88    /// Resolve the load mode for a tool by matching rules in order.
89    pub fn resolve_mode(&self, tool_id: &str) -> ToolLoadMode {
90        for rule in &self.rules {
91            if rule.tool == tool_id || wildcard_match(&rule.tool, tool_id) {
92                return rule.mode;
93            }
94        }
95        self.default_mode
96    }
97
98    /// Check if deferred tools should be enabled.
99    /// `total_savings` is sum of (c_i - c_bar_i) for all deferrable tools.
100    pub fn should_enable(&self, total_savings: f64) -> bool {
101        match self.enabled {
102            Some(forced) => forced,
103            None => total_savings > self.beta_overhead,
104        }
105    }
106
107    /// Get the prior p_i for a tool from agent-level stats.
108    pub fn prior_p(&self, tool_id: &str) -> f64 {
109        self.agent_priors.get(tool_id).copied().unwrap_or(0.01)
110    }
111}
112
113/// [`PluginConfigKey`] binding for deferred tools config in agent specs.
114pub struct DeferredToolsConfigKey;
115
116impl PluginConfigKey for DeferredToolsConfigKey {
117    const KEY: &'static str = "deferred_tools";
118    type Config = DeferredToolsConfig;
119}
awaken_ext_deferred_tools/config.rs

awaken_ext_deferred_tools/
config.rs
