From bf034da624de60c48664e11f37f6f9599a4d1246 Mon Sep 17 00:00:00 2001 From: Aditi Kumari Date: Mon, 22 Jun 2026 23:09:07 +0530 Subject: [PATCH 01/11] feat(governance): add Pydantic AI governance adapter Wraps agent.model with a WrapperModel deriving all four hooks from message parts (UserPromptPart -> BEFORE_MODEL, TextPart -> AFTER_MODEL, ToolCallPart -> TOOL_CALL, ToolReturnPart -> AFTER_TOOL); covers request and request_stream. Self-registers via the uipath.governance.adapters entry point; unit-tested and verified firing through the framework's real execution path. BEFORE/AFTER_AGENT remain owned by the uipath-runtime wrapper. Co-Authored-By: Claude Opus 4.8 --- packages/uipath-pydantic-ai/pyproject.toml | 4 + .../uipath_pydantic_ai/governance/__init__.py | 58 +++ .../uipath_pydantic_ai/governance/adapter.py | 378 ++++++++++++++++++ .../tests/governance/__init__.py | 0 .../tests/governance/test_adapter.py | 272 +++++++++++++ packages/uipath-pydantic-ai/uv.lock | 2 + 6 files changed, 714 insertions(+) create mode 100644 packages/uipath-pydantic-ai/src/uipath_pydantic_ai/governance/__init__.py create mode 100644 packages/uipath-pydantic-ai/src/uipath_pydantic_ai/governance/adapter.py create mode 100644 packages/uipath-pydantic-ai/tests/governance/__init__.py create mode 100644 packages/uipath-pydantic-ai/tests/governance/test_adapter.py 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..e60ea317 --- /dev/null +++ b/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/governance/__init__.py @@ -0,0 +1,58 @@ +"""Governance integration for ``uipath-pydantic-ai``. + +Registers :class:`PydanticAIAdapter` with the global adapter registry in +``uipath.core.adapters`` so ``uipath.runtime.governance.GovernanceRuntime`` 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: + 1. Importing this module triggers registration as a side-effect, so any + caller that does ``import uipath_pydantic_ai.governance`` is opted in. + 2. The package also exposes :func:`register_governance_adapter` as an entry + point under ``uipath.governance.adapters`` so the registry's entry-point + discovery can plug us in without an explicit import. +""" + +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") + + +# Side-effect registration on module import. +register_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..e39a80fe --- /dev/null +++ b/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/governance/adapter.py @@ -0,0 +1,378 @@ +"""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 runtime +wrapper layer in ``uipath-runtime`` 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 self-registers it +with the global adapter registry when ``uipath_pydantic_ai.governance`` is +imported. + +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 if this adapter knows how to hook into the agent.""" + try: + from pydantic_ai import Agent + + if isinstance(agent, Agent): + return True + except ImportError: + pass + + # Duck-typed fallback: a Pydantic AI agent exposes a model slot plus the + # run / iter execution surface. + if hasattr(agent, "model") and hasattr(agent, "run") and hasattr(agent, "iter"): + return True + return False + + 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. + try: + self._callbacks.on_response(stream.get()) + except Exception as e: # noqa: BLE001 - never break on the after-stream check + 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..caa91b6f --- /dev/null +++ b/packages/uipath-pydantic-ai/tests/governance/test_adapter.py @@ -0,0 +1,272 @@ +"""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_plain_object(): + 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." + + +# -------------------------------------------------------------------------- +# 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" }, ] From 2d5d974344dd8abe55a79e4ceb1d5c3fb5be133c Mon Sep 17 00:00:00 2001 From: Aditi Kumari Date: Wed, 24 Jun 2026 17:22:34 +0530 Subject: [PATCH 02/11] chore(governance): apply review feedback (no import-time registration, framework-only can_handle) Mirror radu's LangChain-adapter review across the Pydantic AI adapter: - __init__: drop the import-time registration side-effect; registration only via the uipath.governance.adapters entry point. - can_handle: claim only a real pydantic_ai.Agent; remove the duck-typed (model/run/iter) fallback. - docstring: 'governance host' instead of uipath-runtime internals. - tests: a duck-typed look-alike is now rejected. Co-Authored-By: Claude Opus 4.8 --- .../uipath_pydantic_ai/governance/__init__.py | 22 +++++++------------ .../uipath_pydantic_ai/governance/adapter.py | 18 +++++---------- .../tests/governance/test_adapter.py | 7 +++++- 3 files changed, 19 insertions(+), 28 deletions(-) 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 index e60ea317..0efd4f7e 100644 --- a/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/governance/__init__.py +++ b/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/governance/__init__.py @@ -1,19 +1,17 @@ """Governance integration for ``uipath-pydantic-ai``. -Registers :class:`PydanticAIAdapter` with the global adapter registry in -``uipath.core.adapters`` so ``uipath.runtime.governance.GovernanceRuntime`` can -attach the Pydantic-AI-specific governance (BEFORE_MODEL, AFTER_MODEL, -TOOL_CALL, AFTER_TOOL) when it sees a ``pydantic_ai.Agent``. +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: - 1. Importing this module triggers registration as a side-effect, so any - caller that does ``import uipath_pydantic_ai.governance`` is opted in. - 2. The package also exposes :func:`register_governance_adapter` as an entry - point under ``uipath.governance.adapters`` so the registry's entry-point - discovery can plug us in without an explicit import. +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 @@ -46,10 +44,6 @@ def register_governance_adapter() -> None: logger.debug("Registered uipath-pydantic-ai governance adapter") -# Side-effect registration on module import. -register_governance_adapter() - - __all__ = [ "GovernanceCallbacks", "GovernanceModel", 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 index e39a80fe..de1fea64 100644 --- a/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/governance/adapter.py +++ b/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/governance/adapter.py @@ -20,8 +20,8 @@ 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 runtime -wrapper layer in ``uipath-runtime`` and are intentionally not fired here. +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 self-registers it @@ -70,20 +70,12 @@ def name(self) -> str: return "PydanticAI" def can_handle(self, agent: Any) -> bool: - """Return True if this adapter knows how to hook into the agent.""" + """Return True only for a ``pydantic_ai.Agent``.""" try: from pydantic_ai import Agent - - if isinstance(agent, Agent): - return True except ImportError: - pass - - # Duck-typed fallback: a Pydantic AI agent exposes a model slot plus the - # run / iter execution surface. - if hasattr(agent, "model") and hasattr(agent, "run") and hasattr(agent, "iter"): - return True - return False + return False + return isinstance(agent, Agent) def attach( self, diff --git a/packages/uipath-pydantic-ai/tests/governance/test_adapter.py b/packages/uipath-pydantic-ai/tests/governance/test_adapter.py index caa91b6f..50d720b9 100644 --- a/packages/uipath-pydantic-ai/tests/governance/test_adapter.py +++ b/packages/uipath-pydantic-ai/tests/governance/test_adapter.py @@ -88,7 +88,12 @@ def test_can_handle_agent(): assert PydanticAIAdapter().can_handle(Agent(model=TestModel())) is True -def test_can_handle_rejects_plain_object(): +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 From 7c98ddc18d67fac23697771a41ddada739b0932b Mon Sep 17 00:00:00 2001 From: Aditi Kumari Date: Wed, 24 Jun 2026 18:26:52 +0530 Subject: [PATCH 03/11] fix(governance): propagate GovernanceBlockException from the streaming path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address Copilot review on the Pydantic AI adapter: - request_stream's after-stream check caught all exceptions, swallowing GovernanceBlockException — a DENY during streaming did not abort the run, unlike the non-streaming request() path. Re-raise the block exception; keep swallowing other governance errors. Add a streaming block-propagation test. - Module docstring: registers via the uipath.governance.adapters entry point, not at import time. Co-Authored-By: Claude Opus 4.8 --- .../uipath_pydantic_ai/governance/adapter.py | 13 ++++++---- .../tests/governance/test_adapter.py | 25 +++++++++++++++++++ 2 files changed, 33 insertions(+), 5 deletions(-) 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 index de1fea64..be31c869 100644 --- a/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/governance/adapter.py +++ b/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/governance/adapter.py @@ -24,9 +24,8 @@ 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 self-registers it -with the global adapter registry when ``uipath_pydantic_ai.governance`` is -imported. +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 @@ -153,10 +152,14 @@ async def request_stream( ) 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. + # 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 Exception as e: # noqa: BLE001 - never break on the after-stream check + 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) diff --git a/packages/uipath-pydantic-ai/tests/governance/test_adapter.py b/packages/uipath-pydantic-ai/tests/governance/test_adapter.py index 50d720b9..3495fb05 100644 --- a/packages/uipath-pydantic-ai/tests/governance/test_adapter.py +++ b/packages/uipath-pydantic-ai/tests/governance/test_adapter.py @@ -240,6 +240,31 @@ async def request(self, messages, settings, params): 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 # -------------------------------------------------------------------------- From 6ad2bad80b88a0e2ff8e1e8d3ecce95c59e1396a Mon Sep 17 00:00:00 2001 From: Aditi Kumari Date: Wed, 24 Jun 2026 18:40:49 +0530 Subject: [PATCH 04/11] chore(governance): remove unrelated files bundled from a dirty tree These files were swept into the branch by a broad add; they are unrelated to the governance adapter. Reverting/removing them so the PR contains only governance changes. Co-Authored-By: Claude Opus 4.8 --- SETUP.MD | 141 ------------------ .../docs/llms_and_embeddings.md | 56 +++---- 2 files changed, 19 insertions(+), 178 deletions(-) delete mode 100644 SETUP.MD diff --git a/SETUP.MD b/SETUP.MD deleted file mode 100644 index d7b750a8..00000000 --- a/SETUP.MD +++ /dev/null @@ -1,141 +0,0 @@ -# SETUP.MD - -This file documents how to provision a clean development environment for the five packages in this repo (`uipath-agent-framework`, `uipath-google-adk`, `uipath-llamaindex`, `uipath-openai-agents`, `uipath-pydantic-ai`), run the build, execute the tests, and validate a sample code change end-to-end. It is intended both as a quick reference for human contributors and as a structured guide for automated environment-setup tooling. - -## Prerequisites - -- Python 3.11+ -- [uv](https://docs.astral.sh/uv/) 0.5+ - -### Supported platforms - -`uv` is shell- and OS-agnostic, so the commands below run unchanged on every supported platform: - -- [x] Linux -- [x] Windows -- [x] macOS - -## Environment Variables - -None required for environment setup, build, or unit tests. The suites under the `Test` section run fully offline and require no external authentication. - -> **All commands below must be run from the repository root.** The `uv --directory packages/` invocations resolve each subpackage relative to the current working directory. The first line of `## Setup` enforces this by `cd`-ing to the git root. - -## Setup - -```bash -cd "$(git rev-parse --show-toplevel)" -python3 -m pip install --upgrade uv - -# Sync all five packages (each is independent) -uv --directory packages/uipath-agent-framework sync --all-extras -uv --directory packages/uipath-google-adk sync --all-extras -uv --directory packages/uipath-llamaindex sync --all-extras -uv --directory packages/uipath-openai-agents sync --all-extras -uv --directory packages/uipath-pydantic-ai sync --all-extras -``` - -## Verify Setup - -```bash -uv --version -uv --directory packages/uipath-pydantic-ai run python --version -uv --directory packages/uipath-agent-framework run python -c "import uipath_agent_framework; print('uipath-agent-framework ok')" -uv --directory packages/uipath-google-adk run python -c "import uipath_google_adk; print('uipath-google-adk ok')" -uv --directory packages/uipath-llamaindex run python -c "import uipath_llamaindex; print('uipath-llamaindex ok')" -uv --directory packages/uipath-openai-agents run python -c "import uipath_openai_agents; print('uipath-openai-agents ok')" -uv --directory packages/uipath-pydantic-ai run python -c "import uipath_pydantic_ai; print('uipath-pydantic-ai ok')" -``` - -## Build - -N/A - -## Test - -```bash -uv --directory packages/uipath-agent-framework run pytest -uv --directory packages/uipath-google-adk run pytest -uv --directory packages/uipath-llamaindex run pytest -uv --directory packages/uipath-openai-agents run pytest -uv --directory packages/uipath-pydantic-ai run pytest -``` - -## Sample Code Change - -### The change - -Add a new `agent_count` property to `PydanticAiConfig` in `packages/uipath-pydantic-ai/src/uipath_pydantic_ai/runtime/config.py`, immediately after the existing `entrypoint` property: - -```python -@property -def agent_count(self) -> int: - """Number of agents defined in the configuration.""" - return len(self.agents) -``` - -Then create `packages/uipath-pydantic-ai/tests/test_config_agent_count.py` with two pytest tests: - -```python -"""Tests for PydanticAiConfig.agent_count.""" - -import json -from pathlib import Path - -from uipath_pydantic_ai.runtime.config import PydanticAiConfig - - -def test_agent_count_single(tmp_path: Path) -> None: - config_path = tmp_path / "pydantic_ai.json" - config_path.write_text(json.dumps({"agents": {"main": "main:agent"}})) - cfg = PydanticAiConfig(str(config_path)) - assert cfg.agent_count == 1 - - -def test_agent_count_multiple(tmp_path: Path) -> None: - config_path = tmp_path / "pydantic_ai.json" - config_path.write_text( - json.dumps( - { - "agents": { - "alpha": "alpha:agent", - "beta": "beta:agent", - "gamma": "gamma:agent", - } - } - ) - ) - cfg = PydanticAiConfig(str(config_path)) - assert cfg.agent_count == 3 -``` - -### Verification - -```bash -uv --directory packages/uipath-pydantic-ai run pytest tests/test_config_agent_count.py -v -``` - -## Test with a real UiPath Coded Agent - -The unit tests above are necessary but not sufficient — they don't exercise the package end-to-end through a real agent. The flow below validates changes against a live runtime: - -1. Apply the code changes locally. -2. Run the unit tests (see the `Sample Code Change` section above). -3. Scaffold a coded UiPath agent (PydanticAI / OpenAI / Google ADK / LlamaIndex / Agent Framework, matching the package you changed) that exercises the changed code path. -4. In the downstream project's `pyproject.toml`, add this local library as an editable dependency (substitute the package you changed): - - ```toml - [tool.uv.sources] - uipath-pydantic-ai = { path = "../path/to/uipath-integrations-python/packages/uipath-pydantic-ai", editable = true } - ``` - -5. Exercise the new behavior end-to-end: - - ```bash - uv run uipath run --input '{...}' - ``` - -6. (Optional) Open a PR and apply the `build:dev` label — this publishes the development version to Test PyPI. -7. The PR description is updated automatically with instructions for pointing the downstream agent at the Test PyPI dev version. -8. Push the dev version to UiPath with [`uipath push`](https://uipath.github.io/uipath-python/cli/#push), then deploy it to Orchestrator or Studio Web with [`uipath deploy`](https://uipath.github.io/uipath-python/cli/#deploy), and run it in cloud to confirm the changes behave correctly against the real platform. -9. Once validation is done, close the dev PR — these PRs are not meant to be merged; their only purpose was to publish a Test PyPI build for end-to-end validation. diff --git a/packages/uipath-llamaindex/docs/llms_and_embeddings.md b/packages/uipath-llamaindex/docs/llms_and_embeddings.md index 01416be8..15e25eb4 100644 --- a/packages/uipath-llamaindex/docs/llms_and_embeddings.md +++ b/packages/uipath-llamaindex/docs/llms_and_embeddings.md @@ -1,39 +1,7 @@ # LLMs and Embeddings -UiPath provides pre-configured LLM and embedding classes for several providers (OpenAI via `UiPathOpenAI`, Anthropic on AWS Bedrock via `UiPathChatBedrockConverse`, Google Vertex AI via `UiPathVertex`, and more), plus embeddings via `UiPathOpenAIEmbedding`. These handle authentication, routing, and configuration automatically, allowing you to focus on building your agents. You do not need to add API keys from OpenAI, AWS, or Google, usage of these models will consume `Agent Units` on your account. - -## Available models - -LLM models are served through the UiPath LLM Gateway and are subject to [AI Trust Layer](https://docs.uipath.com/automation-cloud/automation-cloud/latest/admin-guide/about-ai-trust-layer) policies, so the exact set of models available to you depends on your tenant configuration. List the models you can use with the `uipath` CLI: - -```console -$ uipath list-models - Available LLM Models -┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━┓ -┃ AwsBedrock ┃ OpenAi ┃ VertexAi ┃ -┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━┩ -│ anthropic.claude-haiku-4-5-20251001-v1:0 │ gpt-4.1-2025-04-14 │ gemini-2.5-flash │ -│ anthropic.claude-opus-4-7 │ gpt-4.1-mini-2025-04-14 │ gemini-2.5-pro │ -│ ... │ ... │ ... │ -└──────────────────────────────────────────┴─────────────────────────┴──────────────────┘ -``` - -Pick a model id from the relevant provider column and pass it (or the matching enum member) to the matching class: - -```python -from uipath_llamaindex.llms import UiPathOpenAI -from uipath_llamaindex.llms.bedrock import UiPathChatBedrockConverse -from uipath_llamaindex.llms.vertex import UiPathVertex - -# OpenAI models -llm = UiPathOpenAI(model="gpt-4.1-mini-2025-04-14") - -# AWS Bedrock (Anthropic) models -llm = UiPathChatBedrockConverse(model="anthropic.claude-sonnet-4-5-20250929-v1:0") - -# Google Vertex AI (Gemini) models -llm = UiPathVertex(model="gemini-2.5-flash") -``` +UiPath provides pre-configured LLM and embedding classes that handle authentication, routing, and configuration automatically, allowing you to focus on building your agents. +You do not need to add API keys from OpenAI, AWS, or Google, usage of these models will consume `Agent Units` on your account. ## UiPathOpenAI @@ -41,7 +9,17 @@ The `UiPathOpenAI` class is a pre-configured Azure OpenAI client that routes req ### Available Models -The OpenAI models from the `OpenAi` column of [`uipath list-models`](#available-models) can be used here, either as a model string or via the `OpenAIModel` enum. +The following OpenAI models are available through the `OpenAIModel` enum: + +- `GPT_4_1_2025_04_14` +- `GPT_4_1_MINI_2025_04_14` +- `GPT_4_1_NANO_2025_04_14` +- `GPT_4O_2024_05_13` +- `GPT_4O_2024_08_06` +- `GPT_4O_2024_11_20` +- `GPT_4O_MINI_2024_07_18` (default) +- `O3_MINI_2025_01_31` +- `TEXT_DAVINCI_003` ### Basic Usage @@ -165,7 +143,9 @@ from uipath_llamaindex.llms import BedrockModel llm = UiPathChatBedrock(model=BedrockModel.anthropic_claude_sonnet_4) ``` -The available models are the ones in the `AwsBedrock` column of [`uipath list-models`](#available-models). +Currently, the following models can be used (this list can be updated in the future): + +- `anthropic.claude-3-7-sonnet-20250219-v1:0`, `anthropic.claude-sonnet-4-20250514-v1:0`, `anthropic.claude-sonnet-4-5-20250929-v1:0`, `anthropic.claude-haiku-4-5-20251001-v1:0` ## UiPathVertex @@ -204,7 +184,9 @@ response = llm.chat(messages) print(response) ``` -The available models are the ones in the `VertexAi` column of [`uipath list-models`](#available-models). +Currently, the following models can be used (this list can be updated in the future): + +- `gemini-2.0-flash-001`, `gemini-2.5-flash`, `gemini-2.5-pro` ## Integration with LlamaIndex From a814bb2f499bf0cc7f9bf095dbdf1554a83f3d3c Mon Sep 17 00:00:00 2001 From: Aditi Kumari Date: Wed, 24 Jun 2026 19:01:04 +0530 Subject: [PATCH 05/11] chore(governance): restore SETUP.MD and llms doc that belong to main An earlier cleanup commit compared against a stale local main and wrongly removed SETUP.MD and reverted the LlamaIndex docs change. Both files come from main (PRs #352/#356), not this branch. Restore them to the main version so this PR is governance-only with no spurious deletions. Co-Authored-By: Claude Opus 4.8 --- SETUP.MD | 141 ++++++++++++++++++ .../docs/llms_and_embeddings.md | 56 ++++--- 2 files changed, 178 insertions(+), 19 deletions(-) create mode 100644 SETUP.MD diff --git a/SETUP.MD b/SETUP.MD new file mode 100644 index 00000000..d7b750a8 --- /dev/null +++ b/SETUP.MD @@ -0,0 +1,141 @@ +# SETUP.MD + +This file documents how to provision a clean development environment for the five packages in this repo (`uipath-agent-framework`, `uipath-google-adk`, `uipath-llamaindex`, `uipath-openai-agents`, `uipath-pydantic-ai`), run the build, execute the tests, and validate a sample code change end-to-end. It is intended both as a quick reference for human contributors and as a structured guide for automated environment-setup tooling. + +## Prerequisites + +- Python 3.11+ +- [uv](https://docs.astral.sh/uv/) 0.5+ + +### Supported platforms + +`uv` is shell- and OS-agnostic, so the commands below run unchanged on every supported platform: + +- [x] Linux +- [x] Windows +- [x] macOS + +## Environment Variables + +None required for environment setup, build, or unit tests. The suites under the `Test` section run fully offline and require no external authentication. + +> **All commands below must be run from the repository root.** The `uv --directory packages/` invocations resolve each subpackage relative to the current working directory. The first line of `## Setup` enforces this by `cd`-ing to the git root. + +## Setup + +```bash +cd "$(git rev-parse --show-toplevel)" +python3 -m pip install --upgrade uv + +# Sync all five packages (each is independent) +uv --directory packages/uipath-agent-framework sync --all-extras +uv --directory packages/uipath-google-adk sync --all-extras +uv --directory packages/uipath-llamaindex sync --all-extras +uv --directory packages/uipath-openai-agents sync --all-extras +uv --directory packages/uipath-pydantic-ai sync --all-extras +``` + +## Verify Setup + +```bash +uv --version +uv --directory packages/uipath-pydantic-ai run python --version +uv --directory packages/uipath-agent-framework run python -c "import uipath_agent_framework; print('uipath-agent-framework ok')" +uv --directory packages/uipath-google-adk run python -c "import uipath_google_adk; print('uipath-google-adk ok')" +uv --directory packages/uipath-llamaindex run python -c "import uipath_llamaindex; print('uipath-llamaindex ok')" +uv --directory packages/uipath-openai-agents run python -c "import uipath_openai_agents; print('uipath-openai-agents ok')" +uv --directory packages/uipath-pydantic-ai run python -c "import uipath_pydantic_ai; print('uipath-pydantic-ai ok')" +``` + +## Build + +N/A + +## Test + +```bash +uv --directory packages/uipath-agent-framework run pytest +uv --directory packages/uipath-google-adk run pytest +uv --directory packages/uipath-llamaindex run pytest +uv --directory packages/uipath-openai-agents run pytest +uv --directory packages/uipath-pydantic-ai run pytest +``` + +## Sample Code Change + +### The change + +Add a new `agent_count` property to `PydanticAiConfig` in `packages/uipath-pydantic-ai/src/uipath_pydantic_ai/runtime/config.py`, immediately after the existing `entrypoint` property: + +```python +@property +def agent_count(self) -> int: + """Number of agents defined in the configuration.""" + return len(self.agents) +``` + +Then create `packages/uipath-pydantic-ai/tests/test_config_agent_count.py` with two pytest tests: + +```python +"""Tests for PydanticAiConfig.agent_count.""" + +import json +from pathlib import Path + +from uipath_pydantic_ai.runtime.config import PydanticAiConfig + + +def test_agent_count_single(tmp_path: Path) -> None: + config_path = tmp_path / "pydantic_ai.json" + config_path.write_text(json.dumps({"agents": {"main": "main:agent"}})) + cfg = PydanticAiConfig(str(config_path)) + assert cfg.agent_count == 1 + + +def test_agent_count_multiple(tmp_path: Path) -> None: + config_path = tmp_path / "pydantic_ai.json" + config_path.write_text( + json.dumps( + { + "agents": { + "alpha": "alpha:agent", + "beta": "beta:agent", + "gamma": "gamma:agent", + } + } + ) + ) + cfg = PydanticAiConfig(str(config_path)) + assert cfg.agent_count == 3 +``` + +### Verification + +```bash +uv --directory packages/uipath-pydantic-ai run pytest tests/test_config_agent_count.py -v +``` + +## Test with a real UiPath Coded Agent + +The unit tests above are necessary but not sufficient — they don't exercise the package end-to-end through a real agent. The flow below validates changes against a live runtime: + +1. Apply the code changes locally. +2. Run the unit tests (see the `Sample Code Change` section above). +3. Scaffold a coded UiPath agent (PydanticAI / OpenAI / Google ADK / LlamaIndex / Agent Framework, matching the package you changed) that exercises the changed code path. +4. In the downstream project's `pyproject.toml`, add this local library as an editable dependency (substitute the package you changed): + + ```toml + [tool.uv.sources] + uipath-pydantic-ai = { path = "../path/to/uipath-integrations-python/packages/uipath-pydantic-ai", editable = true } + ``` + +5. Exercise the new behavior end-to-end: + + ```bash + uv run uipath run --input '{...}' + ``` + +6. (Optional) Open a PR and apply the `build:dev` label — this publishes the development version to Test PyPI. +7. The PR description is updated automatically with instructions for pointing the downstream agent at the Test PyPI dev version. +8. Push the dev version to UiPath with [`uipath push`](https://uipath.github.io/uipath-python/cli/#push), then deploy it to Orchestrator or Studio Web with [`uipath deploy`](https://uipath.github.io/uipath-python/cli/#deploy), and run it in cloud to confirm the changes behave correctly against the real platform. +9. Once validation is done, close the dev PR — these PRs are not meant to be merged; their only purpose was to publish a Test PyPI build for end-to-end validation. diff --git a/packages/uipath-llamaindex/docs/llms_and_embeddings.md b/packages/uipath-llamaindex/docs/llms_and_embeddings.md index 15e25eb4..01416be8 100644 --- a/packages/uipath-llamaindex/docs/llms_and_embeddings.md +++ b/packages/uipath-llamaindex/docs/llms_and_embeddings.md @@ -1,7 +1,39 @@ # LLMs and Embeddings -UiPath provides pre-configured LLM and embedding classes that handle authentication, routing, and configuration automatically, allowing you to focus on building your agents. -You do not need to add API keys from OpenAI, AWS, or Google, usage of these models will consume `Agent Units` on your account. +UiPath provides pre-configured LLM and embedding classes for several providers (OpenAI via `UiPathOpenAI`, Anthropic on AWS Bedrock via `UiPathChatBedrockConverse`, Google Vertex AI via `UiPathVertex`, and more), plus embeddings via `UiPathOpenAIEmbedding`. These handle authentication, routing, and configuration automatically, allowing you to focus on building your agents. You do not need to add API keys from OpenAI, AWS, or Google, usage of these models will consume `Agent Units` on your account. + +## Available models + +LLM models are served through the UiPath LLM Gateway and are subject to [AI Trust Layer](https://docs.uipath.com/automation-cloud/automation-cloud/latest/admin-guide/about-ai-trust-layer) policies, so the exact set of models available to you depends on your tenant configuration. List the models you can use with the `uipath` CLI: + +```console +$ uipath list-models + Available LLM Models +┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━┓ +┃ AwsBedrock ┃ OpenAi ┃ VertexAi ┃ +┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━┩ +│ anthropic.claude-haiku-4-5-20251001-v1:0 │ gpt-4.1-2025-04-14 │ gemini-2.5-flash │ +│ anthropic.claude-opus-4-7 │ gpt-4.1-mini-2025-04-14 │ gemini-2.5-pro │ +│ ... │ ... │ ... │ +└──────────────────────────────────────────┴─────────────────────────┴──────────────────┘ +``` + +Pick a model id from the relevant provider column and pass it (or the matching enum member) to the matching class: + +```python +from uipath_llamaindex.llms import UiPathOpenAI +from uipath_llamaindex.llms.bedrock import UiPathChatBedrockConverse +from uipath_llamaindex.llms.vertex import UiPathVertex + +# OpenAI models +llm = UiPathOpenAI(model="gpt-4.1-mini-2025-04-14") + +# AWS Bedrock (Anthropic) models +llm = UiPathChatBedrockConverse(model="anthropic.claude-sonnet-4-5-20250929-v1:0") + +# Google Vertex AI (Gemini) models +llm = UiPathVertex(model="gemini-2.5-flash") +``` ## UiPathOpenAI @@ -9,17 +41,7 @@ The `UiPathOpenAI` class is a pre-configured Azure OpenAI client that routes req ### Available Models -The following OpenAI models are available through the `OpenAIModel` enum: - -- `GPT_4_1_2025_04_14` -- `GPT_4_1_MINI_2025_04_14` -- `GPT_4_1_NANO_2025_04_14` -- `GPT_4O_2024_05_13` -- `GPT_4O_2024_08_06` -- `GPT_4O_2024_11_20` -- `GPT_4O_MINI_2024_07_18` (default) -- `O3_MINI_2025_01_31` -- `TEXT_DAVINCI_003` +The OpenAI models from the `OpenAi` column of [`uipath list-models`](#available-models) can be used here, either as a model string or via the `OpenAIModel` enum. ### Basic Usage @@ -143,9 +165,7 @@ from uipath_llamaindex.llms import BedrockModel llm = UiPathChatBedrock(model=BedrockModel.anthropic_claude_sonnet_4) ``` -Currently, the following models can be used (this list can be updated in the future): - -- `anthropic.claude-3-7-sonnet-20250219-v1:0`, `anthropic.claude-sonnet-4-20250514-v1:0`, `anthropic.claude-sonnet-4-5-20250929-v1:0`, `anthropic.claude-haiku-4-5-20251001-v1:0` +The available models are the ones in the `AwsBedrock` column of [`uipath list-models`](#available-models). ## UiPathVertex @@ -184,9 +204,7 @@ response = llm.chat(messages) print(response) ``` -Currently, the following models can be used (this list can be updated in the future): - -- `gemini-2.0-flash-001`, `gemini-2.5-flash`, `gemini-2.5-pro` +The available models are the ones in the `VertexAi` column of [`uipath list-models`](#available-models). ## Integration with LlamaIndex From 09dcec6292fce83e04b41be3d0e44aa3dfd5edc6 Mon Sep 17 00:00:00 2001 From: Aditi Kumari Date: Tue, 30 Jun 2026 16:54:12 +0530 Subject: [PATCH 06/11] refactor(governance): type Pydantic AI part extraction against public message models Replace the duck-typed `_part_kind` string-discriminator with isinstance checks against pydantic_ai.messages public types (UserPromptPart, ToolReturnPart, ToolCallPart, BuiltinToolCallPart, TextPart), matching the typed-extraction pattern from the LangChain adapter review (#899). Drops the `_part_kind` helper and the getattr fallbacks; attribute access is now type-checked. Co-Authored-By: Claude Opus 4.8 --- .../uipath_pydantic_ai/governance/adapter.py | 42 ++++++++----------- 1 file changed, 17 insertions(+), 25 deletions(-) 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 index be31c869..398e854b 100644 --- a/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/governance/adapter.py +++ b/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/governance/adapter.py @@ -41,6 +41,13 @@ from typing import Any, AsyncIterator, Dict, List from uuid import uuid4 +from pydantic_ai.messages import ( + BuiltinToolCallPart, + TextPart, + ToolCallPart, + ToolReturnPart, + UserPromptPart, +) from pydantic_ai.models import Model, ModelRequestParameters, StreamedResponse from pydantic_ai.models.wrapper import WrapperModel from pydantic_ai.settings import ModelSettings @@ -199,11 +206,8 @@ def on_request(self, messages: Any) -> None: 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), - ) + if isinstance(part, ToolReturnPart): + self._after_tool(part.tool_name or "unknown", part.content) # ----- after the model call --------------------------------------- @@ -212,11 +216,8 @@ def on_response(self, response: Any) -> None: 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), - ) + if isinstance(part, (ToolCallPart, BuiltinToolCallPart)): + self._tool_call(part.tool_name or "unknown", part.args) # ----- individual evaluate_* wrappers (block-propagate, else swallow) -- @@ -299,11 +300,10 @@ def _parts_input_text(cls, parts: Any) -> str: """ 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))) + if isinstance(part, UserPromptPart): + collected.append(_content_text(part.content)) + elif isinstance(part, ToolReturnPart): + collected.append(_stringify(part.content)) return "\n".join(p for p in collected if p)[:_BEFORE_MODEL_TEXT_CAP] @classmethod @@ -311,10 +311,8 @@ 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) + if isinstance(part, TextPart) and part.content: + collected.append(part.content) return "\n".join(collected)[:_BEFORE_MODEL_TEXT_CAP] @@ -323,12 +321,6 @@ def _response_text(cls, parts: Any) -> str: # -------------------------------------------------------------------------- -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: From d0ff1f10b8fd4a89a94162f6f4a10c456e06099e Mon Sep 17 00:00:00 2001 From: Aditi Kumari Date: Tue, 30 Jun 2026 19:39:22 +0530 Subject: [PATCH 07/11] refactor(governance): migrate Pydantic AI adapter to factory-evaluator Core PR #1761 removed BaseAdapter from uipath-core. Migrate to the factory-evaluator pattern (matching #899): - governance/adapter.py -> model.py: replace the BaseAdapter subclass (name/can_handle/attach/detach) with module-level install_governance(); keep GovernanceModel + GovernanceCallbacks. File named for its seam (the model wrapper), like LangChain's callbacks.py. - runtime/factory.py: new_runtime reads `evaluator` from kwargs and calls install_governance on the resolved agent. - governance/__init__.py: drop register_governance_adapter + registry import; expose install_governance. No import-time side effects. - pyproject.toml: remove the uipath.governance.adapters entry point. - tests (test_adapter.py -> test_model.py): drop can_handle/attach/ detach; cover install_governance + factory wiring. ruff + mypy clean; 17 governance tests pass. Co-Authored-By: Claude Opus 4.8 --- packages/uipath-pydantic-ai/pyproject.toml | 3 - .../uipath_pydantic_ai/governance/__init__.py | 51 ++------- .../governance/{adapter.py => model.py} | 101 +++++++----------- .../src/uipath_pydantic_ai/runtime/factory.py | 18 +++- .../{test_adapter.py => test_model.py} | 81 ++++++++------ 5 files changed, 111 insertions(+), 143 deletions(-) rename packages/uipath-pydantic-ai/src/uipath_pydantic_ai/governance/{adapter.py => model.py} (80%) rename packages/uipath-pydantic-ai/tests/governance/{test_adapter.py => test_model.py} (83%) diff --git a/packages/uipath-pydantic-ai/pyproject.toml b/packages/uipath-pydantic-ai/pyproject.toml index bcfc47c9..960fc98f 100644 --- a/packages/uipath-pydantic-ai/pyproject.toml +++ b/packages/uipath-pydantic-ai/pyproject.toml @@ -28,9 +28,6 @@ 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 index 0efd4f7e..3c9199ed 100644 --- a/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/governance/__init__.py +++ b/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/governance/__init__.py @@ -1,52 +1,21 @@ """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. +Exposes :func:`install_governance` — wraps a ``pydantic_ai.Agent``'s ``model`` +with a :class:`GovernanceModel` that brackets every model call with governance +(BEFORE_MODEL, AFTER_MODEL, TOOL_CALL, AFTER_TOOL). Wired into a run by passing +an ``evaluator`` to :class:`UiPathPydanticAIRuntimeFactory`; the factory calls +:func:`install_governance` on the resolved agent. + +Importing this module has no side effects: no adapter is registered, no global +state is mutated. """ 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") - +from .model import GovernanceCallbacks, GovernanceModel, install_governance __all__ = [ "GovernanceCallbacks", "GovernanceModel", - "PydanticAIAdapter", - "register_governance_adapter", + "install_governance", ] \ 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/model.py similarity index 80% rename from packages/uipath-pydantic-ai/src/uipath_pydantic_ai/governance/adapter.py rename to packages/uipath-pydantic-ai/src/uipath_pydantic_ai/governance/model.py index 398e854b..f19be66a 100644 --- a/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/governance/adapter.py +++ b/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/governance/model.py @@ -1,4 +1,4 @@ -"""Pydantic AI adapter for UiPath governance. +"""Pydantic AI governance model wrapper for UiPath. 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 @@ -17,15 +17,17 @@ 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. +Because the wrap is installed on ``agent.model`` in place, +:func:`install_governance` returns the **original agent**. 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. +The evaluator protocol comes from ``uipath-core``; this package contributes +only the Pydantic-AI-specific wiring. Governance is installed by the runtime +factory: passing an ``evaluator`` to ``new_runtime`` calls +:func:`install_governance` on the resolved agent. No adapter registry, no +entry point, no import-time side effects. Audit emission and enforcement (raising :class:`GovernanceBlockException` on DENY) are owned by the evaluator. The wrapper only extracts payloads and calls @@ -41,6 +43,7 @@ from typing import Any, AsyncIterator, Dict, List from uuid import uuid4 +from pydantic_ai import Agent from pydantic_ai.messages import ( BuiltinToolCallPart, TextPart, @@ -51,7 +54,7 @@ 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.adapters import EvaluatorProtocol from uipath.core.governance.exceptions import GovernanceBlockException logger = logging.getLogger(__name__) @@ -60,69 +63,39 @@ # 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" +def install_governance( + agent: Agent, + evaluator: EvaluatorProtocol, + *, + agent_name: str, + session_id: str, +) -> Agent: + """Wrap ``agent.model`` with a :class:`GovernanceModel` (mutated in place). -class PydanticAIAdapter(BaseAdapter): - """Adapter for the Pydantic AI framework. + Returns the original ``agent``. Idempotent: an already-wrapped model is + left untouched. If the agent has no concrete ``Model`` bound (the model is + supplied per-run), there is nothing to wrap and a warning is logged. - Detects ``pydantic_ai.Agent`` instances and wraps their ``model`` with a - :class:`GovernanceModel`. + Called by :class:`UiPathPydanticAIRuntimeFactory` when an ``evaluator`` + is supplied to ``new_runtime``. """ - - @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 + model = getattr(agent, "model", None) + if isinstance(model, GovernanceModel): + return agent # idempotent — already governed + if not isinstance(model, Model): + logger.warning( + "install_governance: agent has no bound Model to wrap (got %s); " + "model-layer governance will not fire", + type(model).__name__, ) - 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 + callbacks = GovernanceCallbacks( + evaluator=evaluator, agent_name=agent_name, session_id=session_id + ) + agent.model = GovernanceModel(model, callbacks) + logger.debug("Wrapped Pydantic AI agent model with governance") + return agent class GovernanceModel(WrapperModel): diff --git a/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/runtime/factory.py b/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/runtime/factory.py index c415baef..593291ab 100644 --- a/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/runtime/factory.py +++ b/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/runtime/factory.py @@ -4,6 +4,7 @@ from typing import Any from pydantic_ai import Agent +from uipath.core.adapters import EvaluatorProtocol from uipath.runtime import ( UiPathRuntimeContext, UiPathRuntimeFactorySettings, @@ -12,6 +13,7 @@ ) from uipath.runtime.errors import UiPathErrorCategory +from uipath_pydantic_ai.governance import install_governance from uipath_pydantic_ai.runtime.config import PydanticAiConfig from uipath_pydantic_ai.runtime.errors import ( UiPathPydanticAIErrorCode, @@ -215,6 +217,7 @@ async def _create_runtime_instance( agent: Agent, runtime_id: str, entrypoint: str, + evaluator: EvaluatorProtocol | None = None, ) -> UiPathRuntimeProtocol: """ Create a runtime instance from an agent. @@ -223,10 +226,20 @@ async def _create_runtime_instance( agent: The PydanticAI Agent runtime_id: Unique identifier for the runtime instance entrypoint: Agent entrypoint name + evaluator: When supplied, governance is installed on the agent's + model in place via :func:`install_governance`. Returns: Configured runtime instance """ + if evaluator is not None: + install_governance( + agent, + evaluator, + agent_name=entrypoint, + session_id=runtime_id, + ) + return UiPathPydanticAIRuntime( agent=agent, runtime_id=runtime_id, @@ -242,7 +255,9 @@ async def new_runtime( Args: entrypoint: Agent name from pydantic_ai.json runtime_id: Unique identifier for the runtime instance - **kwargs: Additional keyword arguments (unused) + **kwargs: Forwarded factory kwargs. Recognized: ``evaluator`` + (``EvaluatorProtocol``) — when present, governance is installed + on the agent's model via :func:`install_governance`. Returns: Configured runtime instance with agent @@ -252,6 +267,7 @@ async def new_runtime( return await self._create_runtime_instance( agent=agent, runtime_id=runtime_id, + evaluator=kwargs.get("evaluator"), entrypoint=entrypoint, ) diff --git a/packages/uipath-pydantic-ai/tests/governance/test_adapter.py b/packages/uipath-pydantic-ai/tests/governance/test_model.py similarity index 83% rename from packages/uipath-pydantic-ai/tests/governance/test_adapter.py rename to packages/uipath-pydantic-ai/tests/governance/test_model.py index 3495fb05..7ad90e9d 100644 --- a/packages/uipath-pydantic-ai/tests/governance/test_adapter.py +++ b/packages/uipath-pydantic-ai/tests/governance/test_model.py @@ -1,4 +1,4 @@ -"""Unit tests for the Pydantic AI governance adapter. +"""Unit tests for the Pydantic AI governance model wrapper. These tests use real ``pydantic_ai`` message parts (``UserPromptPart`` etc.) so the part-extraction logic is exercised against the actual types, plus the @@ -12,6 +12,7 @@ from __future__ import annotations import logging +from types import SimpleNamespace from typing import Any, List import pytest @@ -27,12 +28,12 @@ from pydantic_ai.models.test import TestModel from uipath.core.governance.exceptions import GovernanceBlockException -from uipath_pydantic_ai.governance.adapter import ( +from uipath_pydantic_ai.governance.model import ( _BEFORE_MODEL_TEXT_CAP, GovernanceCallbacks, GovernanceModel, - PydanticAIAdapter, _coerce_args, + install_governance, ) # -------------------------------------------------------------------------- @@ -80,56 +81,68 @@ def _hooks(ev: FakeEvaluator) -> List[str]: # -------------------------------------------------------------------------- -# can_handle +# install_governance # -------------------------------------------------------------------------- -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(): +def test_install_governance_wraps_model(): agent = Agent(model=TestModel()) - original = agent.model - adapter = PydanticAIAdapter() - returned = adapter.attach(agent, agent_id="x", session_id="s", evaluator=FakeEvaluator()) + returned = install_governance(agent, FakeEvaluator(), agent_name="x", session_id="s") assert returned is agent assert isinstance(agent.model, GovernanceModel) - adapter.detach(agent) - assert agent.model is original -def test_attach_is_idempotent(): +def test_install_governance_is_idempotent(): agent = Agent(model=TestModel()) - adapter = PydanticAIAdapter() ev = FakeEvaluator() - adapter.attach(agent, agent_id="x", session_id="s", evaluator=ev) + install_governance(agent, ev, agent_name="x", session_id="s") wrapped = agent.model - adapter.attach(agent, agent_id="x", session_id="s", evaluator=ev) + install_governance(agent, ev, agent_name="x", session_id="s") assert agent.model is wrapped # not double-wrapped -def test_attach_warns_when_no_bound_model(caplog): +def test_install_governance_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()) + install_governance(agent, FakeEvaluator(), agent_name="x", session_id="s") assert any("no bound Model" in r.message for r in caplog.records) +# -------------------------------------------------------------------------- +# Factory wiring — the evaluator kwarg drives install_governance +# -------------------------------------------------------------------------- + + +def _factory_without_init(): + """A factory instance that skips __init__ (avoids config/IO).""" + from uipath_pydantic_ai.runtime.factory import UiPathPydanticAIRuntimeFactory + + return UiPathPydanticAIRuntimeFactory.__new__(UiPathPydanticAIRuntimeFactory) + + +async def test_factory_installs_governance_when_evaluator_supplied(monkeypatch): + from uipath_pydantic_ai.runtime import factory as factory_mod + + monkeypatch.setattr(factory_mod, "UiPathPydanticAIRuntime", lambda **kw: SimpleNamespace(**kw)) + agent = Agent(model=TestModel()) + await _factory_without_init()._create_runtime_instance( + agent=agent, runtime_id="r", entrypoint="e", evaluator=FakeEvaluator() + ) + assert isinstance(agent.model, GovernanceModel) + + +async def test_factory_skips_governance_without_evaluator(monkeypatch): + from uipath_pydantic_ai.runtime import factory as factory_mod + + monkeypatch.setattr(factory_mod, "UiPathPydanticAIRuntime", lambda **kw: SimpleNamespace(**kw)) + agent = Agent(model=TestModel()) + original = agent.model + await _factory_without_init()._create_runtime_instance( + agent=agent, runtime_id="r", entrypoint="e" + ) + assert agent.model is original + + # -------------------------------------------------------------------------- # on_request → BEFORE_MODEL + AFTER_TOOL # -------------------------------------------------------------------------- From f3a7cbc8529c9f4d78de1d1cd09831c820f0d748 Mon Sep 17 00:00:00 2001 From: Aditi Kumari Date: Wed, 1 Jul 2026 16:13:34 +0530 Subject: [PATCH 08/11] =?UTF-8?q?fix(pydantic-ai):=20address=20governance?= =?UTF-8?q?=20review=20=E2=80=94=20streaming,=20ModelRequest,=20tool=5Fcal?= =?UTF-8?q?l=5Fid?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review findings (Viswa) for PR #359: - Streaming AFTER_MODEL now runs in a finally *inside* the request_stream async-with, so it fires even if the consumer's `async for` raises partway through, and while the stream context is still open. Documented that streaming governance is inherently post-hoc (tokens are already emitted by the time the response is complete; a DENY aborts the run but cannot un-send them). - _latest_request now scans from the end for the last ModelRequest instead of blindly taking messages[-1] — a history ending in a ModelResponse (mid tool round-trip) would otherwise scan the wrong side of the exchange. - Pass tool_call_id through to evaluate_tool_call / evaluate_after_tool so the evaluator can correlate a call with its result. - Bounded joins (_join_within_cap) for request/response/content text and a capped _stringify, so an oversized tool-return can't build a multi-megabyte string before the final slice. - Drop the per-callbacks uuid trace_id (identical for every call); trace correlation is owned by the layer below, matching LangChain. Requires uipath-core >= 0.5.20 (removed trace_id from EvaluatorProtocol) — bumped. - Count llm/tool calls only after governance passes (no inflation on block). - Documented that agent.model is a property whose setter the agent re-reads per run (no stale ref), and that a per-run model override bypasses the wrap. Tests: added streaming happy-path (finalized response governed), streaming finally-on-consumer-error, ModelRequest-filter, and tool_call_id passthrough. Co-Authored-By: Claude Opus 4.8 --- .../uipath_pydantic_ai/governance/model.py | 189 ++++++++++++------ .../tests/governance/test_model.py | 120 ++++++++++- packages/uipath-pydantic-ai/uv.lock | 6 +- 3 files changed, 247 insertions(+), 68 deletions(-) diff --git a/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/governance/model.py b/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/governance/model.py index f19be66a..23e29a4a 100644 --- a/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/governance/model.py +++ b/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/governance/model.py @@ -40,12 +40,12 @@ import json import logging from contextlib import asynccontextmanager -from typing import Any, AsyncIterator, Dict, List -from uuid import uuid4 +from typing import Any, AsyncIterator, Dict, Iterable, List from pydantic_ai import Agent from pydantic_ai.messages import ( BuiltinToolCallPart, + ModelRequest, TextPart, ToolCallPart, ToolReturnPart, @@ -93,6 +93,11 @@ def install_governance( callbacks = GovernanceCallbacks( evaluator=evaluator, agent_name=agent_name, session_id=session_id ) + # ``agent.model`` is a property whose setter stores ``_model``; the agent + # re-reads ``self.model`` on every run (``_get_model``), so this in-place + # wrap takes effect for all subsequent runs with no stale reference. A model + # supplied per-run (``agent.run(model=...)``) bypasses this wrap — that path + # is governed by whatever model the caller passes, not by us. agent.model = GovernanceModel(model, callbacks) logger.debug("Wrapped Pydantic AI agent model with governance") return agent @@ -130,17 +135,28 @@ async def request_stream( 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) + try: + yield stream + finally: + # Once the caller has consumed the stream the final response is + # assembled — govern it the same as the non-streaming path. This + # runs in a ``finally`` so AFTER_MODEL / TOOL_CALL still fire + # even if the consumer's ``async for`` raised partway through. + # + # Streaming governance is inherently post-hoc: the tokens have + # already been streamed to the caller by the time the response + # is complete, so a DENY here aborts the run but cannot un-send + # what was already emitted. The block exception still + # propagates; any other governance error is logged and swallowed + # so a governance bug can't break the run. + try: + self._callbacks.on_response(stream.get()) + except GovernanceBlockException: + raise + except Exception as e: # noqa: BLE001 + logger.warning( + "after-stream governance check failed (continuing): %s", e + ) class GovernanceCallbacks: @@ -160,7 +176,10 @@ def __init__( self._evaluator = evaluator self._agent_name = agent_name self._session_id = session_id - self._trace_id = str(uuid4()) + # ``trace_id`` is intentionally NOT held here. A single uuid minted at + # install time would be identical for every call. Trace correlation is + # owned by the layer below (OTel span / HTTP resolve at call time), + # matching the LangChain adapter. self._session_state: Dict[str, Any] = {"tool_calls": 0, "llm_calls": 0} # ----- before the model call -------------------------------------- @@ -180,7 +199,11 @@ def on_request(self, messages: Any) -> None: self._before_model(self._parts_input_text(parts)) for part in parts: if isinstance(part, ToolReturnPart): - self._after_tool(part.tool_name or "unknown", part.content) + self._after_tool( + part.tool_name or "unknown", + part.content, + tool_call_id=getattr(part, "tool_call_id", None), + ) # ----- after the model call --------------------------------------- @@ -190,20 +213,25 @@ def on_response(self, response: Any) -> None: self._after_model(self._response_text(parts)) for part in parts: if isinstance(part, (ToolCallPart, BuiltinToolCallPart)): - self._tool_call(part.tool_name or "unknown", part.args) + self._tool_call( + part.tool_name or "unknown", + part.args, + tool_call_id=getattr(part, "tool_call_id", 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, + ) + # Count only calls that passed governance — a DENY raises above, so + # a blocked call must not inflate the counter. + self._session_state["llm_calls"] = ( + self._session_state.get("llm_calls", 0) + 1 ) except GovernanceBlockException: raise @@ -216,39 +244,44 @@ def _after_model(self, text: str) -> None: 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: + def _tool_call( + self, tool_name: str, args: Any, tool_call_id: str | None = None + ) -> 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, + tool_call_id=tool_call_id, + ) + # Count only calls that passed governance; the evaluator saw the + # count of prior tool calls, and a DENY raises before this bump. + self._session_state["tool_calls"] = ( + self._session_state.get("tool_calls", 0) + 1 ) 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: + def _after_tool( + self, tool_name: str, content: Any, tool_call_id: str | None = None + ) -> 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, + tool_call_id=tool_call_id, ) except GovernanceBlockException: raise @@ -259,34 +292,48 @@ def _after_tool(self, tool_name: str, content: Any) -> None: @staticmethod def _latest_request(messages: Any) -> Any: - """Return the most recent message (a ``ModelRequest``) or ``None``.""" + """Return the most recent ``ModelRequest`` message, or ``None``. + + Scans from the end for the last ``ModelRequest`` rather than blindly + taking ``messages[-1]``: the history can end with a ``ModelResponse`` + (e.g. mid tool-call round-trips), and treating that as request input + would scan the wrong side of the exchange. + """ if not messages or not isinstance(messages, (list, tuple)): return None - return messages[-1] + for message in reversed(messages): + if isinstance(message, ModelRequest): + return message + return None @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`. + follow-up turn). Joined with a running cap so an oversized tool-return + can't build a multi-megabyte string before the final slice. """ - collected: List[str] = [] - for part in parts: - if isinstance(part, UserPromptPart): - collected.append(_content_text(part.content)) - elif isinstance(part, ToolReturnPart): - collected.append(_stringify(part.content)) - return "\n".join(p for p in collected if p)[:_BEFORE_MODEL_TEXT_CAP] + + def _pieces() -> Iterable[str]: + for part in parts: + if isinstance(part, UserPromptPart): + yield _content_text(part.content) + elif isinstance(part, ToolReturnPart): + yield _stringify(part.content) + + return _join_within_cap(_pieces(), _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 isinstance(part, TextPart) and part.content: - collected.append(part.content) - return "\n".join(collected)[:_BEFORE_MODEL_TEXT_CAP] + """Join ``TextPart`` content from a model response's parts (running cap).""" + + def _pieces() -> Iterable[str]: + for part in parts: + if isinstance(part, TextPart) and part.content: + yield part.content + + return _join_within_cap(_pieces(), _BEFORE_MODEL_TEXT_CAP) # -------------------------------------------------------------------------- @@ -294,22 +341,42 @@ def _response_text(cls, parts: Any) -> str: # -------------------------------------------------------------------------- +def _join_within_cap(pieces: Iterable[str], cap: int) -> str: + """Join non-empty ``pieces`` with newlines, stopping once ``cap`` is hit. + + Bounds the work (and the allocation) to ``cap`` characters instead of + building the full string and slicing afterwards. + """ + collected: List[str] = [] + remaining = cap + for piece in pieces: + if remaining <= 0: + break + if not piece: + continue + collected.append(piece[:remaining]) + remaining -= len(piece) + 1 + return "\n".join(collected)[:cap] + + 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 + return content[:_BEFORE_MODEL_TEXT_CAP] 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) + + def _pieces() -> Iterable[str]: + for item in content: + if isinstance(item, str): + yield item + else: + text = getattr(item, "text", None) + if isinstance(text, str): + yield text + + return _join_within_cap(_pieces(), _BEFORE_MODEL_TEXT_CAP) return _stringify(content) @@ -328,11 +395,15 @@ def _coerce_args(args: Any) -> Dict[str, Any]: return {} -def _stringify(value: Any) -> str: - """Render a dict / object payload as compact, scannable text.""" +def _stringify(value: Any, cap: int = _BEFORE_MODEL_TEXT_CAP) -> str: + """Render a dict / object payload as compact, scannable text, capped. + + Bounded by ``cap`` so an oversized tool result / return content can't hand + a multi-megabyte string to the evaluator. + """ if isinstance(value, str): - return value + return value[:cap] try: - return json.dumps(value, default=str, ensure_ascii=False) + return json.dumps(value, default=str, ensure_ascii=False)[:cap] except (TypeError, ValueError): - return str(value) \ No newline at end of file + return str(value)[:cap] diff --git a/packages/uipath-pydantic-ai/tests/governance/test_model.py b/packages/uipath-pydantic-ai/tests/governance/test_model.py index 7ad90e9d..dd21f3dc 100644 --- a/packages/uipath-pydantic-ai/tests/governance/test_model.py +++ b/packages/uipath-pydantic-ai/tests/governance/test_model.py @@ -87,7 +87,9 @@ def _hooks(ev: FakeEvaluator) -> List[str]: def test_install_governance_wraps_model(): agent = Agent(model=TestModel()) - returned = install_governance(agent, FakeEvaluator(), agent_name="x", session_id="s") + returned = install_governance( + agent, FakeEvaluator(), agent_name="x", session_id="s" + ) assert returned is agent assert isinstance(agent.model, GovernanceModel) @@ -123,7 +125,9 @@ def _factory_without_init(): async def test_factory_installs_governance_when_evaluator_supplied(monkeypatch): from uipath_pydantic_ai.runtime import factory as factory_mod - monkeypatch.setattr(factory_mod, "UiPathPydanticAIRuntime", lambda **kw: SimpleNamespace(**kw)) + monkeypatch.setattr( + factory_mod, "UiPathPydanticAIRuntime", lambda **kw: SimpleNamespace(**kw) + ) agent = Agent(model=TestModel()) await _factory_without_init()._create_runtime_instance( agent=agent, runtime_id="r", entrypoint="e", evaluator=FakeEvaluator() @@ -134,7 +138,9 @@ async def test_factory_installs_governance_when_evaluator_supplied(monkeypatch): async def test_factory_skips_governance_without_evaluator(monkeypatch): from uipath_pydantic_ai.runtime import factory as factory_mod - monkeypatch.setattr(factory_mod, "UiPathPydanticAIRuntime", lambda **kw: SimpleNamespace(**kw)) + monkeypatch.setattr( + factory_mod, "UiPathPydanticAIRuntime", lambda **kw: SimpleNamespace(**kw) + ) agent = Agent(model=TestModel()) original = agent.model await _factory_without_init()._create_runtime_instance( @@ -165,7 +171,11 @@ def test_on_request_fires_after_tool_for_tool_return(): cb = _make_callbacks(ev) messages = [ ModelRequest( - parts=[ToolReturnPart(tool_name="lookup", content={"balance": "1000"}, tool_call_id="c1")] + parts=[ + ToolReturnPart( + tool_name="lookup", content={"balance": "1000"}, tool_call_id="c1" + ) + ] ) ] cb.on_request(messages) @@ -278,6 +288,102 @@ async def request_stream(self, *_a, **_k): assert stream is not None +async def test_governance_model_request_stream_governs_finalized_response(): + """Streaming happy path: BEFORE_MODEL fires up front, and AFTER_MODEL runs + on the finalized response assembled after the caller consumes the stream.""" + from contextlib import asynccontextmanager + from types import SimpleNamespace + + ev = FakeEvaluator() + cb = _make_callbacks(ev) + final = ModelResponse(parts=[TextPart(content="the streamed answer")]) + + class FakeWrapped: + @asynccontextmanager + async def request_stream(self, *_a, **_k): + yield SimpleNamespace(get=lambda: final) + + gm = GovernanceModel.__new__(GovernanceModel) + gm.wrapped = FakeWrapped() # type: ignore[attr-defined] + gm._callbacks = cb + messages = [ModelRequest(parts=[UserPromptPart(content="the question")])] + + async with gm.request_stream(messages, None, None) as stream: + # BEFORE_MODEL already fired; AFTER_MODEL deferred until the stream + # context exits (final response is assembled). + assert _hooks(ev) == ["before_model"] + assert stream is not None + + assert _hooks(ev) == ["before_model", "after_model"] + assert ev.calls[-1][1]["model_output"] == "the streamed answer" + + +async def test_governance_model_request_stream_governs_even_if_consumer_raises(): + """If the consumer's ``async for`` raises, AFTER_MODEL still fires (finally), + and the consumer's exception propagates.""" + from contextlib import asynccontextmanager + from types import SimpleNamespace + + ev = FakeEvaluator() + cb = _make_callbacks(ev) + final = ModelResponse(parts=[TextPart(content="partial answer")]) + + class FakeWrapped: + @asynccontextmanager + async def request_stream(self, *_a, **_k): + yield SimpleNamespace(get=lambda: final) + + gm = GovernanceModel.__new__(GovernanceModel) + gm.wrapped = FakeWrapped() # type: ignore[attr-defined] + gm._callbacks = cb + messages = [ModelRequest(parts=[UserPromptPart(content="hi")])] + + with pytest.raises(RuntimeError, match="consumer blew up"): + async with gm.request_stream(messages, None, None): + raise RuntimeError("consumer blew up") + + assert "after_model" in _hooks(ev) # ran despite the consumer error + + +def test_latest_request_skips_trailing_model_response(): + """History can end with a ModelResponse (mid tool round-trip); BEFORE_MODEL + must scan the last ModelRequest, not the response.""" + ev = FakeEvaluator() + cb = _make_callbacks(ev) + messages = [ + ModelRequest(parts=[UserPromptPart(content="the real question")]), + ModelResponse(parts=[TextPart(content="assistant reply")]), + ] + cb.on_request(messages) + assert ev.calls[0][1]["model_input"] == "the real question" + + +def test_tool_call_and_after_tool_pass_tool_call_id(): + ev = FakeEvaluator() + cb = _make_callbacks(ev) + cb.on_response( + ModelResponse( + parts=[ToolCallPart(tool_name="t", args={}, tool_call_id="call-42")] + ) + ) + tool_call = [kw for h, kw in ev.calls if h == "tool_call"][0] + assert tool_call["tool_call_id"] == "call-42" + + ev2 = FakeEvaluator() + cb2 = _make_callbacks(ev2) + cb2.on_request( + [ + ModelRequest( + parts=[ + ToolReturnPart(tool_name="t", content="ok", tool_call_id="call-7") + ] + ) + ] + ) + after_tool = [kw for h, kw in ev2.calls if h == "after_tool"][0] + assert after_tool["tool_call_id"] == "call-7" + + # -------------------------------------------------------------------------- # helpers + enforcement # -------------------------------------------------------------------------- @@ -300,7 +406,9 @@ 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")]) + ModelResponse( + parts=[ToolCallPart(tool_name="t", args={}, tool_call_id="c1")] + ) ) @@ -312,4 +420,4 @@ def evaluate_before_model(self, **_: Any) -> None: 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 + assert any("governance check failed" in r.message for r in caplog.records) diff --git a/packages/uipath-pydantic-ai/uv.lock b/packages/uipath-pydantic-ai/uv.lock index 704e7e0c..e02e7469 100644 --- a/packages/uipath-pydantic-ai/uv.lock +++ b/packages/uipath-pydantic-ai/uv.lock @@ -3654,16 +3654,16 @@ wheels = [ [[package]] name = "uipath-core" -version = "0.5.18" +version = "0.5.28" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "opentelemetry-instrumentation" }, { name = "opentelemetry-sdk" }, { name = "pydantic" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/14/b1/d4e555a1a2ccf298195a5f2968e538b0cea8592b3e03f43fc12b178d6c69/uipath_core-0.5.18.tar.gz", hash = "sha256:63ebe8bdb818ca30a4bc9ab0ea8171315680691429931282939359ce039401ab", size = 131988, upload-time = "2026-06-08T14:04:49.688Z" } +sdist = { url = "https://files.pythonhosted.org/packages/4c/f9/8d2f1d98cbebbcf059cf4561f38f34ad4cd58423e4f15cad22bd297a2563/uipath_core-0.5.28.tar.gz", hash = "sha256:942987f6b612c64f93d612ad7b242276ed75f129fdd8f25bc71c24ec8887e388", size = 130578, upload-time = "2026-06-30T14:04:48.841Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/57/de/1a820b33f7bff4565d7649772bc54c88480ac7e70f707097f7da37d05157/uipath_core-0.5.18-py3-none-any.whl", hash = "sha256:351d6faeecfc6a0acea93182e01526f39c04a77e09fa0444be5f4fb580463f5a", size = 54572, upload-time = "2026-06-08T14:04:48.22Z" }, + { url = "https://files.pythonhosted.org/packages/e8/1e/385bb166232a57ebe938cc57ad2717f350bc922bb5d2ce31af84306b7569/uipath_core-0.5.28-py3-none-any.whl", hash = "sha256:b952a46a21710073cbc16d6d5684e9aa645c107f57a636b778cfb94aa81a1e48", size = 54980, upload-time = "2026-06-30T14:04:47.374Z" }, ] [[package]] From 6bca79730871f9cb60cc9995660565565141bfef Mon Sep 17 00:00:00 2001 From: Aditi Kumari Date: Wed, 1 Jul 2026 17:45:42 +0530 Subject: [PATCH 09/11] fix(pydantic-ai): address remaining review minors (page 4) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follow-up to the #359 review pass — the earlier commit covered the blockers + page-3 majors; these are the page-4 minors: - model.py:239 — BuiltinToolReturnPart now fires AFTER_TOOL (symmetric with the built-in TOOL_CALL). A provider-executed built-in tool carries both its call and result inline in the response, so AFTER_TOOL for built-ins fires from on_response, not the next request. - model.py:371 — _coerce_args preserves malformed JSON as {"_raw": args} instead of dropping it to {}, so an arg-based policy can still scan it (a malformed payload must not slip past governance). - model.py:229 — tool_name fallback to "unknown" now logs a warning via a shared _tool_name() helper (used by all three tool-part call sites). Not changed: runtime/factory.py:271 kwargs.get("evaluator") — this is the same magic-string kwarg LangChain #899's factory uses (factory.py:281); it's the host-dispatch contract shared by all adapters, so kept for parity. Tests: added built-in-tool-result AFTER_TOOL coverage; updated _coerce_args variant assertion. 82 pass. Co-Authored-By: Claude Opus 4.8 --- .../uipath_pydantic_ai/governance/model.py | 40 +++++++++++++++++-- .../tests/governance/test_model.py | 30 +++++++++++++- 2 files changed, 65 insertions(+), 5 deletions(-) diff --git a/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/governance/model.py b/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/governance/model.py index 23e29a4a..7b061703 100644 --- a/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/governance/model.py +++ b/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/governance/model.py @@ -45,6 +45,7 @@ from pydantic_ai import Agent from pydantic_ai.messages import ( BuiltinToolCallPart, + BuiltinToolReturnPart, ModelRequest, TextPart, ToolCallPart, @@ -200,7 +201,7 @@ def on_request(self, messages: Any) -> None: for part in parts: if isinstance(part, ToolReturnPart): self._after_tool( - part.tool_name or "unknown", + _tool_name(part), part.content, tool_call_id=getattr(part, "tool_call_id", None), ) @@ -208,16 +209,29 @@ def on_request(self, messages: Any) -> None: # ----- after the model call --------------------------------------- def on_response(self, response: Any) -> None: - """Fire AFTER_MODEL (response text) + TOOL_CALL (each tool-call part).""" + """Fire AFTER_MODEL (response text) + TOOL_CALL / AFTER_TOOL parts. + + A provider-executed **built-in** tool carries both its call + (``BuiltinToolCallPart``) and its result (``BuiltinToolReturnPart``) + inline in the model response, so AFTER_TOOL for built-in tools is fired + here — symmetric with the built-in TOOL_CALL — rather than on the next + request (where only user-tool ``ToolReturnPart``s arrive). + """ parts = getattr(response, "parts", None) or [] self._after_model(self._response_text(parts)) for part in parts: if isinstance(part, (ToolCallPart, BuiltinToolCallPart)): self._tool_call( - part.tool_name or "unknown", + _tool_name(part), part.args, tool_call_id=getattr(part, "tool_call_id", None), ) + elif isinstance(part, BuiltinToolReturnPart): + self._after_tool( + _tool_name(part), + part.content, + tool_call_id=getattr(part, "tool_call_id", None), + ) # ----- individual evaluate_* wrappers (block-propagate, else swallow) -- @@ -380,6 +394,22 @@ def _pieces() -> Iterable[str]: return _stringify(content) +def _tool_name(part: Any) -> str: + """Return ``part.tool_name`` or ``"unknown"``, logging the fallback. + + A missing tool name means TOOL_CALL / AFTER_TOOL can't be attributed to a + real tool, so surface it rather than silently reporting ``"unknown"``. + """ + name = getattr(part, "tool_name", None) + if name: + return name + logger.warning( + "governance: %s carries no tool_name; reporting 'unknown'", + type(part).__name__, + ) + return "unknown" + + def _coerce_args(args: Any) -> Dict[str, Any]: """Normalise ``ToolCallPart.args`` (dict / JSON string / None) to a dict.""" if args is None: @@ -391,7 +421,9 @@ def _coerce_args(args: Any) -> Dict[str, Any]: parsed = json.loads(args) return parsed if isinstance(parsed, dict) else {"_": parsed} except (TypeError, ValueError): - return {} + # Preserve the raw string so an arg-based policy can still scan it; + # a malformed payload must not be a way to slip past governance. + return {"_raw": args} return {} diff --git a/packages/uipath-pydantic-ai/tests/governance/test_model.py b/packages/uipath-pydantic-ai/tests/governance/test_model.py index dd21f3dc..77bff573 100644 --- a/packages/uipath-pydantic-ai/tests/governance/test_model.py +++ b/packages/uipath-pydantic-ai/tests/governance/test_model.py @@ -18,6 +18,8 @@ import pytest from pydantic_ai import Agent from pydantic_ai.messages import ( + BuiltinToolCallPart, + BuiltinToolReturnPart, ModelRequest, ModelResponse, TextPart, @@ -225,6 +227,31 @@ def test_on_response_fires_after_model_and_tool_call(): assert tool_call["session_state"]["tool_calls"] == 1 +def test_on_response_fires_after_tool_for_builtin_tool_result(): + """A provider-executed built-in tool carries its result inline in the + response (BuiltinToolReturnPart); AFTER_TOOL must fire for it — symmetric + with the built-in TOOL_CALL.""" + ev = FakeEvaluator() + cb = _make_callbacks(ev) + response = ModelResponse( + parts=[ + BuiltinToolCallPart( + tool_name="web_search", args={"q": "x"}, tool_call_id="b1" + ), + BuiltinToolReturnPart( + tool_name="web_search", content={"results": "found"}, tool_call_id="b1" + ), + ] + ) + cb.on_response(response) + hooks = _hooks(ev) + assert "tool_call" in hooks and "after_tool" in hooks + after_tool = [kw for h, kw in ev.calls if h == "after_tool"][0] + assert after_tool["tool_name"] == "web_search" + assert after_tool["tool_call_id"] == "b1" + assert "found" in after_tool["tool_result"] + + def test_on_response_coerces_json_string_args(): ev = FakeEvaluator() cb = _make_callbacks(ev) @@ -393,7 +420,8 @@ 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") == {} + # malformed JSON is preserved (not dropped) so arg-based policies can scan it + assert _coerce_args("not json") == {"_raw": "not json"} def test_block_in_before_model_propagates(): From 677b89a854813e1952b3ec97ae0a41f66a0c0be0 Mon Sep 17 00:00:00 2001 From: Aditi Kumari Date: Wed, 1 Jul 2026 19:13:19 +0530 Subject: [PATCH 10/11] test(pydantic-ai): make test_model mypy-clean (fix CI lint) Same pre-existing CI-lint failure as the other packages (mypy runs over tests): - FakeEvaluator evaluate_* -> (self, *args, **kwargs) -> Any so it satisfies EvaluatorProtocol; bare dict -> dict[str, Any]. - corrected the FakeWrapped stub ignore ([attr-defined] -> [assignment]) and added arg-type ignores where None is passed for ModelRequestParameters in the request/request_stream bracket tests. Co-Authored-By: Claude Opus 4.8 --- .../tests/governance/test_model.py | 30 +++++++++---------- 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/packages/uipath-pydantic-ai/tests/governance/test_model.py b/packages/uipath-pydantic-ai/tests/governance/test_model.py index 77bff573..809e3e56 100644 --- a/packages/uipath-pydantic-ai/tests/governance/test_model.py +++ b/packages/uipath-pydantic-ai/tests/governance/test_model.py @@ -48,29 +48,29 @@ class FakeEvaluator: def __init__(self, block_on: str | None = None) -> None: self.block_on = block_on - self.calls: List[tuple[str, dict]] = [] + self.calls: List[tuple[str, dict[str, Any]]] = [] 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: + def evaluate_before_agent(self, *args: Any, **kwargs: Any) -> Any: self._record("before_agent", **kwargs) - def evaluate_after_agent(self, **kwargs: Any) -> None: + def evaluate_after_agent(self, *args: Any, **kwargs: Any) -> Any: self._record("after_agent", **kwargs) - def evaluate_before_model(self, **kwargs: Any) -> None: + def evaluate_before_model(self, *args: Any, **kwargs: Any) -> Any: self._record("before_model", **kwargs) - def evaluate_after_model(self, **kwargs: Any) -> None: + def evaluate_after_model(self, *args: Any, **kwargs: Any) -> Any: self._record("after_model", **kwargs) - def evaluate_tool_call(self, **kwargs: Any) -> None: + def evaluate_tool_call(self, *args: Any, **kwargs: Any) -> Any: self._record("tool_call", **kwargs) - def evaluate_after_tool(self, **kwargs: Any) -> None: + def evaluate_after_tool(self, *args: Any, **kwargs: Any) -> Any: self._record("after_tool", **kwargs) @@ -279,10 +279,10 @@ async def request(self, messages, settings, params): return ModelResponse(parts=[TextPart(content="Your balance is 1000.")]) gm = GovernanceModel.__new__(GovernanceModel) # bypass WrapperModel init - gm.wrapped = FakeWrapped() # type: ignore[attr-defined] + gm.wrapped = FakeWrapped() # type: ignore[assignment] # test double, not a real Model gm._callbacks = cb messages = [ModelRequest(parts=[UserPromptPart(content="What is my balance?")])] - await gm.request(messages, None, None) + await gm.request(messages, None, None) # type: ignore[arg-type] assert order == ["MODEL_CALL"] assert _hooks(ev) == ["before_model", "after_model"] @@ -307,11 +307,11 @@ 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.wrapped = FakeWrapped() # type: ignore[assignment] # test double, not a real Model gm._callbacks = cb messages = [ModelRequest(parts=[UserPromptPart(content="hi")])] with pytest.raises(GovernanceBlockException): - async with gm.request_stream(messages, None, None) as stream: + async with gm.request_stream(messages, None, None) as stream: # type: ignore[arg-type] assert stream is not None @@ -331,11 +331,11 @@ async def request_stream(self, *_a, **_k): yield SimpleNamespace(get=lambda: final) gm = GovernanceModel.__new__(GovernanceModel) - gm.wrapped = FakeWrapped() # type: ignore[attr-defined] + gm.wrapped = FakeWrapped() # type: ignore[assignment] # test double, not a real Model gm._callbacks = cb messages = [ModelRequest(parts=[UserPromptPart(content="the question")])] - async with gm.request_stream(messages, None, None) as stream: + async with gm.request_stream(messages, None, None) as stream: # type: ignore[arg-type] # BEFORE_MODEL already fired; AFTER_MODEL deferred until the stream # context exits (final response is assembled). assert _hooks(ev) == ["before_model"] @@ -361,12 +361,12 @@ async def request_stream(self, *_a, **_k): yield SimpleNamespace(get=lambda: final) gm = GovernanceModel.__new__(GovernanceModel) - gm.wrapped = FakeWrapped() # type: ignore[attr-defined] + gm.wrapped = FakeWrapped() # type: ignore[assignment] # test double, not a real Model gm._callbacks = cb messages = [ModelRequest(parts=[UserPromptPart(content="hi")])] with pytest.raises(RuntimeError, match="consumer blew up"): - async with gm.request_stream(messages, None, None): + async with gm.request_stream(messages, None, None): # type: ignore[arg-type] raise RuntimeError("consumer blew up") assert "after_model" in _hooks(ev) # ran despite the consumer error From 2e9092d58fd77b3fb6b93f64de2cb3cf7e8720f2 Mon Sep 17 00:00:00 2001 From: Aditi Kumari Date: Wed, 1 Jul 2026 19:38:40 +0530 Subject: [PATCH 11/11] test(pydantic-ai): cover swallow/extraction branches (Sonar coverage) New-code coverage ~89% -> lift over the 90% gate. Added: non-block-error swallow on the response/tool paths, on_request with only a ModelResponse (empty-input branch), and _content_text/_stringify/_tool_name helper edges. governance/model.py: 82% -> 96%. Co-Authored-By: Claude Opus 4.8 --- .../tests/governance/test_model.py | 71 +++++++++++++++++++ 1 file changed, 71 insertions(+) diff --git a/packages/uipath-pydantic-ai/tests/governance/test_model.py b/packages/uipath-pydantic-ai/tests/governance/test_model.py index 809e3e56..d8b3cf6c 100644 --- a/packages/uipath-pydantic-ai/tests/governance/test_model.py +++ b/packages/uipath-pydantic-ai/tests/governance/test_model.py @@ -449,3 +449,74 @@ def evaluate_before_model(self, **_: Any) -> None: 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) + + +# -------------------------------------------------------------------------- +# coverage: swallow paths on every hook + extraction/helper edges +# -------------------------------------------------------------------------- + + +class _Boom: + """Evaluator whose every evaluate_* raises a non-block error.""" + + def __getattr__(self, _name: str) -> Any: + def _raise(*_a: Any, **_k: Any) -> None: + raise RuntimeError("evaluator bug") + + return _raise + + +@pytest.mark.parametrize( + "invoke", + [ + lambda cb: cb.on_response( + ModelResponse( + parts=[ + TextPart(content="x"), + ToolCallPart(tool_name="t", args={}, tool_call_id="c"), + ] + ) + ), + lambda cb: cb.on_request( + [ + ModelRequest( + parts=[ToolReturnPart(tool_name="t", content="r", tool_call_id="c")] + ) + ] + ), + ], +) +def test_response_and_tool_paths_swallow_non_block_errors(invoke, caplog): + cb = GovernanceCallbacks(evaluator=_Boom(), agent_name="a", session_id="s") # type: ignore[arg-type] + with caplog.at_level(logging.WARNING): + invoke(cb) # must NOT raise — a governance bug can't break the run + assert any("governance check failed" in r.message for r in caplog.records) + + +def test_on_request_with_only_model_response_scans_empty(): + # history ending with (only) a ModelResponse → no ModelRequest → empty input + ev = FakeEvaluator() + cb = _make_callbacks(ev) + cb.on_request([ModelResponse(parts=[TextPart(content="assistant")])]) + assert ev.calls[0][1]["model_input"] == "" + + +def test_extraction_and_helper_edges(): + from uipath_pydantic_ai.governance.model import ( + _content_text, + _stringify, + _tool_name, + ) + + # _content_text: None / str / list-of-(str|obj) / bare object + assert _content_text(None) == "" + assert _content_text("plain") == "plain" + assert "a" in _content_text(["a", SimpleNamespace(text="b")]) + assert isinstance(_content_text(SimpleNamespace()), str) + # _stringify: str passthrough + circular-ref fallback (no crash) + assert _stringify("hi") == "hi" + circular: dict[str, Any] = {} + circular["self"] = circular + assert isinstance(_stringify(circular), str) + # _tool_name: missing name → warns + "unknown" + assert _tool_name(SimpleNamespace()) == "unknown"