From 7bb7c33d96ab89919669c190f1c7f07036c7b073 Mon Sep 17 00:00:00 2001 From: Viswanath Lekshmanan Date: Fri, 19 Jun 2026 17:11:27 +0530 Subject: [PATCH 01/10] feat: add LangChain governance adapter MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Registers a LangChain/LangGraph adapter with the uipath-core adapter registry so GovernanceRuntime can attach BEFORE_MODEL / AFTER_MODEL / TOOL_CALL / AFTER_TOOL hooks via LangChain's callback system. Exposed as a uipath.governance.adapters entry point and self-registers on import. Bumps uipath-core to 0.6.x and uipath-runtime to 0.11.x to pick up the new adapter contracts. Key behaviour: - BEFORE_MODEL is scoped to the latest message only. Earlier passes concatenated the entire prompt history, which re-fired violations from older turns on every subsequent LLM call. The callback now takes the last entry of the last batched prompt (mirrors the BEFORE_AGENT latest_only contract in uipath-runtime). - List-of-blocks content (multimodal, OpenAI function-call, Anthropic tool_use, Claude extended thinking) is walked via _extract_block_text instead of str(msg.content), so structured shapes produce clean text with no dict-repr noise that breaks regex/field rules. - Both on_chat_model_start and on_llm_start cap extracted input at _BEFORE_MODEL_TEXT_CAP = 64000 (matches the runtime side). - register_governance_adapter is idempotent — uses the AdapterRegistry instance state as the single source of truth (no module-level flag). Packaging: - prerelease = "allow" plus override-dependencies for uipath-runtime under [tool.uv] so the dev pin can satisfy the umbrella's stable-only constraint. Re-locked so uipath-runtime resolves from testpypi rather than the local editable Windows path. Co-Authored-By: Claude Opus 4.7 (1M context) --- pyproject.toml | 17 +- src/uipath_langchain/governance/__init__.py | 58 + src/uipath_langchain/governance/adapter.py | 580 ++++++++++ tests/governance/__init__.py | 0 tests/governance/test_adapter.py | 1065 +++++++++++++++++++ uv.lock | 74 +- 6 files changed, 1783 insertions(+), 11 deletions(-) create mode 100644 src/uipath_langchain/governance/__init__.py create mode 100644 src/uipath_langchain/governance/adapter.py create mode 100644 tests/governance/__init__.py create mode 100644 tests/governance/test_adapter.py diff --git a/pyproject.toml b/pyproject.toml index a0053aa9c..88f17eb86 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -6,9 +6,9 @@ readme = { file = "README.md", content-type = "text/markdown" } requires-python = ">=3.11" dependencies = [ "uipath>=2.10.79, <2.11.0", - "uipath-core>=0.5.17, <0.6.0", + "uipath-core>=0.5.18, <0.6.0", "uipath-platform>=0.1.61, <0.2.0", - "uipath-runtime>=0.11.0, <0.12.0", + "uipath-runtime>=0.11.0.dev1001180449,<0.11.0.dev1001190000", "langgraph>=1.1.8, <2.0.0", "langchain-core>=1.2.11, <2.0.0", "langgraph-checkpoint-sqlite>=3.0.3, <4.0.0", @@ -63,6 +63,9 @@ register = "uipath_langchain.middlewares:register_middleware" [project.entry-points."uipath.runtime.factories"] langgraph = "uipath_langchain.runtime:register_runtime_factory" +[project.entry-points."uipath.governance.adapters"] +langchain = "uipath_langchain.governance:register_governance_adapter" + [project.urls] Homepage = "https://uipath.com" Repository = "https://github.com/UiPath/uipath-langchain-python" @@ -154,6 +157,16 @@ exclude_lines = [ [tool.uv] exclude-newer = "2 days" +# uipath-runtime is pinned to a dev pre-release served from testpypi. +# uipath==2.10.79 transitively pins uipath-runtime>=0.11.0,<0.12.0, +# which excludes pre-releases (PEP 440: 0.11.0.dev* sorts below 0.11.0), +# so we need both the prerelease allowance and an explicit override to +# bypass the umbrella's stable-only constraint. +prerelease = "allow" +override-dependencies = ["uipath-runtime>=0.11.0.dev1001180449,<0.11.0.dev1001190000"] + +[tool.uv.sources] +uipath-runtime = { index = "testpypi" } [tool.uv.exclude-newer-package] uipath = false diff --git a/src/uipath_langchain/governance/__init__.py b/src/uipath_langchain/governance/__init__.py new file mode 100644 index 000000000..71b38b02a --- /dev/null +++ b/src/uipath_langchain/governance/__init__.py @@ -0,0 +1,58 @@ +"""Governance integration for ``uipath-langchain``. + +Registers :class:`LangChainAdapter` with the global adapter registry in +``uipath.core.adapters`` so ``uipath.runtime.governance.GovernanceRuntime`` +can attach the LangChain-specific inner hooks (BEFORE_MODEL, +AFTER_MODEL, TOOL_CALL, AFTER_TOOL) when it sees a LangChain or +LangGraph 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_langchain.governance`` is + opted in. + 2. The package also exposes :func:`register_governance_adapter` as an + entry point under ``uipath.governance.adapters`` so an upstream + discoverer (or ``uipath-core`` if/when it grows 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 ( + GovernanceCallbackHandler, + GovernedLangChainAgent, + LangChainAdapter, +) + +logger = logging.getLogger(__name__) + + +def register_governance_adapter() -> None: + """Register :class:`LangChainAdapter` with the global registry. + + Idempotent — safe to call multiple times. + """ + registry = get_adapter_registry() + if any(a.name == "LangChain" for a in registry.get_all()): + return + registry.register(LangChainAdapter()) + logger.debug("Registered uipath-langchain governance adapter") + + +# Side-effect registration on module import. +register_governance_adapter() + + +__all__ = [ + "GovernanceCallbackHandler", + "GovernedLangChainAgent", + "LangChainAdapter", + "register_governance_adapter", +] diff --git a/src/uipath_langchain/governance/adapter.py b/src/uipath_langchain/governance/adapter.py new file mode 100644 index 000000000..ee97c272e --- /dev/null +++ b/src/uipath_langchain/governance/adapter.py @@ -0,0 +1,580 @@ +"""LangChain / LangGraph adapter for UiPath governance. + +Provides governance for LangChain chains/agents and LangGraph compiled +graphs. Uses LangChain's callback system for deep hooks (model / tool +events) plus a thin proxy that ensures the callback is wired into +``invoke`` / ``ainvoke`` / ``stream`` / ``astream``. + +This adapter intercepts: + +- ``on_llm_start`` / ``on_chat_model_start`` / ``on_llm_end`` → BEFORE_MODEL / AFTER_MODEL +- ``on_tool_start`` / ``on_tool_end`` → TOOL_CALL / AFTER_TOOL + +Chain-level boundaries (BEFORE_AGENT / AFTER_AGENT) are intentionally +*not* fired from here — they are owned by the runtime wrapper layer in +``uipath-runtime`` (``GovernanceRuntime.execute`` / ``.stream``). The +``GovernanceCallbackHandler`` sets ``ignore_chain = True`` so LangChain +skips chain notifications entirely, avoiding duplicate boundary +evaluations and silencing AttributeError noise for the absent methods. + +Contracts and the evaluator protocol come from ``uipath-core``; this +package contributes only the LangChain-specific implementation and +self-registers it with the global adapter registry when +``uipath_langchain.governance`` is imported. + +Audit emission and enforcement (raising :class:`GovernanceBlockException` +on DENY) are owned by the evaluator itself. This module just hooks the +framework callbacks, extracts the data, and calls +``evaluator.evaluate_*``; block exceptions propagate, everything else +is logged and swallowed so a governance bug never breaks an agent run. +""" + +from __future__ import annotations + +import logging +from typing import Any, Dict +from uuid import uuid4 + +from uipath.core.adapters import BaseAdapter, EvaluatorProtocol, GovernedAgentBase +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 (see +# ``_GOVERNANCE_TEXT_CAP`` in ``uipath.runtime.governance.wrapper``) so +# scan-time budgets are consistent across hooks. A long conversation +# history is governed at the LLM layer by scanning only the latest +# message, not the full prompt — see +# :meth:`GovernanceCallbackHandler._latest_message_input`. The same cap +# bounds the AFTER_MODEL ``model_output`` blob so batched or runaway +# responses can't blow scan budgets either. +_BEFORE_MODEL_TEXT_CAP = 64000 + + +def _add_callback( + existing: Any, callback: "GovernanceCallbackHandler" +) -> tuple[Any, bool]: + """Add ``callback`` to a LangChain callbacks container. + + LangChain accepts callbacks as ``None``, a ``list`` of handlers, a + tuple/other iterable, or a ``BaseCallbackManager``. We don't want a + hard import on ``langchain_core`` here, so we duck-type: + + - ``BaseCallbackManager`` exposes ``add_handler(handler, inherit)`` and + a ``handlers`` attribute; mutate it in place so any tracers / + handlers already attached to the manager are preserved. + - Lists are appended in place (callers may rely on identity). + - Anything else (``None``, tuple, generic iterable) is coerced to a + fresh list with the new callback appended. + + Returns ``(container, replaced)`` — ``replaced`` indicates whether + the caller should rebind the attribute / config slot to the returned + container (True for the coerced cases, False when we mutated in + place). + """ + if existing is None: + return [callback], True + if isinstance(existing, list): + if callback not in existing: + existing.append(callback) + return existing, False + # BaseCallbackManager: mutate the manager directly so attached + # tracers stay wired up. + if hasattr(existing, "add_handler") and hasattr(existing, "handlers"): + if callback not in (existing.handlers or []): + existing.add_handler(callback, inherit=True) + return existing, False + # Tuple or other iterable — copy to a list we can mutate. + try: + handlers = list(existing) + except TypeError: + handlers = [] + if callback not in handlers: + handlers.append(callback) + return handlers, True + + +class LangChainAdapter(BaseAdapter): + """Adapter for LangChain / LangGraph frameworks. + + Detects and wraps LangChain chains, agents, and LangGraph + ``CompiledStateGraph`` instances with a governance callback handler. + """ + + @property + def name(self) -> str: + return "LangChain" + + def can_handle(self, agent: Any) -> bool: + """Return True if this adapter knows how to hook into the agent.""" + # LangGraph CompiledStateGraph + try: + from langgraph.graph.state import CompiledStateGraph + + if isinstance(agent, CompiledStateGraph): + return True + except ImportError: + pass + + # LangChain Runnable + try: + from langchain_core.runnables import Runnable + + if isinstance(agent, Runnable): + return True + except ImportError: + pass + + # Duck-typed fallback: anything with invoke + ainvoke that isn't + # claimed by another framework adapter (caller-side ordering also + # ensures more specific adapters resolve first). + if hasattr(agent, "invoke") and hasattr(agent, "ainvoke"): + module = getattr(type(agent), "__module__", "") + if any( + fw in module + for fw in ("autogen", "crewai", "llama_index", "pydantic_ai") + ): + return False + return True + + return False + + def attach( + self, + agent: Any, + agent_id: str, + session_id: str, + evaluator: EvaluatorProtocol, + ) -> "GovernedLangChainAgent": + """Attach governance to a LangChain / LangGraph agent.""" + callback = GovernanceCallbackHandler( + evaluator=evaluator, + agent_name=agent_id, + session_id=session_id, + ) + self._inject_callback(agent, callback) + return GovernedLangChainAgent( + agent=agent, + adapter=self, + agent_id=agent_id, + session_id=session_id, + evaluator=evaluator, + callback=callback, + ) + + def _inject_callback( + self, agent: Any, callback: "GovernanceCallbackHandler" + ) -> None: + """Inject the governance callback into the agent's callback chain.""" + if hasattr(agent, "callbacks"): + container, replaced = _add_callback( + getattr(agent, "callbacks", None), callback + ) + if replaced: + agent.callbacks = container + logger.debug("Injected governance callback via agent.callbacks") + return + + if hasattr(agent, "config"): + config = agent.config or {} + container, replaced = _add_callback(config.get("callbacks"), callback) + if replaced: + config["callbacks"] = container + agent.config = config + logger.debug("Injected governance callback via agent.config") + return + + logger.warning( + "Could not inject governance callback into %s — agent has neither " + "'callbacks' nor 'config' surface; deep hooks will not fire", + type(agent).__name__, + ) + + +class GovernedLangChainAgent(GovernedAgentBase): + """LangChain / LangGraph agent wrapped with governance. + + The callback handler does the actual rule evaluation; this proxy + ensures the handler is present on every ``invoke`` / ``ainvoke`` / + ``stream`` / ``astream`` call's config. + """ + + def __init__( + self, + agent: Any, + adapter: LangChainAdapter, + agent_id: str, + session_id: str, + evaluator: EvaluatorProtocol, + callback: "GovernanceCallbackHandler", + ) -> None: + super().__init__(agent, adapter, agent_id, session_id, evaluator) + self._callback = callback + + def invoke(self, input_data: Any, config: Any = None, **kwargs: Any) -> Any: + config = self._ensure_callback_config(config) + return self._agent.invoke(input_data, config=config, **kwargs) + + async def ainvoke(self, input_data: Any, config: Any = None, **kwargs: Any) -> Any: + config = self._ensure_callback_config(config) + return await self._agent.ainvoke(input_data, config=config, **kwargs) + + def stream(self, input_data: Any, config: Any = None, **kwargs: Any) -> Any: + config = self._ensure_callback_config(config) + yield from self._agent.stream(input_data, config=config, **kwargs) + + async def astream(self, input_data: Any, config: Any = None, **kwargs: Any) -> Any: + config = self._ensure_callback_config(config) + async for chunk in self._agent.astream(input_data, config=config, **kwargs): + yield chunk + + def _ensure_callback_config(self, config: Any) -> Dict[str, Any]: + """Ensure the governance callback is on the config's callback list.""" + if config is None: + config = {} + if isinstance(config, dict): + container, replaced = _add_callback(config.get("callbacks"), self._callback) + if replaced: + config["callbacks"] = container + return config + + +class GovernanceCallbackHandler: + """LangChain callback handler that fires governance evaluation. + + Implements the LangChain ``BaseCallbackHandler`` interface shape + structurally (no formal inheritance — keeps this package free of a + hard ``langchain_core`` import at module load). + + The evaluator owns audit emission and DENY-raising. Each ``on_*`` + callback only extracts the relevant payload and calls the matching + ``evaluate_*`` method; :class:`GovernanceBlockException` is allowed + to propagate, anything else is logged and swallowed. + """ + + # LangChain callback-handler descriptors: + run_inline: bool = True + raise_error: bool = False + ignore_llm: bool = False + # Chain-level events are owned by the runtime wrapper layer + # (BEFORE_AGENT / AFTER_AGENT fire from GovernanceRuntime.execute / + # .stream). Telling LangChain to skip chain callbacks here avoids + # duplicate boundary firings AND silences the AttributeError noise + # LangChain would otherwise log for every chain start/end now that + # we don't define the methods. + ignore_chain: bool = True + ignore_agent: bool = False + ignore_retriever: bool = True + ignore_retry: bool = True + ignore_chat_model: bool = False + ignore_custom_event: bool = True + + 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} + # Tool name lookup keyed by LangChain ``run_id`` so ``on_tool_end`` + # can report the actual tool name to AFTER_TOOL evaluation. + self._tool_runs: Dict[str, str] = {} + + # ----- LLM callbacks --------------------------------------------------- + + def on_llm_start( + self, + serialized: Dict[str, Any], + prompts: list[str], + **kwargs: Any, + ) -> None: + """Evaluate BEFORE_MODEL rules at LLM start (non-chat completion).""" + try: + self._session_state["llm_calls"] = ( + self._session_state.get("llm_calls", 0) + 1 + ) + # Take only the latest prompt. Re-scanning every prompt in a + # batched call would re-fire rules on prior turns' content + # that's still in the prompt for context. + model_input = (prompts[-1] if prompts else "")[:_BEFORE_MODEL_TEXT_CAP] + self._evaluator.evaluate_before_model( + model_input=model_input, + agent_name=self._agent_name, + runtime_id=self._session_id, + trace_id=self._trace_id, + ) + except GovernanceBlockException: + raise + except Exception as e: + logger.warning("on_llm_start governance check failed (continuing): %s", e) + + def on_chat_model_start( + self, + serialized: Dict[str, Any], + messages: list[list[Any]], + **kwargs: Any, + ) -> None: + """Evaluate BEFORE_MODEL rules for chat models. + + Scans only the **latest message** in the prompt — not the full + chat history. The LLM still receives the entire history (this + callback doesn't mutate ``messages``), but the governance + evaluator focuses on the new content the agent is about to + respond to. Without this scoping, a violation in turn 3's user + message would keep re-firing on turns 4, 5, 6 ... because that + text stays in the prompt for context. + + List-of-blocks content (multimodal, function-call, tool_use, + extended thinking) is walked via :meth:`_extract_block_text` so + dict-syntax noise from ``str(list)`` doesn't leak into the + regex-scanned blob. + """ + try: + self._session_state["llm_calls"] = ( + self._session_state.get("llm_calls", 0) + 1 + ) + model_input = self._latest_message_input(messages) + self._evaluator.evaluate_before_model( + model_input=model_input, + agent_name=self._agent_name, + runtime_id=self._session_id, + trace_id=self._trace_id, + ) + except GovernanceBlockException: + raise + except Exception as e: + logger.warning( + "on_chat_model_start governance check failed (continuing): %s", e + ) + + @staticmethod + def _latest_message_input(messages: list[list[Any]]) -> str: + """Extract content from the most-recent message in the prompt. + + ``messages`` is LangChain's nested shape ``list[list[BaseMessage]]`` + — the outer list is for batched calls (rare); the inner list is + the full message stack for one call. We take the last entry of + the last inner list. For string content, that's used directly; + for list-of-blocks content, :meth:`_extract_block_text` pulls + the text / arguments / input / thinking fields cleanly. + + Returns ``""`` (empty) when the message stack is empty or the + last message carries no extractable content. + """ + if not messages: + return "" + last_batch = messages[-1] + if not last_batch: + return "" + last_msg = last_batch[-1] + # BaseMessage exposes ``.content``; dict-shaped messages + # (LangGraph state, raw OpenAI format) carry it under the same + # key. + content = getattr(last_msg, "content", None) + if content is None and isinstance(last_msg, dict): + content = last_msg.get("content") + if isinstance(content, str): + return content[:_BEFORE_MODEL_TEXT_CAP] + if isinstance(content, list): + parts: list[str] = [] + remaining = _BEFORE_MODEL_TEXT_CAP + for block in content: + if remaining <= 0: + break + if not isinstance(block, dict): + continue + piece = GovernanceCallbackHandler._extract_block_text(block) + if piece: + parts.append(piece) + remaining -= len(piece) + 1 + return "\n".join(parts)[:_BEFORE_MODEL_TEXT_CAP] + return "" + + def on_llm_end(self, response: Any, **kwargs: Any) -> None: + """Evaluate AFTER_MODEL rules at LLM end. + + Concatenates text from every generation. The result is capped at + ``_BEFORE_MODEL_TEXT_CAP`` to match the BEFORE_MODEL budget and + the runtime side's ``_GOVERNANCE_TEXT_CAP``, so batched calls or + a runaway single response can't blow scan budgets. + """ + try: + parts: list[str] = [] + remaining = _BEFORE_MODEL_TEXT_CAP + if hasattr(response, "generations"): + for gen_list in response.generations: + for gen in gen_list: + if remaining <= 0: + break + piece = self._extract_generation_text(gen) + if piece: + parts.append(piece) + remaining -= len(piece) + if remaining <= 0: + break + model_output = "".join(parts)[:_BEFORE_MODEL_TEXT_CAP] + self._evaluator.evaluate_after_model( + model_output=model_output, + agent_name=self._agent_name, + runtime_id=self._session_id, + trace_id=self._trace_id, + ) + except GovernanceBlockException: + raise + except Exception as e: + logger.warning("on_llm_end governance check failed (continuing): %s", e) + + @staticmethod + def _extract_generation_text(gen: Any) -> str: + """Return the text payload of a LangChain ``Generation`` / ``ChatGeneration``. + + ``Generation.text`` is set from ``message.content`` only when content + is a plain ``str``. For chat models whose content is a list of + content blocks (multimodal, tool calls, "submit final answer" + function calls, extended thinking) ``.text`` is ``""``. Fall back + to walking ``gen.message.content`` so the governance evaluator + sees the actual assistant text. + """ + text = getattr(gen, "text", "") or "" + if text: + return text + message = getattr(gen, "message", None) + if message is None: + return "" + content = getattr(message, "content", None) + if isinstance(content, str): + return content + if isinstance(content, list): + parts = [ + GovernanceCallbackHandler._extract_block_text(block) + for block in content + if isinstance(block, dict) + ] + return "\n".join(p for p in parts if p) + return "" + + @staticmethod + def _extract_block_text(block: Dict[str, Any]) -> str: + """Return any governance-relevant text from a content block. + + Covers the common block shapes across providers: + + - ``{"type": "text", "text": "..."}`` — plain text block. + - ``{"type": "function_call", "arguments": ""}`` — OpenAI + function call; ``arguments`` is JSON-encoded and routinely + carries the user-visible reply (e.g. ``end_execution(content=...)`` + tools used as a "submit final answer" pattern). + - ``{"type": "tool_use", "input": {...}}`` — Anthropic tool use; + string values in ``input`` are the assistant's outgoing payload. + - ``{"type": "thinking", "thinking": "..."}`` — Claude extended + thinking (governance-relevant: hidden reasoning can also leak + commitments and PII). + + Metadata-only keys (``id``, ``call_id``, ``name``, ``status``, + ``type``, ...) are excluded so the scanned text isn't padded with + opaque identifiers that could false-positive a rule. + """ + parts: list[str] = [] + text_value = block.get("text") + if isinstance(text_value, str): + parts.append(text_value) + arguments_value = block.get("arguments") + if isinstance(arguments_value, str): + parts.append(arguments_value) + thinking_value = block.get("thinking") + if isinstance(thinking_value, str): + parts.append(thinking_value) + input_value = block.get("input") + if isinstance(input_value, dict): + parts.extend(v for v in input_value.values() if isinstance(v, str)) + return "\n".join(p for p in parts if p) + + def on_llm_error(self, error: Exception, **kwargs: Any) -> None: + logger.warning("LLM error in governed session %s: %s", self._session_id, error) + + # ----- Tool callbacks -------------------------------------------------- + + def on_tool_start( + self, + serialized: Dict[str, Any], + input_str: str, + *, + inputs: Dict[str, Any] | None = None, + **kwargs: Any, + ) -> None: + """Evaluate TOOL_CALL rules at tool start. + + ``run_id → tool_name`` is recorded so ``on_tool_end`` / + ``on_tool_error`` can report the actual tool. If the evaluator + BLOCKS, the tool is aborted, ``on_tool_end`` will not fire, and + the mapping is dropped to keep ``_tool_runs`` from growing + unbounded across blocked turns. + """ + run_id = kwargs.get("run_id") + run_id_str = str(run_id) if run_id is not None else None + try: + self._session_state["tool_calls"] = ( + self._session_state.get("tool_calls", 0) + 1 + ) + tool_name = (serialized or {}).get("name", "unknown") + if run_id_str is not None: + self._tool_runs[run_id_str] = tool_name + tool_args = inputs or {"input": input_str} + self._evaluator.evaluate_tool_call( + tool_name=tool_name, + tool_args=tool_args, + agent_name=self._agent_name, + runtime_id=self._session_id, + trace_id=self._trace_id, + session_state=self._session_state, + ) + except GovernanceBlockException: + # Tool will not run → no on_tool_end is coming. Drop the + # mapping so it does not accumulate across blocked turns. + if run_id_str is not None: + self._tool_runs.pop(run_id_str, None) + raise + except Exception as e: + logger.warning("on_tool_start governance check failed (continuing): %s", e) + + def on_tool_end(self, output: Any, **kwargs: Any) -> None: + """Evaluate AFTER_TOOL rules at tool end.""" + try: + run_id = kwargs.get("run_id") + tool_name = "unknown" + if run_id is not None: + tool_name = self._tool_runs.pop(str(run_id), "unknown") + tool_result = str(output) if output is not None else "" + self._evaluator.evaluate_after_tool( + tool_name=tool_name, + tool_result=tool_result, + agent_name=self._agent_name, + runtime_id=self._session_id, + trace_id=self._trace_id, + ) + except GovernanceBlockException: + raise + except Exception as e: + logger.warning("on_tool_end governance check failed (continuing): %s", e) + + def on_tool_error(self, error: Exception, **kwargs: Any) -> None: + # Tool errored out — on_tool_end will not fire. Pop the mapping + # so a session with many failing tool calls does not leak. + run_id = kwargs.get("run_id") + if run_id is not None: + self._tool_runs.pop(str(run_id), None) + logger.warning("Tool error in governed session %s: %s", self._session_id, error) + + # Chain-level callbacks (on_chain_start / on_chain_end / on_chain_error) + # are intentionally NOT implemented here. The runtime wrapper + # (``GovernanceRuntime.execute`` / ``GovernanceRuntime.stream`` in + # ``uipath-runtime``) owns BEFORE_AGENT / AFTER_AGENT — firing them + # here too would duplicate every boundary evaluation. The + # ``ignore_chain = True`` class-level descriptor above tells + # LangChain to skip chain notifications entirely so we don't get + # AttributeError warnings for the absent methods. diff --git a/tests/governance/__init__.py b/tests/governance/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/governance/test_adapter.py b/tests/governance/test_adapter.py new file mode 100644 index 000000000..93aa24611 --- /dev/null +++ b/tests/governance/test_adapter.py @@ -0,0 +1,1065 @@ +"""Tests for the LangChain governance adapter.""" + +from __future__ import annotations + +import sys +from types import SimpleNamespace +from typing import Any, TypedDict +from unittest.mock import MagicMock, patch + +import pytest +from uipath.core.adapters import get_adapter_registry, reset_adapter_registry +from uipath.core.governance.exceptions import GovernanceBlockException + +from uipath_langchain.governance import register_governance_adapter +from uipath_langchain.governance.adapter import ( + GovernanceCallbackHandler, + GovernedLangChainAgent, + LangChainAdapter, + _add_callback, +) + +LOGGER_PATH = "uipath_langchain.governance.adapter.logger" + + +@pytest.fixture +def evaluator() -> MagicMock: + return MagicMock() + + +@pytest.fixture +def adapter() -> LangChainAdapter: + return LangChainAdapter() + + +@pytest.fixture +def handler(evaluator: MagicMock) -> GovernanceCallbackHandler: + return GovernanceCallbackHandler( + evaluator=evaluator, + agent_name="test-agent", + session_id="test-session", + ) + + +class TestCanHandle: + def test_returns_true_for_langgraph_compiled_state_graph( + self, adapter: LangChainAdapter + ) -> None: + from langgraph.graph import StateGraph + + class S(TypedDict): + v: int + + graph = StateGraph(S) + graph.add_node("n", lambda s: s) + graph.set_entry_point("n") + compiled = graph.compile() + assert adapter.can_handle(compiled) is True + + def test_returns_true_for_langchain_runnable( + self, adapter: LangChainAdapter + ) -> None: + from langchain_core.runnables import RunnableLambda + + runnable = RunnableLambda(lambda x: x) + assert adapter.can_handle(runnable) is True + + def test_returns_true_for_duck_typed_invoke_ainvoke( + self, adapter: LangChainAdapter + ) -> None: + class Duck: + def invoke(self, x): + return x + + async def ainvoke(self, x): + return x + + assert adapter.can_handle(Duck()) is True + + @pytest.mark.parametrize( + "module_name", + ["autogen.foo", "crewai.agent", "llama_index.core", "pydantic_ai.bar"], + ) + def test_returns_false_for_excluded_frameworks( + self, adapter: LangChainAdapter, module_name: str + ) -> None: + class Foreign: + def invoke(self, x): + return x + + async def ainvoke(self, x): + return x + + Foreign.__module__ = module_name + assert adapter.can_handle(Foreign()) is False + + def test_returns_false_for_object_without_invoke( + self, adapter: LangChainAdapter + ) -> None: + assert adapter.can_handle(object()) is False + + def test_returns_false_for_object_with_only_invoke( + self, adapter: LangChainAdapter + ) -> None: + class Half: + def invoke(self, x): + return x + + assert adapter.can_handle(Half()) is False + + def test_handles_langgraph_import_failure( + self, monkeypatch: pytest.MonkeyPatch, adapter: LangChainAdapter + ) -> None: + monkeypatch.setitem(sys.modules, "langgraph.graph.state", None) + from langchain_core.runnables import RunnableLambda + + assert adapter.can_handle(RunnableLambda(lambda x: x)) is True + + def test_handles_langchain_core_import_failure( + self, monkeypatch: pytest.MonkeyPatch, adapter: LangChainAdapter + ) -> None: + monkeypatch.setitem(sys.modules, "langgraph.graph.state", None) + monkeypatch.setitem(sys.modules, "langchain_core.runnables", None) + + class Duck: + def invoke(self, x): + return x + + async def ainvoke(self, x): + return x + + assert adapter.can_handle(Duck()) is True + + +class TestAttach: + def test_returns_governed_agent( + self, adapter: LangChainAdapter, evaluator: MagicMock + ) -> None: + class Agent: + callbacks: list[Any] = [] + + a = Agent() + governed = adapter.attach(a, "agent-id", "session-id", evaluator) + assert isinstance(governed, GovernedLangChainAgent) + assert governed.unwrapped is a + + def test_injects_callback_into_existing_callback_list( + self, adapter: LangChainAdapter, evaluator: MagicMock + ) -> None: + prior = object() + + class Agent: + callbacks: list[Any] = [prior] + + a = Agent() + governed = adapter.attach(a, "id", "session", evaluator) + assert a.callbacks[0] is prior + assert a.callbacks[-1] is governed._callback + + def test_replaces_non_list_callbacks_attribute( + self, adapter: LangChainAdapter, evaluator: MagicMock + ) -> None: + sentinel = object() + + class Agent: + callbacks = sentinel # truthy but not a list + + a = Agent() + governed = adapter.attach(a, "id", "session", evaluator) + assert isinstance(a.callbacks, list) + assert a.callbacks == [governed._callback] + + def test_injects_into_empty_config( + self, adapter: LangChainAdapter, evaluator: MagicMock + ) -> None: + class Agent: + config: dict[str, Any] = {} + + a = Agent() + governed = adapter.attach(a, "id", "session", evaluator) + assert a.config["callbacks"] == [governed._callback] + + def test_injects_into_none_config( + self, adapter: LangChainAdapter, evaluator: MagicMock + ) -> None: + class Agent: + config = None + + a = Agent() + governed = adapter.attach(a, "id", "session", evaluator) + assert isinstance(a.config, dict) + assert a.config["callbacks"] == [governed._callback] + + def test_logs_warning_when_no_callback_surface( + self, + adapter: LangChainAdapter, + evaluator: MagicMock, + ) -> None: + class Bare: + pass + + bare = Bare() + with patch(LOGGER_PATH) as mock_logger: + adapter.attach(bare, "id", "session", evaluator) + mock_logger.warning.assert_called_once() + assert "Could not inject" in mock_logger.warning.call_args.args[0] + assert not hasattr(bare, "callbacks") + assert not hasattr(bare, "config") + + def test_callbacks_path_logs_debug( + self, + adapter: LangChainAdapter, + evaluator: MagicMock, + ) -> None: + class Agent: + callbacks: list[Any] = [] + + with patch(LOGGER_PATH) as mock_logger: + adapter.attach(Agent(), "id", "session", evaluator) + debug_msgs = [c.args[0] for c in mock_logger.debug.call_args_list] + assert any("agent.callbacks" in m for m in debug_msgs) + + def test_config_path_logs_debug( + self, + adapter: LangChainAdapter, + evaluator: MagicMock, + ) -> None: + class Agent: + config: dict[str, Any] = {} + + with patch(LOGGER_PATH) as mock_logger: + adapter.attach(Agent(), "id", "session", evaluator) + debug_msgs = [c.args[0] for c in mock_logger.debug.call_args_list] + assert any("agent.config" in m for m in debug_msgs) + + +class TestAdapterMetadata: + def test_name(self, adapter: LangChainAdapter) -> None: + assert adapter.name == "LangChain" + + +class TestGovernedLangChainAgent: + def _governed( + self, adapter: LangChainAdapter, evaluator: MagicMock, agent: object + ) -> GovernedLangChainAgent: + return adapter.attach(agent, "id", "session", evaluator) + + def test_invoke_injects_callback_and_returns_result( + self, adapter: LangChainAdapter, evaluator: MagicMock + ) -> None: + captured: dict[str, Any] = {} + + class Agent: + callbacks: list[Any] = [] + + def invoke(self, x, config=None, **kw): + captured["config"] = config + captured["x"] = x + return "out" + + gov = self._governed(adapter, evaluator, Agent()) + assert gov.invoke("in") == "out" + assert captured["x"] == "in" + assert gov._callback in captured["config"]["callbacks"] + + async def test_ainvoke_injects_callback_and_returns_result( + self, adapter: LangChainAdapter, evaluator: MagicMock + ) -> None: + captured: dict[str, Any] = {} + + class Agent: + callbacks: list[Any] = [] + + async def ainvoke(self, x, config=None, **kw): + captured["config"] = config + return "async-out" + + gov = self._governed(adapter, evaluator, Agent()) + assert await gov.ainvoke("in") == "async-out" + assert gov._callback in captured["config"]["callbacks"] + + def test_stream_yields_chunks_with_callback( + self, adapter: LangChainAdapter, evaluator: MagicMock + ) -> None: + captured: dict[str, Any] = {} + + class Agent: + callbacks: list[Any] = [] + + def stream(self, x, config=None, **kw): + captured["config"] = config + yield "a" + yield "b" + + gov = self._governed(adapter, evaluator, Agent()) + assert list(gov.stream("x")) == ["a", "b"] + assert gov._callback in captured["config"]["callbacks"] + + async def test_astream_yields_chunks_with_callback( + self, adapter: LangChainAdapter, evaluator: MagicMock + ) -> None: + captured: dict[str, Any] = {} + + class Agent: + callbacks: list[Any] = [] + + async def astream(self, x, config=None, **kw): + captured["config"] = config + yield "a" + yield "b" + + gov = self._governed(adapter, evaluator, Agent()) + chunks = [c async for c in gov.astream("x")] + assert chunks == ["a", "b"] + assert gov._callback in captured["config"]["callbacks"] + + def test_ensure_callback_config_with_none( + self, adapter: LangChainAdapter, evaluator: MagicMock + ) -> None: + gov = self._governed(adapter, evaluator, SimpleNamespace(callbacks=[])) + result = gov._ensure_callback_config(None) + assert isinstance(result, dict) + assert result["callbacks"] == [gov._callback] + + def test_ensure_callback_config_preserves_other_keys( + self, adapter: LangChainAdapter, evaluator: MagicMock + ) -> None: + gov = self._governed(adapter, evaluator, SimpleNamespace(callbacks=[])) + config = {"metadata": {"k": "v"}, "callbacks": []} + result = gov._ensure_callback_config(config) + assert result is config + assert result["metadata"] == {"k": "v"} + assert gov._callback in result["callbacks"] + + def test_ensure_callback_config_is_idempotent( + self, adapter: LangChainAdapter, evaluator: MagicMock + ) -> None: + gov = self._governed(adapter, evaluator, SimpleNamespace(callbacks=[])) + first = gov._ensure_callback_config(None) + second = gov._ensure_callback_config(first) + assert second["callbacks"].count(gov._callback) == 1 + + def test_ensure_callback_config_passes_through_non_dict( + self, adapter: LangChainAdapter, evaluator: MagicMock + ) -> None: + gov = self._governed(adapter, evaluator, SimpleNamespace(callbacks=[])) + sentinel = ["not", "a", "dict"] + result = gov._ensure_callback_config(sentinel) + assert result is sentinel + + def test_getattr_forwards_to_wrapped_agent( + self, adapter: LangChainAdapter, evaluator: MagicMock + ) -> None: + class Agent: + callbacks: list[Any] = [] + answer = 42 + + gov = self._governed(adapter, evaluator, Agent()) + assert gov.answer == 42 + + +class TestCallbackHandlerLLM: + def test_on_llm_start_invokes_evaluator_with_latest_prompt( + self, handler: GovernanceCallbackHandler, evaluator: MagicMock + ) -> None: + """Only the latest prompt feeds BEFORE_MODEL — prior prompts in a + batched call would re-fire rules on content the LLM has + already responded to in earlier batches.""" + handler.on_llm_start({"name": "m"}, ["a", "b"]) + evaluator.evaluate_before_model.assert_called_once() + kwargs = evaluator.evaluate_before_model.call_args.kwargs + assert kwargs["model_input"] == "b" + assert kwargs["agent_name"] == "test-agent" + assert kwargs["runtime_id"] == "test-session" + assert kwargs["trace_id"] == handler._trace_id + + def test_on_llm_start_increments_counter( + self, handler: GovernanceCallbackHandler + ) -> None: + handler.on_llm_start({}, ["p"]) + handler.on_llm_start({}, ["p"]) + assert handler._session_state["llm_calls"] == 2 + + def test_on_llm_start_empty_prompts( + self, handler: GovernanceCallbackHandler, evaluator: MagicMock + ) -> None: + handler.on_llm_start({}, []) + assert evaluator.evaluate_before_model.call_args.kwargs["model_input"] == "" + + def test_on_llm_start_propagates_block( + self, handler: GovernanceCallbackHandler, evaluator: MagicMock + ) -> None: + evaluator.evaluate_before_model.side_effect = GovernanceBlockException( + "blocked" + ) + with pytest.raises(GovernanceBlockException): + handler.on_llm_start({}, ["p"]) + + def test_on_llm_start_swallows_other_exceptions( + self, + handler: GovernanceCallbackHandler, + evaluator: MagicMock, + ) -> None: + evaluator.evaluate_before_model.side_effect = RuntimeError("nope") + with patch(LOGGER_PATH) as mock_logger: + handler.on_llm_start({}, ["p"]) # must not raise + mock_logger.warning.assert_called_once() + assert "on_llm_start" in mock_logger.warning.call_args.args[0] + + def test_on_chat_model_start_latest_message_only( + self, handler: GovernanceCallbackHandler, evaluator: MagicMock + ) -> None: + """Only the LAST message in the prompt is scanned. + + Without this scoping, a violation in turn 3's user message + would keep re-firing on every subsequent LLM call because + that text stays in the prompt for context. + """ + handler.on_chat_model_start( + {}, + [[SimpleNamespace(content="hello"), SimpleNamespace(content="world")]], + ) + model_input = evaluator.evaluate_before_model.call_args.kwargs["model_input"] + assert model_input == "world" + assert "hello" not in model_input + + def test_on_chat_model_start_dict_messages_latest_only( + self, handler: GovernanceCallbackHandler, evaluator: MagicMock + ) -> None: + """Dict-shaped (LangGraph state) messages: latest is extracted.""" + handler.on_chat_model_start( + {}, + [[{"content": "from dict"}, {"role": "user", "content": "another"}]], + ) + model_input = evaluator.evaluate_before_model.call_args.kwargs["model_input"] + assert model_input == "another" + assert "from dict" not in model_input + + def test_on_chat_model_start_dict_message_missing_content( + self, handler: GovernanceCallbackHandler, evaluator: MagicMock + ) -> None: + handler.on_chat_model_start({}, [[{"role": "user"}]]) + assert evaluator.evaluate_before_model.call_args.kwargs["model_input"] == "" + + def test_on_chat_model_start_list_of_blocks_content( + self, handler: GovernanceCallbackHandler, evaluator: MagicMock + ) -> None: + """Multi-block content (text + function_call) is extracted cleanly. + + Regression for the prior ``str(msg.content)`` path which produced + ``[{'type': ..., 'text': ...}]`` dict-repr noise instead of + clean text. Field-precise rules can't navigate that shape. + """ + latest = SimpleNamespace( + content=[ + {"type": "text", "text": "Here's the answer:"}, + { + "type": "function_call", + "name": "end_execution", + "arguments": '{"content":"Cost: $1,200"}', + "id": "fc_abc", + }, + ] + ) + handler.on_chat_model_start({}, [[SimpleNamespace(content="old"), latest]]) + model_input = evaluator.evaluate_before_model.call_args.kwargs["model_input"] + assert "Here's the answer:" in model_input + assert "Cost: $1,200" in model_input + # No dict-syntax noise from str(list). + assert "{'type'" not in model_input + + def test_on_chat_model_start_empty_messages( + self, handler: GovernanceCallbackHandler, evaluator: MagicMock + ) -> None: + handler.on_chat_model_start({}, []) + assert evaluator.evaluate_before_model.call_args.kwargs["model_input"] == "" + + def test_on_chat_model_start_empty_inner_batch( + self, handler: GovernanceCallbackHandler, evaluator: MagicMock + ) -> None: + handler.on_chat_model_start({}, [[]]) + assert evaluator.evaluate_before_model.call_args.kwargs["model_input"] == "" + + def test_on_chat_model_start_caps_model_input( + self, handler: GovernanceCallbackHandler, evaluator: MagicMock + ) -> None: + """``model_input`` is bounded so a runaway prompt can't dominate scan time.""" + from uipath_langchain.governance.adapter import _BEFORE_MODEL_TEXT_CAP + + huge = SimpleNamespace(content="x" * (_BEFORE_MODEL_TEXT_CAP + 1000)) + handler.on_chat_model_start({}, [[huge]]) + model_input = evaluator.evaluate_before_model.call_args.kwargs["model_input"] + assert len(model_input) == _BEFORE_MODEL_TEXT_CAP + + def test_on_chat_model_start_block_list_stops_at_remaining_budget( + self, handler: GovernanceCallbackHandler, evaluator: MagicMock + ) -> None: + """The block walk exits early once the per-call cap is exhausted.""" + from uipath_langchain.governance.adapter import _BEFORE_MODEL_TEXT_CAP + + first = "a" * _BEFORE_MODEL_TEXT_CAP # consumes the entire budget + latest = SimpleNamespace( + content=[ + {"type": "text", "text": first}, + {"type": "text", "text": "MUST_NOT_APPEAR"}, + ] + ) + handler.on_chat_model_start({}, [[latest]]) + model_input = evaluator.evaluate_before_model.call_args.kwargs["model_input"] + assert "MUST_NOT_APPEAR" not in model_input + assert len(model_input) == _BEFORE_MODEL_TEXT_CAP + + def test_on_chat_model_start_block_list_skips_non_dict_entries( + self, handler: GovernanceCallbackHandler, evaluator: MagicMock + ) -> None: + """Non-dict entries inside a content list are silently skipped.""" + latest = SimpleNamespace( + content=[ + "ignored-string-block", + {"type": "text", "text": "kept"}, + 42, + None, + ] + ) + handler.on_chat_model_start({}, [[latest]]) + model_input = evaluator.evaluate_before_model.call_args.kwargs["model_input"] + assert model_input == "kept" + + def test_on_chat_model_start_propagates_block( + self, handler: GovernanceCallbackHandler, evaluator: MagicMock + ) -> None: + evaluator.evaluate_before_model.side_effect = GovernanceBlockException("x") + with pytest.raises(GovernanceBlockException): + handler.on_chat_model_start({}, [[SimpleNamespace(content="x")]]) + + def test_on_chat_model_start_swallows_other_exceptions( + self, + handler: GovernanceCallbackHandler, + evaluator: MagicMock, + ) -> None: + evaluator.evaluate_before_model.side_effect = RuntimeError("oops") + with patch(LOGGER_PATH) as mock_logger: + handler.on_chat_model_start({}, [[SimpleNamespace(content="x")]]) + mock_logger.warning.assert_called_once() + assert "on_chat_model_start" in mock_logger.warning.call_args.args[0] + + def test_on_llm_end_extracts_plain_text( + self, handler: GovernanceCallbackHandler, evaluator: MagicMock + ) -> None: + gen = SimpleNamespace(text="output", message=None) + response = SimpleNamespace(generations=[[gen]]) + handler.on_llm_end(response) + kwargs = evaluator.evaluate_after_model.call_args.kwargs + assert kwargs["model_output"] == "output" + + def test_on_llm_end_response_without_generations( + self, handler: GovernanceCallbackHandler, evaluator: MagicMock + ) -> None: + handler.on_llm_end(SimpleNamespace()) + assert evaluator.evaluate_after_model.call_args.kwargs["model_output"] == "" + + def test_on_llm_end_propagates_block( + self, handler: GovernanceCallbackHandler, evaluator: MagicMock + ) -> None: + evaluator.evaluate_after_model.side_effect = GovernanceBlockException("x") + with pytest.raises(GovernanceBlockException): + handler.on_llm_end(SimpleNamespace(generations=[])) + + def test_on_llm_end_caps_model_output( + self, handler: GovernanceCallbackHandler, evaluator: MagicMock + ) -> None: + """A runaway / batched response is capped so the AFTER_MODEL + scan budget matches BEFORE_MODEL and the runtime side's cap. + """ + from uipath_langchain.governance.adapter import _BEFORE_MODEL_TEXT_CAP + + # Many large generations across batched gen_lists. + gen = SimpleNamespace(text="y" * 50_000, message=None) + response = SimpleNamespace(generations=[[gen], [gen, gen]]) + handler.on_llm_end(response) + model_output = evaluator.evaluate_after_model.call_args.kwargs["model_output"] + assert len(model_output) == _BEFORE_MODEL_TEXT_CAP + + def test_on_llm_end_skips_empty_generation_text( + self, handler: GovernanceCallbackHandler, evaluator: MagicMock + ) -> None: + """Generations with no extractable text don't bloat the output.""" + empty = SimpleNamespace(text="", message=None) + keep = SimpleNamespace(text="kept", message=None) + response = SimpleNamespace(generations=[[empty, keep]]) + handler.on_llm_end(response) + assert evaluator.evaluate_after_model.call_args.kwargs["model_output"] == "kept" + + def test_on_llm_end_swallows_other_exceptions( + self, + handler: GovernanceCallbackHandler, + evaluator: MagicMock, + ) -> None: + evaluator.evaluate_after_model.side_effect = RuntimeError("nope") + with patch(LOGGER_PATH) as mock_logger: + handler.on_llm_end(SimpleNamespace()) + mock_logger.warning.assert_called_once() + assert "on_llm_end" in mock_logger.warning.call_args.args[0] + + def test_on_llm_error_logs( + self, + handler: GovernanceCallbackHandler, + ) -> None: + with patch(LOGGER_PATH) as mock_logger: + handler.on_llm_error(RuntimeError("boom")) + mock_logger.warning.assert_called_once() + assert "LLM error" in mock_logger.warning.call_args.args[0] + + +class TestExtractGenerationText: + def test_returns_text_when_present(self) -> None: + gen = SimpleNamespace(text="hello", message=None) + assert GovernanceCallbackHandler._extract_generation_text(gen) == "hello" + + def test_falls_back_to_message_string_content(self) -> None: + gen = SimpleNamespace(text="", message=SimpleNamespace(content="rich")) + assert GovernanceCallbackHandler._extract_generation_text(gen) == "rich" + + def test_returns_empty_when_message_missing(self) -> None: + class G: + text = "" + + assert GovernanceCallbackHandler._extract_generation_text(G()) == "" + + def test_returns_empty_when_message_is_none(self) -> None: + gen = SimpleNamespace(text="", message=None) + assert GovernanceCallbackHandler._extract_generation_text(gen) == "" + + def test_extracts_from_block_list_content(self) -> None: + gen = SimpleNamespace( + text="", + message=SimpleNamespace( + content=[ + {"type": "text", "text": "alpha"}, + {"type": "tool_use", "input": {"q": "beta"}}, + ] + ), + ) + out = GovernanceCallbackHandler._extract_generation_text(gen) + assert "alpha" in out + assert "beta" in out + + def test_block_list_skips_non_dict_entries(self) -> None: + gen = SimpleNamespace( + text="", + message=SimpleNamespace( + content=["string-entry", {"type": "text", "text": "kept"}] + ), + ) + assert GovernanceCallbackHandler._extract_generation_text(gen) == "kept" + + def test_unknown_content_shape_returns_empty(self) -> None: + gen = SimpleNamespace(text="", message=SimpleNamespace(content=123)) + assert GovernanceCallbackHandler._extract_generation_text(gen) == "" + + +class TestExtractBlockText: + def test_plain_text_block(self) -> None: + assert ( + GovernanceCallbackHandler._extract_block_text( + {"type": "text", "text": "hello"} + ) + == "hello" + ) + + def test_function_call_arguments_block(self) -> None: + assert ( + GovernanceCallbackHandler._extract_block_text( + {"type": "function_call", "arguments": '{"a":1}'} + ) + == '{"a":1}' + ) + + def test_thinking_block(self) -> None: + assert ( + GovernanceCallbackHandler._extract_block_text( + {"type": "thinking", "thinking": "step by step"} + ) + == "step by step" + ) + + def test_tool_use_input_extracts_string_values(self) -> None: + result = GovernanceCallbackHandler._extract_block_text( + {"type": "tool_use", "input": {"query": "search", "id": "ignored"}} + ) + assert "search" in result + assert "ignored" in result # both are strings; metadata filtering is by key + + def test_input_ignores_non_string_values(self) -> None: + result = GovernanceCallbackHandler._extract_block_text( + {"input": {"a": 123, "b": ["nested"], "c": "kept"}} + ) + assert result == "kept" + + def test_metadata_only_block_returns_empty(self) -> None: + assert ( + GovernanceCallbackHandler._extract_block_text( + {"type": "tool_use", "id": "abc", "name": "search", "status": "ok"} + ) + == "" + ) + + def test_combined_fields_all_collected(self) -> None: + result = GovernanceCallbackHandler._extract_block_text( + { + "type": "tool_use", + "text": "T", + "arguments": "A", + "thinking": "Th", + "input": {"k": "I"}, + } + ) + for token in ("T", "A", "Th", "I"): + assert token in result + + def test_empty_block(self) -> None: + assert GovernanceCallbackHandler._extract_block_text({}) == "" + + +class TestCallbackHandlerTools: + def test_on_tool_start_with_inputs( + self, handler: GovernanceCallbackHandler, evaluator: MagicMock + ) -> None: + handler.on_tool_start({"name": "search"}, "fallback", inputs={"q": "v"}) + kwargs = evaluator.evaluate_tool_call.call_args.kwargs + assert kwargs["tool_name"] == "search" + assert kwargs["tool_args"] == {"q": "v"} + assert kwargs["session_state"] is handler._session_state + + def test_on_tool_start_without_inputs_uses_input_str( + self, handler: GovernanceCallbackHandler, evaluator: MagicMock + ) -> None: + handler.on_tool_start({"name": "calc"}, "1+2") + kwargs = evaluator.evaluate_tool_call.call_args.kwargs + assert kwargs["tool_args"] == {"input": "1+2"} + + def test_on_tool_start_unknown_name_when_missing( + self, handler: GovernanceCallbackHandler, evaluator: MagicMock + ) -> None: + handler.on_tool_start({}, "x") + assert evaluator.evaluate_tool_call.call_args.kwargs["tool_name"] == "unknown" + + def test_on_tool_start_increments_counter( + self, handler: GovernanceCallbackHandler + ) -> None: + handler.on_tool_start({}, "x") + handler.on_tool_start({}, "y") + assert handler._session_state["tool_calls"] == 2 + + def test_on_tool_start_propagates_block( + self, handler: GovernanceCallbackHandler, evaluator: MagicMock + ) -> None: + evaluator.evaluate_tool_call.side_effect = GovernanceBlockException("no") + with pytest.raises(GovernanceBlockException): + handler.on_tool_start({}, "x") + + def test_on_tool_start_swallows_other_exceptions( + self, + handler: GovernanceCallbackHandler, + evaluator: MagicMock, + ) -> None: + evaluator.evaluate_tool_call.side_effect = RuntimeError("nope") + with patch(LOGGER_PATH) as mock_logger: + handler.on_tool_start({}, "x") + mock_logger.warning.assert_called_once() + assert "on_tool_start" in mock_logger.warning.call_args.args[0] + + def test_on_tool_end_with_output( + self, handler: GovernanceCallbackHandler, evaluator: MagicMock + ) -> None: + handler.on_tool_end({"answer": 42}) + kwargs = evaluator.evaluate_after_tool.call_args.kwargs + assert "42" in kwargs["tool_result"] + assert kwargs["tool_name"] == "unknown" + + def test_on_tool_end_uses_tool_name_from_run_id( + self, handler: GovernanceCallbackHandler, evaluator: MagicMock + ) -> None: + handler.on_tool_start({"name": "search"}, "q", run_id="run-1") + handler.on_tool_end("result", run_id="run-1") + assert evaluator.evaluate_after_tool.call_args.kwargs["tool_name"] == "search" + # The run_id mapping is cleaned up so a stale entry isn't reused. + assert "run-1" not in handler._tool_runs + + def test_on_tool_end_unknown_when_run_id_not_recorded( + self, handler: GovernanceCallbackHandler, evaluator: MagicMock + ) -> None: + handler.on_tool_end("r", run_id="never-started") + assert evaluator.evaluate_after_tool.call_args.kwargs["tool_name"] == "unknown" + + def test_on_tool_start_handles_none_serialized( + self, handler: GovernanceCallbackHandler, evaluator: MagicMock + ) -> None: + handler.on_tool_start(None, "x") # type: ignore[arg-type] + assert evaluator.evaluate_tool_call.call_args.kwargs["tool_name"] == "unknown" + + def test_on_tool_end_with_none_output( + self, handler: GovernanceCallbackHandler, evaluator: MagicMock + ) -> None: + handler.on_tool_end(None) + assert evaluator.evaluate_after_tool.call_args.kwargs["tool_result"] == "" + + def test_on_tool_end_propagates_block( + self, handler: GovernanceCallbackHandler, evaluator: MagicMock + ) -> None: + evaluator.evaluate_after_tool.side_effect = GovernanceBlockException("x") + with pytest.raises(GovernanceBlockException): + handler.on_tool_end("out") + + def test_on_tool_end_swallows_other_exceptions( + self, + handler: GovernanceCallbackHandler, + evaluator: MagicMock, + ) -> None: + evaluator.evaluate_after_tool.side_effect = RuntimeError("err") + with patch(LOGGER_PATH) as mock_logger: + handler.on_tool_end("out") + mock_logger.warning.assert_called_once() + assert "on_tool_end" in mock_logger.warning.call_args.args[0] + + def test_on_tool_error_logs( + self, + handler: GovernanceCallbackHandler, + ) -> None: + with patch(LOGGER_PATH) as mock_logger: + handler.on_tool_error(RuntimeError("broke")) + mock_logger.warning.assert_called_once() + assert "Tool error" in mock_logger.warning.call_args.args[0] + + def test_on_tool_error_pops_run_id_mapping( + self, handler: GovernanceCallbackHandler + ) -> None: + """``on_tool_error`` cleans up ``_tool_runs`` so failed tool calls + don't accumulate over the lifetime of a governed session. + """ + handler.on_tool_start({"name": "search"}, "q", run_id="run-err") + assert handler._tool_runs.get("run-err") == "search" + handler.on_tool_error(RuntimeError("boom"), run_id="run-err") + assert "run-err" not in handler._tool_runs + + def test_on_tool_error_without_run_id_does_not_crash( + self, handler: GovernanceCallbackHandler + ) -> None: + # No run_id kwargs — should still log and not raise. + handler.on_tool_error(RuntimeError("boom")) + assert handler._tool_runs == {} + + def test_on_tool_start_block_pops_run_id_mapping( + self, handler: GovernanceCallbackHandler, evaluator: MagicMock + ) -> None: + """If BEFORE_TOOL evaluation BLOCKS, the recorded mapping is + dropped — the tool never runs and ``on_tool_end`` will not fire. + Leaving the entry would leak across blocked turns. + """ + evaluator.evaluate_tool_call.side_effect = GovernanceBlockException("nope") + with pytest.raises(GovernanceBlockException): + handler.on_tool_start({"name": "search"}, "q", run_id="run-blocked") + assert "run-blocked" not in handler._tool_runs + + def test_on_tool_start_swallowed_error_preserves_mapping( + self, handler: GovernanceCallbackHandler, evaluator: MagicMock + ) -> None: + """When the evaluator raises a non-block exception, we swallow + and the tool still runs — the mapping must survive so + ``on_tool_end`` can resolve the tool name. + """ + evaluator.evaluate_tool_call.side_effect = RuntimeError("flaky") + with patch(LOGGER_PATH): + handler.on_tool_start({"name": "search"}, "q", run_id="run-flaky") + assert handler._tool_runs.get("run-flaky") == "search" + + +class TestCallbackHandlerInit: + def test_session_state_initialized(self, evaluator: MagicMock) -> None: + h = GovernanceCallbackHandler( + evaluator=evaluator, agent_name="a", session_id="s" + ) + assert h._session_state == {"tool_calls": 0, "llm_calls": 0} + assert h._agent_name == "a" + assert h._session_id == "s" + assert h._trace_id # uuid4 string + + +class TestRegisterAdapter: + def setup_method(self, method) -> None: + reset_adapter_registry() + + def teardown_method(self, method) -> None: + reset_adapter_registry() + register_governance_adapter() + + def test_registers_adapter(self) -> None: + register_governance_adapter() + names = [a.name for a in get_adapter_registry().get_all()] + assert names.count("LangChain") == 1 + + def test_is_idempotent(self) -> None: + register_governance_adapter() + register_governance_adapter() + register_governance_adapter() + names = [a.name for a in get_adapter_registry().get_all()] + assert names.count("LangChain") == 1 + + def test_skips_when_adapter_with_same_name_already_present(self) -> None: + registry = get_adapter_registry() + # Drop anything the entry-point discovery may have added. + registry.clear() + + class OtherLang(LangChainAdapter): + pass + + registry.register(OtherLang()) + + register_governance_adapter() + + names = [a.name for a in registry.get_all()] + assert names.count("LangChain") == 1 + # The pre-existing instance was kept (not replaced). + assert isinstance(registry.get_all()[0], OtherLang) + + +class TestAddCallback: + """Direct tests for the ``_add_callback`` shape-normalizer.""" + + def test_none_returns_new_list(self) -> None: + cb = object() + container, replaced = _add_callback(None, cb) # type: ignore[arg-type] + assert container == [cb] + assert replaced is True + + def test_list_is_mutated_in_place(self) -> None: + cb = object() + existing: list[Any] = [] + container, replaced = _add_callback(existing, cb) # type: ignore[arg-type] + assert container is existing + assert existing == [cb] + assert replaced is False + + def test_list_is_idempotent(self) -> None: + cb = object() + existing: list[Any] = [cb] + container, replaced = _add_callback(existing, cb) # type: ignore[arg-type] + assert existing.count(cb) == 1 + assert replaced is False + + def test_tuple_is_coerced_to_list(self) -> None: + prior = object() + cb = object() + container, replaced = _add_callback((prior,), cb) # type: ignore[arg-type] + assert container == [prior, cb] + assert replaced is True + + def test_base_callback_manager_is_mutated_in_place(self) -> None: + """A duck-typed callback manager keeps its identity and tracers.""" + + class FakeManager: + def __init__(self) -> None: + self.handlers: list[Any] = [] + self.inherit_calls: list[bool] = [] + + def add_handler(self, handler: Any, inherit: bool) -> None: + self.handlers.append(handler) + self.inherit_calls.append(inherit) + + manager = FakeManager() + prior = object() + manager.handlers.append(prior) + + cb = object() + container, replaced = _add_callback(manager, cb) # type: ignore[arg-type] + assert container is manager + assert replaced is False + assert manager.handlers == [prior, cb] + assert manager.inherit_calls == [True] + + def test_base_callback_manager_is_idempotent(self) -> None: + class FakeManager: + def __init__(self) -> None: + self.handlers: list[Any] = [] + + def add_handler(self, handler: Any, inherit: bool) -> None: + self.handlers.append(handler) + + cb = object() + manager = FakeManager() + manager.handlers.append(cb) + _add_callback(manager, cb) # type: ignore[arg-type] + assert manager.handlers.count(cb) == 1 + + def test_non_iterable_falls_back_to_new_list(self) -> None: + cb = object() + container, replaced = _add_callback(object(), cb) # type: ignore[arg-type] + assert container == [cb] + assert replaced is True + + +class TestInjectCallbackShapes: + """End-to-end injection across the supported callback container shapes.""" + + def test_injects_into_tuple_callbacks_via_replacement( + self, adapter: LangChainAdapter, evaluator: MagicMock + ) -> None: + prior = object() + + class Agent: + callbacks = (prior,) + + a = Agent() + governed = adapter.attach(a, "id", "session", evaluator) + assert isinstance(a.callbacks, list) + assert a.callbacks == [prior, governed._callback] + + def test_injects_into_base_callback_manager_in_place( + self, adapter: LangChainAdapter, evaluator: MagicMock + ) -> None: + class FakeManager: + def __init__(self) -> None: + self.handlers: list[Any] = [] + + def add_handler(self, handler: Any, inherit: bool) -> None: + self.handlers.append(handler) + + manager = FakeManager() + + class Agent: + callbacks = manager + + a = Agent() + governed = adapter.attach(a, "id", "session", evaluator) + assert a.callbacks is manager + assert manager.handlers == [governed._callback] + + def test_injects_into_config_with_tuple_callbacks( + self, adapter: LangChainAdapter, evaluator: MagicMock + ) -> None: + prior = object() + + class Agent: + config: dict[str, Any] = {"callbacks": (prior,)} + + a = Agent() + governed = adapter.attach(a, "id", "session", evaluator) + assert a.config["callbacks"] == [prior, governed._callback] + + def test_ensure_callback_config_with_base_callback_manager( + self, adapter: LangChainAdapter, evaluator: MagicMock + ) -> None: + class FakeManager: + def __init__(self) -> None: + self.handlers: list[Any] = [] + + def add_handler(self, handler: Any, inherit: bool) -> None: + self.handlers.append(handler) + + manager = FakeManager() + gov = adapter.attach(SimpleNamespace(callbacks=[]), "id", "session", evaluator) + config = {"callbacks": manager} + result = gov._ensure_callback_config(config) + assert result is config + assert result["callbacks"] is manager + assert manager.handlers == [gov._callback] diff --git a/uv.lock b/uv.lock index a425af33d..e12d97448 100644 --- a/uv.lock +++ b/uv.lock @@ -9,6 +9,7 @@ resolution-markers = [ ] [options] +prerelease-mode = "allow" exclude-newer = "0001-01-01T00:00:00Z" # This has no effect and is included for backwards compatibility when using relative exclude-newer values. exclude-newer-span = "P2D" @@ -21,6 +22,9 @@ jsonschema-pydantic-converter = false uipath-langchain-client = false uipath-core = false +[manifest] +overrides = [{ name = "uipath-runtime", specifier = ">=0.11.0.dev1001180449,<0.11.0.dev1001190000", index = "https://test.pypi.org/simple/" }] + [[package]] name = "a2a-sdk" version = "0.3.26" @@ -563,6 +567,43 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/db/3c/33bac158f8ab7f89b2e59426d5fe2e4f63f7ed25df84c036890172b412b5/cfgv-3.5.0-py2.py3-none-any.whl", hash = "sha256:a8dc6b26ad22ff227d2634a65cb388215ce6cc96bbcc5cfde7641ae87e8dacc0", size = 7445, upload-time = "2025-11-19T20:55:50.744Z" }, ] +[[package]] +name = "chardet" +version = "7.4.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/19/b6/9df434a8eeba2e6628c465a1dfa31034228ef79b26f76f46278f4ef7e49d/chardet-7.4.3.tar.gz", hash = "sha256:cc1d4eb92a4ec1c2df3b490836ffa46922e599d34ce0bb75cf41fd2bf6303d56", size = 784800, upload-time = "2026-04-13T21:33:39.803Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/19/52/505c207f334d51e937cbaa27ff95776e16e2d120e13cbe491cd7b3a70b50/chardet-7.4.3-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:25a862cddc6a9ac07023e808aedd297115345fbaabc2690479481ddc0f980e09", size = 870747, upload-time = "2026-04-13T21:32:56.916Z" }, + { url = "https://files.pythonhosted.org/packages/14/4b/d3c79495dee4831b8bebca2790e72cb90f0c5849c940570a7c7e5b70b952/chardet-7.4.3-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:7005c88da26fd95d8abb8acbe6281d833e9a9181b03cf49b4546c4555389bd97", size = 853210, upload-time = "2026-04-13T21:32:58.309Z" }, + { url = "https://files.pythonhosted.org/packages/b9/99/f6a822ad1bde25a4c38dc3e770485e78e0893dfd871cd6e18ed3ea3a795e/chardet-7.4.3-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:dc50f28bad067393cce0af9091052c3b8df7a23115afd8ba7b2e0947f0cef1f8", size = 873625, upload-time = "2026-04-13T21:32:59.606Z" }, + { url = "https://files.pythonhosted.org/packages/b1/10/31932775c94a86814f76b41c4a772b52abfb0e6125324f32c6da1196c297/chardet-7.4.3-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4c3da294de1a681097848ab58bd3f2771a674f8039d2d87a5538b28856b815e9", size = 883436, upload-time = "2026-04-13T21:33:01.351Z" }, + { url = "https://files.pythonhosted.org/packages/6c/63/0f43e3acf2c436fdb32a0f904aeb03a2904d2126eed34a042a194d235926/chardet-7.4.3-cp311-cp311-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:93c45e116dd51b66226a53ade3f9f635e870de5399b90e00ce45dcc311093bf4", size = 876589, upload-time = "2026-04-13T21:33:02.636Z" }, + { url = "https://files.pythonhosted.org/packages/5d/a6/e9b8f8a3e99602792b01fa7d0a731737615ab56d8bfd0b52935a0ef88b85/chardet-7.4.3-cp311-cp311-win_amd64.whl", hash = "sha256:ccc1f83ab4bcfb901cf39e0c4ba6bc6e726fc6264735f10e24ceb5cb47387578", size = 941866, upload-time = "2026-04-13T21:33:04.282Z" }, + { url = "https://files.pythonhosted.org/packages/61/33/29de185079e6675c3f375546e30a559b7ddc75ce972f18d6e566cd9ea4eb/chardet-7.4.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:75d3c65cc16bddf40b8da1fd25ba84fca5f8070f2b14e86083653c1c85aee971", size = 874870, upload-time = "2026-04-13T21:33:05.977Z" }, + { url = "https://files.pythonhosted.org/packages/9c/2f/4c5af01fd1a7506a1d5375403d68925eac70289229492db5aa68b58103d8/chardet-7.4.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:29af5999f654e8729d251f1724a62b538b1262d9292cccaefddf8a02aae1ef6a", size = 854859, upload-time = "2026-04-13T21:33:07.381Z" }, + { url = "https://files.pythonhosted.org/packages/36/21/edb36ad5dfa48d7f8eed97ab43931ecdaa8c15166c21b1d614967e49d681/chardet-7.4.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:626f00299ad62dfe937058a09572beed442ccc7b58f87aa667949b20fd3db235", size = 875032, upload-time = "2026-04-13T21:33:08.741Z" }, + { url = "https://files.pythonhosted.org/packages/e5/59/a32a241d861cf180853a11c8e5a67641cb1b2af13c3a5ccce83ec07e2c9f/chardet-7.4.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9a4904dd5f071b7a7d7f50b4a67a86db3c902d243bf31708f1d5cde2f68239cb", size = 888283, upload-time = "2026-04-13T21:33:10.213Z" }, + { url = "https://files.pythonhosted.org/packages/87/2e/e1ee6a77abf3782c00e05b89c4d4328c8353bf9500661c4348df1dd68614/chardet-7.4.3-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:5d2879598bc220689e8ce509fe9c3f37ad2fca53a36be9c9bd91abdd91dd364f", size = 879974, upload-time = "2026-04-13T21:33:11.448Z" }, + { url = "https://files.pythonhosted.org/packages/32/60/fca69c534602a7ced04280c952a246ad1edde2a6ca3a164f65d32ac41fe7/chardet-7.4.3-cp312-cp312-win_amd64.whl", hash = "sha256:4b2799bd58e7245cfa8d4ab2e8ad1d76a5c3a5b1f32318eb6acca4c69a3e7101", size = 943973, upload-time = "2026-04-13T21:33:12.756Z" }, + { url = "https://files.pythonhosted.org/packages/7c/43/79ac9b4db5bc87020c9dbc419125371d80882d1d197e9c4765ba8682b605/chardet-7.4.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:a9e4486df251b8962e86ea9f139ca235aa6e0542a00f7844c9a04160afb99aa9", size = 873769, upload-time = "2026-04-13T21:33:14.002Z" }, + { url = "https://files.pythonhosted.org/packages/55/5f/25bdec773905bff0ff6cf35ca73b17bd05593b4f87bd8c5fa43705f7167d/chardet-7.4.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:4fbff1907925b0c5a1064cffb5e040cd5e338585c9c552625f30de6bc2f3107a", size = 853991, upload-time = "2026-04-13T21:33:15.564Z" }, + { url = "https://files.pythonhosted.org/packages/b4/07/a29380ee0b215d23d77733b5ad60c5c0c7969650e080c667acdf9462040d/chardet-7.4.3-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:365135eaf37ba65a828f8e668eb0a8c38c479dcbec724dc25f4dfd781049c357", size = 874024, upload-time = "2026-04-13T21:33:16.915Z" }, + { url = "https://files.pythonhosted.org/packages/a8/b1/3338e121cbd4c8a126b8ccb1061170c2ce51a53f678c502793ea49c6fd6d/chardet-7.4.3-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:bfc134b70c846c21ead8e43ada3ae1a805fff732f6922f8abcf2ff27b8f6493d", size = 887410, upload-time = "2026-04-13T21:33:18.368Z" }, + { url = "https://files.pythonhosted.org/packages/63/1c/44a9a9e0c59c185a5d307ceaeee8768afa1558f0a24f7a4b5fa11b67586b/chardet-7.4.3-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:9acd9988a93e09390f3cd231201ea7166c415eb8da1b735928990ffc05cb9fbb", size = 879269, upload-time = "2026-04-13T21:33:20.377Z" }, + { url = "https://files.pythonhosted.org/packages/1b/b3/5d0e77ea774bd3224321c248880ea0c0379000ac5c2bb6d77609549de247/chardet-7.4.3-cp313-cp313-win_amd64.whl", hash = "sha256:e1b98790c284ff813f18f7cf7de5f05ea2435a080030c7f1a8318f3a4f80b131", size = 944155, upload-time = "2026-04-13T21:33:21.694Z" }, + { url = "https://files.pythonhosted.org/packages/70/a8/bf0811d859e13801279a2ae64f37a408027b282f2047bc0001c75dd356ad/chardet-7.4.3-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:d892d3dcd652fdef53e3d6327d39b17c0df40a899dfc919abaeb64c974497531", size = 872887, upload-time = "2026-04-13T21:33:23.328Z" }, + { url = "https://files.pythonhosted.org/packages/51/ac/b9d68ebddfe1b02c77af5bf81120e12b036b4432dc6af7a303d90e2bc38b/chardet-7.4.3-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:acc46d1b8b7d5783216afe15db56d1c179b9a40e5a1558bc13164c4fd20674c4", size = 853964, upload-time = "2026-04-13T21:33:24.724Z" }, + { url = "https://files.pythonhosted.org/packages/2a/81/17fa103ea9caf5d325a5e4051ab2ba65996fd66baa60b81ee41af1f54e10/chardet-7.4.3-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0ac3bf11c645734a1701a3804e43eabd98851838192267d08c353a834ab79fea", size = 876006, upload-time = "2026-04-13T21:33:26.098Z" }, + { url = "https://files.pythonhosted.org/packages/c2/20/193faab46a68ea550587331a698c3dca8099f8901d10937c4443135c7ed9/chardet-7.4.3-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6e3bd9f936e04bae89c254262af08d9e5b98f805175ba1e29d454e6cba3107b7", size = 887680, upload-time = "2026-04-13T21:33:27.49Z" }, + { url = "https://files.pythonhosted.org/packages/40/c6/94a3c673327392652ee8bdea9a45bc8a5f5365197a7387d68f0eed007115/chardet-7.4.3-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:27cc23da03630cdecc9aa81a895aa86629c211f995cd57651f0fbc280717bf93", size = 879865, upload-time = "2026-04-13T21:33:29.052Z" }, + { url = "https://files.pythonhosted.org/packages/b1/2c/cad8b5e3623a987f3c930b68e2bdd06cfc388cd91cd42ed05f1227701b73/chardet-7.4.3-cp314-cp314-win_amd64.whl", hash = "sha256:b95c934b9ad59e2ba8abb9be49df70d3ad1b0d95d864b9fdb7588d4fa8bd921c", size = 939594, upload-time = "2026-04-13T21:33:31.391Z" }, + { url = "https://files.pythonhosted.org/packages/33/e0/d06e42fd6f02a58e5e227e5106587751cb38adcff0aaf949add744b78b6e/chardet-7.4.3-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:c77867f0c1cb8bd819502249fcdc500364aedb07881e11b743726fa2148e7b6e", size = 889714, upload-time = "2026-04-13T21:33:32.772Z" }, + { url = "https://files.pythonhosted.org/packages/d4/ed/40d091954d48abea037baae6be8fb79905e5f78d34d12ea955132c7d8011/chardet-7.4.3-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:cf1efeaf65a6ef2f5b9cc3a1df6f08ba2831b369ccaa4c7018eaf90aa757bb11", size = 872319, upload-time = "2026-04-13T21:33:34.427Z" }, + { url = "https://files.pythonhosted.org/packages/bb/77/82a46821dbfbdfe062710d2bf2ede13426304e3567a23c57d919c0c31630/chardet-7.4.3-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9f3504c139a2ad544077dd2d9e412cd08b01786843d76997cd43bb6de311723c", size = 892021, upload-time = "2026-04-13T21:33:35.766Z" }, + { url = "https://files.pythonhosted.org/packages/49/57/42d30c562bda5b4a839766c1aad8d5856b798ad2a1c3247b72a679afec94/chardet-7.4.3-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:457f619882ba66327d4d8d14c6c342269bdb1e4e1c38e8117df941d14d351b04", size = 902509, upload-time = "2026-04-13T21:33:37.096Z" }, + { url = "https://files.pythonhosted.org/packages/8c/6c/0a40afdb50a0fe041ab95553b835a8160b6cf0e81edf2ae2fe9f5224cbf9/chardet-7.4.3-py3-none-any.whl", hash = "sha256:1173b74051570cf08099d7429d92e4882d375ad4217f92a6e5240ccfb26f231e", size = 626562, upload-time = "2026-04-13T21:33:38.559Z" }, +] + [[package]] name = "charset-normalizer" version = "3.4.7" @@ -4374,16 +4415,16 @@ wheels = [ [[package]] name = "uipath-core" -version = "0.5.17" +version = "0.5.18" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "opentelemetry-instrumentation" }, { name = "opentelemetry-sdk" }, { name = "pydantic" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/e3/80/a626eb3136a6765e0af06c9d5080ac0843c2a72f17b7a2170f1f45da40dd/uipath_core-0.5.17.tar.gz", hash = "sha256:13565e1eba9f059a8221494dfb3239257ddf7f265fc7057199ffe03ed066300a", size = 119023, upload-time = "2026-05-28T21:34:10.903Z" } +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" } wheels = [ - { url = "https://files.pythonhosted.org/packages/74/cf/f4b481970621e2a9aec869302773fa2c7d346aef294a553429626369633f/uipath_core-0.5.17-py3-none-any.whl", hash = "sha256:6e088eec5130bc492ac176ab85d4924d7d4cb07ee290ed7e6a46984e9de8c12b", size = 44957, upload-time = "2026-05-28T21:34:09.534Z" }, + { 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" }, ] [[package]] @@ -4464,7 +4505,7 @@ requires-dist = [ { name = "pydantic-settings", specifier = ">=2.6.0" }, { name = "python-dotenv", specifier = ">=1.0.1" }, { name = "uipath", specifier = ">=2.10.79,<2.11.0" }, - { name = "uipath-core", specifier = ">=0.5.17,<0.6.0" }, + { name = "uipath-core", specifier = ">=0.5.18,<0.6.0" }, { name = "uipath-langchain-client", extras = ["all"], marker = "extra == 'all'", specifier = ">=1.13.0,<1.14.0" }, { name = "uipath-langchain-client", extras = ["anthropic"], marker = "extra == 'anthropic'", specifier = ">=1.13.0,<1.14.0" }, { name = "uipath-langchain-client", extras = ["bedrock"], marker = "extra == 'bedrock'", specifier = ">=1.13.0,<1.14.0" }, @@ -4473,7 +4514,7 @@ requires-dist = [ { name = "uipath-langchain-client", extras = ["openai"], specifier = ">=1.13.0,<1.14.0" }, { name = "uipath-langchain-client", extras = ["vertexai"], marker = "extra == 'vertex'", specifier = ">=1.13.0,<1.14.0" }, { name = "uipath-platform", specifier = ">=0.1.61,<0.2.0" }, - { name = "uipath-runtime", specifier = ">=0.11.0,<0.12.0" }, + { name = "uipath-runtime", specifier = ">=0.11.0.dev1001180449,<0.11.0.dev1001190000", index = "https://test.pypi.org/simple/" }, ] provides-extras = ["anthropic", "vertex", "bedrock", "fireworks", "all"] @@ -4573,14 +4614,17 @@ wheels = [ [[package]] name = "uipath-runtime" -version = "0.11.0" -source = { registry = "https://pypi.org/simple" } +version = "0.11.0.dev1001180449" +source = { registry = "https://test.pypi.org/simple/" } dependencies = [ + { name = "chardet" }, + { name = "pyyaml" }, { name = "uipath-core" }, + { name = "vadersentiment" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/72/8d/4d36d6a5dda4ca5f25e52508bc20dd82cb92fcdf2a36cd0adc4f9832d047/uipath_runtime-0.11.0.tar.gz", hash = "sha256:cc94f2fdab43b593ef678eff904fc6cdd4831963cffe39a83909ffcf9082d76f", size = 143685, upload-time = "2026-05-29T15:13:30.562Z" } +sdist = { url = "https://test-files.pythonhosted.org/packages/42/ec/735ae02437793e5aa8431031e0fe945a2194cf0a6dac45f740522d334f24/uipath_runtime-0.11.0.dev1001180449.tar.gz", hash = "sha256:bc1a65f8803ed85a7f944ac0d805c5a75492b1215741523fcc0a1686da783592", size = 256729, upload-time = "2026-06-12T09:13:05.45Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/e7/08/c7b90851d4544ff5e76ca7c55452597aae1619cf1ebc2c0aa7b098110f14/uipath_runtime-0.11.0-py3-none-any.whl", hash = "sha256:08bf53a0e38bb3d19edc6708d2ecb7d918aa96fdda13e35f3ad0e6f2a6c392b9", size = 43770, upload-time = "2026-05-29T15:13:29.282Z" }, + { url = "https://test-files.pythonhosted.org/packages/d2/6f/766e30a03947a7688e5ce64c6f72995cefc3da7f00d236e200a18bb6c8f2/uipath_runtime-0.11.0.dev1001180449-py3-none-any.whl", hash = "sha256:04de9fbff9faf4f64f0f52cd6cb11546126f8ac5f89fb13eba0b9d5f933b2cd7", size = 110803, upload-time = "2026-06-12T09:13:07.03Z" }, ] [[package]] @@ -4634,6 +4678,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/31/a3/5b1562db76a5a488274b2332a97199b32d0442aca0ed193697fd47786316/uvicorn-0.46.0-py3-none-any.whl", hash = "sha256:bbebbcbed972d162afca128605223022bedd345b7bc7855ce66deb31487a9048", size = 70926, upload-time = "2026-04-23T07:15:58.355Z" }, ] +[[package]] +name = "vadersentiment" +version = "3.3.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "requests" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/77/8c/4a48c10a50f750ae565e341e697d74a38075a3e43ff0df6f1ab72e186902/vaderSentiment-3.3.2.tar.gz", hash = "sha256:5d7c06e027fc8b99238edb0d53d970cf97066ef97654009890b83703849632f9", size = 2466783, upload-time = "2020-05-22T15:06:32.81Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/76/fc/310e16254683c1ed35eeb97386986d6c00bc29df17ce280aed64d55537e9/vaderSentiment-3.3.2-py2.py3-none-any.whl", hash = "sha256:3bf1d243b98b1afad575b9f22bc2cb1e212b94ff89ca74f8a23a588d024ea311", size = 125950, upload-time = "2020-05-22T15:07:00.052Z" }, +] + [[package]] name = "validators" version = "0.35.0" From 4faf5f7c79a55b11623e71bb2d29b1e8f8bc13a6 Mon Sep 17 00:00:00 2001 From: Aditi Kumari Date: Mon, 22 Jun 2026 16:41:22 +0530 Subject: [PATCH 02/10] chore: remove testpypi dev-pin for uipath-runtime Revert uipath-runtime to the stable range (>=0.11.0,<0.12.0) and drop the temporary testpypi source, prerelease allowance, and override-dependencies that were only needed for in-PR testing. Version will be bumped to the governance-bearing release once uipath-runtime goes live publicly. Co-Authored-By: Claude Opus 4.8 --- pyproject.toml | 12 +----------- 1 file changed, 1 insertion(+), 11 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 88f17eb86..668c2624a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,7 +8,7 @@ dependencies = [ "uipath>=2.10.79, <2.11.0", "uipath-core>=0.5.18, <0.6.0", "uipath-platform>=0.1.61, <0.2.0", - "uipath-runtime>=0.11.0.dev1001180449,<0.11.0.dev1001190000", + "uipath-runtime>=0.11.0,<0.12.0", "langgraph>=1.1.8, <2.0.0", "langchain-core>=1.2.11, <2.0.0", "langgraph-checkpoint-sqlite>=3.0.3, <4.0.0", @@ -157,16 +157,6 @@ exclude_lines = [ [tool.uv] exclude-newer = "2 days" -# uipath-runtime is pinned to a dev pre-release served from testpypi. -# uipath==2.10.79 transitively pins uipath-runtime>=0.11.0,<0.12.0, -# which excludes pre-releases (PEP 440: 0.11.0.dev* sorts below 0.11.0), -# so we need both the prerelease allowance and an explicit override to -# bypass the umbrella's stable-only constraint. -prerelease = "allow" -override-dependencies = ["uipath-runtime>=0.11.0.dev1001180449,<0.11.0.dev1001190000"] - -[tool.uv.sources] -uipath-runtime = { index = "testpypi" } [tool.uv.exclude-newer-package] uipath = false From 50f3487a2e78e52ce4fbe103422e9d7c120f7808 Mon Sep 17 00:00:00 2001 From: Aditi Kumari Date: Tue, 23 Jun 2026 18:03:21 +0530 Subject: [PATCH 03/10] =?UTF-8?q?refactor(governance):=20address=20review?= =?UTF-8?q?=20=E2=80=94=20tighten=20adapter=20boundaries?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - can_handle: drop the broad duck-typed invoke/ainvoke fallback; only claim LangChain Runnable / LangGraph CompiledStateGraph objects. - __init__: remove side-effect registration on import; the adapter is registered only through the uipath.governance.adapters entry point / discovery path. Importing the module no longer mutates the registry. - docstrings/comments: stop documenting runtime-internal locations (GovernanceRuntime, uipath.runtime.governance.wrapper); refer to the governance host generically so this package stays focused on the LangChain hook bridge. - tests: update can_handle expectations to match (duck-typed and import-failure cases now return False). Co-Authored-By: Claude Opus 4.8 --- src/uipath_langchain/governance/__init__.py | 25 +++----- src/uipath_langchain/governance/adapter.py | 68 ++++++++------------- tests/governance/test_adapter.py | 11 ++-- 3 files changed, 42 insertions(+), 62 deletions(-) diff --git a/src/uipath_langchain/governance/__init__.py b/src/uipath_langchain/governance/__init__.py index 71b38b02a..b16515760 100644 --- a/src/uipath_langchain/governance/__init__.py +++ b/src/uipath_langchain/governance/__init__.py @@ -1,22 +1,17 @@ """Governance integration for ``uipath-langchain``. -Registers :class:`LangChainAdapter` with the global adapter registry in -``uipath.core.adapters`` so ``uipath.runtime.governance.GovernanceRuntime`` -can attach the LangChain-specific inner hooks (BEFORE_MODEL, -AFTER_MODEL, TOOL_CALL, AFTER_TOOL) when it sees a LangChain or -LangGraph agent. +Registers :class:`LangChainAdapter` with the adapter registry in +``uipath.core.adapters`` so the governance host can attach the +LangChain-specific inner hooks (BEFORE_MODEL, AFTER_MODEL, TOOL_CALL, +AFTER_TOOL) when it sees a LangChain or LangGraph 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_langchain.governance`` is - opted in. - 2. The package also exposes :func:`register_governance_adapter` as an - entry point under ``uipath.governance.adapters`` so an upstream - discoverer (or ``uipath-core`` if/when it grows 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 +41,6 @@ def register_governance_adapter() -> None: logger.debug("Registered uipath-langchain governance adapter") -# Side-effect registration on module import. -register_governance_adapter() - - __all__ = [ "GovernanceCallbackHandler", "GovernedLangChainAgent", diff --git a/src/uipath_langchain/governance/adapter.py b/src/uipath_langchain/governance/adapter.py index ee97c272e..cee6cec28 100644 --- a/src/uipath_langchain/governance/adapter.py +++ b/src/uipath_langchain/governance/adapter.py @@ -11,16 +11,16 @@ - ``on_tool_start`` / ``on_tool_end`` → TOOL_CALL / AFTER_TOOL Chain-level boundaries (BEFORE_AGENT / AFTER_AGENT) are intentionally -*not* fired from here — they are owned by the runtime wrapper layer in -``uipath-runtime`` (``GovernanceRuntime.execute`` / ``.stream``). The -``GovernanceCallbackHandler`` sets ``ignore_chain = True`` so LangChain -skips chain notifications entirely, avoiding duplicate boundary -evaluations and silencing AttributeError noise for the absent methods. +*not* fired from here — they are owned by the governance host that drives +the agent. The ``GovernanceCallbackHandler`` sets ``ignore_chain = True`` +so LangChain skips chain notifications entirely, avoiding duplicate +boundary evaluations and silencing AttributeError noise for the absent +methods. Contracts and the evaluator protocol come from ``uipath-core``; this -package contributes only the LangChain-specific implementation and -self-registers it with the global adapter registry when -``uipath_langchain.governance`` is imported. +package contributes only the LangChain-specific implementation. The +adapter is exposed through the ``uipath.governance.adapters`` entry point +and registered via the governance adapter discovery path. Audit emission and enforcement (raising :class:`GovernanceBlockException` on DENY) are owned by the evaluator itself. This module just hooks the @@ -41,11 +41,9 @@ logger = logging.getLogger(__name__) # Cap on the text blob passed to BEFORE_MODEL / AFTER_MODEL governance -# evaluation. Sized to match the runtime side (see -# ``_GOVERNANCE_TEXT_CAP`` in ``uipath.runtime.governance.wrapper``) so -# scan-time budgets are consistent across hooks. A long conversation -# history is governed at the LLM layer by scanning only the latest -# message, not the full prompt — see +# evaluation, keeping scan-time budgets bounded across hooks. A long +# conversation history is governed at the LLM layer by scanning only the +# latest message, not the full prompt — see # :meth:`GovernanceCallbackHandler._latest_message_input`. The same cap # bounds the AFTER_MODEL ``model_output`` blob so batched or runaway # responses can't blow scan budgets either. @@ -126,18 +124,10 @@ def can_handle(self, agent: Any) -> bool: except ImportError: pass - # Duck-typed fallback: anything with invoke + ainvoke that isn't - # claimed by another framework adapter (caller-side ordering also - # ensures more specific adapters resolve first). - if hasattr(agent, "invoke") and hasattr(agent, "ainvoke"): - module = getattr(type(agent), "__module__", "") - if any( - fw in module - for fw in ("autogen", "crewai", "llama_index", "pydantic_ai") - ): - return False - return True - + # Only LangChain/LangGraph objects are claimed. No duck-typed + # fallback: matching on generic ``invoke`` / ``ainvoke`` would let + # this adapter attach to other frameworks that share those method + # names. return False def attach( @@ -257,12 +247,11 @@ class GovernanceCallbackHandler: run_inline: bool = True raise_error: bool = False ignore_llm: bool = False - # Chain-level events are owned by the runtime wrapper layer - # (BEFORE_AGENT / AFTER_AGENT fire from GovernanceRuntime.execute / - # .stream). Telling LangChain to skip chain callbacks here avoids - # duplicate boundary firings AND silences the AttributeError noise - # LangChain would otherwise log for every chain start/end now that - # we don't define the methods. + # Chain-level events (BEFORE_AGENT / AFTER_AGENT) are owned by the + # governance host, not this handler. Telling LangChain to skip chain + # callbacks here avoids duplicate boundary firings AND silences the + # AttributeError noise LangChain would otherwise log for every chain + # start/end now that we don't define the methods. ignore_chain: bool = True ignore_agent: bool = False ignore_retriever: bool = True @@ -399,9 +388,8 @@ def on_llm_end(self, response: Any, **kwargs: Any) -> None: """Evaluate AFTER_MODEL rules at LLM end. Concatenates text from every generation. The result is capped at - ``_BEFORE_MODEL_TEXT_CAP`` to match the BEFORE_MODEL budget and - the runtime side's ``_GOVERNANCE_TEXT_CAP``, so batched calls or - a runaway single response can't blow scan budgets. + ``_BEFORE_MODEL_TEXT_CAP`` to match the BEFORE_MODEL budget, so + batched calls or a runaway single response can't blow scan budgets. """ try: parts: list[str] = [] @@ -571,10 +559,8 @@ def on_tool_error(self, error: Exception, **kwargs: Any) -> None: logger.warning("Tool error in governed session %s: %s", self._session_id, error) # Chain-level callbacks (on_chain_start / on_chain_end / on_chain_error) - # are intentionally NOT implemented here. The runtime wrapper - # (``GovernanceRuntime.execute`` / ``GovernanceRuntime.stream`` in - # ``uipath-runtime``) owns BEFORE_AGENT / AFTER_AGENT — firing them - # here too would duplicate every boundary evaluation. The - # ``ignore_chain = True`` class-level descriptor above tells - # LangChain to skip chain notifications entirely so we don't get - # AttributeError warnings for the absent methods. + # are intentionally NOT implemented here. The governance host owns + # BEFORE_AGENT / AFTER_AGENT — firing them here too would duplicate + # every boundary evaluation. The ``ignore_chain = True`` class-level + # descriptor above tells LangChain to skip chain notifications entirely + # so we don't get AttributeError warnings for the absent methods. diff --git a/tests/governance/test_adapter.py b/tests/governance/test_adapter.py index 93aa24611..eef162670 100644 --- a/tests/governance/test_adapter.py +++ b/tests/governance/test_adapter.py @@ -64,9 +64,11 @@ def test_returns_true_for_langchain_runnable( runnable = RunnableLambda(lambda x: x) assert adapter.can_handle(runnable) is True - def test_returns_true_for_duck_typed_invoke_ainvoke( + def test_returns_false_for_duck_typed_invoke_ainvoke( self, adapter: LangChainAdapter ) -> None: + # A non-LangChain object that merely exposes invoke/ainvoke must + # NOT be claimed — the adapter only handles LangChain/LangGraph. class Duck: def invoke(self, x): return x @@ -74,13 +76,13 @@ def invoke(self, x): async def ainvoke(self, x): return x - assert adapter.can_handle(Duck()) is True + assert adapter.can_handle(Duck()) is False @pytest.mark.parametrize( "module_name", ["autogen.foo", "crewai.agent", "llama_index.core", "pydantic_ai.bar"], ) - def test_returns_false_for_excluded_frameworks( + def test_returns_false_for_other_frameworks( self, adapter: LangChainAdapter, module_name: str ) -> None: class Foreign: @@ -118,6 +120,7 @@ def test_handles_langgraph_import_failure( def test_handles_langchain_core_import_failure( self, monkeypatch: pytest.MonkeyPatch, adapter: LangChainAdapter ) -> None: + # With both framework imports unavailable, nothing can be claimed. monkeypatch.setitem(sys.modules, "langgraph.graph.state", None) monkeypatch.setitem(sys.modules, "langchain_core.runnables", None) @@ -128,7 +131,7 @@ def invoke(self, x): async def ainvoke(self, x): return x - assert adapter.can_handle(Duck()) is True + assert adapter.can_handle(Duck()) is False class TestAttach: From 7ddde0d67043b52193c2297c555945ec463d1205 Mon Sep 17 00:00:00 2001 From: Aditi Kumari Date: Tue, 23 Jun 2026 19:04:54 +0530 Subject: [PATCH 04/10] refactor(governance): resolve SonarQube code smells - Rename unused LangChain-callback parameter serialized -> _serialized in on_llm_start / on_chat_model_start (S1172). - Reduce cognitive complexity (S3776): extract _blocks_to_text from _latest_message_input and _collect_generations_text from on_llm_end. Behavior unchanged; covered by existing tests. Co-Authored-By: Claude Opus 4.8 --- src/uipath_langchain/governance/adapter.py | 79 ++++++++++++++-------- 1 file changed, 50 insertions(+), 29 deletions(-) diff --git a/src/uipath_langchain/governance/adapter.py b/src/uipath_langchain/governance/adapter.py index cee6cec28..0ba40ecd4 100644 --- a/src/uipath_langchain/governance/adapter.py +++ b/src/uipath_langchain/governance/adapter.py @@ -278,11 +278,15 @@ def __init__( def on_llm_start( self, - serialized: Dict[str, Any], + _serialized: Dict[str, Any], prompts: list[str], **kwargs: Any, ) -> None: - """Evaluate BEFORE_MODEL rules at LLM start (non-chat completion).""" + """Evaluate BEFORE_MODEL rules at LLM start (non-chat completion). + + ``_serialized`` is part of LangChain's callback signature but unused + here. + """ try: self._session_state["llm_calls"] = ( self._session_state.get("llm_calls", 0) + 1 @@ -304,7 +308,7 @@ def on_llm_start( def on_chat_model_start( self, - serialized: Dict[str, Any], + _serialized: Dict[str, Any], messages: list[list[Any]], **kwargs: Any, ) -> None: @@ -370,20 +374,30 @@ def _latest_message_input(messages: list[list[Any]]) -> str: if isinstance(content, str): return content[:_BEFORE_MODEL_TEXT_CAP] if isinstance(content, list): - parts: list[str] = [] - remaining = _BEFORE_MODEL_TEXT_CAP - for block in content: - if remaining <= 0: - break - if not isinstance(block, dict): - continue - piece = GovernanceCallbackHandler._extract_block_text(block) - if piece: - parts.append(piece) - remaining -= len(piece) + 1 - return "\n".join(parts)[:_BEFORE_MODEL_TEXT_CAP] + return GovernanceCallbackHandler._blocks_to_text(content) return "" + @staticmethod + def _blocks_to_text(content: list[Any]) -> str: + """Concatenate governance-relevant text from a list of content blocks. + + Walks list-of-blocks message content (multimodal, function-call, + tool_use, extended thinking) via :meth:`_extract_block_text`, + capping the joined result at ``_BEFORE_MODEL_TEXT_CAP``. + """ + parts: list[str] = [] + remaining = _BEFORE_MODEL_TEXT_CAP + for block in content: + if remaining <= 0: + break + if not isinstance(block, dict): + continue + piece = GovernanceCallbackHandler._extract_block_text(block) + if piece: + parts.append(piece) + remaining -= len(piece) + 1 + return "\n".join(parts)[:_BEFORE_MODEL_TEXT_CAP] + def on_llm_end(self, response: Any, **kwargs: Any) -> None: """Evaluate AFTER_MODEL rules at LLM end. @@ -392,20 +406,7 @@ def on_llm_end(self, response: Any, **kwargs: Any) -> None: batched calls or a runaway single response can't blow scan budgets. """ try: - parts: list[str] = [] - remaining = _BEFORE_MODEL_TEXT_CAP - if hasattr(response, "generations"): - for gen_list in response.generations: - for gen in gen_list: - if remaining <= 0: - break - piece = self._extract_generation_text(gen) - if piece: - parts.append(piece) - remaining -= len(piece) - if remaining <= 0: - break - model_output = "".join(parts)[:_BEFORE_MODEL_TEXT_CAP] + model_output = self._collect_generations_text(response) self._evaluator.evaluate_after_model( model_output=model_output, agent_name=self._agent_name, @@ -417,6 +418,26 @@ def on_llm_end(self, response: Any, **kwargs: Any) -> None: except Exception as e: logger.warning("on_llm_end governance check failed (continuing): %s", e) + def _collect_generations_text(self, response: Any) -> str: + """Concatenate text across all generations, capped at the text budget. + + Returns ``""`` when the response carries no ``generations``. + """ + parts: list[str] = [] + remaining = _BEFORE_MODEL_TEXT_CAP + if hasattr(response, "generations"): + for gen_list in response.generations: + for gen in gen_list: + if remaining <= 0: + break + piece = self._extract_generation_text(gen) + if piece: + parts.append(piece) + remaining -= len(piece) + if remaining <= 0: + break + return "".join(parts)[:_BEFORE_MODEL_TEXT_CAP] + @staticmethod def _extract_generation_text(gen: Any) -> str: """Return the text payload of a LangChain ``Generation`` / ``ChatGeneration``. From 2842f4d4d9017d30aeceb7e94fc667fe0101260e Mon Sep 17 00:00:00 2001 From: Aditi Kumari Date: Tue, 23 Jun 2026 19:28:29 +0530 Subject: [PATCH 05/10] refactor(governance): collapse capped-join loop to clear S3776 Extract the shared "accumulate text up to the scan cap" loop into _join_within_cap, used by both _blocks_to_text and _collect_generations_text. This removes the nested-loop cognitive complexity Sonar flagged on _collect_generations_text (17 -> under 15). Behavior unchanged; existing tests cover both paths. Co-Authored-By: Claude Opus 4.8 --- src/uipath_langchain/governance/adapter.py | 52 ++++++++++++---------- 1 file changed, 29 insertions(+), 23 deletions(-) diff --git a/src/uipath_langchain/governance/adapter.py b/src/uipath_langchain/governance/adapter.py index 0ba40ecd4..c78cfb09b 100644 --- a/src/uipath_langchain/governance/adapter.py +++ b/src/uipath_langchain/governance/adapter.py @@ -32,7 +32,7 @@ from __future__ import annotations import logging -from typing import Any, Dict +from typing import Any, Dict, Iterable from uuid import uuid4 from uipath.core.adapters import BaseAdapter, EvaluatorProtocol, GovernedAgentBase @@ -385,18 +385,30 @@ def _blocks_to_text(content: list[Any]) -> str: tool_use, extended thinking) via :meth:`_extract_block_text`, capping the joined result at ``_BEFORE_MODEL_TEXT_CAP``. """ - parts: list[str] = [] + pieces = ( + GovernanceCallbackHandler._extract_block_text(block) + for block in content + if isinstance(block, dict) + ) + return GovernanceCallbackHandler._join_within_cap(pieces, "\n") + + @staticmethod + def _join_within_cap(pieces: Iterable[str], sep: str) -> str: + """Join non-empty ``pieces`` with ``sep``, stopping at the text cap. + + Shared accumulator for the model-input/output scan blobs: appends + pieces until ``_BEFORE_MODEL_TEXT_CAP`` characters are reached + (counting the separator), then caps the joined result. + """ + out: list[str] = [] remaining = _BEFORE_MODEL_TEXT_CAP - for block in content: + for piece in pieces: if remaining <= 0: break - if not isinstance(block, dict): - continue - piece = GovernanceCallbackHandler._extract_block_text(block) if piece: - parts.append(piece) - remaining -= len(piece) + 1 - return "\n".join(parts)[:_BEFORE_MODEL_TEXT_CAP] + out.append(piece) + remaining -= len(piece) + len(sep) + return sep.join(out)[:_BEFORE_MODEL_TEXT_CAP] def on_llm_end(self, response: Any, **kwargs: Any) -> None: """Evaluate AFTER_MODEL rules at LLM end. @@ -423,20 +435,14 @@ def _collect_generations_text(self, response: Any) -> str: Returns ``""`` when the response carries no ``generations``. """ - parts: list[str] = [] - remaining = _BEFORE_MODEL_TEXT_CAP - if hasattr(response, "generations"): - for gen_list in response.generations: - for gen in gen_list: - if remaining <= 0: - break - piece = self._extract_generation_text(gen) - if piece: - parts.append(piece) - remaining -= len(piece) - if remaining <= 0: - break - return "".join(parts)[:_BEFORE_MODEL_TEXT_CAP] + if not hasattr(response, "generations"): + return "" + pieces = ( + self._extract_generation_text(gen) + for gen_list in response.generations + for gen in gen_list + ) + return self._join_within_cap(pieces, "") @staticmethod def _extract_generation_text(gen: Any) -> str: From 3d260873da4c7c07446ce25fecbb612c0321ef64 Mon Sep 17 00:00:00 2001 From: Aditi Kumari Date: Wed, 24 Jun 2026 16:32:27 +0530 Subject: [PATCH 06/10] chore(governance): trim debt-history comments per review (PR 899) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address radu's review on the LangChain adapter: drop comments that document removed behavior / tech debt and keep code self-explanatory. - can_handle: remove the 'no duck-typed fallback' explanatory comment (the code already only claims LangChain/LangGraph). - ignore_chain descriptor: shorten to the invariant (chain events owned by the governance host → skipped to avoid duplicate boundary evaluations); drop the AttributeError-noise wording. - module docstring: drop the same AttributeError-noise tail. - remove the dangling end-of-file comment re-explaining why chain callbacks aren't implemented (docstring + ignore_chain already cover it). Co-Authored-By: Claude Opus 4.8 --- src/uipath_langchain/governance/adapter.py | 20 +++----------------- 1 file changed, 3 insertions(+), 17 deletions(-) diff --git a/src/uipath_langchain/governance/adapter.py b/src/uipath_langchain/governance/adapter.py index c78cfb09b..d44efeb6f 100644 --- a/src/uipath_langchain/governance/adapter.py +++ b/src/uipath_langchain/governance/adapter.py @@ -14,8 +14,7 @@ *not* fired from here — they are owned by the governance host that drives the agent. The ``GovernanceCallbackHandler`` sets ``ignore_chain = True`` so LangChain skips chain notifications entirely, avoiding duplicate -boundary evaluations and silencing AttributeError noise for the absent -methods. +boundary evaluations. Contracts and the evaluator protocol come from ``uipath-core``; this package contributes only the LangChain-specific implementation. The @@ -124,10 +123,6 @@ def can_handle(self, agent: Any) -> bool: except ImportError: pass - # Only LangChain/LangGraph objects are claimed. No duck-typed - # fallback: matching on generic ``invoke`` / ``ainvoke`` would let - # this adapter attach to other frameworks that share those method - # names. return False def attach( @@ -248,10 +243,8 @@ class GovernanceCallbackHandler: raise_error: bool = False ignore_llm: bool = False # Chain-level events (BEFORE_AGENT / AFTER_AGENT) are owned by the - # governance host, not this handler. Telling LangChain to skip chain - # callbacks here avoids duplicate boundary firings AND silences the - # AttributeError noise LangChain would otherwise log for every chain - # start/end now that we don't define the methods. + # governance host, so this handler skips them to avoid duplicate + # boundary evaluations. ignore_chain: bool = True ignore_agent: bool = False ignore_retriever: bool = True @@ -584,10 +577,3 @@ def on_tool_error(self, error: Exception, **kwargs: Any) -> None: if run_id is not None: self._tool_runs.pop(str(run_id), None) logger.warning("Tool error in governed session %s: %s", self._session_id, error) - - # Chain-level callbacks (on_chain_start / on_chain_end / on_chain_error) - # are intentionally NOT implemented here. The governance host owns - # BEFORE_AGENT / AFTER_AGENT — firing them here too would duplicate - # every boundary evaluation. The ``ignore_chain = True`` class-level - # descriptor above tells LangChain to skip chain notifications entirely - # so we don't get AttributeError warnings for the absent methods. From 820361c95fbf7b257cb4a23a1d580758fe55617d Mon Sep 17 00:00:00 2001 From: Aditi Kumari Date: Wed, 24 Jun 2026 16:42:03 +0530 Subject: [PATCH 07/10] chore(governance): shorten the text-cap comment (PR 899) Co-Authored-By: Claude Opus 4.8 --- src/uipath_langchain/governance/adapter.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/src/uipath_langchain/governance/adapter.py b/src/uipath_langchain/governance/adapter.py index d44efeb6f..efbccecd1 100644 --- a/src/uipath_langchain/governance/adapter.py +++ b/src/uipath_langchain/governance/adapter.py @@ -39,13 +39,8 @@ logger = logging.getLogger(__name__) -# Cap on the text blob passed to BEFORE_MODEL / AFTER_MODEL governance -# evaluation, keeping scan-time budgets bounded across hooks. A long -# conversation history is governed at the LLM layer by scanning only the -# latest message, not the full prompt — see -# :meth:`GovernanceCallbackHandler._latest_message_input`. The same cap -# bounds the AFTER_MODEL ``model_output`` blob so batched or runaway -# responses can't blow scan budgets either. +# Cap on the text scanned per model hook, so a long history / runaway +# response can't blow scan-time budgets. _BEFORE_MODEL_TEXT_CAP = 64000 From f251dc31af16af8cf1d4761a07de50b4ac002de6 Mon Sep 17 00:00:00 2001 From: Aditi Kumari Date: Wed, 24 Jun 2026 16:55:53 +0530 Subject: [PATCH 08/10] chore(governance): drop redundant restatement comments (PR 899) Remove comments that only restate the code: the two isinstance() labels in can_handle and the 'copy to a list' note. Co-Authored-By: Claude Opus 4.8 --- src/uipath_langchain/governance/adapter.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/src/uipath_langchain/governance/adapter.py b/src/uipath_langchain/governance/adapter.py index efbccecd1..e924bc4cf 100644 --- a/src/uipath_langchain/governance/adapter.py +++ b/src/uipath_langchain/governance/adapter.py @@ -77,7 +77,6 @@ def _add_callback( if callback not in (existing.handlers or []): existing.add_handler(callback, inherit=True) return existing, False - # Tuple or other iterable — copy to a list we can mutate. try: handlers = list(existing) except TypeError: @@ -100,7 +99,6 @@ def name(self) -> str: def can_handle(self, agent: Any) -> bool: """Return True if this adapter knows how to hook into the agent.""" - # LangGraph CompiledStateGraph try: from langgraph.graph.state import CompiledStateGraph @@ -109,7 +107,6 @@ def can_handle(self, agent: Any) -> bool: except ImportError: pass - # LangChain Runnable try: from langchain_core.runnables import Runnable From 463156d9465d1828b28b4525129e9494fc2ef639 Mon Sep 17 00:00:00 2001 From: Viswanath Lekshmanan Date: Thu, 25 Jun 2026 10:28:05 +0530 Subject: [PATCH 09/10] refactor(governance): drop adapter; wire handler via factory evaluator kwarg MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per the governance architecture review (uipath-python PR 1761): LangChainAdapter, GovernedLangChainAgent, the _inject_callback proxy, and the uipath.governance.adapters entry-point are removed. The surviving GovernanceCallbackHandler now subclasses langchain_core.callbacks.BaseCallbackHandler directly, and UiPathLangGraphRuntimeFactory.new_runtime gains an evaluator: EvaluatorProtocol | None = None kwarg — when supplied, the factory builds the handler and hands it to UiPathLangGraphRuntime through the existing callbacks constructor arg. Pins uipath-core==0.5.23.dev1017616946 (the PR 1761 dev build) via testpypi, since BaseAdapter and the adapter registry no longer exist in uipath-core. Co-Authored-By: Claude Opus 4.7 (1M context) --- pyproject.toml | 8 +- src/uipath_langchain/governance/__init__.py | 55 +- .../governance/{adapter.py => callbacks.py} | 239 +------- src/uipath_langchain/runtime/factory.py | 30 +- .../{test_adapter.py => test_callbacks.py} | 542 +----------------- tests/runtime/test_factory_governance.py | 116 ++++ uv.lock | 10 +- 7 files changed, 217 insertions(+), 783 deletions(-) rename src/uipath_langchain/governance/{adapter.py => callbacks.py} (64%) rename tests/governance/{test_adapter.py => test_callbacks.py} (56%) create mode 100644 tests/runtime/test_factory_governance.py diff --git a/pyproject.toml b/pyproject.toml index 8ac18a256..5cb82af57 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -6,7 +6,7 @@ readme = { file = "README.md", content-type = "text/markdown" } requires-python = ">=3.11" dependencies = [ "uipath>=2.10.79, <2.12.0", - "uipath-core>=0.5.20, <0.6.0", + "uipath-core==0.5.23.dev1017616946", "uipath-platform>=0.1.71, <0.2.0", "uipath-runtime>=0.11.0, <0.12.0", "langgraph>=1.1.8, <2.0.0", @@ -63,9 +63,6 @@ register = "uipath_langchain.middlewares:register_middleware" [project.entry-points."uipath.runtime.factories"] langgraph = "uipath_langchain.runtime:register_runtime_factory" -[project.entry-points."uipath.governance.adapters"] -langchain = "uipath_langchain.governance:register_governance_adapter" - [project.urls] Homepage = "https://uipath.com" Repository = "https://github.com/UiPath/uipath-langchain-python" @@ -172,3 +169,6 @@ name = "testpypi" url = "https://test.pypi.org/simple/" publish-url = "https://test.pypi.org/legacy/" explicit = true + +[tool.uv.sources] +uipath-core = { index = "testpypi" } diff --git a/src/uipath_langchain/governance/__init__.py b/src/uipath_langchain/governance/__init__.py index b16515760..c7bfd29bc 100644 --- a/src/uipath_langchain/governance/__init__.py +++ b/src/uipath_langchain/governance/__init__.py @@ -1,49 +1,16 @@ """Governance integration for ``uipath-langchain``. -Registers :class:`LangChainAdapter` with the adapter registry in -``uipath.core.adapters`` so the governance host can attach the -LangChain-specific inner hooks (BEFORE_MODEL, AFTER_MODEL, TOOL_CALL, -AFTER_TOOL) when it sees a LangChain or LangGraph 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 :class:`GovernanceCallbackHandler` — a LangChain callback +handler that calls an :class:`~uipath.core.adapters.EvaluatorProtocol` +on the model and tool lifecycle. Wired into a run by passing an +``evaluator`` to :class:`UiPathLangGraphRuntimeFactory`; the factory +builds the handler and hands it to the runtime through the existing +``callbacks`` channel. + +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 ( - GovernanceCallbackHandler, - GovernedLangChainAgent, - LangChainAdapter, -) - -logger = logging.getLogger(__name__) - - -def register_governance_adapter() -> None: - """Register :class:`LangChainAdapter` with the global registry. - - Idempotent — safe to call multiple times. - """ - registry = get_adapter_registry() - if any(a.name == "LangChain" for a in registry.get_all()): - return - registry.register(LangChainAdapter()) - logger.debug("Registered uipath-langchain governance adapter") - +from .callbacks import GovernanceCallbackHandler -__all__ = [ - "GovernanceCallbackHandler", - "GovernedLangChainAgent", - "LangChainAdapter", - "register_governance_adapter", -] +__all__ = ["GovernanceCallbackHandler"] diff --git a/src/uipath_langchain/governance/adapter.py b/src/uipath_langchain/governance/callbacks.py similarity index 64% rename from src/uipath_langchain/governance/adapter.py rename to src/uipath_langchain/governance/callbacks.py index e924bc4cf..5d796eac7 100644 --- a/src/uipath_langchain/governance/adapter.py +++ b/src/uipath_langchain/governance/callbacks.py @@ -1,31 +1,31 @@ -"""LangChain / LangGraph adapter for UiPath governance. +"""LangChain governance callback handler. -Provides governance for LangChain chains/agents and LangGraph compiled -graphs. Uses LangChain's callback system for deep hooks (model / tool -events) plus a thin proxy that ensures the callback is wired into -``invoke`` / ``ainvoke`` / ``stream`` / ``astream``. +A :class:`langchain_core.callbacks.BaseCallbackHandler` that calls a +framework-agnostic :class:`~uipath.core.adapters.EvaluatorProtocol` +on the model and tool lifecycle. -This adapter intercepts: +Wiring lives in :class:`UiPathLangGraphRuntimeFactory`: passing an +``evaluator`` to ``new_runtime`` causes the factory to build this +handler and hand it to :class:`UiPathLangGraphRuntime` through the +existing ``callbacks`` constructor arg. No adapter registry, no global +state, no import-time mutation. + +Intercepts: - ``on_llm_start`` / ``on_chat_model_start`` / ``on_llm_end`` → BEFORE_MODEL / AFTER_MODEL - ``on_tool_start`` / ``on_tool_end`` → TOOL_CALL / AFTER_TOOL Chain-level boundaries (BEFORE_AGENT / AFTER_AGENT) are intentionally -*not* fired from here — they are owned by the governance host that drives -the agent. The ``GovernanceCallbackHandler`` sets ``ignore_chain = True`` -so LangChain skips chain notifications entirely, avoiding duplicate -boundary evaluations. - -Contracts and the evaluator protocol come from ``uipath-core``; this -package contributes only the LangChain-specific implementation. The -adapter is exposed through the ``uipath.governance.adapters`` entry point -and registered via the governance adapter discovery path. - -Audit emission and enforcement (raising :class:`GovernanceBlockException` -on DENY) are owned by the evaluator itself. This module just hooks the -framework callbacks, extracts the data, and calls -``evaluator.evaluate_*``; block exceptions propagate, everything else -is logged and swallowed so a governance bug never breaks an agent run. +*not* fired from here — they are owned by the governance host that +drives the agent. ``ignore_chain = True`` makes LangChain skip chain +notifications entirely, avoiding duplicate boundary evaluations. + +Audit emission and enforcement (raising +:class:`GovernanceBlockException` on DENY) are owned by the evaluator +itself. This module just hooks the framework callbacks, extracts the +data, and calls ``evaluator.evaluate_*``; block exceptions propagate, +everything else is logged and swallowed so a governance bug never +breaks an agent run. """ from __future__ import annotations @@ -34,7 +34,8 @@ from typing import Any, Dict, Iterable from uuid import uuid4 -from uipath.core.adapters import BaseAdapter, EvaluatorProtocol, GovernedAgentBase +from langchain_core.callbacks import BaseCallbackHandler +from uipath.core.adapters import EvaluatorProtocol from uipath.core.governance.exceptions import GovernanceBlockException logger = logging.getLogger(__name__) @@ -44,193 +45,15 @@ _BEFORE_MODEL_TEXT_CAP = 64000 -def _add_callback( - existing: Any, callback: "GovernanceCallbackHandler" -) -> tuple[Any, bool]: - """Add ``callback`` to a LangChain callbacks container. - - LangChain accepts callbacks as ``None``, a ``list`` of handlers, a - tuple/other iterable, or a ``BaseCallbackManager``. We don't want a - hard import on ``langchain_core`` here, so we duck-type: - - - ``BaseCallbackManager`` exposes ``add_handler(handler, inherit)`` and - a ``handlers`` attribute; mutate it in place so any tracers / - handlers already attached to the manager are preserved. - - Lists are appended in place (callers may rely on identity). - - Anything else (``None``, tuple, generic iterable) is coerced to a - fresh list with the new callback appended. - - Returns ``(container, replaced)`` — ``replaced`` indicates whether - the caller should rebind the attribute / config slot to the returned - container (True for the coerced cases, False when we mutated in - place). - """ - if existing is None: - return [callback], True - if isinstance(existing, list): - if callback not in existing: - existing.append(callback) - return existing, False - # BaseCallbackManager: mutate the manager directly so attached - # tracers stay wired up. - if hasattr(existing, "add_handler") and hasattr(existing, "handlers"): - if callback not in (existing.handlers or []): - existing.add_handler(callback, inherit=True) - return existing, False - try: - handlers = list(existing) - except TypeError: - handlers = [] - if callback not in handlers: - handlers.append(callback) - return handlers, True - - -class LangChainAdapter(BaseAdapter): - """Adapter for LangChain / LangGraph frameworks. - - Detects and wraps LangChain chains, agents, and LangGraph - ``CompiledStateGraph`` instances with a governance callback handler. - """ - - @property - def name(self) -> str: - return "LangChain" - - def can_handle(self, agent: Any) -> bool: - """Return True if this adapter knows how to hook into the agent.""" - try: - from langgraph.graph.state import CompiledStateGraph - - if isinstance(agent, CompiledStateGraph): - return True - except ImportError: - pass - - try: - from langchain_core.runnables import Runnable - - if isinstance(agent, Runnable): - return True - except ImportError: - pass - - return False - - def attach( - self, - agent: Any, - agent_id: str, - session_id: str, - evaluator: EvaluatorProtocol, - ) -> "GovernedLangChainAgent": - """Attach governance to a LangChain / LangGraph agent.""" - callback = GovernanceCallbackHandler( - evaluator=evaluator, - agent_name=agent_id, - session_id=session_id, - ) - self._inject_callback(agent, callback) - return GovernedLangChainAgent( - agent=agent, - adapter=self, - agent_id=agent_id, - session_id=session_id, - evaluator=evaluator, - callback=callback, - ) - - def _inject_callback( - self, agent: Any, callback: "GovernanceCallbackHandler" - ) -> None: - """Inject the governance callback into the agent's callback chain.""" - if hasattr(agent, "callbacks"): - container, replaced = _add_callback( - getattr(agent, "callbacks", None), callback - ) - if replaced: - agent.callbacks = container - logger.debug("Injected governance callback via agent.callbacks") - return - - if hasattr(agent, "config"): - config = agent.config or {} - container, replaced = _add_callback(config.get("callbacks"), callback) - if replaced: - config["callbacks"] = container - agent.config = config - logger.debug("Injected governance callback via agent.config") - return - - logger.warning( - "Could not inject governance callback into %s — agent has neither " - "'callbacks' nor 'config' surface; deep hooks will not fire", - type(agent).__name__, - ) - - -class GovernedLangChainAgent(GovernedAgentBase): - """LangChain / LangGraph agent wrapped with governance. - - The callback handler does the actual rule evaluation; this proxy - ensures the handler is present on every ``invoke`` / ``ainvoke`` / - ``stream`` / ``astream`` call's config. - """ - - def __init__( - self, - agent: Any, - adapter: LangChainAdapter, - agent_id: str, - session_id: str, - evaluator: EvaluatorProtocol, - callback: "GovernanceCallbackHandler", - ) -> None: - super().__init__(agent, adapter, agent_id, session_id, evaluator) - self._callback = callback - - def invoke(self, input_data: Any, config: Any = None, **kwargs: Any) -> Any: - config = self._ensure_callback_config(config) - return self._agent.invoke(input_data, config=config, **kwargs) - - async def ainvoke(self, input_data: Any, config: Any = None, **kwargs: Any) -> Any: - config = self._ensure_callback_config(config) - return await self._agent.ainvoke(input_data, config=config, **kwargs) - - def stream(self, input_data: Any, config: Any = None, **kwargs: Any) -> Any: - config = self._ensure_callback_config(config) - yield from self._agent.stream(input_data, config=config, **kwargs) - - async def astream(self, input_data: Any, config: Any = None, **kwargs: Any) -> Any: - config = self._ensure_callback_config(config) - async for chunk in self._agent.astream(input_data, config=config, **kwargs): - yield chunk - - def _ensure_callback_config(self, config: Any) -> Dict[str, Any]: - """Ensure the governance callback is on the config's callback list.""" - if config is None: - config = {} - if isinstance(config, dict): - container, replaced = _add_callback(config.get("callbacks"), self._callback) - if replaced: - config["callbacks"] = container - return config - - -class GovernanceCallbackHandler: +class GovernanceCallbackHandler(BaseCallbackHandler): """LangChain callback handler that fires governance evaluation. - Implements the LangChain ``BaseCallbackHandler`` interface shape - structurally (no formal inheritance — keeps this package free of a - hard ``langchain_core`` import at module load). - The evaluator owns audit emission and DENY-raising. Each ``on_*`` callback only extracts the relevant payload and calls the matching ``evaluate_*`` method; :class:`GovernanceBlockException` is allowed to propagate, anything else is logged and swallowed. """ - # LangChain callback-handler descriptors: run_inline: bool = True raise_error: bool = False ignore_llm: bool = False @@ -263,15 +86,11 @@ def __init__( def on_llm_start( self, - _serialized: Dict[str, Any], + serialized: Dict[str, Any], prompts: list[str], **kwargs: Any, ) -> None: - """Evaluate BEFORE_MODEL rules at LLM start (non-chat completion). - - ``_serialized`` is part of LangChain's callback signature but unused - here. - """ + """Evaluate BEFORE_MODEL rules at LLM start (non-chat completion).""" try: self._session_state["llm_calls"] = ( self._session_state.get("llm_calls", 0) + 1 @@ -293,7 +112,7 @@ def on_llm_start( def on_chat_model_start( self, - _serialized: Dict[str, Any], + serialized: Dict[str, Any], messages: list[list[Any]], **kwargs: Any, ) -> None: @@ -494,7 +313,7 @@ def _extract_block_text(block: Dict[str, Any]) -> str: parts.extend(v for v in input_value.values() if isinstance(v, str)) return "\n".join(p for p in parts if p) - def on_llm_error(self, error: Exception, **kwargs: Any) -> None: + def on_llm_error(self, error: BaseException, **kwargs: Any) -> None: logger.warning("LLM error in governed session %s: %s", self._session_id, error) # ----- Tool callbacks -------------------------------------------------- @@ -562,7 +381,7 @@ def on_tool_end(self, output: Any, **kwargs: Any) -> None: except Exception as e: logger.warning("on_tool_end governance check failed (continuing): %s", e) - def on_tool_error(self, error: Exception, **kwargs: Any) -> None: + def on_tool_error(self, error: BaseException, **kwargs: Any) -> None: # Tool errored out — on_tool_end will not fire. Pop the mapping # so a session with many failing tool calls does not leak. run_id = kwargs.get("run_id") diff --git a/src/uipath_langchain/runtime/factory.py b/src/uipath_langchain/runtime/factory.py index b8f6565f8..fead85d87 100644 --- a/src/uipath_langchain/runtime/factory.py +++ b/src/uipath_langchain/runtime/factory.py @@ -2,6 +2,7 @@ import os from typing import Any, AsyncContextManager +from langchain_core.callbacks import BaseCallbackHandler from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver from langgraph.graph.state import CompiledStateGraph, StateGraph from openinference.instrumentation.langchain import ( @@ -9,6 +10,7 @@ get_ancestor_spans, get_current_span, ) +from uipath.core.adapters import EvaluatorProtocol from uipath.core.tracing import UiPathSpanUtils, UiPathTraceManager from uipath.platform.resume_triggers import ( UiPathResumeTriggerHandler, @@ -23,6 +25,7 @@ from uipath.runtime.errors import UiPathErrorCategory from uipath_langchain._tracing import _instrument_traceable_attributes +from uipath_langchain.governance import GovernanceCallbackHandler from uipath_langchain.runtime.config import LangGraphConfig from uipath_langchain.runtime.errors import LangGraphErrorCode, LangGraphRuntimeError from uipath_langchain.runtime.graph import LangGraphLoader @@ -254,6 +257,7 @@ async def _create_runtime_instance( compiled_graph: CompiledStateGraph[Any, Any, Any, Any], runtime_id: str, entrypoint: str, + evaluator: EvaluatorProtocol | None = None, **kwargs, ) -> UiPathRuntimeProtocol: """ @@ -263,6 +267,9 @@ async def _create_runtime_instance( compiled_graph: The compiled graph runtime_id: Unique identifier for the runtime instance entrypoint: Graph entrypoint name + evaluator: Optional governance evaluator. When supplied, the + factory builds a :class:`GovernanceCallbackHandler` and + hands it to the runtime via its ``callbacks`` arg. Returns: Configured runtime instance @@ -271,10 +278,23 @@ async def _create_runtime_instance( storage = SqliteResumableStorage(memory) trigger_manager = UiPathResumeTriggerHandler() + callbacks: list[BaseCallbackHandler] | None = ( + [ + GovernanceCallbackHandler( + evaluator=evaluator, + agent_name=entrypoint, + session_id=runtime_id, + ) + ] + if evaluator is not None + else None + ) + base_runtime = UiPathLangGraphRuntime( graph=compiled_graph, runtime_id=runtime_id, entrypoint=entrypoint, + callbacks=callbacks, storage=storage, ) @@ -286,7 +306,11 @@ async def _create_runtime_instance( ) async def new_runtime( - self, entrypoint: str, runtime_id: str, **kwargs + self, + entrypoint: str, + runtime_id: str, + evaluator: EvaluatorProtocol | None = None, + **kwargs, ) -> UiPathRuntimeProtocol: """ Create a new LangGraph runtime instance. @@ -294,6 +318,9 @@ async def new_runtime( Args: entrypoint: Graph name from langgraph.json runtime_id: Unique identifier for the runtime instance + evaluator: Optional governance evaluator. When supplied, the + factory wires a :class:`GovernanceCallbackHandler` into + the runtime's callback list. Returns: Configured runtime instance with compiled graph @@ -309,6 +336,7 @@ async def new_runtime( compiled_graph=compiled_graph, runtime_id=runtime_id, entrypoint=entrypoint, + evaluator=evaluator, **kwargs, ) diff --git a/tests/governance/test_adapter.py b/tests/governance/test_callbacks.py similarity index 56% rename from tests/governance/test_adapter.py rename to tests/governance/test_callbacks.py index eef162670..d634548a0 100644 --- a/tests/governance/test_adapter.py +++ b/tests/governance/test_callbacks.py @@ -1,25 +1,18 @@ -"""Tests for the LangChain governance adapter.""" +"""Tests for the LangChain governance callback handler.""" from __future__ import annotations -import sys from types import SimpleNamespace -from typing import Any, TypedDict from unittest.mock import MagicMock, patch import pytest -from uipath.core.adapters import get_adapter_registry, reset_adapter_registry +from langchain_core.callbacks import BaseCallbackHandler from uipath.core.governance.exceptions import GovernanceBlockException -from uipath_langchain.governance import register_governance_adapter -from uipath_langchain.governance.adapter import ( - GovernanceCallbackHandler, - GovernedLangChainAgent, - LangChainAdapter, - _add_callback, -) +from uipath_langchain.governance import GovernanceCallbackHandler +from uipath_langchain.governance.callbacks import _BEFORE_MODEL_TEXT_CAP -LOGGER_PATH = "uipath_langchain.governance.adapter.logger" +LOGGER_PATH = "uipath_langchain.governance.callbacks.logger" @pytest.fixture @@ -27,11 +20,6 @@ def evaluator() -> MagicMock: return MagicMock() -@pytest.fixture -def adapter() -> LangChainAdapter: - return LangChainAdapter() - - @pytest.fixture def handler(evaluator: MagicMock) -> GovernanceCallbackHandler: return GovernanceCallbackHandler( @@ -41,324 +29,26 @@ def handler(evaluator: MagicMock) -> GovernanceCallbackHandler: ) -class TestCanHandle: - def test_returns_true_for_langgraph_compiled_state_graph( - self, adapter: LangChainAdapter - ) -> None: - from langgraph.graph import StateGraph - - class S(TypedDict): - v: int - - graph = StateGraph(S) - graph.add_node("n", lambda s: s) - graph.set_entry_point("n") - compiled = graph.compile() - assert adapter.can_handle(compiled) is True - - def test_returns_true_for_langchain_runnable( - self, adapter: LangChainAdapter - ) -> None: - from langchain_core.runnables import RunnableLambda - - runnable = RunnableLambda(lambda x: x) - assert adapter.can_handle(runnable) is True - - def test_returns_false_for_duck_typed_invoke_ainvoke( - self, adapter: LangChainAdapter - ) -> None: - # A non-LangChain object that merely exposes invoke/ainvoke must - # NOT be claimed — the adapter only handles LangChain/LangGraph. - class Duck: - def invoke(self, x): - return x - - async def ainvoke(self, x): - return x - - assert adapter.can_handle(Duck()) is False - - @pytest.mark.parametrize( - "module_name", - ["autogen.foo", "crewai.agent", "llama_index.core", "pydantic_ai.bar"], - ) - def test_returns_false_for_other_frameworks( - self, adapter: LangChainAdapter, module_name: str - ) -> None: - class Foreign: - def invoke(self, x): - return x - - async def ainvoke(self, x): - return x - - Foreign.__module__ = module_name - assert adapter.can_handle(Foreign()) is False - - def test_returns_false_for_object_without_invoke( - self, adapter: LangChainAdapter - ) -> None: - assert adapter.can_handle(object()) is False - - def test_returns_false_for_object_with_only_invoke( - self, adapter: LangChainAdapter - ) -> None: - class Half: - def invoke(self, x): - return x - - assert adapter.can_handle(Half()) is False - - def test_handles_langgraph_import_failure( - self, monkeypatch: pytest.MonkeyPatch, adapter: LangChainAdapter - ) -> None: - monkeypatch.setitem(sys.modules, "langgraph.graph.state", None) - from langchain_core.runnables import RunnableLambda - - assert adapter.can_handle(RunnableLambda(lambda x: x)) is True - - def test_handles_langchain_core_import_failure( - self, monkeypatch: pytest.MonkeyPatch, adapter: LangChainAdapter - ) -> None: - # With both framework imports unavailable, nothing can be claimed. - monkeypatch.setitem(sys.modules, "langgraph.graph.state", None) - monkeypatch.setitem(sys.modules, "langchain_core.runnables", None) - - class Duck: - def invoke(self, x): - return x - - async def ainvoke(self, x): - return x - - assert adapter.can_handle(Duck()) is False +class TestSubclassesBaseCallbackHandler: + def test_is_base_callback_handler(self, handler: GovernanceCallbackHandler) -> None: + # Closes governance-architecture-review §3.2: the handler must be + # a real LangChain BaseCallbackHandler so LangChain's dispatch / + # tracer wiring treats it natively. + assert isinstance(handler, BaseCallbackHandler) - -class TestAttach: - def test_returns_governed_agent( - self, adapter: LangChainAdapter, evaluator: MagicMock - ) -> None: - class Agent: - callbacks: list[Any] = [] - - a = Agent() - governed = adapter.attach(a, "agent-id", "session-id", evaluator) - assert isinstance(governed, GovernedLangChainAgent) - assert governed.unwrapped is a - - def test_injects_callback_into_existing_callback_list( - self, adapter: LangChainAdapter, evaluator: MagicMock - ) -> None: - prior = object() - - class Agent: - callbacks: list[Any] = [prior] - - a = Agent() - governed = adapter.attach(a, "id", "session", evaluator) - assert a.callbacks[0] is prior - assert a.callbacks[-1] is governed._callback - - def test_replaces_non_list_callbacks_attribute( - self, adapter: LangChainAdapter, evaluator: MagicMock - ) -> None: - sentinel = object() - - class Agent: - callbacks = sentinel # truthy but not a list - - a = Agent() - governed = adapter.attach(a, "id", "session", evaluator) - assert isinstance(a.callbacks, list) - assert a.callbacks == [governed._callback] - - def test_injects_into_empty_config( - self, adapter: LangChainAdapter, evaluator: MagicMock - ) -> None: - class Agent: - config: dict[str, Any] = {} - - a = Agent() - governed = adapter.attach(a, "id", "session", evaluator) - assert a.config["callbacks"] == [governed._callback] - - def test_injects_into_none_config( - self, adapter: LangChainAdapter, evaluator: MagicMock - ) -> None: - class Agent: - config = None - - a = Agent() - governed = adapter.attach(a, "id", "session", evaluator) - assert isinstance(a.config, dict) - assert a.config["callbacks"] == [governed._callback] - - def test_logs_warning_when_no_callback_surface( - self, - adapter: LangChainAdapter, - evaluator: MagicMock, - ) -> None: - class Bare: - pass - - bare = Bare() - with patch(LOGGER_PATH) as mock_logger: - adapter.attach(bare, "id", "session", evaluator) - mock_logger.warning.assert_called_once() - assert "Could not inject" in mock_logger.warning.call_args.args[0] - assert not hasattr(bare, "callbacks") - assert not hasattr(bare, "config") - - def test_callbacks_path_logs_debug( - self, - adapter: LangChainAdapter, - evaluator: MagicMock, - ) -> None: - class Agent: - callbacks: list[Any] = [] - - with patch(LOGGER_PATH) as mock_logger: - adapter.attach(Agent(), "id", "session", evaluator) - debug_msgs = [c.args[0] for c in mock_logger.debug.call_args_list] - assert any("agent.callbacks" in m for m in debug_msgs) - - def test_config_path_logs_debug( - self, - adapter: LangChainAdapter, - evaluator: MagicMock, - ) -> None: - class Agent: - config: dict[str, Any] = {} - - with patch(LOGGER_PATH) as mock_logger: - adapter.attach(Agent(), "id", "session", evaluator) - debug_msgs = [c.args[0] for c in mock_logger.debug.call_args_list] - assert any("agent.config" in m for m in debug_msgs) - - -class TestAdapterMetadata: - def test_name(self, adapter: LangChainAdapter) -> None: - assert adapter.name == "LangChain" - - -class TestGovernedLangChainAgent: - def _governed( - self, adapter: LangChainAdapter, evaluator: MagicMock, agent: object - ) -> GovernedLangChainAgent: - return adapter.attach(agent, "id", "session", evaluator) - - def test_invoke_injects_callback_and_returns_result( - self, adapter: LangChainAdapter, evaluator: MagicMock - ) -> None: - captured: dict[str, Any] = {} - - class Agent: - callbacks: list[Any] = [] - - def invoke(self, x, config=None, **kw): - captured["config"] = config - captured["x"] = x - return "out" - - gov = self._governed(adapter, evaluator, Agent()) - assert gov.invoke("in") == "out" - assert captured["x"] == "in" - assert gov._callback in captured["config"]["callbacks"] - - async def test_ainvoke_injects_callback_and_returns_result( - self, adapter: LangChainAdapter, evaluator: MagicMock - ) -> None: - captured: dict[str, Any] = {} - - class Agent: - callbacks: list[Any] = [] - - async def ainvoke(self, x, config=None, **kw): - captured["config"] = config - return "async-out" - - gov = self._governed(adapter, evaluator, Agent()) - assert await gov.ainvoke("in") == "async-out" - assert gov._callback in captured["config"]["callbacks"] - - def test_stream_yields_chunks_with_callback( - self, adapter: LangChainAdapter, evaluator: MagicMock - ) -> None: - captured: dict[str, Any] = {} - - class Agent: - callbacks: list[Any] = [] - - def stream(self, x, config=None, **kw): - captured["config"] = config - yield "a" - yield "b" - - gov = self._governed(adapter, evaluator, Agent()) - assert list(gov.stream("x")) == ["a", "b"] - assert gov._callback in captured["config"]["callbacks"] - - async def test_astream_yields_chunks_with_callback( - self, adapter: LangChainAdapter, evaluator: MagicMock - ) -> None: - captured: dict[str, Any] = {} - - class Agent: - callbacks: list[Any] = [] - - async def astream(self, x, config=None, **kw): - captured["config"] = config - yield "a" - yield "b" - - gov = self._governed(adapter, evaluator, Agent()) - chunks = [c async for c in gov.astream("x")] - assert chunks == ["a", "b"] - assert gov._callback in captured["config"]["callbacks"] - - def test_ensure_callback_config_with_none( - self, adapter: LangChainAdapter, evaluator: MagicMock - ) -> None: - gov = self._governed(adapter, evaluator, SimpleNamespace(callbacks=[])) - result = gov._ensure_callback_config(None) - assert isinstance(result, dict) - assert result["callbacks"] == [gov._callback] - - def test_ensure_callback_config_preserves_other_keys( - self, adapter: LangChainAdapter, evaluator: MagicMock - ) -> None: - gov = self._governed(adapter, evaluator, SimpleNamespace(callbacks=[])) - config = {"metadata": {"k": "v"}, "callbacks": []} - result = gov._ensure_callback_config(config) - assert result is config - assert result["metadata"] == {"k": "v"} - assert gov._callback in result["callbacks"] - - def test_ensure_callback_config_is_idempotent( - self, adapter: LangChainAdapter, evaluator: MagicMock - ) -> None: - gov = self._governed(adapter, evaluator, SimpleNamespace(callbacks=[])) - first = gov._ensure_callback_config(None) - second = gov._ensure_callback_config(first) - assert second["callbacks"].count(gov._callback) == 1 - - def test_ensure_callback_config_passes_through_non_dict( - self, adapter: LangChainAdapter, evaluator: MagicMock - ) -> None: - gov = self._governed(adapter, evaluator, SimpleNamespace(callbacks=[])) - sentinel = ["not", "a", "dict"] - result = gov._ensure_callback_config(sentinel) - assert result is sentinel - - def test_getattr_forwards_to_wrapped_agent( - self, adapter: LangChainAdapter, evaluator: MagicMock + def test_ignore_flags_override_parent_properties( + self, handler: GovernanceCallbackHandler ) -> None: - class Agent: - callbacks: list[Any] = [] - answer = 42 - - gov = self._governed(adapter, evaluator, Agent()) - assert gov.answer == 42 + # Chain notifications skipped — the governance host owns + # BEFORE_AGENT / AFTER_AGENT and would otherwise double-fire. + assert handler.ignore_chain is True + assert handler.ignore_retriever is True + assert handler.ignore_retry is True + assert handler.ignore_custom_event is True + # LLM / chat model / tool / agent events stay on. + assert handler.ignore_llm is False + assert handler.ignore_chat_model is False + assert handler.ignore_agent is False class TestCallbackHandlerLLM: @@ -487,8 +177,6 @@ def test_on_chat_model_start_caps_model_input( self, handler: GovernanceCallbackHandler, evaluator: MagicMock ) -> None: """``model_input`` is bounded so a runaway prompt can't dominate scan time.""" - from uipath_langchain.governance.adapter import _BEFORE_MODEL_TEXT_CAP - huge = SimpleNamespace(content="x" * (_BEFORE_MODEL_TEXT_CAP + 1000)) handler.on_chat_model_start({}, [[huge]]) model_input = evaluator.evaluate_before_model.call_args.kwargs["model_input"] @@ -498,8 +186,6 @@ def test_on_chat_model_start_block_list_stops_at_remaining_budget( self, handler: GovernanceCallbackHandler, evaluator: MagicMock ) -> None: """The block walk exits early once the per-call cap is exhausted.""" - from uipath_langchain.governance.adapter import _BEFORE_MODEL_TEXT_CAP - first = "a" * _BEFORE_MODEL_TEXT_CAP # consumes the entire budget latest = SimpleNamespace( content=[ @@ -574,8 +260,6 @@ def test_on_llm_end_caps_model_output( """A runaway / batched response is capped so the AFTER_MODEL scan budget matches BEFORE_MODEL and the runtime side's cap. """ - from uipath_langchain.governance.adapter import _BEFORE_MODEL_TEXT_CAP - # Many large generations across batched gen_lists. gen = SimpleNamespace(text="y" * 50_000, message=None) response = SimpleNamespace(generations=[[gen], [gen, gen]]) @@ -886,183 +570,3 @@ def test_session_state_initialized(self, evaluator: MagicMock) -> None: assert h._agent_name == "a" assert h._session_id == "s" assert h._trace_id # uuid4 string - - -class TestRegisterAdapter: - def setup_method(self, method) -> None: - reset_adapter_registry() - - def teardown_method(self, method) -> None: - reset_adapter_registry() - register_governance_adapter() - - def test_registers_adapter(self) -> None: - register_governance_adapter() - names = [a.name for a in get_adapter_registry().get_all()] - assert names.count("LangChain") == 1 - - def test_is_idempotent(self) -> None: - register_governance_adapter() - register_governance_adapter() - register_governance_adapter() - names = [a.name for a in get_adapter_registry().get_all()] - assert names.count("LangChain") == 1 - - def test_skips_when_adapter_with_same_name_already_present(self) -> None: - registry = get_adapter_registry() - # Drop anything the entry-point discovery may have added. - registry.clear() - - class OtherLang(LangChainAdapter): - pass - - registry.register(OtherLang()) - - register_governance_adapter() - - names = [a.name for a in registry.get_all()] - assert names.count("LangChain") == 1 - # The pre-existing instance was kept (not replaced). - assert isinstance(registry.get_all()[0], OtherLang) - - -class TestAddCallback: - """Direct tests for the ``_add_callback`` shape-normalizer.""" - - def test_none_returns_new_list(self) -> None: - cb = object() - container, replaced = _add_callback(None, cb) # type: ignore[arg-type] - assert container == [cb] - assert replaced is True - - def test_list_is_mutated_in_place(self) -> None: - cb = object() - existing: list[Any] = [] - container, replaced = _add_callback(existing, cb) # type: ignore[arg-type] - assert container is existing - assert existing == [cb] - assert replaced is False - - def test_list_is_idempotent(self) -> None: - cb = object() - existing: list[Any] = [cb] - container, replaced = _add_callback(existing, cb) # type: ignore[arg-type] - assert existing.count(cb) == 1 - assert replaced is False - - def test_tuple_is_coerced_to_list(self) -> None: - prior = object() - cb = object() - container, replaced = _add_callback((prior,), cb) # type: ignore[arg-type] - assert container == [prior, cb] - assert replaced is True - - def test_base_callback_manager_is_mutated_in_place(self) -> None: - """A duck-typed callback manager keeps its identity and tracers.""" - - class FakeManager: - def __init__(self) -> None: - self.handlers: list[Any] = [] - self.inherit_calls: list[bool] = [] - - def add_handler(self, handler: Any, inherit: bool) -> None: - self.handlers.append(handler) - self.inherit_calls.append(inherit) - - manager = FakeManager() - prior = object() - manager.handlers.append(prior) - - cb = object() - container, replaced = _add_callback(manager, cb) # type: ignore[arg-type] - assert container is manager - assert replaced is False - assert manager.handlers == [prior, cb] - assert manager.inherit_calls == [True] - - def test_base_callback_manager_is_idempotent(self) -> None: - class FakeManager: - def __init__(self) -> None: - self.handlers: list[Any] = [] - - def add_handler(self, handler: Any, inherit: bool) -> None: - self.handlers.append(handler) - - cb = object() - manager = FakeManager() - manager.handlers.append(cb) - _add_callback(manager, cb) # type: ignore[arg-type] - assert manager.handlers.count(cb) == 1 - - def test_non_iterable_falls_back_to_new_list(self) -> None: - cb = object() - container, replaced = _add_callback(object(), cb) # type: ignore[arg-type] - assert container == [cb] - assert replaced is True - - -class TestInjectCallbackShapes: - """End-to-end injection across the supported callback container shapes.""" - - def test_injects_into_tuple_callbacks_via_replacement( - self, adapter: LangChainAdapter, evaluator: MagicMock - ) -> None: - prior = object() - - class Agent: - callbacks = (prior,) - - a = Agent() - governed = adapter.attach(a, "id", "session", evaluator) - assert isinstance(a.callbacks, list) - assert a.callbacks == [prior, governed._callback] - - def test_injects_into_base_callback_manager_in_place( - self, adapter: LangChainAdapter, evaluator: MagicMock - ) -> None: - class FakeManager: - def __init__(self) -> None: - self.handlers: list[Any] = [] - - def add_handler(self, handler: Any, inherit: bool) -> None: - self.handlers.append(handler) - - manager = FakeManager() - - class Agent: - callbacks = manager - - a = Agent() - governed = adapter.attach(a, "id", "session", evaluator) - assert a.callbacks is manager - assert manager.handlers == [governed._callback] - - def test_injects_into_config_with_tuple_callbacks( - self, adapter: LangChainAdapter, evaluator: MagicMock - ) -> None: - prior = object() - - class Agent: - config: dict[str, Any] = {"callbacks": (prior,)} - - a = Agent() - governed = adapter.attach(a, "id", "session", evaluator) - assert a.config["callbacks"] == [prior, governed._callback] - - def test_ensure_callback_config_with_base_callback_manager( - self, adapter: LangChainAdapter, evaluator: MagicMock - ) -> None: - class FakeManager: - def __init__(self) -> None: - self.handlers: list[Any] = [] - - def add_handler(self, handler: Any, inherit: bool) -> None: - self.handlers.append(handler) - - manager = FakeManager() - gov = adapter.attach(SimpleNamespace(callbacks=[]), "id", "session", evaluator) - config = {"callbacks": manager} - result = gov._ensure_callback_config(config) - assert result is config - assert result["callbacks"] is manager - assert manager.handlers == [gov._callback] diff --git a/tests/runtime/test_factory_governance.py b/tests/runtime/test_factory_governance.py new file mode 100644 index 000000000..5058a10da --- /dev/null +++ b/tests/runtime/test_factory_governance.py @@ -0,0 +1,116 @@ +"""Factory-level governance wiring: evaluator -> callbacks plumbing.""" + +from __future__ import annotations + +import os +import tempfile +from typing import Any, TypedDict +from unittest.mock import MagicMock + +import pytest +from langgraph.graph import END, START, StateGraph +from uipath.core.adapters import EvaluatorProtocol +from uipath.runtime import UiPathRuntimeContext + +from uipath_langchain.governance import GovernanceCallbackHandler +from uipath_langchain.runtime.factory import UiPathLangGraphRuntimeFactory + + +class _State(TypedDict): + v: int + + +def _build_graph() -> StateGraph[Any, Any, Any]: + g = StateGraph(_State) + g.add_node("noop", lambda s: s) + g.add_edge(START, "noop") + g.add_edge("noop", END) + return g + + +@pytest.fixture +def context() -> UiPathRuntimeContext: + tmpdir = tempfile.mkdtemp() + ctx = UiPathRuntimeContext( + runtime_dir=tmpdir, + state_file=os.path.join(tmpdir, "state.db"), + ) + return ctx + + +@pytest.fixture +def factory(context: UiPathRuntimeContext) -> UiPathLangGraphRuntimeFactory: + return UiPathLangGraphRuntimeFactory(context) + + +class TestEvaluatorWiring: + """Passing ``evaluator`` to ``new_runtime`` should attach a + :class:`GovernanceCallbackHandler` to the underlying LangGraph + runtime's callback list. This is the entire surface change — the + previous adapter / register-on-import path is gone. + """ + + async def test_no_evaluator_means_no_callbacks( + self, factory: UiPathLangGraphRuntimeFactory + ) -> None: + compiled = _build_graph().compile() + await factory._get_memory() + runtime = await factory._create_runtime_instance( + compiled_graph=compiled, + runtime_id="rt-1", + entrypoint="ep", + ) + # The resumable runtime wraps the langgraph runtime as ``delegate``. + assert runtime.delegate.callbacks == [] # type: ignore[attr-defined] + await factory.dispose() + + async def test_evaluator_attaches_governance_handler( + self, factory: UiPathLangGraphRuntimeFactory + ) -> None: + evaluator: EvaluatorProtocol = MagicMock(spec=EvaluatorProtocol) + compiled = _build_graph().compile() + await factory._get_memory() # ensure memory is initialized + runtime = await factory._create_runtime_instance( + compiled_graph=compiled, + runtime_id="rt-1", + entrypoint="ep", + evaluator=evaluator, + ) + callbacks = runtime.delegate.callbacks # type: ignore[attr-defined] + assert len(callbacks) == 1 + handler = callbacks[0] + assert isinstance(handler, GovernanceCallbackHandler) + # Identity / session_id / agent_name come from the factory args. + assert handler._evaluator is evaluator + assert handler._agent_name == "ep" + assert handler._session_id == "rt-1" + await factory.dispose() + + async def test_handler_built_per_runtime_instance( + self, factory: UiPathLangGraphRuntimeFactory + ) -> None: + """Two factory calls with the same evaluator yield two distinct + handler instances — each runtime gets its own trace_id and + session_state, so concurrent sessions don't share counters.""" + evaluator: EvaluatorProtocol = MagicMock(spec=EvaluatorProtocol) + compiled = _build_graph().compile() + await factory._get_memory() + first = await factory._create_runtime_instance( + compiled_graph=compiled, + runtime_id="rt-a", + entrypoint="ep", + evaluator=evaluator, + ) + second = await factory._create_runtime_instance( + compiled_graph=compiled, + runtime_id="rt-b", + entrypoint="ep", + evaluator=evaluator, + ) + h1 = first.delegate.callbacks[0] # type: ignore[attr-defined] + h2 = second.delegate.callbacks[0] # type: ignore[attr-defined] + assert h1 is not h2 + assert h1._trace_id != h2._trace_id + assert h1._session_id == "rt-a" + assert h2._session_id == "rt-b" + await factory.dispose() diff --git a/uv.lock b/uv.lock index 51ceb50ae..ff51689ac 100644 --- a/uv.lock +++ b/uv.lock @@ -4374,16 +4374,16 @@ wheels = [ [[package]] name = "uipath-core" -version = "0.5.20" -source = { registry = "https://pypi.org/simple" } +version = "0.5.23.dev1017616946" +source = { registry = "https://test.pypi.org/simple/" } dependencies = [ { name = "opentelemetry-instrumentation" }, { name = "opentelemetry-sdk" }, { name = "pydantic" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/93/da/011ced5af57363caf7ca6d263261fc4b64f19bf7f7b2a5e54132906a36a6/uipath_core-0.5.20.tar.gz", hash = "sha256:2a2430185522869b10c05273128c23a81fbb8c53ee5dd8686c8b5089ea270fa7", size = 132363, upload-time = "2026-06-19T12:01:37.545Z" } +sdist = { url = "https://test-files.pythonhosted.org/packages/49/a7/5683de37c2e12168e1b9d0da80c17411e77ca6e600134de0e1584b36eab6/uipath_core-0.5.23.dev1017616946.tar.gz", hash = "sha256:46b36d45ea984d6ebce134e5ee171b615a4d6f53bff7a0aa2e9c26f0d87da117", size = 129588, upload-time = "2026-06-25T03:15:30.649Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/3c/02/b518bb5569c9c35f4ed19a7f23c744c471352a1fa5233071dd75a4cc9324/uipath_core-0.5.20-py3-none-any.whl", hash = "sha256:2be116ec68d034348ea58fd541675863d73e96649f528648d533206b8c8853cc", size = 55005, upload-time = "2026-06-19T12:01:36.08Z" }, + { url = "https://test-files.pythonhosted.org/packages/85/f9/f3e7bea29d70951167907171117fa4494abf64f9ccfd9b66d22b8604e7d6/uipath_core-0.5.23.dev1017616946-py3-none-any.whl", hash = "sha256:c988ee6b2ca2d7093f2cd7d46b7c9bb3bb504a09a0d76aa3c61d9ce1003dac16", size = 54265, upload-time = "2026-06-25T03:15:29.663Z" }, ] [[package]] @@ -4464,7 +4464,7 @@ requires-dist = [ { name = "pydantic-settings", specifier = ">=2.6.0" }, { name = "python-dotenv", specifier = ">=1.0.1" }, { name = "uipath", specifier = ">=2.10.79,<2.12.0" }, - { name = "uipath-core", specifier = ">=0.5.20,<0.6.0" }, + { name = "uipath-core", specifier = "==0.5.23.dev1017616946", index = "https://test.pypi.org/simple/" }, { name = "uipath-langchain-client", extras = ["all"], marker = "extra == 'all'", specifier = ">=1.14.0,<1.15.0" }, { name = "uipath-langchain-client", extras = ["anthropic"], marker = "extra == 'anthropic'", specifier = ">=1.14.0,<1.15.0" }, { name = "uipath-langchain-client", extras = ["bedrock"], marker = "extra == 'bedrock'", specifier = ">=1.14.0,<1.15.0" }, From 768414cbe27ec9172a2e6a4fbbb724270c079c14 Mon Sep 17 00:00:00 2001 From: Viswanath Lekshmanan Date: Fri, 26 Jun 2026 16:31:21 +0530 Subject: [PATCH 10/10] fix: changes for removing trace_id --- src/uipath_langchain/governance/callbacks.py | 15 ++++++++------- tests/governance/test_callbacks.py | 9 +++++++-- 2 files changed, 15 insertions(+), 9 deletions(-) diff --git a/src/uipath_langchain/governance/callbacks.py b/src/uipath_langchain/governance/callbacks.py index 5d796eac7..5ccbe5f37 100644 --- a/src/uipath_langchain/governance/callbacks.py +++ b/src/uipath_langchain/governance/callbacks.py @@ -32,7 +32,6 @@ import logging from typing import Any, Dict, Iterable -from uuid import uuid4 from langchain_core.callbacks import BaseCallbackHandler from uipath.core.adapters import EvaluatorProtocol @@ -76,7 +75,14 @@ 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. The runtime layer + # is env-free for this code path: the platform-side + # ``resolve_trace_id`` resolves the canonical agent trace id at + # HTTP-call time, and the governance compensator propagates the + # live OTel context across its thread-pool hop via + # ``contextvars.copy_context()``. The callback handler just + # forwards extracted payload to the evaluator and lets the + # plumbing below it resolve the trace id wherever it's needed. self._session_state: Dict[str, Any] = {"tool_calls": 0, "llm_calls": 0} # Tool name lookup keyed by LangChain ``run_id`` so ``on_tool_end`` # can report the actual tool name to AFTER_TOOL evaluation. @@ -103,7 +109,6 @@ def on_llm_start( model_input=model_input, agent_name=self._agent_name, runtime_id=self._session_id, - trace_id=self._trace_id, ) except GovernanceBlockException: raise @@ -140,7 +145,6 @@ def on_chat_model_start( model_input=model_input, agent_name=self._agent_name, runtime_id=self._session_id, - trace_id=self._trace_id, ) except GovernanceBlockException: raise @@ -227,7 +231,6 @@ def on_llm_end(self, response: Any, **kwargs: Any) -> None: model_output=model_output, agent_name=self._agent_name, runtime_id=self._session_id, - trace_id=self._trace_id, ) except GovernanceBlockException: raise @@ -349,7 +352,6 @@ def on_tool_start( tool_args=tool_args, agent_name=self._agent_name, runtime_id=self._session_id, - trace_id=self._trace_id, session_state=self._session_state, ) except GovernanceBlockException: @@ -374,7 +376,6 @@ def on_tool_end(self, output: Any, **kwargs: Any) -> None: tool_result=tool_result, agent_name=self._agent_name, runtime_id=self._session_id, - trace_id=self._trace_id, ) except GovernanceBlockException: raise diff --git a/tests/governance/test_callbacks.py b/tests/governance/test_callbacks.py index d634548a0..775b32a6d 100644 --- a/tests/governance/test_callbacks.py +++ b/tests/governance/test_callbacks.py @@ -64,7 +64,10 @@ def test_on_llm_start_invokes_evaluator_with_latest_prompt( assert kwargs["model_input"] == "b" assert kwargs["agent_name"] == "test-agent" assert kwargs["runtime_id"] == "test-session" - assert kwargs["trace_id"] == handler._trace_id + # ``trace_id`` is intentionally NOT passed — the platform layer + # resolves the canonical trace id at HTTP-call time via + # ``resolve_trace_id``. The callback handler stays env-free. + assert "trace_id" not in kwargs def test_on_llm_start_increments_counter( self, handler: GovernanceCallbackHandler @@ -569,4 +572,6 @@ def test_session_state_initialized(self, evaluator: MagicMock) -> None: assert h._session_state == {"tool_calls": 0, "llm_calls": 0} assert h._agent_name == "a" assert h._session_id == "s" - assert h._trace_id # uuid4 string + # ``trace_id`` is no longer held on the handler — the runtime + # layer is env-free; the platform side resolves at HTTP time. + assert not hasattr(h, "_trace_id")