diff --git a/Cargo.lock b/Cargo.lock index 0ed36255f6..d27ec8c12f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2523,15 +2523,23 @@ name = "goose-llm" version = "1.0.22" dependencies = [ "anyhow", + "async-trait", + "base64 0.21.7", "chrono", - "goose", + "criterion", "include_dir", - "mcp-core", "minijinja", "once_cell", + "regex", + "reqwest 0.12.12", "serde", "serde_json", + "smallvec", + "tempfile", + "thiserror 1.0.69", "tokio", + "tracing", + "url", ] [[package]] @@ -5561,6 +5569,9 @@ name = "smallvec" version = "1.14.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7fcf8323ef1faaee30a44a340193b1ac6814fd9b7b4e88e9d4519a3e4abe1cfd" +dependencies = [ + "serde", +] [[package]] name = "smawk" diff --git a/crates/goose-llm/Cargo.toml b/crates/goose-llm/Cargo.toml index 1d73bff6aa..cdf2d8455b 100644 --- a/crates/goose-llm/Cargo.toml +++ b/crates/goose-llm/Cargo.toml @@ -8,16 +8,38 @@ repository.workspace = true description.workspace = true [dependencies] -goose = { path = "../goose" } -mcp-core = { path = "../mcp-core" } tokio = { version = "1.43", features = ["full"] } serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" anyhow = "1.0" +thiserror = "1.0" minijinja = "2.8.0" include_dir = "0.7.4" once_cell = "1.20.2" chrono = { version = "0.4.38", features = ["serde"] } +reqwest = { version = "0.12.9", features = [ + "rustls-tls-native-roots", + "json", + "cookies", + "gzip", + "brotli", + "deflate", + "zstd", + "charset", + "http2", + "stream" + ], default-features = false } +async-trait = "0.1" +url = "2.5" +base64 = "0.21" +regex = "1.11.1" +tracing = "0.1" +smallvec = { version = "1.13", features = ["serde"] } + +[dev-dependencies] +criterion = "0.5" +tempfile = "3.15.0" + [[example]] name = "simple" diff --git a/crates/goose-llm/examples/simple.rs b/crates/goose-llm/examples/simple.rs index 97e2d830bb..0fbfa568b2 100644 --- a/crates/goose-llm/examples/simple.rs +++ b/crates/goose-llm/examples/simple.rs @@ -1,19 +1,21 @@ use std::vec; use anyhow::Result; -use goose::message::Message; -use goose::model::ModelConfig; -use goose_llm::{completion, CompletionResponse, Extension}; -use mcp_core::tool::Tool; +use goose_llm::{ + completion, + types::completion::{CompletionResponse, ExtensionConfig, ToolApprovalMode, ToolConfig}, + Message, ModelConfig, +}; use serde_json::json; #[tokio::main] async fn main() -> Result<()> { let provider = "databricks"; - let model_name = "goose-claude-3-5-sonnet"; + // let model_name = "goose-claude-3-5-sonnet"; // sequential tool calls + let model_name = "goose-gpt-4-1"; // parallel tool calls let model_config = ModelConfig::new(model_name.to_string()); - let calculator_tool = Tool::new( + let calculator_tool = ToolConfig::new( "calculator", "Perform basic arithmetic operations", json!({ @@ -32,10 +34,10 @@ async fn main() -> Result<()> { } } }), - None, + ToolApprovalMode::Auto, ); - let bash_tool = Tool::new( + let bash_tool = ToolConfig::new( "bash_shell", "Run a shell command", json!({ @@ -48,28 +50,43 @@ async fn main() -> Result<()> { } } }), - None, + ToolApprovalMode::Manual, + ); + + let list_dir_tool = ToolConfig::new( + "list_directory", + "List files in a directory", + json!({ + "type": "object", + "required": ["path"], + "properties": { + "path": { + "type": "string", + "description": "The directory path to list files from" + } + } + }), + ToolApprovalMode::Auto, ); let extensions = vec![ - Extension::new( + ExtensionConfig::new( "calculator_extension".to_string(), Some("This extension provides a calculator tool.".to_string()), vec![calculator_tool], ), - Extension::new( + ExtensionConfig::new( "bash_extension".to_string(), Some("This extension provides a bash shell tool.".to_string()), - vec![bash_tool], + vec![bash_tool, list_dir_tool], ), ]; let system_preamble = "You are a helpful assistant."; for text in [ - "Add 10037 + 23123", - // "Write some random bad words to end of words.txt", - // "List all json files in the current directory and then multiply the count of the files by 7", + "Add 10037 + 23123 using calculator and also run 'date -u' using bash", + "List all files in the current directory", ] { println!("\n---------------\n"); println!("User Input: {text}"); diff --git a/crates/goose-llm/src/completion.rs b/crates/goose-llm/src/completion.rs index 999d7748d1..2a010217d4 100644 --- a/crates/goose-llm/src/completion.rs +++ b/crates/goose-llm/src/completion.rs @@ -1,17 +1,41 @@ +use std::{collections::HashMap, time::Instant}; + use anyhow::Result; use chrono::Utc; use serde_json::Value; -use std::collections::HashMap; -use goose::message::Message; -use goose::model::ModelConfig; -use goose::providers::create; -use goose::providers::errors::ProviderError; +use crate::{ + message::{Message, MessageContent}, + model::ModelConfig, + prompt_template, + providers::{create, errors::ProviderError}, + types::completion::{ + CompletionResponse, ExtensionConfig, RuntimeMetrics, ToolApprovalMode, ToolConfig, + }, +}; -use std::time::Instant; +/// Set `needs_approval` on *every* tool call in the message based on approval mode. +pub fn update_needs_approval_for_tool_calls( + message: &mut Message, + tool_configs: &HashMap, +) { + for content in message.content.iter_mut() { + if let MessageContent::ToolRequest(req) = content { + if let Ok(call) = &mut req.tool_call { + let needs = match tool_configs.get(&call.name) { + Some(cfg) => match cfg.approval_mode { + ToolApprovalMode::Auto => false, + ToolApprovalMode::Manual => true, + ToolApprovalMode::Smart => true, // TODO: implement smart approval later + }, + None => call.needs_approval, // unknown tool: leave flag unchanged + }; -use crate::prompt_template; -use crate::{CompletionResponse, Extension, RuntimeMetrics}; + call.set_needs_approval(needs); + } + } + } +} /// Public API for the Goose LLM completion function pub async fn completion( @@ -19,7 +43,7 @@ pub async fn completion( model_config: ModelConfig, system_preamble: &str, messages: &[Message], - extensions: &[Extension], + extensions: &[ExtensionConfig], ) -> Result { let start_total = Instant::now(); let provider = create(provider, model_config).unwrap(); @@ -32,11 +56,9 @@ pub async fn completion( .collect::>(); let start_provider = Instant::now(); - let (response, usage) = provider.complete(&system_prompt, messages, &tools).await?; + let mut response = provider.complete(&system_prompt, messages, &tools).await?; let total_time_ms_provider = start_provider.elapsed().as_millis(); - let total_time_ms = start_total.elapsed().as_millis(); - - let tokens_per_second = usage.usage.total_tokens.and_then(|toks| { + let tokens_per_second = response.usage.total_tokens.and_then(|toks| { if total_time_ms_provider > 0 { Some(toks as f64 / (total_time_ms_provider as f64 / 1000.0)) } else { @@ -44,15 +66,23 @@ pub async fn completion( } }); - let runtime_metrics = - RuntimeMetrics::new(total_time_ms, total_time_ms_provider, tokens_per_second); + let tool_configs: HashMap = extensions + .iter() + .flat_map(|ext| ext.get_prefixed_tool_configs().into_iter()) + .collect(); - let result = CompletionResponse::new(response.clone(), usage.clone(), runtime_metrics); + update_needs_approval_for_tool_calls(&mut response.message, &tool_configs); - Ok(result) + let total_time_ms = start_total.elapsed().as_millis(); + Ok(CompletionResponse::new( + response.message, + response.model, + response.usage, + RuntimeMetrics::new(total_time_ms, total_time_ms_provider, tokens_per_second), + )) } -fn construct_system_prompt(system_preamble: &str, extensions: &[Extension]) -> String { +fn construct_system_prompt(system_preamble: &str, extensions: &[ExtensionConfig]) -> String { let mut context: HashMap<&str, Value> = HashMap::new(); context.insert( diff --git a/crates/goose-llm/src/lib.rs b/crates/goose-llm/src/lib.rs index 9fbdbc4b44..3798cc9d50 100644 --- a/crates/goose-llm/src/lib.rs +++ b/crates/goose-llm/src/lib.rs @@ -1,6 +1,10 @@ mod completion; +mod message; +mod model; mod prompt_template; -mod types; +mod providers; +pub mod types; pub use completion::completion; -pub use types::{CompletionResponse, Extension, RuntimeMetrics}; +pub use message::Message; +pub use model::ModelConfig; diff --git a/crates/goose-llm/src/message.rs b/crates/goose-llm/src/message.rs new file mode 100644 index 0000000000..4bfa058d24 --- /dev/null +++ b/crates/goose-llm/src/message.rs @@ -0,0 +1,531 @@ +use std::{collections::HashSet, iter::FromIterator, ops::Deref}; + +/// Messages which represent the content sent back and forth to LLM provider +/// +/// We use these messages in the agent code, and interfaces which interact with +/// the agent. That let's us reuse message histories across different interfaces. +/// +/// The content of the messages uses MCP types to avoid additional conversions +/// when interacting with MCP servers. +use chrono::Utc; +use serde::{Deserialize, Serialize}; +use smallvec::SmallVec; + +use crate::types::core::{Content, ImageContent, Role, TextContent, ToolCall, ToolResult}; + +mod tool_result_serde; + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ToolRequest { + pub id: String, + #[serde(with = "tool_result_serde")] + pub tool_call: ToolResult, +} + +impl ToolRequest { + pub fn to_readable_string(&self) -> String { + match &self.tool_call { + Ok(tool_call) => { + format!( + "Tool: {}, Args: {}", + tool_call.name, + serde_json::to_string_pretty(&tool_call.arguments) + .unwrap_or_else(|_| "<>".to_string()) + ) + } + Err(e) => format!("Invalid tool call: {}", e), + } + } +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ToolResponse { + pub id: String, + #[serde(with = "tool_result_serde")] + pub tool_result: ToolResult>, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct ThinkingContent { + pub thinking: String, + pub signature: String, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct RedactedThinkingContent { + pub data: String, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +/// Content passed inside a message, which can be both simple content and tool content +#[serde(tag = "type", rename_all = "camelCase")] +pub enum MessageContent { + Text(TextContent), + Image(ImageContent), + ToolRequest(ToolRequest), + ToolResponse(ToolResponse), + Thinking(ThinkingContent), + RedactedThinking(RedactedThinkingContent), +} + +impl MessageContent { + pub fn text>(text: S) -> Self { + MessageContent::Text(TextContent { text: text.into() }) + } + + pub fn image, T: Into>(data: S, mime_type: T) -> Self { + MessageContent::Image(ImageContent { + data: data.into(), + mime_type: mime_type.into(), + }) + } + + pub fn tool_request>(id: S, tool_call: ToolResult) -> Self { + MessageContent::ToolRequest(ToolRequest { + id: id.into(), + tool_call, + }) + } + + pub fn tool_response>(id: S, tool_result: ToolResult>) -> Self { + MessageContent::ToolResponse(ToolResponse { + id: id.into(), + tool_result, + }) + } + + pub fn thinking, S2: Into>(thinking: S1, signature: S2) -> Self { + MessageContent::Thinking(ThinkingContent { + thinking: thinking.into(), + signature: signature.into(), + }) + } + + pub fn redacted_thinking>(data: S) -> Self { + MessageContent::RedactedThinking(RedactedThinkingContent { data: data.into() }) + } + + pub fn as_tool_request(&self) -> Option<&ToolRequest> { + if let MessageContent::ToolRequest(ref tool_request) = self { + Some(tool_request) + } else { + None + } + } + + pub fn as_tool_response(&self) -> Option<&ToolResponse> { + if let MessageContent::ToolResponse(ref tool_response) = self { + Some(tool_response) + } else { + None + } + } + + pub fn as_tool_response_text(&self) -> Option { + if let Some(tool_response) = self.as_tool_response() { + if let Ok(contents) = &tool_response.tool_result { + let texts: Vec = contents + .iter() + .filter_map(|content| content.as_text().map(String::from)) + .collect(); + if !texts.is_empty() { + return Some(texts.join("\n")); + } + } + } + None + } + + pub fn as_tool_request_id(&self) -> Option<&str> { + if let Self::ToolRequest(r) = self { + Some(&r.id) + } else { + None + } + } + + pub fn as_tool_response_id(&self) -> Option<&str> { + if let Self::ToolResponse(r) = self { + Some(&r.id) + } else { + None + } + } + + /// Get the text content if this is a TextContent variant + pub fn as_text(&self) -> Option<&str> { + match self { + MessageContent::Text(text) => Some(&text.text), + _ => None, + } + } + + /// Get the thinking content if this is a ThinkingContent variant + pub fn as_thinking(&self) -> Option<&ThinkingContent> { + match self { + MessageContent::Thinking(thinking) => Some(thinking), + _ => None, + } + } + + /// Get the redacted thinking content if this is a RedactedThinkingContent variant + pub fn as_redacted_thinking(&self) -> Option<&RedactedThinkingContent> { + match self { + MessageContent::RedactedThinking(redacted) => Some(redacted), + _ => None, + } + } + + pub fn is_text(&self) -> bool { + matches!(self, Self::Text(_)) + } + pub fn is_image(&self) -> bool { + matches!(self, Self::Image(_)) + } + pub fn is_tool_request(&self) -> bool { + matches!(self, Self::ToolRequest(_)) + } + pub fn is_tool_response(&self) -> bool { + matches!(self, Self::ToolResponse(_)) + } +} + +impl From for MessageContent { + fn from(content: Content) -> Self { + match content { + Content::Text(text) => MessageContent::Text(text), + Content::Image(image) => MessageContent::Image(image), + } + } +} + +// ──────────────────────────────────────────────────────────────────────────── +// 2. Contents – a new-type wrapper around SmallVec +// ──────────────────────────────────────────────────────────────────────────── + +/// Holds the heterogeneous fragments that make up one chat message. +/// +/// * Up to two items are stored inline on the stack. +/// * Falls back to a heap allocation only when necessary. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)] +#[serde(transparent)] +pub struct Contents(SmallVec<[MessageContent; 2]>); + +impl Contents { + /*---------------------------------------------------------- + * 1-line ergonomic helpers + *---------------------------------------------------------*/ + + pub fn iter_mut(&mut self) -> std::slice::IterMut<'_, MessageContent> { + self.0.iter_mut() + } + + pub fn push(&mut self, item: impl Into) { + self.0.push(item.into()); + } + + pub fn texts(&self) -> impl Iterator { + self.0.iter().filter_map(|c| c.as_text()) + } + + pub fn concat_text_str(&self) -> String { + self.texts().collect::>().join("\n") + } + + /// Returns `true` if *any* item satisfies the predicate. + pub fn any_is

(&self, pred: P) -> bool + where + P: FnMut(&MessageContent) -> bool, + { + self.iter().any(pred) + } + + /// Returns `true` if *every* item satisfies the predicate. + pub fn all_are

(&self, pred: P) -> bool + where + P: FnMut(&MessageContent) -> bool, + { + self.iter().all(pred) + } +} + +impl From> for Contents { + fn from(v: Vec) -> Self { + Contents(SmallVec::from_vec(v)) + } +} + +impl FromIterator for Contents { + fn from_iter>(iter: I) -> Self { + Contents(SmallVec::from_iter(iter)) + } +} + +/*-------------------------------------------------------------- + * Allow &message.content to behave like a slice of fragments. + *-------------------------------------------------------------*/ +impl Deref for Contents { + type Target = [MessageContent]; + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +/// A message to or from an LLM +#[serde(rename_all = "camelCase")] +pub struct Message { + pub role: Role, + pub created: i64, + pub content: Contents, +} + +impl Message { + pub fn new(role: Role) -> Self { + Self { + role, + created: Utc::now().timestamp(), + content: Contents::default(), + } + } + + /// Create a new user message with the current timestamp + pub fn user() -> Self { + Self::new(Role::User) + } + + /// Create a new assistant message with the current timestamp + pub fn assistant() -> Self { + Self::new(Role::Assistant) + } + + /// Add any item that implements Into to the message + pub fn with_content(mut self, item: impl Into) -> Self { + self.content.push(item); + self + } + + /// Add text content to the message + pub fn with_text>(self, text: S) -> Self { + self.with_content(MessageContent::text(text)) + } + + /// Add image content to the message + pub fn with_image, T: Into>(self, data: S, mime_type: T) -> Self { + self.with_content(MessageContent::image(data, mime_type)) + } + + /// Add a tool request to the message + pub fn with_tool_request>( + self, + id: S, + tool_call: ToolResult, + ) -> Self { + self.with_content(MessageContent::tool_request(id, tool_call)) + } + + /// Add a tool response to the message + pub fn with_tool_response>( + self, + id: S, + result: ToolResult>, + ) -> Self { + self.with_content(MessageContent::tool_response(id, result)) + } + + /// Add thinking content to the message + pub fn with_thinking, S2: Into>( + self, + thinking: S1, + signature: S2, + ) -> Self { + self.with_content(MessageContent::thinking(thinking, signature)) + } + + /// Add redacted thinking content to the message + pub fn with_redacted_thinking>(self, data: S) -> Self { + self.with_content(MessageContent::redacted_thinking(data)) + } + + /// Check if the message is a tool call + pub fn contains_tool_call(&self) -> bool { + self.content.any_is(MessageContent::is_tool_request) + } + + /// Check if the message is a tool response + pub fn contains_tool_response(&self) -> bool { + self.content.any_is(MessageContent::is_tool_response) + } + + /// Check if the message contains only text content + pub fn has_only_text_content(&self) -> bool { + self.content.all_are(MessageContent::is_text) + } + + /// Retrieves all tool `id` from ToolRequest messages + pub fn tool_request_ids(&self) -> HashSet<&str> { + self.content + .iter() + .filter_map(MessageContent::as_tool_request_id) + .collect() + } + + /// Retrieves all tool `id` from ToolResponse messages + pub fn tool_response_ids(&self) -> HashSet<&str> { + self.content + .iter() + .filter_map(MessageContent::as_tool_response_id) + .collect() + } + + /// Retrieves all tool `id` from the message + pub fn tool_ids(&self) -> HashSet<&str> { + self.tool_request_ids() + .into_iter() + .chain(self.tool_response_ids()) + .collect() + } +} + +#[cfg(test)] +mod tests { + use serde_json::{json, Value}; + + use super::*; + use crate::types::core::ToolError; + + #[test] + fn test_message_serialization() { + let message = Message::assistant() + .with_text("Hello, I'll help you with that.") + .with_tool_request( + "tool123", + Ok(ToolCall::new("test_tool", json!({"param": "value"}))), + ); + + let json_str = serde_json::to_string_pretty(&message).unwrap(); + println!("Serialized message: {}", json_str); + + // Parse back to Value to check structure + let value: Value = serde_json::from_str(&json_str).unwrap(); + + // Check top-level fields + assert_eq!(value["role"], "assistant"); + assert!(value["created"].is_i64()); + assert!(value["content"].is_array()); + + // Check content items + let content = &value["content"]; + + // First item should be text + assert_eq!(content[0]["type"], "text"); + assert_eq!(content[0]["text"], "Hello, I'll help you with that."); + + // Second item should be toolRequest + assert_eq!(content[1]["type"], "toolRequest"); + assert_eq!(content[1]["id"], "tool123"); + + // Check tool_call serialization + assert_eq!(content[1]["toolCall"]["status"], "success"); + assert_eq!(content[1]["toolCall"]["value"]["name"], "test_tool"); + assert_eq!( + content[1]["toolCall"]["value"]["arguments"]["param"], + "value" + ); + } + + #[test] + fn test_error_serialization() { + let message = Message::assistant().with_tool_request( + "tool123", + Err(ToolError::ExecutionError( + "Something went wrong".to_string(), + )), + ); + + let json_str = serde_json::to_string_pretty(&message).unwrap(); + println!("Serialized error: {}", json_str); + + // Parse back to Value to check structure + let value: Value = serde_json::from_str(&json_str).unwrap(); + + // Check tool_call serialization with error + let tool_call = &value["content"][0]["toolCall"]; + assert_eq!(tool_call["status"], "error"); + assert_eq!(tool_call["error"], "Execution failed: Something went wrong"); + } + + #[test] + fn test_deserialization() { + // Create a JSON string with our new format + let json_str = r#"{ + "role": "assistant", + "created": 1740171566, + "content": [ + { + "type": "text", + "text": "I'll help you with that." + }, + { + "type": "toolRequest", + "id": "tool123", + "toolCall": { + "status": "success", + "value": { + "name": "test_tool", + "arguments": {"param": "value"}, + "needsApproval": false + } + } + } + ] + }"#; + + let message: Message = serde_json::from_str(json_str).unwrap(); + + assert_eq!(message.role, Role::Assistant); + assert_eq!(message.created, 1740171566); + assert_eq!(message.content.len(), 2); + + // Check first content item + if let MessageContent::Text(text) = &message.content[0] { + assert_eq!(text.text, "I'll help you with that."); + } else { + panic!("Expected Text content"); + } + + // Check second content item + if let MessageContent::ToolRequest(req) = &message.content[1] { + assert_eq!(req.id, "tool123"); + if let Ok(tool_call) = &req.tool_call { + assert_eq!(tool_call.name, "test_tool"); + assert_eq!(tool_call.arguments, json!({"param": "value"})); + } else { + panic!("Expected successful tool call"); + } + } else { + panic!("Expected ToolRequest content"); + } + } + + #[test] + fn test_message_with_text() { + let message = Message::user().with_text("Hello"); + assert_eq!(message.content.concat_text_str(), "Hello"); + } + + #[test] + fn test_message_with_tool_request() { + let tool_call = Ok(ToolCall::new("test_tool", json!({}))); + + let message = Message::assistant().with_tool_request("req1", tool_call); + assert!(message.contains_tool_call()); + assert!(!message.contains_tool_response()); + + let ids = message.tool_ids(); + assert_eq!(ids.len(), 1); + assert!(ids.contains("req1")); + } +} diff --git a/crates/goose-llm/src/message/tool_result_serde.rs b/crates/goose-llm/src/message/tool_result_serde.rs new file mode 100644 index 0000000000..7f1143228d --- /dev/null +++ b/crates/goose-llm/src/message/tool_result_serde.rs @@ -0,0 +1,64 @@ +use serde::{ser::SerializeStruct, Deserialize, Deserializer, Serialize, Serializer}; + +use crate::types::core::{ToolError, ToolResult}; + +pub fn serialize(value: &ToolResult, serializer: S) -> Result +where + T: Serialize, + S: Serializer, +{ + match value { + Ok(val) => { + let mut state = serializer.serialize_struct("ToolResult", 2)?; + state.serialize_field("status", "success")?; + state.serialize_field("value", val)?; + state.end() + } + Err(err) => { + let mut state = serializer.serialize_struct("ToolResult", 2)?; + state.serialize_field("status", "error")?; + state.serialize_field("error", &err.to_string())?; + state.end() + } + } +} + +// For deserialization, let's use a simpler approach that works with the format we're serializing to +pub fn deserialize<'de, T, D>(deserializer: D) -> Result, D::Error> +where + T: Deserialize<'de>, + D: Deserializer<'de>, +{ + // Define a helper enum to handle the two possible formats + #[derive(Deserialize)] + #[serde(untagged)] + enum ResultFormat { + Success { status: String, value: T }, + Error { status: String, error: String }, + } + + let format = ResultFormat::deserialize(deserializer)?; + + match format { + ResultFormat::Success { status, value } => { + if status == "success" { + Ok(Ok(value)) + } else { + Err(serde::de::Error::custom(format!( + "Expected status 'success', got '{}'", + status + ))) + } + } + ResultFormat::Error { status, error } => { + if status == "error" { + Ok(Err(ToolError::ExecutionError(error))) + } else { + Err(serde::de::Error::custom(format!( + "Expected status 'error', got '{}'", + status + ))) + } + } + } +} diff --git a/crates/goose-llm/src/model.rs b/crates/goose-llm/src/model.rs new file mode 100644 index 0000000000..5e0d92d9cb --- /dev/null +++ b/crates/goose-llm/src/model.rs @@ -0,0 +1,118 @@ +use serde::{Deserialize, Serialize}; + +const DEFAULT_CONTEXT_LIMIT: usize = 128_000; + +/// Configuration for model-specific settings and limits +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ModelConfig { + /// The name of the model to use + pub model_name: String, + /// Optional explicit context limit that overrides any defaults + pub context_limit: Option, + /// Optional temperature setting (0.0 - 1.0) + pub temperature: Option, + /// Optional maximum tokens to generate + pub max_tokens: Option, +} + +impl ModelConfig { + /// Create a new ModelConfig with the specified model name + /// + /// The context limit is set with the following precedence: + /// 1. Explicit context_limit if provided in config + /// 2. Model-specific default based on model name + /// 3. Global default (128_000) (in get_context_limit) + pub fn new(model_name: String) -> Self { + let context_limit = Self::get_model_specific_limit(&model_name); + + Self { + model_name, + context_limit, + temperature: None, + max_tokens: None, + } + } + + /// Get model-specific context limit based on model name + fn get_model_specific_limit(model_name: &str) -> Option { + // Implement some sensible defaults + match model_name { + // OpenAI models, https://platform.openai.com/docs/models#models-overview + name if name.contains("gpt-4o") => Some(128_000), + name if name.contains("gpt-4-turbo") => Some(128_000), + + // Anthropic models, https://docs.anthropic.com/en/docs/about-claude/models + name if name.contains("claude-3") => Some(200_000), + + // Meta Llama models, https://github.com/meta-llama/llama-models/tree/main?tab=readme-ov-file#llama-models-1 + name if name.contains("llama3.2") => Some(128_000), + name if name.contains("llama3.3") => Some(128_000), + _ => None, + } + } + + /// Set an explicit context limit + pub fn with_context_limit(mut self, limit: Option) -> Self { + // Default is None and therefore DEFAULT_CONTEXT_LIMIT, only set + // if input is Some to allow passing through with_context_limit in + // configuration cases + if limit.is_some() { + self.context_limit = limit; + } + self + } + + /// Set the temperature + pub fn with_temperature(mut self, temp: Option) -> Self { + self.temperature = temp; + self + } + + /// Set the max tokens + pub fn with_max_tokens(mut self, tokens: Option) -> Self { + self.max_tokens = tokens; + self + } + + /// Get the context_limit for the current model + /// If none are defined, use the DEFAULT_CONTEXT_LIMIT + pub fn context_limit(&self) -> usize { + self.context_limit.unwrap_or(DEFAULT_CONTEXT_LIMIT) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_model_config_context_limits() { + // Test explicit limit + let config = + ModelConfig::new("claude-3-opus".to_string()).with_context_limit(Some(150_000)); + assert_eq!(config.context_limit(), 150_000); + + // Test model-specific defaults + let config = ModelConfig::new("claude-3-opus".to_string()); + assert_eq!(config.context_limit(), 200_000); + + let config = ModelConfig::new("gpt-4-turbo".to_string()); + assert_eq!(config.context_limit(), 128_000); + + // Test fallback to default + let config = ModelConfig::new("unknown-model".to_string()); + assert_eq!(config.context_limit(), DEFAULT_CONTEXT_LIMIT); + } + + #[test] + fn test_model_config_settings() { + let config = ModelConfig::new("test-model".to_string()) + .with_temperature(Some(0.7)) + .with_max_tokens(Some(1000)) + .with_context_limit(Some(50_000)); + + assert_eq!(config.temperature, Some(0.7)); + assert_eq!(config.max_tokens, Some(1000)); + assert_eq!(config.context_limit, Some(50_000)); + } +} diff --git a/crates/goose-llm/src/prompt_template.rs b/crates/goose-llm/src/prompt_template.rs index 8ed1c3da55..eca9facb6e 100644 --- a/crates/goose-llm/src/prompt_template.rs +++ b/crates/goose-llm/src/prompt_template.rs @@ -1,9 +1,12 @@ +use std::{ + path::PathBuf, + sync::{Arc, RwLock}, +}; + use include_dir::{include_dir, Dir}; use minijinja::{Environment, Error as MiniJinjaError, Value as MJValue}; use once_cell::sync::Lazy; use serde::Serialize; -use std::path::PathBuf; -use std::sync::{Arc, RwLock}; /// This directory will be embedded into the final binary. /// Typically used to store "core" or "system" prompts. diff --git a/crates/goose-llm/src/providers/base.rs b/crates/goose-llm/src/providers/base.rs new file mode 100644 index 0000000000..eb580490b7 --- /dev/null +++ b/crates/goose-llm/src/providers/base.rs @@ -0,0 +1,93 @@ +use anyhow::Result; +use async_trait::async_trait; +use serde::{Deserialize, Serialize}; + +use super::errors::ProviderError; +use crate::{message::Message, types::core::Tool}; + +#[derive(Debug, Clone, PartialEq, Default, Serialize, Deserialize)] +pub struct Usage { + pub input_tokens: Option, + pub output_tokens: Option, + pub total_tokens: Option, +} + +impl Usage { + pub fn new( + input_tokens: Option, + output_tokens: Option, + total_tokens: Option, + ) -> Self { + Self { + input_tokens, + output_tokens, + total_tokens, + } + } +} + +#[derive(Debug, Clone)] +pub struct ProviderCompleteResponse { + pub message: Message, + pub model: String, + pub usage: Usage, +} + +impl ProviderCompleteResponse { + pub fn new(message: Message, model: String, usage: Usage) -> Self { + Self { + message, + model, + usage, + } + } +} + +/// Base trait for AI providers (OpenAI, Anthropic, etc) +#[async_trait] +pub trait Provider: Send + Sync { + /// Generate the next message using the configured model and other parameters + /// + /// # Arguments + /// * `system` - The system prompt that guides the model's behavior + /// * `messages` - The conversation history as a sequence of messages + /// * `tools` - Optional list of tools the model can use + /// + /// # Returns + /// A tuple containing the model's response message and provider usage statistics + /// + /// # Errors + /// ProviderError + /// - It's important to raise ContextLengthExceeded correctly since agent handles it + async fn complete( + &self, + system: &str, + messages: &[Message], + tools: &[Tool], + ) -> Result; +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_usage_creation() { + let usage = Usage::new(Some(10), Some(20), Some(30)); + assert_eq!(usage.input_tokens, Some(10)); + assert_eq!(usage.output_tokens, Some(20)); + assert_eq!(usage.total_tokens, Some(30)); + } + + #[test] + fn test_provider_complete_response_creation() { + let message = Message::user().with_text("Hello, world!"); + let usage = Usage::new(Some(10), Some(20), Some(30)); + let response = + ProviderCompleteResponse::new(message.clone(), "test_model".to_string(), usage.clone()); + + assert_eq!(response.message, message); + assert_eq!(response.model, "test_model"); + assert_eq!(response.usage, usage); + } +} diff --git a/crates/goose-llm/src/providers/databricks.rs b/crates/goose-llm/src/providers/databricks.rs new file mode 100644 index 0000000000..013b8a8978 --- /dev/null +++ b/crates/goose-llm/src/providers/databricks.rs @@ -0,0 +1,218 @@ +use std::time::Duration; + +use anyhow::Result; +use async_trait::async_trait; +use reqwest::{Client, StatusCode}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use url::Url; + +use super::{ + errors::ProviderError, + formats::databricks::{create_request, get_usage, response_to_message}, + utils::{get_env, get_model, ImageFormat}, +}; +use crate::{ + message::Message, + model::ModelConfig, + providers::{Provider, ProviderCompleteResponse, Usage}, + types::core::Tool, +}; + +pub const DATABRICKS_DEFAULT_MODEL: &str = "databricks-meta-llama-3-3-70b-instruct"; +// Databricks can passthrough to a wide range of models, we only provide the default +pub const _DATABRICKS_KNOWN_MODELS: &[&str] = &[ + "databricks-meta-llama-3-3-70b-instruct", + "databricks-meta-llama-3-1-405b-instruct", + "databricks-dbrx-instruct", + "databricks-mixtral-8x7b-instruct", +]; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub enum DatabricksAuth { + Token(String), +} + +impl DatabricksAuth { + pub fn token(token: String) -> Self { + Self::Token(token) + } +} + +#[derive(Debug)] +pub struct DatabricksProvider { + client: Client, + host: String, + auth: DatabricksAuth, + model: ModelConfig, + image_format: ImageFormat, +} + +impl Default for DatabricksProvider { + fn default() -> Self { + let model = ModelConfig::new(DATABRICKS_DEFAULT_MODEL.to_string()); + DatabricksProvider::from_env(model).expect("Failed to initialize Databricks provider") + } +} + +impl DatabricksProvider { + pub fn from_env(model: ModelConfig) -> Result { + let host = get_env("DATABRICKS_HOST")?; + let api_key = get_env("DATABRICKS_TOKEN")?; + + let client = Client::builder() + .timeout(Duration::from_secs(600)) + .build()?; + + Ok(Self { + client, + host, + auth: DatabricksAuth::token(api_key), + model, + image_format: ImageFormat::OpenAi, + }) + } + + async fn ensure_auth_header(&self) -> Result { + match &self.auth { + DatabricksAuth::Token(token) => Ok(format!("Bearer {}", token)), + } + } + + async fn post(&self, payload: Value) -> Result { + let base_url = Url::parse(&self.host) + .map_err(|e| ProviderError::RequestFailed(format!("Invalid base URL: {e}")))?; + let path = format!("serving-endpoints/{}/invocations", self.model.model_name); + let url = base_url.join(&path).map_err(|e| { + ProviderError::RequestFailed(format!("Failed to construct endpoint URL: {e}")) + })?; + + let auth_header = self.ensure_auth_header().await?; + let response = self + .client + .post(url) + .header("Authorization", auth_header) + .json(&payload) + .send() + .await?; + + let status = response.status(); + let payload: Option = response.json().await.ok(); + + match status { + StatusCode::OK => payload.ok_or_else(|| { + ProviderError::RequestFailed("Response body is not valid JSON".to_string()) + }), + StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN => { + Err(ProviderError::Authentication(format!( + "Authentication failed. Please ensure your API keys are valid and have the required permissions. \ + Status: {}. Response: {:?}", + status, payload + ))) + } + StatusCode::BAD_REQUEST => { + // Databricks provides a generic 'error' but also includes 'external_model_message' which is provider specific + // We try to extract the error message from the payload and check for phrases that indicate context length exceeded + let payload_str = serde_json::to_string(&payload) + .unwrap_or_default() + .to_lowercase(); + let check_phrases = [ + "too long", + "context length", + "context_length_exceeded", + "reduce the length", + "token count", + "exceeds", + ]; + if check_phrases.iter().any(|c| payload_str.contains(c)) { + return Err(ProviderError::ContextLengthExceeded(payload_str)); + } + + let mut error_msg = "Unknown error".to_string(); + if let Some(payload) = &payload { + // try to convert message to string, if that fails use external_model_message + error_msg = payload + .get("message") + .and_then(|m| m.as_str()) + .or_else(|| { + payload + .get("external_model_message") + .and_then(|ext| ext.get("message")) + .and_then(|m| m.as_str()) + }) + .unwrap_or("Unknown error") + .to_string(); + } + + tracing::debug!( + "{}", + format!( + "Provider request failed with status: {}. Payload: {:?}", + status, payload + ) + ); + Err(ProviderError::RequestFailed(format!( + "Request failed with status: {}. Message: {}", + status, error_msg + ))) + } + StatusCode::TOO_MANY_REQUESTS => { + Err(ProviderError::RateLimitExceeded(format!("{:?}", payload))) + } + StatusCode::INTERNAL_SERVER_ERROR | StatusCode::SERVICE_UNAVAILABLE => { + Err(ProviderError::ServerError(format!("{:?}", payload))) + } + _ => { + tracing::debug!( + "{}", + format!( + "Provider request failed with status: {}. Payload: {:?}", + status, payload + ) + ); + Err(ProviderError::RequestFailed(format!( + "Request failed with status: {}", + status + ))) + } + } + } +} + +#[async_trait] +impl Provider for DatabricksProvider { + #[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 { + 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"); + + let response = self.post(payload.clone()).await?; + + // Parse response + let message = response_to_message(response.clone())?; + let usage = match get_usage(&response) { + Ok(usage) => usage, + Err(ProviderError::UsageError(e)) => { + tracing::debug!("Failed to get usage data: {}", e); + Usage::default() + } + Err(e) => return Err(e), + }; + let model = get_model(&response); + super::utils::emit_debug_trace(&self.model, &payload, &response, &usage); + + Ok(ProviderCompleteResponse::new(message, model, usage)) + } +} diff --git a/crates/goose-llm/src/providers/errors.rs b/crates/goose-llm/src/providers/errors.rs new file mode 100644 index 0000000000..4b31d28629 --- /dev/null +++ b/crates/goose-llm/src/providers/errors.rs @@ -0,0 +1,141 @@ +use thiserror::Error; + +#[derive(Error, Debug)] +pub enum ProviderError { + #[error("Authentication error: {0}")] + Authentication(String), + + #[error("Context length exceeded: {0}")] + ContextLengthExceeded(String), + + #[error("Rate limit exceeded: {0}")] + RateLimitExceeded(String), + + #[error("Server error: {0}")] + ServerError(String), + + #[error("Request failed: {0}")] + RequestFailed(String), + + #[error("Execution error: {0}")] + ExecutionError(String), + + #[error("Usage data error: {0}")] + UsageError(String), +} + +impl From for ProviderError { + fn from(error: anyhow::Error) -> Self { + ProviderError::ExecutionError(error.to_string()) + } +} + +impl From for ProviderError { + fn from(error: reqwest::Error) -> Self { + ProviderError::ExecutionError(error.to_string()) + } +} + +#[derive(serde::Deserialize, Debug)] +pub struct OpenAIError { + #[serde(deserialize_with = "code_as_string")] + pub code: Option, + pub message: Option, + #[serde(rename = "type")] + pub error_type: Option, +} + +fn code_as_string<'de, D>(deserializer: D) -> Result, D::Error> +where + D: serde::Deserializer<'de>, +{ + use std::fmt; + + use serde::de::{self, Visitor}; + + struct CodeVisitor; + + impl<'de> Visitor<'de> for CodeVisitor { + type Value = Option; + + fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result { + formatter.write_str("a string, a number, null, or none for the code field") + } + + fn visit_str(self, value: &str) -> Result + where + E: de::Error, + { + Ok(Some(value.to_string())) + } + + fn visit_u64(self, value: u64) -> Result + where + E: de::Error, + { + Ok(Some(value.to_string())) + } + + fn visit_none(self) -> Result + where + E: de::Error, + { + Ok(None) + } + + fn visit_unit(self) -> Result + where + E: de::Error, + { + Ok(None) + } + + fn visit_some(self, deserializer: D) -> Result + where + D: serde::Deserializer<'de>, + { + deserializer.deserialize_any(CodeVisitor) + } + } + + deserializer.deserialize_option(CodeVisitor) +} + +impl OpenAIError { + pub fn is_context_length_exceeded(&self) -> bool { + if let Some(code) = &self.code { + code == "context_length_exceeded" || code == "string_above_max_length" + } else { + false + } + } +} + +impl std::fmt::Display for OpenAIError { + /// Format the error for display. + /// E.g. {"message": "Invalid API key", "code": "invalid_api_key", "type": "client_error"} + /// would be formatted as "Invalid API key (code: invalid_api_key, type: client_error)" + /// and {"message": "Foo"} as just "Foo", etc. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + if let Some(message) = &self.message { + write!(f, "{}", message)?; + } + let mut in_parenthesis = false; + if let Some(code) = &self.code { + write!(f, " (code: {}", code)?; + in_parenthesis = true; + } + if let Some(typ) = &self.error_type { + if in_parenthesis { + write!(f, ", type: {}", typ)?; + } else { + write!(f, " (type: {}", typ)?; + in_parenthesis = true; + } + } + if in_parenthesis { + write!(f, ")")?; + } + Ok(()) + } +} diff --git a/crates/goose-llm/src/providers/factory.rs b/crates/goose-llm/src/providers/factory.rs new file mode 100644 index 0000000000..220f7ace3a --- /dev/null +++ b/crates/goose-llm/src/providers/factory.rs @@ -0,0 +1,15 @@ +use std::sync::Arc; + +use anyhow::Result; + +use super::{base::Provider, databricks::DatabricksProvider, openai::OpenAiProvider}; +use crate::model::ModelConfig; + +pub fn create(name: &str, model: ModelConfig) -> Result> { + // We use Arc instead of Box to be able to clone for multiple async tasks + match name { + "openai" => Ok(Arc::new(OpenAiProvider::from_env(model)?)), + "databricks" => Ok(Arc::new(DatabricksProvider::from_env(model)?)), + _ => Err(anyhow::anyhow!("Unknown provider: {}", name)), + } +} diff --git a/crates/goose-llm/src/providers/formats/databricks.rs b/crates/goose-llm/src/providers/formats/databricks.rs new file mode 100644 index 0000000000..8aece4d257 --- /dev/null +++ b/crates/goose-llm/src/providers/formats/databricks.rs @@ -0,0 +1,1115 @@ +use anyhow::{anyhow, Error}; +use serde_json::{json, Value}; + +use crate::{ + message::{Message, MessageContent}, + model::ModelConfig, + providers::{ + base::Usage, + errors::ProviderError, + utils::{ + convert_image, detect_image_path, is_valid_function_name, load_image_file, + sanitize_function_name, ImageFormat, + }, + }, + types::core::{Content, Role, Tool, ToolCall, ToolError}, +}; + +/// Convert internal Message format to Databricks' API message specification +/// Databricks is mostly OpenAI compatible, but has some differences (reasoning type, etc) +/// 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 format_messages(messages: &[Message], image_format: &ImageFormat) -> Vec { + let mut result = Vec::new(); + for message in messages { + let mut converted = json!({ + "role": message.role + }); + + let mut content_array = Vec::new(); + let mut has_tool_calls = false; + let mut has_multiple_content = false; + + for content in message.content.iter() { + match content { + MessageContent::Text(text) => { + if !text.text.is_empty() { + // Check for image paths in the text + if let Some(image_path) = detect_image_path(&text.text) { + has_multiple_content = true; + // Try to load and convert the image + if let Ok(image) = load_image_file(image_path) { + content_array.push(json!({ + "type": "text", + "text": text.text + })); + content_array.push(convert_image(&image, image_format)); + } else { + content_array.push(json!({ + "type": "text", + "text": text.text + })); + } + } else { + content_array.push(json!({ + "type": "text", + "text": text.text + })); + } + } + } + MessageContent::Thinking(content) => { + has_multiple_content = true; + content_array.push(json!({ + "type": "reasoning", + "summary": [ + { + "type": "summary_text", + "text": content.thinking, + "signature": content.signature + } + ] + })); + } + MessageContent::RedactedThinking(content) => { + has_multiple_content = true; + content_array.push(json!({ + "type": "reasoning", + "summary": [ + { + "type": "summary_encrypted_text", + "data": content.data + } + ] + })); + } + MessageContent::ToolRequest(request) => { + has_tool_calls = true; + match &request.tool_call { + Ok(tool_call) => { + let sanitized_name = sanitize_function_name(&tool_call.name); + + // Get mutable access to the "tool_calls" field in the converted object + // If "tool_calls" doesn't exist, insert an empty JSON array + let tool_calls = converted + .as_object_mut() + .unwrap() + .entry("tool_calls") + .or_insert(json!([])); + + tool_calls.as_array_mut().unwrap().push(json!({ + "id": request.id, + "type": "function", + "function": { + "name": sanitized_name, + "arguments": tool_call.arguments.to_string(), + } + })); + } + Err(e) => { + content_array.push(json!({ + "type": "text", + "text": format!("Error: {}", e) + })); + } + } + } + MessageContent::ToolResponse(response) => { + match &response.tool_result { + Ok(contents) => { + // Process all content, replacing images with placeholder text + let mut tool_content = Vec::new(); + let mut image_messages = Vec::new(); + + for content in contents { + match content { + Content::Image(image) => { + // Add placeholder text in the tool response + tool_content.push(Content::text("This tool result included an image that is uploaded in the next message.")); + + // Create a separate image message + image_messages.push(json!({ + "role": "user", + "content": [convert_image(image, image_format)] + })); + } + _ => { + tool_content.push(content.clone()); + } + } + } + let tool_response_content: Value = json!(tool_content + .iter() + .map(|content| match content { + Content::Text(text) => text.text.clone(), + _ => String::new(), + }) + .collect::>() + .join(" ")); + + // Add tool response as a separate message + result.push(json!({ + "role": "tool", + "content": tool_response_content, + "tool_call_id": response.id + })); + // Then add any image messages that need to follow + result.extend(image_messages); + } + Err(e) => { + // A tool result error is shown as output so the model can interpret the error message + result.push(json!({ + "role": "tool", + "content": format!("The tool call returned the following error:\n{}", e), + "tool_call_id": response.id + })); + } + } + } + MessageContent::Image(image) => { + // Handle direct image content + content_array.push(json!({ + "type": "image_url", + "image_url": { + "url": convert_image(image, image_format) + } + })); + } + } + } + + if !content_array.is_empty() { + // If we only have a single text content and no other special content, + // use the simple string format + if content_array.len() == 1 + && !has_multiple_content + && content_array[0]["type"] == "text" + { + converted["content"] = json!(content_array[0]["text"]); + } else { + converted["content"] = json!(content_array); + } + } + + if !content_array.is_empty() || has_tool_calls { + result.push(converted); + } + } + + result +} + +/// Convert internal Tool format to OpenAI's API tool specification +pub fn format_tools(tools: &[Tool]) -> anyhow::Result> { + let mut tool_names = std::collections::HashSet::new(); + let mut result = Vec::new(); + + for tool in tools { + if !tool_names.insert(&tool.name) { + return Err(anyhow!("Duplicate tool name: {}", tool.name)); + } + + let mut description = tool.description.clone(); + description.truncate(1024); + + // OpenAI's tool description max str len is 1024 + result.push(json!({ + "type": "function", + "function": { + "name": tool.name, + "description": description, + "parameters": tool.input_schema, + } + })); + } + + Ok(result) +} + +/// Convert Databricks' API response to internal Message format +pub fn response_to_message(response: Value) -> anyhow::Result { + let original = response["choices"][0]["message"].clone(); + let mut content: Vec = Vec::new(); + + // Handle array-based content + if let Some(content_array) = original.get("content").and_then(|c| c.as_array()) { + for content_item in content_array { + match content_item.get("type").and_then(|t| t.as_str()) { + Some("text") => { + if let Some(text) = content_item.get("text").and_then(|t| t.as_str()) { + content.push(MessageContent::text(text)); + } + } + Some("reasoning") => { + if let Some(summary_array) = + content_item.get("summary").and_then(|s| s.as_array()) + { + for summary in summary_array { + match summary.get("type").and_then(|t| t.as_str()) { + Some("summary_text") => { + let text = summary + .get("text") + .and_then(|t| t.as_str()) + .unwrap_or_default(); + let signature = summary + .get("signature") + .and_then(|s| s.as_str()) + .unwrap_or_default(); + content.push(MessageContent::thinking(text, signature)); + } + Some("summary_encrypted_text") => { + if let Some(data) = summary.get("data").and_then(|d| d.as_str()) + { + content.push(MessageContent::redacted_thinking(data)); + } + } + _ => continue, + } + } + } + } + _ => continue, + } + } + } else if let Some(text) = original.get("content").and_then(|t| t.as_str()) { + // Handle legacy single string content + content.push(MessageContent::text(text)); + } + + // Handle tool calls + if let Some(tool_calls) = original.get("tool_calls") { + if let Some(tool_calls_array) = tool_calls.as_array() { + for tool_call in tool_calls_array { + let id = tool_call["id"].as_str().unwrap_or_default().to_string(); + let function_name = tool_call["function"]["name"] + .as_str() + .unwrap_or_default() + .to_string(); + let mut arguments = tool_call["function"]["arguments"] + .as_str() + .unwrap_or_default() + .to_string(); + // If arguments is empty, we will have invalid json parsing error later. + if arguments.is_empty() { + arguments = "{}".to_string(); + } + + if !is_valid_function_name(&function_name) { + let error = ToolError::NotFound(format!( + "The provided function name '{}' had invalid characters, it must match this regex [a-zA-Z0-9_-]+", + function_name + )); + content.push(MessageContent::tool_request(id, Err(error))); + } else { + match serde_json::from_str::(&arguments) { + Ok(params) => { + content.push(MessageContent::tool_request( + id, + Ok(ToolCall::new(&function_name, params)), + )); + } + Err(e) => { + let error = ToolError::InvalidParameters(format!( + "Could not interpret tool use parameters for id {}: {}", + id, e + )); + content.push(MessageContent::tool_request(id, Err(error))); + } + } + } + } + } + } + + Ok(Message { + role: Role::Assistant, + created: chrono::Utc::now().timestamp(), + content: content.into(), + }) +} + +pub fn get_usage(data: &Value) -> Result { + let usage = data + .get("usage") + .ok_or_else(|| ProviderError::UsageError("No usage data in response".to_string()))?; + + let input_tokens = usage + .get("prompt_tokens") + .and_then(|v| v.as_i64()) + .map(|v| v as i32); + + let output_tokens = usage + .get("completion_tokens") + .and_then(|v| v.as_i64()) + .map(|v| v as i32); + + let total_tokens = usage + .get("total_tokens") + .and_then(|v| v.as_i64()) + .map(|v| v as i32) + .or_else(|| match (input_tokens, output_tokens) { + (Some(input), Some(output)) => Some(input + output), + _ => None, + }); + + Ok(Usage::new(input_tokens, output_tokens, total_tokens)) +} + +/// Validates and fixes tool schemas to ensure they have proper parameter structure. +/// If parameters exist, ensures they have properties and required fields, or removes parameters entirely. +pub fn validate_tool_schemas(tools: &mut [Value]) { + for tool in tools.iter_mut() { + if let Some(function) = tool.get_mut("function") { + if let Some(parameters) = function.get_mut("parameters") { + if parameters.is_object() { + ensure_valid_json_schema(parameters); + } + } + } + } +} + +/// Ensures that the given JSON value follows the expected JSON Schema structure. +fn ensure_valid_json_schema(schema: &mut Value) { + if let Some(params_obj) = schema.as_object_mut() { + // Check if this is meant to be an object type schema + let is_object_type = params_obj + .get("type") + .and_then(|t| t.as_str()) + .is_none_or(|t| t == "object"); // Default to true if no type is specified + + // Only apply full schema validation to object types + if is_object_type { + // Ensure required fields exist with default values + params_obj.entry("properties").or_insert_with(|| json!({})); + params_obj.entry("required").or_insert_with(|| json!([])); + params_obj.entry("type").or_insert_with(|| json!("object")); + + // Recursively validate properties if it exists + if let Some(properties) = params_obj.get_mut("properties") { + if let Some(properties_obj) = properties.as_object_mut() { + for (_key, prop) in properties_obj.iter_mut() { + if prop.is_object() + && prop.get("type").and_then(|t| t.as_str()) == Some("object") + { + ensure_valid_json_schema(prop); + } + } + } + } + } + } +} + +pub fn create_request( + model_config: &ModelConfig, + system: &str, + messages: &[Message], + tools: &[Tool], + image_format: &ImageFormat, +) -> anyhow::Result { + if model_config.model_name.starts_with("o1-mini") { + return Err(anyhow!( + "o1-mini model is not currently supported since Goose uses tool calling and o1-mini does not support it. Please use o1 or o3 models instead." + )); + } + + let model_name = model_config.model_name.to_string(); + let is_o1 = model_name.starts_with("o1") || model_name.starts_with("goose-o1"); + let is_o3 = model_name.starts_with("o3") || model_name.starts_with("goose-o3"); + let is_claude_3_7_sonnet = model_name.contains("claude-3-7-sonnet"); // can be goose- or databricks- + + // Only extract reasoning effort for O1/O3 models + let (model_name, reasoning_effort) = if is_o1 || is_o3 { + let parts: Vec<&str> = model_config.model_name.split('-').collect(); + let last_part = parts.last().unwrap(); + + match *last_part { + "low" | "medium" | "high" => { + let base_name = parts[..parts.len() - 1].join("-"); + (base_name, Some(last_part.to_string())) + } + _ => ( + model_config.model_name.to_string(), + Some("medium".to_string()), + ), + } + } else { + // For non-O family models, use the model name as is and no reasoning effort + (model_config.model_name.to_string(), None) + }; + + let system_message = json!({ + "role": if is_o1 || is_o3 { "developer" } else { "system" }, + "content": system + }); + + let messages_spec = format_messages(messages, image_format); + let mut tools_spec = if !tools.is_empty() { + format_tools(tools)? + } else { + vec![] + }; + + // Validate tool schemas + validate_tool_schemas(&mut tools_spec); + + let mut messages_array = vec![system_message]; + messages_array.extend(messages_spec); + + let mut payload = json!({ + "model": model_name, + "messages": messages_array + }); + + if let Some(effort) = reasoning_effort { + payload + .as_object_mut() + .unwrap() + .insert("reasoning_effort".to_string(), json!(effort)); + } + + if !tools_spec.is_empty() { + payload + .as_object_mut() + .unwrap() + .insert("tools".to_string(), json!(tools_spec)); + } + + // Add thinking parameters for Claude 3.7 Sonnet model when requested + let is_thinking_enabled = std::env::var("CLAUDE_THINKING_ENABLED").is_ok(); + if is_claude_3_7_sonnet && is_thinking_enabled { + // Minimum budget_tokens is 1024 + let budget_tokens = std::env::var("CLAUDE_THINKING_BUDGET") + .unwrap_or_else(|_| "16000".to_string()) + .parse() + .unwrap_or(16000); + + // For Claude models with thinking enabled, we need to add max_tokens + budget_tokens + // Default to 8192 (Claude max output) + budget if not specified + let max_completion_tokens = model_config.max_tokens.unwrap_or(8192); + payload.as_object_mut().unwrap().insert( + "max_tokens".to_string(), + json!(max_completion_tokens + budget_tokens), + ); + + payload.as_object_mut().unwrap().insert( + "thinking".to_string(), + json!({ + "type": "enabled", + "budget_tokens": budget_tokens + }), + ); + + // Temperature is fixed to 2 when using claude 3.7 thinking with Databricks + payload + .as_object_mut() + .unwrap() + .insert("temperature".to_string(), json!(2)); + } else { + // o1, o3 models currently don't support temperature + if !is_o1 && !is_o3 { + if let Some(temp) = model_config.temperature { + payload + .as_object_mut() + .unwrap() + .insert("temperature".to_string(), json!(temp)); + } + } + + // o1 models use max_completion_tokens instead of max_tokens + if let Some(tokens) = model_config.max_tokens { + let key = if is_o1 || is_o3 { + "max_completion_tokens" + } else { + "max_tokens" + }; + payload + .as_object_mut() + .unwrap() + .insert(key.to_string(), json!(tokens)); + } + } + + Ok(payload) +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + use crate::types::core::Content; + + #[test] + fn test_validate_tool_schemas() { + // Test case 1: Empty parameters object + // Input JSON with an incomplete parameters object + let mut actual = vec![json!({ + "type": "function", + "function": { + "name": "test_func", + "description": "test description", + "parameters": { + "type": "object" + } + } + })]; + + // Run the function to validate and update schemas + validate_tool_schemas(&mut actual); + + // Expected JSON after validation + let expected = vec![json!({ + "type": "function", + "function": { + "name": "test_func", + "description": "test description", + "parameters": { + "type": "object", + "properties": {}, + "required": [] + } + } + })]; + + // Compare entire JSON structures instead of individual fields + assert_eq!(actual, expected); + + // Test case 2: Missing type field + let mut tools = vec![json!({ + "type": "function", + "function": { + "name": "test_func", + "description": "test description", + "parameters": { + "properties": {} + } + } + })]; + + validate_tool_schemas(&mut tools); + + let params = tools[0]["function"]["parameters"].as_object().unwrap(); + assert_eq!(params["type"], "object"); + + // Test case 3: Complete valid schema should remain unchanged + let original_schema = json!({ + "type": "function", + "function": { + "name": "test_func", + "description": "test description", + "parameters": { + "type": "object", + "properties": { + "location": { + "type": "string", + "description": "City and country" + } + }, + "required": ["location"] + } + } + }); + + let mut tools = vec![original_schema.clone()]; + validate_tool_schemas(&mut tools); + assert_eq!(tools[0], original_schema); + } + + const OPENAI_TOOL_USE_RESPONSE: &str = r#"{ + "choices": [{ + "role": "assistant", + "message": { + "tool_calls": [{ + "id": "1", + "function": { + "name": "example_fn", + "arguments": "{\"param\": \"value\"}" + } + }] + } + }], + "usage": { + "input_tokens": 10, + "output_tokens": 25, + "total_tokens": 35 + } + }"#; + + #[test] + fn test_format_messages() -> anyhow::Result<()> { + let message = Message::user().with_text("Hello"); + let spec = format_messages(&[message], &ImageFormat::OpenAi); + + assert_eq!(spec.len(), 1); + assert_eq!(spec[0]["role"], "user"); + assert_eq!(spec[0]["content"], "Hello"); + Ok(()) + } + + #[test] + fn test_format_tools() -> anyhow::Result<()> { + let tool = Tool::new( + "test_tool", + "A test tool", + json!({ + "type": "object", + "properties": { + "input": { + "type": "string", + "description": "Test parameter" + } + }, + "required": ["input"] + }), + ); + + let spec = format_tools(&[tool])?; + + assert_eq!(spec.len(), 1); + assert_eq!(spec[0]["type"], "function"); + assert_eq!(spec[0]["function"]["name"], "test_tool"); + Ok(()) + } + + #[test] + fn test_format_messages_complex() -> anyhow::Result<()> { + let mut messages = vec![ + Message::assistant().with_text("Hello!"), + Message::user().with_text("How are you?"), + Message::assistant().with_tool_request( + "tool1", + Ok(ToolCall::new("example", json!({"param1": "value1"}))), + ), + ]; + + // Get the ID from the tool request to use in the response + let tool_id = if let MessageContent::ToolRequest(request) = &messages[2].content[0] { + request.id.clone() + } else { + panic!("should be tool request"); + }; + + messages + .push(Message::user().with_tool_response(tool_id, Ok(vec![Content::text("Result")]))); + + let spec = format_messages(&messages, &ImageFormat::OpenAi); + + assert_eq!(spec.len(), 4); + assert_eq!(spec[0]["role"], "assistant"); + assert_eq!(spec[0]["content"], "Hello!"); + assert_eq!(spec[1]["role"], "user"); + assert_eq!(spec[1]["content"], "How are you?"); + assert_eq!(spec[2]["role"], "assistant"); + assert!(spec[2]["tool_calls"].is_array()); + assert_eq!(spec[3]["role"], "tool"); + assert_eq!(spec[3]["content"], "Result"); + assert_eq!(spec[3]["tool_call_id"], spec[2]["tool_calls"][0]["id"]); + + Ok(()) + } + + #[test] + 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"}))), + )]; + + // Get the ID from the tool request to use in the response + let tool_id = if let MessageContent::ToolRequest(request) = &messages[0].content[0] { + request.id.clone() + } else { + panic!("should be tool request"); + }; + + messages + .push(Message::user().with_tool_response(tool_id, Ok(vec![Content::text("Result")]))); + + 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"], "Result"); + assert_eq!(spec[1]["tool_call_id"], spec[0]["tool_calls"][0]["id"]); + + Ok(()) + } + + #[test] + fn test_format_tools_duplicate() -> anyhow::Result<()> { + let tool1 = Tool::new( + "test_tool", + "Test tool", + json!({ + "type": "object", + "properties": { + "input": { + "type": "string", + "description": "Test parameter" + } + }, + "required": ["input"] + }), + ); + + let tool2 = Tool::new( + "test_tool", + "Test tool", + json!({ + "type": "object", + "properties": { + "input": { + "type": "string", + "description": "Test parameter" + } + }, + "required": ["input"] + }), + ); + + let result = format_tools(&[tool1, tool2]); + assert!(result.is_err()); + assert!(result + .unwrap_err() + .to_string() + .contains("Duplicate tool name")); + + Ok(()) + } + + #[test] + fn test_format_tools_empty() -> anyhow::Result<()> { + let spec = format_tools(&[])?; + assert!(spec.is_empty()); + Ok(()) + } + + #[test] + fn test_format_messages_with_image_path() -> anyhow::Result<()> { + // Create a temporary PNG file with valid PNG magic numbers + let temp_dir = tempfile::tempdir()?; + let png_path = temp_dir.path().join("test.png"); + let png_data = [ + 0x89, 0x50, 0x4E, 0x47, // PNG magic number + 0x0D, 0x0A, 0x1A, 0x0A, // PNG header + 0x00, 0x00, 0x00, 0x0D, // Rest of fake PNG data + ]; + std::fs::write(&png_path, &png_data)?; + let png_path_str = png_path.to_str().unwrap(); + + // Create message with image path + let message = Message::user().with_text(format!("Here is an image: {}", png_path_str)); + let spec = format_messages(&[message], &ImageFormat::OpenAi); + + assert_eq!(spec.len(), 1); + assert_eq!(spec[0]["role"], "user"); + + // Content should be an array with text and image + let content = spec[0]["content"].as_array().unwrap(); + assert_eq!(content.len(), 2); + assert_eq!(content[0]["type"], "text"); + assert!(content[0]["text"].as_str().unwrap().contains(png_path_str)); + assert_eq!(content[1]["type"], "image_url"); + assert!(content[1]["image_url"]["url"] + .as_str() + .unwrap() + .starts_with("data:image/png;base64,")); + + Ok(()) + } + + #[test] + fn test_response_to_message_text() -> anyhow::Result<()> { + let response = json!({ + "choices": [{ + "role": "assistant", + "message": { + "content": "Hello from John Cena!" + } + }], + "usage": { + "input_tokens": 10, + "output_tokens": 25, + "total_tokens": 35 + } + }); + + 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!"); + } else { + panic!("Expected Text content"); + } + assert!(matches!(message.role, Role::Assistant)); + + Ok(()) + } + + #[test] + fn test_response_to_message_valid_toolrequest() -> anyhow::Result<()> { + let response: Value = serde_json::from_str(OPENAI_TOOL_USE_RESPONSE)?; + let message = response_to_message(response)?; + + assert_eq!(message.content.len(), 1); + if let MessageContent::ToolRequest(request) = &message.content[0] { + let tool_call = request.tool_call.as_ref().unwrap(); + assert_eq!(tool_call.name, "example_fn"); + assert_eq!(tool_call.arguments, json!({"param": "value"})); + } else { + panic!("Expected ToolRequest content"); + } + + Ok(()) + } + + #[test] + 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 = response_to_message(response)?; + + if let MessageContent::ToolRequest(request) = &message.content[0] { + match &request.tool_call { + Err(ToolError::NotFound(msg)) => { + assert!(msg.starts_with("The provided function name")); + } + _ => panic!("Expected ToolNotFound error"), + } + } else { + panic!("Expected ToolRequest content"); + } + + Ok(()) + } + + #[test] + 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 = response_to_message(response)?; + + if let MessageContent::ToolRequest(request) = &message.content[0] { + match &request.tool_call { + Err(ToolError::InvalidParameters(msg)) => { + assert!(msg.starts_with("Could not interpret tool use parameters")); + } + _ => panic!("Expected InvalidParameters error"), + } + } else { + panic!("Expected ToolRequest content"); + } + + Ok(()) + } + + #[test] + fn test_response_to_message_empty_argument() -> anyhow::Result<()> { + let mut response: Value = serde_json::from_str(OPENAI_TOOL_USE_RESPONSE)?; + response["choices"][0]["message"]["tool_calls"][0]["function"]["arguments"] = + serde_json::Value::String("".to_string()); + + let message = response_to_message(response)?; + + if let MessageContent::ToolRequest(request) = &message.content[0] { + let tool_call = request.tool_call.as_ref().unwrap(); + assert_eq!(tool_call.name, "example_fn"); + assert_eq!(tool_call.arguments, json!({})); + } else { + panic!("Expected ToolRequest content"); + } + + Ok(()) + } + + #[test] + fn test_create_request_gpt_4o() -> anyhow::Result<()> { + // Test default medium reasoning effort for O3 model + let model_config = ModelConfig { + model_name: "gpt-4o".to_string(), + context_limit: Some(4096), + temperature: None, + max_tokens: Some(1024), + }; + let request = create_request(&model_config, "system", &[], &[], &ImageFormat::OpenAi)?; + let obj = request.as_object().unwrap(); + let expected = json!({ + "model": "gpt-4o", + "messages": [ + { + "role": "system", + "content": "system" + } + ], + "max_tokens": 1024 + }); + + for (key, value) in expected.as_object().unwrap() { + assert_eq!(obj.get(key).unwrap(), value); + } + + Ok(()) + } + + #[test] + fn test_create_request_o1_default() -> anyhow::Result<()> { + // Test default medium reasoning effort for O1 model + let model_config = ModelConfig { + model_name: "o1".to_string(), + context_limit: Some(4096), + temperature: None, + max_tokens: Some(1024), + }; + let request = create_request(&model_config, "system", &[], &[], &ImageFormat::OpenAi)?; + let obj = request.as_object().unwrap(); + let expected = json!({ + "model": "o1", + "messages": [ + { + "role": "developer", + "content": "system" + } + ], + "reasoning_effort": "medium", + "max_completion_tokens": 1024 + }); + + for (key, value) in expected.as_object().unwrap() { + assert_eq!(obj.get(key).unwrap(), value); + } + + Ok(()) + } + + #[test] + fn test_create_request_o3_custom_reasoning_effort() -> anyhow::Result<()> { + // Test custom reasoning effort for O3 model + let model_config = ModelConfig { + model_name: "o3-mini-high".to_string(), + context_limit: Some(4096), + temperature: None, + max_tokens: Some(1024), + }; + let request = create_request(&model_config, "system", &[], &[], &ImageFormat::OpenAi)?; + let obj = request.as_object().unwrap(); + let expected = json!({ + "model": "o3-mini", + "messages": [ + { + "role": "developer", + "content": "system" + } + ], + "reasoning_effort": "high", + "max_completion_tokens": 1024 + }); + + for (key, value) in expected.as_object().unwrap() { + assert_eq!(obj.get(key).unwrap(), value); + } + + Ok(()) + } + + #[test] + fn test_response_to_message_claude_thinking() -> anyhow::Result<()> { + let response = json!({ + "model": "us.anthropic.claude-3-7-sonnet-20250219-v1:0", + "choices": [{ + "message": { + "role": "assistant", + "content": [ + { + "type": "reasoning", + "summary": [ + { + "type": "summary_text", + "text": "Test thinking content", + "signature": "test-signature" + } + ] + }, + { + "type": "text", + "text": "Regular text content" + } + ] + }, + "index": 0, + "finish_reason": "stop" + }] + }); + + let message = response_to_message(response)?; + assert_eq!(message.content.len(), 2); + + if let MessageContent::Thinking(thinking) = &message.content[0] { + assert_eq!(thinking.thinking, "Test thinking content"); + assert_eq!(thinking.signature, "test-signature"); + } else { + panic!("Expected Thinking content"); + } + + if let MessageContent::Text(text) = &message.content[1] { + assert_eq!(text.text, "Regular text content"); + } else { + panic!("Expected Text content"); + } + + Ok(()) + } + + #[test] + fn test_response_to_message_claude_encrypted_thinking() -> anyhow::Result<()> { + let response = json!({ + "model": "claude-3-7-sonnet-20250219", + "choices": [{ + "message": { + "role": "assistant", + "content": [ + { + "type": "reasoning", + "summary": [ + { + "type": "summary_encrypted_text", + "data": "E23sQFCkYIARgCKkATCHitsdf327Ber3v4NYUq2" + } + ] + }, + { + "type": "text", + "text": "Regular text content" + } + ] + }, + "index": 0, + "finish_reason": "stop" + }] + }); + + let message = response_to_message(response)?; + assert_eq!(message.content.len(), 2); + + if let MessageContent::RedactedThinking(redacted) = &message.content[0] { + assert_eq!(redacted.data, "E23sQFCkYIARgCKkATCHitsdf327Ber3v4NYUq2"); + } else { + panic!("Expected RedactedThinking content"); + } + + if let MessageContent::Text(text) = &message.content[1] { + assert_eq!(text.text, "Regular text content"); + } else { + panic!("Expected Text content"); + } + + Ok(()) + } +} diff --git a/crates/goose-llm/src/providers/formats/mod.rs b/crates/goose-llm/src/providers/formats/mod.rs new file mode 100644 index 0000000000..cf929f39cc --- /dev/null +++ b/crates/goose-llm/src/providers/formats/mod.rs @@ -0,0 +1,2 @@ +pub mod databricks; +pub mod openai; diff --git a/crates/goose-llm/src/providers/formats/openai.rs b/crates/goose-llm/src/providers/formats/openai.rs new file mode 100644 index 0000000000..b1e96e05cb --- /dev/null +++ b/crates/goose-llm/src/providers/formats/openai.rs @@ -0,0 +1,895 @@ +use anyhow::{anyhow, Error}; +use serde_json::{json, Value}; + +use crate::{ + message::{Message, MessageContent}, + model::ModelConfig, + providers::{ + base::Usage, + errors::ProviderError, + utils::{ + convert_image, detect_image_path, is_valid_function_name, load_image_file, + sanitize_function_name, ImageFormat, + }, + }, + types::core::{Content, Role, Tool, ToolCall, ToolError}, +}; + +/// 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 format_messages(messages: &[Message], image_format: &ImageFormat) -> Vec { + let mut messages_spec = Vec::new(); + for message in messages { + let mut converted = json!({ + "role": message.role + }); + + let mut output = Vec::new(); + + for content in message.content.iter() { + match content { + MessageContent::Text(text) => { + if !text.text.is_empty() { + // Check for image paths in the text + if let Some(image_path) = detect_image_path(&text.text) { + // Try to load and convert the image + if let Ok(image) = load_image_file(image_path) { + converted["content"] = json!([ + {"type": "text", "text": text.text}, + convert_image(&image, image_format) + ]); + } else { + // If image loading fails, just use the text + converted["content"] = json!(text.text); + } + } else { + converted["content"] = json!(text.text); + } + } + } + MessageContent::Thinking(_) => { + // Thinking blocks are not directly used in OpenAI format + continue; + } + MessageContent::RedactedThinking(_) => { + // Redacted thinking blocks are not directly used in OpenAI format + continue; + } + MessageContent::ToolRequest(request) => match &request.tool_call { + Ok(tool_call) => { + let sanitized_name = sanitize_function_name(&tool_call.name); + let tool_calls = converted + .as_object_mut() + .unwrap() + .entry("tool_calls") + .or_insert(json!([])); + + tool_calls.as_array_mut().unwrap().push(json!({ + "id": request.id, + "type": "function", + "function": { + "name": sanitized_name, + "arguments": tool_call.arguments.to_string(), + } + })); + } + Err(e) => { + output.push(json!({ + "role": "tool", + "content": format!("Error: {}", e), + "tool_call_id": request.id + })); + } + }, + MessageContent::ToolResponse(response) => { + match &response.tool_result { + Ok(contents) => { + // Process all content, replacing images with placeholder text + let mut tool_content = Vec::new(); + let mut image_messages = Vec::new(); + + for content in contents { + match content { + Content::Image(image) => { + // Add placeholder text in the tool response + tool_content.push(Content::text("This tool result included an image that is uploaded in the next message.")); + + // Create a separate image message + image_messages.push(json!({ + "role": "user", + "content": [convert_image(image, image_format)] + })); + } + _ => { + tool_content.push(content.clone()); + } + } + } + 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!({ + "role": "tool", + "content": tool_response_content, + "tool_call_id": response.id + })); + // Then add any image messages that need to follow + output.extend(image_messages); + } + Err(e) => { + // A tool result error is shown as output so the model can interpret the error message + output.push(json!({ + "role": "tool", + "content": format!("The tool call returned the following error:\n{}", e), + "tool_call_id": response.id + })); + } + } + } + MessageContent::Image(image) => { + // Handle direct image content + converted["content"] = json!([convert_image(image, image_format)]); + } + } + } + + if converted.get("content").is_some() || converted.get("tool_calls").is_some() { + output.insert(0, converted); + } + messages_spec.extend(output); + } + + messages_spec +} + +/// Convert internal Tool format to OpenAI's API tool specification +pub fn format_tools(tools: &[Tool]) -> anyhow::Result> { + let mut tool_names = std::collections::HashSet::new(); + let mut result = Vec::new(); + + for tool in tools { + if !tool_names.insert(&tool.name) { + return Err(anyhow!("Duplicate tool name: {}", tool.name)); + } + + let mut description = tool.description.clone(); + description.truncate(1024); + + // OpenAI's tool description max str len is 1024 + result.push(json!({ + "type": "function", + "function": { + "name": tool.name, + "description": description, + "parameters": tool.input_schema, + } + })); + } + + Ok(result) +} + +/// Convert OpenAI's API response to internal Message format +pub fn response_to_message(response: Value) -> anyhow::Result { + let original = response["choices"][0]["message"].clone(); + let mut content = Vec::new(); + + if let Some(text) = original.get("content") { + if let Some(text_str) = text.as_str() { + content.push(MessageContent::text(text_str)); + } + } + + if let Some(tool_calls) = original.get("tool_calls") { + if let Some(tool_calls_array) = tool_calls.as_array() { + for tool_call in tool_calls_array { + let id = tool_call["id"].as_str().unwrap_or_default().to_string(); + let function_name = tool_call["function"]["name"] + .as_str() + .unwrap_or_default() + .to_string(); + let mut arguments = tool_call["function"]["arguments"] + .as_str() + .unwrap_or_default() + .to_string(); + // If arguments is empty, we will have invalid json parsing error later. + if arguments.is_empty() { + arguments = "{}".to_string(); + } + + if !is_valid_function_name(&function_name) { + let error = ToolError::NotFound(format!( + "The provided function name '{}' had invalid characters, it must match this regex [a-zA-Z0-9_-]+", + function_name + )); + content.push(MessageContent::tool_request(id, Err(error))); + } else { + match serde_json::from_str::(&arguments) { + Ok(params) => { + content.push(MessageContent::tool_request( + id, + Ok(ToolCall::new(&function_name, params)), + )); + } + Err(e) => { + let error = ToolError::InvalidParameters(format!( + "Could not interpret tool use parameters for id {}: {}", + id, e + )); + content.push(MessageContent::tool_request(id, Err(error))); + } + } + } + } + } + } + + Ok(Message { + role: Role::Assistant, + created: chrono::Utc::now().timestamp(), + content: content.into(), + }) +} + +pub fn get_usage(data: &Value) -> Result { + let usage = data + .get("usage") + .ok_or_else(|| ProviderError::UsageError("No usage data in response".to_string()))?; + + let input_tokens = usage + .get("prompt_tokens") + .and_then(|v| v.as_i64()) + .map(|v| v as i32); + + let output_tokens = usage + .get("completion_tokens") + .and_then(|v| v.as_i64()) + .map(|v| v as i32); + + let total_tokens = usage + .get("total_tokens") + .and_then(|v| v.as_i64()) + .map(|v| v as i32) + .or_else(|| match (input_tokens, output_tokens) { + (Some(input), Some(output)) => Some(input + output), + _ => None, + }); + + Ok(Usage::new(input_tokens, output_tokens, total_tokens)) +} + +/// Validates and fixes tool schemas to ensure they have proper parameter structure. +/// If parameters exist, ensures they have properties and required fields, or removes parameters entirely. +pub fn validate_tool_schemas(tools: &mut [Value]) { + for tool in tools.iter_mut() { + if let Some(function) = tool.get_mut("function") { + if let Some(parameters) = function.get_mut("parameters") { + if parameters.is_object() { + ensure_valid_json_schema(parameters); + } + } + } + } +} + +/// Ensures that the given JSON value follows the expected JSON Schema structure. +fn ensure_valid_json_schema(schema: &mut Value) { + if let Some(params_obj) = schema.as_object_mut() { + // Check if this is meant to be an object type schema + let is_object_type = params_obj + .get("type") + .and_then(|t| t.as_str()) + .is_none_or(|t| t == "object"); // Default to true if no type is specified + + // Only apply full schema validation to object types + if is_object_type { + // Ensure required fields exist with default values + params_obj.entry("properties").or_insert_with(|| json!({})); + params_obj.entry("required").or_insert_with(|| json!([])); + params_obj.entry("type").or_insert_with(|| json!("object")); + + // Recursively validate properties if it exists + if let Some(properties) = params_obj.get_mut("properties") { + if let Some(properties_obj) = properties.as_object_mut() { + for (_key, prop) in properties_obj.iter_mut() { + if prop.is_object() + && prop.get("type").and_then(|t| t.as_str()) == Some("object") + { + ensure_valid_json_schema(prop); + } + } + } + } + } + } +} + +pub fn create_request( + model_config: &ModelConfig, + system: &str, + messages: &[Message], + tools: &[Tool], + image_format: &ImageFormat, +) -> anyhow::Result { + if model_config.model_name.starts_with("o1-mini") { + return Err(anyhow!( + "o1-mini model is not currently supported since Goose uses tool calling and o1-mini does not support it. Please use o1 or o3 models instead." + )); + } + + let is_ox_model = model_config.model_name.starts_with("o"); + + // Only extract reasoning effort for O1/O3 models + let (model_name, reasoning_effort) = if is_ox_model { + let parts: Vec<&str> = model_config.model_name.split('-').collect(); + let last_part = parts.last().unwrap(); + + match *last_part { + "low" | "medium" | "high" => { + let base_name = parts[..parts.len() - 1].join("-"); + (base_name, Some(last_part.to_string())) + } + _ => ( + model_config.model_name.to_string(), + Some("medium".to_string()), + ), + } + } else { + // For non-O family models, use the model name as is and no reasoning effort + (model_config.model_name.to_string(), None) + }; + + let system_message = json!({ + "role": if is_ox_model { "developer" } else { "system" }, + "content": system + }); + + let messages_spec = format_messages(messages, image_format); + let mut tools_spec = if !tools.is_empty() { + format_tools(tools)? + } else { + vec![] + }; + + // Validate tool schemas + validate_tool_schemas(&mut tools_spec); + + let mut messages_array = vec![system_message]; + messages_array.extend(messages_spec); + + let mut payload = json!({ + "model": model_name, + "messages": messages_array + }); + + if let Some(effort) = reasoning_effort { + payload + .as_object_mut() + .unwrap() + .insert("reasoning_effort".to_string(), json!(effort)); + } + + if !tools_spec.is_empty() { + payload + .as_object_mut() + .unwrap() + .insert("tools".to_string(), json!(tools_spec)); + } + // o1, o3 models currently don't support temperature + if !is_ox_model { + if let Some(temp) = model_config.temperature { + payload + .as_object_mut() + .unwrap() + .insert("temperature".to_string(), json!(temp)); + } + } + + // o1 models use max_completion_tokens instead of max_tokens + if let Some(tokens) = model_config.max_tokens { + let key = if is_ox_model { + "max_completion_tokens" + } else { + "max_tokens" + }; + payload + .as_object_mut() + .unwrap() + .insert(key.to_string(), json!(tokens)); + } + Ok(payload) +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + use crate::types::core::Content; + + #[test] + fn test_validate_tool_schemas() { + // Test case 1: Empty parameters object + // Input JSON with an incomplete parameters object + let mut actual = vec![json!({ + "type": "function", + "function": { + "name": "test_func", + "description": "test description", + "parameters": { + "type": "object" + } + } + })]; + + // Run the function to validate and update schemas + validate_tool_schemas(&mut actual); + + // Expected JSON after validation + let expected = vec![json!({ + "type": "function", + "function": { + "name": "test_func", + "description": "test description", + "parameters": { + "type": "object", + "properties": {}, + "required": [] + } + } + })]; + + // Compare entire JSON structures instead of individual fields + assert_eq!(actual, expected); + + // Test case 2: Missing type field + let mut tools = vec![json!({ + "type": "function", + "function": { + "name": "test_func", + "description": "test description", + "parameters": { + "properties": {} + } + } + })]; + + validate_tool_schemas(&mut tools); + + let params = tools[0]["function"]["parameters"].as_object().unwrap(); + assert_eq!(params["type"], "object"); + + // Test case 3: Complete valid schema should remain unchanged + let original_schema = json!({ + "type": "function", + "function": { + "name": "test_func", + "description": "test description", + "parameters": { + "type": "object", + "properties": { + "location": { + "type": "string", + "description": "City and country" + } + }, + "required": ["location"] + } + } + }); + + let mut tools = vec![original_schema.clone()]; + validate_tool_schemas(&mut tools); + assert_eq!(tools[0], original_schema); + } + + const OPENAI_TOOL_USE_RESPONSE: &str = r#"{ + "choices": [{ + "role": "assistant", + "message": { + "tool_calls": [{ + "id": "1", + "function": { + "name": "example_fn", + "arguments": "{\"param\": \"value\"}" + } + }] + } + }], + "usage": { + "input_tokens": 10, + "output_tokens": 25, + "total_tokens": 35 + } + }"#; + + #[test] + fn test_format_messages() -> anyhow::Result<()> { + let message = Message::user().with_text("Hello"); + let spec = format_messages(&[message], &ImageFormat::OpenAi); + + assert_eq!(spec.len(), 1); + assert_eq!(spec[0]["role"], "user"); + assert_eq!(spec[0]["content"], "Hello"); + Ok(()) + } + + #[test] + fn test_format_tools() -> anyhow::Result<()> { + let tool = Tool::new( + "test_tool", + "A test tool", + json!({ + "type": "object", + "properties": { + "input": { + "type": "string", + "description": "Test parameter" + } + }, + "required": ["input"] + }), + ); + + let spec = format_tools(&[tool])?; + + assert_eq!(spec.len(), 1); + assert_eq!(spec[0]["type"], "function"); + assert_eq!(spec[0]["function"]["name"], "test_tool"); + Ok(()) + } + + #[test] + fn test_format_messages_complex() -> anyhow::Result<()> { + let mut messages = vec![ + Message::assistant().with_text("Hello!"), + Message::user().with_text("How are you?"), + Message::assistant().with_tool_request( + "tool1", + Ok(ToolCall::new("example", json!({"param1": "value1"}))), + ), + ]; + + // Get the ID from the tool request to use in the response + let tool_id = if let MessageContent::ToolRequest(request) = &messages[2].content[0] { + request.id.clone() + } else { + panic!("should be tool request"); + }; + + messages + .push(Message::user().with_tool_response(tool_id, Ok(vec![Content::text("Result")]))); + + let spec = format_messages(&messages, &ImageFormat::OpenAi); + + assert_eq!(spec.len(), 4); + assert_eq!(spec[0]["role"], "assistant"); + assert_eq!(spec[0]["content"], "Hello!"); + assert_eq!(spec[1]["role"], "user"); + assert_eq!(spec[1]["content"], "How are you?"); + assert_eq!(spec[2]["role"], "assistant"); + assert!(spec[2]["tool_calls"].is_array()); + assert_eq!(spec[3]["role"], "tool"); + assert_eq!(spec[3]["content"], "Result"); + assert_eq!(spec[3]["tool_call_id"], spec[2]["tool_calls"][0]["id"]); + + Ok(()) + } + + #[test] + 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"}))), + )]; + + // Get the ID from the tool request to use in the response + let tool_id = if let MessageContent::ToolRequest(request) = &messages[0].content[0] { + request.id.clone() + } else { + panic!("should be tool request"); + }; + + messages + .push(Message::user().with_tool_response(tool_id, Ok(vec![Content::text("Result")]))); + + 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"], "Result"); + assert_eq!(spec[1]["tool_call_id"], spec[0]["tool_calls"][0]["id"]); + + Ok(()) + } + + #[test] + fn test_format_tools_duplicate() -> anyhow::Result<()> { + let tool1 = Tool::new( + "test_tool", + "Test tool", + json!({ + "type": "object", + "properties": { + "input": { + "type": "string", + "description": "Test parameter" + } + }, + "required": ["input"] + }), + ); + + let tool2 = Tool::new( + "test_tool", + "Test tool", + json!({ + "type": "object", + "properties": { + "input": { + "type": "string", + "description": "Test parameter" + } + }, + "required": ["input"] + }), + ); + + let result = format_tools(&[tool1, tool2]); + assert!(result.is_err()); + assert!(result + .unwrap_err() + .to_string() + .contains("Duplicate tool name")); + + Ok(()) + } + + #[test] + fn test_format_tools_empty() -> anyhow::Result<()> { + let spec = format_tools(&[])?; + assert!(spec.is_empty()); + Ok(()) + } + + #[test] + fn test_format_messages_with_image_path() -> anyhow::Result<()> { + // Create a temporary PNG file with valid PNG magic numbers + let temp_dir = tempfile::tempdir()?; + let png_path = temp_dir.path().join("test.png"); + let png_data = [ + 0x89, 0x50, 0x4E, 0x47, // PNG magic number + 0x0D, 0x0A, 0x1A, 0x0A, // PNG header + 0x00, 0x00, 0x00, 0x0D, // Rest of fake PNG data + ]; + std::fs::write(&png_path, &png_data)?; + let png_path_str = png_path.to_str().unwrap(); + + // Create message with image path + let message = Message::user().with_text(format!("Here is an image: {}", png_path_str)); + let spec = format_messages(&[message], &ImageFormat::OpenAi); + + assert_eq!(spec.len(), 1); + assert_eq!(spec[0]["role"], "user"); + + // Content should be an array with text and image + let content = spec[0]["content"].as_array().unwrap(); + assert_eq!(content.len(), 2); + assert_eq!(content[0]["type"], "text"); + assert!(content[0]["text"].as_str().unwrap().contains(png_path_str)); + assert_eq!(content[1]["type"], "image_url"); + assert!(content[1]["image_url"]["url"] + .as_str() + .unwrap() + .starts_with("data:image/png;base64,")); + + Ok(()) + } + + #[test] + fn test_response_to_message_text() -> anyhow::Result<()> { + let response = json!({ + "choices": [{ + "role": "assistant", + "message": { + "content": "Hello from John Cena!" + } + }], + "usage": { + "input_tokens": 10, + "output_tokens": 25, + "total_tokens": 35 + } + }); + + 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!"); + } else { + panic!("Expected Text content"); + } + assert!(matches!(message.role, Role::Assistant)); + + Ok(()) + } + + #[test] + fn test_response_to_message_valid_toolrequest() -> anyhow::Result<()> { + let response: Value = serde_json::from_str(OPENAI_TOOL_USE_RESPONSE)?; + let message = response_to_message(response)?; + + assert_eq!(message.content.len(), 1); + if let MessageContent::ToolRequest(request) = &message.content[0] { + let tool_call = request.tool_call.as_ref().unwrap(); + assert_eq!(tool_call.name, "example_fn"); + assert_eq!(tool_call.arguments, json!({"param": "value"})); + } else { + panic!("Expected ToolRequest content"); + } + + Ok(()) + } + + #[test] + 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 = response_to_message(response)?; + + if let MessageContent::ToolRequest(request) = &message.content[0] { + match &request.tool_call { + Err(ToolError::NotFound(msg)) => { + assert!(msg.starts_with("The provided function name")); + } + _ => panic!("Expected ToolNotFound error"), + } + } else { + panic!("Expected ToolRequest content"); + } + + Ok(()) + } + + #[test] + 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 = response_to_message(response)?; + + if let MessageContent::ToolRequest(request) = &message.content[0] { + match &request.tool_call { + Err(ToolError::InvalidParameters(msg)) => { + assert!(msg.starts_with("Could not interpret tool use parameters")); + } + _ => panic!("Expected InvalidParameters error"), + } + } else { + panic!("Expected ToolRequest content"); + } + + Ok(()) + } + + #[test] + fn test_response_to_message_empty_argument() -> anyhow::Result<()> { + let mut response: Value = serde_json::from_str(OPENAI_TOOL_USE_RESPONSE)?; + response["choices"][0]["message"]["tool_calls"][0]["function"]["arguments"] = + serde_json::Value::String("".to_string()); + + let message = response_to_message(response)?; + + if let MessageContent::ToolRequest(request) = &message.content[0] { + let tool_call = request.tool_call.as_ref().unwrap(); + assert_eq!(tool_call.name, "example_fn"); + assert_eq!(tool_call.arguments, json!({})); + } else { + panic!("Expected ToolRequest content"); + } + + Ok(()) + } + + #[test] + fn test_create_request_gpt_4o() -> anyhow::Result<()> { + // Test default medium reasoning effort for O3 model + let model_config = ModelConfig { + model_name: "gpt-4o".to_string(), + context_limit: Some(4096), + temperature: None, + max_tokens: Some(1024), + }; + let request = create_request(&model_config, "system", &[], &[], &ImageFormat::OpenAi)?; + let obj = request.as_object().unwrap(); + let expected = json!({ + "model": "gpt-4o", + "messages": [ + { + "role": "system", + "content": "system" + } + ], + "max_tokens": 1024 + }); + + for (key, value) in expected.as_object().unwrap() { + assert_eq!(obj.get(key).unwrap(), value); + } + + Ok(()) + } + + #[test] + fn test_create_request_o1_default() -> anyhow::Result<()> { + // Test default medium reasoning effort for O1 model + let model_config = ModelConfig { + model_name: "o1".to_string(), + context_limit: Some(4096), + temperature: None, + max_tokens: Some(1024), + }; + let request = create_request(&model_config, "system", &[], &[], &ImageFormat::OpenAi)?; + let obj = request.as_object().unwrap(); + let expected = json!({ + "model": "o1", + "messages": [ + { + "role": "developer", + "content": "system" + } + ], + "reasoning_effort": "medium", + "max_completion_tokens": 1024 + }); + + for (key, value) in expected.as_object().unwrap() { + assert_eq!(obj.get(key).unwrap(), value); + } + + Ok(()) + } + + #[test] + fn test_create_request_o3_custom_reasoning_effort() -> anyhow::Result<()> { + // Test custom reasoning effort for O3 model + let model_config = ModelConfig { + model_name: "o3-mini-high".to_string(), + context_limit: Some(4096), + temperature: None, + max_tokens: Some(1024), + }; + let request = create_request(&model_config, "system", &[], &[], &ImageFormat::OpenAi)?; + let obj = request.as_object().unwrap(); + let expected = json!({ + "model": "o3-mini", + "messages": [ + { + "role": "developer", + "content": "system" + } + ], + "reasoning_effort": "high", + "max_completion_tokens": 1024 + }); + + for (key, value) in expected.as_object().unwrap() { + assert_eq!(obj.get(key).unwrap(), value); + } + + Ok(()) + } +} diff --git a/crates/goose-llm/src/providers/mod.rs b/crates/goose-llm/src/providers/mod.rs new file mode 100644 index 0000000000..c952d2d3b4 --- /dev/null +++ b/crates/goose-llm/src/providers/mod.rs @@ -0,0 +1,10 @@ +pub mod base; +pub mod databricks; +pub mod errors; +mod factory; +pub mod formats; +pub mod openai; +pub mod utils; + +pub use base::{Provider, ProviderCompleteResponse, Usage}; +pub use factory::create; diff --git a/crates/goose-llm/src/providers/openai.rs b/crates/goose-llm/src/providers/openai.rs new file mode 100644 index 0000000000..450dfc4675 --- /dev/null +++ b/crates/goose-llm/src/providers/openai.rs @@ -0,0 +1,160 @@ +use std::{collections::HashMap, time::Duration}; + +use anyhow::Result; +use async_trait::async_trait; +use reqwest::Client; +use serde_json::Value; + +use super::{ + errors::ProviderError, + formats::openai::{create_request, get_usage, response_to_message}, + utils::{emit_debug_trace, get_env, get_model, handle_response_openai_compat, ImageFormat}, +}; +use crate::{ + message::Message, + model::ModelConfig, + providers::{Provider, ProviderCompleteResponse, Usage}, + types::core::Tool, +}; + +pub const OPEN_AI_DEFAULT_MODEL: &str = "gpt-4o"; +pub const _OPEN_AI_KNOWN_MODELS: &[&str] = &[ + "gpt-4o", + "gpt-4o-mini", + "gpt-4-turbo", + "gpt-3.5-turbo", + "o1", + "o3", + "o4-mini", +]; + +#[derive(Debug)] +pub struct OpenAiProvider { + client: Client, + host: String, + base_path: String, + api_key: String, + organization: Option, + project: Option, + model: ModelConfig, + custom_headers: Option>, +} + +impl Default for OpenAiProvider { + fn default() -> Self { + let model = ModelConfig::new(OPEN_AI_DEFAULT_MODEL.to_string()); + OpenAiProvider::from_env(model).expect("Failed to initialize OpenAI provider") + } +} + +impl OpenAiProvider { + pub fn from_env(model: ModelConfig) -> Result { + let api_key: String = get_env("OPENAI_API_KEY")?; + let host: String = + get_env("OPENAI_HOST").unwrap_or_else(|_| "https://api.openai.com".to_string()); + let base_path: String = + get_env("OPENAI_BASE_PATH").unwrap_or_else(|_| "v1/chat/completions".to_string()); + let organization: Option = get_env("OPENAI_ORGANIZATION").ok(); + let project: Option = get_env("OPENAI_PROJECT").ok(); + let custom_headers: Option> = get_env("OPENAI_CUSTOM_HEADERS") + .or_else(|_| get_env("OPENAI_CUSTOM_HEADERS")) + .ok() + .map(parse_custom_headers); + // parse get_env("OPENAI_TIMEOUT") to u64 or set default to 600 + let timeout_secs = get_env("OPENAI_TIMEOUT") + .ok() + .and_then(|s| s.parse::().ok()) + .unwrap_or(600); + let client = Client::builder() + .timeout(Duration::from_secs(timeout_secs)) + .build()?; + + Ok(Self { + client, + host, + base_path, + api_key, + organization, + project, + model, + custom_headers, + }) + } + + async fn post(&self, payload: Value) -> Result { + let base_url = url::Url::parse(&self.host) + .map_err(|e| ProviderError::RequestFailed(format!("Invalid base URL: {e}")))?; + let url = base_url.join(&self.base_path).map_err(|e| { + ProviderError::RequestFailed(format!("Failed to construct endpoint URL: {e}")) + })?; + + let mut request = self + .client + .post(url) + .header("Authorization", format!("Bearer {}", self.api_key)); + + // Add organization header if present + if let Some(org) = &self.organization { + request = request.header("OpenAI-Organization", org); + } + + // Add project header if present + if let Some(project) = &self.project { + request = request.header("OpenAI-Project", project); + } + + if let Some(custom_headers) = &self.custom_headers { + for (key, value) in custom_headers { + request = request.header(key, value); + } + } + + let response = request.json(&payload).send().await?; + + handle_response_openai_compat(response).await + } +} + +#[async_trait] +impl Provider for OpenAiProvider { + #[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 { + let payload = create_request(&self.model, system, messages, tools, &ImageFormat::OpenAi)?; + + // Make request + let response = self.post(payload.clone()).await?; + + // Parse response + let message = response_to_message(response.clone())?; + let usage = match get_usage(&response) { + Ok(usage) => usage, + Err(ProviderError::UsageError(e)) => { + tracing::debug!("Failed to get usage data: {}", e); + Usage::default() + } + Err(e) => return Err(e), + }; + let model = get_model(&response); + emit_debug_trace(&self.model, &payload, &response, &usage); + Ok(ProviderCompleteResponse::new(message, model, usage)) + } +} + +fn parse_custom_headers(s: String) -> HashMap { + s.split(',') + .filter_map(|header| { + let mut parts = header.splitn(2, '='); + let key = parts.next().map(|s| s.trim().to_string())?; + let value = parts.next().map(|s| s.trim().to_string())?; + Some((key, value)) + }) + .collect() +} diff --git a/crates/goose-llm/src/providers/utils.rs b/crates/goose-llm/src/providers/utils.rs new file mode 100644 index 0000000000..3d6f177434 --- /dev/null +++ b/crates/goose-llm/src/providers/utils.rs @@ -0,0 +1,347 @@ +use std::{env, io::Read, path::Path}; + +use anyhow::Result; +use base64::Engine; +use regex::Regex; +use reqwest::{Response, StatusCode}; +use serde::{Deserialize, Serialize}; +use serde_json::{from_value, json, Value}; + +use super::base::Usage; +use crate::{ + model::ModelConfig, + providers::errors::{OpenAIError, ProviderError}, + types::core::ImageContent, +}; + +#[derive(serde::Deserialize)] +struct OpenAIErrorResponse { + error: OpenAIError, +} + +#[derive(Debug, Copy, Clone, Serialize, Deserialize)] +pub enum ImageFormat { + OpenAi, + Anthropic, +} + +/// Convert an image content into an image json based on format +pub fn convert_image(image: &ImageContent, image_format: &ImageFormat) -> Value { + match image_format { + ImageFormat::OpenAi => json!({ + "type": "image_url", + "image_url": { + "url": format!("data:{};base64,{}", image.mime_type, image.data) + } + }), + ImageFormat::Anthropic => json!({ + "type": "image", + "source": { + "type": "base64", + "media_type": image.mime_type, + "data": image.data, + } + }), + } +} + +/// Handle response from OpenAI compatible endpoints +/// Error codes: https://platform.openai.com/docs/guides/error-codes +/// Context window exceeded: https://community.openai.com/t/help-needed-tackling-context-length-limits-in-openai-models/617543 +pub async fn handle_response_openai_compat(response: Response) -> Result { + let status = response.status(); + // Try to parse the response body as JSON (if applicable) + let payload = match response.json::().await { + Ok(json) => json, + Err(e) => return Err(ProviderError::RequestFailed(e.to_string())), + }; + + match status { + StatusCode::OK => Ok(payload), + StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN => { + Err(ProviderError::Authentication(format!( + "Authentication failed. Please ensure your API keys are valid and have the required permissions. \ + Status: {}. Response: {:?}", + status, payload + ))) + } + StatusCode::BAD_REQUEST | StatusCode::NOT_FOUND => { + tracing::debug!( + "{}", + format!( + "Provider request failed with status: {}. Payload: {:?}", + status, payload + ) + ); + if let Ok(err_resp) = from_value::(payload) { + let err = err_resp.error; + if err.is_context_length_exceeded() { + return Err(ProviderError::ContextLengthExceeded( + err.message.unwrap_or("Unknown error".to_string()), + )); + } + return Err(ProviderError::RequestFailed(format!( + "{} (status {})", + err, + status.as_u16() + ))); + } + Err(ProviderError::RequestFailed(format!( + "Unknown error (status {})", + status + ))) + } + StatusCode::TOO_MANY_REQUESTS => { + Err(ProviderError::RateLimitExceeded(format!("{:?}", payload))) + } + StatusCode::INTERNAL_SERVER_ERROR | StatusCode::SERVICE_UNAVAILABLE => { + Err(ProviderError::ServerError(format!("{:?}", payload))) + } + _ => { + tracing::debug!( + "{}", + format!( + "Provider request failed with status: {}. Payload: {:?}", + status, payload + ) + ); + Err(ProviderError::RequestFailed(format!( + "Request failed with status: {}", + status + ))) + } + } +} + +/// Get a secret from environment variables. The secret is expected to be in JSON format. +pub fn get_env(key: &str) -> Result { + // check environment variables (convert to uppercase) + let env_key = key.to_uppercase(); + if let Ok(val) = env::var(&env_key) { + let value: Value = serde_json::from_str(&val).unwrap_or(Value::String(val)); + Ok(serde_json::from_value(value)?) + } else { + Err(anyhow::anyhow!( + "Environment variable {} not found", + env_key + )) + } +} + +pub fn sanitize_function_name(name: &str) -> String { + let re = Regex::new(r"[^a-zA-Z0-9_-]").unwrap(); + re.replace_all(name, "_").to_string() +} + +pub fn is_valid_function_name(name: &str) -> bool { + let re = Regex::new(r"^[a-zA-Z0-9_-]+$").unwrap(); + re.is_match(name) +} + +/// Extract the model name from a JSON object. Common with most providers to have this top level attribute. +pub fn get_model(data: &Value) -> String { + if let Some(model) = data.get("model") { + if let Some(model_str) = model.as_str() { + model_str.to_string() + } else { + "Unknown".to_string() + } + } else { + "Unknown".to_string() + } +} + +/// Check if a file is actually an image by examining its magic bytes +fn is_image_file(path: &Path) -> bool { + if let Ok(mut file) = std::fs::File::open(path) { + let mut buffer = [0u8; 8]; // Large enough for most image magic numbers + if file.read(&mut buffer).is_ok() { + // Check magic numbers for common image formats + return match &buffer[0..4] { + // PNG: 89 50 4E 47 + [0x89, 0x50, 0x4E, 0x47] => true, + // JPEG: FF D8 FF + [0xFF, 0xD8, 0xFF, _] => true, + _ => false, + }; + } + } + false +} + +/// Detect if a string contains a path to an image file +pub fn detect_image_path(text: &str) -> Option<&str> { + // Basic image file extension check + let extensions = [".png", ".jpg", ".jpeg"]; + + // Find any word that ends with an image extension + for word in text.split_whitespace() { + if extensions + .iter() + .any(|ext| word.to_lowercase().ends_with(ext)) + { + let path = Path::new(word); + // Check if it's an absolute path and file exists + if path.is_absolute() && path.is_file() { + // Verify it's actually an image file + if is_image_file(path) { + return Some(word); + } + } + } + } + None +} + +/// Convert a local image file to base64 encoded ImageContent +pub fn load_image_file(path: &str) -> Result { + let path = Path::new(path); + + // Verify it's an image before proceeding + if !is_image_file(path) { + return Err(ProviderError::RequestFailed( + "File is not a valid image".to_string(), + )); + } + + // Read the file + let bytes = std::fs::read(path) + .map_err(|e| ProviderError::RequestFailed(format!("Failed to read image file: {}", e)))?; + + // Detect mime type from extension + let mime_type = match path.extension().and_then(|e| e.to_str()) { + Some(ext) => match ext.to_lowercase().as_str() { + "png" => "image/png", + "jpg" | "jpeg" => "image/jpeg", + _ => { + return Err(ProviderError::RequestFailed( + "Unsupported image format".to_string(), + )); + } + }, + None => { + return Err(ProviderError::RequestFailed( + "Unknown image format".to_string(), + )); + } + }; + + // Convert to base64 + let data = base64::prelude::BASE64_STANDARD.encode(&bytes); + + Ok(ImageContent { + mime_type: mime_type.to_string(), + data, + }) +} + +pub fn emit_debug_trace( + model_config: &ModelConfig, + payload: &Value, + response: &Value, + usage: &Usage, +) { + tracing::debug!( + model_config = %serde_json::to_string_pretty(model_config).unwrap_or_default(), + input = %serde_json::to_string_pretty(payload).unwrap_or_default(), + output = %serde_json::to_string_pretty(response).unwrap_or_default(), + input_tokens = ?usage.input_tokens.unwrap_or_default(), + output_tokens = ?usage.output_tokens.unwrap_or_default(), + total_tokens = ?usage.total_tokens.unwrap_or_default(), + ); +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_detect_image_path() { + // Create a temporary PNG file with valid PNG magic numbers + let temp_dir = tempfile::tempdir().unwrap(); + let png_path = temp_dir.path().join("test.png"); + let png_data = [ + 0x89, 0x50, 0x4E, 0x47, // PNG magic number + 0x0D, 0x0A, 0x1A, 0x0A, // PNG header + 0x00, 0x00, 0x00, 0x0D, // Rest of fake PNG data + ]; + std::fs::write(&png_path, &png_data).unwrap(); + let png_path_str = png_path.to_str().unwrap(); + + // Create a fake PNG (wrong magic numbers) + let fake_png_path = temp_dir.path().join("fake.png"); + std::fs::write(&fake_png_path, b"not a real png").unwrap(); + + // Test with valid PNG file using absolute path + let text = format!("Here is an image {}", png_path_str); + assert_eq!(detect_image_path(&text), Some(png_path_str)); + + // Test with non-image file that has .png extension + let text = format!("Here is a fake image {}", fake_png_path.to_str().unwrap()); + assert_eq!(detect_image_path(&text), None); + + // Test with non-existent file + let text = "Here is a fake.png that doesn't exist"; + assert_eq!(detect_image_path(text), None); + + // Test with non-image file + let text = "Here is a file.txt"; + assert_eq!(detect_image_path(text), None); + + // Test with relative path (should not match) + let text = "Here is a relative/path/image.png"; + assert_eq!(detect_image_path(text), None); + } + + #[test] + fn test_load_image_file() { + // Create a temporary PNG file with valid PNG magic numbers + let temp_dir = tempfile::tempdir().unwrap(); + let png_path = temp_dir.path().join("test.png"); + let png_data = [ + 0x89, 0x50, 0x4E, 0x47, // PNG magic number + 0x0D, 0x0A, 0x1A, 0x0A, // PNG header + 0x00, 0x00, 0x00, 0x0D, // Rest of fake PNG data + ]; + std::fs::write(&png_path, &png_data).unwrap(); + let png_path_str = png_path.to_str().unwrap(); + + // Create a fake PNG (wrong magic numbers) + let fake_png_path = temp_dir.path().join("fake.png"); + std::fs::write(&fake_png_path, b"not a real png").unwrap(); + let fake_png_path_str = fake_png_path.to_str().unwrap(); + + // Test loading valid PNG file + let result = load_image_file(png_path_str); + assert!(result.is_ok()); + let image = result.unwrap(); + assert_eq!(image.mime_type, "image/png"); + + // Test loading fake PNG file + let result = load_image_file(fake_png_path_str); + assert!(result.is_err()); + assert!(result + .unwrap_err() + .to_string() + .contains("not a valid image")); + + // Test non-existent file + let result = load_image_file("nonexistent.png"); + assert!(result.is_err()); + } + + #[test] + fn test_sanitize_function_name() { + assert_eq!(sanitize_function_name("hello-world"), "hello-world"); + assert_eq!(sanitize_function_name("hello world"), "hello_world"); + assert_eq!(sanitize_function_name("hello@world"), "hello_world"); + } + + #[test] + fn test_is_valid_function_name() { + assert!(is_valid_function_name("hello-world")); + assert!(is_valid_function_name("hello_world")); + assert!(!is_valid_function_name("hello world")); + assert!(!is_valid_function_name("hello@world")); + } +} diff --git a/crates/goose-llm/src/types.rs b/crates/goose-llm/src/types.rs deleted file mode 100644 index cdc52232b2..0000000000 --- a/crates/goose-llm/src/types.rs +++ /dev/null @@ -1,70 +0,0 @@ -use goose::message::Message; -use goose::providers::base::ProviderUsage; -use mcp_core::tool::Tool; -use serde::{Deserialize, Serialize}; - -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct CompletionResponse { - message: Message, - usage: ProviderUsage, - runtime_metrics: RuntimeMetrics, -} - -impl CompletionResponse { - pub fn new(message: Message, usage: ProviderUsage, runtime_metrics: RuntimeMetrics) -> Self { - Self { - message, - usage, - runtime_metrics, - } - } -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct RuntimeMetrics { - pub total_time_ms: u128, - pub total_time_ms_provider: u128, - pub tokens_per_second: Option, -} - -impl RuntimeMetrics { - pub fn new( - total_time_ms: u128, - total_time_ms_provider: u128, - tokens_per_second: Option, - ) -> Self { - Self { - total_time_ms, - total_time_ms_provider, - tokens_per_second, - } - } -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct Extension { - name: String, - instructions: Option, - tools: Vec, -} - -impl Extension { - pub fn new(name: String, instructions: Option, tools: Vec) -> Self { - Self { - name, - instructions, - tools, - } - } - - pub fn get_prefixed_tools(&self) -> Vec { - self.tools - .iter() - .map(|tool| { - let mut prefixed_tool = tool.clone(); - prefixed_tool.name = format!("{}__{}", self.name, tool.name); - prefixed_tool - }) - .collect() - } -} diff --git a/crates/goose-llm/src/types/completion.rs b/crates/goose-llm/src/types/completion.rs new file mode 100644 index 0000000000..4f9e33daa2 --- /dev/null +++ b/crates/goose-llm/src/types/completion.rs @@ -0,0 +1,134 @@ +// This file defines types for completion interfaces, including the request and response structures. +// Many of these are adapted based on the Goose Service API: +// https://docs.google.com/document/d/1r5vjSK3nBQU1cIRf0WKysDigqMlzzrzl_bxEE4msOiw/edit?tab=t.0 + +use std::collections::HashMap; + +use serde::{Deserialize, Serialize}; + +use crate::{message::Message, providers::Usage}; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct CompletionResponse { + pub message: Message, + pub model: String, + pub usage: Usage, + pub runtime_metrics: RuntimeMetrics, +} + +impl CompletionResponse { + pub fn new( + message: Message, + model: String, + usage: Usage, + runtime_metrics: RuntimeMetrics, + ) -> Self { + Self { + message, + model, + usage, + runtime_metrics, + } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct RuntimeMetrics { + pub total_time_ms: u128, + pub total_time_ms_provider: u128, + pub tokens_per_second: Option, +} + +impl RuntimeMetrics { + pub fn new( + total_time_ms: u128, + total_time_ms_provider: u128, + tokens_per_second: Option, + ) -> Self { + Self { + total_time_ms, + total_time_ms_provider, + tokens_per_second, + } + } +} + +#[derive(Debug, Clone, PartialEq, Serialize)] +pub enum ToolApprovalMode { + Auto, + Manual, + Smart, +} + +#[derive(Debug, Clone, Serialize)] +pub struct ToolConfig { + pub name: String, + pub description: String, + pub input_schema: serde_json::Value, + pub approval_mode: ToolApprovalMode, +} + +impl ToolConfig { + pub fn new( + name: &str, + description: &str, + input_schema: serde_json::Value, + approval_mode: ToolApprovalMode, + ) -> Self { + Self { + name: name.to_string(), + description: description.to_string(), + input_schema, + approval_mode, + } + } + + /// Convert the tool config to a core tool + pub fn to_core_tool(&self, name: Option<&str>) -> super::core::Tool { + let tool_name = name.unwrap_or(&self.name); + super::core::Tool::new( + tool_name, + self.description.clone(), + self.input_schema.clone(), + ) + } +} + +#[derive(Debug, Clone, Serialize)] +pub struct ExtensionConfig { + name: String, + instructions: Option, + tools: Vec, +} + +impl ExtensionConfig { + pub fn new(name: String, instructions: Option, tools: Vec) -> Self { + Self { + name, + instructions, + tools, + } + } + + /// Convert the tools to core tools with the extension name as a prefix + pub fn get_prefixed_tools(&self) -> Vec { + self.tools + .iter() + .map(|tool| { + let name = format!("{}__{}", self.name, tool.name); + tool.to_core_tool(Some(&name)) + }) + .collect() + } + + /// Get a map of prefixed tool names to their approval modes + pub fn get_prefixed_tool_configs(&self) -> HashMap { + self.tools + .iter() + .map(|tool| { + let name = format!("{}__{}", self.name, tool.name); + (name, tool.clone()) + }) + .collect() + } +} diff --git a/crates/goose-llm/src/types/core.rs b/crates/goose-llm/src/types/core.rs new file mode 100644 index 0000000000..2555c6bb5d --- /dev/null +++ b/crates/goose-llm/src/types/core.rs @@ -0,0 +1,131 @@ +// This file defines core types that require serialization to +// construct payloads for LLM model providers and work with MCPs. + +use serde::{Deserialize, Serialize}; +use thiserror::Error; + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum Role { + User, + Assistant, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(tag = "type", rename_all = "camelCase")] +pub enum Content { + Text(TextContent), + Image(ImageContent), +} + +impl Content { + pub fn text>(text: S) -> Self { + Content::Text(TextContent { text: text.into() }) + } + + pub fn image, T: Into>(data: S, mime_type: T) -> Self { + Content::Image(ImageContent { + data: data.into(), + mime_type: mime_type.into(), + }) + } + + /// Get the text content if this is a TextContent variant + pub fn as_text(&self) -> Option<&str> { + match self { + Content::Text(text) => Some(&text.text), + _ => None, + } + } + + /// Get the image content if this is an ImageContent variant + pub fn as_image(&self) -> Option<(&str, &str)> { + match self { + Content::Image(image) => Some((&image.data, &image.mime_type)), + _ => None, + } + } +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct TextContent { + pub text: String, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ImageContent { + pub data: String, + pub mime_type: String, +} + +/// A tool that can be used by a model. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct Tool { + /// The name of the tool + pub name: String, + /// A description of what the tool does + pub description: String, + /// A JSON Schema object defining the expected parameters for the tool + pub input_schema: serde_json::Value, +} + +impl Tool { + /// Create a new tool with the given name and description + pub fn new(name: N, description: D, input_schema: serde_json::Value) -> Self + where + N: Into, + D: Into, + { + Tool { + name: name.into(), + description: description.into(), + input_schema, + } + } +} + +/// A tool call request that an extension can execute +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ToolCall { + /// The name of the tool to execute + pub name: String, + /// The parameters for the execution + pub arguments: serde_json::Value, + /// Whether the tool call needs approval before execution. Default is false. + pub needs_approval: bool, +} + +impl ToolCall { + /// Create a new ToolUse with the given name and parameters + pub fn new>(name: S, arguments: serde_json::Value) -> Self { + Self { + name: name.into(), + arguments, + needs_approval: false, + } + } + + /// Set needs_approval field + pub fn set_needs_approval(&mut self, flag: bool) { + self.needs_approval = flag; + } +} + +#[non_exhaustive] +#[derive(Error, Debug, Clone, Deserialize, Serialize, PartialEq)] +pub enum ToolError { + #[error("Invalid parameters: {0}")] + InvalidParameters(String), + #[error("Execution failed: {0}")] + ExecutionError(String), + #[error("Schema error: {0}")] + SchemaError(String), + #[error("Tool not found: {0}")] + NotFound(String), +} + +pub type ToolResult = std::result::Result; diff --git a/crates/goose-llm/src/types/mod.rs b/crates/goose-llm/src/types/mod.rs new file mode 100644 index 0000000000..5e2ededc6d --- /dev/null +++ b/crates/goose-llm/src/types/mod.rs @@ -0,0 +1,2 @@ +pub mod completion; +pub mod core;