diff --git a/lib/summarization/strategies/fold_content.rb b/lib/summarization/strategies/fold_content.rb
index 38defd3a..36f76276 100644
--- a/lib/summarization/strategies/fold_content.rb
+++ b/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 XML tags:
+
+
+ #{summaries.join("\n")}
+
+ 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.
diff --git a/spec/lib/modules/summarization/strategies/fold_content_spec.rb b/spec/lib/modules/summarization/strategies/fold_content_spec.rb
index df7f1298..0333dd45 100644
--- a/spec/lib/modules/summarization/strategies/fold_content_spec.rb
+++ b/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