clarifai 11.5.5__py3-none-any.whl → 11.5.6__py3-none-any.whl

clarifai/__init__.py

@@ -1 +1 @@
-__version__ = "11.5.5"
+__version__ = "11.5.6"

clarifai/runners/models/model_builder.py

@@ -550,6 +550,11 @@ class ModelBuilder:
             "inference_compute_info not found in the config file"
         )
         inference_compute_info = self.config.get('inference_compute_info')
+        # Ensure cpu_limit is a string if it exists and is an int
+        if 'cpu_limit' in inference_compute_info and isinstance(
+            inference_compute_info['cpu_limit'], int
+        ):
+            inference_compute_info['cpu_limit'] = str(inference_compute_info['cpu_limit'])
         return json_format.ParseDict(inference_compute_info, resources_pb2.ComputeInfo())
 
     def check_model_exists(self):

clarifai/runners/models/model_class.py

@@ -1,6 +1,7 @@
 import inspect
 import itertools
 import os
+import threading
 import traceback
 from abc import ABC
 from collections import abc
@@ -59,15 +60,22 @@ class ModelClass(ABC):
             yield item.x + ' ' + str(item.y)
     '''
 
+    def __init__(self):
+        super().__init__()
+        self._thread_local = threading.local()
+
     @staticmethod
     def method(func):
         setattr(func, _METHOD_INFO_ATTR, _MethodInfo(func))
         return func
 
     def set_output_context(self, prompt_tokens=None, completion_tokens=None):
-        """This is used to set the prompt and completion tokens in the Output proto"""
-        self._prompt_tokens = prompt_tokens
-        self._completion_tokens = completion_tokens
+        """Set the prompt and completion tokens for the Output proto.
+        In batch mode, call this once per output, in order, before returning each output.
+        """
+        if not hasattr(self._thread_local, 'token_contexts'):
+            self._thread_local.token_contexts = []
+        self._thread_local.token_contexts.append((prompt_tokens, completion_tokens))
 
     def load_model(self):
         """Load the model."""
@@ -213,6 +221,7 @@ class ModelClass(ABC):
         self, request: service_pb2.PostModelOutputsRequest
     ) -> Iterator[service_pb2.MultiOutputResponse]:
         try:
+            assert len(request.inputs) == 1, "Generate requires exactly one input"
             method_name = 'generate'
             if len(request.inputs) > 0 and '_method_name' in request.inputs[0].data.metadata:
                 method_name = request.inputs[0].data.metadata['_method_name']
@@ -385,12 +394,18 @@ class ModelClass(ABC):
             data = DataConverter.convert_output_data_to_old_format(proto.data)
             proto.data.CopyFrom(data)
         proto.status.code = status_code_pb2.SUCCESS
-        if hasattr(self, "_prompt_tokens") and self._prompt_tokens is not None:
-            proto.prompt_tokens = self._prompt_tokens
-        if hasattr(self, "_completion_tokens") and self._completion_tokens is not None:
-            proto.completion_tokens = self._completion_tokens
-        self._prompt_tokens = None
-        self._completion_tokens = None
+        # Per-output token context support
+        token_contexts = getattr(self._thread_local, 'token_contexts', None)
+        prompt_tokens = completion_tokens = None
+        if token_contexts and len(token_contexts) > 0:
+            prompt_tokens, completion_tokens = token_contexts.pop(0)
+            # If this was the last, clean up
+            if len(token_contexts) == 0:
+                del self._thread_local.token_contexts
+        if prompt_tokens is not None:
+            proto.prompt_tokens = prompt_tokens
+        if completion_tokens is not None:
+            proto.completion_tokens = completion_tokens
         return proto
 
     @classmethod

clarifai/runners/models/openai_class.py

@@ -34,6 +34,7 @@ class OpenAIModelClass(ModelClass):
     model = None
 
     def __init__(self) -> None:
+        super().__init__()
         if self.client is None:
             raise NotImplementedError("Subclasses must set the 'client' class attribute")
         if self.model is None:

{clarifai-11.5.5.dist-info → clarifai-11.5.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: clarifai
-Version: 11.5.5
+Version: 11.5.6
 Home-page: https://github.com/Clarifai/clarifai-python
 Author: Clarifai
 Author-email: support@clarifai.com

{clarifai-11.5.5.dist-info → clarifai-11.5.6.dist-info}/RECORD

@@ -1,4 +1,4 @@
-clarifai/__init__.py,sha256=OF6Ppf-NDN6b_03eNnRsnEEmYm8mHwco3nReiqZIKq8,23
+clarifai/__init__.py,sha256=BxuXNC41hShXWv3SgslBYu8ufM3yw21Lbp4jemirGjY,23
 clarifai/cli.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 clarifai/errors.py,sha256=GXa6D4v_L404J83jnRNFPH7s-1V9lk7w6Ws99f1g-AY,2772
 clarifai/versions.py,sha256=ecSuEB_nOL2XSoYHDw2n23XUbm_KPOGjudMXmQrGdS8,224
@@ -75,12 +75,12 @@ clarifai/runners/dockerfile_template/Dockerfile.template,sha256=DUH7F0-uLOV0LTjn
 clarifai/runners/models/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 clarifai/runners/models/dummy_openai_model.py,sha256=pcmAVbqTTGG4J3BLVjKfvM_SQ-GET_XexIUdLcr9Zvo,8373
 clarifai/runners/models/mcp_class.py,sha256=RdKn7rW4vYol0VRDZiLTSMfkqjLhO1ijXAQ0Rq0Jfnw,6647
-clarifai/runners/models/model_builder.py,sha256=QgGYruFVUS-xfva-wdGaZT0zpLbbc3dkYqOoYSM_jyU,63409
-clarifai/runners/models/model_class.py,sha256=2yAXPsRFIjmyo8n9XUu-YaMOxyRpeVg_MJo4IAo0kRU,19365
+clarifai/runners/models/model_builder.py,sha256=lzXf1H4oQ4l672xHoXSaT4bi0z0vTj3YnziT7RRRLIg,63695
+clarifai/runners/models/model_class.py,sha256=Ndh437BNMkpFBo6B108GuKL8sGYaGnSplZ6FxOgd_v8,20010
 clarifai/runners/models/model_run_locally.py,sha256=6-6WjEKc0ba3gAv4wOLdMs2XOzS3b-2bZHJS0wdVqJY,20088
 clarifai/runners/models/model_runner.py,sha256=tZTX1XKMlniJEmd1WMjcwGfej5NCWqv23HZ4xrG8YV8,9153
 clarifai/runners/models/model_servicer.py,sha256=415RqiXCPH-1WgFHrtZMDH1nb8gaQDpVlrQ_tvtbjOg,4523
-clarifai/runners/models/openai_class.py,sha256=aXlk5W6LWkh-A4eZYi74DeLW0i_86_9DYYGxpJHXI0w,6688
+clarifai/runners/models/openai_class.py,sha256=OVYe4dWJPhskyZn53X9yJorzKZA7mjRCqYbu1rzRTFU,6715
 clarifai/runners/models/visual_classifier_class.py,sha256=1ZoLfCT2crrgRbejjTMAIwpTRgQMiH9N9yflOVpFxSg,2721
 clarifai/runners/models/visual_detector_class.py,sha256=ky4oFAkGCKPpGPdgaOso-n6D3HcmnbKee_8hBsNiV8U,2883
 clarifai/runners/pipeline_steps/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -118,9 +118,9 @@ clarifai/workflows/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuF
 clarifai/workflows/export.py,sha256=HvUYG9N_-UZoRR0-_tdGbZ950_AeBqawSppgUxQebR0,1913
 clarifai/workflows/utils.py,sha256=ESL3INcouNcLKCh-nMpfXX-YbtCzX7tz7hT57_RGQ3M,2079
 clarifai/workflows/validate.py,sha256=UhmukyHkfxiMFrPPeBdUTiCOHQT5-shqivlBYEyKTlU,2931
-clarifai-11.5.5.dist-info/licenses/LICENSE,sha256=mUqF_d12-qE2n41g7C5_sq-BMLOcj6CNN-jevr15YHU,555
-clarifai-11.5.5.dist-info/METADATA,sha256=SkHXmknF4MRdxN2h9_8tqq-ZZid1avN-WisUenwU6ak,22737
-clarifai-11.5.5.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-clarifai-11.5.5.dist-info/entry_points.txt,sha256=X9FZ4Z-i_r2Ud1RpZ9sNIFYuu_-9fogzCMCRUD9hyX0,51
-clarifai-11.5.5.dist-info/top_level.txt,sha256=wUMdCQGjkxaynZ6nZ9FAnvBUCgp5RJUVFSy2j-KYo0s,9
-clarifai-11.5.5.dist-info/RECORD,,
+clarifai-11.5.6.dist-info/licenses/LICENSE,sha256=mUqF_d12-qE2n41g7C5_sq-BMLOcj6CNN-jevr15YHU,555
+clarifai-11.5.6.dist-info/METADATA,sha256=dE9Pb314d6-aAeFuNetMEyt6GsbveEuJdqnjTbeln7E,22737
+clarifai-11.5.6.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+clarifai-11.5.6.dist-info/entry_points.txt,sha256=X9FZ4Z-i_r2Ud1RpZ9sNIFYuu_-9fogzCMCRUD9hyX0,51
+clarifai-11.5.6.dist-info/top_level.txt,sha256=wUMdCQGjkxaynZ6nZ9FAnvBUCgp5RJUVFSy2j-KYo0s,9
+clarifai-11.5.6.dist-info/RECORD,,