Single-file nanodjango & llm powered chat app using HTMX and websockets

chat.py

# /// script
# requires-python = ">=3.12"
# dependencies = [
#     "nanodjango",
#     "channels",
#     "daphne",
#     "htpy",
#     "markdown",
#     "markupsafe",
#     "llm"
# ]
# ///

import json
import uuid

from channels.generic.websocket import WebsocketConsumer
from channels.routing import ProtocolTypeRouter, URLRouter
from channels.auth import AuthMiddlewareStack
from django.http import HttpResponse
from django.urls import path
from markupsafe import Markup
from markdown import markdown
from htpy import (
    body,
    button,
    div,
    form,
    h1,
    head,
    html,
    input,
    meta,
    script,
    link,
    title,
    main,
    style,
    fieldset,
    article,
)
from nanodjango import Django
import llm


#  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .
#
#  ┌┬┐┌─┐┌┬┐┌─┐┬  ┌─┐┌┬┐┌─┐
#   │ ├┤ │││├─┘│  ├─┤ │ ├┤ 
#   ┴ └─┘┴ ┴┴  ┴─┘┴ ┴ ┴ └─┘
def html_template():
    return html[
        head[
            meta(charset="utf-8"),
            meta(name="viewport", content="width=device-width, initial-scale=1"),
            title["llm chat"],
            script(src="https://unpkg.com/[email protected]"),
            script(src="https://unpkg.com/[email protected]"),
            script(src="https://cdn.jsdelivr.net/npm/marked/marked.min.js"),
            link(
                rel="stylesheet",
                href="https://cdn.jsdelivr.net/npm/@picocss/pico@2/css/pico.min.css",
            ),
            style[
                Markup("""
                .message { padding: .5rem;  }
                .user-message {
                    border: 1px solid #999;
                    border-radius: 0.5rem;
                    margin-bottom: 0.5rem;
                }
                .response-message { 
                    font-weight: bold; 
                    background-color: #333;
                    border: 1px solid green;
                    border-radius: 0.5rem;
                    margin-bottom: 1rem;
                }
                .markdown-content {
                    display: none;
                }
                """)
            ],
            script[
                Markup("""
                // Create a MutationObserver to watch for content changes in hidden elements
                const observer = new MutationObserver((mutations) => {
                    mutations.forEach((mutation) => {
                        if (mutation.type === 'childList' && mutation.target.classList.contains('markdown-content')) {
                            // Get the visible container (sibling of the hidden content)
                            const visibleContainer = mutation.target.nextElementSibling;
                            if (visibleContainer) {
                                visibleContainer.innerHTML = marked.parse(mutation.target.textContent);
                            }
                        }
                    });
                });

                // Start observing the message list for changes
                document.addEventListener('DOMContentLoaded', () => {
                    const messageList = document.getElementById('message-list');
                    if (messageList) {
                        observer.observe(messageList, {
                            childList: true,
                            subtree: true,
                            characterData: true
                        });
                    }
                });
                """)
            ],
        ],
        body[
            main(class_="container")[
                article[
                    h1["🧢 thinking cap"],
                    div(hx_ext="ws", ws_connect="/ws/echo/")[
                        div("#message-list"),
                        form(ws_send=True)[
                            fieldset(role="group")[
                                input(
                                    name="message",
                                    type="text",
                                    placeholder="Type your message...",
                                    autocomplete="off",
                                ),
                                button(
                                    class_="primary outline",
                                    type="submit",
                                    onclick="setTimeout(() => this.closest('form').querySelector('input[name=message]').value = '', 0)",
                                )["↩"],
                            ]
                        ],
                    ],
                ],
            ]
        ],
    ]

#  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .
#
#   ┌─┐┌─┐┌┬┐┌─┐┌─┐┌┐┌┌─┐┌┐┌┌┬┐┌─┐
#   │  │ ││││├─┘│ ││││├┤ │││ │ └─┐
#   └─┘└─┘┴ ┴┴  └─┘┘└┘└─┘┘└┘ ┴ └─┘
def response_message(message_text, id):
    return div("#message-list", hx_swap_oob=f"beforeend:{id} .markdown-content")[message_text]


def formatted_response_message(message_text, id):
    return div(id, hx_swap_oob="outerHTML")[
        div(data_theme="dark", class_="message response-message")[
            Markup(markdown(message_text, extensions=['fenced_code']))
        ]
    ]


def response_container(id):
    return div("#message-list", hx_swap_oob="beforeend")[
        div(id, class_=["message", "response-message"], data_theme="dark")[
            div(class_="markdown-content")[""],  # Hidden element for raw markdown
            div(class_="rendered-content")[""]   # Visible element for rendered HTML
        ]
    ]


def user_message(message_text):
    return div("#message-list", hx_swap_oob="beforeend")[
        div(class_=["message", "user-message"])[
            message_text
        ]
    ]

#  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .
#
#   ┬  ┬┬┌─┐┬ ┬┌─┐
#   └┐┌┘│├┤ │││└─┐
#    └┘ ┴└─┘└┴┘└─┘
def index(request):
    return HttpResponse(html_template())

#  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .
#
#     ┬ ┬┌─┐┌┐ ┌─┐┌─┐┌─┐┬┌─┌─┐┌┬┐
#     │││├┤ ├┴┐└─┐│ ││  ├┴┐├┤  │ 
#     └┴┘└─┘└─┘└─┘└─┘└─┘┴ ┴└─┘ ┴ 
class EchoConsumer(WebsocketConsumer):
    
    def receive(self, text_data):
        text_data_json = json.loads(text_data)
        message_text = text_data_json.get("message", "")
        if not message_text.strip():
            return
        user_message_html = user_message(message_text)
        self.send(text_data=user_message_html)
        
        response = get_model().prompt(message_text)
        response_container_id = f"#response-message-{str(uuid.uuid4())}"
        response_container_html = response_container(response_container_id)
        self.send(text_data=response_container_html)
        full_response = ""
        for chunk in response:
            full_response += chunk
            echo_message_html = response_message(chunk, response_container_id)
            self.send(text_data=echo_message_html)

#  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .
#
#     ╔═╗╔═╗╔═╗
#     ╠═╣╠═╝╠═╝
#     ╩ ╩╩  ╩  
app = Django(
    # EXTRA_APPS=[
    #     "channels",
    # ],
    #
    # Nanodjango doesn't yet support prepending "priority" apps to INSTALLED_APPS,
    # and `daphne` must be the first app in INSTALLED_APPS.
    INSTALLED_APPS=[
        "daphne",
        "django.contrib.admin",
        "django.contrib.auth",
        "django.contrib.contenttypes",
        "django.contrib.sessions",
        "django.contrib.messages",
        "whitenoise.runserver_nostatic",
        "django.contrib.staticfiles",
        "channels",
    ],
    CHANNEL_LAYERS={
        "default": {
            "BACKEND": "channels.layers.InMemoryChannelLayer",
        },
    },
    ASGI_APPLICATION="__main__.htmx_websocket_interface",
)


#  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .
#
#   ┬ ┬┬─┐┬  ┌─┐
#   │ │├┬┘│  └─┐
#   └─┘┴└─┴─┘└─┘
app.route("/")(index)
websocket_urlpatterns = [
    path("ws/echo/", EchoConsumer.as_asgi()),
]


htmx_websocket_interface = ProtocolTypeRouter(
    {
        "http": app.asgi,
        "websocket": AuthMiddlewareStack(URLRouter(websocket_urlpatterns)),
    }
)


#  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .  .
#
#     ┬  ┬  ┌┬┐
#     │  │  │││
#     ┴─┘┴─┘┴ ┴
_model = None
def get_model():
    global _model
    if _model is None:
        model = llm.get_model()
        _model = model.conversation()
    return _model


# ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ ^^^^ 
#
if __name__ == "__main__":
    app.run()