ignore list-columns from head, can burn tokens

boettiger-lab · Feb 4, 2025 · 8986359 · 8986359
1 parent 64e52f1
commit 8986359
Show file tree

Hide file tree

Showing 4 changed files with 32 additions and 21 deletions.
diff --git a/DESCRIPTION b/DESCRIPTION
@@ -18,7 +18,8 @@ Imports:
     knitr,
     duckdbfs,
     dplyr,
-    utils
+    utils,
+    purrr
 Suggests: 
     testthat (>= 3.0.0)
 Config/testthat/edition: 3
diff --git a/R/agent.R b/R/agent.R
@@ -44,7 +44,8 @@ create_prompt <- function(con = duckdbfs::cached_connection(),
 Pay attention to the schema of the table <x>:
 <schema>
 
-Also pay close attention to how data is represented in each column, as seen by the HEAD of table <x>:
+Also pay close attention to how data is represented in each column,
+as seen by the HEAD (ommitting large list-type columns, if any) of table <x>:
 <head>
 ", .open = "<", .close = ">")
 
@@ -62,11 +63,20 @@ Also pay close attention to how data is represented in each column, as seen by t
 
 # render table info as markdown tables:
 tbl_head_md <- function(table_name, con) {
-  dplyr::tbl(con, table_name) |>
-    utils::head() |>
-    dplyr::collect() |>
-    knitr::kable() |>
-    paste(collapse = "\n")
+    x <- dplyr::tbl(con, table_name)
+
+    # drop non-list types.
+    # backend doesn't support tidyselect predicates
+    types <- dplyr::collect(utils::head(x,1)) |>
+      purrr::map_lgl(function(x) class(x)[[1]] != "list")
+    keep <- names(types)[types]
+
+    x |>
+      dplyr::select(dplyr::all_of(keep)) |>
+      utils::head() |>
+      dplyr::collect() |>
+      knitr::kable() |>
+      paste(collapse = "\n")
 }
 
 tbl_schema_md <- function(table_name, con) {

diff --git a/README.Rmd b/README.Rmd
@@ -45,15 +45,15 @@ system_prompt = create_prompt()
 
 
 ```{r }
-tracts_agent <- ellmer::chat_vllm(
+agent <- ellmer::chat_vllm(
     base_url = "https://llm.cirrus.carlboettiger.info/v1/",
     model = "kosbu/Llama-3.3-70B-Instruct-AWQ",
     api_key = Sys.getenv("CIRRUS_LLM_KEY"),
     system_prompt = system_prompt,
     api_args = list(temperature = 0)
 )
 
-resp <- tracts_agent$chat("Yolo County")
+resp <- agent$chat("Yolo County")
 agent_query(resp)
 ```
 

diff --git a/README.md b/README.md
@@ -34,15 +34,15 @@ system_prompt = create_prompt()
 ```
 
 ``` r
-tracts_agent <- ellmer::chat_vllm(
+agent <- ellmer::chat_vllm(
     base_url = "https://llm.cirrus.carlboettiger.info/v1/",
     model = "kosbu/Llama-3.3-70B-Instruct-AWQ",
     api_key = Sys.getenv("CIRRUS_LLM_KEY"),
     system_prompt = system_prompt,
     api_args = list(temperature = 0)
 )
 
-resp <- tracts_agent$chat("Yolo County")
+resp <- agent$chat("Yolo County")
 #> {
 #> "query": "CREATE OR REPLACE VIEW yolo_county AS SELECT * FROM censustracts 
 #> WHERE COUNTY = 'Yolo County'",
@@ -55,15 +55,15 @@ agent_query(resp)
 #> [90m# Database: DuckDB v1.1.3 [unknown@Linux 6.9.3-76060903-generic:R 4.4.2/:memory:][39m
 #>    STATE      COUNTY      FIPS        h6             
 #>    [3m[90m<chr>[39m[23m      [3m[90m<chr>[39m[23m       [3m[90m<chr>[39m[23m       [3m[90m<chr>[39m[23m          
-#> [90m 1[39m California Yolo County 06113010102 862832b0fffffff
-#> [90m 2[39m California Yolo County 06113010313 862832b07ffffff
-#> [90m 3[39m California Yolo County 06113010401 8628304afffffff
-#> [90m 4[39m California Yolo County 06113010401 8628304c7ffffff
-#> [90m 5[39m California Yolo County 06113010401 8628304d7ffffff
-#> [90m 6[39m California Yolo County 06113010401 8628304dfffffff
-#> [90m 7[39m California Yolo County 06113010401 8628304f7ffffff
-#> [90m 8[39m California Yolo County 06113010401 862830417ffffff
-#> [90m 9[39m California Yolo County 06113010401 86283041fffffff
-#> [90m10[39m California Yolo County 06113010401 862832b2fffffff
+#> [90m 1[39m California Yolo County 06113010102 862832B0FFFFFFF
+#> [90m 2[39m California Yolo County 06113010313 862832B07FFFFFF
+#> [90m 3[39m California Yolo County 06113010401 8628304AFFFFFFF
+#> [90m 4[39m California Yolo County 06113010401 8628304C7FFFFFF
+#> [90m 5[39m California Yolo County 06113010401 8628304D7FFFFFF
+#> [90m 6[39m California Yolo County 06113010401 8628304DFFFFFFF
+#> [90m 7[39m California Yolo County 06113010401 8628304F7FFFFFF
+#> [90m 8[39m California Yolo County 06113010401 862830417FFFFFF
+#> [90m 9[39m California Yolo County 06113010401 86283041FFFFFFF
+#> [90m10[39m California Yolo County 06113010401 862832B2FFFFFFF
 #> [90m# ℹ more rows[39m
 ```