nucliadb 6.5.1.post4504__py3-none-any.whl → 6.5.1.post4507__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-6.5.1.post4504.dist-info → nucliadb-6.5.1.post4507.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nucliadb
-Version: 6.5.1.post4504
+Version: 6.5.1.post4507
 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.post4507
+Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.5.1.post4507
+Requires-Dist: nucliadb-protos>=6.5.1.post4507
+Requires-Dist: nucliadb-models>=6.5.1.post4507
+Requires-Dist: nidx-protos>=6.5.1.post4507
 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.post4507.dist-info}/RECORD

@@ -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.post4507.dist-info/METADATA,sha256=mXI1_xVFcW5gf0mv4h3PdH3RYyaqAUTt6UlvPEU2gb4,4158
+nucliadb-6.5.1.post4507.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+nucliadb-6.5.1.post4507.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
+nucliadb-6.5.1.post4507.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
+nucliadb-6.5.1.post4507.dist-info/RECORD,,