feat: 🎉 Added support for batch mode & batch querying (#302)

feat: 🎉 Added support for batch mode & batch querying

Author: shirbiton

ai21/clients/studio/resources/studio_library.py

@@ -9,7 +9,8 @@
 )
 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, BatchStatusResponse
+from ai21.models.upload_mode import UploadMode
 from ai21.types import NOT_GIVEN, NotGiven
 from ai21.utils.typing import remove_not_given
 
@@ -24,6 +25,7 @@ def __init__(self, client: AI21HTTPClient):
 
 class LibraryFiles(StudioResource):
     _module_name = "library/files"
+    _sub_module_name = "library"
 
     def create(
         self,
@@ -32,10 +34,22 @@ def create(
         path: Optional[str] | NotGiven = NOT_GIVEN,
         labels: Optional[List[str]] | NotGiven = NOT_GIVEN,
         public_url: Optional[str] | NotGiven = NOT_GIVEN,
+        batch_id: Optional[str] | 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,
+                "batch_id": batch_id,
+                **kwargs,
+            }
+        )
+
+        if body.get("batch_id"):
+            body["upload_mode"] = UploadMode.BATCH
 
         raw_response = self._post(path=f"/{self._module_name}", files=files, body=body, response_cls=dict)
 
@@ -44,6 +58,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: str) -> BatchStatusResponse:
+        return self._get(path=f"/{self._sub_module_name}/batches/{batch_id}/status", response_cls=BatchStatusResponse)
+
     def list(
         self,
         *,
@@ -86,6 +103,7 @@ def __init__(self, client: AsyncAI21HTTPClient):
 
 class AsyncLibraryFiles(AsyncStudioResource):
     _module_name = "library/files"
+    _sub_module_name = "library"
 
     async def create(
         self,
@@ -94,10 +112,22 @@ async def create(
         path: Optional[str] | NotGiven = NOT_GIVEN,
         labels: Optional[List[str]] | NotGiven = NOT_GIVEN,
         public_url: Optional[str] | NotGiven = NOT_GIVEN,
+        batch_id: Optional[str] | 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,
+                "batch_id": batch_id,
+                **kwargs,
+            }
+        )
+
+        if body.get("batch_id"):
+            body["upload_mode"] = UploadMode.BATCH
 
         raw_response = await self._post(path=f"/{self._module_name}", files=files, body=body, response_cls=dict)
 
@@ -106,6 +136,11 @@ 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: str) -> BatchStatusResponse:
+        return await self._get(
+            path=f"/{self._sub_module_name}/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 BatchStatusResponse
 
 __all__ = [
     "ChatMessage",
@@ -20,4 +21,5 @@
     "FileResponse",
     "ConversationalRagResponse",
     "ConversationalRagSource",
+    "BatchStatusResponse",
 ]

ai21/models/upload_mode.py

@@ -0,0 +1,19 @@
+from enum import Enum
+from typing import List
+
+from pydantic import BaseModel, Field
+
+
+class UploadMode(str, Enum):
+    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,45 @@
-#!/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
+} && {
+  # 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
+}