diff --git a/Cargo.lock b/Cargo.lock index ed2188565..da8e6db9b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -84,6 +84,12 @@ version = "1.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0" +[[package]] +name = "autocfg" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" + [[package]] name = "axum" version = "0.8.7" @@ -464,6 +470,9 @@ dependencies = [ "opentelemetry_sdk", "pin-project-lite", "prost", + "pyo3", + "pyo3-async-runtimes", + "pythonize", "rand", "reqwest", "rustls", @@ -1000,6 +1009,15 @@ dependencies = [ "hashbrown 0.16.1", ] +[[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" @@ -1135,6 +1153,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" @@ -1572,6 +1599,104 @@ dependencies = [ "pulldown-cmark", ] +[[package]] +name = "pyo3" +version = "0.27.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab53c047fcd1a1d2a8820fe84f05d6be69e9526be40cb03b73f86b6b03e6d87d" +dependencies = [ + "hashbrown 0.16.1", + "indoc", + "libc", + "memoffset", + "once_cell", + "portable-atomic", + "pyo3-build-config", + "pyo3-ffi", + "pyo3-macros", + "serde", + "unindent", +] + +[[package]] +name = "pyo3-async-runtimes" +version = "0.27.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57ddb5b570751e93cc6777e81fee8087e59cd53b5043292f2a6d59d5bd80fdfd" +dependencies = [ + "futures", + "once_cell", + "pin-project-lite", + "pyo3", + "pyo3-async-runtimes-macros", + "tokio", +] + +[[package]] +name = "pyo3-async-runtimes-macros" +version = "0.27.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bcd7d70ee0ca1661c40407e6f84e4463ef2658c90a9e2fbbd4515b2bcdfcaeca" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "pyo3-build-config" +version = "0.27.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b455933107de8642b4487ed26d912c2d899dec6114884214a0b3bb3be9261ea6" +dependencies = [ + "target-lexicon", +] + +[[package]] +name = "pyo3-ffi" +version = "0.27.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1c85c9cbfaddf651b1221594209aed57e9e5cff63c4d11d1feead529b872a089" +dependencies = [ + "libc", + "pyo3-build-config", +] + +[[package]] +name = "pyo3-macros" +version = "0.27.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0a5b10c9bf9888125d917fb4d2ca2d25c8df94c7ab5a52e13313a07e050a3b02" +dependencies = [ + "proc-macro2", + "pyo3-macros-backend", + "quote", + "syn", +] + +[[package]] +name = "pyo3-macros-backend" +version = "0.27.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "03b51720d314836e53327f5871d4c0cfb4fb37cc2c4a11cc71907a86342c40f9" +dependencies = [ + "heck", + "proc-macro2", + "pyo3-build-config", + "quote", + "syn", +] + +[[package]] +name = "pythonize" +version = "0.27.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a3a8f29db331e28c332c63496cfcbb822aca3d7320bc08b655d7fd0c29c50ede" +dependencies = [ + "pyo3", + "serde", +] + [[package]] name = "quinn" version = "0.11.9" @@ -2145,6 +2270,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" @@ -2540,6 +2671,12 @@ version = "1.0.22" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9312f7c4f6ff9069b165498234ce8be658059c6728633667c526e27dc2cf1df5" +[[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 391aa5b0b..5d3300b65 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.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", 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 diff --git a/examples/python_orchestrator.py b/examples/python_orchestrator.py new file mode 100644 index 000000000..0e7b62473 --- /dev/null +++ b/examples/python_orchestrator.py @@ -0,0 +1,81 @@ +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, +) + + +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/local_config.yaml" + +# Showing sync initialization +orch8 = GuardrailsOrchestrator(config_path=CONFIG_FILE, start_up_health_check=False) + +async def detect_content(): + try: + request = TextContentDetectionRequest( + content="This is stupid text.", + detectors= { + "en_syntax_slate.38m.hap": { + } + } + ) + + except Exception as ex: + print(ex) + raise ex + + result = await orch8.content_detection(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_context()) + + asyncio.run(chat_completion_detection()) \ No newline at end of file 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/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 fc60b477d..8e59c35cb 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -28,3 +28,50 @@ pub mod server; pub mod utils; #[allow(unused_imports)] pub(crate) use utils::AsUriExt; + +use pyo3::prelude::*; + +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(()) +} diff --git a/src/models.rs b/src/models.rs index f15e78722..b57cb9df8 100644 --- a/src/models.rs +++ b/src/models.rs @@ -19,6 +19,8 @@ use std::collections::{BTreeMap, HashMap}; +use pyo3::{conversion::FromPyObject, prelude::*, pyclass, types::PyDict}; +use pythonize::depythonize; use serde::{Deserialize, Serialize}; use crate::{ @@ -47,7 +49,7 @@ pub struct InfoParams { /// Parameters relevant to each detector #[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] -pub struct DetectorParams(BTreeMap); +pub struct DetectorParams(pub BTreeMap); pub type Metadata = BTreeMap; @@ -367,16 +369,18 @@ 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> { @@ -396,11 +400,13 @@ 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 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. @@ -558,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 @@ -977,6 +984,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 +1023,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, @@ -1427,3 +1436,51 @@ 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<'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 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() + } +} diff --git a/src/orchestrator.rs b/src/orchestrator.rs index 19ff86ccd..af4830038 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}; @@ -154,3 +155,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> { + +// 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) + +// } diff --git a/src/orchestrator/python_interface.rs b/src/orchestrator/python_interface.rs new file mode 100644 index 000000000..48054bef7 --- /dev/null +++ b/src/orchestrator/python_interface.rs @@ -0,0 +1,5 @@ +pub mod orchestrator; +pub use orchestrator::*; + +pub mod models; +pub use models::*; diff --git a/src/orchestrator/python_interface/models.rs b/src/orchestrator/python_interface/models.rs new file mode 100644 index 000000000..16b0aaa98 --- /dev/null +++ b/src/orchestrator/python_interface/models.rs @@ -0,0 +1,270 @@ +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 tokio::sync::{Mutex, mpsc}; + +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; + + fn extract(obj: Borrowed<'a, 'py, PyAny>) -> PyResult { + // 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)) + } +} + +#[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 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(), + } + } +} + +// 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 new file mode 100644 index 000000000..bd479c255 --- /dev/null +++ b/src/orchestrator/python_interface/orchestrator.rs @@ -0,0 +1,137 @@ +use std::sync::Arc; + +use axum::http::HeaderMap; +use tokio::sync::Mutex; + +use pyo3::{ + exceptions::{PyOSError, PyTypeError}, + prelude::*, +}; +use pyo3_async_runtimes::tokio::future_into_py; + +use crate::{ + clients::openai::ChatCompletionsResponse, + config::OrchestratorConfig, + models::TextContentDetectionHttpRequest, + orchestrator::{ + Orchestrator, + handlers::{Handle, chat_completions_detection, text_content_detection}, + }, +}; + +use crate::{ + orchestrator::python_interface::models::{PyChatCompletionsRequest, PyChatCompletionsStream}, + utils::trace, +}; + +// Define orchestrator +#[pyclass] +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, +} + +#[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), + }) + }) + }) + } + + 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) => Ok(response), + // TODO: Handle errors properly with correct types + Err(error) => Err(PyTypeError::new_err(error.to_string())), + } + }) + } + + 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)) => { + let py_obj = Python::attach(|py| { + let bound_any = Bound::new(py, *response)?.into_any(); + + Ok::, PyErr>(bound_any.unbind()) + }); + // Ok(py_obj) + println!("{:?}", py_obj); + py_obj + } + + Ok(ChatCompletionsResponse::Streaming(receiver_stream)) => { + let py_stream = PyChatCompletionsStream { + receiver: Arc::new(Mutex::new(receiver_stream)), + }; + + let py_obj = Python::attach(|py| { + let bound_any = Bound::new(py, py_stream)?.into_any(); + Ok::, PyErr>(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,