From b6ca594782c6a9b4c55fc65342263f816f306216 Mon Sep 17 00:00:00 2001
From: Tash Boyse <57753415+nboyse@users.noreply.github.com>
Date: Mon, 17 Feb 2025 17:29:14 +0000
Subject: [PATCH] fix: implement exponential backoff (#94)

---
 redbox-core/redbox/graph/nodes/processes.py | 27 ++++++++++++++++++---
 1 file changed, 23 insertions(+), 4 deletions(-)

diff --git a/redbox-core/redbox/graph/nodes/processes.py b/redbox-core/redbox/graph/nodes/processes.py
index 3f3cc6b8..7d620d69 100644
--- a/redbox-core/redbox/graph/nodes/processes.py
+++ b/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"]