From bf034da624de60c48664e11f37f6f9599a4d1246 Mon Sep 17 00:00:00 2001
From: Aditi Kumari <aditi.kumari@uipath.com>
Date: Mon, 22 Jun 2026 23:09:07 +0530
Subject: [PATCH 1/5] 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 <noreply@anthropic.com>
---
 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 <aditi.kumari@uipath.com>
Date: Wed, 24 Jun 2026 17:22:34 +0530
Subject: [PATCH 2/5] 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 <noreply@anthropic.com>
---
 .../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 <aditi.kumari@uipath.com>
Date: Wed, 24 Jun 2026 18:26:52 +0530
Subject: [PATCH 3/5] 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 <noreply@anthropic.com>
---
 .../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 <aditi.kumari@uipath.com>
Date: Wed, 24 Jun 2026 18:40:49 +0530
Subject: [PATCH 4/5] 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 <noreply@anthropic.com>
---
 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/<name>` 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 <agent-name> --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 <aditi.kumari@uipath.com>
Date: Wed, 24 Jun 2026 19:01:04 +0530
Subject: [PATCH 5/5] 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 <noreply@anthropic.com>
---
 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/<name>` 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 <agent-name> --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