hehe_tools/
lib.rs

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}