Add assertions in tests

Author: xrmx

instrumentation/opentelemetry-instrumentation-botocore/src/opentelemetry/instrumentation/botocore/extensions/bedrock.py

@@ -446,12 +446,31 @@ def _invoke_model_on_success(
             if original_body is not None:
                 original_body.close()
 
-    def _on_stream_error_callback(self, span: Span, exception):
+    def _on_stream_error_callback(
+        self,
+        span: Span,
+        exception,
+        instrumentor_context: _BotocoreInstrumentorContext,
+    ):
         span.set_status(Status(StatusCode.ERROR, str(exception)))
         if span.is_recording():
             span.set_attribute(ERROR_TYPE, type(exception).__qualname__)
         span.end()
 
+        metrics = instrumentor_context.metrics
+        metrics_attributes = {
+            **self._extract_metrics_attributes(),
+            ERROR_TYPE: type(exception).__qualname__,
+        }
+        if operation_duration_histogram := metrics.get(
+            GEN_AI_CLIENT_OPERATION_DURATION
+        ):
+            duration = max((default_timer() - self._operation_start), 0)
+            operation_duration_histogram.record(
+                duration,
+                attributes=metrics_attributes,
+            )
+
     def on_success(
         self,
         span: Span,
@@ -475,7 +494,9 @@ def stream_done_callback(response):
                     span.end()
 
                 def stream_error_callback(exception):
-                    self._on_stream_error_callback(span, exception)
+                    self._on_stream_error_callback(
+                        span, exception, instrumentor_context
+                    )
 
                 result["stream"] = ConverseStreamWrapper(
                     result["stream"],
@@ -513,7 +534,9 @@ def invoke_model_stream_done_callback(response):
                     span.end()
 
                 def invoke_model_stream_error_callback(exception):
-                    self._on_stream_error_callback(span, exception)
+                    self._on_stream_error_callback(
+                        span, exception, instrumentor_context
+                    )
 
                 result["body"] = InvokeModelWithResponseStreamWrapper(
                     result["body"],
@@ -716,3 +739,17 @@ def on_error(
 
         if not self.should_end_span_on_exit():
             span.end()
+
+        metrics = instrumentor_context.metrics
+        metrics_attributes = {
+            **self._extract_metrics_attributes(),
+            ERROR_TYPE: type(exception).__qualname__,
+        }
+        if operation_duration_histogram := metrics.get(
+            GEN_AI_CLIENT_OPERATION_DURATION
+        ):
+            duration = max((default_timer() - self._operation_start), 0)
+            operation_duration_histogram.record(
+                duration,
+                attributes=metrics_attributes,
+            )

instrumentation/opentelemetry-instrumentation-botocore/tests/bedrock_utils.py

@@ -19,13 +19,22 @@
 
 from botocore.response import StreamingBody
 
+from opentelemetry.instrumentation.botocore.extensions.bedrock import (
+    _GEN_AI_CLIENT_OPERATION_DURATION_BUCKETS,
+    _GEN_AI_CLIENT_TOKEN_USAGE_BUCKETS,
+)
+from opentelemetry.sdk.metrics._internal.point import ResourceMetrics
 from opentelemetry.sdk.trace import ReadableSpan
 from opentelemetry.semconv._incubating.attributes import (
     event_attributes as EventAttributes,
 )
 from opentelemetry.semconv._incubating.attributes import (
     gen_ai_attributes as GenAIAttributes,
 )
+from opentelemetry.semconv._incubating.metrics.gen_ai_metrics import (
+    GEN_AI_CLIENT_OPERATION_DURATION,
+    GEN_AI_CLIENT_TOKEN_USAGE,
+)
 
 
 # pylint: disable=too-many-branches, too-many-locals
@@ -259,3 +268,96 @@ def assert_message_in_logs(log, event_name, expected_content, parent_span):
             expected_content
         ), dict(log.log_record.body)
     assert_log_parent(log, parent_span)
+
+
+def assert_all_metric_attributes(data_point, operation_name, model):
+    assert GenAIAttributes.GEN_AI_OPERATION_NAME in data_point.attributes
+    assert (
+        data_point.attributes[GenAIAttributes.GEN_AI_OPERATION_NAME]
+        == operation_name
+    )
+    assert GenAIAttributes.GEN_AI_SYSTEM in data_point.attributes
+    assert (
+        data_point.attributes[GenAIAttributes.GEN_AI_SYSTEM]
+        == GenAIAttributes.GenAiSystemValues.AWS_BEDROCK.value
+    )
+    assert GenAIAttributes.GEN_AI_REQUEST_MODEL in data_point.attributes
+    assert data_point.attributes[GenAIAttributes.GEN_AI_REQUEST_MODEL] == model
+
+
+def assert_metrics(
+    resource_metrics: ResourceMetrics,
+    operation_name: str,
+    model: str,
+    input_tokens: float | None = None,
+    output_tokens: float | None = None,
+):
+    assert len(resource_metrics) == 1
+
+    metric_data = resource_metrics[0].scope_metrics[0].metrics
+    if input_tokens is not None or output_tokens is not None:
+        expected_metrics_data_len = 2
+    else:
+        expected_metrics_data_len = 1
+    assert len(metric_data) == expected_metrics_data_len
+
+    duration_metric = next(
+        (m for m in metric_data if m.name == GEN_AI_CLIENT_OPERATION_DURATION),
+        None,
+    )
+    assert duration_metric is not None
+
+    duration_point = duration_metric.data.data_points[0]
+    assert duration_point.sum > 0
+    assert_all_metric_attributes(duration_point, operation_name, model)
+    assert duration_point.explicit_bounds == tuple(
+        _GEN_AI_CLIENT_OPERATION_DURATION_BUCKETS
+    )
+
+    if input_tokens is not None:
+        token_usage_metric = next(
+            (m for m in metric_data if m.name == GEN_AI_CLIENT_TOKEN_USAGE),
+            None,
+        )
+        assert token_usage_metric is not None
+
+        input_token_usage = next(
+            (
+                d
+                for d in token_usage_metric.data.data_points
+                if d.attributes[GenAIAttributes.GEN_AI_TOKEN_TYPE]
+                == GenAIAttributes.GenAiTokenTypeValues.INPUT.value
+            ),
+            None,
+        )
+        assert input_token_usage is not None
+        assert input_token_usage.sum == input_tokens
+
+        assert input_token_usage.explicit_bounds == tuple(
+            _GEN_AI_CLIENT_TOKEN_USAGE_BUCKETS
+        )
+        assert_all_metric_attributes(input_token_usage, operation_name, model)
+
+    if output_tokens is not None:
+        token_usage_metric = next(
+            (m for m in metric_data if m.name == GEN_AI_CLIENT_TOKEN_USAGE),
+            None,
+        )
+        assert token_usage_metric is not None
+
+        output_token_usage = next(
+            (
+                d
+                for d in token_usage_metric.data.data_points
+                if d.attributes[GenAIAttributes.GEN_AI_TOKEN_TYPE]
+                == GenAIAttributes.GenAiTokenTypeValues.COMPLETION.value
+            ),
+            None,
+        )
+        assert output_token_usage is not None
+        assert output_token_usage.sum == output_tokens
+
+        assert output_token_usage.explicit_bounds == tuple(
+            _GEN_AI_CLIENT_TOKEN_USAGE_BUCKETS
+        )
+        assert_all_metric_attributes(output_token_usage, operation_name, model)

instrumentation/opentelemetry-instrumentation-botocore/tests/conftest.py

@@ -17,6 +17,12 @@
     InMemoryLogExporter,
     SimpleLogRecordProcessor,
 )
+from opentelemetry.sdk.metrics import (
+    MeterProvider,
+)
+from opentelemetry.sdk.metrics.export import (
+    InMemoryMetricReader,
+)
 from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace.export import SimpleSpanProcessor
 from opentelemetry.sdk.trace.export.in_memory_span_exporter import (
@@ -36,6 +42,12 @@ def fixture_log_exporter():
     yield exporter
 
 
+@pytest.fixture(scope="function", name="metric_reader")
+def fixture_metric_reader():
+    reader = InMemoryMetricReader()
+    yield reader
+
+
 @pytest.fixture(scope="function", name="tracer_provider")
 def fixture_tracer_provider(span_exporter):
     provider = TracerProvider()
@@ -52,6 +64,15 @@ def fixture_event_logger_provider(log_exporter):
     return event_logger_provider
 
 
+@pytest.fixture(scope="function", name="meter_provider")
+def fixture_meter_provider(metric_reader):
+    meter_provider = MeterProvider(
+        metric_readers=[metric_reader],
+    )
+
+    return meter_provider
+
+
 @pytest.fixture
 def bedrock_runtime_client():
     return boto3.client("bedrock-runtime")
@@ -81,7 +102,9 @@ def vcr_config():
 
 
 @pytest.fixture(scope="function")
-def instrument_no_content(tracer_provider, event_logger_provider):
+def instrument_no_content(
+    tracer_provider, event_logger_provider, meter_provider
+):
     os.environ.update(
         {OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT: "False"}
     )
@@ -90,6 +113,7 @@ def instrument_no_content(tracer_provider, event_logger_provider):
     instrumentor.instrument(
         tracer_provider=tracer_provider,
         event_logger_provider=event_logger_provider,
+        meter_provider=meter_provider,
     )
 
     yield instrumentor
@@ -98,14 +122,17 @@ def instrument_no_content(tracer_provider, event_logger_provider):
 
 
 @pytest.fixture(scope="function")
-def instrument_with_content(tracer_provider, event_logger_provider):
+def instrument_with_content(
+    tracer_provider, event_logger_provider, meter_provider
+):
     os.environ.update(
         {OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT: "True"}
     )
     instrumentor = BotocoreInstrumentor()
     instrumentor.instrument(
         tracer_provider=tracer_provider,
         event_logger_provider=event_logger_provider,
+        meter_provider=meter_provider,
     )
 
     yield instrumentor

instrumentation/opentelemetry-instrumentation-botocore/tests/test_botocore_bedrock.py

@@ -34,6 +34,7 @@
     assert_completion_attributes_from_streaming_body,
     assert_converse_completion_attributes,
     assert_message_in_logs,
+    assert_metrics,
     assert_stream_completion_attributes,
 )
 
@@ -51,6 +52,7 @@ def filter_message_keys(message, keys):
 def test_converse_with_content(
     span_exporter,
     log_exporter,
+    metric_reader,
     bedrock_runtime_client,
     instrument_with_content,
 ):
@@ -95,6 +97,13 @@ def test_converse_with_content(
     }
     assert_message_in_logs(logs[1], "gen_ai.choice", choice_body, span)
 
+    input_tokens = response["usage"]["inputTokens"]
+    output_tokens = response["usage"]["outputTokens"]
+    metrics = metric_reader.get_metrics_data().resource_metrics
+    assert_metrics(
+        metrics, "chat", llm_model_value, input_tokens, output_tokens
+    )
+
 
 @pytest.mark.skipif(
     BOTO3_VERSION < (1, 35, 56), reason="Converse API not available"
@@ -103,6 +112,7 @@ def test_converse_with_content(
 def test_converse_with_content_different_events(
     span_exporter,
     log_exporter,
+    metric_reader,
     bedrock_runtime_client,
     instrument_with_content,
 ):
@@ -150,6 +160,13 @@ def test_converse_with_content_different_events(
     }
     assert_message_in_logs(logs[4], "gen_ai.choice", choice_body, span)
 
+    input_tokens = response["usage"]["inputTokens"]
+    output_tokens = response["usage"]["outputTokens"]
+    metrics = metric_reader.get_metrics_data().resource_metrics
+    assert_metrics(
+        metrics, "chat", llm_model_value, input_tokens, output_tokens
+    )
+
 
 def converse_tool_call(
     span_exporter, log_exporter, bedrock_runtime_client, expect_content
@@ -452,6 +469,7 @@ def test_converse_tool_call_no_content(
 def test_converse_with_invalid_model(
     span_exporter,
     log_exporter,
+    metric_reader,
     bedrock_runtime_client,
     instrument_with_content,
 ):
@@ -479,6 +497,9 @@ def test_converse_with_invalid_model(
     user_content = filter_message_keys(messages[0], ["content"])
     assert_message_in_logs(logs[0], "gen_ai.user.message", user_content, span)
 
+    metrics = metric_reader.get_metrics_data().resource_metrics
+    assert_metrics(metrics, "chat", llm_model_value)
+
 
 @pytest.mark.skipif(
     BOTO3_VERSION < (1, 35, 56), reason="ConverseStream API not available"
@@ -487,6 +508,7 @@ def test_converse_with_invalid_model(
 def test_converse_stream_with_content(
     span_exporter,
     log_exporter,
+    metric_reader,
     bedrock_runtime_client,
     instrument_with_content,
 ):
@@ -553,6 +575,11 @@ def test_converse_stream_with_content(
     }
     assert_message_in_logs(logs[1], "gen_ai.choice", choice_body, span)
 
+    metrics = metric_reader.get_metrics_data().resource_metrics
+    assert_metrics(
+        metrics, "chat", llm_model_value, input_tokens, output_tokens
+    )
+
 
 @pytest.mark.skipif(
     BOTO3_VERSION < (1, 35, 56), reason="ConverseStream API not available"
@@ -561,6 +588,7 @@ def test_converse_stream_with_content(
 def test_converse_stream_with_content_different_events(
     span_exporter,
     log_exporter,
+    metric_reader,
     bedrock_runtime_client,
     instrument_with_content,
 ):
@@ -614,6 +642,9 @@ def test_converse_stream_with_content_different_events(
     }
     assert_message_in_logs(logs[4], "gen_ai.choice", choice_body, span)
 
+    metrics = metric_reader.get_metrics_data().resource_metrics
+    assert_metrics(metrics, "chat", llm_model_value, mock.ANY, mock.ANY)
+
 
 def _rebuild_stream_message(response):
     message = {"content": []}
@@ -986,6 +1017,7 @@ def test_converse_stream_no_content_tool_call(
 def test_converse_stream_handles_event_stream_error(
     span_exporter,
     log_exporter,
+    metric_reader,
     bedrock_runtime_client,
     instrument_with_content,
 ):
@@ -1039,6 +1071,9 @@ def test_converse_stream_handles_event_stream_error(
     user_content = filter_message_keys(messages[0], ["content"])
     assert_message_in_logs(logs[0], "gen_ai.user.message", user_content, span)
 
+    metrics = metric_reader.get_metrics_data().resource_metrics
+    assert_metrics(metrics, "chat", llm_model_value)
+
 
 @pytest.mark.skipif(
     BOTO3_VERSION < (1, 35, 56), reason="ConverseStream API not available"