Done some code refactoring

Author: adharshctr

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

@@ -494,7 +494,6 @@ def wrapper(wrapped, instance, args, kwargs):
 
 @dont_throw
 def _handle_input(span, event_logger, name, instance, args, kwargs):
-    messages = None
     _set_api_attributes(span)
     if "generate" in name:
         set_model_input_attributes(span, instance)
@@ -503,35 +502,19 @@ def _handle_input(span, event_logger, name, instance, args, kwargs):
 
     if "chat" in name:
         messages = kwargs.get("messages")
-        if messages is None and args and isinstance(args[0], list):
-            messages = args[0]
         if messages and span.is_recording():
             for index, msg in enumerate(messages):
-                if not isinstance(msg, dict):
-                    continue
                 role = msg.get("role")
                 content = msg.get("content")
-                if role and isinstance(content, str):
+                if role and content:
                     _set_span_attribute(
                         span,
                         f"{GenAIAttributes.GEN_AI_PROMPT}.{index}.{role}",
                         content.strip(),
                     )
 
     if should_emit_events() and event_logger:
-        if "chat" in name and isinstance(messages, list):
-            for msg in messages:
-                if not isinstance(msg, dict):
-                    continue
-                emit_event(
-                    MessageEvent(
-                        content=msg.get("content"),
-                        role=msg.get("role", "user"),
-                    ),
-                    event_logger,
-                )
-        else:
-            _emit_input_events(args, kwargs, event_logger)
+        _emit_input_events(args, kwargs, event_logger)
 
 
 @dont_throw
@@ -592,30 +575,19 @@ def _handle_chat_response(
     duration_histogram,
     duration,
 ):
-    if not span.is_recording() or not isinstance(response, dict):
+    if not span.is_recording():
         return
 
-    model_id = response.get("model_id") or response.get("model") or "unknown"
-
+    model_id = response.get("model_id") or response.get("model")
     _set_span_attribute(span, GenAIAttributes.GEN_AI_RESPONSE_MODEL, model_id)
 
     # Content
-    choices = response.get("choices") or []
+    choices = response.get("choices", [])
     for index, choice in enumerate(choices):
         message = choice.get("message", {})
         content = message.get("content")
-        finish_reason = choice.get("finish_reason") or "unknown"
-
-        if should_emit_events() and event_logger:
-            emit_event(
-                ChoiceEvent(
-                    index=index,
-                    message=message,
-                    finish_reason=finish_reason or "unknown",
-                ),
-                event_logger,
-            )
-        elif content and should_send_prompts():
+        finish_reason = choice.get("finish_reason")
+        if content and should_send_prompts():
             _set_span_attribute(
                 span,
                 f"{GenAIAttributes.GEN_AI_COMPLETION}.{index}.content",
@@ -630,7 +602,7 @@ def _handle_chat_response(
             response_counter.add(1, attributes=attributes)
 
     # Usage
-    usage = response.get("usage") or {}
+    usage = response.get("usage", {})
     prompt_tokens = usage.get("prompt_tokens", 0)
     completion_tokens = usage.get("completion_tokens", 0)
     total_tokens = usage.get("total_tokens", prompt_tokens + completion_tokens)
@@ -686,7 +658,7 @@ def _wrap(
         },
     )
 
-
+    _handle_input(span, event_logger, name, instance, args, kwargs)
     if "generate" in name or "chat" in name:
         if to_wrap.get("method") == "generate_text_stream":
             if (raw_flag := kwargs.get("raw_response", None)) is None:
@@ -700,7 +672,6 @@ def _wrap(
                 kwargs["messages"] = [
                     {"role": "user", "content": prompt}
                 ]
-    _handle_input(span, event_logger, name, instance, args, kwargs)
 
     try:
         start_time = time.time()