Patch for a streaming ChatGPT app with Django, Channels, and HTMX

chatgpt-channels.patch

diff --git a/apps/chat/consumers.py b/apps/chat/consumers.py
index cf792d9e..a4963491 100644
--- a/apps/chat/consumers.py
+++ b/apps/chat/consumers.py
@@ -1,4 +1,5 @@
 import json
+import uuid
 
 from channels.generic.websocket import WebsocketConsumer
 from django.conf import settings
@@ -105,3 +106,86 @@ def _format_token(token: str) -> str:
     # apply very basic formatting while we're rendering tokens in real-time
     token = token.replace("\n", "<br>")
     return token
+
+
+class ChatConsumerDemo(WebsocketConsumer):
+    def connect(self):
+        self.user = self.scope["user"]
+        self.messages = []
+        if self.user.is_authenticated:
+            self.accept()
+        else:
+            self.close()
+
+    def disconnect(self, close_code):
+        pass
+
+    def receive(self, text_data):
+        text_data_json = json.loads(text_data)
+        message_text = text_data_json["message"]
+
+        # do nothing with empty messages
+        if not message_text.strip():
+            return
+
+        # add to messages
+        self.messages.append(
+            {
+                "role": "user",
+                "content": message_text,
+            }
+        )
+
+        # show user's message
+        user_message_html = render_to_string(
+            "chat/websocket_components/user_message_demo.html",
+            {
+                "message_text": message_text,
+            },
+        )
+        self.send(text_data=user_message_html)
+
+        # render an empty system message where we'll stream our response
+        message_id = uuid.uuid4().hex
+        contents_div_id = f"message-response-{message_id}"
+        system_message_html = render_to_string(
+            "chat/websocket_components/system_message.html",
+            {
+                "contents_div_id": contents_div_id,
+            },
+        )
+        self.send(text_data=system_message_html)
+
+        # call chatgpt api
+        client = OpenAI(api_key=settings.OPENAI_API_KEY)
+        openai_response = client.chat.completions.create(
+            model=settings.OPENAI_MODEL,
+            messages=self.messages,
+            stream=True,
+        )
+        chunks = []
+        for chunk in openai_response:
+            message_chunk = chunk.choices[0].delta.content
+            if message_chunk:
+                chunks.append(message_chunk)
+                # use htmx to insert the next token at the end of our system message.
+                chunk = f'<div hx-swap-oob="beforeend:#{contents_div_id}">{_format_token(message_chunk)}</div>'
+                self.send(text_data=chunk)
+
+        system_message = "".join(chunks)
+        # replace final input with fully rendered version, so we can render markdown, etc.
+        final_message_html = render_to_string(
+            "chat/websocket_components/final_system_message_demo.html",
+            {
+                "contents_div_id": contents_div_id,
+                "message": system_message,
+            },
+        )
+        # add to messages
+        self.messages.append(
+            {
+                "role": "system",
+                "content": system_message,
+            }
+        )
+        self.send(text_data=final_message_html)
diff --git a/apps/chat/routing.py b/apps/chat/routing.py
index 0a2acff6..f925951f 100644
--- a/apps/chat/routing.py
+++ b/apps/chat/routing.py
@@ -5,4 +5,5 @@ from . import consumers
 websocket_urlpatterns = [
     path(r"ws/aichat/", consumers.ChatConsumer.as_asgi(), name="ws_openai_new_chat"),
     path(r"ws/aichat/<slug:chat_id>/", consumers.ChatConsumer.as_asgi(), name="ws_openai_continue_chat"),
+    path(r"ws/ai-demo/", consumers.ChatConsumerDemo.as_asgi(), name="ws_ai_demo_new_chat"),
 ]
diff --git a/apps/chat/urls.py b/apps/chat/urls.py
index 618f3443..071db8ad 100644
--- a/apps/chat/urls.py
+++ b/apps/chat/urls.py
@@ -8,4 +8,5 @@ urlpatterns = [
     path("", views.chat_home, name="chat_home"),
     path("chat/new/", views.new_chat_streaming, name="new_chat"),
     path("chat/<int:chat_id>/", views.single_chat_streaming, name="single_chat"),
+    path("new-chat/", views.new_chat_demo, name="new_chat_demo"),
 ]
diff --git a/apps/chat/views.py b/apps/chat/views.py
index fac761f6..7f302020 100644
--- a/apps/chat/views.py
+++ b/apps/chat/views.py
@@ -46,3 +46,8 @@ def single_chat_streaming(request, chat_id: int):
             "chat": chat,
         },
     )
+
+
+@login_required
+def new_chat_demo(request):
+    return TemplateResponse(request, "chat/single_chat_demo.html")
diff --git a/templates/chat/single_chat_demo.html b/templates/chat/single_chat_demo.html
new file mode 100644
index 00000000..d2bbda84
--- /dev/null
+++ b/templates/chat/single_chat_demo.html
@@ -0,0 +1,35 @@
+{% extends "web/chat/chat_wrapper.html" %}
+{% load i18n %}
+{% load chat_tags %}
+{% block page_head %}
+  <script src="https://unpkg.com/htmx.org/dist/ext/ws.js"></script>
+{% endblock %}
+{% block chat_ui %}
+<div class="pg-chat-wrapper" hx-ext="ws" ws-connect="/ws/ai-demo/">
+  <div id="message-list" class="pg-chat-pane">
+    <div class="pg-chat-message-system">
+      {% include "chat/components/system_icon.html" %}
+      <div class="pg-message-contents">
+        <p>{% translate "Hello, what can I help you with today?" %}</p>
+      </div>
+    </div>
+  </div>
+  <form class="pg-chat-input-bar" ws-send>
+    <input id="chat-message-input" name="message" type="text" placeholder="{% translate 'Type your message...' %}" aria-label="Message" class="pg-control">
+    <button type="submit" class="pg-button-primary mx-2">{% translate "Send" %}</button>
+  </form>
+</div>
+{% endblock %}
+{% block page_js %}
+<script>
+  // clear message input after sending our new message
+  document.body.addEventListener('htmx:wsAfterSend', function(evt) {
+    document.getElementById("chat-message-input").value = "";
+  });
+  // scroll to bottom of chat after every incoming message
+  document.body.addEventListener('htmx:wsAfterMessage', function(evt) {
+    const chatUI = document.getElementById('message-list');
+    chatUI.scrollTop = chatUI.scrollHeight;
+  });
+</script>
+{% endblock %}
diff --git a/templates/chat/websocket_components/final_system_message_demo.html b/templates/chat/websocket_components/final_system_message_demo.html
new file mode 100644
index 00000000..2b047ef3
--- /dev/null
+++ b/templates/chat/websocket_components/final_system_message_demo.html
@@ -0,0 +1,4 @@
+{% load chat_tags %}
+<div id="{{ contents_div_id }}" class="pg-message-contents" hx-swap-oob="true">
+  {{ message|render_markdown }}
+</div>
diff --git a/templates/chat/websocket_components/user_message_demo.html b/templates/chat/websocket_components/user_message_demo.html
new file mode 100644
index 00000000..2fefe1e5
--- /dev/null
+++ b/templates/chat/websocket_components/user_message_demo.html
@@ -0,0 +1,8 @@
+<div id="message-list" hx-swap-oob="beforeend">
+  <div class="pg-chat-message-user">
+    {% include 'chat/components/user_icon.html' %}
+    <div class="pg-message-contents">
+      {{ message_text }}
+    </div>
+  </div>
+</div>