clarifai 9.10.1__py3-none-any.whl → 9.10.3__py3-none-any.whl

clarifai/client/input.py

@@ -25,10 +25,11 @@ class Inputs(Lister, BaseClient):
   """Inputs is a class that provides access to Clarifai API endpoints related to Input information."""
 
   def __init__(self,
-               user_id: str = "",
-               app_id: str = "",
+               user_id: str = None,
+               app_id: str = None,
                logger_level: str = "INFO",
                base_url: str = "https://api.clarifai.com",
+               pat: str = None,
                **kwargs):
     """Initializes an Input object.
 
@@ -43,11 +44,11 @@ class Inputs(Lister, BaseClient):
     self.kwargs = {**kwargs}
     self.input_info = resources_pb2.Input(**self.kwargs)
     self.logger = get_logger(logger_level=logger_level, name=__name__)
-    BaseClient.__init__(self, user_id=self.user_id, app_id=self.app_id, base=base_url)
+    BaseClient.__init__(self, user_id=self.user_id, app_id=self.app_id, base=base_url, pat=pat)
     Lister.__init__(self)
 
-  def _get_proto(self,
-                 input_id: str,
+  @staticmethod
+  def _get_proto(input_id: str,
                  dataset_id: Union[str, None],
                  imagepb: Image = None,
                  video_pb: Video = None,
@@ -106,8 +107,8 @@ class Inputs(Lister, BaseClient):
             concepts=concepts,
             metadata=metadata))
 
-  def get_input_from_url(self,
-                         input_id: str,
+  @staticmethod
+  def get_input_from_url(input_id: str,
                          image_url: str = None,
                          video_url: str = None,
                          audio_url: str = None,
@@ -129,8 +130,7 @@ class Inputs(Lister, BaseClient):
 
     Example:
         >>> from clarifai.client.input import Inputs
-        >>> input_obj = Inputs()
-        >>> input_proto = input_obj.get_input_from_url(input_id = 'demo', image_url='https://samples.clarifai.com/metro-north.jpg')
+        >>> input_proto = Inputs.get_input_from_url(input_id = 'demo', image_url='https://samples.clarifai.com/metro-north.jpg')
     """
     if not any((image_url, video_url, audio_url, text_url)):
       raise ValueError(
@@ -139,7 +139,7 @@ class Inputs(Lister, BaseClient):
     video_pb = resources_pb2.Video(url=video_url) if video_url else None
     audio_pb = resources_pb2.Audio(url=audio_url) if audio_url else None
     text_pb = resources_pb2.Text(url=text_url) if text_url else None
-    return self._get_proto(
+    return Inputs._get_proto(
         input_id=input_id,
         dataset_id=dataset_id,
         imagepb=image_pb,
@@ -148,8 +148,8 @@ class Inputs(Lister, BaseClient):
         text_pb=text_pb,
         **kwargs)
 
-  def get_input_from_file(self,
-                          input_id: str,
+  @staticmethod
+  def get_input_from_file(input_id: str,
                           image_file: str = None,
                           video_file: str = None,
                           audio_file: str = None,
@@ -171,8 +171,7 @@ class Inputs(Lister, BaseClient):
 
     Example:
         >>> from clarifai.client.input import Inputs
-        >>> input_obj = Inputs()
-        >>> input_proto = input_obj.get_input_from_file(input_id = 'demo', video_file='file_path')
+        >>> input_proto = Inputs.get_input_from_file(input_id = 'demo', video_file='file_path')
     """
     if not any((image_file, video_file, audio_file, text_file)):
       raise ValueError(
@@ -181,7 +180,7 @@ class Inputs(Lister, BaseClient):
     video_pb = resources_pb2.Video(base64=open(video_file, 'rb').read()) if video_file else None
     audio_pb = resources_pb2.Audio(base64=open(audio_file, 'rb').read()) if audio_file else None
     text_pb = resources_pb2.Text(raw=open(text_file, 'rb').read()) if text_file else None
-    return self._get_proto(
+    return Inputs._get_proto(
         input_id=input_id,
         dataset_id=dataset_id,
         imagepb=image_pb,
@@ -190,8 +189,8 @@ class Inputs(Lister, BaseClient):
         text_pb=text_pb,
         **kwargs)
 
-  def get_input_from_bytes(self,
-                           input_id: str,
+  @staticmethod
+  def get_input_from_bytes(input_id: str,
                            image_bytes: bytes = None,
                            video_bytes: bytes = None,
                            audio_bytes: bytes = None,
@@ -213,10 +212,9 @@ class Inputs(Lister, BaseClient):
 
     Example:
         >>> from clarifai.client.input import Inputs
-        >>> input_obj = Inputs()
         >>> image = open('demo.jpg', 'rb').read()
         >>> video = open('demo.mp4', 'rb').read()
-        >>> input_proto = input_obj.get_input_from_bytes(input_id = 'demo',image_bytes =image, video_bytes=video)
+        >>> input_proto = Inputs.get_input_from_bytes(input_id = 'demo',image_bytes =image, video_bytes=video)
     """
     if not any((image_bytes, video_bytes, audio_bytes, text_bytes)):
       raise ValueError(
@@ -225,7 +223,7 @@ class Inputs(Lister, BaseClient):
     video_pb = resources_pb2.Video(base64=video_bytes) if video_bytes else None
     audio_pb = resources_pb2.Audio(base64=audio_bytes) if audio_bytes else None
     text_pb = resources_pb2.Text(raw=text_bytes) if text_bytes else None
-    return self._get_proto(
+    return Inputs._get_proto(
         input_id=input_id,
         dataset_id=dataset_id,
         imagepb=image_pb,
@@ -234,9 +232,8 @@ class Inputs(Lister, BaseClient):
         text_pb=text_pb,
         **kwargs)
 
-  def get_image_inputs_from_folder(self,
-                                   folder_path: str,
-                                   dataset_id: str = None,
+  @staticmethod
+  def get_image_inputs_from_folder(folder_path: str, dataset_id: str = None,
                                    labels: bool = False) -> List[Input]:  #image specific
     """Create input protos for image data type from folder.
 
@@ -248,8 +245,7 @@ class Inputs(Lister, BaseClient):
 
     Example:
         >>> from clarifai.client.input import Inputs
-        >>> input_obj = Inputs()
-        >>> input_protos = input_obj.get_image_inputs_from_folder(folder_path='demo_folder')
+        >>> input_protos = Inputs.get_image_inputs_from_folder(folder_path='demo_folder')
     """
     input_protos = []
     labels = [folder_path.split('/')[-1]] if labels else None
@@ -259,11 +255,12 @@ class Inputs(Lister, BaseClient):
       input_id = filename.split('.')[0]
       image_pb = resources_pb2.Image(base64=open(os.path.join(folder_path, filename), 'rb').read())
       input_protos.append(
-          self._get_proto(
+          Inputs._get_proto(
               input_id=input_id, dataset_id=dataset_id, imagepb=image_pb, labels=labels))
     return input_protos
 
-  def get_text_input(self, input_id: str, raw_text: str, dataset_id: str = None,
+  @staticmethod
+  def get_text_input(input_id: str, raw_text: str, dataset_id: str = None,
                      **kwargs) -> Text:  #text specific
     """Create input proto for text data type from rawtext.
 
@@ -278,14 +275,13 @@ class Inputs(Lister, BaseClient):
 
     Example:
         >>> from clarifai.client.input import Inputs
-        >>> input_obj = Inputs()
-        >>> input_protos = input_obj.get_text_input(input_id = 'demo', raw_text = 'This is a test')
+        >>> input_protos = Inputs.get_text_input(input_id = 'demo', raw_text = 'This is a test')
     """
     text_pb = resources_pb2.Text(raw=raw_text)
-    return self._get_proto(input_id=input_id, dataset_id=dataset_id, text_pb=text_pb, **kwargs)
+    return Inputs._get_proto(input_id=input_id, dataset_id=dataset_id, text_pb=text_pb, **kwargs)
 
-  def get_inputs_from_csv(self,
-                          csv_path: str,
+  @staticmethod
+  def get_inputs_from_csv(csv_path: str,
                           input_type: str = 'text',
                           csv_type: str = 'raw',
                           dataset_id: str = None,
@@ -304,8 +300,7 @@ class Inputs(Lister, BaseClient):
 
     Example:
         >>> from clarifai.client.input import Inputs
-        >>> input_obj = Inputs()
-        >>> input_protos = input_obj.get_inputs_from_csv(csv_path='filepath', input_type='text', csv_type='raw')
+        >>> input_protos = Inputs.get_inputs_from_csv(csv_path='filepath', input_type='text', csv_type='raw')
     """
     input_protos = []
     with open(csv_path) as _file:
@@ -356,7 +351,7 @@ class Inputs(Lister, BaseClient):
 
         if csv_type == 'raw':
           input_protos.append(
-              self.get_text_input(
+              Inputs.get_text_input(
                   input_id=input_id,
                   raw_text=text,
                   dataset_id=dataset_id,
@@ -365,7 +360,7 @@ class Inputs(Lister, BaseClient):
                   geo_info=geo_info))
         elif csv_type == 'url':
           input_protos.append(
-              self.get_input_from_url(
+              Inputs.get_input_from_url(
                   input_id=input_id,
                   image_url=image,
                   text_url=text,
@@ -377,7 +372,7 @@ class Inputs(Lister, BaseClient):
                   geo_info=geo_info))
         else:
           input_protos.append(
-              self.get_input_from_file(
+              Inputs.get_input_from_file(
                   input_id=input_id,
                   image_file=image,
                   text_file=text,
@@ -390,9 +385,8 @@ class Inputs(Lister, BaseClient):
 
     return input_protos
 
-  def get_text_inputs_from_folder(self,
-                                  folder_path: str,
-                                  dataset_id: str = None,
+  @staticmethod
+  def get_text_inputs_from_folder(folder_path: str, dataset_id: str = None,
                                   labels: bool = False) -> List[Text]:  #text specific
     """Create input protos for text data type from folder.
 
@@ -404,8 +398,7 @@ class Inputs(Lister, BaseClient):
 
     Example:
         >>> from clarifai.client.input import Inputs
-        >>> input_obj = Inputs()
-        >>> input_protos = input_obj.get_text_inputs_from_folder(folder_path='demo_folder')
+        >>> input_protos = Inputs.get_text_inputs_from_folder(folder_path='demo_folder')
     """
     input_protos = []
     labels = [folder_path.split('/')[-1]] if labels else None
@@ -415,11 +408,12 @@ class Inputs(Lister, BaseClient):
       input_id = filename.split('.')[0]
       text_pb = resources_pb2.Text(raw=open(os.path.join(folder_path, filename), 'rb').read())
       input_protos.append(
-          self._get_proto(
+          Inputs._get_proto(
               input_id=input_id, dataset_id=dataset_id, text_pb=text_pb, labels=labels))
     return input_protos
 
-  def get_annotation_proto(self, input_id: str, label: str, annotations: List) -> Annotation:
+  @staticmethod
+  def get_annotation_proto(input_id: str, label: str, annotations: List) -> Annotation:
     """Create an annotation proto for each bounding box, label input pair.
 
     Args:
@@ -432,8 +426,7 @@ class Inputs(Lister, BaseClient):
 
     Example:
         >>> from clarifai.client.input import Inputs
-        >>> input_obj = Inputs()
-        >>> input_obj.get_annotation_proto(input_id='demo', label='demo', annotations=[x_min, y_min, x_max, y_max])
+        >>> Inputs.get_annotation_proto(input_id='demo', label='demo', annotations=[x_min, y_min, x_max, y_max])
     """
     if not isinstance(annotations, list):
       raise UserError("annotations must be a list of bbox cooridnates")
@@ -458,7 +451,8 @@ class Inputs(Lister, BaseClient):
 
     return input_annot_proto
 
-  def get_mask_proto(self, input_id: str, label: str, polygons: List[List[float]]) -> Annotation:
+  @staticmethod
+  def get_mask_proto(input_id: str, label: str, polygons: List[List[float]]) -> Annotation:
     """Create an annotation proto for each polygon box, label input pair.
 
     Args:
@@ -471,8 +465,7 @@ class Inputs(Lister, BaseClient):
 
     Example:
         >>> from clarifai.client.input import Inputs
-        >>> input_obj = Inputs()
-        >>> input_obj.get_mask_proto(input_id='demo', label='demo', polygons=[[[x,y],...,[x,y]],...])
+        >>> Inputs.get_mask_proto(input_id='demo', label='demo', polygons=[[[x,y],...,[x,y]],...])
     """
     if not isinstance(polygons, list):
       raise UserError("polygons must be a list of points")
@@ -653,8 +646,7 @@ class Inputs(Lister, BaseClient):
     response = self._grpc_request(self.STUB.PostAnnotations, request)
     if response.status.code != status_code_pb2.SUCCESS:
       try:
-        self.logger.warning(
-            f"Post annotations failed, status: {response.annotations[0].status.details}")
+        self.logger.warning(f"Post annotations failed, status: {response.annotations[0].status}")
       except Exception:
         self.logger.warning(f"Post annotations failed, status: {response.status.details}")
       finally:
@@ -788,16 +780,16 @@ class Inputs(Lister, BaseClient):
       annotations_info['id'] = annotations_info.pop('annotation_id')
       yield Annotation(**annotations_info)
 
-  def _bulk_upload(self, inputs: List[Input], chunk_size: int = 128) -> None:
+  def _bulk_upload(self, inputs: List[Input], batch_size: int = 128) -> None:
     """Uploads process for large number of inputs.
 
     Args:
         inputs (List[Input]): input protos
-        chunk_size (int): chunk size for each request
+        batch_size (int): batch size for each request
     """
     num_workers: int = min(10, cpu_count())  # limit max workers to 10
-    chunk_size = min(128, chunk_size)  # limit max protos in a req
-    chunked_inputs = Chunker(inputs, chunk_size).chunk()
+    batch_size = min(128, batch_size)  # limit max protos in a req
+    chunked_inputs = Chunker(inputs, batch_size).chunk()
     with ThreadPoolExecutor(max_workers=num_workers) as executor:
       with tqdm(total=len(chunked_inputs), desc='Uploading inputs') as progress:
         # Submit all jobs to the executor and store the returned futures

clarifai/client/model.py

@@ -27,34 +27,35 @@ class Model(Lister, BaseClient):
   """Model is a class that provides access to Clarifai API endpoints related to Model information."""
 
   def __init__(self,
-               url_init: str = "",
-               model_id: str = "",
+               url: str = None,
+               model_id: str = None,
                model_version: Dict = {'id': ""},
                base_url: str = "https://api.clarifai.com",
+               pat: str = None,
                **kwargs):
     """Initializes a Model object.
 
     Args:
-        url_init (str): The URL to initialize the model object.
+        url (str): The URL to initialize the model object.
         model_id (str): The Model ID to interact with.
         model_version (dict): The Model Version to interact with.
         base_url (str): Base API url. Default "https://api.clarifai.com"
+        pat (str): A personal access token for authentication. Can be set as env var CLARIFAI_PAT
         **kwargs: Additional keyword arguments to be passed to the Model.
     """
-    if url_init != "" and model_id != "":
-      raise UserError("You can only specify one of url_init or model_id.")
-    if url_init == "" and model_id == "":
-      raise UserError("You must specify one of url_init or model_id.")
-    if url_init != "":
-      user_id, app_id, _, model_id, model_version_id = ClarifaiUrlHelper.split_clarifai_url(
-          url_init)
+    if url and model_id:
+      raise UserError("You can only specify one of url or model_id.")
+    if not url and not model_id:
+      raise UserError("You must specify one of url or model_id.")
+    if url:
+      user_id, app_id, _, model_id, model_version_id = ClarifaiUrlHelper.split_clarifai_url(url)
       model_version = {'id': model_version_id}
       kwargs = {'user_id': user_id, 'app_id': app_id}
     self.kwargs = {**kwargs, 'id': model_id, 'model_version': model_version,}
     self.model_info = resources_pb2.Model(**self.kwargs)
     self.logger = get_logger(logger_level="INFO")
     self.training_params = {}
-    BaseClient.__init__(self, user_id=self.user_id, app_id=self.app_id, base=base_url)
+    BaseClient.__init__(self, user_id=self.user_id, app_id=self.app_id, base=base_url, pat=pat)
     Lister.__init__(self)
 
   def list_training_templates(self) -> List[str]:
@@ -306,7 +307,7 @@ class Model(Lister, BaseClient):
 
     Example:
         >>> from clarifai.client.model import Model
-        >>> model = Model("model_url")
+        >>> model = Model("url")
                     or
         >>> model = Model(model_id='model_id', user_id='user_id', app_id='app_id')
         >>> model_version = model.create_version(description='model_version_description')
@@ -330,7 +331,7 @@ class Model(Lister, BaseClient):
     dict_response = MessageToDict(response, preserving_proto_field_name=True)
     kwargs = self.process_response_keys(dict_response['model'], 'model')
 
-    return Model(base_url=self.base, **kwargs)
+    return Model(base_url=self.base, pat=self.pat, **kwargs)
 
   def list_versions(self, page_no: int = None,
                     per_page: int = None) -> Generator['Model', None, None]:
@@ -345,7 +346,7 @@ class Model(Lister, BaseClient):
 
     Example:
         >>> from clarifai.client.model import Model
-        >>> model = Model("model_url") # Example URL: https://clarifai.com/clarifai/main/models/general-image-recognition
+        >>> model = Model("url") # Example URL: https://clarifai.com/clarifai/main/models/general-image-recognition
                     or
         >>> model = Model(model_id='model_id', user_id='user_id', app_id='app_id')
         >>> all_model_versions = list(model.list_versions())
@@ -375,6 +376,7 @@ class Model(Lister, BaseClient):
       yield Model(
           model_id=self.id,
           base_url=self.base,
+          pat=self.pat,
           **dict(self.kwargs, model_version=model_version_info))
 
   def predict(self, inputs: List[Input], inference_params: Dict = {}, output_config: Dict = {}):
@@ -432,7 +434,7 @@ class Model(Lister, BaseClient):
 
     Example:
         >>> from clarifai.client.model import Model
-        >>> model = Model("model_url") # Example URL: https://clarifai.com/clarifai/main/models/general-image-recognition
+        >>> model = Model("url") # Example URL: https://clarifai.com/clarifai/main/models/general-image-recognition
                     or
         >>> model = Model(model_id='model_id', user_id='user_id', app_id='app_id')
         >>> model_prediction = model.predict_by_filepath('/path/to/image.jpg', 'image')
@@ -476,13 +478,13 @@ class Model(Lister, BaseClient):
       raise UserError('Invalid bytes.')
 
     if input_type == "image":
-      input_proto = Inputs().get_input_from_bytes("", image_bytes=input_bytes)
+      input_proto = Inputs.get_input_from_bytes("", image_bytes=input_bytes)
     elif input_type == "text":
-      input_proto = Inputs().get_input_from_bytes("", text_bytes=input_bytes)
+      input_proto = Inputs.get_input_from_bytes("", text_bytes=input_bytes)
     elif input_type == "video":
-      input_proto = Inputs().get_input_from_bytes("", video_bytes=input_bytes)
+      input_proto = Inputs.get_input_from_bytes("", video_bytes=input_bytes)
     elif input_type == "audio":
-      input_proto = Inputs().get_input_from_bytes("", audio_bytes=input_bytes)
+      input_proto = Inputs.get_input_from_bytes("", audio_bytes=input_bytes)
 
     return self.predict(
         inputs=[input_proto], inference_params=inference_params, output_config=output_config)
@@ -505,7 +507,7 @@ class Model(Lister, BaseClient):
 
     Example:
         >>> from clarifai.client.model import Model
-        >>> model = Model("model_url") # Example URL: https://clarifai.com/clarifai/main/models/general-image-recognition
+        >>> model = Model("url") # Example URL: https://clarifai.com/clarifai/main/models/general-image-recognition
                     or
         >>> model = Model(model_id='model_id', user_id='user_id', app_id='app_id')
         >>> model_prediction = model.predict_by_url('url', 'image')
@@ -515,13 +517,13 @@ class Model(Lister, BaseClient):
           f"Got input type {input_type} but expected one of image, text, video, audio.")
 
     if input_type == "image":
-      input_proto = Inputs().get_input_from_url("", image_url=url)
+      input_proto = Inputs.get_input_from_url("", image_url=url)
     elif input_type == "text":
-      input_proto = Inputs().get_input_from_url("", text_url=url)
+      input_proto = Inputs.get_input_from_url("", text_url=url)
     elif input_type == "video":
-      input_proto = Inputs().get_input_from_url("", video_url=url)
+      input_proto = Inputs.get_input_from_url("", video_url=url)
     elif input_type == "audio":
-      input_proto = Inputs().get_input_from_url("", audio_url=url)
+      input_proto = Inputs.get_input_from_url("", audio_url=url)
 
     return self.predict(
         inputs=[input_proto], inference_params=inference_params, output_config=output_config)
@@ -537,8 +539,8 @@ class Model(Lister, BaseClient):
           select_concepts (list[Concept]): The concepts to select.
           sample_ms (int): The number of milliseconds to sample.
     """
+    params = Struct()
     if inference_params is not None:
-      params = Struct()
       params.update(inference_params)
 
     self.model_info.model_version.output_info.CopyFrom(

clarifai/client/module.py

@@ -13,34 +13,35 @@ class Module(Lister, BaseClient):
   """Module is a class that provides access to Clarifai API endpoints related to Module information."""
 
   def __init__(self,
-               url_init: str = "",
-               module_id: str = "",
+               url: str = None,
+               module_id: str = None,
                module_version: Dict = {'id': ""},
                base_url: str = "https://api.clarifai.com",
+               pat: str = None,
                **kwargs):
     """Initializes a Module object.
 
         Args:
-            url_init (str): The URL to initialize the module object.
+            url (str): The URL to initialize the module object.
             module_id (str): The Module ID to interact with.
             module_version (dict): The Module Version to interact with.
             base_url (str): Base API url. Default "https://api.clarifai.com"
+            pat (str): A personal access token for authentication. Can be set as env var CLARIFAI_PAT
             **kwargs: Additional keyword arguments to be passed to the Module.
         """
-    if url_init != "" and module_id != "":
-      raise UserError("You can only specify one of url_init or module_id.")
-    if url_init == "" and module_id == "":
-      raise UserError("You must specify one of url_init or module_id.")
-    if url_init != "":
-      user_id, app_id, module_id, module_version_id = ClarifaiUrlHelper.split_module_ui_url(
-          url_init)
+    if url and module_id:
+      raise UserError("You can only specify one of url or module_id.")
+    if not url and not module_id:
+      raise UserError("You must specify one of url or module_id.")
+    if url:
+      user_id, app_id, module_id, module_version_id = ClarifaiUrlHelper.split_module_ui_url(url)
       module_version = {'id': module_version_id}
       kwargs = {'user_id': user_id, 'app_id': app_id}
 
     self.kwargs = {**kwargs, 'id': module_id, 'module_version': module_version}
     self.module_info = resources_pb2.Module(**self.kwargs)
     self.logger = get_logger(logger_level="INFO")
-    BaseClient.__init__(self, user_id=self.user_id, app_id=self.app_id, base=base_url)
+    BaseClient.__init__(self, user_id=self.user_id, app_id=self.app_id, base=base_url, pat=pat)
     Lister.__init__(self)
 
   def list_versions(self, page_no: int = None,
@@ -80,6 +81,7 @@ class Module(Lister, BaseClient):
       yield Module(
           module_id=self.id,
           base_url=self.base,
+          pat=self.pat,
           **dict(self.kwargs, module_version=module_version_info))
 
   def __getattr__(self, name):

clarifai/client/runner.py

@@ -32,19 +32,21 @@ class Runner(BaseClient):
 
   def __init__(self,
                runner_id: str,
-               user_id: str = "",
+               user_id: str = None,
                check_runner_exists: bool = True,
                base_url: str = "https://api.clarifai.com",
+               pat: str = None,
                **kwargs) -> None:
     """
     Args:
       runner_id (str): the id of the runner to use. Create the runner in the Clarifai API first
       user_id (str): Clarifai User ID
       base_url (str): Base API url. Default "https://api.clarifai.com"
+      pat (str): A personal access token for authentication. Can be set as env var CLARIFAI_PAT
     """
     user_id = user_id or os.environ.get("CLARIFAI_USER_ID", "")
 
-    if user_id == "":
+    if not user_id:
       raise UserError(
           "Set CLARIFAI_USER_ID as environment variables or pass user_id as input arguments")
 
@@ -52,7 +54,7 @@ class Runner(BaseClient):
     self.logger = get_logger("INFO", __name__)
     self.kwargs = {**kwargs, 'id': runner_id, 'user_id': user_id}
     self.runner_info = resources_pb2.Runner(**self.kwargs)
-    BaseClient.__init__(self, user_id=self.user_id, app_id="", base=base_url)
+    BaseClient.__init__(self, user_id=self.user_id, app_id="", base=base_url, pat=pat)
 
     # Check that the runner exists.
     if check_runner_exists:

clarifai/client/search.py

@@ -1,3 +1,4 @@
+from math import ceil
 from typing import Any, Callable, Dict, Generator
 
 from clarifai_grpc.grpc.api import resources_pb2, service_pb2
@@ -20,7 +21,9 @@ class Search(Lister, BaseClient):
                user_id,
                app_id,
                top_k: int = DEFAULT_TOP_K,
-               metric: str = DEFAULT_SEARCH_METRIC):
+               metric: str = DEFAULT_SEARCH_METRIC,
+               base_url: str = "https://api.clarifai.com",
+               pat: str = None):
     """Initialize the Search object.
 
     Args:
@@ -28,6 +31,8 @@ class Search(Lister, BaseClient):
         app_id (str): App ID.
         top_k (int, optional): Top K results to retrieve. Defaults to 10.
         metric (str, optional): Similarity metric (either 'cosine' or 'euclidean'). Defaults to 'cosine'.
+        base_url (str, optional): Base API url. Defaults to "https://api.clarifai.com".
+        pat (str, optional): A personal access token for authentication. Can be set as env var CLARIFAI_PAT
 
     Raises:
         UserError: If the metric is not 'cosine' or 'euclidean'.
@@ -39,11 +44,12 @@ class Search(Lister, BaseClient):
     self.app_id = app_id
     self.metric_distance = dict(cosine="COSINE_DISTANCE", euclidean="EUCLIDEAN_DISTANCE")[metric]
     self.data_proto = resources_pb2.Data()
+    self.top_k = top_k
 
-    self.inputs = Inputs(user_id=self.user_id, app_id=self.app_id)
+    self.inputs = Inputs(user_id=self.user_id, app_id=self.app_id, pat=pat)
     self.rank_filter_schema = get_schema()
-    BaseClient.__init__(self, user_id=self.user_id, app_id=self.app_id)
-    Lister.__init__(self, page_size=top_k)
+    BaseClient.__init__(self, user_id=self.user_id, app_id=self.app_id, base=base_url, pat=pat)
+    Lister.__init__(self, page_size=1000)
 
   def _get_annot_proto(self, **kwargs):
     """Get an Annotation proto message based on keyword arguments.
@@ -118,7 +124,7 @@ class Search(Lister, BaseClient):
             self.data_proto.video.CopyFrom(resources_pb2.Video())
         self.input_proto.data.CopyFrom(self.data_proto)
       elif key == "input_dataset_ids":
-        self.input_proto.dataset_ids = value
+        self.input_proto.dataset_ids.extend(value)
       elif key == "input_status_code":
         self.input_proto.status.code = value
       else:
@@ -154,18 +160,31 @@ class Search(Lister, BaseClient):
         Yields:
             response_dict: The next item in the listing.
         """
+    max_pages = ceil(self.top_k / self.default_page_size)
+    total_hits = 0
     page = 1
-    request_data['pagination'] = service_pb2.Pagination(page=page, per_page=self.default_page_size)
-    while True:
-      request_data['pagination'].page = page
+    while (page <= max_pages):
+      if (page == max_pages):
+        per_page = self.top_k - total_hits
+      else:
+        per_page = self.default_page_size
+      request_data['pagination'] = service_pb2.Pagination(page=page, per_page=per_page)
       response = self._grpc_request(endpoint, proto_message(**request_data))
       dict_response = MessageToDict(response, preserving_proto_field_name=True)
       if response.status.code != status_code_pb2.SUCCESS:
-        raise Exception(f"Listing failed with response {response!r}")
+        if "page * perPage cannot exceed" in str(response.status.details):
+          msg = (f"Your top_k is set to {self.top_k}. "
+                 f"The current pagination settings exceed the limit. Please reach out to "
+                 f"support@clarifai.com to request an increase for your use case.\n"
+                 f"req_id: {response.status.req_id}")
+          raise UserError(msg)
+        else:
+          raise Exception(f"Listing failed with response {response!r}")
 
       if 'hits' not in list(dict_response.keys()):
         break
       page += 1
+      total_hits += per_page
       yield response
 
   def query(self, ranks=[{}], filters=[{}]):
@@ -186,7 +205,7 @@ class Search(Lister, BaseClient):
 
         Vector search over inputs
         >>> from clarifai.client.search import Search
-        >>> search = Search(user_id='user_id', app_id='app_id', top_k=10, metric='cosine')
+        >>> search = Search(user_id='user_id', app_id='app_id', top_k=1, metric='cosine')
         >>> res = search.query(ranks=[{'image_url': 'https://samples.clarifai.com/dog.tiff'}])
 
     Note: For more detailed search examples, please refer to [examples](https://github.com/Clarifai/examples/tree/main/search).