fix: implement exponential backoff to resolve EventStreamError

redbox-core/redbox/graph/nodes/processes.py

@@ -2,11 +2,15 @@
 import logging
 import re
 import textwrap
+import time
+from random import uniform
 from collections.abc import Callable
 from functools import reduce
 from typing import Any, Iterable
 from uuid import uuid4
 
+from botocore.exceptions import EventStreamError
+
 from langchain.schema import StrOutputParser
 from langchain_core.callbacks.manager import dispatch_custom_event
 from langchain_core.documents import Document
@@ -83,11 +87,13 @@ def build_merge_pattern(
 
     When combined with group send, with combine all Documents and use the metadata of the first.
 
-    When used without a send, the first Document receieved defines the metadata.
+    When used without a send, the first Document received defines the metadata.
 
     If tools are supplied, can also set state["tool_calls"].
     """
     tokeniser = get_tokeniser()
+    MAX_RETRIES = 5
+    BACKOFF_FACTOR = 1.5
 
     @RunnableLambda
     def _merge(state: RedboxState) -> dict[str, Any]:
@@ -106,9 +112,22 @@ def _merge(state: RedboxState) -> dict[str, Any]:
             ),
         )
 
-        merge_response = build_llm_chain(
-            prompt_set=prompt_set, llm=llm, final_response_chain=final_response_chain
-        ).invoke(merge_state)
+        retries = 0
+        while retries < MAX_RETRIES:
+            try:
+                merge_response = build_llm_chain(
+                    prompt_set=prompt_set, llm=llm, final_response_chain=final_response_chain
+                ).invoke(merge_state)
+
+                # if reaches a successful citation, exit the loop
+                break
+
+            except EventStreamError as e:
+                retries += 1
+                if retries >= MAX_RETRIES:
+                    raise e
+                wait_time = BACKOFF_FACTOR**retries + uniform(0, 1)
+                time.sleep(wait_time)
 
         merged_document.page_content = merge_response["messages"][-1].content
         request_metadata = merge_response["metadata"]