R/patientChat.R

In PatientGenerator: Generator of Synthetic Patient Data for the OMOP Common Data Model

#' `patientChat()` generates synthetic patients in the OMOP-CDM using an LLM API.
#'
#' @description
#' Requires an OPEN_AI_KEY in ~/.Renviron.
#' After that just sent a prompt and save() the results.
#' The JSON file can be used as an OMOP-CDM patient test set.
#' @details
#' Accepts a prompt as input. Produces a test set using a
#' structured JSON schema. Utilizes tools such as CodelistGenerator or Hecate
#' to look up concept IDs. Accepts subsequent prompts to modify existing test
#' sets that the LLM uses as context.
#'
#' This class allows testing patient sets created by the LLM, prompt engineering,
#' integration of search tools and functionality, and creating a set of patients
#' to test analytical packages.
#' @returns 
#' A JSON response that includes: the natural language answer from the LLM and
#' a JSON with test set patients in accordance to the provided schema.
#' @importFrom ellmer chat_openai
#' @importFrom ellmer type_from_schema
#' @importFrom jsonlite fromJSON
#' @importFrom httr2 request req_headers req_perform resp_body_json
#' @importFrom checkmate assertCharacter assertFileExists assertDirectoryExists
#' @importFrom tools file_path_sans_ext
#' @importFrom R6 R6Class
#' @importFrom cli cli_alert_success cli_progress_step
#' @export
#' @examples
#' \dontrun{
#' generator <- patientChat$new()
#' generator$prompt("Give me 5 patients")
#' generator$save("my_test")
#' }
patientChat <- R6::R6Class(
  classname = "patientChat",
  public = list(

    #' @field chat An ellmer chat instance
    chat = NULL,

    #' @field json_schema_path JSON schema to output structured results
    json_schema_path = NULL,

    #' @field response Ouput from the LLM
    response = NULL,

    #' @field codelist A codelist with details to search for concepts ids
    codelist = NULL,

    #' @description
    #' Create a new chat to create JSON test sets for OMOP-CDM.
    #' @param system_prompt Initial system prompt to impose behaviour to the LLM
    #' @param model Such as "gpt-5.3". For a complete list, call patientChat$availableModels()
    #' @param jsonSchemaPath The JSON schema to structure output from LLM
    #' @param echo How the output will be displayed in the console
    #' @param codelist_data A codelist with details to search for concepts ids
    #'
    #' @return A new `Person` object.
    initialize = function(system_prompt = NULL,
                          model = "gpt-5.4",
                          jsonSchemaPath = NULL,
                          echo = c("none", "output", "all"),
                          codelist_data = NULL) {

      # Check API and available models -----------------------
      private$.api_check(model)

      # check JSON schema file -------------------------------
      private$.json_schema_check(jsonSchemaPath)

      # System propmpt ---------------------------------------
      if (is.null(system_prompt)) {
        system_prompt <- "You are an expert generator of synthetic patient data.
                          Your output must strictly adhere to the provided JSON schema. 
                          All generated data, structure, tables, features and data types 
                          must exclusively conform to the OMOP-CDM v5.4 standard"
      }

      # Create chat ------------------------------------------
      self$chat <- ellmer::chat_openai(
        system_prompt = system_prompt,
        model = model,
        echo = echo
        )
      cli::cli_alert_success("Chat created")

      # Codelist ---------------------------------------------
      if (!is.null(codelist_data)) {
        checkmate::assertDataFrame(codelist_data)
        if (all(c(
          "concept_id",
          "concept_name",
          "domain_id",
          "vocabulary_id",
          "standard_concept"
          ) %in% codelist_data |> 
          names()
          )
          ) {
          self$codelist <- codelist_data
          self$chat$register_tool(
            private$.register_codelist_tool()
            )
          message("Codelist tool added")
        } else {
          stop(
            "Dataframe has missing columns: 'concept_id', 'concept_name', 'domain_id', 'vocabulary_id', 'standard_concept'"
            )
        }
      }

    },

    #' @description
    #' Prompt to request data from LLM API
    #' @param prompt A query in character.
    prompt = function(prompt) {
      checkmate::assertCharacter(prompt)
      cli::cli_progress_step(
        msg = "Generating test patients",
        msg_done = "Test set created successfully",
        msg_failed = "Error or connection lost"
        )
      api_respose <- self$chat$chat_structured(
        prompt,
        type = ellmer::type_from_schema(
          path = self$json_schema_path
        )
      )
      self$response <- api_respose
    },

    #' @description
    #' Output in JSON format
    json_response = function() {
      jsonlite::toJSON(
        self$response,
        dataframe = "rows",
        pretty = TRUE,
        null = "null",
        na = "null",
        auto_unbox = TRUE
        )
    },

    #' @description
    #' Returns the chat object
    output = function() {
      return(self$chat)
    },

    #' @description
    #' Retrieves and filters data from codelist_data
    #' @param concept_label Filters the concept_name in the codelist with details
    #' @param domain Filters the domain in the codelist with details.
    retrieveCodelist = function(concept_label = "Stage 1", domain = "Measurement") {
      checkmate::assertCharacter(concept_label)
      checkmate::assertCharacter(domain)
      domain_names <- self$codelist |>
        dplyr::pull(domain_id) |>
        unique()
      if (!domain %in% domain_names) {
        warning("domain not found in codelist")
      }
      result <- self$codelist %>%
        filter(stringr::str_detect(
          concept_name,
          stringr::regex(
            concept_label,
            ignore_case = TRUE
            )
          )
          ) |>
        dplyr::filter(domain_id == domain) |>
        jsonlite::toJSON(
          dataframe = "rows",
          auto_unbox = TRUE
          )
      return(result)
    },

    #' @description
    #' Saves the JSON test set to disk.
    #' @param name Name of the file
    #' @param path To save the file.
    #' If NULL, the package first tries `testthat::test_path("testCases")`,
    #' then checks `options(PatientGenerator.testSetDir = "...")`, and finally
    #' falls back to the package user data directory.
    save = function(name = "patient-chat-test",
                    path = NULL) {
      if (is.null(path)) {
        path <- testSetDir(create = TRUE)
      } else {
        checkmate::assertCharacter(path)
        checkmate::assertDirectoryExists(path)
      }
      name <- tools::file_path_sans_ext(name)
      test_file_path <- file.path(
        path,
        paste0(
          name,
          ".json"
          )
        )
      write(
        self$json_response(),
        file = test_file_path
        )
      cli::cli_alert_success("Test set saved to {.path {test_file_path} }")
    },

    #' @description
    #' Retrieves available models from the LLM API.
    availableModels = function() {
      private$.check_api_key()
      # Retrieve and check available models
      response <- httr2::request("https://api.openai.com/v1/models") |>
        httr2::req_headers(
          Authorization = paste(
            "Bearer",
            Sys.getenv("OPENAI_API_KEY")
            )
        ) |>
        httr2::req_perform()
      models <- httr2::resp_body_json(response)
      lapply(models$data, function(data) data$id) |>
        unlist()
    }
  ),

  private = list(

    .api_check = function(model) {
      checkmate::assertCharacter(model)
      api_models <- self$availableModels()
      if (!model %in% api_models) {
        stop(
          glue::glue("{model} not available.\n"),
          "\n These are some models that are available to you:\n",
          paste0("  - ", sample(api_models, 10), collapse = "\n"),
          "\n For a complete list, call availableModels()\n",
          call. = FALSE
        )
      } else {
        message(glue::glue("{model} is a valid model"))
      }
    },

    .json_schema_check = function(jsonSchemaPath) {
      if (is.null(jsonSchemaPath)) {
        json_schema_path <- system.file(
          "jsonSchemas",
          "cdm54schema-complete.json",
          package = "PatientGenerator"
          )
        checkmate::assertFileExists(json_schema_path)
        self$json_schema_path <- json_schema_path
        } else {
        checkmate::assertCharacter(jsonSchemaPath)
        checkmate::assertFileExists(jsonSchemaPath)
        is_json <- tryCatch({
          jsonlite::fromJSON(jsonSchemaPath)
          message("The file at 'jsonSchemaPath' can be parsed as JSON")
          self$json_schema_path <- jsonSchemaPath
        }, error = function(e) {
          stop("The file at 'jsonSchemaPath' cannot be parsed as JSON")
        })
      }
    },

    .check_api_key = function(apiKeyName = "OPENAI_API_KEY") {
      key <- Sys.getenv("OPENAI_API_KEY", unset = "")
      if (!nzchar(key)) {
        stop(
          "API key not found.\n",
          "Set it with Sys.setenv(OPENAI_API_KEY = 'your_key') ",
          "and/or add it to your ~/.Renviron.",
          call. = FALSE
        )
      }
    },

    .register_codelist_tool = function() {
      retrieveCodelistTool <- ellmer::tool(
        fun = self$retrieveCodelist,
        description = "Retrieves data for concept_id search by concept_name and domain_id",
        name = "retrieveCodelist",
        arguments = list(
          concept_label = ellmer::type_string(
            description = "A regex to look up for a concept_name; the function uses stringr::regex to filter the concept_name"
            ),
          domain = ellmer::type_string(
            description = "One word typically a domain from the OMOP-CDM, currently only available now: 'Drug', 'Condition', 'Measurement', 'Observation', 'Procedure'"
            )
        )
      )
    }
  )
)