feat(guardrail): integrate into chat handler

src/proxy/handlers/chat_completions/mod.rs

@@ -10,16 +10,22 @@ use span_attributes::{
 };
 pub use types::ChatCompletionError;
 
+use super::FormatHandlerAdapter;
 use crate::{
     gateway::{
+        error::GatewayError,
         formats::OpenAIChatFormat,
-        traits::ProviderCapabilities,
+        traits::{ChatFormat, ProviderCapabilities},
         types::{
             common::Usage,
             openai::{ChatCompletionChunk, ChatCompletionRequest, ChatCompletionResponse},
         },
     },
-    proxy::handlers::format_handler::FormatHandlerAdapter,
+    proxy::guardrails::{
+        input_guardrail_payload_from_chat_messages, input_payload_from_check_payload,
+        input_payload_to_chat_messages, output_guardrail_payload_from_chat_messages,
+        output_payload_from_check_payload, output_payload_to_chat_messages,
+    },
 };
 
 pub(crate) struct ChatCompletionsAdapter;
@@ -86,7 +92,77 @@ impl FormatHandlerAdapter for ChatCompletionsAdapter {
         collector.output_message_span_properties()
     }
 
+    fn guardrail_input_payload(
+        _lifecycle_state: &Self::LifecycleState,
+        request: &Self::Request,
+    ) -> Result<Option<crate::guardrail::traits::GuardrailCheckPayload>, Self::Error> {
+        let (hub_request, _) = OpenAIChatFormat::to_hub(request)?;
+        let payload = input_guardrail_payload_from_chat_messages(&hub_request.messages)
+            .map(crate::guardrail::traits::GuardrailCheckPayload::Input)
+            .map_err(bridge_error)?;
+        Ok(Some(payload))
+    }
+
+    fn apply_input_guardrail_rewrite(
+        _lifecycle_state: &mut Self::LifecycleState,
+        request: &mut Self::Request,
+        rewrite: crate::guardrail::traits::GuardrailCheckPayload,
+    ) -> Result<(), Self::Error> {
+        request.messages = input_payload_to_chat_messages(
+            &input_payload_from_check_payload(rewrite).map_err(bridge_error)?,
+        )
+        .map_err(bridge_error)?;
+        Ok(())
+    }
+
+    fn guardrail_output_payload(
+        _lifecycle_state: &Self::LifecycleState,
+        response: &Self::Response,
+    ) -> Result<Option<crate::guardrail::traits::GuardrailCheckPayload>, Self::Error> {
+        let messages = response
+            .choices
+            .iter()
+            .map(|choice| choice.message.clone())
+            .collect::<Vec<_>>();
+        let payload = output_guardrail_payload_from_chat_messages(&messages)
+            .map(crate::guardrail::traits::GuardrailCheckPayload::Output)
+            .map_err(bridge_error)?;
+        Ok(Some(payload))
+    }
+
+    fn apply_output_guardrail_rewrite(
+        _lifecycle_state: &mut Self::LifecycleState,
+        response: &mut Self::Response,
+        rewrite: crate::guardrail::traits::GuardrailCheckPayload,
+    ) -> Result<(), Self::Error> {
+        let messages = output_payload_to_chat_messages(
+            &output_payload_from_check_payload(rewrite).map_err(bridge_error)?,
+        )
+        .map_err(bridge_error)?;
+
+        if messages.len() != response.choices.len() {
+            return Err(bridge_error(GatewayError::Bridge(format!(
+                "chat completion output guardrail rewrite expected {} messages, got {}",
+                response.choices.len(),
+                messages.len()
+            ))));
+        }
+
+        for (choice, message) in response.choices.iter_mut().zip(messages) {
+            choice.message = message;
+        }
+
+        Ok(())
+    }
+
     fn end_of_stream_event(saw_item: bool) -> Option<SseEvent> {
         saw_item.then(|| SseEvent::default().data("[DONE]"))
     }
 }
+
+fn bridge_error<E>(error: E) -> ChatCompletionError
+where
+    E: std::fmt::Display,
+{
+    ChatCompletionError::GatewayError(GatewayError::Bridge(error.to_string()))
+}

src/proxy/handlers/format_handler.rs

@@ -20,6 +20,7 @@ use tokio::{
 
 use crate::{
     config::entities::{Model, ResourceEntry},
+    guardrail::traits::{GuardrailCheckPayload, GuardrailOutcome},
     gateway::{
         error::GatewayError,
         traits::{ChatFormat, ProviderCapabilities},
@@ -30,6 +31,7 @@ use crate::{
     },
     proxy::{
         AppState,
+        guardrails::{ConfiguredGuardrailRuntime, resolve_model_guardrails},
         hooks::{
             self, RequestContext, authorization::AuthorizationError, rate_limit::RateLimitError,
         },
@@ -85,6 +87,36 @@ pub(crate) trait FormatHandlerAdapter: Send + Sync + 'static {
 
     fn output_message_span_properties(collector: &Self::Collector) -> Vec<(String, String)>;
 
+    fn guardrail_input_payload(
+        _lifecycle_state: &Self::LifecycleState,
+        _request: &Self::Request,
+    ) -> Result<Option<GuardrailCheckPayload>, Self::Error> {
+        Ok(None)
+    }
+
+    fn apply_input_guardrail_rewrite(
+        _lifecycle_state: &mut Self::LifecycleState,
+        _request: &mut Self::Request,
+        _rewrite: GuardrailCheckPayload,
+    ) -> Result<(), Self::Error> {
+        Ok(())
+    }
+
+    fn guardrail_output_payload(
+        _lifecycle_state: &Self::LifecycleState,
+        _response: &Self::Response,
+    ) -> Result<Option<GuardrailCheckPayload>, Self::Error> {
+        Ok(None)
+    }
+
+    fn apply_output_guardrail_rewrite(
+        _lifecycle_state: &mut Self::LifecycleState,
+        _response: &mut Self::Response,
+        _rewrite: GuardrailCheckPayload,
+    ) -> Result<(), Self::Error> {
+        Ok(())
+    }
+
     async fn prepare_lifecycle(
         _state: &AppState,
         _request_ctx: &mut RequestContext,
@@ -183,11 +215,19 @@ where
     let provider = model.provider(resources.as_ref()).ok_or_else(|| {
         GatewayError::Internal(format!("provider {} not found", model.provider_id))
     })?;
+    let configured_guardrails = resolve_model_guardrails(&model, resources.as_ref())?;
     let provider_instance = create_provider_instance(gateway.as_ref(), &provider)?;
     let provider_base_url = provider_instance.effective_base_url().ok();
     let mut lifecycle_state =
         A::prepare_lifecycle(&state, &mut request_ctx, &mut request_data).await?;
 
+    apply_input_guardrails::<A>(
+        &configured_guardrails,
+        &mut lifecycle_state,
+        &mut request_data,
+    )
+    .await?;
+
     let span = Span::enter_with_local_parent(A::span_name());
     apply_span_properties(
         &span,
@@ -214,6 +254,12 @@ where
             mut response,
             usage,
         })) => {
+            apply_output_guardrails::<A>(
+                &configured_guardrails,
+                &mut lifecycle_state,
+                &mut response,
+            )
+            .await?;
             A::handle_complete_response(
                 &state,
                 &mut request_ctx,
@@ -247,6 +293,76 @@ where
     }
 }
 
+async fn apply_input_guardrails<A>(
+    guardrails: &[Box<dyn ConfiguredGuardrailRuntime>],
+    lifecycle_state: &mut A::LifecycleState,
+    request: &mut AdapterRequest<A>,
+) -> Result<(), A::Error>
+where
+    A: FormatHandlerAdapter,
+{
+    for guardrail in guardrails {
+        let Some(payload) = A::guardrail_input_payload(lifecycle_state, request)? else {
+            continue;
+        };
+        let Some(outcome) = guardrail.check(&payload).await? else {
+            continue;
+        };
+
+        match outcome {
+            GuardrailOutcome::Allow => {}
+            GuardrailOutcome::Rewrite(rewrite) => {
+                A::apply_input_guardrail_rewrite(lifecycle_state, request, rewrite)?;
+            }
+            GuardrailOutcome::Block { reason } => {
+                return Err(GatewayError::Validation(format!(
+                    "guardrail {} blocked input: {}",
+                    guardrail.name(),
+                    reason
+                ))
+                .into());
+            }
+        }
+    }
+
+    Ok(())
+}
+
+async fn apply_output_guardrails<A>(
+    guardrails: &[Box<dyn ConfiguredGuardrailRuntime>],
+    lifecycle_state: &mut A::LifecycleState,
+    response: &mut AdapterResponse<A>,
+) -> Result<(), A::Error>
+where
+    A: FormatHandlerAdapter,
+{
+    for guardrail in guardrails {
+        let Some(payload) = A::guardrail_output_payload(lifecycle_state, response)? else {
+            continue;
+        };
+        let Some(outcome) = guardrail.check(&payload).await? else {
+            continue;
+        };
+
+        match outcome {
+            GuardrailOutcome::Allow => {}
+            GuardrailOutcome::Rewrite(rewrite) => {
+                A::apply_output_guardrail_rewrite(lifecycle_state, response, rewrite)?;
+            }
+            GuardrailOutcome::Block { reason } => {
+                return Err(GatewayError::Validation(format!(
+                    "guardrail {} blocked output: {}",
+                    guardrail.name(),
+                    reason
+                ))
+                .into());
+            }
+        }
+    }
+
+    Ok(())
+}
+
 async fn handle_regular_response<A>(
     response: AdapterResponse<A>,
     usage: Usage,