FIX: Restore ability to fold summaries, which was accidentally removed (#700)

lib/summarization/strategies/fold_content.rb

@@ -21,27 +21,122 @@ module DiscourseAi
 
           llm = DiscourseAi::Completions::Llm.proxy(completion_model.model_name)
 
-          summary_content =
-            content[:contents].map { |c| { ids: [c[:id]], summary: format_content_item(c) } }
+          initial_chunks =
+            rebalance_chunks(
+              llm.tokenizer,
+              content[:contents].map { |c| { ids: [c[:id]], summary: format_content_item(c) } },
+            )
 
-          {
-            summary:
-              summarize_single(llm, summary_content.first[:summary], user, opts, &on_partial_blk),
-          }
+          # Special case where we can do all the summarization in one pass.
+          if initial_chunks.length == 1
+            {
+              summary:
+                summarize_single(llm, initial_chunks.first[:summary], user, opts, &on_partial_blk),
+              chunks: [],
+            }
+          else
+            summarize_chunks(llm, initial_chunks, user, opts, &on_partial_blk)
+          end
         end
 
         private
 
+        def summarize_chunks(llm, chunks, user, opts, &on_partial_blk)
+          # Safely assume we always have more than one chunk.
+          summarized_chunks = summarize_in_chunks(llm, chunks, user, opts)
+          total_summaries_size =
+            llm.tokenizer.size(summarized_chunks.map { |s| s[:summary].to_s }.join)
+
+          if total_summaries_size < completion_model.available_tokens
+            # Chunks are small enough, we can concatenate them.
+            {
+              summary:
+                concatenate_summaries(
+                  llm,
+                  summarized_chunks.map { |s| s[:summary] },
+                  user,
+                  &on_partial_blk
+                ),
+              chunks: summarized_chunks,
+            }
+          else
+            # We have summarized chunks but we can't concatenate them yet. Split them into smaller summaries and summarize again.
+            rebalanced_chunks = rebalance_chunks(llm.tokenizer, summarized_chunks)
+
+            summarize_chunks(llm, rebalanced_chunks, user, opts, &on_partial_blk)
+          end
+        end
+
         def format_content_item(item)
           "(#{item[:id]} #{item[:poster]} said: #{item[:text]} "
         end
 
+        def rebalance_chunks(tokenizer, chunks)
+          section = { ids: [], summary: "" }
+
+          chunks =
+            chunks.reduce([]) do |sections, chunk|
+              if tokenizer.can_expand_tokens?(
+                   section[:summary],
+                   chunk[:summary],
+                   completion_model.available_tokens,
+                 )
+                section[:summary] += chunk[:summary]
+                section[:ids] = section[:ids].concat(chunk[:ids])
+              else
+                sections << section
+                section = chunk
+              end
+
+              sections
+            end
+
+          chunks << section if section[:summary].present?
+
+          chunks
+        end
+
         def summarize_single(llm, text, user, opts, &on_partial_blk)
           prompt = summarization_prompt(text, opts)
 
           llm.generate(prompt, user: user, feature_name: "summarize", &on_partial_blk)
         end
 
+        def summarize_in_chunks(llm, chunks, user, opts)
+          chunks.map do |chunk|
+            prompt = summarization_prompt(chunk[:summary], opts)
+
+            chunk[:summary] = llm.generate(
+              prompt,
+              user: user,
+              max_tokens: 300,
+              feature_name: "summarize",
+            )
+            chunk
+          end
+        end
+
+        def concatenate_summaries(llm, summaries, user, &on_partial_blk)
+          prompt = DiscourseAi::Completions::Prompt.new(<<~TEXT.strip)
+            You are a summarization bot that effectively concatenates disjoint summaries, creating a cohesive narrative.
+            The narrative you create is in the form of one or multiple paragraphs.
+            Your reply MUST BE a single concatenated summary using the summaries I'll provide to you.
+            I'm NOT interested in anything other than the concatenated summary, don't include additional text or comments.
+            You understand and generate Discourse forum Markdown.
+            You format the response, including links, using Markdown.
+          TEXT
+
+          prompt.push(type: :user, content: <<~TEXT.strip)
+            THESE are the summaries, each one separated by a newline, all of them inside <input></input> XML tags:
+
+            <input>
+              #{summaries.join("\n")}
+            </input>
+          TEXT
+
+          llm.generate(prompt, user: user, &on_partial_blk)
+        end
+
         def summarization_prompt(input, opts)
           insts = +<<~TEXT
             You are an advanced summarization bot that generates concise, coherent summaries of provided text.

spec/lib/modules/summarization/strategies/fold_content_spec.rb

@@ -32,5 +32,37 @@ RSpec.describe DiscourseAi::Summarization::Strategies::FoldContent do
         expect(result[:summary]).to eq(single_summary)
       end
     end
+
+    context "when the content to summarize doesn't fit in a single call" do
+      it "summarizes each chunk and then concatenates them" do
+        content[:contents] << { poster: "asd2", id: 2, text: summarize_text }
+
+        result =
+          DiscourseAi::Completions::Llm.with_prepared_responses(
+            [single_summary, single_summary, concatenated_summary],
+          ) { |spy| strategy.summarize(content, user).tap { expect(spy.completions).to eq(3) } }
+
+        expect(result[:summary]).to eq(concatenated_summary)
+      end
+
+      it "keeps splitting into chunks until the content fits into a single call to create a cohesive narrative" do
+        content[:contents] << { poster: "asd2", id: 2, text: summarize_text }
+        max_length_response = "(1 asd said: This is a text "
+        chunk_of_chunks = "I'm smol"
+
+        result =
+          DiscourseAi::Completions::Llm.with_prepared_responses(
+            [
+              max_length_response,
+              max_length_response,
+              chunk_of_chunks,
+              chunk_of_chunks,
+              concatenated_summary,
+            ],
+          ) { |spy| strategy.summarize(content, user).tap { expect(spy.completions).to eq(5) } }
+
+        expect(result[:summary]).to eq(concatenated_summary)
+      end
+    end
   end
 end