fix: aws tests

.github/workflows/integration-tests.yaml

@@ -34,7 +34,7 @@ jobs:
           poetry env use ${{ matrix.python-version }}
       - name: Install dependencies
         run: |
-          poetry install --no-root --only dev --all-extras
+          poetry install --only dev --all-extras
       - name: Lint Python (Black)
         run: |
           poetry run inv formatter

ai21/clients/bedrock/resources/bedrock_completion.py

@@ -1,7 +1,7 @@
 from typing import Optional, List
 
-from ai21.resources import Penalty
 from ai21.clients.bedrock.resources.bedrock_resource import BedrockResource
+from ai21.models import Penalty
 from ai21.models.responses.completion_response import CompletionsResponse
 
 

ai21/clients/sagemaker/resources/sagemaker_completion.py

@@ -1,8 +1,8 @@
 from typing import Optional, List
 
-from ai21.resources import Penalty
-from ai21.models.responses.completion_response import CompletionsResponse
 from ai21.clients.sagemaker.resources.sagemaker_resource import SageMakerResource
+from ai21.models import Penalty
+from ai21.models.responses.completion_response import CompletionsResponse
 
 
 class SageMakerCompletion(SageMakerResource):

examples/bedrock/completion.py

@@ -1,5 +1,5 @@
 from ai21 import AI21BedrockClient, BedrockModelID
-from ai21.resources import Penalty
+from ai21.models import Penalty
 
 # Bedrock is currently supported only in us-east-1 region.
 # Either set your profile's region to us-east-1 or uncomment next line