Compute token costs (#418)

Fixes #203

Author: hadley

Diff for: .Rbuildignore

@@ -13,3 +13,4 @@ _cache/
 ^CRAN-SUBMISSION$
 ^[\.]?air\.toml$
 ^\.vscode$
+^data-raw$

Diff for: DESCRIPTION

@@ -16,6 +16,8 @@ Description: Chat with large language models from a range of providers
 License: MIT + file LICENSE
 URL: https://ellmer.tidyverse.org, https://github.com/tidyverse/ellmer
 BugReports: https://github.com/tidyverse/ellmer/issues
+Depends: 
+    R (>= 4.1)
 Imports:
     cli,
     coro (>= 1.1.0),

Diff for: NAMESPACE

@@ -1,6 +1,8 @@
 # Generated by roxygen2: do not edit by hand
 
+S3method(format,ellmer_dollars)
 S3method(print,Chat)
+S3method(print,ellmer_dollars)
 export(Content)
 export(ContentImage)
 export(ContentImageInline)

Diff for: NEWS.md

@@ -1,5 +1,12 @@
 # ellmer (development version)
 
+* ellmer now tracks the cost of input and output tokens. The cost is displayed
+  when you print a `Chat` object, in `tokens_usage()`, and with 
+  `Chat$get_cost()`. This is our best effort at computing the cost, but you 
+  should treat it as an estimate rather than the exact price. Unfortunately LLM APIs
+  currently make it very hard to figure out exactly how much your queries are
+  costing (#203).
+
 * `ContentToolResult` objects now include the error condition in the `error`
   property when a tool call fails (#421, @gadenbuie).
 

Diff for: R/chat.R

@@ -119,32 +119,67 @@ Chat <- R6::R6Class(
       assistant_turns <- keep(turns, function(x) x@role == "assistant")
 
       n <- length(assistant_turns)
-      tokens <- t(vapply(
+      tokens_acc <- t(vapply(
         assistant_turns,
         function(turn) turn@tokens,
         double(2)
       ))
+
+      tokens <- tokens_acc
       if (n > 1) {
         # Compute just the new tokens
         tokens[-1, 1] <- tokens[seq(2, n), 1] -
           (tokens[seq(1, n - 1), 1] + tokens[seq(1, n - 1), 2])
       }
       # collapse into a single vector
       tokens_v <- c(t(tokens))
+      tokens_acc_v <- c(t(tokens_acc))
 
       tokens_df <- data.frame(
         role = rep(c("user", "assistant"), times = n),
-        tokens = tokens_v
+        tokens = tokens_v,
+        tokens_total = tokens_acc_v
       )
 
       if (include_system_prompt && private$has_system_prompt()) {
         # How do we compute this?
-        tokens_df <- rbind(data.frame(role = "system", tokens = 0), tokens_df)
+        tokens_df <- rbind(
+          data.frame(role = "system", tokens = 0, tokens_total = 0),
+          tokens_df
+        )
       }
 
       tokens_df
     },
 
+    #' @description The cost of this chat
+    #' @param include The default, `"all"`, gives the total cumulative cost
+    #'   of this chat. Alternatively, use `"last"` to get the cost of just the
+    #'   most recent turn.
+    get_cost = function(include = c("all", "last")) {
+      include <- arg_match(include)
+
+      turns <- self$get_turns(include_system_prompt = FALSE)
+      assistant_turns <- keep(turns, function(x) x@role == "assistant")
+      n <- length(assistant_turns)
+      tokens <- t(vapply(
+        assistant_turns,
+        function(turn) turn@tokens,
+        double(2)
+      ))
+
+      if (include == "last") {
+        tokens <- tokens[nrow(tokens), , drop = FALSE]
+      }
+
+      get_token_cost(
+        private$provider@name,
+        private$provider@model,
+        input = sum(tokens[, 1]),
+        output = sum(tokens[, 2])
+      )
+    },
+
     #' @description The last turn returned by the assistant.
     #' @param role Optionally, specify a role to find the last turn with
     #'   for the role.
@@ -645,14 +680,22 @@ print.Chat <- function(x, ...) {
   turns <- x$get_turns(include_system_prompt = TRUE)
 
   tokens <- x$tokens(include_system_prompt = TRUE)
-  tokens_user <- sum(tokens$tokens[tokens$role == "user"])
-  tokens_assistant <- sum(tokens$tokens[tokens$role == "assistant"])
+
+  tokens_user <- sum(tokens$tokens_total[tokens$role == "user"])
+  tokens_assistant <- sum(tokens$tokens_total[tokens$role == "assistant"])
+  cost <- x$get_cost()
 
   cat(paste_c(
     "<Chat",
     c(" ", provider@name, "/", provider@model),
     c(" turns=", length(turns)),
-    c(" tokens=", tokens_user, "/", tokens_assistant),
+    c(
+      " tokens=",
+      tokens_user,
+      "/",
+      tokens_assistant
+    ),
+    if (!is.na(cost)) c(" ", format(cost)),
     ">\n"
   ))
 

Diff for: R/provider-bedrock.R

@@ -263,8 +263,11 @@ method(value_turn, ProviderAWSBedrock) <- function(
     }
   })
 
-  tokens <- c(result$usage$inputTokens, result$usage$outputTokens)
-  tokens_log(provider, tokens)
+  tokens <- tokens_log(
+    provider,
+    input = result$usage$inputTokens,
+    output = result$usage$outputTokens
+  )
 
   Turn(result$output$message$role, contents, json = result, tokens = tokens)
 }

Diff for: R/provider-claude.R

@@ -288,8 +288,11 @@ method(value_turn, ProviderAnthropic) <- function(
     }
   })
 
-  tokens <- c(result$usage$input_tokens, result$usage$output_tokens)
-  tokens_log(provider, tokens)
+  tokens <- tokens_log(
+    provider,
+    input = result$usage$input_tokens,
+    output = result$usage$output_tokens
+  )
 
   Turn(result$role, contents, json = result, tokens = tokens)
 }
@@ -401,6 +404,11 @@ method(as_json, list(ProviderAnthropic, ContentThinking)) <- function(
     signature = x@extra$signature
   )
 }
+# Pricing ----------------------------------------------------------------------
+
+method(standardise_model, ProviderAnthropic) <- function(provider, model) {
+  gsub("-(latest|\\d{8})$", "", model)
+}
 
 # Helpers ----------------------------------------------------------------
 

Diff for: R/provider-gemini.R

@@ -203,11 +203,11 @@ method(value_turn, ProviderGoogleGemini) <- function(
   })
   contents <- compact(contents)
   usage <- result$usageMetadata
-  tokens <- c(
-    usage$promptTokenCount %||% NA_integer_,
-    usage$candidatesTokenCount %||% NA_integer_
+  tokens <- tokens_log(
+    provider,
+    input = usage$promptTokenCount,
+    output = usage$candidatesTokenCount
   )
-  tokens_log(provider, tokens)
 
   Turn("assistant", contents, json = result, tokens = tokens)
 }
@@ -562,3 +562,11 @@ default_google_credentials <- function(
     list(Authorization = paste("Bearer", token$credentials$access_token))
   })
 }
+
+# Pricing ----------------------------------------------------------------------
+
+method(standardise_model, ProviderGoogleGemini) <- function(provider, model) {
+  # https://ai.google.dev/gemini-api/docs/models#model-versions
+  # <model>-<generation>-<variation>-...
+  gsub("^([^-]+-[^-]+-[^-]+).*$", "\\1", model)
+}

Diff for: R/provider-openai.R

@@ -237,12 +237,11 @@ method(value_turn, ProviderOpenAI) <- function(
     })
     content <- c(content, calls)
   }
-  tokens <- c(
-    result$usage$prompt_tokens %||% NA_integer_,
-    result$usage$completion_tokens %||% NA_integer_
+  tokens <- tokens_log(
+    provider,
+    input = result$usage$prompt_tokens,
+    output = result$usage$completion_tokens
   )
-  tokens_log(provider, tokens)
-
   Turn(message$role, content, json = result, tokens = tokens)
 }
 

Diff for: R/provider-snowflake.R

@@ -142,11 +142,11 @@ method(value_turn, ProviderSnowflakeCortex) <- function(
 ) {
   deltas <- compact(sapply(result$choices, function(x) x$delta$content))
   content <- list(as_content(paste(deltas, collapse = "")))
-  tokens <- c(
-    result$usage$prompt_tokens %||% NA_integer_,
-    result$usage$completion_tokens %||% NA_integer_
+  tokens <- tokens_log(
+    provider,
+    input = result$usage$prompt_tokens,
+    output = result$usage$completion_tokens
   )
-  tokens_log(provider, tokens)
   Turn(
     # Snowflake's response format seems to omit the role.
     "assistant",

Diff for: R/provider.R

@@ -113,3 +113,17 @@ method(as_json, list(Provider, class_list)) <- function(provider, x) {
 method(as_json, list(Provider, ContentJson)) <- function(provider, x) {
   as_json(provider, ContentText("<structured data/>"))
 }
+
+# Pricing ---------------------------------------------------------------------
+
+standardise_model <- new_generic(
+  "standardise_model",
+  "provider",
+  function(provider, model) {
+    S7_dispatch()
+  }
+)
+
+method(standardise_model, Provider) <- function(provider, model) {
+  model
+}

Diff for: R/sysdata.rda

@@ -1,52 +1,92 @@
-tokens_log <- function(provider, tokens) {
-  # TODO: probably should make this store in a data frame, but will tackle
-  # when implementing token costs.
+on_load(
+  the$tokens <- tokens_row(character(), character(), numeric(), numeric())
+)
 
-  name <- paste0(provider@name, "/", provider@model)
+tokens_log <- function(provider, input = NULL, output = NULL) {
+  input <- input %||% 0
+  output <- output %||% 0
 
-  if (is.null(the$tokens)) {
-    the$tokens <- list()
-  }
-  if (is.null(the$tokens[[name]])) {
-    the$tokens[[name]] <- c(0, 0)
+  model <- standardise_model(provider, provider@model)
+
+  name <- function(provider, model) paste0(provider, "/", model)
+  i <- tokens_match(provider@name, model, the$tokens$provider, the$tokens$model)
+
+  if (is.na(i)) {
+    new_row <- tokens_row(provider@name, model, input, output)
+    the$tokens <- rbind(the$tokens, new_row)
+  } else {
+    the$tokens$input[i] <- the$tokens$input[i] + input
+    the$tokens$output[i] <- the$tokens$output[i] + output
   }
 
-  tokens[is.na(tokens)] <- 0
-  the$tokens[[name]] <- the$tokens[[name]] + tokens
-  invisible()
+  # Returns value to be passed to Turn
+  c(input, output)
+}
+
+tokens_row <- function(provider, model, input, output) {
+  data.frame(provider = provider, model = model, input = input, output = output)
 }
 
+tokens_match <- function(
+  provider_needle,
+  model_needle,
+  provider_haystack,
+  model_haystack
+) {
+  match(
+    paste0(provider_needle, "/", model_needle),
+    paste0(provider_haystack, "/", model_haystack)
+  )
+}
+
+
 local_tokens <- function(frame = parent.frame()) {
   old <- the$tokens
-  the$tokens <- NULL
+  the$tokens <- tokens_row(character(), character(), numeric(), numeric())
 
   defer(the$tokens <- old, env = frame)
 }
 
-tokens_set <- function() {
-  the$tokens <- NULL
-  invisible()
-}
-
 #' Report on token usage in the current session
 #'
 #' Call this function to find out the cumulative number of tokens that you
-#' have sent and recieved in the current session.
+#' have sent and recieved in the current session. The price will be shown
+#' if known.
 #'
 #' @export
 #' @return A data frame
 #' @examples
 #' token_usage()
 token_usage <- function() {
-  if (is.null(the$tokens)) {
+  if (nrow(the$tokens) == 0) {
     cli::cli_inform(c(x = "No recorded usage in this session"))
-    return(invisible(
-      data.frame(name = character(), input = numeric(), output = numeric())
-    ))
+    return(invisible(the$tokens))
   }
 
-  rows <- map2(names(the$tokens), the$tokens, function(name, tokens) {
-    data.frame(name = name, input = tokens[[1]], output = tokens[[2]])
-  })
-  do.call("rbind", rows)
+  out <- the$tokens
+  out$price <- get_token_cost(out$provider, out$model, out$input, out$output)
+  out
+}
+
+# Cost ----------------------------------------------------------------------
+
+get_token_cost <- function(provider, model, input, output) {
+  idx <- tokens_match(provider, model, prices$provider, prices$model)
+
+  input_price <- input * prices$input[idx] / 1e6
+  output_price <- output * prices$output[idx] / 1e6
+  dollars(input_price + output_price)
+}
+
+dollars <- function(x) {
+  structure(x, class = c("ellmer_dollars", "numeric"))
+}
+#' @export
+format.ellmer_dollars <- function(x, ...) {
+  paste0(ifelse(is.na(x), "", "$"), format(unclass(round(x, 2)), nsmall = 2))
+}
+#' @export
+print.ellmer_dollars <- function(x, ...) {
+  print(format(x), quote = FALSE)
+  invisible(x)
 }

Diff for: R/tokens.R

@@ -3,6 +3,9 @@
   S7::methods_register()
 }
 
+# Work around S7 bug
+rm(format)
+
 # enable usage of <S7_object>@name in package code
 #' @rawNamespace if (getRversion() < "4.3.0") importFrom("S7", "@")
 NULL