1pub mod error;
2pub mod traits;
3pub mod registry;
4pub mod executor;
5#[cfg(feature = "builtin")]
6pub mod builtin;
7pub mod sandbox;
8
9pub use error::{Result, ToolError};
10pub use traits::{Artifact, ArtifactData, Tool, ToolOutput};
11pub use registry::ToolRegistry;
12pub use executor::ToolExecutor;
13pub use sandbox::{NativeSandbox, Sandbox, SandboxConfig};
14
15#[cfg(feature = "builtin")]
16pub use builtin::{
17 create_default_registry, register_all,
18 ListDirectoryTool, ReadFileTool, SearchFilesTool, WriteFileTool,
19 GetSystemInfoTool,
20};
21
22#[cfg(all(feature = "builtin", feature = "shell"))]
23pub use builtin::ExecuteShellTool;
24
25#[cfg(all(feature = "builtin", feature = "http"))]
26pub use builtin::HttpRequestTool;
27
28pub mod prelude {
29 pub use crate::error::{Result, ToolError};
30 pub use crate::traits::{Artifact, ArtifactData, Tool, ToolOutput};
31 pub use crate::registry::ToolRegistry;
32 pub use crate::executor::ToolExecutor;
33 pub use crate::sandbox::{NativeSandbox, Sandbox, SandboxConfig};
34
35 #[cfg(feature = "builtin")]
36 pub use crate::builtin::{
37 create_default_registry, register_all,
38 ListDirectoryTool, ReadFileTool, SearchFilesTool, WriteFileTool,
39 GetSystemInfoTool,
40 };
41
42 #[cfg(all(feature = "builtin", feature = "shell"))]
43 pub use crate::builtin::ExecuteShellTool;
44
45 #[cfg(all(feature = "builtin", feature = "http"))]
46 pub use crate::builtin::HttpRequestTool;
47}
48
49#[cfg(test)]
50mod tests {
51 use super::*;
52 use std::sync::Arc;
53
54 #[cfg(feature = "builtin")]
55 #[test]
56 fn test_create_default_registry() {
57 let registry = create_default_registry();
58 assert!(registry.contains("read_file"));
59 assert!(registry.contains("write_file"));
60 assert!(registry.contains("list_directory"));
61 assert!(registry.contains("search_files"));
62 assert!(registry.contains("get_system_info"));
63 }
64
65 #[cfg(all(feature = "builtin", feature = "shell"))]
66 #[test]
67 fn test_registry_with_shell() {
68 let registry = create_default_registry();
69 assert!(registry.contains("execute_shell"));
70 assert!(registry.dangerous_tools().contains(&"execute_shell"));
71 }
72
73 #[cfg(feature = "builtin")]
74 #[tokio::test]
75 async fn test_executor_with_builtin() {
76 use hehe_core::Context;
77
78 let registry = Arc::new(create_default_registry());
79 let executor = ToolExecutor::new(registry);
80 let ctx = Context::new();
81
82 let output = executor
83 .execute(&ctx, "get_system_info", serde_json::Value::Null)
84 .await
85 .unwrap();
86
87 assert!(!output.is_error);
88 assert!(output.content.contains("os"));
89 }
90}