feat!(pkg-r): new data_source plugin API (#123)

* feat!(pkg-r): new data_source plugin API

* `air format` (GitHub Actions)

* Address feedback

* `air format` (GitHub Actions)

* `devtools::document()` (GitHub Actions)

* Apply suggestions from code review

Co-authored-by: Garrick Aden-Buie <[email protected]>

* Remove file committed by mistake

* get_system_prompt -> assemble_system_prompt

* `devtools::document()` (GitHub Actions)

* Missed a renaming

---------

Co-authored-by: cpsievert <[email protected]>
Co-authored-by: Garrick Aden-Buie <[email protected]>

Author: cpsievert

pkg-py/src/querychat/_querychat.py

@@ -49,7 +49,7 @@ def __init__(
 
         self.greeting = greeting.read_text() if isinstance(greeting, Path) else greeting
 
-        prompt = get_system_prompt(
+        prompt = assemble_system_prompt(
             self._data_source,
             data_description=data_description,
             extra_instructions=extra_instructions,
@@ -682,7 +682,7 @@ def as_querychat_client(client: str | chatlas.Chat | None) -> chatlas.Chat:
     return chatlas.ChatAuto(provider_model=client)
 
 
-def get_system_prompt(
+def assemble_system_prompt(
     data_source: DataSource,
     *,
     data_description: Optional[str | Path] = None,

pkg-r/NAMESPACE

@@ -1,22 +1,9 @@
 # Generated by roxygen2: do not edit by hand
 
-S3method(as_querychat_data_source,DBIConnection)
-S3method(as_querychat_data_source,data.frame)
-S3method(cleanup_source,dbi_source)
-S3method(create_system_prompt,querychat_data_source)
-S3method(execute_query,dbi_source)
-S3method(get_db_type,data_frame_source)
-S3method(get_db_type,dbi_source)
-S3method(get_db_type,default)
-S3method(get_schema,dbi_source)
-S3method(test_query,dbi_source)
+export(DBISource)
+export(DataFrameSource)
+export(DataSource)
 export(QueryChat)
-export(as_querychat_data_source)
-export(cleanup_source)
-export(create_system_prompt)
-export(execute_query)
-export(get_db_type)
-export(get_schema)
 export(querychat)
 export(querychat_app)
 export(querychat_data_source)
@@ -25,7 +12,6 @@ export(querychat_init)
 export(querychat_server)
 export(querychat_sidebar)
 export(querychat_ui)
-export(test_query)
 importFrom(R6,R6Class)
 importFrom(bslib,sidebar)
 importFrom(lifecycle,deprecated)

pkg-r/R/QueryChat.R

@@ -171,7 +171,7 @@ QueryChat <- R6::R6Class(
       }
       self$greeting <- greeting
 
-      prompt <- create_system_prompt(
+      prompt <- assemble_system_prompt(
         private$.data_source,
         data_description = data_description,
         categorical_threshold = categorical_threshold,
@@ -522,7 +522,10 @@ QueryChat <- R6::R6Class(
     #'
     #' @return Invisibly returns `NULL`. Resources are cleaned up internally.
     cleanup = function() {
-      cleanup_source(private$.data_source)
+      if (!is.null(private$.data_source)) {
+        private$.data_source$cleanup()
+      }
+      invisible(NULL)
     }
   ),
   active = list(
@@ -686,8 +689,22 @@ querychat_app <- function(
 
 normalize_data_source <- function(data_source, table_name) {
   if (is_data_source(data_source)) {
-    data_source
-  } else {
-    as_querychat_data_source(data_source, table_name)
+    return(data_source)
+  }
+
+  if (is.data.frame(data_source)) {
+    return(DataFrameSource$new(data_source, table_name))
   }
+
+  if (inherits(data_source, "DBIConnection")) {
+    return(DBISource$new(data_source, table_name))
+  }
+
+  cli::cli_abort(
+    paste0(
+      "`data_source` must be a DataSource, data.frame, or DBIConnection. ",
+      "Got: ",
+      class(data_source)[1]
+    )
+  )
 }