flatten LLMInvocation to use attributes instead of dict keys

Author: keith-decker

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

@@ -74,10 +74,20 @@ def start_llm(
     ) -> UUID:
         if run_id is None:
             run_id = uuid.uuid4()
+        provider = attributes.pop("provider", None)
+        response_model_name = attributes.pop("response_model_name", None)
+        response_id = attributes.pop("response_id", None)
+        input_tokens = attributes.pop("input_tokens", None)
+        output_tokens = attributes.pop("output_tokens", None)
         invocation = LLMInvocation(
             request_model=request_model,
             messages=prompts,
             run_id=run_id,
+            provider=provider,
+            response_model_name=response_model_name,
+            response_id=response_id,
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
             attributes=attributes,
         )
         self._llm_registry[invocation.run_id] = invocation
@@ -93,6 +103,19 @@ def stop_llm(
         invocation = self._llm_registry.pop(run_id)
         invocation.end_time = time.time()
         invocation.chat_generations = chat_generations
+        if "provider" in attributes:
+            invocation.provider = attributes.pop("provider")
+        if "response_model_name" in attributes:
+            invocation.response_model_name = attributes.pop(
+                "response_model_name"
+            )
+        if "response_id" in attributes:
+            invocation.response_id = attributes.pop("response_id")
+        if "input_tokens" in attributes:
+            invocation.input_tokens = attributes.pop("input_tokens")
+        if "output_tokens" in attributes:
+            invocation.output_tokens = attributes.pop("output_tokens")
+        # Keep any remaining attributes
         invocation.attributes.update(attributes)
         self._generator.finish(invocation)
         return invocation
@@ -102,6 +125,18 @@ def fail_llm(
     ) -> LLMInvocation:
         invocation = self._llm_registry.pop(run_id)
         invocation.end_time = time.time()
+        if "provider" in attributes:
+            invocation.provider = attributes.pop("provider")
+        if "response_model_name" in attributes:
+            invocation.response_model_name = attributes.pop(
+                "response_model_name"
+            )
+        if "response_id" in attributes:
+            invocation.response_id = attributes.pop("response_id")
+        if "input_tokens" in attributes:
+            invocation.input_tokens = attributes.pop("input_tokens")
+        if "output_tokens" in attributes:
+            invocation.output_tokens = attributes.pop("output_tokens")
         invocation.attributes.update(**attributes)
         self._generator.error(error, invocation)
         return invocation

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

@@ -14,7 +14,7 @@
 
 import json
 from dataclasses import asdict
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List
 
 from opentelemetry.semconv._incubating.attributes import (
     gen_ai_attributes as GenAI,
@@ -31,7 +31,6 @@
     get_content_capturing_mode,
     is_experimental_mode,
 )
-from opentelemetry.util.types import AttributeValue
 
 from .types import Error, InputMessage, LLMInvocation, OutputMessage
 
@@ -43,8 +42,8 @@ def _apply_common_span_attributes(
 
     Returns (genai_attributes) for use with metrics.
     """
-    request_model = invocation.attributes.get("request_model")
-    provider = invocation.attributes.get("provider")
+    request_model = invocation.request_model
+    provider = invocation.provider
 
     span.set_attribute(
         GenAI.GEN_AI_OPERATION_NAME, GenAI.GenAiOperationNameValues.CHAT.value
@@ -63,34 +62,20 @@ def _apply_common_span_attributes(
             GenAI.GEN_AI_RESPONSE_FINISH_REASONS, finish_reasons
         )
 
-    response_model = invocation.attributes.get("response_model_name")
-    response_id = invocation.attributes.get("response_id")
-    prompt_tokens = invocation.attributes.get("input_tokens")
-    completion_tokens = invocation.attributes.get("output_tokens")
-    _set_response_and_usage_attributes(
-        span,
-        response_model,
-        response_id,
-        prompt_tokens,
-        completion_tokens,
-    )
-
-
-def _set_response_and_usage_attributes(
-    span: Span,
-    response_model: Optional[str],
-    response_id: Optional[str],
-    prompt_tokens: Optional[AttributeValue],
-    completion_tokens: Optional[AttributeValue],
-) -> None:
-    if response_model is not None:
-        span.set_attribute(GenAI.GEN_AI_RESPONSE_MODEL, response_model)
-    if response_id is not None:
-        span.set_attribute(GenAI.GEN_AI_RESPONSE_ID, response_id)
-    if isinstance(prompt_tokens, (int, float)):
-        span.set_attribute(GenAI.GEN_AI_USAGE_INPUT_TOKENS, prompt_tokens)
-    if isinstance(completion_tokens, (int, float)):
-        span.set_attribute(GenAI.GEN_AI_USAGE_OUTPUT_TOKENS, completion_tokens)
+    if invocation.response_model_name is not None:
+        span.set_attribute(
+            GenAI.GEN_AI_RESPONSE_MODEL, invocation.response_model_name
+        )
+    if invocation.response_id is not None:
+        span.set_attribute(GenAI.GEN_AI_RESPONSE_ID, invocation.response_id)
+    if isinstance(invocation.input_tokens, (int, float)):
+        span.set_attribute(
+            GenAI.GEN_AI_USAGE_INPUT_TOKENS, invocation.input_tokens
+        )
+    if isinstance(invocation.output_tokens, (int, float)):
+        span.set_attribute(
+            GenAI.GEN_AI_USAGE_OUTPUT_TOKENS, invocation.output_tokens
+        )
 
 
 def _maybe_set_span_messages(

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

@@ -19,6 +19,8 @@
 from typing import Any, Dict, List, Literal, Optional, Type, Union
 from uuid import UUID
 
+from opentelemetry.util.types import AttributeValue
+
 
 class ContentCapturingMode(Enum):
     # Do not capture content (default).
@@ -86,9 +88,12 @@ class LLMInvocation:
     end_time: Optional[float] = None
     messages: List[InputMessage] = field(default_factory=list)
     chat_generations: List[OutputMessage] = field(default_factory=list)
+    provider: Optional[str] = None
+    response_model_name: Optional[str] = None
+    response_id: Optional[str] = None
+    input_tokens: Optional[AttributeValue] = None
+    output_tokens: Optional[AttributeValue] = None
     attributes: Dict[str, Any] = field(default_factory=dict)
-    span_id: int = 0
-    trace_id: int = 0
 
 
 @dataclass