Merge pull request #80 from uktrade/fix-summarisation-backoff-dev

Fix summarisation and exponential backoff for summarisation
uktrade · Feb 12, 2025 · e0d5a6a · e0d5a6a
2 parents 3473b8c + 4f97cb9
commit e0d5a6a
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/redbox-core/redbox/graph/root.py b/redbox-core/redbox/graph/root.py
@@ -35,7 +35,7 @@
 from redbox.models.chat import ChatRoute, ErrorRoute
 from redbox.models.graph import ROUTABLE_KEYWORDS, RedboxActivityEvent
 from redbox.transform import structure_documents_by_file_name, structure_documents_by_group_and_indices
-
+from langgraph.pregel import RetryPolicy
 
 def get_self_route_graph(retriever: VectorStoreRetriever, prompt_set: PromptSet, debug: bool = False):
     builder = StateGraph(RedboxState)
@@ -237,17 +237,20 @@ def get_chat_with_documents_graph(
     builder.add_node(
         "p_summarise_each_document",
         build_merge_pattern(prompt_set=PromptSet.ChatwithDocsMapReduce),
+        retry=RetryPolicy(max_attempts=3)
     )
     builder.add_node(
         "p_summarise_document_by_document",
         build_merge_pattern(prompt_set=PromptSet.ChatwithDocsMapReduce),
+        retry=RetryPolicy(max_attempts=3)
     )
     builder.add_node(
         "p_summarise",
         build_stuff_pattern(
             prompt_set=PromptSet.ChatwithDocs,
             final_response_chain=True,
         ),
+        retry=RetryPolicy(max_attempts=3)
     )
     builder.add_node("p_clear_documents", clear_documents_process)
     builder.add_node(

diff --git a/redbox-core/redbox/models/prompts.py b/redbox-core/redbox/models/prompts.py
@@ -134,6 +134,7 @@
     "2) Avoid repetition,\n"
     "3) Ensure the summary is easy to understand,\n"
     "4) Maintain the original context and meaning.\n"
+    "5) Do not start your answer by saying: here is a summary...Go straight to the point."
 )
 
 REDUCE_SYSTEM_PROMPT = (