From 3f9858a67fa08e2a4c87b8ee4e0865a244e3d845 Mon Sep 17 00:00:00 2001
From: Zach Gollwitzer <zach.gollwitzer@gmail.com>
Date: Wed, 12 Mar 2025 14:06:42 -0400
Subject: [PATCH] Add streaming chat

---
 app/controllers/messages_controller.rb        |  7 +++-
 app/javascript/controllers/chat_controller.js | 32 ++++++++++++++++
 app/jobs/ai_response_job.rb                   |  7 ++++
 app/models/chat.rb                            | 38 +++++++++++++++++++
 app/models/message.rb                         | 25 +++++++++++-
 app/views/chats/show.html.erb                 | 14 ++++---
 app/views/layouts/application.html.erb        |  4 +-
 app/views/messages/_message.html.erb          |  6 +--
 test/jobs/ai_response_job_test.rb             |  7 ++++
 9 files changed, 127 insertions(+), 13 deletions(-)
 create mode 100644 app/javascript/controllers/chat_controller.js
 create mode 100644 app/jobs/ai_response_job.rb
 create mode 100644 test/jobs/ai_response_job_test.rb

diff --git a/app/controllers/messages_controller.rb b/app/controllers/messages_controller.rb
index b0428a58..0921e8a1 100644
--- a/app/controllers/messages_controller.rb
+++ b/app/controllers/messages_controller.rb
@@ -4,7 +4,12 @@ class MessagesController < ApplicationController
   def create
     @message = @chat.messages.create!(message_params.merge(role: "user"))
 
-    redirect_to chat_path(@chat)
+    AiResponseJob.perform_later(@message)
+
+    respond_to do |format|
+      format.turbo_stream
+      format.html { redirect_to chat_path(@chat) }
+    end
   end
 
   private
diff --git a/app/javascript/controllers/chat_controller.js b/app/javascript/controllers/chat_controller.js
new file mode 100644
index 00000000..61a00408
--- /dev/null
+++ b/app/javascript/controllers/chat_controller.js
@@ -0,0 +1,32 @@
+import { Controller } from "@hotwired/stimulus";
+
+// Connects to data-controller="chat-scroll"
+export default class extends Controller {
+  static targets = ["form", "messages"];
+
+  connect() {
+    this.scrollToBottom();
+
+    this.observer = new MutationObserver(() => {
+      this.scrollToBottom();
+      this.clearInput();
+    });
+
+    this.observer.observe(this.messagesTarget, {
+      childList: true,
+      subtree: true,
+    });
+  }
+
+  disconnect() {
+    this.observer.disconnect();
+  }
+
+  scrollToBottom() {
+    this.messagesTarget.scrollTop = this.messagesTarget.scrollHeight;
+  }
+
+  clearInput() {
+    this.formTarget.querySelector("textarea").value = "";
+  }
+}
diff --git a/app/jobs/ai_response_job.rb b/app/jobs/ai_response_job.rb
new file mode 100644
index 00000000..ef73403c
--- /dev/null
+++ b/app/jobs/ai_response_job.rb
@@ -0,0 +1,7 @@
+class AiResponseJob < ApplicationJob
+  queue_as :default
+
+  def perform(message)
+    message.chat.generate_next_ai_response
+  end
+end
diff --git a/app/models/chat.rb b/app/models/chat.rb
index 5b3092c0..8e01601e 100644
--- a/app/models/chat.rb
+++ b/app/models/chat.rb
@@ -20,4 +20,42 @@ class Chat < ApplicationRecord
       )
     end
   end
+
+  def generate_next_ai_response
+    if messages.conversation.ordered.last&.role == "assistant"
+      Rails.logger.info("Skipping response because last message was an assistant message")
+      return
+    end
+
+    openai.chat(
+      parameters: {
+        model: "gpt-4o-mini",
+        stream: streamer,
+        n: 1,
+        messages: messages.conversation.order(:created_at).map do |message|
+          {
+            role: message.role,
+            content: message.content
+          }
+        end
+      }
+    )
+  end
+
+  private
+    def openai
+      OpenAI::Client.new(access_token: ENV["OPENAI_ACCESS_TOKEN"])
+    end
+
+    def streamer
+      message = messages.create!(
+        role: "assistant",
+        content: ""
+      )
+
+      proc do |chunk, _bytesize|
+        new_content = chunk.dig("choices", 0, "delta", "content")
+        message.update(content: message.content + new_content) if new_content
+      end
+    end
 end
diff --git a/app/models/message.rb b/app/models/message.rb
index 933652dd..f8a8a057 100644
--- a/app/models/message.rb
+++ b/app/models/message.rb
@@ -3,8 +3,31 @@ class Message < ApplicationRecord
 
   enum :role, { user: "user", assistant: "assistant", system: "system" }
 
-  validates :content, presence: true
+  validates :content, presence: true, allow_blank: true
   validates :role, presence: true
 
   scope :conversation, -> { where(debug_mode: false, role: [ :user, :assistant ]) }
+  scope :ordered, -> { order(created_at: :asc) }
+
+  after_create_commit :broadcast_to_chat
+  after_update_commit :broadcast_update_to_chat
+
+  private
+    def broadcast_to_chat
+      broadcast_append_to(
+        chat,
+        partial: "messages/message",
+        locals: { message: self },
+        target: "chat_#{chat.id}_messages"
+      )
+    end
+
+    def broadcast_update_to_chat
+      broadcast_update_to(
+        chat,
+        partial: "messages/message",
+        locals: { message: self },
+        target: "message_#{self.id}"
+      )
+    end
 end
diff --git a/app/views/chats/show.html.erb b/app/views/chats/show.html.erb
index 1caa449b..b40e335c 100644
--- a/app/views/chats/show.html.erb
+++ b/app/views/chats/show.html.erb
@@ -1,13 +1,17 @@
 <%= turbo_frame_tag "chat_content" do %>
-  <div class="flex flex-col gap-4 justify-between h-full">
-    <div>
+  <div class="flex flex-col gap-4 justify-between h-full" data-controller="chat">
+    <div class="grow overflow-y-auto" data-chat-target="messages">
       <h2 class="text-2xl font-bold mb-6"><%= @chat.title %></h2>
 
-      <div class="space-y-6">
-        <%= render @chat.messages.conversation %>
+      <%= turbo_stream_from @chat %>
+
+      <div id="<%= dom_id(@chat) %>_messages" class="space-y-6 py-8">
+        <%= render @chat.messages.conversation.ordered %>
       </div>
     </div>
 
-    <%= render "messages/form", chat: @chat %>
+    <div data-chat-target="form">
+      <%= render "messages/form", chat: @chat %>
+    </div>
   </div>
 <% end %>
\ No newline at end of file
diff --git a/app/views/layouts/application.html.erb b/app/views/layouts/application.html.erb
index eb63c56b..7ce593cc 100644
--- a/app/views/layouts/application.html.erb
+++ b/app/views/layouts/application.html.erb
@@ -58,12 +58,12 @@
       <% end %>
     <% end %>
 
-    <%= tag.div class: class_names("py-4 shrink-0 flex flex-col justify-between gap-6 h-full overflow-y-auto transition-all duration-300 w-[375px]") do %>
+    <%= tag.div class: class_names("py-4 shrink-0 flex flex-col justify-between gap-6 h-full transition-all duration-300 w-[375px]") do %>
       <div class="bg-white p-4">
         <%= render "chats/chat_nav" %>
       </div>
 
-      <div id="chat-content" class="grow p-4 bg-white" data-turbo-permanent>
+      <div id="chat-content" class="grow p-4 bg-white overflow-y-auto" data-turbo-permanent>
         <% if Current.user.current_chat.present? %>
           <%= turbo_frame_tag "chat_content", src: chat_path(Current.user.current_chat), loading: "lazy" do %>
             <p>Loading chat...</p>
diff --git a/app/views/messages/_message.html.erb b/app/views/messages/_message.html.erb
index 60c4c144..c8aa6551 100644
--- a/app/views/messages/_message.html.erb
+++ b/app/views/messages/_message.html.erb
@@ -1,8 +1,6 @@
-<div>
+<div id="<%= dom_id(message) %>">
   <p class="text-sm text-gray-500">
     role: <%= message.role %>
   </p>
-  <p class="text-gray-900">
-    <%= message.content %>
-  </p>
+  <p class="text-primary whitespace-pre-wrap"><%= message.content %></p>
 </div>
\ No newline at end of file
diff --git a/test/jobs/ai_response_job_test.rb b/test/jobs/ai_response_job_test.rb
new file mode 100644
index 00000000..48ac8a84
--- /dev/null
+++ b/test/jobs/ai_response_job_test.rb
@@ -0,0 +1,7 @@
+require "test_helper"
+
+class AiResponseJobTest < ActiveJob::TestCase
+  # test "the truth" do
+  #   assert true
+  # end
+end