From dc3be99aa534b5cee7390c021de8a36da6677739 Mon Sep 17 00:00:00 2001 From: Salman Mohammed Date: Sat, 9 Nov 2024 17:07:41 -0500 Subject: [PATCH] Add providers base and implement openai (#229) --- .cursorrules | 27 ++ .goosehints | 27 ++ Cargo.toml | 6 + profiles.yaml | 13 + providers.yaml | 16 + src/lib.rs | 2 +- src/main.rs | 98 +++--- src/providers.rs | 6 + src/providers/base.rs | 79 +++++ src/providers/configs.rs | 3 + src/providers/configs/base.rs | 22 ++ src/providers/configs/databricks.rs | 27 ++ src/providers/configs/openai.rs | 29 ++ src/providers/databricks.rs | 1 + src/providers/openai.rs | 413 +++++++++++++++++++++++++ src/providers/types.rs | 4 + src/providers/types/content.rs | 61 ++++ src/providers/types/message.rs | 294 ++++++++++++++++++ src/providers/types/objectid.rs | 30 ++ src/providers/types/tool.rs | 140 +++++++++ src/providers/utils.rs | 447 ++++++++++++++++++++++++++++ 21 files changed, 1706 insertions(+), 39 deletions(-) create mode 100644 .cursorrules create mode 100644 .goosehints create mode 100644 profiles.yaml create mode 100644 providers.yaml create mode 100644 src/providers.rs create mode 100644 src/providers/base.rs create mode 100644 src/providers/configs.rs create mode 100644 src/providers/configs/base.rs create mode 100644 src/providers/configs/databricks.rs create mode 100644 src/providers/configs/openai.rs create mode 100644 src/providers/databricks.rs create mode 100644 src/providers/openai.rs create mode 100644 src/providers/types.rs create mode 100644 src/providers/types/content.rs create mode 100644 src/providers/types/message.rs create mode 100644 src/providers/types/objectid.rs create mode 100644 src/providers/types/tool.rs create mode 100644 src/providers/utils.rs diff --git a/.cursorrules b/.cursorrules new file mode 100644 index 0000000000..4eaa699918 --- /dev/null +++ b/.cursorrules @@ -0,0 +1,27 @@ +You are an expert programmer in Rust teaching who is teaching another developer who is learning Rust. +The students are familiar with programming in languages such as Python (advanced), Java (novice) and C (novice) so +when possible use analogies from those languages. + +Key Principles +- Write clear, concise, and idiomatic Rust code with accurate examples. +- Use async programming paradigms effectively, leveraging `tokio` for concurrency. +- Prioritize modularity, clean code organization, and efficient resource management. +- Use expressive variable names that convey intent (e.g., `is_ready`, `has_data`). +- Adhere to Rust's naming conventions: snake_case for variables and functions, PascalCase for types and structs. +- Avoid code duplication; use functions and modules to encapsulate reusable logic. +- Write code with safety, concurrency, and performance in mind, embracing Rust's ownership and type system. + +Error Handling and Safety +- Embrace Rust's Result and Option types for error handling. +- Use `?` operator to propagate errors in async functions. +- Implement custom error types using `thiserror` or `anyhow` for more descriptive errors. +- Handle errors and edge cases early, returning errors where appropriate. +- Use `.await` responsibly, ensuring safe points for context switching. + +Key Conventions +1. Structure the application into modules: separate concerns like networking, database, and business logic. +2. Use environment variables for configuration management (e.g., `dotenv` crate). +3. Ensure code is well-documented with inline comments and Rustdoc. +4. Do not use the older style of "MOD/mod.rs" for separing modules and instead use the "MOD.rs" filename convention. + +Refer to "The Rust Programming Language" book (2024 version) and "Command line apps in Rust" documentation for in-depth information on best practices, and advanced features. diff --git a/.goosehints b/.goosehints new file mode 100644 index 0000000000..4eaa699918 --- /dev/null +++ b/.goosehints @@ -0,0 +1,27 @@ +You are an expert programmer in Rust teaching who is teaching another developer who is learning Rust. +The students are familiar with programming in languages such as Python (advanced), Java (novice) and C (novice) so +when possible use analogies from those languages. + +Key Principles +- Write clear, concise, and idiomatic Rust code with accurate examples. +- Use async programming paradigms effectively, leveraging `tokio` for concurrency. +- Prioritize modularity, clean code organization, and efficient resource management. +- Use expressive variable names that convey intent (e.g., `is_ready`, `has_data`). +- Adhere to Rust's naming conventions: snake_case for variables and functions, PascalCase for types and structs. +- Avoid code duplication; use functions and modules to encapsulate reusable logic. +- Write code with safety, concurrency, and performance in mind, embracing Rust's ownership and type system. + +Error Handling and Safety +- Embrace Rust's Result and Option types for error handling. +- Use `?` operator to propagate errors in async functions. +- Implement custom error types using `thiserror` or `anyhow` for more descriptive errors. +- Handle errors and edge cases early, returning errors where appropriate. +- Use `.await` responsibly, ensuring safe points for context switching. + +Key Conventions +1. Structure the application into modules: separate concerns like networking, database, and business logic. +2. Use environment variables for configuration management (e.g., `dotenv` crate). +3. Ensure code is well-documented with inline comments and Rustdoc. +4. Do not use the older style of "MOD/mod.rs" for separing modules and instead use the "MOD.rs" filename convention. + +Refer to "The Rust Programming Language" book (2024 version) and "Command line apps in Rust" documentation for in-depth information on best practices, and advanced features. diff --git a/Cargo.toml b/Cargo.toml index 1ea0179941..cba17eeb35 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -6,6 +6,7 @@ edition = "2021" [dependencies] # Error handling anyhow = "1.0" +thiserror = "1.0" # Requests reqwest = { version = "0.11", features = ["json"] } @@ -20,3 +21,8 @@ clap = { version = "4.4", features = ["derive"] } cliclack = "0.3.5" console = "0.15.8" bat = "0.24.0" +uuid = { version = "1.0", features = ["v4"] } +regex = "1.11.1" + +[dev-dependencies] +mockito = "1.2" diff --git a/profiles.yaml b/profiles.yaml new file mode 100644 index 0000000000..9b1783af50 --- /dev/null +++ b/profiles.yaml @@ -0,0 +1,13 @@ +default: + processor: + provider: databricks + model: claude-3-5-sonnet-2 + accelerator: + provider: openai / ollama + model: gpt-4o-mini / qwen2.5 + systems: + - name: synopsis + type: embedded + - name: qai + type: http + host: localhost:8001 diff --git a/providers.yaml b/providers.yaml new file mode 100644 index 0000000000..b96e58ed8d --- /dev/null +++ b/providers.yaml @@ -0,0 +1,16 @@ +providers: + - name: databricks + config: + host: databricks_host + token: databricks_token + - name: openai + config: + api_key: openai_api_key + - name: ollama + config: + host: "http://localhost:11434/" + - name: bedrock + config: + profile: aws_profile + access_key: key + accesss_secret: secret diff --git a/src/lib.rs b/src/lib.rs index 8b13789179..9ce92fb4d1 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1 +1 @@ - +pub mod providers; diff --git a/src/main.rs b/src/main.rs index 154c13fcb8..17922d04cb 100644 --- a/src/main.rs +++ b/src/main.rs @@ -3,9 +3,15 @@ use bat::PrettyPrinter; use clap::Parser; use cliclack::{input, spinner}; use console::style; -use reqwest::Client; -use serde_json::{json, Value}; use std::env; +use serde_json::json; + +use goose::providers::configs::openai::OpenAiProviderConfig; +use goose::providers::base::Provider; +use goose::providers::openai::OpenAiProvider; +use goose::providers::types::message::Message; +use goose::providers::types::tool::Tool; +use goose::providers::types::content::Content; #[derive(Parser)] #[command(author, version, about, long_about = None)] @@ -29,7 +35,38 @@ async fn main() -> Result<()> { .or_else(|| env::var("OPENAI_API_KEY").ok()) .context("API key must be provided via --api-key or OPENAI_API_KEY environment variable")?; - let client = Client::new(); + // Initialize OpenAI provider + let provider = OpenAiProvider::new(OpenAiProviderConfig { + api_key, + host: "https://api.openai.com".to_string(), + })?; + + // Add word counting tool + let parameters = json!({ + "type": "object", + "properties": { + "text": { + "type": "string", + "description": "The text to count words in" + } + }, + "required": ["text"] + }); + + let word_count_tool = Tool::new( + "count_words".to_string(), + "Count the number of words in text".to_string(), + parameters, + |args| { + let text = args.get("text") + .and_then(|v| v.as_str()) + .unwrap_or_default(); + Ok(json!({ "count": text.split_whitespace().count() })) + } + ); + + let tools = vec![word_count_tool]; + println!( "Example goose CLI {}", style("- type \"exit\" to end the session").dim() @@ -37,25 +74,36 @@ async fn main() -> Result<()> { println!("\n"); loop { - let message: String = input("Message:").placeholder("").multiline().interact()?; + let message_text: String = input("Message:").placeholder("").multiline().interact()?; - if message.trim().eq_ignore_ascii_case("exit") { + if message_text.trim().eq_ignore_ascii_case("exit") { break; } - let future = send_chat_request(&client, &api_key, &cli.model, &message); let spin = spinner(); spin.start("awaiting reply"); - let response = future.await?; + // Create user message and get completion + let user_message = Message::user(&message_text)?; + let (response_message, _usage) = provider.complete( + &cli.model, + "You are a helpful assistant.", + &[user_message], + &tools, // Changed from &[] to &tools + None, // default temperature + None, // default max_tokens + None, // no stop sequences + None, // default top_p + ).await?; // Added .await since complete returns a Future + spin.stop(""); - // Extract the content from the response - let content = response["choices"][0]["message"]["content"] - .as_str() - .context("Failed to get response content")?; + if response_message.has_tool_use() { + render(&Content::ToolUse(response_message.tool_use().first().unwrap().clone()).summary()).await; + } else { + render(&response_message.text()).await; + } - render(content).await; println!("\n"); } Ok(()) @@ -68,29 +116,3 @@ async fn render(content: &str) { .print() .unwrap(); } - -async fn send_chat_request( - client: &Client, - api_key: &str, - model: &str, - message: &str, -) -> Result { - let response = client - .post("https://api.openai.com/v1/chat/completions") - .header("Authorization", format!("Bearer {}", api_key)) - .json(&json!({ - "model": model, - "messages": [ - { - "role": "user", - "content": message - } - ] - })) - .send() - .await? - .json::() - .await?; - - Ok(response) -} diff --git a/src/providers.rs b/src/providers.rs new file mode 100644 index 0000000000..72e2ad0cad --- /dev/null +++ b/src/providers.rs @@ -0,0 +1,6 @@ +pub mod base; +pub mod configs; +pub mod databricks; +pub mod openai; +pub mod types; +pub mod utils; diff --git a/src/providers/base.rs b/src/providers/base.rs new file mode 100644 index 0000000000..31d06cf2e8 --- /dev/null +++ b/src/providers/base.rs @@ -0,0 +1,79 @@ +use anyhow::Result; +use serde::{Deserialize, Serialize}; + +use super::types::{message::Message, tool::Tool}; + +#[derive(Debug, Clone, 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, + } + } +} + +/// Base trait for AI providers (OpenAI, Anthropic, etc) +pub trait Provider: Send + Sync { + /// Create a provider instance from environment variables + fn from_env() -> Result + where + Self: Sized; + + /// Generate the next message using the specified model and other parameters + fn complete( + &self, + model: &str, + system: &str, + messages: &[Message], + tools: &[Tool], + temperature: Option, + max_tokens: Option, + stop_sequences: Option<&[String]>, + top_p: Option, + ) -> impl std::future::Future> + Send; +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[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_usage_serialization() -> Result<()> { + let usage = Usage::new(Some(10), Some(20), Some(30)); + let serialized = serde_json::to_string(&usage)?; + let deserialized: Usage = serde_json::from_str(&serialized)?; + + assert_eq!(usage.input_tokens, deserialized.input_tokens); + assert_eq!(usage.output_tokens, deserialized.output_tokens); + assert_eq!(usage.total_tokens, deserialized.total_tokens); + + // Test JSON structure + let json_value: serde_json::Value = serde_json::from_str(&serialized)?; + assert_eq!(json_value["input_tokens"], json!(10)); + assert_eq!(json_value["output_tokens"], json!(20)); + assert_eq!(json_value["total_tokens"], json!(30)); + + Ok(()) + } +} diff --git a/src/providers/configs.rs b/src/providers/configs.rs new file mode 100644 index 0000000000..2aecf46281 --- /dev/null +++ b/src/providers/configs.rs @@ -0,0 +1,3 @@ +pub mod base; +pub mod databricks; +pub mod openai; diff --git a/src/providers/configs/base.rs b/src/providers/configs/base.rs new file mode 100644 index 0000000000..59067f08e6 --- /dev/null +++ b/src/providers/configs/base.rs @@ -0,0 +1,22 @@ +use anyhow::Result; +use std::env; + +pub trait ProviderConfig { + /// Load configuration from environment variables + fn from_env() -> Result + where + Self: Sized; + + /// Helper function to get environment variables with error handling + fn get_env(key: &str, required: bool, default: Option) -> Result> { + match env::var(key) { + Ok(value) => Ok(Some(value)), + Err(env::VarError::NotPresent) if !required => Ok(default), + Err(env::VarError::NotPresent) => Err(anyhow::anyhow!( + "Environment variable '{}' is required but not set.", + key + )), + Err(e) => Err(e.into()), + } + } +} diff --git a/src/providers/configs/databricks.rs b/src/providers/configs/databricks.rs new file mode 100644 index 0000000000..93ca96c2a3 --- /dev/null +++ b/src/providers/configs/databricks.rs @@ -0,0 +1,27 @@ +use super::base::ProviderConfig; +use anyhow::Result; + +pub struct DatabricksProviderConfig { + pub host: String, + pub token: String, +} + +impl DatabricksProviderConfig { + pub fn new(host: String, token: String) -> Self { + Self { host, token } + } +} + +impl ProviderConfig for DatabricksProviderConfig { + fn from_env() -> Result { + // Get required host + let host = Self::get_env("DATABRICKS_HOST", true, None)? + .ok_or_else(|| anyhow::anyhow!("Databricks host should be present"))?; + + // Get required token + let token = Self::get_env("DATABRICKS_TOKEN", true, None)? + .ok_or_else(|| anyhow::anyhow!("Databricks token should be present"))?; + + Ok(Self::new(host, token)) + } +} diff --git a/src/providers/configs/openai.rs b/src/providers/configs/openai.rs new file mode 100644 index 0000000000..1da3310ff2 --- /dev/null +++ b/src/providers/configs/openai.rs @@ -0,0 +1,29 @@ +use super::base::ProviderConfig; +use anyhow::Result; + +pub struct OpenAiProviderConfig { + pub api_key: String, + pub host: String, +} + +impl OpenAiProviderConfig { + pub fn new(api_key: String, host: String) -> Self { + Self { api_key, host } + } +} + +impl ProviderConfig for OpenAiProviderConfig { + fn from_env() -> Result { + let api_key = Self::get_env("OPENAI_API_KEY", true, None)? + .ok_or_else(|| anyhow::anyhow!("OpenAI API key should be present"))?; + + let host = Self::get_env( + "OPENAI_API_HOST", + false, + Some("https://api.openai.com/".to_string()), + )? + .unwrap_or_else(|| "https://api.openai.com/".to_string()); + + Ok(Self::new(api_key, host)) + } +} diff --git a/src/providers/databricks.rs b/src/providers/databricks.rs new file mode 100644 index 0000000000..8b13789179 --- /dev/null +++ b/src/providers/databricks.rs @@ -0,0 +1 @@ + diff --git a/src/providers/openai.rs b/src/providers/openai.rs new file mode 100644 index 0000000000..849d9b5703 --- /dev/null +++ b/src/providers/openai.rs @@ -0,0 +1,413 @@ +use anyhow::{anyhow, Result}; +use reqwest::Client; // we are using blocking API here to make sync calls +use reqwest::StatusCode; +use serde_json::{json, Value}; +use std::time::Duration; + +use super::{ + base::{Provider, Usage}, + configs::base::ProviderConfig, + configs::openai::OpenAiProviderConfig, + types::{message::Message, tool::Tool}, + utils::{ + check_openai_context_length_error, messages_to_openai_spec, openai_response_to_message, + tools_to_openai_spec, + }, +}; + +pub struct OpenAiProvider { + client: Client, + config: OpenAiProviderConfig, +} + +impl OpenAiProvider { + pub fn new(config: OpenAiProviderConfig) -> Result { + let client = Client::builder() + .timeout(Duration::from_secs(600)) // 10 minutes timeout + .build()?; + + Ok(Self { client, config }) + } + + fn get_usage(data: &Value) -> Result { + let usage = data + .get("usage") + .ok_or_else(|| anyhow!("No usage data in response"))?; + + 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)) + } + + async fn post(&self, payload: Value) -> Result { + let url = format!( + "{}/v1/chat/completions", + self.config.host.trim_end_matches('/') + ); + + let response = self + .client + .post(&url) + .header("Authorization", format!("Bearer {}", self.config.api_key)) + .json(&payload) + .send() + .await?; + + match response.status() { + StatusCode::OK => Ok(response.json().await?), + status if status == StatusCode::TOO_MANY_REQUESTS || status.as_u16() >= 500 => { + // Implement retry logic here if needed + Err(anyhow!("Server error: {}", status)) + } + _ => Err(anyhow!("Request failed: {}", response.status())), + } + } +} + +impl Provider for OpenAiProvider { + fn from_env() -> Result { + let config = OpenAiProviderConfig::from_env()?; + Self::new(config) + } + + async fn complete( + &self, + model: &str, + system: &str, + messages: &[Message], + tools: &[Tool], + temperature: Option, + max_tokens: Option, + stop_sequences: Option<&[String]>, + top_p: Option, + ) -> Result<(Message, Usage)> { + // Not checking for o1 model here since system message is not supported by o1 + let system_message = json!({ + "role": "system", + "content": system + }); + + // Convert messages and tools to OpenAI format + let messages_spec = messages_to_openai_spec(messages); + let tools_spec = if !tools.is_empty() { + tools_to_openai_spec(tools)? + } else { + vec![] + }; + + // Build payload + // create messages array with system message first + let mut messages_array = vec![system_message]; + messages_array.extend(messages_spec); + + let mut payload = json!({ + "model": model, + "messages": messages_array + }); + + // Add optional parameters + if !tools_spec.is_empty() { + payload + .as_object_mut() + .unwrap() + .insert("tools".to_string(), json!(tools_spec)); + } + if let Some(temp) = temperature { + payload + .as_object_mut() + .unwrap() + .insert("temperature".to_string(), json!(temp)); + } + if let Some(tokens) = max_tokens { + payload + .as_object_mut() + .unwrap() + .insert("max_tokens".to_string(), json!(tokens)); + } + if let Some(sequences) = stop_sequences { + payload + .as_object_mut() + .unwrap() + .insert("stop".to_string(), json!(sequences)); + } + if let Some(p) = top_p { + payload + .as_object_mut() + .unwrap() + .insert("top_p".to_string(), json!(p)); + } + + // dbg!(&payload); + + // Make request + let response = self.post(payload).await?; + + // Check for context length error if single message + if let Some(error) = response.get("error") { + if messages.len() == 1 { + if let Some(err) = check_openai_context_length_error(error) { + return Err(err.into()); + } + } + return Err(anyhow!("OpenAI API error: {}", error)); + } + + // Parse response + let message = openai_response_to_message(response.clone())?; + let usage = Self::get_usage(&response)?; + + Ok((message, usage)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_get_usage() { + let response = json!({ + "usage": { + "prompt_tokens": 10, + "completion_tokens": 20, + "total_tokens": 30 + } + }); + + let usage = OpenAiProvider::get_usage(&response).unwrap(); + assert_eq!(usage.input_tokens, Some(10)); + assert_eq!(usage.output_tokens, Some(20)); + assert_eq!(usage.total_tokens, Some(30)); + } + + #[test] + fn test_get_usage_calculated_total() { + let response = json!({ + "usage": { + "prompt_tokens": 10, + "completion_tokens": 20 + } + }); + + let usage = OpenAiProvider::get_usage(&response).unwrap(); + 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_creation() { + std::env::set_var("OPENAI_API_KEY", "test_key"); + + let provider = OpenAiProvider::from_env(); + assert!(provider.is_ok()); + + std::env::remove_var("OPENAI_API_KEY"); + } +} +// fn setup_mock_server(response_body: &str) -> (Mock, OpenAiProvider) { +// let mut server = mockito::Server::new(); +// let mock = server.mock("POST", "/v1/chat/completions") +// .match_header("authorization", "Bearer test_key") +// .with_status(200) +// .with_header("content-type", "application/json") +// .with_body(response_body) +// .create(); + +// let config = OpenAiProviderConfig { +// api_key: "test_key".to_string(), +// host: server.url() +// }; +// let provider = OpenAiProvider::new(config).unwrap(); + +// (mock, provider) +// } + +// fn setup_mock_server_for_complete() -> (Mock, OpenAiProvider) { +// setup_mock_server(r#"{ +// "id": "chatcmpl-123", +// "object": "chat.completion", +// "choices": [{ +// "index": 0, +// "message": { +// "role": "assistant", +// "content": "Hello! How can I assist you today?", +// "tool_calls": null +// }, +// "finish_reason": "stop" +// }], +// "usage": { +// "prompt_tokens": 12, +// "completion_tokens": 15, +// "total_tokens": 27 +// } +// }"#) +// } + +// fn setup_mock_server_for_tools() -> (Mock, OpenAiProvider) { +// setup_mock_server(r#"{ +// "id": "chatcmpl-tool", +// "object": "chat.completion", +// "choices": [{ +// "index": 0, +// "message": { +// "role": "assistant", +// "content": null, +// "tool_calls": [{ +// "id": "call_123", +// "type": "function", +// "function": { +// "name": "get_weather", +// "arguments": "{\"location\":\"San Francisco, CA\"}" +// } +// }] +// }, +// "finish_reason": "tool_calls" +// }], +// "usage": { +// "prompt_tokens": 20, +// "completion_tokens": 15, +// "total_tokens": 35 +// } +// }"#) +// } + +// #[test] +// fn test_basic_completion() -> Result<()> { +// let (_mock, provider) = setup_mock_server_for_complete(); + +// let result = provider.complete( +// "gpt-4", +// "You are a helpful assistant.", +// &[Message::user("Hi")?], +// &[], // no tools +// None, +// None, +// None, +// None, +// )?; + +// let (message, usage) = result; + +// // Verify response content +// assert_eq!(message.text(), "Hello! How can I assist you today?"); +// assert!(message.tool_use().is_empty()); + +// // Verify usage statistics +// assert_eq!(usage.total_tokens, Some(27)); +// assert_eq!(usage.input_tokens, Some(12)); +// assert_eq!(usage.output_tokens, Some(15)); + +// Ok(()) +// } + +// #[test] +// fn test_completion_with_tool_call() -> Result<()> { +// let (_mock, provider) = setup_mock_server_for_tools(); + +// // Create a weather tool +// let mut parameters = HashMap::new(); +// parameters.insert( +// "location".to_string(), +// json!({ +// "type": "string", +// "description": "The city and state" +// }) +// ); + +// let weather_tool = Tool::new( +// "get_weather".to_string(), +// "Get the current weather".to_string(), +// parameters, +// |_| Ok(json!({"temperature": 72})) +// ); + +// let result = provider.complete( +// "gpt-4", +// "You are a helpful assistant.", +// &[Message::user("What's the weather in San Francisco?")?], +// &[weather_tool], +// None, +// None, +// None, +// None, +// )?; + +// let (message, usage) = result; + +// // Verify tool calls +// let tool_uses = message.tool_use(); +// assert_eq!(tool_uses.len(), 1); +// assert_eq!(tool_uses[0].name, "get_weather"); + +// // Verify tool parameters +// let expected_params: Value = serde_json::from_str( +// r#"{"location":"San Francisco, CA"}"# +// )?; +// assert_eq!(tool_uses[0].parameters, expected_params); + +// // Verify usage statistics +// assert_eq!(usage.total_tokens, Some(35)); +// assert_eq!(usage.input_tokens, Some(20)); +// assert_eq!(usage.output_tokens, Some(15)); + +// Ok(()) +// } + +// #[test] +// fn test_completion_with_tool_result() -> Result<()> { +// let (_mock, provider) = setup_mock_server_for_complete(); + +// // Create messages including a tool result +// let messages = vec![ +// Message::user("What's the weather?")?, +// Message::new( +// Role::Assistant, +// vec![Content::ToolResult(ToolResult { +// tool_use_id: "call_123".to_string(), +// output: "The temperature is 72°F".to_string(), +// is_error: false, +// })] +// )? +// ]; + +// let result = provider.complete( +// "gpt-4", +// "You are a helpful assistant.", +// &messages, +// &[], // no tools needed for this response +// None, +// None, +// None, +// None, +// )?; + +// let (message, usage) = result; + +// // Verify response +// assert_eq!(message.text(), "Hello! How can I assist you today?"); +// assert!(message.tool_use().is_empty()); + +// // Verify usage +// assert_eq!(usage.total_tokens, Some(27)); + +// Ok(()) +// } +// } diff --git a/src/providers/types.rs b/src/providers/types.rs new file mode 100644 index 0000000000..698decb28b --- /dev/null +++ b/src/providers/types.rs @@ -0,0 +1,4 @@ +pub mod content; +pub mod message; +pub mod objectid; +pub mod tool; diff --git a/src/providers/types/content.rs b/src/providers/types/content.rs new file mode 100644 index 0000000000..19e337de5d --- /dev/null +++ b/src/providers/types/content.rs @@ -0,0 +1,61 @@ +use serde::{Deserialize, Serialize}; +use serde_json::Value; + + +// Text content +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Text { + pub text: String, +} + + +// Tool use content +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ToolUse { + pub id: String, + pub name: String, + pub parameters: Value, + #[serde(default)] + pub is_error: bool, + #[serde(skip_serializing_if = "Option::is_none")] + pub error_message: Option, +} + + + +// Tool result content +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ToolResult { + pub tool_use_id: String, + pub output: String, + #[serde(default)] + pub is_error: bool, +} + + +// Enum to handle all content types +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(tag = "type")] +pub enum Content { + Text(Text), + ToolUse(ToolUse), + ToolResult(ToolResult), +} + +impl Content { + + pub fn summary(&self) -> String { + match self { + Content::Text(t) => format!("content:text\n{}", t.text), + Content::ToolUse(t) => format!( + "content:tool_use:{}\nparameters:{}", + t.name, + serde_json::to_string(&t.parameters).unwrap_or_default() + ), + Content::ToolResult(t) => format!( + "content:tool_result:error={}\noutput:{}", + t.is_error, t.output + ), + } + } +} diff --git a/src/providers/types/message.rs b/src/providers/types/message.rs new file mode 100644 index 0000000000..2e97a6af7b --- /dev/null +++ b/src/providers/types/message.rs @@ -0,0 +1,294 @@ +use anyhow::{anyhow, Result}; +use serde::{Deserialize, Serialize}; +use std::time::{SystemTime, UNIX_EPOCH}; + +use super::content::{Content, Text, ToolResult, ToolUse}; +use super::objectid::create_object_id; + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum Role { + User, + Assistant, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Message { + pub role: Role, + pub id: String, + pub created: i64, + pub content: Vec, +} + +impl Message { + pub fn new(role: Role, content: Vec) -> Result { + let msg = Self { + role, + id: create_object_id("msg"), + created: SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs() as i64, + content, + }; + msg.validate()?; + Ok(msg) + } + + fn validate(&self) -> Result<()> { + match self.role { + Role::User => { + if !self.has_text() && !self.has_tool_result() { + return Err(anyhow!("User message must include a Text or ToolResult")); + } + if self.has_tool_use() { + return Err(anyhow!("User message does not support ToolUse")); + } + } + Role::Assistant => { + if !self.has_text() && !self.has_tool_use() { + return Err(anyhow!("Assistant message must include a Text or ToolUse")); + } + if self.has_tool_result() { + return Err(anyhow!("Assistant message does not support ToolResult")); + } + } + } + Ok(()) + } + + pub fn text(&self) -> String { + self.content + .iter() + .filter_map(|content| match content { + Content::Text(text) => Some(text.text.clone()), + _ => None, + }) + .collect::>() + .join("\n") + } + + pub fn tool_use(&self) -> Vec { + self.content + .iter() + .filter_map(|content| match content { + Content::ToolUse(tool_use) => Some(tool_use.clone()), + _ => None, + }) + .collect() + } + + pub fn tool_result(&self) -> Vec { + self.content + .iter() + .filter_map(|content| match content { + Content::ToolResult(tool_result) => Some(tool_result.clone()), + _ => None, + }) + .collect() + } + + pub fn has_text(&self) -> bool { + self.content.iter().any(|c| matches!(c, Content::Text(_))) + } + + pub fn has_tool_use(&self) -> bool { + self.content + .iter() + .any(|c| matches!(c, Content::ToolUse(_))) + } + + pub fn has_tool_result(&self) -> bool { + self.content + .iter() + .any(|c| matches!(c, Content::ToolResult(_))) + } + + pub fn user(text: &str) -> Result { + Self::new( + Role::User, + vec![Content::Text(Text { + text: text.to_string(), + })], + ) + } + + pub fn assistant(text: &str) -> Result { + Self::new( + Role::Assistant, + vec![Content::Text(Text { + text: text.to_string(), + })], + ) + } + + pub fn summary(&self) -> String { + let content_summaries: Vec = self.content.iter().map(|c| c.summary()).collect(); + format!("message:{:?}\n{}", self.role, content_summaries.join("\n")) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::{json, Value}; + + #[test] + fn test_user_message() -> Result<()> { + let user_message = Message::user("abcd")?; + assert!(matches!(user_message.role, Role::User)); + assert_eq!(user_message.text(), "abcd"); + Ok(()) + } + + #[test] + fn test_assistant_message() -> Result<()> { + let assistant_message = Message::assistant("abcd")?; + assert!(matches!(assistant_message.role, Role::Assistant)); + assert_eq!(assistant_message.text(), "abcd"); + Ok(()) + } + + #[test] + fn test_message_tool_use() -> Result<()> { + let tu1 = ToolUse { + id: "1".to_string(), + name: "tool".to_string(), + parameters: json!({}), + is_error: false, + error_message: None, + }; + let tu2 = ToolUse { + id: "2".to_string(), + name: "tool".to_string(), + parameters: json!({}), + is_error: false, + error_message: None, + }; + + let message = Message::new( + Role::Assistant, + vec![Content::ToolUse(tu1), Content::ToolUse(tu2)], + )?; + + let tool_uses = message.tool_use(); + assert_eq!(tool_uses.len(), 2); + assert_eq!(tool_uses[0].name, "tool"); + Ok(()) + } + + #[test] + fn test_message_tool_result() -> Result<()> { + let tr1 = ToolResult { + tool_use_id: "1".to_string(), + output: "result".to_string(), + is_error: false, + }; + let tr2 = ToolResult { + tool_use_id: "2".to_string(), + output: "result".to_string(), + is_error: false, + }; + + let message = Message::new( + Role::User, + vec![Content::ToolResult(tr1), Content::ToolResult(tr2)], + )?; + + let tool_results = message.tool_result(); + assert_eq!(tool_results.len(), 2); + assert_eq!(tool_results[0].output, "result"); + Ok(()) + } + + #[test] + fn test_message_validation() -> Result<()> { + // Valid user message + let message = Message::user("Hello")?; + assert_eq!(message.text(), "Hello"); + + // Valid assistant message + let message = Message::assistant("Hello")?; + assert_eq!(message.text(), "Hello"); + + // Invalid message: user with tool_use + let result = Message::new( + Role::User, + vec![ + Content::Text(Text { + text: "".to_string(), + }), + Content::ToolUse(ToolUse { + id: "1".to_string(), + name: "tool".to_string(), + parameters: json!({}), + is_error: false, + error_message: None, + }), + ], + ); + assert!(result.is_err()); + + // Invalid message: assistant with tool_result + let result = Message::new( + Role::Assistant, + vec![ + Content::Text(Text { + text: "".to_string(), + }), + Content::ToolResult(ToolResult { + tool_use_id: "1".to_string(), + output: "result".to_string(), + is_error: false, + }), + ], + ); + assert!(result.is_err()); + Ok(()) + } + + #[test] + fn test_serialization() -> Result<()> { + // Test basic message serialization + let message = Message::user("Hello, world!")?; + let serialized = serde_json::to_string(&message)?; + let deserialized: Message = serde_json::from_str(&serialized)?; + assert_eq!(message.text(), deserialized.text()); + assert!(matches!(deserialized.role, Role::User)); + + // Test complex message with tool use + let tool_use = ToolUse { + id: "test_id".to_string(), + name: "test_tool".to_string(), + parameters: json!({"key": "value"}), + is_error: false, + error_message: None, + }; + + let message = Message::new( + Role::Assistant, + vec![ + Content::Text(Text { + text: "Using tool".to_string(), + }), + Content::ToolUse(tool_use), + ], + )?; + + let serialized = serde_json::to_string(&message)?; + let deserialized: Message = serde_json::from_str(&serialized)?; + + assert_eq!(message.text(), deserialized.text()); + assert_eq!(message.tool_use().len(), deserialized.tool_use().len()); + assert_eq!(message.tool_use()[0].name, deserialized.tool_use()[0].name); + + // Verify JSON structure + let json_value: Value = serde_json::from_str(&serialized)?; + assert!(json_value.get("id").is_some()); + assert!(json_value.get("created").is_some()); + assert!(json_value.get("role").is_some()); + assert!(json_value.get("content").is_some()); + + Ok(()) + } +} diff --git a/src/providers/types/objectid.rs b/src/providers/types/objectid.rs new file mode 100644 index 0000000000..a864914294 --- /dev/null +++ b/src/providers/types/objectid.rs @@ -0,0 +1,30 @@ +use uuid::Uuid; + +pub fn create_object_id(prefix: &str) -> String { + format!("{}_{}", prefix, Uuid::new_v4().simple().to_string()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_create_object_id_format() { + let id = create_object_id("test"); + assert!(id.starts_with("test_"), "ID should start with 'test_'"); + let expected_length = "test_".len() + 32; + assert_eq!( + id.len(), + expected_length, + "ID length should be {} (prefix + '_' + 32 char UUID)", + expected_length + ); + } + + #[test] + fn test_create_object_id_uniqueness() { + let id1 = create_object_id("test"); + let id2 = create_object_id("test"); + assert_ne!(id1, id2, "Generated IDs should be unique"); + } +} diff --git a/src/providers/types/tool.rs b/src/providers/types/tool.rs new file mode 100644 index 0000000000..fdd4f58e8d --- /dev/null +++ b/src/providers/types/tool.rs @@ -0,0 +1,140 @@ +use std::fmt::Debug; + +/// A tool that can be used by a model. +pub struct Tool { + /// The name of the tool + pub name: String, + /// A description of what the tool does + pub description: String, + /// A json schema of the function signature + pub parameters: serde_json::Value, + /// The function that powers the tool + /// Note: We use Box to store callable functions + pub function: + Box anyhow::Result + Send + Sync>, +} + +impl Tool { + pub fn new( + name: String, + description: String, + parameters: serde_json::Value, + function: impl Fn(&serde_json::Value) -> anyhow::Result + + Send + + Sync + + 'static, + ) -> Self { + Tool { + name, + description, + parameters, + function: Box::new(function), + } + } +} + +impl Debug for Tool { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Tool") + .field("name", &self.name) + .field("description", &self.description) + .field("parameters", &self.parameters) + .field("function", &"") + .finish() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + // Helper function to simulate your Python weather function + fn get_current_weather(params: &serde_json::Value) -> anyhow::Result { + // In a real implementation, this would actually fetch weather + Ok(json!({ "location": params["location"], "temperature": 72 })) + } + + #[test] + fn test_basic_tool_creation() { + let parameters = json!({ + "type": "object", + "properties": { + "location": { + "type": "string", + "description": "The city and state, e.g. San Francisco, CA" + } + }, + "required": ["location"] + }); + + let tool = Tool::new( + "get_current_weather".to_string(), + "Get the current weather in a given location".to_string(), + parameters.clone(), + get_current_weather, + ); + + assert_eq!(tool.name, "get_current_weather"); + assert_eq!( + tool.description, + "Get the current weather in a given location" + ); + assert_eq!(tool.parameters, parameters); + + // Test function execution + let result = (tool.function)(&json!({"location": "Chicago, IL"})).unwrap(); + assert_eq!(result.get("temperature").unwrap(), 72); + assert_eq!(result.get("location").unwrap(), "Chicago, IL"); + } + + #[test] + fn test_tool_with_multiple_params() { + let parameters = json!({ + "type": "object", + "properties": { + "param1": { + "type": "integer", + "description": "Description for param1" + }, + "param2": { + "type": "string", + "description": "Description for param2" + } + }, + "required": ["param1", "param2"] + }); + + let tool = Tool::new( + "multi_param_function".to_string(), + "A function with multiple parameters".to_string(), + parameters.clone(), + |params| Ok(json!({ "received": params })), + ); + + assert_eq!(tool.parameters, parameters); + + // Test function execution with multiple parameters + let test_params = json!({ + "param1": 42, + "param2": "test" + }); + let result = (tool.function)(&test_params).unwrap(); + assert_eq!(result.get("received"), Some(&test_params)); + } + + #[test] + fn test_tool_debug_output() { + let tool = Tool::new( + "test_tool".to_string(), + "Test description".to_string(), + json!({}), + |_| Ok(json!({})), + ); + + let debug_output = format!("{:?}", tool); + assert!(debug_output.contains("test_tool")); + assert!(debug_output.contains("Test description")); + assert!(debug_output.contains("")); + } +} diff --git a/src/providers/utils.rs b/src/providers/utils.rs new file mode 100644 index 0000000000..2c3d10c32c --- /dev/null +++ b/src/providers/utils.rs @@ -0,0 +1,447 @@ +use anyhow::{anyhow, Result}; +use regex::Regex; +use serde_json::{json, Value}; + +use super::types::{ + content::{Content, Text, ToolUse}, + message::Message, + tool::Tool, +}; + +/// Convert internal Message format to OpenAI's API message specification +pub fn messages_to_openai_spec(messages: &[Message]) -> 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 { + match content { + Content::Text(Text { text }) => { + converted["content"] = json!(text); + } + Content::ToolUse(tool_use) => { + let sanitized_name = sanitize_function_name(&tool_use.name); + let tool_calls = converted + .as_object_mut() + .unwrap() + .entry("tool_calls") + .or_insert(json!([])); + + tool_calls.as_array_mut().unwrap().push(json!({ + "id": tool_use.id, + "type": "function", + "function": { + "name": sanitized_name, + "arguments": tool_use.parameters.to_string(), + } + })); + } + Content::ToolResult(tool_result) => { + // Handle image results (assuming similar functionality as Python) + if tool_result.output.starts_with("\"image:") { + // TODO: Implement image handling if needed + output.push(json!({ + "role": "tool", + "content": [{ + "type": "text", + "text": "This tool result included an image that is uploaded in the next message." + }], + "tool_call_id": tool_result.tool_use_id + })); + } else { + output.push(json!({ + "role": "tool", + "content": tool_result.output, + "tool_call_id": tool_result.tool_use_id + })); + } + } + } + } + + 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 tools_to_openai_spec(tools: &[Tool]) -> 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)); + } + + result.push(json!({ + "type": "function", + "function": { + "name": tool.name, + "description": tool.description, + "parameters": tool.parameters, + } + })); + } + + Ok(result) +} + +/// Convert OpenAI's API response to internal Message format +pub fn openai_response_to_message(response: Value) -> 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(Content::Text(Text { + text: text_str.to_string(), + })); + } + } + + 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 arguments = tool_call["function"]["arguments"] + .as_str() + .unwrap_or_default() + .to_string(); + + if !is_valid_function_name(&function_name) { + content.push(Content::ToolUse(ToolUse { + id, + name: function_name.clone(), + parameters: json!(arguments), + is_error: true, + error_message: Some(format!( + "The provided function name '{}' had invalid characters, it must match this regex [a-zA-Z0-9_-]+", + function_name + )), + })); + } else { + match serde_json::from_str::(&arguments) { + Ok(params) => { + content.push(Content::ToolUse(ToolUse { + id, + name: function_name, + parameters: params, + is_error: false, + error_message: None, + })); + } + Err(_) => { + content.push(Content::ToolUse(ToolUse { + id: id.clone(), + name: function_name, + parameters: json!(arguments), + is_error: true, + error_message: Some(format!( + "Could not interpret tool use parameters for id {}: {}", + id, arguments + )), + })); + } + } + } + } + } + } + + Message::new(super::types::message::Role::Assistant, content) +} + +fn sanitize_function_name(name: &str) -> String { + let re = Regex::new(r"[^a-zA-Z0-9_-]").unwrap(); + re.replace_all(name, "_").to_string() +} + +fn is_valid_function_name(name: &str) -> bool { + let re = Regex::new(r"^[a-zA-Z0-9_-]+$").unwrap(); + re.is_match(name) +} + +#[derive(Debug, thiserror::Error)] +#[error("Input message too long. Message: {0}")] +pub struct InitialMessageTooLargeError(String); + +pub fn check_openai_context_length_error(error: &Value) -> Option { + let code = error.get("code")?.as_str()?; + if code == "context_length_exceeded" || code == "string_above_max_length" { + let message = error + .get("message") + .and_then(|m| m.as_str()) + .unwrap_or("Unknown error") + .to_string(); + Some(InitialMessageTooLargeError(message)) + } else { + None + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::providers::types::content::ToolResult; + use crate::providers::types::message::Role; + use serde_json::json; + + 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_messages_to_openai_spec() -> Result<()> { + let message = Message::user("Hello")?; + let spec = messages_to_openai_spec(&[message]); + + assert_eq!(spec.len(), 1); + assert_eq!(spec[0]["role"], "user"); + assert_eq!(spec[0]["content"], "Hello"); + Ok(()) + } + + #[test] + fn test_tools_to_openai_spec() -> Result<()> { + let parameters = json!({ + "type": "object" + }); + let tool = Tool::new( + "test_tool".to_string(), + "A test tool".to_string(), + parameters, + |_| Ok(json!({})), + ); + + let spec = tools_to_openai_spec(&[tool])?; + + assert_eq!(spec.len(), 1); + assert_eq!(spec[0]["type"], "function"); + assert_eq!(spec[0]["function"]["name"], "test_tool"); + Ok(()) + } + + #[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")); + } + + #[test] + fn test_messages_to_openai_spec_complex() -> Result<()> { + let messages = vec![ + Message::assistant("Hello!")?, + Message::user("How are you?")?, + Message::new( + Role::Assistant, + vec![Content::ToolUse(ToolUse { + id: "1".to_string(), + name: "tool1".to_string(), + parameters: json!({"param1": "value1"}), + is_error: false, + error_message: None, + })], + )?, + Message::new( + Role::User, + vec![Content::ToolResult(ToolResult { + tool_use_id: "1".to_string(), + output: "Result".to_string(), + is_error: false, + })], + )?, + ]; + + let spec = messages_to_openai_spec(&messages); + + 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"], "1"); + + Ok(()) + } + + #[test] + fn test_tools_to_openai_spec_duplicate() -> Result<()> { + let parameters = json!({ + "type": "object" + }); + + let tool1 = Tool::new( + "test_tool".to_string(), + "Test tool".to_string(), + parameters.clone(), + |_| Ok(json!({})), + ); + + let tool2 = Tool::new( + "test_tool".to_string(), + "Test tool".to_string(), + parameters.clone(), + |_| Ok(json!({})), + ); + + let result = tools_to_openai_spec(&[tool1, tool2]); + assert!(result.is_err()); + assert!(result + .unwrap_err() + .to_string() + .contains("Duplicate tool name")); + + Ok(()) + } + + #[test] + fn test_tools_to_openai_spec_empty() -> Result<()> { + let spec = tools_to_openai_spec(&[])?; + assert!(spec.is_empty()); + Ok(()) + } + + #[test] + fn test_openai_response_to_message_text() -> 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 = openai_response_to_message(response)?; + assert_eq!(message.text(), "Hello from John Cena!"); + assert!(matches!(message.role, Role::Assistant)); + + Ok(()) + } + + #[test] + fn test_openai_response_to_message_valid_tooluse() -> Result<()> { + let response: Value = serde_json::from_str(OPENAI_TOOL_USE_RESPONSE)?; + let message = openai_response_to_message(response)?; + + let tool_uses = message.tool_use(); + assert_eq!(tool_uses.len(), 1); + assert_eq!(tool_uses[0].name, "example_fn"); + assert_eq!(tool_uses[0].parameters, json!({"param": "value"})); + assert!(!tool_uses[0].is_error); + assert!(tool_uses[0].error_message.is_none()); + + Ok(()) + } + + #[test] + fn test_openai_response_to_message_invalid_func_name() -> Result<()> { + let mut response: Value = serde_json::from_str(OPENAI_TOOL_USE_RESPONSE)?; + response["choices"][0]["message"]["tool_calls"][0]["function"]["name"] = + json!("invalid fn"); + + let message = openai_response_to_message(response)?; + let tool_uses = message.tool_use(); + + assert_eq!(tool_uses[0].name, "invalid fn"); + assert!(tool_uses[0].is_error); + assert!(tool_uses[0] + .error_message + .as_ref() + .unwrap() + .starts_with("The provided function name")); + + Ok(()) + } + + #[test] + fn test_openai_response_to_message_json_decode_error() -> Result<()> { + let mut response: Value = serde_json::from_str(OPENAI_TOOL_USE_RESPONSE)?; + response["choices"][0]["message"]["tool_calls"][0]["function"]["arguments"] = + json!("invalid json {"); + + let message = openai_response_to_message(response)?; + let tool_uses = message.tool_use(); + + assert_eq!(tool_uses[0].name, "example_fn"); + assert!(tool_uses[0].is_error); + assert!(tool_uses[0] + .error_message + .as_ref() + .unwrap() + .starts_with("Could not interpret tool use parameters")); + + Ok(()) + } + + #[test] + fn test_check_openai_context_length_error() { + let error = json!({ + "code": "context_length_exceeded", + "message": "This message is too long" + }); + + let result = check_openai_context_length_error(&error); + assert!(result.is_some()); + assert_eq!( + result.unwrap().to_string(), + "Input message too long. Message: This message is too long" + ); + + let error = json!({ + "code": "other_error", + "message": "Some other error" + }); + + let result = check_openai_context_length_error(&error); + assert!(result.is_none()); + } +}