From 70fb365b9b8f0f8efedbd1bcf821ada0a8e19069 Mon Sep 17 00:00:00 2001 From: bzp2010 Date: Tue, 12 May 2026 01:04:43 +0800 Subject: [PATCH 1/2] feat(guardrail): add definitions and traits --- Cargo.lock | 10 + Cargo.toml | 3 +- crates/aisix-guardrail/Cargo.toml | 10 + crates/aisix-guardrail/src/lib.rs | 2 + crates/aisix-guardrail/src/traits.rs | 270 +++++++++++++++++++++++++++ 5 files changed, 294 insertions(+), 1 deletion(-) create mode 100644 crates/aisix-guardrail/Cargo.toml create mode 100644 crates/aisix-guardrail/src/lib.rs create mode 100644 crates/aisix-guardrail/src/traits.rs diff --git a/Cargo.lock b/Cargo.lock index e71e713..d960a85 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -30,6 +30,7 @@ name = "aisix" version = "0.1.0" dependencies = [ "aisix-admin-ui", + "aisix-guardrail", "aisix-llm", "anyhow", "arc-swap", @@ -104,6 +105,15 @@ dependencies = [ "rust-embed", ] +[[package]] +name = "aisix-guardrail" +version = "0.1.0" +dependencies = [ + "async-trait", + "serde", + "utoipa", +] + [[package]] name = "aisix-llm" version = "0.1.0" diff --git a/Cargo.toml b/Cargo.toml index d0011fb..a41e461 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,5 +1,5 @@ [workspace] -members = ["crates/aisix-llm", "crates/admin-ui"] +members = ["crates/aisix-llm", "crates/aisix-guardrail", "crates/admin-ui"] resolver = "3" [workspace.package] @@ -56,6 +56,7 @@ build-ui = ["aisix-admin-ui/build-ui"] [dependencies] aisix-llm = { path = "crates/aisix-llm" } +aisix-guardrail = { path = "crates/aisix-guardrail" } aisix-admin-ui = { path = "crates/admin-ui" } anyhow.workspace = true tokio.workspace = true diff --git a/crates/aisix-guardrail/Cargo.toml b/crates/aisix-guardrail/Cargo.toml new file mode 100644 index 0000000..4a071ca --- /dev/null +++ b/crates/aisix-guardrail/Cargo.toml @@ -0,0 +1,10 @@ +[package] +name = "aisix-guardrail" +version.workspace = true +edition.workspace = true +rust-version.workspace = true + +[dependencies] +async-trait.workspace = true +serde.workspace = true +utoipa.workspace = true diff --git a/crates/aisix-guardrail/src/lib.rs b/crates/aisix-guardrail/src/lib.rs new file mode 100644 index 0000000..d99c24e --- /dev/null +++ b/crates/aisix-guardrail/src/lib.rs @@ -0,0 +1,2 @@ +pub mod guardrails; +pub mod traits; diff --git a/crates/aisix-guardrail/src/traits.rs b/crates/aisix-guardrail/src/traits.rs new file mode 100644 index 0000000..1117b18 --- /dev/null +++ b/crates/aisix-guardrail/src/traits.rs @@ -0,0 +1,270 @@ +use async_trait::async_trait; + +/// Declares metadata about a guardrail implementation. +pub trait GuardrailMeta: Send + Sync + 'static { + fn name(&self) -> &'static str; + + fn supported_stages(&self) -> &'static [GuardrailStage] { + &[GuardrailStage::Input, GuardrailStage::Output] + } + + fn supports_stage(&self, stage: GuardrailStage) -> bool { + self.supported_stages().contains(&stage) + } +} + +/// Request lifecycle stage where a guardrail can run. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum GuardrailStage { + Input, + Output, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum GuardrailRole { + System, + User, + Assistant, + Tool, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum GuardrailMessageContent { + Text(String), + Parts(Vec), +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum GuardrailContentPart { + Text { text: String }, + ImageUrl { image_url: GuardrailImageUrl }, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct GuardrailImageUrl { + pub url: String, + pub detail: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct GuardrailToolCall { + pub id: String, + pub r#type: String, + pub name: String, + pub arguments: String, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct GuardrailMessage { + pub role: GuardrailRole, + pub content: Option, + pub name: Option, + pub tool_calls: Option>, + pub tool_call_id: Option, +} + +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct InputGuardrailPayload { + pub messages: Vec, +} + +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct OutputGuardrailPayload { + pub messages: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum GuardrailCheckPayload { + Input(InputGuardrailPayload), + Output(OutputGuardrailPayload), +} + +impl GuardrailCheckPayload { + pub fn stage(&self) -> GuardrailStage { + match self { + Self::Input(_) => GuardrailStage::Input, + Self::Output(_) => GuardrailStage::Output, + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum GuardrailOutcome { + Allow, + Rewrite(GuardrailCheckPayload), + Block { reason: String }, +} + +/// Runtime contract for message-level guardrail checks. +#[async_trait] +pub trait GuardrailRuntime: GuardrailMeta { + type Error: std::error::Error + Send + Sync + 'static; + + async fn check( + &self, + payload: &GuardrailCheckPayload, + config: &C, + ) -> Result; +} + +#[cfg(test)] +mod tests { + use super::{ + GuardrailCheckPayload, GuardrailContentPart, GuardrailImageUrl, GuardrailMessage, + GuardrailMessageContent, GuardrailMeta, GuardrailOutcome, GuardrailRole, GuardrailStage, + GuardrailToolCall, InputGuardrailPayload, OutputGuardrailPayload, + }; + + struct DummyGuardrailMeta; + + impl GuardrailMeta for DummyGuardrailMeta { + fn name(&self) -> &'static str { + "dummy" + } + } + + #[test] + fn guardrail_meta_should_support_input_and_output_by_default() { + let meta = DummyGuardrailMeta; + + assert!(meta.supports_stage(GuardrailStage::Input)); + assert!(meta.supports_stage(GuardrailStage::Output)); + } + + #[test] + fn guardrail_check_payload_should_report_its_stage() { + let input_payload = GuardrailCheckPayload::Input(InputGuardrailPayload::default()); + let output_payload = GuardrailCheckPayload::Output(OutputGuardrailPayload::default()); + + assert_eq!(input_payload.stage(), GuardrailStage::Input); + assert_eq!(output_payload.stage(), GuardrailStage::Output); + } + + #[test] + fn guardrail_check_payload_should_embed_tool_calls_within_assistant_messages() { + let payload = GuardrailCheckPayload::Output(OutputGuardrailPayload { + messages: vec![GuardrailMessage { + role: GuardrailRole::Assistant, + content: None, + name: None, + tool_calls: Some(vec![GuardrailToolCall { + id: "call_weather_1".into(), + r#type: "function".into(), + name: "get_weather".into(), + arguments: r#"{"city":"Hangzhou"}"#.into(), + }]), + tool_call_id: None, + }], + }); + + let GuardrailCheckPayload::Output(payload) = payload else { + panic!("expected output payload"); + }; + + assert_eq!(payload.messages.len(), 1); + assert_eq!(payload.messages[0].role, GuardrailRole::Assistant); + assert_eq!(payload.messages[0].tool_calls.as_ref().unwrap().len(), 1); + assert_eq!( + payload.messages[0].tool_calls.as_ref().unwrap()[0].name, + "get_weather" + ); + } + + #[test] + fn guardrail_check_payload_should_represent_tool_results_as_tool_messages() { + let payload = GuardrailCheckPayload::Input(InputGuardrailPayload { + messages: vec![GuardrailMessage { + role: GuardrailRole::Tool, + content: Some(GuardrailMessageContent::Text( + r#"{"temperature":23}"#.into(), + )), + name: None, + tool_calls: None, + tool_call_id: Some("call_weather_1".into()), + }], + }); + + let GuardrailCheckPayload::Input(payload) = payload else { + panic!("expected input payload"); + }; + + assert_eq!(payload.messages.len(), 1); + assert_eq!(payload.messages[0].role, GuardrailRole::Tool); + assert_eq!( + payload.messages[0].tool_call_id.as_deref(), + Some("call_weather_1") + ); + assert_eq!( + payload.messages[0].content, + Some(GuardrailMessageContent::Text( + r#"{"temperature":23}"#.into() + )) + ); + } + + #[test] + fn guardrail_message_should_preserve_multimodal_content_parts() { + let payload = GuardrailCheckPayload::Input(InputGuardrailPayload { + messages: vec![GuardrailMessage { + role: GuardrailRole::User, + content: Some(GuardrailMessageContent::Parts(vec![ + GuardrailContentPart::Text { + text: "describe this image".into(), + }, + GuardrailContentPart::ImageUrl { + image_url: GuardrailImageUrl { + url: "https://example.com/cat.png".into(), + detail: Some("high".into()), + }, + }, + ])), + name: Some("alice".into()), + tool_calls: None, + tool_call_id: None, + }], + }); + + let GuardrailCheckPayload::Input(payload) = payload else { + panic!("expected input payload"); + }; + + assert_eq!(payload.messages[0].name.as_deref(), Some("alice")); + assert_eq!( + payload.messages[0].content, + Some(GuardrailMessageContent::Parts(vec![ + GuardrailContentPart::Text { + text: "describe this image".into(), + }, + GuardrailContentPart::ImageUrl { + image_url: GuardrailImageUrl { + url: "https://example.com/cat.png".into(), + detail: Some("high".into()), + }, + }, + ])) + ); + } + + #[test] + fn guardrail_outcome_should_allow_rewriting_full_payloads() { + let outcome = + GuardrailOutcome::Rewrite(GuardrailCheckPayload::Input(InputGuardrailPayload { + messages: vec![GuardrailMessage { + role: GuardrailRole::User, + content: Some(GuardrailMessageContent::Text("hello".into())), + name: None, + tool_calls: None, + tool_call_id: None, + }], + })); + + let GuardrailOutcome::Rewrite(GuardrailCheckPayload::Input(payload)) = outcome else { + panic!("expected rewrite input outcome"); + }; + + assert_eq!( + payload.messages[0].content, + Some(GuardrailMessageContent::Text("hello".into())) + ); + } +} From bdec91ce47299f2dd048e7a2d1ed0d183efa2391 Mon Sep 17 00:00:00 2001 From: bzp2010 Date: Tue, 12 May 2026 01:10:04 +0800 Subject: [PATCH 2/2] fix lint --- crates/aisix-guardrail/src/lib.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/crates/aisix-guardrail/src/lib.rs b/crates/aisix-guardrail/src/lib.rs index d99c24e..f6ac8fc 100644 --- a/crates/aisix-guardrail/src/lib.rs +++ b/crates/aisix-guardrail/src/lib.rs @@ -1,2 +1 @@ -pub mod guardrails; pub mod traits;