fix minor bug by making sure token counts are defined

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

Hide file tree

Showing 2 changed files with 1 addition and 2 deletions.
diff --git a/ddtrace/_trace/trace_handlers.py b/ddtrace/_trace/trace_handlers.py
@@ -642,7 +642,7 @@ def _on_botocore_bedrock_process_response(
     integration = ctx["bedrock_integration"]
     if metadata is not None:
         for k, v in metadata.items():
-            if k in ["usage.completion_tokens", "usage.prompt_tokens"]:
+            if k in ["usage.completion_tokens", "usage.prompt_tokens"] and v:
                 span.set_metric("bedrock.{}".format(k), int(v))
             else:
                 span.set_tag_str("bedrock.{}".format(k), str(v))

diff --git a/tests/contrib/botocore/test_bedrock.py b/tests/contrib/botocore/test_bedrock.py
@@ -373,7 +373,6 @@ 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"