From c6d6de95d41f81b1b329033664623905c654bc9f Mon Sep 17 00:00:00 2001 From: Gaurav Kumbhat Date: Tue, 4 Nov 2025 14:24:49 +0530 Subject: [PATCH 01/10] :sparkles: Initiate python orchestrator interfaces Signed-off-by: Gaurav Kumbhat --- Cargo.lock | 135 +++++++++++++++++++++++++++++++- Cargo.toml | 5 ++ examples/python_orchestrator.py | 57 ++++++++++++++ pyproject.toml | 16 ++++ src/clients/detector.rs | 3 + src/lib.rs | 23 ++++++ src/models.rs | 8 ++ src/orchestrator.rs | 33 ++++++++ 8 files changed, 277 insertions(+), 3 deletions(-) create mode 100644 examples/python_orchestrator.py create mode 100644 pyproject.toml diff --git a/Cargo.lock b/Cargo.lock index 69195e60f..af9782f66 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -470,6 +470,9 @@ dependencies = [ "opentelemetry_sdk", "pin-project-lite", "prost", + "pyo3", + "pyo3-async-runtimes", + "pythonize", "rand", "reqwest", "rustls", @@ -901,7 +904,7 @@ dependencies = [ "libc", "percent-encoding", "pin-project-lite", - "socket2 0.5.10", + "socket2 0.6.0", "system-configuration", "tokio", "tower-service", @@ -1026,6 +1029,15 @@ dependencies = [ "hashbrown 0.16.0", ] +[[package]] +name = "indoc" +version = "2.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "79cf5c93f93228cf8efb3ba362535fb11199ac548a09ce117c9b1adc3030d706" +dependencies = [ + "rustversion", +] + [[package]] name = "ipconfig" version = "0.3.2" @@ -1175,6 +1187,15 @@ version = "2.7.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f52b00d39961fc5b2736ea853c9cc86238e165017a493d1d5c8eac6bdc4cc273" +[[package]] +name = "memoffset" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a" +dependencies = [ + "autocfg", +] + [[package]] name = "mime" version = "0.3.17" @@ -1613,6 +1634,102 @@ dependencies = [ "pulldown-cmark", ] +[[package]] +name = "pyo3" +version = "0.26.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ba0117f4212101ee6544044dae45abe1083d30ce7b29c4b5cbdfa2354e07383" +dependencies = [ + "indoc", + "libc", + "memoffset", + "once_cell", + "portable-atomic", + "pyo3-build-config", + "pyo3-ffi", + "pyo3-macros", + "unindent", +] + +[[package]] +name = "pyo3-async-runtimes" +version = "0.26.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6ee6d4cb3e8d5b925f5cdb38da183e0ff18122eb2048d4041c9e7034d026e23" +dependencies = [ + "futures", + "once_cell", + "pin-project-lite", + "pyo3", + "pyo3-async-runtimes-macros", + "tokio", +] + +[[package]] +name = "pyo3-async-runtimes-macros" +version = "0.26.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c29bc5c673e36a8102d0b9179149c1bb59990d8db4f3ae58bd7dceccab90b951" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "pyo3-build-config" +version = "0.26.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fc6ddaf24947d12a9aa31ac65431fb1b851b8f4365426e182901eabfb87df5f" +dependencies = [ + "target-lexicon", +] + +[[package]] +name = "pyo3-ffi" +version = "0.26.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "025474d3928738efb38ac36d4744a74a400c901c7596199e20e45d98eb194105" +dependencies = [ + "libc", + "pyo3-build-config", +] + +[[package]] +name = "pyo3-macros" +version = "0.26.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e64eb489f22fe1c95911b77c44cc41e7c19f3082fc81cce90f657cdc42ffded" +dependencies = [ + "proc-macro2", + "pyo3-macros-backend", + "quote", + "syn", +] + +[[package]] +name = "pyo3-macros-backend" +version = "0.26.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "100246c0ecf400b475341b8455a9213344569af29a3c841d29270e53102e0fcf" +dependencies = [ + "heck", + "proc-macro2", + "pyo3-build-config", + "quote", + "syn", +] + +[[package]] +name = "pythonize" +version = "0.26.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11e06e4cff9be2bbf2bddf28a486ae619172ea57e79787f856572878c62dcfe2" +dependencies = [ + "pyo3", + "serde", +] + [[package]] name = "quinn" version = "0.11.9" @@ -1626,7 +1743,7 @@ dependencies = [ "quinn-udp", "rustc-hash", "rustls", - "socket2 0.5.10", + "socket2 0.6.0", "thiserror 2.0.17", "tokio", "tracing", @@ -1663,7 +1780,7 @@ dependencies = [ "cfg_aliases", "libc", "once_cell", - "socket2 0.5.10", + "socket2 0.6.0", "tracing", "windows-sys 0.60.2", ] @@ -2192,6 +2309,12 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7b2093cf4c8eb1e67749a6762251bc9cd836b6fc171623bd0a9d324d37af2417" +[[package]] +name = "target-lexicon" +version = "0.13.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df7f62577c25e07834649fc3b39fafdc597c0a3527dc1c60129201ccfcbaa50c" + [[package]] name = "tempfile" version = "3.23.0" @@ -2607,6 +2730,12 @@ version = "1.0.19" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f63a545481291138910575129486daeaf8ac54aee4387fe7906919f7830c7d9d" +[[package]] +name = "unindent" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7264e107f553ccae879d21fbea1d6724ac785e8c3bfc762137959b5802826ef3" + [[package]] name = "untrusted" version = "0.9.0" diff --git a/Cargo.toml b/Cargo.toml index 460f7993a..5cb478a70 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -8,6 +8,8 @@ build = "build.rs" [lib] path = "src/lib.rs" +# "cdylib" is necessary to produce a shared library for Python to import from. +crate-type = ["cdylib"] [[bin]] name = "fms-guardrails-orchestr8" @@ -53,6 +55,9 @@ opentelemetry-otlp = { version = "0.31.0", features = [ opentelemetry_sdk = { version = "0.31.0", features = ["rt-tokio", "metrics"] } pin-project-lite = "0.2.16" prost = "0.14.1" +pyo3 = { version = "0.26.0", features = ["extension-module"]} +pyo3-async-runtimes = { version = "0.26.0", features = ["attributes", "tokio-runtime"] } +pythonize = "0.26.0" reqwest = { version = "0.12.24", features = [ "blocking", "rustls-tls", diff --git a/examples/python_orchestrator.py b/examples/python_orchestrator.py new file mode 100644 index 000000000..5b030116d --- /dev/null +++ b/examples/python_orchestrator.py @@ -0,0 +1,57 @@ +import asyncio +import logging + +from fms_guardrails_orchestr8 import ( + get_guardrails_orchestrator, + ContextDocsHttpRequest, + ContextType, + TextContentDetectionHttpRequest, +) + + +FORMAT = '%(levelname)s %(name)s %(asctime)-15s %(filename)s:%(lineno)d %(message)s' +logging.basicConfig(format=FORMAT) +logging.getLogger().setLevel(logging.DEBUG) + +CONFIG_FILE = "config/test-guardrails-orchestrator.yaml" + +# Showing sync initialization +# orch8 = GuardrailsOrchestrator(config_path=CONFIG_FILE) + +async def detect_content(): + # Showing async initialization + orch8 = await get_guardrails_orchestrator(config_path=CONFIG_FILE, start_up_health_check=False) + + request = TextContentDetectionHttpRequest( + content="This is stupid text.", + detectors= { + "en_syntax_slate.38m.hap": { + "foo": "bar" + } + } + ) + + + result = await orch8.detection_content(request) + print(result) + + +async def detect_context(): + request = ContextDocsHttpRequest( + content="This is a good document", + context_type=PyContextType.DOCUMENT, + context=["Document 1", "Document 2", "Document 3"], + detectors={ + "granite-guardian-context": { + "risk_name": "context_relevance" + } + } + ) + result = await orch8.detect_context_documents(request) + print(result) + + + +if __name__ == "__main__": + asyncio.run(detect_content()) + # asyncio.run(detect_context()) diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 000000000..d20a8bc16 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,16 @@ +[build-system] +requires = ["maturin>=1.8,<2.0"] +build-backend = "maturin" + +[project] +name = "fms_guardrails_orchestr8" +requires-python = ">=3.10,<3.15" +classifiers = [ + "Programming Language :: Rust", + "Programming Language :: Python :: Implementation :: CPython", + "Programming Language :: Python :: Implementation :: PyPy", +] +dynamic = ["version"] + +[tool.maturin] +features = ["pyo3/extension-module"] diff --git a/src/clients/detector.rs b/src/clients/detector.rs index 840f42896..007417e15 100644 --- a/src/clients/detector.rs +++ b/src/clients/detector.rs @@ -39,6 +39,8 @@ use crate::{ models::{DetectionResult, DetectorParams, EvidenceObj, Metadata}, }; +use pyo3::pyclass; + pub const DEFAULT_PORT: u16 = 8080; pub const MODEL_HEADER_NAME: &str = "x-model-name"; pub const DETECTOR_ID_HEADER_NAME: &str = "detector-id"; @@ -271,6 +273,7 @@ pub struct ContextDocsDetectionRequest { } /// Enum representing the context type of a detection +#[pyclass] #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] pub enum ContextType { #[serde(rename = "docs")] diff --git a/src/lib.rs b/src/lib.rs index fc60b477d..239743abc 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -28,3 +28,26 @@ pub mod server; pub mod utils; #[allow(unused_imports)] pub(crate) use utils::AsUriExt; + +use pyo3::prelude::*; + +use crate::{models::{DetectorParams, TextContentDetectionHttpRequest}, orchestrator::Orchestrator}; + + +/// A Python module implemented in Rust. +#[pymodule] +fn fms_guardrails_orchestr8(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { + + // TODO: Replace this with proper implementation + let _ = rustls::crypto::ring::default_provider().install_default(); + + // m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_function(wrap_pyfunction!(orchestrator::get_guardrails_orchestrator, m)?)?; + Ok(()) +} \ No newline at end of file diff --git a/src/models.rs b/src/models.rs index f15e78722..af0ec9236 100644 --- a/src/models.rs +++ b/src/models.rs @@ -20,6 +20,9 @@ use std::collections::{BTreeMap, HashMap}; use serde::{Deserialize, Serialize}; +use pyo3::pyclass; +use pyo3::types::PyDict; + use crate::{ clients::{ @@ -46,6 +49,7 @@ pub struct InfoParams { } /// Parameters relevant to each detector +#[pyclass] #[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] pub struct DetectorParams(BTreeMap); @@ -367,6 +371,7 @@ pub struct ClassifiedGeneratedTextResult { } /// The request format expected in the /api/v2/text/detection/content endpoint. +#[pyclass] #[derive(Clone, Debug, Serialize, Deserialize)] #[serde(deny_unknown_fields)] pub struct TextContentDetectionHttpRequest { @@ -396,6 +401,7 @@ impl TextContentDetectionHttpRequest { } /// The response format of the /api/v2/text/detection/content endpoint +#[pyclass(extends=PyDict)] #[derive(Default, Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct TextContentDetectionResult { /// Detection results @@ -977,6 +983,7 @@ pub enum GuardrailDetection { } /// The request format expected in the /api/v2/text/context endpoint. +#[pyclass] #[derive(Clone, Debug, Serialize, Deserialize)] #[serde(deny_unknown_fields)] pub struct ContextDocsHttpRequest { @@ -1015,6 +1022,7 @@ impl ContextDocsHttpRequest { } /// The response format of the /api/v1/text/task/generation-detection endpoint +#[pyclass] #[derive(Default, Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct ContextDocsResult { pub detections: Vec, diff --git a/src/orchestrator.rs b/src/orchestrator.rs index 19ff86ccd..7360063e8 100644 --- a/src/orchestrator.rs +++ b/src/orchestrator.rs @@ -34,6 +34,11 @@ use crate::{ health::{HealthCheckResult, HealthStatus}, }; +use pyo3::exceptions::{PyOSError}; +use pyo3::prelude::*; +use pyo3::pyclass; +use pyo3_async_runtimes::tokio::future_into_py; + const DEFAULT_MAX_RETRIES: usize = 3; #[cfg_attr(test, derive(Default))] @@ -49,6 +54,7 @@ impl Context { } /// Handles orchestrator tasks. +#[pyclass] #[cfg_attr(test, derive(Default))] pub struct Orchestrator { ctx: Arc, @@ -154,3 +160,30 @@ async fn create_clients(config: &OrchestratorConfig) -> Result } Ok(clients) } + +// Async Factory function to create a guardrails orchestrator +#[pyfunction] +#[pyo3(text_signature = "(config_path=config.yaml, start_up_health_check=True)")] +pub fn get_guardrails_orchestrator<'py>(py: Python<'py>, config_path: String, start_up_health_check: bool) -> PyResult> { + + let guardrails_orch8_future = async move { + let config_result = OrchestratorConfig::load(config_path).await; + + let config = match config_result { + Ok(config) => config, + Err(err) => return Err(PyOSError::new_err(format!("Error loading orchestrator configuration: {}", err))) + }; + + let orchestrator_result = Orchestrator::new(config, start_up_health_check).await; + + match orchestrator_result { + Ok(orchestrator) => Ok(orchestrator), + Err(err) => Err(PyOSError::new_err(format!("Error creating orchestrator: {}", err))) + } + }; + + // Convert the Future into a Python awaitable + future_into_py(py, guardrails_orch8_future) + +} + From 9363729552c272528756aefd4cd1c81270d2ac02 Mon Sep 17 00:00:00 2001 From: Gaurav Kumbhat Date: Wed, 3 Dec 2025 08:55:22 -0600 Subject: [PATCH 02/10] :construction_worker: Update interfaces Signed-off-by: Gaurav Kumbhat --- examples/python_orchestrator.py | 25 +++++---- src/lib.rs | 17 ++++-- src/models.rs | 97 ++++++++++++++++++++++++++++++++- src/orchestrator.rs | 59 ++++++++++++++++++++ 4 files changed, 180 insertions(+), 18 deletions(-) diff --git a/examples/python_orchestrator.py b/examples/python_orchestrator.py index 5b030116d..0fbb65710 100644 --- a/examples/python_orchestrator.py +++ b/examples/python_orchestrator.py @@ -3,9 +3,9 @@ from fms_guardrails_orchestr8 import ( get_guardrails_orchestrator, - ContextDocsHttpRequest, + PyContextDocsHttpRequest, ContextType, - TextContentDetectionHttpRequest, + PyTextContentDetectionHttpRequest, ) @@ -22,22 +22,25 @@ async def detect_content(): # Showing async initialization orch8 = await get_guardrails_orchestrator(config_path=CONFIG_FILE, start_up_health_check=False) - request = TextContentDetectionHttpRequest( - content="This is stupid text.", - detectors= { - "en_syntax_slate.38m.hap": { - "foo": "bar" + try: + request = PyTextContentDetectionHttpRequest( + content="This is stupid text.", + detectors= { + "en_syntax_slate.38m.hap": { + "foo": "bar" + } } - } - ) - + ) + except Exception as ex: + print(ex) + raise ex result = await orch8.detection_content(request) print(result) async def detect_context(): - request = ContextDocsHttpRequest( + request = PyContextDocsHttpRequest( content="This is a good document", context_type=PyContextType.DOCUMENT, context=["Document 1", "Document 2", "Document 3"], diff --git a/src/lib.rs b/src/lib.rs index 239743abc..cfe5d7723 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -31,7 +31,13 @@ pub(crate) use utils::AsUriExt; use pyo3::prelude::*; -use crate::{models::{DetectorParams, TextContentDetectionHttpRequest}, orchestrator::Orchestrator}; +use crate::models::{ + ContextDocsResult, + DetectorParams, + PyTextContentDetectionHttpRequest, + PyContextDocsHttpRequest, + TextContentDetectionResult +}; /// A Python module implemented in Rust. @@ -41,13 +47,12 @@ fn fms_guardrails_orchestr8(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<() // TODO: Replace this with proper implementation let _ = rustls::crypto::ring::default_provider().install_default(); - // m.add_class::()?; - m.add_class::()?; + m.add_class::()?; m.add_class::()?; - m.add_class::()?; + m.add_class::()?; m.add_class::()?; - m.add_class::()?; - m.add_class::()?; + m.add_class::()?; + m.add_class::()?; m.add_function(wrap_pyfunction!(orchestrator::get_guardrails_orchestrator, m)?)?; Ok(()) } \ No newline at end of file diff --git a/src/models.rs b/src/models.rs index af0ec9236..5ed31fe9d 100644 --- a/src/models.rs +++ b/src/models.rs @@ -18,10 +18,14 @@ #![allow(unused_qualifications)] use std::collections::{BTreeMap, HashMap}; +use std::sync::Arc; use serde::{Deserialize, Serialize}; +use pyo3::prelude::*; use pyo3::pyclass; use pyo3::types::PyDict; +use pyo3::conversion::{FromPyObject}; +use pythonize::depythonize; use crate::{ @@ -370,8 +374,8 @@ pub struct ClassifiedGeneratedTextResult { pub input_tokens: Option>, } + /// The request format expected in the /api/v2/text/detection/content endpoint. -#[pyclass] #[derive(Clone, Debug, Serialize, Deserialize)] #[serde(deny_unknown_fields)] pub struct TextContentDetectionHttpRequest { @@ -1435,3 +1439,94 @@ mod tests { Ok(()) } } + + +/// Python Interfaces + +#[pyclass] +#[derive(Serialize, Deserialize)] +pub struct PyDetectorsObj(HashMap); + +// NOTE: If this moves to orchestrator object itself, we don't need to do double conversion +// and we can implement the trait directly there +impl<'py> FromPyObject<'py> for PyDetectorsObj { + fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult { + let dict = ob.downcast::()?; + + let mut result = HashMap::new(); + + for (key, value) in dict.iter() { + let key = key.extract::()?; + let inner_dict = value.downcast::()?; + // let mut inner_result = BTreeMap::new(); + let mut inner_result = DetectorParams::new(); + + for (inner_key, inner_value) in inner_dict.iter() { + let inner_key = inner_key.extract::()?; + let inner_value = depythonize(&inner_value).unwrap(); + inner_result.insert(inner_key, inner_value); + } + result.insert(key, inner_result); + } + Ok(PyDetectorsObj(result)) + } +} + +// Into trait for PyDetectorsObj which translates to +// HashMap +impl Into> for PyDetectorsObj { + fn into(self) -> HashMap { + self.0 + } +} + +impl From<&PyDetectorsObj> for HashMap { + fn from(obj: &PyDetectorsObj) -> Self { + obj.0.clone() + } +} + +#[pyclass] +#[derive(Clone)] +pub struct PyTextContentDetectionHttpRequest { + #[pyo3(get, set)] + pub content: String, + // pub detectors: Arc>, + pub detectors: Arc, +} + +#[pymethods] +impl PyTextContentDetectionHttpRequest { + #[new] + fn new(content: String, detectors: PyDetectorsObj) -> Self { + PyTextContentDetectionHttpRequest { + content, + // detectors: Arc::new(Mutex::new(detectors)), // Wrap HashMap in Arc and Mutex for thread-safe access + detectors: Arc::new(detectors) + } + } + +} + +#[pyclass] +#[derive(Clone)] +pub struct PyContextDocsHttpRequest{ + /// The map of detectors to be used, along with their respective parameters, e.g. thresholds. + pub detectors: Arc, + pub content: String, + pub context_type: ContextType, + pub context: Vec, +} + +#[pymethods] +impl PyContextDocsHttpRequest { + #[new] + fn new(content: String, context_type: ContextType, context: Vec, detectors: PyDetectorsObj) -> Self { + Self { + content, + context, + context_type: context_type.into(), + detectors: detectors.into() + } + } +} \ No newline at end of file diff --git a/src/orchestrator.rs b/src/orchestrator.rs index 7360063e8..35c42fb38 100644 --- a/src/orchestrator.rs +++ b/src/orchestrator.rs @@ -31,7 +31,10 @@ use crate::{ openai::OpenAiClient, }, config::{GenerationProvider, OrchestratorConfig}, + orchestrator::handlers::{TextContentDetectionTask}, health::{HealthCheckResult, HealthStatus}, + models::{PyTextContentDetectionHttpRequest, TextContentDetectionHttpRequest}, + utils::trace }; use pyo3::exceptions::{PyOSError}; @@ -111,6 +114,62 @@ impl Orchestrator { } health } + + fn detection_content<'py>(&self, py: Python<'py>, request: PyTextContentDetectionHttpRequest) -> PyResult> { + let trace_id = trace::current_trace_id(); + // TODO: Replace this with real headers input + let headers = HeaderMap::new(); + let detectors = &request.detectors; + + let content = request.content; + let request = TextContentDetectionHttpRequest { + content, + detectors: detectors.as_ref().into(), + }; + + // TODO: Add request validation here + + let task = TextContentDetectionTask::new(trace_id, request, headers); + + let guardrails_orch = Arc::clone(self); + + pyo3_async_runtimes::tokio::future_into_py(py, async move { + match guardrails_orch.handle(task).await { + Ok(response) => { + // Left it as an example. + // let serialized_response = serde_json::to_string(&response).unwrap(); + // Ok(serialized_response) + Ok(PyTextContentDetectionResult(response)) + }, + // TODO: Handle errors properly with correct types + Err(error) => Err(PyTypeError::new_err(error.to_string())), + } + }) + + } + + + fn detect_context_documents<'py>(&self, py: Python<'py>, request: PyContextDocsHttpRequest) -> PyResult> { + let trace_id = trace::current_trace_id(); + // TODO: Replace this with real headers input + let headers = HeaderMap::new(); + + let request: ContextDocsHttpRequest = request.into(); + + let task = ContextDocsDetectionTask::new(trace_id, request, headers); + let guardrails_orch = Arc::clone(&self.orchestrator); + + pyo3_async_runtimes::tokio::future_into_py(py, async move { + match guardrails_orch.handle(task).await { + Ok(response) => { + Ok(PyContextDocsResult(response)) + }, + // TODO: Handle errors properly with correct types + Err(error) => Err(PyTypeError::new_err(error.to_string())), + } + }) + } + } async fn create_clients(config: &OrchestratorConfig) -> Result { From db0d272dbcc0a686c095dddd28d6a19110ed248d Mon Sep 17 00:00:00 2001 From: gkumbhat Date: Mon, 22 Dec 2025 13:07:24 -0600 Subject: [PATCH 03/10] :construction_worker: WIP Signed-off-by: gkumbhat --- src/lib.rs | 54 +++++----- src/orchestrator.rs | 102 ++++-------------- src/orchestrator/python_interface.rs | 3 + .../python_interface/orchestrator.rs | 50 +++++++++ 4 files changed, 99 insertions(+), 110 deletions(-) create mode 100644 src/orchestrator/python_interface.rs create mode 100644 src/orchestrator/python_interface/orchestrator.rs diff --git a/src/lib.rs b/src/lib.rs index cfe5d7723..23a6bb0a3 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -29,30 +29,30 @@ pub mod utils; #[allow(unused_imports)] pub(crate) use utils::AsUriExt; -use pyo3::prelude::*; - -use crate::models::{ - ContextDocsResult, - DetectorParams, - PyTextContentDetectionHttpRequest, - PyContextDocsHttpRequest, - TextContentDetectionResult -}; - - -/// A Python module implemented in Rust. -#[pymodule] -fn fms_guardrails_orchestr8(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { - - // TODO: Replace this with proper implementation - let _ = rustls::crypto::ring::default_provider().install_default(); - - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - m.add_function(wrap_pyfunction!(orchestrator::get_guardrails_orchestrator, m)?)?; - Ok(()) -} \ No newline at end of file +// use pyo3::prelude::*; + +// use crate::models::{ +// ContextDocsResult, +// DetectorParams, +// PyTextContentDetectionHttpRequest, +// PyContextDocsHttpRequest, +// TextContentDetectionResult +// }; + + +// /// A Python module implemented in Rust. +// #[pymodule] +// fn fms_guardrails_orchestr8(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { + +// // TODO: Replace this with proper implementation +// let _ = rustls::crypto::ring::default_provider().install_default(); + +// m.add_class::()?; +// m.add_class::()?; +// m.add_class::()?; +// m.add_class::()?; +// m.add_class::()?; +// m.add_class::()?; +// // m.add_function(wrap_pyfunction!(orchestrator::get_guardrails_orchestrator, m)?)?; +// Ok(()) +// } \ No newline at end of file diff --git a/src/orchestrator.rs b/src/orchestrator.rs index 35c42fb38..571c2e00d 100644 --- a/src/orchestrator.rs +++ b/src/orchestrator.rs @@ -31,17 +31,9 @@ use crate::{ openai::OpenAiClient, }, config::{GenerationProvider, OrchestratorConfig}, - orchestrator::handlers::{TextContentDetectionTask}, health::{HealthCheckResult, HealthStatus}, - models::{PyTextContentDetectionHttpRequest, TextContentDetectionHttpRequest}, - utils::trace }; -use pyo3::exceptions::{PyOSError}; -use pyo3::prelude::*; -use pyo3::pyclass; -use pyo3_async_runtimes::tokio::future_into_py; - const DEFAULT_MAX_RETRIES: usize = 3; #[cfg_attr(test, derive(Default))] @@ -57,7 +49,6 @@ impl Context { } /// Handles orchestrator tasks. -#[pyclass] #[cfg_attr(test, derive(Default))] pub struct Orchestrator { ctx: Arc, @@ -115,61 +106,6 @@ impl Orchestrator { health } - fn detection_content<'py>(&self, py: Python<'py>, request: PyTextContentDetectionHttpRequest) -> PyResult> { - let trace_id = trace::current_trace_id(); - // TODO: Replace this with real headers input - let headers = HeaderMap::new(); - let detectors = &request.detectors; - - let content = request.content; - let request = TextContentDetectionHttpRequest { - content, - detectors: detectors.as_ref().into(), - }; - - // TODO: Add request validation here - - let task = TextContentDetectionTask::new(trace_id, request, headers); - - let guardrails_orch = Arc::clone(self); - - pyo3_async_runtimes::tokio::future_into_py(py, async move { - match guardrails_orch.handle(task).await { - Ok(response) => { - // Left it as an example. - // let serialized_response = serde_json::to_string(&response).unwrap(); - // Ok(serialized_response) - Ok(PyTextContentDetectionResult(response)) - }, - // TODO: Handle errors properly with correct types - Err(error) => Err(PyTypeError::new_err(error.to_string())), - } - }) - - } - - - fn detect_context_documents<'py>(&self, py: Python<'py>, request: PyContextDocsHttpRequest) -> PyResult> { - let trace_id = trace::current_trace_id(); - // TODO: Replace this with real headers input - let headers = HeaderMap::new(); - - let request: ContextDocsHttpRequest = request.into(); - - let task = ContextDocsDetectionTask::new(trace_id, request, headers); - let guardrails_orch = Arc::clone(&self.orchestrator); - - pyo3_async_runtimes::tokio::future_into_py(py, async move { - match guardrails_orch.handle(task).await { - Ok(response) => { - Ok(PyContextDocsResult(response)) - }, - // TODO: Handle errors properly with correct types - Err(error) => Err(PyTypeError::new_err(error.to_string())), - } - }) - } - } async fn create_clients(config: &OrchestratorConfig) -> Result { @@ -220,29 +156,29 @@ async fn create_clients(config: &OrchestratorConfig) -> Result Ok(clients) } -// Async Factory function to create a guardrails orchestrator -#[pyfunction] -#[pyo3(text_signature = "(config_path=config.yaml, start_up_health_check=True)")] -pub fn get_guardrails_orchestrator<'py>(py: Python<'py>, config_path: String, start_up_health_check: bool) -> PyResult> { +// // Async Factory function to create a guardrails orchestrator +// #[pyfunction] +// #[pyo3(text_signature = "(config_path=config.yaml, start_up_health_check=True)")] +// pub fn get_guardrails_orchestrator<'py>(py: Python<'py>, config_path: String, start_up_health_check: bool) -> PyResult> { - let guardrails_orch8_future = async move { - let config_result = OrchestratorConfig::load(config_path).await; +// let guardrails_orch8_future = async move { +// let config_result = OrchestratorConfig::load(config_path).await; - let config = match config_result { - Ok(config) => config, - Err(err) => return Err(PyOSError::new_err(format!("Error loading orchestrator configuration: {}", err))) - }; +// let config = match config_result { +// Ok(config) => config, +// Err(err) => return Err(PyOSError::new_err(format!("Error loading orchestrator configuration: {}", err))) +// }; - let orchestrator_result = Orchestrator::new(config, start_up_health_check).await; +// let orchestrator_result = Orchestrator::new(config, start_up_health_check).await; - match orchestrator_result { - Ok(orchestrator) => Ok(orchestrator), - Err(err) => Err(PyOSError::new_err(format!("Error creating orchestrator: {}", err))) - } - }; +// match orchestrator_result { +// Ok(orchestrator) => Ok(orchestrator), +// Err(err) => Err(PyOSError::new_err(format!("Error creating orchestrator: {}", err))) +// } +// }; - // Convert the Future into a Python awaitable - future_into_py(py, guardrails_orch8_future) +// // Convert the Future into a Python awaitable +// future_into_py(py, guardrails_orch8_future) -} +// } diff --git a/src/orchestrator/python_interface.rs b/src/orchestrator/python_interface.rs new file mode 100644 index 000000000..7b08c16f8 --- /dev/null +++ b/src/orchestrator/python_interface.rs @@ -0,0 +1,3 @@ + +// pub mod python_interface; +// pub use python_interface::orchestrator::*; \ No newline at end of file diff --git a/src/orchestrator/python_interface/orchestrator.rs b/src/orchestrator/python_interface/orchestrator.rs new file mode 100644 index 000000000..3eefe7d1e --- /dev/null +++ b/src/orchestrator/python_interface/orchestrator.rs @@ -0,0 +1,50 @@ + + +use std::sync::Arc; + +use pyo3::prelude::*; +use pyo3::exceptions::{PyOSError, PyTypeError}; +use pyo3::PyErr; +use pyo3::types::PyDict; +use pyo3::conversion::{FromPyObject}; + +use crate::orchestrator::Orchestrator; + +use crate::config::OrchestratorConfig; + +// Define orchestrator +#[pyclass] +struct GuardrailsOrchestrator { + // Added under Arc to allow safe sharing from multiple threads + // If we need modification capability later, then we might want to make it mutable + orchestrator: Arc, +} + +#[pymethods] +impl GuardrailsOrchestrator { + + #[new] + fn new(config_path: String) -> PyResult { + Python::with_gil(|py| { + pyo3_async_runtimes::tokio::run(py, async move { + let config_result = OrchestratorConfig::load(config_path).await; + + let config = match config_result { + Ok(config) => config, + Err(err) => return Err(PyOSError::new_err(format!("Error loading orchestrator configuration: {}", err))) + }; + + let orchestrator_result = Orchestrator::new(config, STARTUP_HEALTHCHECK).await; + + let orchestrator = match orchestrator_result { + Ok(orchestrator) => orchestrator, + Err(err) => return Err(PyOSError::new_err(format!("Error creating orchestrator: {}", err))) + }; + + Ok(GuardrailsOrchestrator { + orchestrator: Arc::new(orchestrator), + }) + + })}) + } +} \ No newline at end of file From 318c9b360aa0d10c44229e5ce4d0f69732d6e450 Mon Sep 17 00:00:00 2001 From: gkumbhat Date: Mon, 22 Dec 2025 16:12:12 -0600 Subject: [PATCH 04/10] :recycle: Refactor to separate out all python interfaces to separate module and cleanup the code Signed-off-by: gkumbhat --- examples/python_orchestrator.py | 46 +++++----- src/lib.rs | 52 ++++++------ src/models.rs | 57 ++----------- src/orchestrator.rs | 1 + src/orchestrator/python_interface.rs | 7 +- src/orchestrator/python_interface/models.rs | 83 +++++++++++++++++++ .../python_interface/orchestrator.rs | 50 +++++++++-- 7 files changed, 187 insertions(+), 109 deletions(-) create mode 100644 src/orchestrator/python_interface/models.rs diff --git a/examples/python_orchestrator.py b/examples/python_orchestrator.py index 0fbb65710..52b144298 100644 --- a/examples/python_orchestrator.py +++ b/examples/python_orchestrator.py @@ -2,10 +2,10 @@ import logging from fms_guardrails_orchestr8 import ( - get_guardrails_orchestrator, - PyContextDocsHttpRequest, - ContextType, - PyTextContentDetectionHttpRequest, + GuardrailsOrchestrator, + # DetectorParams, + TextContentDetectionRequest, + TextContentDetectionResult, ) @@ -13,45 +13,45 @@ logging.basicConfig(format=FORMAT) logging.getLogger().setLevel(logging.DEBUG) -CONFIG_FILE = "config/test-guardrails-orchestrator.yaml" +CONFIG_FILE = "config/local_config.yaml" # Showing sync initialization -# orch8 = GuardrailsOrchestrator(config_path=CONFIG_FILE) +orch8 = GuardrailsOrchestrator(config_path=CONFIG_FILE, start_up_health_check=False) async def detect_content(): # Showing async initialization - orch8 = await get_guardrails_orchestrator(config_path=CONFIG_FILE, start_up_health_check=False) + # orch8 = await get_guardrails_orchestrator(config_path=CONFIG_FILE, start_up_health_check=False) try: - request = PyTextContentDetectionHttpRequest( + request = TextContentDetectionRequest( content="This is stupid text.", detectors= { "en_syntax_slate.38m.hap": { - "foo": "bar" } } ) + except Exception as ex: print(ex) raise ex - result = await orch8.detection_content(request) + result = await orch8.content_detection(request) print(result) -async def detect_context(): - request = PyContextDocsHttpRequest( - content="This is a good document", - context_type=PyContextType.DOCUMENT, - context=["Document 1", "Document 2", "Document 3"], - detectors={ - "granite-guardian-context": { - "risk_name": "context_relevance" - } - } - ) - result = await orch8.detect_context_documents(request) - print(result) +# async def detect_context(): +# request = PyContextDocsHttpRequest( +# content="This is a good document", +# context_type=PyContextType.DOCUMENT, +# context=["Document 1", "Document 2", "Document 3"], +# detectors={ +# "granite-guardian-context": { +# "risk_name": "context_relevance" +# } +# } +# ) +# result = await orch8.detect_context_documents(request) +# print(result) diff --git a/src/lib.rs b/src/lib.rs index 23a6bb0a3..c2c7b43c5 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -29,30 +29,28 @@ pub mod utils; #[allow(unused_imports)] pub(crate) use utils::AsUriExt; -// use pyo3::prelude::*; - -// use crate::models::{ -// ContextDocsResult, -// DetectorParams, -// PyTextContentDetectionHttpRequest, -// PyContextDocsHttpRequest, -// TextContentDetectionResult -// }; - - -// /// A Python module implemented in Rust. -// #[pymodule] -// fn fms_guardrails_orchestr8(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { - -// // TODO: Replace this with proper implementation -// let _ = rustls::crypto::ring::default_provider().install_default(); - -// m.add_class::()?; -// m.add_class::()?; -// m.add_class::()?; -// m.add_class::()?; -// m.add_class::()?; -// m.add_class::()?; -// // m.add_function(wrap_pyfunction!(orchestrator::get_guardrails_orchestrator, m)?)?; -// Ok(()) -// } \ No newline at end of file +use pyo3::prelude::*; + +use crate::models::{ + ContextDocsResult, + ContextDocsHttpRequest, + TextContentDetectionHttpRequest, + TextContentDetectionResult +}; + + +/// A Python module implemented in Rust. +#[pymodule] +fn fms_guardrails_orchestr8(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { + + // TODO: Replace this with proper implementation + let _ = rustls::crypto::ring::default_provider().install_default(); + + m.add_class::()?; + m.add_class::()?; // TextContentDetectionRequest + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + Ok(()) +} \ No newline at end of file diff --git a/src/models.rs b/src/models.rs index 5ed31fe9d..25ad8f476 100644 --- a/src/models.rs +++ b/src/models.rs @@ -18,7 +18,6 @@ #![allow(unused_qualifications)] use std::collections::{BTreeMap, HashMap}; -use std::sync::Arc; use serde::{Deserialize, Serialize}; use pyo3::prelude::*; @@ -53,9 +52,9 @@ pub struct InfoParams { } /// Parameters relevant to each detector -#[pyclass] #[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] -pub struct DetectorParams(BTreeMap); +pub struct DetectorParams(pub BTreeMap); + pub type Metadata = BTreeMap; @@ -376,17 +375,21 @@ pub struct ClassifiedGeneratedTextResult { /// The request format expected in the /api/v2/text/detection/content endpoint. +#[pyclass(name = "TextContentDetectionRequest")] #[derive(Clone, Debug, Serialize, Deserialize)] #[serde(deny_unknown_fields)] pub struct TextContentDetectionHttpRequest { + /// The content to run detectors on + #[pyo3(set)] pub content: String, /// The map of detectors to be used, along with their respective parameters, e.g. thresholds. + #[pyo3(set)] pub detectors: HashMap, } - impl TextContentDetectionHttpRequest { + /// Upfront validation of user request pub fn validate(&self) -> Result<(), ValidationError> { // Validate required parameters @@ -411,6 +414,7 @@ pub struct TextContentDetectionResult { /// Detection results pub detections: Vec, } + /// Streaming classification result on text produced by a text generation model, containing /// information from the original text generation output as well as the result of /// classification on the generated text. Also indicates where in stream is processed. @@ -1484,49 +1488,4 @@ impl From<&PyDetectorsObj> for HashMap { fn from(obj: &PyDetectorsObj) -> Self { obj.0.clone() } -} - -#[pyclass] -#[derive(Clone)] -pub struct PyTextContentDetectionHttpRequest { - #[pyo3(get, set)] - pub content: String, - // pub detectors: Arc>, - pub detectors: Arc, -} - -#[pymethods] -impl PyTextContentDetectionHttpRequest { - #[new] - fn new(content: String, detectors: PyDetectorsObj) -> Self { - PyTextContentDetectionHttpRequest { - content, - // detectors: Arc::new(Mutex::new(detectors)), // Wrap HashMap in Arc and Mutex for thread-safe access - detectors: Arc::new(detectors) - } - } - -} - -#[pyclass] -#[derive(Clone)] -pub struct PyContextDocsHttpRequest{ - /// The map of detectors to be used, along with their respective parameters, e.g. thresholds. - pub detectors: Arc, - pub content: String, - pub context_type: ContextType, - pub context: Vec, -} - -#[pymethods] -impl PyContextDocsHttpRequest { - #[new] - fn new(content: String, context_type: ContextType, context: Vec, detectors: PyDetectorsObj) -> Self { - Self { - content, - context, - context_type: context_type.into(), - detectors: detectors.into() - } - } } \ No newline at end of file diff --git a/src/orchestrator.rs b/src/orchestrator.rs index 571c2e00d..da58454f3 100644 --- a/src/orchestrator.rs +++ b/src/orchestrator.rs @@ -18,6 +18,7 @@ pub mod errors; pub use errors::Error; pub mod common; pub mod handlers; +pub mod python_interface; pub mod types; use std::{collections::HashMap, sync::Arc}; diff --git a/src/orchestrator/python_interface.rs b/src/orchestrator/python_interface.rs index 7b08c16f8..eac62754c 100644 --- a/src/orchestrator/python_interface.rs +++ b/src/orchestrator/python_interface.rs @@ -1,3 +1,6 @@ -// pub mod python_interface; -// pub use python_interface::orchestrator::*; \ No newline at end of file +pub mod orchestrator; +pub use orchestrator::*; + +pub mod models; +pub use models::*; \ No newline at end of file diff --git a/src/orchestrator/python_interface/models.rs b/src/orchestrator/python_interface/models.rs new file mode 100644 index 000000000..c0251abab --- /dev/null +++ b/src/orchestrator/python_interface/models.rs @@ -0,0 +1,83 @@ +use std::collections::{BTreeMap, HashMap}; + +use pyo3::prelude::*; +use pyo3::types::PyDict; +use pythonize::{depythonize}; + + +use crate::clients::detector::{ContextType}; + +use crate::models::{ + ContextDocsHttpRequest, + DetectorParams, + TextContentDetectionHttpRequest, + TextContentDetectionResult, +}; + + + +impl<'py> FromPyObject<'py> for DetectorParams { + fn extract_bound(obj: &Bound<'py, PyAny>) -> PyResult { + println!("{:?} reached here", obj); + // Use depythonize to convert the Python object directly to BTreeMap + let map: BTreeMap = depythonize(obj) + .map_err(|e| PyErr::new::(e.to_string()))?; + + Ok(DetectorParams(map)) + } +} + +// Text Content Detection + +#[pymethods] +impl TextContentDetectionHttpRequest { + + #[new] + #[pyo3(signature = (content, detectors = HashMap::new()))] + fn new(content: String, detectors: HashMap) -> Self { + Self { content, detectors } + } + +} + +impl<'py> IntoPyObject<'py> for TextContentDetectionResult { + type Target = PyAny; + type Output = Bound<'py, Self::Target>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> Result { + let mut detections = Vec::with_capacity(self.detections.len()); + // Note below will get simplified once we move all of this declaration to Orchestrator itself. + for detection in self.detections { + let content_response = PyDict::new(py); + content_response.set_item("start", detection.start)?; + content_response.set_item("end", detection.end)?; + content_response.set_item("text", detection.text)?; + content_response.set_item("detection", detection.detection)?; + content_response.set_item("detection_type", detection.detection_type)?; + content_response.set_item("detector_id", detection.detector_id)?; + content_response.set_item("score", detection.score)?; + detections.push(content_response) + } + detections.into_pyobject(py) + } +} + + +// Context Doc + + +#[pymethods] +impl ContextDocsHttpRequest { + #[new] + #[pyo3(signature = (content, context_type, context, detectors = HashMap::new()))] + fn new(content: String, context_type: ContextType, context: Vec, detectors: HashMap) -> Self { + Self { + content, + context, + context_type: context_type.into(), + detectors: detectors.into() + } + } +} + diff --git a/src/orchestrator/python_interface/orchestrator.rs b/src/orchestrator/python_interface/orchestrator.rs index 3eefe7d1e..fa23b521e 100644 --- a/src/orchestrator/python_interface/orchestrator.rs +++ b/src/orchestrator/python_interface/orchestrator.rs @@ -2,19 +2,27 @@ use std::sync::Arc; +use axum::http::HeaderMap; + use pyo3::prelude::*; use pyo3::exceptions::{PyOSError, PyTypeError}; -use pyo3::PyErr; -use pyo3::types::PyDict; -use pyo3::conversion::{FromPyObject}; +use pyo3_async_runtimes::tokio::future_into_py; + +use crate::models::{ + TextContentDetectionHttpRequest, +}; -use crate::orchestrator::Orchestrator; +use crate::orchestrator::{ + Orchestrator, + handlers::{text_content_detection, Handle}, +}; +use crate::utils::trace; use crate::config::OrchestratorConfig; // Define orchestrator #[pyclass] -struct GuardrailsOrchestrator { +pub struct GuardrailsOrchestrator { // Added under Arc to allow safe sharing from multiple threads // If we need modification capability later, then we might want to make it mutable orchestrator: Arc, @@ -24,8 +32,8 @@ struct GuardrailsOrchestrator { impl GuardrailsOrchestrator { #[new] - fn new(config_path: String) -> PyResult { - Python::with_gil(|py| { + fn new(config_path: String, start_up_health_check: bool) -> PyResult { + Python::attach(|py| { pyo3_async_runtimes::tokio::run(py, async move { let config_result = OrchestratorConfig::load(config_path).await; @@ -34,7 +42,7 @@ impl GuardrailsOrchestrator { Err(err) => return Err(PyOSError::new_err(format!("Error loading orchestrator configuration: {}", err))) }; - let orchestrator_result = Orchestrator::new(config, STARTUP_HEALTHCHECK).await; + let orchestrator_result = Orchestrator::new(config, start_up_health_check).await; let orchestrator = match orchestrator_result { Ok(orchestrator) => orchestrator, @@ -47,4 +55,30 @@ impl GuardrailsOrchestrator { })}) } + + pub fn content_detection<'py>( + &self, + py: Python<'py>, + request: TextContentDetectionHttpRequest, + ) -> PyResult> { + + let headers = HeaderMap::new(); + let trace_id = trace::current_trace_id(); + + let task = text_content_detection::TextContentDetectionTask::new(trace_id, request, headers); + + let orchestrator = Arc::clone(&self.orchestrator); + + future_into_py(py, async move { + match orchestrator.handle(task).await { + Ok(response) => { + println!("Response: {:?}", response); + Ok(response) + }, + // TODO: Handle errors properly with correct types + Err(error) => Err(PyTypeError::new_err(error.to_string())), + } + }) + } + } \ No newline at end of file From bd4d6014183f40a286124fd550ebeb571318609b Mon Sep 17 00:00:00 2001 From: gkumbhat Date: Wed, 24 Dec 2025 13:58:47 -0600 Subject: [PATCH 05/10] :package: Update pyo3 to 0.27.2 Signed-off-by: gkumbhat --- Cargo.lock | 34 ++++++++++++++++++---------------- Cargo.toml | 6 +++--- 2 files changed, 21 insertions(+), 19 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 6bcd3e3dc..da8e6db9b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1601,10 +1601,11 @@ dependencies = [ [[package]] name = "pyo3" -version = "0.26.0" +version = "0.27.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7ba0117f4212101ee6544044dae45abe1083d30ce7b29c4b5cbdfa2354e07383" +checksum = "ab53c047fcd1a1d2a8820fe84f05d6be69e9526be40cb03b73f86b6b03e6d87d" dependencies = [ + "hashbrown 0.16.1", "indoc", "libc", "memoffset", @@ -1613,14 +1614,15 @@ dependencies = [ "pyo3-build-config", "pyo3-ffi", "pyo3-macros", + "serde", "unindent", ] [[package]] name = "pyo3-async-runtimes" -version = "0.26.0" +version = "0.27.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e6ee6d4cb3e8d5b925f5cdb38da183e0ff18122eb2048d4041c9e7034d026e23" +checksum = "57ddb5b570751e93cc6777e81fee8087e59cd53b5043292f2a6d59d5bd80fdfd" dependencies = [ "futures", "once_cell", @@ -1632,9 +1634,9 @@ dependencies = [ [[package]] name = "pyo3-async-runtimes-macros" -version = "0.26.0" +version = "0.27.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c29bc5c673e36a8102d0b9179149c1bb59990d8db4f3ae58bd7dceccab90b951" +checksum = "bcd7d70ee0ca1661c40407e6f84e4463ef2658c90a9e2fbbd4515b2bcdfcaeca" dependencies = [ "proc-macro2", "quote", @@ -1643,18 +1645,18 @@ dependencies = [ [[package]] name = "pyo3-build-config" -version = "0.26.0" +version = "0.27.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4fc6ddaf24947d12a9aa31ac65431fb1b851b8f4365426e182901eabfb87df5f" +checksum = "b455933107de8642b4487ed26d912c2d899dec6114884214a0b3bb3be9261ea6" dependencies = [ "target-lexicon", ] [[package]] name = "pyo3-ffi" -version = "0.26.0" +version = "0.27.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "025474d3928738efb38ac36d4744a74a400c901c7596199e20e45d98eb194105" +checksum = "1c85c9cbfaddf651b1221594209aed57e9e5cff63c4d11d1feead529b872a089" dependencies = [ "libc", "pyo3-build-config", @@ -1662,9 +1664,9 @@ dependencies = [ [[package]] name = "pyo3-macros" -version = "0.26.0" +version = "0.27.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2e64eb489f22fe1c95911b77c44cc41e7c19f3082fc81cce90f657cdc42ffded" +checksum = "0a5b10c9bf9888125d917fb4d2ca2d25c8df94c7ab5a52e13313a07e050a3b02" dependencies = [ "proc-macro2", "pyo3-macros-backend", @@ -1674,9 +1676,9 @@ dependencies = [ [[package]] name = "pyo3-macros-backend" -version = "0.26.0" +version = "0.27.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "100246c0ecf400b475341b8455a9213344569af29a3c841d29270e53102e0fcf" +checksum = "03b51720d314836e53327f5871d4c0cfb4fb37cc2c4a11cc71907a86342c40f9" dependencies = [ "heck", "proc-macro2", @@ -1687,9 +1689,9 @@ dependencies = [ [[package]] name = "pythonize" -version = "0.26.0" +version = "0.27.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "11e06e4cff9be2bbf2bddf28a486ae619172ea57e79787f856572878c62dcfe2" +checksum = "a3a8f29db331e28c332c63496cfcbb822aca3d7320bc08b655d7fd0c29c50ede" dependencies = [ "pyo3", "serde", diff --git a/Cargo.toml b/Cargo.toml index 39e35c88d..5d3300b65 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -55,9 +55,9 @@ opentelemetry-otlp = { version = "0.31.0", features = [ opentelemetry_sdk = { version = "0.31.0", features = ["rt-tokio", "metrics"] } pin-project-lite = "0.2.16" prost = "0.14.1" -pyo3 = { version = "0.26.0", features = ["extension-module"]} -pyo3-async-runtimes = { version = "0.26.0", features = ["attributes", "tokio-runtime"] } -pythonize = "0.26.0" +pyo3 = { version = "0.27.2", features = ["extension-module", "serde", "hashbrown"]} +pyo3-async-runtimes = { version = "0.27.0", features = ["attributes", "tokio-runtime"] } +pythonize = "0.27.0" reqwest = { version = "0.12.25", features = [ "blocking", "rustls-tls", From 81a6c971e27db4df20a515e6a41c92005550068e Mon Sep 17 00:00:00 2001 From: gkumbhat Date: Wed, 24 Dec 2025 14:22:55 -0600 Subject: [PATCH 06/10] :recycle::sparkles: Refactor and implement chat completion detection Signed-off-by: gkumbhat --- src/clients/openai.rs | 86 ++++++- src/lib.rs | 37 ++- src/models.rs | 25 +- src/orchestrator.rs | 2 - src/orchestrator/python_interface.rs | 3 +- src/orchestrator/python_interface/models.rs | 233 ++++++++++++++++-- .../python_interface/orchestrator.rs | 168 ++++++++++--- src/orchestrator/types/detection.rs | 13 + 8 files changed, 475 insertions(+), 92 deletions(-) diff --git a/src/clients/openai.rs b/src/clients/openai.rs index 6d25482a4..062b36c45 100644 --- a/src/clients/openai.rs +++ b/src/clients/openai.rs @@ -27,6 +27,9 @@ use serde_json::{Map, Value}; use tokio::sync::mpsc; use url::Url; +use pyo3::{prelude::*, pyclass}; +use pythonize::pythonize; + use super::{ Client, Error, HttpClient, create_http_client, http::{HttpClientExt, RequestBody}, @@ -343,6 +346,7 @@ impl ChatCompletionsRequest { /// This is to avoid tracking and updating OpenAI and vLLM /// parameter additions/changes. Full validation is delegated to /// the downstream server implementation. +#[pyclass(name = "CompletionsRequest")] #[derive(Debug, Default, Clone, PartialEq, Serialize, Deserialize)] pub struct CompletionsRequest { /// Detector config. @@ -416,6 +420,7 @@ impl TokenizeRequest { } /// Detector config. +#[pyclass(name = "DetectorConfig")] #[derive(Default, Debug, Clone, PartialEq, Serialize, Deserialize)] #[serde(deny_unknown_fields)] pub struct DetectorConfig { @@ -426,6 +431,7 @@ pub struct DetectorConfig { } /// Response format. +#[pyclass(name = "ResponseFormat")] #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ResponseFormat { /// The type of response format being defined. @@ -436,6 +442,7 @@ pub struct ResponseFormat { } /// Tool. +#[pyclass(get_all)] #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] #[serde(untagged)] pub enum Tool { @@ -444,6 +451,7 @@ pub enum Tool { } /// Function tool. +#[pyclass(get_all)] #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct FunctionTool { /// The type of the tool. Always `function`. @@ -477,6 +485,7 @@ impl From for Tool { } /// Custom tool. +#[pyclass(get_all)] #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct CustomTool { /// The type of the tool. Always `custom`. @@ -510,6 +519,7 @@ impl From for Tool { } /// A function tool that can be used to generate a response. +#[pyclass] #[derive(Debug, Default, Clone, PartialEq, Serialize, Deserialize)] pub struct FunctionDefinition { /// The name of the function to be called. @@ -527,19 +537,33 @@ pub struct FunctionDefinition { } /// A custom tool that processes input using a specified format. +#[pyclass] #[derive(Debug, Default, Clone, PartialEq, Serialize, Deserialize)] pub struct CustomDefinition { /// The name of the custom tool, used to identify it in tool calls. + #[pyo3(get)] pub name: String, /// Optional description of the custom tool, used to provide more context. #[serde(skip_serializing_if = "Option::is_none")] + #[pyo3(get)] pub description: Option, /// The input format for the custom tool. Default is unconstrained text. #[serde(skip_serializing_if = "Option::is_none")] pub format: Option, } +#[pymethods] +impl CustomDefinition { + #[getter(format)] + fn get_format<'py>(&self, py: Python<'py>) -> PyResult> { + // Use pythonize to convert serde_json::Value to a Python dict/list + pythonize(py, &self.format) + .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string())) + } +} + /// Stream options. +#[pyclass(get_all)] #[derive(Debug, Clone, Serialize, Deserialize)] pub struct StreamOptions { /// If set, an additional chunk will be streamed before the data: [DONE] message. @@ -551,6 +575,7 @@ pub struct StreamOptions { } /// Role. +#[pyclass(get_all)] #[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)] #[serde(rename_all = "lowercase")] pub enum Role { @@ -563,6 +588,7 @@ pub enum Role { } /// Message. +#[pyclass(name = "ChatCompletionMessage", get_all, set_all)] #[derive(Debug, Default, Clone, PartialEq, Serialize, Deserialize)] pub struct Message { /// The role of the author of this message. @@ -606,12 +632,15 @@ impl Message { } /// Content. +#[pyclass] #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] #[serde(untagged)] pub enum Content { /// The text contents of the message. + #[pyo3(constructor = (_0))] // Makes the value accessible via content._0 Text(String), - /// Array of content parts. + /// Array of content parts. // Makes the value accessible via content._0 + #[pyo3(constructor = (_0))] Array(Vec), } @@ -661,6 +690,7 @@ impl From> for Content { } /// Content type. +#[pyclass(get_all)] #[derive(Debug, Default, Clone, Serialize, Deserialize, PartialEq)] pub enum ContentType { #[serde(rename = "text")] @@ -671,6 +701,7 @@ pub enum ContentType { } /// Content part. +#[pyclass(get_all)] #[derive(Debug, Default, Clone, Serialize, Deserialize, PartialEq)] pub struct ContentPart { /// The type of the content part. @@ -701,6 +732,7 @@ impl ContentPart { } /// Image url. +#[pyclass(get_all)] #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] pub struct ImageUrl { /// Either a URL of the image or the base64 encoded image data. @@ -711,6 +743,7 @@ pub struct ImageUrl { } /// Tool call. +#[pyclass(get_all)] #[derive(Debug, Default, Clone, Serialize, Deserialize, PartialEq)] pub struct ToolCall { /// Index (streaming) @@ -731,6 +764,7 @@ pub struct ToolCall { } /// Function. +#[pyclass(get_all)] #[derive(Debug, Default, Clone, Serialize, Deserialize, PartialEq)] pub struct Function { /// The name of the function to call. @@ -742,6 +776,7 @@ pub struct Function { } /// Custom. +#[pyclass(get_all)] #[derive(Debug, Default, Clone, Serialize, Deserialize, PartialEq)] pub struct Custom { /// The name of the custom tool to call. @@ -753,38 +788,50 @@ pub struct Custom { } /// Chat completion response. +#[pyclass] #[derive(Debug, Default, Clone, Serialize, Deserialize, PartialEq)] pub struct ChatCompletion { /// A unique identifier for the chat completion. + #[pyo3(get)] pub id: String, /// The object type, which is always `chat.completion`. + #[pyo3(get)] pub object: String, /// The Unix timestamp (in seconds) of when the chat completion was created. + #[pyo3(get)] pub created: i64, /// The model used for the chat completion. + #[pyo3(get)] pub model: String, /// A list of chat completion choices. Can be more than one if n is greater than 1. + #[pyo3(get)] pub choices: Vec, /// Usage statistics for the completion request. + #[pyo3(get)] pub usage: Usage, /// Prompt logprobs. + #[pyo3(get)] pub prompt_logprobs: Option>>>, /// This fingerprint represents the backend configuration that the model runs with. #[serde(skip_serializing_if = "Option::is_none")] + #[pyo3(get)] pub system_fingerprint: Option, /// The service tier used for processing the request. /// This field is only included if the `service_tier` parameter is specified in the request. #[serde(skip_serializing_if = "Option::is_none")] + #[pyo3(get)] pub service_tier: Option, /// Detections #[serde(skip_serializing_if = "Option::is_none")] pub detections: Option, /// Warnings #[serde(default, skip_serializing_if = "Vec::is_empty")] + #[pyo3(get)] pub warnings: Vec, } /// Helper to accept both string and integer for stop_reason. +#[pyclass(get_all)] #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] #[serde(untagged)] pub enum StopReason { @@ -793,6 +840,7 @@ pub enum StopReason { } /// Chat completion choice. +#[pyclass(get_all)] #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] pub struct ChatCompletionChoice { /// The index of the choice in the list of choices. @@ -808,6 +856,7 @@ pub struct ChatCompletionChoice { } /// Chat completion logprobs. +#[pyclass(get_all)] #[derive(Debug, Default, Clone, Deserialize, Serialize, PartialEq)] pub struct ChatCompletionLogprobs { /// A list of message content tokens with log probability information. @@ -819,6 +868,7 @@ pub struct ChatCompletionLogprobs { } /// Chat completion logprob. +#[pyclass(get_all)] #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] pub struct ChatCompletionLogprob { /// The token. @@ -832,6 +882,7 @@ pub struct ChatCompletionLogprob { } /// Chat completion top logprob. +#[pyclass(get_all)] #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] pub struct ChatCompletionTopLogprob { /// The token. @@ -843,32 +894,42 @@ pub struct ChatCompletionTopLogprob { } /// Streaming chat completion chunk. +#[pyclass] #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct ChatCompletionChunk { /// A unique identifier for the chat completion. Each chunk has the same ID. + #[pyo3(get)] pub id: String, /// The object type, which is always `chat.completion.chunk`. + #[pyo3(get)] pub object: String, /// The Unix timestamp (in seconds) of when the chat completion was created. Each chunk has the same timestamp. + #[pyo3(get)] pub created: i64, /// The model to generate the completion. + #[pyo3(get)] pub model: String, /// This fingerprint represents the backend configuration that the model runs with. #[serde(skip_serializing_if = "Option::is_none")] + #[pyo3(get)] pub system_fingerprint: Option, /// A list of chat completion choices. + #[pyo3(get)] pub choices: Vec, /// The service tier used for processing the request. /// This field is only included if the service_tier parameter is specified in the request. #[serde(skip_serializing_if = "Option::is_none")] + #[pyo3(get)] pub service_tier: Option, #[serde(skip_serializing_if = "Option::is_none")] + #[pyo3(get)] pub usage: Option, /// Detections #[serde(skip_serializing_if = "Option::is_none")] pub detections: Option, /// Warnings #[serde(default, skip_serializing_if = "Vec::is_empty")] + #[pyo3(get)] pub warnings: Vec, } @@ -890,6 +951,7 @@ impl Default for ChatCompletionChunk { } /// Streaming chat completion chunk choice. +#[pyclass(get_all)] #[derive(Debug, Default, Clone, PartialEq, Serialize, Deserialize)] pub struct ChatCompletionChunkChoice { /// The index of the choice in the list of choices. @@ -905,6 +967,7 @@ pub struct ChatCompletionChunkChoice { } /// Streaming chat completion delta. +#[pyclass(get_all)] #[derive(Debug, Default, Clone, PartialEq, Serialize, Deserialize)] pub struct ChatCompletionDelta { /// The role of the author of this message. @@ -922,28 +985,37 @@ pub struct ChatCompletionDelta { } /// Completion (legacy) response. Also used for streaming. +#[pyclass] #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] pub struct Completion { /// A unique identifier for the completion. + #[pyo3(get)] pub id: String, /// The object type, which is always `text_completion`. + #[pyo3(get)] pub object: String, /// The Unix timestamp (in seconds) of when the chat completion was created. + #[pyo3(get)] pub created: i64, /// The model used for the completion. + #[pyo3(get)] pub model: String, /// A list of completion choices. Can be more than one if n is greater than 1. + #[pyo3(get)] pub choices: Vec, /// Usage statistics for the completion request. + #[pyo3(get)] pub usage: Option, /// This fingerprint represents the backend configuration that the model runs with. #[serde(skip_serializing_if = "Option::is_none")] + #[pyo3(get)] pub system_fingerprint: Option, /// Detections #[serde(skip_serializing_if = "Option::is_none")] pub detections: Option, /// Warnings #[serde(default, skip_serializing_if = "Vec::is_empty")] + #[pyo3(get)] pub warnings: Vec, } @@ -964,6 +1036,7 @@ impl Default for Completion { } /// Completion (legacy) choice. +#[pyclass(get_all)] #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)] pub struct CompletionChoice { /// The index of the choice in the list of choices. @@ -982,6 +1055,7 @@ pub struct CompletionChoice { } /// Completion logprobs. +#[pyclass(get_all)] #[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Default)] pub struct CompletionLogprobs { /// Tokens generated by the model. @@ -995,6 +1069,7 @@ pub struct CompletionLogprobs { } /// Logprob. +#[pyclass(get_all)] #[derive(Debug, Clone, Deserialize, Serialize, PartialEq)] pub struct Logprob { /// The logprob of the chosen token @@ -1006,6 +1081,7 @@ pub struct Logprob { } /// Completion usage statistics. +#[pyclass(get_all)] #[derive(Debug, Default, Clone, Serialize, Deserialize, PartialEq)] pub struct Usage { /// Number of tokens in the prompt. @@ -1023,6 +1099,7 @@ pub struct Usage { } /// Completion token details. +#[pyclass(get_all)] #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] pub struct CompletionTokenDetails { pub audio_tokens: u32, @@ -1030,6 +1107,7 @@ pub struct CompletionTokenDetails { } /// Prompt token details. +#[pyclass(get_all)] #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] pub struct PromptTokenDetails { pub audio_tokens: u32, @@ -1037,6 +1115,7 @@ pub struct PromptTokenDetails { } /// Stop tokens. +#[pyclass(get_all)] #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(untagged)] pub enum StopTokens { @@ -1045,6 +1124,7 @@ pub enum StopTokens { } /// Error response v1, for backwards compatability. +#[pyclass(get_all)] #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ErrorResponseV1 { pub object: String, @@ -1056,12 +1136,14 @@ pub struct ErrorResponseV1 { } /// Error response v2. vLLM >= v0.10.1. +#[pyclass(get_all)] #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ErrorResponseV2 { pub error: ErrorInfo, } /// Error info. +#[pyclass(get_all)] #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ErrorInfo { pub message: String, @@ -1072,6 +1154,7 @@ pub struct ErrorInfo { } /// Error response. +#[pyclass(get_all)] #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(untagged)] pub enum ErrorResponse { @@ -1169,6 +1252,7 @@ pub struct CompletionOutputDetections { } /// Guardrails completion detection warning. +#[pyclass(get_all)] #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] pub struct CompletionDetectionWarning { r#type: DetectionWarningReason, diff --git a/src/lib.rs b/src/lib.rs index c2c7b43c5..8e59c35cb 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -31,26 +31,47 @@ pub(crate) use utils::AsUriExt; use pyo3::prelude::*; -use crate::models::{ - ContextDocsResult, - ContextDocsHttpRequest, - TextContentDetectionHttpRequest, - TextContentDetectionResult +use crate::{ + clients::openai::{DetectorConfig, Message}, + models::{ + ContextDocsHttpRequest, ContextDocsResult, TextContentDetectionHttpRequest, + TextContentDetectionResult, + }, }; - /// A Python module implemented in Rust. #[pymodule] fn fms_guardrails_orchestr8(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { - // TODO: Replace this with proper implementation let _ = rustls::crypto::ring::default_provider().install_default(); m.add_class::()?; + + // Text Content Detection m.add_class::()?; // TextContentDetectionRequest m.add_class::()?; + + // Context Detection m.add_class::()?; m.add_class::()?; m.add_class::()?; + + // Chat Completion Detection + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + + // Completion Detection + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + Ok(()) -} \ No newline at end of file +} diff --git a/src/models.rs b/src/models.rs index 25ad8f476..b57cb9df8 100644 --- a/src/models.rs +++ b/src/models.rs @@ -19,13 +19,9 @@ use std::collections::{BTreeMap, HashMap}; -use serde::{Deserialize, Serialize}; -use pyo3::prelude::*; -use pyo3::pyclass; -use pyo3::types::PyDict; -use pyo3::conversion::{FromPyObject}; +use pyo3::{conversion::FromPyObject, prelude::*, pyclass, types::PyDict}; use pythonize::depythonize; - +use serde::{Deserialize, Serialize}; use crate::{ clients::{ @@ -55,7 +51,6 @@ pub struct InfoParams { #[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] pub struct DetectorParams(pub BTreeMap); - pub type Metadata = BTreeMap; impl DetectorParams { @@ -373,13 +368,11 @@ pub struct ClassifiedGeneratedTextResult { pub input_tokens: Option>, } - /// The request format expected in the /api/v2/text/detection/content endpoint. #[pyclass(name = "TextContentDetectionRequest")] #[derive(Clone, Debug, Serialize, Deserialize)] #[serde(deny_unknown_fields)] pub struct TextContentDetectionHttpRequest { - /// The content to run detectors on #[pyo3(set)] pub content: String, @@ -389,7 +382,6 @@ pub struct TextContentDetectionHttpRequest { pub detectors: HashMap, } impl TextContentDetectionHttpRequest { - /// Upfront validation of user request pub fn validate(&self) -> Result<(), ValidationError> { // Validate required parameters @@ -572,6 +564,7 @@ impl DetectionWarning { /// Since this enum's variants do not hold data, we can easily define them as `#[repr(C)]` /// which helps with FFI. #[repr(C)] +#[pyclass] #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] pub enum DetectionWarningReason { /// Unsuitable text detected on input @@ -1444,7 +1437,6 @@ mod tests { } } - /// Python Interfaces #[pyclass] @@ -1453,15 +1445,18 @@ pub struct PyDetectorsObj(HashMap); // NOTE: If this moves to orchestrator object itself, we don't need to do double conversion // and we can implement the trait directly there -impl<'py> FromPyObject<'py> for PyDetectorsObj { - fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult { +impl<'a, 'py> FromPyObject<'a, 'py> for PyDetectorsObj { + // Add an associated Error type + type Error = PyErr; + + fn extract(ob: Borrowed<'a, 'py, PyAny>) -> PyResult { let dict = ob.downcast::()?; let mut result = HashMap::new(); for (key, value) in dict.iter() { let key = key.extract::()?; - let inner_dict = value.downcast::()?; + let inner_dict = value.downcast::()?; // let mut inner_result = BTreeMap::new(); let mut inner_result = DetectorParams::new(); @@ -1488,4 +1483,4 @@ impl From<&PyDetectorsObj> for HashMap { fn from(obj: &PyDetectorsObj) -> Self { obj.0.clone() } -} \ No newline at end of file +} diff --git a/src/orchestrator.rs b/src/orchestrator.rs index da58454f3..af4830038 100644 --- a/src/orchestrator.rs +++ b/src/orchestrator.rs @@ -106,7 +106,6 @@ impl Orchestrator { } health } - } async fn create_clients(config: &OrchestratorConfig) -> Result { @@ -182,4 +181,3 @@ async fn create_clients(config: &OrchestratorConfig) -> Result // future_into_py(py, guardrails_orch8_future) // } - diff --git a/src/orchestrator/python_interface.rs b/src/orchestrator/python_interface.rs index eac62754c..48054bef7 100644 --- a/src/orchestrator/python_interface.rs +++ b/src/orchestrator/python_interface.rs @@ -1,6 +1,5 @@ - pub mod orchestrator; pub use orchestrator::*; pub mod models; -pub use models::*; \ No newline at end of file +pub use models::*; diff --git a/src/orchestrator/python_interface/models.rs b/src/orchestrator/python_interface/models.rs index c0251abab..16b0aaa98 100644 --- a/src/orchestrator/python_interface/models.rs +++ b/src/orchestrator/python_interface/models.rs @@ -1,48 +1,69 @@ -use std::collections::{BTreeMap, HashMap}; - -use pyo3::prelude::*; -use pyo3::types::PyDict; -use pythonize::{depythonize}; +use serde_json::{Map, Value}; +use std::{ + collections::{BTreeMap, HashMap}, + sync::Arc, +}; +use pyo3::{prelude::*, types::PyDict}; +use pyo3_async_runtimes::tokio::future_into_py; +use pythonize::{depythonize, pythonize}; -use crate::clients::detector::{ContextType}; +use tokio::sync::{Mutex, mpsc}; -use crate::models::{ - ContextDocsHttpRequest, - DetectorParams, - TextContentDetectionHttpRequest, - TextContentDetectionResult, +use crate::{ + clients::{detector::ContextType, openai}, + models::{ + ContextDocsHttpRequest, DetectorParams, TextContentDetectionHttpRequest, + TextContentDetectionResult, + }, + orchestrator, + orchestrator::types::Detection, }; +impl<'a, 'py> FromPyObject<'a, 'py> for DetectorParams { + // Add an associated Error type + type Error = PyErr; - -impl<'py> FromPyObject<'py> for DetectorParams { - fn extract_bound(obj: &Bound<'py, PyAny>) -> PyResult { - println!("{:?} reached here", obj); + fn extract(obj: Borrowed<'a, 'py, PyAny>) -> PyResult { // Use depythonize to convert the Python object directly to BTreeMap - let map: BTreeMap = depythonize(obj) + let map: BTreeMap = depythonize(&obj) .map_err(|e| PyErr::new::(e.to_string()))?; Ok(DetectorParams(map)) } } +#[pymethods] +impl Detection { + #[getter] + fn metadata<'py>(&self, py: Python<'py>) -> PyResult> { + // pythonize converts the BTreeMap and all nested Values into a Python dict + pythonize(py, &self.metadata) + .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string())) + } +} + +#[pymethods] +impl openai::ChatCompletionChoice { + pub fn to_dict<'py>(&self, py: Python<'py>) -> PyResult> { + pythonize(py, self).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string())) + } +} + // Text Content Detection #[pymethods] impl TextContentDetectionHttpRequest { - #[new] #[pyo3(signature = (content, detectors = HashMap::new()))] fn new(content: String, detectors: HashMap) -> Self { Self { content, detectors } } - } impl<'py> IntoPyObject<'py> for TextContentDetectionResult { type Target = PyAny; - type Output = Bound<'py, Self::Target>; + type Output = Bound<'py, Self::Target>; type Error = PyErr; fn into_pyobject(self, py: Python<'py>) -> Result { @@ -63,21 +84,187 @@ impl<'py> IntoPyObject<'py> for TextContentDetectionResult { } } - // Context Doc - #[pymethods] impl ContextDocsHttpRequest { #[new] #[pyo3(signature = (content, context_type, context, detectors = HashMap::new()))] - fn new(content: String, context_type: ContextType, context: Vec, detectors: HashMap) -> Self { + fn new( + content: String, + context_type: ContextType, + context: Vec, + detectors: HashMap, + ) -> Self { Self { content, context, context_type: context_type.into(), - detectors: detectors.into() + detectors: detectors.into(), } } } +// Completion Detection + +#[pymethods] +impl openai::Completion { + // Implementing manual getter to override the default get behavior. Reasons: + // 1. Default getter gives error because of serde_value (eventually) + // 2. Adds way too much nesting to be able to get to final answer + // 3. This provides an easy way to get the response in python dict, allowing + // easier access and possibly conversion in python. + #[getter] + fn detections<'py>(&self, py: Python<'py>) -> PyResult> { + // This converts the Rust struct into a native Python dict automatically + pythonize(py, &self.detections) + .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string())) + } +} + +// Chat Completion Detection + +#[pymethods] +impl openai::ChatCompletion { + #[getter] + fn detections<'py>(&self, py: Python<'py>) -> PyResult> { + // This converts the Rust struct into a native Python dict automatically + pythonize(py, &self.detections) + .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string())) + } +} + +#[pymethods] +impl openai::ChatCompletionChunk { + #[getter] + fn detections<'py>(&self, py: Python<'py>) -> PyResult> { + // This converts the Rust struct into a native Python dict automatically + pythonize(py, &self.detections) + .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string())) + } +} + +#[pymethods] +impl openai::Message { + #[new] + #[pyo3(signature = (data = None, **kwargs))] + fn new( + py: Python<'_>, + data: Option>, + kwargs: Option>, + ) -> PyResult { + // Identify the source: prioritize positional 'data' then 'kwargs' + let source = if let Some(datum) = data { + datum + } else if let Some(k) = kwargs { + k.into_any() + } else { + return Err(pyo3::exceptions::PyTypeError::new_err( + "Missing message data", + )); + }; + + // Convert from Python object/dict to Rust struct via depythonize + // This works for: + // - A dict: {"role": "user", "content": "..."} + let message: openai::Message = depythonize(&source).map_err(|e| { + pyo3::exceptions::PyTypeError::new_err(format!("Invalid message format: {}", e)) + })?; + + Ok(message) + } +} + +#[pymethods] +impl openai::DetectorConfig { + // TODO: Allow both dictionary as well as struct input + #[new] + #[pyo3(signature = (input = HashMap::new(), output = HashMap::new()))] + fn new( + input: HashMap, + output: HashMap, + ) -> Self { + Self { input, output } + } +} + +#[pyclass(name = "ChatCompletionsRequest")] +#[derive(Clone)] +pub struct PyChatCompletionsRequest(pub openai::ChatCompletionsRequest); + +#[pymethods] +impl PyChatCompletionsRequest { + #[new] + #[pyo3(signature = (model, detectors, messages, stream = None, extra = HashMap::new(), tools = None ))] + fn new( + py: Python<'_>, + model: String, + detectors: openai::DetectorConfig, + messages: Bound<'_, PyAny>, + stream: Option, + extra: HashMap>, + tools: Option>, + ) -> PyResult { + let mut json_map = Map::new(); + + for (key, py_obj) in extra { + // Bind the Py to the current GIL to get a Bound<'_, PyAny> + let bound_obj = py_obj.bind(py); + + // Convert Python object to serde_json::Value + let value: Value = depythonize(bound_obj) + .map_err(|e| PyErr::new::(e.to_string()))?; + + json_map.insert(key, value); + } + let new_messages: Vec = depythonize(&messages).map_err(|e| { + pyo3::exceptions::PyTypeError::new_err(format!( + "Argument 'messages' must be a list of valid message objects. Error: {}", + e + )) + })?; + + Ok(Self(openai::ChatCompletionsRequest { + stream, + detectors, + model, + messages: new_messages, + tools, + extra: json_map, + })) + } +} + +#[pyclass(name = "ChatCompletionsStream")] +pub struct PyChatCompletionsStream { + pub receiver: Arc< + Mutex, orchestrator::Error>>>, + >, +} + +#[pymethods] +impl PyChatCompletionsStream { + fn __aiter__(slf: PyRef) -> PyRef { + slf + } + + fn __anext__<'py>(&self, py: Python<'py>) -> PyResult> { + let receiver = self.receiver.clone(); + + future_into_py(py, async move { + let mut receiver_stream = receiver.lock().await; + + match receiver_stream.recv().await { + Some(Ok(Some(chunk))) => Ok(chunk), + + Some(Ok(None)) => Err(pyo3::exceptions::PyStopAsyncIteration::new_err(())), + + Some(Err(e)) => Err(PyErr::new::( + e.to_string(), + )), + + None => Err(pyo3::exceptions::PyStopAsyncIteration::new_err(())), + } + }) + } +} diff --git a/src/orchestrator/python_interface/orchestrator.rs b/src/orchestrator/python_interface/orchestrator.rs index fa23b521e..c2c583cdd 100644 --- a/src/orchestrator/python_interface/orchestrator.rs +++ b/src/orchestrator/python_interface/orchestrator.rs @@ -1,24 +1,28 @@ - - use std::sync::Arc; use axum::http::HeaderMap; +use tokio::sync::Mutex; -use pyo3::prelude::*; -use pyo3::exceptions::{PyOSError, PyTypeError}; +use pyo3::{ + exceptions::{PyOSError, PyTypeError}, + prelude::*, +}; use pyo3_async_runtimes::tokio::future_into_py; -use crate::models::{ - TextContentDetectionHttpRequest, +use crate::{ + clients::openai::ChatCompletionsResponse, + config::OrchestratorConfig, + models::TextContentDetectionHttpRequest, + orchestrator::{ + Orchestrator, + handlers::{Handle, chat_completions_detection, text_content_detection}, + }, }; -use crate::orchestrator::{ - Orchestrator, - handlers::{text_content_detection, Handle}, +use crate::{ + orchestrator::python_interface::models::{PyChatCompletionsRequest, PyChatCompletionsStream}, + utils::trace, }; -use crate::utils::trace; - -use crate::config::OrchestratorConfig; // Define orchestrator #[pyclass] @@ -30,55 +34,137 @@ pub struct GuardrailsOrchestrator { #[pymethods] impl GuardrailsOrchestrator { - #[new] fn new(config_path: String, start_up_health_check: bool) -> PyResult { Python::attach(|py| { pyo3_async_runtimes::tokio::run(py, async move { - let config_result = OrchestratorConfig::load(config_path).await; - - let config = match config_result { - Ok(config) => config, - Err(err) => return Err(PyOSError::new_err(format!("Error loading orchestrator configuration: {}", err))) - }; - - let orchestrator_result = Orchestrator::new(config, start_up_health_check).await; - - let orchestrator = match orchestrator_result { - Ok(orchestrator) => orchestrator, - Err(err) => return Err(PyOSError::new_err(format!("Error creating orchestrator: {}", err))) - }; - - Ok(GuardrailsOrchestrator { - orchestrator: Arc::new(orchestrator), + let config_result = OrchestratorConfig::load(config_path).await; + + let config = match config_result { + Ok(config) => config, + Err(err) => { + return Err(PyOSError::new_err(format!( + "Error loading orchestrator configuration: {}", + err + ))); + } + }; + + let orchestrator_result = Orchestrator::new(config, start_up_health_check).await; + + let orchestrator = match orchestrator_result { + Ok(orchestrator) => orchestrator, + Err(err) => { + return Err(PyOSError::new_err(format!( + "Error creating orchestrator: {}", + err + ))); + } + }; + + Ok(GuardrailsOrchestrator { + orchestrator: Arc::new(orchestrator), + }) }) - - })}) + }) } pub fn content_detection<'py>( &self, py: Python<'py>, request: TextContentDetectionHttpRequest, - ) -> PyResult> { - + ) -> PyResult> { let headers = HeaderMap::new(); let trace_id = trace::current_trace_id(); - let task = text_content_detection::TextContentDetectionTask::new(trace_id, request, headers); + let task = + text_content_detection::TextContentDetectionTask::new(trace_id, request, headers); let orchestrator = Arc::clone(&self.orchestrator); future_into_py(py, async move { match orchestrator.handle(task).await { - Ok(response) => { - println!("Response: {:?}", response); - Ok(response) - }, - // TODO: Handle errors properly with correct types - Err(error) => Err(PyTypeError::new_err(error.to_string())), + Ok(response) => { + println!("Response: {:?}", response); + Ok(response) } + // TODO: Handle errors properly with correct types + Err(error) => Err(PyTypeError::new_err(error.to_string())), + } }) } -} \ No newline at end of file + pub fn chat_completions_detection<'py>( + &self, + py: Python<'py>, + request: PyChatCompletionsRequest, + ) -> PyResult> { + let headers = HeaderMap::new(); + let trace_id = trace::current_trace_id(); + + let task = chat_completions_detection::ChatCompletionsDetectionTask::new( + trace_id, request.0, headers, + ); + + let orchestrator = Arc::clone(&self.orchestrator); + + // Convert the async move { } block to a Python awaitable + future_into_py(py, async move { + match orchestrator.handle(task).await { + Ok(ChatCompletionsResponse::Unary(response)) => { + println!("Response: {:?}", response); + // response.into_pyobject(py) + // response.into_pyobject(py).map(|obj| obj.into()) + // Ok(Python::attach(|py| pythonize(py, &*response))) + // Ok(Python::attach(|py| response.into_py_any(py))) + // let result_py_object: Py = Python::attach(|py| { + // let bound_any = pythonize(py, &response) + // .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?; + + // // 3. .unbind() converts Bound<'py, PyAny> to Py (aka PyObject) + // // This "erases" the lifetime so it can safely leave the closure + // Ok::, PyErr>(bound_any.unbind()) + // })?; + + // Ok(result_py_object) + let py_obj = Python::attach(|py| { + let bound_any = Bound::new(py, *response)?.into_any(); + + Ok::(bound_any.unbind()) + // bound_any.unbind() + // bound_any + }); + // Ok(py_obj) + println!("{:?}", py_obj); + py_obj + } + + Ok(ChatCompletionsResponse::Streaming(receiver_stream)) => { + let py_stream = PyChatCompletionsStream { + receiver: Arc::new(Mutex::new(receiver_stream)), + }; + // Ok(Python::attach(|py| pythonize(py, py_stream))) + // Ok(Python::attach(|py| py_stream.into_py_any(py))) + // py_stream.into_pyobject(py) + // py_stream.into_pyobject(py).map(|obj| obj.into()) + // let result_py_object: Py = Python::attach(|py| { + // let bound_any = pythonize(py, &py_stream) + // .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?; + + // // 3. .unbind() converts Bound<'py, PyAny> to Py (aka PyObject) + // // This "erases" the lifetime so it can safely leave the closure + // Ok::, PyErr>(bound_any.unbind()) + // })?; + // Ok(result_py_object) + + let py_obj = Python::attach(|py| { + let bound_any = Bound::new(py, py_stream)?.into_any(); + Ok::(bound_any.unbind()) + }); + py_obj + } + Err(e) => Err(PyTypeError::new_err(e.to_string())), + } + }) + } +} diff --git a/src/orchestrator/types/detection.rs b/src/orchestrator/types/detection.rs index 82ead2481..388f56390 100644 --- a/src/orchestrator/types/detection.rs +++ b/src/orchestrator/types/detection.rs @@ -18,30 +18,41 @@ use std::collections::BTreeMap; use serde::{Deserialize, Serialize}; +use pyo3::prelude::*; + use crate::{clients::detector, models}; /// A detection. +#[pyclass] #[derive(Default, Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct Detection { /// Start index of the detection #[serde(skip_serializing_if = "Option::is_none")] + #[pyo3(get)] pub start: Option, /// End index of the detection #[serde(skip_serializing_if = "Option::is_none")] + #[pyo3(get)] pub end: Option, /// Text corresponding to the detection #[serde(skip_serializing_if = "Option::is_none")] + #[pyo3(get)] pub text: Option, /// ID of the detector + #[pyo3(get)] pub detector_id: Option, /// Type of detection + #[pyo3(get)] pub detection_type: String, /// Detection class + #[pyo3(get)] pub detection: String, /// Confidence level of the detection class + #[pyo3(get)] pub score: f64, /// Detection evidence #[serde(default, skip_serializing_if = "Vec::is_empty")] + #[pyo3(get)] pub evidence: Vec, /// Detection metadata #[serde(default, skip_serializing_if = "BTreeMap::is_empty")] @@ -49,6 +60,7 @@ pub struct Detection { } /// Detection evidence. +#[pyclass[get_all]] #[derive(Default, Clone, Debug, PartialEq, Serialize, Deserialize)] pub struct DetectionEvidence { /// Evidence name @@ -65,6 +77,7 @@ pub struct DetectionEvidence { } /// Additional detection evidence. +#[pyclass[get_all]] #[derive(Default, Clone, Debug, PartialEq, Serialize, Deserialize)] pub struct Evidence { pub name: String, From c70c2b560e35ea813f0f3524288ae99bd3788f1f Mon Sep 17 00:00:00 2001 From: gkumbhat Date: Wed, 24 Dec 2025 14:24:42 -0600 Subject: [PATCH 07/10] :iphone: Add python example Signed-off-by: gkumbhat --- examples/python_orchestrator.py | 55 +++++++++++++++++++++++---------- 1 file changed, 38 insertions(+), 17 deletions(-) diff --git a/examples/python_orchestrator.py b/examples/python_orchestrator.py index 52b144298..0e7b62473 100644 --- a/examples/python_orchestrator.py +++ b/examples/python_orchestrator.py @@ -1,9 +1,14 @@ import asyncio import logging +from openai.types.chat import ChatCompletionMessage from fms_guardrails_orchestr8 import ( + ChatCompletionsRequest, + ChatCompletion, + DetectorConfig, GuardrailsOrchestrator, # DetectorParams, + # ChatCompletionMessage, # Can be imported from openai.types TextContentDetectionRequest, TextContentDetectionResult, ) @@ -19,9 +24,6 @@ orch8 = GuardrailsOrchestrator(config_path=CONFIG_FILE, start_up_health_check=False) async def detect_content(): - # Showing async initialization - # orch8 = await get_guardrails_orchestrator(config_path=CONFIG_FILE, start_up_health_check=False) - try: request = TextContentDetectionRequest( content="This is stupid text.", @@ -39,22 +41,41 @@ async def detect_content(): print(result) -# async def detect_context(): -# request = PyContextDocsHttpRequest( -# content="This is a good document", -# context_type=PyContextType.DOCUMENT, -# context=["Document 1", "Document 2", "Document 3"], -# detectors={ -# "granite-guardian-context": { -# "risk_name": "context_relevance" -# } -# } -# ) -# result = await orch8.detect_context_documents(request) -# print(result) +async def chat_completion_detection(): + stream = False + try: + request = ChatCompletionsRequest( + model="meta-llama/Llama-3.2-1B-Instruct", + messages=[ChatCompletionMessage( + content="This is completely stupid", + role="assistant" + ).dict()], # Note: Currently we need to convert these to dict form or pass dict directly or use our own typ + detectors=DetectorConfig( + input = { + "en_syntax_slate.38m.hap": {} + }, + ), + stream=stream, + ) + if not stream: + result = await orch8.chat_completions_detection(request) + print(result) + breakpoint() + else: + async for response in await orch8.chat_completions_detection(request): + if len(response.choices) > 0: + print(response.choices[0].delta.content) + if len(response.detections) > 0: + print(response.detections) + # breakpoint() + except Exception as ex: + print(ex) + raise ex if __name__ == "__main__": - asyncio.run(detect_content()) + # asyncio.run(detect_content()) # asyncio.run(detect_context()) + + asyncio.run(chat_completion_detection()) \ No newline at end of file From 8c9ea906e103e57824e69e7cb43dd883f6438043 Mon Sep 17 00:00:00 2001 From: gkumbhat Date: Mon, 29 Dec 2025 10:59:52 -0700 Subject: [PATCH 08/10] :wastebasket: Cleanup Signed-off-by: gkumbhat --- .../python_interface/orchestrator.rs | 36 ++----------------- 1 file changed, 3 insertions(+), 33 deletions(-) diff --git a/src/orchestrator/python_interface/orchestrator.rs b/src/orchestrator/python_interface/orchestrator.rs index c2c583cdd..13609bf8b 100644 --- a/src/orchestrator/python_interface/orchestrator.rs +++ b/src/orchestrator/python_interface/orchestrator.rs @@ -85,7 +85,6 @@ impl GuardrailsOrchestrator { future_into_py(py, async move { match orchestrator.handle(task).await { Ok(response) => { - println!("Response: {:?}", response); Ok(response) } // TODO: Handle errors properly with correct types @@ -112,27 +111,11 @@ impl GuardrailsOrchestrator { future_into_py(py, async move { match orchestrator.handle(task).await { Ok(ChatCompletionsResponse::Unary(response)) => { - println!("Response: {:?}", response); - // response.into_pyobject(py) - // response.into_pyobject(py).map(|obj| obj.into()) - // Ok(Python::attach(|py| pythonize(py, &*response))) - // Ok(Python::attach(|py| response.into_py_any(py))) - // let result_py_object: Py = Python::attach(|py| { - // let bound_any = pythonize(py, &response) - // .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?; - - // // 3. .unbind() converts Bound<'py, PyAny> to Py (aka PyObject) - // // This "erases" the lifetime so it can safely leave the closure - // Ok::, PyErr>(bound_any.unbind()) - // })?; - - // Ok(result_py_object) + let py_obj = Python::attach(|py| { let bound_any = Bound::new(py, *response)?.into_any(); - Ok::(bound_any.unbind()) - // bound_any.unbind() - // bound_any + Ok::, PyErr>(bound_any.unbind()) }); // Ok(py_obj) println!("{:?}", py_obj); @@ -143,23 +126,10 @@ impl GuardrailsOrchestrator { let py_stream = PyChatCompletionsStream { receiver: Arc::new(Mutex::new(receiver_stream)), }; - // Ok(Python::attach(|py| pythonize(py, py_stream))) - // Ok(Python::attach(|py| py_stream.into_py_any(py))) - // py_stream.into_pyobject(py) - // py_stream.into_pyobject(py).map(|obj| obj.into()) - // let result_py_object: Py = Python::attach(|py| { - // let bound_any = pythonize(py, &py_stream) - // .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?; - - // // 3. .unbind() converts Bound<'py, PyAny> to Py (aka PyObject) - // // This "erases" the lifetime so it can safely leave the closure - // Ok::, PyErr>(bound_any.unbind()) - // })?; - // Ok(result_py_object) let py_obj = Python::attach(|py| { let bound_any = Bound::new(py, py_stream)?.into_any(); - Ok::(bound_any.unbind()) + Ok::, PyErr>(bound_any.unbind()) }); py_obj } From a613282d86004dc5ad8abcef42326fe534412358 Mon Sep 17 00:00:00 2001 From: gkumbhat Date: Fri, 2 Jan 2026 13:13:15 -0600 Subject: [PATCH 09/10] :art: Fix linting Signed-off-by: gkumbhat --- src/orchestrator/python_interface/orchestrator.rs | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/src/orchestrator/python_interface/orchestrator.rs b/src/orchestrator/python_interface/orchestrator.rs index 13609bf8b..bd479c255 100644 --- a/src/orchestrator/python_interface/orchestrator.rs +++ b/src/orchestrator/python_interface/orchestrator.rs @@ -84,9 +84,7 @@ impl GuardrailsOrchestrator { future_into_py(py, async move { match orchestrator.handle(task).await { - Ok(response) => { - Ok(response) - } + Ok(response) => Ok(response), // TODO: Handle errors properly with correct types Err(error) => Err(PyTypeError::new_err(error.to_string())), } @@ -111,7 +109,6 @@ impl GuardrailsOrchestrator { future_into_py(py, async move { match orchestrator.handle(task).await { Ok(ChatCompletionsResponse::Unary(response)) => { - let py_obj = Python::attach(|py| { let bound_any = Bound::new(py, *response)?.into_any(); From 3107873ac20fcb98bbe8a57a7b9c0d883e7bb7db Mon Sep 17 00:00:00 2001 From: gkumbhat Date: Fri, 2 Jan 2026 13:13:49 -0600 Subject: [PATCH 10/10] :sparkles: Add langchain example Signed-off-by: gkumbhat --- examples/langchain/__init__.py | 0 .../langchain/content_detect_middleware.py | 365 ++++++++++++++++++ examples/langchain/run.py | 54 +++ 3 files changed, 419 insertions(+) create mode 100644 examples/langchain/__init__.py create mode 100644 examples/langchain/content_detect_middleware.py create mode 100644 examples/langchain/run.py diff --git a/examples/langchain/__init__.py b/examples/langchain/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/examples/langchain/content_detect_middleware.py b/examples/langchain/content_detect_middleware.py new file mode 100644 index 000000000..2cc0d8d69 --- /dev/null +++ b/examples/langchain/content_detect_middleware.py @@ -0,0 +1,365 @@ +""""Custom Middleware for langchain + +Supported Detectors: Content Detector +""" +from __future__ import annotations + +import asyncio +from typing import TYPE_CHECKING, Any, Dict, List, Literal + +from langchain_core.messages import AIMessage, AnyMessage, HumanMessage, ToolMessage +from langchain.agents.middleware.types import AgentMiddleware, AgentState, hook_config + +if TYPE_CHECKING: + from langgraph.runtime import Runtime + +## Guardrails orchestrator deps +from fms_guardrails_orchestr8 import ( + GuardrailsOrchestrator, + TextContentDetectionRequest, +) + + +class ContentDetectionError(Exception): + """Raised when configured to block on detected harmful content""" + + def __init__(self, detections: List[Any], *args): + count = len(detections) + msg = f"Detected {count} instances(s) of harm in the content" + super().__init__(msg, *args) + + +class ContentDetectionMiddleware(AgentMiddleware): + """ + + This middleware runs content detections and applies configurable strategies + to handle them. It can run detections on both user input and agent output. + + Strategies: + + - `block`: Raise an exception when harmful content is detected + - `mask`: Partially mask harmful content (e.g., `****-****-****-1234` for credit card) + + Strategy Selection Guide: + + | Strategy | Preserves Identity? | Best For | + | -------- | ------------------- | --------------------------------------- | + | `block` | N/A | Avoid harmful content completely | + | `mask` | No | Human readability, customer service UIs | + + Example: + ```python + from langchain.agents.middleware import ContentDetectionMiddleware + from langchain.agents import create_agent + + # Redact all harmful content in user input + agent = create_agent( + "openai:gpt-5", + middleware=[ + ContentDetectionMiddleware(detectors=["en_syntax_slate.38m.hap"], strategy="mask", orchestrator_config="config.yaml"), + ], + ) + ``` + """ + + def __init__( + self, + *, + strategy: Literal["block", "mask"] = "mask", + detectors: List[str] | Dict[str, Any], + orchestrator_config: str | GuardrailsOrchestrator, + apply_to_input: bool = True, + apply_to_output: bool = False, + apply_to_tool_results: bool = False, + ) -> None: + """Initialize the Content detection middleware. + + Args: + strategy: How to handle detections. + + Options: + + * `block`: Raise error when harmful content is detected + * `mask`: Partially mask detected content (show last few characters) + + detectors: List of name of detectors. + * If `List[str]`: List of names of the content detection models + * If `List[Dir]`: List of detector names along with their individual config + orchestrator_config: Path to orchestrator config or initialized orchestrator instance + apply_to_input: Whether to check user messages before model call. + apply_to_output: Whether to check AI messages after model call. + apply_to_tool_results: Whether to check tool result messages after tool execution. + + Raises: + ContentDetectionError: When hamful content is detected and strategy is to block + """ + super().__init__() + + self.apply_to_input = apply_to_input + self.apply_to_output = apply_to_output + self.apply_to_tool_results = apply_to_tool_results + + self.strategy = strategy + self.detectors = detectors + if isinstance(orchestrator_config, GuardrailsOrchestrator): + self.orchestrator = orchestrator_config + else: + self.orchestrator = GuardrailsOrchestrator(config_path=orchestrator_config, start_up_health_check=False) + + @property + def name(self) -> str: + """Name of the middleware.""" + return f"{self.__class__.__name__}" + + async def _process_content(self, content: str) -> tuple[str, List[Dict[str, Any]]]: + """Apply the configured redaction rule to the provided content.""" + + if isinstance(self.detectors, dict): + detectors = self.detectors + else: + detectors = { detector_name: {} for detector_name in self.detectors } + + request = TextContentDetectionRequest ( + content=content, + detectors= detectors + ) + detections = await self.orchestrator.content_detection(request) + + if detections: + sanitized = self._apply_strategy(content, detections, self.strategy) + return sanitized, detections + else: + return content, detections + + + def _apply_strategy(self, content, detections, strategy): + """Function to apply strategy for detections""" + + if strategy == "block": + raise ContentDetectionError(detections) + + if strategy == "mask": + return self._apply_mask_strategy(content, detections) + + raise ValueError("Unknown strategy") + + + def _apply_mask_strategy(self, content, detections, unmasked_char = 3): + """Function to mask content for every detections + + Args: + content: str + detections: List containing detection dicts, which contains span and text + unmasked_char: Number of characters to keep as-is and mask all others + + Returns: + Masked content + """ + masked_content = content + # Note: we assume orchestrator is already returning sorted + # detections here based on spans (which it does automatically) + for detection in detections: + start = detection["start"] + end = detection["end"] + text = detection["text"] + # TODO: Figure out thresholding here + masked_text = f"{"*" * (len(text) - unmasked_char)}{text[-unmasked_char: ]}" + masked_content = masked_content[:start] + masked_text + masked_content[end:] + + return masked_content + + + @hook_config(can_jump_to=["end"]) + def before_model( + self, + state: AgentState, + runtime: Runtime, + ) -> dict[str, Any] | None: + """Check user messages and tool results for detection before model invocation. + + Args: + state: The current agent state. + runtime: The langgraph runtime. + + Returns: + Updated state with harmful detection handled according to strategy, or `None` if no + harmful content is detected. + + Raises: + ContentDetectionError: If harm is detected and strategy is `'block'`. + """ + return asyncio.run(self.abefore_model(state, runtime)) + + @hook_config(can_jump_to=["end"]) + async def abefore_model( + self, + state: AgentState, + runtime: Runtime, + ) -> dict[str, Any] | None: + """Async check user messages and tool results for detection before model invocation. + + Args: + state: The current agent state. + runtime: The langgraph runtime. + + Returns: + Updated state with harmful detection handled according to strategy, or `None` if no + harmful content is detected. + + Raises: + ContentDetectionError: If harm is detected and strategy is `'block'`. + """ + if not self.apply_to_input and not self.apply_to_tool_results: + return None + + messages = state["messages"] + if not messages: + return None + + new_messages = list(messages) + any_modified = False + + # Check user input if enabled + if self.apply_to_input: + # Get last user message + last_user_msg = None + last_user_idx = None + for i in range(len(messages) - 1, -1, -1): + if isinstance(messages[i], HumanMessage): + last_user_msg = messages[i] + last_user_idx = i + break + + if last_user_idx is not None and last_user_msg and last_user_msg.content: + # Detect harmful content in message + content = str(last_user_msg.content) + new_content, matches = await self._process_content(content) + + if matches: + updated_message: AnyMessage = HumanMessage( + content=new_content, + id=last_user_msg.id, + name=last_user_msg.name, + ) + + new_messages[last_user_idx] = updated_message + any_modified = True + + # Check tool results if enabled + if self.apply_to_tool_results: + # Find the last AIMessage, then process all `ToolMessage` objects after it + last_ai_idx = None + for i in range(len(messages) - 1, -1, -1): + if isinstance(messages[i], AIMessage): + last_ai_idx = i + break + + if last_ai_idx is not None: + # Get all tool messages after the last AI message + for i in range(last_ai_idx + 1, len(messages)): + msg = messages[i] + if isinstance(msg, ToolMessage): + tool_msg = msg + if not tool_msg.content: + continue + + content = str(tool_msg.content) + new_content, matches = await self._process_content(content) + + if not matches: + continue + + # Create updated tool message + updated_message = ToolMessage( + content=new_content, + id=tool_msg.id, + name=tool_msg.name, + tool_call_id=tool_msg.tool_call_id, + ) + + new_messages[i] = updated_message + any_modified = True + + if any_modified: + return {"messages": new_messages} + + return None + + def after_model( + self, + state: AgentState, + runtime: Runtime, + ) -> dict[str, Any] | None: + """Check AI messages for harmful content after model invocation. + + Args: + state: The current agent state. + runtime: The langgraph runtime. + + Returns: + Updated state with harmful content handled according to strategy, or `None` if no + harmful content is detected. + + Raises: + ContentDetectionError: If harm is detected and strategy is `'block'`. + """ + return asyncio.run(self.aafter_model(state, runtime)) + + async def aafter_model( + self, + state: AgentState, + runtime: Runtime, + ) -> dict[str, Any] | None: + """Async check AI messages for harmful content after model invocation. + + Args: + state: The current agent state. + runtime: The langgraph runtime. + + Returns: + Updated state with harmful content handled according to strategy, or `None` if no + harmful content is detected. + + Raises: + ContentDetectionError: If harm is detected and strategy is `'block'`. + """ + if not self.apply_to_output: + return None + + messages = state["messages"] + if not messages: + return None + + # Get last AI message + last_ai_msg = None + last_ai_idx = None + for i in range(len(messages) - 1, -1, -1): + msg = messages[i] + if isinstance(msg, AIMessage): + last_ai_msg = msg + last_ai_idx = i + break + + if last_ai_idx is None or not last_ai_msg or not last_ai_msg.content: + return None + + # Detect any harmful content in message + content = str(last_ai_msg.content) + new_content, matches = await self._process_content(content) + + if not matches: + return None + + # Create updated message + updated_message = AIMessage( + content=new_content, + id=last_ai_msg.id, + name=last_ai_msg.name, + tool_calls=last_ai_msg.tool_calls, + ) + + # Return updated messages + new_messages = list(messages) + new_messages[last_ai_idx] = updated_message + + return {"messages": new_messages} diff --git a/examples/langchain/run.py b/examples/langchain/run.py new file mode 100644 index 000000000..4448cdbf8 --- /dev/null +++ b/examples/langchain/run.py @@ -0,0 +1,54 @@ +"""Example showing how to use detector middleware +""" + + +import os +from langchain.agents import create_agent +from langchain.chat_models import init_chat_model + +from content_detect_middleware import ContentDetectionMiddleware +from fms_guardrails_orchestr8 import ( + GuardrailsOrchestrator, +) + + + +CONFIG_FILE = "config/local_config.yaml" + +# Showing sync initialization +orch8 = GuardrailsOrchestrator(config_path=CONFIG_FILE, start_up_health_check=False) + +model = init_chat_model( + model="meta-llama/Llama-3.2-1B-Instruct", + model_provider="openai", + base_url="http://localhost:3001/v1/", + api_key="DUMMY_API_KEY" +) + +agent = create_agent( + model=model, + # Custom planning instructions can be added via middleware + middleware=[ + ContentDetectionMiddleware( + detectors=["en_syntax_slate.38m.hap"], + strategy="block", + orchestrator_config=orch8, + apply_to_input=False, + apply_to_output=True, + ), + ], +) + +############ Non-Streaming Example ############ +query = "What is the weather in Tokyo?" +response = agent.invoke( + {"messages": [{"role": "user", "content": query}]} +) + +print("Response from agent: \n", response) + +############ Streaming Example ############ + +query = "What is the weather in Tokyo?" +for chunk in agent.stream({"messages": [{"role": "user", "content": query}]}): + print("Response from agent: \n", chunk) \ No newline at end of file