refactor: completion test refactor

Author: etang-ai21

ai21/clients/bedrock/resources/bedrock_completion.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from typing import List
 
 from ai21.clients.bedrock.resources.bedrock_resource import BedrockResource

tests/integration_tests/clients/bedrock/test_completion.py

@@ -40,28 +40,24 @@ def test_completion__when_no_penalties__should_return_response(
     frequency_penalty: Optional[Penalty], presence_penalty: Optional[Penalty], count_penalty: Optional[Penalty]
 ):
     client = AI21BedrockClient()
+    completion_args = dict(
+        prompt=_PROMPT,
+        max_tokens=64,
+        model_id=BedrockModelID.J2_MID_V1,
+        temperature=0,
+        top_p=1,
+        top_k_return=0,
+    )
 
-    if frequency_penalty is None and presence_penalty is None and count_penalty is None:
-        response = client.completion.create(
-            prompt=_PROMPT,
-            max_tokens=64,
-            model_id=BedrockModelID.J2_MID_V1,
-            temperature=0,
-            top_p=1,
-            top_k_return=0,
-        )
-    else:
-        response = client.completion.create(
-            prompt=_PROMPT,
-            max_tokens=64,
-            model_id=BedrockModelID.J2_MID_V1,
-            temperature=0,
-            top_p=1,
-            top_k_return=0,
-            frequency_penalty=frequency_penalty,
-            presence_penalty=presence_penalty,
-            count_penalty=count_penalty,
-        )
+    for arg_name, penalty in [
+        ("frequency_penalty", frequency_penalty),
+        ("presence_penalty", presence_penalty),
+        ("count_penalty", count_penalty),
+    ]:
+        if penalty:
+            completion_args[arg_name] = penalty
+
+    response = client.completion.create(**completion_args)
 
     assert response.prompt.text == _PROMPT
     assert len(response.completions) == 1