No longer use Redis to help with cancelled jobs (#324)

app/controllers/messages_controller.rb

@@ -24,7 +24,7 @@ def index
     @streaming_message = Message.where(
       content_text: [nil, ""],
       cancelled_at: nil
-    ).find_by(id: redis.get("conversation-#{@conversation.id}-latest-assistant_message-id"))
+    ).find_by(id: @conversation.last_assistant_message_id)
   end
 
   def show
@@ -43,7 +43,7 @@ def create
 
     if @message.save
       after_create_assistant_reply = @message.conversation.latest_message_for_version(@message.version)
-      GetNextAIMessageJob.perform_later(after_create_assistant_reply.id, @assistant.id)
+      GetNextAIMessageJob.perform_later(current_user.id, after_create_assistant_reply.id, @assistant.id)
       redirect_to conversation_messages_path(@message.conversation, version: @message.version)
     else
       # what's the right flow for a failed message create? it's not this, but hacking it so tests pass until we have a plan
@@ -110,8 +110,4 @@ def message_params
     end
     modified_params
   end
-
-  def redis
-    RedisConnection.client
-  end
 end

app/jobs/get_next_ai_message_job.rb

@@ -12,12 +12,13 @@ def ai_backend
     end
   end
 
-  def perform(message_id, assistant_id, attempt = 1)
-    puts "\n### GetNextAIMessageJob.perform(#{message_id}, #{assistant_id}, #{attempt})" unless Rails.env.test?
+  def perform(user_id, message_id, assistant_id, attempt = 1)
+    puts "\n### GetNextAIMessageJob.perform(#{user_id}, #{message_id}, #{assistant_id}, #{attempt})" unless Rails.env.test?
 
-    @message      = Message.find_by(id: message_id)
+    @user         = User.find(user_id)
+    @message      = Message.find(message_id)
     @conversation = @message.conversation
-    @assistant    = Assistant.find_by(id: assistant_id)
+    @assistant    = Assistant.find(assistant_id)
     @prev_message = @conversation.messages.assistant.for_conversation_version(@message.version).find_by(index: @message.index-1)
 
     return false          if generation_was_cancelled? || message_is_populated?
@@ -153,19 +154,15 @@ def generation_was_cancelled?
 
   def message_cancelled?
     @message.cancelled? ||
-      (@cancel_counter > 1 && @message.id == redis.get("message-cancelled-id")&.to_i)
+      (@cancel_counter > 1 && @message.id == @user.reload.last_cancelled_message_id)
   end
 
   def newer_messages_in_conversation?
     @message != @conversation.latest_message_for_version(@message.version) ||
-      (@cancel_counter > 1 && @message.id != redis.get("conversation-#{@conversation.id}-latest-assistant_message-id")&.to_i)
+      (@cancel_counter > 1 && @message.id != @conversation.reload.last_assistant_message_id)
   end
 
   def message_is_populated?
     @message.content_text.present?
   end
-
-  def redis
-    RedisConnection.client
-  end
 end

app/models/conversation.rb

@@ -7,6 +7,7 @@ class Conversation < ApplicationRecord
   has_many :messages, dependent: :destroy
   has_many :runs, dependent: :destroy
   has_many :steps, dependent: :destroy
+  belongs_to :last_assistant_message, class_name: "Message", optional: true
 
   after_touch :set_title_async, if: -> { title.blank? && messages.count >= 2 }
 

app/models/message.rb

@@ -5,6 +5,7 @@ class Message < ApplicationRecord
   belongs_to :conversation
   belongs_to :content_document, class_name: "Document", optional: true
   belongs_to :run, optional: true
+  has_one :latest_assistant_message_for, class_name: "Conversation", foreign_key: :last_assistant_message_id, dependent: :nullify
 
   enum role: %w[user assistant].index_by(&:to_sym)
 
@@ -20,7 +21,7 @@ class Message < ApplicationRecord
   validate  :validate_assistant,     if: -> { assistant.present? && Current.user }
 
   after_create :start_assistant_reply, if: :user?
-  after_create :set_latest_assistant_message, if: :assistant?
+  after_create :set_last_assistant_message, if: :assistant?
   after_save :update_assistant_on_conversation, if: -> { assistant.present? && conversation.present? }
 
   scope :ordered, -> { latest_version_for_conversation }
@@ -49,8 +50,8 @@ def start_assistant_reply
     )
   end
 
-  def set_latest_assistant_message
-    redis.set("conversation-#{conversation_id}-latest-assistant_message-id", id)
+  def set_last_assistant_message
+    conversation.update!(last_assistant_message: self)
   end
 
   def update_assistant_on_conversation

app/models/message/cancellable.rb

@@ -2,16 +2,14 @@ module Message::Cancellable
   extend ActiveSupport::Concern
 
   included do
-    after_save :save_cancelled_id_to_redis, if: :saved_change_to_cancelled_at?
+    has_one :cancelled_by, class_name: "User", foreign_key: :last_cancelled_message_id, dependent: :nullify
+
+    after_save :set_cancelled_by, if: :saved_change_to_cancelled_at?
   end
 
   private
 
-  def save_cancelled_id_to_redis
-    redis.set("message-cancelled-id", id)
-  end
-
-  def redis
-    RedisConnection.client
+  def set_cancelled_by
+    Current.user&.update!(last_cancelled_message: self)
   end
 end

app/models/user.rb

@@ -11,6 +11,7 @@ class User < ApplicationRecord
 
   has_many :assistants, dependent: :destroy
   has_many :conversations, dependent: :destroy
+  belongs_to :last_cancelled_message, class_name: "Message", optional: true
 
   serialize :preferences, coder: JsonSerializer
 

db/migrate/20240504212248_add_last_message_id_cancelled_to_users.rb

@@ -0,0 +1,5 @@
+class AddLastMessageIdCancelledToUsers < ActiveRecord::Migration[7.1]
+  def change
+    add_reference :users, :last_cancelled_message, null: true, foreign_key: { to_table: :messages }
+  end
+end

db/migrate/20240504212345_add_last_assistant_message_id_to_conversations.rb

@@ -0,0 +1,5 @@
+class AddLastAssistantMessageIdToConversations < ActiveRecord::Migration[7.1]
+  def change
+    add_reference :conversations, :last_assistant_message, null: true, foreign_key: { to_table: :messages }
+  end
+end

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: 2024_05_01_215824) do
+ActiveRecord::Schema[7.1].define(version: 2024_05_04_212345) do
   # These are extensions that must be enabled in order to support this database
   enable_extension "plpgsql"
 
@@ -75,7 +75,9 @@
     t.string "title"
     t.datetime "created_at", null: false
     t.datetime "updated_at", null: false
+    t.bigint "last_assistant_message_id"
     t.index ["assistant_id"], name: "index_conversations_on_assistant_id"
+    t.index ["last_assistant_message_id"], name: "index_conversations_on_last_assistant_message_id"
     t.index ["updated_at"], name: "index_conversations_on_updated_at"
     t.index ["user_id"], name: "index_conversations_on_user_id"
   end
@@ -292,13 +294,16 @@
     t.string "openai_key"
     t.string "anthropic_key"
     t.jsonb "preferences"
+    t.bigint "last_cancelled_message_id"
+    t.index ["last_cancelled_message_id"], name: "index_users_on_last_cancelled_message_id"
   end
 
   add_foreign_key "active_storage_attachments", "active_storage_blobs", column: "blob_id"
   add_foreign_key "active_storage_variant_records", "active_storage_blobs", column: "blob_id"
   add_foreign_key "assistants", "users"
   add_foreign_key "chats", "users"
   add_foreign_key "conversations", "assistants"
+  add_foreign_key "conversations", "messages", column: "last_assistant_message_id"
   add_foreign_key "conversations", "users"
   add_foreign_key "documents", "assistants"
   add_foreign_key "documents", "messages"
@@ -319,4 +324,5 @@
   add_foreign_key "steps", "assistants"
   add_foreign_key "steps", "conversations"
   add_foreign_key "steps", "runs"
+  add_foreign_key "users", "messages", column: "last_cancelled_message_id"
 end

test/fixtures/conversations.yml

@@ -2,36 +2,43 @@ greeting:
   user: keith
   assistant: samantha
   title: Meeting Samantha
+  last_assistant_message: im_a_bot
 
 attachment:
   user: keith
   assistant: samantha
   title: Testing one attachment
+  last_assistant_message: not_human
 
 attachments:
   user: keith
   assistant: samantha
   title: Testing attachments
+  last_assistant_message: yes_cat_and_dog
 
 javascript:
   user: keith
   assistant: keith_gpt4
   title: Javascript popstate
+  last_assistant_message: popstate_event
 
 ruby_version:
   user: keith
   assistant: samantha
   title: Ruby version
+  last_assistant_message: latest_ruby_version
 
 hello_claude:
   user: keith
   assistant: keith_claude3
   title: Meeting Claude
+  last_assistant_message: claude_age_replying
 
 debugging:
   user: rob
   assistant: rob_gpt4
   title: Verifying Ruby syntax
+  last_assistant_message: filter_map_example
 
 #  0.1
 #  1.1
@@ -47,6 +54,7 @@ versioned:
   user: keith
   assistant: samantha
   title: Versioned conversation
+  last_assistant_message: message5_v2
 
 #  0.1
 #  1.1
@@ -60,3 +68,4 @@ versioned2:
   user: keith
   assistant: samantha
   title: Versioned v2 conversation
+  last_assistant_message: msg2_v5

test/fixtures/messages.yml

@@ -68,6 +68,7 @@ dont_know_day:
   content_document:
   run: what_day_response
   created_at: 2023-12-30 1:05:00
+  cancelled_at: 2023-12-30 1:05:03
   index: 5
   version: 1
 

test/fixtures/users.yml

@@ -5,6 +5,7 @@ keith:
   registered_at: 2023-12-19 08:40:05
   openai_key: abc123
   anthropic_key: def789
+  last_cancelled_message: dont_know_day
   preferences: {}
 
 rob:

test/jobs/get_next_ai_message_job_anthropic_test.rb

@@ -3,49 +3,50 @@
 class GetNextAIMessageJobAnthropicTest < ActiveJob::TestCase
   setup do
     @conversation = conversations(:hello_claude)
+    @user = @conversation.user
     @conversation.messages.create! role: :user, content_text: "Still there?", assistant: @conversation.assistant
     @message = @conversation.latest_message_for_version(:latest)
     @test_client = TestClients::Anthropic.new(access_token: 'abc')
   end
 
   test "populates the latest message from the assistant" do
     assert_no_difference "@conversation.messages.reload.length" do
-      assert GetNextAIMessageJob.perform_now(@message.id, @conversation.assistant.id)
+      assert GetNextAIMessageJob.perform_now(@user.id, @message.id, @conversation.assistant.id)
     end
 
     message_text = @test_client.messages
     assert_equal message_text, @conversation.latest_message_for_version(:latest).content_text
   end
 
   test "returns early if the message id was invalid" do
-    refute GetNextAIMessageJob.perform_now(0, @conversation.assistant.id)
+    refute GetNextAIMessageJob.perform_now(@user.id, 0, @conversation.assistant.id)
   end
 
   test "returns early if the assistant id was invalid" do
-    refute GetNextAIMessageJob.perform_now(@message.id, 0)
+    refute GetNextAIMessageJob.perform_now(@user.id, @message.id, 0)
   end
 
   test "returns early if the message was already generated" do
     @message.update!(content_text: "Hello")
-    refute GetNextAIMessageJob.perform_now(@message.id, @conversation.assistant.id)
+    refute GetNextAIMessageJob.perform_now(@user.id, @message.id, @conversation.assistant.id)
   end
 
   test "returns early if the user has replied after this" do
     @conversation.messages.create! role: :user, content_text: "Ignore that, new question:", assistant: @conversation.assistant
-    refute GetNextAIMessageJob.perform_now(@message.id, @conversation.assistant.id)
+    refute GetNextAIMessageJob.perform_now(@user.id, @message.id, @conversation.assistant.id)
   end
 
   test "when anthropic key is blank, a nice error message is displayed" do
     user = conversations(:greeting).user
     user.update!(anthropic_key: "")
 
-    assert GetNextAIMessageJob.perform_now(@message.id, @conversation.assistant.id)
+    assert GetNextAIMessageJob.perform_now(@user.id, @message.id, @conversation.assistant.id)
     assert_includes @conversation.latest_message_for_version(:latest).content_text, "need to enter a valid API key for Anthropic"
   end
 
   test "when API response key is, a nice error message is displayed" do
     TestClients::Anthropic.stub :text, "" do
-      assert GetNextAIMessageJob.perform_now(@message.id, @conversation.assistant.id)
+      assert GetNextAIMessageJob.perform_now(@user.id, @message.id, @conversation.assistant.id)
       assert_includes @conversation.latest_message_for_version(:latest).content_text, "a blank response"
     end
   end

test/jobs/get_next_ai_message_job_openai_test.rb

@@ -3,49 +3,50 @@
 class GetNextAIMessageJobOpenaiTest < ActiveJob::TestCase
   setup do
     @conversation = conversations(:greeting)
+    @user = @conversation.user
     @conversation.messages.create! role: :user, content_text: "Still there?", assistant: @conversation.assistant
     @message = @conversation.latest_message_for_version(:latest)
     @test_client = TestClients::OpenAI.new(access_token: 'abc')
   end
 
   test "populates the latest message from the assistant" do
     assert_no_difference "@conversation.messages.reload.length" do
-      assert GetNextAIMessageJob.perform_now(@message.id, @conversation.assistant.id)
+      assert GetNextAIMessageJob.perform_now(@user.id, @message.id, @conversation.assistant.id)
     end
 
     message_text = @test_client.chat
     assert_equal message_text, @conversation.latest_message_for_version(:latest).content_text
   end
 
   test "returns early if the message id was invalid" do
-    refute GetNextAIMessageJob.perform_now(0, @conversation.assistant.id)
+    refute GetNextAIMessageJob.perform_now(@user.id, 0, @conversation.assistant.id)
   end
 
   test "returns early if the assistant id was invalid" do
-    refute GetNextAIMessageJob.perform_now(@message.id, 0)
+    refute GetNextAIMessageJob.perform_now(@user.id, @message.id, 0)
   end
 
   test "returns early if the message was already generated" do
     @message.update!(content_text: "Hello")
-    refute GetNextAIMessageJob.perform_now(@message.id, @conversation.assistant.id)
+    refute GetNextAIMessageJob.perform_now(@user.id, @message.id, @conversation.assistant.id)
   end
 
   test "returns early if the user has replied after this" do
     @conversation.messages.create! role: :user, content_text: "Ignore that, new question:", assistant: @conversation.assistant
-    refute GetNextAIMessageJob.perform_now(@message.id, @conversation.assistant.id)
+    refute GetNextAIMessageJob.perform_now(@user.id, @message.id, @conversation.assistant.id)
   end
 
   test "when openai key is blank, a nice error message is displayed" do
     user = @conversation.user
     user.update!(openai_key: "")
 
-    assert GetNextAIMessageJob.perform_now(@message.id, @conversation.assistant.id)
+    assert GetNextAIMessageJob.perform_now(@user.id, @message.id, @conversation.assistant.id)
     assert_includes @conversation.latest_message_for_version(:latest).content_text, "need to enter a valid API key for OpenAI"
   end
 
   test "when API response key is, a nice error message is displayed" do
     TestClients::OpenAI.stub :text, "" do
-      assert GetNextAIMessageJob.perform_now(@message.id, @conversation.assistant.id)
+      assert GetNextAIMessageJob.perform_now(@user.id, @message.id, @conversation.assistant.id)
       assert_includes @conversation.latest_message_for_version(:latest).content_text, "a blank response"
     end
   end