From 87425350244e5ec911c7b132a6fbf3cfed5acb59 Mon Sep 17 00:00:00 2001
From: Rafael dos Santos Silva <xfalcox@gmail.com>
Date: Tue, 13 Jun 2023 15:23:48 -0300
Subject: [PATCH] FEATURE: Allow using large context OpenAI models for
 summarization (#86)

---
 lib/modules/summarization/entry_point.rb        | 2 ++
 lib/modules/summarization/strategies/open_ai.rb | 7 ++++++-
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/lib/modules/summarization/entry_point.rb b/lib/modules/summarization/entry_point.rb
index c4be1baf..93b27ea0 100644
--- a/lib/modules/summarization/entry_point.rb
+++ b/lib/modules/summarization/entry_point.rb
@@ -12,7 +12,9 @@ module DiscourseAi
       def inject_into(plugin)
         [
           Strategies::OpenAi.new("gpt-4"),
+          Strategies::OpenAi.new("gpt-4-32k"),
           Strategies::OpenAi.new("gpt-3.5-turbo"),
+          Strategies::OpenAi.new("gpt-3.5-turbo-16k"),
           Strategies::DiscourseAi.new("bart-large-cnn-samsum"),
           Strategies::DiscourseAi.new("flan-t5-base-samsum"),
           Strategies::DiscourseAi.new("long-t5-tglobal-base-16384-book-summary"),
diff --git a/lib/modules/summarization/strategies/open_ai.rb b/lib/modules/summarization/strategies/open_ai.rb
index 2da7b508..35011b66 100644
--- a/lib/modules/summarization/strategies/open_ai.rb
+++ b/lib/modules/summarization/strategies/open_ai.rb
@@ -41,7 +41,12 @@ module DiscourseAi
         private
 
         def max_length
-          lengths = { "gpt-3.5-turbo" => 4096, "gpt-4" => 8192 }
+          lengths = {
+            "gpt-3.5-turbo" => 4096,
+            "gpt-4" => 8192,
+            "gpt-3.5-turbo-16k" => 16_384,
+            "gpt-4-32k" => 32_768,
+          }
 
           lengths[model]
         end