diff --git a/.cargo/config.toml b/.cargo/config.toml new file mode 100644 index 00000000..cc01b33f --- /dev/null +++ b/.cargo/config.toml @@ -0,0 +1,7 @@ +# RPATH $ORIGIN so libgraphbit.so finds libguardrail_ffi.so in the same dir at runtime +# (used when bundling libguardrail_ffi.so in the wheel via python-src for manylinux) +[target.x86_64-unknown-linux-gnu] +rustflags = ["-C", "link-arg=-Wl,-rpath,$ORIGIN"] + +[target.aarch64-unknown-linux-gnu] +rustflags = ["-C", "link-arg=-Wl,-rpath,$ORIGIN"] diff --git a/.github/workflows/build-artifacts-only.yml b/.github/workflows/build-artifacts-only.yml index 49b04741..5939dc47 100644 --- a/.github/workflows/build-artifacts-only.yml +++ b/.github/workflows/build-artifacts-only.yml @@ -89,21 +89,63 @@ jobs: ~/.cargo/registry/cache/ ~/.cargo/git/db/ target/ - key: ${{ runner.os }}-cargo-build-${{ matrix.platform.target }}-${{ hashFiles('**/Cargo.lock') }} + key: ${{ runner.os }}-cargo-build-v2-${{ matrix.platform.target }}-${{ hashFiles('**/Cargo.lock') }} restore-keys: | - ${{ runner.os }}-cargo-build-${{ matrix.platform.target }}- - ${{ runner.os }}-cargo-build- - ${{ runner.os }}-cargo- + ${{ runner.os }}-cargo-build-v2-${{ matrix.platform.target }}- + ${{ runner.os }}-cargo-build-v2- + ${{ runner.os }}-cargo-v2- + + - name: Set SONAME on guardrail .so (Linux only) + if: matrix.platform.runner == 'ubuntu-22.04' + shell: bash + run: | + sudo apt-get install -y patchelf + if [[ "${{ matrix.platform.target }}" == "x86_64" ]]; then + patchelf --set-soname libguardrail_ffi.so \ + vendor/guardrail/x86_64-unknown-linux-gnu/libguardrail_ffi.so + elif [[ "${{ matrix.platform.target }}" == "aarch64" ]]; then + patchelf --set-soname libguardrail_ffi.so \ + vendor/guardrail/aarch64-unknown-linux-gnu/libguardrail_ffi.so + fi + + - name: Copy guardrail .so for auditwheel (Linux only) + if: matrix.platform.runner == 'ubuntu-22.04' + shell: bash + run: | + if [[ "${{ matrix.platform.target }}" == "x86_64" ]]; then + cp vendor/guardrail/x86_64-unknown-linux-gnu/libguardrail_ffi.so \ + vendor/guardrail/libguardrail_ffi.so + cp vendor/guardrail/x86_64-unknown-linux-gnu/libguardrail_ffi.so \ + python/python-src/graphbit/libguardrail_ffi.so + elif [[ "${{ matrix.platform.target }}" == "aarch64" ]]; then + cp vendor/guardrail/aarch64-unknown-linux-gnu/libguardrail_ffi.so \ + vendor/guardrail/libguardrail_ffi.so + cp vendor/guardrail/aarch64-unknown-linux-gnu/libguardrail_ffi.so \ + python/python-src/graphbit/libguardrail_ffi.so + fi + + - name: Debug vendor contents + if: matrix.platform.runner == 'ubuntu-22.04' + shell: bash + run: | + echo "=== vendor/guardrail/ ===" + ls -la vendor/guardrail/ + echo "=== arch of copied .so ===" + file vendor/guardrail/libguardrail_ffi.so + echo "=== checking SONAME ===" + readelf -d vendor/guardrail/libguardrail_ffi.so | grep -E "SONAME|NEEDED" || echo "readelf not available" - name: Build wheels uses: messense/maturin-action@v1 + env: + LD_LIBRARY_PATH: ${{ github.workspace }}/vendor/guardrail with: target: ${{ matrix.platform.target }} - args: --release --out dist --find-interpreter ${{ matrix.platform.runner == 'ubuntu-22.04' && '--zig' || '' }} + args: --release --out dist --find-interpreter ${{ matrix.platform.runner == 'ubuntu-22.04' && '--zig' || '' }} ${{ matrix.platform.runner == 'ubuntu-22.04' && '--skip-auditwheel' || '' }} sccache: ${{ matrix.platform.runner != 'ubuntu-22.04' }} manylinux: ${{ matrix.platform.manylinux }} working-directory: python - + docker-options: '' - name: Verify wheel contents shell: bash diff --git a/.gitignore b/.gitignore index 57e1eff5..311893ef 100644 --- a/.gitignore +++ b/.gitignore @@ -187,3 +187,6 @@ $RECYCLE.BIN/ #.idea/ .flake8 .zed/ + +# allow vendor so +!vendor/**/*.so \ No newline at end of file diff --git a/Cargo.lock b/Cargo.lock index 4021a21a..9bbfcd84 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -840,6 +840,7 @@ dependencies = [ "csv", "docx-rs", "futures", + "guardrail_ffi", "jemallocator", "lopdf 0.32.0", "mimalloc", @@ -880,6 +881,14 @@ dependencies = [ "tokio", ] +[[package]] +name = "guardrail_ffi" +version = "0.1.0" +dependencies = [ + "reqwest", + "serde_json", +] + [[package]] name = "half" version = "2.7.1" @@ -2085,6 +2094,7 @@ checksum = "ab3f43e3283ab1488b624b44b0e988d0acea0b3214e694730a055cb6b2efa801" dependencies = [ "base64", "bytes", + "futures-channel", "futures-core", "futures-util", "http", diff --git a/Cargo.toml b/Cargo.toml index 263ae092..a754d9bc 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -41,7 +41,7 @@ version.workspace = true [profile.bench] debug = false incremental = false -lto = "fat" +lto = false opt-level = 3 # Workspace profiles (applied to all members) @@ -55,7 +55,8 @@ overflow-checks = true codegen-units = 1 debug = false incremental = false -lto = "fat" +# LTO disabled: prebuilt libguardrail_ffi.a has no bitcode; fat LTO would fail at link. +lto = false opt-level = 3 panic = "abort" strip = "symbols" @@ -64,7 +65,8 @@ strip = "symbols" [profile.release-python] codegen-units = 1 inherits = "release" -lto = "fat" +# LTO follows release (false) so we can link prebuilt libguardrail_ffi.a +lto = false opt-level = "s" # Optimize for size (important for Python extensions) panic = "abort" strip = "symbols" @@ -80,6 +82,7 @@ jemallocator.workspace = true [workspace] members = [ "core", + "guardrail_ffi", "python" ] resolver = "2" diff --git a/README.md b/README.md index ad84db85..20cf88fa 100644 --- a/README.md +++ b/README.md @@ -156,7 +156,7 @@ export ANTHROPIC_API_KEY=your_anthropic_api_key_here ```python import os -from graphbit import LlmConfig, Executor, Workflow, Node, tool +from graphbit import LlmConfig, Executor, Workflow, Node, tool, GuardRailPolicyConfig # Initialize and configure config = LlmConfig.openai(os.getenv("OPENAI_API_KEY"), "gpt-4o-mini") @@ -200,7 +200,9 @@ id1 = workflow.add_node(smart_agent) id2 = workflow.add_node(processor) workflow.connect(id1, id2) +# Run (optionally with a guardrail policy for PII masking/mapping) result = executor.execute(workflow) +# Or with policy: result = executor.execute(workflow, policy=GuardRailPolicyConfig.from_json('{"guardrail_policy": {"pii_rules": [...]}}')) print(f"Workflow completed: {result.is_success()}") print("\nSmart Agent Output: \n", result.get_node_output("Smart Agent")) print("\nData Processor Output: \n", result.get_node_output("Data Processor")) diff --git a/core/Cargo.toml b/core/Cargo.toml index 9ee69d4d..d2e0389a 100644 --- a/core/Cargo.toml +++ b/core/Cargo.toml @@ -1,5 +1,7 @@ [dependencies] pyo3 = {workspace = true, optional = true} +# GuardRail: prebuilt libguardrail_ffi.a only (see vendor/guardrail/README.md). +guardrail_ffi = { path = "../guardrail_ffi" } anyhow.workspace = true async-trait.workspace = true calamine.workspace = true diff --git a/core/src/lib.rs b/core/src/lib.rs index 4c20d0b0..96e4ba6e 100644 --- a/core/src/lib.rs +++ b/core/src/lib.rs @@ -61,6 +61,9 @@ pub use types::{ pub use validation::ValidationResult; pub use workflow::{Workflow, WorkflowBuilder, WorkflowExecutor}; +// Re-export guardrail types (from prebuilt libguardrail_ffi.a via guardrail_ffi crate) +pub use guardrail_ffi::{DecodeContext, EncodeContext, EncodeResult, DecodeResult, Enforcer, GuardRail, GuardRailConfig}; + /// Version information pub const VERSION: &str = env!("CARGO_PKG_VERSION"); diff --git a/core/src/workflow.rs b/core/src/workflow.rs index a0498bdb..d4e1f647 100644 --- a/core/src/workflow.rs +++ b/core/src/workflow.rs @@ -7,6 +7,7 @@ use crate::agents::AgentTrait; use crate::document_loader::DocumentLoader; use crate::errors::{GraphBitError, GraphBitResult}; use crate::graph::{NodeType, WorkflowGraph, WorkflowNode}; +use crate::{DecodeContext, EncodeContext, Enforcer}; use crate::types::{ AgentId, AgentMessage, CircuitBreaker, CircuitBreakerConfig, ConcurrencyConfig, ConcurrencyManager, ConcurrencyStats, MessageContent, NodeExecutionResult, NodeId, RetryConfig, @@ -68,6 +69,7 @@ impl Workflow { /// Validate the workflow pub fn validate(&self) -> GraphBitResult<()> { + tracing::debug!("Workflow '{:#?}' validated successfully", self.graph); self.graph.validate() } @@ -280,8 +282,15 @@ impl WorkflowExecutor { self.concurrency_manager.get_available_permits().await } - /// Execute a workflow with enhanced performance monitoring - pub async fn execute(&self, workflow: Workflow) -> GraphBitResult { + /// Execute a workflow with enhanced performance monitoring. + /// + /// When `guardrail_enforcer` is `Some`, PII is encoded before each LLM call and + /// decoded on LLM output; tool-call boundaries are handled by the executor layer. + pub async fn execute( + &self, + workflow: Workflow, + guardrail_enforcer: Option>, + ) -> GraphBitResult { let start_time = std::time::Instant::now(); // Initialize workflow context with simple constructor @@ -472,6 +481,7 @@ impl WorkflowExecutor { let circuit_breaker_config = self.circuit_breaker_config.clone(); let retry_config = self.default_retry_config.clone(); let concurrency_manager = self.concurrency_manager.clone(); + let guardrail_enforcer = guardrail_enforcer.clone(); // Use lightweight task spawning without unnecessary permit acquisition overhead let task = tokio::spawn(async move { @@ -503,6 +513,7 @@ impl WorkflowExecutor { circuit_breakers_clone, circuit_breaker_config, retry_config, + guardrail_enforcer, ) .await }); @@ -633,6 +644,7 @@ impl WorkflowExecutor { circuit_breakers: Arc>>, circuit_breaker_config: CircuitBreakerConfig, retry_config: Option, + guardrail_enforcer: Option>, ) -> GraphBitResult { let start_time = std::time::Instant::now(); let mut attempt = 0; @@ -678,6 +690,7 @@ impl WorkflowExecutor { &node.config, context.clone(), agents.clone(), + guardrail_enforcer.clone(), ) .await } @@ -785,7 +798,8 @@ impl WorkflowExecutor { } } - /// Execute an agent node (static version) + /// Execute an agent node (static version). + /// When `guardrail_enforcer` is `Some`, encodes prompt before LLM and decodes response after. async fn execute_agent_node_static( current_node_id: &NodeId, agent_id: &crate::types::AgentId, @@ -793,6 +807,7 @@ impl WorkflowExecutor { node_config: &std::collections::HashMap, context: Arc>, agents: Arc>>>, + guardrail_enforcer: Option>, ) -> GraphBitResult { // Use read lock for better performance let agents_guard = agents.read().await; @@ -959,6 +974,7 @@ impl WorkflowExecutor { current_node_id, &node_name, context.clone(), + guardrail_enforcer.clone(), ) .await; tracing::info!("Agent with tools execution result: {:?}", result); @@ -967,9 +983,27 @@ impl WorkflowExecutor { // Execute agent without tools (original behavior) tracing::info!("NO TOOLS DETECTED - using standard agent execution"); + // Guardrail: encode prompt before sending to LLM; combine injection text + payload + let prompt_for_llm = if let Some(ref enforcer) = guardrail_enforcer { + tracing::debug!("Guardrail: encoding prompt before LLM call (sensitive data will be masked)"); + let result = enforcer.encode( + serde_json::Value::String(resolved_prompt.clone()), + EncodeContext::Llm, + ); + tracing::debug!("Guardrail: prompt encoded for LLM (payload only): {}", result.payload.as_str().unwrap_or("")); + tracing::debug!("[GuardRail] encoded prompt (sent to LLM, payload only): {}", result.payload.as_str().unwrap_or("")); + format!( + "{}{}", + result.signature_injection_text, + result.payload.as_str().unwrap_or("") + ) + } else { + resolved_prompt.clone() + }; + // Call LLM provider directly to capture metadata use crate::llm::LlmRequest; - let mut request = LlmRequest::new(resolved_prompt.clone()); + let mut request = LlmRequest::new(prompt_for_llm.clone()); // Apply node-level configuration overrides (temperature, max_tokens, etc.) if let Some(temp_value) = node_config.get("temperature") { @@ -990,7 +1024,46 @@ impl WorkflowExecutor { let llm_response = agent.llm_provider().complete(request).await?; let llm_duration_ms = llm_start.elapsed().as_secs_f64() * 1000.0; + if guardrail_enforcer.is_some() { + tracing::debug!( + "[GuardRail] raw LLM response (before decode): content={:?}, tool_calls={:?}", + llm_response.content, + llm_response.tool_calls + ); + } + + // Guardrail: decode LLM output before storing in context + let llm_response = if let Some(ref enforcer) = guardrail_enforcer { + tracing::debug!("Guardrail: decoding LLM response (rehydrating for context) llm_response.content: {}", llm_response.content); + tracing::debug!("Guardrail: decoding LLM response (rehydrating for context)"); + let payload = serde_json::json!({ + "content": llm_response.content, + "tool_calls": llm_response.tool_calls + }); + let decoded_result = enforcer.decode(payload, DecodeContext::LlmResponse); + tracing::debug!("Guardrail: LLM response decoded"); + let content = decoded_result + .payload + .get("content") + .and_then(|v| v.as_str()) + .map(String::from) + .unwrap_or_else(|| llm_response.content.clone()); + let tool_calls = decoded_result + .payload + .get("tool_calls") + .and_then(|v| serde_json::from_value(v.clone()).ok()) + .unwrap_or_else(|| llm_response.tool_calls.clone()); + crate::llm::LlmResponse { + content, + tool_calls, + ..llm_response + } + } else { + llm_response + }; + // Store LLM response metadata AND request prompt in context for observability + // When GuardRail is on, store the encoded prompt so PII is never written to metadata. { // First, get the node name before mutable borrow let node_name = { @@ -1004,14 +1077,18 @@ impl WorkflowExecutor { .unwrap_or_else(|| "unknown".to_string()) }; + let prompt_for_metadata = guardrail_enforcer + .as_ref() + .map(|_| prompt_for_llm.clone()) + .unwrap_or_else(|| resolved_prompt.clone()); + // Now store the metadata let mut ctx = context.lock().await; if let Ok(mut response_metadata) = serde_json::to_value(&llm_response) { - // Add the request prompt to the metadata if let Some(obj) = response_metadata.as_object_mut() { obj.insert( "prompt".to_string(), - serde_json::Value::String(resolved_prompt.clone()), + serde_json::Value::String(prompt_for_metadata), ); // Add LLM call duration for accurate latency tracking obj.insert( @@ -1046,7 +1123,8 @@ impl WorkflowExecutor { } } - /// Execute an agent with tool calling orchestration + /// Execute an agent with tool calling orchestration. + /// When `guardrail_enforcer` is `Some`, encodes prompt before LLM and decodes response after. async fn execute_agent_with_tools( _agent_id: &crate::types::AgentId, prompt: &str, @@ -1055,10 +1133,29 @@ impl WorkflowExecutor { node_id: &NodeId, node_name: &str, context: Arc>, + guardrail_enforcer: Option>, ) -> GraphBitResult { tracing::info!("Starting execute_agent_with_tools for agent: {_agent_id}"); use crate::llm::{LlmRequest, LlmTool}; + // Guardrail: encode prompt before sending to LLM; combine injection text + payload + let prompt_for_llm = if let Some(ref enforcer) = guardrail_enforcer { + tracing::debug!("Guardrail: encoding prompt before LLM call (tool path; sensitive data masked)"); + let result = enforcer.encode( + serde_json::Value::String(prompt.to_string()), + EncodeContext::Llm, + ); + tracing::debug!("Guardrail: prompt encoded for LLM (payload only): {}", result.payload.as_str().unwrap_or_default()); + tracing::debug!("[GuardRail] encoded prompt (sent to LLM, payload only): {}", result.payload.as_str().unwrap_or_default()); + format!( + "{}{}", + result.signature_injection_text, + result.payload.as_str().unwrap_or_default() + ) + } else { + prompt.to_string() + }; + // Extract tool schemas from node config let tool_schemas = node_config .get("tool_schemas") @@ -1079,8 +1176,8 @@ impl WorkflowExecutor { } } - // Create initial LLM request with tools - let mut request = LlmRequest::new(prompt); + // Create initial LLM request with tools (using encoded prompt when guardrail is active) + let mut request = LlmRequest::new(prompt_for_llm.clone()); for tool in &tools { request = request.with_tool(tool.clone()); } @@ -1125,37 +1222,65 @@ impl WorkflowExecutor { // Measure LLM call duration and capture execution timestamp let execution_timestamp = chrono::Utc::now(); let llm_start = std::time::Instant::now(); - let llm_response = agent.llm_provider().complete(request).await?; + let mut llm_response = agent.llm_provider().complete(request).await?; let llm_duration_ms = llm_start.elapsed().as_secs_f64() * 1000.0; + if guardrail_enforcer.is_some() { + tracing::debug!( + "[GuardRail] raw LLM response (before decode): content={:?}, tool_calls={:?}", + llm_response.content, + llm_response.tool_calls + ); + } + + // Guardrail: decode LLM output before storing in context + if let Some(ref enforcer) = guardrail_enforcer { + tracing::debug!("Guardrail: decoding LLM response (tool path; rehydrating for context) llm_response.content: {}", llm_response.content); + tracing::debug!("Guardrail: decoding LLM response (tool path; rehydrating for context)"); + let payload = serde_json::json!({ + "content": llm_response.content, + "tool_calls": llm_response.tool_calls + }); + let decoded_result = enforcer.decode(payload, DecodeContext::LlmResponse); + tracing::debug!("Guardrail: LLM response decoded"); + if let Some(c) = decoded_result.payload.get("content").and_then(|v| v.as_str()) { + llm_response.content = c.to_string(); + } + if let Some(tc) = decoded_result.payload.get("tool_calls") { + if let Ok(parsed) = serde_json::from_value(tc.clone()) { + llm_response.tool_calls = parsed; + } + } + } + // Store LLM response metadata AND request prompt in context for observability + // When GuardRail is on, store the encoded prompt so PII is never written to metadata. { let mut ctx = context.lock().await; if let Ok(mut response_metadata) = serde_json::to_value(&llm_response) { - // Add the request prompt to the metadata if let Some(obj) = response_metadata.as_object_mut() { + let prompt_for_metadata = guardrail_enforcer + .as_ref() + .map(|_| prompt_for_llm.clone()) + .unwrap_or_else(|| prompt.to_string()); obj.insert( "prompt".to_string(), - serde_json::Value::String(prompt.to_string()), + serde_json::Value::String(prompt_for_metadata), ); - // Add LLM call duration for accurate latency tracking obj.insert( "duration_ms".to_string(), serde_json::json!(llm_duration_ms), ); - // Add execution timestamp for chronological ordering obj.insert( "execution_timestamp".to_string(), serde_json::json!(execution_timestamp.to_rfc3339()), ); } - // Store by node ID ctx.metadata.insert( format!("node_response_{node_id}"), response_metadata.clone(), ); - // Store by node name ctx.metadata .insert(format!("node_response_{node_name}"), response_metadata); } @@ -1188,13 +1313,17 @@ impl WorkflowExecutor { GraphBitError::workflow_execution(format!("Failed to serialize tool calls: {e}")) })?; - // Return a structured response that the Python layer can interpret - // Include token usage for budget tracking + // Return a structured response that the Python layer can interpret. + // When GuardRail is on, pass the encoded prompt so the final LLM never sees raw PII. + let original_prompt_for_response = guardrail_enforcer + .as_ref() + .map(|_| prompt_for_llm.clone()) + .unwrap_or_else(|| prompt.to_string()); Ok(serde_json::json!({ "type": "tool_calls_required", "content": llm_response.content, "tool_calls": tool_calls_json, - "original_prompt": prompt, + "original_prompt": original_prompt_for_response, "initial_tokens_used": llm_response.usage.completion_tokens, "max_tokens_configured": node_config.get("max_tokens").and_then(|v| v.as_u64()), "message": "Tool execution should be handled by Python layer with proper tool registry" diff --git a/examples/guardrail_financial/README.md b/examples/guardrail_financial/README.md new file mode 100644 index 00000000..79794f40 --- /dev/null +++ b/examples/guardrail_financial/README.md @@ -0,0 +1,146 @@ +# GuardRail Examples + +This directory contains examples demonstrating GraphBit's GuardRail feature for protecting Personally Identifiable Information (PII) in LLM workflows. + +## How GuardRail Works + +GuardRail provides intelligent masking of sensitive data: +- **LLM sees**: Masked tokens (e.g., `[CREDIT_CARD_1]`, `[EMAIL_1]`) +- **Tools receive**: Real unmasked values for accurate processing +- **Automatic handling**: Encode/decode at LLM and tool boundaries + +## Examples + +### 1. Phone Number Sum (`guardrail_phone/`) + +**Purpose**: Demonstrate GuardRail protecting phone numbers while allowing tools to process them correctly. + +**Pattern**: +- Policy masks phone numbers matching pattern `\d{3}-\d{4}` (e.g., `123-4567`) +- Tool `sum_digits_in_phone()` receives the real unmasked number +- LLM only sees masked token (e.g., `[PHONE_NUMBER_1]`) + +**Files**: +- `guardrail_phone_policy.json` - Policy defining phone number masking rules +- `run_guardrail_phone.py` - Example workflow with tool calling + +**Run**: +```bash +.venv/bin/python examples/guardrail_phone/run_guardrail_phone.py +``` + +### 2. Financial Payment Processing (`guardrail_financial/`) + +**Purpose**: Demonstrate GuardRail protecting multiple types of sensitive financial data (credit cards, emails, SSNs) in a payment processing workflow. + +**Pattern**: +- Policy masks: + - Credit cards: `\d{4}-\d{4}-\d{4}-\d{4}` (e.g., `4532-1234-5678-9010`) + - Emails: Standard email pattern (e.g., `customer@example.com`) + - SSNs: `\d{3}-\d{2}-\d{4}` (e.g., `123-45-6789`) +- Three tools with different data requirements: + - `validate_credit_card()` - Validates the card and returns last 4 digits + - `calculate_transaction_fee()` - Computes 2% fee on amount + - `send_payment_confirmation()` - Sends confirmation to recipient email +- LLM sees only masked tokens, tools get real values + +**Files**: +- `guardrail_financial_policy.json` - Policy defining credit card, email, and SSN masking +- `run_guardrail_financial.py` - Complete payment processing workflow + +**Run**: +```bash +.venv/bin/python examples/guardrail_financial/run_guardrail_financial.py +``` + +## Key Implementation Details + +### Policy Definition (JSON) + +```json +{ + "policy_name": "example_policy", + "policy_version": "1.0.0", + "active": true, + "guardrail_policy": { + "pii_rules": [ + { + "type": "regex", + "name": "RULE_NAME", + "pattern": "regex_pattern_here" + } + ], + "masking": true, + "mapping": true + } +} +``` + +### Tool Definition + +```python +from graphbit import tool + +@tool(_description="Description shown to LLM") +def my_tool(param: str) -> str: + """ + When GuardRail is enabled, this tool receives DECODED values. + """ + print(f"[Tool received] param = {param!r}") + # Process the real unmasked value + return result +``` + +### Workflow Execution with GuardRail + +```python +from graphbit import Executor, GuardRailPolicyConfig + +# Load policy +policy = GuardRailPolicyConfig.from_file("policy.json") + +# Execute with policy - LLM gets masked, tools get real data +result = executor.execute(workflow, policy=policy) +``` + +## Testing Without OpenAI + +These examples use OpenAI (GPT-4o-mini) if `OPENAI_API_KEY` is set, otherwise fallback to Ollama. + +**Using Ollama**: +```bash +# Ensure Ollama is running +ollama pull llama3.2 +ollama serve + +# In another terminal +.venv/bin/python examples/guardrail_financial/run_guardrail_financial.py +``` + +**Using OpenAI**: +```bash +export OPENAI_API_KEY="sk-..." +.venv/bin/python examples/guardrail_financial/run_guardrail_financial.py +``` + +## Debug Output + +Run with debug logging to see GuardRail in action: +- `Guardrail: encoding prompt before LLM` - Shows data being masked +- `Guardrail: decoding tool call parameters` - Shows unmasking at tool boundary +- `[Tool received]` - Shows the real unmasked values tools receive + +Look for masked tokens in logs like: +- `` +- `` +- `` + +## Creating Your Own Example + +Follow this pattern: + +1. **Define a policy** (JSON with PII patterns) +2. **Create tools** (Python functions marked with `@tool`) +3. **Build workflow** (Node.agent with tools parameter) +4. **Execute with policy** (executor.execute(workflow, policy=policy)) +5. **Verify** (check tool received real values in debug output) diff --git a/examples/guardrail_financial/guardrail_financial_policy.json b/examples/guardrail_financial/guardrail_financial_policy.json new file mode 100644 index 00000000..46030400 --- /dev/null +++ b/examples/guardrail_financial/guardrail_financial_policy.json @@ -0,0 +1,26 @@ +{ + "policy_name": "financial_pii_policy", + "policy_version": "1.0.0", + "active": true, + "guardrail_policy": { + "pii_rules": [ + { + "type": "regex", + "name": "CREDIT_CARD", + "pattern": "\\d{4}-\\d{4}-\\d{4}-\\d{4}" + }, + { + "type": "regex", + "name": "EMAIL", + "pattern": "[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\\.[a-zA-Z]{2,}" + }, + { + "type": "regex", + "name": "SSN", + "pattern": "\\d{3}-\\d{2}-\\d{4}" + } + ], + "masking": true, + "mapping": true + } +} diff --git a/examples/guardrail_financial/run_guardrail_financial.py b/examples/guardrail_financial/run_guardrail_financial.py new file mode 100644 index 00000000..3dc05984 --- /dev/null +++ b/examples/guardrail_financial/run_guardrail_financial.py @@ -0,0 +1,153 @@ +""" +GuardRail Financial Example: Secure Payment Processing with PII Masking + +This example demonstrates how GuardRail masks sensitive financial information +(credit card numbers, emails, SSNs) from the LLM while ensuring tools receive +the actual decoded values for proper processing. + +Pattern: +- LLM sees masked tokens (e.g., [CREDIT_CARD_1], [EMAIL_1]) +- Tools receive the real unmasked values for accurate computation +- GuardRail automatically handles encode/decode boundaries +""" + +import os +import sys + +import graphbit +from graphbit import ( + Executor, + GuardRailPolicyConfig, + LlmConfig, + Node, + Workflow, + tool, + init, +) + +# using our policy +POLICY_DIR = os.path.dirname(os.path.abspath(__file__)) +POLICY_PATH = os.path.join(POLICY_DIR, "guardrail_financial_policy.json") + + +@tool(_description="Validate a credit card number and return the last 4 digits safely. Input is a full CC number in format XXXX-XXXX-XXXX-XXXX.") +def validate_credit_card(card_number: str) -> str: + """ + Validate a credit card by checking digit count and return last 4 digits. + When GuardRail is enabled, this tool receives the DECODED (real) card number. + """ + digits_only = card_number.replace("-", "") + is_valid = len(digits_only) == 16 and digits_only.isdigit() + last_four = digits_only[-4:] if len(digits_only) >= 4 else "INVALID" + + print(f"[Tool received] card_number = {card_number!r}") + print(f"[Tool validates] Valid: {is_valid}, Last 4 digits: {last_four}") + + return f"Card valid: {is_valid}, Last 4: {last_four}" + + +@tool(_description="Send a payment confirmation email. Recipient email must be in format user@domain.com") +def send_payment_confirmation(recipient_email: str, amount: str) -> str: + """ + Send payment confirmation to recipient. + When GuardRail is enabled, this tool receives the DECODED (real) email address. + """ + print(f"[Tool received] recipient_email = {recipient_email!r}, amount = {amount!r}") + + # Simulate email validation + if "@" in recipient_email and "." in recipient_email.split("@")[1]: + result = f"Confirmation email sent to {recipient_email} for amount ${amount}" + print(f"[Tool result] {result}") + return result + else: + return "Error: Invalid email format" + + +@tool(_description="Calculate transaction fee based on amount. Amount should be a number like 100.50") +def calculate_transaction_fee(amount: str) -> str: + """ + Calculate the transaction fee (2% of amount). + """ + try: + amt = float(amount) + fee = amt * 0.02 + print(f"[Tool received] amount = {amount!r} -> fee = ${fee:.2f}") + return f"Transaction fee: ${fee:.2f}" + except ValueError: + return "Error: Amount must be a number" + + +def main(): + init(enable_tracing=True, log_level="debug") + + print("=" * 80) + print("GuardRail Financial Example: Secure Payment Processing") + print("=" * 80) + print("\nScenario: Process a payment with sensitive information") + print("- LLM should NOT see real credit card, email, or SSN") + print("- Tools MUST receive actual values for proper processing\n") + + # Prefer OpenAI for reliable tool calling; fallback to Ollama + llm_config = LlmConfig.openai(os.getenv("OPENAI_API_KEY"), "gpt-4o-mini") + + executor = Executor(llm_config) + + # Load policy that masks credit cards, emails, and SSNs + if not os.path.isfile(POLICY_PATH): + print(f"Policy file not found: {POLICY_PATH}") + sys.exit(1) + + policy = GuardRailPolicyConfig.from_file(POLICY_PATH) + print(f"Loaded policy: {policy.policy_name()}, active={policy.is_active()}\n") + + workflow = Workflow("Secure Payment Processing [GuardRail]") + + payment_agent = Node.agent( + name="Payment Processor", + prompt=( + "I want to send money using my bank with this credit card and sender info:\n" + "Credit Card: 4532-1234-5678-9010\n" + "Recipient Email: customer@example.com\n" + "Amount: 250.00\n" + "Customer SSN (for verification): 123-45-6789\n" + + ), + system_prompt=( + "You are a secure payment processor. Use the available tools to validate " + "payments, calculate fees, and send confirmations. Always use the tools " + "with the exact information provided." + ), + tools=[validate_credit_card, calculate_transaction_fee, send_payment_confirmation], + ) + + workflow.add_node(payment_agent) + workflow.validate() + + # Execute WITH policy: LLM sees masked data; tools receive decoded data + print("Executing workflow with GuardRail policy...\n") + result = executor.execute(workflow, policy=policy) + + print("\n" + "=" * 80) + print("--- Result ---") + print("=" * 80) + + if result.is_success(): + out = result.get_node_output("Payment Processor") + print(f"\nAgent Output:\n{out}") + print("\n" + "-" * 80) + print("Verification Notes:") + print("- Above '[Tool received]' entries should show REAL values (unmasked)") + print(" - Credit card: 4532-1234-5678-9010") + print(" - Email: customer@example.com") + print(" - SSN: 123-45-6789") + print("\n- In debug logs you should see:") + print(" - 'Guardrail: encoding prompt before LLM'") + print(" - 'Guardrail: decoding tool call parameters before execution'") + print(" - Masked tokens like [CREDIT_CARD_1], [EMAIL_1], [SSN_1]") + print("-" * 80) + else: + print(f"Workflow failed: {result.state()}") + + +if __name__ == "__main__": + main() diff --git a/examples/guardrail_phone/guardrail_phone_policy.json b/examples/guardrail_phone/guardrail_phone_policy.json new file mode 100644 index 00000000..0831bb2f --- /dev/null +++ b/examples/guardrail_phone/guardrail_phone_policy.json @@ -0,0 +1,16 @@ +{ + "policy_name": "phone_mask_policy", + "policy_version": "1.0.0", + "active": true, + "guardrail_policy": { + "pii_rules": [ + { + "type": "regex", + "name": "PHONE_NUMBER", + "pattern": "\\d{3}-\\d{4}" + } + ], + "masking": true, + "mapping": true + } +} diff --git a/examples/guardrail_phone/run_guardrail_phone.py b/examples/guardrail_phone/run_guardrail_phone.py new file mode 100644 index 00000000..8b49a1d3 --- /dev/null +++ b/examples/guardrail_phone/run_guardrail_phone.py @@ -0,0 +1,83 @@ +import os +import sys + +import graphbit +from graphbit import ( + Executor, + GuardRailPolicyConfig, + LlmConfig, + Node, + Workflow, + tool, + init, +) + +# using our policy +POLICY_DIR = os.path.dirname(os.path.abspath(__file__)) +POLICY_PATH = os.path.join(POLICY_DIR, "guardrail_phone_policy.json") + + +@tool(_description="Sum all digits in a phone number. Input is a string that may look like 555-5555 or a token.") +def sum_digits_in_phone(phone_number: str) -> str: + """ + Compute the sum of all digits in the given phone number. + When GuardRail is enabled, this tool receives the DECODED (real) number so it can compute correctly. + """ + digits = [int(c) for c in phone_number if c.isdigit()] + total = sum(digits) + print(f"[Tool received] phone_number = {phone_number!r} -> sum of digits = {total}") + return str(total) + + +def main(): + init(enable_tracing=True, log_level="debug") + + print("GuardRail phone example: LLM should never see 123-4567; tool should always receive it.\n") + + # Prefer OpenAI for reliable tool calling; fallback to Ollama + if os.getenv("OPENAI_API_KEY"): + llm_config = LlmConfig.openai(os.getenv("OPENAI_API_KEY"), "gpt-5") + else: + print("No OPENAI_API_KEY set. Using Ollama (ollama run llama3.2).") + llm_config = LlmConfig.ollama("llama3.2") + + executor = Executor(llm_config) + + # Load policy that masks 123-4567-style numbers + if not os.path.isfile(POLICY_PATH): + print(f"Policy file not found: {POLICY_PATH}") + sys.exit(1) + policy = GuardRailPolicyConfig.from_file(POLICY_PATH) + print(f"Loaded policy: {policy.policy_name()}, active={policy.is_active()}\n") + + workflow = Workflow("Phone digits sum [GuardRail]") + agent = Node.agent( + name="Phone Agent", + prompt=( + "The user's phone number is 123-4567. " + "Use the sum_digits_in_phone tool to compute the sum of all digits in that phone number, " + "then reply with the result." + ), + system_prompt="You have a tool to sum digits in a phone number. Use it when asked.", + tools=[sum_digits_in_phone], + max_tokens=1000, + ) + workflow.add_node(agent) + workflow.validate() + + # Execute WITH policy: LLM sees masked data; tool receives decoded data + result = executor.execute(workflow, policy=policy) + + print("\n--- Result ---") + + if result.is_success(): + out = result.get_node_output("Phone Agent") + print(f"Agent output: {out}") + print("\nVerify: above '[Tool received]' should show phone_number = '123-4567' (decoded).") + print("In debug logs you should see 'Guardrail: encoding prompt before LLM' and 'decoding tool call parameters'.") + else: + print(f"Workflow failed: {result.state()}") + print("\n--- Node Response Metadata ---") + print(result.get_all_node_response_metadata()) +if __name__ == "__main__": + main() diff --git a/guardrail_ffi/Cargo.toml b/guardrail_ffi/Cargo.toml new file mode 100644 index 00000000..12048ee4 --- /dev/null +++ b/guardrail_ffi/Cargo.toml @@ -0,0 +1,12 @@ +[package] +name = "guardrail_ffi" +version = "0.1.0" +edition = "2024" +description = "GraphBit wrapper for prebuilt GuardRail C ABI (libguardrail_ffi.a)" + +[dependencies] +serde_json = "1.0" +reqwest = { version = "0.13", default-features = false, features = ["json", "rustls", "blocking"] } + +[lints.rust] +unsafe_code = "allow" diff --git a/guardrail_ffi/build.rs b/guardrail_ffi/build.rs new file mode 100644 index 00000000..a386c6ac --- /dev/null +++ b/guardrail_ffi/build.rs @@ -0,0 +1,63 @@ +use std::env; +use std::path::Path; + +fn main() { + let manifest_dir = env::var("CARGO_MANIFEST_DIR").expect("CARGO_MANIFEST_DIR"); + let target_os = env::var("CARGO_CFG_TARGET_OS").unwrap_or_default(); + let target_arch = env::var("CARGO_CFG_TARGET_ARCH").unwrap_or_default(); + + let default_lib_dir = { + let base = Path::new(&manifest_dir).join("../vendor/guardrail"); + if target_os == "linux" { + let triple = match target_arch.as_str() { + "x86_64" => "x86_64-unknown-linux-gnu", + "aarch64" => "aarch64-unknown-linux-gnu", + other => panic!("unsupported Linux arch: {other}"), + }; + base.join(triple).to_string_lossy().into_owned() + } else { + base.to_string_lossy().into_owned() + } + }; + + let lib_dir = env::var("GUARDRAIL_LIB_DIR").unwrap_or(default_lib_dir); + let lib_path = Path::new(&lib_dir); + + if !lib_path.exists() { + eprintln!( + "cargo:warning=GuardRail lib dir not found: {} (set GUARDRAIL_LIB_DIR or populate vendor/guardrail/)", + lib_dir + ); + return; + } + + println!("cargo:rerun-if-changed=../vendor/guardrail/"); + + match target_os.as_str() { + "windows" => { + println!("cargo:rustc-link-search=native={}", lib_path.display()); + println!("cargo:rustc-link-lib=dylib=guardrail_ffi"); + } + "linux" => { + // Use path relative to workspace root so no absolute build path gets embedded + // in the binary (avoids runtime "cannot open shared object file" on other machines). + let link_search = Path::new("vendor") + .join("guardrail") + .join(match target_arch.as_str() { + "x86_64" => "x86_64-unknown-linux-gnu", + "aarch64" => "aarch64-unknown-linux-gnu", + other => panic!("unsupported Linux arch: {other}"), + }); + println!( + "cargo:rustc-link-search=native={}", + link_search.display() + ); + println!("cargo:rustc-link-lib=dylib=guardrail_ffi"); + } + _ => { + // macOS + println!("cargo:rustc-link-search=native={}", lib_path.display()); + println!("cargo:rustc-link-lib=static=guardrail_ffi"); + } + } +} \ No newline at end of file diff --git a/guardrail_ffi/src/lib.rs b/guardrail_ffi/src/lib.rs new file mode 100644 index 00000000..652a1f3a --- /dev/null +++ b/guardrail_ffi/src/lib.rs @@ -0,0 +1,397 @@ +//! GuardRail FFI wrapper — links the prebuilt `libguardrail_ffi.a` only (no guardrail source). + +use std::ffi::{CStr, CString}; +use std::os::raw::{c_char, c_uint}; +use std::path::Path; +use std::sync::Arc; + +unsafe extern "C" { + fn guardrail_config_from_json(json_ptr: *const c_char, json_len: usize) -> *mut std::ffi::c_void; + fn guardrail_config_default() -> *mut std::ffi::c_void; + fn guardrail_config_clone(handle: *mut std::ffi::c_void) -> *mut std::ffi::c_void; + fn guardrail_config_drop(handle: *mut std::ffi::c_void); + fn guardrail_config_policy_name(handle: *mut std::ffi::c_void) -> *mut c_char; + fn guardrail_config_policy_version(handle: *mut std::ffi::c_void) -> *mut c_char; + fn guardrail_config_active(handle: *mut std::ffi::c_void) -> bool; + + fn guardrail_enforcer_create( + config_handle: *mut std::ffi::c_void, + workflow_id_ptr: *const c_char, + workflow_id_len: usize, + ) -> *mut std::ffi::c_void; + fn guardrail_enforcer_drop(handle: *mut std::ffi::c_void); + + fn guardrail_encode( + enforcer_handle: *mut std::ffi::c_void, + json_ptr: *const c_char, + json_len: usize, + encode_context: c_uint, + ) -> *mut c_char; + fn guardrail_decode( + enforcer_handle: *mut std::ffi::c_void, + json_ptr: *const c_char, + json_len: usize, + context: c_uint, + ) -> *mut c_char; + fn guardrail_free(ptr: *mut c_char); +} + +const CONTEXT_TOOL_BOUNDARY: c_uint = 0; +const CONTEXT_LLM_RESPONSE: c_uint = 1; +const CONTEXT_MANUAL_CALL: c_uint = 2; + +/// Encode context: Llm adds signature and instruction text; Manual does not. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum EncodeContext { + /// No signature or instruction (e.g. manual / logging). + Manual, + /// Add 3-digit signature to tokens and prepend instruction text for LLM. + Llm, +} + +const ENCODE_CONTEXT_MANUAL: c_uint = 0; +const ENCODE_CONTEXT_LLM: c_uint = 1; + +/// Decode context for rehydration. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum DecodeContext { + /// Rehydrate at tool boundary so the tool receives real PII. + ToolBoundary, + /// Rehydrate LLM output for context. + LlmResponse, + /// Explicit host decode. + ManualCall, +} + +/// Opaque config handle (refcounted via clone/drop). +pub struct GuardRailConfigInner { + pub(crate) handle: *mut std::ffi::c_void, +} + +impl Drop for GuardRailConfigInner { + fn drop(&mut self) { + if !self.handle.is_null() { + unsafe { guardrail_config_drop(self.handle) }; + self.handle = std::ptr::null_mut(); + } + } +} + +impl Clone for GuardRailConfigInner { + fn clone(&self) -> Self { + let handle = if self.handle.is_null() { + std::ptr::null_mut() + } else { + unsafe { guardrail_config_clone(self.handle) } + }; + Self { handle } + } +} + +// Opaque FFI handle: safe to Send/Sync as the C library manages thread safety. +unsafe impl Send for GuardRailConfigInner {} +unsafe impl Sync for GuardRailConfigInner {} + +/// GuardRail policy configuration. +#[derive(Clone)] +pub struct GuardRailConfig { + pub(crate) inner: Arc, +} + +impl GuardRailConfig { + /// Load from JSON string. + pub fn new(json: &str) -> Result { + let c_str = CString::new(json).map_err(|e| e.to_string())?; + let handle = unsafe { guardrail_config_from_json(c_str.as_ptr(), c_str.as_bytes().len()) }; + if handle.is_null() { + return Err("GuardRail config from_json failed".into()); + } + Ok(Self { + inner: Arc::new(GuardRailConfigInner { handle }), + }) + } + + /// Load from file. + pub fn from_file(path: &Path) -> Result { + let json = std::fs::read_to_string(path).map_err(|e| e.to_string())?; + Self::new(&json) + } + + /// Load from URL (blocking GET). + pub fn from_url(url: &str) -> Result { + let client = reqwest::blocking::Client::new(); + let json = client + .get(url) + .send() + .map_err(|e| e.to_string())? + .text() + .map_err(|e| e.to_string())?; + Self::new(&json) + } + + /// Default inactive config. + pub fn default_config() -> Self { + let handle = unsafe { guardrail_config_default() }; + assert!(!handle.is_null(), "guardrail_config_default failed"); + Self { + inner: Arc::new(GuardRailConfigInner { handle }), + } + } + + pub(crate) fn ptr(&self) -> *mut std::ffi::c_void { + self.inner.handle + } + + /// Policy name. + pub fn policy_name(&self) -> String { + if self.ptr().is_null() { + return String::new(); + } + let p = unsafe { guardrail_config_policy_name(self.ptr()) }; + if p.is_null() { + return String::new(); + } + let s = unsafe { CStr::from_ptr(p).to_string_lossy().into_owned() }; + unsafe { guardrail_free(p) }; + s + } + + /// Policy version. + pub fn policy_version(&self) -> String { + if self.ptr().is_null() { + return String::new(); + } + let p = unsafe { guardrail_config_policy_version(self.ptr()) }; + if p.is_null() { + return String::new(); + } + let s = unsafe { CStr::from_ptr(p).to_string_lossy().into_owned() }; + unsafe { guardrail_free(p) }; + s + } + + /// Whether the policy is active. + pub fn active(&self) -> bool { + if self.ptr().is_null() { + return false; + } + unsafe { guardrail_config_active(self.ptr()) } + } +} + +/// Enforcer for one workflow (encode/decode). +pub struct Enforcer { + pub(crate) handle: *mut std::ffi::c_void, +} + +impl Drop for Enforcer { + fn drop(&mut self) { + if !self.handle.is_null() { + unsafe { guardrail_enforcer_drop(self.handle) }; + self.handle = std::ptr::null_mut(); + } + } +} + +impl std::fmt::Debug for Enforcer { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Enforcer").finish_non_exhaustive() + } +} + +unsafe impl Send for Enforcer {} +unsafe impl Sync for Enforcer {} + +/// Result of encode: payload (masked only) plus optional injection text and metadata. +#[derive(Debug, Clone)] +pub struct EncodeResult { + pub payload: serde_json::Value, + /// Rule text to prepend when sending to LLM; empty when not applicable. Caller concatenates with payload. + pub signature_injection_text: String, + pub rules_applied_count: u32, + pub rule_names: Vec, + pub policy_name: String, +} + +/// Result of decode: payload plus metadata. +#[derive(Debug, Clone)] +pub struct DecodeResult { + pub payload: serde_json::Value, + pub rules_applied_count: u32, + pub rule_names: Vec, + pub policy_name: String, +} + +impl Enforcer { + /// Encode payload (mask PII). When context is Llm, tokens get a 3-digit signature and instruction text is prepended. + pub fn encode(&self, payload: serde_json::Value, context: EncodeContext) -> EncodeResult { + let default_result = EncodeResult { + payload: payload.clone(), + signature_injection_text: String::new(), + rules_applied_count: 0, + rule_names: Vec::new(), + policy_name: String::new(), + }; + if self.handle.is_null() { + return default_result; + } + let json = match serde_json::to_string(&payload) { + Ok(s) => s, + Err(_) => return default_result, + }; + let c_str = match CString::new(json.as_bytes()) { + Ok(c) => c, + Err(_) => return default_result, + }; + let enc_ctx = match context { + EncodeContext::Llm => ENCODE_CONTEXT_LLM, + EncodeContext::Manual => ENCODE_CONTEXT_MANUAL, + }; + let out = unsafe { + guardrail_encode( + self.handle, + c_str.as_ptr(), + c_str.as_bytes().len(), + enc_ctx, + ) + }; + if out.is_null() { + return default_result; + } + let out_slice = unsafe { CStr::from_ptr(out).to_bytes() }; + let out_str = String::from_utf8_lossy(out_slice).into_owned(); + unsafe { guardrail_free(out) }; + parse_encode_result(&out_str) + .or_else(|| parse_encode_result_legacy(&out_str)) + .unwrap_or(default_result) + } + + /// Decode payload (rehydrate PII). + pub fn decode(&self, payload: serde_json::Value, context: DecodeContext) -> DecodeResult { + let default_result = DecodeResult { + payload: payload.clone(), + rules_applied_count: 0, + rule_names: Vec::new(), + policy_name: String::new(), + }; + if self.handle.is_null() { + return default_result; + } + let json = match serde_json::to_string(&payload) { + Ok(s) => s, + Err(_) => return default_result, + }; + let c_str = match CString::new(json.as_bytes()) { + Ok(c) => c, + Err(_) => return default_result, + }; + let ctx = match context { + DecodeContext::ToolBoundary => CONTEXT_TOOL_BOUNDARY, + DecodeContext::LlmResponse => CONTEXT_LLM_RESPONSE, + DecodeContext::ManualCall => CONTEXT_MANUAL_CALL, + }; + let out = + unsafe { guardrail_decode(self.handle, c_str.as_ptr(), c_str.as_bytes().len(), ctx) }; + if out.is_null() { + return default_result; + } + let out_slice = unsafe { CStr::from_ptr(out).to_bytes() }; + let out_str = String::from_utf8_lossy(out_slice).into_owned(); + unsafe { guardrail_free(out) }; + parse_decode_result(&out_str) + .or_else(|| parse_decode_result_legacy(&out_str)) + .unwrap_or(default_result) + } +} + +/// Legacy FFI return: raw encoded payload as JSON (old guardrail_encode returned Value serialized). +fn parse_encode_result_legacy(s: &str) -> Option { + let payload: serde_json::Value = serde_json::from_str(s).ok()?; + Some(EncodeResult { + payload, + signature_injection_text: String::new(), + rules_applied_count: 0, + rule_names: Vec::new(), + policy_name: String::new(), + }) +} + +/// Legacy FFI return: raw decoded payload as JSON (old guardrail_decode returned Value serialized). +fn parse_decode_result_legacy(s: &str) -> Option { + let payload: serde_json::Value = serde_json::from_str(s).ok()?; + Some(DecodeResult { + payload, + rules_applied_count: 0, + rule_names: Vec::new(), + policy_name: String::new(), + }) +} + +fn parse_encode_result(s: &str) -> Option { + let v: serde_json::Value = serde_json::from_str(s).ok()?; + let payload = v.get("payload")?.clone(); + let signature_injection_text = v + .get("signature_injection_text") + .and_then(|x| x.as_str()) + .unwrap_or("") + .to_string(); + let rules_applied_count = v.get("rules_applied_count")?.as_u64()? as u32; + let rule_names: Vec = v + .get("rule_names")? + .as_array()? + .iter() + .filter_map(|x| x.as_str().map(String::from)) + .collect(); + let policy_name = v.get("policy_name")?.as_str()?.to_string(); + Some(EncodeResult { + payload, + signature_injection_text, + rules_applied_count, + rule_names, + policy_name, + }) +} + +fn parse_decode_result(s: &str) -> Option { + let v: serde_json::Value = serde_json::from_str(s).ok()?; + let payload = v.get("payload")?.clone(); + let rules_applied_count = v.get("rules_applied_count")?.as_u64()? as u32; + let rule_names: Vec = v + .get("rule_names")? + .as_array()? + .iter() + .filter_map(|x| x.as_str().map(String::from)) + .collect(); + let policy_name = v.get("policy_name")?.as_str()?.to_string(); + Some(DecodeResult { + payload, + rules_applied_count, + rule_names, + policy_name, + }) +} + +/// Singleton entry point to create enforcers. +#[derive(Debug, Clone)] +pub struct GuardRail; + +impl GuardRail { + /// Initialize (no-op for FFI; state is inside the lib). + #[must_use] + pub fn init() -> Self { + Self + } + + /// Create an enforcer for this workflow. + pub fn enforcer_for(config: Arc, workflow_id: impl Into) -> Enforcer { + let workflow_id = workflow_id.into(); + let (ptr, len) = if workflow_id.is_empty() { + (std::ptr::null(), 0) + } else { + (workflow_id.as_ptr() as *const c_char, workflow_id.len()) + }; + let handle = unsafe { guardrail_enforcer_create(config.ptr(), ptr, len) }; + assert!(!handle.is_null(), "guardrail_enforcer_create failed"); + Enforcer { handle } + } +} diff --git a/python/python-src/graphbit/__init__.py b/python/python-src/graphbit/__init__.py index 0cb9cbd8..fa08787a 100644 --- a/python/python-src/graphbit/__init__.py +++ b/python/python-src/graphbit/__init__.py @@ -10,7 +10,6 @@ configure_runtime, shutdown, ) - # Document loader classes from .graphbit import ( DocumentLoaderConfig, @@ -30,6 +29,7 @@ # Workflow classes from .graphbit import ( + GuardRailPolicyConfig, Node, Workflow, WorkflowContext, @@ -99,6 +99,8 @@ "FinishReason", "LlmToolCall", "LlmResponse", + # GuardRail + "GuardRailPolicyConfig", # Workflow "Node", "Workflow", diff --git a/python/src/guardrail.rs b/python/src/guardrail.rs new file mode 100644 index 00000000..2706f933 --- /dev/null +++ b/python/src/guardrail.rs @@ -0,0 +1,90 @@ +//! GuardRail policy config exposed to Python as `GuardRailPolicyConfig`. +//! Used with `Executor.execute(workflow, policy=...)` for PII masking/mapping. + +use graphbit_core::GuardRailConfig; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use std::sync::Arc; + +/// Python-facing guardrail policy configuration. +/// +/// Create via `GuardRailPolicyConfig.from_json(...)`, `from_file(...)`, or `from_url(...)`, +/// then pass to `executor.execute(workflow, policy=config)`. +#[pyclass] +#[derive(Clone)] +pub struct GuardRailPolicyConfig { + pub(crate) inner: Arc, +} + +#[pymethods] +impl GuardRailPolicyConfig { + /// Create a config from a JSON string. + /// + /// # Errors + /// Raises `ValueError` if the JSON is invalid or validation fails. + #[staticmethod] + pub fn from_json(json_str: &str) -> PyResult { + let config = GuardRailConfig::new(json_str) + .map_err(|e| PyValueError::new_err(format!("GuardRail config error: {}", e)))?; + Ok(Self { + inner: Arc::new(config), + }) + } + + /// Create a config from a local file path. + /// + /// # Errors + /// Raises `ValueError` if the file cannot be read or validation fails. + #[staticmethod] + pub fn from_file(path: &str) -> PyResult { + let config = GuardRailConfig::from_file(std::path::Path::new(path)) + .map_err(|e| PyValueError::new_err(format!("GuardRail config error: {}", e)))?; + Ok(Self { + inner: Arc::new(config), + }) + } + + /// Create a config from a remote URL (HTTP GET). + /// + /// # Errors + /// Raises `ValueError` if the URL cannot be fetched or validation fails. + #[staticmethod] + pub fn from_url(url: &str) -> PyResult { + let config = GuardRailConfig::from_url(url) + .map_err(|e| PyValueError::new_err(format!("GuardRail config error: {}", e)))?; + Ok(Self { + inner: Arc::new(config), + }) + } + + /// Default (inactive) policy — no masking or mapping. + #[staticmethod] + pub fn default_config() -> Self { + Self { + inner: Arc::new(GuardRailConfig::default_config()), + } + } + + /// Policy name. + pub fn policy_name(&self) -> String { + self.inner.policy_name() + } + + /// Policy version. + pub fn policy_version(&self) -> String { + self.inner.policy_version() + } + + /// Whether the policy is active. + pub fn is_active(&self) -> bool { + self.inner.active() + } +} + +impl GuardRailPolicyConfig { + /// Return the inner config for use by the executor (same-crate only). Not exposed to Python. + #[inline] + pub(crate) fn get_inner(&self) -> Arc { + Arc::clone(&self.inner) + } +} diff --git a/python/src/lib.rs b/python/src/lib.rs index 458d7e38..23ce3995 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -33,6 +33,7 @@ use tracing::{error, info, warn}; mod document_loader; mod embeddings; mod errors; +mod guardrail; mod llm; mod runtime; mod text_splitter; @@ -49,6 +50,7 @@ pub use text_splitter::{ TokenSplitter, }; pub use tools::{ToolDecorator, ToolExecutor, ToolRegistry, ToolResult}; +pub use guardrail::GuardRailPolicyConfig; pub use workflow::{Executor, Node, Workflow, WorkflowContext, WorkflowResult}; /// Global initialization flag to ensure init is called only once @@ -386,6 +388,9 @@ fn graphbit(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; + // GuardRail policy config (optional for executor.execute(workflow, policy=...)) + m.add_class::()?; + // Workflow classes m.add_class::()?; m.add_class::()?; diff --git a/python/src/workflow/executor.rs b/python/src/workflow/executor.rs index 7997531d..5569f45e 100644 --- a/python/src/workflow/executor.rs +++ b/python/src/workflow/executor.rs @@ -8,13 +8,16 @@ //! - Graceful error handling and recovery use graphbit_core::workflow::WorkflowExecutor as CoreWorkflowExecutor; +use graphbit_core::{DecodeContext, EncodeContext, GuardRail, Enforcer}; use pyo3::prelude::*; use pyo3::types::{PyAny, PyDict}; +use std::sync::Arc; use std::time::{Duration, Instant}; use tracing::{debug, error, info, instrument, warn}; use super::{result::WorkflowResult, workflow::Workflow}; use crate::errors::{timeout_error, to_py_runtime_error, validation_error}; +use crate::guardrail::GuardRailPolicyConfig; use crate::llm::config::LlmConfig; use crate::runtime::get_runtime; @@ -90,10 +93,11 @@ pub struct Executor { #[pymethods] impl Executor { #[new] - #[pyo3(signature = (config, _lightweight_mode=None, timeout_seconds=None, debug=None))] + #[pyo3(signature = (config, lightweight_mode=None, timeout_seconds=None, debug=None))] + #[allow(unused_variables)] fn new( config: LlmConfig, - _lightweight_mode: Option, + lightweight_mode: Option, timeout_seconds: Option, debug: Option, ) -> PyResult { @@ -132,9 +136,18 @@ impl Executor { }) } - /// Execute a workflow with comprehensive error handling and monitoring - #[instrument(skip(self, py, workflow), fields(workflow_name = %workflow.inner.name))] - fn execute(&mut self, py: Python<'_>, workflow: &Workflow) -> PyResult { + /// Execute a workflow with comprehensive error handling and monitoring. + /// + /// `policy` is optional. When provided: encode before every LLM call, decode after every LLM call; + /// before tool usage decode (so tools see real PII); after tool usage do nothing (no encode). + #[instrument(skip(self, py, workflow, policy), fields(workflow_name = %workflow.inner.name))] + #[pyo3(signature = (workflow, policy=None))] + fn execute( + &mut self, + py: Python<'_>, + workflow: &Workflow, + policy: Option<&Bound<'_, GuardRailPolicyConfig>>, + ) -> PyResult { let start_time = Instant::now(); // Validate workflow @@ -161,6 +174,14 @@ impl Executor { let timeout_duration = config.timeout; let debug = config.enable_tracing; // Capture debug flag + // Build optional guardrail enforcer from policy (for encode/decode at LLM and tool boundaries) + let guardrail_enforcer = policy.map(|p| { + let config = p.borrow().get_inner(); + Arc::new( + GuardRail::enforcer_for(config, workflow_clone.id.to_string()), + ) + }); + if debug { debug!("Starting workflow execution with mode: {:?}", config.mode); } @@ -171,7 +192,13 @@ impl Executor { get_runtime().block_on(async move { // Apply timeout to the entire execution tokio::time::timeout(timeout_duration, async move { - Self::execute_workflow_internal(llm_config, workflow_clone, config).await + Self::execute_workflow_internal( + llm_config, + workflow_clone, + config, + guardrail_enforcer, + ) + .await }) .await }) @@ -210,8 +237,14 @@ impl Executor { } /// Async execution with enhanced performance optimizations - #[instrument(skip(self, workflow, py), fields(workflow_name = %workflow.inner.name))] - fn run_async<'a>(&mut self, workflow: &Workflow, py: Python<'a>) -> PyResult> { + #[instrument(skip(self, workflow, py, policy), fields(workflow_name = %workflow.inner.name))] + #[pyo3(signature = (workflow, policy=None))] + fn run_async<'a>( + &mut self, + workflow: &Workflow, + py: Python<'a>, + policy: Option<&Bound<'_, GuardRailPolicyConfig>>, + ) -> PyResult> { // Validate workflow if let Err(e) = workflow.inner.validate() { return Err(validation_error( @@ -226,7 +259,13 @@ impl Executor { let config = self.config.clone(); let timeout_duration = config.timeout; let start_time = Instant::now(); - let debug = config.enable_tracing; // Capture debug flag + let debug = config.enable_tracing; + let guardrail_enforcer = policy.map(|p| { + let config = p.borrow().get_inner(); + Arc::new( + GuardRail::enforcer_for(config, workflow_clone.id.to_string()), + ) + }); if debug { debug!( @@ -236,9 +275,14 @@ impl Executor { } pyo3_async_runtimes::tokio::future_into_py(py, async move { - // Apply timeout to the entire execution let result = tokio::time::timeout(timeout_duration, async move { - Self::execute_workflow_internal(llm_config, workflow_clone, config).await + Self::execute_workflow_internal( + llm_config, + workflow_clone, + config, + guardrail_enforcer, + ) + .await }) .await; @@ -381,19 +425,24 @@ impl Executor { } impl Executor { - /// Internal workflow execution with mode-specific optimizations and tool call handling + /// Internal workflow execution with mode-specific optimizations and tool call handling. + /// When `guardrail_enforcer` is `Some`, the core encodes before LLM and decodes after LLM; + /// we decode before tool usage only (no encode after tool). async fn execute_workflow_internal( llm_config: graphbit_core::llm::LlmConfig, workflow: graphbit_core::workflow::Workflow, config: ExecutionConfig, + guardrail_enforcer: Option>, ) -> Result { let executor = match config.mode { ExecutionMode::Balanced => CoreWorkflowExecutor::new() .with_default_llm_config(llm_config.clone()), }; - // Execute the workflow - let mut context = executor.execute(workflow.clone()).await?; + // Execute the workflow (core applies encode before LLM, decode after LLM when enforcer is Some) + let mut context = executor + .execute(workflow.clone(), guardrail_enforcer.clone()) + .await?; // Store LLM config in context metadata for tool call handling if let Ok(llm_config_json) = serde_json::to_value(&llm_config) { @@ -403,16 +452,25 @@ impl Executor { } // Check if any node outputs contain tool_calls_required responses and handle them - context = Self::handle_tool_calls_in_context(context, &workflow).await?; + context = Self::handle_tool_calls_in_context( + context, + &workflow, + guardrail_enforcer.as_ref().map(|arc| arc.as_ref()), + ) + .await?; Ok(context) } - /// Handle tool calls in workflow context by executing them and updating the context + /// Handle tool calls in workflow context by executing them and updating the context. + /// When `guardrail_enforcer` is `Some`, decodes tool-call parameters before execution only; + /// after tool execution we do nothing (no encode of tool results). async fn handle_tool_calls_in_context( mut context: graphbit_core::types::WorkflowContext, workflow: &graphbit_core::workflow::Workflow, + guardrail_enforcer: Option<&Enforcer>, ) -> Result { + use graphbit_core::DecodeContext; use crate::workflow::node::execute_production_tool_calls; use graphbit_core::llm::{LlmProvider, LlmRequest}; @@ -447,22 +505,35 @@ impl Executor { }) .unwrap_or_default(); - // Convert tool calls to the format expected by Python layer + // Convert tool calls to the format expected by Python layer. + // Guardrail: decode parameters before tool execution so tools see real PII. + if guardrail_enforcer.is_some() { + tracing::debug!( + "[GuardRail] tool call parameters from LLM (before decode): {:?}", + tool_calls + ); + } let python_tool_calls: Vec = if let Some(tool_calls_array) = tool_calls.as_array() { tool_calls_array .iter() .map(|tc| { - // Extract name and parameters from the tool call object let name = tc .get("name") .and_then(|v| v.as_str()) .unwrap_or("unknown"); - let parameters = tc + let mut parameters = tc .get("parameters") .cloned() .unwrap_or(serde_json::json!({})); - + if let Some(enforcer) = guardrail_enforcer { + tracing::debug!( + "Guardrail: decoding tool call parameters (tool boundary — tool will receive real PII)" + ); + let decoded_result = + enforcer.decode(parameters, DecodeContext::ToolBoundary); + parameters = decoded_result.payload; + } serde_json::json!({ "tool_name": name, "parameters": parameters @@ -524,11 +595,30 @@ impl Executor { summary_lines.join("\n") }; - // Create final prompt with tool results summary + // Guardrail: before tool we decode; after tool we do nothing (no encode of results). + let summary_for_llm = tool_results_summary.clone(); + + // Build final prompt; when GuardRail is active encode it and debug-print. let final_prompt = format!( "{}\n\nTool execution results:\n{}\n\nPlease provide a comprehensive response based on the tool results.", - original_prompt, tool_results_summary + original_prompt, summary_for_llm ); + let prompt_for_final_llm = if let Some(ref enforcer) = guardrail_enforcer { + tracing::info!("[GuardRail] final prompt (before encode): {}", final_prompt); + let result = enforcer.encode( + serde_json::Value::String(final_prompt.clone()), + EncodeContext::Llm, + ); + let encoded_str = format!( + "{}{}", + result.signature_injection_text, + result.payload.as_str().unwrap_or_default() + ); + tracing::info!("[GuardRail] final prompt (after encode, sent to LLM, payload only): {}", result.payload.as_str().unwrap_or_default()); + encoded_str + } else { + final_prompt.clone() + }; // Get LLM provider from node configuration and make final call if let graphbit_core::graph::NodeType::Agent { .. } = @@ -553,8 +643,8 @@ impl Executor { let llm_provider = LlmProvider::new(provider_trait, llm_config); - // Create final request and apply node configuration parameters - let mut final_request = LlmRequest::new(final_prompt.clone()); + // Create final request (with encoded prompt when GuardRail is on) + let mut final_request = LlmRequest::new(prompt_for_final_llm); // CUMULATIVE TOKEN BUDGET TRACKING // Extract initial tokens used and max_tokens to calculate remaining budget @@ -616,6 +706,11 @@ impl Executor { match llm_provider.complete(final_request).await { Ok(final_response) => { + tracing::info!( + "[GuardRail] final LLM response (GuardRail active={}); before decode: {}", + guardrail_enforcer.is_some(), + final_response.content + ); tracing::debug!( "Final LLM response received - content: '{}', tokens: {}, finish_reason: {:?}", final_response.content, @@ -623,9 +718,24 @@ impl Executor { final_response.finish_reason ); - // Clone the content to avoid borrow checker issues - let response_content = - final_response.content.clone(); + // Guardrail: decode after every LLM call so user sees rehydrated content + let response_content = if let Some(ref enforcer) = guardrail_enforcer { + let payload = serde_json::json!({ + "content": final_response.content + }); + let decoded_result = + enforcer.decode(payload, DecodeContext::LlmResponse); + let content = decoded_result + .payload + .get("content") + .and_then(|v| v.as_str()) + .map(String::from) + .unwrap_or_else(|| final_response.content.clone()); + tracing::info!("[GuardRail] final LLM response (after decode): {}", content); + content + } else { + final_response.content.clone() + }; // Store full LLM response metadata in context // This enables observability tools to capture complete LLM metadata diff --git a/tests/rust_integration_tests/full_workflow_tests.rs b/tests/rust_integration_tests/full_workflow_tests.rs index 49653de5..4fb5dec6 100644 --- a/tests/rust_integration_tests/full_workflow_tests.rs +++ b/tests/rust_integration_tests/full_workflow_tests.rs @@ -593,7 +593,7 @@ async fn test_real_llm_workflow_execution() { let workflow = builder.build().expect("Failed to build workflow"); // Try to execute the workflow - let result = executor.execute(workflow).await; + let result = executor.execute(workflow, None).await; match result { Ok(context) => { println!("Real LLM workflow executed successfully"); @@ -721,7 +721,7 @@ async fn test_workflow_error_propagation() { // Try to execute - should fail gracefully let executor = WorkflowExecutor::new(); - let result = executor.execute(workflow).await; + let result = executor.execute(workflow, None).await; // Should handle gracefully (may succeed or fail depending on implementation) match result { @@ -802,7 +802,7 @@ async fn test_multi_provider_workflow_execution() { let workflow = builder.build().expect("Failed to build workflow"); - let result = executor.execute(workflow).await; + let result = executor.execute(workflow, None).await; match result { Ok(context) => { println!("{provider_name} workflow executed successfully"); @@ -936,7 +936,7 @@ async fn test_comprehensive_real_api_workflow() { .expect("Failed to build workflow"); // Execute workflow - let result = executor.execute(workflow).await; + let result = executor.execute(workflow, None).await; match result { Ok(context) => { println!("Comprehensive real API workflow executed successfully"); @@ -981,7 +981,7 @@ async fn test_workflow_timeout_handling() { // Execute with timeout let executor = WorkflowExecutor::new().with_max_node_execution_time(2000); // 2 second max let start = std::time::Instant::now(); - let result = executor.execute(workflow).await; + let result = executor.execute(workflow, None).await; let duration = start.elapsed(); // Should complete quickly due to timeout, not wait full 10 seconds diff --git a/tests/rust_unit_tests/python_bindings_tests.rs b/tests/rust_unit_tests/python_bindings_tests.rs index aafef5ab..b91fb61f 100644 --- a/tests/rust_unit_tests/python_bindings_tests.rs +++ b/tests/rust_unit_tests/python_bindings_tests.rs @@ -82,7 +82,7 @@ async fn test_workflow_executor() { assert!(workflow.metadata.contains_key("test_key")); let executor = WorkflowExecutor::new(); - let result = executor.execute(workflow).await; + let result = executor.execute(workflow, None).await; // Empty workflows may fail execution, which is expected behavior // Let's check what the actual result is @@ -109,7 +109,7 @@ async fn test_workflow_integration() { .unwrap(); let executor = WorkflowExecutor::new(); - let result = executor.execute(workflow).await; + let result = executor.execute(workflow, None).await; // Empty workflows may fail execution, which is expected behavior match result { diff --git a/tests/rust_unit_tests/workflow_tests.rs b/tests/rust_unit_tests/workflow_tests.rs index fa92719d..db06808d 100644 --- a/tests/rust_unit_tests/workflow_tests.rs +++ b/tests/rust_unit_tests/workflow_tests.rs @@ -191,7 +191,7 @@ async fn test_workflow_execute_with_dummy_agent_success() { let exec = WorkflowExecutor::new(); exec.register_agent(agent).await; - let ctx = exec.execute(wf).await.expect("workflow should execute"); + let ctx = exec.execute(wf, None).await.expect("workflow should execute"); assert!(matches!(ctx.state, WorkflowState::Completed)); let stats = ctx.stats.expect("stats present"); assert!(stats.total_nodes >= 3); @@ -233,7 +233,7 @@ async fn test_workflow_execute_fail_fast_on_error() { exec.register_agent(agent).await; let ctx = exec - .execute(wf) + .execute(wf, None) .await .expect("execution should return context"); // Current executor records node failure but continues; ensure at least one failed node counted @@ -454,7 +454,7 @@ async fn test_workflow_with_llm() { .expect("Failed to build agent"); executor.register_agent(Arc::new(agent)).await; - let result = executor.execute(workflow).await; + let result = executor.execute(workflow, None).await; assert!(result.is_ok()); } @@ -497,7 +497,7 @@ async fn test_workflow_with_anthropic() { .expect("Failed to build agent"); executor.register_agent(Arc::new(agent)).await; - let result = executor.execute(workflow).await; + let result = executor.execute(workflow, None).await; assert!(result.is_ok()); } @@ -542,7 +542,7 @@ async fn test_workflow_with_ollama() { .expect("Failed to build agent"); executor.register_agent(Arc::new(agent)).await; - let result = executor.execute(workflow).await; + let result = executor.execute(workflow, None).await; assert!(result.is_ok()); } diff --git a/vendor/guardrail/aarch64-unknown-linux-gnu/libguardrail_ffi.so b/vendor/guardrail/aarch64-unknown-linux-gnu/libguardrail_ffi.so new file mode 100644 index 00000000..6ff9b079 Binary files /dev/null and b/vendor/guardrail/aarch64-unknown-linux-gnu/libguardrail_ffi.so differ diff --git a/vendor/guardrail/guardrail_ffi.dll b/vendor/guardrail/guardrail_ffi.dll new file mode 100644 index 00000000..ff4dfe4e Binary files /dev/null and b/vendor/guardrail/guardrail_ffi.dll differ diff --git a/vendor/guardrail/guardrail_ffi.lib b/vendor/guardrail/guardrail_ffi.lib new file mode 100644 index 00000000..7513e2a9 Binary files /dev/null and b/vendor/guardrail/guardrail_ffi.lib differ diff --git a/vendor/guardrail/libguardrail_ffi.a b/vendor/guardrail/libguardrail_ffi.a new file mode 100644 index 00000000..7741f353 Binary files /dev/null and b/vendor/guardrail/libguardrail_ffi.a differ diff --git a/vendor/guardrail/x86_64-unknown-linux-gnu/libguardrail_ffi.so b/vendor/guardrail/x86_64-unknown-linux-gnu/libguardrail_ffi.so new file mode 100644 index 00000000..ebe95b97 Binary files /dev/null and b/vendor/guardrail/x86_64-unknown-linux-gnu/libguardrail_ffi.so differ