Rafael dos Santos Silva
8fcba12fae
This allows admins to configure services with multiple backends using DNS SRV records. This PR also adds support for shared secret auth via headers for TEI and vLLM endpoints, so they are inline with the other ones.
97 lines
2.7 KiB
Ruby
97 lines
2.7 KiB
Ruby
# frozen_string_literal: true
|
|
|
|
module DiscourseAi
|
|
module Nsfw
|
|
class Classification
|
|
def type
|
|
:nsfw
|
|
end
|
|
|
|
def can_classify?(target)
|
|
content_of(target).present?
|
|
end
|
|
|
|
def get_verdicts(classification_data)
|
|
classification_data
|
|
.map do |model_name, classifications|
|
|
verdict =
|
|
classifications.values.any? do |data|
|
|
send("#{model_name}_verdict?", data.except(:neutral, :target_classified_type))
|
|
end
|
|
|
|
[model_name, verdict]
|
|
end
|
|
.to_h
|
|
end
|
|
|
|
def should_flag_based_on?(verdicts)
|
|
return false if !SiteSetting.ai_nsfw_flag_automatically
|
|
|
|
verdicts.values.any?
|
|
end
|
|
|
|
def request(target_to_classify)
|
|
uploads_to_classify = content_of(target_to_classify)
|
|
|
|
available_models.reduce({}) do |memo, model|
|
|
memo[model] = uploads_to_classify.reduce({}) do |upl_memo, upload|
|
|
classification =
|
|
evaluate_with_model(model, upload).merge(target_classified_type: upload.class.name)
|
|
|
|
# 415 denotes that the image is not supported by the model, so we skip it
|
|
upl_memo[upload.id] = classification if classification.dig(:status) != 415
|
|
|
|
upl_memo
|
|
end
|
|
|
|
memo
|
|
end
|
|
end
|
|
|
|
private
|
|
|
|
def evaluate_with_model(model, upload)
|
|
upload_url = Discourse.store.cdn_url(upload.url)
|
|
upload_url = "#{Discourse.base_url_no_prefix}#{upload_url}" if upload_url.starts_with?("/")
|
|
|
|
DiscourseAi::Inference::DiscourseClassifier.perform!(
|
|
"#{endpoint}/api/v1/classify",
|
|
model,
|
|
upload_url,
|
|
SiteSetting.ai_nsfw_inference_service_api_key,
|
|
)
|
|
end
|
|
|
|
def available_models
|
|
SiteSetting.ai_nsfw_models.split("|")
|
|
end
|
|
|
|
def content_of(target_to_classify)
|
|
target_to_classify.uploads.to_a.select { |u| FileHelper.is_supported_image?(u.url) }
|
|
end
|
|
|
|
def opennsfw2_verdict?(clasification)
|
|
clasification.values.first.to_i >= SiteSetting.ai_nsfw_flag_threshold_general
|
|
end
|
|
|
|
def nsfw_detector_verdict?(classification)
|
|
classification.any? do |key, value|
|
|
value.to_i >= SiteSetting.send("ai_nsfw_flag_threshold_#{key}")
|
|
end
|
|
end
|
|
|
|
def endpoint
|
|
if SiteSetting.ai_nsfw_inference_service_api_endpoint_srv.present?
|
|
service =
|
|
DiscourseAi::Utils::DnsSrv.lookup(
|
|
SiteSetting.ai_nsfw_inference_service_api_endpoint_srv,
|
|
)
|
|
"https://#{service.target}:#{service.port}"
|
|
else
|
|
SiteSetting.ai_nsfw_inference_service_api_endpoint
|
|
end
|
|
end
|
|
end
|
|
end
|
|
end
|