Inline model default for embedding functions in closures. (#52)

* Store argument defaults for embedding functions in closures.

* refactor `embed` processing

- only inline `model` arg
- refactor recursive ast walker
- process `embed` before computing `embedding_size`

* add tests

* preserve `...`, ensure `model` arg always named

---------

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

Author: dfalbel

R/store.R

@@ -71,7 +71,9 @@ ragnar_store_create <- function(
   if (is_motherduck_location(location)) {
     con <- motherduck_connection(location, create = TRUE, overwrite)
   } else {
-    if (any(file.exists(c(location, location.wal <- paste0(location, ".wal"))))) {
+    if (
+      any(file.exists(c(location, location.wal <- paste0(location, ".wal"))))
+    ) {
       if (overwrite) {
         unlink(c(location, location.wal), force = TRUE)
       } else {
@@ -91,14 +93,10 @@ ragnar_store_create <- function(
   if (is.null(embed)) {
     embedding_size <- NULL
   } else {
+    embed <- process_embed_func(embed)
     check_number_whole(embedding_size, min = 0)
     embedding_size <- as.integer(embedding_size)
 
-    if (!inherits(embed, "crate")) {
-      environment(embed) <- baseenv()
-      embed <- rlang::zap_srcref(embed)
-    }
-
     default_schema$embedding <- matrix(
       numeric(0),
       nrow = 0,
@@ -232,18 +230,80 @@ motherduck_connection <- function(location, create = FALSE, overwrite = FALSE) {
   if (create) {
     if (dbName %in% DBI::dbGetQuery(con, "SHOW DATABASES")$database_name) {
       if (overwrite) {
-        DBI::dbExecute(con, glue::glue_sql(.con = con, "DROP DATABASE {`dbName`}"))
+        DBI::dbExecute(
+          con,
+          glue::glue_sql(.con = con, "DROP DATABASE {`dbName`}")
+        )
       } else {
         stop("Database already exists: ", dbName)
       }
     }
-    DBI::dbExecute(con, glue::glue_sql(.con = con, "CREATE DATABASE {`dbName`}"))
+    DBI::dbExecute(
+      con,
+      glue::glue_sql(.con = con, "CREATE DATABASE {`dbName`}")
+    )
   }
 
   DBI::dbExecute(con, glue::glue_sql(.con = con, "USE {`dbName`}"))
   con
 }
 
+process_embed_func <- function(embed) {
+  if (inherits(embed, "crate")) {
+    return(embed)
+  }
+  environment(embed) <- baseenv()
+  embed <- rlang::zap_srcref(embed)
+
+  embed_func_names <- grep(
+    "^embed_",
+    getNamespaceExports("ragnar"),
+    value = TRUE
+  )
+
+  walker <- function(x) {
+    switch(
+      typeof(x),
+      list = {
+        x <- lapply(x, walker)
+      },
+      language = {
+        if (rlang::is_call(x, embed_func_names, ns = c("", "ragnar"))) {
+          name <- rlang::call_name(x)
+          fn <- get(name)
+          ox <- x
+          x <- rlang::call_match(x, fn, defaults = FALSE, dots_expand = FALSE)
+          x <- as.list(x)
+
+          # ensure 'model' is explicit arg embedded in call
+          if (!"model" %in% names(x)) {
+            x["model"] <- formals(fn)["model"]
+          }
+
+          # preserve `...` if they were present in the call (call_match() removes them)
+          if (any(map_lgl(as.list(ox), identical, quote(...)))) {
+            x <- c(x, quote(...))
+          }
+          x <- as.call(x)
+
+          # ensure the call is namespaced
+          if (is.null(rlang::call_ns(x))) {
+            x[[1L]] <- call("::", quote(ragnar), as.symbol(name))
+          }
+        } else {
+          x <- as.call(lapply(x, walker))
+        }
+      },
+      x
+    )
+    x
+  }
+
+  body(embed) <- walker(body(embed))
+  embed
+}
+
+
 #' Connect to `RagnarStore`
 #'
 #' @param location string, a filepath location.

tests/testthat/_snaps/store.md

@@ -0,0 +1,45 @@
+# embed functions get the defaults stored
+
+    Code
+      store@embed
+    Output
+      function (x) 
+      ragnar::embed_openai(x = x, model = "text-embedding-3-small")
+      <environment: base>
+
+---
+
+    Code
+      store@embed
+    Output
+      function (x) 
+      ragnar::embed_openai(x = x, model = "text-embedding-3-small")
+      <environment: base>
+
+---
+
+    Code
+      store@embed
+    Output
+      function (x, ...) 
+      ragnar::embed_openai(x = x, ..., model = "text-embedding-3-small")
+      <environment: base>
+
+---
+
+    Code
+      store@embed
+    Output
+      function (x) 
+      ragnar::embed_openai(x = x, model = "text-embedding-3-small")
+      <environment: base>
+
+---
+
+    Code
+      store@embed
+    Output
+      function (x) 
+      ragnar::embed_ollama(x = x, model = "all-minilm")
+      <environment: base>
+

tests/testthat/test-store.R

@@ -168,7 +168,6 @@ test_that("Allow a NULL embedding function", {
 })
 
 test_that("works with MotherDuck", {
-
   testthat::skip_if(Sys.getenv("motherduck_token", "") == "")
 
   store <- ragnar_store_create(
@@ -184,7 +183,10 @@ test_that("works with MotherDuck", {
   )
 
   expect_error(ragnar_store_insert(store, chunks), regexp = NA)
-  expect_warning(ragnar_store_build_index(store), regexp = "MotherDuck does not support")
+  expect_warning(
+    ragnar_store_build_index(store),
+    regexp = "MotherDuck does not support"
+  )
   expect_error(ragnar_retrieve(store, "hello"), regexp = NA)
 
   # Since we used insert, there's no checking if the hash is the same
@@ -198,3 +200,26 @@ test_that("works with MotherDuck", {
   val <- dbGetQuery(store@.con, "select origin, hash, text from chunks")
   expect_equal(nrow(val), 1)
 })
+
+test_that("embed functions get the defaults stored", {
+  store <- ragnar_store_create(embed = function(x) ragnar::embed_openai(x))
+  expect_snapshot(store@embed)
+
+  # here embed_openai is implicitly obtained from ragnar::embed_openai
+  store <- ragnar_store_create(embed = function(x) embed_openai(x))
+  expect_snapshot(store@embed)
+
+  # if the embed function takes ..., they're preserved
+  store <- ragnar_store_create(
+    embed = function(x, ...) ragnar::embed_openai(x, ...)
+  )
+  expect_snapshot(store@embed)
+
+  # when using the partialized version, we should also add the defaults
+  store <- ragnar_store_create(embed = embed_openai())
+  expect_snapshot(store@embed)
+
+  # test other embed funcs
+  store <- ragnar_store_create(embed = function(x) ragnar::embed_ollama(x))
+  expect_snapshot(store@embed)
+})