Add Open Telemetry instrumentation

DESCRIPTION

@@ -30,7 +30,7 @@ Imports:
     jsonlite,
     later (>= 1.4.0),
     lifecycle,
-    promises (>= 1.3.1),
+    promises (>= 1.5.0),
     R6,
     rlang (>= 1.1.0),
     S7 (>= 0.2.0),
@@ -45,6 +45,8 @@ Suggests:
     knitr,
     magick,
     openssl,
+    otel (>= 0.2.0),
+    otelsdk (>= 0.2.0),
     paws.common,
     png,
     rmarkdown,
@@ -82,11 +84,13 @@ Collate:
     'content-pdf.R'
     'content-replay.R'
     'httr2.R'
+    'import-standalone-defer.R'
     'import-standalone-obj-type.R'
     'import-standalone-purrr.R'
     'import-standalone-types-check.R'
     'interpolate.R'
     'live.R'
+    'otel.R'
     'parallel-chat.R'
     'params.R'
     'provider-any.R'

R/chat-tools.R

@@ -33,7 +33,8 @@ on_load({
     echo = "none",
     on_tool_request = function(request) invisible(),
     on_tool_result = function(result) invisible(),
-    yield_request = FALSE
+    yield_request = FALSE,
+    otel_span = NULL
   ) {
     tool_requests <- extract_tool_requests(turn)
 
@@ -51,7 +52,7 @@ on_load({
         next
       }
 
-      result <- invoke_tool(request)
+      result <- invoke_tool(request, otel_span = otel_span)
 
       if (promises::is.promise(result@value)) {
         cli::cli_abort(
@@ -78,7 +79,8 @@ on_load({
     echo = "none",
     on_tool_request = function(request) invisible(),
     on_tool_result = function(result) invisible(),
-    yield_request = FALSE
+    yield_request = FALSE,
+    otel_span = NULL
   ) {
     tool_requests <- extract_tool_requests(turn)
 
@@ -94,7 +96,7 @@ on_load({
         return(rejected)
       }
 
-      result <- coro::await(invoke_tool_async(request))
+      result <- coro::await(invoke_tool_async(request, otel_span = otel_span))
 
       maybe_echo_tool(result, echo = echo)
       on_tool_result(result)
@@ -144,7 +146,7 @@ new_tool_result <- function(request, result = NULL, error = NULL) {
 }
 
 # Also need to handle edge cases: https://platform.openai.com/docs/guides/function-calling/edge-cases
-invoke_tool <- function(request) {
+invoke_tool <- function(request, otel_span = NULL) {
   if (is.null(request@tool)) {
     return(new_tool_result(request, error = "Unknown tool"))
   }
@@ -155,19 +157,25 @@ invoke_tool <- function(request) {
     return(args)
   }
 
+  tool_span <- local_tool_otel_span(
+    request,
+    parent = otel_span
+  )
+
   tryCatch(
     {
       result <- do.call(request@tool, args)
       new_tool_result(request, result)
     },
     error = function(e) {
+      record_tool_otel_span_error(tool_span, e)
       new_tool_result(request, error = e)
     }
   )
 }
 
 on_load(
-  invoke_tool_async <- coro::async(function(request) {
+  invoke_tool_async <- coro::async(function(request, otel_span = NULL) {
     if (is.null(request@tool)) {
       return(new_tool_result(request, error = "Unknown tool"))
     }
@@ -178,12 +186,18 @@ on_load(
       return(args)
     }
 
+    tool_span <- local_tool_otel_span(
+      request,
+      parent = otel_span
+    )
+
     tryCatch(
       {
         result <- await(do.call(request@tool, args))
         new_tool_result(request, result)
       },
       error = function(e) {
+        record_tool_otel_span_error(tool_span, e)
         new_tool_result(request, error = e)
       }
     )

R/chat.R

@@ -456,14 +456,17 @@ Chat <- R6::R6Class(
       yield_as_content = FALSE
     ) {
       tool_errors <- list()
-      withr::defer(warn_tool_errors(tool_errors))
+      defer(warn_tool_errors(tool_errors))
+
+      agent_span <- local_agent_otel_span(private$provider, activate = FALSE)
 
       while (!is.null(user_turn)) {
         assistant_chunks <- private$submit_turns(
           user_turn,
           stream = stream,
           echo = echo,
-          yield_as_content = yield_as_content
+          yield_as_content = yield_as_content,
+          otel_span = agent_span
         )
         for (chunk in assistant_chunks) {
           yield(chunk)
@@ -478,7 +481,8 @@ Chat <- R6::R6Class(
             echo = echo,
             on_tool_request = private$callback_on_tool_request$invoke,
             on_tool_result = private$callback_on_tool_result$invoke,
-            yield_request = yield_as_content
+            yield_request = yield_as_content,
+            otel_span = agent_span
           )
 
           tool_results <- list()
@@ -515,14 +519,17 @@ Chat <- R6::R6Class(
       yield_as_content = FALSE
     ) {
       tool_errors <- list()
-      withr::defer(warn_tool_errors(tool_errors))
+      defer(warn_tool_errors(tool_errors))
+
+      agent_span <- local_agent_otel_span(private$provider, activate = FALSE)
 
       while (!is.null(user_turn)) {
         assistant_chunks <- private$submit_turns_async(
           user_turn,
           stream = stream,
           echo = echo,
-          yield_as_content = yield_as_content
+          yield_as_content = yield_as_content,
+          otel_span = agent_span
         )
         for (chunk in await_each(assistant_chunks)) {
           yield(chunk)
@@ -537,11 +544,12 @@ Chat <- R6::R6Class(
             echo = echo,
             on_tool_request = private$callback_on_tool_request$invoke_async,
             on_tool_result = private$callback_on_tool_result$invoke_async,
-            yield_request = yield_as_content
+            yield_request = yield_as_content,
+            otel_span = agent_span
           )
           if (tool_mode == "sequential") {
             tool_results <- list()
-            for (tool_step in coro::await_each(tool_calls)) {
+            for (tool_step in await_each(tool_calls)) {
               if (yield_as_content) {
                 yield(tool_step)
               }
@@ -587,19 +595,27 @@ Chat <- R6::R6Class(
       stream,
       echo,
       type = NULL,
-      yield_as_content = FALSE
+      yield_as_content = FALSE,
+      otel_span = NULL
     ) {
       if (echo == "all") {
         cat_line(format(user_turn), prefix = "> ")
       }
 
+      chat_span <- local_chat_otel_span(
+        private$provider,
+        parent = otel_span
+      )
+
       response <- chat_perform(
         provider = private$provider,
         mode = if (stream) "stream" else "value",
         turns = c(private$.turns, list(user_turn)),
         tools = if (is.null(type)) private$tools,
-        type = type
+        type = type,
+        otel_span = chat_span
       )
+
       emit <- emitter(echo)
       any_text <- FALSE
 
@@ -619,9 +635,15 @@ Chat <- R6::R6Class(
 
           result <- stream_merge_chunks(private$provider, result, chunk)
         }
-        turn <- value_turn(private$provider, result, has_type = !is.null(type))
+        record_chat_otel_span_status(chat_span, result)
+        turn <- value_turn(
+          private$provider,
+          result,
+          has_type = !is.null(type)
+        )
         turn <- match_tools(turn, private$tools)
       } else {
+        record_chat_otel_span_status(chat_span, response)
         turn <- value_turn(
           private$provider,
           resp_body_json(response),
@@ -673,15 +695,23 @@ Chat <- R6::R6Class(
       stream,
       echo,
       type = NULL,
-      yield_as_content = FALSE
+      yield_as_content = FALSE,
+      otel_span = NULL
     ) {
+      chat_span <- local_chat_otel_span(
+        private$provider,
+        parent = otel_span
+      )
+
       response <- chat_perform(
         provider = private$provider,
         mode = if (stream) "async-stream" else "async-value",
         turns = c(private$.turns, list(user_turn)),
         tools = if (is.null(type)) private$tools,
-        type = type
+        type = type,
+        otel_span = chat_span
       )
+
       emit <- emitter(echo)
       any_text <- FALSE
 
@@ -701,10 +731,16 @@ Chat <- R6::R6Class(
 
           result <- stream_merge_chunks(private$provider, result, chunk)
         }
-        turn <- value_turn(private$provider, result, has_type = !is.null(type))
+        record_chat_otel_span_status(chat_span, result)
+        turn <- value_turn(
+          private$provider,
+          result,
+          has_type = !is.null(type)
+        )
       } else {
         result <- await(response)
 
+        record_chat_otel_span_status(chat_span, result)
         turn <- value_turn(
           private$provider,
           resp_body_json(result),

R/httr2.R

@@ -6,12 +6,15 @@ chat_perform <- function(
   mode = c("value", "stream", "async-stream", "async-value"),
   turns,
   tools = NULL,
-  type = NULL
+  type = NULL,
+  otel_span = NULL
 ) {
   mode <- arg_match(mode)
   stream <- mode %in% c("stream", "async-stream")
   tools <- tools %||% list()
 
+  setup_active_promise_otel_span(otel_span)
+
   req <- chat_request(
     provider = provider,
     turns = turns,
@@ -23,14 +26,28 @@ chat_perform <- function(
   switch(
     mode,
     "value" = req_perform(req),
-    "stream" = chat_perform_stream(provider, req),
+    "stream" = chat_perform_stream(
+      provider,
+      req,
+      otel_span = otel_span
+    ),
     "async-value" = req_perform_promise(req),
-    "async-stream" = chat_perform_async_stream(provider, req)
+    "async-stream" = chat_perform_async_stream(
+      provider,
+      req,
+      otel_span = otel_span
+    )
   )
 }
 
 on_load(
-  chat_perform_stream <- coro::generator(function(provider, req) {
+  chat_perform_stream <- coro::generator(function(
+    provider,
+    req,
+    otel_span = NULL
+  ) {
+    setup_active_promise_otel_span(otel_span)
+
     resp <- req_perform_connection(req)
     on.exit(close(resp))
 
@@ -47,7 +64,13 @@ on_load(
 )
 
 on_load(
-  chat_perform_async_stream <- coro::async_generator(function(provider, req) {
+  chat_perform_async_stream <- coro::async_generator(function(
+    provider,
+    req,
+    otel_span = NULL
+  ) {
+    setup_active_promise_otel_span(otel_span)
+
     resp <- req_perform_connection(req, blocking = FALSE)
     on.exit(close(resp))
 

R/import-standalone-defer.R

@@ -0,0 +1,35 @@
+# Standalone file: do not edit by hand
+# Source: https://github.com/r-lib/withr/blob/HEAD/R/standalone-defer.R
+# Generated by: usethis::use_standalone("r-lib/withr", "defer")
+# ----------------------------------------------------------------------
+#
+# ---
+# repo: r-lib/withr
+# file: standalone-defer.R
+# last-updated: 2024-01-15
+# license: https://unlicense.org
+# ---
+#
+# `defer()` is similar to `on.exit()` but with a better default for
+# `add` (hardcoded to `TRUE`) and `after` (`FALSE` by default).
+# It also supports adding handlers to other frames which is useful
+# to implement `local_` functions.
+#
+#
+# ## Changelog
+#
+# 2024-01-15:
+# * Rewritten to be pure base R.
+#
+# nocov start
+
+defer <- function(expr, envir = parent.frame(), after = FALSE) {
+  thunk <- as.call(list(function() expr))
+  do.call(
+    on.exit,
+    list(thunk, add = TRUE, after = after),
+    envir = envir
+  )
+}
+
+# nocov end