Add wait and fixes types (#410)

* Add wait and fixes types

* Fix types

* Fix imports

* Fix type

src/cohere/client.py

@@ -4,6 +4,8 @@
 
 from .base_client import BaseCohere, AsyncBaseCohere
 from .environment import ClientEnvironment
+from .utils import wait, async_wait
+
 
 # Use NoReturn as Never type for compatibility
 Never = typing.NoReturn
@@ -25,6 +27,7 @@ def throw_if_stream_is_true(*args, **kwargs) -> None:
             "Since python sdk cohere==5.0.0, you must now use chat_stream(...) instead of chat(stream=True, ...)"
         )
 
+
 def moved_function(fn_name: str, new_fn_name: str) -> typing.Any:
     """
     This method is moved. Please update usage.
@@ -56,7 +59,7 @@ def fn(*args, **kwargs):
 class Client(BaseCohere):
     def __init__(
             self,
-            api_key: typing.Union[str, typing.Callable[[], str]],
+            api_key: typing.Optional[typing.Union[str, typing.Callable[[], str]]] = None,
             *,
             base_url: typing.Optional[str] = None,
             environment: ClientEnvironment = ClientEnvironment.PRODUCTION,
@@ -76,6 +79,8 @@ def __init__(
 
         validate_args(self, "chat", throw_if_stream_is_true)
 
+    wait = wait
+
     """
     The following methods have been moved or deprecated in cohere==5.0.0. Please update your usage.
     Issues may be filed in https://github.com/cohere-ai/cohere-python/issues.
@@ -125,7 +130,7 @@ def __init__(
 class AsyncClient(AsyncBaseCohere):
     def __init__(
             self,
-            api_key: typing.Union[str, typing.Callable[[], str]],
+            api_key: typing.Optional[typing.Union[str, typing.Callable[[], str]]] = None,
             *,
             base_url: typing.Optional[str] = None,
             environment: ClientEnvironment = ClientEnvironment.PRODUCTION,
@@ -145,6 +150,8 @@ def __init__(
 
         validate_args(self, "chat", throw_if_stream_is_true)
 
+    wait = async_wait
+
     """
     The following methods have been moved or deprecated in cohere==5.0.0. Please update your usage.
     Issues may be filed in https://github.com/cohere-ai/cohere-python/issues.

src/cohere/utils.py

@@ -0,0 +1,152 @@
+import asyncio
+import time
+import typing
+from typing import Optional
+
+from .types import EmbedJob, CreateEmbedJobResponse
+from .datasets import DatasetsCreateResponse, DatasetsGetResponse
+
+
+def get_terminal_states():
+    return get_success_states() | get_failed_states()
+
+
+def get_success_states():
+    return {"complete", "validated"}
+
+
+def get_failed_states():
+    return {"unknown", "failed", "skipped", "cancelled", "failed"}
+
+
+def get_id(
+        awaitable: typing.Union[CreateEmbedJobResponse, DatasetsCreateResponse, EmbedJob, DatasetsGetResponse]):
+    return getattr(awaitable, "job_id", None) or getattr(awaitable, "id", None) or getattr(
+        getattr(awaitable, "dataset", None), "id", None)
+
+
+def get_validation_status(awaitable: typing.Union[EmbedJob, DatasetsGetResponse]):
+    return getattr(awaitable, "status", None) or getattr(getattr(awaitable, "dataset", None), "validation_status", None)
+
+
+def get_job(cohere: typing.Any,
+            awaitable: typing.Union[CreateEmbedJobResponse, DatasetsCreateResponse, EmbedJob, DatasetsGetResponse]) -> \
+typing.Union[
+    EmbedJob, DatasetsGetResponse]:
+    if awaitable.__class__.__name__ == "EmbedJob" or awaitable.__class__.__name__ == "CreateEmbedJobResponse":
+        return cohere.embed_jobs.get(id=get_id(awaitable))
+    elif awaitable.__class__.__name__ == "DatasetsGetResponse" or awaitable.__class__.__name__ == "DatasetsCreateResponse":
+        return cohere.datasets.get(id=get_id(awaitable))
+    else:
+        raise ValueError(f"Unexpected awaitable type {awaitable}")
+
+
+async def async_get_job(cohere: typing.Any, awaitable: typing.Union[CreateEmbedJobResponse, DatasetsCreateResponse]) -> \
+        typing.Union[
+            EmbedJob, DatasetsGetResponse]:
+    if awaitable.__class__.__name__ == "EmbedJob" or awaitable.__class__.__name__ == "CreateEmbedJobResponse":
+        return await cohere.embed_jobs.get(id=get_id(awaitable))
+    elif awaitable.__class__.__name__ == "DatasetsGetResponse" or awaitable.__class__.__name__ == "DatasetsCreateResponse":
+        return await cohere.datasets.get(id=get_id(awaitable))
+    else:
+        raise ValueError(f"Unexpected awaitable type {awaitable}")
+
+
+def get_failure_reason(job: typing.Union[EmbedJob, DatasetsGetResponse]) -> Optional[str]:
+    if isinstance(job, EmbedJob):
+        return f"Embed job {job.job_id} failed with status {job.status}"
+    elif isinstance(job, DatasetsGetResponse):
+        return f"Dataset creation {job.dataset.validation_status} failed with status {job.dataset.validation_status}"
+    return None
+
+
+@typing.overload
+def wait(
+        cohere: typing.Any,
+        awaitable: CreateEmbedJobResponse,
+        timeout: Optional[float] = None,
+        interval: float = 10,
+) -> EmbedJob:
+    ...
+
+
+@typing.overload
+def wait(
+        cohere: typing.Any,
+        awaitable: DatasetsCreateResponse,
+        timeout: Optional[float] = None,
+        interval: float = 10,
+) -> DatasetsGetResponse:
+    ...
+
+
+def wait(
+        cohere: typing.Any,
+        awaitable: typing.Union[CreateEmbedJobResponse, DatasetsCreateResponse],
+        timeout: Optional[float] = None,
+        interval: float = 2,
+) -> typing.Union[EmbedJob, DatasetsGetResponse]:
+    start_time = time.time()
+    terminal_states = get_terminal_states()
+    failed_states = get_failed_states()
+
+    job = get_job(cohere, awaitable)
+    while get_validation_status(job) not in terminal_states:
+        if timeout is not None and time.time() - start_time > timeout:
+            raise TimeoutError(f"wait timed out after {timeout} seconds")
+
+        time.sleep(interval)
+        print("...")
+
+        job = get_job(cohere, awaitable)
+
+    if get_validation_status(job) in failed_states:
+        raise Exception(get_failure_reason(job))
+
+    return job
+
+
+@typing.overload
+async def async_wait(
+        cohere: typing.Any,
+        awaitable: CreateEmbedJobResponse,
+        timeout: Optional[float] = None,
+        interval: float = 10,
+) -> EmbedJob:
+    ...
+
+
+@typing.overload
+async def async_wait(
+        cohere: typing.Any,
+        awaitable: DatasetsCreateResponse,
+        timeout: Optional[float] = None,
+        interval: float = 10,
+) -> DatasetsGetResponse:
+    ...
+
+
+async def async_wait(
+        cohere: typing.Any,
+        awaitable: typing.Union[CreateEmbedJobResponse, DatasetsCreateResponse],
+        timeout: Optional[float] = None,
+        interval: float = 10,
+) -> typing.Union[EmbedJob, DatasetsGetResponse]:
+    start_time = time.time()
+    terminal_states = get_terminal_states()
+    failed_states = get_failed_states()
+
+    job = await async_get_job(cohere, awaitable)
+    while get_validation_status(job) not in terminal_states:
+        if timeout is not None and time.time() - start_time > timeout:
+            raise TimeoutError(f"wait timed out after {timeout} seconds")
+
+        await asyncio.sleep(interval)
+        print("...")
+
+        job = await async_get_job(cohere, awaitable)
+
+    if get_validation_status(job) in failed_states:
+        raise Exception(get_failure_reason(job))
+
+    return job