Summarize tool calls

Author: dfalbel

R/ragnar-chat.R

@@ -118,6 +118,19 @@ RagnarChat <- R6::R6Class(
       )
     },
 
+    stream_async = function(..., tool_mode = c("concurrent", "sequential"), stream = c("text", "content")) {
+      result <- private$callback_user_turn(...)
+      do.call(
+        super$stream_async,
+        append(result, list(tool_mode = tool_mode, stream = stream))
+      )
+    },
+
+    stream = function(..., stream = c("text", "content")) {
+      result <- private$callback_user_turn(...)
+      do.call(super$stream, append(result, list(stream = stream)))
+    },
+
     #' @field ragnar_tool A function that retrieves relevant chunks from the store.
     #'  This is the function that is registered as a tool in the chat.
     ragnar_tool = function(query) {
@@ -413,6 +426,34 @@ RagnarChat <- R6::R6Class(
         ...,
         ContentRagnarDocuments(text = documents)
       )
+    },
+
+    #' @description
+    #' Summarizes the tools calls in the chat history. The assistant tool
+    #' call request is redacted, and the chunks are summarized. 
+    #' 
+    #' The ellmer::ToolCallRequest becomes a `ellmer::ContentText` with:
+    #' ```
+    #' <Redacted by summarization tool: query={}>
+    #' ```
+    #' 
+    #' The `ellmer::ContentToolResult` becomes a `ellmer::ContentText` with
+    #' a summary of the chunks given by the `summarize_chunks` callback.
+    #' `summarize_chunks` takes a list of `chunks` as argument.
+    #' 
+    turns_summarize_tool_calls = function(summarize_chunks) {
+      turns <- self$get_turns() |> 
+        turns_modify_tool_calls(function(assistant_turn, user_turn) {
+          summarize_tool_call(
+            assistant_turn = assistant_turn,
+            user_turn = user_turn, 
+            summarize_chunks, 
+            tool_name = self$ragnar_tool_def@name
+          )
+        },
+        tool_name = self$ragnar_tool_def@name
+      )
+      self$set_turns(turns)
     }
   ),
 
@@ -457,3 +498,60 @@ content_set_chunks <- function(x, chunks) {
   }
   x
 }
+
+#' Applies a function to pairs of turns that represent tool calls.
+#' Otherwise keep them unchanged.
+#' @noRd
+turns_modify_tool_calls <- function(turns, fn, tool_name = NULL) {
+  i <- 1
+  while (i < length(turns)) {
+    if (turns[[i]]@role == "user") {
+      i <- i + 1; next
+    };
+    if (!is_tool_call(turns[[i]], tool_name)) {
+      i <- i + 1; next;
+    }
+
+    turns[c(i, i+1)] <- fn(turns[[i]], turns[[i+1]])
+    i <- i + 2
+  }
+  turns
+}
+
+is_tool_call <- function(x, tool_name = NULL) {
+  for (content in x@contents) {
+    if (S7::S7_inherits(content, ellmer::ContentToolRequest)) {
+      if (length(x@contents) > 1) {
+        cli::cli_warn(
+          "Tool call request found in turn with multiple contents, this may lead to unexpected results."
+        )
+      }
+      if (!is.null(tool_name) && content@tool@name != tool_name) {
+        return(FALSE)
+      }
+
+      return(TRUE)
+    }
+  }
+  FALSE
+}
+
+summarize_tool_call <- function(assistant_turn, user_turn, summarize, tool_name) {
+  assistant_turn@contents <- list(
+    ellmer::ContentText(
+      text = paste0(
+        "<Redacted by summarization tool: query=",
+        user_turn@contents[[1]]@request@arguments$query, ">"
+      )
+    )
+  )
+
+  summary <- lapply(user_turn@contents, \(x) content_get_chunks(x, tool_name)) |> 
+    unlist(recursive = FALSE) |> 
+    summarize()
+
+  user_turn@contents <- list(ellmer::ContentText(text = summary))
+
+  list(assistant_turn, user_turn)
+}
+

tests/testthat/test-ragnar-chat.R

@@ -131,3 +131,31 @@ test_that("Can insert chunks premptively in the user chat", {
   expect_equal(length(chat$get_turns()), 2)
   expect_equal(length(chat$get_turns()[[1]]@contents), 1)
 })
+
+test_that("Can summarize tool calls", {
+
+  store <- test_store()
+  chat <- chat_ragnar(
+    \() ellmer::chat_openai(model = "gpt-4.1-nano"), 
+    store = store
+  )
+
+  x <- chat$chat("functional")
+
+  summarize_chunks <- function(chunks) {
+    json <- jsonlite::toJSON(chunks)
+    ellmer::chat_openai(
+      model = "gpt-4.1-nano",
+      system_prompt = "Summarize the following chunks of text.",
+      echo = FALSE
+    )$chat(json, echo = FALSE)
+  }
+
+  chat$turns_summarize_tool_calls(summarize_chunks)
+
+  turns <- chat$get_turns()
+  expect_equal(length(turns), 4) # 2 user + 2 assistant
+  expect_equal(turns[[2]]@role, "assistant")
+  expect_true(grepl("Redacted by summarization tool", turns[[2]]@contents[[1]]@text))
+  expect_true(S7::S7_inherits(turns[[3]]@contents[[1]], ellmer::ContentText))
+})

vignettes/articles/ragnar-chat.Rmd

@@ -68,9 +68,9 @@ Now we create a `ragnar_chat` object that uses the `query_rewrite` function to r
 
 ```{r}
 chat <- chat_ragnar(
-  ellmer::chat_openai,
-  .store = store,
-  .retrieve = function(self, query) {
+  \() ellmer::chat_openai(),
+  store = store,
+  retrieve = function(self, query) {
     queries <- query_rewrite(query)
     cli::cli_inform(
       i = "Rewriten queries:",
@@ -104,13 +104,13 @@ Here's an example of how to implement this callback:
 
 ```{r}
 chat <- chat_ragnar(
-  ellmer::chat_openai, 
-  .store = store,
+  \() ellmer::chat_openai(), 
+  store = store,
   # We explicitly avoid registering the store so the LLM is not capable of making tool
   # calls to the store. Instead, it relies on the documents inserted at the top of
   # the chat history to answer the question.
-  .register_store = FALSE,
-  .on_user_turn = function(self, ...) {
+  register_store = FALSE,
+  on_user_turn = function(self, ...) {
     self$turns_prune_chunks()
     # Inserts documents relevant to the query at the top of the chat history.
     self$turns_insert_documents(
@@ -125,3 +125,31 @@ chat <- chat_ragnar(
 cat(chat$chat("What is the difference between a vector and a list in R?"))
 ```
 
+## Summarize tool calls
+
+To save tokens, it's common to summarize the results of the tool calls after the chunks
+were used to produce the LLM response. With this, the next message in the chat won't need
+to use the full text of the chunks, but rather a summary of the results.
+
+This can be acomplished by customizing the `on_user_turn` callback to summarise the
+results of the previous tool calls. For example:
+
+```{r}
+summarize_chunks <- function(chunks) {
+  json <- jsonlite::toJSON(chunks)
+  ellmer::chat_openai(
+    model = "gpt-4.1-nano",
+    system_prompt = "Summarize the following chunks of text.",
+    echo = FALSE
+  )$chat(json, echo = FALSE)
+}
+
+chat <- chat_ragnar(
+  \() ellmer::chat_openai(), 
+  store = store,
+  on_user_turn = function(self, ...) {
+    self$turns_summarize_tool_calls(summarize_chunks)
+    self$turns_insert_tool_call_request(..., query = paste(..., collapse = " "))
+  }
+)
+```