From 67873de681058713aa24128589c0f11f3cbf565c Mon Sep 17 00:00:00 2001 From: Dinmukhamed Mailibay <47117969+dinmukhamedm@users.noreply.github.com> Date: Thu, 23 Jan 2025 12:34:07 +0500 Subject: [PATCH] fix(bedrock): add response id attribute (#2544) Co-authored-by: Nir Gazit --- .../instrumentation/bedrock/__init__.py | 13 +++++++++++-- .../tests/traces/test_ai21.py | 3 +++ .../tests/traces/test_anthropic.py | 5 +++++ .../tests/traces/test_cohere.py | 1 + .../tests/traces/test_meta.py | 2 ++ .../tests/traces/test_titan.py | 3 +++ 6 files changed, 25 insertions(+), 2 deletions(-) diff --git a/packages/opentelemetry-instrumentation-bedrock/opentelemetry/instrumentation/bedrock/__init__.py b/packages/opentelemetry-instrumentation-bedrock/opentelemetry/instrumentation/bedrock/__init__.py index ecb3be18e..3833140a4 100644 --- a/packages/opentelemetry-instrumentation-bedrock/opentelemetry/instrumentation/bedrock/__init__.py +++ b/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) diff --git a/packages/opentelemetry-instrumentation-bedrock/tests/traces/test_ai21.py b/packages/opentelemetry-instrumentation-bedrock/tests/traces/test_ai21.py index 9d33a9e22..bbce97a29 100644 --- a/packages/opentelemetry-instrumentation-bedrock/tests/traces/test_ai21.py +++ b/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 diff --git a/packages/opentelemetry-instrumentation-bedrock/tests/traces/test_anthropic.py b/packages/opentelemetry-instrumentation-bedrock/tests/traces/test_anthropic.py index 5a0689625..c3a09fec2 100644 --- a/packages/opentelemetry-instrumentation-bedrock/tests/traces/test_anthropic.py +++ b/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" diff --git a/packages/opentelemetry-instrumentation-bedrock/tests/traces/test_cohere.py b/packages/opentelemetry-instrumentation-bedrock/tests/traces/test_cohere.py index f35374c9e..63992e3db 100644 --- a/packages/opentelemetry-instrumentation-bedrock/tests/traces/test_cohere.py +++ b/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 diff --git a/packages/opentelemetry-instrumentation-bedrock/tests/traces/test_meta.py b/packages/opentelemetry-instrumentation-bedrock/tests/traces/test_meta.py index 0b8d516b3..82e799444 100644 --- a/packages/opentelemetry-instrumentation-bedrock/tests/traces/test_meta.py +++ b/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 diff --git a/packages/opentelemetry-instrumentation-bedrock/tests/traces/test_titan.py b/packages/opentelemetry-instrumentation-bedrock/tests/traces/test_titan.py index 577f11459..d9b28e7da 100644 --- a/packages/opentelemetry-instrumentation-bedrock/tests/traces/test_titan.py +++ b/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