1+ """LlamaIndex adapter for UiPath governance.
2+
3+ Provides governance for LlamaIndex agents/workflows. Unlike the ADK / OpenAI /
4+ Agent-Framework adapters — which install per-agent callbacks or middleware —
5+ LlamaIndex routes everything (LLM calls, tool calls) through its global
6+ **instrumentation dispatcher** (the same mechanism the package already uses for
7+ OpenInference tracing). So this adapter governs by registering a
8+ :class:`GovernanceEventHandler` on the **root dispatcher**, which receives every
9+ event propagated from child dispatchers:
10+
11+ - ``LLMChatStartEvent`` → BEFORE_MODEL (scans the latest input message)
12+ - ``LLMChatEndEvent`` → AFTER_MODEL (scans the response)
13+ - ``AgentToolCallEvent`` → TOOL_CALL (tool name + arguments)
14+
15+ The dispatcher is process-global, so registration is process-wide — which fits
16+ the coded-agent model (one workflow per process). :meth:`attach` therefore
17+ returns the ``agent`` unchanged (nothing is mutated on it); the wiring lives on
18+ the dispatcher. :meth:`detach` removes the handler.
19+
20+ LlamaIndex does **not** emit a tool-*end* instrumentation event, so AFTER_TOOL
21+ is not wired here; a tool's result is governed at the next ``LLMChatStartEvent``
22+ where it is fed back to the model as input (analogous to how the OpenAI adapter
23+ handles its missing tool-args).
24+
25+ Chain-level boundaries (BEFORE_AGENT / AFTER_AGENT) are owned by the runtime
26+ wrapper layer in ``uipath-runtime`` and are intentionally not fired here.
27+
28+ Contracts and the evaluator protocol come from ``uipath-core``; this package
29+ contributes only the LlamaIndex-specific implementation and self-registers it
30+ with the global adapter registry when ``uipath_llamaindex.governance`` is
31+ imported.
32+
33+ Audit emission and enforcement (raising :class:`GovernanceBlockException` on
34+ DENY) are owned by the evaluator. The handler only extracts payloads and calls
35+ the matching ``evaluate_*`` method; :class:`GovernanceBlockException` propagates
36+ (aborting the run), anything else is logged and swallowed.
37+ """
38+
39+ from __future__ import annotations
40+
41+ import json
42+ import logging
43+ from typing import Any , Dict , List
44+ from uuid import uuid4
45+
46+ from llama_index .core .instrumentation import (
47+ get_dispatcher , # type: ignore[attr-defined]
48+ )
49+ from llama_index .core .instrumentation .event_handlers .base import ( # type: ignore[attr-defined]
50+ BaseEventHandler ,
51+ )
52+ from llama_index .core .instrumentation .events .agent import AgentToolCallEvent
53+ from llama_index .core .instrumentation .events .llm import (
54+ LLMChatEndEvent ,
55+ LLMChatStartEvent ,
56+ )
57+ from pydantic import PrivateAttr
58+ from uipath .core .adapters import BaseAdapter , EvaluatorProtocol
59+ from uipath .core .governance .exceptions import GovernanceBlockException
60+
61+ logger = logging .getLogger (__name__ )
62+
63+ # Cap on the text blob passed to BEFORE_MODEL / AFTER_MODEL governance
64+ # evaluation. Sized to match the runtime side and the other adapters.
65+ _BEFORE_MODEL_TEXT_CAP = 64000
66+
67+
68+ class LlamaIndexAdapter (BaseAdapter ):
69+ """Adapter for the LlamaIndex framework.
70+
71+ Detects LlamaIndex workflows/agents and governs them by registering a
72+ :class:`GovernanceEventHandler` on the root instrumentation dispatcher.
73+ """
74+
75+ @property
76+ def name (self ) -> str :
77+ return "LlamaIndex"
78+
79+ def can_handle (self , agent : Any ) -> bool :
80+ """Return True if this looks like a LlamaIndex workflow/agent."""
81+ try :
82+ from workflows import Workflow
83+
84+ if isinstance (agent , Workflow ):
85+ return True
86+ except ImportError :
87+ pass
88+
89+ # Duck-typed fallback: a workflow/agent exposes ``run`` and a workflow
90+ # step surface (``_get_steps`` / ``steps``) or a Workflow-shaped name.
91+ if hasattr (agent , "run" ) and (
92+ hasattr (agent , "_get_steps" )
93+ or hasattr (agent , "steps" )
94+ or type (agent ).__name__ .endswith (("Workflow" , "Agent" ))
95+ ):
96+ return True
97+ return False
98+
99+ def attach (
100+ self ,
101+ agent : Any ,
102+ agent_id : str ,
103+ session_id : str ,
104+ evaluator : EvaluatorProtocol ,
105+ ) -> Any :
106+ """Register the governance event handler on the root dispatcher.
107+
108+ Returns the ``agent`` unchanged — LlamaIndex governance is wired on the
109+ process-global dispatcher, not on the agent object. Idempotent: a
110+ second attach is a no-op while a handler is already registered.
111+ """
112+ dispatcher = get_dispatcher ()
113+ if any (isinstance (h , GovernanceEventHandler ) for h in dispatcher .event_handlers ):
114+ return agent # idempotent — already governed
115+ callbacks = GovernanceCallbacks (
116+ evaluator = evaluator , agent_name = agent_id , session_id = session_id
117+ )
118+ dispatcher .add_event_handler (GovernanceEventHandler (callbacks = callbacks ))
119+ logger .debug ("Registered governance event handler on LlamaIndex dispatcher" )
120+ return agent
121+
122+ def detach (self , governed : Any ) -> Any :
123+ """Remove the governance event handler from the root dispatcher."""
124+ dispatcher = get_dispatcher ()
125+ dispatcher .event_handlers = [
126+ h
127+ for h in dispatcher .event_handlers
128+ if not isinstance (h , GovernanceEventHandler )
129+ ]
130+ return governed
131+
132+
133+ class GovernanceEventHandler (BaseEventHandler ):
134+ """Routes LlamaIndex instrumentation events to a governance evaluator.
135+
136+ A pydantic model (``BaseEventHandler`` is one), so the evaluator + state
137+ are held in a private attribute. ``handle`` is called synchronously by the
138+ dispatcher for every event; we dispatch the three governance-relevant
139+ types and ignore the rest.
140+ """
141+
142+ _callbacks : "GovernanceCallbacks" = PrivateAttr ()
143+
144+ def __init__ (self , callbacks : "GovernanceCallbacks" , ** data : Any ) -> None :
145+ super ().__init__ (** data )
146+ self ._callbacks = callbacks
147+
148+ @classmethod
149+ def class_name (cls ) -> str :
150+ return "GovernanceEventHandler"
151+
152+ def handle (self , event : Any , ** kwargs : Any ) -> Any :
153+ if isinstance (event , LLMChatStartEvent ):
154+ self ._callbacks .before_model (event .messages )
155+ elif isinstance (event , LLMChatEndEvent ):
156+ self ._callbacks .after_model (event .response )
157+ elif isinstance (event , AgentToolCallEvent ):
158+ self ._callbacks .tool_call (event .tool , event .arguments )
159+ return None
160+
161+
162+ class GovernanceCallbacks :
163+ """Holds the evaluator + per-attach state, called by the event handler.
164+
165+ :class:`GovernanceBlockException` is re-raised (it aborts the run);
166+ anything else is logged and swallowed so a governance bug never breaks an
167+ agent run.
168+ """
169+
170+ def __init__ (
171+ self ,
172+ evaluator : EvaluatorProtocol ,
173+ agent_name : str ,
174+ session_id : str ,
175+ ) -> None :
176+ self ._evaluator = evaluator
177+ self ._agent_name = agent_name
178+ self ._session_id = session_id
179+ self ._trace_id = str (uuid4 ())
180+ self ._session_state : Dict [str , Any ] = {"tool_calls" : 0 , "llm_calls" : 0 }
181+
182+ def before_model (self , messages : Any ) -> None :
183+ """Evaluate BEFORE_MODEL on the latest input message (see ADK rationale)."""
184+ try :
185+ self ._session_state ["llm_calls" ] = (
186+ self ._session_state .get ("llm_calls" , 0 ) + 1
187+ )
188+ self ._evaluator .evaluate_before_model (
189+ model_input = _latest_message_text (messages ),
190+ agent_name = self ._agent_name ,
191+ runtime_id = self ._session_id ,
192+ trace_id = self ._trace_id ,
193+ )
194+ except GovernanceBlockException :
195+ raise
196+ except Exception as e : # noqa: BLE001 - governance must not break the run
197+ logger .warning ("before_model governance check failed (continuing): %s" , e )
198+
199+ def after_model (self , response : Any ) -> None :
200+ """Evaluate AFTER_MODEL on the chat response text."""
201+ try :
202+ self ._evaluator .evaluate_after_model (
203+ model_output = _response_text (response ),
204+ agent_name = self ._agent_name ,
205+ runtime_id = self ._session_id ,
206+ trace_id = self ._trace_id ,
207+ )
208+ except GovernanceBlockException :
209+ raise
210+ except Exception as e : # noqa: BLE001
211+ logger .warning ("after_model governance check failed (continuing): %s" , e )
212+
213+ def tool_call (self , tool : Any , arguments : Any ) -> None :
214+ """Evaluate TOOL_CALL with the tool name + arguments."""
215+ try :
216+ self ._session_state ["tool_calls" ] = (
217+ self ._session_state .get ("tool_calls" , 0 ) + 1
218+ )
219+ self ._evaluator .evaluate_tool_call (
220+ tool_name = getattr (tool , "name" , None ) or "unknown" ,
221+ tool_args = _coerce_args (arguments ),
222+ agent_name = self ._agent_name ,
223+ runtime_id = self ._session_id ,
224+ trace_id = self ._trace_id ,
225+ session_state = self ._session_state ,
226+ )
227+ except GovernanceBlockException :
228+ raise
229+ except Exception as e : # noqa: BLE001
230+ logger .warning ("tool_call governance check failed (continuing): %s" , e )
231+
232+
233+ # --------------------------------------------------------------------------
234+ # Text / argument extraction
235+ # --------------------------------------------------------------------------
236+
237+
238+ def _latest_message_text (messages : Any ) -> str :
239+ """Text of the most-recent message in a chat request."""
240+ if not messages :
241+ return ""
242+ if isinstance (messages , (list , tuple )):
243+ return _message_text (messages [- 1 ])
244+ return _message_text (messages )
245+
246+
247+ def _message_text (message : Any ) -> str :
248+ """Pull text from a ``ChatMessage`` (``.content``) or a bare string."""
249+ if message is None :
250+ return ""
251+ if isinstance (message , str ):
252+ return message [:_BEFORE_MODEL_TEXT_CAP ]
253+ content = getattr (message , "content" , None )
254+ if isinstance (content , str ) and content :
255+ return content [:_BEFORE_MODEL_TEXT_CAP ]
256+ # Newer ChatMessage carries typed blocks; fall back to str().
257+ return str (message )[:_BEFORE_MODEL_TEXT_CAP ]
258+
259+
260+ def _response_text (response : Any ) -> str :
261+ """Pull assistant text from a ``ChatResponse`` (``.message.content``)."""
262+ if response is None :
263+ return ""
264+ message = getattr (response , "message" , None )
265+ if message is not None :
266+ return _message_text (message )
267+ text = getattr (response , "text" , None )
268+ if isinstance (text , str ):
269+ return text [:_BEFORE_MODEL_TEXT_CAP ]
270+ return str (response )[:_BEFORE_MODEL_TEXT_CAP ]
271+
272+
273+ def _coerce_args (arguments : Any ) -> Dict [str , Any ]:
274+ """Normalise tool arguments (JSON string / Mapping / None) to a dict.
275+
276+ ``AgentToolCallEvent.arguments`` is a JSON-encoded string; other call
277+ sites may hand a dict directly.
278+ """
279+ if arguments is None :
280+ return {}
281+ if isinstance (arguments , dict ):
282+ return arguments
283+ if isinstance (arguments , str ):
284+ try :
285+ parsed = json .loads (arguments )
286+ return parsed if isinstance (parsed , dict ) else {"_" : parsed }
287+ except (TypeError , ValueError ):
288+ return {}
289+ return {}
290+
291+
292+ __all__ : List [str ] = [
293+ "GovernanceCallbacks" ,
294+ "GovernanceEventHandler" ,
295+ "LlamaIndexAdapter" ,
296+ ]
0 commit comments