From ab0cdf6c3ed05bcb1475f043ccf75d599bb2c389 Mon Sep 17 00:00:00 2001 From: paramthakkar123 Date: Sun, 19 Jul 2026 09:34:02 +0530 Subject: [PATCH 1/5] feat: add Prompt module for constructing grounded RAG prompts Introduces a Prompt module with: - FUNSQL system prompt and customizable PromptTemplate - Context formatting from retrieved chunks - Build prompts for LLM query construction Includes full test coverage and documentation page. --- docs/make.jl | 1 + docs/src/index.md | 3 +- docs/src/querying.md | 206 ++++++++++++++++++++++++++++++ src/HealthLLM.jl | 5 + src/prompt.jl | 292 +++++++++++++++++++++++++++++++++++++++++++ test/PromptTest.jl | 79 ++++++++++++ test/runtests.jl | 1 + 7 files changed, 586 insertions(+), 1 deletion(-) create mode 100644 docs/src/querying.md create mode 100644 src/prompt.jl create mode 100644 test/PromptTest.jl diff --git a/docs/make.jl b/docs/make.jl index 634c9a8..09972e5 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -17,6 +17,7 @@ makedocs(; "Getting Started" => "getting-started.md", "Document Ingestion" => "ingestion.md", "Building Embeddings" => "embeddings.md", + "Querying the RAG System" => "querying.md", ], ) diff --git a/docs/src/index.md b/docs/src/index.md index 28b4b04..2ed31c2 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -13,6 +13,7 @@ The package centers on five areas: - collecting source files and writing combined corpora - ingesting curated docs and web-search results into an index (see [Document Ingestion](ingestion.md)) - building retrieval indexes through `RAGTools` +- constructing grounded prompts from retrieved chunks (see [Querying the RAG System](querying.md)) - generating retrieval-backed answers for query construction - building and validating embeddings across Ollama and HuggingFace (see [Building Embeddings](embeddings.md)) - storing embeddings in a local file, PostgreSQL/`pgvector`, or FAISS @@ -32,5 +33,5 @@ index = build_index_rag(RAGTools.SimpleIndexer(), files) More detailed setup, testing commands, and the end-to-end walkthrough are in [Getting Started](getting-started.md). ```@autodocs -Modules = [HealthLLM, HealthLLM.Utils, HealthLLM.Database, HealthLLM.Query, HealthLLM.Ingestion, HealthLLM.Embeddings, HealthLLM.Storage] +Modules = [HealthLLM, HealthLLM.Utils, HealthLLM.Database, HealthLLM.Prompt, HealthLLM.Query, HealthLLM.Ingestion, HealthLLM.Embeddings, HealthLLM.Storage] ``` diff --git a/docs/src/querying.md b/docs/src/querying.md new file mode 100644 index 0000000..bad5c9d --- /dev/null +++ b/docs/src/querying.md @@ -0,0 +1,206 @@ +```@meta +CurrentModule = HealthLLM +``` + +# Querying the RAG System + +Once documents are [ingested](ingestion.md) and [embedded](embeddings.md) into a +vector store, querying is a three-step loop: **retrieve** the chunks relevant to +a question, **construct** a grounded prompt from them, and **generate** a FunSQL +answer with a chat model. + +``` +question ─▶ retrieve ─▶ ranked chunks ─┐ + ├─▶ build_prompt ─▶ (; system, user, prompt) ─▶ chat model ─▶ FunSQL + (schema + docs + examples) ─┘ +``` + +The middle step — [`build_prompt`](@ref) — is what keeps generations honest: it +injects the retrieved schema/doc chunks alongside the natural-language question +and pairs them with a system message that forbids the model from using any table +or column name it cannot see in the context. + +## From question to prompt + +Assume you already have a populated store (see [Building Embeddings](embeddings.md)): + +```julia +using HealthLLM + +store = load(LocalVectorStore, "omop_index.jls") +question = "How many distinct patients had a diabetes diagnosis in 2020?" +``` + +Retrieve the most relevant chunks, then build the prompt from them: + +```julia +hits = retrieve(store, question, 6) # top-6 nearest chunks +p = build_prompt(question, hits) # (; system, user, prompt) +``` + +[`build_prompt`](@ref) returns a NamedTuple with three fields so it fits either +prompting style: + +| Field | What it is | Use it for | +|----------|-------------------------------------------------------|---------------------------------------------| +| `system` | The FunSQL/OMOP system message ([`FUNSQL_SYSTEM_PROMPT`](@ref)) | a proper system/user chat conversation | +| `user` | Context block + question + answer cue | the user turn of that conversation | +| `prompt` | `system` and `user` joined into one string | single-string / `{input_query}`-style calls | + +The `user` turn interleaves the retrieved context with the question: + +``` +# Retrieved context + +[Source 1] https://ohdsi.github.io/CommonDataModel/ › condition_occurrence +# condition_occurrence +condition_occurrence_id, person_id, condition_concept_id, condition_start_date, ... + +[Source 2] FunSQL-examples › patients with a condition +Query: patients with a condition +FunSQL: From(condition_occurrence) |> Group(Get.person_id) + +# Analytical question + +How many distinct patients had a diabetes diagnosis in 2020? + +# FunSQL query +``` + +## Generating the answer + +Feed the prompt to a registered chat model through `PromptingTools`. The split +`system`/`user` form is preferred — it puts the grounding rules in the system +role where models weight them most heavily: + +```julia +using PromptingTools + +register_models("llama3.2", "nomic-embed-text") # once per session + +msg = aigenerate( + [PromptingTools.SystemMessage(p.system), + PromptingTools.UserMessage(p.user)]; + model = "llama3.2", +) +println(msg.content) # the FunSQL query + a short explanation +``` + +For a store-free index built with RAGTools, [`generate_funsql_query`](@ref) +retrieves and generates in one call against a RAG `index` (see +[Document Ingestion](ingestion.md)): + +```julia +index = ingest_to_index(; sources = ["FunSQL.jl", "OMOP CDM"], + query = "OMOP condition_occurrence columns") +answer = generate_funsql_query(index, "nomic-embed-text", "llama3.2", + "Context: {input_query}", question) +``` + +## Grounding: why the model stays on-schema + +The single biggest failure mode for text-to-query over a large schema like OMOP +is **hallucinated identifiers** — the model confidently writes `patient_id` or +`diagnosis_date` because those names are plausible, even though the CDM uses +`person_id` and `condition_start_date`. [`FUNSQL_SYSTEM_PROMPT`](@ref) counters +this with strict grounding rules that make the retrieved context the *only* +admissible source of names and syntax: + +- **Names must appear verbatim in the context.** Every table and column the model + references has to be present in a retrieved chunk. Memory of the OMOP CDM is + explicitly overridden, and inventing, guessing, or "correcting" an identifier is + forbidden. +- **FunSQL constructs are copied, not recalled.** The model may use `From`, + `Where`, `Join`, `Group`, `Select`, `Get`, `Agg`, `Fun`, and `|>` only as the + retrieved examples demonstrate them — no raw SQL strings, no undemonstrated + features. +- **Missing schema is surfaced, not faked.** When the context lacks a table or + column the question needs, the model is told to say what is missing and generate + the closest correct query from what *is* available — never to paper over the gap + with a fabricated name. + +Because grounding is only as good as the context, retrieval quality matters: if +`retrieve` does not surface the `condition_occurrence` schema chunk, no prompt +rule can make the model use it correctly. Two levers help: + +1. **Retrieve enough chunks.** A larger `k` (or a larger `max_chunks`) raises the + chance the needed table definition is present — at the cost of a longer prompt. +2. **Index schema at table granularity.** The [`HeaderChunk`](@ref) strategy keeps + each OMOP table's full column list in one chunk, so a single hit carries the + whole definition rather than a fragment. + +## Tuning the prompt + +[`build_prompt`](@ref) is driven by a [`PromptTemplate`](@ref) — pass one to +change the system message, the section headers, per-chunk formatting, or the +context budget. [`DEFAULT_FUNSQL_TEMPLATE`](@ref) is the OMOP/FunSQL default +(provenance on, scores off, up to 8 chunks or 8000 characters). + +```julia +tmpl = PromptTemplate( + include_scores = true, # show each chunk's retrieval score + max_chunks = 5, # keep at most 5 chunks + max_context_chars = 6000, # ...and cap the context at 6000 chars +) + +p = build_prompt(question, hits; template = tmpl) +``` + +The two budgets protect the context window: chunks are admitted highest-ranked +first, and admission stops at a **chunk boundary** once either limit is reached — +a chunk is included whole or skipped, never cut mid-text (the sole exception is a +single chunk that exceeds the character budget on its own, which is truncated with +an ellipsis). This means the prompt never silently overflows the model's window, +and lower-ranked chunks are the ones dropped when space runs out. + +You can render just the context block — for inspection or a custom prompt layout — +with [`format_context`](@ref): + +```julia +println(format_context(DEFAULT_FUNSQL_TEMPLATE, hits)) +``` + +### Flexible chunk inputs + +`build_prompt` and `format_context` accept whatever your retrieval step produces, +so the same call works across store backends and index paths: + +| Input | Where it comes from | +|--------------------------------|----------------------------------------------| +| `search`/`retrieve` hit NamedTuples | [`LocalVectorStore`](@ref) / [`FaissVectorStore`](@ref) (`.chunk`, `.score`) and [`PgVectorStore`](@ref) (`.chunk`, `.distance`) | +| plain `String`s | any ad-hoc list of context snippets | +| [`Chunk`](@ref)s | the ingestion layer, carrying `:heading`/`:source`/`:url` provenance | + +When a chunk carries provenance (a URL or source plus its parent heading), it is +shown in the `[Source i]` header so the model — and you — can trace each fact back +to its origin. + +## End to end + +```julia +using HealthLLM, PromptingTools + +register_models("llama3.2", "nomic-embed-text") + +# 1. Load (or build) the index. +store = load(LocalVectorStore, "omop_index.jls") + +# 2. Retrieve context for the question. +question = "What is the average age of patients with a hypertension diagnosis?" +hits = retrieve(store, question, 6) + +# 3. Construct the grounded prompt. +p = build_prompt(question, hits) + +# 4. Generate the FunSQL query. +msg = aigenerate( + [PromptingTools.SystemMessage(p.system), + PromptingTools.UserMessage(p.user)]; + model = "llama3.2", +) +println(msg.content) +``` + +See the [API reference](index.md) for full docstrings of [`build_prompt`](@ref), +[`format_context`](@ref), [`PromptTemplate`](@ref), and [`FUNSQL_SYSTEM_PROMPT`](@ref). +``` diff --git a/src/HealthLLM.jl b/src/HealthLLM.jl index 10e5fa0..bf21079 100644 --- a/src/HealthLLM.jl +++ b/src/HealthLLM.jl @@ -11,6 +11,7 @@ include("utils.jl") include("database.jl") include("embeddings.jl") include("storage.jl") +include("prompt.jl") include("query.jl") include("ingestion.jl") @@ -22,6 +23,8 @@ import .Embeddings: EmbeddingModel, EMBEDDING_MODELS, DEFAULT_EMBEDDING_MODEL, validate_embeddings, embedding_sanity_check import .Storage: AbstractVectorStore, LocalVectorStore, PgVectorStore, FaissVectorStore, add!, search, retrieve, save, load +import .Prompt: FUNSQL_SYSTEM_PROMPT, PromptTemplate, DEFAULT_FUNSQL_TEMPLATE, + format_context, build_prompt import .Query: generate_funsql_query import .Ingestion: SourceDocument, SearchResult, AbstractSearchProvider, DuckDuckGoProvider, @@ -30,6 +33,8 @@ import .Ingestion: SourceDocument, SearchResult, ingest, ingest_to_index export PromptingTools, RAGTools +export FUNSQL_SYSTEM_PROMPT, PromptTemplate, DEFAULT_FUNSQL_TEMPLATE, + format_context, build_prompt export collect_files_with_extensions, write_combined_file, generate_funsql_query, build_index_rag, store_embeddings_pgvector, search_embeddings_pgvector, validate_embeddings_inputs, diff --git a/src/prompt.jl b/src/prompt.jl new file mode 100644 index 0000000..347f94d --- /dev/null +++ b/src/prompt.jl @@ -0,0 +1,292 @@ +""" + 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`](@ref)/[`search`](@ref) + (or raw strings / [`Chunk`](@ref)s), 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`](@ref) — the reusable shape (system message, section headers, + per-chunk formatting, and context budgets). [`DEFAULT_FUNSQL_TEMPLATE`](@ref) is + the OMOP/FunSQL default. +- [`format_context`](@ref)`(template, hits)` — render retrieved chunks to a string. +- [`build_prompt`](@ref)`(question, hits; template)` — the entry point; returns + `(; system, user, prompt)`. + +## Example + +```julia +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 +``` +""" +module Prompt + +export FUNSQL_SYSTEM_PROMPT, PromptTemplate, DEFAULT_FUNSQL_TEMPLATE, + format_context, build_prompt + +""" + 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. +""" +const FUNSQL_SYSTEM_PROMPT = """ +You are an expert data analyst for the OHDSI OMOP Common Data Model (CDM), \ +writing analytical queries with the Julia package FunSQL.jl. + +Your task: translate the user's natural-language analytical question into a \ +correct FunSQL query against the OMOP CDM, using ONLY the retrieved context below. + +Grounding (strict — this overrides your prior knowledge): +- The retrieved context is your ONLY source of truth for table names, column \ + names, and FunSQL syntax. Every table and column you reference MUST appear \ + verbatim in the context. Do not rely on memory of the OMOP CDM or of FunSQL, \ + and do not invent, guess, pluralise, or "correct" any identifier. +- Before using a table or column, confirm it is present in the context. If a name \ + you need is not there, DO NOT substitute a similar-looking one — treat it as \ + unavailable. +- Reproduce FunSQL constructs (`From`, `Where`, `Join`, `Group`, `Select`, \ + `Get`, `Agg`, `Fun`, `|>`) only as they appear in the retrieved examples. Do \ + not use SQL string syntax or FunSQL features you cannot see demonstrated in the \ + context. +- If the context lacks a table, column, or construct required to answer the \ + question, state exactly what is missing, then generate the closest correct \ + query you can from what IS available, labelling any assumption. Never paper over \ + a gap with a fabricated identifier. + +Style: +- Follow the idioms of the retrieved FunSQL examples. +- Join through the standard OMOP keys shown in the context (e.g. `person_id`) and \ + filter on the `*_concept_id` columns present in the context where the question \ + implies a clinical concept. + +Output: +- Return the query as ONE fenced code block tagged `julia`, followed by a one- or \ + two-sentence explanation that names the context tables/columns you used. Do not \ + pad the answer with unrelated commentary. +""" + +""" + 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`](@ref). +- `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 `"[