feat: add submit_input method for thread run (#250)

* feat: add submit_input method for thread run

* fix: change gt to gte operator in poll_for_status

* fix: add =None to optionals

Author: guyl25

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

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import List
+from typing import Any, List
 
 from ai21.models.assistant.assistant import Optimization
 from ai21.models.assistant.run import ToolOutput
@@ -85,3 +85,19 @@ def create_and_poll(
         **kwargs,
     ) -> RunResponse:
         pass
+
+    @abstractmethod
+    def submit_input(self, *, thread_id: str, run_id: str, input: Any) -> RunResponse:
+        pass
+
+    @abstractmethod
+    def submit_input_and_poll(
+        self,
+        *,
+        thread_id: str,
+        run_id: str,
+        input: Any,
+        poll_interval_sec: float,
+        poll_timeout_sec: float,
+    ) -> RunResponse:
+        pass

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

@@ -2,7 +2,7 @@
 
 import asyncio
 import time
-from typing import List
+from typing import Any, List
 
 from ai21.clients.common.beta.assistant.runs import BaseRuns
 from ai21.clients.studio.resources.studio_resource import StudioResource, AsyncStudioResource
@@ -73,7 +73,7 @@ def _poll_for_status(
             if run.status in TERMINATED_RUN_STATUSES:
                 return run
 
-            if (time.time() - start_time) > poll_timeout:
+            if (time.time() - start_time) >= poll_timeout:
                 return run
 
             time.sleep(poll_interval)
@@ -97,6 +97,30 @@ def create_and_poll(
             thread_id=thread_id, run_id=run.id, poll_interval=poll_interval_sec, poll_timeout=poll_timeout_sec
         )
 
+    def submit_input(self, *, thread_id: str, run_id: str, input: Any) -> RunResponse:
+        body = dict(input=input)
+
+        return self._post(
+            path=f"/threads/{thread_id}/{self._module_name}/{run_id}/submit_input",
+            body=body,
+            response_cls=RunResponse,
+        )
+
+    def submit_input_and_poll(
+        self,
+        *,
+        thread_id: str,
+        run_id: str,
+        input: Any,
+        poll_interval_sec: float = DEFAULT_RUN_POLL_INTERVAL,
+        poll_timeout_sec: float = DEFAULT_RUN_POLL_TIMEOUT,
+    ) -> RunResponse:
+        run = self.submit_input(thread_id=thread_id, run_id=run_id, input=input)
+
+        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(
@@ -156,7 +180,7 @@ async def _poll_for_status(
             if run.status in TERMINATED_RUN_STATUSES:
                 return run
 
-            if (time.time() - start_time) > poll_timeout:
+            if (time.time() - start_time) >= poll_timeout:
                 return run
 
             await asyncio.sleep(poll_interval)
@@ -179,3 +203,27 @@ async def create_and_poll(
         return await self._poll_for_status(
             thread_id=thread_id, run_id=run.id, poll_interval=poll_interval_sec, poll_timeout=poll_timeout_sec
         )
+
+    async def submit_input(self, *, thread_id: str, run_id: str, input: Any) -> RunResponse:
+        body = dict(input=input)
+
+        return await self._post(
+            path=f"/threads/{thread_id}/{self._module_name}/{run_id}/submit_inputs",
+            body=body,
+            response_cls=RunResponse,
+        )
+
+    async def submit_input_and_poll(
+        self,
+        *,
+        thread_id: str,
+        run_id: str,
+        input: Any,
+        poll_interval_sec: float = DEFAULT_RUN_POLL_INTERVAL,
+        poll_timeout_sec: float = DEFAULT_RUN_POLL_TIMEOUT,
+    ) -> RunResponse:
+        run = await self.submit_input(thread_id=thread_id, run_id=run_id, input=input)
+
+        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, Set
+from typing import Dict, Literal, Any, List, Set, Optional
 
 from typing_extensions import TypedDict
 
@@ -40,6 +40,12 @@ class SubmitToolCallOutputs(TypedDict):
     tool_calls: List[ToolOutput]
 
 
+class SubmitInput(TypedDict):
+    event_name: str
+    data: Dict[str, Any]
+
+
 class RequiredAction(TypedDict):
-    type: Literal["submit_tool_outputs"]
-    submit_tool_outputs: SubmitToolCallOutputs
+    type: Literal["submit_tool_outputs", "submit_input"]
+    submit_tool_outputs: Optional[SubmitToolCallOutputs] = None
+    submit_input: Optional[SubmitInput] = None