diff --git a/agentops/instrumentation/providers/openai/stream_wrapper.py b/agentops/instrumentation/providers/openai/stream_wrapper.py index 4d23c5b35..0eabad2d5 100644 --- a/agentops/instrumentation/providers/openai/stream_wrapper.py +++ b/agentops/instrumentation/providers/openai/stream_wrapper.py @@ -51,6 +51,7 @@ def __init__(self, stream: Any, span: Span, request_kwargs: dict): self._usage = None self._tool_calls = {} self._current_tool_call_index = None + self._finished = False # Make sure the span is attached to the current context current_context = context_api.get_current() @@ -62,6 +63,8 @@ def __iter__(self) -> Iterator[Any]: def __next__(self) -> Any: """Process the next chunk from the stream.""" + if self._finished: + raise StopIteration try: chunk = next(self._stream) self._process_chunk(chunk) @@ -69,6 +72,11 @@ def __next__(self) -> Any: except StopIteration: self._finalize_stream() raise + except Exception as e: + self._span.record_exception(e) + self._span.set_status(Status(StatusCode.ERROR, str(e))) + self._end_stream() + raise def __enter__(self): """Support context manager protocol.""" @@ -76,14 +84,23 @@ def __enter__(self): def __exit__(self, exc_type, exc_val, exc_tb): """Clean up on context manager exit.""" - if exc_type is not None: + if not self._finished and exc_type is not None: self._span.record_exception(exc_val) self._span.set_status(Status(StatusCode.ERROR, str(exc_val))) - self._span.end() - context_api.detach(self._token) + self._end_stream() return False + def _end_stream(self) -> None: + """End the span and detach its context once across terminal paths.""" + if self._finished: + return + self._finished = True + try: + self._span.end() + finally: + context_api.detach(self._token) + def _process_chunk(self, chunk: Any) -> None: """Process a single chunk from the stream. @@ -228,8 +245,7 @@ def _finalize_stream(self) -> None: # Finalize span and context self._span.set_status(Status(StatusCode.OK)) - self._span.end() - context_api.detach(self._token) + self._end_stream() class OpenAIAsyncStreamWrapper: diff --git a/tests/unit/instrumentation/openai_core/test_sync_stream_lifecycle.py b/tests/unit/instrumentation/openai_core/test_sync_stream_lifecycle.py new file mode 100644 index 000000000..6ea89deea --- /dev/null +++ b/tests/unit/instrumentation/openai_core/test_sync_stream_lifecycle.py @@ -0,0 +1,161 @@ +"""Synchronous chat streams must release tracing state on every terminal path.""" + +from contextlib import nullcontext +import json + +import httpx +import openai +import pytest +from opentelemetry import context, trace +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.trace import StatusCode +from wrapt import wrap_function_wrapper + +from agentops.instrumentation.providers.openai.stream_wrapper import OpenaiStreamWrapper, chat_completion_stream_wrapper + + +@pytest.fixture +def stream_tracing(): + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + tracer = provider.get_tracer(__name__) + saved = context.attach(context.get_current()) + original = openai.resources.chat.completions.Completions.create + wrap_function_wrapper( + "openai.resources.chat.completions", "Completions.create", chat_completion_stream_wrapper(tracer) + ) + try: + yield tracer, exporter + finally: + openai.resources.chat.completions.Completions.create = original + context.detach(saved) + provider.shutdown() + + +def stream_client(fail): + chunk = { + "id": "synthetic-completion", + "object": "chat.completion.chunk", + "created": 1, + "model": "synthetic-model", + "choices": [{"index": 0, "delta": {"content": "first"}, "finish_reason": None}], + } + terminal = ( + json.dumps({"error": {"message": "synthetic stream failure", "type": "server_error"}}) if fail else "[DONE]" + ) + payload = ("data: " + json.dumps(chunk) + "\n\ndata: " + terminal + "\n\n").encode() + transport = httpx.MockTransport( + lambda request: httpx.Response(200, headers={"content-type": "text/event-stream"}, content=payload) + ) + return openai.OpenAI(api_key="synthetic-placeholder", max_retries=0, http_client=httpx.Client(transport=transport)) + + +def create_stream(client): + return client.chat.completions.create( + model="synthetic-model", messages=[{"role": "user", "content": "synthetic input"}], stream=True + ) + + +@pytest.mark.parametrize("managed", [False, True]) +@pytest.mark.parametrize("fail", [False, True]) +def test_sync_sdk_stream_ends_once_and_restores_request(stream_tracing, mocker, managed, fail): + tracer, exporter = stream_tracing + with tracer.start_as_current_span("request") as parent, stream_client(fail) as client: + stream = create_stream(client) + span = trace.get_current_span() + end = mocker.spy(span, "end") + try: + with pytest.raises(openai.APIError, match="synthetic stream failure") if fail else nullcontext(): + with stream if managed else nullcontext(stream): + assert next(stream).choices[0].delta.content == "first" + if fail: + next(stream) + else: + assert list(stream) == [] + assert trace.get_current_span() is parent + spans = exporter.get_finished_spans() + assert len(spans) == 1 + assert spans[0].status.status_code == (StatusCode.ERROR if fail else StatusCode.OK) + assert sum(event.name == "exception" for event in spans[0].events) == (1 if fail else 0) + if not fail: + assert spans[0].attributes["llm.openai.stream.chunk_count"] == 1 + for _ in range(2): + with pytest.raises(StopIteration): + next(stream) + end.assert_called_once() + with tracer.start_as_current_span("following") as following: + assert following.parent.span_id == parent.get_span_context().span_id + finally: + if span.is_recording(): + span.end() + + +def test_caught_stream_error_restores_context_before_manager_exit(stream_tracing): + tracer, exporter = stream_tracing + with tracer.start_as_current_span("request") as parent, stream_client(True) as client: + with create_stream(client) as stream: + next(stream) + with pytest.raises(openai.APIError): + next(stream) + assert trace.get_current_span() is parent + assert len(exporter.get_finished_spans()) == 1 + with tracer.start_as_current_span("following") as following: + assert following.parent.span_id == parent.get_span_context().span_id + assert trace.get_current_span() is parent + + +def test_iterator_error_preserves_original_exception_and_chunk(stream_tracing, mocker): + tracer, exporter = stream_tracing + error = ValueError("synthetic iterator failure") + chunk = openai.types.chat.ChatCompletionChunk( + id="synthetic-completion", + object="chat.completion.chunk", + created=1, + model="synthetic-model", + choices=[{"index": 0, "delta": {"content": "first"}}], + ) + + def source(): + yield chunk + raise error + + with tracer.start_as_current_span("request") as parent: + span = tracer.start_span("stream") + end = mocker.spy(span, "end") + stream = OpenaiStreamWrapper(source(), span, {}) + try: + assert next(stream) is chunk + with pytest.raises(ValueError) as caught: + next(stream) + assert caught.value is error + assert trace.get_current_span() is parent + spans = exporter.get_finished_spans() + assert len(spans) == 1 + assert spans[0].status.status_code == StatusCode.ERROR + assert sum(event.name == "exception" for event in spans[0].events) == 1 + end.assert_called_once() + finally: + if span.is_recording(): + span.end() + + +def test_context_manager_records_caller_error_before_exhaustion(stream_tracing, mocker): + tracer, exporter = stream_tracing + error = ValueError("synthetic caller failure") + with tracer.start_as_current_span("request") as parent, stream_client(False) as client: + stream = create_stream(client) + end = mocker.spy(trace.get_current_span(), "end") + with pytest.raises(ValueError) as caught: + with stream: + next(stream) + raise error + assert caught.value is error + assert trace.get_current_span() is parent + spans = exporter.get_finished_spans() + assert len(spans) == 1 + assert spans[0].status.status_code == StatusCode.ERROR + assert sum(event.name == "exception" for event in spans[0].events) == 1 + end.assert_called_once()