FEATURE: Try to automatically handle larger embedding indexes (#403)

* FEATURE: Try to automatically handle larger embedding indexes

* linteeeeeeeer
This commit is contained in:
Rafael dos Santos Silva 2024-01-05 09:56:28 -03:00 committed by GitHub
parent dd42a4e47b
commit 6fc1c9f7a6
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 24 additions and 11 deletions

View File

@ -80,20 +80,33 @@ module DiscourseAi
end end
def create_index!(table_name, memory, lists, probes) def create_index!(table_name, memory, lists, probes)
tries = 0
index_name = index_name(table_name) index_name = index_name(table_name)
DB.exec("SET work_mem TO '#{memory}';") DB.exec("SET work_mem TO '#{memory}';")
DB.exec("SET maintenance_work_mem TO '#{memory}';") DB.exec("SET maintenance_work_mem TO '#{memory}';")
DB.exec(<<~SQL) begin
DROP INDEX IF EXISTS #{index_name}; DB.exec(<<~SQL)
CREATE INDEX IF NOT EXISTS DROP INDEX IF EXISTS #{index_name};
#{index_name} CREATE INDEX IF NOT EXISTS
ON #{index_name}
#{table_name} ON
USING #{table_name}
ivfflat (embeddings #{pg_index_type}) USING
WITH ivfflat (embeddings #{pg_index_type})
(lists = #{lists}); WITH
SQL (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("COMMENT ON INDEX #{index_name} IS '#{Time.now.to_i}';")
DB.exec("RESET work_mem;") DB.exec("RESET work_mem;")
DB.exec("RESET maintenance_work_mem;") DB.exec("RESET maintenance_work_mem;")