clarifai 11.5.1__py3-none-any.whl → 11.5.3__py3-none-any.whl

clarifai/client/auth/stub.py

@@ -1,3 +1,4 @@
+import asyncio
 import itertools
 import time
 from concurrent.futures import ThreadPoolExecutor
@@ -31,6 +32,10 @@ def validate_response(response, attempt, max_attempts):
             else:
                 return response
 
+        # Check if response is an async iterator
+        if hasattr(response, '__aiter__'):
+            return response  # Return async iterator directly for handling in _async_call__
+
     # Check if the response is an instance of a gRPC streaming call
     if isinstance(response, grpc._channel._MultiThreadedRendezvous):
         try:
@@ -49,7 +54,9 @@ def validate_response(response, attempt, max_attempts):
         return handle_simple_response(response)
 
 
-def create_stub(auth_helper: ClarifaiAuthHelper = None, max_retry_attempts: int = 10) -> V2Stub:
+def create_stub(
+    auth_helper: ClarifaiAuthHelper = None, max_retry_attempts: int = 10, is_async: bool = False
+) -> V2Stub:
     """
     Create client stub that handles authorization and basic retries for
     unavailable or throttled connections.
@@ -58,41 +65,48 @@ def create_stub(auth_helper: ClarifaiAuthHelper = None, max_retry_attempts: int
       auth_helper:  ClarifaiAuthHelper to use for auth metadata (default: from env)
       max_retry_attempts:  max attempts to retry rpcs with retryable failures
     """
-    stub = AuthorizedStub(auth_helper)
+    stub = AuthorizedStub(auth_helper, is_async=is_async)
     if max_retry_attempts > 0:
-        return RetryStub(stub, max_retry_attempts)
+        return RetryStub(stub, max_retry_attempts, is_async=is_async)
     return stub
 
 
 class AuthorizedStub(V2Stub):
     """V2Stub proxy that inserts metadata authorization in rpc calls."""
 
-    def __init__(self, auth_helper: ClarifaiAuthHelper = None):
+    def __init__(self, auth_helper: ClarifaiAuthHelper = None, is_async: bool = False):
         if auth_helper is None:
             auth_helper = ClarifaiAuthHelper.from_env()
-        self.stub = auth_helper.get_stub()
+        self.is_async = is_async
+        self.stub = auth_helper.get_async_stub() if is_async else auth_helper.get_stub()
         self.metadata = auth_helper.metadata
 
     def __getattr__(self, name):
         value = getattr(self.stub, name)
         if isinstance(value, RpcCallable):
-            value = _AuthorizedRpcCallable(value, self.metadata)
+            value = _AuthorizedRpcCallable(value, self.metadata, self.is_async)
         return value
 
 
 class _AuthorizedRpcCallable(RpcCallable):
     """Adds metadata(authorization header) to rpc calls"""
 
-    def __init__(self, func, metadata):
+    def __init__(self, func, metadata, is_async):
         self.f = func
         self.metadata = metadata
+        self.is_async = is_async or asyncio.iscoroutinefunction(func)
 
     def __repr__(self):
         return repr(self.f)
 
     def __call__(self, *args, **kwargs):
         metadata = kwargs.pop('metadata', self.metadata)
-        return self.f(*args, **kwargs, metadata=metadata)
+
+        return self.f(
+            *args,
+            **kwargs,
+            metadata=metadata,
+        )
 
     def future(self, *args, **kwargs):
         metadata = kwargs.pop('metadata', self.metadata)
@@ -107,30 +121,109 @@ class RetryStub(V2Stub):
     V2Stub proxy that retries requests (currently on unavailable server or throttle codes)
     """
 
-    def __init__(self, stub, max_attempts=10, backoff_time=5):
+    def __init__(self, stub, max_attempts=10, backoff_time=5, is_async=False):
         self.stub = stub
         self.max_attempts = max_attempts
         self.backoff_time = backoff_time
+        self.is_async = is_async
 
     def __getattr__(self, name):
         value = getattr(self.stub, name)
         if isinstance(value, RpcCallable):
-            value = _RetryRpcCallable(value, self.max_attempts, self.backoff_time)
+            value = _RetryRpcCallable(value, self.max_attempts, self.backoff_time, self.is_async)
         return value
 
 
 class _RetryRpcCallable(RpcCallable):
     """Retries rpc calls on unavailable server or throttle codes"""
 
-    def __init__(self, func, max_attempts, backoff_time):
+    def __init__(self, func, max_attempts, backoff_time, is_async=False):
         self.f = func
         self.max_attempts = max_attempts
         self.backoff_time = backoff_time
+        self.is_async = is_async or asyncio.iscoroutinefunction(func)
 
     def __repr__(self):
         return repr(self.f)
 
-    def __call__(self, *args, **kwargs):
+    async def _async_call__(self, *args, **kwargs):
+        """Handle async RPC calls with retries and validation"""
+        for attempt in range(1, self.max_attempts + 1):
+            if attempt != 1:
+                await asyncio.sleep(self.backoff_time)
+
+            try:
+                response = self.f(*args, **kwargs)
+
+                # Handle streaming response
+                if hasattr(response, '__aiter__'):
+                    return await self._handle_streaming_response(response, attempt)
+
+                # Handle regular async response
+                validated_response = await self._handle_regular_response(response, attempt)
+                if validated_response is not None:
+                    return validated_response
+
+            except grpc.RpcError as e:
+                if not self._should_retry(e, attempt):
+                    raise
+                logger.debug(
+                    f'Retrying after error {e.code()} (attempt {attempt}/{self.max_attempts})'
+                )
+
+        raise Exception(f'Max attempts reached ({self.max_attempts}) without success')
+
+    async def _handle_streaming_response(self, response, attempt):
+        """Handle streaming response validation and processing"""
+
+        async def validated_stream():
+            try:
+                async for item in response:
+                    if not self._is_valid_response(item):
+                        if attempt < self.max_attempts:
+                            yield None  # Signal for retry
+                        raise Exception(
+                            f'Validation failed on streaming response (attempt {attempt})'
+                        )
+                    yield item
+            except grpc.RpcError as e:
+                if not self._should_retry(e, attempt):
+                    raise
+                yield None  # Signal for retry
+
+        return validated_stream()
+
+    async def _handle_regular_response(self, response, attempt):
+        """Handle regular async response validation and processing"""
+        try:
+            result = await response
+            if not self._is_valid_response(result):
+                if attempt < self.max_attempts:
+                    return None  # Signal for retry
+                raise Exception(f'Validation failed on response (attempt {attempt})')
+            return result
+        except grpc.RpcError as e:
+            if not self._should_retry(e, attempt):
+                raise
+            return None  # Signal for retry
+
+    def _is_valid_response(self, response):
+        """Check if response status is valid"""
+        return not (
+            hasattr(response, 'status')
+            and hasattr(response.status, 'code')
+            and response.status.code in throttle_status_codes
+        )
+
+    def _should_retry(self, error, attempt):
+        """Determine if we should retry based on error and attempt count"""
+        return (
+            isinstance(error, grpc.RpcError)
+            and error.code() in retry_codes_grpc
+            and attempt < self.max_attempts
+        )
+
+    def _sync_call__(self, *args, **kwargs):
         attempt = 0
         while attempt < self.max_attempts:
             attempt += 1
@@ -147,8 +240,19 @@ class _RetryRpcCallable(RpcCallable):
                 else:
                     raise
 
+    def __call__(self, *args, **kwargs):
+        if self.is_async:
+            return self._async_call__(*args, **kwargs)
+        return self._sync_call__(*args, **kwargs)
+
+    async def __call_async__(self, *args, **kwargs):
+        """Explicit async call method"""
+        return await self._async_call(*args, **kwargs)
+
     def future(self, *args, **kwargs):
         # TODO use single result event loop thread with asyncio
+        if self.is_async:
+            return asyncio.create_task(self._async_call(*args, **kwargs))
         return _threadpool.submit(self, *args, **kwargs)
 
     def __getattr__(self, name):

clarifai/client/base.py

@@ -52,6 +52,7 @@ class BaseClient:
 
         self.auth_helper = ClarifaiAuthHelper(**kwargs, validate=False)
         self.STUB = create_stub(self.auth_helper)
+        self._async_stub = None
         self.metadata = self.auth_helper.metadata
         self.pat = self.auth_helper.pat
         self.token = self.auth_helper._token
@@ -59,6 +60,14 @@ class BaseClient:
         self.base = self.auth_helper.base
         self.root_certificates_path = self.auth_helper._root_certificates_path
 
+    @property
+    def async_stub(self):
+        """Returns the asynchronous gRPC stub for the API interaction.
+        Lazy initialization of async stub"""
+        if self._async_stub is None:
+            self._async_stub = create_stub(self.auth_helper, is_async=True)
+        return self._async_stub
+
     @classmethod
     def from_env(cls, validate: bool = False):
         auth = ClarifaiAuthHelper.from_env(validate=validate)

clarifai/client/model.py

@@ -104,11 +104,6 @@ class Model(Lister, BaseClient):
         self.input_types = None
         self._client = None
         self._added_methods = False
-        self._set_runner_selector(
-            compute_cluster_id=compute_cluster_id,
-            nodepool_id=nodepool_id,
-            deployment_id=deployment_id,
-        )
         BaseClient.__init__(
             self,
             user_id=self.user_id,
@@ -120,6 +115,12 @@ class Model(Lister, BaseClient):
         )
         Lister.__init__(self)
 
+        self._set_runner_selector(
+            compute_cluster_id=compute_cluster_id,
+            nodepool_id=nodepool_id,
+            deployment_id=deployment_id,
+        )
+
     @classmethod
     def from_current_context(cls, **kwargs) -> 'Model':
         from clarifai.urls.helper import ClarifaiUrlHelper
@@ -502,7 +503,9 @@ class Model(Lister, BaseClient):
                 model=self.model_info,
                 runner_selector=self._runner_selector,
             )
-            self._client = ModelClient(self.STUB, request_template=request_template)
+            self._client = ModelClient(
+                stub=self.STUB, async_stub=self.async_stub, request_template=request_template
+            )
         return self._client
 
     def predict(self, *args, **kwargs):
@@ -530,6 +533,35 @@ class Model(Lister, BaseClient):
 
         return self.client.predict(*args, **kwargs)
 
+    async def async_predict(self, *args, **kwargs):
+        """
+        Calls the model's async predict() method with the given arguments.
+
+        If passed in request_pb2.PostModelOutputsRequest values, will send the model the raw
+        protos directly for compatibility with previous versions of the SDK.
+        """
+        inputs = None
+        if 'inputs' in kwargs:
+            inputs = kwargs['inputs']
+        elif args:
+            inputs = args[0]
+        if inputs and isinstance(inputs, list) and isinstance(inputs[0], resources_pb2.Input):
+            assert len(args) <= 1, (
+                "Cannot pass in raw protos and additional arguments at the same time."
+            )
+            inference_params = kwargs.get('inference_params', {})
+            output_config = kwargs.get('output_config', {})
+            return await self.client._async_predict_by_proto(
+                inputs=inputs, inference_params=inference_params, output_config=output_config
+            )
+
+        # Adding try-except, since the await works differently with jupyter kernels and in regular python scripts.
+        try:
+            return await self.client.predict(*args, **kwargs)
+        except TypeError:
+            # In jupyter, it returns a str object instead of a co-routine.
+            return self.client.predict(*args, **kwargs)
+
     def __getattr__(self, name):
         try:
             return getattr(self.model_info, name)
@@ -603,7 +635,11 @@ class Model(Lister, BaseClient):
         if any([deployment_id, nodepool_id, compute_cluster_id]):
             from clarifai.client.user import User
 
-            user_id = User().get_user_info(user_id='me').user.id
+            user_id = (
+                User(pat=self.auth_helper.pat, token=self.auth_helper._token)
+                .get_user_info(user_id='me')
+                .user.id
+            )
 
         runner_selector = None
         if deployment_id and (compute_cluster_id or nodepool_id):
@@ -761,6 +797,30 @@ class Model(Lister, BaseClient):
 
         return self.client.generate(*args, **kwargs)
 
+    async def async_generate(self, *args, **kwargs):
+        """
+        Calls the model's async generate() method with the given arguments.
+
+        If passed in request_pb2.PostModelOutputsRequest values, will send the model the raw
+        protos directly for compatibility with previous versions of the SDK.
+        """
+        inputs = None
+        if 'inputs' in kwargs:
+            inputs = kwargs['inputs']
+        elif args:
+            inputs = args[0]
+        if inputs and isinstance(inputs, list) and isinstance(inputs[0], resources_pb2.Input):
+            assert len(args) <= 1, (
+                "Cannot pass in raw protos and additional arguments at the same time."
+            )
+            inference_params = kwargs.get('inference_params', {})
+            output_config = kwargs.get('output_config', {})
+            return self.client._async_generate_by_proto(
+                inputs=inputs, inference_params=inference_params, output_config=output_config
+            )
+
+        return self.client.generate(*args, **kwargs)
+
     def generate_by_filepath(
         self,
         filepath: str,
@@ -926,6 +986,50 @@ class Model(Lister, BaseClient):
 
         return self.client.stream(*args, **kwargs)
 
+    async def async_stream(self, *args, **kwargs):
+        """
+        Calls the model's async stream() method with the given arguments.
+
+        If passed in request_pb2.PostModelOutputsRequest values, will send the model the raw
+        protos directly for compatibility with previous versions of the SDK.
+        """
+
+        use_proto_call = False
+        inputs = None
+        if 'inputs' in kwargs:
+            inputs = kwargs['inputs']
+        elif args:
+            inputs = args[0]
+        if inputs and isinstance(inputs, Iterable):
+            inputs_iter = inputs
+            try:
+                peek = next(inputs_iter)
+            except StopIteration:
+                pass
+            else:
+                use_proto_call = (
+                    peek and isinstance(peek, list) and isinstance(peek[0], resources_pb2.Input)
+                )
+                # put back the peeked value
+                if inputs_iter is inputs:
+                    inputs = itertools.chain([peek], inputs_iter)
+                    if 'inputs' in kwargs:
+                        kwargs['inputs'] = inputs
+                    else:
+                        args = (inputs,) + args[1:]
+
+            if use_proto_call:
+                assert len(args) <= 1, (
+                    "Cannot pass in raw protos and additional arguments at the same time."
+                )
+                inference_params = kwargs.get('inference_params', {})
+                output_config = kwargs.get('output_config', {})
+                return self.client._async_stream_by_proto(
+                    inputs=inputs, inference_params=inference_params, output_config=output_config
+                )
+
+            return self.client.async_stream(*args, **kwargs)
+
     def stream_by_filepath(
         self,
         filepath: str,