nucliadb 6.5.1.post4504__py3-none-any.whl → 6.5.1.post4510__py3-none-any.whl

nucliadb/search/search/chat/prompt.py

@@ -100,33 +100,58 @@ class CappedPromptContext:
         self.output: PromptContext = {}
         self.images: PromptContextImages = {}
         self.max_size = max_size
-        self._size = 0
 
     def __setitem__(self, key: str, value: str) -> None:
-        prev_value_len = len(self.output.get(key, ""))
-        if self.max_size is None:
-            # Unbounded size context
-            to_add = value
-        else:
-            # Make sure we don't exceed the max size
-            size_available = max(self.max_size - self._size + prev_value_len, 0)
-            to_add = value[:size_available]
-        self.output[key] = to_add
-        self._size = self._size - prev_value_len + len(to_add)
+        self.output.__setitem__(key, value)
 
     def __getitem__(self, key: str) -> str:
         return self.output.__getitem__(key)
 
     def __delitem__(self, key: str) -> None:
-        value = self.output.pop(key, "")
-        self._size -= len(value)
+        try:
+            self.output.__delitem__(key)
+        except KeyError:
+            pass
 
     def text_block_ids(self) -> list[str]:
         return list(self.output.keys())
 
     @property
     def size(self) -> int:
-        return self._size
+        """
+        Returns the total size of the context in characters.
+        """
+        return sum(len(text) for text in self.output.values())
+
+    def cap(self) -> dict[str, str]:
+        """
+        This method will trim the context to the maximum size if it exceeds it.
+        It will remove text from the most recent entries first, until the size is below the limit.
+        """
+        if self.max_size is None:
+            return self.output
+
+        if self.size <= self.max_size:
+            return self.output
+
+        logger.info("Removing text from context to fit within the max size limit")
+        # Iterate the dictionary in reverse order of insertion
+        for key in reversed(list(self.output.keys())):
+            current_size = self.size
+            if current_size <= self.max_size:
+                break
+            # Remove text from the value
+            text = self.output[key]
+            # If removing the whole text still keeps the total size above the limit, remove it
+            if current_size - len(text) >= self.max_size:
+                del self.output[key]
+            else:
+                # Otherwise, trim the text to fit within the limit
+                excess_size = current_size - self.max_size
+                if excess_size > 0:
+                    trimmed_text = text[:-excess_size]
+                    self.output[key] = trimmed_text
+        return self.output
 
 
 async def get_next_conversation_messages(
@@ -1001,12 +1026,11 @@ class PromptContextBuilder:
         self,
     ) -> tuple[PromptContext, PromptContextOrder, PromptContextImages, AugmentedContext]:
         ccontext = CappedPromptContext(max_size=self.max_context_characters)
-        print(".......................")
         self.prepend_user_context(ccontext)
         await self._build_context(ccontext)
         if self.visual_llm:
             await self._build_context_images(ccontext)
-        context = ccontext.output
+        context = ccontext.cap()
         context_images = ccontext.images
         context_order = {text_block_id: order for order, text_block_id in enumerate(context.keys())}
         return context, context_order, context_images, self.augmented_context

nucliadb/search/search/predict_proxy.py

@@ -171,14 +171,13 @@ async def chat_streaming_generator(
     user_query: str,
     is_json: bool,
 ):
-    stream = predict_response.content.iter_any()
     first = True
     status_code = AnswerStatusCode.ERROR.value
     text_answer = ""
     json_object = None
     metrics = AskMetrics()
     with metrics.time(PREDICT_ANSWER_METRIC):
-        async for chunk in stream:
+        async for chunk in predict_response.content:
             if first:
                 metrics.record_first_chunk_yielded()
                 first = False

{nucliadb-6.5.1.post4504.dist-info → nucliadb-6.5.1.post4510.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nucliadb
-Version: 6.5.1.post4504
+Version: 6.5.1.post4510
 Summary: NucliaDB
 Author-email: Nuclia <nucliadb@nuclia.com>
 License-Expression: AGPL-3.0-or-later
@@ -19,11 +19,11 @@ Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3 :: Only
 Requires-Python: <4,>=3.9
 Description-Content-Type: text/markdown
-Requires-Dist: nucliadb-telemetry[all]>=6.5.1.post4504
-Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.5.1.post4504
-Requires-Dist: nucliadb-protos>=6.5.1.post4504
-Requires-Dist: nucliadb-models>=6.5.1.post4504
-Requires-Dist: nidx-protos>=6.5.1.post4504
+Requires-Dist: nucliadb-telemetry[all]>=6.5.1.post4510
+Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.5.1.post4510
+Requires-Dist: nucliadb-protos>=6.5.1.post4510
+Requires-Dist: nucliadb-models>=6.5.1.post4510
+Requires-Dist: nidx-protos>=6.5.1.post4510
 Requires-Dist: nucliadb-admin-assets>=1.0.0.post1224
 Requires-Dist: nuclia-models>=0.24.2
 Requires-Dist: uvicorn[standard]

{nucliadb-6.5.1.post4504.dist-info → nucliadb-6.5.1.post4510.dist-info}/RECORD

@@ -255,7 +255,7 @@ nucliadb/search/search/merge.py,sha256=XiRBsxhYPshPV7lZXD-9E259KZOPIf4I2tKosY0lP
 nucliadb/search/search/metrics.py,sha256=3I6IN0qDSmqIvUaWJmT3rt-Jyjs6LcvnKI8ZqCiuJPY,3501
 nucliadb/search/search/paragraphs.py,sha256=pNAEiYqJGGUVcEf7xf-PFMVqz0PX4Qb-WNG-_zPGN2o,7799
 nucliadb/search/search/pgcatalog.py,sha256=O_nRjSJf1Qc-XorVwcNlsDOftzy_zQLLfagkjU4YmSA,16718
-nucliadb/search/search/predict_proxy.py,sha256=cuD_sfM3RLdEoQaanRz0CflO6nKVGGKPzoFA17shb_w,8647
+nucliadb/search/search/predict_proxy.py,sha256=rbclJeMBW7Yd6PpSAhp9fiH6YpPrVnUb-qgIGQTUB90,8616
 nucliadb/search/search/query.py,sha256=0qIQdt548L3jtKOyKo06aGJ73SLBxAW3N38_Hc1M3Uw,11528
 nucliadb/search/search/rank_fusion.py,sha256=xZtXhbmKb_56gs73u6KkFm2efvTATOSMmpOV2wrAIqE,9613
 nucliadb/search/search/rerankers.py,sha256=E2J1QdKAojqbhHM3KAyaOXKf6tJyETUxKs4tf_BEyqk,7472
@@ -266,7 +266,7 @@ nucliadb/search/search/chat/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn
 nucliadb/search/search/chat/ask.py,sha256=0sgfiCbNaCZrTvYaRGtf5xL6VnzRgzofINiEP4IvhWs,38278
 nucliadb/search/search/chat/exceptions.py,sha256=Siy4GXW2L7oPhIR86H3WHBhE9lkV4A4YaAszuGGUf54,1356
 nucliadb/search/search/chat/images.py,sha256=PA8VWxT5_HUGfW1ULhKTK46UBsVyINtWWqEM1ulzX1E,3095
-nucliadb/search/search/chat/prompt.py,sha256=TKyIblYvZ3EiyfwVFbEeYJ66ZNDtbeJySSa45HptCOk,51837
+nucliadb/search/search/chat/prompt.py,sha256=SNsCtB9mZTODjnUMAH8YfPxn05Kjl2d5xTIteNxyVcI,52783
 nucliadb/search/search/chat/query.py,sha256=3jMPNbiFEOoS0ydMOPYkSx1qVlvAv51npzadWXDwkMs,16650
 nucliadb/search/search/query_parser/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
 nucliadb/search/search/query_parser/exceptions.py,sha256=sVl9gRNzhE-s480LBBVkiXzNRbKhYRQN5F3it5tNNp8,939
@@ -375,8 +375,8 @@ nucliadb/writer/tus/local.py,sha256=7jYa_w9b-N90jWgN2sQKkNcomqn6JMVBOVeDOVYJHto,
 nucliadb/writer/tus/s3.py,sha256=vF0NkFTXiXhXq3bCVXXVV-ED38ECVoUeeYViP8uMqcU,8357
 nucliadb/writer/tus/storage.py,sha256=ToqwjoYnjI4oIcwzkhha_MPxi-k4Jk3Lt55zRwaC1SM,2903
 nucliadb/writer/tus/utils.py,sha256=MSdVbRsRSZVdkaum69_0wku7X3p5wlZf4nr6E0GMKbw,2556
-nucliadb-6.5.1.post4504.dist-info/METADATA,sha256=NgbKcJmI6-SK9h2awP4bpcbndCDv-XWCqmh6rGtBREg,4158
-nucliadb-6.5.1.post4504.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-nucliadb-6.5.1.post4504.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
-nucliadb-6.5.1.post4504.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
-nucliadb-6.5.1.post4504.dist-info/RECORD,,
+nucliadb-6.5.1.post4510.dist-info/METADATA,sha256=lLAN-xpgGXv6qlz-HyosNkSlKRX5qV6Be8-Xj59N8F0,4158
+nucliadb-6.5.1.post4510.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+nucliadb-6.5.1.post4510.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
+nucliadb-6.5.1.post4510.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
+nucliadb-6.5.1.post4510.dist-info/RECORD,,