ci: Add ready label to automate tests

.github/workflows/integration-tests.yaml

@@ -4,12 +4,6 @@ on:
   push:
     branches:
       - main
-  workflow_dispatch:
-    inputs:
-      commit-to-test:
-        description: "Commit SHA to test"
-        required: true
-        type: string
 
 env:
   POETRY_VERSION: "1.8"
@@ -26,21 +20,6 @@ jobs:
     steps:
       - name: Checkout
         uses: actions/checkout@v3
-        with:
-          ref: ${{ inputs.commit-to-test }}
-      - name: Set status to pending
-        if: ${{ github.event_name == 'workflow_dispatch' }}
-        uses: actions/github-script@v7
-        with:
-          script: |
-            github.rest.repos.createCommitStatus({
-              owner: context.repo.owner,
-              repo: context.repo.repo,
-              sha: '${{ inputs.commit-to-test }}',
-              state: 'pending',
-              context: 'Integration Tests',
-              description: 'Integration tests are in progress'
-            })
       - name: Install Poetry
         run: |
           pipx install poetry==1.8
@@ -81,29 +60,3 @@ jobs:
           path: junit/test-results-${{ matrix.python-version }}.xml
         # Use always() to always run this step to publish test results when there are test failures
         if: ${{ always() }}
-      - name: Set status to success
-        if: ${{ github.event_name == 'workflow_dispatch' && success() }}
-        uses: actions/github-script@v7
-        with:
-          script: |
-            github.rest.repos.createCommitStatus({
-              owner: context.repo.owner,
-              repo: context.repo.repo,
-              sha: '${{ inputs.commit-to-test }}',
-              state: 'success',
-              context: 'Integration Tests',
-              description: 'Integration tests passed'
-            })
-      - name: Set status to failure
-        if: ${{ github.event_name == 'workflow_dispatch' && failure() }}
-        uses: actions/github-script@v7
-        with:
-          script: |
-            github.rest.repos.createCommitStatus({
-              owner: context.repo.owner,
-              repo: context.repo.repo,
-              sha: '${{ inputs.commit-to-test }}',
-              state: 'failure',
-              context: 'Integration Tests',
-              description: 'Integration tests failed'
-            })

.github/workflows/labeler.yml

@@ -13,6 +13,7 @@ jobs:
       contents: read
       pull-requests: write
     runs-on: ubuntu-latest
+    if: github.event.action == 'synchronize'
     steps:
       - name: Checkout repository
         uses: actions/checkout@v4
@@ -112,16 +113,20 @@ jobs:
           github_token: ${{ secrets.GITHUB_TOKEN }}
           labels: ${{ env.label }}
 
-  add_lgtm_label:
-    runs-on: ubuntu-latest
-    if: github.event.review.state == 'APPROVED'
-
-    steps:
-      - name: Checkout repository
-        uses: actions/checkout@v3
-
-      - name: Add LGTM label
-        uses: actions-ecosystem/action-add-labels@v1
+      - name: Remove lgtm label on new commits
+        uses: actions/github-script@v7
         with:
-          github_token: ${{ secrets.GITHUB_TOKEN }}
-          labels: lgtm
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+          script: |
+            try {
+              await github.rest.issues.removeLabel({
+                owner: context.repo.owner,
+                repo: context.repo.repo,
+                issue_number: context.issue.number,
+                name: 'lgtm'
+              });
+            } catch (error) {
+              if (error.status !== 404) {
+                throw error;
+              }
+            }

.github/workflows/pr-approval.yml

@@ -0,0 +1,118 @@
+name: PR Approval Handler
+
+on:
+  pull_request_review:
+    types: [submitted]
+
+env:
+  POETRY_VERSION: "1.8"
+  POETRY_URL: https://install.python-poetry.org
+
+jobs:
+  handle-approval:
+    if: github.event.review.state == 'APPROVED'
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+      pull-requests: write
+      statuses: write
+
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event.pull_request.head.sha }}
+
+      - name: Add lgtm label
+        uses: actions-ecosystem/action-add-labels@v1
+        with:
+          github_token: ${{ secrets.GITHUB_TOKEN }}
+          labels: lgtm
+
+      - name: Install Poetry
+        run: |
+          pipx install poetry==1.8
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+          cache: poetry
+          cache-dependency-path: poetry.lock
+
+      - name: Set Poetry environment
+        run: |
+          poetry env use 3.11
+          poetry cache clear --all pypi
+
+      - name: Install dependencies
+        run: |
+          poetry install --all-extras
+
+      - name: Set status to pending
+        uses: actions/github-script@v7
+        with:
+          script: |
+            github.rest.repos.createCommitStatus({
+              owner: context.repo.owner,
+              repo: context.repo.repo,
+              sha: '${{ github.event.pull_request.head.sha }}',
+              state: 'pending',
+              context: 'Integration Tests',
+              description: 'Integration tests are running after approval'
+            })
+
+      - name: Run Integration Tests
+        env:
+          AI21_API_KEY: ${{ secrets.AI21_API_KEY }}
+          AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
+          AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
+        run: |
+          poetry run pytest tests/integration_tests/
+
+      - name: Set status to success
+        if: success()
+        uses: actions/github-script@v7
+        with:
+          script: |
+            github.rest.repos.createCommitStatus({
+              owner: context.repo.owner,
+              repo: context.repo.repo,
+              sha: '${{ github.event.pull_request.head.sha }}',
+              state: 'success',
+              context: 'Integration Tests',
+              description: 'Integration tests passed after approval'
+            })
+
+      - name: Set status to failure
+        if: failure()
+        uses: actions/github-script@v7
+        with:
+          script: |
+            github.rest.repos.createCommitStatus({
+              owner: context.repo.owner,
+              repo: context.repo.repo,
+              sha: '${{ github.event.pull_request.head.sha }}',
+              state: 'failure',
+              context: 'Integration Tests',
+              description: 'Integration tests failed after approval'
+            })
+
+      - name: Upload pytest integration tests results
+        uses: actions/upload-artifact@v4
+        with:
+          name: pytest-results-pr-approval
+          path: junit/test-results-*.xml
+        if: always()
+
+      - name: Auto-merge PR
+        if: success()
+        uses: actions/github-script@v7
+        with:
+          script: |
+            await github.rest.pulls.merge({
+              owner: context.repo.owner,
+              repo: context.repo.repo,
+              pull_number: context.issue.number,
+              merge_method: 'squash'
+            })

ai21/models/ai21_base_model.py

@@ -1,11 +1,14 @@
 import warnings
 from typing import Any, Dict
 
-from pydantic import BaseModel, ConfigDict
+from pydantic import BaseModel
 from typing_extensions import Self
 
 from ai21.models._pydantic_compatibility import _to_dict, _to_json, _from_dict, _from_json, IS_PYDANTIC_V2
 
+if IS_PYDANTIC_V2:
+    from pydantic import ConfigDict
+
 
 class AI21BaseModel(BaseModel):
     if IS_PYDANTIC_V2:

tests/integration_tests/clients/studio/test_maestro.py

@@ -5,7 +5,7 @@
 
 
 @pytest.mark.asyncio
-async def test_maestro__when_upload__should_return_data_sources():  # file_in_library: str):
+async def test_maestro__when_upload__should_return_data_sources(file_in_library: str):
     client = AsyncAI21Client()
     run = await client.beta.maestro.runs.create_and_poll(
         input="When did Einstein receive a Nobel Prize?",