feat: ✨ add support for status polling for Run

ai21/clients/common/beta/assistant/runs.py

@@ -65,3 +65,23 @@ def cancel(
     @abstractmethod
     def submit_tool_outputs(self, *, thread_id: str, run_id: str, tool_outputs: List[ToolOutput]) -> RunResponse:
         pass
+
+    @abstractmethod
+    def _poll_for_status(
+        self, *, thread_id: str, run_id: str, poll_interval: float, poll_timeout: float
+    ) -> RunResponse:
+        pass
+
+    @abstractmethod
+    def create_and_poll(
+        self,
+        *,
+        thread_id: str,
+        assistant_id: str,
+        description: str | NotGiven,
+        optimization: Optimization | NotGiven,
+        poll_interval_sec: float,
+        poll_timeout_sec: float,
+        **kwargs,
+    ) -> RunResponse:
+        pass

ai21/clients/studio/resources/beta/assistant/thread_runs.py

@@ -1,11 +1,18 @@
 from __future__ import annotations
 
+import asyncio
+import time
 from typing import List
 
 from ai21.clients.common.beta.assistant.runs import BaseRuns
 from ai21.clients.studio.resources.studio_resource import StudioResource, AsyncStudioResource
 from ai21.models.assistant.assistant import Optimization
-from ai21.models.assistant.run import ToolOutput
+from ai21.models.assistant.run import (
+    ToolOutput,
+    TERMINATED_RUN_STATUSES,
+    DEFAULT_RUN_POLL_INTERVAL,
+    DEFAULT_RUN_POLL_TIMEOUT,
+)
 from ai21.models.responses.run_response import RunResponse
 from ai21.types import NotGiven, NOT_GIVEN
 
@@ -55,6 +62,41 @@ def submit_tool_outputs(self, *, thread_id: str, run_id: str, tool_outputs: List
             response_cls=RunResponse,
         )
 
+    def _poll_for_status(
+        self, *, thread_id: str, run_id: str, poll_interval: float, poll_timeout: float
+    ) -> RunResponse:
+        start_time = time.time()
+
+        while True:
+            run = self.retrieve(thread_id=thread_id, run_id=run_id)
+
+            if run.status in TERMINATED_RUN_STATUSES:
+                return run
+
+            if (time.time() - start_time) > poll_timeout:
+                return run
+
+            time.sleep(poll_interval)
+
+    def create_and_poll(
+        self,
+        *,
+        thread_id: str,
+        assistant_id: str,
+        description: str | NotGiven = NOT_GIVEN,
+        optimization: Optimization | NotGiven = NOT_GIVEN,
+        poll_interval_sec: float = DEFAULT_RUN_POLL_INTERVAL,
+        poll_timeout_sec: float = DEFAULT_RUN_POLL_TIMEOUT,
+        **kwargs,
+    ) -> RunResponse:
+        run = self.create(
+            thread_id=thread_id, assistant_id=assistant_id, description=description, optimization=optimization, **kwargs
+        )
+
+        return self._poll_for_status(
+            thread_id=thread_id, run_id=run.id, poll_interval=poll_interval_sec, poll_timeout=poll_timeout_sec
+        )
+
 
 class AsyncThreadRuns(AsyncStudioResource, BaseRuns):
     async def create(
@@ -102,3 +144,38 @@ async def submit_tool_outputs(self, *, thread_id: str, run_id: str, tool_outputs
             body=body,
             response_cls=RunResponse,
         )
+
+    async def _poll_for_status(
+        self, *, thread_id: str, run_id: str, poll_interval: float, poll_timeout: float
+    ) -> RunResponse:
+        start_time = time.time()
+
+        while True:
+            run = await self.retrieve(thread_id=thread_id, run_id=run_id)
+
+            if run.status in TERMINATED_RUN_STATUSES:
+                return run
+
+            if (time.time() - start_time) > poll_timeout:
+                return run
+
+            await asyncio.sleep(poll_interval)
+
+    async def create_and_poll(
+        self,
+        *,
+        thread_id: str,
+        assistant_id: str,
+        description: str | NotGiven = NOT_GIVEN,
+        optimization: Optimization | NotGiven = NOT_GIVEN,
+        poll_interval_sec: float = DEFAULT_RUN_POLL_INTERVAL,
+        poll_timeout_sec: float = DEFAULT_RUN_POLL_TIMEOUT,
+        **kwargs,
+    ) -> RunResponse:
+        run = await self.create(
+            thread_id=thread_id, assistant_id=assistant_id, description=description, optimization=optimization, **kwargs
+        )
+
+        return await self._poll_for_status(
+            thread_id=thread_id, run_id=run.id, poll_interval=poll_interval_sec, poll_timeout=poll_timeout_sec
+        )

ai21/models/assistant/run.py

@@ -1,4 +1,4 @@
-from typing import Literal, Any, List
+from typing import Literal, Any, List, Set
 
 from typing_extensions import TypedDict
 
@@ -15,6 +15,10 @@
     "requires_action",
 ]
 
+TERMINATED_RUN_STATUSES: Set[RunStatus] = {"completed", "failed", "expired", "cancelled", "requires_action"}
+DEFAULT_RUN_POLL_INTERVAL: float = 1  # seconds
+DEFAULT_RUN_POLL_TIMEOUT: float = 60  # seconds
+
 
 class ToolOutput(TypedDict):
     tool_call_id: str

examples/studio/assistant/assistant.py

@@ -1,9 +1,5 @@
-import time
-
 from ai21 import AI21Client
 
-TIMEOUT = 20
-
 
 def main():
     ai21_client = AI21Client()
@@ -19,25 +15,17 @@ def main():
         ]
     )
 
-    run = ai21_client.beta.threads.runs.create(
+    run = ai21_client.beta.threads.runs.create_and_poll(
         thread_id=thread.id,
         assistant_id=assistant.id,
     )
 
-    start = time.time()
-
-    while run.status == "in_progress":
-        run = ai21_client.beta.threads.runs.retrieve(thread_id=thread.id, run_id=run.id)
-        if time.time() - start > TIMEOUT:
-            break
-        time.sleep(1)
-
     if run.status == "completed":
         messages = ai21_client.beta.threads.messages.list(thread_id=thread.id)
         print("Messages:")
         print("\n".join(f"{msg.role}: {msg.content['text']}" for msg in messages.results))
     else:
-        raise Exception(f"Run failed. Status: {run.status}")
+        print(f"Run status: {run.status}")
 
 
 if __name__ == "__main__":

examples/studio/assistant/async_assistant.py

@@ -1,12 +1,8 @@
 import asyncio
-import time
 
 from ai21 import AsyncAI21Client
 
 
-TIMEOUT = 20
-
-
 async def main():
     ai21_client = AsyncAI21Client()
 
@@ -21,25 +17,19 @@ async def main():
         ]
     )
 
-    run = await ai21_client.beta.threads.runs.create(
+    run = await ai21_client.beta.threads.runs.create_and_poll(
         thread_id=thread.id,
         assistant_id=assistant.id,
     )
 
-    start = time.time()
-
-    while run.status == "in_progress":
-        run = await ai21_client.beta.threads.runs.retrieve(thread_id=thread.id, run_id=run.id)
-        if time.time() - start > TIMEOUT:
-            break
-        time.sleep(1)
-
     if run.status == "completed":
         messages = await ai21_client.beta.threads.messages.list(thread_id=thread.id)
         print("Messages:")
         print("\n".join(f"{msg.role}: {msg.content['text']}" for msg in messages.results))
-    else:
+    elif run.status == "failed":
         raise Exception(f"Run failed. Status: {run.status}")
+    else:
+        print(f"Run status: {run.status}")
 
 
 if __name__ == "__main__":