fix(bedrock): add response id attribute (#2544)

Co-authored-by: Nir Gazit <[email protected]>

packages/opentelemetry-instrumentation-bedrock/opentelemetry/instrumentation/bedrock/__init__.py

@@ -25,6 +25,7 @@
     unwrap,
 )
 
+from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import GEN_AI_RESPONSE_ID
 from opentelemetry.semconv_ai import (
     SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY,
     SpanAttributes,
@@ -203,9 +204,13 @@ def stream_done(response_body):
         metric_params.model = model
         metric_params.is_stream = True
 
+        response_model = response_body.get("model")
+        response_id = response_body.get("id")
+
         _set_span_attribute(span, SpanAttributes.LLM_SYSTEM, vendor)
         _set_span_attribute(span, SpanAttributes.LLM_REQUEST_MODEL, model)
-        _set_span_attribute(span, SpanAttributes.LLM_RESPONSE_MODEL, model)
+        _set_span_attribute(span, SpanAttributes.LLM_RESPONSE_MODEL, response_model)
+        _set_span_attribute(span, GEN_AI_RESPONSE_ID, response_id)
 
         if vendor == "cohere":
             _set_cohere_span_attributes(
@@ -248,9 +253,13 @@ def _handle_call(span, kwargs, response, metric_params):
     metric_params.model = model
     metric_params.is_stream = False
 
+    response_model = response_body.get("model")
+    response_id = response_body.get("id")
+
     _set_span_attribute(span, SpanAttributes.LLM_SYSTEM, vendor)
     _set_span_attribute(span, SpanAttributes.LLM_REQUEST_MODEL, model)
-    _set_span_attribute(span, SpanAttributes.LLM_RESPONSE_MODEL, model)
+    _set_span_attribute(span, SpanAttributes.LLM_RESPONSE_MODEL, response_model)
+    _set_span_attribute(span, GEN_AI_RESPONSE_ID, response_id)
 
     if vendor == "cohere":
         _set_cohere_span_attributes(span, request_body, response_body, metric_params)

packages/opentelemetry-instrumentation-bedrock/tests/traces/test_ai21.py

@@ -42,3 +42,6 @@ def test_ai21_j2_completion_string_content(test_context, brt):
         == meta_span.attributes[SpanAttributes.LLM_USAGE_PROMPT_TOKENS]
         + meta_span.attributes[SpanAttributes.LLM_USAGE_COMPLETION_TOKENS]
     )
+    # It is apparently always 1234, but for the sake of consistency,
+    # we should not assert on it.
+    assert meta_span.attributes.get("gen_ai.response.id") == 1234

packages/opentelemetry-instrumentation-bedrock/tests/traces/test_anthropic.py

@@ -46,6 +46,8 @@ def test_anthropic_2_completion(test_context, brt):
     ) == anthropic_span.attributes.get(
         SpanAttributes.LLM_USAGE_TOTAL_TOKENS
     )
+    # Bedrock does not return the response id for claude-v2:1
+    assert anthropic_span.attributes.get("gen_ai.response.id") is None
 
 
 @pytest.mark.vcr
@@ -102,6 +104,7 @@ def test_anthropic_3_completion_complex_content(test_context, brt):
     ) == anthropic_span.attributes.get(
         SpanAttributes.LLM_USAGE_TOTAL_TOKENS
     )
+    assert anthropic_span.attributes.get("gen_ai.response.id") == "msg_bdrk_01Q6Z4xmUkMigo9K4qd1fshW"
 
 
 @pytest.mark.vcr
@@ -165,6 +168,7 @@ def test_anthropic_3_completion_streaming(test_context, brt):
     ) == anthropic_span.attributes.get(
         SpanAttributes.LLM_USAGE_TOTAL_TOKENS
     )
+    assert anthropic_span.attributes.get("gen_ai.response.id") == "msg_bdrk_014eJfxWXNnxFKhmuiT8FYf7"
 
 
 @pytest.mark.vcr
@@ -218,3 +222,4 @@ def test_anthropic_3_completion_string_content(test_context, brt):
     ) == anthropic_span.attributes.get(
         SpanAttributes.LLM_USAGE_TOTAL_TOKENS
     )
+    assert anthropic_span.attributes.get("gen_ai.response.id") == "msg_bdrk_01WR9VHqpyBzBhzgwCDapaQD"

packages/opentelemetry-instrumentation-bedrock/tests/traces/test_cohere.py

@@ -52,6 +52,7 @@ def test_cohere_completion(test_context, brt):
         bedrock_span.attributes[f"{SpanAttributes.LLM_COMPLETIONS}.0.content"]
         == generated_text
     )
+    assert bedrock_span.attributes.get("gen_ai.response.id") == "3266ca30-473c-4491-b6ef-5b1f033798d2"
 
     # Assert on other request parameters
     assert bedrock_span.attributes[SpanAttributes.LLM_REQUEST_MAX_TOKENS] == 200

packages/opentelemetry-instrumentation-bedrock/tests/traces/test_meta.py

@@ -43,6 +43,7 @@ def test_meta_llama2_completion_string_content(test_context, brt):
         meta_span.attributes[SpanAttributes.LLM_USAGE_TOTAL_TOKENS]
         == response_body["generation_token_count"] + response_body["prompt_token_count"]
     )
+    assert meta_span.attributes.get("gen_ai.response.id") is None
 
 
 @pytest.mark.vcr
@@ -80,3 +81,4 @@ def test_meta_llama3_completion(test_context, brt):
         meta_span.attributes[f"{SpanAttributes.LLM_COMPLETIONS}.0.content"]
         == response_body["generation"]
     )
+    assert meta_span.attributes.get("gen_ai.response.id") is None

packages/opentelemetry-instrumentation-bedrock/tests/traces/test_titan.py

@@ -68,3 +68,6 @@ def test_titan_completion(test_context, brt):
     assert bedrock_span.attributes[SpanAttributes.LLM_REQUEST_MAX_TOKENS] == 200
     assert bedrock_span.attributes[SpanAttributes.LLM_REQUEST_TEMPERATURE] == 0.5
     assert bedrock_span.attributes[SpanAttributes.LLM_REQUEST_TOP_P] == 0.5
+    # There is no response id for Amazon Titan models in the response body,
+    # only request id in the response.
+    assert bedrock_span.attributes.get("gen_ai.response.id") is None