From ac87a1f626460c2c1bc2753b9745d35c9543d436 Mon Sep 17 00:00:00 2001 From: JINSONG WANG Date: Thu, 17 Oct 2024 22:13:36 -0700 Subject: [PATCH 1/2] fix(watsonx): use with statement to automatically handle exception --- .../instrumentation/watsonx/__init__.py | 101 +++++++++--------- 1 file changed, 50 insertions(+), 51 deletions(-) diff --git a/packages/opentelemetry-instrumentation-watsonx/opentelemetry/instrumentation/watsonx/__init__.py b/packages/opentelemetry-instrumentation-watsonx/opentelemetry/instrumentation/watsonx/__init__.py index f6402e572..d6a72359c 100644 --- a/packages/opentelemetry-instrumentation-watsonx/opentelemetry/instrumentation/watsonx/__init__.py +++ b/packages/opentelemetry-instrumentation-watsonx/opentelemetry/instrumentation/watsonx/__init__.py @@ -445,67 +445,66 @@ def _wrap( name = to_wrap.get("span_name") - span = tracer.start_span( + with tracer.start_span( name, kind=SpanKind.CLIENT, attributes={ SpanAttributes.LLM_SYSTEM: "Watsonx", SpanAttributes.LLM_REQUEST_TYPE: LLMRequestTypeValues.COMPLETION.value, }, - ) - - _set_api_attributes(span) - if "generate" in name: - _set_input_attributes(span, instance, kwargs) - if to_wrap.get("method") == "generate_text_stream": - if (raw_flag := kwargs.get("raw_response", None)) is None: - kwargs = {**kwargs, "raw_response": True} - elif raw_flag is False: - kwargs["raw_response"] = True - - try: - start_time = time.time() - response = wrapped(*args, **kwargs) - end_time = time.time() - except Exception as e: - end_time = time.time() - duration = end_time - start_time if "start_time" in locals() else 0 - - attributes = { - "error.type": e.__class__.__name__, - } + ) as span: + + _set_api_attributes(span) + if "generate" in name: + _set_input_attributes(span, instance, kwargs) + if to_wrap.get("method") == "generate_text_stream": + if (raw_flag := kwargs.get("raw_response", None)) is None: + kwargs = {**kwargs, "raw_response": True} + elif raw_flag is False: + kwargs["raw_response"] = True + + try: + start_time = time.time() + response = wrapped(*args, **kwargs) + end_time = time.time() + except Exception as e: + end_time = time.time() + duration = end_time - start_time if "start_time" in locals() else 0 + + attributes = { + "error.type": e.__class__.__name__, + } - if duration > 0 and duration_histogram: - duration_histogram.record(duration, attributes=attributes) - if exception_counter: - exception_counter.add(1, attributes=attributes) + if duration > 0 and duration_histogram: + duration_histogram.record(duration, attributes=attributes) + if exception_counter: + exception_counter.add(1, attributes=attributes) - raise e + raise e - if "generate" in name: - if isinstance(response, types.GeneratorType): - return _build_and_set_stream_response( - span, - response, - raw_flag, - token_histogram, - response_counter, - duration_histogram, - start_time, - ) - else: - duration = end_time - start_time - _set_response_attributes( - span, - response, - token_histogram, - response_counter, - duration_histogram, - duration, - ) + if "generate" in name: + if isinstance(response, types.GeneratorType): + return _build_and_set_stream_response( + span, + response, + raw_flag, + token_histogram, + response_counter, + duration_histogram, + start_time, + ) + else: + duration = end_time - start_time + _set_response_attributes( + span, + response, + token_histogram, + response_counter, + duration_histogram, + duration, + ) - span.end() - return response + return response class WatsonxSpanAttributes: From 74eaaa0766cefe9e2d6c4f2b6173968936ad750a Mon Sep 17 00:00:00 2001 From: JINSONG WANG Date: Fri, 18 Oct 2024 23:36:17 -0700 Subject: [PATCH 2/2] fix(watsonx): explicitly call span.__exit__ for exception record --- .../instrumentation/watsonx/__init__.py | 105 +++++++++--------- 1 file changed, 55 insertions(+), 50 deletions(-) diff --git a/packages/opentelemetry-instrumentation-watsonx/opentelemetry/instrumentation/watsonx/__init__.py b/packages/opentelemetry-instrumentation-watsonx/opentelemetry/instrumentation/watsonx/__init__.py index d6a72359c..44fb8ad30 100644 --- a/packages/opentelemetry-instrumentation-watsonx/opentelemetry/instrumentation/watsonx/__init__.py +++ b/packages/opentelemetry-instrumentation-watsonx/opentelemetry/instrumentation/watsonx/__init__.py @@ -2,6 +2,7 @@ import logging import os +import sys import types import time from typing import Collection, Optional @@ -445,66 +446,70 @@ def _wrap( name = to_wrap.get("span_name") - with tracer.start_span( + span = tracer.start_span( name, kind=SpanKind.CLIENT, attributes={ SpanAttributes.LLM_SYSTEM: "Watsonx", SpanAttributes.LLM_REQUEST_TYPE: LLMRequestTypeValues.COMPLETION.value, }, - ) as span: - - _set_api_attributes(span) - if "generate" in name: - _set_input_attributes(span, instance, kwargs) - if to_wrap.get("method") == "generate_text_stream": - if (raw_flag := kwargs.get("raw_response", None)) is None: - kwargs = {**kwargs, "raw_response": True} - elif raw_flag is False: - kwargs["raw_response"] = True - - try: - start_time = time.time() - response = wrapped(*args, **kwargs) - end_time = time.time() - except Exception as e: - end_time = time.time() - duration = end_time - start_time if "start_time" in locals() else 0 - - attributes = { - "error.type": e.__class__.__name__, - } + ) - if duration > 0 and duration_histogram: - duration_histogram.record(duration, attributes=attributes) - if exception_counter: - exception_counter.add(1, attributes=attributes) + _set_api_attributes(span) + if "generate" in name: + _set_input_attributes(span, instance, kwargs) + if to_wrap.get("method") == "generate_text_stream": + if (raw_flag := kwargs.get("raw_response", None)) is None: + kwargs = {**kwargs, "raw_response": True} + elif raw_flag is False: + kwargs["raw_response"] = True + + try: + start_time = time.time() + response = wrapped(*args, **kwargs) + end_time = time.time() + except Exception as e: + end_time = time.time() + duration = end_time - start_time if "start_time" in locals() else 0 + + attributes = { + "error.type": e.__class__.__name__, + } - raise e + if duration > 0 and duration_histogram: + duration_histogram.record(duration, attributes=attributes) + if exception_counter: + exception_counter.add(1, attributes=attributes) - if "generate" in name: - if isinstance(response, types.GeneratorType): - return _build_and_set_stream_response( - span, - response, - raw_flag, - token_histogram, - response_counter, - duration_histogram, - start_time, - ) - else: - duration = end_time - start_time - _set_response_attributes( - span, - response, - token_histogram, - response_counter, - duration_histogram, - duration, - ) + exc_type, exc_val, exc_tb = sys.exc_info() + span.__exit__(exc_type, exc_val, exc_tb) + + raise e - return response + if "generate" in name: + if isinstance(response, types.GeneratorType): + return _build_and_set_stream_response( + span, + response, + raw_flag, + token_histogram, + response_counter, + duration_histogram, + start_time, + ) + else: + duration = end_time - start_time + _set_response_attributes( + span, + response, + token_histogram, + response_counter, + duration_histogram, + duration, + ) + + span.end() + return response class WatsonxSpanAttributes: