From a91aab4b3c1eae9c11107383cd98fc28153b5c11 Mon Sep 17 00:00:00 2001
From: Johannes Wesch <johannes.wesch@aleph-alpha.com>
Date: Tue, 6 Feb 2024 09:28:03 +0100
Subject: [PATCH] added break for recursive summary when number of partial
 summary doesn't change

---
 .../summarize/recursive_summarize.py          |  8 +++++---
 .../summarize/test_recursive_summarize.py     | 20 +++++++++++++++++++
 2 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/src/intelligence_layer/use_cases/summarize/recursive_summarize.py b/src/intelligence_layer/use_cases/summarize/recursive_summarize.py
index 80cebabdd..0c2e4e7ab 100644
--- a/src/intelligence_layer/use_cases/summarize/recursive_summarize.py
+++ b/src/intelligence_layer/use_cases/summarize/recursive_summarize.py
@@ -27,20 +27,22 @@ def __init__(
     def do_run(
         self, input: LongContextSummarizeInput, task_span: TaskSpan
     ) -> SummarizeOutput:
+        num_partial_summaries = 0
         text = input.text
-        loop_count = 0
         while True:
             summarize_output = self.long_context_summarize_task.run(
                 LongContextSummarizeInput(text=text, language=input.language), task_span
             )
+            if num_partial_summaries == len(summarize_output.partial_summaries):
+                break
+            num_partial_summaries = len(summarize_output.partial_summaries)
+
             num_generated_tokens = 0
             text = ""
             for partial_summary in summarize_output.partial_summaries:
                 num_generated_tokens += partial_summary.generated_tokens
                 text += partial_summary.summary + "\n"
 
-            loop_count += 1
-
             if len(summarize_output.partial_summaries) == 1:
                 break
 
diff --git a/tests/use_cases/summarize/test_recursive_summarize.py b/tests/use_cases/summarize/test_recursive_summarize.py
index b0e95df1e..c8cb401c3 100644
--- a/tests/use_cases/summarize/test_recursive_summarize.py
+++ b/tests/use_cases/summarize/test_recursive_summarize.py
@@ -3,6 +3,9 @@
 
 from aleph_alpha_client import Client, CompletionRequest, CompletionResponse
 from pytest import fixture
+from intelligence_layer.connectors.limited_concurrency_client import (
+    AlephAlphaClientProtocol,
+)
 
 from intelligence_layer.core import NoOpTracer
 from intelligence_layer.use_cases import (
@@ -10,6 +13,9 @@
     LongContextSummarizeInput,
     RecursiveSummarize,
 )
+from intelligence_layer.use_cases.summarize.steerable_long_context_summarize import (
+    SteerableLongContextSummarize,
+)
 
 
 class RecursiveCountingClient(Client):
@@ -52,6 +58,20 @@ def test_recursive_summarize_stops_when_hitting_max_tokens(
     assert "new orleans" in output.summary.lower()
 
 
+def test_recursive_summarize_stops_when_num_partial_summaries_stays_same(
+    client: AlephAlphaClientProtocol,
+) -> None:
+    max_tokens = None
+    long_context_high_compression_summarize = SteerableLongContextSummarize(
+        client, model="luminous-base", max_generated_tokens=75, max_tokens_per_chunk=145
+    )
+    input = LongContextSummarizeInput(text=short_text, max_tokens=max_tokens)
+    task = RecursiveSummarize(long_context_high_compression_summarize)
+    output = task.run(input, NoOpTracer())
+
+    assert output.generated_tokens > 145
+
+
 def test_recursive_summarize_stops_after_one_chunk(
     recursive_counting_client: RecursiveCountingClient,
 ) -> None: