diff --git a/db/migrate/20260714230000_replace_captain_assistant_response_embedding_index_with_cosine.rb b/db/migrate/20260714230000_replace_captain_assistant_response_embedding_index_with_cosine.rb new file mode 100644 index 000000000..be2f4b449 --- /dev/null +++ b/db/migrate/20260714230000_replace_captain_assistant_response_embedding_index_with_cosine.rb @@ -0,0 +1,26 @@ +class ReplaceCaptainAssistantResponseEmbeddingIndexWithCosine < ActiveRecord::Migration[7.1] + disable_ddl_transaction! + + OLD_INDEX_NAME = 'vector_idx_knowledge_entries_embedding'.freeze + NEW_INDEX_NAME = 'vector_idx_captain_assistant_responses_embedding_cosine'.freeze + + def up + add_index :captain_assistant_responses, + :embedding, + using: :ivfflat, + opclass: :vector_cosine_ops, + name: NEW_INDEX_NAME, + algorithm: :concurrently + remove_index :captain_assistant_responses, name: OLD_INDEX_NAME, algorithm: :concurrently + end + + def down + add_index :captain_assistant_responses, + :embedding, + using: :ivfflat, + opclass: :vector_l2_ops, + name: OLD_INDEX_NAME, + algorithm: :concurrently + remove_index :captain_assistant_responses, name: NEW_INDEX_NAME, algorithm: :concurrently + end +end diff --git a/db/schema.rb b/db/schema.rb index 43e7135b9..e616f2b4f 100644 --- a/db/schema.rb +++ b/db/schema.rb @@ -10,7 +10,7 @@ # # It's strongly recommended that you check this file into your version control system. -ActiveRecord::Schema[7.1].define(version: 2026_07_13_184351) do +ActiveRecord::Schema[7.1].define(version: 2026_07_14_230000) do # These extensions should be enabled to support this database enable_extension "pg_stat_statements" enable_extension "pg_trgm" @@ -361,7 +361,7 @@ ActiveRecord::Schema[7.1].define(version: 2026_07_13_184351) do t.index ["account_id"], name: "index_captain_assistant_responses_on_account_id" t.index ["assistant_id"], name: "index_captain_assistant_responses_on_assistant_id" t.index ["documentable_id", "documentable_type"], name: "idx_cap_asst_resp_on_documentable" - t.index ["embedding"], name: "vector_idx_knowledge_entries_embedding", using: :ivfflat + t.index ["embedding"], name: "vector_idx_captain_assistant_responses_embedding_cosine", opclass: :vector_cosine_ops, using: :ivfflat t.index ["status"], name: "index_captain_assistant_responses_on_status" end diff --git a/enterprise/app/models/captain/assistant_response.rb b/enterprise/app/models/captain/assistant_response.rb index db2ac058a..7289601c8 100644 --- a/enterprise/app/models/captain/assistant_response.rb +++ b/enterprise/app/models/captain/assistant_response.rb @@ -17,11 +17,11 @@ # # Indexes # -# idx_cap_asst_resp_on_documentable (documentable_id,documentable_type) -# index_captain_assistant_responses_on_account_id (account_id) -# index_captain_assistant_responses_on_assistant_id (assistant_id) -# index_captain_assistant_responses_on_status (status) -# vector_idx_knowledge_entries_embedding (embedding) USING ivfflat +# idx_cap_asst_resp_on_documentable (documentable_id,documentable_type) +# index_captain_assistant_responses_on_account_id (account_id) +# index_captain_assistant_responses_on_assistant_id (assistant_id) +# index_captain_assistant_responses_on_status (status) +# vector_idx_captain_assistant_responses_embedding_cosine (embedding) USING ivfflat # class Captain::AssistantResponse < ApplicationRecord self.table_name = 'captain_assistant_responses' @@ -48,7 +48,11 @@ class Captain::AssistantResponse < ApplicationRecord def self.search(query, account_id: nil) embedding = Captain::Llm::EmbeddingService.new(account_id: account_id).get_embedding(query) - nearest_neighbors(:embedding, embedding, distance: 'cosine').limit(5) + + transaction do + connection.execute("SET LOCAL ivfflat.iterative_scan = 'relaxed_order'") + nearest_neighbors(:embedding, embedding, distance: 'cosine').limit(5).load + end end private