formatting

Author: liustve

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

@@ -29,7 +29,7 @@
 from amazon.opentelemetry.distro.sampler.aws_xray_remote_sampler import AwsXRayRemoteSampler
 from amazon.opentelemetry.distro.scope_based_exporter import ScopeBasedPeriodicExportingMetricReader
 from amazon.opentelemetry.distro.scope_based_filtering_view import ScopeBasedRetainingView
-from opentelemetry._logs import set_logger_provider, get_logger_provider
+from opentelemetry._logs import get_logger_provider, set_logger_provider
 from opentelemetry.exporter.otlp.proto.http._log_exporter import OTLPLogExporter
 from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter as OTLPHttpOTLPMetricExporter
 from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/exporter/otlp/aws/logs/otlp_aws_logs_exporter.py

@@ -71,7 +71,8 @@ def export(self, batch: Sequence[LogData]) -> LogExportResult:
 
         1. Always compresses the serialized data into gzip before sending.
 
-        2. If self._gen_ai_flag is enabled, the log data is > 1 MB and the assumption is that the log is a normalized gen.ai LogEvent.
+        2. If self._gen_ai_flag is enabled, the log data is > 1 MB a
+           and the assumption is that the log is a normalized gen.ai LogEvent.
             - inject the 'x-aws-log-semantics' flag into the header.
 
         3. Retry behavior is now the following:

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/exporter/otlp/aws/traces/otlp_aws_span_exporter.py

@@ -3,12 +3,12 @@
 
 from typing import Dict, Optional, Sequence
 
+from amazon.opentelemetry.distro._utils import is_agent_observability_enabled
 from amazon.opentelemetry.distro.exporter.otlp.aws.common.aws_auth_session import AwsAuthSession
 from amazon.opentelemetry.distro.llo_handler import LLOHandler
-from amazon.opentelemetry.distro._utils import is_agent_observability_enabled
-from opentelemetry.sdk._logs import LoggerProvider
 from opentelemetry.exporter.otlp.proto.http import Compression
 from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from opentelemetry.sdk._logs import LoggerProvider
 from opentelemetry.sdk.trace import ReadableSpan
 from opentelemetry.sdk.trace.export import SpanExportResult
 
@@ -25,7 +25,7 @@ def __init__(
         headers: Optional[Dict[str, str]] = None,
         timeout: Optional[int] = None,
         compression: Optional[Compression] = None,
-        logger_provider: Optional[LoggerProvider] = None
+        logger_provider: Optional[LoggerProvider] = None,
     ):
         self._aws_region = None
 
@@ -52,4 +52,4 @@ def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
             llo_processed_spans = self._llo_handler.process_spans(spans)
             return super().export(llo_processed_spans)
 
-        return super().export(spans)
+        return super().export(spans)

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

@@ -1,13 +1,13 @@
 import logging
 import re
-
 from typing import Any, Dict, List, Optional, Sequence
 
-from opentelemetry.attributes import BoundedAttributes
 from opentelemetry._events import Event
-from opentelemetry.sdk._logs import LoggerProvider
+from opentelemetry.attributes import BoundedAttributes
 from opentelemetry.sdk._events import EventLoggerProvider
-from opentelemetry.sdk.trace import ReadableSpan, Event as SpanEvent
+from opentelemetry.sdk._logs import LoggerProvider
+from opentelemetry.sdk.trace import Event as SpanEvent
+from opentelemetry.sdk.trace import ReadableSpan
 
 # Message event types
 GEN_AI_SYSTEM_MESSAGE = "gen_ai.system.message"
@@ -678,4 +678,4 @@ def _get_gen_ai_event(self, name, span_ctx, timestamp, attributes, body):
             trace_id=span_ctx.trace_id,
             span_id=span_ctx.span_id,
             trace_flags=span_ctx.trace_flags,
-            )
+        )

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_bedrock_patches.py

@@ -220,7 +220,7 @@ def extract_attributes(self, attributes: _AttributeMapT):
         knowledge_base_id = self._call_context.params.get(_KNOWLEDGE_BASE_ID)
         if knowledge_base_id:
             attributes[AWS_BEDROCK_KNOWLEDGE_BASE_ID] = knowledge_base_id
-            
+
     def on_success(self, span: Span, result: _BotoResultT, instrumentor_context=None):
         # Currently no attributes to extract from the result
         pass

aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_llo_handler.py

@@ -1,5 +1,5 @@
 from unittest import TestCase
-from unittest.mock import MagicMock, patch, call
+from unittest.mock import MagicMock, call, patch
 
 from amazon.opentelemetry.distro.llo_handler import LLOHandler
 from opentelemetry._events import Event
@@ -67,21 +67,21 @@ def test_is_llo_attribute_no_match(self):
         self.assertFalse(self.llo_handler._is_llo_attribute("some.other.attribute"))
 
     def test_is_llo_attribute_traceloop_match(self):
-      """
-      Test _is_llo_attribute method with Traceloop patterns
-      """
-      # Test exact matches for Traceloop attributes
-      self.assertTrue(self.llo_handler._is_llo_attribute("traceloop.entity.input"))
-      self.assertTrue(self.llo_handler._is_llo_attribute("traceloop.entity.output"))
+        """
+        Test _is_llo_attribute method with Traceloop patterns
+        """
+        # Test exact matches for Traceloop attributes
+        self.assertTrue(self.llo_handler._is_llo_attribute("traceloop.entity.input"))
+        self.assertTrue(self.llo_handler._is_llo_attribute("traceloop.entity.output"))
 
     def test_is_llo_attribute_openlit_match(self):
-      """
-      Test _is_llo_attribute method with OpenLit patterns
-      """
-      # Test exact matches for direct OpenLit attributes
-      self.assertTrue(self.llo_handler._is_llo_attribute("gen_ai.prompt"))
-      self.assertTrue(self.llo_handler._is_llo_attribute("gen_ai.completion"))
-      self.assertTrue(self.llo_handler._is_llo_attribute("gen_ai.content.revised_prompt"))
+        """
+        Test _is_llo_attribute method with OpenLit patterns
+        """
+        # Test exact matches for direct OpenLit attributes
+        self.assertTrue(self.llo_handler._is_llo_attribute("gen_ai.prompt"))
+        self.assertTrue(self.llo_handler._is_llo_attribute("gen_ai.completion"))
+        self.assertTrue(self.llo_handler._is_llo_attribute("gen_ai.content.revised_prompt"))
 
     def test_is_llo_attribute_openinference_match(self):
         """
@@ -296,105 +296,96 @@ def test_extract_traceloop_events(self):
         self.assertEqual(output_event.timestamp, 1234567899)  # end_time
 
     def test_extract_openlit_direct_prompt(self):
-      """
-      Test _extract_openlit_span_event_attributes with direct prompt attribute
-      """
-      attributes = {
-          "gen_ai.prompt": "user direct prompt",
-          "gen_ai.system": "openlit"
-      }
-
-      span = self._create_mock_span(attributes)
-
-      events = self.llo_handler._extract_openlit_span_event_attributes(span, attributes)
-
-      self.assertEqual(len(events), 1)
-      event = events[0]
-      self.assertEqual(event.name, "gen_ai.user.message")
-      self.assertEqual(event.body["content"], "user direct prompt")
-      self.assertEqual(event.body["role"], "user")
-      self.assertEqual(event.attributes["gen_ai.system"], "openlit")
-      self.assertEqual(event.attributes["original_attribute"], "gen_ai.prompt")
-      self.assertEqual(event.timestamp, 1234567890)  # start_time
+        """
+        Test _extract_openlit_span_event_attributes with direct prompt attribute
+        """
+        attributes = {"gen_ai.prompt": "user direct prompt", "gen_ai.system": "openlit"}
+
+        span = self._create_mock_span(attributes)
+
+        events = self.llo_handler._extract_openlit_span_event_attributes(span, attributes)
+
+        self.assertEqual(len(events), 1)
+        event = events[0]
+        self.assertEqual(event.name, "gen_ai.user.message")
+        self.assertEqual(event.body["content"], "user direct prompt")
+        self.assertEqual(event.body["role"], "user")
+        self.assertEqual(event.attributes["gen_ai.system"], "openlit")
+        self.assertEqual(event.attributes["original_attribute"], "gen_ai.prompt")
+        self.assertEqual(event.timestamp, 1234567890)  # start_time
 
     def test_extract_openlit_direct_completion(self):
-      """
-      Test _extract_openlit_span_event_attributes with direct completion attribute
-      """
-      attributes = {
-          "gen_ai.completion": "assistant direct completion",
-          "gen_ai.system": "openlit"
-      }
-
-      span = self._create_mock_span(attributes)
-      span.end_time = 1234567899
-
-      events = self.llo_handler._extract_openlit_span_event_attributes(span, attributes)
-
-      self.assertEqual(len(events), 1)
-      event = events[0]
-      self.assertEqual(event.name, "gen_ai.assistant.message")
-      self.assertEqual(event.body["content"], "assistant direct completion")
-      self.assertEqual(event.body["role"], "assistant")
-      self.assertEqual(event.attributes["gen_ai.system"], "openlit")
-      self.assertEqual(event.attributes["original_attribute"], "gen_ai.completion")
-      self.assertEqual(event.timestamp, 1234567899)  # end_time
+        """
+        Test _extract_openlit_span_event_attributes with direct completion attribute
+        """
+        attributes = {"gen_ai.completion": "assistant direct completion", "gen_ai.system": "openlit"}
+
+        span = self._create_mock_span(attributes)
+        span.end_time = 1234567899
+
+        events = self.llo_handler._extract_openlit_span_event_attributes(span, attributes)
+
+        self.assertEqual(len(events), 1)
+        event = events[0]
+        self.assertEqual(event.name, "gen_ai.assistant.message")
+        self.assertEqual(event.body["content"], "assistant direct completion")
+        self.assertEqual(event.body["role"], "assistant")
+        self.assertEqual(event.attributes["gen_ai.system"], "openlit")
+        self.assertEqual(event.attributes["original_attribute"], "gen_ai.completion")
+        self.assertEqual(event.timestamp, 1234567899)  # end_time
 
     def test_extract_openlit_all_attributes(self):
-      """
-      Test _extract_openlit_span_event_attributes with all OpenLit attributes
-      """
-      attributes = {
-          "gen_ai.prompt": "user prompt",
-          "gen_ai.completion": "assistant response",
-          "gen_ai.content.revised_prompt": "revised prompt",
-          "gen_ai.system": "langchain"
-      }
-
-      span = self._create_mock_span(attributes)
-      span.end_time = 1234567899
-
-      events = self.llo_handler._extract_openlit_span_event_attributes(span, attributes)
-
-      self.assertEqual(len(events), 3)
-
-      # Check that all events have the correct system
-      for event in events:
-          self.assertEqual(event.attributes["gen_ai.system"], "langchain")
-
-      # Check we have the expected event types
-      event_types = {event.name for event in events}
-      self.assertIn("gen_ai.user.message", event_types)
-      self.assertIn("gen_ai.assistant.message", event_types)
-      self.assertIn("gen_ai.system.message", event_types)
-
-      # Check original attributes
-      original_attrs = {event.attributes["original_attribute"] for event in events}
-      self.assertIn("gen_ai.prompt", original_attrs)
-      self.assertIn("gen_ai.completion", original_attrs)
-      self.assertIn("gen_ai.content.revised_prompt", original_attrs)
+        """
+        Test _extract_openlit_span_event_attributes with all OpenLit attributes
+        """
+        attributes = {
+            "gen_ai.prompt": "user prompt",
+            "gen_ai.completion": "assistant response",
+            "gen_ai.content.revised_prompt": "revised prompt",
+            "gen_ai.system": "langchain",
+        }
+
+        span = self._create_mock_span(attributes)
+        span.end_time = 1234567899
+
+        events = self.llo_handler._extract_openlit_span_event_attributes(span, attributes)
+
+        self.assertEqual(len(events), 3)
+
+        # Check that all events have the correct system
+        for event in events:
+            self.assertEqual(event.attributes["gen_ai.system"], "langchain")
+
+        # Check we have the expected event types
+        event_types = {event.name for event in events}
+        self.assertIn("gen_ai.user.message", event_types)
+        self.assertIn("gen_ai.assistant.message", event_types)
+        self.assertIn("gen_ai.system.message", event_types)
+
+        # Check original attributes
+        original_attrs = {event.attributes["original_attribute"] for event in events}
+        self.assertIn("gen_ai.prompt", original_attrs)
+        self.assertIn("gen_ai.completion", original_attrs)
+        self.assertIn("gen_ai.content.revised_prompt", original_attrs)
 
     def test_extract_openlit_revised_prompt(self):
-      """
-      Test _extract_openlit_span_event_attributes with revised prompt attribute
-      """
-      attributes = {
-          "gen_ai.content.revised_prompt": "revised system prompt",
-          "gen_ai.system": "openlit"
-      }
-
-      span = self._create_mock_span(attributes)
-
-      events = self.llo_handler._extract_openlit_span_event_attributes(span, attributes)
-
-      self.assertEqual(len(events), 1)
-      event = events[0]
-      self.assertEqual(event.name, "gen_ai.system.message")
-      self.assertEqual(event.body["content"], "revised system prompt")
-      self.assertEqual(event.body["role"], "system")
-      self.assertEqual(event.attributes["gen_ai.system"], "openlit")
-      self.assertEqual(event.attributes["original_attribute"], "gen_ai.content.revised_prompt")
-      self.assertEqual(event.timestamp, 1234567890)  # start_time
+        """
+        Test _extract_openlit_span_event_attributes with revised prompt attribute
+        """
+        attributes = {"gen_ai.content.revised_prompt": "revised system prompt", "gen_ai.system": "openlit"}
+
+        span = self._create_mock_span(attributes)
+
+        events = self.llo_handler._extract_openlit_span_event_attributes(span, attributes)
+
+        self.assertEqual(len(events), 1)
+        event = events[0]
+        self.assertEqual(event.name, "gen_ai.system.message")
+        self.assertEqual(event.body["content"], "revised system prompt")
+        self.assertEqual(event.body["role"], "system")
+        self.assertEqual(event.attributes["gen_ai.system"], "openlit")
+        self.assertEqual(event.attributes["original_attribute"], "gen_ai.content.revised_prompt")
+        self.assertEqual(event.timestamp, 1234567890)  # start_time
 
     def test_extract_openinference_direct_attributes(self):
         """