diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 4c04a6e424..c96f4d80f7 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -40,13 +40,6 @@ jobs: sudo apt update -y sudo apt install -y libdbus-1-dev gnome-keyring libxcb1-dev - - name: Start gnome-keyring - # run gnome-keyring with 'foobar' as password for the login keyring - # this will create a new login keyring and unlock it - # the login password doesn't matter, but the keyring must be unlocked for the tests to work - run: | - gnome-keyring-daemon --components=secrets --daemonize --unlock <<< 'foobar' - - name: Setup Rust uses: dtolnay/rust-toolchain@stable with: @@ -76,27 +69,11 @@ jobs: restore-keys: | ${{ runner.os }}-cargo-build- - - name: Install Ollama - run: curl -fsSL https://ollama.com/install.sh | sh - - - name: Start Ollama + - name: Build and Test run: | - # Run the background, in a way that survives to the next step - nohup ollama serve > ollama.log 2>&1 & - # Block using the ready endpoint - time curl --retry 5 --retry-connrefused --retry-delay 1 -sf http://localhost:11434 + gnome-keyring-daemon --components=secrets --daemonize --unlock <<< 'foobar' + cargo test - - name: Test Ollama Model - run: ollama run qwen2.5 hello || cat ollama.log - - - name: Build Rust Project - run: cargo build - - - name: Run Tests - run: cargo test --verbose - env: - OLLAMA_MODEL: "qwen2.5" - ## TODO: Need to decide if we wanna error out on clippy warnings. It was not being used before. # - name: Run Cargo Clippy (Lint) # run: cargo clippy -- -D warnings diff --git a/crates/goose-cli/src/log_usage.rs b/crates/goose-cli/src/log_usage.rs index 862078efa3..a7235b7f76 100644 --- a/crates/goose-cli/src/log_usage.rs +++ b/crates/goose-cli/src/log_usage.rs @@ -50,7 +50,6 @@ pub fn log_usage(session_file: String, usage: Vec) { #[cfg(test)] mod tests { use goose::providers::base::{ProviderUsage, Usage}; - use rust_decimal_macros::dec; use crate::{ log_usage::{log_usage, SessionLog}, @@ -72,7 +71,6 @@ mod tests { vec![ProviderUsage::new( "model".to_string(), Usage::new(Some(10), Some(20), Some(30)), - Some(dec!(0.5)), )], ); @@ -87,7 +85,6 @@ mod tests { assert_eq!(log.usage[0].usage.output_tokens, Some(20)); assert_eq!(log.usage[0].usage.total_tokens, Some(30)); assert_eq!(log.usage[0].model, "model"); - assert_eq!(log.usage[0].cost, Some(dec!(0.5))); // Remove the log file after test std::fs::remove_file(&log_file).ok(); diff --git a/crates/goose-server/src/routes/reply.rs b/crates/goose-server/src/routes/reply.rs index 05310711af..902e2fbc05 100644 --- a/crates/goose-server/src/routes/reply.rs +++ b/crates/goose-server/src/routes/reply.rs @@ -9,7 +9,7 @@ use axum::{ use bytes::Bytes; use futures::{stream::StreamExt, Stream}; use goose::message::{Message, MessageContent}; -use goose::providers::base::ModerationError; + use mcp_core::{content::Content, role::Role}; use serde::Deserialize; use serde_json::{json, Value}; @@ -166,18 +166,6 @@ impl ProtocolFormatter { format!("3:{}\n", encoded_error) } - fn format_moderation_error(error: &ModerationError) -> String { - let error_part = match error { - ModerationError::ContentFlagged { categories, .. } => { - format!( - "Content was flagged by moderation in the following categories: {}", - categories - ) - } - }; - format!("3:\"{}\"\n", error_part) - } - fn format_finish(reason: &str) -> String { // Finish messages start with "d:" let finish = json!({ @@ -324,19 +312,9 @@ async fn handler( Ok(stream) => stream, Err(e) => { tracing::error!("Failed to start reply stream: {}", e); - // Check if it's a moderation error - if let Some(moderation_error) = e.downcast_ref::() { - let _ = tx - .send(ProtocolFormatter::format_moderation_error(moderation_error)) - .await; - // Kill the stream since we encountered a moderation error - } else { - // Send a generic error message - let _ = tx - .send(ProtocolFormatter::format_error(&e.to_string())) - .await; - } - // Send a finish message with error as the reason + let _ = tx + .send(ProtocolFormatter::format_error(&e.to_string())) + .await; let _ = tx.send(ProtocolFormatter::format_finish("error")).await; return; } @@ -355,12 +333,7 @@ async fn handler( } Ok(Some(Err(e))) => { tracing::error!("Error processing message: {}", e); - // Check if it's a moderation error - if let Some(moderation_error) = e.downcast_ref::() { - let _ = tx.send(ProtocolFormatter::format_moderation_error(moderation_error)).await; - } else { - let _ = tx.send(ProtocolFormatter::format_error(&e.to_string())).await; - } + let _ = tx.send(ProtocolFormatter::format_error(&e.to_string())).await; break; } Ok(None) => { @@ -467,7 +440,7 @@ mod tests { use super::*; use goose::{ agents::DefaultAgent as Agent, - providers::base::{Moderation, ModerationResult, Provider, ProviderUsage, Usage}, + providers::base::{Provider, ProviderUsage, Usage}, providers::configs::ModelConfig, }; use mcp_core::tool::Tool; @@ -480,37 +453,27 @@ mod tests { #[async_trait::async_trait] impl Provider for MockProvider { - async fn complete_internal( - &self, - _system_prompt: &str, - _messages: &[Message], - _tools: &[Tool], - ) -> Result<(Message, ProviderUsage), anyhow::Error> { - Ok(( - Message::assistant().with_text("Mock response"), - ProviderUsage::new("mock".to_string(), Usage::default(), None), - )) - } - fn get_model_config(&self) -> &ModelConfig { &self.model_config } + async fn complete( + &self, + _system: &str, + _messages: &[Message], + _tools: &[Tool], + ) -> anyhow::Result<(Message, ProviderUsage)> { + Ok(( + Message::assistant().with_text("Mock response"), + ProviderUsage::new("mock".to_string(), Usage::default()), + )) + } + fn get_usage(&self, _data: &Value) -> anyhow::Result { Ok(Usage::new(None, None, None)) } } - #[async_trait::async_trait] - impl Moderation for MockProvider { - async fn moderate_content( - &self, - _content: &str, - ) -> Result { - Ok(ModerationResult::new(false, None, None)) - } - } - #[test] fn test_convert_messages_user_only() { let incoming = vec![IncomingMessage { @@ -584,21 +547,6 @@ mod tests { assert!(formatted.starts_with("3:")); assert!(formatted.contains("Test error")); - // Test moderation error formatting - let moderation_error = ModerationError::ContentFlagged { - categories: "hate, violence".to_string(), - category_scores: Some(json!({ - "hate": 0.9, - "violence": 0.8 - })), - }; - let formatted = ProtocolFormatter::format_moderation_error(&moderation_error); - println!("{}", formatted); - assert!(formatted.starts_with("3:")); - assert!( - formatted.contains("Content was flagged by moderation in the following categories:") - ); - // Test finish formatting let formatted = ProtocolFormatter::format_finish("stop"); assert!(formatted.starts_with("d:")); diff --git a/crates/goose/.env.example b/crates/goose/.env.example deleted file mode 100644 index 1b6ab55798..0000000000 --- a/crates/goose/.env.example +++ /dev/null @@ -1,9 +0,0 @@ -# OpenAI Configuration -OPENAI_API_KEY=your_api_key_here - -# Databricks Configuration -DATABRICKS_HOST=your_databricks_host -DATABRICKS_TOKEN=your_databricks_token - -# Test Configuration -INTEGRATION_TESTS=true # Set to true to run integration tests \ No newline at end of file diff --git a/crates/goose/Cargo.toml b/crates/goose/Cargo.toml index a65bf285e6..4546efb9f6 100644 --- a/crates/goose/Cargo.toml +++ b/crates/goose/Cargo.toml @@ -96,4 +96,4 @@ path = "examples/databricks_oauth.rs" [[bench]] name = "tokenization_benchmark" -harness = false \ No newline at end of file +harness = false diff --git a/crates/goose/src/agents/capabilities.rs b/crates/goose/src/agents/capabilities.rs index 2cd6a95c36..fba442e786 100644 --- a/crates/goose/src/agents/capabilities.rs +++ b/crates/goose/src/agents/capabilities.rs @@ -1,6 +1,5 @@ use chrono::{DateTime, TimeZone, Utc}; use mcp_client::McpService; -use rust_decimal_macros::dec; use std::collections::HashMap; use std::sync::Arc; use std::sync::LazyLock; @@ -193,11 +192,6 @@ impl Capabilities { e.usage.total_tokens = Some( e.usage.total_tokens.unwrap_or(0) + usage.usage.total_tokens.unwrap_or(0), ); - if e.cost.is_none() || usage.cost.is_none() { - e.cost = None; // Pricing is not available for all models - } else { - e.cost = Some(e.cost.unwrap_or(dec!(0)) + usage.cost.unwrap_or(dec!(0))); - } }) .or_insert_with(|| usage.clone()); }); diff --git a/crates/goose/src/providers/anthropic.rs b/crates/goose/src/providers/anthropic.rs index 0f65e1ac59..617fb38a29 100644 --- a/crates/goose/src/providers/anthropic.rs +++ b/crates/goose/src/providers/anthropic.rs @@ -1,20 +1,16 @@ -use anyhow::{anyhow, Result}; +use anyhow::Result; use async_trait::async_trait; use reqwest::Client; use reqwest::StatusCode; -use serde_json::{json, Value}; -use std::collections::HashSet; +use serde_json::Value; use std::time::Duration; -use super::base::{Moderation, ModerationResult, Provider, ProviderUsage, Usage}; +use super::base::{Provider, ProviderUsage, Usage}; use super::configs::ModelConfig; -use super::model_pricing::cost; -use super::model_pricing::model_pricing_for; +use super::formats::anthropic::{create_request, get_usage, response_to_message}; use super::utils::{emit_debug_trace, get_model, non_ok_response_to_provider_error}; -use crate::message::{Message, MessageContent}; -use mcp_core::content::Content; -use mcp_core::role::Role; -use mcp_core::tool::{Tool, ToolCall}; +use crate::message::Message; +use mcp_core::tool::Tool; pub const ANTHROPIC_DEFAULT_MODEL: &str = "claude-3-5-sonnet-latest"; @@ -48,176 +44,6 @@ impl AnthropicProvider { }) } - fn tools_to_anthropic_spec(tools: &[Tool]) -> Vec { - let mut unique_tools = HashSet::new(); - let mut tool_specs = Vec::new(); - - for tool in tools { - if unique_tools.insert(tool.name.clone()) { - tool_specs.push(json!({ - "name": tool.name, - "description": tool.description, - "input_schema": tool.input_schema - })); - } - } - - // Add "cache_control" to the last tool spec, if any. This means that all tool definitions, - // will be cached as a single prefix. - if let Some(last_tool) = tool_specs.last_mut() { - last_tool - .as_object_mut() - .unwrap() - .insert("cache_control".to_string(), json!({ "type": "ephemeral" })); - } - - tool_specs - } - - fn system_to_anthropic_spec(system: &str) -> Value { - json!([{ - "type": "text", - "text": system, - "cache_control": { "type": "ephemeral" } - }]) - } - - fn messages_to_anthropic_spec(messages: &[Message]) -> Vec { - let mut anthropic_messages = Vec::new(); - - // Convert messages to Anthropic format - for message in messages { - let role = match message.role { - Role::User => "user", - Role::Assistant => "assistant", - }; - - let mut content = Vec::new(); - for msg_content in &message.content { - match msg_content { - MessageContent::Text(text) => { - content.push(json!({ - "type": "text", - "text": text.text - })); - } - MessageContent::ToolRequest(tool_request) => { - if let Ok(tool_call) = &tool_request.tool_call { - content.push(json!({ - "type": "tool_use", - "id": tool_request.id, - "name": tool_call.name, - "input": tool_call.arguments - })); - } - } - 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"); - - content.push(json!({ - "type": "tool_result", - "tool_use_id": tool_response.id, - "content": text - })); - } - } - MessageContent::Image(_) => continue, // Anthropic doesn't support image content yet - } - } - - // Skip messages with empty content - if !content.is_empty() { - anthropic_messages.push(json!({ - "role": role, - "content": content - })); - } - } - - // If no messages, add a default one - if anthropic_messages.is_empty() { - anthropic_messages.push(json!({ - "role": "user", - "content": [{ - "type": "text", - "text": "Ignore" - }] - })); - } - - // Add "cache_control" to the last and second-to-last "user" messages. - // During each turn, we mark the final message with cache_control so the conversation can be - // incrementally cached. The second-to-last user message is also marked for caching with the - // cache_control parameter, so that this checkpoint can read from the previous cache. - let mut user_count = 0; - for message in anthropic_messages.iter_mut().rev() { - if message.get("role") == Some(&json!("user")) { - if let Some(content) = message.get_mut("content") { - if let Some(content_array) = content.as_array_mut() { - if let Some(last_content) = content_array.last_mut() { - last_content.as_object_mut().unwrap().insert( - "cache_control".to_string(), - json!({ "type": "ephemeral" }), - ); - } - } - } - user_count += 1; - if user_count >= 2 { - break; - } - } - } - - anthropic_messages - } - - fn parse_anthropic_response(response: Value) -> Result { - let content_blocks = response - .get("content") - .and_then(|c| c.as_array()) - .ok_or_else(|| anyhow!("Invalid response format: missing content array"))?; - - let mut message = Message::assistant(); - - for block in content_blocks { - match block.get("type").and_then(|t| t.as_str()) { - Some("text") => { - if let Some(text) = block.get("text").and_then(|t| t.as_str()) { - message = message.with_text(text.to_string()); - } - } - Some("tool_use") => { - let id = block - .get("id") - .and_then(|i| i.as_str()) - .ok_or_else(|| anyhow!("Missing tool_use id"))?; - let name = block - .get("name") - .and_then(|n| n.as_str()) - .ok_or_else(|| anyhow!("Missing tool_use name"))?; - let input = block - .get("input") - .ok_or_else(|| anyhow!("Missing tool_use input"))?; - - let tool_call = ToolCall::new(name, input.clone()); - message = message.with_tool_request(id, Ok(tool_call)); - } - _ => continue, - } - } - - Ok(message) - } - async fn post(&self, payload: Value) -> Result { let url = format!("{}/v1/messages", self.host.trim_end_matches('/')); @@ -248,275 +74,28 @@ impl Provider for AnthropicProvider { #[tracing::instrument( skip(self, system, messages, tools), - fields( - model_config, - input, - output, - input_tokens, - output_tokens, - total_tokens, - cost - ) + fields(model_config, input, output, input_tokens, output_tokens, total_tokens) )] - async fn complete_internal( + async fn complete( &self, system: &str, messages: &[Message], tools: &[Tool], ) -> Result<(Message, ProviderUsage)> { - let anthropic_messages = Self::messages_to_anthropic_spec(messages); - let tool_specs = Self::tools_to_anthropic_spec(tools); - let system_spec = Self::system_to_anthropic_spec(system); - - // Check if we have any messages to send - if anthropic_messages.is_empty() { - return Err(anyhow!("No valid messages to send to Anthropic API")); - } - - let mut payload = json!({ - "model": self.model.model_name, - "messages": anthropic_messages, - "max_tokens": self.model.max_tokens.unwrap_or(4096) - }); - - // Add system message if present - if !system.is_empty() { - payload - .as_object_mut() - .unwrap() - .insert("system".to_string(), json!(system_spec)); - } - - // Add tools if present - if !tool_specs.is_empty() { - payload - .as_object_mut() - .unwrap() - .insert("tools".to_string(), json!(tool_specs)); - } - - // Add temperature if specified - if let Some(temp) = self.model.temperature { - payload - .as_object_mut() - .unwrap() - .insert("temperature".to_string(), json!(temp)); - } + let payload = create_request(&self.model, system, messages, tools)?; // Make request let response = self.post(payload.clone()).await?; // Parse response - let message = Self::parse_anthropic_response(response.clone())?; + let message = response_to_message(response.clone())?; let usage = self.get_usage(&response)?; let model = get_model(&response); - let cost = cost(&usage, &model_pricing_for(&model)); - emit_debug_trace(self, &payload, &response, &usage, cost); - Ok((message, ProviderUsage::new(model, usage, cost))) + emit_debug_trace(self, &payload, &response, &usage); + Ok((message, ProviderUsage::new(model, usage))) } fn get_usage(&self, 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(i), Some(o)) => Some(i + o), - _ => None, - }; - - Ok(Usage::new(input_tokens, output_tokens, total_tokens)) - } else { - // If no usage data, return None for all values - Ok(Usage::new(None, None, None)) - } - } -} - -#[async_trait] -impl Moderation for AnthropicProvider { - async fn moderate_content(&self, _content: &str) -> Result { - Ok(ModerationResult::new(false, None, None)) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use rust_decimal_macros::dec; - use serde_json::json; - use wiremock::matchers::{header, method, path}; - use wiremock::{Mock, MockServer, ResponseTemplate}; - - async fn setup_mock_server(response_body: Value) -> (MockServer, AnthropicProvider) { - let mock_server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/v1/messages")) - .and(header("x-api-key", "test_api_key")) - .respond_with(ResponseTemplate::new(200).set_body_json(response_body)) - .mount(&mock_server) - .await; - - let provider = AnthropicProvider { - client: Client::builder().build().unwrap(), - host: mock_server.uri(), - api_key: "test_api_key".to_string(), - model: ModelConfig::new("claude-3-5-sonnet-latest".to_string()) - .with_temperature(Some(0.7)) - .with_context_limit(Some(200_000)), - }; - - (mock_server, provider) - } - - #[tokio::test] - async fn test_complete_basic() -> Result<()> { - let response_body = json!({ - "id": "msg_123", - "type": "message", - "role": "assistant", - "content": [{ - "type": "text", - "text": "Hello! How can I assist you today?" - }], - "model": "claude-3-5-sonnet-latest", - "stop_reason": "end_turn", - "stop_sequence": null, - "usage": { - "input_tokens": 12, - "output_tokens": 15, - "cache_creation_input_tokens": 12, - "cache_read_input_tokens": 0 - } - }); - - let (_, provider) = setup_mock_server(response_body).await; - - let messages = vec![Message::user().with_text("Hello?")]; - - let (message, usage) = provider - .complete_internal("You are a helpful assistant.", &messages, &[]) - .await?; - - 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.usage.input_tokens, Some(12)); - assert_eq!(usage.usage.output_tokens, Some(15)); - assert_eq!(usage.usage.total_tokens, Some(27)); - assert_eq!(usage.model, "claude-3-5-sonnet-latest"); - assert_eq!(usage.cost, Some(dec!(0.000261))); - - Ok(()) - } - - #[tokio::test] - async fn test_complete_with_tools() -> Result<()> { - let response_body = json!({ - "id": "msg_123", - "type": "message", - "role": "assistant", - "content": [{ - "type": "tool_use", - "id": "tool_1", - "name": "calculator", - "input": { - "expression": "2 + 2" - } - }], - "model": "claude-3-5-sonnet-latest", - "stop_reason": "end_turn", - "stop_sequence": null, - "usage": { - "input_tokens": 15, - "output_tokens": 20, - "cache_creation_input_tokens": 15, - "cache_read_input_tokens": 0, - } - }); - - let (_, provider) = setup_mock_server(response_body).await; - - let messages = vec![Message::user().with_text("What is 2 + 2?")]; - let tool = Tool::new( - "calculator", - "Calculate mathematical expressions", - json!({ - "type": "object", - "properties": { - "expression": { - "type": "string", - "description": "The mathematical expression to evaluate" - } - } - }), - ); - - let (message, usage) = provider - .complete_internal("You are a helpful assistant.", &messages, &[tool]) - .await?; - - 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.usage.input_tokens, Some(15)); - assert_eq!(usage.usage.output_tokens, Some(20)); - assert_eq!(usage.usage.total_tokens, Some(35)); - - Ok(()) - } - - #[tokio::test] - async fn test_empty_messages() -> Result<()> { - let response_body = json!({ - "id": "msg_123", - "type": "message", - "role": "assistant", - "content": [{ - "type": "text", - "text": "Hello!" - }], - "model": "claude-3-5-sonnet-latest", - "stop_reason": "end_turn", - "stop_sequence": null, - "usage": { - "input_tokens": 12, - "output_tokens": 15 - } - }); - - let (_, provider) = setup_mock_server(response_body).await; - - // Create a message with empty content - let messages = vec![ - Message::user().with_text(""), - Message::user().with_text("Hello"), - ]; - - let (message, _) = provider - .complete("You are a helpful assistant.", &messages, &[]) - .await?; - - if let MessageContent::Text(text) = &message.content[0] { - assert_eq!(text.text, "Hello!"); - } else { - panic!("Expected Text content"); - } - - Ok(()) + get_usage(data) } } diff --git a/crates/goose/src/providers/base.rs b/crates/goose/src/providers/base.rs index 1734bba013..182ca57a72 100644 --- a/crates/goose/src/providers/base.rs +++ b/crates/goose/src/providers/base.rs @@ -1,45 +1,23 @@ use anyhow::Result; -use lazy_static::lazy_static; -use rust_decimal::Decimal; use serde::{Deserialize, Serialize}; -use std::collections::HashMap; -use std::sync::Arc; -use thiserror::Error; -use tokio::sync::RwLock; +use serde_json::Value; use super::configs::ModelConfig; use crate::message::Message; use mcp_core::tool::Tool; -#[derive(Error, Debug)] -pub enum ModerationError { - #[error("Content was flagged for moderation in categories: {categories}")] - ContentFlagged { - categories: String, - category_scores: Option, - }, -} - #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ProviderUsage { pub model: String, pub usage: Usage, - pub cost: Option, } impl ProviderUsage { - pub fn new(model: String, usage: Usage, cost: Option) -> Self { - Self { model, usage, cost } + pub fn new(model: String, usage: Usage) -> Self { + Self { model, usage } } } -#[derive(Debug, Clone)] -pub struct Pricing { - /// Prices are per million tokens. - pub input_token_price: Decimal, - pub output_token_price: Decimal, -} - #[derive(Debug, Clone, Serialize, Deserialize, Default)] pub struct Usage { pub input_tokens: Option, @@ -61,101 +39,11 @@ impl Usage { } } -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct ModerationResult { - /// Whether the content was flagged as inappropriate - pub flagged: bool, - /// Optional categories that were flagged (provider specific) - pub categories: Option>, - /// Optional scores for each category (provider specific) - pub category_scores: Option, -} - -impl ModerationResult { - pub fn new( - flagged: bool, - categories: Option>, - category_scores: Option, - ) -> Self { - Self { - flagged, - categories, - category_scores, - } - } -} - -#[derive(Debug, Clone, Default)] -pub struct ModerationCache { - cache: Arc>>, -} - -impl ModerationCache { - pub fn new() -> Self { - Self { - cache: Arc::new(RwLock::new(HashMap::new())), - } - } - - pub async fn get(&self, content: &str) -> Option { - let cache = self.cache.read().await; - cache.get(content).cloned() - } - - pub async fn set(&self, content: String, result: ModerationResult) { - let mut cache = self.cache.write().await; - cache.insert(content, result); - } -} - -lazy_static! { - static ref DEFAULT_CACHE: ModerationCache = ModerationCache::new(); -} - use async_trait::async_trait; -use serde_json::Value; - -/// Trait for handling content moderation -#[async_trait] -pub trait Moderation: Send + Sync { - /// Get the moderation cache - fn moderation_cache(&self) -> &ModerationCache { - &DEFAULT_CACHE - } - - /// Internal moderation method to be implemented by providers - async fn moderate_content_internal(&self, _content: &str) -> Result { - Ok(ModerationResult::new(false, None, None)) - } - - /// Moderate the given content - /// - /// # Arguments - /// * `content` - The text content to moderate - /// - /// # Returns - /// A ModerationResult containing the moderation decision and details - async fn moderate_content(&self, content: &str) -> Result { - // Check cache first - if let Some(cached) = self.moderation_cache().get(content).await { - return Ok(cached); - } - - // If not in cache, do moderation - let result = self.moderate_content_internal(content).await?; - - // Cache the result - self.moderation_cache() - .set(content.to_string(), result.clone()) - .await; - - Ok(result) - } -} /// Base trait for AI providers (OpenAI, Anthropic, etc) #[async_trait] -pub trait Provider: Send + Sync + Moderation { +pub trait Provider: Send + Sync { /// Get the model configuration fn get_model_config(&self) -> &ModelConfig; @@ -173,74 +61,6 @@ pub trait Provider: Send + Sync + Moderation { system: &str, messages: &[Message], tools: &[Tool], - ) -> Result<(Message, ProviderUsage)> { - self.complete_internal(system, messages, tools).await - - // Get the latest user message - //let latest_user_msg = messages - // .iter() - // .rev() - // .find(|msg| { - // msg.role == Role::User - // && msg - // .content - // .iter() - // .any(|content| matches!(content, MessageContent::Text(_))) - // }) - // .ok_or_else(|| anyhow::anyhow!("No user message with text content found in history"))?; - // - //// Get the content to moderate - //let content = latest_user_msg.content.first().unwrap().as_text().unwrap(); - - // Start completion and moderation immediately - //let moderation_fut = self.moderate_content(content); - //tokio::pin!(completion_fut); - //tokio::pin!(moderation_fut); - - // Run moderation and completion concurrently - //select! { - // moderation = &mut moderation_fut => { - // let result = moderation?; - // - // if result.flagged { - // let categories = result.categories - // .unwrap_or_else(|| vec!["unknown".to_string()]) - // .join(", "); - // return Err(ModerationError::ContentFlagged { - // categories, - // category_scores: result.category_scores, - // }.into()); - // } - // - // // Moderation passed, wait for completion - // Ok(completion_fut.await?) - // } - // completion = &mut completion_fut => { - // // Completion finished first, still need to check moderation - // let completion_result = completion?; - // let moderation_result = moderation_fut.await?; - // - // if moderation_result.flagged { - // let categories = moderation_result.categories - // .unwrap_or_else(|| vec!["unknown".to_string()]) - // .join(", "); - // return Err(ModerationError::ContentFlagged { - // categories, - // category_scores: moderation_result.category_scores, - // }.into()); - // } - // - // Ok(completion_result) - // } - //} - } - - /// Internal completion method to be implemented by providers - async fn complete_internal( - &self, - system: &str, - messages: &[Message], - tools: &[Tool], ) -> Result<(Message, ProviderUsage)>; fn get_usage(&self, data: &Value) -> Result; diff --git a/crates/goose/src/providers/databricks.rs b/crates/goose/src/providers/databricks.rs index 87f371780a..e5ef1b90d0 100644 --- a/crates/goose/src/providers/databricks.rs +++ b/crates/goose/src/providers/databricks.rs @@ -2,26 +2,23 @@ use anyhow::{anyhow, Result}; use async_trait::async_trait; use reqwest::Client; use serde::{Deserialize, Serialize}; -use serde_json::{json, Value}; +use serde_json::Value; use std::time::Duration; -use super::base::{Moderation, ModerationResult, Provider, ProviderUsage, Usage}; +use super::base::{Provider, ProviderUsage, Usage}; use super::configs::ModelConfig; -use super::model_pricing::{cost, model_pricing_for}; -use super::oauth; -use super::utils::{check_bedrock_context_length_error, get_model, handle_response, ImageFormat}; -use crate::message::Message; -use crate::providers::openai_utils::{ - check_openai_context_length_error, get_openai_usage, messages_to_openai_spec, - openai_response_to_message, tools_to_openai_spec, +use super::formats::openai::{ + create_request, get_usage, is_context_length_error, response_to_message, }; +use super::oauth; +use super::utils::{get_model, handle_response, ImageFormat}; +use crate::message::Message; use mcp_core::tool::Tool; const DEFAULT_CLIENT_ID: &str = "databricks-cli"; const DEFAULT_REDIRECT_URL: &str = "http://localhost:8020"; const DEFAULT_SCOPES: &[&str] = &["all-apis"]; pub const DATABRICKS_DEFAULT_MODEL: &str = "claude-3-5-sonnet-2"; -pub const INPUT_GUARDRAIL: &str = "input_guardrail"; #[derive(Debug, Clone, Serialize, Deserialize)] pub enum DatabricksAuth { @@ -129,35 +126,6 @@ impl DatabricksProvider { handle_response(payload, response).await } - - async fn handle_moderation_response(&self, response: reqwest::Response) -> Result { - match response.status() { - reqwest::StatusCode::OK => { - let payload = response.json().await?; - Ok(payload) - } - reqwest::StatusCode::BAD_REQUEST => { - let error_body: Value = response.json().await?; - - // Check if this is a moderation error - if let Some(finish_reason) = error_body.get("finishReason") { - if finish_reason == "input_guardrail_triggered" { - return Ok(error_body); - } - } - // Not a moderation error, return the original error - Err(anyhow::anyhow!("Bad request: {}", error_body)) - } - status => { - let error_body: Value = response.json().await?; - Err(anyhow::anyhow!( - "Moderation request failed with status: {}\nPayload {}", - status, - error_body - )) - } - } - } } #[async_trait] @@ -168,372 +136,41 @@ impl Provider for DatabricksProvider { #[tracing::instrument( skip(self, system, messages, tools), - fields( - model_config, - input, - output, - input_tokens, - output_tokens, - total_tokens, - cost - ) + fields(model_config, input, output, input_tokens, output_tokens, total_tokens) )] - async fn complete_internal( + async fn complete( &self, system: &str, messages: &[Message], tools: &[Tool], ) -> Result<(Message, ProviderUsage)> { - // Prepare messages and tools - let concat_tool_response_contents = false; - let messages_spec = - messages_to_openai_spec(messages, &self.image_format, concat_tool_response_contents); - let tools_spec = if !tools.is_empty() { - tools_to_openai_spec(tools)? - } else { - vec![] - }; + let mut payload = create_request(&self.model, system, messages, tools, &self.image_format)?; + // Remove the model key which is part of the url with databricks + payload + .as_object_mut() + .expect("payload should have model key") + .remove("model"); - // Build payload with system message - let mut messages_array = vec![json!({ "role": "system", "content": system })]; - messages_array.extend(messages_spec); - - let mut payload = json!({ "messages": messages_array }); - - // Add optional parameters - if !tools_spec.is_empty() { - payload["tools"] = json!(tools_spec); - } - if let Some(temp) = self.model.temperature { - payload["temperature"] = json!(temp); - } - if let Some(tokens) = self.model.max_tokens { - payload["max_tokens"] = json!(tokens); - } - - // Remove null values - let payload = serde_json::Value::Object( - payload - .as_object() - .unwrap() - .iter() - .filter(|&(_, v)| !v.is_null()) - .map(|(k, v)| (k.clone(), v.clone())) - .collect(), - ); - - // Make request let response = self.post(payload.clone()).await?; // Raise specific error if context length is exceeded if let Some(error) = response.get("error") { - if let Some(err) = check_openai_context_length_error(error) { - return Err(err.into()); - } else if let Some(err) = check_bedrock_context_length_error(error) { + if let Some(err) = is_context_length_error(error) { return Err(err.into()); } return Err(anyhow!("Databricks API error: {}", error)); } // Parse response - let message = openai_response_to_message(response.clone())?; + let message = response_to_message(response.clone())?; let usage = self.get_usage(&response)?; let model = get_model(&response); - let cost = cost(&usage, &model_pricing_for(&model)); - super::utils::emit_debug_trace(self, &payload, &response, &usage, cost); + super::utils::emit_debug_trace(self, &payload, &response, &usage); - Ok((message, ProviderUsage::new(model, usage, cost))) + Ok((message, ProviderUsage::new(model, usage))) } fn get_usage(&self, data: &Value) -> Result { - get_openai_usage(data) - } -} - -#[async_trait] -impl Moderation for DatabricksProvider { - async fn moderate_content_internal(&self, content: &str) -> Result { - let url = format!( - "{}/serving-endpoints/moderation/invocations", - self.host.trim_end_matches('/') - ); - - let auth_header = self.ensure_auth_header().await?; - let payload = json!({ - "messages": [ - { - "role": "user", - "content": content - } - ] - }); - - let response = self - .client - .post(&url) - .header("Authorization", auth_header) - .json(&payload) - .send() - .await?; - - // let response: Value = response.json().await?; - // let response = handle_response(payload, response).await??; - let response = self.handle_moderation_response(response).await?; - - // Check if we got a moderation result - if let Some(input_guardrail) = response.get(INPUT_GUARDRAIL) { - if let Some(first_result) = input_guardrail.as_array().and_then(|arr| arr.first()) { - if let Some(flagged) = first_result.get("flagged").and_then(|f| f.as_bool()) { - // Extract categories if they exist and if content is flagged - let categories = if flagged { - first_result - .get("categories") - .and_then(|cats| cats.as_object()) - .map(|cats| { - cats.iter() - .filter(|(_, v)| v.as_bool().unwrap_or(false)) - .map(|(k, _)| k.to_string()) - .collect::>() - }) - } else { - None - }; - - return Ok(ModerationResult::new(flagged, categories, None)); - } - } - } - - // If we get here, there was no moderation result, so the content is considered safe - Ok(ModerationResult::new(false, None, None)) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::message::MessageContent; - use crate::providers::mock_server::{ - create_mock_open_ai_response, TEST_INPUT_TOKENS, TEST_OUTPUT_TOKENS, TEST_TOTAL_TOKENS, - }; - - use serde_json::json; - use wiremock::matchers::{body_json, header, method, path}; - use wiremock::{Mock, MockServer, ResponseTemplate}; - - fn create_test_provider(mock_server: &MockServer) -> DatabricksProvider { - DatabricksProvider { - client: Client::builder().build().unwrap(), - host: mock_server.uri(), - auth: DatabricksAuth::Token("test_token".to_string()), - model: ModelConfig::new("my-databricks-model".to_string()), - image_format: ImageFormat::Anthropic, - } - } - - #[tokio::test] - async fn test_moderation_flagged_content() -> Result<()> { - // Start a mock server - let mock_server = MockServer::start().await; - - // Mock response for moderation with flagged content - let mock_response = json!({ - "usage": { - "prompt_tokens": 199, - "total_tokens": 199 - }, - "input_guardrail": [{ - "flagged": true, - "categories": { - "violent-crimes": false, - "non-violent-crimes": true, - "sex-crimes": false, - "child-exploitation": false, - "specialized-advice": false, - "privacy": false, - "intellectual-property": false, - "indiscriminate-weapons": false, - "hate": false, - "self-harm": false, - "sexual-content": false - } - }], - "finishReason": "input_guardrail_triggered" - }); - - // Set up the mock - Mock::given(method("POST")) - .and(path("/serving-endpoints/moderation/invocations")) - .and(header("Authorization", "Bearer test_token")) - .respond_with(ResponseTemplate::new(200).set_body_json(mock_response)) - .expect(1) - .mount(&mock_server) - .await; - - let provider = create_test_provider(&mock_server); - - // Test moderation - let result = provider.moderate_content("test content").await?; - - assert!(result.flagged); - assert_eq!(result.categories.unwrap(), vec!["non-violent-crimes"]); - assert!(result.category_scores.is_none()); - - Ok(()) - } - - #[tokio::test] - async fn test_moderation_safe_content() -> Result<()> { - // Start a mock server - let mock_server = MockServer::start().await; - - // Mock response for safe content (regular chat response) - let mock_response = json!({ - "usage": { - "prompt_tokens": 10, - "completion_tokens": 20, - "total_tokens": 30 - }, - "choices": [{ - "message": { - "role": "assistant", - "content": "This is a safe response" - }, - "finish_reason": "stop" - }] - }); - - // Set up the mock - Mock::given(method("POST")) - .and(path("/serving-endpoints/moderation/invocations")) - .and(header("Authorization", "Bearer test_token")) - .respond_with(ResponseTemplate::new(200).set_body_json(mock_response)) - .expect(1) - .mount(&mock_server) - .await; - - let provider = create_test_provider(&mock_server); - - // Test moderation - let result = provider.moderate_content("safe content").await?; - - assert!(!result.flagged); - assert!(result.categories.is_none()); - assert!(result.category_scores.is_none()); - - Ok(()) - } - - #[tokio::test] - async fn test_moderation_explicit_safe() -> Result<()> { - // Start a mock server - let mock_server = MockServer::start().await; - - // Mock response for explicitly safe content - let mock_response = json!({ - "usage": { - "prompt_tokens": 199, - "total_tokens": 199 - }, - "input_guardrail": [{ - "flagged": false, - "categories": { - "violent-crimes": false, - "non-violent-crimes": false, - "sex-crimes": false, - "child-exploitation": false, - "specialized-advice": false, - "privacy": false, - "intellectual-property": false, - "indiscriminate-weapons": false, - "hate": false, - "self-harm": false, - "sexual-content": false - } - }] - }); - - // Set up the mock - Mock::given(method("POST")) - .and(path("/serving-endpoints/moderation/invocations")) - .and(header("Authorization", "Bearer test_token")) - .respond_with(ResponseTemplate::new(200).set_body_json(mock_response)) - .expect(1) - .mount(&mock_server) - .await; - - let provider = create_test_provider(&mock_server); - - // Test moderation - let result = provider.moderate_content("explicitly safe content").await?; - - assert!(!result.flagged); - assert!(result.categories.is_none()); - assert!(result.category_scores.is_none()); - - Ok(()) - } - - #[tokio::test] - async fn test_databricks_completion_with_token() -> Result<()> { - // Start a mock server - let mock_server = MockServer::start().await; - - // Mock response for completion - let mock_response = create_mock_open_ai_response("my-databricks-model", "Hello!"); - - let _moderator_mock_response = json!({ - "usage": { - "prompt_tokens": 10, - "completion_tokens": 20, - "total_tokens": 30 - }, - "choices": [{ - "message": { - "role": "assistant", - "content": "This is a safe response" - }, - "finish_reason": "stop" - }] - }); - - // Expected request body - let system = "You are a helpful assistant."; - let expected_request_body = json!({ - "messages": [ - {"role": "system", "content": system}, - {"role": "user", "content": "Hello"} - ] - }); - // Set up the mock to intercept the request and respond with the mocked response - Mock::given(method("POST")) - .and(path("/serving-endpoints/my-databricks-model/invocations")) - .and(header("Authorization", "Bearer test_token")) - .and(body_json(expected_request_body.clone())) - .respond_with(ResponseTemplate::new(200).set_body_json(mock_response)) - .expect(1) // Expect exactly one matching request - .mount(&mock_server) - .await; - - let provider = create_test_provider(&mock_server); - - // Prepare input - let messages = vec![Message::user().with_text("Hello")]; - let tools = vec![]; // Empty tools list - - // Call the complete method - let (reply_message, reply_usage) = provider.complete(system, &messages, &tools).await?; - - // Assert the response - if let MessageContent::Text(text) = &reply_message.content[0] { - assert_eq!(text.text, "Hello!"); - } else { - panic!("Expected Text content"); - } - assert_eq!(reply_usage.usage.input_tokens, Some(TEST_INPUT_TOKENS)); - assert_eq!(reply_usage.usage.output_tokens, Some(TEST_OUTPUT_TOKENS)); - assert_eq!(reply_usage.usage.total_tokens, Some(TEST_TOTAL_TOKENS)); - - Ok(()) + get_usage(data) } } diff --git a/crates/goose/src/providers/formats/anthropic.rs b/crates/goose/src/providers/formats/anthropic.rs new file mode 100644 index 0000000000..bb92e4bd63 --- /dev/null +++ b/crates/goose/src/providers/formats/anthropic.rs @@ -0,0 +1,418 @@ +use crate::message::{Message, MessageContent}; +use crate::providers::base::Usage; +use crate::providers::configs::ModelConfig; +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 Anthropic's API message specification +pub fn format_messages(messages: &[Message]) -> Vec { + let mut anthropic_messages = Vec::new(); + + // Convert messages to Anthropic format + for message in messages { + let role = match message.role { + Role::User => "user", + Role::Assistant => "assistant", + }; + + let mut content = Vec::new(); + for msg_content in &message.content { + match msg_content { + MessageContent::Text(text) => { + content.push(json!({ + "type": "text", + "text": text.text + })); + } + MessageContent::ToolRequest(tool_request) => { + if let Ok(tool_call) = &tool_request.tool_call { + content.push(json!({ + "type": "tool_use", + "id": tool_request.id, + "name": tool_call.name, + "input": tool_call.arguments + })); + } + } + 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"); + + content.push(json!({ + "type": "tool_result", + "tool_use_id": tool_response.id, + "content": text + })); + } + } + MessageContent::Image(_) => continue, // Anthropic doesn't support image content yet + } + } + + // Skip messages with empty content + if !content.is_empty() { + anthropic_messages.push(json!({ + "role": role, + "content": content + })); + } + } + + // If no messages, add a default one + if anthropic_messages.is_empty() { + anthropic_messages.push(json!({ + "role": "user", + "content": [{ + "type": "text", + "text": "Ignore" + }] + })); + } + + // Add "cache_control" to the last and second-to-last "user" messages. + // During each turn, we mark the final message with cache_control so the conversation can be + // incrementally cached. The second-to-last user message is also marked for caching with the + // cache_control parameter, so that this checkpoint can read from the previous cache. + let mut user_count = 0; + for message in anthropic_messages.iter_mut().rev() { + if message.get("role") == Some(&json!("user")) { + if let Some(content) = message.get_mut("content") { + if let Some(content_array) = content.as_array_mut() { + if let Some(last_content) = content_array.last_mut() { + last_content + .as_object_mut() + .unwrap() + .insert("cache_control".to_string(), json!({ "type": "ephemeral" })); + } + } + } + user_count += 1; + if user_count >= 2 { + break; + } + } + } + + anthropic_messages +} + +/// Convert internal Tool format to Anthropic'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 { + if unique_tools.insert(tool.name.clone()) { + tool_specs.push(json!({ + "name": tool.name, + "description": tool.description, + "input_schema": tool.input_schema + })); + } + } + + // Add "cache_control" to the last tool spec, if any. This means that all tool definitions, + // will be cached as a single prefix. + if let Some(last_tool) = tool_specs.last_mut() { + last_tool + .as_object_mut() + .unwrap() + .insert("cache_control".to_string(), json!({ "type": "ephemeral" })); + } + + tool_specs +} + +/// Convert system message to Anthropic's API system specification +pub fn format_system(system: &str) -> Value { + json!([{ + "type": "text", + "text": system, + "cache_control": { "type": "ephemeral" } + }]) +} + +/// Convert Anthropic's API response to internal Message format +pub fn response_to_message(response: Value) -> Result { + let content_blocks = response + .get("content") + .and_then(|c| c.as_array()) + .ok_or_else(|| anyhow!("Invalid response format: missing content array"))?; + + let mut message = Message::assistant(); + + for block in content_blocks { + match block.get("type").and_then(|t| t.as_str()) { + Some("text") => { + if let Some(text) = block.get("text").and_then(|t| t.as_str()) { + message = message.with_text(text.to_string()); + } + } + Some("tool_use") => { + let id = block + .get("id") + .and_then(|i| i.as_str()) + .ok_or_else(|| anyhow!("Missing tool_use id"))?; + let name = block + .get("name") + .and_then(|n| n.as_str()) + .ok_or_else(|| anyhow!("Missing tool_use name"))?; + let input = block + .get("input") + .ok_or_else(|| anyhow!("Missing tool_use input"))?; + + let tool_call = ToolCall::new(name, input.clone()); + message = message.with_tool_request(id, Ok(tool_call)); + } + _ => continue, + } + } + + Ok(message) +} + +/// Extract usage information from Anthropic'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(i), Some(o)) => Some(i + o), + _ => None, + }; + + Ok(Usage::new(input_tokens, output_tokens, total_tokens)) + } else { + // If no usage data, return None for all values + Ok(Usage::new(None, None, None)) + } +} + +/// Create a complete request payload for Anthropic's API +pub fn create_request( + model_config: &ModelConfig, + system: &str, + messages: &[Message], + tools: &[Tool], +) -> Result { + let anthropic_messages = format_messages(messages); + let tool_specs = format_tools(tools); + let system_spec = format_system(system); + + // Check if we have any messages to send + if anthropic_messages.is_empty() { + return Err(anyhow!("No valid messages to send to Anthropic API")); + } + + let mut payload = json!({ + "model": model_config.model_name, + "messages": anthropic_messages, + "max_tokens": model_config.max_tokens.unwrap_or(4096) + }); + + // Add system message if present + if !system.is_empty() { + payload + .as_object_mut() + .unwrap() + .insert("system".to_string(), json!(system_spec)); + } + + // Add tools if present + if !tool_specs.is_empty() { + payload + .as_object_mut() + .unwrap() + .insert("tools".to_string(), json!(tool_specs)); + } + + // Add temperature if specified + if let Some(temp) = model_config.temperature { + payload + .as_object_mut() + .unwrap() + .insert("temperature".to_string(), json!(temp)); + } + + 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": [{ + "type": "text", + "text": "Hello! How can I assist you today?" + }], + "model": "claude-3-5-sonnet-latest", + "stop_reason": "end_turn", + "stop_sequence": null, + "usage": { + "input_tokens": 12, + "output_tokens": 15, + "cache_creation_input_tokens": 12, + "cache_read_input_tokens": 0 + } + }); + + 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)); + + Ok(()) + } + + #[test] + fn test_parse_tool_response() -> Result<()> { + let response = json!({ + "id": "msg_123", + "type": "message", + "role": "assistant", + "content": [{ + "type": "tool_use", + "id": "tool_1", + "name": "calculator", + "input": { + "expression": "2 + 2" + } + }], + "model": "claude-3-sonnet-20240229", + "stop_reason": "end_turn", + "stop_sequence": null, + "usage": { + "input_tokens": 15, + "output_tokens": 20, + "cache_creation_input_tokens": 15, + "cache_read_input_tokens": 0, + } + }); + + 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)); + + Ok(()) + } + + #[test] + fn test_message_to_anthropic_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"][0]["type"], "text"); + assert_eq!(spec[0]["content"][0]["text"], "Hello"); + assert_eq!(spec[1]["role"], "assistant"); + assert_eq!(spec[1]["content"][0]["text"], "Hi there"); + assert_eq!(spec[2]["role"], "user"); + assert_eq!(spec[2]["content"][0]["text"], "How are you?"); + } + + #[test] + fn test_tools_to_anthropic_spec() { + let tools = vec![ + Tool::new( + "calculator", + "Calculate mathematical expressions", + json!({ + "type": "object", + "properties": { + "expression": { + "type": "string", + "description": "The mathematical expression to evaluate" + } + } + }), + ), + Tool::new( + "weather", + "Get weather information", + json!({ + "type": "object", + "properties": { + "location": { + "type": "string", + "description": "The location to get weather for" + } + } + }), + ), + ]; + + let spec = format_tools(&tools); + + assert_eq!(spec.len(), 2); + assert_eq!(spec[0]["name"], "calculator"); + assert_eq!(spec[0]["description"], "Calculate mathematical expressions"); + assert_eq!(spec[1]["name"], "weather"); + assert_eq!(spec[1]["description"], "Get weather information"); + + // Verify cache control is added to last tool + assert!(spec[1].get("cache_control").is_some()); + } + + #[test] + fn test_system_to_anthropic_spec() { + let system = "You are a helpful assistant."; + let spec = format_system(system); + + assert!(spec.is_array()); + let spec_array = spec.as_array().unwrap(); + assert_eq!(spec_array.len(), 1); + assert_eq!(spec_array[0]["type"], "text"); + assert_eq!(spec_array[0]["text"], system); + assert!(spec_array[0].get("cache_control").is_some()); + } +} diff --git a/crates/goose/src/providers/formats/google.rs b/crates/goose/src/providers/formats/google.rs new file mode 100644 index 0000000000..9500379ac2 --- /dev/null +++ b/crates/goose/src/providers/formats/google.rs @@ -0,0 +1,531 @@ +use crate::message::{Message, MessageContent}; +use crate::providers::base::Usage; +use crate::providers::configs::ModelConfig; +use crate::providers::utils::{is_valid_function_name, sanitize_function_name}; +use anyhow::Result; +use mcp_core::content::Content; +use mcp_core::role::Role; +use mcp_core::tool::{Tool, ToolCall}; +use serde_json::{json, Map, Value}; + +/// Convert internal Message format to Google's API message specification +pub fn format_messages(messages: &[Message]) -> Vec { + messages + .iter() + .map(|message| { + let role = if message.role == Role::User { + "user" + } else { + "model" + }; + let mut parts = Vec::new(); + for message_content in message.content.iter() { + match message_content { + MessageContent::Text(text) => { + if !text.text.is_empty() { + parts.push(json!({"text": text.text})); + } + } + MessageContent::ToolRequest(request) => match &request.tool_call { + Ok(tool_call) => { + let mut function_call_part = Map::new(); + function_call_part.insert( + "name".to_string(), + json!(sanitize_function_name(&tool_call.name)), + ); + if tool_call.arguments.is_object() + && !tool_call.arguments.as_object().unwrap().is_empty() + { + function_call_part + .insert("args".to_string(), tool_call.arguments.clone()); + } + parts.push(json!({ + "functionCall": function_call_part + })); + } + Err(e) => { + parts.push(json!({"text":format!("Error: {}", e)})); + } + }, + MessageContent::ToolResponse(response) => { + match &response.tool_result { + Ok(contents) => { + // Send only contents with no audience or with Assistant in the audience + let abridged: Vec<_> = contents + .iter() + .filter(|content| { + content.audience().is_none_or(|audience| { + audience.contains(&Role::Assistant) + }) + }) + .map(|content| content.unannotated()) + .collect(); + + for content in abridged { + match content { + Content::Image(image) => { + parts.push(json!({ + "inline_data": { + "mime_type": image.mime_type, + "data": image.data, + } + })); + } + _ => { + parts.push(json!({ + "functionResponse": { + "name": response.id, + "response": {"content": content}, + }} + )); + } + } + } + } + Err(e) => { + parts.push(json!({"text":format!("Error: {}", e)})); + } + } + } + + _ => {} + } + } + json!({"role": role, "parts": parts}) + }) + .collect() +} + +/// Convert internal Tool format to Google's API tool specification +pub fn format_tools(tools: &[Tool]) -> Vec { + tools + .iter() + .map(|tool| { + let mut parameters = Map::new(); + parameters.insert("name".to_string(), json!(tool.name)); + parameters.insert("description".to_string(), json!(tool.description)); + let tool_input_schema = tool.input_schema.as_object().unwrap(); + let tool_input_schema_properties = tool_input_schema + .get("properties") + .unwrap_or(&json!({})) + .as_object() + .unwrap() + .clone(); + if !tool_input_schema_properties.is_empty() { + let accepted_tool_schema_attributes = vec![ + "type".to_string(), + "format".to_string(), + "description".to_string(), + "nullable".to_string(), + "enum".to_string(), + "maxItems".to_string(), + "minItems".to_string(), + "properties".to_string(), + "required".to_string(), + "items".to_string(), + ]; + parameters.insert( + "parameters".to_string(), + json!(process_map( + tool_input_schema, + &accepted_tool_schema_attributes, + None + )), + ); + } + json!(parameters) + }) + .collect() +} + +/// Process a JSON map to filter out unsupported attributes +fn process_map( + map: &Map, + accepted_keys: &[String], + parent_key: Option<&str>, +) -> Value { + let mut filtered_map: Map = map + .iter() + .filter_map(|(key, value)| { + let should_remove = !accepted_keys.contains(key) && parent_key != Some("properties"); + if should_remove { + return None; + } + // Process nested maps recursively + let filtered_value = match value { + Value::Object(nested_map) => process_map( + &nested_map + .iter() + .map(|(k, v)| (k.clone(), v.clone())) + .collect(), + accepted_keys, + Some(key), + ), + _ => value.clone(), + }; + + Some((key.clone(), filtered_value)) + }) + .collect(); + if parent_key != Some("properties") && !filtered_map.contains_key("type") { + filtered_map.insert("type".to_string(), Value::String("string".to_string())); + } + + Value::Object(filtered_map) +} + +/// Convert Google's API response to internal Message format +pub fn response_to_message(response: Value) -> Result { + let mut content = Vec::new(); + let binding = vec![]; + let candidates: &Vec = response + .get("candidates") + .and_then(|v| v.as_array()) + .unwrap_or(&binding); + let candidate = candidates.first(); + let role = Role::Assistant; + let created = chrono::Utc::now().timestamp(); + if candidate.is_none() { + return Ok(Message { + role, + created, + content, + }); + } + let candidate = candidate.unwrap(); + let parts = candidate + .get("content") + .and_then(|content| content.get("parts")) + .and_then(|parts| parts.as_array()) + .unwrap_or(&binding); + for part in parts { + if let Some(text) = part.get("text").and_then(|v| v.as_str()) { + content.push(MessageContent::text(text.to_string())); + } else if let Some(function_call) = part.get("functionCall") { + let id = function_call["name"] + .as_str() + .unwrap_or_default() + .to_string(); + let name = function_call["name"] + .as_str() + .unwrap_or_default() + .to_string(); + if !is_valid_function_name(&name) { + let error = mcp_core::ToolError::NotFound(format!( + "The provided function name '{}' had invalid characters, it must match this regex [a-zA-Z0-9_-]+", + name + )); + content.push(MessageContent::tool_request(id, Err(error))); + } else { + let parameters = function_call.get("args"); + if let Some(params) = parameters { + content.push(MessageContent::tool_request( + id, + Ok(ToolCall::new(&name, params.clone())), + )); + } + } + } + } + Ok(Message { + role, + created, + content, + }) +} + +/// Extract usage information from Google's API response +pub fn get_usage(data: &Value) -> Result { + if let Some(usage_meta_data) = data.get("usageMetadata") { + let input_tokens = usage_meta_data + .get("promptTokenCount") + .and_then(|v| v.as_u64()) + .map(|v| v as i32); + let output_tokens = usage_meta_data + .get("candidatesTokenCount") + .and_then(|v| v.as_u64()) + .map(|v| v as i32); + let total_tokens = usage_meta_data + .get("totalTokenCount") + .and_then(|v| v.as_u64()) + .map(|v| v as i32); + Ok(Usage::new(input_tokens, output_tokens, total_tokens)) + } else { + // If no usage data, return None for all values + Ok(Usage::new(None, None, None)) + } +} + +/// Create a complete request payload for Google's API +pub fn create_request( + model_config: &ModelConfig, + system: &str, + messages: &[Message], + tools: &[Tool], +) -> Result { + let mut payload = Map::new(); + payload.insert( + "system_instruction".to_string(), + json!({"parts": [{"text": system}]}), + ); + payload.insert("contents".to_string(), json!(format_messages(messages))); + if !tools.is_empty() { + payload.insert( + "tools".to_string(), + json!({"functionDeclarations": format_tools(tools)}), + ); + } + let mut generation_config = Map::new(); + if let Some(temp) = model_config.temperature { + generation_config.insert("temperature".to_string(), json!(temp)); + } + if let Some(tokens) = model_config.max_tokens { + generation_config.insert("maxOutputTokens".to_string(), json!(tokens)); + } + if !generation_config.is_empty() { + payload.insert("generationConfig".to_string(), json!(generation_config)); + } + + Ok(Value::Object(payload)) +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + fn set_up_text_message(text: &str, role: Role) -> Message { + Message { + role, + created: 0, + content: vec![MessageContent::text(text.to_string())], + } + } + + fn set_up_tool_request_message(id: &str, tool_call: ToolCall) -> Message { + Message { + role: Role::User, + created: 0, + content: vec![MessageContent::tool_request(id.to_string(), Ok(tool_call))], + } + } + + fn set_up_tool_response_message(id: &str, tool_response: Vec) -> Message { + Message { + role: Role::Assistant, + created: 0, + content: vec![MessageContent::tool_response( + id.to_string(), + Ok(tool_response), + )], + } + } + + fn set_up_tool(name: &str, description: &str, params: Value) -> Tool { + Tool { + name: name.to_string(), + description: description.to_string(), + input_schema: json!({ + "properties": params + }), + } + } + + #[test] + fn test_get_usage() { + let data = json!({ + "usageMetadata": { + "promptTokenCount": 1, + "candidatesTokenCount": 2, + "totalTokenCount": 3 + } + }); + let usage = get_usage(&data).unwrap(); + assert_eq!(usage.input_tokens, Some(1)); + assert_eq!(usage.output_tokens, Some(2)); + assert_eq!(usage.total_tokens, Some(3)); + } + + #[test] + fn test_message_to_google_spec_text_message() { + let messages = vec![ + set_up_text_message("Hello", Role::User), + set_up_text_message("World", Role::Assistant), + ]; + let payload = format_messages(&messages); + assert_eq!(payload.len(), 2); + assert_eq!(payload[0]["role"], "user"); + assert_eq!(payload[0]["parts"][0]["text"], "Hello"); + assert_eq!(payload[1]["role"], "model"); + assert_eq!(payload[1]["parts"][0]["text"], "World"); + } + + #[test] + fn test_message_to_google_spec_tool_request_message() { + let arguments = json!({ + "param1": "value1" + }); + let messages = vec![set_up_tool_request_message( + "id", + ToolCall::new("tool_name", json!(arguments)), + )]; + let payload = format_messages(&messages); + assert_eq!(payload.len(), 1); + assert_eq!(payload[0]["role"], "user"); + assert_eq!(payload[0]["parts"][0]["functionCall"]["args"], arguments); + } + + #[test] + fn test_message_to_google_spec_tool_result_message() { + let tool_result: Vec = vec![Content::text("Hello")]; + let messages = vec![set_up_tool_response_message("response_id", tool_result)]; + let payload = format_messages(&messages); + assert_eq!(payload.len(), 1); + assert_eq!(payload[0]["role"], "model"); + assert_eq!( + payload[0]["parts"][0]["functionResponse"]["name"], + "response_id" + ); + assert_eq!( + payload[0]["parts"][0]["functionResponse"]["response"]["content"]["text"], + "Hello" + ); + } + + #[test] + fn test_tools_to_google_spec_with_valid_tools() { + let params1 = json!({ + "param1": { + "type": "string", + "description": "A parameter", + "field_does_not_accept": ["value1", "value2"] + } + }); + let params2 = json!({ + "param2": { + "type": "string", + "description": "B parameter", + } + }); + let tools = vec![ + set_up_tool("tool1", "description1", params1), + set_up_tool("tool2", "description2", params2), + ]; + let result = format_tools(&tools); + assert_eq!(result.len(), 2); + assert_eq!(result[0]["name"], "tool1"); + assert_eq!(result[0]["description"], "description1"); + assert_eq!( + result[0]["parameters"]["properties"], + json!({"param1": json!({ + "type": "string", + "description": "A parameter" + })}) + ); + assert_eq!(result[1]["name"], "tool2"); + assert_eq!(result[1]["description"], "description2"); + assert_eq!( + result[1]["parameters"]["properties"], + json!({"param2": json!({ + "type": "string", + "description": "B parameter" + })}) + ); + } + + #[test] + fn test_tools_to_google_spec_with_empty_properties() { + let tools = vec![Tool { + name: "tool1".to_string(), + description: "description1".to_string(), + input_schema: json!({ + "properties": {} + }), + }]; + let result = format_tools(&tools); + assert_eq!(result.len(), 1); + assert_eq!(result[0]["name"], "tool1"); + assert_eq!(result[0]["description"], "description1"); + assert!(result[0]["parameters"].get("properties").is_none()); + } + + #[test] + fn test_response_to_message_with_no_candidates() { + let response = json!({}); + let message = response_to_message(response).unwrap(); + assert_eq!(message.role, Role::Assistant); + assert!(message.content.is_empty()); + } + + #[test] + fn test_response_to_message_with_text_part() { + let response = json!({ + "candidates": [{ + "content": { + "parts": [{ + "text": "Hello, world!" + }] + } + }] + }); + let message = response_to_message(response).unwrap(); + assert_eq!(message.role, Role::Assistant); + assert_eq!(message.content.len(), 1); + if let MessageContent::Text(text) = &message.content[0] { + assert_eq!(text.text, "Hello, world!"); + } else { + panic!("Expected text content"); + } + } + + #[test] + fn test_response_to_message_with_invalid_function_name() { + let response = json!({ + "candidates": [{ + "content": { + "parts": [{ + "functionCall": { + "name": "invalid name!", + "args": {} + } + }] + } + }] + }); + let message = response_to_message(response).unwrap(); + assert_eq!(message.role, Role::Assistant); + assert_eq!(message.content.len(), 1); + if let Err(error) = &message.content[0].as_tool_request().unwrap().tool_call { + assert!(matches!(error, mcp_core::ToolError::NotFound(_))); + } else { + panic!("Expected tool request error"); + } + } + + #[test] + fn test_response_to_message_with_valid_function_call() { + let response = json!({ + "candidates": [{ + "content": { + "parts": [{ + "functionCall": { + "name": "valid_name", + "args": { + "param": "value" + } + } + }] + } + }] + }); + let message = response_to_message(response).unwrap(); + assert_eq!(message.role, Role::Assistant); + assert_eq!(message.content.len(), 1); + if let Ok(tool_call) = &message.content[0].as_tool_request().unwrap().tool_call { + assert_eq!(tool_call.name, "valid_name"); + assert_eq!(tool_call.arguments["param"], "value"); + } else { + panic!("Expected valid tool request"); + } + } +} diff --git a/crates/goose/src/providers/formats/mod.rs b/crates/goose/src/providers/formats/mod.rs new file mode 100644 index 0000000000..7134682858 --- /dev/null +++ b/crates/goose/src/providers/formats/mod.rs @@ -0,0 +1,3 @@ +pub mod anthropic; +pub mod google; +pub mod openai; diff --git a/crates/goose/src/providers/openai_utils.rs b/crates/goose/src/providers/formats/openai.rs similarity index 83% rename from crates/goose/src/providers/openai_utils.rs rename to crates/goose/src/providers/formats/openai.rs index 0d8ce78d1a..5bc9feba40 100644 --- a/crates/goose/src/providers/openai_utils.rs +++ b/crates/goose/src/providers/formats/openai.rs @@ -13,11 +13,7 @@ use serde_json::{json, Value}; /// Convert internal Message format to OpenAI's API message specification /// some openai compatible endpoints use the anthropic image spec at the content level /// even though the message structure is otherwise following openai, the enum switches this -pub fn messages_to_openai_spec( - messages: &[Message], - image_format: &ImageFormat, - concat_tool_response_contents: bool, -) -> Vec { +pub fn format_messages(messages: &[Message], image_format: &ImageFormat) -> Vec { let mut messages_spec = Vec::new(); for message in messages { let mut converted = json!({ @@ -97,19 +93,14 @@ pub fn messages_to_openai_spec( } } } - let tool_response_content: Value = match concat_tool_response_contents { - true => { - json!(tool_content - .iter() - .map(|content| match content { - Content::Text(text) => text.text.clone(), - _ => String::new(), - }) - .collect::>() - .join(" ")) - } - false => json!(tool_content), - }; + let tool_response_content: Value = json!(tool_content + .iter() + .map(|content| match content { + Content::Text(text) => text.text.clone(), + _ => String::new(), + }) + .collect::>() + .join(" ")); // First add the tool response with all content output.push(json!({ @@ -147,7 +138,7 @@ pub fn messages_to_openai_spec( } /// Convert internal Tool format to OpenAI's API tool specification -pub fn tools_to_openai_spec(tools: &[Tool]) -> anyhow::Result> { +pub fn format_tools(tools: &[Tool]) -> anyhow::Result> { let mut tool_names = std::collections::HashSet::new(); let mut result = Vec::new(); @@ -170,7 +161,7 @@ pub fn tools_to_openai_spec(tools: &[Tool]) -> anyhow::Result> { } /// Convert OpenAI's API response to internal Message format -pub fn openai_response_to_message(response: Value) -> anyhow::Result { +pub fn response_to_message(response: Value) -> anyhow::Result { let original = response["choices"][0]["message"].clone(); let mut content = Vec::new(); @@ -227,7 +218,7 @@ pub fn openai_response_to_message(response: Value) -> anyhow::Result { }) } -pub fn get_openai_usage(data: &Value) -> anyhow::Result { +pub fn get_usage(data: &Value) -> anyhow::Result { let usage = data .get("usage") .ok_or_else(|| anyhow!("No usage data in response"))?; @@ -254,40 +245,21 @@ pub fn get_openai_usage(data: &Value) -> anyhow::Result { Ok(Usage::new(input_tokens, output_tokens, total_tokens)) } -pub fn create_openai_request_payload( +pub fn create_request( model_config: &ModelConfig, system: &str, messages: &[Message], tools: &[Tool], -) -> anyhow::Result { - create_openai_request_payload_handling_concat_response_content( - model_config, - system, - messages, - tools, - false, - ) -} - -fn create_openai_request_payload_handling_concat_response_content( - model_config: &ModelConfig, - system: &str, - messages: &[Message], - tools: &[Tool], - concat_tool_response_contents: bool, + image_format: &ImageFormat, ) -> anyhow::Result { let system_message = json!({ "role": "system", "content": system }); - let messages_spec = messages_to_openai_spec( - messages, - &ImageFormat::OpenAi, - concat_tool_response_contents, - ); + let messages_spec = format_messages(messages, image_format); let tools_spec = if !tools.is_empty() { - tools_to_openai_spec(tools)? + format_tools(tools)? } else { vec![] }; @@ -321,22 +293,7 @@ fn create_openai_request_payload_handling_concat_response_content( Ok(payload) } -pub fn create_openai_request_payload_with_concat_response_content( - model_config: &ModelConfig, - system: &str, - messages: &[Message], - tools: &[Tool], -) -> anyhow::Result { - create_openai_request_payload_handling_concat_response_content( - model_config, - system, - messages, - tools, - true, - ) -} - -pub fn check_openai_context_length_error(error: &Value) -> Option { +pub fn is_context_length_error(error: &Value) -> Option { let code = error.get("code")?.as_str()?; if code == "context_length_exceeded" || code == "string_above_max_length" { let message = error @@ -377,9 +334,9 @@ mod tests { }"#; #[test] - fn test_messages_to_openai_spec() -> anyhow::Result<()> { + fn test_format_messages() -> anyhow::Result<()> { let message = Message::user().with_text("Hello"); - let spec = messages_to_openai_spec(&[message], &ImageFormat::OpenAi, false); + let spec = format_messages(&[message], &ImageFormat::OpenAi); assert_eq!(spec.len(), 1); assert_eq!(spec[0]["role"], "user"); @@ -388,7 +345,7 @@ mod tests { } #[test] - fn test_tools_to_openai_spec() -> anyhow::Result<()> { + fn test_format_tools() -> anyhow::Result<()> { let tool = Tool::new( "test_tool", "A test tool", @@ -404,7 +361,7 @@ mod tests { }), ); - let spec = tools_to_openai_spec(&[tool])?; + let spec = format_tools(&[tool])?; assert_eq!(spec.len(), 1); assert_eq!(spec[0]["type"], "function"); @@ -413,7 +370,7 @@ mod tests { } #[test] - fn test_messages_to_openai_spec_complex() -> anyhow::Result<()> { + fn test_format_messages_complex() -> anyhow::Result<()> { let mut messages = vec![ Message::assistant().with_text("Hello!"), Message::user().with_text("How are you?"), @@ -433,7 +390,7 @@ mod tests { messages .push(Message::user().with_tool_response(tool_id, Ok(vec![Content::text("Result")]))); - let spec = messages_to_openai_spec(&messages, &ImageFormat::OpenAi, true); + let spec = format_messages(&messages, &ImageFormat::OpenAi); assert_eq!(spec.len(), 4); assert_eq!(spec[0]["role"], "assistant"); @@ -450,7 +407,7 @@ mod tests { } #[test] - fn test_messages_to_openai_spec_not_concat_tool_response_content() -> anyhow::Result<()> { + fn test_format_messages_multiple_content() -> anyhow::Result<()> { let mut messages = vec![Message::assistant().with_tool_request( "tool1", Ok(ToolCall::new("example", json!({"param1": "value1"}))), @@ -466,20 +423,20 @@ mod tests { messages .push(Message::user().with_tool_response(tool_id, Ok(vec![Content::text("Result")]))); - let spec = messages_to_openai_spec(&messages, &ImageFormat::OpenAi, false); + let spec = format_messages(&messages, &ImageFormat::OpenAi); assert_eq!(spec.len(), 2); assert_eq!(spec[0]["role"], "assistant"); assert!(spec[0]["tool_calls"].is_array()); assert_eq!(spec[1]["role"], "tool"); - assert_eq!(spec[1]["content"][0]["text"], "Result"); + assert_eq!(spec[1]["content"], "Result"); assert_eq!(spec[1]["tool_call_id"], spec[0]["tool_calls"][0]["id"]); Ok(()) } #[test] - fn test_tools_to_openai_spec_duplicate() -> anyhow::Result<()> { + fn test_format_tools_duplicate() -> anyhow::Result<()> { let tool1 = Tool::new( "test_tool", "Test tool", @@ -510,7 +467,7 @@ mod tests { }), ); - let result = tools_to_openai_spec(&[tool1, tool2]); + let result = format_tools(&[tool1, tool2]); assert!(result.is_err()); assert!(result .unwrap_err() @@ -521,14 +478,14 @@ mod tests { } #[test] - fn test_tools_to_openai_spec_empty() -> anyhow::Result<()> { - let spec = tools_to_openai_spec(&[])?; + fn test_format_tools_empty() -> anyhow::Result<()> { + let spec = format_tools(&[])?; assert!(spec.is_empty()); Ok(()) } #[test] - fn test_openai_response_to_message_text() -> anyhow::Result<()> { + fn test_response_to_message_text() -> anyhow::Result<()> { let response = json!({ "choices": [{ "role": "assistant", @@ -543,7 +500,7 @@ mod tests { } }); - let message = openai_response_to_message(response)?; + let message = response_to_message(response)?; assert_eq!(message.content.len(), 1); if let MessageContent::Text(text) = &message.content[0] { assert_eq!(text.text, "Hello from John Cena!"); @@ -556,9 +513,9 @@ mod tests { } #[test] - fn test_openai_response_to_message_valid_toolrequest() -> anyhow::Result<()> { + fn test_response_to_message_valid_toolrequest() -> anyhow::Result<()> { let response: Value = serde_json::from_str(OPENAI_TOOL_USE_RESPONSE)?; - let message = openai_response_to_message(response)?; + let message = response_to_message(response)?; assert_eq!(message.content.len(), 1); if let MessageContent::ToolRequest(request) = &message.content[0] { @@ -573,12 +530,12 @@ mod tests { } #[test] - fn test_openai_response_to_message_invalid_func_name() -> anyhow::Result<()> { + fn test_response_to_message_invalid_func_name() -> anyhow::Result<()> { let mut response: Value = serde_json::from_str(OPENAI_TOOL_USE_RESPONSE)?; response["choices"][0]["message"]["tool_calls"][0]["function"]["name"] = json!("invalid fn"); - let message = openai_response_to_message(response)?; + let message = response_to_message(response)?; if let MessageContent::ToolRequest(request) = &message.content[0] { match &request.tool_call { @@ -595,12 +552,12 @@ mod tests { } #[test] - fn test_openai_response_to_message_json_decode_error() -> anyhow::Result<()> { + fn test_response_to_message_json_decode_error() -> anyhow::Result<()> { let mut response: Value = serde_json::from_str(OPENAI_TOOL_USE_RESPONSE)?; response["choices"][0]["message"]["tool_calls"][0]["function"]["arguments"] = json!("invalid json {"); - let message = openai_response_to_message(response)?; + let message = response_to_message(response)?; if let MessageContent::ToolRequest(request) = &message.content[0] { match &request.tool_call { @@ -617,13 +574,13 @@ mod tests { } #[test] - fn test_check_openai_context_length_error() { + fn test_is_context_length_error() { let error = json!({ "code": "context_length_exceeded", "message": "This message is too long" }); - let result = check_openai_context_length_error(&error); + let result = is_context_length_error(&error); assert!(result.is_some()); assert_eq!( result.unwrap().to_string(), @@ -635,7 +592,7 @@ mod tests { "message": "Some other error" }); - let result = check_openai_context_length_error(&error); + let result = is_context_length_error(&error); assert!(result.is_none()); } } diff --git a/crates/goose/src/providers/google.rs b/crates/goose/src/providers/google.rs index 95256e112e..5705362d56 100644 --- a/crates/goose/src/providers/google.rs +++ b/crates/goose/src/providers/google.rs @@ -1,19 +1,17 @@ -use crate::message::{Message, MessageContent}; -use crate::providers::base::{Moderation, ModerationResult, Provider, ProviderUsage, Usage}; +use crate::message::Message; +use crate::providers::base::{Provider, ProviderUsage, Usage}; use crate::providers::configs::ModelConfig; -use crate::providers::utils::{ - handle_response, is_valid_function_name, sanitize_function_name, unescape_json_values, -}; +use crate::providers::formats::google::{create_request, get_usage, response_to_message}; +use crate::providers::utils::{emit_debug_trace, handle_response, unescape_json_values}; use anyhow::Result; use async_trait::async_trait; -use mcp_core::ToolError; -use mcp_core::{Content, Role, Tool, ToolCall}; +use mcp_core::tool::Tool; use reqwest::Client; -use serde_json::{json, Map, Value}; +use serde_json::Value; use std::time::Duration; pub const GOOGLE_API_HOST: &str = "https://generativelanguage.googleapis.com"; -pub const GOOGLE_DEFAULT_MODEL: &str = "gemini-1.5-flash"; +pub const GOOGLE_DEFAULT_MODEL: &str = "gemini-2.0-flash-exp"; #[derive(Debug, serde::Serialize)] pub struct GoogleProvider { @@ -43,7 +41,7 @@ impl GoogleProvider { }) } - async fn post(&self, payload: Value) -> anyhow::Result { + async fn post(&self, payload: Value) -> Result { let url = format!( "{}/v1beta/models/{}:generateContent?key={}", self.host.trim_end_matches('/'), @@ -61,228 +59,6 @@ impl GoogleProvider { handle_response(payload, response).await } - - fn messages_to_google_spec(&self, messages: &[Message]) -> Vec { - messages - .iter() - .map(|message| { - let role = if message.role == Role::User { - "user" - } else { - "model" - }; - let mut parts = Vec::new(); - for message_content in message.content.iter() { - match message_content { - MessageContent::Text(text) => { - if !text.text.is_empty() { - parts.push(json!({"text": text.text})); - } - } - MessageContent::ToolRequest(request) => match &request.tool_call { - Ok(tool_call) => { - let mut function_call_part = Map::new(); - function_call_part.insert( - "name".to_string(), - json!(sanitize_function_name(&tool_call.name)), - ); - if tool_call.arguments.is_object() - && !tool_call.arguments.as_object().unwrap().is_empty() - { - function_call_part - .insert("args".to_string(), tool_call.arguments.clone()); - } - parts.push(json!({ - "functionCall": function_call_part - })); - } - Err(e) => { - parts.push(json!({"text":format!("Error: {}", e)})); - } - }, - MessageContent::ToolResponse(response) => { - match &response.tool_result { - Ok(contents) => { - // Send only contents with no audience or with Assistant in the audience - let abridged: Vec<_> = contents - .iter() - .filter(|content| { - content.audience().is_none_or(|audience| { - audience.contains(&Role::Assistant) - }) - }) - .map(|content| content.unannotated()) - .collect(); - - for content in abridged { - match content { - Content::Image(image) => { - parts.push(json!({ - "inline_data": { - "mime_type": image.mime_type, - "data": image.data, - } - })); - } - _ => { - parts.push(json!({ - "functionResponse": { - "name": response.id, - "response": {"content": content}, - }} - )); - } - } - } - } - Err(e) => { - parts.push(json!({"text":format!("Error: {}", e)})); - } - } - } - - _ => {} - } - } - json!({"role": role, "parts": parts}) - }) - .collect() - } - - fn tools_to_google_spec(&self, tools: &[Tool]) -> Vec { - tools - .iter() - .map(|tool| { - let mut parameters = Map::new(); - parameters.insert("name".to_string(), json!(tool.name)); - parameters.insert("description".to_string(), json!(tool.description)); - let tool_input_schema = tool.input_schema.as_object().unwrap(); - let tool_input_schema_properties = tool_input_schema - .get("properties") - .unwrap_or(&json!({})) - .as_object() - .unwrap() - .clone(); - if !tool_input_schema_properties.is_empty() { - let accepted_tool_schema_attributes = vec![ - "type".to_string(), - "format".to_string(), - "description".to_string(), - "nullable".to_string(), - "enum".to_string(), - "maxItems".to_string(), - "minItems".to_string(), - "properties".to_string(), - "required".to_string(), - "items".to_string(), - ]; - parameters.insert( - "parameters".to_string(), - json!(process_map( - tool_input_schema, - &accepted_tool_schema_attributes, - None - )), - ); - } - json!(parameters) - }) - .collect() - } - - fn google_response_to_message(&self, response: Value) -> anyhow::Result { - let mut content = Vec::new(); - let binding = vec![]; - let candidates: &Vec = response - .get("candidates") - .and_then(|v| v.as_array()) - .unwrap_or(&binding); - let candidate = candidates.first(); - let role = Role::Assistant; - let created = chrono::Utc::now().timestamp(); - if candidate.is_none() { - return Ok(Message { - role, - created, - content, - }); - } - let candidate = candidate.unwrap(); - let parts = candidate - .get("content") - .and_then(|content| content.get("parts")) - .and_then(|parts| parts.as_array()) - .unwrap_or(&binding); - for part in parts { - if let Some(text) = part.get("text").and_then(|v| v.as_str()) { - content.push(MessageContent::text(text.to_string())); - } else if let Some(function_call) = part.get("functionCall") { - let id = function_call["name"] - .as_str() - .unwrap_or_default() - .to_string(); - let name = function_call["name"] - .as_str() - .unwrap_or_default() - .to_string(); - if !is_valid_function_name(&name) { - let error = ToolError::NotFound(format!( - "The provided function name '{}' had invalid characters, it must match this regex [a-zA-Z0-9_-]+", - name - )); - content.push(MessageContent::tool_request(id, Err(error))); - } else { - let parameters = function_call.get("args"); - if let Some(params) = parameters { - content.push(MessageContent::tool_request( - id, - Ok(ToolCall::new(&name, params.clone())), - )); - } - } - } - } - Ok(Message { - role, - created, - content, - }) - } -} - -fn process_map( - map: &Map, - accepted_keys: &[String], - parent_key: Option<&str>, // Track the parent key -) -> Value { - let mut filtered_map: Map = map - .iter() - .filter_map(|(key, value)| { - let should_remove = !accepted_keys.contains(key) && parent_key != Some("properties"); - if should_remove { - return None; - } - // Process nested maps recursively - let filtered_value = match value { - Value::Object(nested_map) => process_map( - &nested_map - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(), - accepted_keys, - Some(key), - ), - _ => value.clone(), - }; - - Some((key.clone(), filtered_value)) - }) - .collect(); - if parent_key != Some("properties") && !filtered_map.contains_key("type") { - filtered_map.insert("type".to_string(), Value::String("string".to_string())); - } - - Value::Object(filtered_map) } #[async_trait] @@ -293,441 +69,32 @@ impl Provider for GoogleProvider { #[tracing::instrument( skip(self, system, messages, tools), - fields( - model_config, - input, - output, - input_tokens, - output_tokens, - total_tokens, - cost - ) + fields(model_config, input, output, input_tokens, output_tokens, total_tokens) )] - async fn complete_internal( + async fn complete( &self, system: &str, messages: &[Message], tools: &[Tool], - ) -> anyhow::Result<(Message, ProviderUsage)> { - let mut payload = Map::new(); - payload.insert( - "system_instruction".to_string(), - json!({"parts": [{"text": system}]}), - ); - payload.insert( - "contents".to_string(), - json!(self.messages_to_google_spec(messages)), - ); - if !tools.is_empty() { - payload.insert( - "tools".to_string(), - json!({"functionDeclarations": self.tools_to_google_spec(tools)}), - ); - } - let mut generation_config = Map::new(); - if let Some(temp) = self.model.temperature { - generation_config.insert("temperature".to_string(), json!(temp)); - } - if let Some(tokens) = self.model.max_tokens { - generation_config.insert("maxOutputTokens".to_string(), json!(tokens)); - } - if !generation_config.is_empty() { - payload.insert("generationConfig".to_string(), json!(generation_config)); - } + ) -> Result<(Message, ProviderUsage)> { + let payload = create_request(&self.model, system, messages, tools)?; // Make request - let response = self.post(Value::Object(payload.clone())).await?; + let response = self.post(payload.clone()).await?; + // Parse response - let message = self.google_response_to_message(unescape_json_values(&response))?; + let message = response_to_message(unescape_json_values(&response))?; let usage = self.get_usage(&response)?; let model = match response.get("modelVersion") { Some(model_version) => model_version.as_str().unwrap_or_default().to_string(), None => self.model.model_name.clone(), }; - super::utils::emit_debug_trace(self, &payload, &response, &usage, None); - let provider_usage = ProviderUsage::new(model, usage, None); + emit_debug_trace(self, &payload, &response, &usage); + let provider_usage = ProviderUsage::new(model, usage); Ok((message, provider_usage)) } - fn get_usage(&self, data: &Value) -> anyhow::Result { - if let Some(usage_meta_data) = data.get("usageMetadata") { - let input_tokens = usage_meta_data - .get("promptTokenCount") - .and_then(|v| v.as_u64()) - .map(|v| v as i32); - let output_tokens = usage_meta_data - .get("candidatesTokenCount") - .and_then(|v| v.as_u64()) - .map(|v| v as i32); - let total_tokens = usage_meta_data - .get("totalTokenCount") - .and_then(|v| v.as_u64()) - .map(|v| v as i32); - Ok(Usage::new(input_tokens, output_tokens, total_tokens)) - } else { - // If no usage data, return None for all values - Ok(Usage::new(None, None, None)) - } - } -} - -#[async_trait] -impl Moderation for GoogleProvider { - async fn moderate_content(&self, _content: &str) -> Result { - Ok(ModerationResult::new(false, None, None)) - } -} - -#[cfg(test)] // Only compiles this module when running tests -mod tests { - use super::*; - - use crate::providers::mock_server::{ - create_mock_google_ai_response_with_tools, create_test_tool, - get_expected_function_call_arguments, setup_mock_server, TEST_INPUT_TOKENS, - TEST_OUTPUT_TOKENS, TEST_TOOL_FUNCTION_NAME, TEST_TOTAL_TOKENS, - }; - use wiremock::MockServer; - - fn set_up_provider() -> GoogleProvider { - GoogleProvider { - client: Client::builder().build().unwrap(), - host: "dummy_host".to_string(), - api_key: "dummy_key".to_string(), - model: ModelConfig::new("dummy_model".to_string()), - } - } - - fn set_up_text_message(text: &str, role: Role) -> Message { - Message { - role, - created: 0, - content: vec![MessageContent::text(text.to_string())], - } - } - - fn set_up_tool_request_message(id: &str, tool_call: ToolCall) -> Message { - Message { - role: Role::User, - created: 0, - content: vec![MessageContent::tool_request(id.to_string(), Ok(tool_call))], - } - } - - fn set_up_tool_response_message(id: &str, tool_response: Vec) -> Message { - Message { - role: Role::Assistant, - created: 0, - content: vec![MessageContent::tool_response( - id.to_string(), - Ok(tool_response), - )], - } - } - - fn set_up_tool(name: &str, description: &str, params: Value) -> Tool { - Tool { - name: name.to_string(), - description: description.to_string(), - input_schema: json!({ - "properties": params - }), - } - } - - #[test] - fn test_get_usage() { - let provider = set_up_provider(); - let data = json!({ - "usageMetadata": { - "promptTokenCount": 1, - "candidatesTokenCount": 2, - "totalTokenCount": 3 - } - }); - let usage = provider.get_usage(&data).unwrap(); - assert_eq!(usage.input_tokens, Some(1)); - assert_eq!(usage.output_tokens, Some(2)); - assert_eq!(usage.total_tokens, Some(3)); - } - - #[test] - fn test_message_to_google_spec_text_message() { - let provider = set_up_provider(); - let messages = vec![ - set_up_text_message("Hello", Role::User), - set_up_text_message("World", Role::Assistant), - ]; - let payload = provider.messages_to_google_spec(&messages); - assert_eq!(payload.len(), 2); - assert_eq!(payload[0]["role"], "user"); - assert_eq!(payload[0]["parts"][0]["text"], "Hello"); - assert_eq!(payload[1]["role"], "model"); - assert_eq!(payload[1]["parts"][0]["text"], "World"); - } - - #[test] - fn test_message_to_google_spec_tool_request_message() { - let provider = set_up_provider(); - let arguments = json!({ - "param1": "value1" - }); - let messages = vec![set_up_tool_request_message( - "id", - ToolCall::new("tool_name", json!(arguments)), - )]; - let payload = provider.messages_to_google_spec(&messages); - assert_eq!(payload.len(), 1); - assert_eq!(payload[0]["role"], "user"); - assert_eq!(payload[0]["parts"][0]["functionCall"]["args"], arguments); - } - - #[test] - fn test_message_to_google_spec_tool_result_message() { - let provider = set_up_provider(); - let tool_result: Vec = vec![Content::text("Hello")]; - let messages = vec![set_up_tool_response_message("response_id", tool_result)]; - let payload = provider.messages_to_google_spec(&messages); - assert_eq!(payload.len(), 1); - assert_eq!(payload[0]["role"], "model"); - assert_eq!( - payload[0]["parts"][0]["functionResponse"]["name"], - "response_id" - ); - assert_eq!( - payload[0]["parts"][0]["functionResponse"]["response"]["content"]["text"], - "Hello" - ); - } - - #[test] - fn tools_to_google_spec_with_valid_tools() { - let provider = set_up_provider(); - let params1 = json!({ - "param1": { - "type": "string", - "description": "A parameter", - "field_does_not_accept": ["value1", "value2"] - } - }); - let params2 = json!({ - "param2": { - "type": "string", - "description": "B parameter", - } - }); - let tools = vec![ - set_up_tool("tool1", "description1", params1), - set_up_tool("tool2", "description2", params2), - ]; - let result = provider.tools_to_google_spec(&tools); - assert_eq!(result.len(), 2); - assert_eq!(result[0]["name"], "tool1"); - assert_eq!(result[0]["description"], "description1"); - assert_eq!( - result[0]["parameters"]["properties"], - json!({"param1": json!({ - "type": "string", - "description": "A parameter" - })}) - ); - assert_eq!(result[1]["name"], "tool2"); - assert_eq!(result[1]["description"], "description2"); - assert_eq!( - result[1]["parameters"]["properties"], - json!({"param2": json!({ - "type": "string", - "description": "B parameter" - })}) - ); - } - - #[test] - fn tools_to_google_spec_with_empty_properties() { - let provider = set_up_provider(); - let tools = vec![Tool { - name: "tool1".to_string(), - description: "description1".to_string(), - input_schema: json!({ - "properties": {} - }), - }]; - let result = provider.tools_to_google_spec(&tools); - assert_eq!(result.len(), 1); - assert_eq!(result[0]["name"], "tool1"); - assert_eq!(result[0]["description"], "description1"); - assert!(result[0]["parameters"].get("properties").is_none()); - } - - #[test] - fn google_response_to_message_with_no_candidates() { - let provider = set_up_provider(); - let response = json!({}); - let message = provider.google_response_to_message(response).unwrap(); - assert_eq!(message.role, Role::Assistant); - assert!(message.content.is_empty()); - } - - #[test] - fn google_response_to_message_with_text_part() { - let provider = set_up_provider(); - let response = json!({ - "candidates": [{ - "content": { - "parts": [{ - "text": "Hello, world!" - }] - } - }] - }); - let message = provider.google_response_to_message(response).unwrap(); - assert_eq!(message.role, Role::Assistant); - assert_eq!(message.content.len(), 1); - if let MessageContent::Text(text) = &message.content[0] { - assert_eq!(text.text, "Hello, world!"); - } else { - panic!("Expected text content"); - } - } - - #[test] - fn google_response_to_message_with_invalid_function_name() { - let provider = set_up_provider(); - let response = json!({ - "candidates": [{ - "content": { - "parts": [{ - "functionCall": { - "name": "invalid name!", - "args": {} - } - }] - } - }] - }); - let message = provider.google_response_to_message(response).unwrap(); - assert_eq!(message.role, Role::Assistant); - assert_eq!(message.content.len(), 1); - if let Err(error) = &message.content[0].as_tool_request().unwrap().tool_call { - assert!(matches!(error, ToolError::NotFound(_))); - } else { - panic!("Expected tool request error"); - } - } - - #[test] - fn google_response_to_message_with_valid_function_call() { - let provider = set_up_provider(); - let response = json!({ - "candidates": [{ - "content": { - "parts": [{ - "functionCall": { - "name": "valid_name", - "args": { - "param": "value" - } - } - }] - } - }] - }); - let message = provider.google_response_to_message(response).unwrap(); - assert_eq!(message.role, Role::Assistant); - assert_eq!(message.content.len(), 1); - if let Ok(tool_call) = &message.content[0].as_tool_request().unwrap().tool_call { - assert_eq!(tool_call.name, "valid_name"); - assert_eq!(tool_call.arguments["param"], "value"); - } else { - panic!("Expected valid tool request"); - } - } - - async fn _setup_mock_server( - model_name: &str, - response_body: Value, - ) -> (MockServer, GoogleProvider) { - let path_url = format!("/v1beta/models/{}:generateContent", model_name); - let mock_server = setup_mock_server(&path_url, response_body).await; - - let provider = GoogleProvider { - client: Client::builder().build().unwrap(), - host: mock_server.uri(), - api_key: "test_api_key".to_string(), - model: ModelConfig::new(GOOGLE_DEFAULT_MODEL.to_string()), - }; - - (mock_server, provider) - } - - // TODO Fix this test, it's failing in CI, but not locally - // #[tokio::test] - // async fn test_complete_basic() -> anyhow::Result<()> { - // let model_name = "gemini-1.5-flash"; - // // Mock response for normal completion - // let response_body = - // create_mock_google_ai_response(model_name, "Hello! How can I assist you today?"); - - // let (_, provider) = _setup_mock_server(model_name, response_body).await; - - // // Prepare input messages - // let messages = vec![Message::user().with_text("Hello?")]; - - // // Call the complete method - // let (message, usage) = provider - // .complete_internal("You are a helpful assistant.", &messages, &[]) - // .await?; - - // // Assert the response - // if let MessageContent::Text(text) = &message.content[0] { - // println!("text: {:?}", text); - // println!("text: {:?}", text.text); - // assert_eq!(text.text, "Hello! How can I assist you today?"); - // } else { - // panic!("Expected Text content"); - // } - // assert_eq!(usage.usage.input_tokens, Some(TEST_INPUT_TOKENS)); - // assert_eq!(usage.usage.output_tokens, Some(TEST_OUTPUT_TOKENS)); - // assert_eq!(usage.usage.total_tokens, Some(TEST_TOTAL_TOKENS)); - // assert_eq!(usage.model, model_name); - // assert_eq!(usage.cost, None); - - // Ok(()) - // } - - #[tokio::test] - async fn test_complete_tool_request() -> anyhow::Result<()> { - let model_name = "gemini-1.5-flash"; - // Mock response for tool calling - let response_body = create_mock_google_ai_response_with_tools("gpt-4o"); - - let (_, provider) = _setup_mock_server(model_name, response_body).await; - - // Input messages - let messages = vec![Message::user().with_text("What's the weather in San Francisco?")]; - - // Call the complete method - let (message, usage) = provider - .complete_internal( - "You are a helpful assistant.", - &messages, - &[create_test_tool()], - ) - .await?; - - // Assert the 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, TEST_TOOL_FUNCTION_NAME); - assert_eq!(tool_call.arguments, get_expected_function_call_arguments()); - } else { - panic!("Expected ToolCall content"); - } - - assert_eq!(usage.usage.input_tokens, Some(TEST_INPUT_TOKENS)); - assert_eq!(usage.usage.output_tokens, Some(TEST_OUTPUT_TOKENS)); - assert_eq!(usage.usage.total_tokens, Some(TEST_TOTAL_TOKENS)); - - Ok(()) + fn get_usage(&self, data: &Value) -> Result { + get_usage(data) } } diff --git a/crates/goose/src/providers/groq.rs b/crates/goose/src/providers/groq.rs index f768dce114..c84c79ddaf 100644 --- a/crates/goose/src/providers/groq.rs +++ b/crates/goose/src/providers/groq.rs @@ -1,10 +1,7 @@ use crate::message::Message; -use crate::providers::base::{Moderation, ModerationResult, Provider, ProviderUsage, Usage}; +use crate::providers::base::{Provider, ProviderUsage, Usage}; use crate::providers::configs::ModelConfig; -use crate::providers::openai_utils::{ - create_openai_request_payload_with_concat_response_content, get_openai_usage, - openai_response_to_message, -}; +use crate::providers::formats::openai::{create_request, get_usage, response_to_message}; use crate::providers::utils::{get_model, handle_response}; use anyhow::Result; use async_trait::async_trait; @@ -70,161 +67,32 @@ impl Provider for GroqProvider { #[tracing::instrument( skip(self, system, messages, tools), - fields( - model_config, - input, - output, - input_tokens, - output_tokens, - total_tokens, - cost - ) + fields(model_config, input, output, input_tokens, output_tokens, total_tokens) )] - async fn complete_internal( + async fn complete( &self, system: &str, messages: &[Message], tools: &[Tool], ) -> anyhow::Result<(Message, ProviderUsage)> { - let payload = create_openai_request_payload_with_concat_response_content( + let payload = create_request( &self.model, system, messages, tools, + &super::utils::ImageFormat::OpenAi, )?; let response = self.post(payload.clone()).await?; - let message = openai_response_to_message(response.clone())?; + let message = response_to_message(response.clone())?; let usage = self.get_usage(&response)?; let model = get_model(&response); - super::utils::emit_debug_trace(self, &payload, &response, &usage, None); - Ok((message, ProviderUsage::new(model, usage, None))) + super::utils::emit_debug_trace(self, &payload, &response, &usage); + Ok((message, ProviderUsage::new(model, usage))) } fn get_usage(&self, data: &Value) -> anyhow::Result { - get_openai_usage(data) - } -} - -#[async_trait] -impl Moderation for GroqProvider { - async fn moderate_content(&self, _content: &str) -> Result { - Ok(ModerationResult::new(false, None, None)) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::message::MessageContent; - use crate::providers::mock_server::{ - create_mock_open_ai_response, create_mock_open_ai_response_with_tools, create_test_tool, - get_expected_function_call_arguments, setup_mock_server, - setup_mock_server_with_response_code, TEST_INPUT_TOKENS, TEST_OUTPUT_TOKENS, - TEST_TOOL_FUNCTION_NAME, TEST_TOTAL_TOKENS, - }; - use wiremock::MockServer; - - async fn _setup_mock_server(response_body: Value) -> (MockServer, GroqProvider) { - let mock_server = setup_mock_server("/openai/v1/chat/completions", response_body).await; - - let provider = GroqProvider { - client: Client::builder().build().unwrap(), - host: mock_server.uri(), - api_key: "test_api_key".to_string(), - model: ModelConfig::new(GROQ_DEFAULT_MODEL.to_string()), - }; - - (mock_server, provider) - } - - #[tokio::test] - async fn test_complete_basic() -> anyhow::Result<()> { - let model_name = "gpt-4o"; - let expected_response = "Hello! How can I assist you today?"; - let response_body = create_mock_open_ai_response(model_name, expected_response); - - let (mock_server, provider) = _setup_mock_server(response_body).await; - - let messages = vec![Message::user().with_text("Hello?")]; - let (message, usage) = provider - .complete("You are a helpful assistant.", &messages, &[]) - .await?; - - assert!( - !message.content.is_empty(), - "Message content should not be empty" - ); - match &message.content[0] { - MessageContent::Text(text) => { - assert_eq!( - text.text, expected_response, - "Response text does not match expected" - ); - } - other => panic!("Expected Text content, got {:?}", other), - } - - assert_eq!(usage.usage.input_tokens, Some(TEST_INPUT_TOKENS)); - assert_eq!(usage.usage.output_tokens, Some(TEST_OUTPUT_TOKENS)); - assert_eq!(usage.usage.total_tokens, Some(TEST_TOTAL_TOKENS)); - assert_eq!(usage.model, model_name); - assert_eq!(usage.cost, None); - - mock_server.verify().await; - Ok(()) - } - - #[tokio::test] - async fn test_complete_tool_request() -> anyhow::Result<()> { - let response_body = create_mock_open_ai_response_with_tools("gpt-4o"); - let (mock_server, provider) = _setup_mock_server(response_body).await; - - let messages = vec![Message::user().with_text("What's the weather in San Francisco?")]; - let (message, usage) = provider - .complete( - "You are a helpful assistant.", - &messages, - &[create_test_tool()], - ) - .await?; - - if let MessageContent::ToolRequest(tool_request) = &message.content[0] { - let tool_call = tool_request.tool_call.as_ref().unwrap(); - assert_eq!(tool_call.name, TEST_TOOL_FUNCTION_NAME); - assert_eq!(tool_call.arguments, get_expected_function_call_arguments()); - } else { - panic!("Expected ToolCall content"); - } - - assert_eq!(usage.usage.input_tokens, Some(TEST_INPUT_TOKENS)); - assert_eq!(usage.usage.output_tokens, Some(TEST_OUTPUT_TOKENS)); - assert_eq!(usage.usage.total_tokens, Some(TEST_TOTAL_TOKENS)); - - mock_server.verify().await; - Ok(()) - } - - #[tokio::test] - async fn test_server_error() -> anyhow::Result<()> { - let mock_server = - setup_mock_server_with_response_code("/openai/v1/chat/completions", 500).await; - - let provider = GroqProvider { - client: Client::builder().build().unwrap(), - host: mock_server.uri(), - api_key: "test_api_key".to_string(), - model: ModelConfig::new(GROQ_DEFAULT_MODEL.to_string()), - }; - let messages = vec![Message::user().with_text("Hello?")]; - let result = provider - .complete("You are a helpful assistant.", &messages, &[]) - .await; - - assert!(result.is_err()); - assert!(result.unwrap_err().to_string().starts_with("Server error")); - - Ok(()) + get_usage(data) } } diff --git a/crates/goose/src/providers/mock.rs b/crates/goose/src/providers/mock.rs deleted file mode 100644 index e37ec8aa78..0000000000 --- a/crates/goose/src/providers/mock.rs +++ /dev/null @@ -1,75 +0,0 @@ -use super::base::{Moderation, ModerationResult, ProviderUsage}; -use crate::message::Message; -use crate::providers::base::{Provider, Usage}; -use crate::providers::configs::ModelConfig; -use anyhow::Result; -use async_trait::async_trait; -use mcp_core::tool::Tool; -use rust_decimal_macros::dec; -use serde_json::Value; -use std::sync::Arc; -use std::sync::Mutex; - -/// A mock provider that returns pre-configured responses for testing -pub struct MockProvider { - responses: Arc>>, - model_config: ModelConfig, -} - -impl MockProvider { - /// Create a new mock provider with a sequence of responses - pub fn new(responses: Vec) -> Self { - Self { - responses: Arc::new(Mutex::new(responses)), - model_config: ModelConfig::new("mock".to_string()), - } - } - - /// Create a new mock provider with specific responses and model config - pub fn with_config(responses: Vec, model_config: ModelConfig) -> Self { - Self { - responses: Arc::new(Mutex::new(responses)), - model_config, - } - } -} - -#[async_trait] -impl Provider for MockProvider { - fn get_model_config(&self) -> &ModelConfig { - &self.model_config - } - - async fn complete_internal( - &self, - _system_prompt: &str, - _messages: &[Message], - _tools: &[Tool], - ) -> Result<(Message, ProviderUsage)> { - let mut responses = self.responses.lock().unwrap(); - let usage = Usage::new(Some(1), Some(1), Some(2)); - if responses.is_empty() { - // Return empty response if no more pre-configured responses - Ok(( - Message::assistant().with_text(""), - ProviderUsage::new("mock".to_string(), usage, Some(dec!(1))), - )) - } else { - Ok(( - responses.remove(0), - ProviderUsage::new("mock".to_string(), usage, Some(dec!(1))), - )) - } - } - - fn get_usage(&self, _data: &Value) -> Result { - Ok(Usage::new(None, None, None)) - } -} - -#[async_trait] -impl Moderation for MockProvider { - async fn moderate_content(&self, _content: &str) -> Result { - Ok(ModerationResult::new(false, None, None)) - } -} diff --git a/crates/goose/src/providers/mock_server.rs b/crates/goose/src/providers/mock_server.rs deleted file mode 100644 index 8712cb8635..0000000000 --- a/crates/goose/src/providers/mock_server.rs +++ /dev/null @@ -1,152 +0,0 @@ -use mcp_core::Tool; -use serde_json::{json, Value}; -use wiremock::matchers::{method, path}; -use wiremock::{Mock, MockServer, ResponseTemplate}; - -pub const TEST_INPUT_TOKENS: i32 = 12; -pub const TEST_OUTPUT_TOKENS: i32 = 15; -pub const TEST_TOTAL_TOKENS: i32 = 27; -pub const TEST_TOOL_FUNCTION_NAME: &str = "get_weather"; -pub const TEST_TOOL_FUNCTION_ARGUMENTS: &str = "{\"location\":\"San Francisco, CA\"}"; - -pub async fn setup_mock_server(path_url: &str, response_body: Value) -> MockServer { - let mock_server = MockServer::start().await; - Mock::given(method("POST")) - .and(path(path_url)) - .respond_with(ResponseTemplate::new(200).set_body_json(response_body)) - .mount(&mock_server) - .await; - mock_server -} - -pub async fn setup_mock_server_with_response_code( - path_url: &str, - response_code: u16, -) -> MockServer { - let mock_server = MockServer::start().await; - Mock::given(method("POST")) - .and(path(path_url)) - .respond_with(ResponseTemplate::new(response_code)) - .mount(&mock_server) - .await; - mock_server -} -pub fn create_mock_open_ai_response_with_tools(model_name: &str) -> Value { - json!({ - "id": "chatcmpl-123", - "object": "chat.completion", - "choices": [{ - "index": 0, - "message": { - "role": "assistant", - "content": null, - "tool_calls": [{ - "id": "call_123", - "type": "function", - "function": { - "name": TEST_TOOL_FUNCTION_NAME, - "arguments": TEST_TOOL_FUNCTION_ARGUMENTS - } - }] - }, - "finish_reason": "stop" - }], - "usage": { - "prompt_tokens": TEST_INPUT_TOKENS, - "completion_tokens": TEST_OUTPUT_TOKENS, - "total_tokens": TEST_TOTAL_TOKENS - }, - "model": model_name - }) -} - -pub fn create_mock_google_ai_response_with_tools(model_name: &str) -> Value { - json!({ - "candidates": [{ - "content": { - "parts": [{ - "functionCall": { - "name": TEST_TOOL_FUNCTION_NAME, - "args":{ - "location": "San Francisco, CA" - } - - } - }], - "role": "model" - }, - "finishReason": "STOP" - }], - "modelVersion": model_name, - "usageMetadata": { - "candidatesTokenCount": TEST_OUTPUT_TOKENS, - "promptTokenCount": TEST_INPUT_TOKENS, - "totalTokenCount": TEST_TOTAL_TOKENS - } - }) -} - -pub fn create_mock_google_ai_response(model_name: &str, content: &str) -> Value { - json!({ - "candidates": [{ - "content": { - "parts": [{ - "text": content - }], - "role": "model" - }, - "finishReason": "STOP" - }], - "modelVersion": model_name, - "usageMetadata": { - "candidatesTokenCount": TEST_OUTPUT_TOKENS, - "promptTokenCount": TEST_INPUT_TOKENS, - "totalTokenCount": TEST_TOTAL_TOKENS - } - }) -} - -pub fn create_mock_open_ai_response(model_name: &str, content: &str) -> Value { - json!({ - "id": "chatcmpl-123", - "object": "chat.completion", - "choices": [{ - "index": 0, - "message": { - "role": "assistant", - "content": content, - "tool_calls": null - }, - "finish_reason": "stop" - }], - "usage": { - "prompt_tokens": TEST_INPUT_TOKENS, - "completion_tokens": TEST_OUTPUT_TOKENS, - "total_tokens": TEST_TOTAL_TOKENS - }, - "model": model_name - }) -} - -pub fn create_test_tool() -> Tool { - Tool::new( - "get_weather", - "Gets the current weather for a location", - json!({ - "type": "object", - "properties": { - "location": { - "type": "string", - "description": "The city and state, e.g. New York, NY" - } - }, - "required": ["location"] - }), - ) -} - -pub fn get_expected_function_call_arguments() -> Value { - json!({ - "location": "San Francisco, CA" - }) -} diff --git a/crates/goose/src/providers.rs b/crates/goose/src/providers/mod.rs similarity index 72% rename from crates/goose/src/providers.rs rename to crates/goose/src/providers/mod.rs index 29f8d669bb..53c39520b4 100644 --- a/crates/goose/src/providers.rs +++ b/crates/goose/src/providers/mod.rs @@ -4,19 +4,13 @@ pub mod configs; pub mod databricks; pub mod errors; pub mod factory; -pub mod mock; -pub mod model_pricing; +pub mod formats; +pub mod google; +pub mod groq; pub mod oauth; pub mod ollama; pub mod openai; -pub mod openai_utils; +pub mod openrouter; pub mod utils; -pub mod google; -pub mod groq; -pub mod openrouter; - pub use factory::get_provider; - -#[cfg(test)] -pub mod mock_server; diff --git a/crates/goose/src/providers/model_pricing.rs b/crates/goose/src/providers/model_pricing.rs deleted file mode 100644 index fc5b074a5c..0000000000 --- a/crates/goose/src/providers/model_pricing.rs +++ /dev/null @@ -1,138 +0,0 @@ -use std::collections::HashMap; - -use rust_decimal::Decimal; -use rust_decimal_macros::dec; - -use super::base::{Pricing, Usage}; - -lazy_static::lazy_static! { - static ref MODEL_PRICING: HashMap = { - let mut m = HashMap::new(); - // Anthropic - m.insert("claude-3-5-sonnet-latest".to_string(), Pricing { - input_token_price: dec!(3), - output_token_price: dec!(15), - }); - m.insert("claude-3-5-sonnet-20241022".to_string(), Pricing { - input_token_price: dec!(3), - output_token_price: dec!(15), - }); - m.insert("anthropic.claude-3-5-sonnet-20241022-v2:0".to_string(), Pricing { - input_token_price: dec!(3), - output_token_price: dec!(15), - }); - m.insert("claude-3-5-sonnet-20241022-v2:0".to_string(), Pricing { - input_token_price: dec!(3), - output_token_price: dec!(15), - }); - m.insert("claude-3-5-sonnet-v2@20241022".to_string(), Pricing { - input_token_price: dec!(3), - output_token_price: dec!(15), - }); - m.insert("claude-3-5-haiku-latest".to_string(), Pricing { - input_token_price: dec!(0.8), - output_token_price: dec!(4), - }); - m.insert("claude-3-5-haiku-20241022".to_string(), Pricing { - input_token_price: dec!(0.8), - output_token_price: dec!(4), - }); - m.insert("anthropic.claude-3-5-haiku-20241022-v1:0".to_string(), Pricing { - input_token_price: dec!(0.8), - output_token_price: dec!(4), - }); - m.insert("claude-3-5-haiku@20241022".to_string(), Pricing { - input_token_price: dec!(0.8), - output_token_price: dec!(4), - }); - m.insert("claude-3-opus-latest".to_string(), Pricing { - input_token_price: dec!(15.00), - output_token_price: dec!(75.00), - }); - m.insert("claude-3-opus-20240229".to_string(), Pricing { - input_token_price: dec!(15.00), - output_token_price: dec!(75.00), - }); - m.insert("anthropic.claude-3-opus-20240229-v1:0".to_string(), Pricing { - input_token_price: dec!(15.00), - output_token_price: dec!(75.00), - }); - m.insert("claude-3-opus@20240229".to_string(), Pricing { - input_token_price: dec!(15.00), - output_token_price: dec!(75.00), - }); - // OpenRouter Models - m.insert("anthropic/claude-3-sonnet".to_string(), Pricing { - input_token_price: dec!(3.00), - output_token_price: dec!(15.00), - }); - m.insert("claude-3-sonnet".to_string(), Pricing { - input_token_price: dec!(3.00), - output_token_price: dec!(15.00), - }); - // OpenAI - m.insert("gpt-4o".to_string(), Pricing { - input_token_price: dec!(2.50), - output_token_price: dec!(10.00), - }); - m.insert("gpt-4".to_string(), Pricing { - input_token_price: dec!(2.50), - output_token_price: dec!(10.00), - }); - m.insert("gpt-4o-2024-11-20".to_string(), Pricing { - input_token_price: dec!(2.50), - output_token_price: dec!(10.00), - }); - m.insert("gpt-4o-2024-08-06".to_string(), Pricing { - input_token_price: dec!(2.50), - output_token_price: dec!(10.00), - }); - m.insert("gpt-4o-2024-05-13".to_string(), Pricing { - input_token_price: dec!(5.00), - output_token_price: dec!(15.00), - }); - m.insert("gpt-4o-mini".to_string(), Pricing { - input_token_price: dec!(0.150), - output_token_price: dec!(0.600), - }); - m.insert("gpt-4o-mini-2024-07-18".to_string(), Pricing { - input_token_price: dec!(0.150), - output_token_price: dec!(0.600), - }); - m.insert("o1-preview".to_string(), Pricing { - input_token_price: dec!(15.00), - output_token_price: dec!(60.00), - }); - m.insert("o1-preview-2024-09-12".to_string(), Pricing { - input_token_price: dec!(15.00), - output_token_price: dec!(60.00), - }); - m.insert("o1-mini".to_string(), Pricing { - input_token_price: dec!(3.00), - output_token_price: dec!(12.00), - }); - m.insert("o1-mini-2024-09-12".to_string(), Pricing { - input_token_price: dec!(3.00), - output_token_price: dec!(12.00), - }); - m - }; -} - -pub fn model_pricing_for(model: &str) -> Option { - MODEL_PRICING.get(model).cloned() -} - -pub fn cost(usage: &Usage, model_pricing: &Option) -> Option { - if let Some(model_pricing) = model_pricing { - let input_price = Decimal::from(usage.input_tokens.unwrap_or(0)) - * model_pricing.input_token_price - / Decimal::from(1_000_000); - let output_price = Decimal::from(usage.output_tokens.unwrap_or(0)) - * model_pricing.output_token_price - / Decimal::from(1_000_000); - Some(input_price + output_price) - } else { - None - } -} diff --git a/crates/goose/src/providers/ollama.rs b/crates/goose/src/providers/ollama.rs index a7a94d98fa..f49832e58e 100644 --- a/crates/goose/src/providers/ollama.rs +++ b/crates/goose/src/providers/ollama.rs @@ -1,10 +1,8 @@ -use super::base::{Moderation, ModerationResult, Provider, ProviderUsage, Usage}; +use super::base::{Provider, ProviderUsage, Usage}; use super::configs::ModelConfig; use super::utils::{get_model, handle_response}; use crate::message::Message; -use crate::providers::openai_utils::{ - create_openai_request_payload, get_openai_usage, openai_response_to_message, -}; +use crate::providers::formats::openai::{create_request, get_usage, response_to_message}; use anyhow::Result; use async_trait::async_trait; use mcp_core::tool::Tool; @@ -60,181 +58,33 @@ impl Provider for OllamaProvider { #[tracing::instrument( skip(self, system, messages, tools), - fields( - model_config, - input, - output, - input_tokens, - output_tokens, - total_tokens, - cost - ) + fields(model_config, input, output, input_tokens, output_tokens, total_tokens) )] - async fn complete_internal( + async fn complete( &self, system: &str, messages: &[Message], tools: &[Tool], ) -> Result<(Message, ProviderUsage)> { - let payload = create_openai_request_payload(&self.model, system, messages, tools)?; + let payload = create_request( + &self.model, + system, + messages, + tools, + &super::utils::ImageFormat::OpenAi, + )?; let response = self.post(payload.clone()).await?; // Parse response - let message = openai_response_to_message(response.clone())?; + let message = response_to_message(response.clone())?; let usage = self.get_usage(&response)?; let model = get_model(&response); - let cost = None; - super::utils::emit_debug_trace(self, &payload, &response, &usage, cost); - Ok((message, ProviderUsage::new(model, usage, cost))) + super::utils::emit_debug_trace(self, &payload, &response, &usage); + Ok((message, ProviderUsage::new(model, usage))) } fn get_usage(&self, data: &Value) -> Result { - get_openai_usage(data) - } -} - -#[async_trait] -impl Moderation for OllamaProvider { - async fn moderate_content(&self, _content: &str) -> Result { - Ok(ModerationResult::new(false, None, None)) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::message::MessageContent; - use crate::providers::mock_server::{ - create_mock_open_ai_response, create_mock_open_ai_response_with_tools, create_test_tool, - get_expected_function_call_arguments, setup_mock_server, - setup_mock_server_with_response_code, TEST_INPUT_TOKENS, TEST_OUTPUT_TOKENS, - TEST_TOOL_FUNCTION_NAME, TEST_TOTAL_TOKENS, - }; - use wiremock::MockServer; - - async fn _setup_mock_server(response_body: Value) -> (MockServer, OllamaProvider) { - let mock_server = setup_mock_server("/v1/chat/completions", response_body).await; - - let provider = OllamaProvider { - client: Client::builder().build().unwrap(), - host: mock_server.uri(), - model: ModelConfig::new(OLLAMA_MODEL.to_string()), - }; - - (mock_server, provider) - } - - #[tokio::test] - async fn test_complete_basic() -> Result<()> { - let model_name = "gpt-4o"; - let expected_response = "Hello! How can I assist you today?"; - // Mock response for normal completion - let response_body = create_mock_open_ai_response(model_name, expected_response); - - let (mock_server, provider) = _setup_mock_server(response_body).await; - - // Prepare input messages - let messages = vec![Message::user().with_text("Hello?")]; - - // Call the complete method - let (message, usage) = provider - .complete("You are a helpful assistant.", &messages, &[]) - .await?; - - // Assert the response - assert!( - !message.content.is_empty(), - "Message content should not be empty" - ); - match &message.content[0] { - MessageContent::Text(text) => { - assert_eq!( - text.text, expected_response, - "Response text does not match expected" - ); - } - other => panic!("Expected Text content, got {:?}", other), - } - - // Verify usage metrics - assert_eq!( - usage.usage.input_tokens, - Some(TEST_INPUT_TOKENS), - "Input tokens mismatch" - ); - assert_eq!( - usage.usage.output_tokens, - Some(TEST_OUTPUT_TOKENS), - "Output tokens mismatch" - ); - assert_eq!( - usage.usage.total_tokens, - Some(TEST_TOTAL_TOKENS), - "Total tokens mismatch" - ); - assert_eq!(usage.model, model_name, "Model name mismatch"); - assert_eq!(usage.cost, None, "Cost should be None"); - - // Ensure mock server handled the request - mock_server.verify().await; - - Ok(()) - } - - #[tokio::test] - async fn test_complete_tool_request() -> Result<()> { - // Mock response for tool calling - let response_body = create_mock_open_ai_response_with_tools("gpt-4o"); - - let (mock_server, provider) = _setup_mock_server(response_body).await; - - // Input messages - let messages = vec![Message::user().with_text("What's the weather in San Francisco?")]; - - // Call the complete method - let (message, usage) = provider - .complete( - "You are a helpful assistant.", - &messages, - &[create_test_tool()], - ) - .await?; - - // Assert the 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, TEST_TOOL_FUNCTION_NAME); - assert_eq!(tool_call.arguments, get_expected_function_call_arguments()); - } else { - panic!("Expected ToolCall content"); - } - - assert_eq!(usage.usage.input_tokens, Some(TEST_INPUT_TOKENS)); - assert_eq!(usage.usage.output_tokens, Some(TEST_OUTPUT_TOKENS)); - assert_eq!(usage.usage.total_tokens, Some(TEST_TOTAL_TOKENS)); - - mock_server.verify().await; - Ok(()) - } - - #[tokio::test] - async fn test_server_error() -> Result<()> { - let mock_server = setup_mock_server_with_response_code("/v1/chat/completions", 500).await; - - let provider = OllamaProvider { - client: Client::builder().build().unwrap(), - host: mock_server.uri(), - model: ModelConfig::new(OLLAMA_MODEL.to_string()), - }; - let messages = vec![Message::user().with_text("Hello?")]; - let result = provider - .complete("You are a helpful assistant.", &messages, &[]) - .await; - - assert!(result.is_err()); - assert!(result.unwrap_err().to_string().starts_with("Server error")); - - Ok(()) + get_usage(data) } } diff --git a/crates/goose/src/providers/openai.rs b/crates/goose/src/providers/openai.rs index 19d1c93155..6df238e746 100644 --- a/crates/goose/src/providers/openai.rs +++ b/crates/goose/src/providers/openai.rs @@ -4,22 +4,16 @@ use reqwest::Client; use serde_json::Value; use std::time::Duration; -use super::base::{Moderation, ModerationResult, ProviderUsage}; -use super::base::{Provider, Usage}; +use super::base::{Provider, ProviderUsage, Usage}; use super::configs::ModelConfig; -use super::model_pricing::cost; -use super::model_pricing::model_pricing_for; -use super::utils::{emit_debug_trace, get_model, handle_response}; -use crate::message::Message; -use crate::providers::openai_utils::{ - check_openai_context_length_error, create_openai_request_payload, get_openai_usage, - openai_response_to_message, +use super::formats::openai::{ + create_request, get_usage, is_context_length_error, response_to_message, }; +use super::utils::{emit_debug_trace, get_model, handle_response, ImageFormat}; +use crate::message::Message; use mcp_core::tool::Tool; -use serde::Serialize; pub const OPEN_AI_DEFAULT_MODEL: &str = "gpt-4o"; -pub const OPEN_AI_MODERATION_MODEL: &str = "omni-moderation-latest"; #[derive(Debug, serde::Serialize)] pub struct OpenAiProvider { @@ -30,18 +24,6 @@ pub struct OpenAiProvider { model: ModelConfig, } -#[derive(Serialize)] -struct OpenAiModerationRequest { - input: String, - model: String, -} - -impl OpenAiModerationRequest { - pub fn new(input: String, model: String) -> Self { - Self { input, model } - } -} - impl OpenAiProvider { pub fn from_env() -> Result { let api_key = crate::key_manager::get_keyring_secret("OPENAI_API_KEY", Default::default())?; @@ -85,185 +67,54 @@ impl Provider for OpenAiProvider { #[tracing::instrument( skip(self, system, messages, tools), - fields( - model_config, - input, - output, - input_tokens, - output_tokens, - total_tokens, - cost - ) + fields(model_config, input, output, input_tokens, output_tokens, total_tokens) )] - async fn complete_internal( + async fn complete( &self, system: &str, messages: &[Message], tools: &[Tool], ) -> Result<(Message, ProviderUsage)> { - // Not checking for o1 model here since system message is not supported by o1 - let payload = create_openai_request_payload(&self.model, system, messages, tools)?; + let payload = create_request(&self.model, system, messages, tools, &ImageFormat::OpenAi)?; // Make request let response = self.post(payload.clone()).await?; // Raise specific error if context length is exceeded if let Some(error) = response.get("error") { - if let Some(err) = check_openai_context_length_error(error) { + if let Some(err) = is_context_length_error(error) { return Err(err.into()); } return Err(anyhow!("OpenAI API error: {}", error)); } // Parse response - let message = openai_response_to_message(response.clone())?; + let message = response_to_message(response.clone())?; let usage = self.get_usage(&response)?; let model = get_model(&response); - let cost = cost(&usage, &model_pricing_for(&model)); - emit_debug_trace(self, &payload, &response, &usage, cost); - Ok((message, ProviderUsage::new(model, usage, cost))) + emit_debug_trace(self, &payload, &response, &usage); + Ok((message, ProviderUsage::new(model, usage))) } fn get_usage(&self, data: &Value) -> Result { - get_openai_usage(data) - } -} - -#[async_trait] -impl Moderation for OpenAiProvider { - async fn moderate_content_internal(&self, content: &str) -> Result { - let url = format!("{}/v1/moderations", self.host.trim_end_matches('/')); - - let request = - OpenAiModerationRequest::new(content.to_string(), OPEN_AI_MODERATION_MODEL.to_string()); - - let response = self - .client - .post(&url) - .header("Authorization", format!("Bearer {}", self.api_key)) - .json(&request) - .send() - .await?; - - let response_json = handle_response(serde_json::to_value(&request)?, response).await?; - - let flagged = response_json["results"][0]["flagged"] - .as_bool() - .unwrap_or(false); - if flagged { - let categories = response_json["results"][0]["categories"] - .as_object() - .unwrap(); - let category_scores = response_json["results"][0]["category_scores"].clone(); - return Ok(ModerationResult::new( - flagged, - Some( - categories - .iter() - .filter(|(_, value)| value.as_bool().unwrap_or(false)) - .map(|(key, _)| key.to_string()) - .collect(), - ), - Some(category_scores), - )); - } else { - return Ok(ModerationResult::new(flagged, None, None)); - } + get_usage(data) } } #[cfg(test)] mod tests { use super::*; - use crate::message::MessageContent; - use crate::providers::mock_server::{ - create_mock_open_ai_response, create_mock_open_ai_response_with_tools, create_test_tool, - get_expected_function_call_arguments, setup_mock_server, TEST_INPUT_TOKENS, - TEST_OUTPUT_TOKENS, TEST_TOOL_FUNCTION_NAME, TEST_TOTAL_TOKENS, - }; - use rust_decimal_macros::dec; - use wiremock::MockServer; - async fn _setup_mock_response(response_body: Value) -> (MockServer, OpenAiProvider) { - let mock_server = setup_mock_server("/v1/chat/completions", response_body).await; + #[test] + fn test_provider_construction() -> Result<()> { + std::env::set_var("OPENAI_API_KEY", "test-key"); + std::env::set_var("OPENAI_HOST", "https://test.openai.com"); + std::env::set_var("OPENAI_MODEL", "gpt-4o"); - // Create the OpenAiProvider with the mock server's URL as the host - let provider = OpenAiProvider { - client: Client::builder().build().unwrap(), - host: mock_server.uri(), - api_key: "test_api_key".to_string(), - model: ModelConfig::new("gpt-4o".to_string()).with_temperature(Some(0.7)), - }; + let provider = OpenAiProvider::from_env()?; + assert_eq!(provider.host, "https://test.openai.com"); + assert_eq!(provider.model.model_name, "gpt-4o"); - (mock_server, provider) - } - - #[tokio::test] - async fn test_complete_basic() -> Result<()> { - let model_name = "gpt-4o"; - // Mock response for normal completion - let response_body = - create_mock_open_ai_response(model_name, "Hello! How can I assist you today?"); - - let (mock_server, provider) = _setup_mock_response(response_body).await; - - // Prepare input messages - let messages = vec![Message::user().with_text("Hello?")]; - - // Call the complete method - let (message, usage) = provider - .complete_internal("You are a helpful assistant.", &messages, &[]) - .await?; - - // Assert the 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.usage.input_tokens, Some(TEST_INPUT_TOKENS)); - assert_eq!(usage.usage.output_tokens, Some(TEST_OUTPUT_TOKENS)); - assert_eq!(usage.usage.total_tokens, Some(TEST_TOTAL_TOKENS)); - assert_eq!(usage.model, model_name); - assert_eq!(usage.cost, Some(dec!(0.00018))); - - mock_server.verify().await; - Ok(()) - } - - #[tokio::test] - async fn test_complete_tool_request() -> Result<()> { - // Mock response for tool calling - let response_body = create_mock_open_ai_response_with_tools("gpt-4o"); - - let (mock_server, provider) = _setup_mock_response(response_body).await; - - // Input messages - let messages = vec![Message::user().with_text("What's the weather in San Francisco?")]; - - // Call the complete method - let (message, usage) = provider - .complete_internal( - "You are a helpful assistant.", - &messages, - &[create_test_tool()], - ) - .await?; - - // Assert the 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, TEST_TOOL_FUNCTION_NAME); - assert_eq!(tool_call.arguments, get_expected_function_call_arguments()); - } else { - panic!("Expected ToolCall content"); - } - - assert_eq!(usage.usage.input_tokens, Some(TEST_INPUT_TOKENS)); - assert_eq!(usage.usage.output_tokens, Some(TEST_OUTPUT_TOKENS)); - assert_eq!(usage.usage.total_tokens, Some(TEST_TOTAL_TOKENS)); - - mock_server.verify().await; Ok(()) } } diff --git a/crates/goose/src/providers/openrouter.rs b/crates/goose/src/providers/openrouter.rs index 9c33df3650..36c415cfe1 100644 --- a/crates/goose/src/providers/openrouter.rs +++ b/crates/goose/src/providers/openrouter.rs @@ -4,15 +4,12 @@ use reqwest::Client; use serde_json::Value; use std::time::Duration; -use super::base::{Moderation, ModerationResult, Provider, ProviderUsage, Usage}; +use super::base::{Provider, ProviderUsage, Usage}; use super::configs::ModelConfig; -use super::model_pricing::cost; -use super::model_pricing::model_pricing_for; use super::utils::{emit_debug_trace, get_model, handle_response}; use crate::message::Message; -use crate::providers::openai_utils::{ - check_openai_context_length_error, create_openai_request_payload_with_concat_response_content, - get_openai_usage, openai_response_to_message, +use crate::providers::formats::openai::{ + create_request, get_usage, is_context_length_error, response_to_message, }; use mcp_core::tool::Tool; @@ -77,28 +74,21 @@ impl Provider for OpenRouterProvider { #[tracing::instrument( skip(self, system, messages, tools), - fields( - model_config, - input, - output, - input_tokens, - output_tokens, - total_tokens, - cost - ) + fields(model_config, input, output, input_tokens, output_tokens, total_tokens) )] - async fn complete_internal( + async fn complete( &self, system: &str, messages: &[Message], tools: &[Tool], ) -> Result<(Message, ProviderUsage)> { // Create the base payload - let payload = create_openai_request_payload_with_concat_response_content( + let payload = create_request( &self.model, system, messages, tools, + &super::utils::ImageFormat::OpenAi, )?; // Make request @@ -106,123 +96,21 @@ impl Provider for OpenRouterProvider { // Raise specific error if context length is exceeded if let Some(error) = response.get("error") { - if let Some(err) = check_openai_context_length_error(error) { + if let Some(err) = is_context_length_error(error) { return Err(err.into()); } return Err(anyhow!("OpenRouter API error: {}", error)); } // Parse response - let message = openai_response_to_message(response.clone())?; + let message = response_to_message(response.clone())?; let usage = self.get_usage(&response)?; let model = get_model(&response); - let cost = cost(&usage, &model_pricing_for(&model)); - emit_debug_trace(self, &payload, &response, &usage, cost); - Ok((message, ProviderUsage::new(model, usage, cost))) + emit_debug_trace(self, &payload, &response, &usage); + Ok((message, ProviderUsage::new(model, usage))) } fn get_usage(&self, data: &Value) -> Result { - get_openai_usage(data) - } -} - -#[async_trait] -impl Moderation for OpenRouterProvider { - async fn moderate_content(&self, _content: &str) -> Result { - Ok(ModerationResult::new(false, None, None)) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::message::MessageContent; - use crate::providers::configs::ModelConfig; - use crate::providers::mock_server::{ - create_mock_open_ai_response, create_mock_open_ai_response_with_tools, create_test_tool, - get_expected_function_call_arguments, setup_mock_server, TEST_INPUT_TOKENS, - TEST_OUTPUT_TOKENS, TEST_TOOL_FUNCTION_NAME, TEST_TOTAL_TOKENS, - }; - use rust_decimal_macros::dec; - use wiremock::MockServer; - - async fn _setup_mock_response(response_body: Value) -> (MockServer, OpenRouterProvider) { - let mock_server = setup_mock_server("/api/v1/chat/completions", response_body).await; - - let provider = OpenRouterProvider { - client: Client::builder().build().unwrap(), - host: mock_server.uri(), - api_key: "test_api_key".to_string(), - model: ModelConfig::new("gpt-3.5-turbo".to_string()).with_temperature(Some(0.7)), - }; - - (mock_server, provider) - } - - #[tokio::test] - async fn test_complete_basic() -> Result<()> { - let model_name = "gpt-4"; - // Mock response for normal completion - let response_body = - create_mock_open_ai_response(model_name, "Hello! How can I assist you today?"); - - let (_, provider) = _setup_mock_response(response_body).await; - - // Prepare input messages - let messages = vec![Message::user().with_text("Hello?")]; - - // Call the complete method - let (message, usage) = provider - .complete("You are a helpful assistant.", &messages, &[]) - .await?; - - // Assert the 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.usage.input_tokens, Some(TEST_INPUT_TOKENS)); - assert_eq!(usage.usage.output_tokens, Some(TEST_OUTPUT_TOKENS)); - assert_eq!(usage.usage.total_tokens, Some(TEST_TOTAL_TOKENS)); - assert_eq!(usage.model, model_name); - assert_eq!(usage.cost, Some(dec!(0.00018))); - - Ok(()) - } - - #[tokio::test] - async fn test_complete_tool_request() -> Result<()> { - // Mock response for tool calling - let response_body = create_mock_open_ai_response_with_tools("gpt-4"); - - let (_, provider) = _setup_mock_response(response_body).await; - - // Input messages - let messages = vec![Message::user().with_text("What's the weather in San Francisco?")]; - - // Call the complete method - let (message, usage) = provider - .complete( - "You are a helpful assistant.", - &messages, - &[create_test_tool()], - ) - .await?; - - // Assert the 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, TEST_TOOL_FUNCTION_NAME); - assert_eq!(tool_call.arguments, get_expected_function_call_arguments()); - } else { - panic!("Expected ToolCall content"); - } - - assert_eq!(usage.usage.input_tokens, Some(TEST_INPUT_TOKENS)); - assert_eq!(usage.usage.output_tokens, Some(TEST_OUTPUT_TOKENS)); - assert_eq!(usage.usage.total_tokens, Some(TEST_TOTAL_TOKENS)); - - Ok(()) + get_usage(data) } } diff --git a/crates/goose/src/providers/utils.rs b/crates/goose/src/providers/utils.rs index b80d5c3fca..a63fcd9421 100644 --- a/crates/goose/src/providers/utils.rs +++ b/crates/goose/src/providers/utils.rs @@ -2,7 +2,6 @@ use super::base::Usage; use anyhow::{Error, Result}; use regex::Regex; use reqwest::{Response, StatusCode}; -use rust_decimal::Decimal; use serde::{Deserialize, Serialize}; use serde_json::{json, Map, Value}; use tracing::debug; @@ -137,7 +136,6 @@ pub fn emit_debug_trace( payload: &impl serde::Serialize, response: &Value, usage: &Usage, - cost: Option, ) { // Handle both Map and Value payload types let payload_str = match serde_json::to_value(payload) { @@ -152,7 +150,6 @@ pub fn emit_debug_trace( input_tokens = ?usage.input_tokens.unwrap_or_default(), output_tokens = ?usage.output_tokens.unwrap_or_default(), total_tokens = ?usage.total_tokens.unwrap_or_default(), - cost = ?cost.unwrap_or_default() ); } diff --git a/crates/goose/tests/providers.rs b/crates/goose/tests/providers.rs index 2451ab8c1a..7e01375b0f 100644 --- a/crates/goose/tests/providers.rs +++ b/crates/goose/tests/providers.rs @@ -2,18 +2,87 @@ use anyhow::Result; use dotenv::dotenv; use goose::message::{Message, MessageContent}; use goose::providers::base::Provider; -use goose::providers::{databricks, ollama, openai}; +use goose::providers::{anthropic, databricks, google, groq, ollama, openai, openrouter}; +use mcp_core::content::Content; use mcp_core::tool::Tool; +use std::collections::HashMap; +use std::sync::Arc; +use std::sync::Mutex; + +#[derive(Debug, Clone, Copy)] +enum TestStatus { + Passed, + Skipped, + Failed, +} + +impl std::fmt::Display for TestStatus { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + TestStatus::Passed => write!(f, "✅"), + TestStatus::Skipped => write!(f, "⏭️"), + TestStatus::Failed => write!(f, "❌"), + } + } +} + +struct TestReport { + results: Mutex>, +} + +impl TestReport { + fn new() -> Arc { + Arc::new(Self { + results: Mutex::new(HashMap::new()), + }) + } + + fn record_status(&self, provider: &str, status: TestStatus) { + let mut results = self.results.lock().unwrap(); + results.insert(provider.to_string(), status); + } + + fn record_pass(&self, provider: &str) { + self.record_status(provider, TestStatus::Passed); + } + + fn record_skip(&self, provider: &str) { + self.record_status(provider, TestStatus::Skipped); + } + + fn record_fail(&self, provider: &str) { + self.record_status(provider, TestStatus::Failed); + } + + fn print_summary(&self) { + println!("\n============== Providers =============="); + let results = self.results.lock().unwrap(); + let mut providers: Vec<_> = results.iter().collect(); + providers.sort_by(|a, b| a.0.cmp(b.0)); + + for (provider, status) in providers { + println!("{} {}", status, provider); + } + println!("=======================================\n"); + } +} + +lazy_static::lazy_static! { + static ref TEST_REPORT: Arc = TestReport::new(); + static ref ENV_LOCK: Mutex<()> = Mutex::new(()); +} /// Generic test harness for any Provider implementation struct ProviderTester { provider: Box, + name: String, } impl ProviderTester { - fn new(provider: T) -> Self { + fn new(provider: T, name: String) -> Self { Self { provider: Box::new(provider), + name, } } @@ -59,24 +128,72 @@ impl ProviderTester { let message = Message::user().with_text("What's the weather like in San Francisco?"); - let (response, _) = self + let (response1, _) = self .provider .complete( "You are a helpful weather assistant.", - &[message], - &[weather_tool], + &[message.clone()], + &[weather_tool.clone()], ) .await?; + println!("=== {}::reponse1 ===", self.name); + dbg!(&response1); + println!("==================="); + // Verify we got a tool request assert!( - response + response1 .content .iter() .any(|content| matches!(content, MessageContent::ToolRequest(_))), "Expected tool request in response" ); + let id = &response1 + .content + .iter() + .filter_map(|message| message.as_tool_request()) + .last() + .expect("got tool request") + .id; + + let weather = Message::user().with_tool_response( + id, + Ok(vec![Content::text( + " + 50°F°C + Precipitation: 0% + Humidity: 84% + Wind: 2 mph + Weather + Saturday 9:00 PM + Clear", + )]), + ); + + // Verify we construct a valid payload including the request/response pair for the next inference + let (response2, _) = self + .provider + .complete( + "You are a helpful weather assistant.", + &[message, response1, weather], + &[weather_tool], + ) + .await?; + + println!("=== {}::reponse2 ===", self.name); + dbg!(&response2); + println!("==================="); + + assert!( + response2 + .content + .iter() + .any(|content| matches!(content, MessageContent::Text(_))), + "Expected text for final response" + ); + Ok(()) } @@ -94,67 +211,187 @@ fn load_env() { } } -#[tokio::test] -async fn test_openai_provider() -> Result<()> { +/// Helper function to run a provider test with proper error handling and reporting +async fn test_provider( + name: &str, + required_vars: &[&str], + env_modifications: Option>>, + provider_fn: F, +) -> Result<()> +where + F: FnOnce() -> Result, + T: Provider + Send + Sync + 'static, +{ + // We start off as failed, so that if the process panics it is seen as a failure + TEST_REPORT.record_fail(name); + + // Take exclusive access to environment modifications + let lock = ENV_LOCK.lock().unwrap(); + load_env(); - // Skip if credentials aren't available - if std::env::var("OPENAI_API_KEY").is_err() { - println!("Skipping OpenAI tests - credentials not configured"); + // Save current environment state for required vars and modified vars + let mut original_env = HashMap::new(); + for &var in required_vars { + if let Ok(val) = std::env::var(var) { + original_env.insert(var, val); + } + } + if let Some(mods) = &env_modifications { + for &var in mods.keys() { + if let Ok(val) = std::env::var(var) { + original_env.insert(var, val); + } + } + } + + // Apply any environment modifications + if let Some(mods) = &env_modifications { + for (&var, value) in mods.iter() { + match value { + Some(val) => std::env::set_var(var, val), + None => std::env::remove_var(var), + } + } + } + + // Setup the provider + let missing_vars = required_vars.iter().any(|var| std::env::var(var).is_err()); + let provider = provider_fn(); + + // Restore original environment + for (&var, value) in original_env.iter() { + std::env::set_var(var, value); + } + if let Some(mods) = env_modifications { + for &var in mods.keys() { + if !original_env.contains_key(var) { + std::env::remove_var(var); + } + } + } + + std::mem::drop(lock); + + if missing_vars { + println!("Skipping {} tests - credentials not configured", name); + TEST_REPORT.record_skip(name); return Ok(()); } - let provider = openai::OpenAiProvider::from_env()?; - let tester = ProviderTester::new(provider); - tester.run_test_suite().await?; + if provider.is_err() { + println!("Could not setup {} from env", name); + TEST_REPORT.record_fail(name); + return Err(provider.err().expect("is error")); + } - Ok(()) + let tester = ProviderTester::new(provider.expect("already checked"), name.to_string()); + match tester.run_test_suite().await { + Ok(_) => { + TEST_REPORT.record_pass(name); + Ok(()) + } + Err(e) => { + println!("{} test failed: {}", name, e); + TEST_REPORT.record_fail(name); + Err(e) + } + } +} + +#[tokio::test] +async fn test_openai_provider() -> Result<()> { + test_provider( + "OpenAI", + &["OPENAI_API_KEY"], + None, + openai::OpenAiProvider::from_env, + ) + .await } #[tokio::test] async fn test_databricks_provider() -> Result<()> { - load_env(); - - // Skip if credentials aren't available - if std::env::var("DATABRICKS_HOST").is_err() || std::env::var("DATABRICKS_TOKEN").is_err() { - println!("Skipping Databricks tests - credentials not configured"); - return Ok(()); - } - - let provider = databricks::DatabricksProvider::from_env()?; - let tester = ProviderTester::new(provider); - tester.run_test_suite().await?; - - Ok(()) + test_provider( + "Databricks", + &["DATABRICKS_HOST", "DATABRICKS_TOKEN"], + None, + databricks::DatabricksProvider::from_env, + ) + .await } #[tokio::test] async fn test_databricks_provider_oauth() -> Result<()> { - load_env(); + let mut env_mods = HashMap::new(); + env_mods.insert("DATABRICKS_TOKEN", None); - // Skip if credentials aren't available - if std::env::var("DATABRICKS_HOST").is_err() { - println!("Skipping Databricks OAuth tests - credentials not configured"); - return Ok(()); - } - - // Clear any token to force OAuth - std::env::remove_var("DATABRICKS_TOKEN"); - let provider = databricks::DatabricksProvider::from_env()?; - let tester = ProviderTester::new(provider); - tester.run_test_suite().await?; - - Ok(()) + test_provider( + "Databricks OAuth", + &["DATABRICKS_HOST"], + Some(env_mods), + databricks::DatabricksProvider::from_env, + ) + .await } -// Integration tests that run against a real Ollama server #[tokio::test] async fn test_ollama_provider() -> Result<()> { - load_env(); - - let provider = ollama::OllamaProvider::from_env()?; - let tester = ProviderTester::new(provider); - tester.run_test_suite().await?; - - Ok(()) + test_provider( + "Ollama", + &["OLLAMA_HOST"], + None, + ollama::OllamaProvider::from_env, + ) + .await +} + +#[tokio::test] +async fn test_groq_provider() -> Result<()> { + test_provider( + "Groq", + &["GROQ_API_KEY"], + None, + groq::GroqProvider::from_env, + ) + .await +} + +#[tokio::test] +async fn test_anthropic_provider() -> Result<()> { + test_provider( + "Anthropic", + &["ANTHROPIC_API_KEY"], + None, + anthropic::AnthropicProvider::from_env, + ) + .await +} + +#[tokio::test] +async fn test_openrouter_provider() -> Result<()> { + test_provider( + "OpenRouter", + &["OPENROUTER_API_KEY"], + None, + openrouter::OpenRouterProvider::from_env, + ) + .await +} + +#[tokio::test] +async fn test_google_provider() -> Result<()> { + test_provider( + "Google", + &["GOOGLE_API_KEY"], + None, + google::GoogleProvider::from_env, + ) + .await +} + +// Print the final test report +#[ctor::dtor] +fn print_test_report() { + TEST_REPORT.print_summary(); }