From 6fc1c9f7a68c3242ab6b94bdc1d30bcc266501ea Mon Sep 17 00:00:00 2001
From: Rafael dos Santos Silva <xfalcox@gmail.com>
Date: Fri, 5 Jan 2024 09:56:28 -0300
Subject: [PATCH] FEATURE: Try to automatically handle larger embedding indexes
 (#403)

* FEATURE: Try to automatically handle larger embedding indexes

* linteeeeeeeer
---
 lib/embeddings/vector_representations/base.rb | 35 +++++++++++++------
 1 file changed, 24 insertions(+), 11 deletions(-)

diff --git a/lib/embeddings/vector_representations/base.rb b/lib/embeddings/vector_representations/base.rb
index fc17bafe..6b7f8ff0 100644
--- a/lib/embeddings/vector_representations/base.rb
+++ b/lib/embeddings/vector_representations/base.rb
@@ -80,20 +80,33 @@ module DiscourseAi
         end
 
         def create_index!(table_name, memory, lists, probes)
+          tries = 0
           index_name = index_name(table_name)
           DB.exec("SET work_mem TO '#{memory}';")
           DB.exec("SET maintenance_work_mem TO '#{memory}';")
-          DB.exec(<<~SQL)
-            DROP INDEX IF EXISTS #{index_name};
-            CREATE INDEX IF NOT EXISTS
-              #{index_name}
-            ON
-              #{table_name}
-            USING
-              ivfflat (embeddings #{pg_index_type})
-            WITH
-              (lists = #{lists});
-          SQL
+          begin
+            DB.exec(<<~SQL)
+              DROP INDEX IF EXISTS #{index_name};
+              CREATE INDEX IF NOT EXISTS
+                #{index_name}
+              ON
+                #{table_name}
+              USING
+                ivfflat (embeddings #{pg_index_type})
+              WITH
+                (lists = #{lists});
+            SQL
+          rescue PG::ProgramLimitExceeded => e
+            parsed_error = e.message.match(/memory required is (\d+ [A-Z]{2}), ([a-z_]+)/)
+            if parsed_error[1].present? && parsed_error[2].present?
+              DB.exec("SET #{parsed_error[2]} TO '#{parsed_error[1].tr(" ", "")}';")
+              tries += 1
+              retry if tries < 3
+            else
+              raise e
+            end
+          end
+
           DB.exec("COMMENT ON INDEX #{index_name} IS '#{Time.now.to_i}';")
           DB.exec("RESET work_mem;")
           DB.exec("RESET maintenance_work_mem;")