From da719641dce663389277e8c5e351f9b9be4adf83 Mon Sep 17 00:00:00 2001 From: damienrj Date: Tue, 27 May 2025 13:02:05 -0700 Subject: [PATCH 1/4] Sagemaker TGI provider --- Cargo.lock | 26 ++ crates/goose/Cargo.toml | 3 + crates/goose/src/providers/factory.rs | 3 + crates/goose/src/providers/mod.rs | 1 + crates/goose/src/providers/sagemaker_tgi.rs | 486 ++++++++++++++++++++ 5 files changed, 519 insertions(+) create mode 100644 crates/goose/src/providers/sagemaker_tgi.rs diff --git a/Cargo.lock b/Cargo.lock index 02967400a288..2322aa5c6d6a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -704,6 +704,7 @@ dependencies = [ "aws-credential-types", "aws-sigv4", "aws-smithy-async", + "aws-smithy-eventstream", "aws-smithy-http 0.60.12", "aws-smithy-runtime", "aws-smithy-runtime-api", @@ -767,6 +768,29 @@ dependencies = [ "tracing", ] +[[package]] +name = "aws-sdk-sagemakerruntime" +version = "1.63.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8c3188bb9f962a9e1781c917dbe7f016ab9430e4bd81ba7daf422e58d86a3595" +dependencies = [ + "aws-credential-types", + "aws-runtime", + "aws-smithy-async", + "aws-smithy-eventstream", + "aws-smithy-http 0.61.1", + "aws-smithy-json", + "aws-smithy-runtime", + "aws-smithy-runtime-api", + "aws-smithy-types", + "aws-types", + "bytes", + "http 0.2.12", + "once_cell", + "regex-lite", + "tracing", +] + [[package]] name = "aws-sdk-sso" version = "1.61.0" @@ -841,6 +865,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9bfe75fad52793ce6dec0dc3d4b1f388f038b5eb866c8d4d7f3a8e21b5ea5051" dependencies = [ "aws-credential-types", + "aws-smithy-eventstream", "aws-smithy-http 0.60.12", "aws-smithy-runtime-api", "aws-smithy-types", @@ -3394,6 +3419,7 @@ dependencies = [ "async-trait", "aws-config", "aws-sdk-bedrockruntime", + "aws-sdk-sagemakerruntime", "aws-smithy-types", "axum", "base64 0.21.7", diff --git a/crates/goose/Cargo.toml b/crates/goose/Cargo.toml index 4aa3f12975da..bf2a23d96387 100644 --- a/crates/goose/Cargo.toml +++ b/crates/goose/Cargo.toml @@ -68,6 +68,9 @@ aws-config = { version = "1.5.16", features = ["behavior-version-latest"] } aws-smithy-types = "1.2.13" aws-sdk-bedrockruntime = "1.74.0" +# For SageMaker TGI provider +aws-sdk-sagemakerruntime = "1.62.0" + # For GCP Vertex AI provider auth jsonwebtoken = "9.3.1" diff --git a/crates/goose/src/providers/factory.rs b/crates/goose/src/providers/factory.rs index 22bdaa95ddef..e3893e8046a5 100644 --- a/crates/goose/src/providers/factory.rs +++ b/crates/goose/src/providers/factory.rs @@ -15,6 +15,7 @@ use super::{ openai::OpenAiProvider, openrouter::OpenRouterProvider, snowflake::SnowflakeProvider, + sagemaker_tgi::SageMakerTgiProvider, venice::VeniceProvider, }; use crate::model::ModelConfig; @@ -48,6 +49,7 @@ pub fn providers() -> Vec { OllamaProvider::metadata(), OpenAiProvider::metadata(), OpenRouterProvider::metadata(), + SageMakerTgiProvider::metadata(), VeniceProvider::metadata(), SnowflakeProvider::metadata(), ] @@ -122,6 +124,7 @@ fn create_provider(name: &str, model: ModelConfig) -> Result> "openrouter" => Ok(Arc::new(OpenRouterProvider::from_env(model)?)), "gcp_vertex_ai" => Ok(Arc::new(GcpVertexAIProvider::from_env(model)?)), "google" => Ok(Arc::new(GoogleProvider::from_env(model)?)), + "sagemaker_tgi" => Ok(Arc::new(SageMakerTgiProvider::from_env(model)?)), "venice" => Ok(Arc::new(VeniceProvider::from_env(model)?)), "snowflake" => Ok(Arc::new(SnowflakeProvider::from_env(model)?)), "github_copilot" => Ok(Arc::new(GithubCopilotProvider::from_env(model)?)), diff --git a/crates/goose/src/providers/mod.rs b/crates/goose/src/providers/mod.rs index 2f6f1f87bd2c..16a0c921e09e 100644 --- a/crates/goose/src/providers/mod.rs +++ b/crates/goose/src/providers/mod.rs @@ -19,6 +19,7 @@ pub mod ollama; pub mod openai; pub mod openrouter; pub mod snowflake; +pub mod sagemaker_tgi; pub mod toolshim; pub mod utils; pub mod utils_universal_openai_stream; diff --git a/crates/goose/src/providers/sagemaker_tgi.rs b/crates/goose/src/providers/sagemaker_tgi.rs new file mode 100644 index 000000000000..08f9453a4734 --- /dev/null +++ b/crates/goose/src/providers/sagemaker_tgi.rs @@ -0,0 +1,486 @@ +use std::collections::HashMap; +use std::time::Duration; + +use anyhow::Result; +use async_trait::async_trait; +use aws_config; +use aws_sdk_bedrockruntime::config::ProvideCredentials; +use aws_sdk_sagemakerruntime::Client as SageMakerClient; +use mcp_core::Tool; +use serde_json::{json, Value}; +use tokio::time::sleep; + +use super::base::{ConfigKey, Provider, ProviderMetadata, ProviderUsage, Usage}; +use super::errors::ProviderError; +use super::utils::emit_debug_trace; +use crate::message::{Message, MessageContent}; +use crate::model::ModelConfig; +use mcp_core::content::TextContent; +use mcp_core::role::Role; +use chrono::Utc; + +pub const SAGEMAKER_TGI_DOC_LINK: &str = + "https://docs.aws.amazon.com/sagemaker/latest/dg/realtime-endpoints.html"; + +pub const SAGEMAKER_TGI_DEFAULT_MODEL: &str = "sagemaker-tgi-endpoint"; + +#[derive(Debug, serde::Serialize)] +pub struct SageMakerTgiProvider { + #[serde(skip)] + sagemaker_client: SageMakerClient, + endpoint_name: String, + model: ModelConfig, +} + +impl SageMakerTgiProvider { + pub fn from_env(model: ModelConfig) -> Result { + let config = crate::config::Config::global(); + + // Get SageMaker endpoint name (just the name, not full URL) + let endpoint_name: String = config.get_param("SAGEMAKER_ENDPOINT_NAME") + .map_err(|_| anyhow::anyhow!("SAGEMAKER_ENDPOINT_NAME is required for SageMaker TGI provider"))?; + + // Attempt to load config and secrets to get AWS_ prefixed keys + let set_aws_env_vars = |res: Result, _>| { + if let Ok(map) = res { + map.into_iter() + .filter(|(key, _)| key.starts_with("AWS_")) + .filter_map(|(key, value)| value.as_str().map(|s| (key, s.to_string()))) + .for_each(|(key, s)| std::env::set_var(key, s)); + } + }; + + set_aws_env_vars(config.load_values()); + set_aws_env_vars(config.load_secrets()); + + let aws_config = futures::executor::block_on(aws_config::load_from_env()); + + // Validate credentials + futures::executor::block_on( + aws_config + .credentials_provider() + .unwrap() + .provide_credentials(), + )?; + + // Create client with longer timeout for model initialization + let timeout_config = aws_config::timeout::TimeoutConfig::builder() + .operation_timeout(Duration::from_secs(300)) // 5 minutes for cold starts + .build(); + + let config_with_timeout = aws_config.into_builder() + .timeout_config(timeout_config) + .build(); + + let sagemaker_client = SageMakerClient::new(&config_with_timeout); + + Ok(Self { + sagemaker_client, + endpoint_name, + model, + }) + } + + fn create_tgi_request(&self, system: &str, messages: &[Message], tools: &[Tool]) -> Result { + // Check if we should use tool calling format + let use_tool_calling = self.supports_tool_calling(); + + if use_tool_calling && !tools.is_empty() { + // Use OpenAI-compatible format for models that support it + return self.create_tool_calling_request(system, messages, tools); + } + + // Create a simplified prompt for basic TGI models + // Skip the complex system prompt and tool descriptions that cause the model to mimic tool formats + let mut prompt = String::new(); + + // Use a very simple system prompt if provided + if !system.is_empty() && !system.contains("Available tools") && system.len() < 200 { + prompt.push_str(&format!("System: {}\n\n", system)); + } else { + // Use a minimal system prompt for TGI + prompt.push_str("System: You are a helpful AI assistant.\n\n"); + } + + // Only include the most recent user messages to avoid overwhelming the model + let recent_messages: Vec<_> = messages.iter().rev().take(3).collect(); + for message in recent_messages.iter().rev() { + match &message.role { + Role::User => { + prompt.push_str("User: "); + for content in &message.content { + if let MessageContent::Text(text) = content { + prompt.push_str(&text.text); + } + } + prompt.push_str("\n\n"); + } + Role::Assistant => { + prompt.push_str("Assistant: "); + for content in &message.content { + if let MessageContent::Text(text) = content { + // Skip responses that look like tool descriptions + if !text.text.contains("__") && !text.text.contains("Available tools") { + prompt.push_str(&text.text); + } + } + } + prompt.push_str("\n\n"); + } + } + } + + prompt.push_str("Assistant: "); + + // Skip tool descriptions entirely for TGI models to avoid confusion + // TGI models don't support tools natively and including tool descriptions + // causes them to mimic that format in their responses + + // Build TGI request with reasonable parameters + let request = json!({ + "inputs": prompt, + "parameters": { + "max_new_tokens": self.model.max_tokens.unwrap_or(150), + "temperature": self.model.temperature.unwrap_or(0.7), + "do_sample": true, + "return_full_text": false + } + }); + + Ok(request) + } + + fn supports_tool_calling(&self) -> bool { + // Check if the model supports tool calling + // You can configure this via environment variable or model name detection + std::env::var("GOOSE_SAGEMAKER_SUPPORTS_TOOLS") + .map(|v| v == "true" || v == "1") + .unwrap_or(false) + } + + fn create_tool_calling_request(&self, system: &str, messages: &[Message], tools: &[Tool]) -> Result { + // Convert messages to OpenAI format + let mut openai_messages = Vec::new(); + + if !system.is_empty() { + openai_messages.push(json!({ + "role": "system", + "content": system + })); + } + + for message in messages { + match &message.role { + Role::User => { + let content = message.content.iter() + .filter_map(|c| { + if let MessageContent::Text(text) = c { + Some(text.text.clone()) + } else { + None + } + }) + .collect::>() + .join("\n"); + + openai_messages.push(json!({ + "role": "user", + "content": content + })); + } + Role::Assistant => { + let content = message.content.iter() + .filter_map(|c| { + if let MessageContent::Text(text) = c { + Some(text.text.clone()) + } else { + None + } + }) + .collect::>() + .join("\n"); + + if !content.is_empty() { + openai_messages.push(json!({ + "role": "assistant", + "content": content + })); + } + } + } + } + + // Convert tools to OpenAI format, cleaning up unsupported schema formats + let openai_tools: Vec = tools.iter().map(|tool| { + let cleaned_schema = self.clean_schema_for_tgi(&tool.input_schema); + json!({ + "type": "function", + "function": { + "name": tool.name, + "description": tool.description, + "parameters": cleaned_schema + } + }) + }).collect(); + + // Create OpenAI-compatible request + let request = json!({ + "messages": openai_messages, + "tools": openai_tools, + "tool_choice": "auto", + "max_tokens": self.model.max_tokens.unwrap_or(150), + "temperature": self.model.temperature.unwrap_or(0.7) + }); + + Ok(request) + } + + fn clean_schema_for_tgi(&self, schema: &Value) -> Value { + // Recursively clean JSON schema to remove formats unsupported by TGI/Outlines + match schema { + Value::Object(obj) => { + let mut cleaned = serde_json::Map::new(); + for (key, value) in obj { + if key == "format" { + // Skip format constraints that TGI doesn't support + if let Some(format_str) = value.as_str() { + match format_str { + "uri" | "email" | "date" | "date-time" | "uuid" => { + // Skip these unsupported formats + continue; + } + _ => { + // Keep supported formats + cleaned.insert(key.clone(), value.clone()); + } + } + } + } else { + // Recursively clean nested objects/arrays + cleaned.insert(key.clone(), self.clean_schema_for_tgi(value)); + } + } + Value::Object(cleaned) + } + Value::Array(arr) => { + Value::Array(arr.iter().map(|v| self.clean_schema_for_tgi(v)).collect()) + } + _ => schema.clone(), + } + } + + async fn invoke_endpoint(&self, payload: Value) -> Result { + let body = serde_json::to_string(&payload) + .map_err(|e| ProviderError::RequestFailed(format!("Failed to serialize request: {}", e)))?; + + let response = self.sagemaker_client + .invoke_endpoint() + .endpoint_name(&self.endpoint_name) + .content_type("application/json") + .body(body.into_bytes().into()) + .send() + .await + .map_err(|e| ProviderError::RequestFailed(format!("SageMaker invoke failed: {}", e)))?; + + let response_body = response.body.as_ref() + .ok_or_else(|| ProviderError::RequestFailed("Empty response body".to_string()))?; + let response_text = std::str::from_utf8(response_body.as_ref()) + .map_err(|e| ProviderError::RequestFailed(format!("Failed to decode response: {}", e)))?; + + serde_json::from_str(response_text) + .map_err(|e| ProviderError::RequestFailed(format!("Failed to parse response JSON: {}", e))) + } + + fn parse_tgi_response(&self, response: Value) -> Result { + // Check if this is an OpenAI-compatible response (tool calling format) + if response.get("choices").is_some() { + return self.parse_openai_response(response); + } + + // Handle standard TGI response: [{"generated_text": "..."}] + let response_array = response.as_array() + .ok_or_else(|| ProviderError::RequestFailed("Expected array response".to_string()))?; + + if response_array.is_empty() { + return Err(ProviderError::RequestFailed("Empty response array".to_string())); + } + + let first_result = &response_array[0]; + let generated_text = first_result.get("generated_text") + .and_then(|v| v.as_str()) + .ok_or_else(|| ProviderError::RequestFailed("No generated_text in response".to_string()))?; + + Ok(Message { + role: Role::Assistant, + created: Utc::now().timestamp(), + content: vec![MessageContent::Text(TextContent { + text: generated_text.to_string(), + annotations: None, + })], + }) + } + + fn parse_openai_response(&self, response: Value) -> Result { + let choices = response.get("choices") + .and_then(|c| c.as_array()) + .ok_or_else(|| ProviderError::RequestFailed("No choices in OpenAI response".to_string()))?; + + if choices.is_empty() { + return Err(ProviderError::RequestFailed("Empty choices array".to_string())); + } + + let choice = &choices[0]; + let message = choice.get("message") + .ok_or_else(|| ProviderError::RequestFailed("No message in choice".to_string()))?; + + let mut content = Vec::new(); + + // Handle text content + if let Some(text_content) = message.get("content").and_then(|c| c.as_str()) { + if !text_content.is_empty() { + content.push(MessageContent::Text(TextContent { + text: text_content.to_string(), + annotations: None, + })); + } + } + + // Handle tool calls + if let Some(tool_calls) = message.get("tool_calls").and_then(|tc| tc.as_array()) { + for tool_call in tool_calls { + if let (Some(id), Some(function)) = ( + tool_call.get("id").and_then(|i| i.as_str()), + tool_call.get("function") + ) { + if let (Some(name), Some(arguments)) = ( + function.get("name").and_then(|n| n.as_str()), + function.get("arguments").and_then(|a| a.as_str()) + ) { + // Parse the tool call arguments + let args: Value = serde_json::from_str(arguments) + .unwrap_or_else(|_| json!({})); + + let tool_call = mcp_core::tool::ToolCall { + name: name.to_string(), + arguments: args, + }; + + content.push(MessageContent::tool_request(id, Ok(tool_call))); + } + } + } + } + + Ok(Message { + role: Role::Assistant, + created: Utc::now().timestamp(), + content, + }) + } +} + +impl Default for SageMakerTgiProvider { + fn default() -> Self { + let model = ModelConfig::new(SageMakerTgiProvider::metadata().default_model); + SageMakerTgiProvider::from_env(model).expect("Failed to initialize SageMaker TGI provider") + } +} + +#[async_trait] +impl Provider for SageMakerTgiProvider { + fn metadata() -> ProviderMetadata { + ProviderMetadata::new( + "sagemaker_tgi", + "Amazon SageMaker TGI", + "Run Text Generation Inference models through Amazon SageMaker endpoints. Requires AWS credentials and a SageMaker endpoint URL.", + SAGEMAKER_TGI_DEFAULT_MODEL, + vec![SAGEMAKER_TGI_DEFAULT_MODEL], + SAGEMAKER_TGI_DOC_LINK, + vec![ + ConfigKey::new("SAGEMAKER_ENDPOINT_NAME", false, false, None), + ConfigKey::new("AWS_REGION", true, false, Some("us-east-1")), + ConfigKey::new("AWS_PROFILE", true, false, Some("default")), + ConfigKey::new("GOOSE_SAGEMAKER_SUPPORTS_TOOLS", true, false, Some("false")), + ], + ) + } + + fn get_model_config(&self) -> ModelConfig { + self.model.clone() + } + + #[tracing::instrument( + skip(self, system, messages, tools), + fields(model_config, input, output, input_tokens, output_tokens, total_tokens) + )] + async fn complete( + &self, + system: &str, + messages: &[Message], + tools: &[Tool], + ) -> Result<(Message, ProviderUsage), ProviderError> { + let model_name = &self.model.model_name; + + let request_payload = self.create_tgi_request(system, messages, tools) + .map_err(|e| ProviderError::RequestFailed(format!("Failed to create request: {}", e)))?; + + // Retry configuration + const MAX_RETRIES: u32 = 3; + const INITIAL_BACKOFF_MS: u64 = 1000; // 1 second + const MAX_BACKOFF_MS: u64 = 30000; // 30 seconds + + let mut attempts = 0; + let mut backoff_ms = INITIAL_BACKOFF_MS; + + loop { + attempts += 1; + + match self.invoke_endpoint(request_payload.clone()).await { + Ok(response) => { + let message = self.parse_tgi_response(response)?; + + // TGI doesn't provide usage statistics, so we estimate + let usage = Usage { + input_tokens: Some(0), // Would need to tokenize input to get accurate count + output_tokens: Some(0), // Would need to tokenize output to get accurate count + total_tokens: Some(0), + }; + + // Add debug trace + let debug_payload = serde_json::json!({ + "system": system, + "messages": messages, + "tools": tools + }); + emit_debug_trace( + &self.model, + &debug_payload, + &serde_json::to_value(&message).unwrap_or_default(), + &usage, + ); + + let provider_usage = ProviderUsage::new(model_name.to_string(), usage); + return Ok((message, provider_usage)); + } + Err(err) => { + if attempts > MAX_RETRIES { + return Err(err); + } + + // Log retry attempt + tracing::warn!( + "SageMaker TGI request failed (attempt {}/{}), retrying in {} ms: {:?}", + attempts, + MAX_RETRIES, + backoff_ms, + err + ); + + // Wait before retry + sleep(Duration::from_millis(backoff_ms)).await; + backoff_ms = (backoff_ms * 2).min(MAX_BACKOFF_MS); + } + } + } + } +} \ No newline at end of file From 257c0347c03e2583f6ce389353cb4be9056b3188 Mon Sep 17 00:00:00 2001 From: damienrj Date: Fri, 30 May 2025 12:27:42 -0700 Subject: [PATCH 2/4] remove Goose tool flag --- crates/goose/src/providers/sagemaker_tgi.rs | 227 +++----------------- 1 file changed, 33 insertions(+), 194 deletions(-) diff --git a/crates/goose/src/providers/sagemaker_tgi.rs b/crates/goose/src/providers/sagemaker_tgi.rs index 08f9453a4734..c767db50bab0 100644 --- a/crates/goose/src/providers/sagemaker_tgi.rs +++ b/crates/goose/src/providers/sagemaker_tgi.rs @@ -81,25 +81,17 @@ impl SageMakerTgiProvider { }) } - fn create_tgi_request(&self, system: &str, messages: &[Message], tools: &[Tool]) -> Result { - // Check if we should use tool calling format - let use_tool_calling = self.supports_tool_calling(); - - if use_tool_calling && !tools.is_empty() { - // Use OpenAI-compatible format for models that support it - return self.create_tool_calling_request(system, messages, tools); - } - + fn create_tgi_request(&self, system: &str, messages: &[Message], _tools: &[Tool]) -> Result { // Create a simplified prompt for basic TGI models // Skip the complex system prompt and tool descriptions that cause the model to mimic tool formats let mut prompt = String::new(); - // Use a very simple system prompt if provided - if !system.is_empty() && !system.contains("Available tools") && system.len() < 200 { + // Use a very simple system prompt if provided, but ensure it doesn't contain HTML instructions + if !system.is_empty() && !system.contains("Available tools") && system.len() < 200 && !system.contains("HTML") && !system.contains("markdown") { prompt.push_str(&format!("System: {}\n\n", system)); } else { - // Use a minimal system prompt for TGI - prompt.push_str("System: You are a helpful AI assistant.\n\n"); + // Use a minimal system prompt for TGI that explicitly avoids HTML + prompt.push_str("System: You are a helpful AI assistant. Provide responses in plain text only. Do not use HTML tags, markup, or formatting.\n\n"); } // Only include the most recent user messages to avoid overwhelming the model @@ -119,8 +111,8 @@ impl SageMakerTgiProvider { prompt.push_str("Assistant: "); for content in &message.content { if let MessageContent::Text(text) = content { - // Skip responses that look like tool descriptions - if !text.text.contains("__") && !text.text.contains("Available tools") { + // Skip responses that look like tool descriptions or contain HTML + if !text.text.contains("__") && !text.text.contains("Available tools") && !text.text.contains("<") { prompt.push_str(&text.text); } } @@ -150,125 +142,6 @@ impl SageMakerTgiProvider { Ok(request) } - fn supports_tool_calling(&self) -> bool { - // Check if the model supports tool calling - // You can configure this via environment variable or model name detection - std::env::var("GOOSE_SAGEMAKER_SUPPORTS_TOOLS") - .map(|v| v == "true" || v == "1") - .unwrap_or(false) - } - - fn create_tool_calling_request(&self, system: &str, messages: &[Message], tools: &[Tool]) -> Result { - // Convert messages to OpenAI format - let mut openai_messages = Vec::new(); - - if !system.is_empty() { - openai_messages.push(json!({ - "role": "system", - "content": system - })); - } - - for message in messages { - match &message.role { - Role::User => { - let content = message.content.iter() - .filter_map(|c| { - if let MessageContent::Text(text) = c { - Some(text.text.clone()) - } else { - None - } - }) - .collect::>() - .join("\n"); - - openai_messages.push(json!({ - "role": "user", - "content": content - })); - } - Role::Assistant => { - let content = message.content.iter() - .filter_map(|c| { - if let MessageContent::Text(text) = c { - Some(text.text.clone()) - } else { - None - } - }) - .collect::>() - .join("\n"); - - if !content.is_empty() { - openai_messages.push(json!({ - "role": "assistant", - "content": content - })); - } - } - } - } - - // Convert tools to OpenAI format, cleaning up unsupported schema formats - let openai_tools: Vec = tools.iter().map(|tool| { - let cleaned_schema = self.clean_schema_for_tgi(&tool.input_schema); - json!({ - "type": "function", - "function": { - "name": tool.name, - "description": tool.description, - "parameters": cleaned_schema - } - }) - }).collect(); - - // Create OpenAI-compatible request - let request = json!({ - "messages": openai_messages, - "tools": openai_tools, - "tool_choice": "auto", - "max_tokens": self.model.max_tokens.unwrap_or(150), - "temperature": self.model.temperature.unwrap_or(0.7) - }); - - Ok(request) - } - - fn clean_schema_for_tgi(&self, schema: &Value) -> Value { - // Recursively clean JSON schema to remove formats unsupported by TGI/Outlines - match schema { - Value::Object(obj) => { - let mut cleaned = serde_json::Map::new(); - for (key, value) in obj { - if key == "format" { - // Skip format constraints that TGI doesn't support - if let Some(format_str) = value.as_str() { - match format_str { - "uri" | "email" | "date" | "date-time" | "uuid" => { - // Skip these unsupported formats - continue; - } - _ => { - // Keep supported formats - cleaned.insert(key.clone(), value.clone()); - } - } - } - } else { - // Recursively clean nested objects/arrays - cleaned.insert(key.clone(), self.clean_schema_for_tgi(value)); - } - } - Value::Object(cleaned) - } - Value::Array(arr) => { - Value::Array(arr.iter().map(|v| self.clean_schema_for_tgi(v)).collect()) - } - _ => schema.clone(), - } - } - async fn invoke_endpoint(&self, payload: Value) -> Result { let body = serde_json::to_string(&payload) .map_err(|e| ProviderError::RequestFailed(format!("Failed to serialize request: {}", e)))?; @@ -292,11 +165,6 @@ impl SageMakerTgiProvider { } fn parse_tgi_response(&self, response: Value) -> Result { - // Check if this is an OpenAI-compatible response (tool calling format) - if response.get("choices").is_some() { - return self.parse_openai_response(response); - } - // Handle standard TGI response: [{"generated_text": "..."}] let response_array = response.as_array() .ok_or_else(|| ProviderError::RequestFailed("Expected array response".to_string()))?; @@ -310,72 +178,44 @@ impl SageMakerTgiProvider { .and_then(|v| v.as_str()) .ok_or_else(|| ProviderError::RequestFailed("No generated_text in response".to_string()))?; + // Strip any HTML tags that might have been generated + let clean_text = self.strip_html_tags(generated_text); + Ok(Message { role: Role::Assistant, created: Utc::now().timestamp(), content: vec![MessageContent::Text(TextContent { - text: generated_text.to_string(), + text: clean_text, annotations: None, })], }) } - fn parse_openai_response(&self, response: Value) -> Result { - let choices = response.get("choices") - .and_then(|c| c.as_array()) - .ok_or_else(|| ProviderError::RequestFailed("No choices in OpenAI response".to_string()))?; - - if choices.is_empty() { - return Err(ProviderError::RequestFailed("Empty choices array".to_string())); - } - - let choice = &choices[0]; - let message = choice.get("message") - .ok_or_else(|| ProviderError::RequestFailed("No message in choice".to_string()))?; - - let mut content = Vec::new(); - - // Handle text content - if let Some(text_content) = message.get("content").and_then(|c| c.as_str()) { - if !text_content.is_empty() { - content.push(MessageContent::Text(TextContent { - text: text_content.to_string(), - annotations: None, - })); - } + /// Strip HTML tags from text to ensure clean output + fn strip_html_tags(&self, text: &str) -> String { + // Simple regex-free approach to strip common HTML tags + let mut result = text.to_string(); + + // Remove common HTML tags like , , , , etc. + let tags_to_remove = ["", "", "", "", "", "", + "", "", "", "", "
", "
", + "

", "

", "
", "
", "", ""]; + + for tag in &tags_to_remove { + result = result.replace(tag, ""); } - - // Handle tool calls - if let Some(tool_calls) = message.get("tool_calls").and_then(|tc| tc.as_array()) { - for tool_call in tool_calls { - if let (Some(id), Some(function)) = ( - tool_call.get("id").and_then(|i| i.as_str()), - tool_call.get("function") - ) { - if let (Some(name), Some(arguments)) = ( - function.get("name").and_then(|n| n.as_str()), - function.get("arguments").and_then(|a| a.as_str()) - ) { - // Parse the tool call arguments - let args: Value = serde_json::from_str(arguments) - .unwrap_or_else(|_| json!({})); - - let tool_call = mcp_core::tool::ToolCall { - name: name.to_string(), - arguments: args, - }; - - content.push(MessageContent::tool_request(id, Ok(tool_call))); - } - } + + // Remove any remaining HTML-like tags using a simple pattern + // This is a basic implementation - for production use, consider using a proper HTML parser + while let Some(start) = result.find('<') { + if let Some(end) = result[start..].find('>') { + result.replace_range(start..start + end + 1, ""); + } else { + break; } } - - Ok(Message { - role: Role::Assistant, - created: Utc::now().timestamp(), - content, - }) + + result.trim().to_string() } } @@ -400,7 +240,6 @@ impl Provider for SageMakerTgiProvider { ConfigKey::new("SAGEMAKER_ENDPOINT_NAME", false, false, None), ConfigKey::new("AWS_REGION", true, false, Some("us-east-1")), ConfigKey::new("AWS_PROFILE", true, false, Some("default")), - ConfigKey::new("GOOSE_SAGEMAKER_SUPPORTS_TOOLS", true, false, Some("false")), ], ) } From 163cf99f06982ae06bed8da8645c568e9f56d04e Mon Sep 17 00:00:00 2001 From: damienrj Date: Thu, 12 Jun 2025 11:09:48 -0700 Subject: [PATCH 3/4] fix formatting --- Justfile | 19 ++++ crates/goose/src/providers/factory.rs | 2 +- crates/goose/src/providers/mod.rs | 2 +- crates/goose/src/providers/sagemaker_tgi.rs | 117 ++++++++++++++------ 4 files changed, 103 insertions(+), 37 deletions(-) diff --git a/Justfile b/Justfile index 4888fcd2b049..d7b27706b492 100644 --- a/Justfile +++ b/Justfile @@ -388,4 +388,23 @@ kotlin-example: cd bindings/kotlin/ && java \ -Djna.library.path=$HOME/Development/goose/target/debug \ -classpath "example.jar:libs/kotlin-stdlib-1.9.0.jar:libs/kotlinx-coroutines-core-jvm-1.7.3.jar:libs/jna-5.13.0.jar" \ +<<<<<<< Updated upstream UsageKt +||||||| Stash base + UsageKt +======= + UsageKt + + + huggingface-pytorch-tgi-inference-2025-05-30-16-42-23-042 + + export GOOSE_PROVIDER=sagemaker_tgi + export SAGEMAKER_ENDPOINT_NAME="huggingface-pytorch-tgi-inference-2025-05-30-16-42-23-042" + export AWS_REGION="us-east-2" + export AWS_PROFILE="osllm-app-staging--admin" + export RUST_LOG=info + + + cargo build + cargo run --bin goose -- session +>>>>>>> Stashed changes diff --git a/crates/goose/src/providers/factory.rs b/crates/goose/src/providers/factory.rs index e3893e8046a5..af4fc2e34f99 100644 --- a/crates/goose/src/providers/factory.rs +++ b/crates/goose/src/providers/factory.rs @@ -14,8 +14,8 @@ use super::{ ollama::OllamaProvider, openai::OpenAiProvider, openrouter::OpenRouterProvider, - snowflake::SnowflakeProvider, sagemaker_tgi::SageMakerTgiProvider, + snowflake::SnowflakeProvider, venice::VeniceProvider, }; use crate::model::ModelConfig; diff --git a/crates/goose/src/providers/mod.rs b/crates/goose/src/providers/mod.rs index 16a0c921e09e..ccda29a431d4 100644 --- a/crates/goose/src/providers/mod.rs +++ b/crates/goose/src/providers/mod.rs @@ -18,8 +18,8 @@ pub mod oauth; pub mod ollama; pub mod openai; pub mod openrouter; -pub mod snowflake; pub mod sagemaker_tgi; +pub mod snowflake; pub mod toolshim; pub mod utils; pub mod utils_universal_openai_stream; diff --git a/crates/goose/src/providers/sagemaker_tgi.rs b/crates/goose/src/providers/sagemaker_tgi.rs index c767db50bab0..214a24f31d47 100644 --- a/crates/goose/src/providers/sagemaker_tgi.rs +++ b/crates/goose/src/providers/sagemaker_tgi.rs @@ -15,11 +15,11 @@ use super::errors::ProviderError; use super::utils::emit_debug_trace; use crate::message::{Message, MessageContent}; use crate::model::ModelConfig; +use chrono::Utc; use mcp_core::content::TextContent; use mcp_core::role::Role; -use chrono::Utc; -pub const SAGEMAKER_TGI_DOC_LINK: &str = +pub const SAGEMAKER_TGI_DOC_LINK: &str = "https://docs.aws.amazon.com/sagemaker/latest/dg/realtime-endpoints.html"; pub const SAGEMAKER_TGI_DEFAULT_MODEL: &str = "sagemaker-tgi-endpoint"; @@ -35,10 +35,11 @@ pub struct SageMakerTgiProvider { impl SageMakerTgiProvider { pub fn from_env(model: ModelConfig) -> Result { let config = crate::config::Config::global(); - + // Get SageMaker endpoint name (just the name, not full URL) - let endpoint_name: String = config.get_param("SAGEMAKER_ENDPOINT_NAME") - .map_err(|_| anyhow::anyhow!("SAGEMAKER_ENDPOINT_NAME is required for SageMaker TGI provider"))?; + let endpoint_name: String = config.get_param("SAGEMAKER_ENDPOINT_NAME").map_err(|_| { + anyhow::anyhow!("SAGEMAKER_ENDPOINT_NAME is required for SageMaker TGI provider") + })?; // Attempt to load config and secrets to get AWS_ prefixed keys let set_aws_env_vars = |res: Result, _>| { @@ -67,11 +68,12 @@ impl SageMakerTgiProvider { let timeout_config = aws_config::timeout::TimeoutConfig::builder() .operation_timeout(Duration::from_secs(300)) // 5 minutes for cold starts .build(); - - let config_with_timeout = aws_config.into_builder() + + let config_with_timeout = aws_config + .into_builder() .timeout_config(timeout_config) .build(); - + let sagemaker_client = SageMakerClient::new(&config_with_timeout); Ok(Self { @@ -81,13 +83,23 @@ impl SageMakerTgiProvider { }) } - fn create_tgi_request(&self, system: &str, messages: &[Message], _tools: &[Tool]) -> Result { + fn create_tgi_request( + &self, + system: &str, + messages: &[Message], + _tools: &[Tool], + ) -> Result { // Create a simplified prompt for basic TGI models // Skip the complex system prompt and tool descriptions that cause the model to mimic tool formats let mut prompt = String::new(); - + // Use a very simple system prompt if provided, but ensure it doesn't contain HTML instructions - if !system.is_empty() && !system.contains("Available tools") && system.len() < 200 && !system.contains("HTML") && !system.contains("markdown") { + if !system.is_empty() + && !system.contains("Available tools") + && system.len() < 200 + && !system.contains("HTML") + && !system.contains("markdown") + { prompt.push_str(&format!("System: {}\n\n", system)); } else { // Use a minimal system prompt for TGI that explicitly avoids HTML @@ -112,7 +124,10 @@ impl SageMakerTgiProvider { for content in &message.content { if let MessageContent::Text(text) = content { // Skip responses that look like tool descriptions or contain HTML - if !text.text.contains("__") && !text.text.contains("Available tools") && !text.text.contains("<") { + if !text.text.contains("__") + && !text.text.contains("Available tools") + && !text.text.contains("<") + { prompt.push_str(&text.text); } } @@ -143,10 +158,12 @@ impl SageMakerTgiProvider { } async fn invoke_endpoint(&self, payload: Value) -> Result { - let body = serde_json::to_string(&payload) - .map_err(|e| ProviderError::RequestFailed(format!("Failed to serialize request: {}", e)))?; + let body = serde_json::to_string(&payload).map_err(|e| { + ProviderError::RequestFailed(format!("Failed to serialize request: {}", e)) + })?; - let response = self.sagemaker_client + let response = self + .sagemaker_client .invoke_endpoint() .endpoint_name(&self.endpoint_name) .content_type("application/json") @@ -155,28 +172,38 @@ impl SageMakerTgiProvider { .await .map_err(|e| ProviderError::RequestFailed(format!("SageMaker invoke failed: {}", e)))?; - let response_body = response.body.as_ref() + let response_body = response + .body + .as_ref() .ok_or_else(|| ProviderError::RequestFailed("Empty response body".to_string()))?; - let response_text = std::str::from_utf8(response_body.as_ref()) - .map_err(|e| ProviderError::RequestFailed(format!("Failed to decode response: {}", e)))?; + let response_text = std::str::from_utf8(response_body.as_ref()).map_err(|e| { + ProviderError::RequestFailed(format!("Failed to decode response: {}", e)) + })?; - serde_json::from_str(response_text) - .map_err(|e| ProviderError::RequestFailed(format!("Failed to parse response JSON: {}", e))) + serde_json::from_str(response_text).map_err(|e| { + ProviderError::RequestFailed(format!("Failed to parse response JSON: {}", e)) + }) } fn parse_tgi_response(&self, response: Value) -> Result { // Handle standard TGI response: [{"generated_text": "..."}] - let response_array = response.as_array() + let response_array = response + .as_array() .ok_or_else(|| ProviderError::RequestFailed("Expected array response".to_string()))?; if response_array.is_empty() { - return Err(ProviderError::RequestFailed("Empty response array".to_string())); + return Err(ProviderError::RequestFailed( + "Empty response array".to_string(), + )); } let first_result = &response_array[0]; - let generated_text = first_result.get("generated_text") + let generated_text = first_result + .get("generated_text") .and_then(|v| v.as_str()) - .ok_or_else(|| ProviderError::RequestFailed("No generated_text in response".to_string()))?; + .ok_or_else(|| { + ProviderError::RequestFailed("No generated_text in response".to_string()) + })?; // Strip any HTML tags that might have been generated let clean_text = self.strip_html_tags(generated_text); @@ -195,16 +222,33 @@ impl SageMakerTgiProvider { fn strip_html_tags(&self, text: &str) -> String { // Simple regex-free approach to strip common HTML tags let mut result = text.to_string(); - + // Remove common HTML tags like , , , , etc. - let tags_to_remove = ["", "", "", "", "", "", - "", "", "", "", "
", "
", - "

", "

", "
", "
", "", ""]; - + let tags_to_remove = [ + "", + "", + "", + "", + "", + "", + "", + "", + "", + "", + "
", + "
", + "

", + "

", + "
", + "
", + "", + "", + ]; + for tag in &tags_to_remove { result = result.replace(tag, ""); } - + // Remove any remaining HTML-like tags using a simple pattern // This is a basic implementation - for production use, consider using a proper HTML parser while let Some(start) = result.find('<') { @@ -214,7 +258,7 @@ impl SageMakerTgiProvider { break; } } - + result.trim().to_string() } } @@ -260,8 +304,11 @@ impl Provider for SageMakerTgiProvider { ) -> Result<(Message, ProviderUsage), ProviderError> { let model_name = &self.model.model_name; - let request_payload = self.create_tgi_request(system, messages, tools) - .map_err(|e| ProviderError::RequestFailed(format!("Failed to create request: {}", e)))?; + let request_payload = self + .create_tgi_request(system, messages, tools) + .map_err(|e| { + ProviderError::RequestFailed(format!("Failed to create request: {}", e)) + })?; // Retry configuration const MAX_RETRIES: u32 = 3; @@ -280,7 +327,7 @@ impl Provider for SageMakerTgiProvider { // TGI doesn't provide usage statistics, so we estimate let usage = Usage { - input_tokens: Some(0), // Would need to tokenize input to get accurate count + input_tokens: Some(0), // Would need to tokenize input to get accurate count output_tokens: Some(0), // Would need to tokenize output to get accurate count total_tokens: Some(0), }; @@ -322,4 +369,4 @@ impl Provider for SageMakerTgiProvider { } } } -} \ No newline at end of file +} From 2845d9a94e532cb7489239c3ff7a7a58164703a2 Mon Sep 17 00:00:00 2001 From: damienrj Date: Thu, 12 Jun 2025 14:24:00 -0700 Subject: [PATCH 4/4] Clean code and add a tes --- Justfile | 19 ------------------- crates/goose/src/providers/sagemaker_tgi.rs | 19 ++++++------------- crates/goose/tests/providers.rs | 11 +++++++++++ 3 files changed, 17 insertions(+), 32 deletions(-) diff --git a/Justfile b/Justfile index d7b27706b492..4888fcd2b049 100644 --- a/Justfile +++ b/Justfile @@ -388,23 +388,4 @@ kotlin-example: cd bindings/kotlin/ && java \ -Djna.library.path=$HOME/Development/goose/target/debug \ -classpath "example.jar:libs/kotlin-stdlib-1.9.0.jar:libs/kotlinx-coroutines-core-jvm-1.7.3.jar:libs/jna-5.13.0.jar" \ -<<<<<<< Updated upstream UsageKt -||||||| Stash base - UsageKt -======= - UsageKt - - - huggingface-pytorch-tgi-inference-2025-05-30-16-42-23-042 - - export GOOSE_PROVIDER=sagemaker_tgi - export SAGEMAKER_ENDPOINT_NAME="huggingface-pytorch-tgi-inference-2025-05-30-16-42-23-042" - export AWS_REGION="us-east-2" - export AWS_PROFILE="osllm-app-staging--admin" - export RUST_LOG=info - - - cargo build - cargo run --bin goose -- session ->>>>>>> Stashed changes diff --git a/crates/goose/src/providers/sagemaker_tgi.rs b/crates/goose/src/providers/sagemaker_tgi.rs index 214a24f31d47..19b95deb83f6 100644 --- a/crates/goose/src/providers/sagemaker_tgi.rs +++ b/crates/goose/src/providers/sagemaker_tgi.rs @@ -83,14 +83,9 @@ impl SageMakerTgiProvider { }) } - fn create_tgi_request( - &self, - system: &str, - messages: &[Message], - _tools: &[Tool], - ) -> Result { - // Create a simplified prompt for basic TGI models - // Skip the complex system prompt and tool descriptions that cause the model to mimic tool formats + fn create_tgi_request(&self, system: &str, messages: &[Message]) -> Result { + // Create a simplified prompt for TGI models using recent user and assistant messages. + // Uses a minimal system prompt and avoids HTML or tool-related formatting. let mut prompt = String::new(); // Use a very simple system prompt if provided, but ensure it doesn't contain HTML instructions @@ -304,11 +299,9 @@ impl Provider for SageMakerTgiProvider { ) -> Result<(Message, ProviderUsage), ProviderError> { let model_name = &self.model.model_name; - let request_payload = self - .create_tgi_request(system, messages, tools) - .map_err(|e| { - ProviderError::RequestFailed(format!("Failed to create request: {}", e)) - })?; + let request_payload = self.create_tgi_request(system, messages).map_err(|e| { + ProviderError::RequestFailed(format!("Failed to create request: {}", e)) + })?; // Retry configuration const MAX_RETRIES: u32 = 3; diff --git a/crates/goose/tests/providers.rs b/crates/goose/tests/providers.rs index 4d41251fc440..842f53789963 100644 --- a/crates/goose/tests/providers.rs +++ b/crates/goose/tests/providers.rs @@ -490,6 +490,17 @@ async fn test_snowflake_provider() -> Result<()> { .await } +#[tokio::test] +async fn test_sagemaker_tgi_provider() -> Result<()> { + test_provider( + "SageMakerTgi", + &["SAGEMAKER_ENDPOINT_NAME"], + None, + goose::providers::sagemaker_tgi::SageMakerTgiProvider::default, + ) + .await +} + // Print the final test report #[ctor::dtor] fn print_test_report() {