fix: Pass model id in Bedrock client init (#71)

* feat: Moved model_id in bedrock to init

* test: Added tests to check model id

* fix: bedrock example

* refactor: rename test params

Author: asafgardin

ai21/clients/bedrock/ai21_bedrock_client.py

@@ -9,15 +9,16 @@
 
 class AI21BedrockClient:
     """
+    :param model_id: The model ID to use for the client.
     :param session: An optional boto3 session to use for the client.
     """
 
     def __init__(
         self,
+        model_id: str,
         session: Optional[boto3.Session] = None,
         region: Optional[str] = None,
         env_config: _AI21EnvConfig = AI21EnvConfig,
     ):
-
         self._bedrock_session = BedrockSession(session=session, region=region or env_config.aws_region)
-        self.completion = BedrockCompletion(self._bedrock_session)
+        self.completion = BedrockCompletion(model_id=model_id, bedrock_session=self._bedrock_session)

ai21/clients/bedrock/resources/bedrock_completion.py

@@ -7,7 +7,6 @@
 class BedrockCompletion(BedrockResource):
     def create(
         self,
-        model_id: str,
         prompt: str,
         *,
         max_tokens: Optional[int] = None,
@@ -42,6 +41,8 @@ def create(
         if count_penalty is not None:
             body["countPenalty"] = count_penalty.to_dict()
 
+        model_id = kwargs.get("model_id", self._model_id)
+
         raw_response = self._invoke(model_id=model_id, body=body)
 
         return CompletionsResponse.from_dict(raw_response)

ai21/clients/bedrock/resources/bedrock_resource.py

@@ -6,7 +6,8 @@
 
 
 class BedrockResource(ABC):
-    def __init__(self, bedrock_session: BedrockSession):
+    def __init__(self, model_id: str, bedrock_session: BedrockSession):
+        self._model_id = model_id
         self._bedrock_session = bedrock_session
 
     def _invoke(self, model_id: str, body: Dict[str, Any]) -> Dict[str, Any]:

examples/bedrock/completion.py

@@ -38,10 +38,9 @@
     "User: Hi, I have a question for you"
 )
 
-response = AI21BedrockClient().completion.create(
+response = AI21BedrockClient(model_id=BedrockModelID.J2_MID_V1).completion.create(
     prompt=prompt,
     max_tokens=1000,
-    model_id=BedrockModelID.J2_MID_V1,
     temperature=0,
     top_p=1,
     top_k_return=0,

tests/integration_tests/clients/bedrock/test_completion.py

@@ -50,11 +50,10 @@
 def test_completion__when_no_penalties__should_return_response(
     frequency_penalty: Optional[Penalty], presence_penalty: Optional[Penalty], count_penalty: Optional[Penalty]
 ):
-    client = AI21BedrockClient()
+    client = AI21BedrockClient(model_id=BedrockModelID.J2_MID_V1)
     response = client.completion.create(
         prompt=_PROMPT,
         max_tokens=64,
-        model_id=BedrockModelID.J2_MID_V1,
         temperature=0,
         top_p=1,
         top_k_return=0,

tests/unittests/clients/bedrock/__init__.py

@@ -0,0 +1,11 @@
+from unittest.mock import Mock
+
+import pytest
+from pytest_mock import MockerFixture
+
+from ai21.clients.bedrock.bedrock_session import BedrockSession
+
+
+@pytest.fixture
+def mock_bedrock_session(mocker: MockerFixture) -> Mock:
+    return mocker.MagicMock(spec=BedrockSession)

tests/unittests/clients/bedrock/resources/__init__.py

@@ -0,0 +1,45 @@
+from typing import Optional
+
+import pytest
+from pytest_mock import MockerFixture
+from unittest.mock import Mock, ANY
+
+from ai21.clients.bedrock.resources.bedrock_completion import BedrockCompletion
+from ai21.models import Prompt
+
+_CTOR_PROVIDED_MODEL_ID = "constructor_provided_model_id"
+_INVOCATION_MODEL_ID = "invocation_model_id"
+
+
+@pytest.mark.parametrize(
+    ids=[
+        "when_model_id_not_passed_in_create__should_use_model_id_from_init",
+        "when_model_id_passed_in_create__should_use_model_id_from_create",
+    ],
+    argnames=["invocation_model_id", "expected_model_id"],
+    argvalues=[
+        (None, _CTOR_PROVIDED_MODEL_ID),
+        (_INVOCATION_MODEL_ID, _INVOCATION_MODEL_ID),
+    ],
+)
+def test__when_model_id_create_and_init__should_use_one_from_create(
+    invocation_model_id: Optional[str],
+    expected_model_id: str,
+    mock_bedrock_session: Mock,
+    mocker: MockerFixture,
+):
+    mock_bedrock_session.invoke_model.return_value = {
+        "id": expected_model_id,
+        "prompt": mocker.MagicMock(spec=Prompt),
+        "completions": [],
+    }
+
+    client = BedrockCompletion(model_id=_CTOR_PROVIDED_MODEL_ID, bedrock_session=mock_bedrock_session)
+
+    # We can not pass "None" explicitly to the create method, so we have to use the if else statement
+    if invocation_model_id is None:
+        client.create(prompt="test")
+    else:
+        client.create(model_id=invocation_model_id, prompt="test")
+
+    mock_bedrock_session.invoke_model.assert_called_once_with(model_id=expected_model_id, input_json=ANY)