diff --git a/agentops/instrumentation/common/streaming.py b/agentops/instrumentation/common/streaming.py index eaf320558..53c9a2b5f 100644 --- a/agentops/instrumentation/common/streaming.py +++ b/agentops/instrumentation/common/streaming.py @@ -4,6 +4,7 @@ in a consistent way across different providers. """ +import asyncio from typing import Optional, Any, Dict, Callable from abc import ABC import time @@ -67,8 +68,12 @@ def _process_chunk(self, chunk: Any): self.token_usage.completion_tokens or 0 ) + chunk_usage.completion_tokens - def _finalize(self): - """Finalize the stream processing.""" + def _finalize(self, *, set_success_status: bool = True): + """Finalize the stream processing. + + Args: + set_success_status: Whether to mark the span as successful after setting final attributes. + """ try: # Set final content final_content = "".join(self.accumulated_content) @@ -87,7 +92,8 @@ def _finalize(self): for attr_name, value in self.token_usage.to_attributes().items(): safe_set_attribute(self.span, attr_name, value) - self.span.set_status(Status(StatusCode.OK)) + if set_success_status: + self.span.set_status(Status(StatusCode.OK)) except Exception as e: logger.error(f"Error finalizing stream: {e}") self.span.set_status(Status(StatusCode.ERROR, str(e))) @@ -100,32 +106,41 @@ class SyncStreamWrapper(BaseStreamWrapper): """Wrapper for synchronous streaming responses.""" def __iter__(self): + failed = False try: for chunk in self.stream: self._process_chunk(chunk) yield chunk except Exception as e: + failed = True self.span.set_status(Status(StatusCode.ERROR, str(e))) self.span.record_exception(e) raise finally: - self._finalize() + self._finalize(set_success_status=not failed) class AsyncStreamWrapper(BaseStreamWrapper): """Wrapper for asynchronous streaming responses.""" async def __aiter__(self): + failed = False try: async for chunk in self.stream: self._process_chunk(chunk) yield chunk + except asyncio.CancelledError as e: + failed = True + self.span.set_status(Status(StatusCode.ERROR, str(e))) + self.span.record_exception(e) + raise except Exception as e: + failed = True self.span.set_status(Status(StatusCode.ERROR, str(e))) self.span.record_exception(e) raise finally: - self._finalize() + self._finalize(set_success_status=not failed) def create_stream_wrapper_factory( diff --git a/tests/unit/instrumentation/common/test_streaming.py b/tests/unit/instrumentation/common/test_streaming.py index 31fabc683..ea920bb1c 100644 --- a/tests/unit/instrumentation/common/test_streaming.py +++ b/tests/unit/instrumentation/common/test_streaming.py @@ -1,7 +1,11 @@ +import asyncio + import pytest from unittest.mock import Mock, patch from types import SimpleNamespace +from opentelemetry.trace import StatusCode + from agentops.instrumentation.common.streaming import ( BaseStreamWrapper, SyncStreamWrapper, @@ -248,7 +252,8 @@ def failing_stream(): with pytest.raises(ValueError, match="Test error"): list(wrapper) - mock_span.set_status.assert_called() + statuses = [call.args[0].status_code for call in mock_span.set_status.call_args_list] + assert statuses == [StatusCode.ERROR] mock_span.record_exception.assert_called_once() mock_span.end.assert_called_once() @@ -294,10 +299,32 @@ async def failing_async_stream(): async for chunk in wrapper: pass - mock_span.set_status.assert_called() + statuses = [call.args[0].status_code for call in mock_span.set_status.call_args_list] + assert statuses == [StatusCode.ERROR] mock_span.record_exception.assert_called_once() mock_span.end.assert_called_once() + @pytest.mark.asyncio + async def test_async_iteration_with_cancellation(self): + """Test cancellation is recorded without overwriting it with success.""" + + async def cancelled_async_stream(): + yield "chunk1" + raise asyncio.CancelledError("stream cancelled") + + mock_span = Mock() + extract_content = lambda x: x + wrapper = AsyncStreamWrapper(cancelled_async_stream(), mock_span, extract_content) + + with pytest.raises(asyncio.CancelledError, match="stream cancelled") as exc_info: + async for chunk in wrapper: + pass + + statuses = [call.args[0].status_code for call in mock_span.set_status.call_args_list] + assert statuses == [StatusCode.ERROR] + mock_span.record_exception.assert_called_once_with(exc_info.value) + mock_span.end.assert_called_once() + class TestCreateStreamWrapperFactory: """Test the create_stream_wrapper_factory function."""