Move choice event creation to _Choice and add types

Author: xrmx

instrumentation/opentelemetry-instrumentation-botocore/src/opentelemetry/instrumentation/botocore/extensions/bedrock.py

@@ -32,7 +32,6 @@
     _Choice,
     genai_capture_message_content,
     message_to_event,
-    to_choice_event,
 )
 from opentelemetry.instrumentation.botocore.extensions.types import (
     _AttributeMapT,
@@ -293,8 +292,7 @@ def _converse_on_success(
         # context so need to add the span context manually
         span_ctx = span.get_span_context()
         event_logger.emit(
-            to_choice_event(
-                choice,
+            choice.to_choice_event(
                 trace_id=span_ctx.trace_id,
                 span_id=span_ctx.span_id,
                 trace_flags=span_ctx.trace_flags,
@@ -444,7 +442,7 @@ def _handle_amazon_titan_response(
             choice = _Choice.from_invoke_amazon_titan(
                 response_body, capture_content
             )
-            event_logger.emit(to_choice_event(choice))
+            event_logger.emit(choice.to_choice_event())
 
     # pylint: disable=no-self-use
     def _handle_amazon_nova_response(
@@ -471,7 +469,7 @@ def _handle_amazon_nova_response(
 
         event_logger = instrumentor_context.event_logger
         choice = _Choice.from_converse(response_body, capture_content)
-        event_logger.emit(to_choice_event(choice))
+        event_logger.emit(choice.to_choice_event())
 
     # pylint: disable=no-self-use
     def _handle_anthropic_claude_response(
@@ -499,7 +497,7 @@ def _handle_anthropic_claude_response(
         choice = _Choice.from_invoke_anthropic_claude(
             response_body, capture_content
         )
-        event_logger.emit(to_choice_event(choice))
+        event_logger.emit(choice.to_choice_event())
 
     def on_error(
         self,

instrumentation/opentelemetry-instrumentation-botocore/src/opentelemetry/instrumentation/botocore/extensions/bedrock_utils.py

@@ -16,7 +16,7 @@
 
 import json
 from os import environ
-from typing import Callable, Dict, Union
+from typing import Any, Callable, Dict, Union
 
 from botocore.eventstream import EventStream, EventStreamError
 from wrapt import ObjectProxy
@@ -297,7 +297,7 @@ def genai_capture_message_content() -> bool:
     return capture_content.lower() == "true"
 
 
-def message_to_event(message, capture_content):
+def message_to_event(message: dict[str, Any], capture_content: bool) -> Event:
     attributes = {GEN_AI_SYSTEM: GenAiSystemValues.AWS_BEDROCK.value}
     role = message.get("role")
     content = message.get("content")
@@ -314,13 +314,17 @@ def message_to_event(message, capture_content):
 
 
 class _Choice:
-    def __init__(self, message, finish_reason, index):
+    def __init__(
+        self, message: dict[str, Any], finish_reason: str, index: int
+    ):
         self.message = message
         self.finish_reason = finish_reason
         self.index = index
 
     @classmethod
-    def from_converse(cls, response, capture_content):
+    def from_converse(
+        cls, response: dict[str, Any], capture_content: bool
+    ) -> _Choice:
         orig_message = response["output"]["message"]
         if role := orig_message.get("role"):
             message = {"role": role}
@@ -332,7 +336,9 @@ def from_converse(cls, response, capture_content):
         return cls(message, response["stopReason"], index=0)
 
     @classmethod
-    def from_invoke_amazon_titan(cls, response, capture_content):
+    def from_invoke_amazon_titan(
+        cls, response: dict[str, Any], capture_content: bool
+    ) -> _Choice:
         result = response["results"][0]
         if capture_content:
             message = {"content": result["outputText"]}
@@ -341,7 +347,9 @@ def from_invoke_amazon_titan(cls, response, capture_content):
         return cls(message, result["completionReason"], index=0)
 
     @classmethod
-    def from_invoke_anthropic_claude(cls, response, capture_content):
+    def from_invoke_anthropic_claude(
+        cls, response: dict[str, Any], capture_content: bool
+    ) -> _Choice:
         if capture_content:
             message = {
                 "content": response["content"],
@@ -352,19 +360,18 @@ def from_invoke_anthropic_claude(cls, response, capture_content):
 
         return cls(message, response["stop_reason"], index=0)
 
-    def to_body_dict(self):
+    def _to_body_dict(self) -> dict[str, Any]:
         return {
             "finish_reason": self.finish_reason,
             "index": self.index,
             "message": self.message,
         }
 
-
-def to_choice_event(choice: _Choice, **event_kwargs):
-    attributes = {GEN_AI_SYSTEM: GenAiSystemValues.AWS_BEDROCK.value}
-    return Event(
-        name="gen_ai.choice",
-        attributes=attributes,
-        body=choice.to_body_dict(),
-        **event_kwargs,
-    )
+    def to_choice_event(self, **event_kwargs) -> Event:
+        attributes = {GEN_AI_SYSTEM: GenAiSystemValues.AWS_BEDROCK.value}
+        return Event(
+            name="gen_ai.choice",
+            attributes=attributes,
+            body=self._to_body_dict(),
+            **event_kwargs,
+        )