diff --git a/Cargo.lock b/Cargo.lock index 697d048898..6fa1bbedd2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2710,6 +2710,7 @@ dependencies = [ "regex", "rmcp", "rustyline", + "schemars", "serde", "serde_json", "serde_yaml", diff --git a/crates/goose-cli/Cargo.toml b/crates/goose-cli/Cargo.toml index 83ebb8644a..ffa535df3c 100644 --- a/crates/goose-cli/Cargo.toml +++ b/crates/goose-cli/Cargo.toml @@ -59,6 +59,7 @@ is-terminal = "0.4.16" anstream = "0.6.18" url = "2.5.7" open = "5.3.2" +schemars = "1.0" [target.'cfg(target_os = "windows")'.dependencies] winapi = { version = "0.3", features = ["wincred"] } diff --git a/crates/goose-cli/src/scenario_tests/mock_client.rs b/crates/goose-cli/src/scenario_tests/mock_client.rs index a042c55a93..bd78417204 100644 --- a/crates/goose-cli/src/scenario_tests/mock_client.rs +++ b/crates/goose-cli/src/scenario_tests/mock_client.rs @@ -3,11 +3,13 @@ use goose::agents::mcp_client::{Error, McpClientTrait}; use rmcp::{ + handler::server::wrapper::Parameters, model::{ CallToolResult, Content, ErrorData, GetPromptResult, ListPromptsResult, ListResourcesResult, ListToolsResult, ReadResourceResult, ServerNotification, Tool, }, - object, + schemars::JsonSchema, + tool, }; use serde_json::Value; use std::collections::HashMap; @@ -17,7 +19,7 @@ use tokio_util::sync::CancellationToken; type Handler = Box Result, ErrorData> + Send + Sync>; pub struct MockClient { - tools: HashMap, + pub tools: HashMap, handlers: HashMap, } @@ -137,22 +139,15 @@ impl McpClientTrait for MockClient { pub const WEATHER_TYPE: &str = "cloudy"; pub fn weather_client() -> MockClient { - let weather_tool = Tool::new( - "get_weather", - "Get the weather for a location", - object!({ - "type": "object", - "required": ["location"], - "properties": { - "location": { - "type": "string", - "description": "The city and state, e.g. San Francisco, CA" - } - } - }), - ); + #[derive(JsonSchema)] + struct WeatherParams { + location: String, + } - let mock_client = MockClient::new().add_tool(weather_tool, |args| { + #[tool(name = "get_weather", description = "Get the weather for a location")] + async fn weather_tool(param: Parameters) {} + + let mock_client = MockClient::new().add_tool(weather_tool_tool_attr(), |args| { let location = args .get("location") .and_then(|v| v.as_str()) diff --git a/crates/goose-cli/src/scenario_tests/scenarios.rs b/crates/goose-cli/src/scenario_tests/scenarios.rs index 6b6c0ef999..263c9b3ed4 100644 --- a/crates/goose-cli/src/scenario_tests/scenarios.rs +++ b/crates/goose-cli/src/scenario_tests/scenarios.rs @@ -4,7 +4,7 @@ #[cfg(test)] mod tests { use crate::scenario_tests::message_generator::{image, text}; - use crate::scenario_tests::mock_client::WEATHER_TYPE; + use crate::scenario_tests::mock_client::{weather_client, WEATHER_TYPE}; use crate::scenario_tests::scenario_runner::run_scenario; use anyhow::Result; use goose::conversation::message::Message; @@ -102,4 +102,11 @@ mod tests { ) .await } + + #[tokio::test] + async fn test_weather_tool_schema() -> Result<()> { + let client = weather_client(); + eprintln!("{:?}", client.tools); + Ok(()) + } }