chore(llmobs): remove llm io from bedrock spans (#14057)

Remove potentially sensitive i/o data from apm spans. This way, prompt
and completion data will only appear on the llm obs spans, which
are/will be subject to data access controls.

Mostly, this just removes io tag sets. A few things (mostly metrics)
have llmobs tags dependent on span tags, so there is a bit more
refactoring there.

Let me know if I removed anything that should really stay, or if I
missed something that should be restricted.

## Checklist
- [x] PR author has checked that all the criteria below are met
- The PR description includes an overview of the change
- The PR description articulates the motivation for the change
- The change includes tests OR the PR description describes a testing
strategy
- The PR description notes risks associated with the change, if any
- Newly-added code is easy to change
- The change follows the [library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html)
- The change includes or references documentation updates if necessary
- Backport labels are set (if
[applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting))

## Reviewer Checklist
- [x] Reviewer has checked that all the criteria below are met 
- Title is accurate
- All changes are related to the pull request's stated goal
- Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes
- Testing strategy adequately addresses listed risks
- Newly-added code is easy to change
- Release note makes sense to a user of the library
- If necessary, author has acknowledged and discussed the performance
implications of this PR as reported in the benchmarks PR comment
- Backport labels are set in a manner that is consistent with the
[release branch maintenance
policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)

Author: jsimpher

ddtrace/_trace/trace_handlers.py

@@ -661,13 +661,9 @@ def _on_botocore_patched_bedrock_api_call_started(ctx, request_params):
 
     span.set_tag_str("bedrock.request.model_provider", ctx["model_provider"])
     span.set_tag_str("bedrock.request.model", ctx["model_name"])
-    for k, v in request_params.items():
-        if k == "prompt":
-            if integration.is_pc_sampled_span(span):
-                v = integration.trunc(str(v))
-        span.set_tag_str("bedrock.request.{}".format(k), str(v))
-        if k == "n":
-            ctx.set_item("num_generations", str(v))
+
+    if "n" in request_params:
+        ctx.set_item("num_generations", str(request_params["n"]))
 
 
 def _on_botocore_patched_bedrock_api_call_exception(ctx, exc_info):
@@ -680,16 +676,6 @@ def _on_botocore_patched_bedrock_api_call_exception(ctx, exc_info):
     span.finish()
 
 
-def _on_botocore_patched_bedrock_api_call_success(ctx, reqid, latency, input_token_count, output_token_count):
-    span = ctx.span
-    span.set_tag_str("bedrock.response.id", reqid)
-    span.set_tag_str("bedrock.response.duration", latency)
-    if input_token_count:
-        span.set_metric("bedrock.response.usage.prompt_tokens", int(input_token_count))
-    if output_token_count:
-        span.set_metric("bedrock.response.usage.completion_tokens", int(output_token_count))
-
-
 def _propagate_context(ctx, headers):
     distributed_tracing_enabled = ctx["integration_config"].distributed_tracing_enabled
     span = ctx.span
@@ -731,38 +717,13 @@ def _on_botocore_bedrock_process_response_converse(
 def _on_botocore_bedrock_process_response(
     ctx: core.ExecutionContext,
     formatted_response: Dict[str, Any],
-    metadata: Dict[str, Any],
-    body: Dict[str, List[Dict]],
-    should_set_choice_ids: bool,
 ) -> None:
-    text = formatted_response["text"]
-    span = ctx.span
-    model_name = ctx["model_name"]
-    if should_set_choice_ids:
-        for i in range(len(text)):
-            span.set_tag_str("bedrock.response.choices.{}.id".format(i), str(body["generations"][i]["id"]))
-    integration = ctx["bedrock_integration"]
-    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.response.{}".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()
-        return
-    for i in range(len(formatted_response["text"])):
-        if integration.is_pc_sampled_span(span):
-            span.set_tag_str(
-                "bedrock.response.choices.{}.text".format(i),
-                integration.trunc(str(formatted_response["text"][i])),
-            )
-        span.set_tag_str(
-            "bedrock.response.choices.{}.finish_reason".format(i), str(formatted_response["finish_reason"][i])
-        )
-    integration.llmobs_set_tags(span, args=[ctx], kwargs={}, response=formatted_response)
-    span.finish()
+    with ctx.span as span:
+        model_name = ctx["model_name"]
+        integration = ctx["bedrock_integration"]
+        if "embed" in model_name:
+            return
+        integration.llmobs_set_tags(span, args=[ctx], kwargs={}, response=formatted_response)
 
 
 def _on_botocore_sqs_recvmessage_post(
@@ -931,7 +892,6 @@ def listen():
     core.on("botocore.client_context.update_messages", _on_botocore_update_messages)
     core.on("botocore.patched_bedrock_api_call.started", _on_botocore_patched_bedrock_api_call_started)
     core.on("botocore.patched_bedrock_api_call.exception", _on_botocore_patched_bedrock_api_call_exception)
-    core.on("botocore.patched_bedrock_api_call.success", _on_botocore_patched_bedrock_api_call_success)
     core.on("botocore.bedrock.process_response", _on_botocore_bedrock_process_response)
     core.on("botocore.bedrock.process_response_converse", _on_botocore_bedrock_process_response_converse)
     core.on("botocore.sqs.ReceiveMessage.post", _on_botocore_sqs_recvmessage_post)

ddtrace/contrib/internal/botocore/services/bedrock.py

@@ -48,12 +48,9 @@ def read(self, amt=None):
             self._body.append(json.loads(body))
             if self.__wrapped__.tell() == int(self.__wrapped__._content_length):
                 formatted_response = _extract_text_and_response_reason(self._execution_ctx, self._body[0])
-                model_provider = self._execution_ctx["model_provider"]
-                model_name = self._execution_ctx["model_name"]
-                should_set_choice_ids = model_provider == _COHERE and "embed" not in model_name
                 core.dispatch(
                     "botocore.bedrock.process_response",
-                    [self._execution_ctx, formatted_response, None, self._body[0], should_set_choice_ids],
+                    [self._execution_ctx, formatted_response],
                 )
             return body
         except Exception:
@@ -67,12 +64,9 @@ def readlines(self):
             for line in lines:
                 self._body.append(json.loads(line))
             formatted_response = _extract_text_and_response_reason(self._execution_ctx, self._body[0])
-            model_provider = self._execution_ctx["model_provider"]
-            model_name = self._execution_ctx["model_name"]
-            should_set_choice_ids = model_provider == _COHERE and "embed" not in model_name
             core.dispatch(
                 "botocore.bedrock.process_response",
-                [self._execution_ctx, formatted_response, None, self._body[0], should_set_choice_ids],
+                [self._execution_ctx, formatted_response],
             )
             return lines
         except Exception:
@@ -93,16 +87,10 @@ def __iter__(self):
         finally:
             if exception_raised:
                 return
-            metadata = _extract_streamed_response_metadata(self._execution_ctx, self._body)
             formatted_response = _extract_streamed_response(self._execution_ctx, self._body)
-            model_provider = self._execution_ctx["model_provider"]
-            model_name = self._execution_ctx["model_name"]
-            should_set_choice_ids = (
-                model_provider == _COHERE and "is_finished" not in self._body[0] and "embed" not in model_name
-            )
             core.dispatch(
                 "botocore.bedrock.process_response",
-                [self._execution_ctx, formatted_response, metadata, self._body, should_set_choice_ids],
+                [self._execution_ctx, formatted_response],
             )
 
 
@@ -443,18 +431,6 @@ def handle_bedrock_response(
         safe_token_count(cache_write_tokens),
     )
 
-    # for both converse & invoke, dispatch success event to store basic metrics
-    core.dispatch(
-        "botocore.patched_bedrock_api_call.success",
-        [
-            ctx,
-            str(metadata.get("RequestId", "")),
-            request_latency,
-            str(input_tokens),
-            str(output_tokens),
-        ],
-    )
-
     if ctx["resource"] == "Converse":
         core.dispatch("botocore.bedrock.process_response_converse", [ctx, result])
         return result

releasenotes/notes/remove-io-data-from-apm-span-bedrock-integration-4269a599d64a5756.yaml

@@ -0,0 +1,3 @@
+other:
+  - |
+    botocore: Removes I/O tags, including llm prompts, completions, and token usage, from APM spans for Bedrock LLM requests and responses, which is duplicated in LLM Observability.

tests/contrib/botocore/test_bedrock.py

@@ -252,4 +252,3 @@ def test_span_finishes_after_generator_exit(bedrock_client, request_vcr, mock_tr
     assert span is not None
     assert span.name == "bedrock-runtime.command"
     assert span.resource == "InvokeModelWithResponseStream"
-    assert span.get_tag("bedrock.response.choices.0.text").startswith("Hobb")

tests/contrib/botocore/test_bedrock_llmobs.py

@@ -25,41 +25,41 @@
 )
 class TestLLMObsBedrock:
     @staticmethod
-    def expected_llmobs_span_event(span, n_output, message=False):
-        prompt_tokens = span.get_metric("bedrock.response.usage.prompt_tokens")
-        completion_tokens = span.get_metric("bedrock.response.usage.completion_tokens")
-        token_metrics = {}
-        if prompt_tokens is not None:
-            token_metrics["input_tokens"] = prompt_tokens
-        if completion_tokens is not None:
-            token_metrics["output_tokens"] = completion_tokens
-        if prompt_tokens is not None and completion_tokens is not None:
-            token_metrics["total_tokens"] = prompt_tokens + completion_tokens
-
-        if span.get_tag("bedrock.request.temperature"):
-            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"))
-
+    def expected_llmobs_span_event(span, n_output, message=False, metadata=None, token_metrics=None):
         expected_input = [{"content": mock.ANY}]
         if message:
             expected_input = [{"content": mock.ANY, "role": "user"}]
-        return _expected_llmobs_llm_span_event(
+
+        # Use empty dicts as defaults for _expected_llmobs_llm_span_event to avoid None issues
+        expected_parameters = metadata if metadata is not None else {}
+        expected_token_metrics = token_metrics if token_metrics is not None else None
+
+        expected_event = _expected_llmobs_llm_span_event(
             span,
             model_name=span.get_tag("bedrock.request.model"),
             model_provider=span.get_tag("bedrock.request.model_provider"),
             input_messages=expected_input,
             output_messages=[{"content": mock.ANY} for _ in range(n_output)],
             metadata=expected_parameters,
-            token_metrics=token_metrics,
+            token_metrics=expected_token_metrics,
             tags={"service": "aws.bedrock-runtime", "ml_app": "<ml-app-name>"},
         )
 
+        # If parameters were not explicitly provided, use mock.ANY to match anything
+        if metadata is None:
+            expected_event["meta"]["metadata"] = mock.ANY
+        if token_metrics is None:
+            expected_event["metrics"] = mock.ANY
+
+        return expected_event
+
     @classmethod
     def _test_llmobs_invoke(cls, provider, bedrock_client, mock_tracer, llmobs_events, cassette_name=None, n_output=1):
         if cassette_name is None:
             cassette_name = "%s_invoke.yaml" % provider
         body = _REQUEST_BODIES[provider]
+        expected_metadata = None
+
         if provider == "cohere":
             body = {
                 "prompt": "\n\nHuman: %s\n\nAssistant: Can you explain what a LLM chain is?",
@@ -71,6 +71,8 @@ def _test_llmobs_invoke(cls, provider, bedrock_client, mock_tracer, llmobs_event
                 "stream": False,
                 "num_generations": n_output,
             }
+            expected_metadata = {"temperature": 0.9, "max_tokens": 10}
+
         with get_request_vcr().use_cassette(cassette_name):
             body, model = json.dumps(body), _MODELS[provider]
             if provider == "anthropic_message":
@@ -82,7 +84,9 @@ def _test_llmobs_invoke(cls, provider, bedrock_client, mock_tracer, llmobs_event
         span = mock_tracer.pop_traces()[0][0]
 
         assert len(llmobs_events) == 1
-        assert llmobs_events[0] == cls.expected_llmobs_span_event(span, n_output, message="message" in provider)
+        assert llmobs_events[0] == cls.expected_llmobs_span_event(
+            span, n_output, message="message" in provider, metadata=expected_metadata
+        )
         LLMObs.disable()
 
     @classmethod
@@ -92,6 +96,8 @@ def _test_llmobs_invoke_stream(
         if cassette_name is None:
             cassette_name = "%s_invoke_stream.yaml" % provider
         body = _REQUEST_BODIES[provider]
+        expected_metadata = None
+
         if provider == "cohere":
             body = {
                 "prompt": "\n\nHuman: %s\n\nAssistant: Can you explain what a LLM chain is?",
@@ -103,6 +109,8 @@ def _test_llmobs_invoke_stream(
                 "stream": True,
                 "num_generations": n_output,
             }
+            expected_metadata = {"temperature": 0.9, "max_tokens": 10}
+
         with get_request_vcr().use_cassette(cassette_name):
             body, model = json.dumps(body), _MODELS[provider]
             response = bedrock_client.invoke_model_with_response_stream(body=body, modelId=model)
@@ -111,7 +119,9 @@ def _test_llmobs_invoke_stream(
         span = mock_tracer.pop_traces()[0][0]
 
         assert len(llmobs_events) == 1
-        assert llmobs_events[0] == cls.expected_llmobs_span_event(span, n_output, message="message" in provider)
+        assert llmobs_events[0] == cls.expected_llmobs_span_event(
+            span, n_output, message="message" in provider, metadata=expected_metadata
+        )
 
     def test_llmobs_ai21_invoke(self, ddtrace_global_config, bedrock_client, mock_tracer, llmobs_events):
         self._test_llmobs_invoke("ai21", bedrock_client, mock_tracer, llmobs_events)
@@ -216,16 +226,15 @@ def test_llmobs_error(self, ddtrace_global_config, bedrock_client, mock_tracer,
                 json.loads(response.get("body").read())
         span = mock_tracer.pop_traces()[0][0]
 
+        metadata = mock.ANY
+
         assert len(llmobs_events) == 1
         assert llmobs_events[0] == _expected_llmobs_llm_span_event(
             span,
             model_name=span.get_tag("bedrock.request.model"),
             model_provider=span.get_tag("bedrock.request.model_provider"),
             input_messages=[{"content": mock.ANY}],
-            metadata={
-                "temperature": float(span.get_tag("bedrock.request.temperature")),
-                "max_tokens": int(span.get_tag("bedrock.request.max_tokens")),
-            },
+            metadata=metadata,
             output_messages=[{"content": ""}],
             error=span.get_tag("error.type"),
             error_message=span.get_tag("error.message"),
@@ -647,11 +656,8 @@ def test_llmobs_converse_tool_result_json_non_text_or_json(
 )
 class TestLLMObsBedrockProxy:
     @staticmethod
-    def expected_llmobs_span_event_proxy(span, n_output, message=False):
-        if span.get_tag("bedrock.request.temperature"):
-            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"))
+    def expected_llmobs_span_event_proxy(span, n_output, message=False, metadata=None):
+        expected_parameters = metadata if metadata is not None else mock.ANY
         return _expected_llmobs_non_llm_span_event(
             span,
             span_kind="workflow",

tests/snapshots/tests.contrib.botocore.test_bedrock.test_ai21_invoke.json

@@ -11,29 +11,18 @@
     "meta": {
       "_dd.base_service": "tests.contrib.botocore",
       "_dd.p.dm": "-0",
-      "_dd.p.tid": "659c346700000000",
-      "bedrock.request.max_tokens": "10",
+      "_dd.p.tid": "687f905d00000000",
       "bedrock.request.model": "j2-mid-v1",
       "bedrock.request.model_provider": "ai21",
-      "bedrock.request.prompt": "Explain like I'm a five-year old: what is a neural network?",
-      "bedrock.request.stop_sequences": "[]",
-      "bedrock.request.temperature": "0.9",
-      "bedrock.request.top_p": "1.0",
-      "bedrock.response.choices.0.finish_reason": "{'reason': 'length', 'length': 10}",
-      "bedrock.response.choices.0.text": "\\nA neural network is like a secret recipe that a computer uses to learn how to",
-      "bedrock.response.duration": "319",
-      "bedrock.response.id": "1de3312e-48d1-4d7f-8694-733c1c1ea20f",
       "language": "python",
-      "runtime-id": "3dd17f1c810946349e47a84acb56402a"
+      "runtime-id": "f9bd434e80c74a2e87a2796368dac486"
     },
     "metrics": {
       "_dd.top_level": 1,
       "_dd.tracer_kr": 1.0,
       "_sampling_priority_v1": 1,
-      "bedrock.response.usage.completion_tokens": 10,
-      "bedrock.response.usage.prompt_tokens": 10,
-      "process_id": 7458
+      "process_id": 25551
     },
-    "duration": 2112000,
-    "start": 1704735847029810000
+    "duration": 2443000,
+    "start": 1753190493493215000
   }]]

tests/snapshots/tests.contrib.botocore.test_bedrock.test_amazon_embedding.json

@@ -11,23 +11,18 @@
     "meta": {
       "_dd.base_service": "tests.contrib.botocore",
       "_dd.p.dm": "-0",
-      "_dd.p.tid": "662820e400000000",
+      "_dd.p.tid": "687f909300000000",
       "bedrock.request.model": "titan-embed-text-v1",
       "bedrock.request.model_provider": "amazon",
-      "bedrock.request.prompt": "Hello World!",
-      "bedrock.response.duration": "311",
-      "bedrock.response.id": "1fd884e0-c9e8-44fa-b736-d31e2f607d54",
       "language": "python",
-      "runtime-id": "a7bb6456241740dea419398d37aa13d2"
+      "runtime-id": "f9bd434e80c74a2e87a2796368dac486"
     },
     "metrics": {
       "_dd.top_level": 1,
       "_dd.tracer_kr": 1.0,
       "_sampling_priority_v1": 1,
-      "bedrock.response.embedding_length": 1536,
-      "bedrock.response.usage.prompt_tokens": 3,
-      "process_id": 60939
+      "process_id": 25551
     },
-    "duration": 6739000,
-    "start": 1713905892539987000
+    "duration": 2984000,
+    "start": 1753190547196841000
   }]]