Completely refactor the app for better shiny practices.

app.R

@@ -25,99 +25,38 @@ ui <- page_fillable(
     heights_equal = "row",
     width = 1,
     fillable = FALSE,
-    card_body(id = "messages", gap = 5, fillable = FALSE)
+    uiOutput("messages")
   ),
   layout_column_wrap(
     width = 1/2,
     textInput("prompt", label = NULL, width="100%"),
-    actionButton("send", "Send", width = "100%")
+    uiOutput("sendButton")
   )
 )
 
 server <- function(input, output, session) {
-  idxs <- reactiveVal()
-  n_tokens <- reactiveVal(value = 0)
-  msg_id <- reactiveVal(value = 0)
-  
-  observeEvent(input$send, {
-    if (is.null(input$prompt) || input$prompt == "") {
-      return()
-    }
-    shinyjs::disable("send")
-    updateActionButton(inputId = "send", label = "Waiting for model...")
-    insert_message(msg_id, as.character(glue::glue("🤗: {input$prompt}")))  
-    
-    if (is.null(idxs())) {
-      current_idxs <- sess$tok$encode(system_prompt)$ids
-    } else {
-      current_idxs <- idxs()
-    }
-
-    new_idxs <- paste0("<|USER|>", input$prompt, "<|ASSISTANT|>")
-    new_idxs <- sess$tok$encode(new_idxs)$ids
-
-    # we modify the prompt to trigger the 'next_token' reactive
-    idxs(c(current_idxs, new_idxs)) 
-  })
-  
-  next_token <- eventReactive(idxs(), ignoreInit = TRUE, {
-    idxs() %>% 
-      sess$generate() %>% 
-      promises::then(
-        onFulfilled = function(x) {x},
-        onRejected = function(x) {
-          insert_message(msg_id, paste0("😭 Error generating token.", as.character(x)))
-          updateActionButton(inputId = "send", label = "Failing generation. Contact admin.")
-          NULL
-        }
-      )
-  })
-  
-  observeEvent(next_token(), {
-    tok <- next_token()
-    n_tokens(n_tokens() + 1)
-    tok %>% promises::then(function(tok) {
-      tok_dec <- sess$tok$decode(tok)
-      if (n_tokens() == 1) {
-        insert_message(msg_id, paste0("🤖: ", tok_dec), append = FALSE)
-      } else {
-        insert_message(msg_id, tok_dec, append = TRUE)
-      }
-      
-      if ((!tok %in% c(50278L, 50279L, 50277L, 1L, 0L)) &&
-          n_tokens() < max_n_tokens) {
-        idxs(c(idxs(), tok))
-      } else {
-        shinyjs::enable("send")
-        updateActionButton(inputId = "send", label = "Send")
-        n_tokens(0)
-      }
-    })
-  })
-  
+  # context for the observers that load the model in the background session
+  # it also handles reloads
+  loading <- reactiveValues(
+    model = NULL,
+    reload = NULL
+  )
   # Observer used at app startup time to allow using the 'Send' button once the
   # model has been loaded.
-  model_loaded <- reactiveVal()
-  event_reload <- reactiveVal(val = 0)
-  observeEvent(event_reload(), ignoreNULL=FALSE, {
-    
-    # the model is already loaded, we just make sure the send button is enabled
+  observeEvent(loading$reload, ignoreInit = FALSE, ignoreNULL = FALSE, priority = 0, {
+
+    # the model is already loaded, we just make sure that we propagate this
+    # by setting generating to FALSE
     if (!is.null(sess$is_loaded) && sess$is_loaded) {
-      shinyjs::enable("send")
-      updateActionButton(inputId = "send", label = "Send")
+      context$generating <- FALSE
       return()
     } 
     
-    # the model isn't loaded, this we disable the send button and
-    # show that we are loading the model
-    shinyjs::disable("send")
-    updateActionButton(inputId = "send", label = "Loading the model...")
-    
     # the model isn't loaded and no task is trying to load it, so we start a new
     # task to load it
     if (is.null(sess$is_loaded)) {
       cat("Started loading model ....", "\n")
-      model_loaded(sess$load_model(repo))
+      loading$model <- sess$load_model(repo)
       sess$is_loaded <- FALSE # not yet loaded, but loading
     } else {
       # the model is loading, but this is handled by another session. We should
@@ -129,51 +68,172 @@ server <- function(input, output, session) {
     
     # this runs for the cases where sess$is_loaded was NULL
     # ie there was no model currently loading.
-    m <- model_loaded() %>% 
-      promises::then(onFulfilled = function(x) {
-        cat("Model has been loaded!", "\n")
-        shinyjs::enable("send")
-        updateActionButton(inputId = "send", label = "Send")
-        sess$is_loaded <- TRUE
-        TRUE
-      }, onRejected = function(x) {
-        shinyjs::disable("send")
-        insert_message(msg_id, paste0("😭 Error loading the model:\n", as.character(x)))
-        sess$is_loaded <- NULL # means failure!
-        sess$sess <- NULL
-        if (event_reload() < 10) {
-          Sys.sleep(5)
-          event_reload(event_reload() + 1)
-        }
-        FALSE
-      })
-    model_loaded(m)
+    m <- loading$model %>% 
+      promises::then(
+        onFulfilled = function(x) {
+          cat("Model has been loaded!", "\n")
+          context$generating <- FALSE
+          sess$is_loaded <- TRUE
+          TRUE
+        }, 
+        onRejected = function(x) {
+          context$generating <- "error"
+          msg <- list(
+            role = "error", 
+            content = paste0("Error loading the model:\n", as.character(x))
+          )
+          context$messages <- append(context$messages, list(msg))
+          
+          # setup for retry!
+          sess$is_loaded <- NULL # means failure!
+          sess$sess <- NULL
+          if (loading$reload < 10) {
+            Sys.sleep(5)
+            loading$reload <- loading$reload + 1
+          }
+          FALSE
+        })
+    loading$model <- m
   })
-}
-
-insert_message <- function(message_id, msg, append = FALSE) {
-  if (!append) {
-    id <- message_id() + 1
-    message_id(id)
+  
+  
+  # context for generating messages
+  context <- reactiveValues(
+    generating = "loading", # a flag indicating if we are still generating tokens
+    idxs = NULL, # the current sequence of tokens
+    n_tokens = 0, # number of tokens already generated
+    messages = list()
+  )
+  
+  observeEvent(input$send, ignoreInit = TRUE, {
+    # the is the observer for send message action button that triggers the rest 
+    # of the reactions.
+    # if the prompt is empty, there's nothing to do
+    if (is.null(input$prompt) || input$prompt == "") {
+      return()
+    }
+    
+    # the user clicked 'send' and the prompt is not empty:
+    # we will enter in generation mode
+    context$generating <- TRUE
     
-    insertUI(
-      "#messages", 
-      "beforeEnd", 
-      immediate = TRUE,
-      ui = card(style="margin-bottom:5px;", card_body(
-        p(id = paste0("msg-",id), msg)
-      ))
+    # we add the user message into the messages list
+    context$messages <- append(
+      context$messages, 
+      list(list(role = "user", content = input$prompt))
     )
-  } else {
-    id <- message_id()
-    shinyjs::runjs(glue::glue(
-      "document.getElementById('msg-{id}').textContent += '{msg}'"
-    ))
-  }
-  # scroll to bottom
-  shinyjs::runjs("var elem = document.getElementById('messages'); elem.scrollTop = elem.scrollHeight;")
-  id
+    # ... and the start of the assistant message
+    context$messages <- append(
+      context$messages, 
+      list(list(role = "assistant", content = ""))
+    )
+    
+    # we also update the idxs context value with the newly added prompt
+    # in case, this is the first send call, we also need to add the system
+    # prompt
+    if (is.null(context$idxs)) {
+      context$idxs <- sess$tok$encode(system_prompt)$ids
+    }
+    
+    # we now append the prompt. the prompt is wrapped around special tokens
+    # for generation:
+    prompt <- paste0("<|USER|>", input$prompt, "<|ASSISTANT|>")
+    context$idxs <- c(context$idxs, sess$tok$encode(prompt)$ids)
+    cat("Tokens in context: ", length(context$idxs), "\n")
+  })
+  
+  observeEvent(context$generating, priority = 10, {
+    # this controls the state of the send button.
+    # if generating is TRUE we want it to be disabled, otherwise it's enabled
+    # if generating is `NULL`, then the model is not yet loaded
+    btn <- if (is.null(context$generating) || context$generating == "loading") {
+      btn <- list(class = "btn-secondary disabled", label = "Loading model ...")
+      insertUI(
+        "#sendButton", 
+        ui = actionButton("send", width = "100%", label = btn$label, class = btn$class), 
+        immediate = TRUE
+      )
+      btn
+    } else if (context$generating == "error") {
+      list(class = "btn-secondary disabled", label = "Generating error ...")
+    } else if (context$generating) {
+      list(class = "btn-secondary disabled", label = "Generating response ...")
+    } else {
+      list(class = "btn-primary", label = "Send")
+    }
+    
+    output$sendButton <- renderUI({
+      actionButton("send", width = "100%", label = btn$label, class = btn$class)
+    })
+  })
+  
+  observeEvent(context$messages, priority = 10, {
+    # this observer generates and updates the messages list
+    msg_cards <- context$messages %>% 
+      lapply(function(msg) {
+        emoji <- if (msg$role == "user") "🤗" else "🤖"
+        card(style="margin-bottom:5px;", card_body(
+          p(paste0(emoji, ":",  msg$content))
+        ))
+      })
+    
+    output$messages <- renderUI({
+      rlang::exec(card_body, !!!msg_cards, gap = 5, fillable = FALSE)
+    })
+  })
+  
+  observeEvent(context$idxs, {
+    # this observer is responsible for actually generating text by calling the
+    # model that is loaded in `sess`. it takes the context$idxs to generate new
+    # text and updates it once it's done. It also appends the last message with
+    # the newly generated token
+    context$idxs %>% 
+      sess$generate() %>% 
+      promises::then(
+        onFulfilled = function(id) {
+          if (id %in% c(50278L, 50279L, 50277L, 1L, 0L)) {
+            context$generating <- FALSE
+            context$n_tokens <- 0
+            return() # special tokens that stop generation.
+          }
+          
+          # update last message with the newly generated token
+          messages <- context$messages
+          new_msg <- paste0(
+            messages[[length(messages)]]$content, 
+            sess$tok$decode(id)
+          )
+    
+          messages[[length(messages)]]$content <- new_msg
+          context$messages <- messages
+          
+          # update the token counter
+          context$n_tokens <- context$n_tokens + 1
+          
+          if (context$n_tokens > max_n_tokens) {
+            context$generating <- FALSE
+            context$n_tokens <- 0
+            return() # we already generated enough tokens
+          }
+          
+          context$idxs <- c(context$idxs, id)  
+        },
+        onRejected = function(x) {
+          # if there was a generation error, we post a message in the message
+          # list with a error role
+          msg <- list(
+            role = "error", 
+            content = paste0("Error generating token.", x)
+          )
+          context$messages <- append(context$messages, list(msg))
+          
+          # we also say that we are no longer generating, by setting another
+          # value for the `generating` "flag"
+          context$generating <- "error"
+        }
+      )
+    NULL
+  })
 }
 
-
 shinyApp(ui, server)