DEV: Robust check for embeddings enabled (#1116)

This commit is contained in:
Roman Rizzi 2025-02-06 12:18:55 -03:00 committed by GitHub
parent 83c7919856
commit e52045ebdc
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
13 changed files with 32 additions and 23 deletions

View File

@ -20,7 +20,7 @@ module DiscourseAi
def upload_file
file = params[:file] || params[:files].first
if !SiteSetting.ai_embeddings_enabled?
if !DiscourseAi::Embeddings.enabled?
raise Discourse::InvalidAccess.new("Embeddings not enabled")
end

View File

@ -5,7 +5,7 @@ module Jobs
sidekiq_options queue: "low"
def execute(args)
return unless SiteSetting.ai_embeddings_enabled
return unless DiscourseAi::Embeddings.enabled?
return if args[:target_type].blank? || args[:target_id].blank?
target = args[:target_type].constantize.find_by_id(args[:target_id])
return if target.nil? || target.deleted_at.present?

View File

@ -7,7 +7,7 @@ module Jobs
cluster_concurrency 1
def execute(args)
return unless SiteSetting.ai_embeddings_enabled
return unless DiscourseAi::Embeddings.enabled?
limit = SiteSetting.ai_embeddings_backfill_batch_size

View File

@ -11,7 +11,7 @@ class RagDocumentFragment < ActiveRecord::Base
def link_target_and_uploads(target, upload_ids)
return if target.blank?
return if upload_ids.blank?
return if !SiteSetting.ai_embeddings_enabled?
return if !DiscourseAi::Embeddings.enabled?
UploadReference.ensure_exist!(upload_ids: upload_ids, target: target)
@ -27,7 +27,7 @@ class RagDocumentFragment < ActiveRecord::Base
def update_target_uploads(target, upload_ids)
return if target.blank?
return if !SiteSetting.ai_embeddings_enabled?
return if !DiscourseAi::Embeddings.enabled?
if upload_ids.blank?
RagDocumentFragment.where(target: target).destroy_all

View File

@ -196,7 +196,7 @@ module DiscourseAi
)
plugin.on(:site_setting_changed) do |name, old_value, new_value|
if name == :ai_embeddings_selected_model && SiteSetting.ai_embeddings_enabled? &&
if name == :ai_embeddings_selected_model && DiscourseAi::Embeddings.enabled? &&
new_value != old_value
RagDocumentFragment.delete_all
UploadReference

View File

@ -295,7 +295,7 @@ module DiscourseAi
upload_refs =
UploadReference.where(target_id: id, target_type: "AiPersona").pluck(:upload_id)
return nil if !SiteSetting.ai_embeddings_enabled?
return nil if !DiscourseAi::Embeddings.enabled?
return nil if conversation_context.blank? || upload_refs.blank?
latest_interactions =

View File

@ -9,7 +9,7 @@ module DiscourseAi
def categories
return [] if @text.blank?
return [] unless SiteSetting.ai_embeddings_enabled
return [] if !DiscourseAi::Embeddings.enabled?
candidates = nearest_neighbors(limit: 100)
return [] if candidates.empty?
@ -51,7 +51,7 @@ module DiscourseAi
def tags
return [] if @text.blank?
return [] unless SiteSetting.ai_embeddings_enabled
return [] if !DiscourseAi::Embeddings.enabled?
candidates = nearest_neighbors(limit: 100)
return [] if candidates.empty?

View File

@ -8,20 +8,11 @@ module DiscourseAi
end
def valid_value?(val)
if val.blank?
@module_enabled = SiteSetting.ai_embeddings_enabled
!@module_enabled
else
EmbeddingDefinition.exists?(id: val).tap { |def_exists| @invalid_option = !def_exists }
end
val.present? || !SiteSetting.ai_embeddings_enabled
end
def error_message
return I18n.t("discourse_ai.embeddings.configuration.disable_embeddings") if @module_enabled
return I18n.t("discourse_ai.embeddings.configuration.invalid_config") if @invalid_option
""
I18n.t("discourse_ai.embeddings.configuration.disable_embeddings")
end
end
end

10
lib/embeddings.rb Normal file
View File

@ -0,0 +1,10 @@
# frozen_string_literal: true
module DiscourseAi
module Embeddings
def self.enabled?
SiteSetting.ai_embeddings_enabled && SiteSetting.ai_embeddings_selected_model.present? &&
EmbeddingDefinition.exists?(id: SiteSetting.ai_embeddings_selected_model)
end
end
end

View File

@ -56,7 +56,7 @@ module DiscourseAi
# embeddings generation.
callback =
Proc.new do |target|
if SiteSetting.ai_embeddings_enabled &&
if DiscourseAi::Embeddings.enabled? &&
(target.is_a?(Topic) || SiteSetting.ai_embeddings_per_post_enabled)
Jobs.enqueue(
:generate_embeddings,

View File

@ -3,6 +3,10 @@
describe DiscourseAi::Embeddings::EntryPoint do
fab!(:user) { Fabricate(:user, refresh_auto_groups: true) }
fab!(:embedding_definition)
before { SiteSetting.ai_embeddings_selected_model = embedding_definition.id }
describe "registering event callbacks" do
context "when creating a topic" do
let(:creator) do

View File

@ -6,7 +6,10 @@ RSpec.describe RagDocumentFragment do
fab!(:upload_2) { Fabricate(:upload) }
fab!(:vector_def) { Fabricate(:embedding_definition) }
before { SiteSetting.ai_embeddings_enabled = true }
before do
SiteSetting.ai_embeddings_selected_model = vector_def.id
SiteSetting.ai_embeddings_enabled = true
end
describe ".link_uploads_and_persona" do
it "does nothing if there is no persona" do

View File

@ -3,10 +3,11 @@
RSpec.describe DiscourseAi::Admin::AiPersonasController do
fab!(:admin)
fab!(:ai_persona)
fab!(:embedding_definition)
before do
sign_in(admin)
SiteSetting.ai_embeddings_selected_model = embedding_definition.id
SiteSetting.ai_embeddings_enabled = true
end