agent_diva_nano/
tool_assembly.rs

1//! Thin wrapper around the shared agent-diva tool assembly.
2
3use agent_diva_agent::tool_config::network::{
4    NetworkToolConfig, WebFetchRuntimeConfig, WebRuntimeConfig, WebSearchRuntimeConfig,
5};
6pub use agent_diva_agent::{BuiltInToolsConfig, SubagentSpawner};
7use agent_diva_core::config::MCPServerConfig;
8#[cfg(feature = "files")]
9use agent_diva_files::FileManager;
10use agent_diva_tooling::{Tool, ToolRegistry};
11use std::collections::HashMap;
12use std::path::PathBuf;
13use std::sync::Arc;
14
15/// Shell tool configuration.
16#[derive(Debug, Clone)]
17pub struct ShellToolConfig {
18    pub timeout_secs: u64,
19    pub working_dir: Option<PathBuf>,
20    pub restrict_to_workspace: bool,
21}
22
23impl Default for ShellToolConfig {
24    fn default() -> Self {
25        Self {
26            timeout_secs: 60,
27            working_dir: None,
28            restrict_to_workspace: true,
29        }
30    }
31}
32
33/// Web tool configuration.
34#[derive(Debug, Clone)]
35pub struct WebToolConfig {
36    pub search_enabled: bool,
37    pub fetch_enabled: bool,
38    pub search_provider: String,
39    pub search_api_key: Option<String>,
40    pub max_results: u32,
41}
42
43impl Default for WebToolConfig {
44    fn default() -> Self {
45        Self {
46            search_enabled: true,
47            fetch_enabled: true,
48            search_provider: "bocha".to_string(),
49            search_api_key: None,
50            max_results: 5,
51        }
52    }
53}
54
55/// Builder for assembling a ToolRegistry with fine-grained control.
56pub struct ToolAssembly {
57    workspace: PathBuf,
58    builtin_config: BuiltInToolsConfig,
59    shell_config: ShellToolConfig,
60    web_config: WebToolConfig,
61    custom_tools: Vec<Arc<dyn Tool>>,
62    mcp_servers: HashMap<String, MCPServerConfig>,
63    subagent_spawner: Option<Arc<dyn SubagentSpawner>>,
64    #[cfg(feature = "files")]
65    file_manager: Option<Arc<FileManager>>,
66}
67
68impl ToolAssembly {
69    pub fn new(workspace: PathBuf) -> Self {
70        Self {
71            workspace,
72            builtin_config: BuiltInToolsConfig::default(),
73            shell_config: ShellToolConfig::default(),
74            web_config: WebToolConfig::default(),
75            custom_tools: Vec::new(),
76            mcp_servers: HashMap::new(),
77            subagent_spawner: None,
78            #[cfg(feature = "files")]
79            file_manager: None,
80        }
81    }
82
83    #[cfg(feature = "files")]
84    pub fn with_file_manager(mut self, manager: Arc<FileManager>) -> Self {
85        self.file_manager = Some(manager);
86        self
87    }
88
89    pub fn builtin(mut self, config: BuiltInToolsConfig) -> Self {
90        self.builtin_config = config;
91        self
92    }
93
94    pub fn filesystem(mut self, enabled: bool) -> Self {
95        self.builtin_config.filesystem = enabled;
96        self
97    }
98
99    pub fn shell(mut self, enabled: bool) -> Self {
100        self.builtin_config.shell = enabled;
101        self
102    }
103
104    pub fn shell_config(mut self, config: ShellToolConfig) -> Self {
105        self.shell_config = config;
106        self.builtin_config.shell = true;
107        self
108    }
109
110    pub fn web(mut self, enabled: bool) -> Self {
111        self.builtin_config.web_search = enabled;
112        self.builtin_config.web_fetch = enabled;
113        self
114    }
115
116    pub fn web_config(mut self, config: WebToolConfig) -> Self {
117        self.web_config = config;
118        self.builtin_config.web_search = true;
119        self.builtin_config.web_fetch = true;
120        self
121    }
122
123    pub fn spawn(mut self, enabled: bool) -> Self {
124        self.builtin_config.spawn = enabled;
125        self
126    }
127
128    pub fn with_subagent_spawner(mut self, spawner: Arc<dyn SubagentSpawner>) -> Self {
129        self.subagent_spawner = Some(spawner);
130        self.builtin_config.spawn = true;
131        self
132    }
133
134    pub fn cron(mut self, enabled: bool) -> Self {
135        self.builtin_config.cron = enabled;
136        self
137    }
138
139    pub fn mcp(mut self, enabled: bool) -> Self {
140        self.builtin_config.mcp = enabled;
141        self
142    }
143
144    pub fn add_mcp_server(mut self, name: String, config: MCPServerConfig) -> Self {
145        self.mcp_servers.insert(name, config);
146        self.builtin_config.mcp = true;
147        self
148    }
149
150    pub fn mcp_servers(mut self, servers: HashMap<String, MCPServerConfig>) -> Self {
151        self.mcp_servers = servers;
152        self
153    }
154
155    pub fn attachment(mut self, enabled: bool) -> Self {
156        self.builtin_config.attachment = enabled;
157        self
158    }
159
160    pub fn restrict_to_workspace(mut self, restrict: bool) -> Self {
161        self.shell_config.restrict_to_workspace = restrict;
162        self
163    }
164
165    pub fn with_tool(mut self, tool: Arc<dyn Tool>) -> Self {
166        self.custom_tools.push(tool);
167        self
168    }
169
170    pub fn with_tools(mut self, tools: Vec<Arc<dyn Tool>>) -> Self {
171        self.custom_tools.extend(tools);
172        self
173    }
174
175    pub fn build(self) -> ToolRegistry {
176        self.into_shared().build()
177    }
178
179    fn into_shared(self) -> agent_diva_agent::ToolAssembly {
180        let network = NetworkToolConfig {
181            web: WebRuntimeConfig {
182                search: WebSearchRuntimeConfig {
183                    provider: self.web_config.search_provider,
184                    enabled: self.web_config.search_enabled,
185                    api_key: self.web_config.search_api_key,
186                    max_results: self.web_config.max_results,
187                },
188                fetch: WebFetchRuntimeConfig {
189                    enabled: self.web_config.fetch_enabled,
190                },
191            },
192        };
193
194        let mut assembly = agent_diva_agent::ToolAssembly::new(self.workspace)
195            .builtin(self.builtin_config)
196            .with_network_config(network)
197            .with_exec_timeout(self.shell_config.timeout_secs)
198            .restrict_to_workspace(self.shell_config.restrict_to_workspace)
199            .mcp_servers(self.mcp_servers)
200            .with_tools(self.custom_tools);
201        if let Some(spawner) = self.subagent_spawner {
202            assembly = assembly.with_subagent_spawner(spawner);
203        }
204        #[cfg(feature = "files")]
205        if let Some(file_manager) = self.file_manager {
206            assembly = assembly.with_file_manager(file_manager);
207        }
208        assembly
209    }
210}
agent_diva_nano/tool_assembly.rs

agent_diva_nano/
tool_assembly.rs
