feat(bedrock): add support for ARN and cross region endpoint (#2785)

Author: thisthat

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

@@ -261,7 +261,7 @@ def with_instrumentation(*args, **kwargs):
 @dont_throw
 def _handle_stream_call(span, kwargs, response, metric_params):
 
-    (vendor, model) = kwargs.get("modelId").split(".")
+    (vendor, model) = _get_vendor_model(kwargs.get("modelId"))
     request_body = json.loads(kwargs.get("body"))
 
     headers = {}
@@ -299,12 +299,7 @@ def _handle_call(span, kwargs, response, metric_params):
     if "ResponseMetadata" in response:
         headers = response.get("ResponseMetadata").get("HTTPHeaders", {})
 
-    modelId = kwargs.get("modelId")
-    if modelId is not None and "." in modelId:
-        (vendor, model) = modelId.split(".")
-    else:
-        vendor = "imported_model"
-        model = kwargs.get("modelId")
+    (vendor, model) = _get_vendor_model(kwargs.get("modelId"))
     metric_params.vendor = vendor
     metric_params.model = model
     metric_params.is_stream = False
@@ -318,7 +313,7 @@ def _handle_call(span, kwargs, response, metric_params):
 
 @dont_throw
 def _handle_converse(span, kwargs, response, metric_params):
-    (vendor, model) = kwargs.get("modelId").split(".")
+    (vendor, model) = _get_vendor_model(kwargs.get("modelId"))
     guardrail_converse(response, vendor, model, metric_params)
 
     _set_span_attribute(span, SpanAttributes.LLM_SYSTEM, vendor)
@@ -359,7 +354,7 @@ def _handle_converse(span, kwargs, response, metric_params):
 
 @dont_throw
 def _handle_converse_stream(span, kwargs, response, metric_params):
-    (vendor, model) = kwargs.get("modelId").split(".")
+    (vendor, model) = _get_vendor_model(kwargs.get("modelId"))
 
     _set_span_attribute(span, SpanAttributes.LLM_SYSTEM, vendor)
     _set_span_attribute(span, SpanAttributes.LLM_REQUEST_MODEL, model)
@@ -417,6 +412,37 @@ def wrap(*args, **kwargs):
         stream._parse_event = handler(stream._parse_event)
 
 
+def _get_vendor_model(modelId):
+    # Docs:
+    # https://docs.aws.amazon.com/bedrock/latest/userguide/inference-profiles-support.html#inference-profiles-support-system
+    vendor = "imported_model"
+    model = modelId
+
+    if modelId is not None and modelId.startswith("arn"):
+        components = modelId.split(":")
+        if len(components) > 5:
+            inf_profile = components[5].split("/")
+            if len(inf_profile) == 2:
+                if "." in inf_profile[1]:
+                    (vendor, model) = _cross_region_check(inf_profile[1])
+    elif modelId is not None and "." in modelId:
+        (vendor, model) = _cross_region_check(modelId)
+
+    return vendor, model
+
+
+def _cross_region_check(value):
+    prefixes = ["us", "us-gov", "eu", "apac"]
+    if any(value.startswith(prefix + ".") for prefix in prefixes):
+        parts = value.split(".")
+        if len(parts) > 2:
+            parts.pop(0)
+        return parts[0], parts[1]
+    else:
+        (vendor, model) = value.split(".")
+    return vendor, model
+
+
 def _report_converse_input_prompt(kwargs, span):
     prompt_idx = 0
     if "system" in kwargs:

packages/opentelemetry-instrumentation-bedrock/tests/traces/cassettes/test_anthropic/test_anthropic_cross_region.yaml

@@ -0,0 +1,50 @@
+interactions:
+  - request:
+      body: '{"messages": [{"role": "user", "content": [{"text": "Human: Tell me a joke
+      about opentelemetry Assistant:"}]}], "guardrailConfig": {"guardrailIdentifier":
+      "v9kpg6yrwhs2", "guardrailVersion": "DRAFT", "trace": "enabled"}, "inferenceConfig":
+      {"temperature": 0.5}}'
+      headers:
+        Content-Length:
+          - '261'
+        Content-Type:
+          - !!binary |
+            YXBwbGljYXRpb24vanNvbg==
+        User-Agent:
+          - !!binary |
+            Qm90bzMvMS4zNy4xMyBtZC9Cb3RvY29yZSMxLjM3LjEzIHVhLzIuMSBvcy9tYWNvcyMyNC4zLjAg
+            bWQvYXJjaCNhcm02NCBsYW5nL3B5dGhvbiMzLjkuNiBtZC9weWltcGwjQ1B5dGhvbiBjZmcvcmV0
+            cnktbW9kZSNsZWdhY3kgQm90b2NvcmUvMS4zNy4xMw==
+        X-Amz-Date:
+          - !!binary |
+            MjAyNTAzMTdUMTMzMjAzWg==
+        amz-sdk-invocation-id:
+          - !!binary |
+            YWRkNDMwM2EtMDhmZS00MjdjLWE0OWEtMzU0NzkzYzE5OWIz
+        amz-sdk-request:
+          - !!binary |
+            YXR0ZW1wdD0x
+      method: POST
+      uri: https://bedrock-runtime.us-east-1.amazonaws.com/model/arn%3Aaws%3Abedrock%3Aus-east-1%3A012345678901%3Ainference-profile%2Fus.anthropic.claude-3-7-sonnet-20250219-v1%3A0/converse
+    response:
+      body:
+        string: '{"metrics":{"latencyMs":2900},"output":{"message":{"content":[{"text":"Why
+        did the developer start using OpenTelemetry?\n\nBecause they wanted to \"trace\"
+        their steps when their application was running slow, but got lost in a \"span\"
+        of metrics, logs, and traces!\n\nTurns out debugging is much easier when you''re
+        not completely in the dark about what your distributed system is doing!"}],"role":"assistant"}},"stopReason":"end_turn","trace":{"guardrail":{"inputAssessment":{"v9kpg6yrwhs2":{"invocationMetrics":{"guardrailCoverage":{"textCharacters":{"guarded":52,"total":52}},"guardrailProcessingLatency":296,"usage":{"contentPolicyUnits":1,"contextualGroundingPolicyUnits":0,"sensitiveInformationPolicyFreeUnits":0,"sensitiveInformationPolicyUnits":1,"topicPolicyUnits":1,"wordPolicyUnits":1}}}},"outputAssessments":{"v9kpg6yrwhs2":[{"invocationMetrics":{"guardrailCoverage":{"textCharacters":{"guarded":308,"total":308}},"guardrailProcessingLatency":236,"usage":{"contentPolicyUnits":1,"contextualGroundingPolicyUnits":0,"sensitiveInformationPolicyFreeUnits":1,"sensitiveInformationPolicyUnits":1,"topicPolicyUnits":1,"wordPolicyUnits":1}}}]}}},"usage":{"cacheReadInputTokenCount":0,"cacheReadInputTokens":0,"cacheWriteInputTokenCount":0,"cacheWriteInputTokens":0,"inputTokens":20,"outputTokens":72,"totalTokens":92}}'
+      headers:
+        Connection:
+          - keep-alive
+        Content-Length:
+          - '1322'
+        Content-Type:
+          - application/json
+        Date:
+          - Mon, 17 Mar 2025 13:32:07 GMT
+        x-amzn-RequestId:
+          - 48a1e323-e366-4401-93a3-eb326790a4db
+      status:
+        code: 200
+        message: OK
+version: 1

packages/opentelemetry-instrumentation-bedrock/tests/traces/cassettes/test_nova/test_nova_cross_region_invoke.yaml

@@ -0,0 +1,55 @@
+interactions:
+- request:
+    body: '{"messages": [{"role": "user", "content": [{"text": "Tell me a joke about
+      OpenTelemetry"}]}], "inferenceConfig": {"maxTokens": 500, "topP": 0.9, "topK":
+      20, "temperature": 0.7}}'
+    headers:
+      Content-Length:
+      - '177'
+      User-Agent:
+      - !!binary |
+        Qm90bzMvMS4zNy4xMyBtZC9Cb3RvY29yZSMxLjM3LjEzIHVhLzIuMSBvcy9tYWNvcyMyNC4zLjAg
+        bWQvYXJjaCNhcm02NCBsYW5nL3B5dGhvbiMzLjkuNiBtZC9weWltcGwjQ1B5dGhvbiBjZmcvcmV0
+        cnktbW9kZSNsZWdhY3kgQm90b2NvcmUvMS4zNy4xMw==
+      X-Amz-Date:
+      - !!binary |
+        MjAyNTAzMTdUMTQxOTIyWg==
+      amz-sdk-invocation-id:
+      - !!binary |
+        MTQ1ZTJiMTktYzRiMC00OTQzLTkxNzEtNjdiZjA3NDRhYmFj
+      amz-sdk-request:
+      - !!binary |
+        YXR0ZW1wdD0x
+    method: POST
+    uri: https://bedrock-runtime.us-east-1.amazonaws.com/model/us.amazon.nova-lite-v1%3A0/invoke
+  response:
+    body:
+      string: '{"output":{"message":{"content":[{"text":"Sure, here''s a joke for
+        you:\n\nWhy did the OpenTelemetry developer bring a ladder to work?\n\nBecause
+        they heard the company was going to great lengths to trace their performance
+        issues to the \"root cause\"!\n\nHope that brought a smile to your face!"}],"role":"assistant"}},"stopReason":"end_turn","usage":{"inputTokens":7,"outputTokens":53,"totalTokens":60,"cacheReadInputTokenCount":0,"cacheWriteInputTokenCount":0}}'
+    headers:
+      Connection:
+      - keep-alive
+      Content-Length:
+      - '463'
+      Content-Type:
+      - application/json
+      Date:
+      - Mon, 17 Mar 2025 14:19:23 GMT
+      X-Amzn-Bedrock-Cache-Read-Input-Token-Count:
+      - '0'
+      X-Amzn-Bedrock-Cache-Write-Input-Token-Count:
+      - '0'
+      X-Amzn-Bedrock-Input-Token-Count:
+      - '7'
+      X-Amzn-Bedrock-Invocation-Latency:
+      - '596'
+      X-Amzn-Bedrock-Output-Token-Count:
+      - '53'
+      x-amzn-RequestId:
+      - ae14e85a-2f22-46b0-8540-277b810fab77
+    status:
+      code: 200
+      message: OK
+version: 1

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

@@ -232,3 +232,55 @@ def test_anthropic_3_completion_string_content(test_context, brt):
         anthropic_span.attributes.get("gen_ai.response.id")
         == "msg_bdrk_01WR9VHqpyBzBhzgwCDapaQD"
     )
+
+
+@pytest.mark.vcr
+def test_anthropic_cross_region(test_context, brt):
+    inference_config = {"temperature": 0.5}
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"text": "Human: Tell me a joke about opentelemetry Assistant:"},
+            ],
+        },
+    ]
+    response = brt.converse(
+        modelId="arn:aws:bedrock:us-east-1:012345678901:inference-profile/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
+        messages=messages,
+        inferenceConfig=inference_config,
+    )
+
+    completion = response["output"]["message"]["content"][0]['text']
+
+    exporter, _, _ = test_context
+    spans = exporter.get_finished_spans()
+
+    assert len(spans) == 1
+
+    anthropic_span = spans[0]
+    assert anthropic_span.name == "bedrock.converse"
+
+    # Assert on model name and vendor
+    assert anthropic_span.attributes[SpanAttributes.LLM_REQUEST_MODEL] == "claude-3-7-sonnet-20250219-v1"
+    assert anthropic_span.attributes[SpanAttributes.LLM_SYSTEM] == "anthropic"
+
+    assert (
+            anthropic_span.attributes[f"{SpanAttributes.LLM_PROMPTS}.0.content"]
+            == json.dumps(messages[0]["content"])
+    )
+    assert (
+            anthropic_span.attributes.get(f"{SpanAttributes.LLM_COMPLETIONS}.0.content")
+            == completion
+    )
+
+    assert anthropic_span.attributes.get(SpanAttributes.LLM_USAGE_PROMPT_TOKENS) == 20
+    assert anthropic_span.attributes.get(
+        SpanAttributes.LLM_USAGE_COMPLETION_TOKENS
+    ) + anthropic_span.attributes.get(
+        SpanAttributes.LLM_USAGE_PROMPT_TOKENS
+    ) == 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

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

@@ -374,3 +374,67 @@ def test_nova_converse_stream(test_context, brt):
         bedrock_span.attributes[SpanAttributes.LLM_USAGE_TOTAL_TOKENS]
         == inputTokens + outputTokens
     )
+
+
+@pytest.mark.vcr
+def test_nova_cross_region_invoke(test_context, brt):
+
+    message_list = [{"role": "user", "content": [{"text": "Tell me a joke about OpenTelemetry"}]}]
+    inf_params = {"maxTokens": 500, "topP": 0.9, "topK": 20, "temperature": 0.7}
+    request_body = {
+        "messages": message_list,
+        "inferenceConfig": inf_params,
+    }
+
+    modelId = "us.amazon.nova-lite-v1:0"
+    accept = "application/json"
+    contentType = "application/json"
+
+    response = brt.invoke_model(
+        body=json.dumps(request_body),
+        modelId=modelId,
+        accept=accept,
+        contentType=contentType,
+    )
+
+    response_body = json.loads(response.get("body").read())
+
+    exporter, _, _ = test_context
+    spans = exporter.get_finished_spans()
+    assert len(spans) == 1
+    assert spans[0].name == "bedrock.completion"
+
+    bedrock_span = spans[0]
+
+    # Assert on model name and vendor
+    assert bedrock_span.attributes[SpanAttributes.LLM_REQUEST_MODEL] == "nova-lite-v1:0"
+    assert bedrock_span.attributes[SpanAttributes.LLM_SYSTEM] == "amazon"
+
+    # Assert on vendor
+    assert bedrock_span.attributes[SpanAttributes.LLM_SYSTEM] == "amazon"
+
+    # Assert on request type
+    assert bedrock_span.attributes[SpanAttributes.LLM_REQUEST_TYPE] == "completion"
+
+    # Assert on prompt
+    assert bedrock_span.attributes[f"{SpanAttributes.LLM_PROMPTS}.0.role"] == "user"
+
+    assert bedrock_span.attributes[
+               f"{SpanAttributes.LLM_PROMPTS}.0.content"
+           ] == json.dumps(message_list[0].get("content"), default=str)
+
+    # Assert on response
+    generated_text = response_body["output"]["message"]["content"]
+    for i in range(0, len(generated_text)):
+        assert (
+                bedrock_span.attributes[f"{SpanAttributes.LLM_COMPLETIONS}.{i}.content"]
+                == generated_text[i]["text"]
+        )
+
+    # Assert on other request parameters
+    assert bedrock_span.attributes[SpanAttributes.LLM_REQUEST_MAX_TOKENS] == 500
+    assert bedrock_span.attributes[SpanAttributes.LLM_REQUEST_TEMPERATURE] == 0.7
+    assert bedrock_span.attributes[SpanAttributes.LLM_REQUEST_TOP_P] == 0.9
+    # 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