fix(watsonx): refactor event emission

Author: LuizDMM

packages/opentelemetry-instrumentation-watsonx/opentelemetry/instrumentation/watsonx/__init__.py

@@ -4,27 +4,27 @@
 import os
 import time
 import types
-from typing import Collection, Optional
+from typing import Collection, Optional, Union
 
 from opentelemetry import context as context_api
-from opentelemetry._events import get_event_logger
+from opentelemetry._events import EventLogger, get_event_logger
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
 from opentelemetry.instrumentation.utils import (
     _SUPPRESS_INSTRUMENTATION_KEY,
     unwrap,
 )
 from opentelemetry.instrumentation.watsonx.config import Config
-from opentelemetry.instrumentation.watsonx.event_handler import (
-    ChoiceEvent,
-    MessageEvent,
+from opentelemetry.instrumentation.watsonx.event_emitter import (
     emit_event,
 )
-from opentelemetry.instrumentation.watsonx.utils import dont_throw
+from opentelemetry.instrumentation.watsonx.event_models import ChoiceEvent, MessageEvent
+from opentelemetry.instrumentation.watsonx.utils import (
+    dont_throw,
+    should_emit_events,
+    should_send_prompts,
+)
 from opentelemetry.instrumentation.watsonx.version import __version__
 from opentelemetry.metrics import Counter, Histogram, get_meter
-from opentelemetry.semconv._incubating.attributes import (
-    gen_ai_attributes as GenAIAttributes,
-)
 from opentelemetry.semconv_ai import (
     SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY,
     LLMRequestTypeValues,
@@ -107,6 +107,8 @@ def _set_span_attribute(span, name, value):
 
 
 def _set_api_attributes(span):
+    if not span.is_recording():
+        return
     _set_span_attribute(
         span,
         WatsonxSpanAttributes.WATSONX_API_BASE,
@@ -115,20 +117,15 @@ def _set_api_attributes(span):
     _set_span_attribute(span, WatsonxSpanAttributes.WATSONX_API_TYPE, "watsonx.ai")
     _set_span_attribute(span, WatsonxSpanAttributes.WATSONX_API_VERSION, "1.0")
 
-    return
-
-
-def should_send_prompts():
-    return (
-        os.getenv("TRACELOOP_TRACE_CONTENT") or "true"
-    ).lower() == "true" or context_api.get_value("override_enable_content_tracing")
-
 
 def is_metrics_enabled() -> bool:
     return (os.getenv("TRACELOOP_METRICS_ENABLED") or "true").lower() == "true"
 
 
 def _set_input_attributes(span, instance, kwargs):
+    if not span.is_recording():
+        return
+
     if should_send_prompts() and kwargs is not None and len(kwargs) > 0:
         prompt = kwargs.get("prompt")
         if isinstance(prompt, list):
@@ -145,6 +142,11 @@ def _set_input_attributes(span, instance, kwargs):
                 prompt,
             )
 
+
+def set_model_input_attributes(span, instance):
+    if not span.is_recording():
+        return
+
     _set_span_attribute(span, SpanAttributes.LLM_REQUEST_MODEL, instance.model_id)
     # Set other attributes
     modelParameters = instance.params
@@ -186,10 +188,20 @@ def _set_input_attributes(span, instance, kwargs):
             span, SpanAttributes.LLM_REQUEST_TOP_P, modelParameters.get("top_p", None)
         )
 
-    return
-
 
 def _set_stream_response_attributes(span, stream_response):
+    if not span.is_recording():
+        return
+    _set_span_attribute(
+        span,
+        f"{SpanAttributes.LLM_COMPLETIONS}.0.content",
+        stream_response.get("generated_text"),
+    )
+
+
+def _set_model_stream_response_attributes(span, stream_response):
+    if not span.is_recording():
+        return
     _set_span_attribute(
         span, SpanAttributes.LLM_RESPONSE_MODEL, stream_response.get("model_id")
     )
@@ -211,11 +223,6 @@ def _set_stream_response_attributes(span, stream_response):
         SpanAttributes.LLM_USAGE_TOTAL_TOKENS,
         total_token,
     )
-    _set_span_attribute(
-        span,
-        f"{SpanAttributes.LLM_COMPLETIONS}.0.content",
-        stream_response.get("generated_text"),
-    )
 
 
 def _set_completion_content_attributes(
@@ -263,7 +270,7 @@ def _token_usage_count(responses):
 def _set_response_attributes(
     span, responses, token_histogram, response_counter, duration_histogram, duration
 ):
-    if not isinstance(responses, (list, dict)):
+    if not isinstance(responses, (list, dict)) or not span.is_recording():
         return
 
     if isinstance(responses, list):
@@ -283,6 +290,32 @@ def _set_response_attributes(
         return
     _set_span_attribute(span, SpanAttributes.LLM_RESPONSE_MODEL, model_id)
 
+    shared_attributes = _metric_shared_attributes(response_model=model_id)
+
+    prompt_token, completion_token = _token_usage_count(responses)
+
+    if token_histogram:
+        attributes_with_token_type = {
+            **shared_attributes,
+            SpanAttributes.LLM_TOKEN_TYPE: "output",
+        }
+        token_histogram.record(completion_token, attributes=attributes_with_token_type)
+        attributes_with_token_type = {
+            **shared_attributes,
+            SpanAttributes.LLM_TOKEN_TYPE: "input",
+        }
+        token_histogram.record(prompt_token, attributes=attributes_with_token_type)
+
+    if duration and isinstance(duration, (float, int)) and duration_histogram:
+        duration_histogram.record(duration, attributes=shared_attributes)
+
+
+def set_model_response_attributes(
+    span, responses, token_histogram, duration_histogram, duration
+):
+    if not span.is_recording():
+        return
+
     prompt_token, completion_token = _token_usage_count(responses)
     if (prompt_token + completion_token) != 0:
         _set_span_attribute(
@@ -301,35 +334,16 @@ def _set_response_attributes(
             prompt_token + completion_token,
         )
 
-        shared_attributes = _metric_shared_attributes(response_model=model_id)
 
-        if token_histogram:
-            attributes_with_token_type = {
-                **shared_attributes,
-                SpanAttributes.LLM_TOKEN_TYPE: "output",
-            }
-            token_histogram.record(
-                completion_token, attributes=attributes_with_token_type
-            )
-            attributes_with_token_type = {
-                **shared_attributes,
-                SpanAttributes.LLM_TOKEN_TYPE: "input",
-            }
-            token_histogram.record(prompt_token, attributes=attributes_with_token_type)
-
-    if duration and isinstance(duration, (float, int)) and duration_histogram:
-        duration_histogram.record(duration, attributes=shared_attributes)
-
-
-def _emit_input_events(args, kwargs):
+def _emit_input_events(args, kwargs, event_logger):
     prompt = kwargs.get("prompt") or args[0]
 
     if isinstance(prompt, list):
         for message in prompt:
-            emit_event(MessageEvent(content=message, role="user"))
+            emit_event(MessageEvent(content=message, role="user"), event_logger)
 
     elif isinstance(prompt, str):
-        emit_event(MessageEvent(content=prompt, role="user"))
+        emit_event(MessageEvent(content=prompt, role="user"), event_logger)
 
 
 def _emit_response_events(response: dict):
@@ -345,6 +359,7 @@ def _emit_response_events(response: dict):
 
 def _build_and_set_stream_response(
     span,
+    event_logger,
     response,
     raw_flag,
     token_histogram,
@@ -378,7 +393,9 @@ def _build_and_set_stream_response(
         "generated_token_count": stream_generated_token_count,
         "input_token_count": stream_input_token_count,
     }
-    _set_stream_response_attributes(span, stream_response)
+    _handle_stream_response(
+        span, event_logger, stream_response, stream_generated_text, stream_stop_reason
+    )
     # response counter
     if response_counter:
         attributes_with_reason = {
@@ -412,16 +429,6 @@ def _build_and_set_stream_response(
     if duration and isinstance(duration, (float, int)) and duration_histogram:
         duration_histogram.record(duration, attributes=shared_attributes)
 
-    _emit_response_events(
-        {
-            "results": [
-                {
-                    "stop_reason": stream_stop_reason,
-                    "generated_text": stream_generated_text,
-                }
-            ]
-        },
-    )
     span.set_status(Status(StatusCode.OK))
     span.end()
 
@@ -444,6 +451,7 @@ def _with_tracer(
         response_counter,
         duration_histogram,
         exception_counter,
+        event_logger,
     ):
         def wrapper(wrapped, instance, args, kwargs):
             return func(
@@ -453,6 +461,7 @@ def wrapper(wrapped, instance, args, kwargs):
                 response_counter,
                 duration_histogram,
                 exception_counter,
+                event_logger,
                 wrapped,
                 instance,
                 args,
@@ -464,6 +473,67 @@ def wrapper(wrapped, instance, args, kwargs):
     return _with_tracer
 
 
+@dont_throw
+def _handle_input(span, event_logger, name, instance, response_counter, args, kwargs):
+    _set_api_attributes(span)
+
+    if "generate" in name:
+        set_model_input_attributes(span, instance)
+
+    if should_emit_events() and event_logger:
+        _emit_input_events(args, kwargs, event_logger)
+    elif "generate" in name:
+        _set_input_attributes(span, instance, kwargs)
+
+
+@dont_throw
+def _handle_response(
+    span,
+    event_logger,
+    responses,
+    response_counter,
+    token_histogram,
+    duration_histogram,
+    duration,
+):
+    set_model_response_attributes(
+        span, responses, token_histogram, duration_histogram, duration
+    )
+
+    if should_emit_events() and event_logger:
+        _emit_response_events(responses, event_logger)
+    else:
+        _set_response_attributes(
+            span,
+            responses,
+            token_histogram,
+            response_counter,
+            duration_histogram,
+            duration,
+        )
+
+
+@dont_throw
+def _handle_stream_response(
+    span, event_logger, stream_response, stream_generated_text, stream_stop_reason
+):
+    _set_model_stream_response_attributes(span, stream_response)
+
+    if should_emit_events() and event_logger:
+        _emit_response_events(
+            {
+                "results": [
+                    {
+                        "stop_reason": stream_stop_reason,
+                        "generated_text": stream_generated_text,
+                    }
+                ]
+            },
+        )
+    else:
+        _set_stream_response_attributes(span, stream_response)
+
+
 @_with_tracer_wrapper
 def _wrap(
     tracer,
@@ -472,6 +542,7 @@ def _wrap(
     response_counter: Counter,
     duration_histogram: Histogram,
     exception_counter: Counter,
+    event_logger: Union[EventLogger, None],
     wrapped,
     instance,
     args,
@@ -494,17 +565,15 @@ def _wrap(
         },
     )
 
-    _set_api_attributes(span)
+    _handle_input(span, event_logger, name, instance, args, kwargs)
+
     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
 
-    _emit_input_events(args, kwargs)
-
     try:
         start_time = time.time()
         response = wrapped(*args, **kwargs)
@@ -528,6 +597,7 @@ def _wrap(
         if isinstance(response, types.GeneratorType):
             return _build_and_set_stream_response(
                 span,
+                event_logger,
                 response,
                 raw_flag,
                 token_histogram,
@@ -537,17 +607,15 @@ def _wrap(
             )
         else:
             duration = end_time - start_time
-            _set_response_attributes(
+            _handle_response(
                 span,
+                event_logger,
                 response,
-                token_histogram,
                 response_counter,
+                token_histogram,
                 duration_histogram,
                 duration,
             )
-
-            _emit_response_events(response)
-
     span.end()
     return response
 
@@ -613,9 +681,11 @@ def _instrument(self, **kwargs):
                 None,
             )
 
+        event_logger = None
+
         if not Config.use_legacy_attributes:
             event_logger_provider = kwargs.get("event_logger_provider")
-            Config.event_logger = get_event_logger(
+            event_logger = get_event_logger(
                 __name__, __version__, event_logger_provider=event_logger_provider
             )
 
@@ -634,6 +704,7 @@ def _instrument(self, **kwargs):
                         response_counter,
                         duration_histogram,
                         exception_counter,
+                        event_logger,
                     ),
                 )