diff --git a/Cargo.lock b/Cargo.lock index 065e56a5..13b491a1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2265,6 +2265,26 @@ dependencies = [ "yup-oauth2", ] +[[package]] +name = "google-sheets4" +version = "6.0.0+20240621" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4f8ccfc6418e81d1e2ed66fad49d0487526281505b8a0ed8ee770dc7d6bb1e5" +dependencies = [ + "chrono", + "google-apis-common", + "hyper 1.6.0", + "hyper-rustls 0.27.5", + "hyper-util", + "mime", + "serde", + "serde_json", + "serde_with", + "tokio", + "url", + "yup-oauth2", +] + [[package]] name = "goose" version = "1.0.13" @@ -2388,6 +2408,7 @@ dependencies = [ "docx-rs", "etcetera", "google-drive3", + "google-sheets4", "http-body-util", "ignore", "image 0.24.9", diff --git a/crates/goose-mcp/Cargo.toml b/crates/goose-mcp/Cargo.toml index be12dc2b..602b071f 100644 --- a/crates/goose-mcp/Cargo.toml +++ b/crates/goose-mcp/Cargo.toml @@ -33,6 +33,7 @@ etcetera = "0.8.0" tempfile = "3.8" include_dir = "0.7.4" google-drive3 = "6.0.0" +google-sheets4 = "6.0.0" webbrowser = "0.8" http-body-util = "0.1.2" regex = "1.11.1" diff --git a/crates/goose-mcp/src/google_drive/mod.rs b/crates/goose-mcp/src/google_drive/mod.rs index e3dcfec2..61028d7a 100644 --- a/crates/goose-mcp/src/google_drive/mod.rs +++ b/crates/goose-mcp/src/google_drive/mod.rs @@ -29,6 +29,8 @@ use google_drive3::{ DriveHub, }; +use google_sheets4::{self, Sheets}; + use http_body_util::BodyExt; /// async function to be pinned by the `present_user_url` method of the trait @@ -67,10 +69,14 @@ pub struct GoogleDriveRouter { tools: Vec, instructions: String, drive: DriveHub>, + sheets: Sheets>, } impl GoogleDriveRouter { - async fn google_auth() -> DriveHub> { + async fn google_auth() -> ( + DriveHub>, + Sheets>, + ) { let oauth_config = env::var("GOOGLE_DRIVE_OAUTH_CONFIG"); let keyfile_path_str = env::var("GOOGLE_DRIVE_OAUTH_PATH") .unwrap_or_else(|_| "./gcp-oauth.keys.json".to_string()); @@ -129,11 +135,14 @@ impl GoogleDriveRouter { .build(), ); - DriveHub::new(client, auth) + let drive_hub = DriveHub::new(client.clone(), auth.clone()); + let sheets_hub = Sheets::new(client, auth); + + (drive_hub, sheets_hub) } pub async fn new() -> Self { - let drive = Self::google_auth().await; + let (drive, sheets) = Self::google_auth().await; // handle auth let search_tool = Tool::new( @@ -185,13 +194,49 @@ impl GoogleDriveRouter { }), ); + let sheets_tool = Tool::new( + "sheets_tool".to_string(), + indoc! {r#" + Work with Google Sheets data using various operations. + Supports operations: + - list_sheets: List all sheets in a spreadsheet + - get_columns: Get column headers from a specific sheet + - get_values: Get values from a range + "#} + .to_string(), + json!({ + "type": "object", + "properties": { + "spreadsheetId": { + "type": "string", + "description": "The ID of the spreadsheet to work with", + }, + "operation": { + "type": "string", + "enum": ["list_sheets", "get_columns", "get_values"], + "description": "The operation to perform on the spreadsheet", + }, + "sheetName": { + "type": "string", + "description": "The name of the sheet to work with (optional for some operations)", + }, + "range": { + "type": "string", + "description": "The A1 notation of the range to retrieve values (e.g., 'Sheet1!A1:D10')", + } + }, + "required": ["spreadsheetId", "operation"], + }), + ); + let instructions = indoc::formatdoc! {r#" Google Drive MCP Server Instructions ## Overview - The Google Drive MCP server provides two main tools for interacting with Google Drive files: + The Google Drive MCP server provides tools for interacting with Google Drive files and Google Sheets: 1. search - Find files in your Google Drive 2. read - Read file contents directly using a uri in the `gdrive:///uri` format + 3. sheets_tool - Work with Google Sheets data using various operations ## Available Tools @@ -210,6 +255,18 @@ impl GoogleDriveRouter { Limitations: Google Sheets exporting only supports reading the first sheet. This is an important limitation that should be communicated to the user whenever dealing with a Google Sheet (mimeType: application/vnd.google-apps.spreadsheet). + ### 3. Sheets Tool + Work with Google Sheets data using various operations: + - list_sheets: List all sheets in a spreadsheet + - get_columns: Get column headers from a specific sheet + - get_values: Get values from a range + + Parameters: + - spreadsheetId: The ID of the spreadsheet (can be obtained from search results) + - operation: The operation to perform (one of the operations listed above) + - sheetName: The name of the sheet to work with (optional for some operations) + - range: The A1 notation of the range to retrieve values (e.g., 'Sheet1!A1:D10') + ## File Format Handling The server automatically handles different file types: - Google Docs → Markdown @@ -222,6 +279,7 @@ impl GoogleDriveRouter { 1. First, search for the file you want to read, searching by name. 2. Then, use the file URI from the search results to read its contents. + 3. For Google Sheets, use the sheets_tool with the appropriate operation. ## Best Practices 1. Always use search first to find the correct file URI @@ -240,9 +298,10 @@ impl GoogleDriveRouter { "#}; Self { - tools: vec![search_tool, read_tool], + tools: vec![search_tool, read_tool, sheets_tool], instructions, drive, + sheets, } } @@ -506,6 +565,152 @@ impl GoogleDriveRouter { } } + // Implement sheets_tool functionality + async fn sheets_tool(&self, params: Value) -> Result, ToolError> { + let spreadsheet_id = params.get("spreadsheetId").and_then(|q| q.as_str()).ok_or( + ToolError::InvalidParameters("The spreadsheetId is required".to_string()), + )?; + + let operation = params.get("operation").and_then(|q| q.as_str()).ok_or( + ToolError::InvalidParameters("The operation is required".to_string()), + )?; + + match operation { + "list_sheets" => { + // Get spreadsheet metadata to list all sheets + let result = self + .sheets + .spreadsheets() + .get(spreadsheet_id) + .clear_scopes() + .add_scope(Scope::Readonly) + .doit() + .await; + + match result { + Err(e) => Err(ToolError::ExecutionError(format!( + "Failed to execute Google Sheets get query, {}.", + e + ))), + Ok(r) => { + let spreadsheet = r.1; + let sheets = spreadsheet.sheets.unwrap_or_default(); + let sheets_info = sheets + .into_iter() + .filter_map(|sheet| { + let properties = sheet.properties?; + let title = properties.title?; + let sheet_id = properties.sheet_id?; + let grid_properties = properties.grid_properties?; + Some(format!( + "Sheet: {} (ID: {}, Rows: {}, Columns: {})", + title, + sheet_id, + grid_properties.row_count.unwrap_or(0), + grid_properties.column_count.unwrap_or(0) + )) + }) + .collect::>() + .join("\n"); + + Ok(vec![Content::text(sheets_info).with_priority(0.1)]) + } + } + }, + "get_columns" => { + // Get the sheet name if provided, otherwise we'll use the first sheet + let sheet_name = params + .get("sheetName") + .and_then(|q| q.as_str()) + .map(|s| format!("{}!1:1", s)) + .unwrap_or_else(|| "1:1".to_string()); // Default to first row of first sheet + + let result = self + .sheets + .spreadsheets() + .values_get(spreadsheet_id, &sheet_name) + .clear_scopes() + .add_scope(Scope::Readonly) + .doit() + .await; + + match result { + Err(e) => Err(ToolError::ExecutionError(format!( + "Failed to execute Google Sheets get_columns query, {}.", + e + ))), + Ok(r) => { + let value_range = r.1; + // Extract just the headers (first row) + let headers = match value_range.values { + Some(mut values) if !values.is_empty() => { + // Take the first row only + let headers = values.remove(0); + let header_values: Vec = headers + .into_iter() + .map(|cell| cell.as_str().unwrap_or_default().to_string()) + .collect(); + header_values.join(", ") + } + _ => "No headers found".to_string(), + }; + + Ok(vec![Content::text(headers).with_priority(0.1)]) + } + } + }, + "get_values" => { + let range = params + .get("range") + .and_then(|q| q.as_str()) + .ok_or(ToolError::InvalidParameters( + "The range is required for get_values operation".to_string(), + ))?; + + let result = self + .sheets + .spreadsheets() + .values_get(spreadsheet_id, range) + .clear_scopes() + .add_scope(Scope::Readonly) + .doit() + .await; + + match result { + Err(e) => Err(ToolError::ExecutionError(format!( + "Failed to execute Google Sheets values_get query, {}.", + e + ))), + Ok(r) => { + let value_range = r.1; + // Convert the values to a CSV string + let csv_content = match value_range.values { + Some(values) => { + let mut csv_string = String::new(); + for row in values { + let row_values: Vec = row + .into_iter() + .map(|cell| cell.as_str().unwrap_or_default().to_string()) + .collect(); + csv_string.push_str(&row_values.join(",")); + csv_string.push('\n'); + } + csv_string + } + None => "No data found".to_string(), + }; + + Ok(vec![Content::text(csv_content).with_priority(0.1)]) + } + } + }, + _ => Err(ToolError::InvalidParameters(format!( + "Invalid operation: {}. Supported operations are: list_sheets, get_columns, get_values", + operation + ))), + } + } + async fn read_google_resource(&self, uri: String) -> Result { self.read(json!({"uri": uri})) .await @@ -599,6 +804,7 @@ impl Router for GoogleDriveRouter { match tool_name.as_str() { "search" => this.search(arguments).await, "read" => this.read(arguments).await, + "sheets_tool" => this.sheets_tool(arguments).await, _ => Err(ToolError::NotFound(format!("Tool {} not found", tool_name))), } }) @@ -644,6 +850,7 @@ impl Clone for GoogleDriveRouter { tools: self.tools.clone(), instructions: self.instructions.clone(), drive: self.drive.clone(), + sheets: self.sheets.clone(), } } }