diff --git a/instrumentation/opentelemetry-instrumentation-genai-smolagents/.changelog/1006.fixed b/instrumentation/opentelemetry-instrumentation-genai-smolagents/.changelog/1006.fixed new file mode 100644 index 000000000..5e5ebedbe --- /dev/null +++ b/instrumentation/opentelemetry-instrumentation-genai-smolagents/.changelog/1006.fixed @@ -0,0 +1 @@ +Restore the caller's context while a model or agent run stream is unconsumed, so caller work is not parented under the `chat` or `invoke_agent` span. diff --git a/instrumentation/opentelemetry-instrumentation-genai-smolagents/pyproject.toml b/instrumentation/opentelemetry-instrumentation-genai-smolagents/pyproject.toml index 6d5f60794..b5574268c 100644 --- a/instrumentation/opentelemetry-instrumentation-genai-smolagents/pyproject.toml +++ b/instrumentation/opentelemetry-instrumentation-genai-smolagents/pyproject.toml @@ -28,7 +28,7 @@ dependencies = [ "opentelemetry-api ~= 1.43", "opentelemetry-instrumentation >= 0.64b0, <1", "opentelemetry-semantic-conventions >= 0.64b0, <1", - "opentelemetry-util-genai >= 1.2b0, <2", + "opentelemetry-util-genai >= 1.3b0, <2", ] [project.optional-dependencies] diff --git a/instrumentation/opentelemetry-instrumentation-genai-smolagents/src/opentelemetry/instrumentation/genai/smolagents/patch.py b/instrumentation/opentelemetry-instrumentation-genai-smolagents/src/opentelemetry/instrumentation/genai/smolagents/patch.py index 6ef5fa234..c1aa13273 100644 --- a/instrumentation/opentelemetry-instrumentation-genai-smolagents/src/opentelemetry/instrumentation/genai/smolagents/patch.py +++ b/instrumentation/opentelemetry-instrumentation-genai-smolagents/src/opentelemetry/instrumentation/genai/smolagents/patch.py @@ -20,7 +20,7 @@ import logging from collections.abc import Callable, Generator, Mapping -from contextlib import ExitStack +from contextlib import AbstractContextManager, ExitStack from typing import TYPE_CHECKING, Any, TypeAlias, TypeVar from smolagents import AgentMaxStepsError @@ -297,6 +297,7 @@ def __init__( ) -> None: super().__init__(stream, invocation=invocation) self._self_inference = invocation + invocation.suspend() self._self_capture_content = handler.should_capture_content() self._self_content: list[str] = [] self._self_input_tokens = 0 @@ -305,7 +306,8 @@ def __init__( def __del__(self) -> None: try: - self._finalize_success() + with self._execution_context(): + self._finalize_success() except BaseException: # pylint: disable=broad-except pass @@ -424,15 +426,23 @@ def __init__( # invocation to the base class. super().__init__(stream) self._self_agent_invocation = invocation + invocation.suspend() self._self_agent = agent self._self_capture_content = capture_content self._self_final_output: object = None self._self_saw_final = False self._self_finished = False + def _execution_context(self) -> AbstractContextManager[None]: + return self._self_agent_invocation.activate() + def __del__(self) -> None: try: - self._finalize_success() + with self._execution_context(): + try: + self._self_stream.close() + finally: + self._finalize_success() except BaseException: # pylint: disable=broad-except pass diff --git a/instrumentation/opentelemetry-instrumentation-genai-smolagents/tests/test_agents.py b/instrumentation/opentelemetry-instrumentation-genai-smolagents/tests/test_agents.py index e4199b592..d34dbaa54 100644 --- a/instrumentation/opentelemetry-instrumentation-genai-smolagents/tests/test_agents.py +++ b/instrumentation/opentelemetry-instrumentation-genai-smolagents/tests/test_agents.py @@ -7,6 +7,7 @@ import inspect import json import pathlib +import sys import tempfile import weakref from collections.abc import Generator @@ -36,7 +37,7 @@ ) from opentelemetry.semconv._incubating.metrics import gen_ai_metrics from opentelemetry.semconv.attributes import error_attributes -from opentelemetry.trace import StatusCode +from opentelemetry.trace import INVALID_SPAN, StatusCode, get_current_span from .test_utils import ( BrokenTool, @@ -309,6 +310,18 @@ def test_abandoned_streaming_run_does_not_leak_context( assert unrelated.parent is None +def test_abandoned_streaming_run_closes_the_run( + instrument_with_content, span_exporter, monkeypatch +) -> None: + unraisable: list[Any] = [] + monkeypatch.setattr(sys, "unraisablehook", unraisable.append) + agent = CodeAgent(tools=[], model=FakeCodeModel(), max_steps=3) + _first_step_only(agent) + gc.collect() + + assert unraisable == [] + + def test_code_agent_streaming_is_lazy_until_drained( instrument_with_content, span_exporter ) -> None: @@ -330,6 +343,42 @@ def test_code_agent_streaming_is_lazy_until_drained( assert attr(agent_span, error_attributes.ERROR_TYPE) is None +class _SpanningCodeModel(FakeCodeModel): + def __init__(self, tracer: Any) -> None: + super().__init__() + self._tracer = tracer + + def generate(self, *args: Any, **kwargs: Any) -> Any: + with self._tracer.start_as_current_span("model work"): + return super().generate(*args, **kwargs) + + +def test_streaming_run_is_current_only_while_read( + instrument_with_content, span_exporter, tracer_provider +) -> None: + tracer = tracer_provider.get_tracer("test") + agent = CodeAgent(tools=[], model=_SpanningCodeModel(tracer), max_steps=3) + stream = agent.run("Test question", stream=True) + assert get_current_span() is INVALID_SPAN + + next(stream) + assert get_current_span() is INVALID_SPAN + with tracer.start_as_current_span("caller"): + pass + list(stream) + + spans = span_exporter.get_finished_spans() + (agent_span,) = spans_by_operation(spans, "invoke_agent") + (caller,) = [span for span in spans if span.name == "caller"] + model_work = [span for span in spans if span.name == "model work"] + assert caller.parent is None + assert model_work + assert all( + span.parent.span_id == agent_span.context.span_id + for span in model_work + ) + + def test_streaming_run_stays_a_generator( instrument_with_content, span_exporter ) -> None: diff --git a/instrumentation/opentelemetry-instrumentation-genai-smolagents/tests/test_models.py b/instrumentation/opentelemetry-instrumentation-genai-smolagents/tests/test_models.py index 9bd35de6e..5c24847aa 100644 --- a/instrumentation/opentelemetry-instrumentation-genai-smolagents/tests/test_models.py +++ b/instrumentation/opentelemetry-instrumentation-genai-smolagents/tests/test_models.py @@ -44,7 +44,7 @@ error_attributes, server_attributes, ) -from opentelemetry.trace import StatusCode +from opentelemetry.trace import INVALID_SPAN, StatusCode, get_current_span from opentelemetry.util.genai.types import BlobPart, TextPart, UriPart from .test_utils import ( @@ -652,6 +652,45 @@ def test_generate_stream_is_lazy_and_records_the_drained_response( assert attr(span, GenAI.GEN_AI_RESPONSE_FINISH_REASONS) is None +def test_generate_stream_is_current_only_while_read( + instrument_with_content, span_exporter, tracer_provider +) -> None: + model = transformers_model(stream_chunks=["In ", "Paris"]) + stream = model.generate_stream(messages=MESSAGES) + assert get_current_span() is INVALID_SPAN + + next(stream) + assert get_current_span() is INVALID_SPAN + with tracer_provider.get_tracer("test").start_as_current_span("caller"): + pass + list(stream) + + (caller,) = [ + span + for span in span_exporter.get_finished_spans() + if span.name == "caller" + ] + assert caller.parent is None + + +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 + + def test_generate_stream_stays_a_generator( instrument_with_content, span_exporter ) -> None: