only set token metrics if defined

DataDog · Feb 4, 2025 · 7baffdf · 7baffdf
1 parent 8416367
commit 7baffdf
Show file tree

Hide file tree

Showing 20 changed files with 42 additions and 40 deletions.
diff --git a/ddtrace/_trace/trace_handlers.py b/ddtrace/_trace/trace_handlers.py
@@ -595,8 +595,10 @@ def _on_botocore_patched_bedrock_api_call_success(ctx, reqid, latency, input_tok
     span = ctx.span
     span.set_tag_str("bedrock.response.id", reqid)
     span.set_tag_str("bedrock.response.duration", latency)
-    span.set_metric("bedrock.usage.prompt_tokens", int(input_token_count))
-    span.set_metric("bedrock.usage.completion_tokens", int(output_token_count))
+    if input_token_count:
+        span.set_metric("bedrock.usage.prompt_tokens", int(input_token_count))
+    if output_token_count:
+        span.set_metric("bedrock.usage.completion_tokens", int(output_token_count))
 
 
 def _propagate_context(ctx, headers):
@@ -641,8 +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"]:
-                v = int(v)
-            span.set_metric("bedrock.{}".format(k), v)
+                span.set_metric("bedrock.{}".format(k), int(v))
+            else:
+                span.set_tag_str("bedrock.{}".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/tests/contrib/botocore/test_bedrock.py b/tests/contrib/botocore/test_bedrock.py
@@ -373,6 +373,7 @@ def test_span_finishes_after_generator_exit(bedrock_client, request_vcr, mock_tr
                 if i >= 6:
                     raise GeneratorExit
                 i += 1
+    breakpoint()
     span = mock_tracer.pop_traces()[0][0]
     assert span is not None
     assert span.name == "bedrock-runtime.command"

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.usage.completion_tokens": 10,
+      "bedrock.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,8 +25,7 @@
       "_dd.tracer_kr": 1.0,
       "_sampling_priority_v1": 1,
       "bedrock.response.embedding_length": 1536,
-      "bedrock.usage.completion_tokens": "",
-      "bedrock.usage.prompt_tokens": "3",
+      "bedrock.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.usage.completion_tokens": 50,
+      "bedrock.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.usage.completion_tokens": 51,
+      "bedrock.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.usage.completion_tokens": 32,
+      "bedrock.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.usage.completion_tokens": 4,
+      "bedrock.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.usage.completion_tokens": 22,
+      "bedrock.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.usage.completion_tokens": 22,
+      "bedrock.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,8 +27,7 @@
       "_dd.tracer_kr": 1.0,
       "_sampling_priority_v1": 1,
       "bedrock.response.embedding_length": 1024,
-      "bedrock.usage.completion_tokens": "",
-      "bedrock.usage.prompt_tokens": "7",
+      "bedrock.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.usage.completion_tokens": 20,
+      "bedrock.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.usage.completion_tokens": 10,
+      "bedrock.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.usage.completion_tokens": 20,
+      "bedrock.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.usage.completion_tokens": 10,
+      "bedrock.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.usage.completion_tokens": 50,
+      "bedrock.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.usage.completion_tokens": 60,
+      "bedrock.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.usage.completion_tokens": 60,
+      "bedrock.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.usage.completion_tokens": 60,
+      "bedrock.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.usage.completion_tokens": 60,
+      "bedrock.usage.prompt_tokens": 10,
       "process_id": 42139
     },
     "duration": 3064000,