fix(botocore): accept AWS ARN bedrock model IDs [backport 2.18] (#11972)

Backport e8d68ae from #11944 to 2.18. This PR fixes the model ID parsing in the bedrock integration to additionally accept model IDs in AWS ARN format. Previously bedrock only worked with base/foundation models which were of the format `{model_provider}.{model_name}`. However now [bedrock](https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_InvokeModel.html#API_runtime_InvokeModel_RequestSyntax) has expanded to allow custom/provisioned/imported/prompt/inference/sagemaker models, following the typical AWS ARN format: ``` Base model: "{model_provider}.{model_name}" Cross-region model: "{region}.{model_provider}.{model_name}" AWS ARNs: Prefixed by "arn:aws{+region?}:bedrock:{region}:{account-id}:" a. Foundation model: ARN prefix + "foundation-model/{region?}.{model_provider}.{model_name}" b. Custom model: ARN prefix + "custom-model/{model_provider}.{model_name}" c. Provisioned model: ARN prefix + "provisioned-model/{model-id}" d. Imported model: ARN prefix + "imported-module/{model-id}" e. Prompt management: ARN prefix + "prompt/{prompt-id}" f. Sagemaker: ARN prefix + "endpoint/{model-id}" g. Inference profile: ARN prefix + "{application-?}inference-profile/{model-id}" h. Default prompt router: ARN prefix + "default-prompt-router/{prompt-id}" ``` Currently if an AWS ARN gets submitted as the model ID (for example, prompt management or a provisioned model) then because we attempt to split the string into a 2/3-part period-separated list, our code breaks when there is no period delimitter in the string. This PR makes a best effort attempt to check each case and extract a model provider and name from the model ID. However, due to the nature of some formats, we default model provider to "custom" and just return the model/prompt/resource ID from the rest of the AWS ARN. ## 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) Co-authored-by: Yun Kim <[email protected]>
DataDog · Jan 15, 2025 · 0e3612f · 0e3612f
1 parent 75a0205
commit 0e3612f
Show file tree

Hide file tree

Showing 6 changed files with 153 additions and 5 deletions.
diff --git a/ddtrace/contrib/internal/botocore/services/bedrock.py b/ddtrace/contrib/internal/botocore/services/bedrock.py
@@ -24,6 +24,17 @@
 _META = "meta"
 _STABILITY = "stability"
 
+_MODEL_TYPE_IDENTIFIERS = (
+    "foundation-model/",
+    "custom-model/",
+    "provisioned-model/",
+    "imported-model/",
+    "prompt/",
+    "endpoint/",
+    "inference-profile/",
+    "default-prompt-router/",
+)
+
 
 class TracedBotocoreStreamingBody(wrapt.ObjectProxy):
     """
@@ -320,14 +331,45 @@ def handle_bedrock_response(
     return result
 
 
+def _parse_model_id(model_id: str):
+    """Best effort to extract the model provider and model name from the bedrock model ID.
+    model_id can be a 1/2 period-separated string or a full AWS ARN, based on the following formats:
+    1. Base model: "{model_provider}.{model_name}"
+    2. Cross-region model: "{region}.{model_provider}.{model_name}"
+    3. Other: Prefixed by AWS ARN "arn:aws{+region?}:bedrock:{region}:{account-id}:"
+        a. Foundation model: ARN prefix + "foundation-model/{region?}.{model_provider}.{model_name}"
+        b. Custom model: ARN prefix + "custom-model/{model_provider}.{model_name}"
+        c. Provisioned model: ARN prefix + "provisioned-model/{model-id}"
+        d. Imported model: ARN prefix + "imported-module/{model-id}"
+        e. Prompt management: ARN prefix + "prompt/{prompt-id}"
+        f. Sagemaker: ARN prefix + "endpoint/{model-id}"
+        g. Inference profile: ARN prefix + "{application-?}inference-profile/{model-id}"
+        h. Default prompt router: ARN prefix + "default-prompt-router/{prompt-id}"
+    If model provider cannot be inferred from the model_id formatting, then default to "custom"
+    """
+    if not model_id.startswith("arn:aws"):
+        model_meta = model_id.split(".")
+        if len(model_meta) < 2:
+            return "custom", model_meta[0]
+        return model_meta[-2], model_meta[-1]
+    for identifier in _MODEL_TYPE_IDENTIFIERS:
+        if identifier not in model_id:
+            continue
+        model_id = model_id.rsplit(identifier, 1)[-1]
+        if identifier in ("foundation-model/", "custom-model/"):
+            model_meta = model_id.split(".")
+            if len(model_meta) < 2:
+                return "custom", model_id
+            return model_meta[-2], model_meta[-1]
+        return "custom", model_id
+    return "custom", "custom"
+
+
 def patched_bedrock_api_call(original_func, instance, args, kwargs, function_vars):
     params = function_vars.get("params")
     pin = function_vars.get("pin")
-    model_meta = params.get("modelId").split(".")
-    if len(model_meta) == 2:
-        model_provider, model_name = model_meta
-    else:
-        _, model_provider, model_name = model_meta  # cross-region inference
+    model_id = params.get("modelId")
+    model_provider, model_name = _parse_model_id(model_id)
     integration = function_vars.get("integration")
     submit_to_llmobs = integration.llmobs_enabled and "embed" not in model_name
     with core.context_with_data(

diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt
@@ -38,6 +38,8 @@ attrs
 autodetected
 autopatching
 aws
+AWS
+ARN
 backend
 backends
 backport

diff --git a/releasenotes/notes/fix-bedrock-model-id-parsing-611aea2ca2e00656.yaml b/releasenotes/notes/fix-bedrock-model-id-parsing-611aea2ca2e00656.yaml
@@ -0,0 +1,4 @@
+---
+fixes:
+  - |
+    botocore: Resolves formatting errors in the bedrock integration when parsing request model IDs, which can now accept AWS ARNs. 
diff --git a/tests/contrib/botocore/bedrock_cassettes/amazon_invoke_model_arn.yaml b/tests/contrib/botocore/bedrock_cassettes/amazon_invoke_model_arn.yaml
@@ -0,0 +1,52 @@
+interactions:
+- request:
+    body: '{"inputText": "Command: can you explain what Datadog is to someone not
+      in the tech industry?", "textGenerationConfig": {"maxTokenCount": 50, "stopSequences":
+      [], "temperature": 0, "topP": 0.9}}'
+    headers:
+      Content-Length:
+      - '193'
+      User-Agent:
+      - !!binary |
+        Qm90bzMvMS4zNC40OSBtZC9Cb3RvY29yZSMxLjM0LjQ5IHVhLzIuMCBvcy9tYWNvcyMyNC4yLjAg
+        bWQvYXJjaCNhcm02NCBsYW5nL3B5dGhvbiMzLjEwLjUgbWQvcHlpbXBsI0NQeXRob24gY2ZnL3Jl
+        dHJ5LW1vZGUjbGVnYWN5IEJvdG9jb3JlLzEuMzQuNDk=
+      X-Amz-Date:
+      - !!binary |
+        MjAyNTAxMTRUMjIwNDAyWg==
+      amz-sdk-invocation-id:
+      - !!binary |
+        YjY5NGZlNDgtNDBmNy00YTJlLWI1YTgtYjRiZGVhZTU5MjQ0
+      amz-sdk-request:
+      - !!binary |
+        YXR0ZW1wdD0x
+    method: POST
+    uri: https://bedrock-runtime.us-east-1.amazonaws.com/model/arn%3Aaws%3Abedrock%3Aus-east-1%3A%3Afoundation-model%2Famazon.titan-tg1-large/invoke
+  response:
+    body:
+      string: '{"inputTextTokenCount":18,"results":[{"tokenCount":50,"outputText":"\n\nDatadog
+        is a monitoring and analytics platform for IT operations, DevOps, and software
+        development teams. It provides real-time monitoring of infrastructure, applications,
+        and services, allowing users to identify and resolve issues quickly. Datadog
+        collects","completionReason":"LENGTH"}]}'
+    headers:
+      Connection:
+      - keep-alive
+      Content-Length:
+      - '361'
+      Content-Type:
+      - application/json
+      Date:
+      - Tue, 14 Jan 2025 22:04:05 GMT
+      X-Amzn-Bedrock-Input-Token-Count:
+      - '18'
+      X-Amzn-Bedrock-Invocation-Latency:
+      - '2646'
+      X-Amzn-Bedrock-Output-Token-Count:
+      - '50'
+      x-amzn-RequestId:
+      - b2d0fd44-c29a-4cd4-a97a-6901a48f6264
+    status:
+      code: 200
+      message: OK
+version: 1
diff --git a/tests/contrib/botocore/test_bedrock.py b/tests/contrib/botocore/test_bedrock.py
@@ -222,6 +222,15 @@ def test_meta_invoke(bedrock_client, request_vcr):
         json.loads(response.get("body").read())
 
 
+@pytest.mark.snapshot
+def test_invoke_model_using_aws_arn_model_id(bedrock_client, request_vcr):
+    body = json.dumps(_REQUEST_BODIES["amazon"])
+    model = "arn:aws:bedrock:us-east-1::foundation-model/amazon.titan-tg1-large"
+    with request_vcr.use_cassette("amazon_invoke_model_arn.yaml"):
+        response = bedrock_client.invoke_model(body=body, modelId=model)
+        json.loads(response.get("body").read())
+
+
 @pytest.mark.snapshot
 def test_amazon_invoke_stream(bedrock_client, request_vcr):
     body, model = json.dumps(_REQUEST_BODIES["amazon"]), _MODELS["amazon"]

diff --git a/...apshots/tests.contrib.botocore.test_bedrock.test_invoke_model_using_aws_arn_model_id.json b/...apshots/tests.contrib.botocore.test_bedrock.test_invoke_model_using_aws_arn_model_id.json
@@ -0,0 +1,39 @@
+[[
+  {
+    "name": "bedrock-runtime.command",
+    "service": "aws.bedrock-runtime",
+    "resource": "InvokeModel",
+    "trace_id": 0,
+    "span_id": 1,
+    "parent_id": 0,
+    "type": "",
+    "error": 0,
+    "meta": {
+      "_dd.base_service": "tests.contrib.botocore",
+      "_dd.p.dm": "-0",
+      "_dd.p.tid": "6786dfda00000000",
+      "bedrock.request.max_tokens": "50",
+      "bedrock.request.model": "titan-tg1-large",
+      "bedrock.request.model_provider": "amazon",
+      "bedrock.request.prompt": "Command: can you explain what Datadog is to someone not in the tech industry?",
+      "bedrock.request.stop_sequences": "[]",
+      "bedrock.request.temperature": "0",
+      "bedrock.request.top_p": "0.9",
+      "bedrock.response.choices.0.finish_reason": "LENGTH",
+      "bedrock.response.choices.0.text": "\\n\\nDatadog is a monitoring and analytics platform for IT operations, DevOps, and software development teams. It provides real-t...",
+      "bedrock.response.duration": "2646",
+      "bedrock.response.id": "b2d0fd44-c29a-4cd4-a97a-6901a48f6264",
+      "bedrock.usage.completion_tokens": "50",
+      "bedrock.usage.prompt_tokens": "18",
+      "language": "python",
+      "runtime-id": "cf8ef38d3504475ba71634071f15d00f"
+    },
+    "metrics": {
+      "_dd.top_level": 1,
+      "_dd.tracer_kr": 1.0,
+      "_sampling_priority_v1": 1,
+      "process_id": 96028
+    },
+    "duration": 2318000,
+    "start": 1736892378210317000
+  }]]