Skip to content

feat: expose event sanitizers in Python [3/6] #373

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
rapids-bot merged 4 commits into from
Jul 9, 2026
Merged

feat: expose event sanitizers in Python [3/6] #373

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
c37d47f
feat: expose event sanitizers in Python
willkill07 Jul 7, 2026
f224e4b
fix: address event sanitizer review feedback
willkill07 Jul 9, 2026
c7c3e5e
refactor: centralize event sanitizer handlers
willkill07 Jul 9, 2026
34283fc
Merge branch 'main' into feat/relay-409-python-event-sanitizers
willkill07 Jul 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
122 changes: 122 additions & 0 deletions crates/python/src/py_api/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -875,6 +875,44 @@ fn llm_stream_call_execute<'py>(
// Guardrail registrations (macro-generated)
// ---------------------------------------------------------------------------

macro_rules! py_event_guardrail_api {
Comment thread
willkill07 marked this conversation as resolved.
($register_name:ident, $deregister_name:ident, $core_register:path, $core_deregister:path) => {
#[pyfunction]
fn $register_name(name: &str, priority: i32, guardrail: Py<PyAny>) -> PyResult<()> {
$core_register(
name,
priority,
py_callable::wrap_py_event_sanitize_fn(guardrail),
)
.map_err(to_py_err)
}

#[pyfunction]
fn $deregister_name(name: &str) -> PyResult<bool> {
$core_deregister(name).map_err(to_py_err)
}
};
}

py_event_guardrail_api!(
register_mark_sanitize_guardrail,
deregister_mark_sanitize_guardrail,
core_registry_api::register_mark_sanitize_guardrail,
core_registry_api::deregister_mark_sanitize_guardrail
);
py_event_guardrail_api!(
register_scope_sanitize_start_guardrail,
deregister_scope_sanitize_start_guardrail,
core_registry_api::register_scope_sanitize_start_guardrail,
core_registry_api::deregister_scope_sanitize_start_guardrail
);
py_event_guardrail_api!(
register_scope_sanitize_end_guardrail,
deregister_scope_sanitize_end_guardrail,
core_registry_api::register_scope_sanitize_end_guardrail,
core_registry_api::deregister_scope_sanitize_end_guardrail
);

/// Macro that generates a register/deregister pair for tool guardrails
/// whose callback signature is `(tool_name: str, json: Any) -> Any`.
macro_rules! py_guardrail_tool_api {
Expand Down Expand Up @@ -1272,6 +1310,52 @@ fn parse_uuid(scope_uuid: &str) -> PyResult<Uuid> {
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(format!("invalid UUID: {e}")))
}

macro_rules! py_scope_event_guardrail_api {
($register_name:ident, $deregister_name:ident, $core_register:path, $core_deregister:path) => {
#[pyfunction]
fn $register_name(
scope_uuid: &str,
name: &str,
priority: i32,
guardrail: Py<PyAny>,
) -> PyResult<()> {
let uuid = parse_uuid(scope_uuid)?;
$core_register(
&uuid,
name,
priority,
py_callable::wrap_py_event_sanitize_fn(guardrail),
)
.map_err(to_py_err)
}

#[pyfunction]
fn $deregister_name(scope_uuid: &str, name: &str) -> PyResult<bool> {
let uuid = parse_uuid(scope_uuid)?;
$core_deregister(&uuid, name).map_err(to_py_err)
}
};
}

py_scope_event_guardrail_api!(
scope_register_mark_sanitize_guardrail,
scope_deregister_mark_sanitize_guardrail,
core_registry_api::scope_register_mark_sanitize_guardrail,
core_registry_api::scope_deregister_mark_sanitize_guardrail
);
py_scope_event_guardrail_api!(
scope_register_scope_sanitize_start_guardrail,
scope_deregister_scope_sanitize_start_guardrail,
core_registry_api::scope_register_scope_sanitize_start_guardrail,
core_registry_api::scope_deregister_scope_sanitize_start_guardrail
);
py_scope_event_guardrail_api!(
scope_register_scope_sanitize_end_guardrail,
scope_deregister_scope_sanitize_end_guardrail,
core_registry_api::scope_register_scope_sanitize_end_guardrail,
core_registry_api::scope_deregister_scope_sanitize_end_guardrail
);

/// Macro that generates a scope-local register/deregister pair for guardrails
/// whose callback signature is `(tool_name: str, json: Any) -> Any`.
macro_rules! py_scope_local_guardrail_tool_api {
Expand Down Expand Up @@ -1651,6 +1735,23 @@ pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
m
)?)?;

// Mark and scope event guardrails
m.add_function(wrap_pyfunction!(register_mark_sanitize_guardrail, m)?)?;
m.add_function(wrap_pyfunction!(deregister_mark_sanitize_guardrail, m)?)?;
m.add_function(wrap_pyfunction!(
register_scope_sanitize_start_guardrail,
m
)?)?;
m.add_function(wrap_pyfunction!(
deregister_scope_sanitize_start_guardrail,
m
)?)?;
m.add_function(wrap_pyfunction!(register_scope_sanitize_end_guardrail, m)?)?;
m.add_function(wrap_pyfunction!(
deregister_scope_sanitize_end_guardrail,
m
)?)?;

// Tool intercepts
m.add_function(wrap_pyfunction!(register_tool_request_intercept, m)?)?;
m.add_function(wrap_pyfunction!(deregister_tool_request_intercept, m)?)?;
Expand Down Expand Up @@ -1727,6 +1828,27 @@ pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
scope_deregister_tool_conditional_execution_guardrail,
m
)?)?;
m.add_function(wrap_pyfunction!(scope_register_mark_sanitize_guardrail, m)?)?;
m.add_function(wrap_pyfunction!(
scope_deregister_mark_sanitize_guardrail,
m
)?)?;
m.add_function(wrap_pyfunction!(
scope_register_scope_sanitize_start_guardrail,
m
)?)?;
m.add_function(wrap_pyfunction!(
scope_deregister_scope_sanitize_start_guardrail,
m
)?)?;
m.add_function(wrap_pyfunction!(
scope_register_scope_sanitize_end_guardrail,
m
)?)?;
m.add_function(wrap_pyfunction!(
scope_deregister_scope_sanitize_end_guardrail,
m
)?)?;

// Scope-local tool intercepts
m.add_function(wrap_pyfunction!(scope_register_tool_request_intercept, m)?)?;
Expand Down
69 changes: 65 additions & 4 deletions crates/python/src/py_callable.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,16 +25,16 @@ use std::pin::Pin;
use std::sync::Arc;

use nemo_relay::api::runtime::{
EventSubscriberFn, LlmConditionalFn, LlmExecutionNextFn, LlmRequestInterceptFn,
LlmSanitizeRequestFn, LlmSanitizeResponseFn, LlmStreamExecutionNextFn, ToolConditionalFn,
ToolExecutionNextFn, ToolInterceptFn, ToolSanitizeFn,
EventSanitizeFn, EventSubscriberFn, LlmConditionalFn, LlmExecutionNextFn,
LlmRequestInterceptFn, LlmSanitizeRequestFn, LlmSanitizeResponseFn, LlmStreamExecutionNextFn,
ToolConditionalFn, ToolExecutionNextFn, ToolInterceptFn, ToolSanitizeFn,
};
use nemo_relay::error::{FlowError, Result as FlowResult};
use pyo3::prelude::*;
use serde_json::Value as Json;
use tokio_stream::Stream;

use nemo_relay::api::event::Event;
use nemo_relay::api::event::{Event, EventSanitizeFields};
use nemo_relay::api::llm::{LlmRequest, LlmRequestInterceptOutcome};
use nemo_relay::api::tool::ToolExecutionInterceptOutcome;
use nemo_relay::codec::request::AnnotatedLlmRequest as AnnotatedLLMRequest;
Expand Down Expand Up @@ -854,6 +854,67 @@ pub fn wrap_py_event_subscriber(py_fn: Py<PyAny>) -> EventSubscriberFn {
})
}

/// Wrap a Python callable ``(Event, EventSanitizeFields) -> EventSanitizeFields``.
pub fn wrap_py_event_sanitize_fn(py_fn: Py<PyAny>) -> EventSanitizeFn {
Arc::new(move |event: &Event, fields: EventSanitizeFields| {
Python::attach(|py| {
let py_event = match event {
Event::Scope(inner) => Py::new(
py,
crate::py_types::PyScopeEvent {
inner: inner.clone(),
},
)
.map(|value| value.into_any()),
Event::Mark(inner) => Py::new(
py,
crate::py_types::PyMarkEvent {
inner: inner.clone(),
},
)
.map(|value| value.into_any()),
};
let py_event = match py_event {
Ok(value) => value,
Err(error) => {
eprintln!("nemo_relay: failed to convert event sanitizer context: {error}");
return fields.clone();
}
};
let fields_json = match serde_json::to_value(&fields) {
Ok(value) => value,
Err(error) => {
eprintln!("nemo_relay: failed to serialize event sanitizer fields: {error}");
return fields.clone();
}
};
let py_fields = match json_to_py(py, &fields_json) {
Ok(value) => value,
Err(error) => {
eprintln!("nemo_relay: failed to convert event sanitizer fields: {error}");
return fields.clone();
}
};
let result = match py_fn.call1(py, (py_event, py_fields)) {
Ok(value) => value,
Err(error) => {
eprintln!("nemo_relay: Python event sanitizer callable failed: {error}");
return fields.clone();
}
};
py_to_json(result.bind(py))
.ok()
.and_then(|value| serde_json::from_value(value).ok())
.unwrap_or_else(|| {
eprintln!(
"nemo_relay: event sanitizer must return data, category_profile, and metadata fields"
);
fields.clone()
})
})
})
}

// ---------------------------------------------------------------------------
// LLM Codec wrapper
// ---------------------------------------------------------------------------
Expand Down
Loading
Loading