fix: Pass default model_id in bedrock client (#72)

* fix: Pass default model_id in bedrock client

* fix: Added validation for model_id

* fix: Added validation for model_id

Author: asafgardin

ai21/clients/bedrock/ai21_bedrock_client.py

@@ -15,7 +15,7 @@ class AI21BedrockClient:
 
     def __init__(
         self,
-        model_id: str,
+        model_id: Optional[str] = None,
         session: Optional[boto3.Session] = None,
         region: Optional[str] = None,
         env_config: _AI21EnvConfig = AI21EnvConfig,

ai21/clients/bedrock/resources/bedrock_completion.py

@@ -43,6 +43,9 @@ def create(
 
         model_id = kwargs.get("model_id", self._model_id)
 
+        if model_id is None:
+            raise ValueError("model_id should be provided in either the constructor or the 'create' method call")
+
         raw_response = self._invoke(model_id=model_id, body=body)
 
         return CompletionsResponse.from_dict(raw_response)

tests/integration_tests/clients/bedrock/test_completion.py

@@ -68,3 +68,12 @@ def test_completion__when_no_penalties__should_return_response(
     assert len([completion.data.text for completion in response.completions]) == 1
     for completion in response.completions:
         assert isinstance(completion.data.text, str)
+
+
+@pytest.mark.skipif(should_skip_bedrock_integration_tests(), reason="No keys supplied for AWS. Skipping.")
+def test_completion__when_no_model_id__should_raise_exception():
+    with pytest.raises(ValueError) as e:
+        client = AI21BedrockClient()
+        client.completion.create(prompt=_PROMPT)
+
+    assert e.value.args[0] == "model_id should be provided in either the constructor or the 'create' method call"