Refatoring and adding input token metric for Embedding invocation

Author: shuningc

util/opentelemetry-util-genai/src/opentelemetry/util/genai/metrics.py

@@ -18,7 +18,11 @@
     create_duration_histogram,
     create_token_histogram,
 )
-from opentelemetry.util.genai.types import GenAIInvocation, LLMInvocation
+from opentelemetry.util.genai.types import (
+    EmbeddingInvocation,
+    GenAIInvocation,
+    LLMInvocation,
+)
 from opentelemetry.util.types import AttributeValue
 
 
@@ -37,18 +41,15 @@ def _build_attributes(
         """Build metric attributes from an invocation."""
         attributes: Dict[str, AttributeValue] = {}
 
-        # Set attributes using getattr for fields that may not exist on base class
-        operation_name = getattr(invocation, "operation_name", None)
-        if operation_name:
-            attributes[GenAI.GEN_AI_OPERATION_NAME] = operation_name
+        if invocation.operation_name:
+            attributes[GenAI.GEN_AI_OPERATION_NAME] = invocation.operation_name
 
         request_model = getattr(invocation, "request_model", None)
         if request_model:
             attributes[GenAI.GEN_AI_REQUEST_MODEL] = request_model
 
-        provider = getattr(invocation, "provider", None)
-        if provider:
-            attributes[GenAI.GEN_AI_PROVIDER_NAME] = provider
+        if invocation.provider:
+            attributes[GenAI.GEN_AI_PROVIDER_NAME] = invocation.provider
 
         response_model_name = getattr(invocation, "response_model_name", None)
         if response_model_name:
@@ -62,9 +63,8 @@ def _build_attributes(
         if server_port is not None:
             attributes[server_attributes.SERVER_PORT] = server_port
 
-        metric_attributes = getattr(invocation, "metric_attributes", None)
-        if metric_attributes:
-            attributes.update(metric_attributes)
+        if invocation.metric_attributes:
+            attributes.update(invocation.metric_attributes)
 
         if error_type:
             attributes[error_attributes.ERROR_TYPE] = error_type
@@ -108,29 +108,27 @@ def record(
                 context=span_context,
             )
 
-        # Only record token metrics for LLMInvocation
-        if isinstance(invocation, LLMInvocation):
-            token_counts: list[tuple[int, str]] = []
+        # Record token metrics for LLMInvocation and EmbeddingInvocation
+        if isinstance(invocation, (LLMInvocation, EmbeddingInvocation)):
             if invocation.input_tokens is not None:
-                token_counts.append(
-                    (
-                        invocation.input_tokens,
-                        GenAI.GenAiTokenTypeValues.INPUT.value,
-                    )
-                )
-            if invocation.output_tokens is not None:
-                token_counts.append(
-                    (
-                        invocation.output_tokens,
-                        GenAI.GenAiTokenTypeValues.OUTPUT.value,
-                    )
+                self._token_histogram.record(
+                    invocation.input_tokens,
+                    attributes=attributes
+                    | {
+                        GenAI.GEN_AI_TOKEN_TYPE: GenAI.GenAiTokenTypeValues.INPUT.value
+                    },
+                    context=span_context,
                 )
 
-            for token_count, token_type in token_counts:
+        # Only LLMInvocation has output tokens
+        if isinstance(invocation, LLMInvocation):
+            if invocation.output_tokens is not None:
                 self._token_histogram.record(
-                    token_count,
+                    invocation.output_tokens,
                     attributes=attributes
-                    | {GenAI.GEN_AI_TOKEN_TYPE: token_type},
+                    | {
+                        GenAI.GEN_AI_TOKEN_TYPE: GenAI.GenAiTokenTypeValues.OUTPUT.value
+                    },
                     context=span_context,
                 )
 

util/opentelemetry-util-genai/src/opentelemetry/util/genai/types.py

@@ -256,9 +256,18 @@ def _new_str_any_dict() -> dict[str, Any]:
 
 @dataclass
 class GenAIInvocation:
+    operation_name: str = ""
+    provider: str | None = None
     context_token: ContextToken | None = None
     span: Span | None = None
     attributes: dict[str, Any] = field(default_factory=_new_str_any_dict)
+    metric_attributes: dict[str, Any] = field(
+        default_factory=_new_str_any_dict
+    )
+    """
+    Additional attributes to set on metrics. Must be of a low cardinality.
+    These attributes will not be set on spans or events.
+    """
     error_type: str | None = None
 
     monotonic_start_s: float | None = None
@@ -319,13 +328,6 @@ class LLMInvocation(GenAIInvocation):
     Additional attributes to set on spans and/or events. These attributes
     will not be set on metrics.
     """
-    metric_attributes: dict[str, Any] = field(
-        default_factory=_new_str_any_dict
-    )
-    """
-    Additional attributes to set on metrics. Must be of a low cardinality.
-    These attributes will not be set on spans or events.
-    """
     temperature: float | None = None
     top_p: float | None = None
     frequency_penalty: float | None = None
@@ -364,14 +366,6 @@ class EmbeddingInvocation(GenAIInvocation):
     will not be set on metrics.
     """
 
-    metric_attributes: dict[str, Any] = field(
-        default_factory=_new_str_any_dict
-    )
-    """
-    Additional attributes to set on metrics. Must be of a low cardinality.
-    These attributes will not be set on spans or events.
-    """
-
 
 @dataclass()
 class ToolCall(GenAIInvocation):
@@ -399,6 +393,8 @@ class ToolCall(GenAIInvocation):
     - error.type: Error type if operation failed (Conditionally Required)
     """
 
+    operation_name: str = GenAI.GenAiOperationNameValues.EXECUTE_TOOL.value
+
     # Message identification fields (same as ToolCallRequest)
     # Note: These are required fields but must have defaults due to dataclass inheritance
     name: str = ""

util/opentelemetry-util-genai/tests/test_handler_metrics.py

@@ -186,16 +186,16 @@ def _assert_metric_scope_schema_urls(
                     scope_metric.scope.schema_url, expected_schema_url
                 )
 
-    def test_stop_embedding_records_duration_only(self) -> None:
-        """Verify embedding invocations record duration but not token metrics."""
+    def test_stop_embedding_records_duration_and_tokens(self) -> None:
+        """Verify embedding invocations record duration and input token metrics."""
         handler = TelemetryHandler(
             tracer_provider=self.tracer_provider,
             meter_provider=self.meter_provider,
         )
         invocation = EmbeddingInvocation(
             request_model="embed-model", provider="embed-prov"
         )
-        invocation.input_tokens = 100  # Should NOT be recorded as token metric
+        invocation.input_tokens = 100
         # Patch default_timer during start to ensure monotonic_start_s
         with patch("timeit.default_timer", return_value=1000.0):
             handler.start(invocation)
@@ -225,8 +225,16 @@ def test_stop_embedding_records_duration_only(self) -> None:
         )
         self.assertAlmostEqual(duration_point.sum, 1.5, places=3)
 
-        # Token metrics should NOT be recorded for embedding
-        self.assertNotIn("gen_ai.client.token.usage", metrics)
+        # Token metrics should be recorded for embedding (input only)
+        self.assertIn("gen_ai.client.token.usage", metrics)
+        token_points = metrics["gen_ai.client.token.usage"]
+        self.assertEqual(len(token_points), 1)  # Only input tokens
+        token_point = token_points[0]
+        self.assertEqual(
+            token_point.attributes[GenAI.GEN_AI_TOKEN_TYPE],
+            GenAI.GenAiTokenTypeValues.INPUT.value,
+        )
+        self.assertAlmostEqual(token_point.sum, 100.0, places=3)
 
     def test_stop_embedding_records_duration_with_additional_attributes(
         self,
@@ -300,5 +308,26 @@ def test_fail_embedding_records_error_and_duration(self) -> None:
         )
         self.assertAlmostEqual(duration_point.sum, 2.5, places=3)
 
-        # Token metrics should NOT be recorded for embedding even on failure
+        # Token metrics should NOT be recorded when input_tokens is not set
+        self.assertNotIn("gen_ai.client.token.usage", metrics)
+
+    def test_stop_embedding_without_tokens(self) -> None:
+        """Verify embedding without input_tokens does not record token metrics."""
+        handler = TelemetryHandler(
+            tracer_provider=self.tracer_provider,
+            meter_provider=self.meter_provider,
+        )
+        invocation = EmbeddingInvocation(
+            request_model="embed-model", provider="embed-prov"
+        )
+        # input_tokens is not set
+        handler.start(invocation)
+        handler.stop(invocation)
+
+        metrics = self._harvest_metrics()
+
+        # Duration should be recorded
+        self.assertIn("gen_ai.client.operation.duration", metrics)
+
+        # Token metrics should NOT be recorded when input_tokens is not set
         self.assertNotIn("gen_ai.client.token.usage", metrics)