Skip to content

feat(sdk): LlamaIndex callback adapter - #125

Merged
prashar32 merged 1 commit into
mainfrom
feat/llamaindex-adapter
Jun 14, 2026
Merged

feat(sdk): LlamaIndex callback adapter#125
prashar32 merged 1 commit into
mainfrom
feat/llamaindex-adapter

Conversation

@prashar32

Copy link
Copy Markdown
Owner

Adds a LlamaIndex integration to the Python SDK so a LlamaIndex query or agent is governed by a RiskKernel run with ~no code change — the LlamaIndex analog of the existing LangChain adapter.

What it does

RiskKernelCallbackHandler (in riskkernel/adapters/llama_index.py) is a LlamaIndex BaseCallbackHandler:

  • 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. The next step over budget raises BudgetExceeded, halting the query/agent.
  • The halt surfaces, not swallowed: LlamaIndex's CallbackManager.on_event_start doesn't wrap handler calls in a try/except, so the exception propagates to the caller — no extra flag needed (unlike LangChain's raise_error). I verified this in the LlamaIndex source and with a real integration test.
  • Optional tool gating (gate_tools=True) routes CBEventType.FUNCTION_CALL events through the approval gate, mirroring the LangChain adapter's gate_tools option.
  • llama-index-core is imported lazily with an object fallback base class, so the module imports and the SDK installs without llama-index present — it is not a dependency of the SDK. Pinned to the llama-index-core >= 0.10 callback protocol (documented in the module docstring).

Register it on Settings.callback_manager and governance is invisible until the budget bites:

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

Callback events hooked

  • CBEventType.LLM ("llm") — on_event_start ticks run.step() (loop/time enforcement).
  • CBEventType.FUNCTION_CALL ("function_call") — on_event_start gated through the approval gate when gate_tools=True.
  • on_event_end / start_trace / end_trace — implemented as no-ops (the protocol requires them); on_event_start returns the event_id so the callback manager can correlate the matching end event.

Tests

tests/test_llama_index.py exercises the governance path on stdlib alone against a fake Run: per-LLM-call step tick, loop-budget halt propagation, non-LLM events not ticking a step, tool-gating allow/deny, and event-type (CBEventType-like) normalization. A skip-unless-installed integration test proves a runaway loop of real LlamaIndex LLM calls actually stops at the budget. I ran the integration test against llama-index-core 0.14.22 in a throwaway venv and it passes; the full suite passes stdlib-only (25 tests, 3 framework-gated skips):

python -m unittest discover -s tests -t . -v
Ran 25 tests in 0.545s
OK (skipped=3)

README and CHANGELOG updated.

Closes #86.

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.
@prashar32
prashar32 merged commit 789cd1b into main Jun 14, 2026
7 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

LlamaIndex callback adapter

1 participant