wip, refactor chat generation to output message

Author: keith-decker

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

@@ -58,9 +58,8 @@
 from opentelemetry.trace.status import Status, StatusCode
 from opentelemetry.util.types import AttributeValue
 
-from .data import ChatGeneration, Error
 from .instruments import Instruments
-from .types import InputMessage, LLMInvocation
+from .types import Error, InputMessage, LLMInvocation, OutputMessage, Text
 
 
 @dataclass
@@ -111,7 +110,7 @@ def _message_to_log_record(
 
 
 def _chat_generation_to_log_record(
-    chat_generation: ChatGeneration,
+    chat_generation: OutputMessage,
     index: int,
     provider_name: Optional[str],
     framework: Optional[str],
@@ -133,11 +132,18 @@ def _chat_generation_to_log_record(
         "event.name": "gen_ai.choice",
     }
 
+    # Extract message content from parts (first Text part if available)
+    # TODO: use dataclass to dict
+    content: Optional[str] = None
+    for part in chat_generation.parts:
+        if isinstance(part, Text):
+            content = part.content
+            break
     message = {
-        "type": chat_generation.type,
+        "type": chat_generation.role,
     }
-    if capture_content and chat_generation.content:
-        message["content"] = chat_generation.content
+    if capture_content and content is not None:
+        message["content"] = content
 
     body = {
         "index": index,
@@ -218,7 +224,7 @@ def _set_response_and_usage_attributes(
 
 def _emit_chat_generation_logs(
     logger: Optional[Logger],
-    generations: List[ChatGeneration],
+    generations: List[OutputMessage],
     provider_name: Optional[str],
     framework: Optional[str],
     capture_content: bool,
@@ -234,16 +240,14 @@ def _emit_chat_generation_logs(
         )
         if log and logger:
             logger.emit(log)
-        if chat_generation.finish_reason is not None:
-            finish_reasons.append(chat_generation.finish_reason)
+        finish_reasons.append(chat_generation.finish_reason)
     return finish_reasons
 
 
-def _collect_finish_reasons(generations: List[ChatGeneration]) -> List[str]:
+def _collect_finish_reasons(generations: List[OutputMessage]) -> List[str]:
     finish_reasons: List[str] = []
     for gen in generations:
-        if gen.finish_reason is not None:
-            finish_reasons.append(gen.finish_reason)
+        finish_reasons.append(gen.finish_reason)
     return finish_reasons
 
 
@@ -260,15 +264,20 @@ def _maybe_set_input_messages(
 
 
 def _set_chat_generation_attrs(
-    span: Span, generations: List[ChatGeneration]
+    span: Span, generations: List[OutputMessage]
 ) -> None:
     for index, chat_generation in enumerate(generations):
+        # Extract content
+        # TODO: use dataclass to dict - Handle multiple responses
+        content: Optional[str] = None
+        for part in chat_generation.parts:
+            if isinstance(part, Text):
+                content = part.content
+                break
         # Upcoming semconv fields
+        span.set_attribute(f"gen_ai.completion.{index}.content", content or "")
         span.set_attribute(
-            f"gen_ai.completion.{index}.content", chat_generation.content
-        )
-        span.set_attribute(
-            f"gen_ai.completion.{index}.role", chat_generation.type
+            f"gen_ai.completion.{index}.role", chat_generation.role
         )
 
 

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

@@ -45,9 +45,8 @@
 from opentelemetry.semconv.schemas import Schemas
 from opentelemetry.trace import get_tracer
 
-from .data import ChatGeneration, Error
 from .generators import SpanMetricEventGenerator, SpanMetricGenerator
-from .types import InputMessage, LLMInvocation
+from .types import Error, InputMessage, LLMInvocation, OutputMessage
 
 # TODO: Get the tool version for emitting spans, use GenAI Utils for now
 from .version import __version__
@@ -140,7 +139,7 @@ def start_llm(
     def stop_llm(
         self,
         run_id: UUID,
-        chat_generations: List[ChatGeneration],
+        chat_generations: List[OutputMessage],
         **attributes: Any,
     ) -> LLMInvocation:
         with self._lock:
@@ -192,7 +191,7 @@ def llm_start(
 
 
 def llm_stop(
-    run_id: UUID, chat_generations: List[ChatGeneration], **attributes: Any
+    run_id: UUID, chat_generations: List[OutputMessage], **attributes: Any
 ) -> LLMInvocation:
     return get_telemetry_handler().stop_llm(
         run_id=run_id, chat_generations=chat_generations, **attributes

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

@@ -16,11 +16,9 @@
 import time
 from dataclasses import dataclass, field
 from enum import Enum
-from typing import Any, Dict, List, Literal, Optional, Union
+from typing import Any, Dict, List, Literal, Optional, Type, Union
 from uuid import UUID
 
-from .data import ChatGeneration
-
 
 class ContentCapturingMode(Enum):
     # Do not capture content (default).
@@ -86,7 +84,14 @@ class LLMInvocation:
     start_time: float = field(default_factory=time.time)
     end_time: Optional[float] = None
     messages: List[InputMessage] = field(default_factory=list)
-    chat_generations: List[ChatGeneration] = field(default_factory=list)
+    chat_generations: List[OutputMessage] = field(default_factory=list)
     attributes: Dict[str, Any] = field(default_factory=dict)
     span_id: int = 0
     trace_id: int = 0
+
+
+# TODO: Do we need this?
+@dataclass
+class Error:
+    message: str
+    type: Type[BaseException]

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

@@ -17,7 +17,11 @@
     llm_start,
     llm_stop,
 )
-from opentelemetry.util.genai.types import ChatGeneration, InputMessage, Text
+from opentelemetry.util.genai.types import (
+    InputMessage,
+    OutputMessage,
+    Text,
+)
 
 
 class TestTelemetryHandler(unittest.TestCase):
@@ -42,7 +46,9 @@ def test_llm_start_and_stop_creates_span(self):
         message = InputMessage(
             role="Human", parts=[Text(content="hello world")]
         )
-        chat_generation = ChatGeneration(content="hello back", type="AI")
+        chat_generation = OutputMessage(
+            role="AI", parts=[Text(content="hello back")], finish_reason="stop"
+        )
 
         # Start and stop LLM invocation
         llm_start(
@@ -98,9 +104,9 @@ def test_structured_logs_emitted(self):
         message = InputMessage(
             role="user", parts=[Text(content="hello world")]
         )
-        generation = ChatGeneration(
-            content="hello back",
-            type="assistant",
+        generation = OutputMessage(
+            role="assistant",
+            parts=[Text(content="hello back")],
             finish_reason="stop",
         )