diff --git a/crates/goose/src/providers/factory.rs b/crates/goose/src/providers/factory.rs index da65c9ef91..42d7e69b9d 100644 --- a/crates/goose/src/providers/factory.rs +++ b/crates/goose/src/providers/factory.rs @@ -13,6 +13,7 @@ use super::{ ollama::OllamaProvider, openai::OpenAiProvider, openrouter::OpenRouterProvider, + snowflake::SnowflakeProvider, venice::VeniceProvider, }; use crate::model::ModelConfig; @@ -32,6 +33,7 @@ pub fn providers() -> Vec { OpenAiProvider::metadata(), OpenRouterProvider::metadata(), VeniceProvider::metadata(), + SnowflakeProvider::metadata(), ] } @@ -49,6 +51,7 @@ pub fn create(name: &str, model: ModelConfig) -> Result> { "gcp_vertex_ai" => Ok(Arc::new(GcpVertexAIProvider::from_env(model)?)), "google" => Ok(Arc::new(GoogleProvider::from_env(model)?)), "venice" => Ok(Arc::new(VeniceProvider::from_env(model)?)), + "snowflake" => Ok(Arc::new(SnowflakeProvider::from_env(model)?)), "github_copilot" => Ok(Arc::new(GithubCopilotProvider::from_env(model)?)), _ => Err(anyhow::anyhow!("Unknown provider: {}", name)), } diff --git a/crates/goose/src/providers/formats/mod.rs b/crates/goose/src/providers/formats/mod.rs index fc767e2724..6f3df3d0bf 100644 --- a/crates/goose/src/providers/formats/mod.rs +++ b/crates/goose/src/providers/formats/mod.rs @@ -4,3 +4,4 @@ pub mod databricks; pub mod gcpvertexai; pub mod google; pub mod openai; +pub mod snowflake; diff --git a/crates/goose/src/providers/formats/snowflake.rs b/crates/goose/src/providers/formats/snowflake.rs new file mode 100644 index 0000000000..c7cfe5924b --- /dev/null +++ b/crates/goose/src/providers/formats/snowflake.rs @@ -0,0 +1,716 @@ +use crate::message::{Message, MessageContent}; +use crate::model::ModelConfig; +use crate::providers::base::Usage; +use crate::providers::errors::ProviderError; +use anyhow::{anyhow, Result}; +use mcp_core::content::Content; +use mcp_core::role::Role; +use mcp_core::tool::{Tool, ToolCall}; +use serde_json::{json, Value}; +use std::collections::HashSet; + +/// Convert internal Message format to Snowflake's API message specification +pub fn format_messages(messages: &[Message]) -> Vec { + let mut snowflake_messages = Vec::new(); + + // Convert messages to Snowflake format + for message in messages { + let role = match message.role { + Role::User => "user", + Role::Assistant => "assistant", + }; + + let mut text_content = String::new(); + + for msg_content in &message.content { + match msg_content { + MessageContent::Text(text) => { + if !text_content.is_empty() { + text_content.push('\n'); + } + text_content.push_str(&text.text); + } + MessageContent::ToolRequest(_tool_request) => { + // Skip tool requests in message formatting - tools are handled separately + // through the tools parameter in the API request + continue; + } + MessageContent::ToolResponse(tool_response) => { + if let Ok(result) = &tool_response.tool_result { + let text = result + .iter() + .filter_map(|c| match c { + Content::Text(t) => Some(t.text.clone()), + _ => None, + }) + .collect::>() + .join("\n"); + + if !text_content.is_empty() { + text_content.push('\n'); + } + if !text.is_empty() { + text_content.push_str(&format!("Tool result: {}", text)); + } + } + } + MessageContent::ToolConfirmationRequest(_) => { + // Skip tool confirmation requests + } + MessageContent::ContextLengthExceeded(_) => { + // Skip + } + MessageContent::SummarizationRequested(_) => { + // Skip + } + MessageContent::Thinking(_thinking) => { + // Skip thinking for now + } + MessageContent::RedactedThinking(_redacted) => { + // Skip redacted thinking for now + } + MessageContent::Image(_) => continue, // Snowflake doesn't support image content yet + MessageContent::FrontendToolRequest(_tool_request) => { + // Skip frontend tool requests + } + } + } + + // Add message if it has text content + if !text_content.is_empty() { + snowflake_messages.push(json!({ + "role": role, + "content": text_content + })); + } + } + + // Only add default message if we truly have no messages at all + // This should be rare and only for edge cases + if snowflake_messages.is_empty() { + snowflake_messages.push(json!({ + "role": "user", + "content": "Continue the conversation" + })); + } + + snowflake_messages +} + +/// Convert internal Tool format to Snowflake's API tool specification +pub fn format_tools(tools: &[Tool]) -> Vec { + let mut unique_tools = HashSet::new(); + let mut tool_specs = Vec::new(); + + for tool in tools.iter() { + if unique_tools.insert(tool.name.clone()) { + let tool_spec = json!({ + "type": "generic", + "name": tool.name, + "description": tool.description, + "input_schema": tool.input_schema + }); + + tool_specs.push(json!({"tool_spec": tool_spec})); + } + } + + tool_specs +} + +/// Convert system message to Snowflake's API system specification +pub fn format_system(system: &str) -> Value { + json!({ + "role": "system", + "content": system, + }) +} + +/// Convert Snowflake's streaming API response to internal Message format +pub fn parse_streaming_response(sse_data: &str) -> Result { + let mut message = Message::assistant(); + let mut accumulated_text = String::new(); + let mut tool_use_id: Option = None; + let mut tool_name: Option = None; + let mut tool_input = String::new(); + + // Parse each SSE event + for line in sse_data.lines() { + if !line.starts_with("data: ") { + continue; + } + + let json_str = &line[6..]; // Remove "data: " prefix + if json_str.trim().is_empty() || json_str.trim() == "[DONE]" { + continue; + } + + let event: Value = match serde_json::from_str(json_str) { + Ok(v) => v, + Err(_) => { + continue; + } + }; + + if let Some(choices) = event.get("choices").and_then(|c| c.as_array()) { + if let Some(choice) = choices.first() { + if let Some(delta) = choice.get("delta") { + match delta.get("type").and_then(|t| t.as_str()) { + Some("text") => { + if let Some(content) = delta.get("content").and_then(|c| c.as_str()) { + accumulated_text.push_str(content); + } + } + Some("tool_use") => { + if let Some(id) = delta.get("tool_use_id").and_then(|i| i.as_str()) { + tool_use_id = Some(id.to_string()); + } + if let Some(name) = delta.get("name").and_then(|n| n.as_str()) { + tool_name = Some(name.to_string()); + } + if let Some(input) = delta.get("input").and_then(|i| i.as_str()) { + tool_input.push_str(input); + } + } + _ => {} + } + } + } + } + } + + // Add accumulated text if any + if !accumulated_text.is_empty() { + message = message.with_text(accumulated_text); + } + + // Add tool use if complete + if let (Some(id), Some(name)) = (&tool_use_id, &tool_name) { + if !tool_input.is_empty() { + let input_value = serde_json::from_str::(&tool_input) + .unwrap_or_else(|_| Value::String(tool_input.clone())); + let tool_call = ToolCall::new(name, input_value); + message = message.with_tool_request(id, Ok(tool_call)); + } else if tool_name.is_some() { + // Tool with no input - use empty object + let tool_call = ToolCall::new(name, Value::Object(serde_json::Map::new())); + message = message.with_tool_request(id, Ok(tool_call)); + } + } + + Ok(message) +} + +/// Convert Snowflake's API response to internal Message format +pub fn response_to_message(response: Value) -> Result { + let mut message = Message::assistant(); + + let content_list = response.get("content_list").and_then(|cl| cl.as_array()); + + // Handle case where content_list is missing or empty + let content_list = match content_list { + Some(list) if !list.is_empty() => list, + _ => { + // If no content_list or empty, check if there's a direct content field + if let Some(direct_content) = response.get("content").and_then(|c| c.as_str()) { + if !direct_content.is_empty() { + message = message.with_text(direct_content.to_string()); + } + return Ok(message); + } else { + // Return empty assistant message for empty responses + return Ok(message); + } + } + }; + + // Process all content items in the list + for content in content_list { + match content.get("type").and_then(|t| t.as_str()) { + Some("text") => { + if let Some(text) = content.get("text").and_then(|t| t.as_str()) { + if !text.is_empty() { + message = message.with_text(text.to_string()); + } + } + } + Some("tool_use") => { + let id = content + .get("tool_use_id") + .and_then(|i| i.as_str()) + .ok_or_else(|| anyhow!("Missing tool_use id"))?; + let name = content + .get("name") + .and_then(|n| n.as_str()) + .ok_or_else(|| anyhow!("Missing tool_use name"))?; + + let input = content + .get("input") + .ok_or_else(|| anyhow!("Missing tool input"))? + .clone(); + + let tool_call = ToolCall::new(name, input); + message = message.with_tool_request(id, Ok(tool_call)); + } + Some("thinking") => { + let thinking = content + .get("thinking") + .and_then(|t| t.as_str()) + .ok_or_else(|| anyhow!("Missing thinking content"))?; + let signature = content + .get("signature") + .and_then(|s| s.as_str()) + .ok_or_else(|| anyhow!("Missing thinking signature"))?; + message = message.with_thinking(thinking, signature); + } + Some("redacted_thinking") => { + let data = content + .get("data") + .and_then(|d| d.as_str()) + .ok_or_else(|| anyhow!("Missing redacted_thinking data"))?; + message = message.with_redacted_thinking(data); + } + _ => { + // Ignore unrecognized content types + } + } + } + + Ok(message) +} + +/// Extract usage information from Snowflake's API response +pub fn get_usage(data: &Value) -> Result { + // Extract usage data if available + if let Some(usage) = data.get("usage") { + let input_tokens = usage + .get("input_tokens") + .and_then(|v| v.as_u64()) + .map(|v| v as i32); + + let output_tokens = usage + .get("output_tokens") + .and_then(|v| v.as_u64()) + .map(|v| v as i32); + + let total_tokens = match (input_tokens, output_tokens) { + (Some(input), Some(output)) => Some(input + output), + _ => None, + }; + + Ok(Usage::new(input_tokens, output_tokens, total_tokens)) + } else { + tracing::debug!( + "Failed to get usage data: {}", + ProviderError::UsageError("No usage data found in response".to_string()) + ); + // If no usage data, return None for all values + Ok(Usage::new(None, None, None)) + } +} + +/// Create a complete request payload for Snowflake's API +pub fn create_request( + model_config: &ModelConfig, + system: &str, + messages: &[Message], + tools: &[Tool], +) -> Result { + let mut snowflake_messages = format_messages(messages); + let system_spec = format_system(system); + + // Add system message to the beginning of the messages + snowflake_messages.insert(0, system_spec); + + // Check if we have any messages to send + if snowflake_messages.is_empty() { + return Err(anyhow!("No valid messages to send to Snowflake API")); + } + + // Detect description generation requests and exclude tools to prevent interference + // with normal tool execution flow + let is_description_request = + system.contains("Reply with only a description in four words or less"); + + let tool_specs = if is_description_request { + // For description generation, don't include any tools to avoid confusion + format_tools(&[]) + } else { + format_tools(tools) + }; + + let max_tokens = model_config.max_tokens.unwrap_or(4096); + let mut payload = json!({ + "model": model_config.model_name, + "messages": snowflake_messages, + "max_tokens": max_tokens, + }); + + // Add tools if present and not a description request + if !tool_specs.is_empty() { + if let Some(obj) = payload.as_object_mut() { + obj.insert("tools".to_string(), json!(tool_specs)); + } else { + return Err(anyhow!( + "Failed to create request payload: payload is not a JSON object" + )); + } + } + + Ok(payload) +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn test_parse_text_response() -> Result<()> { + let response = json!({ + "id": "msg_123", + "type": "message", + "role": "assistant", + "content_list": [{ + "type": "text", + "text": "Hello! How can I assist you today?" + }], + "model": "claude-3-5-sonnet", + "stop_reason": "end_turn", + "stop_sequence": null, + "usage": { + "input_tokens": 12, + "output_tokens": 15 + } + }); + + let message = response_to_message(response.clone())?; + let usage = get_usage(&response)?; + + if let MessageContent::Text(text) = &message.content[0] { + assert_eq!(text.text, "Hello! How can I assist you today?"); + } else { + panic!("Expected Text content"); + } + + assert_eq!(usage.input_tokens, Some(12)); + assert_eq!(usage.output_tokens, Some(15)); + assert_eq!(usage.total_tokens, Some(27)); // 12 + 15 + + Ok(()) + } + + #[test] + fn test_parse_tool_response() -> Result<()> { + let response = json!({ + "id": "msg_123", + "type": "message", + "role": "assistant", + "content_list": [{ + "type": "tool_use", + "tool_use_id": "tool_1", + "name": "calculator", + "input": {"expression": "2 + 2"} + }], + "model": "claude-3-5-sonnet", + "stop_reason": "end_turn", + "stop_sequence": null, + "usage": { + "input_tokens": 15, + "output_tokens": 20 + } + }); + + let message = response_to_message(response.clone())?; + let usage = get_usage(&response)?; + + if let MessageContent::ToolRequest(tool_request) = &message.content[0] { + let tool_call = tool_request.tool_call.as_ref().unwrap(); + assert_eq!(tool_call.name, "calculator"); + assert_eq!(tool_call.arguments, json!({"expression": "2 + 2"})); + } else { + panic!("Expected ToolRequest content"); + } + + assert_eq!(usage.input_tokens, Some(15)); + assert_eq!(usage.output_tokens, Some(20)); + assert_eq!(usage.total_tokens, Some(35)); // 15 + 20 + + Ok(()) + } + + #[test] + fn test_message_to_snowflake_spec() { + let messages = vec![ + Message::user().with_text("Hello"), + Message::assistant().with_text("Hi there"), + Message::user().with_text("How are you?"), + ]; + + let spec = format_messages(&messages); + + assert_eq!(spec.len(), 3); + assert_eq!(spec[0]["role"], "user"); + assert_eq!(spec[0]["content"], "Hello"); + assert_eq!(spec[1]["role"], "assistant"); + assert_eq!(spec[1]["content"], "Hi there"); + assert_eq!(spec[2]["role"], "user"); + assert_eq!(spec[2]["content"], "How are you?"); + } + + #[test] + fn test_tools_to_snowflake_spec() { + let tools = vec![ + Tool::new( + "calculator", + "Calculate mathematical expressions", + json!({ + "type": "object", + "properties": { + "expression": { + "type": "string", + "description": "The mathematical expression to evaluate" + } + } + }), + None, + ), + Tool::new( + "weather", + "Get weather information", + json!({ + "type": "object", + "properties": { + "location": { + "type": "string", + "description": "The location to get weather for" + } + } + }), + None, + ), + ]; + + let spec = format_tools(&tools); + + assert_eq!(spec.len(), 2); + assert_eq!(spec[0]["tool_spec"]["name"], "calculator"); + assert_eq!( + spec[0]["tool_spec"]["description"], + "Calculate mathematical expressions" + ); + assert_eq!(spec[1]["tool_spec"]["name"], "weather"); + assert_eq!( + spec[1]["tool_spec"]["description"], + "Get weather information" + ); + } + + #[test] + fn test_system_to_snowflake_spec() { + let system = "You are a helpful assistant."; + let spec = format_system(system); + + assert_eq!(spec["role"], "system"); + assert_eq!(spec["content"], system); + } + + #[test] + fn test_parse_streaming_response() -> Result<()> { + let sse_data = r#"data: {"id":"a9537c2c-2017-4906-9817-2456168d89fa","model":"claude-3-5-sonnet","choices":[{"delta":{"type":"text","content":"I","content_list":[{"type":"text","text":"I"}],"text":"I"}}],"usage":{}} + +data: {"id":"a9537c2c-2017-4906-9817-2456168d89fa","model":"claude-3-5-sonnet","choices":[{"delta":{"type":"text","content":"'ll help you check Nvidia's current","content_list":[{"type":"text","text":"'ll help you check Nvidia's current"}],"text":"'ll help you check Nvidia's current"}}],"usage":{}} + +data: {"id":"a9537c2c-2017-4906-9817-2456168d89fa","model":"claude-3-5-sonnet","choices":[{"delta":{"type":"tool_use","tool_use_id":"tooluse_FB_nOElDTAOKa-YnVWI5Uw","name":"get_stock_price","content_list":[{"tool_use_id":"tooluse_FB_nOElDTAOKa-YnVWI5Uw","name":"get_stock_price"}],"text":""}}],"usage":{}} + +data: {"id":"a9537c2c-2017-4906-9817-2456168d89fa","model":"claude-3-5-sonnet","choices":[{"delta":{"type":"tool_use","input":"{\"symbol\":\"NVDA\"}","content_list":[{"input":"{\"symbol\":\"NVDA\"}"}],"text":""}}],"usage":{"prompt_tokens":397,"completion_tokens":65,"total_tokens":462}} +"#; + + let message = parse_streaming_response(sse_data)?; + + // Should have both text and tool request + assert_eq!(message.content.len(), 2); + + if let MessageContent::Text(text) = &message.content[0] { + assert!(text.text.contains("I'll help you check Nvidia's current")); + } else { + panic!("Expected Text content first"); + } + + if let MessageContent::ToolRequest(tool_request) = &message.content[1] { + let tool_call = tool_request.tool_call.as_ref().unwrap(); + assert_eq!(tool_call.name, "get_stock_price"); + assert_eq!(tool_call.arguments, json!({"symbol": "NVDA"})); + assert_eq!(tool_request.id, "tooluse_FB_nOElDTAOKa-YnVWI5Uw"); + } else { + panic!("Expected ToolRequest content second"); + } + + Ok(()) + } + + #[test] + fn test_create_request_format() -> Result<()> { + use crate::model::ModelConfig; + + let model_config = ModelConfig::new("claude-3-5-sonnet".to_string()); + + let system = "You are a helpful assistant that can use tools to get information."; + let messages = vec![Message::user().with_text("What is the stock price of Nvidia?")]; + + let tools = vec![Tool::new( + "get_stock_price", + "Get stock price information", + json!({ + "type": "object", + "properties": { + "symbol": { + "type": "string", + "description": "The symbol for the stock ticker, e.g. Snowflake = SNOW" + } + }, + "required": ["symbol"] + }), + None, + )]; + + let request = create_request(&model_config, system, &messages, &tools)?; + + // Check basic structure + assert_eq!(request["model"], "claude-3-5-sonnet"); + + let messages_array = request["messages"].as_array().unwrap(); + assert_eq!(messages_array.len(), 2); // system + user message + + // First message should be system with simple content + assert_eq!(messages_array[0]["role"], "system"); + assert_eq!( + messages_array[0]["content"], + "You are a helpful assistant that can use tools to get information." + ); + + // Second message should be user with simple content + assert_eq!(messages_array[1]["role"], "user"); + assert_eq!( + messages_array[1]["content"], + "What is the stock price of Nvidia?" + ); + + // Tools should have tool_spec wrapper + let tools_array = request["tools"].as_array().unwrap(); + assert_eq!(tools_array[0]["tool_spec"]["name"], "get_stock_price"); + + Ok(()) + } + + #[test] + fn test_parse_mixed_text_and_tool_response() -> Result<()> { + let response = json!({ + "id": "msg_123", + "type": "message", + "role": "assistant", + "content": "I'll help you with that calculation.", + "content_list": [ + { + "type": "text", + "text": "I'll help you with that calculation." + }, + { + "type": "tool_use", + "tool_use_id": "tool_1", + "name": "calculator", + "input": {"expression": "2 + 2"} + } + ], + "model": "claude-3-5-sonnet", + "usage": { + "input_tokens": 10, + "output_tokens": 15 + } + }); + + let message = response_to_message(response.clone())?; + + // Should have both text and tool request content + assert_eq!(message.content.len(), 2); + + if let MessageContent::Text(text) = &message.content[0] { + assert_eq!(text.text, "I'll help you with that calculation."); + } else { + panic!("Expected Text content first"); + } + + if let MessageContent::ToolRequest(tool_request) = &message.content[1] { + let tool_call = tool_request.tool_call.as_ref().unwrap(); + assert_eq!(tool_call.name, "calculator"); + assert_eq!(tool_request.id, "tool_1"); + } else { + panic!("Expected ToolRequest content second"); + } + + Ok(()) + } + + #[test] + fn test_empty_tools_array() { + let tools: Vec = vec![]; + let spec = format_tools(&tools); + assert_eq!(spec.len(), 0); + } + + #[test] + fn test_create_request_excludes_tools_for_description() -> Result<()> { + use crate::model::ModelConfig; + + let model_config = ModelConfig::new("claude-3-5-sonnet".to_string()); + let system = "Reply with only a description in four words or less"; + let messages = vec![Message::user().with_text("Test message")]; + let tools = vec![Tool::new( + "test_tool", + "Test tool", + json!({"type": "object", "properties": {}}), + None, + )]; + + let request = create_request(&model_config, system, &messages, &tools)?; + + // Should not include tools for description requests + assert!(request.get("tools").is_none()); + + Ok(()) + } + + #[test] + fn test_message_formatting_skips_tool_requests() { + use mcp_core::tool::ToolCall; + + // Create a conversation with text, tool requests, and tool responses + let tool_call = ToolCall::new("calculator", json!({"expression": "2 + 2"})); + + let messages = vec![ + Message::user().with_text("Calculate 2 + 2"), + Message::assistant() + .with_text("I'll help you calculate that.") + .with_tool_request("tool_1", Ok(tool_call)), + Message::user().with_text("Thanks!"), + ]; + + let spec = format_messages(&messages); + + // Should only have 3 messages - the tool request should be skipped + assert_eq!(spec.len(), 3); + assert_eq!(spec[0]["role"], "user"); + assert_eq!(spec[0]["content"], "Calculate 2 + 2"); + assert_eq!(spec[1]["role"], "assistant"); + assert_eq!(spec[1]["content"], "I'll help you calculate that."); + assert_eq!(spec[2]["role"], "user"); + assert_eq!(spec[2]["content"], "Thanks!"); + + // Verify no tool request content is in the message history + for message in &spec { + let content = message["content"].as_str().unwrap(); + assert!(!content.contains("Using tool:")); + assert!(!content.contains("calculator")); + } + } +} diff --git a/crates/goose/src/providers/mod.rs b/crates/goose/src/providers/mod.rs index c91e43c626..2c1a4571b6 100644 --- a/crates/goose/src/providers/mod.rs +++ b/crates/goose/src/providers/mod.rs @@ -17,6 +17,7 @@ pub mod oauth; pub mod ollama; pub mod openai; pub mod openrouter; +pub mod snowflake; pub mod toolshim; pub mod utils; pub mod utils_universal_openai_stream; diff --git a/crates/goose/src/providers/snowflake.rs b/crates/goose/src/providers/snowflake.rs new file mode 100644 index 0000000000..dcd0ac444b --- /dev/null +++ b/crates/goose/src/providers/snowflake.rs @@ -0,0 +1,434 @@ +use anyhow::Result; +use async_trait::async_trait; +use reqwest::{Client, StatusCode}; +use serde::{Deserialize, Serialize}; +use serde_json::{json, Value}; +use std::time::Duration; + +use super::base::{ConfigKey, Provider, ProviderMetadata, ProviderUsage}; +use super::errors::ProviderError; +use super::formats::snowflake::{create_request, get_usage, response_to_message}; +use super::utils::{get_model, ImageFormat}; +use crate::config::ConfigError; +use crate::message::Message; +use crate::model::ModelConfig; +use mcp_core::tool::Tool; +use url::Url; + +pub const SNOWFLAKE_DEFAULT_MODEL: &str = "claude-3-7-sonnet"; +pub const SNOWFLAKE_KNOWN_MODELS: &[&str] = &["claude-3-7-sonnet", "claude-3-5-sonnet"]; + +pub const SNOWFLAKE_DOC_URL: &str = + "https://docs.snowflake.com/en/user-guide/snowflake-cortex/llm-functions#choosing-a-model"; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub enum SnowflakeAuth { + Token(String), +} + +impl SnowflakeAuth { + pub fn token(token: String) -> Self { + Self::Token(token) + } +} + +#[derive(Debug, serde::Serialize)] +pub struct SnowflakeProvider { + #[serde(skip)] + client: Client, + host: String, + auth: SnowflakeAuth, + model: ModelConfig, + image_format: ImageFormat, +} + +impl Default for SnowflakeProvider { + fn default() -> Self { + let model = ModelConfig::new(SnowflakeProvider::metadata().default_model); + SnowflakeProvider::from_env(model).expect("Failed to initialize Snowflake provider") + } +} + +impl SnowflakeProvider { + pub fn from_env(model: ModelConfig) -> Result { + let config = crate::config::Config::global(); + let mut host: Result = config.get_param("SNOWFLAKE_HOST"); + if host.is_err() { + host = config.get_secret("SNOWFLAKE_HOST") + } + if host.is_err() { + return Err(ConfigError::NotFound( + "Did not find SNOWFLAKE_HOST in either config file or keyring".to_string(), + ) + .into()); + } + + let mut host = host?; + + // Convert host to lowercase + host = host.to_lowercase(); + + // Ensure host ends with snowflakecomputing.com + if !host.ends_with("snowflakecomputing.com") { + host = format!("{}.snowflakecomputing.com", host); + } + + let mut token: Result = config.get_param("SNOWFLAKE_TOKEN"); + + if token.is_err() { + token = config.get_secret("SNOWFLAKE_TOKEN") + } + + if token.is_err() { + return Err(ConfigError::NotFound( + "Did not find SNOWFLAKE_TOKEN in either config file or keyring".to_string(), + ) + .into()); + } + + let client = Client::builder() + .timeout(Duration::from_secs(600)) + .build()?; + + // Use token-based authentication + let api_key = token?; + Ok(Self { + client, + host, + auth: SnowflakeAuth::token(api_key), + model, + image_format: ImageFormat::OpenAi, + }) + } + + async fn ensure_auth_header(&self) -> Result { + match &self.auth { + // https://docs.snowflake.com/en/developer-guide/snowflake-rest-api/authentication#using-a-programmatic-access-token-pat + SnowflakeAuth::Token(token) => Ok(format!("Bearer {}", token)), + } + } + + async fn post(&self, payload: Value) -> Result { + let base_url_str = + if !self.host.starts_with("https://") && !self.host.starts_with("http://") { + format!("https://{}", self.host) + } else { + self.host.clone() + }; + let base_url = Url::parse(&base_url_str) + .map_err(|e| ProviderError::RequestFailed(format!("Invalid base URL: {e}")))?; + let path = "api/v2/cortex/inference:complete"; + let url = base_url.join(path).map_err(|e| { + ProviderError::RequestFailed(format!("Failed to construct endpoint URL: {e}")) + })?; + + let auth_header = self.ensure_auth_header().await?; + let response = self + .client + .post(url) + .header("Authorization", auth_header) + .header("User-Agent", "Goose") + .json(&payload) + .send() + .await?; + + let status = response.status(); + + let payload_text: String = response.text().await.ok().unwrap_or_default(); + + if status == StatusCode::OK { + if let Ok(payload) = serde_json::from_str::(&payload_text) { + if payload.get("code").is_some() { + let code = payload + .get("code") + .and_then(|c| c.as_str()) + .unwrap_or("Unknown code"); + let message = payload + .get("message") + .and_then(|m| m.as_str()) + .unwrap_or("Unknown message"); + return Err(ProviderError::RequestFailed(format!( + "{} - {}", + code, message + ))); + } + } + } + + let lines = payload_text.lines().collect::>(); + + let mut text = String::new(); + let mut tool_name = String::new(); + let mut tool_input = String::new(); + let mut tool_use_id = String::new(); + for line in lines.iter() { + if line.is_empty() { + continue; + } + + let json_str = match line.strip_prefix("data: ") { + Some(s) => s, + None => continue, + }; + + if let Ok(json_line) = serde_json::from_str::(json_str) { + let choices = match json_line.get("choices").and_then(|c| c.as_array()) { + Some(choices) => choices, + None => { + continue; + } + }; + + let choice = match choices.first() { + Some(choice) => choice, + None => { + continue; + } + }; + + let delta = match choice.get("delta") { + Some(delta) => delta, + None => { + continue; + } + }; + + // Track if we found text in content_list to avoid duplication + let mut found_text_in_content_list = false; + + // Handle content_list array first + if let Some(content_list) = delta.get("content_list").and_then(|cl| cl.as_array()) { + for content_item in content_list { + match content_item.get("type").and_then(|t| t.as_str()) { + Some("text") => { + if let Some(text_content) = + content_item.get("text").and_then(|t| t.as_str()) + { + text.push_str(text_content); + found_text_in_content_list = true; + } + } + Some("tool_use") => { + if let Some(tool_id) = + content_item.get("tool_use_id").and_then(|id| id.as_str()) + { + tool_use_id.push_str(tool_id); + } + if let Some(name) = + content_item.get("name").and_then(|n| n.as_str()) + { + tool_name.push_str(name); + } + if let Some(input) = + content_item.get("input").and_then(|i| i.as_str()) + { + tool_input.push_str(input); + } + } + _ => { + // Handle content items without explicit type but with tool information + if let Some(name) = + content_item.get("name").and_then(|n| n.as_str()) + { + tool_name.push_str(name); + } + if let Some(tool_id) = + content_item.get("tool_use_id").and_then(|id| id.as_str()) + { + tool_use_id.push_str(tool_id); + } + if let Some(input) = + content_item.get("input").and_then(|i| i.as_str()) + { + tool_input.push_str(input); + } + } + } + } + } + + // Handle direct content field (for text) only if we didn't find text in content_list + if !found_text_in_content_list { + if let Some(content) = delta.get("content").and_then(|c| c.as_str()) { + text.push_str(content); + } + } + } + } + + // Build the appropriate response structure + let mut content_list = Vec::new(); + + // Add text content if available + if !text.is_empty() { + content_list.push(json!({ + "type": "text", + "text": text + })); + } + + // Add tool use content only if we have complete tool information + if !tool_use_id.is_empty() && !tool_name.is_empty() { + // Parse tool input as JSON if it's not empty + let parsed_input = if tool_input.is_empty() { + json!({}) + } else { + serde_json::from_str::(&tool_input) + .unwrap_or_else(|_| json!({"raw_input": tool_input})) + }; + + content_list.push(json!({ + "type": "tool_use", + "tool_use_id": tool_use_id, + "name": tool_name, + "input": parsed_input + })); + } + + // Ensure we always have at least some content + if content_list.is_empty() { + content_list.push(json!({ + "type": "text", + "text": "" + })); + } + + let answer_payload = json!({ + "role": "assistant", + "content": text, + "content_list": content_list + }); + + match status { + StatusCode::OK => Ok(answer_payload), + StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN => { + // Extract a clean error message from the response if available + let error_msg = payload_text + .lines() + .find(|line| line.contains("\"message\"")) + .and_then(|line| { + let json_str = line.strip_prefix("data: ").unwrap_or(line); + serde_json::from_str::(json_str).ok() + }) + .and_then(|json| { + json.get("message") + .and_then(|m| m.as_str()) + .map(|s| s.to_string()) + }) + .unwrap_or_else(|| "Invalid credentials".to_string()); + + Err(ProviderError::Authentication(format!( + "Authentication failed. Please check your SNOWFLAKE_TOKEN and SNOWFLAKE_HOST configuration. Error: {}", + error_msg + ))) + } + StatusCode::BAD_REQUEST => { + // Snowflake provides a generic 'error' but also includes 'external_model_message' which is provider specific + // We try to extract the error message from the payload and check for phrases that indicate context length exceeded + let payload_str = payload_text.to_lowercase(); + let check_phrases = [ + "too long", + "context length", + "context_length_exceeded", + "reduce the length", + "token count", + "exceeds", + ]; + if check_phrases.iter().any(|c| payload_str.contains(c)) { + return Err(ProviderError::ContextLengthExceeded("Request exceeds maximum context length. Please reduce the number of messages or content size.".to_string())); + } + + // Try to parse a clean error message from the response + let error_msg = if let Ok(json) = serde_json::from_str::(&payload_text) { + json.get("message") + .and_then(|m| m.as_str()) + .map(|s| s.to_string()) + .or_else(|| { + json.get("external_model_message") + .and_then(|ext| ext.get("message")) + .and_then(|m| m.as_str()) + .map(|s| s.to_string()) + }) + .unwrap_or_else(|| "Bad request".to_string()) + } else { + "Bad request".to_string() + }; + + tracing::debug!( + "Provider request failed with status: {}. Response: {}", + status, + payload_text + ); + Err(ProviderError::RequestFailed(format!( + "Request failed: {}", + error_msg + ))) + } + StatusCode::TOO_MANY_REQUESTS => Err(ProviderError::RateLimitExceeded( + "Rate limit exceeded. Please try again later.".to_string(), + )), + StatusCode::INTERNAL_SERVER_ERROR | StatusCode::SERVICE_UNAVAILABLE => { + Err(ProviderError::ServerError( + "Snowflake service is temporarily unavailable. Please try again later." + .to_string(), + )) + } + _ => { + tracing::debug!( + "Provider request failed with status: {}. Response: {}", + status, + payload_text + ); + Err(ProviderError::RequestFailed(format!( + "Request failed with status: {}", + status + ))) + } + } + } +} + +#[async_trait] +impl Provider for SnowflakeProvider { + fn metadata() -> ProviderMetadata { + ProviderMetadata::new( + "snowflake", + "Snowflake", + "Access several models using Snowflake Cortex services.", + SNOWFLAKE_DEFAULT_MODEL, + SNOWFLAKE_KNOWN_MODELS.to_vec(), + SNOWFLAKE_DOC_URL, + vec![ + ConfigKey::new("SNOWFLAKE_HOST", true, false, None), + ConfigKey::new("SNOWFLAKE_TOKEN", true, true, None), + ], + ) + } + + fn get_model_config(&self) -> ModelConfig { + self.model.clone() + } + + #[tracing::instrument( + skip(self, system, messages, tools), + fields(model_config, input, output, input_tokens, output_tokens, total_tokens) + )] + async fn complete( + &self, + system: &str, + messages: &[Message], + tools: &[Tool], + ) -> Result<(Message, ProviderUsage), ProviderError> { + let payload = create_request(&self.model, system, messages, tools)?; + + let response = self.post(payload.clone()).await?; + + // Parse response + let message = response_to_message(response.clone())?; + let usage = get_usage(&response)?; + let model = get_model(&response); + super::utils::emit_debug_trace(&self.model, &payload, &response, &usage); + + Ok((message, ProviderUsage::new(model, usage))) + } +} diff --git a/crates/goose/tests/providers.rs b/crates/goose/tests/providers.rs index e65aff6618..4d41251fc4 100644 --- a/crates/goose/tests/providers.rs +++ b/crates/goose/tests/providers.rs @@ -4,7 +4,7 @@ use goose::message::{Message, MessageContent}; use goose::providers::base::Provider; use goose::providers::errors::ProviderError; use goose::providers::{ - anthropic, azure, bedrock, databricks, google, groq, ollama, openai, openrouter, + anthropic, azure, bedrock, databricks, google, groq, ollama, openai, openrouter, snowflake, }; use mcp_core::content::Content; use mcp_core::tool::Tool; @@ -479,6 +479,17 @@ async fn test_google_provider() -> Result<()> { .await } +#[tokio::test] +async fn test_snowflake_provider() -> Result<()> { + test_provider( + "Snowflake", + &["SNOWFLAKE_HOST", "SNOWFLAKE_TOKEN"], + None, + snowflake::SnowflakeProvider::default, + ) + .await +} + // Print the final test report #[ctor::dtor] fn print_test_report() { diff --git a/ui/desktop/src/components/settings/models/hardcoded_stuff.tsx b/ui/desktop/src/components/settings/models/hardcoded_stuff.tsx index 20b464fbb8..1b885096b8 100644 --- a/ui/desktop/src/components/settings/models/hardcoded_stuff.tsx +++ b/ui/desktop/src/components/settings/models/hardcoded_stuff.tsx @@ -9,6 +9,7 @@ export const default_models = { azure_openai: 'gpt-4o', gcp_vertex_ai: 'gemini-2.0-flash-001', aws_bedrock: 'us.anthropic.claude-3-7-sonnet-20250219-v1:0', + snowflake: 'claude-3-5-sonnet', venice: 'llama-3.3-70b', }; @@ -24,6 +25,7 @@ export const required_keys = { Ollama: ['OLLAMA_HOST'], Google: ['GOOGLE_API_KEY'], OpenRouter: ['OPENROUTER_API_KEY'], + Snowflake: ['SNOWFLAKE_HOST', 'SNOWFLAKE_TOKEN'], 'Azure OpenAI': [ 'AZURE_OPENAI_ENDPOINT', 'AZURE_OPENAI_DEPLOYMENT_NAME', @@ -53,6 +55,7 @@ export const supported_providers = [ 'Azure OpenAI', 'GCP Vertex AI', 'AWS Bedrock', + 'Snowflake', ]; export const model_docs_link = [ @@ -68,6 +71,10 @@ export const model_docs_link = [ { name: 'Ollama', href: 'https://ollama.com/library' }, { name: 'GCP Vertex AI', href: 'https://cloud.google.com/vertex-ai' }, { name: 'AWS Bedrock', href: 'https://console.aws.amazon.com/bedrock/home#/model-catalog' }, + { + name: 'Snowflake', + href: 'https://docs.snowflake.com/user-guide/snowflake-cortex/llm-functions#availability', + }, ]; export const provider_aliases = [ @@ -81,5 +88,6 @@ export const provider_aliases = [ { provider: 'Azure OpenAI', alias: 'azure_openai' }, { provider: 'GCP Vertex AI', alias: 'gcp_vertex_ai' }, { provider: 'AWS Bedrock', alias: 'aws_bedrock' }, + { provider: 'Snowflake', alias: 'snowflake' }, { provider: 'Venice', alias: 'venice' }, ]; diff --git a/ui/desktop/src/components/settings_v2/providers/ProviderRegistry.tsx b/ui/desktop/src/components/settings_v2/providers/ProviderRegistry.tsx index 5f5b1fd8ae..1e55f5c2f8 100644 --- a/ui/desktop/src/components/settings_v2/providers/ProviderRegistry.tsx +++ b/ui/desktop/src/components/settings_v2/providers/ProviderRegistry.tsx @@ -216,4 +216,22 @@ export const PROVIDER_REGISTRY: ProviderRegistry[] = [ ], }, }, + { + name: 'Snowflake', + details: { + id: 'snowflake', + name: 'Snowflake', + description: 'Access Cortex models hosted on your Snowflake account', + parameters: [ + { + name: 'SNOWFLAKE_HOST', + is_secret: false, + }, + { + name: 'SNOWFLAKE_TOKEN', + is_secret: true, + }, + ], + }, + }, ]; diff --git a/ui/desktop/src/components/settings_v2/providers/modal/subcomponents/ProviderLogo.tsx b/ui/desktop/src/components/settings_v2/providers/modal/subcomponents/ProviderLogo.tsx index 465546cc27..585b436c0d 100644 --- a/ui/desktop/src/components/settings_v2/providers/modal/subcomponents/ProviderLogo.tsx +++ b/ui/desktop/src/components/settings_v2/providers/modal/subcomponents/ProviderLogo.tsx @@ -5,6 +5,7 @@ import GroqLogo from './icons/groq@3x.png'; import OllamaLogo from './icons/ollama@3x.png'; import DatabricksLogo from './icons/databricks@3x.png'; import OpenRouterLogo from './icons/openrouter@3x.png'; +import SnowflakeLogo from './icons/snowflake@3x.png'; import DefaultLogo from './icons/default@3x.png'; // Map provider names to their logos @@ -16,6 +17,7 @@ const providerLogos: Record = { ollama: OllamaLogo, databricks: DatabricksLogo, openrouter: OpenRouterLogo, + snowflake: SnowflakeLogo, default: DefaultLogo, }; diff --git a/ui/desktop/src/components/settings_v2/providers/modal/subcomponents/icons/snowflake.png b/ui/desktop/src/components/settings_v2/providers/modal/subcomponents/icons/snowflake.png new file mode 100644 index 0000000000..07971ea060 Binary files /dev/null and b/ui/desktop/src/components/settings_v2/providers/modal/subcomponents/icons/snowflake.png differ diff --git a/ui/desktop/src/components/settings_v2/providers/modal/subcomponents/icons/snowflake@2x.png b/ui/desktop/src/components/settings_v2/providers/modal/subcomponents/icons/snowflake@2x.png new file mode 100644 index 0000000000..5414984ed4 Binary files /dev/null and b/ui/desktop/src/components/settings_v2/providers/modal/subcomponents/icons/snowflake@2x.png differ diff --git a/ui/desktop/src/components/settings_v2/providers/modal/subcomponents/icons/snowflake@3x.png b/ui/desktop/src/components/settings_v2/providers/modal/subcomponents/icons/snowflake@3x.png new file mode 100644 index 0000000000..b08d4e6514 Binary files /dev/null and b/ui/desktop/src/components/settings_v2/providers/modal/subcomponents/icons/snowflake@3x.png differ