feat: decorator guardrail implementation [AL-288] by apetraru-uipath · Pull Request #699 · UiPath/uipath-langchain-python

apetraru-uipath · 2026-03-17T13:48:31Z

What changed?

This pull request adds guardrails as decorators so you can protect agents, language models, and tools without using middleware. Guardrails run when you call the agent or its parts, and you can mix different kinds of protection (PII detection, prompt injection, and custom rules) at the level that fits your app.

In short:

You can attach PII detection and prompt injection guardrails to the LLM and to the whole agent.
You can attach PII and custom rules (e.g. block or filter certain content) to individual tools.
Blocking guardrails stop execution and report a clear error when a violation is found; logging guardrails record a warning and continue.
A new sample (joke-agent-decorator) shows how to use only decorators (no middleware) and how to verify that guardrails behave as expected.
The existing joke-agent sample stays as before, using the middleware-based guardrails.

No existing APIs are removed or changed in a breaking way.

How has this been tested?

Dependencies were installed and the decorator sample was run with different inputs (normal topic, filtered word, PII in the topic) to confirm guardrails behave correctly.
Static type checking (mypy) was run and passes.
The middleware-based joke-agent sample was checked to ensure it still matches the main branch except for the intended API usage.

Are there any breaking changes?

uv.lock

src/uipath_langchain/guardrails/decorators/pii.py

valentinabojan · 2026-03-19T13:36:37Z

src/uipath_langchain/guardrails/decorators/pii_detection.py

+        config={
+            "entities": list(entities),
+            "action": action,
+            "stage": GuardrailExecutionStage.PRE_AND_POST,


Should we allow the developer to specify the stage/stages where he want to enforce the guardrail?

In the future yes, I would leave it as a new feature. Now the behaviour is pre&post

I see this was done in this PR, right?

src/uipath_langchain/guardrails/decorators/prompt_injection.py

src/uipath_langchain/guardrails/decorators/deterministic.py

valentinabojan · 2026-03-19T13:56:40Z

src/uipath_langchain/guardrails/decorators/_base.py

+            ):
+                return GuardrailScope.AGENT
+
+    if callable(obj):


Is this a correct assumption?
Also, do we know if all Tools, Llms, agents are instances of BaseTool, BaseChatModel and StateGraph?

I'm thinking, what if I want to create subgraph in my loop, then I will have a StateGraph, right? Will this be seen as an agent scope?

Right, added a more strict rule

Looks like now a subgraph is a candidate for the guardrails to get applied? Am I understanding right?
If so, is this what we want to achieve? What's the value in having guardrails for subgraphs?

src/uipath_langchain/guardrails/decorators/_base.py

samples/joke-agent-decorator/graph.py

Add stage parameter to pii_detection_guardrail and prompt_injection_guardrail decorators, matching the flexibility already available in deterministic_guardrail. pii_detection_guardrail defaults to PRE_AND_POST; prompt_injection_guardrail defaults to PRE and raises pydantic.ValidationError for POST or PRE_AND_POST (prompt injection is input-only). Fix _apply_llm_input_guardrail and _apply_guardrail_to_message_list: previously concatenated the full conversation history and tried to replace the joined text inside a single message, which silently failed in any multi-turn conversation. Now evaluates only the last HumanMessage (PRE/input) or last AIMessage (POST/output) — semantically correct and replacement works reliably. Replace _extract_text_from_messages with focused helpers: _get_last_human_message, _get_last_ai_message, _extract_message_text, _apply_message_text_modification. Add target_type parameter to _apply_guardrail_to_message_list so input and output graph-scope wrappers target the correct message type. PII decorator, joke-agent-decorator sample, BlockAction uses AgentRuntimeError. Fix _wrap_llm_with_guardrail: factory functions returning BaseChatModel were not wrapped (fell through StateGraph/dict branches). Also fix Pydantic setattr block on UiPathChat by using __class__ swap to a dynamic subclass instead of monkey-patching invoke/ainvoke. Fix BlockAction swallowed by bare except: split try/except so only guardrail API errors are suppressed; action exceptions (AgentRuntimeError) now propagate. Fix CompiledStateGraph not recognised: add _wrap_compiled_graph_with_guardrail and handle CompiledStateGraph return type from factory functions. Fix mypy errors in decorators.py: typed list[BaseMessage], CompiledStateGraph type params, type: ignore[valid-type, misc] for dynamic subclass, and type: ignore[method-assign] for CompiledStateGraph method patching. Add prompt_injection_guardrail decorator: _create_prompt_injection_guardrail, _apply_prompt_injection_guardrail, public prompt_injection_guardrail function; exported from guardrails/__init__.py; stacked on create_llm() in joke-agent-decorator/graph.py with BlockAction to block on detection. Reformat decorators.py for consistency. Middleware cleanup: delete monolithic middleware.py (duplicate of middlewares/ split files); update guardrails/__init__.py to import from .middlewares; update joke-agent/graph.py to use new split-file API (tool_names -> tools, optional scopes on PromptInjection, unconditional POST filter comment). Revert renames: restore LoggingSeverityLevel as proper int Enum (ERROR, INFO, WARNING, DEBUG) in actions.py; remove PromptInjectionValidatorType from enums.py; fix pii_detection.py docstring to use Entity+PIIDetectionEntity; export LoggingSeverityLevel from guardrails/__init__.py; update both samples to use LoggingSeverityLevel instead of AgentGuardrailSeverityLevel. Manual refinements: updates to actions.py, decorators.py, enums.py, models.py, middlewares/pii_detection.py, guardrails/__init__.py, and joke-agent/graph.py. Remove joke-agent/.agent/REQUIRED_STRUCTURE.md; further manual edits to joke-agent/graph.py. Refactor decorators.py into decorators/ package (pii.py, prompt_injection.py, deterministic.py, _base.py) with tool-level guardrail support: - Split monolithic decorators.py into decorators/ subpackage - Add _wrap_tool_with_guardrail using __class__ swap (Pydantic-safe) - Add deterministic_guardrail decorator (TOOL scope, local rules, no API call) - Extend pii_guardrail to support BaseTool and optional tools= kwarg - Extend _detect_scope to return GuardrailScope.TOOL for BaseTool instances - Export deterministic_guardrail and RuleFunction from guardrails/__init__.py - Update joke-agent-decorator/graph.py to demonstrate all three decorator types on analyze_joke_syntax tool (3x @deterministic_guardrail + @pii_guardrail) - Add local CustomFilterAction to joke-agent-decorator sample Tool guardrail fixes and sample updates: - _base.py: unwrap LangGraph tool-call envelope (args) for rule evaluation; rewrap modified args so super().invoke() receives valid input; handle ToolMessage/Command in _extract_output for POST-stage deterministic rules. - joke-agent-decorator: Agent PII uses LogAction(WARNING) with custom message; README aligned with current guardrails and verification scenarios. Deps and lockfiles: pyproject.toml updates (root and joke-agent-decorator); remove samples/joke-agent-decorator/uv.lock and samples/joke-agent/uv.lock; uv.lock at repo root updated. Refactor _base.py: remove unnecessary casts in _evaluate_rules; catch only ValueError (JSONDecodeError is subclass); extract _apply_guardrail_to_message_list, _apply_guardrail_to_input_messages, _apply_guardrail_to_output_messages and use them in _wrap_stategraph_with_guardrail, _wrap_compiled_graph_with_guardrail, and _wrap_function_with_guardrail to reduce cognitive complexity. Made-with: Cursor Enable enabled_for_evals override across decorators and middlewares for PII, prompt injection, and deterministic guardrails (default true, user-overridable), plus docs/sample updates for the new parameter. Fix mypy errors in pii_detection.py and prompt_injection.py: rebind guardrail to a typed non-optional local variable so mypy can narrow the type inside nested class closures. Made-with: Cursor Fix _wrap_compiled_graph_with_guardrail: output guardrail (POST stage) was never applied — invoke/ainvoke discarded the graph output without running _apply_guardrail_to_output_messages. Now captures the output and evaluates it, matching the behaviour of _wrap_stategraph_with_guardrail. Update pyproject.toml and uv.lock. Fix double POST guardrail application on StructuredTool: remove ainvoke override from _GuardedTool in deterministic.py and pii_detection.py. StructuredTool.ainvoke delegates to self.invoke via run_in_executor, so the guardrail chain in invoke already runs once; the ainvoke override caused a second POST application, producing "words++++" instead of "words++". Fix LogAction double-logging: replace print()+logger.log() with a single logger.log() call using the guardrail name as context prefix. Fix _evaluate_rules violation message: include guardrail name instead of positional index so errors read "Rule <name> detected violation" rather than "Rule 1 detected violation". Pass guardrail_name through _apply_pre/_apply_post call sites in deterministic.py. Fix AgentRuntimeError swallowed in guardrail middleware/decorator except handlers: add explicit except AgentRuntimeError: raise before the generic except Exception in pii_detection decorator (_apply_pre, _apply_post), pii_detection middleware (_wrap_tool_call_func, _check_messages), and prompt_injection middleware (_check_messages) so BlockAction errors propagate instead of being silently logged.

valentinabojan · 2026-03-27T08:12:40Z

src/uipath_langchain/guardrails/decorators/pii_detection.py

+    """
+    action = metadata.config["action"]
+    guardrail_name = metadata.name
+    api_guardrail = metadata.guardrail


nit: can we rename this to something else? api_guardrail was confusing the first time I read it

valentinabojan · 2026-03-27T08:12:58Z

src/uipath_langchain/guardrails/decorators/pii_detection.py

+    Returns:
+        The same LLM with wrapped invoke/ainvoke.
+    """
+    _guardrail_opt = metadata.guardrail


nit: same here - bit of a weird name

valentinabojan · 2026-03-27T08:18:40Z

src/uipath_langchain/guardrails/decorators/prompt_injection.py

+    Returns:
+        The same LLM with wrapped invoke/ainvoke.
+    """
+    _guardrail_opt = metadata.guardrail


nit: rename var name

valentinabojan · 2026-03-27T08:21:50Z

src/uipath_langchain/guardrails/decorators/prompt_injection.py

+    name: str = "Prompt Injection Detection",
+    description: str | None = None,
+    enabled_for_evals: bool = True,
+    stage: GuardrailExecutionStage = GuardrailExecutionStage.PRE,


Why expose this if we only allow PRE? Can't it be hidden from the users?

valentinabojan · 2026-03-27T08:23:06Z

src/uipath_langchain/guardrails/decorators/prompt_injection.py

+        raise ValueError("enabled_for_evals must be a boolean")
+
+    # Validate stage — prompt injection is an input-only concern
+    if stage != GuardrailExecutionStage.PRE:


If we hide the stage from the user, we might not need to do this stage validation here, right? Or at least check it and simply raise the error. Why have this _StageValidator in place?

valentinabojan · 2026-03-27T08:39:23Z

src/uipath_langchain/guardrails/decorators/_base.py

+    guardrail: BuiltInValidatorGuardrail,
+    uipath: UiPath,
+    action: GuardrailAction,
+    guardrail_name: str,


Why having guardrail_name separate from guardrail? I see BuiltInValidatorGuardrail has a guardrail name as field. Why not using that name?
I see this pattern in multiple files in this PR. Can we use a single param for the guardrail name?

apetraru-uipath force-pushed the feat/guardrails_decorators branch 21 times, most recently from 00d035f to 30ea969 Compare March 18, 2026 18:49

apetraru-uipath requested a review from cristipufu March 18, 2026 20:21