diff --git a/pkg-py/src/querychat/_querychat.py b/pkg-py/src/querychat/_querychat.py
index a32ef934..5292061d 100644
--- a/pkg-py/src/querychat/_querychat.py
+++ b/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,
diff --git a/pkg-r/NAMESPACE b/pkg-r/NAMESPACE
index fa321d5b..bc89efb0 100644
--- a/pkg-r/NAMESPACE
+++ b/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)
diff --git a/pkg-r/R/QueryChat.R b/pkg-r/R/QueryChat.R
index 274ebba4..4bcf3612 100644
--- a/pkg-r/R/QueryChat.R
+++ b/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]
+ )
+ )
}
diff --git a/pkg-r/R/data_source.R b/pkg-r/R/data_source.R
index 4f8e6633..1efbe73a 100644
--- a/pkg-r/R/data_source.R
+++ b/pkg-r/R/data_source.R
@@ -1,282 +1,395 @@
-#' Create a data source for querychat
+#' Data Source Base Class
#'
-#' An entrypoint for developers to create custom data sources for use with
-#' querychat. Most users shouldn't use this function directly; instead, they
-#' should pass their data to `QueryChat$new()`.
+#' @description
+#' An abstract R6 class defining the interface that custom QueryChat data
+#' sources must implement. This class should not be instantiated directly;
+#' instead, use one of its concrete implementations like [DataFrameSource] or
+#' [DBISource].
#'
-#' @param x A data frame or DBI connection
-#' @param table_name The name to use for the table in the data source. Can be:
-#' - A character string (e.g., "table_name")
-#' - Or, for tables contained within catalogs or schemas, a [DBI::Id()] object (e.g., `DBI::Id(schema = "schema_name", table = "table_name")`)
-#' @return A querychat_data_source object
-#' @keywords internal
-#' @export
-as_querychat_data_source <- function(x, table_name, ...) {
- UseMethod("as_querychat_data_source")
-}
-
-#' @export
-as_querychat_data_source.data.frame <- function(x, table_name, ...) {
- is_table_name_ok <- is.character(table_name) &&
- length(table_name) == 1 &&
- grepl("^[a-zA-Z][a-zA-Z0-9_]*$", table_name, perl = TRUE)
- if (!is_table_name_ok) {
- cli::cli_abort(
- "`table_name` argument must be a string containing alphanumeric characters and underscores, starting with a letter."
- )
- }
-
- # Create duckdb connection
- conn <- DBI::dbConnect(duckdb::duckdb(), dbdir = ":memory:")
- duckdb::duckdb_register(conn, table_name, x, experimental = FALSE)
-
- structure(
- list(conn = conn, table_name = table_name),
- class = c("data_frame_source", "dbi_source", "querychat_data_source")
- )
-}
-
#' @export
-as_querychat_data_source.DBIConnection <- function(x, table_name, ...) {
- # Handle different types of table_name inputs
- if (inherits(table_name, "Id")) {
- # DBI::Id object - keep as is
- } else if (is.character(table_name) && length(table_name) == 1) {
- # Character string - keep as is
- } else {
- # Invalid input
- cli::cli_abort(
- "`table_name` must be a single character string or a DBI::Id object"
- )
- }
-
- # Check if table exists
- if (!DBI::dbExistsTable(x, table_name)) {
- cli::cli_abort(c(
- "Table {DBI::dbQuoteIdentifier(x, table_name)} not found in database.",
- "i" = "If you're using a table in a catalog or schema, pass a DBI::Id object to `table_name`"
- ))
- }
-
- structure(
- list(conn = x, table_name = table_name),
- class = c("dbi_source", "querychat_data_source")
+DataSource <- R6::R6Class(
+ "DataSource",
+ public = list(
+ #' @field table_name Name of the table to be used in SQL queries
+ table_name = NULL,
+
+ #' @description
+ #' Get the database type
+ #'
+ #' @return A string describing the database type (e.g., "DuckDB", "SQLite")
+ get_db_type = function() {
+ cli::cli_abort(
+ "get_db_type() must be implemented by subclass",
+ class = "not_implemented_error"
+ )
+ },
+
+ #' @description
+ #' Get schema information about the table
+ #'
+ #' @param categorical_threshold Maximum number of unique values for a text
+ #' column to be considered categorical
+ #' @return A string containing schema information formatted for LLM prompts
+ get_schema = function(categorical_threshold = 20) {
+ cli::cli_abort(
+ "get_schema() must be implemented by subclass",
+ class = "not_implemented_error"
+ )
+ },
+
+ #' @description
+ #' Execute a SQL query and return results
+ #'
+ #' @param query SQL query string to execute
+ #' @return A data frame containing query results
+ execute_query = function(query) {
+ cli::cli_abort(
+ "execute_query() must be implemented by subclass",
+ class = "not_implemented_error"
+ )
+ },
+
+ #' @description
+ #' Test a SQL query by fetching only one row
+ #'
+ #' @param query SQL query string to test
+ #' @return A data frame containing one row of results (or empty if no matches)
+ test_query = function(query) {
+ cli::cli_abort(
+ "test_query() must be implemented by subclass",
+ class = "not_implemented_error"
+ )
+ },
+
+ #' @description
+ #' Get the unfiltered data as a data frame
+ #'
+ #' @return A data frame containing all data from the table
+ get_data = function() {
+ cli::cli_abort(
+ "get_data() must be implemented by subclass",
+ class = "not_implemented_error"
+ )
+ },
+
+ #' @description
+ #' Clean up resources (close connections, etc.)
+ #'
+ #' @return NULL (invisibly)
+ cleanup = function() {
+ cli::cli_abort(
+ "cleanup() must be implemented by subclass",
+ class = "not_implemented_error"
+ )
+ }
)
-}
+)
-is_data_source <- function(x) {
- inherits(x, "querychat_data_source")
-}
-#' Execute an SQL query on a data source
+#' Data Frame Source
#'
-#' An entrypoint for developers to create custom data source objects for use
-#' with querychat. Most users shouldn't use this function directly; instead,
-#' they call the `$sql()` method on the [QueryChat] object to run queries.
+#' @description
+#' A DataSource implementation that wraps a data frame using DuckDB for SQL
+#' query execution.
#'
-#' @param source A querychat_data_source object
-#' @param query SQL query string
-#' @param ... Additional arguments passed to methods
-#' @return Result of the query as a data frame
-#' @keywords internal
-#' @export
-execute_query <- function(source, query, ...) {
- UseMethod("execute_query")
-}
-
-#' @export
-execute_query.dbi_source <- function(source, query, ...) {
- if (is.null(query) || query == "") {
- # For a null or empty query, default to returning the whole table (ie SELECT *)
- query <- paste0(
- "SELECT * FROM ",
- DBI::dbQuoteIdentifier(source$conn, source$table_name)
- )
- }
- # Execute the query directly
- DBI::dbGetQuery(source$conn, query)
-}
-
-#' Test a SQL query on a data source.
-#'
-#' An entrypoint for developers to create custom data sources for use with
-#' querychat. Most users shouldn't use this function directly; instead, they
-#' should call the `$sql()` method on the [QueryChat] object to run queries.
+#' @details
+#' This class creates an in-memory DuckDB connection and registers the provided
+#' data frame as a table. All SQL queries are executed against this DuckDB table.
#'
-#' @param source A querychat_data_source object
-#' @param query SQL query string
-#' @param ... Additional arguments passed to methods
-#' @return Result of the query, limited to one row of data.
-#' @keywords internal
-#' @export
-test_query <- function(source, query, ...) {
- UseMethod("test_query")
-}
-
#' @export
-test_query.dbi_source <- function(source, query, ...) {
- rs <- DBI::dbSendQuery(source$conn, query)
- df <- DBI::dbFetch(rs, n = 1)
- DBI::dbClearResult(rs)
- df
-}
-
-
-#' Get type information for a data source
+#' @examples
+#' \dontrun{
+#' # Create a data frame source
+#' df_source <- DataFrameSource$new(mtcars, "mtcars")
#'
-#' An entrypoint for developers to create custom data sources for use with
-#' querychat. Most users shouldn't use this function directly; instead, they
-#' should call the `$set_system_prompt()` method on the [QueryChat] object.
+#' # Get database type
+#' df_source$get_db_type() # Returns "DuckDB"
#'
-#' @param source A querychat_data_source object
-#' @param ... Additional arguments passed to methods
-#' @return A character string containing the type information
-#' @keywords internal
-#' @export
-get_db_type <- function(source, ...) {
- UseMethod("get_db_type")
-}
-
-#' @export
-get_db_type.default <- function(source, ...) {
- "standard"
-}
-
-#' @export
-get_db_type.data_frame_source <- function(source, ...) {
- # Local dataframes are always duckdb!
- "DuckDB"
-}
-
-#' @export
-get_db_type.dbi_source <- function(source, ...) {
- conn <- source$conn
+#' # Execute a query
+#' result <- df_source$execute_query("SELECT * FROM mtcars WHERE mpg > 25")
+#'
+#' # Clean up when done
+#' df_source$cleanup()
+#' }
+DataFrameSource <- R6::R6Class(
+ "DataFrameSource",
+ inherit = DataSource,
+ private = list(
+ conn = NULL
+ ),
+ public = list(
+ #' @description
+ #' Create a new DataFrameSource
+ #'
+ #' @param df A data frame.
+ #' @param table_name Name to use for the table in SQL queries. Must be a
+ #' valid table name (start with letter, contain only letters, numbers,
+ #' and underscores)
+ #' @return A new DataFrameSource object
+ #' @examples
+ #' \dontrun{
+ #' source <- DataFrameSource$new(iris, "iris")
+ #' }
+ initialize = function(df, table_name) {
+ if (!is.data.frame(df)) {
+ cli::cli_abort("`df` must be a data frame")
+ }
- # Special handling for known database types
- if (inherits(conn, "duckdb_connection")) {
- return("DuckDB")
- }
- if (inherits(conn, "SQLiteConnection")) {
- return("SQLite")
- }
+ # Validate table name
+ is_table_name_ok <- is.character(table_name) &&
+ length(table_name) == 1 &&
+ grepl("^[a-zA-Z][a-zA-Z0-9_]*$", table_name, perl = TRUE)
+ if (!is_table_name_ok) {
+ cli::cli_abort(
+ "`table_name` argument must be a string containing alphanumeric characters and underscores, starting with a letter."
+ )
+ }
- # default to 'POSIX' if dbms name not found
- conn_info <- DBI::dbGetInfo(conn)
- dbms_name <- purrr::pluck(conn_info, "dbms.name", .default = "POSIX")
+ self$table_name <- table_name
- # remove ' SQL', if exists (SQL is already in the prompt)
- gsub(" SQL", "", dbms_name)
-}
+ # Create DuckDB connection and register the data frame
+ private$conn <- DBI::dbConnect(duckdb::duckdb(), dbdir = ":memory:")
+ duckdb::duckdb_register(
+ private$conn,
+ table_name,
+ df,
+ experimental = FALSE
+ )
+ },
+
+ #' @description Get the database type
+ #' @return The string "DuckDB"
+ get_db_type = function() {
+ "DuckDB"
+ },
+
+ #' @description
+ #' Get schema information for the data frame
+ #'
+ #' @param categorical_threshold Maximum number of unique values for a text
+ #' column to be considered categorical (default: 20)
+ #' @return A string describing the schema
+ get_schema = function(categorical_threshold = 20) {
+ get_schema_impl(private$conn, self$table_name, categorical_threshold)
+ },
+
+ #' @description
+ #' Execute a SQL query
+ #'
+ #' @param query SQL query string. If NULL or empty, returns all data
+ #' @return A data frame with query results
+ execute_query = function(query) {
+ if (is.null(query) || query == "") {
+ query <- paste0(
+ "SELECT * FROM ",
+ DBI::dbQuoteIdentifier(private$conn, self$table_name)
+ )
+ }
+ DBI::dbGetQuery(private$conn, query)
+ },
+
+ #' @description
+ #' Test a SQL query by fetching only one row
+ #'
+ #' @param query SQL query string
+ #' @return A data frame with one row of results
+ test_query = function(query) {
+ rs <- DBI::dbSendQuery(private$conn, query)
+ df <- DBI::dbFetch(rs, n = 1)
+ DBI::dbClearResult(rs)
+ df
+ },
+
+ #' @description
+ #' Get all data from the table
+ #'
+ #' @return A data frame containing all data
+ get_data = function() {
+ self$execute_query(NULL)
+ },
+
+ #' @description
+ #' Close the DuckDB connection
+ #'
+ #' @return NULL (invisibly)
+ cleanup = function() {
+ if (!is.null(private$conn) && DBI::dbIsValid(private$conn)) {
+ DBI::dbDisconnect(private$conn)
+ }
+ invisible(NULL)
+ }
+ )
+)
-#' Create a system prompt for the data source
+#' DBI Source
#'
-#' An entrypoint for developers to create custom data sources for use with
-#' querychat. Most users shouldn't use this function directly; instead, they
-#' should call the `$set_system_prompt()` method on the [QueryChat] object.
+#' @description
+#' A DataSource implementation for DBI database connections (SQLite, PostgreSQL,
+#' MySQL, etc.).
+#'
+#' @details
+#' This class wraps a DBI connection and provides SQL query execution against
+#' a specified table in the database.
#'
-#' @param source A querychat_data_source object
-#' @param data_description Optional description of the data
-#' @param extra_instructions Optional additional instructions
-#' @param categorical_threshold For text columns, the maximum number of unique
-#' values to consider as a categorical variable
-#' @param ... Additional arguments passed to methods
-#' @return A string with the system prompt
-#' @keywords internal
#' @export
-create_system_prompt <- function(
- source,
- data_description = NULL,
- extra_instructions = NULL,
- categorical_threshold = 20,
- ...
-) {
- UseMethod("create_system_prompt")
-}
+#' @examples
+#' \dontrun{
+#' # Connect to a database
+#' conn <- DBI::dbConnect(RSQLite::SQLite(), ":memory:")
+#' DBI::dbWriteTable(conn, "mtcars", mtcars)
+#'
+#' # Create a DBI source
+#' db_source <- DBISource$new(conn, "mtcars")
+#'
+#' # Get database type
+#' db_source$get_db_type() # Returns "SQLite"
+#'
+#' # Execute a query
+#' result <- db_source$execute_query("SELECT * FROM mtcars WHERE mpg > 25")
+#'
+#' # Note: cleanup() will disconnect the connection
+#' # If you want to keep the connection open, don't call cleanup()
+#' }
+DBISource <- R6::R6Class(
+ "DBISource",
+ inherit = DataSource,
+ private = list(
+ conn = NULL
+ ),
+ public = list(
+ #' @description
+ #' Create a new DBISource
+ #'
+ #' @param conn A DBI connection object
+ #' @param table_name Name of the table in the database. Can be a character
+ #' string or a [DBI::Id()] object for tables in catalogs/schemas
+ #' @return A new DBISource object
+ #' @examples
+ #' \dontrun{
+ #' conn <- DBI::dbConnect(RSQLite::SQLite(), ":memory:")
+ #' DBI::dbWriteTable(conn, "iris", iris)
+ #' source <- DBISource$new(conn, "iris")
+ #' }
+ initialize = function(conn, table_name) {
+ if (!inherits(conn, "DBIConnection")) {
+ cli::cli_abort("`conn` must be a DBI connection")
+ }
-#' @export
-create_system_prompt.querychat_data_source <- function(
- source,
- data_description = NULL,
- extra_instructions = NULL,
- categorical_threshold = 20,
- ...
-) {
- if (!is.null(data_description)) {
- data_description <- paste(data_description, collapse = "\n")
- }
- if (!is.null(extra_instructions)) {
- extra_instructions <- paste(extra_instructions, collapse = "\n")
- }
+ # Validate table_name type
+ if (inherits(table_name, "Id")) {
+ # DBI::Id object - keep as is
+ } else if (is.character(table_name) && length(table_name) == 1) {
+ # Character string - keep as is
+ } else {
+ cli::cli_abort(
+ "`table_name` must be a single character string or a DBI::Id object"
+ )
+ }
- # Read the prompt file
- prompt_path <- system.file("prompts", "prompt.md", package = "querychat")
- prompt_content <- readLines(prompt_path, warn = FALSE)
- prompt_text <- paste(prompt_content, collapse = "\n")
+ # Check if table exists
+ if (!DBI::dbExistsTable(conn, table_name)) {
+ cli::cli_abort(c(
+ "Table {DBI::dbQuoteIdentifier(x, table_name)} not found in database.",
+ "i" = "If you're using a table in a catalog or schema, pass a DBI::Id object to `table_name`"
+ ))
+ }
- # Get schema for the data source
- schema <- get_schema(source, categorical_threshold = categorical_threshold)
+ private$conn <- conn
+ self$table_name <- table_name
+ },
- # Examine the data source and get the type for the prompt
- db_type <- get_db_type(source)
+ #' @description Get the database type
+ #' @return A string identifying the database type
+ get_db_type = function() {
+ # Special handling for known database types
+ if (inherits(private$conn, "duckdb_connection")) {
+ return("DuckDB")
+ }
+ if (inherits(private$conn, "SQLiteConnection")) {
+ return("SQLite")
+ }
- whisker::whisker.render(
- prompt_text,
- list(
- schema = schema,
- data_description = data_description,
- extra_instructions = extra_instructions,
- db_type = db_type,
- is_duck_db = identical(db_type, "DuckDB")
- )
+ # Default to 'POSIX' if dbms name not found
+ conn_info <- DBI::dbGetInfo(private$conn)
+ dbms_name <- purrr::pluck(conn_info, "dbms.name", .default = "POSIX")
+
+ # Remove ' SQL', if exists (SQL is already in the prompt)
+ gsub(" SQL", "", dbms_name)
+ },
+
+ #' @description
+ #' Get schema information for the database table
+ #'
+ #' @param categorical_threshold Maximum number of unique values for a text
+ #' column to be considered categorical (default: 20)
+ #' @return A string describing the schema
+ get_schema = function(categorical_threshold = 20) {
+ get_schema_impl(private$conn, self$table_name, categorical_threshold)
+ },
+
+ #' @description
+ #' Execute a SQL query
+ #'
+ #' @param query SQL query string. If NULL or empty, returns all data
+ #' @return A data frame with query results
+ execute_query = function(query) {
+ if (is.null(query) || query == "") {
+ query <- paste0(
+ "SELECT * FROM ",
+ DBI::dbQuoteIdentifier(private$conn, self$table_name)
+ )
+ }
+ DBI::dbGetQuery(private$conn, query)
+ },
+
+ #' @description
+ #' Test a SQL query by fetching only one row
+ #'
+ #' @param query SQL query string
+ #' @return A data frame with one row of results
+ test_query = function(query) {
+ rs <- DBI::dbSendQuery(private$conn, query)
+ df <- DBI::dbFetch(rs, n = 1)
+ DBI::dbClearResult(rs)
+ df
+ },
+
+ #' @description
+ #' Get all data from the table
+ #'
+ #' @return A data frame containing all data
+ get_data = function() {
+ self$execute_query(NULL)
+ },
+
+ #' @description
+ #' Disconnect from the database
+ #'
+ #' @return NULL (invisibly)
+ cleanup = function() {
+ if (!is.null(private$conn) && DBI::dbIsValid(private$conn)) {
+ DBI::dbDisconnect(private$conn)
+ }
+ invisible(NULL)
+ }
)
-}
+)
-#' Clean up a data source (close connections, etc.)
-#'
-#' An entrypoint for developers to create custom data sources for use with
-#' querychat. Most users shouldn't use this function directly; instead, they
-#' should call the `$cleanup()` method on the [QueryChat] object.
-#'
-#' @param source A querychat_data_source object
-#' @param ... Additional arguments passed to methods
-#' @return NULL (invisibly)
-#' @keywords internal
-#' @export
-cleanup_source <- function(source, ...) {
- UseMethod("cleanup_source")
-}
-
-#' @export
-cleanup_source.dbi_source <- function(source, ...) {
- if (!is.null(source$conn) && DBI::dbIsValid(source$conn)) {
- DBI::dbDisconnect(source$conn)
- }
- invisible(NULL)
-}
+# Helper Functions -------------------------------------------------------------
-#' Get schema for a data source
-#'
-#' An entrypoint for developers to create custom data sources for use with
-#' querychat. Most users shouldn't use this function directly; instead, they
-#' should call the `$set_system_prompt()` method on the [QueryChat] object.
+#' Check if object is a DataSource
#'
-#' @param source A querychat_data_source object
-#' @param categorical_threshold For text columns, the maximum number of unique values to consider as a categorical variable
-#' @param ... Additional arguments passed to methods
-#' @return A character string describing the schema
+#' @param x Object to check
+#' @return TRUE if x is a DataSource, FALSE otherwise
#' @keywords internal
-#' @export
-get_schema <- function(source, categorical_threshold = 20, ...) {
- UseMethod("get_schema")
+is_data_source <- function(x) {
+ inherits(x, "DataSource")
}
-#' @export
-get_schema.dbi_source <- function(source, categorical_threshold = 20, ...) {
- conn <- source$conn
- table_name <- source$table_name
+get_schema_impl <- function(conn, table_name, categorical_threshold = 20) {
# Get column information
columns <- DBI::dbListFields(conn, table_name)
@@ -448,7 +561,7 @@ get_schema.dbi_source <- function(source, categorical_threshold = 20, ...) {
}
-# Helper function to map R classes to SQL types
+# Map R classes to SQL types
r_class_to_sql_type <- function(r_class) {
switch(
r_class,
@@ -464,3 +577,51 @@ r_class_to_sql_type <- function(r_class) {
"TEXT" # default
)
}
+
+
+assemble_system_prompt <- function(
+ source,
+ data_description = NULL,
+ extra_instructions = NULL,
+ categorical_threshold = 20,
+ prompt_template = NULL
+) {
+ if (!is_data_source(source)) {
+ cli::cli_abort("`source` must be a DataSource object")
+ }
+
+ prompt_text <- read_text(
+ prompt_template %||%
+ system.file("prompts", "prompt.md", package = "querychat")
+ )
+
+ if (!is.null(data_description)) {
+ data_description <- read_text(data_description)
+ }
+ if (!is.null(extra_instructions)) {
+ extra_instructions <- read_text(extra_instructions)
+ }
+
+ schema <- source$get_schema(categorical_threshold = categorical_threshold)
+ db_type <- source$get_db_type()
+
+ whisker::whisker.render(
+ prompt_text,
+ list(
+ schema = schema,
+ data_description = data_description,
+ extra_instructions = extra_instructions,
+ db_type = db_type,
+ is_duck_db = identical(db_type, "DuckDB")
+ )
+ )
+}
+
+
+read_text <- function(x) {
+ if (file.exists(x)) {
+ x <- readLines(x, warn = FALSE)
+ }
+
+ paste(x, collapse = "\n")
+}
diff --git a/pkg-r/R/querychat-package.R b/pkg-r/R/querychat-package.R
index b38f2bb3..3ed5f093 100644
--- a/pkg-r/R/querychat-package.R
+++ b/pkg-r/R/querychat-package.R
@@ -41,7 +41,7 @@
#'
#' @section Main Components:
#' - [QueryChat]: The main R6 class for creating chat interfaces
-#' - [as_querychat_data_source()]: (Advanced) Create custom data source objects
+#' - [DataSource], [DataFrameSource], [DBISource]: R6 classes for data sources
#'
#' @section Examples:
#' To see examples included with the package, run:
diff --git a/pkg-r/R/querychat_tools.R b/pkg-r/R/querychat_tools.R
index e37a7a7b..6b3275b4 100644
--- a/pkg-r/R/querychat_tools.R
+++ b/pkg-r/R/querychat_tools.R
@@ -8,7 +8,7 @@ tool_update_dashboard <- function(
current_query,
current_title
) {
- db_type <- get_db_type(data_source)
+ db_type <- data_source$get_db_type()
ellmer::tool(
tool_update_dashboard_impl(data_source, current_query, current_title),
@@ -82,7 +82,7 @@ tool_reset_dashboard <- function(reset_fn) {
# @return The results of the query as a data frame.
tool_query <- function(data_source) {
force(data_source)
- db_type <- get_db_type(data_source)
+ db_type <- data_source$get_db_type()
ellmer::tool(
function(query, `_intent` = "") {
@@ -125,10 +125,10 @@ querychat_tool_result <- function(
switch(
action,
update = {
- test_query(data_source, query)
+ data_source$test_query(query)
NULL
},
- query = execute_query(data_source, query),
+ query = data_source$execute_query(query),
reset = "The dashboard has been reset to show all data."
),
error = function(err) err
diff --git a/pkg-r/man/DBISource.Rd b/pkg-r/man/DBISource.Rd
new file mode 100644
index 00000000..14a001af
--- /dev/null
+++ b/pkg-r/man/DBISource.Rd
@@ -0,0 +1,211 @@
+% Generated by roxygen2: do not edit by hand
+% Please edit documentation in R/data_source.R
+\name{DBISource}
+\alias{DBISource}
+\title{DBI Source}
+\description{
+A DataSource implementation for DBI database connections (SQLite, PostgreSQL,
+MySQL, etc.).
+}
+\details{
+This class wraps a DBI connection and provides SQL query execution against
+a specified table in the database.
+}
+\examples{
+\dontrun{
+# Connect to a database
+conn <- DBI::dbConnect(RSQLite::SQLite(), ":memory:")
+DBI::dbWriteTable(conn, "mtcars", mtcars)
+
+# Create a DBI source
+db_source <- DBISource$new(conn, "mtcars")
+
+# Get database type
+db_source$get_db_type() # Returns "SQLite"
+
+# Execute a query
+result <- db_source$execute_query("SELECT * FROM mtcars WHERE mpg > 25")
+
+# Note: cleanup() will disconnect the connection
+# If you want to keep the connection open, don't call cleanup()
+}
+
+## ------------------------------------------------
+## Method `DBISource$new`
+## ------------------------------------------------
+
+\dontrun{
+conn <- DBI::dbConnect(RSQLite::SQLite(), ":memory:")
+DBI::dbWriteTable(conn, "iris", iris)
+source <- DBISource$new(conn, "iris")
+}
+}
+\section{Super class}{
+\code{\link[querychat:DataSource]{querychat::DataSource}} -> \code{DBISource}
+}
+\section{Methods}{
+\subsection{Public methods}{
+\itemize{
+\item \href{#method-DBISource-new}{\code{DBISource$new()}}
+\item \href{#method-DBISource-get_db_type}{\code{DBISource$get_db_type()}}
+\item \href{#method-DBISource-get_schema}{\code{DBISource$get_schema()}}
+\item \href{#method-DBISource-execute_query}{\code{DBISource$execute_query()}}
+\item \href{#method-DBISource-test_query}{\code{DBISource$test_query()}}
+\item \href{#method-DBISource-get_data}{\code{DBISource$get_data()}}
+\item \href{#method-DBISource-cleanup}{\code{DBISource$cleanup()}}
+\item \href{#method-DBISource-clone}{\code{DBISource$clone()}}
+}
+}
+\if{html}{\out{
}}
+\if{html}{\out{}}
+\if{latex}{\out{\hypertarget{method-DBISource-new}{}}}
+\subsection{Method \code{new()}}{
+Create a new DBISource
+\subsection{Usage}{
+\if{html}{\out{}}\preformatted{DBISource$new(conn, table_name)}\if{html}{\out{
}}
+}
+
+\subsection{Arguments}{
+\if{html}{\out{}}
+\describe{
+\item{\code{conn}}{A DBI connection object}
+
+\item{\code{table_name}}{Name of the table in the database. Can be a character
+string or a \code{\link[DBI:Id]{DBI::Id()}} object for tables in catalogs/schemas}
+}
+\if{html}{\out{
}}
+}
+\subsection{Returns}{
+A new DBISource object
+}
+\subsection{Examples}{
+\if{html}{\out{}}
+\preformatted{\dontrun{
+conn <- DBI::dbConnect(RSQLite::SQLite(), ":memory:")
+DBI::dbWriteTable(conn, "iris", iris)
+source <- DBISource$new(conn, "iris")
+}
+}
+\if{html}{\out{
}}
+
+}
+
+}
+\if{html}{\out{
}}
+\if{html}{\out{}}
+\if{latex}{\out{\hypertarget{method-DBISource-get_db_type}{}}}
+\subsection{Method \code{get_db_type()}}{
+Get the database type
+\subsection{Usage}{
+\if{html}{\out{}}\preformatted{DBISource$get_db_type()}\if{html}{\out{
}}
+}
+
+\subsection{Returns}{
+A string identifying the database type
+}
+}
+\if{html}{\out{
}}
+\if{html}{\out{}}
+\if{latex}{\out{\hypertarget{method-DBISource-get_schema}{}}}
+\subsection{Method \code{get_schema()}}{
+Get schema information for the database table
+\subsection{Usage}{
+\if{html}{\out{}}\preformatted{DBISource$get_schema(categorical_threshold = 20)}\if{html}{\out{
}}
+}
+
+\subsection{Arguments}{
+\if{html}{\out{}}
+\describe{
+\item{\code{categorical_threshold}}{Maximum number of unique values for a text
+column to be considered categorical (default: 20)}
+}
+\if{html}{\out{
}}
+}
+\subsection{Returns}{
+A string describing the schema
+}
+}
+\if{html}{\out{
}}
+\if{html}{\out{}}
+\if{latex}{\out{\hypertarget{method-DBISource-execute_query}{}}}
+\subsection{Method \code{execute_query()}}{
+Execute a SQL query
+\subsection{Usage}{
+\if{html}{\out{}}\preformatted{DBISource$execute_query(query)}\if{html}{\out{
}}
+}
+
+\subsection{Arguments}{
+\if{html}{\out{}}
+\describe{
+\item{\code{query}}{SQL query string. If NULL or empty, returns all data}
+}
+\if{html}{\out{
}}
+}
+\subsection{Returns}{
+A data frame with query results
+}
+}
+\if{html}{\out{
}}
+\if{html}{\out{}}
+\if{latex}{\out{\hypertarget{method-DBISource-test_query}{}}}
+\subsection{Method \code{test_query()}}{
+Test a SQL query by fetching only one row
+\subsection{Usage}{
+\if{html}{\out{}}\preformatted{DBISource$test_query(query)}\if{html}{\out{
}}
+}
+
+\subsection{Arguments}{
+\if{html}{\out{}}
+\describe{
+\item{\code{query}}{SQL query string}
+}
+\if{html}{\out{
}}
+}
+\subsection{Returns}{
+A data frame with one row of results
+}
+}
+\if{html}{\out{
}}
+\if{html}{\out{}}
+\if{latex}{\out{\hypertarget{method-DBISource-get_data}{}}}
+\subsection{Method \code{get_data()}}{
+Get all data from the table
+\subsection{Usage}{
+\if{html}{\out{}}\preformatted{DBISource$get_data()}\if{html}{\out{
}}
+}
+
+\subsection{Returns}{
+A data frame containing all data
+}
+}
+\if{html}{\out{
}}
+\if{html}{\out{}}
+\if{latex}{\out{\hypertarget{method-DBISource-cleanup}{}}}
+\subsection{Method \code{cleanup()}}{
+Disconnect from the database
+\subsection{Usage}{
+\if{html}{\out{}}\preformatted{DBISource$cleanup()}\if{html}{\out{
}}
+}
+
+\subsection{Returns}{
+NULL (invisibly)
+}
+}
+\if{html}{\out{
}}
+\if{html}{\out{}}
+\if{latex}{\out{\hypertarget{method-DBISource-clone}{}}}
+\subsection{Method \code{clone()}}{
+The objects of this class are cloneable with this method.
+\subsection{Usage}{
+\if{html}{\out{}}\preformatted{DBISource$clone(deep = FALSE)}\if{html}{\out{
}}
+}
+
+\subsection{Arguments}{
+\if{html}{\out{}}
+\describe{
+\item{\code{deep}}{Whether to make a deep clone.}
+}
+\if{html}{\out{
}}
+}
+}
+}
diff --git a/pkg-r/man/DataFrameSource.Rd b/pkg-r/man/DataFrameSource.Rd
new file mode 100644
index 00000000..6885bf27
--- /dev/null
+++ b/pkg-r/man/DataFrameSource.Rd
@@ -0,0 +1,204 @@
+% Generated by roxygen2: do not edit by hand
+% Please edit documentation in R/data_source.R
+\name{DataFrameSource}
+\alias{DataFrameSource}
+\title{Data Frame Source}
+\description{
+A DataSource implementation that wraps a data frame using DuckDB for SQL
+query execution.
+}
+\details{
+This class creates an in-memory DuckDB connection and registers the provided
+data frame as a table. All SQL queries are executed against this DuckDB table.
+}
+\examples{
+\dontrun{
+# Create a data frame source
+df_source <- DataFrameSource$new(mtcars, "mtcars")
+
+# Get database type
+df_source$get_db_type() # Returns "DuckDB"
+
+# Execute a query
+result <- df_source$execute_query("SELECT * FROM mtcars WHERE mpg > 25")
+
+# Clean up when done
+df_source$cleanup()
+}
+
+## ------------------------------------------------
+## Method `DataFrameSource$new`
+## ------------------------------------------------
+
+\dontrun{
+source <- DataFrameSource$new(iris, "iris")
+}
+}
+\section{Super class}{
+\code{\link[querychat:DataSource]{querychat::DataSource}} -> \code{DataFrameSource}
+}
+\section{Methods}{
+\subsection{Public methods}{
+\itemize{
+\item \href{#method-DataFrameSource-new}{\code{DataFrameSource$new()}}
+\item \href{#method-DataFrameSource-get_db_type}{\code{DataFrameSource$get_db_type()}}
+\item \href{#method-DataFrameSource-get_schema}{\code{DataFrameSource$get_schema()}}
+\item \href{#method-DataFrameSource-execute_query}{\code{DataFrameSource$execute_query()}}
+\item \href{#method-DataFrameSource-test_query}{\code{DataFrameSource$test_query()}}
+\item \href{#method-DataFrameSource-get_data}{\code{DataFrameSource$get_data()}}
+\item \href{#method-DataFrameSource-cleanup}{\code{DataFrameSource$cleanup()}}
+\item \href{#method-DataFrameSource-clone}{\code{DataFrameSource$clone()}}
+}
+}
+\if{html}{\out{
}}
+\if{html}{\out{}}
+\if{latex}{\out{\hypertarget{method-DataFrameSource-new}{}}}
+\subsection{Method \code{new()}}{
+Create a new DataFrameSource
+\subsection{Usage}{
+\if{html}{\out{}}\preformatted{DataFrameSource$new(df, table_name)}\if{html}{\out{
}}
+}
+
+\subsection{Arguments}{
+\if{html}{\out{}}
+\describe{
+\item{\code{df}}{A data frame.}
+
+\item{\code{table_name}}{Name to use for the table in SQL queries. Must be a
+valid table name (start with letter, contain only letters, numbers,
+and underscores)}
+}
+\if{html}{\out{
}}
+}
+\subsection{Returns}{
+A new DataFrameSource object
+}
+\subsection{Examples}{
+\if{html}{\out{}}
+\preformatted{\dontrun{
+source <- DataFrameSource$new(iris, "iris")
+}
+}
+\if{html}{\out{
}}
+
+}
+
+}
+\if{html}{\out{
}}
+\if{html}{\out{}}
+\if{latex}{\out{\hypertarget{method-DataFrameSource-get_db_type}{}}}
+\subsection{Method \code{get_db_type()}}{
+Get the database type
+\subsection{Usage}{
+\if{html}{\out{}}\preformatted{DataFrameSource$get_db_type()}\if{html}{\out{
}}
+}
+
+\subsection{Returns}{
+The string "DuckDB"
+}
+}
+\if{html}{\out{
}}
+\if{html}{\out{}}
+\if{latex}{\out{\hypertarget{method-DataFrameSource-get_schema}{}}}
+\subsection{Method \code{get_schema()}}{
+Get schema information for the data frame
+\subsection{Usage}{
+\if{html}{\out{}}\preformatted{DataFrameSource$get_schema(categorical_threshold = 20)}\if{html}{\out{
}}
+}
+
+\subsection{Arguments}{
+\if{html}{\out{}}
+\describe{
+\item{\code{categorical_threshold}}{Maximum number of unique values for a text
+column to be considered categorical (default: 20)}
+}
+\if{html}{\out{
}}
+}
+\subsection{Returns}{
+A string describing the schema
+}
+}
+\if{html}{\out{
}}
+\if{html}{\out{}}
+\if{latex}{\out{\hypertarget{method-DataFrameSource-execute_query}{}}}
+\subsection{Method \code{execute_query()}}{
+Execute a SQL query
+\subsection{Usage}{
+\if{html}{\out{}}\preformatted{DataFrameSource$execute_query(query)}\if{html}{\out{
}}
+}
+
+\subsection{Arguments}{
+\if{html}{\out{}}
+\describe{
+\item{\code{query}}{SQL query string. If NULL or empty, returns all data}
+}
+\if{html}{\out{
}}
+}
+\subsection{Returns}{
+A data frame with query results
+}
+}
+\if{html}{\out{
}}
+\if{html}{\out{}}
+\if{latex}{\out{\hypertarget{method-DataFrameSource-test_query}{}}}
+\subsection{Method \code{test_query()}}{
+Test a SQL query by fetching only one row
+\subsection{Usage}{
+\if{html}{\out{}}\preformatted{DataFrameSource$test_query(query)}\if{html}{\out{
}}
+}
+
+\subsection{Arguments}{
+\if{html}{\out{}}
+\describe{
+\item{\code{query}}{SQL query string}
+}
+\if{html}{\out{
}}
+}
+\subsection{Returns}{
+A data frame with one row of results
+}
+}
+\if{html}{\out{
}}
+\if{html}{\out{}}
+\if{latex}{\out{\hypertarget{method-DataFrameSource-get_data}{}}}
+\subsection{Method \code{get_data()}}{
+Get all data from the table
+\subsection{Usage}{
+\if{html}{\out{}}\preformatted{DataFrameSource$get_data()}\if{html}{\out{
}}
+}
+
+\subsection{Returns}{
+A data frame containing all data
+}
+}
+\if{html}{\out{
}}
+\if{html}{\out{}}
+\if{latex}{\out{\hypertarget{method-DataFrameSource-cleanup}{}}}
+\subsection{Method \code{cleanup()}}{
+Close the DuckDB connection
+\subsection{Usage}{
+\if{html}{\out{}}\preformatted{DataFrameSource$cleanup()}\if{html}{\out{
}}
+}
+
+\subsection{Returns}{
+NULL (invisibly)
+}
+}
+\if{html}{\out{
}}
+\if{html}{\out{}}
+\if{latex}{\out{\hypertarget{method-DataFrameSource-clone}{}}}
+\subsection{Method \code{clone()}}{
+The objects of this class are cloneable with this method.
+\subsection{Usage}{
+\if{html}{\out{}}\preformatted{DataFrameSource$clone(deep = FALSE)}\if{html}{\out{
}}
+}
+
+\subsection{Arguments}{
+\if{html}{\out{}}
+\describe{
+\item{\code{deep}}{Whether to make a deep clone.}
+}
+\if{html}{\out{
}}
+}
+}
+}
diff --git a/pkg-r/man/DataSource.Rd b/pkg-r/man/DataSource.Rd
new file mode 100644
index 00000000..f9bb2350
--- /dev/null
+++ b/pkg-r/man/DataSource.Rd
@@ -0,0 +1,148 @@
+% Generated by roxygen2: do not edit by hand
+% Please edit documentation in R/data_source.R
+\name{DataSource}
+\alias{DataSource}
+\title{Data Source Base Class}
+\description{
+An abstract R6 class defining the interface that custom QueryChat data
+sources must implement. This class should not be instantiated directly;
+instead, use one of its concrete implementations like \link{DataFrameSource} or
+\link{DBISource}.
+}
+\section{Public fields}{
+\if{html}{\out{}}
+\describe{
+\item{\code{table_name}}{Name of the table to be used in SQL queries}
+}
+\if{html}{\out{
}}
+}
+\section{Methods}{
+\subsection{Public methods}{
+\itemize{
+\item \href{#method-DataSource-get_db_type}{\code{DataSource$get_db_type()}}
+\item \href{#method-DataSource-get_schema}{\code{DataSource$get_schema()}}
+\item \href{#method-DataSource-execute_query}{\code{DataSource$execute_query()}}
+\item \href{#method-DataSource-test_query}{\code{DataSource$test_query()}}
+\item \href{#method-DataSource-get_data}{\code{DataSource$get_data()}}
+\item \href{#method-DataSource-cleanup}{\code{DataSource$cleanup()}}
+\item \href{#method-DataSource-clone}{\code{DataSource$clone()}}
+}
+}
+\if{html}{\out{
}}
+\if{html}{\out{}}
+\if{latex}{\out{\hypertarget{method-DataSource-get_db_type}{}}}
+\subsection{Method \code{get_db_type()}}{
+Get the database type
+\subsection{Usage}{
+\if{html}{\out{}}\preformatted{DataSource$get_db_type()}\if{html}{\out{
}}
+}
+
+\subsection{Returns}{
+A string describing the database type (e.g., "DuckDB", "SQLite")
+}
+}
+\if{html}{\out{
}}
+\if{html}{\out{}}
+\if{latex}{\out{\hypertarget{method-DataSource-get_schema}{}}}
+\subsection{Method \code{get_schema()}}{
+Get schema information about the table
+\subsection{Usage}{
+\if{html}{\out{}}\preformatted{DataSource$get_schema(categorical_threshold = 20)}\if{html}{\out{
}}
+}
+
+\subsection{Arguments}{
+\if{html}{\out{}}
+\describe{
+\item{\code{categorical_threshold}}{Maximum number of unique values for a text
+column to be considered categorical}
+}
+\if{html}{\out{
}}
+}
+\subsection{Returns}{
+A string containing schema information formatted for LLM prompts
+}
+}
+\if{html}{\out{
}}
+\if{html}{\out{}}
+\if{latex}{\out{\hypertarget{method-DataSource-execute_query}{}}}
+\subsection{Method \code{execute_query()}}{
+Execute a SQL query and return results
+\subsection{Usage}{
+\if{html}{\out{}}\preformatted{DataSource$execute_query(query)}\if{html}{\out{
}}
+}
+
+\subsection{Arguments}{
+\if{html}{\out{}}
+\describe{
+\item{\code{query}}{SQL query string to execute}
+}
+\if{html}{\out{
}}
+}
+\subsection{Returns}{
+A data frame containing query results
+}
+}
+\if{html}{\out{
}}
+\if{html}{\out{}}
+\if{latex}{\out{\hypertarget{method-DataSource-test_query}{}}}
+\subsection{Method \code{test_query()}}{
+Test a SQL query by fetching only one row
+\subsection{Usage}{
+\if{html}{\out{}}\preformatted{DataSource$test_query(query)}\if{html}{\out{
}}
+}
+
+\subsection{Arguments}{
+\if{html}{\out{}}
+\describe{
+\item{\code{query}}{SQL query string to test}
+}
+\if{html}{\out{
}}
+}
+\subsection{Returns}{
+A data frame containing one row of results (or empty if no matches)
+}
+}
+\if{html}{\out{
}}
+\if{html}{\out{}}
+\if{latex}{\out{\hypertarget{method-DataSource-get_data}{}}}
+\subsection{Method \code{get_data()}}{
+Get the unfiltered data as a data frame
+\subsection{Usage}{
+\if{html}{\out{}}\preformatted{DataSource$get_data()}\if{html}{\out{
}}
+}
+
+\subsection{Returns}{
+A data frame containing all data from the table
+}
+}
+\if{html}{\out{
}}
+\if{html}{\out{}}
+\if{latex}{\out{\hypertarget{method-DataSource-cleanup}{}}}
+\subsection{Method \code{cleanup()}}{
+Clean up resources (close connections, etc.)
+\subsection{Usage}{
+\if{html}{\out{}}\preformatted{DataSource$cleanup()}\if{html}{\out{
}}
+}
+
+\subsection{Returns}{
+NULL (invisibly)
+}
+}
+\if{html}{\out{
}}
+\if{html}{\out{}}
+\if{latex}{\out{\hypertarget{method-DataSource-clone}{}}}
+\subsection{Method \code{clone()}}{
+The objects of this class are cloneable with this method.
+\subsection{Usage}{
+\if{html}{\out{}}\preformatted{DataSource$clone(deep = FALSE)}\if{html}{\out{
}}
+}
+
+\subsection{Arguments}{
+\if{html}{\out{}}
+\describe{
+\item{\code{deep}}{Whether to make a deep clone.}
+}
+\if{html}{\out{
}}
+}
+}
+}
diff --git a/pkg-r/man/as_querychat_data_source.Rd b/pkg-r/man/as_querychat_data_source.Rd
deleted file mode 100644
index a9bc4cec..00000000
--- a/pkg-r/man/as_querychat_data_source.Rd
+++ /dev/null
@@ -1,26 +0,0 @@
-% Generated by roxygen2: do not edit by hand
-% Please edit documentation in R/data_source.R
-\name{as_querychat_data_source}
-\alias{as_querychat_data_source}
-\title{Create a data source for querychat}
-\usage{
-as_querychat_data_source(x, table_name, ...)
-}
-\arguments{
-\item{x}{A data frame or DBI connection}
-
-\item{table_name}{The name to use for the table in the data source. Can be:
-\itemize{
-\item A character string (e.g., "table_name")
-\item Or, for tables contained within catalogs or schemas, a \code{\link[DBI:Id]{DBI::Id()}} object (e.g., \code{DBI::Id(schema = "schema_name", table = "table_name")})
-}}
-}
-\value{
-A querychat_data_source object
-}
-\description{
-An entrypoint for developers to create custom data sources for use with
-querychat. Most users shouldn't use this function directly; instead, they
-should pass their data to \code{QueryChat$new()}.
-}
-\keyword{internal}
diff --git a/pkg-r/man/cleanup_source.Rd b/pkg-r/man/cleanup_source.Rd
deleted file mode 100644
index 938f585b..00000000
--- a/pkg-r/man/cleanup_source.Rd
+++ /dev/null
@@ -1,22 +0,0 @@
-% Generated by roxygen2: do not edit by hand
-% Please edit documentation in R/data_source.R
-\name{cleanup_source}
-\alias{cleanup_source}
-\title{Clean up a data source (close connections, etc.)}
-\usage{
-cleanup_source(source, ...)
-}
-\arguments{
-\item{source}{A querychat_data_source object}
-
-\item{...}{Additional arguments passed to methods}
-}
-\value{
-NULL (invisibly)
-}
-\description{
-An entrypoint for developers to create custom data sources for use with
-querychat. Most users shouldn't use this function directly; instead, they
-should call the \verb{$cleanup()} method on the \link{QueryChat} object.
-}
-\keyword{internal}
diff --git a/pkg-r/man/create_system_prompt.Rd b/pkg-r/man/create_system_prompt.Rd
deleted file mode 100644
index 46e93c4b..00000000
--- a/pkg-r/man/create_system_prompt.Rd
+++ /dev/null
@@ -1,35 +0,0 @@
-% Generated by roxygen2: do not edit by hand
-% Please edit documentation in R/data_source.R
-\name{create_system_prompt}
-\alias{create_system_prompt}
-\title{Create a system prompt for the data source}
-\usage{
-create_system_prompt(
- source,
- data_description = NULL,
- extra_instructions = NULL,
- categorical_threshold = 20,
- ...
-)
-}
-\arguments{
-\item{source}{A querychat_data_source object}
-
-\item{data_description}{Optional description of the data}
-
-\item{extra_instructions}{Optional additional instructions}
-
-\item{categorical_threshold}{For text columns, the maximum number of unique
-values to consider as a categorical variable}
-
-\item{...}{Additional arguments passed to methods}
-}
-\value{
-A string with the system prompt
-}
-\description{
-An entrypoint for developers to create custom data sources for use with
-querychat. Most users shouldn't use this function directly; instead, they
-should call the \verb{$set_system_prompt()} method on the \link{QueryChat} object.
-}
-\keyword{internal}
diff --git a/pkg-r/man/execute_query.Rd b/pkg-r/man/execute_query.Rd
deleted file mode 100644
index ab8fd054..00000000
--- a/pkg-r/man/execute_query.Rd
+++ /dev/null
@@ -1,24 +0,0 @@
-% Generated by roxygen2: do not edit by hand
-% Please edit documentation in R/data_source.R
-\name{execute_query}
-\alias{execute_query}
-\title{Execute an SQL query on a data source}
-\usage{
-execute_query(source, query, ...)
-}
-\arguments{
-\item{source}{A querychat_data_source object}
-
-\item{query}{SQL query string}
-
-\item{...}{Additional arguments passed to methods}
-}
-\value{
-Result of the query as a data frame
-}
-\description{
-An entrypoint for developers to create custom data source objects for use
-with querychat. Most users shouldn't use this function directly; instead,
-they call the \verb{$sql()} method on the \link{QueryChat} object to run queries.
-}
-\keyword{internal}
diff --git a/pkg-r/man/get_db_type.Rd b/pkg-r/man/get_db_type.Rd
deleted file mode 100644
index 70d4e01b..00000000
--- a/pkg-r/man/get_db_type.Rd
+++ /dev/null
@@ -1,22 +0,0 @@
-% Generated by roxygen2: do not edit by hand
-% Please edit documentation in R/data_source.R
-\name{get_db_type}
-\alias{get_db_type}
-\title{Get type information for a data source}
-\usage{
-get_db_type(source, ...)
-}
-\arguments{
-\item{source}{A querychat_data_source object}
-
-\item{...}{Additional arguments passed to methods}
-}
-\value{
-A character string containing the type information
-}
-\description{
-An entrypoint for developers to create custom data sources for use with
-querychat. Most users shouldn't use this function directly; instead, they
-should call the \verb{$set_system_prompt()} method on the \link{QueryChat} object.
-}
-\keyword{internal}
diff --git a/pkg-r/man/get_schema.Rd b/pkg-r/man/get_schema.Rd
deleted file mode 100644
index 2b4c9ded..00000000
--- a/pkg-r/man/get_schema.Rd
+++ /dev/null
@@ -1,24 +0,0 @@
-% Generated by roxygen2: do not edit by hand
-% Please edit documentation in R/data_source.R
-\name{get_schema}
-\alias{get_schema}
-\title{Get schema for a data source}
-\usage{
-get_schema(source, categorical_threshold = 20, ...)
-}
-\arguments{
-\item{source}{A querychat_data_source object}
-
-\item{categorical_threshold}{For text columns, the maximum number of unique values to consider as a categorical variable}
-
-\item{...}{Additional arguments passed to methods}
-}
-\value{
-A character string describing the schema
-}
-\description{
-An entrypoint for developers to create custom data sources for use with
-querychat. Most users shouldn't use this function directly; instead, they
-should call the \verb{$set_system_prompt()} method on the \link{QueryChat} object.
-}
-\keyword{internal}
diff --git a/pkg-r/man/is_data_source.Rd b/pkg-r/man/is_data_source.Rd
new file mode 100644
index 00000000..6e7348b2
--- /dev/null
+++ b/pkg-r/man/is_data_source.Rd
@@ -0,0 +1,18 @@
+% Generated by roxygen2: do not edit by hand
+% Please edit documentation in R/data_source.R
+\name{is_data_source}
+\alias{is_data_source}
+\title{Check if object is a DataSource}
+\usage{
+is_data_source(x)
+}
+\arguments{
+\item{x}{Object to check}
+}
+\value{
+TRUE if x is a DataSource, FALSE otherwise
+}
+\description{
+Check if object is a DataSource
+}
+\keyword{internal}
diff --git a/pkg-r/man/querychat-package.Rd b/pkg-r/man/querychat-package.Rd
index 67a92081..63598eda 100644
--- a/pkg-r/man/querychat-package.Rd
+++ b/pkg-r/man/querychat-package.Rd
@@ -52,7 +52,7 @@ shinyApp(ui, server)
\itemize{
\item \link{QueryChat}: The main R6 class for creating chat interfaces
-\item \code{\link[=as_querychat_data_source]{as_querychat_data_source()}}: (Advanced) Create custom data source objects
+\item \link{DataSource}, \link{DataFrameSource}, \link{DBISource}: R6 classes for data sources
}
}
diff --git a/pkg-r/man/test_query.Rd b/pkg-r/man/test_query.Rd
deleted file mode 100644
index d8307c6f..00000000
--- a/pkg-r/man/test_query.Rd
+++ /dev/null
@@ -1,24 +0,0 @@
-% Generated by roxygen2: do not edit by hand
-% Please edit documentation in R/data_source.R
-\name{test_query}
-\alias{test_query}
-\title{Test a SQL query on a data source.}
-\usage{
-test_query(source, query, ...)
-}
-\arguments{
-\item{source}{A querychat_data_source object}
-
-\item{query}{SQL query string}
-
-\item{...}{Additional arguments passed to methods}
-}
-\value{
-Result of the query, limited to one row of data.
-}
-\description{
-An entrypoint for developers to create custom data sources for use with
-querychat. Most users shouldn't use this function directly; instead, they
-should call the \verb{$sql()} method on the \link{QueryChat} object to run queries.
-}
-\keyword{internal}
diff --git a/pkg-r/tests/testthat/test-data-source.R b/pkg-r/tests/testthat/test-data-source.R
index 343f336a..de28dcc9 100644
--- a/pkg-r/tests/testthat/test-data-source.R
+++ b/pkg-r/tests/testthat/test-data-source.R
@@ -3,7 +3,7 @@ library(DBI)
library(RSQLite)
library(querychat)
-test_that("as_querychat_data_source.data.frame creates proper S3 object", {
+test_that("DataFrameSource creates proper R6 object", {
# Create a simple data frame
test_df <- data.frame(
id = 1:5,
@@ -13,16 +13,15 @@ test_that("as_querychat_data_source.data.frame creates proper S3 object", {
)
# Test with explicit table name
- source <- as_querychat_data_source(test_df, table_name = "test_table")
- withr::defer(cleanup_source(source))
+ source <- DataFrameSource$new(test_df, "test_table")
+ withr::defer(source$cleanup())
- expect_s3_class(source, "data_frame_source")
- expect_s3_class(source, "querychat_data_source")
+ expect_s3_class(source, "DataFrameSource")
+ expect_s3_class(source, "DataSource")
expect_equal(source$table_name, "test_table")
- expect_true(inherits(source$conn, "DBIConnection"))
})
-test_that("as_querychat_data_source.DBIConnection creates proper S3 object", {
+test_that("DBISource creates proper R6 object", {
# Create temporary SQLite database
temp_db <- withr::local_tempfile(fileext = ".db")
conn <- dbConnect(RSQLite::SQLite(), temp_db)
@@ -39,9 +38,9 @@ test_that("as_querychat_data_source.DBIConnection creates proper S3 object", {
dbWriteTable(conn, "users", test_data, overwrite = TRUE)
# Test DBI source creation
- db_source <- as_querychat_data_source(conn, "users")
- expect_s3_class(db_source, "dbi_source")
- expect_s3_class(db_source, "querychat_data_source")
+ db_source <- DBISource$new(conn, "users")
+ expect_s3_class(db_source, "DBISource")
+ expect_s3_class(db_source, "DataSource")
expect_equal(db_source$table_name, "users")
})
@@ -54,10 +53,10 @@ test_that("get_schema methods return proper schema", {
stringsAsFactors = FALSE
)
- df_source <- as_querychat_data_source(test_df, table_name = "test_table")
- withr::defer(cleanup_source(df_source))
+ df_source <- DataFrameSource$new(test_df, "test_table")
+ withr::defer(df_source$cleanup())
- schema <- get_schema(df_source)
+ schema <- df_source$get_schema()
expect_type(schema, "character")
expect_match(schema, "Table: test_table")
expect_match(schema, "id \\(INTEGER\\)")
@@ -75,8 +74,8 @@ test_that("get_schema methods return proper schema", {
dbWriteTable(conn, "test_table", test_df, overwrite = TRUE)
- dbi_source <- as_querychat_data_source(conn, "test_table")
- schema <- get_schema(dbi_source)
+ dbi_source <- DBISource$new(conn, "test_table")
+ schema <- dbi_source$get_schema()
expect_type(schema, "character")
expect_match(schema, "Table: `test_table`")
expect_match(schema, "id \\(INTEGER\\)")
@@ -94,12 +93,9 @@ test_that("execute_query works for both source types", {
stringsAsFactors = FALSE
)
- df_source <- as_querychat_data_source(test_df, table_name = "test_table")
- withr::defer(cleanup_source(df_source))
- result <- execute_query(
- df_source,
- "SELECT * FROM test_table WHERE value > 25"
- )
+ df_source <- DataFrameSource$new(test_df, "test_table")
+ withr::defer(df_source$cleanup())
+ result <- df_source$execute_query("SELECT * FROM test_table WHERE value > 25")
expect_s3_class(result, "data.frame")
expect_equal(nrow(result), 3) # Should return 3 rows (30, 40, 50)
@@ -109,9 +105,8 @@ test_that("execute_query works for both source types", {
withr::defer(dbDisconnect(conn))
dbWriteTable(conn, "test_table", test_df, overwrite = TRUE)
- dbi_source <- as_querychat_data_source(conn, "test_table")
- result <- execute_query(
- dbi_source,
+ dbi_source <- DBISource$new(conn, "test_table")
+ result <- dbi_source$execute_query(
"SELECT * FROM test_table WHERE value > 25"
)
expect_s3_class(result, "data.frame")
@@ -126,17 +121,17 @@ test_that("execute_query works with empty/null queries", {
stringsAsFactors = FALSE
)
- df_source <- as_querychat_data_source(test_df, table_name = "test_table")
- withr::defer(cleanup_source(df_source))
+ df_source <- DataFrameSource$new(test_df, "test_table")
+ withr::defer(df_source$cleanup())
# Test with NULL query
- result_null <- execute_query(df_source, NULL)
+ result_null <- df_source$execute_query(NULL)
expect_s3_class(result_null, "data.frame")
expect_equal(nrow(result_null), 5) # Should return all rows
expect_equal(ncol(result_null), 2) # Should return all columns
# Test with empty string query
- result_empty <- execute_query(df_source, "")
+ result_empty <- df_source$execute_query("")
expect_s3_class(result_empty, "data.frame")
expect_equal(nrow(result_empty), 5) # Should return all rows
expect_equal(ncol(result_empty), 2) # Should return all columns
@@ -148,16 +143,16 @@ test_that("execute_query works with empty/null queries", {
dbWriteTable(conn, "test_table", test_df, overwrite = TRUE)
- dbi_source <- as_querychat_data_source(conn, "test_table")
+ dbi_source <- DBISource$new(conn, "test_table")
# Test with NULL query
- result_null <- execute_query(dbi_source, NULL)
+ result_null <- dbi_source$execute_query(NULL)
expect_s3_class(result_null, "data.frame")
expect_equal(nrow(result_null), 5) # Should return all rows
expect_equal(ncol(result_null), 2) # Should return all columns
# Test with empty string query
- result_empty <- execute_query(dbi_source, "")
+ result_empty <- dbi_source$execute_query("")
expect_s3_class(result_empty, "data.frame")
expect_equal(nrow(result_empty), 5) # Should return all rows
expect_equal(ncol(result_empty), 2) # Should return all columns
@@ -173,9 +168,9 @@ test_that("get_schema correctly reports min/max values for numeric columns", {
stringsAsFactors = FALSE
)
- df_source <- as_querychat_data_source(test_df, table_name = "test_metrics")
- withr::defer(cleanup_source(df_source))
- schema <- get_schema(df_source)
+ df_source <- DataFrameSource$new(test_df, "test_metrics")
+ withr::defer(df_source$cleanup())
+ schema <- df_source$get_schema()
# Check that each numeric column has the correct min/max values
expect_match(schema, "- id \\(INTEGER\\)\\n Range: 1 to 5")
@@ -184,17 +179,17 @@ test_that("get_schema correctly reports min/max values for numeric columns", {
expect_match(schema, "- count \\(FLOAT\\)\\n Range: 50 to 200")
})
-test_that("create_system_prompt generates appropriate system prompt", {
+test_that("assemble_system_prompt generates appropriate system prompt", {
test_df <- data.frame(
id = 1:3,
name = c("A", "B", "C"),
stringsAsFactors = FALSE
)
- df_source <- as_querychat_data_source(test_df, table_name = "test_table")
- withr::defer(cleanup_source(df_source))
+ df_source <- DataFrameSource$new(test_df, "test_table")
+ withr::defer(df_source$cleanup())
- prompt <- create_system_prompt(
+ prompt <- assemble_system_prompt(
df_source,
data_description = "A test dataframe"
)
@@ -216,19 +211,19 @@ test_that("QueryChat$new() automatically handles data.frame inputs", {
)
withr::defer(qc$cleanup())
- expect_s3_class(qc$data_source, "querychat_data_source")
- expect_s3_class(qc$data_source, "data_frame_source")
+ expect_s3_class(qc$data_source, "DataSource")
+ expect_s3_class(qc$data_source, "DataFrameSource")
# Should work with proper data source too
- df_source <- as_querychat_data_source(test_df, table_name = "test_table")
- withr::defer(cleanup_source(df_source))
+ df_source <- DataFrameSource$new(test_df, "test_table")
+ withr::defer(df_source$cleanup())
qc2 <- QueryChat$new(
data_source = df_source,
table_name = "test_table",
greeting = "Test greeting"
)
- expect_s3_class(qc2$data_source, "querychat_data_source")
+ expect_s3_class(qc2$data_source, "DataSource")
})
test_that("QueryChat$new() works with both source types", {
@@ -240,8 +235,8 @@ test_that("QueryChat$new() works with both source types", {
)
# Create data source and test with QueryChat$new()
- df_source <- as_querychat_data_source(test_df, table_name = "test_source")
- withr::defer(cleanup_source(df_source))
+ df_source <- DataFrameSource$new(test_df, "test_source")
+ withr::defer(df_source$cleanup())
qc <- QueryChat$new(
data_source = df_source,
@@ -249,7 +244,7 @@ test_that("QueryChat$new() works with both source types", {
greeting = "Test greeting"
)
- expect_s3_class(qc$data_source, "data_frame_source")
+ expect_s3_class(qc$data_source, "DataFrameSource")
expect_equal(qc$data_source$table_name, "test_source")
# Test with database connection
@@ -259,12 +254,41 @@ test_that("QueryChat$new() works with both source types", {
dbWriteTable(conn, "test_table", test_df, overwrite = TRUE)
- dbi_source <- as_querychat_data_source(conn, "test_table")
+ dbi_source <- DBISource$new(conn, "test_table")
qc2 <- QueryChat$new(
data_source = dbi_source,
table_name = "test_table",
greeting = "Test greeting"
)
- expect_s3_class(qc2$data_source, "dbi_source")
+ expect_s3_class(qc2$data_source, "DBISource")
expect_equal(qc2$data_source$table_name, "test_table")
})
+
+test_that("get_data returns all data", {
+ # Test with data frame source
+ test_df <- data.frame(
+ id = 1:5,
+ value = c(10, 20, 30, 40, 50),
+ stringsAsFactors = FALSE
+ )
+
+ df_source <- DataFrameSource$new(test_df, "test_table")
+ withr::defer(df_source$cleanup())
+
+ result <- df_source$get_data()
+ expect_s3_class(result, "data.frame")
+ expect_equal(nrow(result), 5)
+ expect_equal(ncol(result), 2)
+
+ # Test with DBI source
+ temp_db <- withr::local_tempfile(fileext = ".db")
+ conn <- dbConnect(RSQLite::SQLite(), temp_db)
+ withr::defer(dbDisconnect(conn))
+ dbWriteTable(conn, "test_table", test_df, overwrite = TRUE)
+
+ dbi_source <- DBISource$new(conn, "test_table")
+ result <- dbi_source$get_data()
+ expect_s3_class(result, "data.frame")
+ expect_equal(nrow(result), 5)
+ expect_equal(ncol(result), 2)
+})
diff --git a/pkg-r/tests/testthat/test-db-type.R b/pkg-r/tests/testthat/test-db-type.R
index 7894d6e8..765be28e 100644
--- a/pkg-r/tests/testthat/test-db-type.R
+++ b/pkg-r/tests/testthat/test-db-type.R
@@ -1,15 +1,16 @@
library(testthat)
-test_that("get_db_type returns correct type for data_frame_source", {
+test_that("get_db_type returns correct type for DataFrameSource", {
# Create a simple data frame source
df <- data.frame(x = 1:5, y = letters[1:5])
- df_source <- as_querychat_data_source(df, "test_table")
+ df_source <- DataFrameSource$new(df, "test_table")
+ withr::defer(df_source$cleanup())
# Test that get_db_type returns "DuckDB"
- expect_equal(get_db_type(df_source), "DuckDB")
+ expect_equal(df_source$get_db_type(), "DuckDB")
})
-test_that("get_db_type returns correct type for dbi_source with SQLite", {
+test_that("get_db_type returns correct type for DBISource with SQLite", {
skip_if_not_installed("RSQLite")
# Create a SQLite database source
@@ -17,19 +18,20 @@ test_that("get_db_type returns correct type for dbi_source with SQLite", {
conn <- DBI::dbConnect(RSQLite::SQLite(), temp_db)
withr::defer(DBI::dbDisconnect(conn))
DBI::dbWriteTable(conn, "test_table", data.frame(x = 1:5, y = letters[1:5]))
- db_source <- as_querychat_data_source(conn, "test_table")
+ db_source <- DBISource$new(conn, "test_table")
# Test that get_db_type returns the correct database type
- expect_equal(get_db_type(db_source), "SQLite")
+ expect_equal(db_source$get_db_type(), "SQLite")
})
-test_that("get_db_type is correctly used in create_system_prompt", {
+test_that("get_db_type is correctly used in assemble_system_prompt", {
# Create a simple data frame source
df <- data.frame(x = 1:5, y = letters[1:5])
- df_source <- as_querychat_data_source(df, "test_table")
+ df_source <- DataFrameSource$new(df, "test_table")
+ withr::defer(df_source$cleanup())
# Generate system prompt
- sys_prompt <- create_system_prompt(df_source)
+ sys_prompt <- assemble_system_prompt(df_source)
# Check that "DuckDB" appears in the prompt content
expect_true(grepl("DuckDB SQL", sys_prompt, fixed = TRUE))
@@ -38,10 +40,11 @@ test_that("get_db_type is correctly used in create_system_prompt", {
test_that("get_db_type is used to customize prompt template", {
# Create a simple data frame source
df <- data.frame(x = 1:5, y = letters[1:5])
- df_source <- as_querychat_data_source(df, "test_table")
+ df_source <- DataFrameSource$new(df, "test_table")
+ withr::defer(df_source$cleanup())
# Get the db_type
- db_type <- get_db_type(df_source)
+ db_type <- df_source$get_db_type()
# Check that the db_type is correctly returned
expect_equal(db_type, "DuckDB")
@@ -49,6 +52,6 @@ test_that("get_db_type is used to customize prompt template", {
# Verify the value is used in the system prompt
# This is an indirect test that doesn't need mocking
# We just check that the string appears somewhere in the system prompt
- prompt <- create_system_prompt(df_source)
+ prompt <- assemble_system_prompt(df_source)
expect_true(grepl(db_type, prompt, fixed = TRUE))
})
diff --git a/pkg-r/tests/testthat/test-querychat-server.R b/pkg-r/tests/testthat/test-querychat-server.R
index 503620d0..2c4e3a9f 100644
--- a/pkg-r/tests/testthat/test-querychat-server.R
+++ b/pkg-r/tests/testthat/test-querychat-server.R
@@ -20,23 +20,22 @@ test_that("database source query functionality", {
dbWriteTable(conn, "users", test_data, overwrite = TRUE)
# Create database source
- db_source <- as_querychat_data_source(conn, "users")
+ db_source <- DBISource$new(conn, "users")
# Test that we can execute queries
- result <- execute_query(db_source, "SELECT * FROM users WHERE age > 30")
+ result <- db_source$execute_query("SELECT * FROM users WHERE age > 30")
expect_s3_class(result, "data.frame")
expect_equal(nrow(result), 2) # Charlie and Eve
expect_equal(result$name, c("Charlie", "Eve"))
# Test that we can get all data
- all_data <- execute_query(db_source, NULL)
+ all_data <- db_source$execute_query(NULL)
expect_s3_class(all_data, "data.frame")
expect_equal(nrow(all_data), 5)
expect_equal(ncol(all_data), 3)
# Test ordering works
- ordered_result <- execute_query(
- db_source,
+ ordered_result <- db_source$execute_query(
"SELECT * FROM users ORDER BY age DESC"
)
expect_equal(ordered_result$name[1], "Charlie") # Oldest first
diff --git a/pkg-r/tests/testthat/test-shiny-app.R b/pkg-r/tests/testthat/test-shiny-app.R
index b4b7cf3c..589981e0 100644
--- a/pkg-r/tests/testthat/test-shiny-app.R
+++ b/pkg-r/tests/testthat/test-shiny-app.R
@@ -35,7 +35,7 @@ test_that("database reactive functionality works correctly", {
db_conn <- dbConnect(RSQLite::SQLite(), temp_db)
withr::defer(dbDisconnect(db_conn))
- iris_source <- as_querychat_data_source(db_conn, "iris")
+ iris_source <- DBISource$new(db_conn, "iris")
# Mock chat function
mock_client <- ellmer::chat_openai(api_key = "boop")
@@ -48,18 +48,17 @@ test_that("database reactive functionality works correctly", {
client = mock_client
)
- expect_s3_class(qc$data_source, "dbi_source")
- expect_s3_class(qc$data_source, "querychat_data_source")
+ expect_s3_class(qc$data_source, "DBISource")
+ expect_s3_class(qc$data_source, "DataSource")
# Test that we can get all data
- result_data <- execute_query(qc$data_source, NULL)
+ result_data <- qc$data_source$execute_query(NULL)
expect_s3_class(result_data, "data.frame")
expect_equal(nrow(result_data), 150)
expect_equal(ncol(result_data), 5)
# Test with a specific query
- query_result <- execute_query(
- qc$data_source,
+ query_result <- qc$data_source$execute_query(
"SELECT \"Sepal.Length\", \"Sepal.Width\" FROM iris WHERE \"Species\" = 'setosa'"
)
expect_s3_class(query_result, "data.frame")
diff --git a/pkg-r/tests/testthat/test-sql-comments.R b/pkg-r/tests/testthat/test-sql-comments.R
index 737093ca..bf056879 100644
--- a/pkg-r/tests/testthat/test-sql-comments.R
+++ b/pkg-r/tests/testthat/test-sql-comments.R
@@ -12,8 +12,8 @@ test_that("execute_query handles SQL with inline comments", {
)
# Create data source
- df_source <- as_querychat_data_source(test_df, table_name = "test_table")
- withr::defer(cleanup_source(df_source))
+ df_source <- DataFrameSource$new(test_df, "test_table")
+ withr::defer(df_source$cleanup())
# Test with inline comments
inline_comment_query <- "
@@ -22,7 +22,7 @@ test_that("execute_query handles SQL with inline comments", {
WHERE value > 25 -- Filter for higher values
"
- result <- execute_query(df_source, inline_comment_query)
+ result <- df_source$execute_query(inline_comment_query)
expect_s3_class(result, "data.frame")
expect_equal(nrow(result), 3) # Should return 3 rows (30, 40, 50)
expect_equal(ncol(result), 2)
@@ -36,7 +36,7 @@ test_that("execute_query handles SQL with inline comments", {
WHERE value > 25 -- Only higher values
"
- result <- execute_query(df_source, multiple_comments_query)
+ result <- df_source$execute_query(multiple_comments_query)
expect_s3_class(result, "data.frame")
expect_equal(nrow(result), 3)
expect_equal(ncol(result), 2)
@@ -51,8 +51,8 @@ test_that("execute_query handles SQL with multiline comments", {
)
# Create data source
- df_source <- as_querychat_data_source(test_df, table_name = "test_table")
- withr::defer(cleanup_source(df_source))
+ df_source <- DataFrameSource$new(test_df, "test_table")
+ withr::defer(df_source$cleanup())
# Test with multiline comments
multiline_comment_query <- "
@@ -65,7 +65,7 @@ test_that("execute_query handles SQL with multiline comments", {
WHERE value > 25
"
- result <- execute_query(df_source, multiline_comment_query)
+ result <- df_source$execute_query(multiline_comment_query)
expect_s3_class(result, "data.frame")
expect_equal(nrow(result), 3)
expect_equal(ncol(result), 2)
@@ -80,7 +80,7 @@ test_that("execute_query handles SQL with multiline comments", {
WHERE value /* another comment */ > 25
"
- result <- execute_query(df_source, embedded_multiline_query)
+ result <- df_source$execute_query(embedded_multiline_query)
expect_s3_class(result, "data.frame")
expect_equal(nrow(result), 3)
expect_equal(ncol(result), 2)
@@ -95,8 +95,8 @@ test_that("execute_query handles SQL with trailing semicolons", {
)
# Create data source
- df_source <- as_querychat_data_source(test_df, table_name = "test_table")
- withr::defer(cleanup_source(df_source))
+ df_source <- DataFrameSource$new(test_df, "test_table")
+ withr::defer(df_source$cleanup())
# Test with trailing semicolon
query_with_semicolon <- "
@@ -105,7 +105,7 @@ test_that("execute_query handles SQL with trailing semicolons", {
WHERE value > 25;
"
- result <- execute_query(df_source, query_with_semicolon)
+ result <- df_source$execute_query(query_with_semicolon)
expect_s3_class(result, "data.frame")
expect_equal(nrow(result), 3)
expect_equal(ncol(result), 2)
@@ -117,7 +117,7 @@ test_that("execute_query handles SQL with trailing semicolons", {
WHERE value > 25;;;;
"
- result <- execute_query(df_source, query_with_multiple_semicolons)
+ result <- df_source$execute_query(query_with_multiple_semicolons)
expect_s3_class(result, "data.frame")
expect_equal(nrow(result), 3)
expect_equal(ncol(result), 2)
@@ -132,8 +132,8 @@ test_that("execute_query handles SQL with mixed comments and semicolons", {
)
# Create data source
- df_source <- as_querychat_data_source(test_df, table_name = "test_table")
- withr::defer(cleanup_source(df_source))
+ df_source <- DataFrameSource$new(test_df, "test_table")
+ withr::defer(df_source$cleanup())
# Test with a mix of comment styles and semicolons
complex_query <- "
@@ -150,7 +150,7 @@ test_that("execute_query handles SQL with mixed comments and semicolons", {
value > 25; -- End of query
"
- result <- execute_query(df_source, complex_query)
+ result <- df_source$execute_query(complex_query)
expect_s3_class(result, "data.frame")
expect_equal(nrow(result), 3)
expect_equal(ncol(result), 2)
@@ -165,7 +165,7 @@ test_that("execute_query handles SQL with mixed comments and semicolons", {
WHERE value > 25 -- WHERE id = 'value; DROP TABLE test;'
"
- result <- execute_query(df_source, tricky_comment_query)
+ result <- df_source$execute_query(tricky_comment_query)
expect_s3_class(result, "data.frame")
expect_equal(nrow(result), 3)
expect_equal(ncol(result), 2)
@@ -180,8 +180,8 @@ test_that("execute_query handles SQL with unusual whitespace patterns", {
)
# Create data source
- df_source <- as_querychat_data_source(test_df, table_name = "test_table")
- withr::defer(cleanup_source(df_source))
+ df_source <- DataFrameSource$new(test_df, "test_table")
+ withr::defer(df_source$cleanup())
# Test with unusual whitespace patterns (which LLMs might generate)
unusual_whitespace_query <- "
@@ -194,7 +194,7 @@ test_that("execute_query handles SQL with unusual whitespace patterns", {
"
- result <- execute_query(df_source, unusual_whitespace_query)
+ result <- df_source$execute_query(unusual_whitespace_query)
expect_s3_class(result, "data.frame")
expect_equal(nrow(result), 3)
expect_equal(ncol(result), 2)
diff --git a/pkg-r/tests/testthat/test-test-query.R b/pkg-r/tests/testthat/test-test-query.R
index afdb0f54..b08d2766 100644
--- a/pkg-r/tests/testthat/test-test-query.R
+++ b/pkg-r/tests/testthat/test-test-query.R
@@ -3,7 +3,7 @@ library(DBI)
library(RSQLite)
library(querychat)
-test_that("test_query.dbi_source correctly retrieves one row of data", {
+test_that("test_query correctly retrieves one row of data", {
# Create a simple data frame
test_df <- data.frame(
id = 1:5,
@@ -18,24 +18,23 @@ test_that("test_query.dbi_source correctly retrieves one row of data", {
dbWriteTable(conn, "test_table", test_df, overwrite = TRUE)
- dbi_source <- as_querychat_data_source(conn, "test_table")
- withr::defer(cleanup_source(dbi_source))
+ dbi_source <- DBISource$new(conn, "test_table")
+ withr::defer(dbi_source$cleanup())
# Test basic query - should only return one row
- result <- test_query(dbi_source, "SELECT * FROM test_table")
+ result <- dbi_source$test_query("SELECT * FROM test_table")
expect_s3_class(result, "data.frame")
expect_equal(nrow(result), 1) # Should only return 1 row
expect_equal(result$id, 1) # Should be first row
# Test with WHERE clause
- result <- test_query(dbi_source, "SELECT * FROM test_table WHERE value > 25")
+ result <- dbi_source$test_query("SELECT * FROM test_table WHERE value > 25")
expect_s3_class(result, "data.frame")
expect_equal(nrow(result), 1) # Should only return 1 row
expect_equal(result$value, 30) # Should return first row with value > 25
# Test with ORDER BY - should get the highest value
- result <- test_query(
- dbi_source,
+ result <- dbi_source$test_query(
"SELECT * FROM test_table ORDER BY value DESC"
)
expect_s3_class(result, "data.frame")
@@ -43,12 +42,12 @@ test_that("test_query.dbi_source correctly retrieves one row of data", {
expect_equal(result$value, 50) # Should be the highest value
# Test with query returning no results
- result <- test_query(dbi_source, "SELECT * FROM test_table WHERE value > 100")
+ result <- dbi_source$test_query("SELECT * FROM test_table WHERE value > 100")
expect_s3_class(result, "data.frame")
expect_equal(nrow(result), 0) # Should return empty data frame
})
-test_that("test_query.dbi_source handles errors correctly", {
+test_that("test_query handles errors correctly", {
# Setup DBI source
temp_db <- withr::local_tempfile(fileext = ".db")
conn <- dbConnect(RSQLite::SQLite(), temp_db)
@@ -62,23 +61,21 @@ test_that("test_query.dbi_source handles errors correctly", {
)
dbWriteTable(conn, "test_table", test_df, overwrite = TRUE)
- dbi_source <- as_querychat_data_source(conn, "test_table")
- withr::defer(cleanup_source(dbi_source), priority = "last")
+ dbi_source <- DBISource$new(conn, "test_table")
# Test with invalid SQL
- expect_error(test_query(dbi_source, "SELECT * WRONG SYNTAX"))
+ expect_error(dbi_source$test_query("SELECT * WRONG SYNTAX"))
# Test with non-existent table
- expect_error(test_query(dbi_source, "SELECT * FROM non_existent_table"))
+ expect_error(dbi_source$test_query("SELECT * FROM non_existent_table"))
# Test with non-existent column
- expect_error(test_query(
- dbi_source,
+ expect_error(dbi_source$test_query(
"SELECT non_existent_column FROM test_table"
))
})
-test_that("test_query.dbi_source works with different data types", {
+test_that("test_query works with different data types", {
# Create a data frame with different data types
test_df <- data.frame(
id = 1:3,
@@ -96,11 +93,10 @@ test_that("test_query.dbi_source works with different data types", {
dbWriteTable(conn, "types_table", test_df, overwrite = TRUE)
- dbi_source <- as_querychat_data_source(conn, "types_table")
- withr::defer(cleanup_source(dbi_source), priority = "last")
+ dbi_source <- DBISource$new(conn, "types_table")
# Test query with different column types
- result <- test_query(dbi_source, "SELECT * FROM types_table")
+ result <- dbi_source$test_query("SELECT * FROM types_table")
expect_s3_class(result, "data.frame")
expect_equal(nrow(result), 1)
expect_type(result$text_col, "character")