1use serde::{Deserialize, Serialize};
19use std::collections::{HashMap, HashSet};
20
21use super::microagent::SubtaskOutput;
22
23#[derive(Clone, Debug, Serialize, Deserialize)]
28pub struct ToolSchema {
29 pub name: String,
31 pub description: String,
33 #[serde(default)]
35 pub parameters: HashMap<String, String>,
36 #[serde(default)]
38 pub required: Vec<String>,
39 #[serde(skip_serializing_if = "Option::is_none")]
41 pub category: Option<ToolCategory>,
42}
43
44impl From<brainwires_core::Tool> for ToolSchema {
45 fn from(tool: brainwires_core::Tool) -> Self {
46 let mut schema = Self::new(&tool.name, &tool.description);
47
48 if let Some(props) = &tool.input_schema.properties {
50 for (name, value) in props {
51 let desc = value
52 .get("description")
53 .and_then(|v| v.as_str())
54 .unwrap_or("No description")
55 .to_string();
56 schema.parameters.insert(name.clone(), desc);
57 }
58 }
59
60 if let Some(required) = &tool.input_schema.required {
61 schema.required = required.clone();
62 }
63
64 schema
65 }
66}
67
68impl ToolSchema {
69 pub fn new(name: impl Into<String>, description: impl Into<String>) -> Self {
71 Self {
72 name: name.into(),
73 description: description.into(),
74 parameters: HashMap::new(),
75 required: Vec::new(),
76 category: None,
77 }
78 }
79
80 pub fn with_param(mut self, name: impl Into<String>, description: impl Into<String>) -> Self {
82 self.parameters.insert(name.into(), description.into());
83 self
84 }
85
86 pub fn with_required_param(
88 mut self,
89 name: impl Into<String>,
90 description: impl Into<String>,
91 ) -> Self {
92 let name = name.into();
93 self.parameters.insert(name.clone(), description.into());
94 self.required.push(name);
95 self
96 }
97
98 pub fn with_category(mut self, category: ToolCategory) -> Self {
100 self.category = Some(category);
101 self
102 }
103
104 pub fn to_prompt_format(&self) -> String {
106 let mut result = format!("- **{}**: {}\n", self.name, self.description);
107 if !self.parameters.is_empty() {
108 result.push_str(" Parameters:\n");
109 for (name, desc) in &self.parameters {
110 let required = if self.required.contains(name) {
111 " (required)"
112 } else {
113 ""
114 };
115 result.push_str(&format!(" - {}{}: {}\n", name, required, desc));
116 }
117 }
118 result
119 }
120}
121
122#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
127pub struct ToolIntent {
128 pub tool_name: String,
130 #[serde(default)]
132 pub arguments: serde_json::Value,
133 #[serde(skip_serializing_if = "Option::is_none")]
135 pub rationale: Option<String>,
136}
137
138impl ToolIntent {
139 pub fn new(tool_name: impl Into<String>, arguments: serde_json::Value) -> Self {
141 Self {
142 tool_name: tool_name.into(),
143 arguments,
144 rationale: None,
145 }
146 }
147
148 pub fn with_rationale(mut self, rationale: impl Into<String>) -> Self {
150 self.rationale = Some(rationale.into());
151 self
152 }
153
154 pub fn matches_category(&self, category: &ToolCategory) -> bool {
156 category.contains_tool(&self.tool_name)
157 }
158}
159
160#[derive(Clone, Debug, Serialize, Deserialize)]
165pub struct SubtaskOutputWithIntent {
166 pub output: SubtaskOutput,
168 #[serde(skip_serializing_if = "Option::is_none")]
170 pub tool_intent: Option<ToolIntent>,
171 #[serde(default)]
173 pub awaiting_tool_result: bool,
174}
175
176impl SubtaskOutputWithIntent {
177 pub fn from_output(output: SubtaskOutput) -> Self {
179 Self {
180 output,
181 tool_intent: None,
182 awaiting_tool_result: false,
183 }
184 }
185
186 pub fn with_tool_intent(output: SubtaskOutput, intent: ToolIntent) -> Self {
188 Self {
189 output,
190 tool_intent: Some(intent),
191 awaiting_tool_result: true,
192 }
193 }
194
195 pub fn has_tool_intent(&self) -> bool {
197 self.tool_intent.is_some()
198 }
199
200 pub fn mark_tool_complete(mut self) -> Self {
202 self.awaiting_tool_result = false;
203 self
204 }
205}
206
207#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
212pub enum ToolCategory {
213 FileRead,
215 FileWrite,
217 Search,
219 SemanticSearch,
221 Bash,
223 Git,
225 Web,
227 AgentPool,
229 TaskManager,
231 Mcp,
233 Custom(String),
235}
236
237impl ToolCategory {
238 pub fn contains_tool(&self, tool_name: &str) -> bool {
240 match self {
241 ToolCategory::FileRead => {
242 matches!(tool_name, "read_file" | "file_read" | "get_file_contents")
243 }
244 ToolCategory::FileWrite => matches!(
245 tool_name,
246 "write_file" | "edit_file" | "delete_file" | "create_directory" | "file_write"
247 ),
248 ToolCategory::Search => matches!(
249 tool_name,
250 "search_files" | "grep" | "find_files" | "glob" | "file_search"
251 ),
252 ToolCategory::SemanticSearch => matches!(
253 tool_name,
254 "semantic_search" | "query_codebase" | "rag_search"
255 ),
256 ToolCategory::Bash => matches!(
257 tool_name,
258 "bash" | "execute_command" | "shell" | "run_command"
259 ),
260 ToolCategory::Git => matches!(
261 tool_name,
262 "git" | "git_status" | "git_diff" | "git_commit" | "git_log"
263 ),
264 ToolCategory::Web => matches!(
265 tool_name,
266 "web_search" | "fetch_url" | "browse" | "http_request"
267 ),
268 ToolCategory::AgentPool => {
269 matches!(tool_name, "spawn_agent" | "agent_pool" | "create_agent")
270 }
271 ToolCategory::TaskManager => {
272 matches!(tool_name, "create_task" | "update_task" | "task_manager")
273 }
274 ToolCategory::Mcp => tool_name.starts_with("mcp_") || tool_name.starts_with("mcp__"),
275 ToolCategory::Custom(prefix) => tool_name.starts_with(prefix),
276 }
277 }
278
279 pub fn read_only_categories() -> HashSet<ToolCategory> {
281 HashSet::from([
282 ToolCategory::FileRead,
283 ToolCategory::Search,
284 ToolCategory::SemanticSearch,
285 ])
286 }
287
288 pub fn side_effect_categories() -> HashSet<ToolCategory> {
290 HashSet::from([
291 ToolCategory::FileWrite,
292 ToolCategory::Bash,
293 ToolCategory::Git,
294 ToolCategory::Web,
295 ToolCategory::AgentPool,
296 ToolCategory::TaskManager,
297 ])
298 }
299}
300
301#[derive(Clone, Debug)]
303pub enum IntentParseResult {
304 NoIntent(SubtaskOutput),
306 WithIntent(SubtaskOutputWithIntent),
308 ParseError(String),
310}
311
312pub fn parse_tool_intent(subtask_id: &str, response_text: &str) -> IntentParseResult {
316 if let Some(intent) = extract_tool_intent_json(response_text) {
318 match serde_json::from_value::<ToolIntent>(intent.clone()) {
319 Ok(tool_intent) => {
320 let output_text = remove_json_block(response_text);
322 let output = SubtaskOutput::new(
323 subtask_id,
324 serde_json::json!({
325 "text": output_text.trim(),
326 "awaiting_tool": true,
327 }),
328 );
329 IntentParseResult::WithIntent(SubtaskOutputWithIntent::with_tool_intent(
330 output,
331 tool_intent,
332 ))
333 }
334 Err(e) => IntentParseResult::ParseError(format!("Failed to parse tool intent: {}", e)),
335 }
336 } else {
337 let output = SubtaskOutput::new(subtask_id, serde_json::json!({ "text": response_text }));
339 IntentParseResult::NoIntent(output)
340 }
341}
342
343fn extract_tool_intent_json(text: &str) -> Option<serde_json::Value> {
345 if let Some(json_block) = extract_json_code_block(text)
347 && let Ok(value) = serde_json::from_str::<serde_json::Value>(&json_block)
348 {
349 if value.get("tool_intent").is_some() {
350 return value.get("tool_intent").cloned();
351 }
352 if value.get("tool_name").is_some() {
354 return Some(value);
355 }
356 }
357
358 for line in text.lines() {
360 let trimmed = line.trim();
361 if trimmed.starts_with('{')
362 && trimmed.ends_with('}')
363 && let Ok(value) = serde_json::from_str::<serde_json::Value>(trimmed)
364 {
365 if value.get("tool_intent").is_some() {
366 return value.get("tool_intent").cloned();
367 }
368 if value.get("tool_name").is_some() {
369 return Some(value);
370 }
371 }
372 }
373
374 None
375}
376
377fn extract_json_code_block(text: &str) -> Option<String> {
379 let start_markers = ["```json", "```JSON"];
380 let end_marker = "```";
381
382 for start in start_markers {
383 if let Some(start_idx) = text.find(start) {
384 let content_start = start_idx + start.len();
385 if let Some(end_idx) = text[content_start..].find(end_marker) {
386 return Some(
387 text[content_start..content_start + end_idx]
388 .trim()
389 .to_string(),
390 );
391 }
392 }
393 }
394
395 None
396}
397
398fn remove_json_block(text: &str) -> String {
400 let start_markers = ["```json", "```JSON"];
401 let end_marker = "```";
402
403 let mut result = text.to_string();
404
405 for start in start_markers {
406 if let Some(start_idx) = result.find(start) {
407 let content_start = start_idx + start.len();
408 if let Some(end_idx) = result[content_start..].find(end_marker) {
409 let block_end = content_start + end_idx + end_marker.len();
410 result = format!("{}{}", &result[..start_idx], &result[block_end..]);
411 }
412 }
413 }
414
415 result
416}
417
418#[cfg(test)]
419mod tests {
420 use super::*;
421
422 #[test]
423 fn test_tool_intent_creation() {
424 let intent = ToolIntent::new("read_file", serde_json::json!({"path": "/test.txt"}))
425 .with_rationale("Need to read configuration");
426
427 assert_eq!(intent.tool_name, "read_file");
428 assert_eq!(
429 intent.rationale,
430 Some("Need to read configuration".to_string())
431 );
432 }
433
434 #[test]
435 fn test_tool_category_matching() {
436 assert!(ToolCategory::FileRead.contains_tool("read_file"));
437 assert!(ToolCategory::FileWrite.contains_tool("write_file"));
438 assert!(ToolCategory::Search.contains_tool("grep"));
439 assert!(ToolCategory::Mcp.contains_tool("mcp__brainwires-rag__query"));
440 assert!(!ToolCategory::FileRead.contains_tool("bash"));
441 }
442
443 #[test]
444 fn test_parse_tool_intent_with_json_block() {
445 let response = r#"I need to read a file first.
446
447```json
448{
449 "tool_name": "read_file",
450 "arguments": {"path": "/test.txt"},
451 "rationale": "Check contents"
452}
453```
454"#;
455
456 match parse_tool_intent("task-1", response) {
457 IntentParseResult::WithIntent(output) => {
458 assert!(output.has_tool_intent());
459 let intent = output.tool_intent.unwrap();
460 assert_eq!(intent.tool_name, "read_file");
461 }
462 _ => panic!("Expected WithIntent result"),
463 }
464 }
465
466 #[test]
467 fn test_parse_no_intent() {
468 let response = "This is just a regular response without any tool calls.";
469
470 match parse_tool_intent("task-1", response) {
471 IntentParseResult::NoIntent(output) => {
472 assert_eq!(output.subtask_id, "task-1");
473 }
474 _ => panic!("Expected NoIntent result"),
475 }
476 }
477
478 #[test]
479 fn test_read_only_categories() {
480 let read_only = ToolCategory::read_only_categories();
481 assert!(read_only.contains(&ToolCategory::FileRead));
482 assert!(read_only.contains(&ToolCategory::Search));
483 assert!(!read_only.contains(&ToolCategory::FileWrite));
484 assert!(!read_only.contains(&ToolCategory::Bash));
485 }
486}