diff --git a/app/models/reviewable_ai_chat_message.rb b/app/models/reviewable_ai_chat_message.rb
index 4751c31e..2639b3b4 100644
--- a/app/models/reviewable_ai_chat_message.rb
+++ b/app/models/reviewable_ai_chat_message.rb
@@ -17,7 +17,7 @@ class ReviewableAiChatMessage < Reviewable
   end
 
   def chat_message
-    @chat_message ||= (target || ChatMessage.with_deleted.find_by(id: target_id))
+    @chat_message ||= (target || Chat::Message.with_deleted.find_by(id: target_id))
   end
 
   def chat_message_creator
diff --git a/lib/modules/toxicity/jobs/regular/toxicity_classify_chat_message.rb b/lib/modules/toxicity/jobs/regular/toxicity_classify_chat_message.rb
index 7f538ee5..0ea5927a 100644
--- a/lib/modules/toxicity/jobs/regular/toxicity_classify_chat_message.rb
+++ b/lib/modules/toxicity/jobs/regular/toxicity_classify_chat_message.rb
@@ -7,7 +7,7 @@ module ::Jobs
 
       return if (chat_message_id = args[:chat_message_id]).blank?
 
-      chat_message = ChatMessage.find_by(id: chat_message_id)
+      chat_message = ::Chat::Message.find_by(id: chat_message_id)
       return if chat_message&.message.blank?
 
       DiscourseAi::ChatMessageClassificator.new(
diff --git a/lib/modules/toxicity/toxicity_classification.rb b/lib/modules/toxicity/toxicity_classification.rb
index 604fee4e..4425e63e 100644
--- a/lib/modules/toxicity/toxicity_classification.rb
+++ b/lib/modules/toxicity/toxicity_classification.rb
@@ -59,7 +59,7 @@ module DiscourseAi
       end
 
       def content_of(target_to_classify)
-        return target_to_classify.message if target_to_classify.is_a?(ChatMessage)
+        return target_to_classify.message if target_to_classify.is_a?(Chat::Message)
 
         if target_to_classify.post_number == 1
           "#{target_to_classify.topic.title}\n#{target_to_classify.raw}"
diff --git a/lib/shared/classificator.rb b/lib/shared/classificator.rb
index f56b893f..890f28b8 100644
--- a/lib/shared/classificator.rb
+++ b/lib/shared/classificator.rb
@@ -59,7 +59,7 @@ module ::DiscourseAi
           {
             model_used: model_name,
             target_id: target.id,
-            target_type: target.class.name,
+            target_type: target.class.sti_name,
             classification_type: classification_model.type,
             classification: classifications,
             updated_at: DateTime.now,
diff --git a/spec/lib/modules/toxicity/entry_point_spec.rb b/spec/lib/modules/toxicity/entry_point_spec.rb
index c5a6a609..75009139 100644
--- a/spec/lib/modules/toxicity/entry_point_spec.rb
+++ b/spec/lib/modules/toxicity/entry_point_spec.rb
@@ -39,7 +39,7 @@ describe DiscourseAi::Toxicity::EntryPoint do
     context "when creating a chat message" do
       let(:public_chat_channel) { Fabricate(:chat_channel) }
       let(:creator) do
-        Chat::ChatMessageCreator.new(
+        Chat::MessageCreator.new(
           chat_channel: public_chat_channel,
           user: user,
           content: "This is my new test",
@@ -54,7 +54,7 @@ describe DiscourseAi::Toxicity::EntryPoint do
     context "when editing a chat message" do
       let(:chat_message) { Fabricate(:chat_message) }
       let(:updater) do
-        Chat::ChatMessageUpdater.new(
+        Chat::MessageUpdater.new(
           guardian: Guardian.new(chat_message.user),
           chat_message: chat_message,
           new_content: "This is my updated message",
diff --git a/spec/models/reviewable_ai_chat_message_spec.rb b/spec/models/reviewable_ai_chat_message_spec.rb
index a4715219..98aa9a0f 100644
--- a/spec/models/reviewable_ai_chat_message_spec.rb
+++ b/spec/models/reviewable_ai_chat_message_spec.rb
@@ -21,7 +21,7 @@ RSpec.describe ReviewableAiChatMessage, type: :model do
     reviewable.perform(moderator, :agree_and_delete)
 
     expect(reviewable).to be_approved
-    expect(ChatMessage.with_deleted.find_by(id: chat_message_id).deleted_at).to be_present
+    expect(Chat::Message.with_deleted.find_by(id: chat_message_id).deleted_at).to be_present
   end
 
   it "agree_and_restore agrees with the flag and restores the message" do