Repository navigation
fix(smolagents): suspend stream invocations between reads - #1006
ishan-1010 wants to merge 5 commits into
Conversation
Pull request dashboard statusWaiting on maintainers · refreshed 2026-11-01 10:17 UTC Merge when ready. Status above doesn't look right?
|
There was a problem hiding this comment.
🟡 Changes recommended
The dependency minimum permits util-genai versions lacking required stream-context reactivation, and both stream destructors finalize outside their execution contexts, risking lost context and incorrect span parenting.
2 open issues
What changed in this PR
This PR addresses the smolagents portion of #872 by restoring caller context between stream reads while keeping streamed work under its model or agent span.
Changes:
- Suspend model and agent invocations before returning streams, with explicit reactivation for agent reads.
- Add regression tests for caller isolation and nested model-span parenting.
- Document the context fix in a changelog fragment.
| File | Description |
|---|---|
instrumentation/opentelemetry-instrumentation-genai-smolagents/tests/test_models.py |
Tests caller context between model-stream reads. |
instrumentation/opentelemetry-instrumentation-genai-smolagents/tests/test_agents.py |
Tests caller isolation and model-work parenting during agent streaming. |
instrumentation/opentelemetry-instrumentation-genai-smolagents/src/opentelemetry/instrumentation/genai/smolagents/patch.py |
Suspends stream invocations and scopes agent execution context. |
instrumentation/opentelemetry-instrumentation-genai-smolagents/.changelog/1006.fixed |
Records the stream-context correction. |
Review effort: Balanced
| try: | ||
| self._finalize_success() | ||
| with self._execution_context(): | ||
| self._finalize_success() |
There was a problem hiding this comment.
Could we also close the wrapped stream inside this context? On abandonment, the generator’s finally runs after __del__ returns, so it still executes outside the agent span. The current tests surface this as generator ignored GeneratorExit warnings.
There was a problem hiding this comment.
done in 9512eb2, the GeneratorExit warning is gone as well.
thanks!
| for span in span_exporter.get_finished_spans() | ||
| if span.name == "caller" | ||
| ] | ||
| assert caller.parent is None |
There was a problem hiding this comment.
This covers the caller side, but nothing checks that chat is still current while the model generator runs (the agent test has _SpanningCodeModel for that). Something like this would catch a regression in the per-read activation:
def test_generate_stream_reads_under_the_chat_span(
instrument_with_content, span_exporter
) -> None:
model = transformers_model()
seen = []
def chunks():
for text in ["In ", "Paris"]:
seen.append(get_current_span().get_span_context().span_id)
yield text
model.streamer = chunks()
list(model.generate_stream(messages=MESSAGES))
(chat,) = spans_by_operation(span_exporter.get_finished_spans(), "chat")
assert seen == [chat.context.span_id] * 2
Part of #872 (smolagents).
Both smolagents stream wrappers left their invocation attached after handing the stream back, so whatever the caller did before draining it was parented under the
chatorinvoke_agentspan. The model stream and the agent run stream now suspend the invocation when they are returned, following the OpenAI wrappers. The agent run wrapper does not pass its invocation to the base class, so it overrides_execution_contextto reactivate it around each read, which keeps model calls made during the run under the agent span.