Aws sut factory (#1499)

* aws sut factory

* fix matching

* rename nova sut

* fix ID parsing, only include active models

* tests

* Fix query

* get region name from env

Author: bkorycki

src/modelgauge/sut_factory.py

@@ -7,6 +7,7 @@
 from modelgauge.sut_definition import SUTDefinition
 from modelgauge.sut_registry import SUTS
 from modelgauge.suts.anthropic_sut_factory import AnthropicSUTFactory
+from modelgauge.suts.aws_bedrock_sut_factory import AWSBedrockSUTFactory
 from modelgauge.suts.google_sut_factory import GoogleSUTFactory
 from modelgauge.suts.huggingface_sut_factory import HuggingFaceSUTFactory
 from modelgauge.suts.indirect_sut import IndirectSUTFactory
@@ -37,6 +38,7 @@ class SUTType(Enum):
 # that can be used to create a dynamic sut
 DYNAMIC_SUT_FACTORIES: dict = {
     "anthropic": AnthropicSUTFactory,
+    "aws": AWSBedrockSUTFactory,
     "google": GoogleSUTFactory,
     "hf": HuggingFaceSUTFactory,
     "hfrelay": HuggingFaceSUTFactory,

src/modelgauge/suts/aws_bedrock_client.py

@@ -1,6 +1,7 @@
 # as defined here:
 # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/bedrock-runtime/client/converse.html
 
+import os
 from typing import Any, Dict, List, Optional
 
 import boto3
@@ -106,7 +107,7 @@ class BedrockResponseUsage(BaseModel):
 
 
 @modelgauge_sut(capabilities=[AcceptsTextPrompt])
-class AmazonNovaSut(PromptResponseSUT):
+class AmazonBedrockSut(PromptResponseSUT):
 
     def __init__(self, uid: str, model_id: str, access_key_id: AwsAccessKeyId, secret_access_key: AwsSecretAccessKey):
         super().__init__(uid)
@@ -118,7 +119,7 @@ def __init__(self, uid: str, model_id: str, access_key_id: AwsAccessKeyId, secre
     def _load_client(self):
         return boto3.client(
             service_name="bedrock-runtime",
-            region_name="us-east-1",
+            region_name=os.getenv("AWS_REGION", "us-east-1"),
             aws_access_key_id=self.access_key_id,
             aws_secret_access_key=self.secret_access_key,
         )
@@ -159,7 +160,7 @@ def translate_response(self, request: BedrockRequest, response: BedrockResponse)
 
 for model in BEDROCK_MODELS:
     SUTS.register(
-        AmazonNovaSut,
+        AmazonBedrockSut,
         f"amazon-nova-1.0-{model}",
         f"amazon.nova-{model}-v1:0",
         InjectSecret(AwsAccessKeyId),
@@ -169,7 +170,7 @@ def translate_response(self, request: BedrockRequest, response: BedrockResponse)
 BEDROCK_INFERENCE_PROFILES = ["premier"]
 for model in BEDROCK_INFERENCE_PROFILES:
     SUTS.register(
-        AmazonNovaSut,
+        AmazonBedrockSut,
         f"amazon-nova-1.0-{model}",
         f"us.amazon.nova-{model}-v1:0",
         InjectSecret(AwsAccessKeyId),

src/modelgauge/suts/aws_bedrock_sut_factory.py

@@ -0,0 +1,60 @@
+import os
+
+import boto3
+
+from modelgauge.dynamic_sut_factory import DynamicSUTFactory, ModelNotSupportedError
+from modelgauge.secret_values import InjectSecret, RawSecrets
+from modelgauge.sut import SUT
+from modelgauge.sut_definition import SUTDefinition
+from modelgauge.suts.aws_bedrock_client import AmazonBedrockSut, AwsAccessKeyId, AwsSecretAccessKey
+
+
+class AWSBedrockSUTFactory(DynamicSUTFactory):
+    DRIVER_NAME = "aws"
+
+    def __init__(self, raw_secrets: RawSecrets):
+        super().__init__(raw_secrets)
+        self._client = None  # Lazy load.
+
+    @property
+    def client(self):
+        if self._client is None:
+            self._client = boto3.client(
+                service_name="bedrock",
+                region_name=os.getenv("AWS_REGION", "us-east-1"),
+                aws_access_key_id=self.injected_secrets()[0].value,
+                aws_secret_access_key=self.injected_secrets()[1].value,
+            )
+        return self._client
+
+    def _convert_model_id(self, model_id: str) -> SUTDefinition:
+        """Convert AWS model IDs (maker.model[:version?]) to our standard format."""
+        maker, model_name = model_id.split(".", maxsplit=1)
+        model_name = model_name.replace(":", ".")
+        return SUTDefinition({"maker": maker, "model": model_name, "driver": self.DRIVER_NAME})
+
+    def _get_available_models(self, maker: str):
+        response = self.client.list_foundation_models()
+        models = {}
+        for m in response["modelSummaries"]:
+            if m.get("modelLifecycle", {}).get("status") != "ACTIVE":
+                continue
+            models[m["modelId"]] = self._convert_model_id(m["modelId"])
+        return models
+
+    def _get_model_id(self, sut_definition: SUTDefinition):
+        models = self._get_available_models(sut_definition.to_dynamic_sut_metadata().maker)
+        for model_id, model_definition in models.items():
+            if str(model_definition.to_dynamic_sut_metadata()) == str(sut_definition.to_dynamic_sut_metadata()):
+                return model_id
+        supported_models = [model_def.to_dynamic_sut_metadata().external_model_name() for model_def in models.values()]
+        raise ModelNotSupportedError(
+            f"Model {sut_definition.external_model_name()} not found among AWS Bedrock models. AWS carries the following models from maker {sut_definition.get("maker")}: {supported_models} "
+        )
+
+    def get_secrets(self) -> list[InjectSecret]:
+        return [InjectSecret(AwsAccessKeyId), InjectSecret(AwsSecretAccessKey)]
+
+    def make_sut(self, sut_definition: SUTDefinition) -> SUT:
+        model_id = self._get_model_id(sut_definition)
+        return AmazonBedrockSut(sut_definition.dynamic_uid, model_id, *self.injected_secrets())

tests/modelgauge_tests/sut_tests/test_aws_bedrock_client.py

@@ -7,7 +7,7 @@
 from modelgauge.typed_data import is_typeable
 
 from modelgauge.suts.aws_bedrock_client import (
-    AmazonNovaSut,
+    AmazonBedrockSut,
     AwsAccessKeyId,
     AwsSecretAccessKey,
     BedrockRequest,
@@ -19,7 +19,7 @@
 
 @pytest.fixture
 def fake_sut():
-    return AmazonNovaSut(
+    return AmazonBedrockSut(
         "fake-sut", FAKE_MODEL_ID, AwsAccessKeyId("fake-api-key"), AwsSecretAccessKey("fake-secret-key")
     )
 

tests/modelgauge_tests/sut_tests/test_aws_bedrock_sut_factory.py

@@ -0,0 +1,68 @@
+import pytest
+from unittest.mock import patch
+
+from modelgauge.dynamic_sut_factory import ModelNotSupportedError
+from modelgauge.sut_definition import SUTDefinition
+from modelgauge.suts.aws_bedrock_client import AmazonBedrockSut
+from modelgauge.suts.aws_bedrock_sut_factory import AWSBedrockSUTFactory
+
+
+@pytest.fixture
+def factory():
+    return AWSBedrockSUTFactory({"aws": {"access_key_id": "value", "secret_access_key": "value"}})
+
+
+@pytest.fixture
+def mock_list_foundation_models():
+    models = {
+        "modelSummaries": [
+            {"modelId": "amazon.nova-1.0-micro-v1:0", "modelLifecycle": {"status": "ACTIVE"}},
+            {"modelId": "old_model", "modelLifecycle": {"status": "LEGACY"}},
+        ]
+    }
+
+    with patch("boto3.client") as mock_client:
+
+        mock_client.return_value.list_foundation_models.return_value = models
+
+        yield mock_client
+
+
+def test_convert_model_id(factory):
+    definition = factory._convert_model_id("amazon.nova-v1")
+    assert definition.get("maker") == "amazon"
+    assert definition.get("model") == "nova-v1"
+    assert definition.get("driver") == "aws"
+
+    # Sometimes they have colons
+    definition = factory._convert_model_id("amazon.nova-v1:0")
+    assert definition.get("maker") == "amazon"
+    assert definition.get("model") == "nova-v1.0"
+    assert definition.get("driver") == "aws"
+
+    # "." in the model name
+    definition = factory._convert_model_id("moonshotai.kimi-k2.5")
+    assert definition.get("maker") == "moonshotai"
+    assert definition.get("model") == "kimi-k2.5"
+    assert definition.get("driver") == "aws"
+
+
+def test_make_sut(factory, mock_list_foundation_models):
+    sut_definition = SUTDefinition(model="nova-1.0-micro-v1.0", maker="amazon", driver="aws")
+    sut = factory.make_sut(sut_definition)
+
+    assert isinstance(sut, AmazonBedrockSut)
+    assert sut.uid == "amazon/nova-1.0-micro-v1.0:aws"
+    assert sut.model_id == "amazon.nova-1.0-micro-v1:0"
+
+
+def test_make_sut_no_model(factory, mock_list_foundation_models):
+    sut_definition = SUTDefinition(model="unknown", maker="amazon", driver="aws")
+    with pytest.raises(ModelNotSupportedError):
+        factory.make_sut(sut_definition)
+
+
+def test_make_sut_legacy_model(factory, mock_list_foundation_models):
+    sut_definition = SUTDefinition(model="old_model", maker="amazon", driver="aws")
+    with pytest.raises(ModelNotSupportedError):
+        factory.make_sut(sut_definition)