diff --git a/agentops/instrumentation/common/streaming.py b/agentops/instrumentation/common/streaming.py index eaf320558..e99f1a5b1 100644 --- a/agentops/instrumentation/common/streaming.py +++ b/agentops/instrumentation/common/streaming.py @@ -67,8 +67,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 +91,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 +105,36 @@ 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 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..a466218eb 100644 --- a/tests/unit/instrumentation/common/test_streaming.py +++ b/tests/unit/instrumentation/common/test_streaming.py @@ -2,6 +2,8 @@ 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 +250,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,7 +297,8 @@ 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()