diff --git a/packages/opentelemetry-instrumentation-cohere/opentelemetry/instrumentation/cohere/streaming.py b/packages/opentelemetry-instrumentation-cohere/opentelemetry/instrumentation/cohere/streaming.py index f0ad2d8d24..e0c85ec5bd 100644 --- a/packages/opentelemetry-instrumentation-cohere/opentelemetry/instrumentation/cohere/streaming.py +++ b/packages/opentelemetry-instrumentation-cohere/opentelemetry/instrumentation/cohere/streaming.py @@ -122,25 +122,32 @@ async def aprocess_chat_v2_streaming_response(span, event_logger, llm_request_ty "type": "function", "function": {"name": "", "arguments": "", "description": ""}, } - async for item in response: - span.add_event(name=f"{SpanAttributes.LLM_CONTENT_COMPLETION_CHUNK}") - item_to_yield = item - try: - _accumulate_stream_item(item, current_content_item, current_tool_call_item, final_response) - except Exception: - pass - yield item_to_yield - - set_span_response_attributes(span, final_response) - if should_emit_events(): - emit_response_events(event_logger, llm_request_type, final_response) - elif should_send_prompts(): - _set_span_chat_response(span, final_response) - if final_response.get("error"): - span.set_status(Status(StatusCode.ERROR, final_response.get("error"))) - span.record_exception(final_response.get("error")) - else: - span.set_status(Status(StatusCode.OK)) + try: + async for item in response: + span.add_event(name=f"{SpanAttributes.LLM_CONTENT_COMPLETION_CHUNK}") + item_to_yield = item + try: + _accumulate_stream_item(item, current_content_item, current_tool_call_item, final_response) + except Exception: + pass + yield item_to_yield + + set_span_response_attributes(span, final_response) + if should_emit_events(): + emit_response_events(event_logger, llm_request_type, final_response) + elif should_send_prompts(): + _set_span_chat_response(span, final_response) + if final_response.get("error"): + span.set_status(Status(StatusCode.ERROR, final_response.get("error"))) + span.record_exception(final_response.get("error")) + else: + span.set_status(Status(StatusCode.OK)) + except Exception as e: + if span.is_recording(): + span.set_status(Status(StatusCode.ERROR, str(e))) + span.record_exception(e) + raise + finally: span.end() diff --git a/packages/opentelemetry-instrumentation-cohere/tests/test_streaming_error_handling.py b/packages/opentelemetry-instrumentation-cohere/tests/test_streaming_error_handling.py new file mode 100644 index 0000000000..0f1a79d5da --- /dev/null +++ b/packages/opentelemetry-instrumentation-cohere/tests/test_streaming_error_handling.py @@ -0,0 +1,30 @@ +import pytest +from opentelemetry.instrumentation.cohere.streaming import aprocess_chat_v2_streaming_response +from opentelemetry.semconv_ai import LLMRequestTypeValues +from opentelemetry.trace import StatusCode + + +@pytest.mark.asyncio +async def test_async_v2_streaming_exception_ends_span(span_exporter, tracer_provider): + tracer = tracer_provider.get_tracer(__name__) + span = tracer.start_span("cohere.chat") + + async def failing_stream(): + yield {"type": "message-start", "delta": {"message": {"role": "assistant"}}, "id": "response-id"} + raise RuntimeError("cohere stream failure") + + with pytest.raises(RuntimeError, match="cohere stream failure"): + async for _ in aprocess_chat_v2_streaming_response( + span, None, LLMRequestTypeValues.CHAT, failing_stream() + ): + pass + + spans = span_exporter.get_finished_spans() + assert len(spans) == 1 + assert spans[0].status.status_code == StatusCode.ERROR + assert "cohere stream failure" in spans[0].status.description + assert any( + event.name == "exception" + and "cohere stream failure" in event.attributes.get("exception.message", "") + for event in spans[0].events + )