From 299e3452210c8da9d64cc77969900150584ca007 Mon Sep 17 00:00:00 2001 From: Adarsh Prashar Date: Sat, 13 Jun 2026 23:59:38 +0530 Subject: [PATCH] feat(sdk): LlamaIndex callback adapter MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a LlamaIndex BaseCallbackHandler so a LlamaIndex query or agent is governed by a RiskKernel run with ~no code change — the LlamaIndex analog of the LangChain adapter. RiskKernelCallbackHandler ticks one governed step per LLM call (CBEventType.LLM on_event_start), so a run's loop/time budget is enforced over the outer LlamaIndex loop and a halt surfaces as BudgetExceeded. LlamaIndex's CallbackManager doesn't wrap handler calls in a try/except, so the halt propagates to the caller without any extra flag (unlike LangChain's raise_error). With gate_tools=True, FUNCTION_CALL events route through the approval gate. llama-index-core is imported lazily with an object fallback, so the module imports and the SDK installs without llama-index present (it is not a dependency). Pinned to the llama-index-core >= 0.10 callback protocol. Tests exercise the governance path on stdlib alone against a fake Run (step tick, loop-budget halt, tool gating allow/deny, event-type normalization), plus a skip-unless-installed integration test that proves a runaway loop of real LlamaIndex LLM calls actually stops. Full suite passes; README and CHANGELOG updated. --- CHANGELOG.md | 10 ++ sdks/python/README.md | 6 + sdks/python/riskkernel/adapters/__init__.py | 1 + .../python/riskkernel/adapters/llama_index.py | 148 +++++++++++++++++ sdks/python/tests/test_llama_index.py | 152 ++++++++++++++++++ 5 files changed, 317 insertions(+) create mode 100644 sdks/python/riskkernel/adapters/llama_index.py create mode 100644 sdks/python/tests/test_llama_index.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 4924658..904aaab 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,16 @@ surface is governed by [`COMPATIBILITY.md`](COMPATIBILITY.md). ## [Unreleased] ### Added +- **Python SDK: LlamaIndex adapter.** `RiskKernelCallbackHandler` (from + `riskkernel.adapters.llama_index`) is a LlamaIndex `BaseCallbackHandler` that ticks + one governed step per LLM call (`CBEventType.LLM`), so a run's loop/time budget is + enforced over a LlamaIndex query or agent and a halt surfaces as `BudgetExceeded` — + the LlamaIndex analog of the LangChain and OpenAI-Agents adapters. Register it on + `Settings.callback_manager` and the governance is invisible until the budget bites; + LlamaIndex doesn't swallow handler exceptions, so no extra flag is needed. Pass + `gate_tools=True` to route `CBEventType.FUNCTION_CALL` through the approval gate. + `llama-index-core` is lazily imported, so the SDK stays dependency-free; pinned to + the `llama-index-core` >= 0.10 callback protocol. - **Streaming proxy.** `POST /v1/chat/completions` now supports `stream:true`: the budget is enforced before the stream opens, the OpenAI provider's SSE is forwarded to the client verbatim (authentic chunks, no translation) while token usage is diff --git a/sdks/python/README.md b/sdks/python/README.md index 62cea55..7c6c577 100644 --- a/sdks/python/README.md +++ b/sdks/python/README.md @@ -81,6 +81,12 @@ Lazy-imported, so you only pay for what you use: from riskkernel.adapters.langchain import RiskKernelCallbackHandler llm.invoke(prompt, config={"callbacks": [RiskKernelCallbackHandler(run)]}) +# LlamaIndex — a CallbackHandler that enforces loop/time budgets per LLM call +from llama_index.core import Settings +from llama_index.core.callbacks import CallbackManager +from riskkernel.adapters.llama_index import RiskKernelCallbackHandler +Settings.callback_manager = CallbackManager([RiskKernelCallbackHandler(run)]) + # Claude Agent SDK — PreToolUse approval hook from riskkernel.adapters.claude_agent import make_pre_tool_use_hook hook = make_pre_tool_use_hook(run, side_effect_for={"Bash": "exec", "Write": "write"}) diff --git a/sdks/python/riskkernel/adapters/__init__.py b/sdks/python/riskkernel/adapters/__init__.py index 39fc6e4..8ea671e 100644 --- a/sdks/python/riskkernel/adapters/__init__.py +++ b/sdks/python/riskkernel/adapters/__init__.py @@ -3,6 +3,7 @@ third-party dependencies and you only pay for what you use. - ``langchain`` — a CallbackHandler (loop/time enforcement per LLM call). +- ``llama_index`` — a CallbackHandler (loop/time enforcement per LLM call). - ``claude_agent`` — a PreToolUse hook for the Claude Agent SDK (approval gate). - ``openai_agents`` — RunHooks for the OpenAI Agents SDK (steps + approval gate). """ diff --git a/sdks/python/riskkernel/adapters/llama_index.py b/sdks/python/riskkernel/adapters/llama_index.py new file mode 100644 index 0000000..8514988 --- /dev/null +++ b/sdks/python/riskkernel/adapters/llama_index.py @@ -0,0 +1,148 @@ +"""LlamaIndex adapter: a callback handler that enforces a governed run's loop and +time budgets, ticking one step per LLM call. Register it on your LlamaIndex +``CallbackManager`` (or ``Settings.callback_manager``) and point your LlamaIndex +LLM at the governing proxy (``run.proxy_config()``) for token/cost/budget metering; +this handler adds the outer-loop enforcement the proxy can't see. + + from llama_index.core.callbacks import CallbackManager + from llama_index.core import Settings + from riskkernel.adapters.llama_index import RiskKernelCallbackHandler + + Settings.callback_manager = CallbackManager([RiskKernelCallbackHandler(run)]) + +A BudgetExceeded raised here propagates out of the LlamaIndex call, halting the +query/agent — exactly when the run is out of budget. Unlike LangChain, LlamaIndex's +``CallbackManager.on_event_start`` does NOT wrap handler calls in a try/except, so a +budget halt surfaces to the caller without any extra flag. + +Supported API: LlamaIndex's ``BaseCallbackHandler`` callback protocol +(``llama-index-core`` >= 0.10, where the package split landed). We tick a step on +``CBEventType.LLM`` (``"llm"``) ``on_event_start`` — one LLM call == one governed +step — and, when ``gate_tools`` is set, gate ``CBEventType.FUNCTION_CALL`` +(``"function_call"``) through the approval gate. +""" + +from __future__ import annotations + +from typing import Any, Optional + +from ..approval import ApprovalGate +from ..runtime import Run + +# LlamaIndex event-type string values (CBEventType.LLM / .FUNCTION_CALL). We match +# on the string value so we don't need the enum imported when LlamaIndex is absent. +_LLM_EVENT = "llm" +_FUNCTION_CALL_EVENT = "function_call" + + +def _base_handler(): + # Inherit the real base class when LlamaIndex is installed (best integration); + # otherwise fall back to object so the module still imports and the SDK installs + # without llama-index present (it is NOT a dependency of the SDK). + try: + from llama_index.core.callbacks.base_handler import ( # type: ignore + BaseCallbackHandler, + ) + return BaseCallbackHandler + except Exception: + return object + + +class RiskKernelCallbackHandler(_base_handler()): # type: ignore[misc] + """Enforces loop/time budgets and (optionally) gates tools on approval. + + Args: + run: the governed Run. + gate_tools: if True, every tool/function call must pass the approval gate. + tool_side_effect: side-effect label reported for gated tools. + timeout: max seconds to await a human decision on a gated tool. + """ + + def __init__(self, run: Run, gate_tools: bool = False, + tool_side_effect: str = "tool", timeout: Optional[float] = None): + self.run = run + self.gate_tools = gate_tools + self.tool_side_effect = tool_side_effect + self.timeout = timeout + self._gate = ApprovalGate(run) + # BaseCallbackHandler.__init__ requires the ignore lists. Call it only when + # we actually inherit it (not the object fallback), so the module works + # whether or not LlamaIndex is installed. + base = type(self).__mro__[1] + if base is not object: + base.__init__(self, event_starts_to_ignore=[], event_ends_to_ignore=[]) + + # One LLM call == one governed step; FUNCTION_CALL starts are gated when asked. + # Returns the event_id unchanged — the callback manager uses it to correlate the + # matching on_event_end, so we must not drop it. + def on_event_start(self, event_type: Any, payload: Optional[dict] = None, + event_id: str = "", parent_id: str = "", + **kwargs: Any) -> str: + et = _event_value(event_type) + if et == _LLM_EVENT: + self.run.step() # raises BudgetExceeded when the loop/time budget is spent + elif et == _FUNCTION_CALL_EVENT and self.gate_tools: + name = _tool_name(payload) + self._gate.require(name or "tool", side_effect=self.tool_side_effect, + arguments={"payload": _stringify(payload)}, + timeout=self.timeout) + return event_id + + def on_event_end(self, event_type: Any, payload: Optional[dict] = None, + event_id: str = "", **kwargs: Any) -> None: + return None + + def start_trace(self, trace_id: Optional[str] = None) -> None: + return None + + def end_trace(self, trace_id: Optional[str] = None, + trace_map: Optional[dict] = None) -> None: + return None + + +def _event_value(event_type: Any) -> str: + """Normalize a CBEventType (or its string value) to its lowercase string.""" + value = getattr(event_type, "value", event_type) + return str(value).lower() + + +def _tool_name(payload: Optional[dict]) -> str: + """Best-effort extraction of the tool name from a FUNCTION_CALL payload across + LlamaIndex versions. The payload keys are EventPayload enum members whose string + values are ``"tool"`` and ``"function_call"``; EventPayload.TOOL is typically a + ToolMetadata with a ``.name``.""" + if not isinstance(payload, dict): + return "" + # Resolve by enum-string key, matching however the enum stringifies as a dict key. + tool = None + fn = None + for key, val in payload.items(): + k = _event_value(key) + if k == "tool": + tool = val + elif k == "function_call": + fn = val + if tool is not None: + name = getattr(tool, "name", None) + if name: + return str(name) + if isinstance(tool, dict): + n = tool.get("name") + if n: + return str(n) + if isinstance(tool, str): + return tool + if isinstance(fn, dict): + n = fn.get("name") or fn.get("tool_name") + if n: + return str(n) + return "" + + +def _stringify(v: Any) -> Any: + try: + import json + json.dumps(v) + return v + except Exception: + return repr(v) diff --git a/sdks/python/tests/test_llama_index.py b/sdks/python/tests/test_llama_index.py new file mode 100644 index 0000000..91441ea --- /dev/null +++ b/sdks/python/tests/test_llama_index.py @@ -0,0 +1,152 @@ +"""LlamaIndex adapter tests — governance behavior against a fake Run, no daemon, +no third-party deps. The handler imports even without llama-index installed (lazy +base class), so these exercise the enforcement path on stdlib alone; a real +LlamaIndex integration test is gated behind skipUnless. +""" + +import unittest + +from riskkernel.adapters.llama_index import RiskKernelCallbackHandler +from riskkernel.errors import ApprovalDenied, BudgetExceeded + + +def _has_llama_index() -> bool: + try: + import llama_index.core # noqa: F401 + + return True + except Exception: + return False + + +class FakeRun: + """A stand-in for runtime.Run: counts steps and halts past a loop budget, + mirroring how the daemon's BeginStep raises BudgetExceeded when the budget is + spent. No HTTP, no daemon.""" + + def __init__(self, loop_budget=None): + self.steps = 0 + self.loop_budget = loop_budget + + def step(self): + self.steps += 1 + if self.loop_budget is not None and self.steps > self.loop_budget: + raise BudgetExceeded("loop_budget_exceeded") + return self.steps + + +class _FakeGate: + """Records gate calls; can be set to deny like a human denying a tool.""" + + def __init__(self, deny=False): + self.deny = deny + self.calls = [] + + def require(self, tool, side_effect="", arguments=None, timeout=None): + self.calls.append({"tool": tool, "side_effect": side_effect, + "arguments": arguments, "timeout": timeout}) + if self.deny: + raise ApprovalDenied(tool, "no") + + +class LlamaIndexAdapterTest(unittest.TestCase): + def test_module_imports_without_llama_index(self): + # The lazy base-class fallback must let the SDK import the adapter even with + # no llama-index installed (it is not a dependency). + self.assertTrue(callable(RiskKernelCallbackHandler)) + + def test_llm_event_ticks_one_step(self): + run = FakeRun() + h = RiskKernelCallbackHandler(run) + # on_event_start must return the event_id so the callback manager can + # correlate the matching on_event_end. + self.assertEqual(h.on_event_start("llm", event_id="e1"), "e1") + self.assertEqual(run.steps, 1) + + def test_llm_event_enforces_loop_budget(self): + run = FakeRun(loop_budget=2) + h = RiskKernelCallbackHandler(run) + h.on_event_start("llm", event_id="e1") # step 1 + h.on_event_start("llm", event_id="e2") # step 2 + with self.assertRaises(BudgetExceeded) as cm: + h.on_event_start("llm", event_id="e3") # step 3 -> over budget + self.assertEqual(cm.exception.reason, "loop_budget_exceeded") + # The halt is surfaced, not swallowed; the step counter stopped at the cap+1. + self.assertEqual(run.steps, 3) + + def test_non_llm_events_do_not_tick_a_step(self): + run = FakeRun() + h = RiskKernelCallbackHandler(run) + for et in ("retrieve", "embedding", "query", "synthesize"): + h.on_event_start(et, event_id="e") + self.assertEqual(run.steps, 0) + + def test_function_call_not_gated_by_default(self): + run = FakeRun() + h = RiskKernelCallbackHandler(run) + gate = _FakeGate() + h._gate = gate + h.on_event_start("function_call", event_id="e", + payload={"tool": _ToolMeta("deploy")}) + self.assertEqual(gate.calls, []) # gate_tools defaults off + self.assertEqual(run.steps, 0) # and it isn't an LLM step + + def test_function_call_gated_passes_tool_name(self): + run = FakeRun() + h = RiskKernelCallbackHandler(run, gate_tools=True, tool_side_effect="exec") + gate = _FakeGate() + h._gate = gate + h.on_event_start("function_call", event_id="e", + payload={"tool": _ToolMeta("deploy")}) + self.assertEqual(len(gate.calls), 1) + self.assertEqual(gate.calls[0]["tool"], "deploy") + self.assertEqual(gate.calls[0]["side_effect"], "exec") + + def test_function_call_gated_denied_raises(self): + run = FakeRun() + h = RiskKernelCallbackHandler(run, gate_tools=True) + h._gate = _FakeGate(deny=True) + with self.assertRaises(ApprovalDenied): + h.on_event_start("function_call", event_id="e", + payload={"function_call": {"name": "deploy"}}) + + def test_enum_like_event_type_is_normalized(self): + # CBEventType has a .value; the handler matches on the string value, so an + # enum-like object with .value == "llm" must still tick a step. + run = FakeRun() + h = RiskKernelCallbackHandler(run) + + class _EventType: + value = "LLM" + + h.on_event_start(_EventType(), event_id="e") + self.assertEqual(run.steps, 1) + + @unittest.skipUnless(_has_llama_index(), "llama-index-core not installed") + def test_llama_index_integration_stops_runaway_loop(self): + # The real proof: a runaway loop of LlamaIndex LLM calls must actually stop. + # LlamaIndex's CallbackManager doesn't swallow handler exceptions, so the + # BudgetExceeded raised on the 3rd LLM call propagates out of the call. + from llama_index.core.callbacks import CallbackManager + from llama_index.core.llms import MockLLM + + run = FakeRun(loop_budget=2) + cm = CallbackManager([RiskKernelCallbackHandler(run)]) + llm = MockLLM(callback_manager=cm) + calls = 0 + with self.assertRaises(BudgetExceeded): + while True: # a deliberately runaway loop + llm.complete("step") + calls += 1 + self.assertEqual(calls, 2) # 2 allowed; the 3rd halted + + +class _ToolMeta: + """A minimal ToolMetadata stand-in: has a .name like LlamaIndex's real one.""" + + def __init__(self, name): + self.name = name + + +if __name__ == "__main__": + unittest.main()