Add snowflake support (#148)

* add snowflake support

* two tweaks for correct results with PAT

* actually keep the static oauth and add snowflake_pat support

* fix recall

* Better surface error messages

* update docs

* Add NEWS

---------

Co-authored-by: Tomasz Kalinowski <[email protected]>
Co-authored-by: Tomasz Kalinowski <[email protected]>

Author: dfalbel

DESCRIPTION

@@ -60,7 +60,9 @@ Suggests:
     shiny,
     stringr,
     testthat (>= 3.0.0),
-    tibble
+    tibble,
+    jose,
+    openssl
 VignetteBuilder:
     knitr
 Config/Needs/website: tidyverse/tidytemplate, rmarkdown

NAMESPACE

@@ -12,6 +12,7 @@ export(embed_google_vertex)
 export(embed_lm_studio)
 export(embed_ollama)
 export(embed_openai)
+export(embed_snowflake)
 export(markdown_chunk)
 export(markdown_frame)
 export(markdown_segment)

NEWS.md

@@ -3,6 +3,9 @@
 -   New `embed_azure_openai()` helper for generating embeddings from
     Azure AI Foundry (#144).
 
+-   New `embed_snowflake()` helper for generating embeddings with the
+    Snowflake Cortex Embedding API (#148).
+
 -   `ragnar_retrieve()` (and the corresponding ellmer retrieve tool) now
     accept a vector of queries (#150).
 

R/embed-snowflake.R

@@ -0,0 +1,256 @@
+#' Generate embeddings using Snowflake
+#'
+#' Uses the [Cortex API `EMBED`](https://docs.snowflake.com/en/release-notes/2025/other/2025-04-14-cortex-offers-embed-rest-api)
+#' functions to generate embeddings.
+#'
+#' See [complete documentation](https://docs.snowflake.com/en/user-guide/snowflake-cortex/cortex-rest-api#label-cortex-llm-embed-function).
+#'
+#' @section Authentication:
+#'
+#' - a *Programmatic Access Token* (PAT) defined via the SNOWFLAKE_PAT environment variable.
+#' - A static OAuth token defined via the SNOWFLAKE_TOKEN environment variable.
+#' - Key-pair authentication credentials defined via the SNOWFLAKE_USER and SNOWFLAKE_PRIVATE_KEY (which can be a PEM-encoded private key or a path to one) environment variables.
+#' - Posit Workbench-managed Snowflake credentials for the corresponding account.
+#' - Viewer-based credentials on Posit Connect. Requires the connectcreds package.
+#'
+#' @inheritParams embed_ollama
+#' @inheritParams ellmer::chat_snowflake
+#' @export
+embed_snowflake <- function(
+  x,
+  account = snowflake_account(),
+  credentials = NULL,
+  model = "snowflake-arctic-embed-m-v1.5",
+  api_args = list(),
+  batch_size = 512L
+) {
+  if (is.data.frame(x)) {
+    x[["embedding"]] <- Recall(
+      x[["text"]],
+      account = account,
+      credentials = credentials,
+      model = model,
+      api_args = api_args,
+      batch_size = batch_size
+    )
+    return(x)
+  }
+
+  text <- x
+  check_character(text)
+  if (!length(text)) {
+    # ideally we'd return a 0-row matrix, but currently the correct
+    # embedding_size is not convenient to access in this context
+    return(NULL)
+  }
+  check_string(model, allow_empty = FALSE)
+  if (!is.list(api_args)) {
+    cli::cli_abort("`api_args` must be a list.")
+  }
+
+  auth_headers <- function() {
+    if (is.null(credentials)) {
+      default_snowflake_credentials(account)()
+    } else if (is.function(credentials)) {
+      credentials()
+    } else {
+      credentials
+    }
+  }
+
+  headers <- auth_headers()
+
+  base_req <- account |>
+    snowflake_url() |>
+    httr2::request() |>
+    embed_req_retry() |>
+    httr2::req_url_path_append("api/v2/cortex/inference:embed") |>
+    httr2::req_headers_redacted(!!!headers) |>
+    httr2::req_headers(
+      "Content-Type" = "application/json",
+      "Accept" = "application/json"
+    ) |>
+    httr2::req_user_agent(ragnar_user_agent()) |>
+    httr2::req_error(body = function(resp) {
+      tryCatch({
+        json <- httr2::resp_body_json(resp, check_type = FALSE)
+        json$message
+      }, error = function(e) {
+        "Unknown error"
+      })
+    })
+
+  out <- vector("list", length(text))
+  base_body <- rlang::list2(model = model, !!!api_args)
+
+  for (indices in chunk_list(seq_along(text), batch_size)) {
+    body <- base_body
+    body$text <- as.list(text[indices])
+
+    resp <- base_req |>
+      httr2::req_body_json(body) |>
+      httr2::req_perform() |>
+      httr2::resp_body_json()
+
+    out[indices] <- lapply(resp$data, \(x) x$embedding)
+  }
+
+  matrix(unlist(out), nrow = length(text), byrow = TRUE)
+}
+
+# Snowflake utilities copied from ellmer ------------
+# Handling Snowflake credentials and authentication in workbench + multiple other scnearios.
+
+snowflake_account <- function() {
+  val <- Sys.getenv("SNOWFLAKE_ACCOUNT")
+  if (!identical(val, "")) {
+    val
+  } else {
+    cli::cli_abort("SNOWFLAKE_ACCOUNT environment variable is not set.")
+  }
+}
+
+snowflake_url <- function(account) {
+  paste0("https://", account, ".snowflakecomputing.com")
+}
+
+default_snowflake_credentials <- function(account = snowflake_account()) {
+  # Detect viewer-based credentials from Posit Connect.
+  url <- snowflake_url(account)
+  if (is_installed("connectcreds") && connectcreds::has_viewer_token(url)) {
+    return(function() {
+      token <- connectcreds::connect_viewer_token(url)
+      list(
+        Authorization = paste("Bearer", token$access_token),
+        `X-Snowflake-Authorization-Token-Type` = "OAUTH"
+      )
+    })
+  }
+
+  token <- Sys.getenv("SNOWFLAKE_TOKEN")
+  if (nchar(token) != 0) {
+    return(function() {
+      list(
+        Authorization = paste("Bearer", token),
+        # See: https://docs.snowflake.com/en/developer-guide/snowflake-rest-api/authentication#using-oauth
+        `X-Snowflake-Authorization-Token-Type` = "OAUTH"
+      )
+    })
+  }
+
+  token <- Sys.getenv("SNOWFLAKE_PAT")
+  if (nchar(token) != 0) {
+    return(function() {
+      list(
+        Authorization = paste("Bearer", token),
+        # See https://docs.snowflake.com/en/user-guide/programmatic-access-tokens
+        `X-Snowflake-Authorization-Token-Type` = "PROGRAMMATIC_ACCESS_TOKEN"
+      )
+    })
+  }
+
+  # Support for Snowflake key-pair authentication.
+  # See: https://docs.snowflake.com/en/developer-guide/snowflake-rest-api/authentication#generate-a-jwt-token
+  user <- Sys.getenv("SNOWFLAKE_USER")
+  private_key <- Sys.getenv("SNOWFLAKE_PRIVATE_KEY")
+  if (nchar(user) != 0 && nchar(private_key) != 0) {
+    check_installed(c("jose", "openssl"), "for key-pair authentication")
+    key <- openssl::read_key(private_key)
+    return(function() {
+      token <- snowflake_keypair_token(account, user, key)
+      list(
+        Authorization = paste("Bearer", token),
+        `X-Snowflake-Authorization-Token-Type` = "KEYPAIR_JWT"
+      )
+    })
+  }
+
+  # Check for Workbench-managed credentials.
+  sf_home <- Sys.getenv("SNOWFLAKE_HOME")
+  if (grepl("posit-workbench", sf_home, fixed = TRUE)) {
+    token <- workbench_snowflake_token(account, sf_home)
+    if (!is.null(token)) {
+      return(function() {
+        # Ensure we get an up-to-date token.
+        token <- workbench_snowflake_token(account, sf_home)
+        list(
+          Authorization = paste("Bearer", token),
+          `X-Snowflake-Authorization-Token-Type` = "OAUTH"
+        )
+      })
+    }
+  }
+
+  cli::cli_abort("No Snowflake credentials are available.")
+}
+
+snowflake_keypair_token <- function(
+  account,
+  user,
+  key,
+  cache = snowflake_keypair_cache(account, key),
+  lifetime = 600L,
+  reauth = FALSE
+) {
+  # Producing a signed JWT is a fairly expensive operation (in the order of
+  # ~10ms), but adding a cache speeds this up approximately 500x.
+  creds <- cache$get()
+  if (reauth || is.null(creds) || creds$expiry < Sys.time()) {
+    cache$clear()
+    expiry <- Sys.time() + lifetime
+    # We can't use openssl::fingerprint() here because it uses a different
+    # algorithm.
+    fp <- openssl::base64_encode(
+      openssl::sha256(openssl::write_der(key$pubkey))
+    )
+    if (grepl(".+\\.privatelink$", account)) {
+      # account identifier is everything up to the first period
+      account <- gsub("^([^.]*).+", "\\1", account)
+    }
+    sub <- toupper(paste0(account, ".", user))
+    iss <- paste0(sub, ".SHA256:", fp)
+    # Note: Snowflake employs a malformed issuer claim, so we have to inject it
+    # manually after jose's validation phase.
+    claim <- jose::jwt_claim("dummy", sub, exp = as.integer(expiry))
+    claim$iss <- iss
+    creds <- list(expiry = expiry, token = jose::jwt_encode_sig(claim, key))
+    cache$set(creds)
+  }
+  creds$token
+}
+
+snowflake_keypair_cache <- function(account, key) {
+  credentials_cache(key = hash(c("sf", account, openssl::fingerprint(key))))
+}
+
+snowflake_credentials_exist <- function(...) {
+  tryCatch(
+    is_list(default_snowflake_credentials(...)),
+    error = function(e) FALSE
+  )
+}
+
+# Reads Posit Workbench-managed Snowflake credentials from a
+# $SNOWFLAKE_HOME/connections.toml file, as used by the Snowflake Connector for
+# Python implementation. The file will look as follows:
+#
+# [workbench]
+# account = "account-id"
+# token = "token"
+# authenticator = "oauth"
+workbench_snowflake_token <- function(account, sf_home) {
+  cfg <- readLines(file.path(sf_home, "connections.toml"))
+  # We don't attempt a full parse of the TOML syntax, instead relying on the
+  # fact that this file will always contain only one section.
+  if (!any(grepl(account, cfg, fixed = TRUE))) {
+    # The configuration doesn't actually apply to this account.
+    return(NULL)
+  }
+  line <- grepl("token = ", cfg, fixed = TRUE)
+  token <- gsub("token = ", "", cfg[line])
+  if (nchar(token) == 0) {
+    return(NULL)
+  }
+  # Drop enclosing quotes.
+  gsub("\"", "", token)
+}

man/embed_snowflake.Rd

@@ -0,0 +1,40 @@
+test_that("snowflake embeddings works", {
+  account <- Sys.getenv("SNOWFLAKE_ACCOUNT")
+  testthat::skip_if(account == "")
+  testthat::skip_if(!snowflake_credentials_exist(account))
+
+  model <- "snowflake-arctic-embed-m-v1.5"
+  text <- c("hello world", "another hello world")
+
+  embs_single_1 <- embed_snowflake(
+    text[1],
+    account = account,
+    model = model
+  )
+
+  embs_single_2 <- embed_snowflake(
+    text[2],
+    account = account,
+    model = model
+  )
+
+  embs_batch <- embed_snowflake(
+    text,
+    account = account,
+    model = model,
+    batch_size = 1L
+  )
+
+  expect_equal(embs_single_1[1, ], embs_batch[1, ])
+  expect_equal(embs_single_2[1, ], embs_batch[2, ])
+})
+
+test_that("embed_snowflake returns NULL for empty inputs", {
+  expect_null(
+    embed_snowflake(
+      character(),
+      account = "placeholder",
+      credentials = function() list()
+    )
+  )
+})