From f99616d1a25427686375a7d51da323c047aa1e78 Mon Sep 17 00:00:00 2001 From: Andrei Petraru Date: Tue, 17 Mar 2026 16:53:39 +0200 Subject: [PATCH] feat: decorator guardrail implementation [AL-288] Refactor guardrail decorator layer to use adapter/registry pattern: - Move validator logic (PIIValidator, PromptInjectionValidator, CustomValidator), actions (LogAction, BlockAction), and @guardrail decorator core into uipath-platform; thin re-exports remain in uipath-langchain for backward compat - Add _langchain_adapter.py with LangChainGuardrailAdapter implementing the GuardrailTargetAdapter Protocol; wraps BaseTool, BaseChatModel, StateGraph, CompiledStateGraph; auto-registered on import of uipath_langchain.guardrails - BlockAction now raises GuardrailBlockException (platform); adapter converts to AgentRuntimeError at wrapper boundaries - Update exception guards in middlewares/pii_detection.py and prompt_injection.py to re-raise GuardrailBlockException alongside AgentRuntimeError - Delete decorators/_base.py, decorators/guardrail.py, decorators/validators/* (logic lives in platform); update thin re-exports in decorators/__init__.py Made-with: Cursor Co-Authored-By: Claude Sonnet 4.6 --- pyproject.toml | 2 +- samples/joke-agent-decorator/README.md | 197 +++++++ samples/joke-agent-decorator/graph.py | 245 ++++++++ samples/joke-agent-decorator/langgraph.json | 7 + samples/joke-agent-decorator/pyproject.toml | 18 + samples/joke-agent/graph.py | 6 +- src/uipath_langchain/guardrails/__init__.py | 53 +- .../guardrails/_langchain_adapter.py | 537 ++++++++++++++++++ src/uipath_langchain/guardrails/actions.py | 115 +--- .../guardrails/decorators/__init__.py | 19 + src/uipath_langchain/guardrails/enums.py | 47 +- .../guardrails/middlewares/deterministic.py | 7 + .../guardrails/middlewares/pii_detection.py | 17 +- .../middlewares/prompt_injection.py | 14 +- src/uipath_langchain/guardrails/models.py | 89 +-- uv.lock | 2 +- 16 files changed, 1124 insertions(+), 251 deletions(-) create mode 100644 samples/joke-agent-decorator/README.md create mode 100644 samples/joke-agent-decorator/graph.py create mode 100644 samples/joke-agent-decorator/langgraph.json create mode 100644 samples/joke-agent-decorator/pyproject.toml create mode 100644 src/uipath_langchain/guardrails/_langchain_adapter.py create mode 100644 src/uipath_langchain/guardrails/decorators/__init__.py diff --git a/pyproject.toml b/pyproject.toml index 35df3ec58..351f0ca74 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "uipath-langchain" -version = "0.9.12" +version = "0.9.13" description = "Python SDK that enables developers to build and deploy LangGraph agents to the UiPath Cloud Platform" readme = { file = "README.md", content-type = "text/markdown" } requires-python = ">=3.11" diff --git a/samples/joke-agent-decorator/README.md b/samples/joke-agent-decorator/README.md new file mode 100644 index 000000000..603f50a30 --- /dev/null +++ b/samples/joke-agent-decorator/README.md @@ -0,0 +1,197 @@ +# Joke Agent (Decorator-based Guardrails) + +A simple LangGraph agent that generates family-friendly jokes based on a given topic using UiPath's LLM. This sample demonstrates all three guardrail decorator types — PII, Prompt Injection, and Deterministic — applied directly to the LLM, agent, and tool without a middleware stack. + +## Requirements + +- Python 3.11+ + +## Installation + +```bash +uv venv -p 3.11 .venv +source .venv/bin/activate # On Windows: .venv\Scripts\activate +uv sync +``` + +## Usage + +Run the joke agent: + +```bash +uv run uipath run agent '{"topic": "banana"}' +``` + +### Input Format + +```json +{ + "topic": "banana" +} +``` + +### Output Format + +```json +{ + "joke": "Why did the banana go to the doctor? Because it wasn't peeling well!" +} +``` + +## Guardrails Overview + +This sample achieves full parity with the middleware-based `joke-agent` sample using only decorators. The table below shows which scope each guardrail covers: + +| Decorator | Target | Scope | Action | +|---|---|---|---| +| `@prompt_injection_guardrail` | `create_llm` factory | LLM | `BlockAction` — blocks on detection | +| `@pii_detection_guardrail` | `create_llm` factory | LLM | `LogAction(WARNING)` — logs and continues | +| `@pii_detection_guardrail` | `analyze_joke_syntax` tool | TOOL | `LogAction(WARNING)` — logs email/phone | +| `@deterministic_guardrail` | `analyze_joke_syntax` tool | TOOL (PRE) | `CustomFilterAction` — replaces "donkey" with "[censored]" | +| `@deterministic_guardrail` | `analyze_joke_syntax` tool | TOOL (PRE) | `BlockAction` — blocks jokes > 1000 chars | +| `@deterministic_guardrail` | `analyze_joke_syntax` tool | TOOL (POST) | `CustomFilterAction` — always-on output transform | +| `@pii_detection_guardrail` | `create_joke_agent` factory | AGENT | `LogAction(WARNING)` — logs agent-level PII | + +## Guardrail Decorators + +### LLM-level guardrails + +Stacked decorators on a factory function. The outermost decorator runs first: + +```python +@prompt_injection_guardrail( + threshold=0.5, + action=BlockAction(), + name="LLM Prompt Injection Detection", + enabled_for_evals=False, # default is True +) +@pii_detection_guardrail( + entities=[PIIDetectionEntity(PIIDetectionEntityType.EMAIL, 0.5)], + action=LogAction(severity_level=LoggingSeverityLevel.WARNING), + name="LLM PII Detection", +) +def create_llm(): + return UiPathChat(model="gpt-4o-2024-08-06", temperature=0.7) + +llm = create_llm() +``` + +### Tool-level guardrails + +`@deterministic_guardrail` applies local rule functions — no UiPath API call. Rules receive the tool input dict and return `True` to signal a violation. `@pii_detection_guardrail` at TOOL scope evaluates via the UiPath guardrails API. + +```python +@deterministic_guardrail( + rules=[lambda args: "donkey" in args.get("joke", "").lower()], + action=CustomFilterAction(word_to_filter="donkey", replacement="[censored]"), + stage=GuardrailExecutionStage.PRE, + name="Joke Content Word Filter", + enabled_for_evals=False, # default is True +) +@deterministic_guardrail( + rules=[lambda args: len(args.get("joke", "")) > 1000], + action=BlockAction(), + stage=GuardrailExecutionStage.PRE, + name="Joke Content Length Limiter", +) +@deterministic_guardrail( + rules=[], # empty rules = always apply (unconditional transform) + action=CustomFilterAction(word_to_filter="words", replacement="words++"), + stage=GuardrailExecutionStage.POST, + name="Joke Content Always Filter", +) +@pii_detection_guardrail( + entities=[ + PIIDetectionEntity(PIIDetectionEntityType.EMAIL, 0.5), + PIIDetectionEntity(PIIDetectionEntityType.PHONE_NUMBER, 0.5), + ], + action=LogAction(severity_level=LoggingSeverityLevel.WARNING), + name="Tool PII Detection", +) +@tool +def analyze_joke_syntax(joke: str) -> str: + ... +``` + +### Agent-level guardrail + +```python +@pii_detection_guardrail( + entities=[PIIDetectionEntity(PIIDetectionEntityType.EMAIL, 0.5)], + action=LogAction( + severity_level=LoggingSeverityLevel.WARNING, + message="PII detected from agent guardrails decorator", + ), + name="Agent PII Detection", + enabled_for_evals=False, # default is True +) +def create_joke_agent(): + return create_agent(model=llm, tools=[analyze_joke_syntax], ...) + +agent = create_joke_agent() +``` + +### Custom action + +`CustomFilterAction` (defined locally in `graph.py`) demonstrates how to implement a custom `GuardrailAction`. When a violation is detected it replaces the offending word in the tool input dict or string, logs the change, then returns the modified data so execution continues with the sanitised input: + +```python +@dataclass +class CustomFilterAction(GuardrailAction): + word_to_filter: str + replacement: str = "***" + + def handle_validation_result(self, result, data, guardrail_name): + # filter word from dict/str and return modified data + ... +``` + +## Rule semantics (`@deterministic_guardrail`) + +- A rule with **1 parameter** receives the tool input dict (`PRE` stage). +- A rule with **2 parameters** receives `(input_dict, output_dict)` (`POST` stage). +- A rule returns `True` to signal a **violation**, `False` to **pass**. +- **All** rules must detect a violation for the guardrail to trigger. If any rule passes, the guardrail passes. +- **Empty `rules=[]`** always triggers the action (useful for unconditional transforms). + +## `enabled_for_evals` override + +All decorator guardrails accept `enabled_for_evals` (default `True`). Set it to `False` +when you want runtime guardrail behavior but do not want that guardrail enabled for eval scenarios. + +## Verification + +To manually verify each guardrail fires, run from this directory: + +```bash +uv run uipath run agent '{"topic": "donkey"}' +``` + +**Scenario 1 — word filter (PRE):** the LLM includes "donkey" in the joke passed to `analyze_joke_syntax`. `CustomFilterAction` replaces it with `[censored]` before the tool executes. Look for `[FILTER][Joke Content Word Filter]` in stdout. + +**Scenario 2 — length limiter (PRE):** if the generated joke exceeds 1000 characters, `BlockAction` raises `AgentRuntimeError(TERMINATION_GUARDRAIL_VIOLATION)` before the tool is called. + +**Scenario 3 — PII at tool and agent scope:** supply a topic containing an email address: + +```bash +uv run uipath run agent '{"topic": "donkey, test@example.com"}' +``` + +Both the agent-scope and LLM-scope `@pii_detection_guardrail` decorators log a `WARNING` when the email is detected. The tool-scope `@pii_detection_guardrail` logs when the email reaches the tool input. + +## Differences from the Middleware Approach (`joke-agent`) + +| Aspect | Middleware (`joke-agent`) | Decorator (`joke-agent-decorator`) | +|---|---|---| +| Configuration | Middleware class instances passed to `create_agent(middleware=[...])` | `@decorator` stacked on the target object | +| Scope | Explicit `scopes=[...]` list | Inferred automatically from the decorated object | +| Tool guardrails | `UiPathDeterministicGuardrailMiddleware(tools=[...])` | `@deterministic_guardrail` directly on the `@tool` | +| Custom loops | Not supported (requires `create_agent`) | Works in any custom LangChain loop | +| API calls | Via middleware stack | Direct `uipath.guardrails.evaluate_guardrail()` | + +## Example Topics + +- `"banana"` — normal run, all guardrails pass +- `"donkey"` — triggers the word filter on `analyze_joke_syntax` +- `"donkey, test@example.com"` — triggers word filter + PII guardrails at all scopes +- `"computer"`, `"coffee"`, `"pizza"`, `"weather"` diff --git a/samples/joke-agent-decorator/graph.py b/samples/joke-agent-decorator/graph.py new file mode 100644 index 000000000..32ee4afed --- /dev/null +++ b/samples/joke-agent-decorator/graph.py @@ -0,0 +1,245 @@ +"""Joke generating agent that creates family-friendly jokes based on a topic.""" + +import logging +import re +from dataclasses import dataclass +from typing import Any + +from langchain.agents import create_agent +from langchain_core.messages import HumanMessage +from langchain_core.tools import tool +from langgraph.constants import END, START +from langgraph.graph import StateGraph +from pydantic import BaseModel +from uipath.core.guardrails import ( + GuardrailValidationResult, + GuardrailValidationResultType, +) + +from uipath_langchain.chat import UiPathChat +from uipath_langchain.guardrails import ( + BlockAction, + CustomValidator, + GuardrailAction, + GuardrailExecutionStage, + LogAction, + LoggingSeverityLevel, + PIIDetectionEntity, + PIIValidator, + PromptInjectionValidator, + guardrail, +) +from uipath_langchain.guardrails.enums import PIIDetectionEntityType + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Custom filter action (defined locally) +# --------------------------------------------------------------------------- + +@dataclass +class CustomFilterAction(GuardrailAction): + """Filters/replaces a word in tool input when a violation is detected.""" + + word_to_filter: str + replacement: str = "***" + + def _filter(self, text: str) -> str: + return re.sub(re.escape(self.word_to_filter), self.replacement, text, flags=re.IGNORECASE) + + def handle_validation_result( + self, + result: GuardrailValidationResult, + data: str | dict[str, Any], + guardrail_name: str, + ) -> str | dict[str, Any] | None: + if result.result != GuardrailValidationResultType.VALIDATION_FAILED: + return None + if isinstance(data, str): + filtered = self._filter(data) + print(f"[FILTER][{guardrail_name}] '{self.word_to_filter}' replaced → '{filtered[:80]}'") + return filtered + if isinstance(data, dict): + filtered_data = data.copy() + for key in ["joke", "text", "content", "message", "input", "output"]: + if key in filtered_data and isinstance(filtered_data[key], str): + filtered_data[key] = self._filter(filtered_data[key]) + print(f"[FILTER][{guardrail_name}] dict filtered") + return filtered_data + return data + + +# --------------------------------------------------------------------------- +# Input / Output schemas +# --------------------------------------------------------------------------- + +class Input(BaseModel): + """Input schema for the joke agent.""" + topic: str + + +class Output(BaseModel): + """Output schema for the joke agent.""" + joke: str + + +# --------------------------------------------------------------------------- +# Reusable validators (declared once, used in multiple @guardrail decorators) +# --------------------------------------------------------------------------- + +pii_email = PIIValidator( + entities=[PIIDetectionEntity(PIIDetectionEntityType.EMAIL, 0.5)], +) + +pii_email_phone = PIIValidator( + entities=[ + PIIDetectionEntity(PIIDetectionEntityType.EMAIL, 0.5), + PIIDetectionEntity(PIIDetectionEntityType.PHONE_NUMBER, 0.5), + ], +) + + +# --------------------------------------------------------------------------- +# LLM with guardrails (prompt injection + PII at LLM scope) +# --------------------------------------------------------------------------- + +@guardrail( + validator=PromptInjectionValidator(threshold=0.5), + action=BlockAction(), + name="LLM Prompt Injection Detection", + stage=GuardrailExecutionStage.PRE, +) +@guardrail( + validator=pii_email, + action=LogAction(severity_level=LoggingSeverityLevel.WARNING), + name="LLM PII Detection", + stage=GuardrailExecutionStage.PRE, +) +def create_llm(): + """Create LLM instance with guardrails.""" + return UiPathChat(model="gpt-4o-2024-08-06", temperature=0.7) + + +llm = create_llm() + + +# --------------------------------------------------------------------------- +# Tool with guardrails (deterministic + PII at TOOL scope) +# --------------------------------------------------------------------------- + +@guardrail( + validator=CustomValidator(lambda args: "donkey" in args.get("joke", "").lower()), + action=CustomFilterAction(word_to_filter="donkey", replacement="[censored]"), + stage=GuardrailExecutionStage.PRE, + name="Joke Content Word Filter", +) +@guardrail( + validator=CustomValidator(lambda args:len(args.get("joke", "")) > 1000), + action=BlockAction(title="Joke is too long", detail="The generated joke is too long"), + stage=GuardrailExecutionStage.PRE, + name="Joke Content Length Limiter", +) +@guardrail( + validator=CustomValidator(lambda args:True), + action=CustomFilterAction(word_to_filter="words", replacement="words++"), + stage=GuardrailExecutionStage.POST, + name="Joke Content Always Filter", +) +@guardrail( + validator=pii_email_phone, + action=LogAction( + severity_level=LoggingSeverityLevel.WARNING, + message="Email or phone number detected", + ), + name="Tool PII Detection", + stage=GuardrailExecutionStage.PRE, +) +@tool +def analyze_joke_syntax(joke: str) -> str: + """Analyze the syntax of a joke by counting words and letters. + + Args: + joke: The joke text to analyze + + Returns: + A string with the analysis results showing word count and letter count + """ + words = joke.split() + word_count = len(words) + letter_count = sum(1 for char in joke if char.isalpha()) + return f"Words number: {word_count}\nLetters: {letter_count}" + + +# --------------------------------------------------------------------------- +# System prompt +# --------------------------------------------------------------------------- + +SYSTEM_PROMPT = """You are an AI assistant designed to generate family-friendly jokes. Your process is as follows: + +1. Generate a family-friendly joke based on the given topic. +2. Use the analyze_joke_syntax tool to analyze the joke's syntax (word count and letter count). +3. Ensure your output includes the joke. + +When creating jokes, ensure they are: + +1. Appropriate for children +2. Free from offensive language or themes +3. Clever and entertaining +4. Not based on stereotypes or sensitive topics + +If you're unable to generate a suitable joke for any reason, politely explain why and offer to try again with a different topic. + +Example joke: Topic: "banana" Joke: "Why did the banana go to the doctor? Because it wasn't peeling well!" + +Remember to always include the 'joke' property in your output to match the required schema.""" + + +# --------------------------------------------------------------------------- +# Agent with PII guardrail at AGENT scope +# --------------------------------------------------------------------------- + +@guardrail( + validator=PIIValidator( + entities=[PIIDetectionEntity(PIIDetectionEntityType.PERSON, 0.5)], + ), + action=BlockAction( + title="Person name detection", + detail="Person name detected and is not allowed", + ), + name="Agent PII Detection", + stage=GuardrailExecutionStage.PRE, +) +def create_joke_agent(): + """Create the joke agent with guardrails.""" + return create_agent( + model=llm, + tools=[analyze_joke_syntax], + system_prompt=SYSTEM_PROMPT, + ) + + +agent = create_joke_agent() + + +# --------------------------------------------------------------------------- +# Wrapper graph node +# --------------------------------------------------------------------------- + +async def joke_node(state: Input) -> Output: + """Convert topic to messages, call agent, and extract joke.""" + messages = [ + HumanMessage(content=f"Generate a family-friendly joke based on the topic: {state.topic}") + ] + result = await agent.ainvoke({"messages": messages}) + joke = result["messages"][-1].content + return Output(joke=joke) + + +# Build wrapper graph with custom input/output schemas +builder = StateGraph(Input, input=Input, output=Output) +builder.add_node("joke", joke_node) +builder.add_edge(START, "joke") +builder.add_edge("joke", END) + +graph = builder.compile() diff --git a/samples/joke-agent-decorator/langgraph.json b/samples/joke-agent-decorator/langgraph.json new file mode 100644 index 000000000..c465a881b --- /dev/null +++ b/samples/joke-agent-decorator/langgraph.json @@ -0,0 +1,7 @@ +{ + "dependencies": ["."], + "graphs": { + "agent": "./graph.py:graph" + }, + "env": ".env" +} diff --git a/samples/joke-agent-decorator/pyproject.toml b/samples/joke-agent-decorator/pyproject.toml new file mode 100644 index 000000000..e0a580628 --- /dev/null +++ b/samples/joke-agent-decorator/pyproject.toml @@ -0,0 +1,18 @@ +[project] +name = "joke-agent-decorator" +version = "0.0.1" +description = "Joke generating agent that creates family-friendly jokes based on a topic - using decorator-based guardrails" +authors = [{ name = "Andrei Petraru", email = "andrei.petraru@uipath.com" }] +requires-python = ">=3.11" +dependencies = [ + "uipath-langchain>=0.8.28", + "uipath>2.7.0", +] + +[dependency-groups] +dev = [ + "uipath-dev>=0.0.14", +] + +[tool.uv.sources] +uipath-langchain = { path = "../..", editable = true } diff --git a/samples/joke-agent/graph.py b/samples/joke-agent/graph.py index 85025cae3..31acb3251 100644 --- a/samples/joke-agent/graph.py +++ b/samples/joke-agent/graph.py @@ -15,12 +15,12 @@ PIIDetectionEntity, GuardrailExecutionStage, LogAction, - PIIDetectionEntityType, UiPathDeterministicGuardrailMiddleware, UiPathPIIDetectionMiddleware, UiPathPromptInjectionMiddleware, ) from uipath_langchain.guardrails.actions import LoggingSeverityLevel +from uipath_langchain.guardrails.enums import PIIDetectionEntityType # Define input schema for the agent @@ -102,12 +102,13 @@ def analyze_joke_syntax(joke: str) -> str: PIIDetectionEntity(PIIDetectionEntityType.PHONE_NUMBER, 0.5), ], tools=[analyze_joke_syntax], + enabled_for_evals=False, ), *UiPathPromptInjectionMiddleware( name="Prompt Injection Detection", - scopes=[GuardrailScope.LLM], action=BlockAction(), threshold=0.5, + enabled_for_evals=False, ), # Custom FilterAction example: demonstrates how developers can implement their own actions *UiPathDeterministicGuardrailMiddleware( @@ -121,6 +122,7 @@ def analyze_joke_syntax(joke: str) -> str: ), stage=GuardrailExecutionStage.PRE, name="Joke Content Validator", + enabled_for_evals=False, ), *UiPathDeterministicGuardrailMiddleware( tools=[analyze_joke_syntax], diff --git a/src/uipath_langchain/guardrails/__init__.py b/src/uipath_langchain/guardrails/__init__.py index efd50e194..3be8db389 100644 --- a/src/uipath_langchain/guardrails/__init__.py +++ b/src/uipath_langchain/guardrails/__init__.py @@ -1,31 +1,68 @@ -"""UiPath Guardrails middleware for LangChain agents. +"""UiPath Guardrails for LangChain agents. -This module provides a developer-friendly API for configuring guardrails -that integrate with UiPath's guardrails service. +Platform guardrail decorators plus LangChain/LangGraph adapter auto-registration. """ from uipath.agent.models.agent import AgentGuardrailSeverityLevel from uipath.core.guardrails import GuardrailScope +from uipath.platform.guardrails.decorators import ( + BlockAction, + CustomValidator, + GuardrailAction, + GuardrailBlockException, + GuardrailExecutionStage, + GuardrailTargetAdapter, + GuardrailValidatorBase, + LogAction, + LoggingSeverityLevel, + PIIDetectionEntity, + PIIDetectionEntityType, + PIIValidator, + PromptInjectionValidator, + RuleFunction, + guardrail, + register_guardrail_adapter, +) -from .actions import BlockAction, LogAction -from .enums import GuardrailExecutionStage, PIIDetectionEntityType +from ._langchain_adapter import LangChainGuardrailAdapter from .middlewares import ( UiPathDeterministicGuardrailMiddleware, UiPathPIIDetectionMiddleware, UiPathPromptInjectionMiddleware, ) -from .models import GuardrailAction, PIIDetectionEntity + +# Auto-register the LangChain adapter so @guardrail knows how to wrap +# BaseTool, BaseChatModel, StateGraph, and CompiledStateGraph. +register_guardrail_adapter(LangChainGuardrailAdapter()) __all__ = [ + # Decorator + "guardrail", + # Validators + "GuardrailValidatorBase", + "PIIValidator", + "PromptInjectionValidator", + "CustomValidator", + "RuleFunction", + # Models & enums + "PIIDetectionEntity", "PIIDetectionEntityType", "GuardrailExecutionStage", "GuardrailScope", - "PIIDetectionEntity", "GuardrailAction", + # Actions "LogAction", "BlockAction", + "LoggingSeverityLevel", + # Exception + "GuardrailBlockException", + # Adapter registry + "GuardrailTargetAdapter", + "register_guardrail_adapter", + # Middlewares (unchanged) "UiPathPIIDetectionMiddleware", "UiPathPromptInjectionMiddleware", "UiPathDeterministicGuardrailMiddleware", - "AgentGuardrailSeverityLevel", # Re-export for convenience + # Re-exports for backwards compat + "AgentGuardrailSeverityLevel", ] diff --git a/src/uipath_langchain/guardrails/_langchain_adapter.py b/src/uipath_langchain/guardrails/_langchain_adapter.py new file mode 100644 index 000000000..326117468 --- /dev/null +++ b/src/uipath_langchain/guardrails/_langchain_adapter.py @@ -0,0 +1,537 @@ +"""LangChain/LangGraph adapter for the ``@guardrail`` decorator. + +Registers a :class:`LangChainGuardrailAdapter` that teaches the platform +``@guardrail`` decorator how to wrap LangChain and LangGraph objects: + +- ``BaseTool`` → TOOL scope +- ``BaseChatModel`` → LLM scope +- ``StateGraph`` / ``CompiledStateGraph`` → AGENT scope + +The adapter is auto-registered when ``uipath_langchain.guardrails`` is imported. +""" + +import logging +from functools import wraps +from typing import Any + +from langchain_core.language_models import BaseChatModel +from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, ToolMessage +from langchain_core.tools import BaseTool +from langgraph.graph import StateGraph +from langgraph.graph.state import CompiledStateGraph +from langgraph.types import Command +from uipath.core.guardrails import GuardrailScope, GuardrailValidationResultType +from uipath.platform.guardrails.decorators._core import ( + _EvaluatorFn, + _extract_input, + _extract_output, + _rewrap_input, +) +from uipath.platform.guardrails.decorators._enums import GuardrailExecutionStage +from uipath.platform.guardrails.decorators._exceptions import GuardrailBlockException +from uipath.platform.guardrails.decorators._models import GuardrailAction +from uipath.runtime.errors import UiPathErrorCategory + +from uipath_langchain.agent.exceptions import AgentRuntimeError, AgentRuntimeErrorCode +from uipath_langchain.guardrails.middlewares._utils import create_modified_tool_result + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Exception conversion helper +# --------------------------------------------------------------------------- + + +def _convert_block_exception(exc: GuardrailBlockException) -> AgentRuntimeError: + """Convert a :class:`GuardrailBlockException` to :class:`AgentRuntimeError`.""" + return AgentRuntimeError( + code=AgentRuntimeErrorCode.TERMINATION_GUARDRAIL_VIOLATION, + title=exc.title, + detail=exc.detail, + category=UiPathErrorCategory.USER, + ) + + +# --------------------------------------------------------------------------- +# Message helpers +# --------------------------------------------------------------------------- + + +def _get_last_human_message(messages: list[BaseMessage]) -> HumanMessage | None: + """Return the last HumanMessage in a list, or None if absent.""" + for msg in reversed(messages): + if isinstance(msg, HumanMessage): + return msg + return None + + +def _get_last_ai_message(messages: list[BaseMessage]) -> AIMessage | None: + """Return the last AIMessage in a list, or None if absent.""" + for msg in reversed(messages): + if isinstance(msg, AIMessage): + return msg + return None + + +def _extract_message_text(msg: BaseMessage) -> str: + """Extract plain text content from a message.""" + if isinstance(msg.content, str): + return msg.content + if isinstance(msg.content, list): + parts = [ + part.get("text", "") + for part in msg.content + if isinstance(part, dict) and part.get("type") == "text" + ] + return "\n".join(filter(None, parts)) + return "" + + +def _apply_message_text_modification(msg: BaseMessage, modified: str) -> None: + """Apply a modified text string back to a message in-place.""" + if isinstance(msg.content, str): + msg.content = modified + elif isinstance(msg.content, list): + for part in msg.content: + if isinstance(part, dict) and part.get("type") == "text": + part["text"] = modified + break + + +# --------------------------------------------------------------------------- +# LangChain-aware output extraction +# --------------------------------------------------------------------------- + + +def _lc_extract_output(result: Any) -> dict[str, Any]: + """Normalise tool output to a dict, handling LangGraph ToolMessage / Command. + + Unwraps ``ToolMessage`` and ``Command`` envelopes first, then delegates to + the platform's pure-Python ``_extract_output`` for JSON / literal-eval parsing. + """ + content: Any = result + if isinstance(result, Command): + update = result.update if hasattr(result, "update") else {} + messages = update.get("messages", []) if isinstance(update, dict) else [] + if messages and isinstance(messages[0], ToolMessage): + content = messages[0].content + else: + return {} + elif isinstance(result, ToolMessage): + content = result.content + + if content is not result: + # Delegate to the pure-Python extractor with pre-processed content. + return _extract_output(content) + return _extract_output(result) + + +# --------------------------------------------------------------------------- +# Tool wrapper +# --------------------------------------------------------------------------- + + +def _wrap_tool_with_guardrail( + tool: BaseTool, + evaluator: _EvaluatorFn, + action: GuardrailAction, + name: str, + stage: GuardrailExecutionStage, +) -> BaseTool: + """Wrap a ``BaseTool`` using Pydantic ``__class__`` swapping.""" + _stage = stage + + def _apply_pre(tool_input: Any) -> Any: + input_data = _extract_input(tool_input) + try: + result = evaluator( + input_data, GuardrailExecutionStage.PRE, input_data, None + ) + except Exception as exc: + logger.error( + "Error evaluating guardrail (pre) for tool %r: %s", + tool.name, + exc, + exc_info=True, + ) + return tool_input + modified = None + if result.result == GuardrailValidationResultType.VALIDATION_FAILED: + try: + modified = action.handle_validation_result(result, input_data, name) + except GuardrailBlockException as exc: + raise _convert_block_exception(exc) from exc + if modified is not None and isinstance(modified, dict): + return _rewrap_input(tool_input, modified) + return tool_input + + def _apply_post(tool_input: Any, raw_result: Any) -> Any: + input_data = _extract_input(tool_input) + output_data = _lc_extract_output(raw_result) + try: + result = evaluator( + output_data, GuardrailExecutionStage.POST, input_data, output_data + ) + except Exception as exc: + logger.error( + "Error evaluating guardrail (post) for tool %r: %s", + tool.name, + exc, + exc_info=True, + ) + return raw_result + modified = None + if result.result == GuardrailValidationResultType.VALIDATION_FAILED: + try: + modified = action.handle_validation_result(result, output_data, name) + except GuardrailBlockException as exc: + raise _convert_block_exception(exc) from exc + if modified is not None: + if isinstance(raw_result, (ToolMessage, Command)): + return create_modified_tool_result(raw_result, modified) + return modified + return raw_result + + ConcreteToolType = type(tool) + + class _GuardedTool(ConcreteToolType): # type: ignore[valid-type, misc] + def invoke(self, tool_input: Any, config: Any = None, **kwargs: Any) -> Any: + guarded_input = tool_input + if _stage in ( + GuardrailExecutionStage.PRE, + GuardrailExecutionStage.PRE_AND_POST, + ): + guarded_input = _apply_pre(tool_input) + result = super().invoke(guarded_input, config, **kwargs) + if _stage in ( + GuardrailExecutionStage.POST, + GuardrailExecutionStage.PRE_AND_POST, + ): + result = _apply_post(guarded_input, result) + return result + + # ainvoke is intentionally NOT overridden here. + # StructuredTool.ainvoke (for sync tools without a coroutine) delegates to + # self.invoke via run_in_executor. Overriding ainvoke would cause the POST + # guardrail to fire twice: once inside self.invoke and once after + # super().ainvoke() returns. + + tool.__class__ = _GuardedTool # type: ignore[assignment] + return tool + + +# --------------------------------------------------------------------------- +# LLM wrapper +# --------------------------------------------------------------------------- + + +def _apply_llm_pre( + messages: list[BaseMessage], + evaluator: _EvaluatorFn, + action: GuardrailAction, + name: str, +) -> None: + """Evaluate the last HumanMessage in-place (PRE stage, LLM scope).""" + msg = _get_last_human_message(messages) + if msg is None: + return + text = _extract_message_text(msg) + if not text: + return + try: + result = evaluator(text, GuardrailExecutionStage.PRE, None, None) + except Exception: + return + modified = None + if result.result == GuardrailValidationResultType.VALIDATION_FAILED: + try: + modified = action.handle_validation_result(result, text, name) + except GuardrailBlockException as exc: + raise _convert_block_exception(exc) from exc + if isinstance(modified, str) and modified != text: + _apply_message_text_modification(msg, modified) + + +def _apply_llm_post( + response: AIMessage, + evaluator: _EvaluatorFn, + action: GuardrailAction, + name: str, +) -> None: + """Evaluate the AIMessage content in-place (POST stage, LLM scope).""" + if not isinstance(response.content, str) or not response.content: + return + try: + result = evaluator(response.content, GuardrailExecutionStage.POST, None, None) + except Exception: + return + modified = None + if result.result == GuardrailValidationResultType.VALIDATION_FAILED: + try: + modified = action.handle_validation_result(result, response.content, name) + except GuardrailBlockException as exc: + raise _convert_block_exception(exc) from exc + if isinstance(modified, str) and modified != response.content: + response.content = modified + + +def _wrap_llm_with_guardrail( + llm: BaseChatModel, + evaluator: _EvaluatorFn, + action: GuardrailAction, + name: str, + stage: GuardrailExecutionStage, +) -> BaseChatModel: + """Wrap a ``BaseChatModel`` using Pydantic ``__class__`` swapping.""" + _stage = stage + ConcreteType = type(llm) + + class _GuardedLLM(ConcreteType): # type: ignore[valid-type, misc] + def invoke(self, messages: Any, config: Any = None, **kwargs: Any) -> Any: + if isinstance(messages, list) and _stage in ( + GuardrailExecutionStage.PRE, + GuardrailExecutionStage.PRE_AND_POST, + ): + _apply_llm_pre(messages, evaluator, action, name) + response = super().invoke(messages, config, **kwargs) + if isinstance(response, AIMessage) and _stage in ( + GuardrailExecutionStage.POST, + GuardrailExecutionStage.PRE_AND_POST, + ): + _apply_llm_post(response, evaluator, action, name) + return response + + async def ainvoke( + self, messages: Any, config: Any = None, **kwargs: Any + ) -> Any: + if isinstance(messages, list) and _stage in ( + GuardrailExecutionStage.PRE, + GuardrailExecutionStage.PRE_AND_POST, + ): + _apply_llm_pre(messages, evaluator, action, name) + response = await super().ainvoke(messages, config, **kwargs) + if isinstance(response, AIMessage) and _stage in ( + GuardrailExecutionStage.POST, + GuardrailExecutionStage.PRE_AND_POST, + ): + _apply_llm_post(response, evaluator, action, name) + return response + + llm.__class__ = _GuardedLLM # type: ignore[assignment] + return llm + + +# --------------------------------------------------------------------------- +# StateGraph / CompiledStateGraph wrappers (AGENT scope) +# --------------------------------------------------------------------------- + + +def _apply_agent_input_guardrail( + input: Any, + evaluator: _EvaluatorFn, + action: GuardrailAction, + name: str, +) -> None: + """Evaluate the last HumanMessage from agent input in-place.""" + if not isinstance(input, dict) or "messages" not in input: + return + messages = input["messages"] + if not isinstance(messages, list): + return + msg = _get_last_human_message(messages) + if msg is None: + return + text = _extract_message_text(msg) + if not text: + return + try: + result = evaluator(text, GuardrailExecutionStage.PRE, None, None) + except Exception: + return + modified = None + if result.result == GuardrailValidationResultType.VALIDATION_FAILED: + try: + modified = action.handle_validation_result(result, text, name) + except GuardrailBlockException as exc: + raise _convert_block_exception(exc) from exc + if isinstance(modified, str) and modified != text: + _apply_message_text_modification(msg, modified) + + +def _apply_agent_output_guardrail( + output: Any, + evaluator: _EvaluatorFn, + action: GuardrailAction, + name: str, +) -> None: + """Evaluate the last AIMessage from agent output in-place.""" + if not isinstance(output, dict) or "messages" not in output: + return + messages = output["messages"] + if not isinstance(messages, list): + return + msg = _get_last_ai_message(messages) + if msg is None: + return + text = _extract_message_text(msg) + if not text: + return + try: + result = evaluator(text, GuardrailExecutionStage.POST, None, None) + except Exception: + return + modified = None + if result.result == GuardrailValidationResultType.VALIDATION_FAILED: + try: + modified = action.handle_validation_result(result, text, name) + except GuardrailBlockException as exc: + raise _convert_block_exception(exc) from exc + if isinstance(modified, str) and modified != text: + _apply_message_text_modification(msg, modified) + + +def _wrap_stategraph_with_guardrail( + graph: "StateGraph[Any, Any]", + evaluator: _EvaluatorFn, + action: GuardrailAction, + name: str, + stage: GuardrailExecutionStage, +) -> "StateGraph[Any, Any]": + """Wrap a ``StateGraph``'s invoke/ainvoke to apply the guardrail.""" + if hasattr(graph, "invoke"): + original_invoke = graph.invoke + + @wraps(original_invoke) + def wrapped_invoke(input: Any, config: Any = None, **kwargs: Any) -> Any: + if stage in ( + GuardrailExecutionStage.PRE, + GuardrailExecutionStage.PRE_AND_POST, + ): + _apply_agent_input_guardrail(input, evaluator, action, name) + output = original_invoke(input, config, **kwargs) + if stage in ( + GuardrailExecutionStage.POST, + GuardrailExecutionStage.PRE_AND_POST, + ): + _apply_agent_output_guardrail(output, evaluator, action, name) + return output + + graph.invoke = wrapped_invoke + + if hasattr(graph, "ainvoke"): + original_ainvoke = graph.ainvoke + + @wraps(original_ainvoke) + async def wrapped_ainvoke(input: Any, config: Any = None, **kwargs: Any) -> Any: + if stage in ( + GuardrailExecutionStage.PRE, + GuardrailExecutionStage.PRE_AND_POST, + ): + _apply_agent_input_guardrail(input, evaluator, action, name) + output = await original_ainvoke(input, config, **kwargs) + if stage in ( + GuardrailExecutionStage.POST, + GuardrailExecutionStage.PRE_AND_POST, + ): + _apply_agent_output_guardrail(output, evaluator, action, name) + return output + + graph.ainvoke = wrapped_ainvoke + + return graph + + +def _wrap_compiled_graph_with_guardrail( + graph: "CompiledStateGraph[Any, Any, Any]", + evaluator: _EvaluatorFn, + action: GuardrailAction, + name: str, + stage: GuardrailExecutionStage, +) -> "CompiledStateGraph[Any, Any, Any]": + """Wrap a ``CompiledStateGraph``'s invoke/ainvoke to apply the guardrail.""" + original_invoke = graph.invoke + original_ainvoke = graph.ainvoke + + @wraps(original_invoke) + def wrapped_invoke(input: Any, config: Any = None, **kwargs: Any) -> Any: + if stage in ( + GuardrailExecutionStage.PRE, + GuardrailExecutionStage.PRE_AND_POST, + ): + _apply_agent_input_guardrail(input, evaluator, action, name) + output = original_invoke(input, config, **kwargs) + if stage in ( + GuardrailExecutionStage.POST, + GuardrailExecutionStage.PRE_AND_POST, + ): + _apply_agent_output_guardrail(output, evaluator, action, name) + return output + + @wraps(original_ainvoke) + async def wrapped_ainvoke(input: Any, config: Any = None, **kwargs: Any) -> Any: + if stage in ( + GuardrailExecutionStage.PRE, + GuardrailExecutionStage.PRE_AND_POST, + ): + _apply_agent_input_guardrail(input, evaluator, action, name) + output = await original_ainvoke(input, config, **kwargs) + if stage in ( + GuardrailExecutionStage.POST, + GuardrailExecutionStage.PRE_AND_POST, + ): + _apply_agent_output_guardrail(output, evaluator, action, name) + return output + + graph.invoke = wrapped_invoke # type: ignore[method-assign] + graph.ainvoke = wrapped_ainvoke # type: ignore[method-assign] + return graph + + +# --------------------------------------------------------------------------- +# Adapter implementation +# --------------------------------------------------------------------------- + + +class LangChainGuardrailAdapter: + """Framework adapter for LangChain/LangGraph objects. + + Implements :class:`~uipath.platform.guardrails.decorators.GuardrailTargetAdapter` + for ``BaseTool``, ``BaseChatModel``, ``StateGraph``, and + ``CompiledStateGraph``. + + Auto-registered when ``uipath_langchain.guardrails`` is imported. + """ + + def detect_scope(self, target: Any) -> GuardrailScope | None: + """Return the guardrail scope for a LangChain/LangGraph object.""" + if isinstance(target, BaseTool): + return GuardrailScope.TOOL + if isinstance(target, BaseChatModel): + return GuardrailScope.LLM + if isinstance(target, (StateGraph, CompiledStateGraph)): + return GuardrailScope.AGENT + return None + + def wrap( + self, + target: Any, + evaluator: _EvaluatorFn, + action: GuardrailAction, + name: str, + stage: GuardrailExecutionStage, + ) -> Any: + """Wrap a LangChain/LangGraph object with guardrail enforcement.""" + if isinstance(target, BaseTool): + return _wrap_tool_with_guardrail(target, evaluator, action, name, stage) + if isinstance(target, BaseChatModel): + return _wrap_llm_with_guardrail(target, evaluator, action, name, stage) + if isinstance(target, CompiledStateGraph): + return _wrap_compiled_graph_with_guardrail( + target, evaluator, action, name, stage + ) + if isinstance(target, StateGraph): + return _wrap_stategraph_with_guardrail( + target, evaluator, action, name, stage + ) + return target diff --git a/src/uipath_langchain/guardrails/actions.py b/src/uipath_langchain/guardrails/actions.py index d4fcd5ec7..49aa5993f 100644 --- a/src/uipath_langchain/guardrails/actions.py +++ b/src/uipath_langchain/guardrails/actions.py @@ -1,112 +1,9 @@ -"""Example implementations of GuardrailAction. +"""Action implementations for UiPath guardrails.""" -This module provides example implementations of the GuardrailAction interface. -These are part of the SDK and demonstrate common use cases for guardrail actions. -""" - -import logging -from dataclasses import dataclass -from enum import Enum -from typing import Any, Optional - -from uipath.agent.models.agent import AgentGuardrailSeverityLevel -from uipath.core.guardrails import ( - GuardrailValidationResult, - GuardrailValidationResultType, -) -from uipath.runtime.errors import UiPathErrorCategory - -from uipath_langchain.agent.exceptions import ( - AgentRuntimeError, - AgentRuntimeErrorCode, +from uipath.platform.guardrails.decorators import ( + BlockAction, + LogAction, + LoggingSeverityLevel, ) -from .models import GuardrailAction - - -def _severity_to_log_level(severity: AgentGuardrailSeverityLevel) -> int: - """Convert AgentGuardrailSeverityLevel to Python logging level.""" - mapping = { - AgentGuardrailSeverityLevel.ERROR: logging.ERROR, - AgentGuardrailSeverityLevel.WARNING: logging.WARNING, - AgentGuardrailSeverityLevel.INFO: logging.INFO, - } - return mapping.get(severity, logging.WARNING) - - -# here! -class LoggingSeverityLevel(int, Enum): - """Severity level enumeration.""" - - ERROR = logging.ERROR - INFO = logging.INFO - WARNING = logging.WARNING - DEBUG = logging.DEBUG - - -@dataclass -class LogAction(GuardrailAction): - """Example implementation: Log action for guardrails. - - This action logs guardrail violations at a specified severity level. - - Args: - severity_level: Severity level for logging (Error, Warning, Info) - message: Optional custom message to log. If not provided, a default - message will be generated. - """ - - severity_level: LoggingSeverityLevel = LoggingSeverityLevel.WARNING - message: Optional[str] = None - - def handle_validation_result( - self, - result: GuardrailValidationResult, - data: str | dict[str, Any], - guardrail_name: str, - ) -> str | dict[str, Any] | None: - """Handle validation result by logging it.""" - if result.result == GuardrailValidationResultType.VALIDATION_FAILED: - log_level = self.severity_level - log_level_name = logging.getLevelName(log_level) - message = self.message or f"Failed: {result.reason}" - logger = logging.getLogger(__name__) - logger.log(log_level, message) - print(f"[{log_level_name}][GUARDRAIL] [{guardrail_name}] {message}") - return None - - -@dataclass -class BlockAction(GuardrailAction): - """Example implementation: Block action for guardrails. - - This action blocks execution by raising an AgentRuntimeError when - a guardrail validation fails. - - Args: - title: Optional custom title for the termination exception. - If not provided, a default title will be generated. - detail: Optional custom detail message for the termination exception. - If not provided, the guardrail validation reason will be used. - """ - - title: Optional[str] = None - detail: Optional[str] = None - - def handle_validation_result( - self, - result: GuardrailValidationResult, - data: str | dict[str, Any], - guardrail_name: str, - ) -> str | dict[str, Any] | None: - """Handle validation result by blocking execution.""" - if result.result == GuardrailValidationResultType.VALIDATION_FAILED: - title = self.title or f"Guardrail [{guardrail_name}] blocked execution" - detail = self.detail or result.reason or "Guardrail validation failed" - raise AgentRuntimeError( - code=AgentRuntimeErrorCode.TERMINATION_GUARDRAIL_VIOLATION, - title=title, - detail=detail, - category=UiPathErrorCategory.USER, - ) - return None +__all__ = ["LoggingSeverityLevel", "LogAction", "BlockAction"] diff --git a/src/uipath_langchain/guardrails/decorators/__init__.py b/src/uipath_langchain/guardrails/decorators/__init__.py new file mode 100644 index 000000000..c81c5c7d3 --- /dev/null +++ b/src/uipath_langchain/guardrails/decorators/__init__.py @@ -0,0 +1,19 @@ +"""Guardrail decorators package.""" + +from uipath.platform.guardrails.decorators import ( + CustomValidator, + GuardrailValidatorBase, + PIIValidator, + PromptInjectionValidator, + RuleFunction, + guardrail, +) + +__all__ = [ + "guardrail", + "GuardrailValidatorBase", + "PIIValidator", + "PromptInjectionValidator", + "CustomValidator", + "RuleFunction", +] diff --git a/src/uipath_langchain/guardrails/enums.py b/src/uipath_langchain/guardrails/enums.py index c13ecf3a9..3bfb3624a 100644 --- a/src/uipath_langchain/guardrails/enums.py +++ b/src/uipath_langchain/guardrails/enums.py @@ -1,44 +1,9 @@ """Enums for UiPath guardrails.""" -from enum import Enum +from uipath.core.guardrails import GuardrailScope +from uipath.platform.guardrails.decorators import ( + GuardrailExecutionStage, + PIIDetectionEntityType, +) -from uipath.core.guardrails import GuardrailScope as CoreGuardrailScope - -# Re-export GuardrailScope from core for convenience -GuardrailScope = CoreGuardrailScope - - -class PIIDetectionEntityType(str, Enum): - """PII detection entity types supported by UiPath guardrails. - - These entities match the available options from the UiPath guardrails service backend. - The enum values correspond to the exact strings expected by the backend API. - """ - - PERSON = "Person" - ADDRESS = "Address" - DATE = "Date" - PHONE_NUMBER = "PhoneNumber" - EUGPS_COORDINATES = "EugpsCoordinates" - EMAIL = "Email" - CREDIT_CARD_NUMBER = "CreditCardNumber" - INTERNATIONAL_BANKING_ACCOUNT_NUMBER = "InternationalBankingAccountNumber" - SWIFT_CODE = "SwiftCode" - ABA_ROUTING_NUMBER = "ABARoutingNumber" - US_DRIVERS_LICENSE_NUMBER = "USDriversLicenseNumber" - UK_DRIVERS_LICENSE_NUMBER = "UKDriversLicenseNumber" - US_INDIVIDUAL_TAXPAYER_IDENTIFICATION = "USIndividualTaxpayerIdentification" - UK_UNIQUE_TAXPAYER_NUMBER = "UKUniqueTaxpayerNumber" - US_BANK_ACCOUNT_NUMBER = "USBankAccountNumber" - US_SOCIAL_SECURITY_NUMBER = "USSocialSecurityNumber" - USUK_PASSPORT_NUMBER = "UsukPassportNumber" - URL = "URL" - IP_ADDRESS = "IPAddress" - - -class GuardrailExecutionStage(str, Enum): - """Execution stage for deterministic guardrails.""" - - PRE = "pre" # Pre-execution only - POST = "post" # Post-execution only - PRE_AND_POST = "pre&post" # Both pre and post execution +__all__ = ["GuardrailScope", "PIIDetectionEntityType", "GuardrailExecutionStage"] diff --git a/src/uipath_langchain/guardrails/middlewares/deterministic.py b/src/uipath_langchain/guardrails/middlewares/deterministic.py index 606d4d5b7..fdaaceab7 100644 --- a/src/uipath_langchain/guardrails/middlewares/deterministic.py +++ b/src/uipath_langchain/guardrails/middlewares/deterministic.py @@ -65,6 +65,7 @@ class UiPathDeterministicGuardrailMiddleware: rules=[], action=CustomFilterAction(...), stage=GuardrailExecutionStage.POST, + enabled_for_evals=False, ) agent = create_agent( @@ -91,6 +92,8 @@ class UiPathDeterministicGuardrailMiddleware: - GuardrailExecutionStage.PRE_AND_POST: Validate both input and output name: Optional name for the guardrail (defaults to "Deterministic Guardrail") description: Optional description for the guardrail + enabled_for_evals: Whether this guardrail is enabled for evaluation scenarios. + Defaults to True. """ def __init__( @@ -102,6 +105,7 @@ def __init__( *, name: str = "Deterministic Guardrail", description: str | None = None, + enabled_for_evals: bool = True, ): """Initialize deterministic guardrail middleware.""" if not tools: @@ -112,6 +116,8 @@ def __init__( raise ValueError( f"stage must be an instance of GuardrailExecutionStage, got {type(stage)}" ) + if not isinstance(enabled_for_evals, bool): + raise ValueError("enabled_for_evals must be a boolean") for i, rule in enumerate(rules): if not callable(rule): @@ -139,6 +145,7 @@ def __init__( self.action = action self._stage = stage self._name = name + self.enabled_for_evals = enabled_for_evals self._description = description or "Deterministic guardrail with custom rules" self._middleware_instances = self._create_middleware_instances() diff --git a/src/uipath_langchain/guardrails/middlewares/pii_detection.py b/src/uipath_langchain/guardrails/middlewares/pii_detection.py index be5a95334..0ee2c395e 100644 --- a/src/uipath_langchain/guardrails/middlewares/pii_detection.py +++ b/src/uipath_langchain/guardrails/middlewares/pii_detection.py @@ -30,6 +30,9 @@ GuardrailScope, MapEnumParameterValue, ) +from uipath.platform.guardrails.decorators._exceptions import GuardrailBlockException + +from uipath_langchain.agent.exceptions import AgentRuntimeError from ..models import GuardrailAction, PIIDetectionEntity from ._utils import ( @@ -70,6 +73,7 @@ def analyze_joke_syntax(joke: str) -> str: PIIDetectionEntity(PIIDetectionEntityType.EMAIL, 0.5), PIIDetectionEntity(PIIDetectionEntityType.ADDRESS, 0.7), ], + enabled_for_evals=True, ) # PII detection for specific tools (using tool reference directly) @@ -78,6 +82,7 @@ def analyze_joke_syntax(joke: str) -> str: action=LogAction(severity_level=LoggingSeverityLevel.WARNING), entities=[PIIDetectionEntity(PIIDetectionEntityType.EMAIL, 0.5)], tools=[analyze_joke_syntax], + enabled_for_evals=False, ) agent = create_agent( @@ -97,6 +102,8 @@ def analyze_joke_syntax(joke: str) -> str: If TOOL scope is not specified, this parameter is ignored. name: Optional name for the guardrail (defaults to "PII Detection") description: Optional description for the guardrail + enabled_for_evals: Whether this guardrail is enabled for evaluation scenarios. + Defaults to True. """ def __init__( @@ -108,6 +115,7 @@ def __init__( tools: Sequence[str | BaseTool] | None = None, name: str = "PII Detection", description: str | None = None, + enabled_for_evals: bool = True, ): """Initialize PII detection guardrail middleware.""" if not scopes: @@ -116,6 +124,8 @@ def __init__( raise ValueError("At least one entity must be specified") if not isinstance(action, GuardrailAction): raise ValueError("action must be an instance of GuardrailAction") + if not isinstance(enabled_for_evals, bool): + raise ValueError("enabled_for_evals must be a boolean") self._tool_names: list[str] | None = None if tools is not None: @@ -144,6 +154,7 @@ def __init__( self.action = action self.entities = list(entities) self._name = name + self.enabled_for_evals = enabled_for_evals self._description = ( description or f"Detects PII entities: {', '.join(e.name for e in entities)}" @@ -230,6 +241,8 @@ async def _wrap_tool_call_func( ) if modified_input is not None and isinstance(modified_input, dict): request = create_modified_tool_request(request, modified_input) + except (AgentRuntimeError, GuardrailBlockException): + raise except Exception as e: logger.error( f"Error evaluating PII guardrail for tool '{tool_name}': {e}", @@ -274,7 +287,7 @@ def _create_guardrail(self) -> BuiltInValidatorGuardrail: id=str(uuid4()), name=self._name, description=self._description, - enabled_for_evals=True, + enabled_for_evals=self.enabled_for_evals, selector=GuardrailSelector(**selector_kwargs), guardrail_type="builtInValidator", validator_type="pii_detection", @@ -334,5 +347,7 @@ def _check_messages(self, messages: list[BaseMessage]) -> None: if isinstance(msg.content, str) and text in msg.content: msg.content = msg.content.replace(text, modified_text, 1) break + except (AgentRuntimeError, GuardrailBlockException): + raise except Exception as e: logger.error(f"Error evaluating PII guardrail: {e}", exc_info=True) diff --git a/src/uipath_langchain/guardrails/middlewares/prompt_injection.py b/src/uipath_langchain/guardrails/middlewares/prompt_injection.py index 787fd10bb..8e8964e2b 100644 --- a/src/uipath_langchain/guardrails/middlewares/prompt_injection.py +++ b/src/uipath_langchain/guardrails/middlewares/prompt_injection.py @@ -14,8 +14,11 @@ ) from uipath.platform import UiPath from uipath.platform.guardrails import BuiltInValidatorGuardrail, GuardrailScope +from uipath.platform.guardrails.decorators._exceptions import GuardrailBlockException from uipath.platform.guardrails.guardrails import NumberParameterValue +from uipath_langchain.agent.exceptions import AgentRuntimeError + from ..models import GuardrailAction from ._utils import extract_text_from_messages @@ -37,6 +40,7 @@ class UiPathPromptInjectionMiddleware: middleware = UiPathPromptInjectionMiddleware( action=LogAction(severity_level=LoggingSeverityLevel.WARNING), threshold=0.5, + enabled_for_evals=True, ) ``` @@ -47,6 +51,8 @@ class UiPathPromptInjectionMiddleware: threshold: Detection threshold (0.0 to 1.0) name: Optional name for the guardrail (defaults to "Prompt Injection Detection") description: Optional description for the guardrail + enabled_for_evals: Whether this guardrail is enabled for evaluation scenarios. + Defaults to True. """ def __init__( @@ -57,12 +63,15 @@ def __init__( scopes: Sequence[GuardrailScope] | None = None, name: str = "Prompt Injection Detection", description: str | None = None, + enabled_for_evals: bool = True, ): """Initialize prompt injection detection guardrail middleware.""" if not isinstance(action, GuardrailAction): raise ValueError("action must be an instance of GuardrailAction") if not 0.0 <= threshold <= 1.0: raise ValueError(f"Threshold must be between 0.0 and 1.0, got {threshold}") + if not isinstance(enabled_for_evals, bool): + raise ValueError("enabled_for_evals must be a boolean") scopes_list = list(scopes) if scopes is not None else [GuardrailScope.LLM] if scopes_list != [GuardrailScope.LLM]: @@ -75,6 +84,7 @@ def __init__( self.action = action self.threshold = threshold self._name = name + self.enabled_for_evals = enabled_for_evals self._description = ( description or f"Detects prompt injection attempts with threshold {threshold}" @@ -118,7 +128,7 @@ def _create_guardrail(self) -> BuiltInValidatorGuardrail: id=str(uuid4()), name=self._name, description=self._description, - enabled_for_evals=True, + enabled_for_evals=self.enabled_for_evals, selector=GuardrailSelector(scopes=self.scopes), guardrail_type="builtInValidator", validator_type="prompt_injection", @@ -168,6 +178,8 @@ def _check_messages(self, messages: list[BaseMessage]) -> None: if isinstance(msg.content, str) and text in msg.content: msg.content = msg.content.replace(text, modified_text, 1) break + except (AgentRuntimeError, GuardrailBlockException): + raise except Exception as e: logger.error( f"Error evaluating prompt injection guardrail: {e}", exc_info=True diff --git a/src/uipath_langchain/guardrails/models.py b/src/uipath_langchain/guardrails/models.py index c5eb8e722..dc7a07572 100644 --- a/src/uipath_langchain/guardrails/models.py +++ b/src/uipath_langchain/guardrails/models.py @@ -1,90 +1,5 @@ """Models for UiPath guardrails configuration.""" -from abc import ABC, abstractmethod -from dataclasses import dataclass -from typing import Any +from uipath.platform.guardrails.decorators import GuardrailAction, PIIDetectionEntity -from uipath.core.guardrails import GuardrailValidationResult - - -@dataclass -class PIIDetectionEntity: - """PII entity configuration with threshold. - - Args: - name: The entity type (e.g., PIIDetectionEntity.EMAIL) - threshold: Confidence threshold (0.0 to 1.0) for detection - """ - - name: str - threshold: float = 0.5 - - def __post_init__(self): - """Validate threshold range.""" - if not 0.0 <= self.threshold <= 1.0: - raise ValueError( - f"Threshold must be between 0.0 and 1.0, got {self.threshold}" - ) - - -class GuardrailAction(ABC): - """Interface for defining custom actions when guardrails are triggered. - - Extend this interface to implement custom behavior when guardrail validation fails. - Common use cases include: - - Logging violations (see actions.LogAction) - - Blocking execution (see actions.BlockAction) - - Escalating to monitoring systems - - Sending alerts or notifications - - Collecting metrics or analytics - - Example: - ```python - from uipath_langchain.guardrails import GuardrailAction - from uipath.core.guardrails import GuardrailValidationResult, GuardrailValidationResultType - - class CustomAction(GuardrailAction): - def handle_validation_result( - self, - result: GuardrailValidationResult, - input_data: str | dict[str, Any], - guardrail_name: str, - ) -> str | dict[str, Any] | None: - if result.result == GuardrailValidationResultType.VALIDATION_FAILED: - # Your custom logic here - # Return modified data or None - return None - ``` - """ - - @abstractmethod - def handle_validation_result( - self, - result: GuardrailValidationResult, - data: str | dict[str, Any], - guardrail_name: str, - ) -> str | dict[str, Any] | None: - """Handle a guardrail validation result. - - This method is called when a guardrail validation fails. - Actions can optionally return modified data to sanitize/filter - the validated data before execution continues. - - Args: - result: The validation result from the guardrails service - data: The data that was validated (string or dictionary). - This can be tool input (arguments), tool output (result), - or message content depending on the guardrail scope. - guardrail_name: The name of the guardrail that triggered - - Returns: - Modified data if the action wants to sanitize/filter the validated data, - or None if no modification is needed. If None is returned, original - data is used. If a value is returned, it replaces the original data. - - Note: The returned data type should match the data type: - - For tool input: return dict[str, Any] (tool arguments) - - For tool output: return dict[str, Any] (tool result) - - For messages: return str (message content) - """ - pass +__all__ = ["PIIDetectionEntity", "GuardrailAction"] diff --git a/uv.lock b/uv.lock index 4a41c373e..6d6ff018d 100644 --- a/uv.lock +++ b/uv.lock @@ -3333,7 +3333,7 @@ wheels = [ [[package]] name = "uipath-langchain" -version = "0.9.12" +version = "0.9.13" source = { editable = "." } dependencies = [ { name = "httpx" },