feat(sdk): LlamaIndex callback adapter

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

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"})

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).
 """

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)

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()