update token metric name and tests accordingly

DataDog · Feb 5, 2025 · 8931ffa · 8931ffa
1 parent 42289dd
commit 8931ffa
Show file tree

Hide file tree

Showing 21 changed files with 40 additions and 53 deletions.
diff --git a/ddtrace/_trace/trace_handlers.py b/ddtrace/_trace/trace_handlers.py
@@ -596,9 +596,9 @@ def _on_botocore_patched_bedrock_api_call_success(ctx, reqid, latency, input_tok
     span.set_tag_str("bedrock.response.id", reqid)
     span.set_tag_str("bedrock.response.duration", latency)
     if input_token_count:
-        span.set_metric("bedrock.usage.prompt_tokens", int(input_token_count))
+        span.set_metric("bedrock.response.usage.prompt_tokens", int(input_token_count))
     if output_token_count:
-        span.set_metric("bedrock.usage.completion_tokens", int(output_token_count))
+        span.set_metric("bedrock.response.usage.completion_tokens", int(output_token_count))
 
 
 def _propagate_context(ctx, headers):
@@ -643,9 +643,9 @@ def _on_botocore_bedrock_process_response(
     if metadata is not None:
         for k, v in metadata.items():
             if k in ["usage.completion_tokens", "usage.prompt_tokens"] and v:
-                span.set_metric("bedrock.{}".format(k), int(v))
+                span.set_metric("bedrock.response{}".format(k), int(v))
             else:
-                span.set_tag_str("bedrock.{}".format(k), str(v))
+                span.set_tag_str("bedrock.response{}".format(k), str(v))
     if "embed" in model_name:
         span.set_metric("bedrock.response.embedding_length", len(formatted_response["text"][0]))
         span.finish()

diff --git a/ddtrace/llmobs/_integrations/utils.py b/ddtrace/llmobs/_integrations/utils.py
@@ -121,19 +121,6 @@ def extract_message_from_part_google(part, role=None):
 def get_llmobs_metrics_tags(integration_name, span):
     usage = {}
 
-    # bedrock integration tags usage under meta instead of metrics
-    if integration_name == "bedrock":
-        input_tokens = int(span.get_tag("bedrock.usage.prompt_tokens") or 0)
-        output_tokens = int(span.get_tag("bedrock.usage.completion_tokens") or 0)
-        total_tokens = input_tokens + output_tokens
-        if input_tokens:
-            usage[INPUT_TOKENS_METRIC_KEY] = input_tokens
-        if output_tokens:
-            usage[OUTPUT_TOKENS_METRIC_KEY] = output_tokens
-        if total_tokens:
-            usage[TOTAL_TOKENS_METRIC_KEY] = total_tokens
-        return usage
-
     # check for both prompt / completion or input / output tokens
     input_tokens = span.get_metric("%s.response.usage.prompt_tokens" % integration_name) or span.get_metric(
         "%s.response.usage.input_tokens" % integration_name

diff --git a/tests/contrib/botocore/test_bedrock_llmobs.py b/tests/contrib/botocore/test_bedrock_llmobs.py
@@ -80,8 +80,8 @@ def mock_llmobs_span_writer():
 class TestLLMObsBedrock:
     @staticmethod
     def expected_llmobs_span_event(span, n_output, message=False):
-        prompt_tokens = int(span.get_metric("bedrock.usage.prompt_tokens"))
-        completion_tokens = int(span.get_metric("bedrock.usage.completion_tokens"))
+        prompt_tokens = int(span.get_metric("bedrock.response.usage.prompt_tokens"))
+        completion_tokens = int(span.get_metric("bedrock.response.usage.completion_tokens"))
         expected_parameters = {"temperature": float(span.get_tag("bedrock.request.temperature"))}
         if span.get_tag("bedrock.request.max_tokens"):
             expected_parameters["max_tokens"] = int(span.get_tag("bedrock.request.max_tokens"))

diff --git a/tests/snapshots/tests.contrib.botocore.test_bedrock.test_ai21_invoke.json b/tests/snapshots/tests.contrib.botocore.test_bedrock.test_ai21_invoke.json
@@ -30,8 +30,8 @@
       "_dd.top_level": 1,
       "_dd.tracer_kr": 1.0,
       "_sampling_priority_v1": 1,
-      "bedrock.usage.completion_tokens": 10,
-      "bedrock.usage.prompt_tokens": 10,
+      "bedrock.response.usage.completion_tokens": 10,
+      "bedrock.response.usage.prompt_tokens": 10,
       "process_id": 7458
     },
     "duration": 2112000,

diff --git a/tests/snapshots/tests.contrib.botocore.test_bedrock.test_amazon_embedding.json b/tests/snapshots/tests.contrib.botocore.test_bedrock.test_amazon_embedding.json
@@ -25,7 +25,7 @@
       "_dd.tracer_kr": 1.0,
       "_sampling_priority_v1": 1,
       "bedrock.response.embedding_length": 1536,
-      "bedrock.usage.prompt_tokens": 3,
+      "bedrock.response.usage.prompt_tokens": 3,
       "process_id": 60939
     },
     "duration": 6739000,

diff --git a/tests/snapshots/tests.contrib.botocore.test_bedrock.test_amazon_invoke.json b/tests/snapshots/tests.contrib.botocore.test_bedrock.test_amazon_invoke.json
@@ -30,8 +30,8 @@
       "_dd.top_level": 1,
       "_dd.tracer_kr": 1.0,
       "_sampling_priority_v1": 1,
-      "bedrock.usage.completion_tokens": 50,
-      "bedrock.usage.prompt_tokens": 18,
+      "bedrock.response.usage.completion_tokens": 50,
+      "bedrock.response.usage.prompt_tokens": 18,
       "process_id": 14088
     },
     "duration": 2147082000,

diff --git a/tests/snapshots/tests.contrib.botocore.test_bedrock.test_amazon_invoke_stream.json b/tests/snapshots/tests.contrib.botocore.test_bedrock.test_amazon_invoke_stream.json
@@ -30,8 +30,8 @@
       "_dd.top_level": 1,
       "_dd.tracer_kr": 1.0,
       "_sampling_priority_v1": 1,
-      "bedrock.usage.completion_tokens": 51,
-      "bedrock.usage.prompt_tokens": 18,
+      "bedrock.response.usage.completion_tokens": 51,
+      "bedrock.response.usage.prompt_tokens": 18,
       "process_id": 14088
     },
     "duration": 2185710000,

diff --git a/tests/snapshots/tests.contrib.botocore.test_bedrock.test_anthropic_invoke.json b/tests/snapshots/tests.contrib.botocore.test_bedrock.test_anthropic_invoke.json
@@ -31,8 +31,8 @@
       "_dd.top_level": 1,
       "_dd.tracer_kr": 1.0,
       "_sampling_priority_v1": 1,
-      "bedrock.usage.completion_tokens": 32,
-      "bedrock.usage.prompt_tokens": 23,
+      "bedrock.response.usage.completion_tokens": 32,
+      "bedrock.response.usage.prompt_tokens": 23,
       "process_id": 7272
     },
     "duration": 2434000,

diff --git a/tests/snapshots/tests.contrib.botocore.test_bedrock.test_anthropic_invoke_stream.json b/tests/snapshots/tests.contrib.botocore.test_bedrock.test_anthropic_invoke_stream.json
@@ -31,8 +31,8 @@
       "_dd.top_level": 1,
       "_dd.tracer_kr": 1.0,
       "_sampling_priority_v1": 1,
-      "bedrock.usage.completion_tokens": 4,
-      "bedrock.usage.prompt_tokens": 25,
+      "bedrock.response.usage.completion_tokens": 4,
+      "bedrock.response.usage.prompt_tokens": 25,
       "process_id": 13707
     },
     "duration": 624710000,

diff --git a/tests/snapshots/tests.contrib.botocore.test_bedrock.test_anthropic_message_invoke.json b/tests/snapshots/tests.contrib.botocore.test_bedrock.test_anthropic_message_invoke.json
@@ -31,8 +31,8 @@
       "_dd.top_level": 1,
       "_dd.tracer_kr": 1.0,
       "_sampling_priority_v1": 1,
-      "bedrock.usage.completion_tokens": 22,
-      "bedrock.usage.prompt_tokens": 21,
+      "bedrock.response.usage.completion_tokens": 22,
+      "bedrock.response.usage.prompt_tokens": 21,
       "process_id": 40705
     },
     "duration": 2160000,

diff --git a/...s/snapshots/tests.contrib.botocore.test_bedrock.test_anthropic_message_invoke_stream.json b/...s/snapshots/tests.contrib.botocore.test_bedrock.test_anthropic_message_invoke_stream.json
@@ -31,8 +31,8 @@
       "_dd.top_level": 1,
       "_dd.tracer_kr": 1.0,
       "_sampling_priority_v1": 1,
-      "bedrock.usage.completion_tokens": 22,
-      "bedrock.usage.prompt_tokens": 21,
+      "bedrock.response.usage.completion_tokens": 22,
+      "bedrock.response.usage.prompt_tokens": 21,
       "process_id": 40896
     },
     "duration": 2950000,

diff --git a/tests/snapshots/tests.contrib.botocore.test_bedrock.test_cohere_embedding.json b/tests/snapshots/tests.contrib.botocore.test_bedrock.test_cohere_embedding.json
@@ -27,7 +27,7 @@
       "_dd.tracer_kr": 1.0,
       "_sampling_priority_v1": 1,
       "bedrock.response.embedding_length": 1024,
-      "bedrock.usage.prompt_tokens": 7,
+      "bedrock.response.usage.prompt_tokens": 7,
       "process_id": 61336
     },
     "duration": 630192000,

diff --git a/tests/snapshots/tests.contrib.botocore.test_bedrock.test_cohere_invoke_multi_output.json b/tests/snapshots/tests.contrib.botocore.test_bedrock.test_cohere_invoke_multi_output.json
@@ -37,8 +37,8 @@
       "_dd.top_level": 1,
       "_dd.tracer_kr": 1.0,
       "_sampling_priority_v1": 1,
-      "bedrock.usage.completion_tokens": 20,
-      "bedrock.usage.prompt_tokens": 40,
+      "bedrock.response.usage.completion_tokens": 20,
+      "bedrock.response.usage.prompt_tokens": 40,
       "process_id": 3568
     },
     "duration": 810213000,

diff --git a/tests/snapshots/tests.contrib.botocore.test_bedrock.test_cohere_invoke_single_output.json b/tests/snapshots/tests.contrib.botocore.test_bedrock.test_cohere_invoke_single_output.json
@@ -34,8 +34,8 @@
       "_dd.top_level": 1,
       "_dd.tracer_kr": 1.0,
       "_sampling_priority_v1": 1,
-      "bedrock.usage.completion_tokens": 10,
-      "bedrock.usage.prompt_tokens": 20,
+      "bedrock.response.usage.completion_tokens": 10,
+      "bedrock.response.usage.prompt_tokens": 20,
       "process_id": 13549
     },
     "duration": 659370000,

diff --git a/...snapshots/tests.contrib.botocore.test_bedrock.test_cohere_invoke_stream_multi_output.json b/...snapshots/tests.contrib.botocore.test_bedrock.test_cohere_invoke_stream_multi_output.json
@@ -35,8 +35,8 @@
       "_dd.top_level": 1,
       "_dd.tracer_kr": 1.0,
       "_sampling_priority_v1": 1,
-      "bedrock.usage.completion_tokens": 20,
-      "bedrock.usage.prompt_tokens": 40,
+      "bedrock.response.usage.completion_tokens": 20,
+      "bedrock.response.usage.prompt_tokens": 40,
       "process_id": 21816
     },
     "duration": 980170000,

diff --git a/...napshots/tests.contrib.botocore.test_bedrock.test_cohere_invoke_stream_single_output.json b/...napshots/tests.contrib.botocore.test_bedrock.test_cohere_invoke_stream_single_output.json
@@ -33,8 +33,8 @@
       "_dd.top_level": 1,
       "_dd.tracer_kr": 1.0,
       "_sampling_priority_v1": 1,
-      "bedrock.usage.completion_tokens": 10,
-      "bedrock.usage.prompt_tokens": 20,
+      "bedrock.response.usage.completion_tokens": 10,
+      "bedrock.response.usage.prompt_tokens": 20,
       "process_id": 21816
     },
     "duration": 630536000,

diff --git a/...apshots/tests.contrib.botocore.test_bedrock.test_invoke_model_using_aws_arn_model_id.json b/...apshots/tests.contrib.botocore.test_bedrock.test_invoke_model_using_aws_arn_model_id.json
@@ -30,8 +30,8 @@
       "_dd.top_level": 1,
       "_dd.tracer_kr": 1.0,
       "_sampling_priority_v1": 1,
-      "bedrock.usage.completion_tokens": 50,
-      "bedrock.usage.prompt_tokens": 18,
+      "bedrock.response.usage.completion_tokens": 50,
+      "bedrock.response.usage.prompt_tokens": 18,
       "process_id": 96028
     },
     "duration": 2318000,

diff --git a/tests/snapshots/tests.contrib.botocore.test_bedrock.test_meta_invoke.json b/tests/snapshots/tests.contrib.botocore.test_bedrock.test_meta_invoke.json
@@ -29,8 +29,8 @@
       "_dd.top_level": 1,
       "_dd.tracer_kr": 1.0,
       "_sampling_priority_v1": 1,
-      "bedrock.usage.completion_tokens": 60,
-      "bedrock.usage.prompt_tokens": 10,
+      "bedrock.response.usage.completion_tokens": 60,
+      "bedrock.response.usage.prompt_tokens": 10,
       "process_id": 10703
     },
     "duration": 2120703000,

diff --git a/tests/snapshots/tests.contrib.botocore.test_bedrock.test_meta_invoke_stream.json b/tests/snapshots/tests.contrib.botocore.test_bedrock.test_meta_invoke_stream.json
@@ -29,8 +29,8 @@
       "_dd.top_level": 1,
       "_dd.tracer_kr": 1.0,
       "_sampling_priority_v1": 1,
-      "bedrock.usage.completion_tokens": 60,
-      "bedrock.usage.prompt_tokens": 10,
+      "bedrock.response.usage.completion_tokens": 60,
+      "bedrock.response.usage.prompt_tokens": 10,
       "process_id": 2664
     },
     "duration": 3795000,

diff --git a/tests/snapshots/tests.contrib.botocore.test_bedrock.test_read_error.json b/tests/snapshots/tests.contrib.botocore.test_bedrock.test_read_error.json
@@ -30,8 +30,8 @@
       "_dd.top_level": 1,
       "_dd.tracer_kr": 1.0,
       "_sampling_priority_v1": 1,
-      "bedrock.usage.completion_tokens": 60,
-      "bedrock.usage.prompt_tokens": 10,
+      "bedrock.response.usage.completion_tokens": 60,
+      "bedrock.response.usage.prompt_tokens": 10,
       "process_id": 42139
     },
     "duration": 2505000,

diff --git a/tests/snapshots/tests.contrib.botocore.test_bedrock.test_readlines_error.json b/tests/snapshots/tests.contrib.botocore.test_bedrock.test_readlines_error.json
@@ -30,8 +30,8 @@
       "_dd.top_level": 1,
       "_dd.tracer_kr": 1.0,
       "_sampling_priority_v1": 1,
-      "bedrock.usage.completion_tokens": 60,
-      "bedrock.usage.prompt_tokens": 10,
+      "bedrock.response.usage.completion_tokens": 60,
+      "bedrock.response.usage.prompt_tokens": 10,
       "process_id": 42139
     },
     "duration": 3064000,