From 18cfadae620de5edc28599c8fcdb7dd5350b34b2 Mon Sep 17 00:00:00 2001
From: Angelo D'Ambrosio <bakaburg1@gmail.com>
Date: Wed, 22 May 2024 10:48:48 +0200
Subject: [PATCH] feat: manage outputs longer than token limit

The changes improve the handling of LLM responses that exceed the output token limit. Instead of stopping the process, the user is now prompted to decide how to proceed. They can choose to let the LLM complete the answer, keep the incomplete answer, or stop the process. The incomplete answer is also saved to a file for reference.
---
 R/LLM_calls.R | 54 +++++++++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 50 insertions(+), 4 deletions(-)

diff --git a/R/LLM_calls.R b/R/LLM_calls.R
index d4e728f..853cc76 100644
--- a/R/LLM_calls.R
+++ b/R/LLM_calls.R
@@ -255,13 +255,59 @@ prompt_llm <- function(
 
   # Return the response
   purrr::imap_chr(parsed$choices, \(ans, i) {
+    ans_content <- ans$message$content
+
+    # Manage the case when the answer is cut off due to exceeding the
+    # output token limit
     if (ans$finish_reason == "length") {
       i <- if (length(parsed$choices) > 1) paste0(" ", i, " ") else " "
 
-      stop("Answer", i, "exhausted the context window!")
-    }
-
-    ans$message$content
+      warning("Answer", i, "exhausted the context window!")
+
+      file_name <- paste0("output_", Sys.time(), ".txt")
+
+      warning(
+        "Answer", i, "exhausted the context window!\n",
+        "The answer has been saved to a file: ", file_name
+      )
+
+      readr::write_lines(ans_content, file_name)
+
+      choice <- utils::menu(
+        c(
+          "Try to complete the answer",
+          "Keep the incomplete answer",
+          "Stop the process"),
+        title = "How do you want to proceed?"
+      )
+
+      if (choice == 1) {
+        # Ask the model to continue the answer
+        messages_new <- c(
+          messages,
+          list(list(
+            role = "assistant",
+            content = ans_content
+          )),
+          list(list(
+            role = "user",
+            content = "continue"
+          ))
+        )
+
+        ans_new <- prompt_llm(
+          messages_new, provider = provider, params = params,
+          force_json = force_json,
+          log_request = log_request, ...
+        )
+
+        return(paste0(ans_content, ans_new))
+      } else if (choice == 2) {
+        return(ans_content)
+      } else {
+        stop("The process has been stopped.")
+      }
+    } else ans_content
   })
 }