feat: 🎉 Added support for batch mode & batch querying

Author: shirbiton

ai21/clients/studio/resources/studio_library.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import uuid
 from os import PathLike
 from typing import List, Optional
 
@@ -9,7 +10,7 @@
 )
 from ai21.http_client.async_http_client import AsyncAI21HTTPClient
 from ai21.http_client.http_client import AI21HTTPClient
-from ai21.models import FileResponse
+from ai21.models import FileResponse, UploadMode, BatchStatusResponse
 from ai21.types import NOT_GIVEN, NotGiven
 from ai21.utils.typing import remove_not_given
 
@@ -32,10 +33,21 @@ def create(
         path: Optional[str] | NotGiven = NOT_GIVEN,
         labels: Optional[List[str]] | NotGiven = NOT_GIVEN,
         public_url: Optional[str] | NotGiven = NOT_GIVEN,
+        upload_mode: Optional[UploadMode] | NotGiven = NOT_GIVEN,
+        batch_id: Optional[uuid.UUID] | NotGiven = NOT_GIVEN,
         **kwargs,
     ) -> str:
         files = {"file": open(file_path, "rb")}
-        body = remove_not_given({"path": path, "labels": labels, "publicUrl": public_url, **kwargs})
+        body = remove_not_given(
+            {
+                "path": path,
+                "labels": labels,
+                "publicUrl": public_url,
+                "upload_mode": upload_mode,
+                "batch_id": batch_id,
+                **kwargs,
+            }
+        )
 
         raw_response = self._post(path=f"/{self._module_name}", files=files, body=body, response_cls=dict)
 
@@ -44,6 +56,9 @@ def create(
     def get(self, file_id: str) -> FileResponse:
         return self._get(path=f"/{self._module_name}/{file_id}", response_cls=FileResponse)
 
+    def get_batch_status(self, batch_id: uuid.UUID) -> BatchStatusResponse:
+        return self._get(path=f"/library/batches/{batch_id}/status", response_cls=BatchStatusResponse)
+
     def list(
         self,
         *,
@@ -94,10 +109,21 @@ async def create(
         path: Optional[str] | NotGiven = NOT_GIVEN,
         labels: Optional[List[str]] | NotGiven = NOT_GIVEN,
         public_url: Optional[str] | NotGiven = NOT_GIVEN,
+        upload_mode: Optional[UploadMode] | NotGiven = NOT_GIVEN,
+        batch_id: Optional[uuid.UUID] | NotGiven = NOT_GIVEN,
         **kwargs,
     ) -> str:
         files = {"file": open(file_path, "rb")}
-        body = remove_not_given({"path": path, "labels": labels, "publicUrl": public_url, **kwargs})
+        body = remove_not_given(
+            {
+                "path": path,
+                "labels": labels,
+                "publicUrl": public_url,
+                "upload_mode": upload_mode,
+                "batch_id": batch_id,
+                **kwargs,
+            }
+        )
 
         raw_response = await self._post(path=f"/{self._module_name}", files=files, body=body, response_cls=dict)
 
@@ -106,6 +132,9 @@ async def create(
     async def get(self, file_id: str) -> FileResponse:
         return await self._get(path=f"/{self._module_name}/{file_id}", response_cls=FileResponse)
 
+    async def get_batch_status(self, batch_id: uuid.UUID) -> BatchStatusResponse:
+        return await self._get(path=f"/library/batches/{batch_id}/status", response_cls=BatchStatusResponse)
+
     async def list(
         self,
         *,

ai21/models/__init__.py

@@ -8,6 +8,7 @@
     ConversationalRagSource,
 )
 from ai21.models.responses.file_response import FileResponse
+from ai21.models.upload_mode import UploadMode, BatchStatusResponse
 
 __all__ = [
     "ChatMessage",
@@ -20,4 +21,6 @@
     "FileResponse",
     "ConversationalRagResponse",
     "ConversationalRagSource",
+    "UploadMode",
+    "BatchStatusResponse",
 ]

ai21/models/upload_mode.py

@@ -0,0 +1,20 @@
+from enum import StrEnum
+from typing import List
+
+from pydantic import BaseModel, Field
+
+
+class UploadMode(StrEnum):
+    IMMEDIATE = "immediate"
+    BATCH = "batch"
+
+
+class BatchStatusCount(BaseModel):
+    status: str
+    count: int
+
+
+class BatchStatusResponse(BaseModel):
+    batch_id: str = Field(description="The UUID of the batch")
+    total_documents: int = Field(description="Total number of documents in the batch")
+    statuses: List[BatchStatusCount] = Field(description="List of document counts by status")

init.sh

@@ -1,38 +1,48 @@
-#!/bin/bash
-
-cd "$(dirname "$0")" || exit
+#!/usr/bin/env bash
 
 # create .git folder
 if [[ ! -d .git ]]; then
   git init
 fi
 
-# install the python version specified in .python-version, if not already installed
-if pyenv --version; then
-  pyenv install --skip-existing
-fi
-
-# install poetry if not already installed
-if ! poetry --version; then
-  brew install poetry
-fi
-
-# poetry needs to create the venv with the same python version
-poetry env use "$(cat .python-version)"
-
-# update lock file
-poetry lock --no-update
-
-# install dependencies
-poetry install
+PYTHON_VERSION=$(cat .python-version)
 
-# install pre-commit if not already installed
-if ! pre-commit --version; then
-  brew install pre-commit
+# install the python version specified in .python-version, if not already installed
+if ! pyenv versions --bare | grep -q "^${PYTHON_VERSION}"; then
+  pyenv install "${PYTHON_VERSION}" --skip-existing
 fi
 
-# install pre-commit hooks
-pre-commit install --install-hooks -t pre-commit -t commit-msg
-
-# shellcheck source=/dev/null
-source .venv/bin/activate
+{ [[ -d .venv ]] || {
+  echo 'creating virtualenv...'
+  python -m venv .venv
+}; } && {
+  # shellcheck disable=SC1091
+  . .venv/bin/activate
+} && {
+  # install poetry if not already installed, or upgrade to version 2.x
+  poetry_version_installed=$(poetry --version 2> /dev/null || true)
+  poetry_major_version=$(echo "$poetry_version_installed" | awk '{print $3}' | cut -d '.' -f1)
+
+  # Check if poetry_major_version is a number and if it's >= 2
+  if [[ -z $poetry_version_installed || ! $poetry_major_version =~ ^[0-9]+$ || $poetry_major_version -lt 2 ]]; then
+    echo "Poetry >=2 is not installed. Installing/upgrading..."
+    pip install --upgrade poetry || brew install poetry
+  fi
+} && {
+  # install keyring
+  poetry self add keyrings-google-artifactregistry-auth@1.1.2
+} && {
+  # update lock file
+  poetry lock
+} && {
+  # install dependencies
+  poetry install --no-root && poetry sync
+}
+
+{
+  # install pre-commit if not already installed
+  pre-commit --version || brew install pre-commit
+} && {
+  # install pre-commit hooks
+  pre-commit install --install-hooks -t pre-commit -t pre-push -t commit-msg
+}