diff --git a/packages/uipath-pydantic-ai/pyproject.toml b/packages/uipath-pydantic-ai/pyproject.toml index 5b51024b..bcfc47c9 100644 --- a/packages/uipath-pydantic-ai/pyproject.toml +++ b/packages/uipath-pydantic-ai/pyproject.toml @@ -8,6 +8,7 @@ dependencies = [ "pydantic-ai>=1.63.0, <2.0.0", "openinference-instrumentation-pydantic-ai>=0.1.12", "uipath>=2.10.2, <2.11.0", + "uipath-core>=0.5.18, <0.7.0", "uipath-runtime>=0.11.0, <0.12.0", ] classifiers = [ @@ -27,6 +28,9 @@ register = "uipath_pydantic_ai.middlewares:register_middleware" [project.entry-points."uipath.runtime.factories"] pydantic-ai = "uipath_pydantic_ai.runtime:register_runtime_factory" +[project.entry-points."uipath.governance.adapters"] +pydantic-ai = "uipath_pydantic_ai.governance:register_governance_adapter" + [project.urls] Homepage = "https://uipath.com" Repository = "https://github.com/UiPath/uipath-integrations-python" diff --git a/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/governance/__init__.py b/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/governance/__init__.py new file mode 100644 index 00000000..0efd4f7e --- /dev/null +++ b/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/governance/__init__.py @@ -0,0 +1,52 @@ +"""Governance integration for ``uipath-pydantic-ai``. + +Registers :class:`PydanticAIAdapter` with the adapter registry in +``uipath.core.adapters`` so the governance host can attach the +Pydantic-AI-specific governance (BEFORE_MODEL, AFTER_MODEL, TOOL_CALL, +AFTER_TOOL) when it sees a ``pydantic_ai.Agent``. + +Registration is **idempotent**: calling :func:`register_governance_adapter` +twice is a no-op on the second call. + +Wiring: the package exposes :func:`register_governance_adapter` as an entry +point under ``uipath.governance.adapters``. The governance adapter discovery +path calls it to register the adapter. Importing this module does not, by +itself, mutate the global registry. +""" + +from __future__ import annotations + +import logging + +from uipath.core.adapters import get_adapter_registry + +from .adapter import GovernanceCallbacks, GovernanceModel, PydanticAIAdapter + +logger = logging.getLogger(__name__) + +_registered: bool = False + + +def register_governance_adapter() -> None: + """Register :class:`PydanticAIAdapter` with the global registry. + + Idempotent — safe to call multiple times. + """ + global _registered + if _registered: + return + registry = get_adapter_registry() + if any(a.name == "PydanticAI" for a in registry.get_all()): + _registered = True + return + registry.register(PydanticAIAdapter()) + _registered = True + logger.debug("Registered uipath-pydantic-ai governance adapter") + + +__all__ = [ + "GovernanceCallbacks", + "GovernanceModel", + "PydanticAIAdapter", + "register_governance_adapter", +] \ No newline at end of file diff --git a/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/governance/adapter.py b/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/governance/adapter.py new file mode 100644 index 00000000..be31c869 --- /dev/null +++ b/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/governance/adapter.py @@ -0,0 +1,373 @@ +"""Pydantic AI adapter for UiPath governance. + +Pydantic AI has the thinnest hook surface of the supported frameworks — there +is no per-agent callback or middleware system. But *everything* an agent does +flows through its ``Model``: the LLM request, the model's tool-call requests +(``ToolCallPart`` in the response), and the tool results fed back on the next +turn (``ToolReturnPart`` in the request). So this adapter governs by wrapping +``agent.model`` with a :class:`GovernanceModel` (a ``pydantic_ai`` ``WrapperModel``) +that brackets every model call: + +- BEFORE_MODEL — the latest request message's text (user prompt or tool result + being fed back), before delegating to the wrapped model. +- AFTER_TOOL — any ``ToolReturnPart`` in that latest request message. +- AFTER_MODEL — the ``TextPart`` content of the model's response. +- TOOL_CALL — each ``ToolCallPart`` the model emits (tool name + arguments). + +Both the non-streaming ``request`` and the streaming ``request_stream`` paths +are covered (the runtime uses ``agent.run`` and ``agent.iter`` respectively). + +Because the wrap is installed on ``agent.model`` in place, :meth:`attach` +returns the **original agent**; :meth:`detach` restores the original model. + +Chain-level boundaries (BEFORE_AGENT / AFTER_AGENT) are owned by the +governance host and are intentionally not fired here. + +Contracts and the evaluator protocol come from ``uipath-core``; this package +contributes only the Pydantic-AI-specific implementation and registers it with +the adapter registry via the ``uipath.governance.adapters`` entry point. + +Audit emission and enforcement (raising :class:`GovernanceBlockException` on +DENY) are owned by the evaluator. The wrapper only extracts payloads and calls +the matching ``evaluate_*`` method; :class:`GovernanceBlockException` propagates +(aborting the run), anything else is logged and swallowed. +""" + +from __future__ import annotations + +import json +import logging +from contextlib import asynccontextmanager +from typing import Any, AsyncIterator, Dict, List +from uuid import uuid4 + +from pydantic_ai.models import Model, ModelRequestParameters, StreamedResponse +from pydantic_ai.models.wrapper import WrapperModel +from pydantic_ai.settings import ModelSettings +from uipath.core.adapters import BaseAdapter, EvaluatorProtocol +from uipath.core.governance.exceptions import GovernanceBlockException + +logger = logging.getLogger(__name__) + +# Cap on the text blob passed to BEFORE_MODEL / AFTER_MODEL governance +# evaluation. Sized to match the runtime side and the other adapters. +_BEFORE_MODEL_TEXT_CAP = 64000 + +# Attribute used to stash the original (unwrapped) model so detach can restore it. +_ORIGINAL_MODEL_ATTR = "_uipath_governance_original_model" + + +class PydanticAIAdapter(BaseAdapter): + """Adapter for the Pydantic AI framework. + + Detects ``pydantic_ai.Agent`` instances and wraps their ``model`` with a + :class:`GovernanceModel`. + """ + + @property + def name(self) -> str: + return "PydanticAI" + + def can_handle(self, agent: Any) -> bool: + """Return True only for a ``pydantic_ai.Agent``.""" + try: + from pydantic_ai import Agent + except ImportError: + return False + return isinstance(agent, Agent) + + def attach( + self, + agent: Any, + agent_id: str, + session_id: str, + evaluator: EvaluatorProtocol, + ) -> Any: + """Wrap ``agent.model`` with governance (mutated in place). + + Returns the original ``agent``. If the agent has no concrete ``Model`` + bound (the model is supplied per-run), there is nothing to wrap and a + warning is logged. + """ + model = getattr(agent, "model", None) + if isinstance(model, GovernanceModel): + return agent # idempotent — already governed + if not isinstance(model, Model): + logger.warning( + "PydanticAIAdapter: agent has no bound Model to wrap (got %s); " + "model-layer governance will not fire", + type(model).__name__, + ) + return agent + callbacks = GovernanceCallbacks( + evaluator=evaluator, agent_name=agent_id, session_id=session_id + ) + setattr(agent, _ORIGINAL_MODEL_ATTR, model) + agent.model = GovernanceModel(model, callbacks) + logger.debug("Wrapped Pydantic AI agent model with governance") + return agent + + def detach(self, governed: Any) -> Any: + """Restore the agent's original (unwrapped) model and return it.""" + if isinstance(getattr(governed, "model", None), GovernanceModel): + original = getattr(governed, _ORIGINAL_MODEL_ATTR, None) + if original is not None: + governed.model = original + if hasattr(governed, _ORIGINAL_MODEL_ATTR): + delattr(governed, _ORIGINAL_MODEL_ATTR) + return governed + + +class GovernanceModel(WrapperModel): + """A ``WrapperModel`` that brackets every model call with governance.""" + + def __init__(self, wrapped: Model, callbacks: "GovernanceCallbacks") -> None: + super().__init__(wrapped) + self._callbacks = callbacks + + async def request( + self, + messages: List[Any], + model_settings: ModelSettings | None, + model_request_parameters: ModelRequestParameters, + ) -> Any: + self._callbacks.on_request(messages) + response = await super().request( + messages, model_settings, model_request_parameters + ) + self._callbacks.on_response(response) + return response + + @asynccontextmanager + async def request_stream( + self, + messages: List[Any], + model_settings: ModelSettings | None, + model_request_parameters: ModelRequestParameters, + run_context: Any = None, + ) -> AsyncIterator[StreamedResponse]: + self._callbacks.on_request(messages) + async with super().request_stream( + messages, model_settings, model_request_parameters, run_context + ) as stream: + yield stream + # After the caller has consumed the stream, the final response is + # assembled — govern it the same as the non-streaming path. A DENY + # decision must still abort the run, so the block exception propagates; + # any other governance error is logged and swallowed. + try: + self._callbacks.on_response(stream.get()) + except GovernanceBlockException: + raise + except Exception as e: # noqa: BLE001 - a governance bug must not break the run + logger.warning("after-stream governance check failed (continuing): %s", e) + + +class GovernanceCallbacks: + """Holds the evaluator + per-attach state, called by :class:`GovernanceModel`. + + :class:`GovernanceBlockException` is re-raised (it aborts the run); + anything else is logged and swallowed so a governance bug never breaks an + agent run. + """ + + def __init__( + self, + evaluator: EvaluatorProtocol, + agent_name: str, + session_id: str, + ) -> None: + self._evaluator = evaluator + self._agent_name = agent_name + self._session_id = session_id + self._trace_id = str(uuid4()) + self._session_state: Dict[str, Any] = {"tool_calls": 0, "llm_calls": 0} + + # ----- before the model call -------------------------------------- + + def on_request(self, messages: Any) -> None: + """Fire BEFORE_MODEL (latest message text) + AFTER_TOOL (tool returns). + + Only the latest request message is scanned, so a tool result / prompt + is not re-evaluated on every subsequent model call (the full history is + re-sent each turn for context). + """ + latest = self._latest_request(messages) + if latest is None: + self._before_model("") + return + parts = getattr(latest, "parts", None) or [] + self._before_model(self._parts_input_text(parts)) + for part in parts: + if _part_kind(part) == "tool-return": + self._after_tool( + getattr(part, "tool_name", None) or "unknown", + getattr(part, "content", None), + ) + + # ----- after the model call --------------------------------------- + + def on_response(self, response: Any) -> None: + """Fire AFTER_MODEL (response text) + TOOL_CALL (each tool-call part).""" + parts = getattr(response, "parts", None) or [] + self._after_model(self._response_text(parts)) + for part in parts: + if _part_kind(part) in ("tool-call", "builtin-tool-call"): + self._tool_call( + getattr(part, "tool_name", None) or "unknown", + getattr(part, "args", None), + ) + + # ----- individual evaluate_* wrappers (block-propagate, else swallow) -- + + def _before_model(self, text: str) -> None: + try: + self._session_state["llm_calls"] = ( + self._session_state.get("llm_calls", 0) + 1 + ) + self._evaluator.evaluate_before_model( + model_input=text, + agent_name=self._agent_name, + runtime_id=self._session_id, + trace_id=self._trace_id, + ) + except GovernanceBlockException: + raise + except Exception as e: # noqa: BLE001 + logger.warning("before_model governance check failed (continuing): %s", e) + + def _after_model(self, text: str) -> None: + try: + self._evaluator.evaluate_after_model( + model_output=text, + agent_name=self._agent_name, + runtime_id=self._session_id, + trace_id=self._trace_id, + ) + except GovernanceBlockException: + raise + except Exception as e: # noqa: BLE001 + logger.warning("after_model governance check failed (continuing): %s", e) + + def _tool_call(self, tool_name: str, args: Any) -> None: + try: + self._session_state["tool_calls"] = ( + self._session_state.get("tool_calls", 0) + 1 + ) + self._evaluator.evaluate_tool_call( + tool_name=tool_name, + tool_args=_coerce_args(args), + agent_name=self._agent_name, + runtime_id=self._session_id, + trace_id=self._trace_id, + session_state=self._session_state, + ) + except GovernanceBlockException: + raise + except Exception as e: # noqa: BLE001 + logger.warning("tool_call governance check failed (continuing): %s", e) + + def _after_tool(self, tool_name: str, content: Any) -> None: + try: + self._evaluator.evaluate_after_tool( + tool_name=tool_name, + tool_result="" if content is None else _stringify(content), + agent_name=self._agent_name, + runtime_id=self._session_id, + trace_id=self._trace_id, + ) + except GovernanceBlockException: + raise + except Exception as e: # noqa: BLE001 + logger.warning("after_tool governance check failed (continuing): %s", e) + + # ----- text extraction -------------------------------------------- + + @staticmethod + def _latest_request(messages: Any) -> Any: + """Return the most recent message (a ``ModelRequest``) or ``None``.""" + if not messages or not isinstance(messages, (list, tuple)): + return None + return messages[-1] + + @classmethod + def _parts_input_text(cls, parts: Any) -> str: + """Join governance-relevant input text from a request message's parts. + + Covers user prompts and tool-return content (the model's input on a + follow-up turn). Capped at :data:`_BEFORE_MODEL_TEXT_CAP`. + """ + collected: List[str] = [] + for part in parts: + kind = _part_kind(part) + if kind == "user-prompt": + collected.append(_content_text(getattr(part, "content", None))) + elif kind == "tool-return": + collected.append(_stringify(getattr(part, "content", None))) + return "\n".join(p for p in collected if p)[:_BEFORE_MODEL_TEXT_CAP] + + @classmethod + def _response_text(cls, parts: Any) -> str: + """Join ``TextPart`` content from a model response's parts.""" + collected: List[str] = [] + for part in parts: + if _part_kind(part) == "text": + text = getattr(part, "content", None) + if isinstance(text, str) and text: + collected.append(text) + return "\n".join(collected)[:_BEFORE_MODEL_TEXT_CAP] + + +# -------------------------------------------------------------------------- +# Helpers +# -------------------------------------------------------------------------- + + +def _part_kind(part: Any) -> str: + """Return a message part's discriminator (``part_kind``), or ``""``.""" + kind = getattr(part, "part_kind", None) + return kind if isinstance(kind, str) else "" + + +def _content_text(content: Any) -> str: + """Render a ``UserPromptPart.content`` (str or list of items) as text.""" + if content is None: + return "" + if isinstance(content, str): + return content + if isinstance(content, (list, tuple)): + out: List[str] = [] + for item in content: + if isinstance(item, str): + out.append(item) + else: + text = getattr(item, "text", None) + if isinstance(text, str): + out.append(text) + return "\n".join(out) + return _stringify(content) + + +def _coerce_args(args: Any) -> Dict[str, Any]: + """Normalise ``ToolCallPart.args`` (dict / JSON string / None) to a dict.""" + if args is None: + return {} + if isinstance(args, dict): + return args + if isinstance(args, str): + try: + parsed = json.loads(args) + return parsed if isinstance(parsed, dict) else {"_": parsed} + except (TypeError, ValueError): + return {} + return {} + + +def _stringify(value: Any) -> str: + """Render a dict / object payload as compact, scannable text.""" + if isinstance(value, str): + return value + try: + return json.dumps(value, default=str, ensure_ascii=False) + except (TypeError, ValueError): + return str(value) \ No newline at end of file diff --git a/packages/uipath-pydantic-ai/tests/governance/__init__.py b/packages/uipath-pydantic-ai/tests/governance/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/packages/uipath-pydantic-ai/tests/governance/test_adapter.py b/packages/uipath-pydantic-ai/tests/governance/test_adapter.py new file mode 100644 index 00000000..3495fb05 --- /dev/null +++ b/packages/uipath-pydantic-ai/tests/governance/test_adapter.py @@ -0,0 +1,302 @@ +"""Unit tests for the Pydantic AI governance adapter. + +These tests use real ``pydantic_ai`` message parts (``UserPromptPart`` etc.) +so the part-extraction logic is exercised against the actual types, plus the +adapter's model-wrapping attach/detach against a real ``Agent`` (driven by the +offline ``TestModel``). + +The package is configured with ``asyncio_mode = "auto"``, so ``async def`` +tests run without an explicit marker. +""" + +from __future__ import annotations + +import logging +from typing import Any, List + +import pytest +from pydantic_ai import Agent +from pydantic_ai.messages import ( + ModelRequest, + ModelResponse, + TextPart, + ToolCallPart, + ToolReturnPart, + UserPromptPart, +) +from pydantic_ai.models.test import TestModel +from uipath.core.governance.exceptions import GovernanceBlockException + +from uipath_pydantic_ai.governance.adapter import ( + _BEFORE_MODEL_TEXT_CAP, + GovernanceCallbacks, + GovernanceModel, + PydanticAIAdapter, + _coerce_args, +) + +# -------------------------------------------------------------------------- +# Fakes +# -------------------------------------------------------------------------- + + +class FakeEvaluator: + """Records evaluate_* calls; optionally BLOCKs on a named hook.""" + + def __init__(self, block_on: str | None = None) -> None: + self.block_on = block_on + self.calls: List[tuple[str, dict]] = [] + + def _record(self, hook: str, **kwargs: Any) -> None: + self.calls.append((hook, kwargs)) + if self.block_on == hook: + raise GovernanceBlockException("blocked") # type: ignore[call-arg] + + def evaluate_before_agent(self, **kwargs: Any) -> None: + self._record("before_agent", **kwargs) + + def evaluate_after_agent(self, **kwargs: Any) -> None: + self._record("after_agent", **kwargs) + + def evaluate_before_model(self, **kwargs: Any) -> None: + self._record("before_model", **kwargs) + + def evaluate_after_model(self, **kwargs: Any) -> None: + self._record("after_model", **kwargs) + + def evaluate_tool_call(self, **kwargs: Any) -> None: + self._record("tool_call", **kwargs) + + def evaluate_after_tool(self, **kwargs: Any) -> None: + self._record("after_tool", **kwargs) + + +def _make_callbacks(ev: FakeEvaluator) -> GovernanceCallbacks: + return GovernanceCallbacks(evaluator=ev, agent_name="agent-1", session_id="sess-1") + + +def _hooks(ev: FakeEvaluator) -> List[str]: + return [h for h, _ in ev.calls] + + +# -------------------------------------------------------------------------- +# can_handle +# -------------------------------------------------------------------------- + + +def test_can_handle_agent(): + assert PydanticAIAdapter().can_handle(Agent(model=TestModel())) is True + + +def test_can_handle_rejects_non_agent(): + from types import SimpleNamespace + + # A duck-typed look-alike (model/run/iter) must NOT be claimed — only a real Agent. + look_alike = SimpleNamespace(model=object(), run=lambda: None, iter=lambda: None) + assert PydanticAIAdapter().can_handle(look_alike) is False + assert PydanticAIAdapter().can_handle(object()) is False + + +# -------------------------------------------------------------------------- +# attach / detach +# -------------------------------------------------------------------------- + + +def test_attach_wraps_model_and_detach_restores(): + agent = Agent(model=TestModel()) + original = agent.model + adapter = PydanticAIAdapter() + returned = adapter.attach(agent, agent_id="x", session_id="s", evaluator=FakeEvaluator()) + assert returned is agent + assert isinstance(agent.model, GovernanceModel) + adapter.detach(agent) + assert agent.model is original + + +def test_attach_is_idempotent(): + agent = Agent(model=TestModel()) + adapter = PydanticAIAdapter() + ev = FakeEvaluator() + adapter.attach(agent, agent_id="x", session_id="s", evaluator=ev) + wrapped = agent.model + adapter.attach(agent, agent_id="x", session_id="s", evaluator=ev) + assert agent.model is wrapped # not double-wrapped + + +def test_attach_warns_when_no_bound_model(caplog): + agent = Agent() # no model bound + with caplog.at_level(logging.WARNING): + PydanticAIAdapter().attach(agent, agent_id="x", session_id="s", evaluator=FakeEvaluator()) + assert any("no bound Model" in r.message for r in caplog.records) + + +# -------------------------------------------------------------------------- +# on_request → BEFORE_MODEL + AFTER_TOOL +# -------------------------------------------------------------------------- + + +def test_on_request_fires_before_model_with_latest_user_prompt(): + ev = FakeEvaluator() + cb = _make_callbacks(ev) + messages = [ + ModelRequest(parts=[UserPromptPart(content="old turn")]), + ModelRequest(parts=[UserPromptPart(content="the question")]), + ] + cb.on_request(messages) + assert _hooks(ev) == ["before_model"] + assert ev.calls[0][1]["model_input"] == "the question" + + +def test_on_request_fires_after_tool_for_tool_return(): + ev = FakeEvaluator() + cb = _make_callbacks(ev) + messages = [ + ModelRequest( + parts=[ToolReturnPart(tool_name="lookup", content={"balance": "1000"}, tool_call_id="c1")] + ) + ] + cb.on_request(messages) + # both BEFORE_MODEL (tool result is the model's new input) and AFTER_TOOL fire + assert "before_model" in _hooks(ev) + after_tool = [kw for h, kw in ev.calls if h == "after_tool"] + assert after_tool and after_tool[0]["tool_name"] == "lookup" + assert "1000" in after_tool[0]["tool_result"] + + +def test_on_request_caps_text(): + ev = FakeEvaluator() + cb = _make_callbacks(ev) + huge = "x" * (_BEFORE_MODEL_TEXT_CAP + 5000) + cb.on_request([ModelRequest(parts=[UserPromptPart(content=huge)])]) + assert len(ev.calls[0][1]["model_input"]) <= _BEFORE_MODEL_TEXT_CAP + + +def test_on_request_empty(): + ev = FakeEvaluator() + cb = _make_callbacks(ev) + cb.on_request([]) + assert ev.calls[0][1]["model_input"] == "" + + +# -------------------------------------------------------------------------- +# on_response → AFTER_MODEL + TOOL_CALL +# -------------------------------------------------------------------------- + + +def test_on_response_fires_after_model_and_tool_call(): + ev = FakeEvaluator() + cb = _make_callbacks(ev) + response = ModelResponse( + parts=[ + TextPart(content="thinking out loud"), + ToolCallPart(tool_name="transfer", args={"amount": 50}, tool_call_id="c1"), + ] + ) + cb.on_response(response) + assert "after_model" in _hooks(ev) and "tool_call" in _hooks(ev) + after_model = [kw for h, kw in ev.calls if h == "after_model"][0] + assert after_model["model_output"] == "thinking out loud" + tool_call = [kw for h, kw in ev.calls if h == "tool_call"][0] + assert tool_call["tool_name"] == "transfer" + assert tool_call["tool_args"] == {"amount": 50} + assert tool_call["session_state"]["tool_calls"] == 1 + + +def test_on_response_coerces_json_string_args(): + ev = FakeEvaluator() + cb = _make_callbacks(ev) + response = ModelResponse( + parts=[ToolCallPart(tool_name="t", args='{"x": 1}', tool_call_id="c1")] + ) + cb.on_response(response) + tool_call = [kw for h, kw in ev.calls if h == "tool_call"][0] + assert tool_call["tool_args"] == {"x": 1} + + +# -------------------------------------------------------------------------- +# GovernanceModel.request brackets a wrapped model +# -------------------------------------------------------------------------- + + +async def test_governance_model_request_brackets_call(): + ev = FakeEvaluator() + cb = _make_callbacks(ev) + order: List[str] = [] + + class FakeWrapped: + async def request(self, messages, settings, params): + order.append("MODEL_CALL") + return ModelResponse(parts=[TextPart(content="Your balance is 1000.")]) + + gm = GovernanceModel.__new__(GovernanceModel) # bypass WrapperModel init + gm.wrapped = FakeWrapped() # type: ignore[attr-defined] + gm._callbacks = cb + messages = [ModelRequest(parts=[UserPromptPart(content="What is my balance?")])] + await gm.request(messages, None, None) + + assert order == ["MODEL_CALL"] + assert _hooks(ev) == ["before_model", "after_model"] + assert ev.calls[0][1]["model_input"] == "What is my balance?" + assert ev.calls[1][1]["model_output"] == "Your balance is 1000." + + +async def test_governance_model_request_stream_block_propagates(): + # A DENY during the after-stream check must abort the run, exactly like the + # non-streaming request() path — it must not be swallowed by the catch-all. + from contextlib import asynccontextmanager + from types import SimpleNamespace + + cb = _make_callbacks(FakeEvaluator(block_on="tool_call")) + denied = ModelResponse( + parts=[ToolCallPart(tool_name="t", args={}, tool_call_id="c1")] + ) + + class FakeWrapped: + @asynccontextmanager + async def request_stream(self, *_a, **_k): + yield SimpleNamespace(get=lambda: denied) + + gm = GovernanceModel.__new__(GovernanceModel) # bypass WrapperModel init + gm.wrapped = FakeWrapped() # type: ignore[attr-defined] + gm._callbacks = cb + messages = [ModelRequest(parts=[UserPromptPart(content="hi")])] + with pytest.raises(GovernanceBlockException): + async with gm.request_stream(messages, None, None) as stream: + assert stream is not None + + +# -------------------------------------------------------------------------- +# helpers + enforcement +# -------------------------------------------------------------------------- + + +def test_coerce_args_variants(): + assert _coerce_args({"a": 1}) == {"a": 1} + assert _coerce_args('{"a": 1}') == {"a": 1} + assert _coerce_args(None) == {} + assert _coerce_args("not json") == {} + + +def test_block_in_before_model_propagates(): + cb = _make_callbacks(FakeEvaluator(block_on="before_model")) + with pytest.raises(GovernanceBlockException): + cb.on_request([ModelRequest(parts=[UserPromptPart(content="hi")])]) + + +def test_block_in_tool_call_propagates(): + cb = _make_callbacks(FakeEvaluator(block_on="tool_call")) + with pytest.raises(GovernanceBlockException): + cb.on_response( + ModelResponse(parts=[ToolCallPart(tool_name="t", args={}, tool_call_id="c1")]) + ) + + +def test_non_block_exception_is_swallowed(caplog): + class Boom: + def evaluate_before_model(self, **_: Any) -> None: + raise RuntimeError("evaluator bug") + + cb = GovernanceCallbacks(evaluator=Boom(), agent_name="a", session_id="s") # type: ignore[arg-type] + with caplog.at_level(logging.WARNING): + cb.on_request([ModelRequest(parts=[UserPromptPart(content="x")])]) + assert any("governance check failed" in r.message for r in caplog.records) \ No newline at end of file diff --git a/packages/uipath-pydantic-ai/uv.lock b/packages/uipath-pydantic-ai/uv.lock index 64691c63..704e7e0c 100644 --- a/packages/uipath-pydantic-ai/uv.lock +++ b/packages/uipath-pydantic-ai/uv.lock @@ -3691,6 +3691,7 @@ dependencies = [ { name = "openinference-instrumentation-pydantic-ai" }, { name = "pydantic-ai" }, { name = "uipath" }, + { name = "uipath-core" }, { name = "uipath-runtime" }, ] @@ -3710,6 +3711,7 @@ requires-dist = [ { name = "openinference-instrumentation-pydantic-ai", specifier = ">=0.1.12" }, { name = "pydantic-ai", specifier = ">=1.63.0,<2.0.0" }, { name = "uipath", specifier = ">=2.10.2,<2.11.0" }, + { name = "uipath-core", specifier = ">=0.5.18,<0.7.0" }, { name = "uipath-runtime", specifier = ">=0.11.0,<0.12.0" }, ]