Discource-C
/
discourse-ai
mirror of https://github.com/discourse/discourse-ai.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
lib/embeddings/vector_representations/base.rb
View File

@ -80,9 +80,11 @@ 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}';")
begin
DB.exec(<<~SQL) DB.exec(<<~SQL)
DROP INDEX IF EXISTS #{index_name}; DROP INDEX IF EXISTS #{index_name};
CREATE INDEX IF NOT EXISTS CREATE INDEX IF NOT EXISTS
@ -94,6 +96,17 @@ module DiscourseAi
WITH WITH
(lists = #{lists}); (lists = #{lists});
SQL 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;")