custom_tool/
custom_tool.rs

1use async_trait::async_trait;
2use helios_engine::{Agent, Config, Tool, ToolParameter, ToolResult};
3use serde_json::Value;
4use std::collections::HashMap;
5
6// Define a custom tool
7struct WeatherTool;
8
9#[async_trait]
10impl Tool for WeatherTool {
11    fn name(&self) -> &str {
12        "get_weather"
13    }
14
15    fn description(&self) -> &str {
16        "Get the current weather for a location"
17    }
18
19    fn parameters(&self) -> HashMap<String, ToolParameter> {
20        let mut params = HashMap::new();
21        params.insert(
22            "location".to_string(),
23            ToolParameter {
24                param_type: "string".to_string(),
25                description: "The city and state, e.g. San Francisco, CA".to_string(),
26                required: Some(true),
27            },
28        );
29        params.insert(
30            "unit".to_string(),
31            ToolParameter {
32                param_type: "string".to_string(),
33                description: "Temperature unit: 'celsius' or 'fahrenheit'".to_string(),
34                required: Some(false),
35            },
36        );
37        params
38    }
39
40    async fn execute(&self, args: Value) -> helios_engine::Result<ToolResult> {
41        let location = args
42            .get("location")
43            .and_then(|v| v.as_str())
44            .unwrap_or("Unknown");
45        
46        let unit = args
47            .get("unit")
48            .and_then(|v| v.as_str())
49            .unwrap_or("fahrenheit");
50
51        // Simulate weather data (in a real implementation, call a weather API)
52        let temp = if unit == "celsius" { "22" } else { "72" };
53        let weather = format!(
54            "The weather in {} is sunny with a temperature of {}°{}",
55            location,
56            temp,
57            if unit == "celsius" { "C" } else { "F" }
58        );
59
60        Ok(ToolResult::success(weather))
61    }
62}
63
64#[tokio::main]
65async fn main() -> helios_engine::Result<()> {
66    // Load configuration
67    let config = Config::from_file("config.toml")?;
68
69    // Create an agent with custom tool
70    let mut agent = Agent::builder("WeatherAgent")
71        .config(config)
72        .system_prompt("You are a helpful weather assistant. Use the weather tool to answer questions about weather.")
73        .tool(Box::new(WeatherTool))
74        .build()?;
75
76    // Ask about weather
77    let response = agent.chat("What's the weather like in New York?").await?;
78    println!("Agent: {}\n", response);
79
80    let response = agent.chat("How about in London, but in celsius?").await?;
81    println!("Agent: {}\n", response);
82
83    Ok(())
84}