Repository navigation
[NA] [SDK] fix: drop LlamaIndex event outputs once their trace has ended - #9556
anishmehta24 wants to merge 3 commits into
Conversation
LlamaIndexCallbackHandler stored the output of every event in _map_event_id_to_output so that end_trace could read the last one, but never removed anything. A handler set once on Settings.callback_manager kept every LLM response, retrieved node list and chunk list for the life of the process. Outputs are now stored per root and dropped in end_trace after it has read the one it needs. An event that ends after its root (the last event of a streamed call) is not stored; the pending root update still gets its output.
github-actions
Bot
added
python
Codex Review Summary: ✅ CompletedThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
|
![baz-reviewer[bot]](https://avatars.githubusercontent.com/in/933528?s=60&v=4){kind=small}
| current_root = _llama_root.get() | ||
| if current_root is not None: | ||
| self._map_event_id_to_output[event_id] = span_output | ||
| self._map_root_id_to_event_ids.setdefault(current_root.id, []).append( | ||
| event_id | ||
| ) |
There was a problem hiding this comment.
Events arriving after end_trace cleans up their root leak one entry per query by repopulating _map_event_id_to_output and _map_root_id_to_event_ids, either through late streamed LLM history tasks or a race between concurrent on_event_end calls and cleanup. Reject events whose root has ended, update and send late output through _pending_root_output_updates without persistent map entries, and protect registration and cleanup with a shared lock or atomic ended check.
Want Baz to fix this for you? Activate Fixer
Other fix methods
Prompt for AI Agents
Before applying, verify this suggestion against the current code. In
`sdks/python/src/opik/integrations/llama_index/callback.py` around lines 260-265, update
`LlamaIndexCallbackHandler.on_event_end` so it cannot leave entries in
`_map_event_id_to_output` or `_map_root_id_to_event_ids` after the root has ended. Two
causes need handling: (1) late streamed events whose history-writing task inherits a
copied ContextVar containing an already-ended root, and (2) concurrent `on_event_end`
delivery racing with the cleanup in `end_trace` around lines 161-162. Check whether the
event's root is already in `_pending_root_output_updates` or otherwise ended before
recording bookkeeping; for that late-event path, update and send the pending root
directly without adding persistent map entries. Add and use one shared lock (or an
equivalent atomic root-ended check) around both event registration and root cleanup in
`end_trace`, so an event is rejected or cleaned up if the root has already ended.
Preserve tracking and existing streaming behavior for events that end before
`end_trace`.
There was a problem hiding this comment.
Fixed in f4c9e00. start_trace now registers the root and end_trace removes it, and on_event_end only stores an output while the root is registered, under a lock. A late event in a copied context stores nothing; there's a test for that case.
A root's entry is now created in start_trace and removed in end_trace, and on_event_end only stores an output while that entry exists. An event ending later, even in a task whose context still holds the ended root, stores nothing. A lock covers the registration and the cleanup.
| ) | ||
| from llama_index.core.llms import MockLLM | ||
|
|
||
| Settings.llm = MockLLM(max_tokens=4) |
There was a problem hiding this comment.
Global LLM setting contaminates later tests
This test leaves the process-wide Settings.llm set to MockLLM for subsequent tests, and failures skip cleanup; use monkeypatch or a restoring fixture instead.
Want Baz to fix this for you? Activate Fixer
Other fix methods
Prompt for AI Agents
Before applying, verify this suggestion against the current code. In
`sdks/python/tests/library_integration/llama_index/test_llama_index.py` around lines
1446-1446, update
`test_llama_index__event_ends_after_end_trace_in_copied_context__output_is_not_kept` so
the process-wide `Settings.llm` assignment is automatically restored. Add a
`monkeypatch` parameter and set `Settings.llm` through it, or use an existing fixture
that restores the setting after the test, ensuring cleanup also occurs when assertions
fail.
There was a problem hiding this comment.
Fixed in 28154cc, the test now sets it through monkeypatch so it's restored afterwards.
Details
LlamaIndexCallbackHandlerkeeps every event's output in_map_event_id_to_outputso thatend_tracecan read the last one, but nothing ever removes them. With the handler set once onSettings.callback_manager, every query keeps its retrieved nodes, chunk lists and LLM responses in memory for the life of the process. A small query engine with MockLLM left 8 more entries per query (11, 19, 27, 35, 43 after five queries).Outputs are now tracked per root and dropped in
end_traceafter it reads the one it needs. An event that ends after its root has ended (the last event of a streamed call) is no longer stored; the pending root update still gets its output. Same kind of fix as #7566 for the LangChain tracer.Change checklist
Issues
No existing issue.
AI-WATERMARK
AI-WATERMARK: yes
Testing
pytest tests/library_integration/llama_index/test_llama_index.py -k "event_outputs_are_released"(insdks/python, Windows, local process with thefake_backendfixture): three new tests, repeated queries with and withoutskip_index_construction_traceand astream_chat. They use MockLLM/MockEmbedding, so no API key is needed. All 3 fail on main and pass with this change.OPENAI_API_KEYand weren't run. Instead I compared the logged trace trees on main and this branch for query, streaming query, chat, stream_chat and complete, with and withoutskip_index_construction_traceand withevent_starts_to_ignore: identical. 16 concurrent calls through one shared handler still give 16 complete traces.ruff check,ruff format --check(0.14.14) and mypy 1.19.1 pass.Documentation
No docs change needed.