add gen ai completion extraction and traceloop extraction + black formatting

Author: yiyuan-he

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/llo_handler.py

@@ -12,9 +12,12 @@
 GEN_AI_SYSTEM_MESSAGE = "gen_ai.system.message"
 GEN_AI_USER_MESSAGE = "gen_ai.user.message"
 GEN_AI_ASSISTANT_MESSAGE = "gen_ai.assistant.message"
+TRACELOOP_ENTITY_INPUT = "traceloop.entity.input"
+TRACELOOP_ENTITY_OUTPUT = "traceloop.entity.output"
 
 _logger = logging.getLogger(__name__)
 
+
 class LLOHandler:
     """
     Utility class for handling Large Language Objects (LLO) in OpenTelemetry spans.
@@ -24,6 +27,7 @@ class LLOHandler:
     2. Extracts and transforms these attributes into OpenTelemetry Gen AI Events
     3. Filters LLO from spans
     """
+
     def __init__(self, logger_provider: LoggerProvider):
         """
         Initialize an LLOHandler with the specified logger provider.
@@ -37,12 +41,15 @@ def __init__(self, logger_provider: LoggerProvider):
         self._event_logger_provider = EventLoggerProvider(logger_provider=self._logger_provider)
         self._event_logger = self._event_logger_provider.get_event_logger("gen_ai.events")
 
-        self._exact_match_patterns = []
+        self._exact_match_patterns = [
+            "traceloop.entity.input",
+            "traceloop.entity.output",
+        ]
         self._regex_match_patterns = [
-            r"^gen_ai\.prompt\.\d+\.content$"
+            r"^gen_ai\.prompt\.\d+\.content$",
+            r"^gen_ai\.completion\.\d+\.content$",
         ]
 
-
     def process_spans(self, spans: Sequence[ReadableSpan]) -> List[ReadableSpan]:
         """
         Performs LLO processing for each span:
@@ -66,7 +73,7 @@ def process_spans(self, spans: Sequence[ReadableSpan]) -> List[ReadableSpan]:
                     maxlen=span.attributes.maxlen,
                     attributes=updated_attributes,
                     immutable=span.attributes._immutable,
-                    max_value_len=span.attributes.max_value_len
+                    max_value_len=span.attributes.max_value_len,
                 )
             else:
                 span._attributes = updated_attributes
@@ -75,7 +82,6 @@ def process_spans(self, spans: Sequence[ReadableSpan]) -> List[ReadableSpan]:
 
         return modified_spans
 
-
     def _emit_llo_attributes(self, span: ReadableSpan, attributes: Dict[str, Any]) -> None:
         """
         Collects the Gen AI Events for each LLO attribute in the span and emits them
@@ -90,12 +96,13 @@ def _emit_llo_attributes(self, span: ReadableSpan, attributes: Dict[str, Any]) -
         """
         all_events = []
         all_events.extend(self._extract_gen_ai_prompt_events(span, attributes))
+        all_events.extend(self._extract_gen_ai_completion_events(span, attributes))
+        all_events.extend(self._extract_traceloop_events(span, attributes))
 
         for event in all_events:
             self._event_logger.emit(event)
             _logger.debug(f"Emitted Gen AI Event: {event.name}")
 
-
     def _filter_attributes(self, attributes: Dict[str, Any]) -> Dict[str, Any]:
         """
         Filter out attributes that contain LLO from the span's attributes. This
@@ -116,7 +123,6 @@ def _filter_attributes(self, attributes: Dict[str, Any]) -> Dict[str, Any]:
 
         return filtered_attributes
 
-
     def _is_llo_attribute(self, key: str) -> bool:
         """
         Determine if a span attribute contains an LLO based on its key name.
@@ -131,12 +137,10 @@ def _is_llo_attribute(self, key: str) -> bool:
         Returns:
             bool: True if the key matches an LLO pattern, False otherwise
         """
-        return (
-            any(pattern == key for pattern in self._exact_match_patterns) or
-            any(re.match(pattern, key) for pattern in self._regex_match_patterns)
+        return any(pattern == key for pattern in self._exact_match_patterns) or any(
+            re.match(pattern, key) for pattern in self._regex_match_patterns
         )
 
-
     def _extract_gen_ai_prompt_events(self, span: ReadableSpan, attributes: Dict[str, Any]) -> List[Event]:
         """
         Extract gen_ai prompt events from attributes. Each item `gen_ai.prompt.{n}.content`
@@ -173,15 +177,9 @@ def _extract_gen_ai_prompt_events(self, span: ReadableSpan, attributes: Dict[str
             role_key = f"gen_ai.prompt.{index}.role"
             role = attributes.get(role_key, "unknown")
 
-            event_attributes = {
-                "gen_ai.system": gen_ai_system,
-                "original_attribute": key
-            }
+            event_attributes = {"gen_ai.system": gen_ai_system, "original_attribute": key}
 
-            body = {
-                "content": value,
-                "role": role
-            }
+            body = {"content": value, "role": role}
 
             event = None
             if role == "system":
@@ -190,46 +188,104 @@ def _extract_gen_ai_prompt_events(self, span: ReadableSpan, attributes: Dict[str
                     span_ctx=span_ctx,
                     timestamp=prompt_timestamp,
                     attributes=event_attributes,
-                    body=body
+                    body=body,
                 )
             elif role == "user":
                 event = self._get_gen_ai_event(
                     name=GEN_AI_USER_MESSAGE,
                     span_ctx=span_ctx,
                     timestamp=prompt_timestamp,
                     attributes=event_attributes,
-                    body=body
+                    body=body,
                 )
             elif role == "assistant":
                 event = self._get_gen_ai_event(
                     name=GEN_AI_ASSISTANT_MESSAGE,
                     span_ctx=span_ctx,
                     timestamp=prompt_timestamp,
                     attributes=event_attributes,
-                    body=body
+                    body=body,
                 )
             elif role in ["function", "unknown"]:
                 event = self._get_gen_ai_event(
                     name=f"gen_ai.{gen_ai_system}.message",
                     span_ctx=span_ctx,
                     timestamp=prompt_timestamp,
                     attributes=event_attributes,
-                    body=body
+                    body=body,
                 )
 
             if event:
                 events.append(event)
 
         return events
 
-    def _get_gen_ai_event(
-        self,
-        name,
-        span_ctx,
-        timestamp,
-        attributes,
-        body
-    ):
+    def _extract_gen_ai_completion_events(self, span: ReadableSpan, attributes: Dict[str, Any]) -> List[Event]:
+        events = []
+        span_ctx = span.context
+        gen_ai_system = span.attributes.get("gen_ai.system", "unknown")
+
+        completion_timestamp = span.end_time
+
+        completion_content_pattern = re.compile(r"^gen_ai\.completion\.(\d+)\.content$")
+
+        for key, value in attributes.items():
+            match = completion_content_pattern.match(key)
+            if not match:
+                continue
+
+            index = match.group(1)
+            role_key = f"gen_ai.completion.{index}.role"
+            role = attributes.get(role_key, "unknown")
+
+            event_attributes = {"gen_ai.system": gen_ai_system, "original_attribute": key}
+
+            body = {"content": value, "role": role}
+
+            event = None
+            if role == "assistant":
+                event = self._get_gen_ai_event(
+                    name=GEN_AI_ASSISTANT_MESSAGE,
+                    span_ctx=span_ctx,
+                    timestamp=completion_timestamp,
+                    attributes=event_attributes,
+                    body=body,
+                )
+            else:
+                event = self._get_gen_ai_event(
+                    name=f"gen_ai.{gen_ai_system}.message",
+                    span_ctx=span_ctx,
+                    timestamp=completion_timestamp,
+                    attributes=event_attributes,
+                    body=body,
+                )
+
+            if event:
+                events.append(event)
+
+        return events
+
+    def _extract_traceloop_events(self, span: ReadableSpan, attributes: Dict[str, Any]) -> List[Event]:
+        events = []
+        span_ctx = span.context
+        gen_ai_system = span.attributes.get("traceloop.entity.name", "unknown")
+
+        traceloop_attrs = [(TRACELOOP_ENTITY_INPUT, span.start_time), (TRACELOOP_ENTITY_OUTPUT, span.end_time)]
+
+        for attr_key, timestamp in traceloop_attrs:
+            if attr_key in attributes:
+                event = self._get_gen_ai_event(
+                    name=f"gen_ai.{gen_ai_system}.message",
+                    span_ctx=span_ctx,
+                    timestamp=timestamp,
+                    attributes={"gen_ai.system": gen_ai_system, "original_attribute": attr_key},
+                    body={"content": attributes[attr_key]},
+                )
+                events.append(event)
+
+        return events
+
+    def _get_gen_ai_event(self, name, span_ctx, timestamp, attributes, body):
         """
         Create and return a Gen AI Event with the provided parameters.
 
@@ -250,5 +306,5 @@ def _get_gen_ai_event(
             body=body,
             trace_id=span_ctx.trace_id,
             span_id=span_ctx.span_id,
-            trace_flags=span_ctx.trace_flags
+            trace_flags=span_ctx.trace_flags,
         )