ai-space/goose
3
0
Fork 0
mirror of https://github.com/aaif-goose/goose.git synced 2026-07-03 14:10:03 +02:00
Code Issues Packages Projects Releases Wiki Activity

Add Databricks moderation (#540)

Browse Source
This commit is contained in:
Zaki Ali
2025-01-13 10:57:25 -08:00
committed by GitHub
parent d70fc344eb
commit ec4672d3ab
5 changed files with 402 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-server/src/routes/reply.rs
+81 -17
View File
@@ -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::<ModerationError>() {
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::<ModerationError>() {
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:"));
crates/goose/src/providers/base.rs
+22 -16
View File
@@ -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<serde_json::Value>,
},
}
#[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::<ModerationError>().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::<ModerationError>().is_some());
}
#[tokio::test]
crates/goose/src/providers/databricks.rs
+283 -1
View File
@@ -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<Value> {
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<ModerationResult> {
async fn moderate_content_internal(&self, content: &str) -> Result<ModerationResult> {
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::<Vec<_>>()
})
} 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"))
crates/goose/src/providers/google.rs
+3 -3
View File
@@ -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;
ui/desktop/src/ChatWindow.tsx
+13 -13
View File
@@ -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() {
)}
</div>
);
}
}