HealthLLM

HealthLLM.PUBLIC_MODULES

The submodules whose exported names make up the public API of HealthLLM. Everything a submodule exports is re-exported from the package, so each symbol is declared in exactly one place — the export list of the module that defines it. Adding a name to that list is all it takes to publish it; there is no second mirrored list here to fall out of sync.

HuggingFace

HuggingFace backend for PromptingTools.

PromptingTools ships schemas for a long list of OpenAI-compatible providers (Groq, Together, Fireworks, DeepSeek, Mistral, ...) but — as of v0.94, the latest release — none for HuggingFace. This module supplies the missing one, built the same way PromptingTools builds its own provider schemas: a marker type under AbstractOpenAISchema plus create_chat/create_embeddings methods that point the shared OpenAI transport at HuggingFace's router.

Because HuggingFaceOpenAISchema is an AbstractOpenAISchema, everything PromptingTools already does for OpenAI — aigenerate, aiembed, message rendering, streaming, retries — works against HuggingFace unchanged.

Credentials

The token is read from huggingface_api_key: an explicitly passed api_key wins, then a key set with set_huggingface_api_key!, then the first of HF_API_TOKEN, HF_TOKEN, HUGGINGFACE_API_KEY, or HUGGING_FACE_HUB_TOKEN found in the environment.

Example

using HealthLLM, PromptingTools

set_huggingface_api_key!(ENV["HF_TOKEN"])

msg = PromptingTools.aigenerate(HuggingFaceOpenAISchema(), "Say hi";
                                model = "meta-llama/Llama-3.1-8B-Instruct")

emb = PromptingTools.aiembed(HuggingFaceOpenAISchema(), ["some text"];
                             model = "BAAI/bge-m3")

# or point embeddings at a TEI / Inference Endpoint that speaks OpenAI
emb = PromptingTools.aiembed(HuggingFaceOpenAISchema(), ["some text"];
                             model = "BAAI/bge-m3",
                             api_kwargs = (; url = "https://my-endpoint.hf.space/v1"))

Endpoints

Chat and embeddings use different HuggingFace surfaces, because the router's OpenAI-compatible API covers chat only:

HUGGINGFACE_EMBED_TIMEOUT

Read timeout in seconds (300) used for feature-extraction when the caller did not choose one. A HuggingFace model that is not already warm loads while holding the connection open — measured at ~55s for BAAI/bge-m3 — so PromptingTools' 120s aiembed default times out on the first call to a cold large model.

HUGGINGFACE_INFERENCE_URL

Base URL for HuggingFace pipeline inference (https://router.huggingface.co/hf-inference/models). Embeddings go here rather than through HUGGINGFACE_ROUTER_URL: the router's OpenAI-compatible surface covers chat completions only and answers /v1/embeddings with a 404.

HUGGINGFACE_ROUTER_URL

Base URL of the HuggingFace Inference Providers router (https://router.huggingface.co/v1), the OpenAI-compatible entry point used when no url is supplied.

HuggingFaceOpenAISchema

Schema for HuggingFace models served over an OpenAI-compatible API — the HuggingFace Inference Providers router by default, or any Text Embeddings Inference / Inference Endpoint deployment you point it at.

A subtype of PromptingTools.AbstractOpenAISchema, so it plugs into aigenerate and aiembed exactly like the built-in provider schemas. Override the endpoint per call with api_kwargs = (; url = "...").

Example

PromptingTools.aigenerate(HuggingFaceOpenAISchema(), "Hello";
                          model = "meta-llama/Llama-3.1-8B-Instruct")

hf_model_id(model) -> String

Strip the "hf:" prefix this package uses to mark HuggingFace references (see embedding_ref), leaving the bare repo id the API expects.

Example

hf_model_id("hf:BAAI/bge-m3")   # "BAAI/bge-m3"
hf_model_id("BAAI/bge-m3")      # "BAAI/bge-m3"

huggingface_api_key() -> String

Return the HuggingFace token: the one set by set_huggingface_api_key! if any, else the first of HF_API_TOKEN, HF_TOKEN, HUGGINGFACE_API_KEY, HUGGING_FACE_HUB_TOKEN present in the environment, else "".

An empty result is not an error here — some endpoints are open — but the request will fail with a 401 if the model requires authentication.

huggingface_providers(model; api_key=huggingface_api_key()) -> Vector{String}

Names of the inference providers currently serving model, i.e. those whose mapping status is "live". Returns an empty vector when nothing serves it.

The router only auto-routes to providers enabled on your account, so a model can be live somewhere and still be refused. Use this to see the options, then pin one by appending it to the model name.

Example

huggingface_providers("Qwen/Qwen2.5-7B-Instruct")   # ["featherless-ai"]
# then: model = "hf:Qwen/Qwen2.5-7B-Instruct:featherless-ai"

set_huggingface_api_key!(key) -> String

Set the HuggingFace token for the current session, taking precedence over the environment. Pass "" to clear it and fall back to the environment again.

split_provider(model) -> (repo, provider)

Split a "org/repo:provider" reference into its parts, returning an empty provider when the model is not pinned. HuggingFace repo ids never contain :, so the last colon unambiguously marks a provider pin.

Example

split_provider("Qwen/Qwen2.5-7B-Instruct:featherless-ai")
# ("Qwen/Qwen2.5-7B-Instruct", "featherless-ai")

OpenAI.create_embeddings(schema::HuggingFaceOpenAISchema, api_key, docs, model; url="", kwargs...)

Embed docs with a HuggingFace model.

With no url, this calls the feature-extraction pipeline under HUGGINGFACE_INFERENCE_URL and reshapes the reply into the OpenAI embeddings response aiembed expects, so the HuggingFace backend behaves like any other from the caller's side. Token-level output (a matrix per input, from models that do not pool internally) is mean-pooled into one vector per input.

Pass api_kwargs = (; url = "https://.../v1") to target a deployment that does serve an OpenAI-compatible /embeddings route — a Text Embeddings Inference container or a dedicated Inference Endpoint — and the request is forwarded there unchanged instead.

Utils

Cross-cutting helpers shared by the rest of HealthLLM: model/schema resolution, HuggingFace model loading, provenance rendering, optional-dependency lookup, and small file utilities.

HuggingFaceLoadResult

Result of a HuggingFace model load operation.

Fields

• path::Union{String,Nothing}: Local path to the downloaded model snapshot, or nothing if not downloaded.
• info::Any: Additional metadata from HuggingFaceHub (e.g. model info), or nothing.
• downloaded::Bool: Whether the model was successfully downloaded.

build_index_rag(cfg, files; embedder_kwargs=NamedTuple())

Build a RAG index from source files using RAGTools.

Arguments

• cfg: A RAGTools indexer configuration (e.g. RAGTools.SimpleIndexer()).
• files: Vector of file paths or text sources to index.

Keywords

• embedder_kwargs=NamedTuple(): Additional keyword arguments passed to the embedder (e.g. (model="nomic-embed-text",)).

Returns

A RAG index object suitable for querying with generate_funsql_query.

Example

index = build_index_rag(RAGTools.SimpleIndexer(), ["doc1.md", "doc2.md"])

check_dims(embeddings, chunks=nothing, expected_dim=nothing) -> (dim, n)

Shape check for a dim × n embedding matrix whose columns are chunks — the convention used throughout embedding, storage, and database code. Verifies the row count against expected_dim and the column count against length(chunks) when either is supplied, then returns size(embeddings).

Throws DimensionMismatch on the first mismatch. Every layer that accepts an embedding matrix alongside its texts routes its shape check through here, so the same misuse reports the same way everywhere.

check_k(k) -> Int

Validate a nearest-neighbour count, throwing ArgumentError unless k > 0.

collect_files_with_extensions(directory, extensions)

Recursively collect file paths under directory whose extension is in extensions. Extensions are matched case-insensitively (e.g. ".jl").

Arguments

• directory::AbstractString: Root directory to search recursively.
• extensions::AbstractVector{<:AbstractString}: File extensions to match (e.g. [".jl", ".md"]).

Returns

• Vector{String}: List of full file paths matching the given extensions.

Example

files = collect_files_with_extensions("src", [".jl", ".md"])

get_schema(provider::Symbol)
get_schema(schema_name=nothing, model=nothing)

Return a PromptingTools schema instance.

The Symbol method is the explicit form: :ollama and :huggingface map directly onto their schemas. The two-argument method is the heuristic form used when only user-supplied strings are available — it tries <Name>Schema from schema_name, then sniffs model for a HuggingFace marker ("hf:", ...), and finally falls back to PromptingTools.OllamaSchema().

HuggingFace resolves to HuggingFaceOpenAISchema, which this package defines: PromptingTools has no HuggingFace schema of its own.

Example

get_schema(:ollama)                          # -> OllamaSchema()
get_schema(:huggingface)                     # -> HuggingFaceOpenAISchema()
get_schema("HuggingFace")                    # -> HuggingFaceOpenAISchema()
get_schema(nothing, "hf:facebook/opt-350m")  # -> HuggingFaceOpenAISchema()

load_huggingface_model(model; token=nothing)

Download and load a HuggingFace model by name using HuggingFaceHub.jl.

The download entry point is probed at run time (snapshot, then repo_download, then info/file_download) because HuggingFaceHub.jl has moved it across versions. Any failure — no network, gated repo, unknown model — is reported as a warning and returned as a non-downloaded result rather than thrown.

Arguments

• model::String: HuggingFace model identifier (e.g. "gpt2", "facebook/opt-350m").

Keywords

• token::Union{Nothing,String}: HuggingFace API token for private/gated models.

Returns

• HuggingFaceLoadResult: Contains path, info, and downloaded fields.

Example

res = load_huggingface_model("gpt2")
if res.downloaded
    println("Model downloaded to $(res.path)")
end

register_models(model_name, model_embedding; schema_name=nothing)

Register chat and embedding models in PromptingTools and set them as active defaults.

Arguments

• model_name::String: Name of the chat/generation model (e.g. "llama3.2", "hf:facebook/opt-350m").
• model_embedding::String: Name of the embedding model (e.g. "nomic-embed-text", "hf:all-MiniLM-L6-v2").

Keywords

• schema_name::Union{Nothing,String}: Optional explicit schema name (e.g. "Ollama", "HuggingFace"). If nothing, the schema is inferred from model names.

Returns

nothing.

Example

register_models("llama3.2", "nomic-embed-text")
register_models("hf:facebook/opt-350m", "hf:all-MiniLM-L6-v2"; schema_name="HuggingFace")

render_provenance(md::AbstractDict; separator=" › ", maxlen=512) -> String

Render grounding metadata as a single provenance string, "<url-or-source> › <heading-or-group>", truncated to maxlen characters.

Either half may be missing: with no parent heading only the base is returned, and with no base only the parent. Both the retrieval-index provenance recorded by chunk_provenance and the per-chunk tag rendered into a prompt by format_context go through here, so the two cannot drift apart.

Example

render_provenance(Dict{Symbol,Any}(:url => "http://cdm", :heading => "person"))
# "http://cdm › person"

require_main_module(name::Symbol, hint::AbstractString) -> Module

Return the module name if the driver session has loaded it into Main, otherwise throw an ArgumentError carrying hint.

Optional backends (FAISS, FunSQL) are deliberately not hard dependencies of HealthLLM: they are heavy, and only some workflows need them. This is the single place that resolves such a module, so every optional backend fails the same way with an actionable message instead of a MethodError deep inside a call.

Example

FunSQL = require_main_module(:FunSQL, "Install FunSQL.jl and run `using FunSQL`.")

write_combined_file(files, output_file)

Concatenate the provided text files into a single output file and include a # File: <path> header before each file body.

Arguments

• files::AbstractVector{<:AbstractString}: List of file paths to concatenate.
• output_file::AbstractString: Path to the output file.

Returns

• String: The path to the output file.

Example

files = collect_files_with_extensions("data", [".jl", ".md"])
combined = write_combined_file(files, "combined.txt")

search_embeddings_pgvector(conn, query, k; table="embeddings", metric=:cosine)

Return the k nearest stored chunks to query from a pgvector table, ordered most-similar first. metric selects the distance operator: :cosine (<=>), :dot (<#>), or :l2 (<->). Each result is a (; id, chunk, distance) named tuple.

Example

hits = search_embeddings_pgvector(conn, rand(384), 5)
for h in hits
    println(h.distance, "  ", first(h.chunk, 60))
end

store_embeddings_pgvector(conn, embeddings, chunks, embedding_dimension)

Create an embeddings table in PostgreSQL with pgvector extension and insert all chunk embeddings in a transaction.

Arguments

• conn::LibPQ.Connection: A connection to the PostgreSQL database.
• embeddings::AbstractMatrix: Matrix where rows are embedding dimensions and columns are chunks.
• chunks::AbstractVector: Vector of text chunks to store with their embeddings.
• embedding_dimension::Int: Dimension of the embedding vectors.

Returns

nothing.

Example

conn = LibPQ.Connection("postgresql://user:pass@localhost/db")
embeddings = rand(384, 10)
chunks = ["chunk $i" for i in 1:10]
store_embeddings_pgvector(conn, embeddings, chunks, 384)

validate_embeddings_inputs(embeddings, chunks, embedding_dimension)

Validate shape consistency between an embedding matrix and associated text chunks.

Arguments

• embeddings::AbstractMatrix: Matrix where rows are embedding dimensions and columns are chunks.
• chunks::AbstractVector: Vector of text chunks corresponding to embedding columns.
• embedding_dimension::Integer: Expected embedding dimension.

Returns

nothing if validation passes.

Throws

• DimensionMismatch: If embedding dimensions or chunk count don't match.

Example

embeddings = rand(384, 10)
chunks = ["chunk $i" for i in 1:10]
validate_embeddings_inputs(embeddings, chunks, 384)  # passes

Prompt

Prompt construction for retrieval-augmented FunSQL generation. This module sits between retrieval and the language model: it takes the chunks a vector store returned for a question — OMOP schema definitions, FunSQL examples, prose docs — and the user's natural-language analytical question, and renders one grounded prompt asking the model to translate the question into a FunSQL query.

The two moving parts injected into every prompt are:

1. Retrieved context — the ranked chunks from retrieve/search (or raw strings / Chunks), formatted as numbered, optionally provenance-tagged blocks so the model can ground its answer and cite sources.
2. The analytical question — the user's natural-language request, verbatim.

Interface

• PromptTemplate — the reusable shape (system message, section headers, per-chunk formatting, and context budgets). DEFAULT_FUNSQL_TEMPLATE is the OMOP/FunSQL default.
• format_context(template, hits) — render retrieved chunks to a string.
• build_prompt(question, hits; template) — the entry point; returns (; system, user, prompt).

Example

store = LocalVectorStore(embedding_dimension())
add!(store, embed(chunks), chunks)

hits = retrieve(store, "How many distinct patients had a diabetes diagnosis in 2020?", 6)
p = build_prompt("How many distinct patients had a diabetes diagnosis in 2020?", hits)

# split form, for a proper system/user conversation:
msg = PromptingTools.aigenerate(p.system * "\n\n" * p.user; model="llama3.2")
# or the single joined string, for the {input_query}-style path:
p.prompt

DEFAULT_FUNSQL_TEMPLATE

The default PromptTemplate: FunSQL/OMOP system message, provenance on, scores off, up to 8 chunks or 8000 characters of context.

FUNSQL_SYSTEM_PROMPT

Default system message for FunSQL generation over the OMOP CDM. It fixes the model's role and the rules that keep generations grounded: build queries only from tables/columns that appear in the retrieved context, follow the retrieved examples' FunSQL style, flag missing schema instead of inventing it, and return the query as a single fenced julia code block.

The Grounding rules are the load-bearing part: they instruct the model to treat the retrieved context as the only admissible source of schema (table/column) names and FunSQL syntax, and to refuse rather than hallucinate when the context is insufficient. Weaken them and generations drift toward plausible-but-nonexistent OMOP columns.

PromptTemplate(; kwargs...)

Reusable shape for a retrieval-augmented prompt. A template is data — the system message, the section headers, how each retrieved chunk is rendered, and how much context to admit — so the same construction logic serves different models and tasks by swapping fields.

Fields

• system::String: System message. Defaults to FUNSQL_SYSTEM_PROMPT.
• context_header::String: Heading printed above the retrieved chunks.
• question_header::String: Heading printed above the user's question.
• answer_cue::String: Trailing line that cues the model to answer (e.g. a # FunSQL query header). Empty to omit.
• empty_context_note::String: Placeholder used when no chunks are supplied, so the model is told the context is empty rather than seeing a blank section.
• chunk_label::String: Per-chunk label prefix; numbered as "[<label> i]".
• include_provenance::Bool: Append each chunk's provenance (url/source › heading) when available.
• include_scores::Bool: Append each chunk's retrieval score when available.
• max_chunks::Int: Keep at most this many chunks (highest-ranked first).
• max_context_chars::Int: Stop admitting chunks once the rendered context would exceed this many characters. Whole chunks only — a chunk is admitted or skipped, never cut mid-text — except a single chunk longer than the budget on its own, which is truncated with an ellipsis.
• separator::String: Placed between rendered chunks.

build_prompt(question, hits; template=DEFAULT_FUNSQL_TEMPLATE)
    -> (; system, user, prompt)

Construct a retrieval-augmented prompt from the user's natural-language analytical question and the retrieved hits, using template.

Returns a NamedTuple with three fields, so both prompting styles are covered:

• system::String — the template's system message, for a proper system/user conversation (e.g. PromptingTools.aigenerate with a system role).
• user::String — the user turn: the formatted context block, the question, and the answer cue.
• prompt::String — system and user joined, for single-string callers.

hits accepts plain strings, search/retrieve hit NamedTuples, or Chunks (see format_context). Throws ArgumentError on an empty question.

Example

hits = retrieve(store, "Average age of patients with hypertension?", 6)
p = build_prompt("Average age of patients with hypertension?", hits)
println(p.prompt)

build_prompt(question; template=DEFAULT_FUNSQL_TEMPLATE) -> (; system, user, prompt)

Convenience method for a question with no retrieved context — renders the template's empty_context_note in the context slot.

format_context(template::PromptTemplate, hits) -> String

Render the retrieved hits into a numbered context block per template. hits is any iterable of retrieval results — plain strings, search/retrieve hit NamedTuples, or Chunks. Blank chunks are dropped; the first template.max_chunks are kept, and chunks are admitted only while the running size stays within template.max_context_chars (a single over-budget chunk is truncated). Returns template.empty_context_note when nothing survives.

Execution

Turn a grounded prompt into a FunSQL query and check that the query the model produced is actually usable. This closes the loop opened by the Prompt module:

build_prompt ─▶ generate_funsql ─▶ FunSQLGeneration ─▶ sanity_check_funsql ─▶ FunSQLCheck
                (call the model)    (extracted code)     (parse / build / render / run)

Two responsibilities:

1. Generation — generate_funsql sends the grounded prompt to a chat model and extracts the FunSQL code block from the reply (extract_funsql).
2. Execution-based sanity check — sanity_check_funsql confirms the generated code parses as Julia, builds into a FunSQL.SQLNode, and renders to SQL — optionally against a schema catalog, which is what catches a query that references a table or column the OMOP CDM does not have. With a connection it can also run the SQL.

Optional dependency

Building and rendering FunSQL needs the FunSQL package loaded in Main (it is a test/driver-side dependency, not a hard dependency of HealthLLM, mirroring how FaissVectorStore treats Faiss). The parse stage works without it; the later stages raise a clear error until using FunSQL has run.

FunSQLCheck

Outcome of sanity_check_funsql. Records how far the generated query got through the parse → build → render → run pipeline and why it stopped.

Fields

• ok::Bool: Reached the deepest requested stage with no error.
• stage::Symbol: Deepest stage reached — :parsed, :built, :rendered, or :executed. On failure this is the last stage that succeeded.
• parsed::Bool: The code parsed as Julia.
• built::Bool: Evaluating it produced a FunSQL.SQLNode.
• rendered::Bool: The node rendered to SQL (against catalog, when supplied).
• executed::Bool: The SQL ran via executor (only when one was given).
• sql::Union{String,Nothing}: The rendered SQL, when rendering succeeded.
• error::Union{String,Nothing}: Message from the failing stage, or nothing.

FunSQLGeneration

Result of generate_funsql.

Fields

• funsql::String: The FunSQL code extracted from the model's reply — the code block only, ready to hand to sanity_check_funsql.
• answer::String: The model's full reply (code block plus any explanation).

extract_funsql(text::AbstractString) -> String

Pull the FunSQL code out of a model reply. Prefers the first fenced block tagged $julia ` / `$jl, falls back to the first fenced block of any language, and finally to the whole stripped text when no fence is present. The returned code is stripped of surrounding whitespace.

generate_funsql(prompt; model=PromptingTools.MODEL_CHAT, schema=nothing,
                generator=<aigenerate>, kwargs...) -> FunSQLGeneration

Generate a FunSQL query from a grounded prompt and extract the code block from the reply. prompt is typically the NamedTuple returned by build_prompt — its system/user split is sent as a proper chat conversation so the grounding rules land in the system role — but a plain prompt string works too.

model is the chat model name (defaults to the active PromptingTools.MODEL_CHAT). schema is an optional PromptingTools schema; when given it is passed to aigenerate as the first argument. Extra kwargs are forwarded to the generator.

generator is injectable for testing: it is called as generator(prompt; model, schema, kwargs...) and must return the reply text. The default calls PromptingTools.aigenerate.

Example

hits = retrieve(store, question, 6)
p    = build_prompt(question, hits)
gen  = generate_funsql(p; model = "llama3.2")
gen.funsql     # the extracted FunSQL code

sanity_check_funsql(code::AbstractString; catalog=nothing, dialect=:duckdb,
                    mod=nothing, executor=nothing) -> FunSQLCheck
sanity_check_funsql(gen::FunSQLGeneration; kwargs...) -> FunSQLCheck

Check that generated FunSQL code is usable, without trusting the model's word for it. Runs as far through the pipeline as the inputs allow and reports the result as a FunSQLCheck:

1. parse — code parses as Julia. Needs nothing beyond Julia itself.
2. build — evaluating it yields a FunSQL.SQLNode. Requires FunSQL loaded in Main (see the module note). By default the code is evaluated in a built-in sandbox module that already binds the FunSQL node constructors (From, Get, Agg, ...), so bare-name generations resolve without importing anything into Main. Pass mod to evaluate in your own module instead.
3. render — the node serialises to SQL. With catalog (a FunSQL.SQLCatalog) the query is resolved against that schema, so a reference to a table or column the schema does not contain fails here — this is the check that catches a hallucinated OMOP name that slipped past the prompt. Without a catalog the render is structural, using dialect (a dialect name like :duckdb or a FunSQL.SQLDialect).
4. run — when executor is supplied, the rendered SQL string is passed to it (e.g. sql -> DuckDB.execute(conn, sql)), confirming the database accepts it.

The check never throws for a bad query: any stage failure is captured in the returned error with ok = false. It only throws for a genuine misuse (build/ render requested without FunSQL available).

Examples

using FunSQL
catalog = FunSQL.SQLCatalog(
    FunSQL.SQLTable(:person, columns = [:person_id, :year_of_birth]),
    dialect = :duckdb,
)

sanity_check_funsql("From(:person) |> Group() |> Select(:n => Agg.count())";
                    catalog = catalog)
# FunSQLCheck(ok @ :rendered)

sanity_check_funsql("From(:person) |> Select(Get.made_up_column)"; catalog = catalog)
# FunSQLCheck(failed @ :built)   -> error names the unresolved column

Query

End-to-end retrieval-augmented querying via RAGTools.airag: retrieval, generation, and the prompt wrapper that carries the FunSQL grounding rules into the question.

Unlike the Prompt/Execution modules — which format retrieved chunks themselves and call the model directly — this module hands the whole loop to RAGTools.airag and only decorates the question. Both paths share FUNSQL_SYSTEM_PROMPT as their single source of grounding rules.

DEFAULT_QUERY_TEMPLATE

Default template for generate_funsql_query: the shared FUNSQL_SYSTEM_PROMPT grounding rules followed by the question slot. Using the same system prompt as build_prompt keeps the two generation paths from drifting apart in what they tell the model.

QUERY_PLACEHOLDER

The substring a query template substitutes the user's question into: "{input_query}".

generate_funsql_query(index, model_embedding, model_name, question;
                      template=DEFAULT_QUERY_TEMPLATE, schema=nothing, kwargs...)
generate_funsql_query(index, model_embedding, model_name, template, question; kwargs...)

Run retrieval-augmented generation against index and return the generated answer.

The question is substituted into template at QUERY_PLACEHOLDER ("{input_query}") and the result is handed to RAGTools.airag, which does its own retrieval and generation.

Arguments

• index: The RAG index to query against.
• model_embedding::AbstractString: Embedding model used to retrieve relevant chunks.
• model_name::AbstractString: Chat model used to generate the answer.
• question::AbstractString: The user's question.

Keywords

• template::AbstractString=DEFAULT_QUERY_TEMPLATE: Wrapper containing "{input_query}". Defaults to the shared FunSQL grounding rules.
• schema=nothing: PromptingTools schema for both the generator and the embedder. When nothing (the default) a schema is inferred from each model name.
• kwargs...: Forwarded to RAGTools.airag.

Returns

The generated answer from the RAG query.

Example

index  = build_index_rag(RAGTools.SimpleIndexer(), files)
answer = generate_funsql_query(index, "nomic-embed-text", "llama3.2",
                               "How many distinct patients had a diabetes diagnosis?")

Ingestion

Pull documentation into a RAG index from two kinds of sources:

1. Curated docs — a fixed, reviewed list of high-quality pages (JuliaHealth, OMOP CDM, OHDSI, FunSQL.jl). See CURATED_SOURCES.
2. Web search — a live query against a pluggable search backend for anything not in the curated set. See AbstractSearchProvider.

Data flow

CURATED_SOURCES ─┐
                 ├─▶ fetch_url ─▶ SourceDocument ─┐
web_search ──────┘                                ├─▶ ingest ─▶ ingest_to_index ─▶ RAG index
                                                   │
                      (each hit's URL fetched) ────┘

Layout (one concern per file, in include order)

Quickstart

using HealthLLM

register_models("llama3.2", "nomic-embed-text")
index = ingest_to_index(; sources=["FunSQL.jl", "OMOP CDM"],
                          query="OMOP condition_occurrence table",
                          embedder_kwargs=(model="nomic-embed-text",))
answer = generate_funsql_query(index, "nomic-embed-text", "llama3.2",
                               "Context: {input_query}", "How do I join person and visit?")

CURATED_SOURCES

Registry of curated documentation sources as name => Vector{url}.

Each entry lists one or more URLs for a documentation source. Raw Markdown endpoints (e.g. raw.githubusercontent.com) are preferred because they need no HTML cleaning; rendered pages are also accepted and stripped to text on fetch.

URLs that fail to fetch (404, timeout, network error) are skipped, so entries may safely list several candidate locations for robustness. Extend or override this registry to point at your own curated set.

GROUNDING_DIR

Directory of the local grounding documents shipped with HealthLLM (data/grounding/): an OMOP CDM schema reference, a FunSQL.jl reference, and SQL → FunSQL translation patterns. Every FunSQL snippet in them is checked to render against the documented schema by the test suite, so they are safe to use as the model's only source of syntax. Load them with load_grounding_docs.

AbstractChunkStrategy

Supertype for chunking strategies. A strategy turns one document's text into a vector of Chunks via chunk(strategy, text).

Concrete strategies:

• RecursiveChunk — size-bounded recursive split (default for prose).
• HeaderChunk — one chunk per Markdown heading section (OMOP tables).
• RecordChunk — one chunk per JSONL record (FunSQL examples).
• FixedSizeChunk — fixed character windows with overlap.

AbstractSearchProvider

Interface for web-search backends. Implement web_search for a concrete provider to plug it into ingestion. Only DuckDuckGoProvider is wired up for now; add more providers as subtypes later.

Chunk

One retrieval chunk plus grounding metadata.

Fields

• text::String: The chunk text.
• metadata::Dict{Symbol,Any}: Grounding info such as :heading (the parent Markdown/section heading, e.g. the OMOP table a chunk describes), :group (FunSQL example group), :source, :url, and :title. Populated by the strategy and enriched with the document's provenance by chunk_document.

DuckDuckGoProvider()

Keyless web search via the DuckDuckGo HTML endpoint. No API key required. The endpoint is rate-limited and results are best-effort scraped from the HTML.

FixedSizeChunk(; size=512, overlap=64)

Split into fixed windows of size characters that overlap their neighbours by overlap characters. Character-based and structure-blind (it will split code mid-expression) — provided as a baseline for comparison; prefer RecursiveChunk for real prose.

HeaderChunk(; min_level=1, max_level=6, max_length=4096)

Split Markdown on ATX heading boundaries (#..######), emitting one chunk per heading together with the body beneath it, up to the next heading of the same or higher level. This keeps an OMOP table definition (its name heading plus the full column/foreign-key list) in a single chunk, tagged with :heading.

A section longer than max_length characters is further split with the same code-preserving splitter as RecursiveChunk, so no single chunk blows past embedding limits and no code block is broken. Text before the first heading becomes its own leading chunk. Heading lines inside fenced code blocks are ignored. min_level/max_level bound which heading depths start a new section (e.g. min_level=2 ignores the document title # and breaks on ## headings).

RecordChunk(; fields=["query", "description", "sql_query", "response"],
              labels=["Query", "Description", "SQL", "FunSQL"])

Treat the input as JSON Lines (one JSON object per line) and emit exactly one chunk per record — the atomic retrieval unit for the FunSQL example dataset in FunSQLQueries/.

fields are pulled from each record in order and rendered as "<label>: <value>" lines. The natural-language query/description lead so the embedding emphasises the question side, while the sql_query/response travel in the same chunk as the retrievable answer. Each chunk keeps the record's group and query in :group/:heading metadata for grounding. Missing fields are skipped; lines that fail to parse as JSON are ignored.

RecursiveChunk(; max_length=1024, separators=_RECURSIVE_SEPARATORS)

Split text into heading sections, then split each section to at or below max_length characters, preferring the earliest separator that fits. Fenced code blocks ($```$/~~~) are kept whole regardless of length so a FunSQL snippet is never cut mid-expression. Each resulting chunk records its parent :heading. The right default for prose docs (FunSQL guide, OHDSI book, JuliaHealth), where a size cut degrades gracefully.

SearchResult

A single web-search hit returned by an AbstractSearchProvider.

Fields

• title::String: Result title.
• url::String: Result URL.
• content::String: Snippet or extracted content (may be empty depending on provider).
• score::Float64: Provider-supplied relevance score, or NaN if unavailable.

SourceDocument

A single ingested document with provenance.

Fields

• source::String: Logical source name (e.g. "FunSQL.jl", "web-search").
• url::String: URL the content came from.
• title::String: Human-readable title (best effort; may be empty).
• content::String: Cleaned plain-text content.

chunk(strategy::AbstractChunkStrategy, text::AbstractString) -> Vector{Chunk}

Split text into atomic Chunks according to strategy. Fenced code blocks are never split mid-expression, and each chunk carries the parent heading it came from where the strategy can determine one. Empty/whitespace-only chunks are dropped.

chunk_document(doc::SourceDocument; strategy=default_strategy(doc)) -> Vector{Chunk}

Chunk one SourceDocument at the granularity appropriate to its content type (or an explicit strategy), enriching each chunk's metadata with the document's :source, :url, and :title for later grounding.

chunk_provenance(c::Chunk) -> String

Render a chunk's grounding as a single provenance string ("<url-or-source> › <parent>", truncated to 512 chars) suitable for RAGTools chunk sources.

Delegates to render_provenance, which is also what the prompt layer uses to tag a retrieved chunk — so what is stored alongside a vector and what the model is shown are the same string.

default_search_provider() -> AbstractSearchProvider

Return the default search backend. Currently DuckDuckGoProvider (keyless).

default_strategy(doc::SourceDocument) -> AbstractChunkStrategy

Pick a chunking strategy from a document's source/content:

Detection is by source name first (e.g. the "FunSQL-examples" source produced by load_funsql_examples), then by a content sniff so ad-hoc documents are still routed sensibly.

fetch_curated(names=keys(CURATED_SOURCES); timeout=30, min_length=200) -> Vector{SourceDocument}

Fetch the curated documentation sources named in names from CURATED_SOURCES.

For each source, its candidate URLs are tried in order and every one that fetches successfully and yields at least min_length characters is kept. URLs that error are skipped with a warning. Returns a flat vector of SourceDocument.

Example

docs = fetch_curated(["FunSQL.jl", "OMOP CDM"])

fetch_url(url; timeout=30, max_bytes=5_000_000) -> String

Fetch url over HTTP(S) and return cleaned text.

Content served as HTML is run through html_to_text; Markdown/plain text is returned as-is. Responses larger than max_bytes are truncated. Throws on network errors or non-2xx status.

Keywords

• timeout=30: Per-request timeout in seconds.
• max_bytes=5_000_000: Byte cap on the response body before cleaning.

html_to_text(html) -> String

Strip HTML markup to readable plain text without an external parser.

Removes <script>/<style>/<head> blocks and HTML comments, converts a few block-level tags to newlines, drops all remaining tags, decodes common HTML entities, and collapses excess whitespace.

Example

html_to_text("<p>Hello <b>world</b></p>")  # -> "Hello world"

ingest(; sources=keys(CURATED_SOURCES), query=nothing, provider=default_search_provider(),
         max_results=5, fetch_search_content=true, timeout=30, min_length=200) -> Vector{SourceDocument}

Gather documents from curated sources and/or a live web search into a single vector of SourceDocuments.

Keywords

• sources: Curated source names to fetch (see CURATED_SOURCES). Pass String[]/nothing to skip curated fetching and search only.
• query: If given, also run a web search for query via provider.
• provider: Search backend (default: default_search_provider).
• max_results: Max web-search hits to keep.
• fetch_search_content: When true, fetch each result URL for full text; when false, use only the provider-supplied snippet/content.
• timeout, min_length: Passed through to fetching/filtering.

Example

# Curated docs plus a targeted web search
docs = ingest(; sources=["FunSQL.jl"], query="FunSQL From Where Select example")

ingest_to_index(cfg=RAGTools.SimpleIndexer(); sources=keys(CURATED_SOURCES), query=nothing,
                provider=default_search_provider(), max_results=5, fetch_search_content=true,
                timeout=30, min_length=200, strategy=nothing,
                embedder_kwargs=NamedTuple(), docs=nothing) -> index

Ingest curated docs and/or web-search results and build a RAG index directly from the in-memory text — no intermediate files.

Each document is chunked at the granularity appropriate to its content type via chunk_document: FunSQL example records stay whole (RecordChunk), OMOP/Markdown table sections break on headings (HeaderChunk), and prose is split recursively (RecursiveChunk); fenced code blocks are never split mid-expression. Pass an explicit strategy (an AbstractChunkStrategy) to force one for every document. Each chunk's provenance — its source/URL plus the parent heading or table it came from — is recorded via chunk_provenance.

Chunking happens here in-pipeline; RAGTools.TextChunker() is then used only as a passthrough (a large max_length so it does not re-split the already-atomic chunks). Pass a precomputed docs vector (from ingest) to reuse it instead of fetching again. The returned index is ready for generate_funsql_query.

Example

register_models("llama3.2", "nomic-embed-text")
index = ingest_to_index(; sources=["FunSQL.jl", "OMOP CDM"],
                          query="OMOP condition_occurrence table",
                          embedder_kwargs=(model="nomic-embed-text",))
answer = generate_funsql_query(index, "nomic-embed-text", "llama3.2",
                               "Context: {input_query}", "How do I join person and visit?")

load_funsql_examples(path; source="FunSQL-examples") -> Vector{SourceDocument}

Load a local JSON Lines file of FunSQL examples (e.g. data/benchmark/funsql_tasks.jsonl, whose FunSQL is verified on DuckDB) into a single SourceDocument whose content is the raw JSONL. Routed through RecordChunk at index time, it yields one chunk per example. The file path is recorded as provenance.

load_grounding_docs(dir=GROUNDING_DIR) -> Vector{SourceDocument}

Load every Markdown file in dir as a SourceDocument, with source "grounding:<file name>" and the file's first heading as its title. They are Markdown with one heading per table or pattern, so default_strategy chunks them with HeaderChunk and each OMOP table definition or translation pattern becomes one retrievable chunk.

Unlike fetch_curated this needs no network, which makes it the reproducible grounding source for benchmarks.

Example

examples = joinpath(pkgdir(HealthLLM), "data", "benchmark", "funsql_tasks.jsonl")
docs = vcat(load_grounding_docs(), load_funsql_examples(examples))
index = ingest_to_index(; docs=docs, embedder_kwargs=(model="nomic-embed-text",))

web_search(query; max_results=5) -> Vector{SearchResult}
web_search(provider, query; max_results=5) -> Vector{SearchResult}

Run query against provider and return up to max_results SearchResults. provider defaults to default_search_provider.

Example

hits = web_search("OMOP CDM person table columns"; max_results=3)

Embeddings

Embedding-model selection, generation, and validation for the RAG pipeline.

The default model, DEFAULT_EMBEDDING_MODEL, is chosen to be usable from both a local Ollama server and HuggingFace, so the same index can be reproduced in either environment. See EMBEDDING_MODELS for the registry.

Layout

DEFAULT_EMBEDDING_MODEL

Canonical name of the recommended embedding model: "nomic-embed-text". It is available on both backends, produces 768-dimensional vectors with strong retrieval quality, and is already the model used throughout the ingestion docs. For a lighter/faster option use "all-minilm" (384-d).

EMBEDDING_MODELS

Registry of embedding models reachable from both Ollama and HuggingFace, keyed by canonical name. Every entry can be pulled with ollama pull <tag> and also exists as a HuggingFace repo, so an index is reproducible across the two backends.

EmbeddingModel

Describes one embedding model and how to reach it from each backend.

Fields

• name::String: Canonical name (registry key).
• ollama::String: Ollama model tag (e.g. "nomic-embed-text").
• huggingface::String: HuggingFace repo id (e.g. "nomic-ai/nomic-embed-text-v1.5").
• dim::Int: Output embedding dimension.
• normalized::Bool: Whether the model emits unit-norm vectors by default.

cosine_similarity(a, b) -> Float64

Cosine similarity between two vectors. Returns 0.0 if either vector is zero.

embed(texts, name=DEFAULT_EMBEDDING_MODEL; provider=:ollama, kwargs...) -> Matrix{Float32}

Embed one or more texts with the named model and return a dim × n matrix whose columns are the embeddings (the row=dimension / column=chunk convention used by the storage layer and validate_embeddings).

provider picks the backend (:ollama or :huggingface); extra kwargs are forwarded to PromptingTools.aiembed. Requires the corresponding backend to be reachable (a running Ollama server, or HuggingFace access).

Example

E = embed(["OMOP person table", "FunSQL From clause"])   # 768 × 2
size(E, 1) == embedding_dimension()                       # true

embedding_dimension(name=DEFAULT_EMBEDDING_MODEL) -> Int

Return the output dimension of the named embedding model.

embedding_model(name=DEFAULT_EMBEDDING_MODEL) -> EmbeddingModel

Look up an EmbeddingModel by canonical name, erroring with the list of known names if it is not registered.

embedding_ref(name=DEFAULT_EMBEDDING_MODEL; provider=:ollama) -> String

Return the backend-specific model reference for name: the Ollama tag for provider=:ollama, or the "hf:<repo>"-prefixed HuggingFace id for provider=:huggingface.

embedding_sanity_check(name=DEFAULT_EMBEDDING_MODEL;
                       provider=:ollama,
                       anchor="How do I count patients in OMOP?",
                       near="Query the number of persons in the CDM",
                       far="A recipe for chocolate cake",
                       kwargs...) -> NamedTuple

Live sanity check that the model embeds meaning, not noise: it embeds an anchor, a semantically near sentence, and an unrelated far sentence, then verifies cos(anchor, near) > cos(anchor, far). Returns (; near, far, passed, dim). Requires the backend to be reachable.

similarity_matrix(embeddings) -> Matrix{Float64}

Pairwise cosine-similarity matrix over the columns of a dim × n embedding matrix. The diagonal is ~1 for any non-degenerate embedding set.

validate_embeddings(embeddings; expected_dim=nothing, chunks=nothing, atol=1e-4) -> NamedTuple

Validate an embedding matrix (dim × n, columns = chunks) and return a (; dim, n, norms) summary. Checks, in order:

1. non-empty (at least one column);
2. dimension equals expected_dim when given;
3. column count equals length(chunks) when chunks is given;
4. all entries finite (no NaN/Inf);
5. no zero-vector columns (degenerate embeddings);
6. self-similarity: each column's cosine similarity with itself is ≈ 1 (atol).

Throws ArgumentError/DimensionMismatch on the first failure.

Example

E = rand(Float32, 384, 12)
validate_embeddings(E; expected_dim=384)   # (; dim=384, n=12, norms=[...])

Storage

Vector-store backends for the RAG pipeline. Every backend shares one interface so embeddings produced by embed can be persisted and queried the same way regardless of where they live:

Interface

Each AbstractVectorStore supports:

• add!(store, embeddings, chunks) — append dim × n embeddings with their texts;
• search(store, query, k) — return the k nearest Hits;
• Base.length(store) — number of stored vectors;
• save / load — persist and restore where the backend supports it.

AbstractVectorStore

Supertype for all vector-store backends. Concrete stores implement add!, search, and Base.length.

FaissVectorStore(dim; metric=:cosine)

FAISS-backed store for large local indexes. FAISS is an optional dependency: the store only works when a Faiss module is loaded into Main (e.g. using Faiss in the driver session) — otherwise construction raises a clear error. Chunks are kept alongside the index so search can return their text.

Fields

• dim::Int: Embedding dimension.
• metric::Symbol: :cosine/:dot (inner-product index) or :l2 (L2 index).
• index::Any: The underlying FAISS index object.
• chunks::Vector{String}: Text aligned with the vectors added to index.
• normalize::Bool: L2-normalise vectors (so inner product == cosine).

Example

using Faiss                       # optional, must be loaded first
store = FaissVectorStore(768)
add!(store, embed(chunks), chunks)
hits = search(store, embed("count patients"), 5)   # (; index, chunk, score)

Hit(index, chunk, score, distance=nothing)

One retrieval result, in the shape every backend returns. A uniform hit means downstream code — prompt building above all — never has to know which store the chunk came from.

Fields

• index::Int: Position of the vector in the store, or the database row id for PgVectorStore.
• chunk::String: The stored text.
• score::Float64: Similarity, larger is always more similar, regardless of the backend's native metric. Cosine backends report cosine similarity in [-1, 1]; distance-based ones report a negated distance.
• distance::Union{Float64,Nothing}: The backend's raw distance where it has one (pgvector), else nothing.

LocalVectorStore(dim; normalize=true)

In-memory, file-backed vector store. Embeddings are held as a dim × n matrix and searched with cosine similarity — no external service required. Use save and load to round-trip the store through Serialization.

Fields

• dim::Int: Embedding dimension every added vector must match.
• embeddings::Matrix{Float32}: dim × n matrix, columns = stored vectors.
• chunks::Vector{String}: Text for each column, aligned with embeddings.
• normalize::Bool: Store L2-normalised vectors (makes cosine search a dot product).

Example

store = LocalVectorStore(768)
add!(store, embed(["OMOP person table", "FunSQL From clause"]),
             ["OMOP person table", "FunSQL From clause"])
hits = search(store, embed("count patients"), 1)
hits[1].chunk        # "OMOP person table"
save(store, "index.jls")
store2 = load(LocalVectorStore, "index.jls")

PgVectorStore(conn, dim; table="embeddings", metric=:cosine)

PostgreSQL/pgvector-backed store. Delegates persistence to store_embeddings_pgvector and retrieval to search_embeddings_pgvector; vectors live in table in the database conn points at.

Fields

• conn::LibPQ.Connection: Open connection to a pgvector-enabled database.
• dim::Int: Embedding dimension for the table's VECTOR(dim) column.
• table::String: Destination table name (validated as a SQL identifier).
• metric::Symbol: Distance metric for search — :cosine, :dot, or :l2.

Example

using LibPQ
conn  = LibPQ.Connection("postgresql://localhost/health")
store = PgVectorStore(conn, 768)
add!(store, embed(chunks), chunks)
hits  = search(store, embed("count patients"), 5)   # (; id, chunk, distance)

add!(store::FaissVectorStore, embeddings, chunks) -> store

Add dim × n embeddings (normalised when store.normalize) and their chunks to the FAISS index. FAISS expects one row per vector, so the dim × n matrix is transposed to n × dim before being handed to Faiss.add.

add!(store::LocalVectorStore, embeddings, chunks) -> store

Append a dim × n embedding matrix and its n text chunks to store. Vectors are L2-normalised first when store.normalize is set. Errors if the dimension or the column/chunk counts do not line up.

add!(store::PgVectorStore, embeddings, chunks) -> store

Create store.table if needed and insert the dim × n embeddings with their chunks, in a single transaction. See store_embeddings_pgvector.

load(::Type{LocalVectorStore}, path) -> LocalVectorStore

Load a store previously written with save.

retrieve(store, query::AbstractString, k=5;
         model=DEFAULT_EMBEDDING_MODEL, provider=:ollama, embedder=embed, kwargs...)

Per-query retrieval: embed a natural-language query with model and return the k nearest stored chunks from store. This is the text-in entry point that pairs embed with a store's search — whereas search takes a pre-computed query vector, retrieve takes the string and handles embedding.

Works with any AbstractVectorStore; the returned hits are exactly what the backend's search yields — a Vector{Hit} whatever the backend. provider and extra kwargs are forwarded to the embedder. Pass embedder to substitute the embedding function (useful for tests or a cached embedder); it is called as embedder(query, model; provider, kwargs...) and must return the query embedding as a dim × 1 matrix or a length-dim vector.

Example

store = LocalVectorStore(embedding_dimension())
add!(store, embed(chunks), chunks)

hits = retrieve(store, "How do I count patients in OMOP?", 5)
for h in hits
    println(round(h.score, digits=3), "  ", first(h.chunk, 80))
end

save(store::LocalVectorStore, path) -> path

Serialise store to path (via Serialization). Restore with load(LocalVectorStore, path).

search(store::FaissVectorStore, query, k=5) -> Vector{Hit}

Return the k nearest chunks to query as Hits. score is the raw FAISS score (inner product for :cosine/:dot, negative L2 distance for :l2, so larger is always more similar).

search(store::PgVectorStore, query, k=5) -> Vector{Hit}

Return the k nearest stored chunks to query as Hits, ordered most-similar first using store.metric. The row id lands in index and the raw pgvector distance in distance; score is that distance flipped to the larger-is-more-similar convention every backend shares (1 - d for cosine, -d for dot and L2). See search_embeddings_pgvector.

search(store::LocalVectorStore, query, k=5) -> Vector{Hit}

Return the k most cosine-similar chunks to query (a length-dim vector), as Hits ordered most-similar first. score is the cosine similarity in [-1, 1].

Evaluation

A benchmark harness for natural-language → FunSQL generation that measures whether retrieval grounding actually helps. Every task is run twice through the same model — once zero-shot (no retrieved context) and once RAG (context from a vector store) — and each output is scored by executing it, not by string similarity to a reference.

BenchmarkTask ─▶ prompt (zero-shot | RAG) ─▶ generate_funsql ─▶ sanity_check_funsql
                                                                   │
                       gold SQL ─▶ executor ─▶ ResultSet ◀─ executor ◀┘
                                                   │
                                          results_match ─▶ TaskResult

Scoring

Each TaskResult records the deepest stage the generated query reached and why it stopped (outcome):

Execution accuracy — the headline metric — is the fraction of scorable tasks whose generated query returns the same result as the gold SQL (see results_match for what "the same" means). compare_modes pairs the two runs task-by-task and reports the discordant counts plus an exact McNemar test, so a RAG gain can be told apart from noise.

Leakage

The benchmark questions come from the same FunSQL example set the RAG index can draw on. Indexing a task's own gold answer and then asking that question measures lookup, not grounding. split_tasks holds the evaluated tasks out, and examples_document turns only the remaining tasks into retrievable examples.

Dependencies

The harness itself depends on nothing beyond HealthLLM: the database is reached through an injected executor(sql) -> ResultSet, the model through generate_funsql's injectable generator, and retrieval through an injected retriever(question) -> hits. Building and rendering the generated FunSQL requires FunSQL loaded in Main, as for sanity_check_funsql. scripts/benchmark/run_benchmark.jl wires all of this to DuckDB and Ollama.

OUTCOMES

Every value TaskResult.outcome can take, from best to worst. See the module docstring for their meaning.

TIERS

Benchmark difficulty tiers, in increasing order of what the model must know:

• :simple — one OMOP table: filter, aggregate, group, order.
• :compositional — several clauses composed: joins, subqueries, CTEs, set or window operations.
• :schema_aware — needs the OMOP vocabulary (concept, concept_ancestor, ...) to resolve a clinical idea to concept ids, i.e. domain-model knowledge rather than query syntax alone.

ZERO_SHOT_SYSTEM_PROMPT

System message for the zero-shot baseline. It asks for the same thing as FUNSQL_SYSTEM_PROMPT — one FunSQL query over the OMOP CDM in a fenced julia block — but drops the "use only the retrieved context" grounding rules, which would otherwise tell the model to refuse when there is no context at all. The difference between the two runs is then the retrieved context, not a stricter instruction.

BenchmarkTask

One natural-language → FunSQL benchmark task.

Fields

• id::String: Stable identifier, used to pair zero-shot and RAG runs.
• question::String: The natural-language analytical question.
• group::String: The OMOP area the task belongs to (e.g. "drug_exposure").
• tier::Symbol: One of TIERS.
• gold_sql::String: Reference SQL whose result defines the correct answer.
• gold_funsql::String: Reference FunSQL (used as a retrievable example, never as the scoring reference).
• description::String: Longer task description (may be empty).

ResultSet(columns, rows)

A query result in a backend-neutral form, so the harness does not depend on any particular database or table package.

Fields

• columns::Vector{String}: Column names, in order.
• rows::Vector{Vector{Any}}: One vector of values per row, aligned with columns.

A DataFrame converts as ResultSet(names(df), [collect(Any, r) for r in eachrow(df)]).

TaskResult

The outcome of running one BenchmarkTask in one mode.

Fields

• task_id::String, question::String, group::String, tier::Symbol: Copied from the task.
• mode::Symbol: :zero_shot or :rag.
• outcome::Symbol: One of OUTCOMES.
• funsql::String: The FunSQL code extracted from the model's reply.
• answer::String: The model's full reply.
• sql::Union{String,Nothing}: The SQL the generated query rendered to.
• error::Union{String,Nothing}: Message from the failing stage.
• context::Vector{String}: Provenance (or a text excerpt) of each retrieved chunk, in rank order. Empty for zero-shot.
• scores::Vector{Union{Float64,Nothing}}: Retrieval score of each chunk, or nothing where the retriever gives none.
• seconds::Float64: Wall time for retrieval, generation, and checking.

classify_tier(sql::AbstractString) -> Symbol

Assign a TIERS tier from a task's gold SQL: :schema_aware when it reads an OMOP vocabulary table, else :compositional when it joins, nests a subquery, defines a CTE, or uses a set/window operation, else :simple.

A heuristic over the reference query, so it labels what the answer needs rather than how the question is phrased. Override it by giving a task an explicit "tier" field in its JSONL record.

compare_modes(results; baseline=:zero_shot, treatment=:rag) -> NamedTuple

Pair baseline and treatment results by task and compare their correctness. Tasks excluded in either mode (:gold_error) are dropped. Returns the paired task count, both accuracies and their difference, the 2×2 agreement counts (both_correct, treatment_only, baseline_only, neither), the exact McNemar p-value on the discordant pairs, and the ids of the tasks each mode alone got right — the starting point for error analysis.

examples_document(tasks; source="FunSQL-examples") -> SourceDocument

Turn the tasks that have a reference FunSQL into a single JSONL SourceDocument that chunk_document splits one example per chunk (RecordChunk). Pass only the examples half of split_tasks so no evaluated task can retrieve its own answer.

load_tasks(path; tiers=TIERS) -> Vector{BenchmarkTask}

Load benchmark tasks from a JSON Lines file in the FunSQL example format (query, group, description, sql_query, response, plus optional id and tier). Records missing a question or gold SQL are skipped. A record without an id gets "task-<line>"; one without a tier is classified with classify_tier. Only tasks whose tier is in tiers are returned.

markdown_report(results; meta=(;)) -> String

Render a benchmark run as a Markdown report: the headline comparison, per-mode stage success rates, accuracy by tier, the error-outcome breakdown, and the tasks each mode alone solved. Intended to be pasted into the evaluation write-up.

mcnemar_exact(b, c) -> Float64

Two-sided exact McNemar p-value for paired binary outcomes, from the discordant counts b (only the first condition correct) and c (only the second correct): the binomial probability of a split at least as uneven as min(b, c) out of b + c under no difference. Concordant pairs carry no information and are not inputs. Returns 1.0 when there are no discordant pairs.

results_match(gold::ResultSet, pred::ResultSet; ordered=false, atol=1e-6, rtol=1e-4) -> Bool

Execution-accuracy comparison: does pred return the same answer as gold?

• Columns are compared by content, not by name — generated queries rarely reproduce the reference's aliases. Both results must have the same number of columns. When the column names match as a set (case-insensitively) they are aligned by name; otherwise every column ordering of pred is tried (up to six columns), so Select(:n, :gender) matches Select(:gender, :n).
• Rows are compared as a multiset unless ordered=true: SQL row order is only defined under ORDER BY, and the questions rarely require one.
• Values: numbers compare with isapprox(; atol, rtol) (so COUNT as Int64 matches Float64, and AVG tolerates rounding), missing/nothing equal each other, and anything else compares by value or, failing that, by its string form (so a Date matches its ISO string).

run_benchmark(tasks; modes=(:zero_shot, :rag), progress=true, kwargs...) -> Vector{TaskResult}

Run every task in every mode with run_task (kwargs are forwarded to it — model, executor, catalog, and, for :rag, retriever are required). Each task's gold SQL is executed once and its result shared by both modes, so the modes are compared on identical references. Results are ordered by task, then mode.

run_task(task, mode; model, executor, catalog, retriever=nothing,
         template=DEFAULT_FUNSQL_TEMPLATE, gold=nothing, generator=nothing,
         match_kwargs=(;), kwargs...) -> TaskResult

Run one task in mode (:zero_shot or :rag) and score it.

1. Build the prompt: zero_shot_prompt or, for :rag, build_prompt(question, retriever(question); template).
2. Generate with generate_funsql (model, generator, and kwargs are forwarded to it).
3. Parse, build, and render the result against catalog (a FunSQL.SQLCatalog of the target database) with sanity_check_funsql.
4. Run the SQL with executor(sql) -> ResultSet and compare it to the gold result with results_match.

gold is the task's gold ResultSet, or an exception if the gold SQL failed; when nothing it is computed with executor. run_benchmark computes it once per task and shares it between modes.

save_tasks(path, tasks) -> path

Write tasks as JSON Lines in the format load_tasks reads, with explicit id and tier fields so a saved benchmark reloads identically.

split_tasks(tasks; test_fraction=0.5, seed=0) -> (; examples, test)

Split tasks into a held-out test set to evaluate on and an examples set the RAG index may retrieve from. The split is stratified by tier so each tier is represented in both halves.

Tasks are ordered by a fixed hash of seed and their id (FNV-1a, not Base.hash), so the split is identical across Julia versions and platforms and does not change when unrelated tasks are added or removed.

summarize(results) -> Dict{Symbol,Any}

Per-mode metrics over results. For each mode: the overall stats — tasks scored, execution accuracy (correct ÷ scorable), the success rate of each pipeline stage (parse → build → render → execute), the outcome histogram, and mean seconds per task — and the same broken down by_tier and by_group. Tasks whose gold SQL failed are reported as excluded and left out of every rate.

write_results(path, results; meta=(;)) -> path

Save a benchmark run as JSON: meta (model names, dataset, seed, ...), the summarize metrics, the compare_modes comparison when both modes ran, and every TaskResult so a run can be re-analysed without re-querying the model.

zero_shot_prompt(question) -> (; system, user, prompt)

The no-retrieval baseline prompt, in the same (; system, user, prompt) shape build_prompt returns, so both modes go through generate_funsql identically.