- enabled backoff for async chunk processing

Author: markhofm

Diff for: src/hierarchical.py

@@ -13,8 +13,9 @@
 
 from .types import TYPE_SPECIFIC_PROMPTS, DocumentType
 import asyncio
-import functools
 import nest_asyncio
+import backoff
+import functools
 
 # Load the LLMDataModel class from the summarize-lib module
 LLMDataModel = Import.load_expression("ExtensityAI/primitives-extend", "LLMDataModel")
@@ -183,7 +184,6 @@ def prompt(self):
         """
         )
         prompt_text += Summary.instruct_llm()
-
         return prompt_text
 
     @property
@@ -284,6 +284,18 @@ def chunk_by_token_count(self, text, chunk_size, include_context=False):
         return chunks
 
     async def summarize_chunks(self, chunks):
+        @backoff.on_exception(
+            backoff.expo,
+            Exception,
+            max_tries=10,
+            factor=2,
+            on_backoff=lambda details: logger.warning(
+                f"Retrying summarization due to error: {details}"
+            ),
+            on_giveup=lambda details: logger.error(
+                f"Failed to summarize chunk after {details['tries']} attempts: {details}"
+            ),
+        )
         async def summarize_chunk(chunk):
             loop = asyncio.get_event_loop()
             forward_fn = functools.partial(

Diff for: src/hierarchical_OLD_BASIC.py

@@ -10,8 +10,9 @@
 import tempfile
 
 import asyncio
-import functools
 import nest_asyncio
+import backoff
+import functools
 
 
 def always_get_an_event_loop() -> asyncio.AbstractEventLoop:
@@ -230,6 +231,18 @@ def chunk_by_token_count(self, text, chunk_size, include_context=False):
         return chunks
 
     async def summarize_chunks(self, chunks):
+        @backoff.on_exception(
+            backoff.expo,
+            Exception,
+            max_tries=10,
+            factor=2,
+            on_backoff=lambda details: logger.warning(
+                f"Retrying summarization due to error: {details}"
+            ),
+            on_giveup=lambda details: logger.error(
+                f"Failed to summarize chunk after {details['tries']} attempts: {details}"
+            ),
+        )
         async def summarize_chunk(chunk):
             loop = asyncio.get_event_loop()
             forward_fn = functools.partial(