From ec4672d3ab2ae96b5f1e2d013b3a16e978d06d0b Mon Sep 17 00:00:00 2001 From: Zaki Ali Date: Mon, 13 Jan 2025 10:57:25 -0800 Subject: [PATCH] Add Databricks moderation (#540) --- crates/goose-server/src/routes/reply.rs | 98 ++++++-- crates/goose/src/providers/base.rs | 38 +-- crates/goose/src/providers/databricks.rs | 284 ++++++++++++++++++++++- crates/goose/src/providers/google.rs | 6 +- ui/desktop/src/ChatWindow.tsx | 26 +-- 5 files changed, 402 insertions(+), 50 deletions(-) diff --git a/crates/goose-server/src/routes/reply.rs b/crates/goose-server/src/routes/reply.rs index 141a584f88..d15b3b485b 100644 --- a/crates/goose-server/src/routes/reply.rs +++ b/crates/goose-server/src/routes/reply.rs @@ -9,6 +9,7 @@ use axum::{ use bytes::Bytes; use futures::{stream::StreamExt, Stream}; use goose::message::{Message, MessageContent}; +use goose::providers::base::ModerationError; use mcp_core::{content::Content, role::Role}; use serde::Deserialize; use serde_json::{json, Value}; @@ -159,6 +160,23 @@ impl ProtocolFormatter { format!("a:{}\n", response) } + fn format_error(error: &str) -> String { + // Error messages start with "3:" in the new protocol. + format!("3:{}\n", error) + } + + fn format_moderation_error(error: &ModerationError) -> String { + let error_part = match error { + ModerationError::ContentFlagged { categories, .. } => { + format!( + "Content was flagged by moderation in the following categories: {}", + categories + ) + } + }; + format!("3:\"{}\"\n", error_part) + } + fn format_finish(reason: &str) -> String { // Finish messages start with "d:" let finish = json!({ @@ -193,8 +211,12 @@ async fn stream_message( .await?; } Err(err) => { + // Send an error message first + tx.send(ProtocolFormatter::format_error(&err.to_string())) + .await?; + // Then send an empty tool response to maintain the protocol let result = - vec![Content::text(format!("Error {}", err)).with_priority(0.0)]; + vec![Content::text(format!("Error: {}", err)).with_priority(0.0)]; tx.send(ProtocolFormatter::format_tool_response( &response.id, &result, @@ -209,22 +231,24 @@ async fn stream_message( for content in message.content { match content { MessageContent::ToolRequest(request) => { - if let Ok(tool_call) = request.tool_call { - tx.send(ProtocolFormatter::format_tool_call( - &request.id, - &tool_call.name, - &tool_call.arguments, - )) - .await?; - } else { - // if the llm generates an invalid object tool call, we still have - // to include it in the history. It always comes with a response indicating the error - tx.send(ProtocolFormatter::format_tool_call( - &request.id, - "invalid name", - &json!({}), - )) - .await?; + match request.tool_call { + Ok(tool_call) => { + tx.send(ProtocolFormatter::format_tool_call( + &request.id, + &tool_call.name, + &tool_call.arguments, + )) + .await?; + } + Err(err) => { + // Send a placeholder tool call to maintain protocol + tx.send(ProtocolFormatter::format_tool_call( + &request.id, + "invalid_tool", + &json!({"error": err.to_string()}), + )) + .await?; + } } } MessageContent::Text(text) => { @@ -278,6 +302,18 @@ async fn handler( Ok(stream) => stream, Err(e) => { tracing::error!("Failed to start reply stream: {}", e); + // Check if it's a moderation error + if let Some(moderation_error) = e.downcast_ref::() { + let _ = tx + .send(ProtocolFormatter::format_moderation_error(moderation_error)) + .await; + // Kill the stream since we encountered a moderation error + } else { + // Send a generic error message + let _ = tx + .send(ProtocolFormatter::format_error(&e.to_string())) + .await; + } // Send a finish message with error as the reason let _ = tx.send(ProtocolFormatter::format_finish("error")).await; return; @@ -291,11 +327,18 @@ async fn handler( Ok(Some(Ok(message))) => { if let Err(e) = stream_message(message, &tx).await { tracing::error!("Error sending message through channel: {}", e); + let _ = tx.send(ProtocolFormatter::format_error(&e.to_string())).await; break; } } Ok(Some(Err(e))) => { tracing::error!("Error processing message: {}", e); + // Check if it's a moderation error + if let Some(moderation_error) = e.downcast_ref::() { + let _ = tx.send(ProtocolFormatter::format_moderation_error(moderation_error)).await; + } else { + let _ = tx.send(ProtocolFormatter::format_error(&e.to_string())).await; + } break; } Ok(None) => { @@ -503,6 +546,27 @@ mod tests { assert!(formatted.starts_with("a:")); assert!(formatted.contains("\"toolCallId\":\"123\"")); + // Test error formatting + let formatted = ProtocolFormatter::format_error("Test error"); + println!("Formatted error: {}", formatted); + assert!(formatted.starts_with("3:")); + assert!(formatted.contains("Test error")); + + // Test moderation error formatting + let moderation_error = ModerationError::ContentFlagged { + categories: "hate, violence".to_string(), + category_scores: Some(json!({ + "hate": 0.9, + "violence": 0.8 + })), + }; + let formatted = ProtocolFormatter::format_moderation_error(&moderation_error); + println!("{}", formatted); + assert!(formatted.starts_with("3:")); + assert!( + formatted.contains("Content was flagged by moderation in the following categories:") + ); + // Test finish formatting let formatted = ProtocolFormatter::format_finish("stop"); assert!(formatted.starts_with("d:")); diff --git a/crates/goose/src/providers/base.rs b/crates/goose/src/providers/base.rs index e00ac93388..a9e0f2e9d0 100644 --- a/crates/goose/src/providers/base.rs +++ b/crates/goose/src/providers/base.rs @@ -4,6 +4,7 @@ use rust_decimal::Decimal; use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::sync::Arc; +use thiserror::Error; use tokio::select; use tokio::sync::RwLock; @@ -12,6 +13,15 @@ use crate::message::{Message, MessageContent}; use mcp_core::role::Role; use mcp_core::tool::Tool; +#[derive(Error, Debug)] +pub enum ModerationError { + #[error("Content was flagged for moderation in categories: {categories}")] + ContentFlagged { + categories: String, + category_scores: Option, + }, +} + #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ProviderUsage { pub model: String, @@ -197,10 +207,10 @@ pub trait Provider: Send + Sync + Moderation { let categories = result.categories .unwrap_or_else(|| vec!["unknown".to_string()]) .join(", "); - return Err(anyhow::anyhow!( - "Content was flagged for moderation in categories: {}", - categories - )); + return Err(ModerationError::ContentFlagged { + categories, + category_scores: result.category_scores, + }.into()); } // Moderation passed, wait for completion @@ -215,10 +225,10 @@ pub trait Provider: Send + Sync + Moderation { let categories = moderation_result.categories .unwrap_or_else(|| vec!["unknown".to_string()]) .join(", "); - return Err(anyhow::anyhow!( - "Content was flagged for moderation in categories: {}", - categories - )); + return Err(ModerationError::ContentFlagged { + categories, + category_scores: moderation_result.category_scores, + }.into()); } Ok(completion_result) @@ -338,10 +348,8 @@ mod tests { let result = provider.complete("system", &[test_message], &[]).await; assert!(result.is_err()); - assert!(result - .unwrap_err() - .to_string() - .contains("Content was flagged")); + let err = result.unwrap_err(); + assert!(err.downcast_ref::().is_some()); } #[tokio::test] @@ -407,10 +415,8 @@ mod tests { let result = provider.complete("system", &[test_message], &[]).await; assert!(result.is_err()); - assert!(result - .unwrap_err() - .to_string() - .contains("Content was flagged")); + let err = result.unwrap_err(); + assert!(err.downcast_ref::().is_some()); } #[tokio::test] diff --git a/crates/goose/src/providers/databricks.rs b/crates/goose/src/providers/databricks.rs index 40f22f177f..e3139ebd93 100644 --- a/crates/goose/src/providers/databricks.rs +++ b/crates/goose/src/providers/databricks.rs @@ -17,7 +17,9 @@ use crate::providers::openai_utils::{ use mcp_core::tool::Tool; pub const DATABRICKS_DEFAULT_MODEL: &str = "claude-3-5-sonnet-2"; +pub const INPUT_GUARDRAIL: &str = "input_guardrail"; +#[derive(Debug)] pub struct DatabricksProvider { client: Client, config: DatabricksProviderConfig, @@ -66,6 +68,35 @@ impl DatabricksProvider { handle_response(payload, response).await? } + + async fn handle_moderation_response(&self, response: reqwest::Response) -> Result { + match response.status() { + reqwest::StatusCode::OK => { + let payload = response.json().await?; + Ok(payload) + } + reqwest::StatusCode::BAD_REQUEST => { + let error_body: Value = response.json().await?; + + // Check if this is a moderation error + if let Some(finish_reason) = error_body.get("finishReason") { + if finish_reason == "input_guardrail_triggered" { + return Ok(error_body); + } + } + // Not a moderation error, return the original error + Err(anyhow::anyhow!("Bad request: {}", error_body)) + } + status => { + let error_body: Value = response.json().await?; + Err(anyhow::anyhow!( + "Moderation request failed with status: {}\nPayload {}", + status, + error_body + )) + } + } + } } #[async_trait] @@ -133,6 +164,7 @@ impl Provider for DatabricksProvider { .collect(), ); + // Make request let response = self.post(payload.clone()).await?; // Raise specific error if context length is exceeded @@ -151,6 +183,7 @@ impl Provider for DatabricksProvider { let model = get_model(&response); let cost = cost(&usage, &model_pricing_for(&model)); super::utils::emit_debug_trace(&self.config, &payload, &response, &usage, cost); + Ok((message, ProviderUsage::new(model, usage, cost))) } @@ -161,7 +194,59 @@ impl Provider for DatabricksProvider { #[async_trait] impl Moderation for DatabricksProvider { - async fn moderate_content(&self, _content: &str) -> Result { + async fn moderate_content_internal(&self, content: &str) -> Result { + let url = format!( + "{}/serving-endpoints/moderation/invocations", + self.config.host.trim_end_matches('/') + ); + + let auth_header = self.ensure_auth_header().await?; + let payload = json!({ + "messages": [ + { + "role": "user", + "content": content + } + ] + }); + + let response = self + .client + .post(&url) + .header("Authorization", auth_header) + .json(&payload) + .send() + .await?; + + // let response: Value = response.json().await?; + // let response = handle_response(payload, response).await??; + let response = self.handle_moderation_response(response).await?; + + // Check if we got a moderation result + if let Some(input_guardrail) = response.get(INPUT_GUARDRAIL) { + if let Some(first_result) = input_guardrail.as_array().and_then(|arr| arr.first()) { + if let Some(flagged) = first_result.get("flagged").and_then(|f| f.as_bool()) { + // Extract categories if they exist and if content is flagged + let categories = if flagged { + first_result + .get("categories") + .and_then(|cats| cats.as_object()) + .map(|cats| { + cats.iter() + .filter(|(_, v)| v.as_bool().unwrap_or(false)) + .map(|(k, _)| k.to_string()) + .collect::>() + }) + } else { + None + }; + + return Ok(ModerationResult::new(flagged, categories, None)); + } + } + } + + // If we get here, there was no moderation result, so the content is considered safe Ok(ModerationResult::new(false, None, None)) } } @@ -174,9 +259,177 @@ mod tests { use crate::providers::mock_server::{ create_mock_open_ai_response, TEST_INPUT_TOKENS, TEST_OUTPUT_TOKENS, TEST_TOTAL_TOKENS, }; + use serde_json::json; use wiremock::matchers::{body_json, header, method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; + #[tokio::test] + async fn test_moderation_flagged_content() -> Result<()> { + // Start a mock server + let mock_server = MockServer::start().await; + + // Mock response for moderation with flagged content + let mock_response = json!({ + "usage": { + "prompt_tokens": 199, + "total_tokens": 199 + }, + "input_guardrail": [{ + "flagged": true, + "categories": { + "violent-crimes": false, + "non-violent-crimes": true, + "sex-crimes": false, + "child-exploitation": false, + "specialized-advice": false, + "privacy": false, + "intellectual-property": false, + "indiscriminate-weapons": false, + "hate": false, + "self-harm": false, + "sexual-content": false + } + }], + "finishReason": "input_guardrail_triggered" + }); + + // Set up the mock + Mock::given(method("POST")) + .and(path("/serving-endpoints/moderation/invocations")) + .and(header("Authorization", "Bearer test_token")) + .respond_with(ResponseTemplate::new(200).set_body_json(mock_response)) + .expect(1) + .mount(&mock_server) + .await; + + // Create the DatabricksProvider + let config = DatabricksProviderConfig { + host: mock_server.uri(), + auth: DatabricksAuth::Token("test_token".to_string()), + model: ModelConfig::new("my-databricks-model".to_string()), + image_format: crate::providers::utils::ImageFormat::Anthropic, + }; + + let provider = DatabricksProvider::new(config)?; + + // Test moderation + let result = provider.moderate_content("test content").await?; + + assert!(result.flagged); + assert_eq!(result.categories.unwrap(), vec!["non-violent-crimes"]); + assert!(result.category_scores.is_none()); + + Ok(()) + } + + #[tokio::test] + async fn test_moderation_safe_content() -> Result<()> { + // Start a mock server + let mock_server = MockServer::start().await; + + // Mock response for safe content (regular chat response) + let mock_response = json!({ + "usage": { + "prompt_tokens": 10, + "completion_tokens": 20, + "total_tokens": 30 + }, + "choices": [{ + "message": { + "role": "assistant", + "content": "This is a safe response" + }, + "finish_reason": "stop" + }] + }); + + // Set up the mock + Mock::given(method("POST")) + .and(path("/serving-endpoints/moderation/invocations")) + .and(header("Authorization", "Bearer test_token")) + .respond_with(ResponseTemplate::new(200).set_body_json(mock_response)) + .expect(1) + .mount(&mock_server) + .await; + + // Create the DatabricksProvider + let config = DatabricksProviderConfig { + host: mock_server.uri(), + auth: DatabricksAuth::Token("test_token".to_string()), + model: ModelConfig::new("my-databricks-model".to_string()), + image_format: crate::providers::utils::ImageFormat::Anthropic, + }; + + let provider = DatabricksProvider::new(config)?; + + // Test moderation + let result = provider.moderate_content("safe content").await?; + + assert!(!result.flagged); + assert!(result.categories.is_none()); + assert!(result.category_scores.is_none()); + + Ok(()) + } + + #[tokio::test] + async fn test_moderation_explicit_safe() -> Result<()> { + // Start a mock server + let mock_server = MockServer::start().await; + + // Mock response for explicitly safe content + let mock_response = json!({ + "usage": { + "prompt_tokens": 199, + "total_tokens": 199 + }, + "input_guardrail": [{ + "flagged": false, + "categories": { + "violent-crimes": false, + "non-violent-crimes": false, + "sex-crimes": false, + "child-exploitation": false, + "specialized-advice": false, + "privacy": false, + "intellectual-property": false, + "indiscriminate-weapons": false, + "hate": false, + "self-harm": false, + "sexual-content": false + } + }] + }); + + // Set up the mock + Mock::given(method("POST")) + .and(path("/serving-endpoints/moderation/invocations")) + .and(header("Authorization", "Bearer test_token")) + .respond_with(ResponseTemplate::new(200).set_body_json(mock_response)) + .expect(1) + .mount(&mock_server) + .await; + + // Create the DatabricksProvider + let config = DatabricksProviderConfig { + host: mock_server.uri(), + auth: DatabricksAuth::Token("test_token".to_string()), + model: ModelConfig::new("my-databricks-model".to_string()), + image_format: crate::providers::utils::ImageFormat::Anthropic, + }; + + let provider = DatabricksProvider::new(config)?; + + // Test moderation + let result = provider.moderate_content("explicitly safe content").await?; + + assert!(!result.flagged); + assert!(result.categories.is_none()); + assert!(result.category_scores.is_none()); + + Ok(()) + } + #[tokio::test] async fn test_databricks_completion_with_token() -> Result<()> { // Start a mock server @@ -185,6 +438,21 @@ mod tests { // Mock response for completion let mock_response = create_mock_open_ai_response("my-databricks-model", "Hello!"); + let moderator_mock_response = json!({ + "usage": { + "prompt_tokens": 10, + "completion_tokens": 20, + "total_tokens": 30 + }, + "choices": [{ + "message": { + "role": "assistant", + "content": "This is a safe response" + }, + "finish_reason": "stop" + }] + }); + // Expected request body let system = "You are a helpful assistant."; let expected_request_body = json!({ @@ -193,6 +461,20 @@ mod tests { {"role": "user", "content": "Hello"} ] }); + let expected_moderation_request_body = json!({ + "messages": [ + {"role": "user", "content": "Hello"} + ] + }); + // Set up the mock to intercept the request and respond with the mocked response + Mock::given(method("POST")) + .and(path("/serving-endpoints/moderation/invocations")) + .and(header("Authorization", "Bearer test_token")) + .and(body_json(expected_moderation_request_body.clone())) + .respond_with(ResponseTemplate::new(200).set_body_json(moderator_mock_response)) + .expect(1) // Expect exactly one matching request + .mount(&mock_server) + .await; // Set up the mock to intercept the request and respond with the mocked response Mock::given(method("POST")) diff --git a/crates/goose/src/providers/google.rs b/crates/goose/src/providers/google.rs index 0a83229929..3187c7d511 100644 --- a/crates/goose/src/providers/google.rs +++ b/crates/goose/src/providers/google.rs @@ -369,9 +369,9 @@ mod tests { use super::*; use crate::providers::mock_server::{ - create_mock_google_ai_response, create_mock_google_ai_response_with_tools, - create_test_tool, get_expected_function_call_arguments, setup_mock_server, - TEST_INPUT_TOKENS, TEST_OUTPUT_TOKENS, TEST_TOOL_FUNCTION_NAME, TEST_TOTAL_TOKENS, + create_mock_google_ai_response_with_tools, create_test_tool, + get_expected_function_call_arguments, setup_mock_server, TEST_INPUT_TOKENS, + TEST_OUTPUT_TOKENS, TEST_TOOL_FUNCTION_NAME, TEST_TOTAL_TOKENS, }; use wiremock::MockServer; diff --git a/ui/desktop/src/ChatWindow.tsx b/ui/desktop/src/ChatWindow.tsx index 124fcfa609..c391e92b86 100644 --- a/ui/desktop/src/ChatWindow.tsx +++ b/ui/desktop/src/ChatWindow.tsx @@ -1,21 +1,21 @@ import React, { useEffect, useRef, useState } from 'react'; +import { Navigate, Route, Routes } from 'react-router-dom'; import { Message, useChat } from './ai-sdk-fork/useChat'; -import { Route, Routes, Navigate } from 'react-router-dom'; -import { getApiUrl } from './config'; +import { ApiKeyWarning } from './components/ApiKeyWarning'; +import BottomMenu from './components/BottomMenu'; +import FlappyGoose from './components/FlappyGoose'; +import GooseMessage from './components/GooseMessage'; +import Input from './components/Input'; +import LoadingGoose from './components/LoadingGoose'; +import MoreMenu from './components/MoreMenu'; +import Splash from './components/Splash'; import { Card } from './components/ui/card'; import { ScrollArea } from './components/ui/scroll-area'; -import Splash from './components/Splash'; -import GooseMessage from './components/GooseMessage'; import UserMessage from './components/UserMessage'; -import Input from './components/Input'; -import MoreMenu from './components/MoreMenu'; -import BottomMenu from './components/BottomMenu'; -import LoadingGoose from './components/LoadingGoose'; -import { ApiKeyWarning } from './components/ApiKeyWarning'; -import { askAi } from './utils/askAI'; -import WingToWing, { Working } from './components/WingToWing'; import { WelcomeScreen } from './components/WelcomeScreen'; -import FlappyGoose from './components/FlappyGoose'; +import WingToWing, { Working } from './components/WingToWing'; +import { getApiUrl } from './config'; +import { askAi } from './utils/askAI'; // update this when you want to show the welcome screen again - doesn't have to be an actual version, just anything woudln't have been seen before const CURRENT_VERSION = '0.0.0'; @@ -462,4 +462,4 @@ export default function ChatWindow() { )} ); -} +} \ No newline at end of file