@sjcrh/proteinpaint-server 2.174.1 → 2.175.0

package/routes/termdb.chat.js

@@ -7,6 +7,81 @@ import { mayLog } from "#src/helpers.ts";
 import Database from "better-sqlite3";
 import { formatElapsedTime } from "#shared";
 const num_filter_cutoff = 3;
+const FILTER_TERM_DEFINITIONS = {
+  FilterTerm: {
+    anyOf: [{ $ref: "#/definitions/CategoricalFilterTerm" }, { $ref: "#/definitions/NumericFilterTerm" }]
+  },
+  CategoricalFilterTerm: {
+    type: "object",
+    properties: {
+      term: { type: "string", description: "Name of categorical term" },
+      category: { type: "string", description: "The category of the term" },
+      join: {
+        type: "string",
+        enum: ["and", "or"],
+        description: "join term to be used only when there is more than one filter term and should be placed from the 2nd filter term onwards describing how it connects to the previous term"
+      }
+    },
+    required: ["term", "category"],
+    additionalProperties: false
+  },
+  NumericFilterTerm: {
+    type: "object",
+    properties: {
+      term: { type: "string", description: "Name of numeric term" },
+      start: { type: "number", description: "start position (or lower limit) of numeric term" },
+      stop: { type: "number", description: "stop position (or upper limit) of numeric term" },
+      join: {
+        type: "string",
+        enum: ["and", "or"],
+        description: "join term to be used only when there is more than one filter term and should be placed from the 2nd filter term onwards describing how it connects to the previous term"
+      }
+    },
+    required: ["term"],
+    additionalProperties: false
+  }
+};
+function formatTrainingExamples(trainingData) {
+  return trainingData.map(
+    (td, i) => "Example question" + (i + 1).toString() + ": " + td.question + " Example answer" + (i + 1).toString() + ":" + JSON.stringify(td.answer)
+  ).join(" ");
+}
+const FILTER_DESCRIPTION = 'There are two kinds of filter variables: "Categorical" and "Numeric". "Categorical" variables are those variables which can have a fixed set of values e.g. gender, race. They are defined by the "CategoricalFilterTerm" which consists of "term" (a field from the sqlite3 db) and "category" (a value of the field from the sqlite db). "Numeric" variables are those which can have any numeric value. They are defined by "NumericFilterTerm" and contain the subfields "term" (a field from the sqlite3 db), "start" an optional filter which is defined when a lower cutoff is defined in the user input for the numeric variable and "stop" an optional filter which is defined when a higher cutoff is defined in the user input for the numeric variable. ';
+function extractGenesFromPrompt(prompt, genes_list) {
+  const words = prompt.replace(/[^a-zA-Z0-9\s]/g, "").split(/\s+/).map((str) => str.toLowerCase());
+  return words.filter((item) => genes_list.includes(item));
+}
+const CHILD_TYPE_DEFAULTS = {
+  "categorical:undefined": "barchart",
+  "numeric:undefined": "violin",
+  "categorical:categorical": "barchart",
+  "numeric:categorical": "violin",
+  "categorical:numeric": "violin",
+  "numeric:numeric": "sampleScatter"
+};
+const CHILD_TYPE_INVALID = {
+  "categorical:undefined": /* @__PURE__ */ new Set(["violin", "boxplot", "sampleScatter"]),
+  "categorical:categorical": /* @__PURE__ */ new Set(["violin", "boxplot", "sampleScatter"])
+};
+function resolveChildType(cat1, cat2, llmChildType) {
+  const norm1 = cat1 == "float" || cat1 == "integer" ? "numeric" : cat1 || "undefined";
+  const norm2 = cat2 == "float" || cat2 == "integer" ? "numeric" : cat2 || "undefined";
+  const key = norm1 + ":" + norm2;
+  const defaultType = CHILD_TYPE_DEFAULTS[key];
+  if (!defaultType) {
+    return { childType: "barchart" };
+  }
+  const invalid = CHILD_TYPE_INVALID[key];
+  if (llmChildType && invalid && invalid.has(llmChildType)) {
+    return {
+      error: "Invalid plot type supplied by the user: " + llmChildType + ". For " + key.replace(":", " and ") + " variables the plot type should always be " + defaultType
+    };
+  }
+  return {
+    childType: llmChildType || defaultType,
+    bothNumeric: norm1 == "numeric" && norm2 == "numeric"
+  };
+}
 const api = {
   endpoint: "termdb/chat",
   methods: {
@@ -32,72 +107,25 @@ function init({ genomes }) {
       if (!serverconfig_ds_entries.aifiles) {
         throw "aifiles are missing for chatbot to work";
       }
-      let apilink;
-      let comp_model_name;
-      if (serverconfig.llm_backend == "SJ") {
-        apilink = serverconfig.sj_apilink;
-        comp_model_name = serverconfig.sj_comp_model_name;
-      } else if (serverconfig.llm_backend == "ollama") {
-        apilink = serverconfig.ollama_apilink;
-        comp_model_name = serverconfig.ollama_comp_model_name;
-      } else {
-        throw "llm_backend either needs to be 'SJ' or 'ollama'";
+      const llm = serverconfig.llm;
+      if (!llm) throw "serverconfig.llm is not configured";
+      if (llm.provider !== "SJ" && llm.provider !== "ollama") {
+        throw "llm.provider must be 'SJ' or 'ollama'";
       }
       const dataset_db = serverconfig.tpmasterdir + "/" + ds.cohort.db.file;
       const genedb = serverconfig.tpmasterdir + "/" + g.genedb.dbfile;
       const dataset_json = await readJSONFile(serverconfig_ds_entries.aifiles);
-      const time1 = (/* @__PURE__ */ new Date()).valueOf();
-      const class_response = await classify_query_by_dataset_type(
+      const testing = false;
+      const ai_output_json = await run_chat_pipeline(
         q.prompt,
-        comp_model_name,
-        serverconfig.llm_backend,
-        apilink,
+        llm,
         serverconfig.aiRoute,
-        dataset_json
+        dataset_json,
+        testing,
+        dataset_db,
+        genedb,
+        ds
       );
-      let ai_output_json;
-      mayLog("Time taken for classification:", formatElapsedTime(Date.now() - time1));
-      if (class_response.type == "html") {
-        ai_output_json = class_response;
-      } else if (class_response.type == "plot") {
-        const classResult = class_response.plot;
-        mayLog("classResult:", classResult);
-        if (classResult == "summary") {
-          const time12 = (/* @__PURE__ */ new Date()).valueOf();
-          ai_output_json = await extract_summary_terms(
-            q.prompt,
-            serverconfig.llm_backend,
-            comp_model_name,
-            apilink,
-            dataset_db,
-            dataset_json,
-            genedb,
-            ds
-          );
-          mayLog("Time taken for summary agent:", formatElapsedTime(Date.now() - time12));
-        } else if (classResult == "dge") {
-          const time12 = (/* @__PURE__ */ new Date()).valueOf();
-          ai_output_json = await extract_DE_search_terms_from_query(
-            q.prompt,
-            serverconfig.llm_backend,
-            comp_model_name,
-            apilink,
-            dataset_db,
-            dataset_json,
-            ds
-          );
-          mayLog("Time taken for DE agent:", formatElapsedTime(Date.now() - time12));
-        } else if (classResult == "survival") {
-          ai_output_json = { type: "html", html: "survival agent has not been implemented yet" };
-        } else {
-          ai_output_json = { type: "html", html: "Unknown classification value" };
-        }
-      } else {
-        ai_output_json = {
-          type: "html",
-          html: "Unknown classification type"
-        };
-      }
       res.send(ai_output_json);
     } catch (e) {
       if (e.stack) mayLog(e.stack);
@@ -105,6 +133,72 @@ function init({ genomes }) {
     }
   };
 }
+async function run_chat_pipeline(user_prompt, llm, aiRoute, dataset_json, testing, dataset_db, genedb, ds) {
+  const time1 = (/* @__PURE__ */ new Date()).valueOf();
+  const class_response = await classify_query_by_dataset_type(
+    user_prompt,
+    llm,
+    aiRoute,
+    dataset_json,
+    testing
+  );
+  let ai_output_json;
+  mayLog("Time taken for classification:", formatElapsedTime(Date.now() - time1));
+  if (class_response.type == "html") {
+    ai_output_json = class_response;
+  } else if (class_response.type == "plot") {
+    const classResult = class_response.plot;
+    mayLog("classResult:", classResult);
+    const dataset_db_output = await parse_dataset_db(dataset_db);
+    const genes_list = dataset_json.hasGeneExpression ? await parse_geneset_db(genedb) : [];
+    if (classResult == "summary") {
+      const time12 = (/* @__PURE__ */ new Date()).valueOf();
+      ai_output_json = await extract_summary_terms(
+        user_prompt,
+        llm,
+        dataset_db_output,
+        dataset_json,
+        genes_list,
+        ds,
+        testing
+      );
+      mayLog("Time taken for summary agent:", formatElapsedTime(Date.now() - time12));
+    } else if (classResult == "dge") {
+      const time12 = (/* @__PURE__ */ new Date()).valueOf();
+      ai_output_json = await extract_DE_search_terms_from_query(
+        user_prompt,
+        llm,
+        dataset_db_output,
+        dataset_json,
+        ds,
+        testing
+      );
+      mayLog("Time taken for DE agent:", formatElapsedTime(Date.now() - time12));
+    } else if (classResult == "survival") {
+      ai_output_json = { type: "html", html: "survival agent has not been implemented yet" };
+    } else if (classResult == "matrix") {
+      const time12 = (/* @__PURE__ */ new Date()).valueOf();
+      ai_output_json = await extract_matrix_search_terms_from_query(
+        user_prompt,
+        llm,
+        dataset_db_output,
+        dataset_json,
+        genes_list,
+        ds,
+        testing
+      );
+      mayLog("Time taken for matrix agent:", formatElapsedTime(Date.now() - time12));
+    } else {
+      ai_output_json = { type: "html", html: "Unknown classification value" };
+    }
+  } else {
+    ai_output_json = {
+      type: "html",
+      html: "Unknown classification type"
+    };
+  }
+  return ai_output_json;
+}
 async function call_ollama(prompt, model_name, apilink) {
   const temperature = 0.01;
   const top_p = 0.95;
@@ -177,6 +271,17 @@ async function call_sj_llm(prompt, model_name, apilink) {
     throw "SJ API request failed:" + error;
   }
 }
+async function route_to_appropriate_llm_provider(template, llm) {
+  let response;
+  if (llm.provider == "SJ") {
+    response = await call_sj_llm(template, llm.modelName, llm.api);
+  } else if (llm.provider == "ollama") {
+    response = await call_ollama(template, llm.modelName, llm.api);
+  } else {
+    throw "Unknown LLM provider";
+  }
+  return response;
+}
 function checkField(sentence) {
   if (!sentence) return "";
   else return sentence;
@@ -185,7 +290,7 @@ async function readJSONFile(file) {
   const json_file = await fs.promises.readFile(file);
   return JSON.parse(json_file.toString());
 }
-async function classify_query_by_dataset_type(user_prompt, comp_model_name, llm_backend_type, apilink, aiRoute, dataset_json) {
+async function classify_query_by_dataset_type(user_prompt, llm, aiRoute, dataset_json, testing) {
   const data = await readJSONFile(aiRoute);
   let contents = data["general"];
   for (const key of Object.keys(data)) {
@@ -193,32 +298,24 @@ async function classify_query_by_dataset_type(user_prompt, comp_model_name, llm_
       contents += data[key];
     }
   }
-  const classification_ds = dataset_json.charts.filter((chart) => chart.type == "Classification");
-  if (classification_ds.length == 0) throw "Classification information is not present in the dataset file.";
-  if (classification_ds[0].TrainingData.length == 0) throw "No training data is provided for the classification agent.";
-  let train_iter = 0;
+  const classification_ds = dataset_json.charts.find((chart) => chart.type == "Classification");
+  if (!classification_ds) throw "Classification information is not present in the dataset file.";
+  if (classification_ds.TrainingData.length == 0) throw "No training data is provided for the classification agent.";
   let training_data = "";
-  if (classification_ds.length > 0 && classification_ds[0].TrainingData.length > 0) {
-    contents += checkField(dataset_json.DatasetPrompt) + checkField(classification_ds[0].SystemPrompt);
-    for (const train_data of classification_ds[0].TrainingData) {
-      train_iter += 1;
-      training_data += "Example question" + train_iter.toString() + ": " + train_data.question + " Example answer" + train_iter.toString() + ":" + JSON.stringify(train_data.answer) + " ";
-    }
+  if (classification_ds && classification_ds.TrainingData.length > 0) {
+    contents += checkField(dataset_json.DatasetPrompt) + checkField(classification_ds.SystemPrompt);
+    training_data = formatTrainingExamples(classification_ds.TrainingData);
   }
   const template = contents + " training data is as follows:" + training_data + " Question: {" + user_prompt + "} Answer: {answer}";
-  let response;
-  if (llm_backend_type == "SJ") {
-    response = await call_sj_llm(template, comp_model_name, apilink);
-  } else if (llm_backend_type == "ollama") {
-    response = await call_ollama(template, comp_model_name, apilink);
+  const response = await route_to_appropriate_llm_provider(template, llm);
+  if (testing) {
+    return { action: "html", response: JSON.parse(response) };
   } else {
-    throw "Unknown LLM backend";
+    return JSON.parse(response);
   }
-  return JSON.parse(response);
 }
-async function extract_DE_search_terms_from_query(prompt, llm_backend_type, comp_model_name, apilink, dataset_db, dataset_json, ds) {
+async function extract_DE_search_terms_from_query(prompt, llm, dataset_db_output, dataset_json, ds, testing) {
   if (dataset_json.hasDE) {
-    const dataset_db_output = await parse_dataset_db(dataset_db);
     const Schema = {
       $schema: "http://json-schema.org/draft-07/schema#",
       $ref: "#/definitions/DEType",
@@ -245,60 +342,20 @@ async function extract_DE_search_terms_from_query(prompt, llm_backend_type, comp
           required: ["group1", "group2"],
           additionalProperties: false
         },
-        FilterTerm: {
-          anyOf: [{ $ref: "#/definitions/CategoricalFilterTerm" }, { $ref: "#/definitions/NumericFilterTerm" }]
-        },
-        CategoricalFilterTerm: {
-          type: "object",
-          properties: {
-            term: { type: "string", description: "Name of categorical term" },
-            category: { type: "string", description: "The category of the term" },
-            join: {
-              type: "string",
-              enum: ["and", "or"],
-              description: "join term to be used only when there is more than one filter term and should be placed from the 2nd filter term onwards describing how it connects to the previous term"
-            }
-          },
-          required: ["term", "category"],
-          additionalProperties: false
-        },
-        NumericFilterTerm: {
-          type: "object",
-          properties: {
-            term: { type: "string", description: "Name of numeric term" },
-            start: { type: "number", description: "start position (or lower limit) of numeric term" },
-            stop: { type: "number", description: "stop position (or upper limit) of numeric term" },
-            join: {
-              type: "string",
-              enum: ["and", "or"],
-              description: "join term to be used only when there is more than one filter term and should be placed from the 2nd filter term onwards describing how it connects to the previous term"
-            }
-          },
-          required: ["term"],
-          additionalProperties: false
-        }
+        ...FILTER_TERM_DEFINITIONS
       }
     };
-    const DE_ds = dataset_json.charts.filter((chart) => chart.type == "DE");
-    if (DE_ds.length == 0) throw "DE information is not present in the dataset file.";
-    if (DE_ds[0].TrainingData.length == 0) throw "No training data is provided for the DE agent.";
-    let train_iter = 0;
-    let training_data = "";
-    for (const train_data of DE_ds[0].TrainingData) {
-      train_iter += 1;
-      training_data += "Example question" + train_iter.toString() + ": " + train_data.question + " Example answer" + train_iter.toString() + ":" + JSON.stringify(train_data.answer) + " ";
-    }
-    const system_prompt = "I am an assistant that extracts the groups from the user prompt to carry out differential gene expression. The final output must be in the following JSON with NO extra comments. The schema is as follows: " + JSON.stringify(Schema) + ' . "group1" and "group2" fields are compulsory. Both "group1" and "group2" consist of an array of filter variables. There are two kinds of filter variables: "Categorical" and "Numeric". "Categorical" variables are those variables which can have a fixed set of values e.g. gender, race. They are defined by the "CategoricalFilterTerm" which consists of "term" (a field from the sqlite3 db)  and "category" (a value of the field from the sqlite db).  "Numeric" variables are those which can have any numeric value. They are defined by "NumericFilterTerm" and contain  the subfields "term" (a field from the sqlite3 db), "start" an optional filter which is defined when a lower cutoff is defined in the user input for the numeric variable and "stop" an optional filter which is defined when a higher cutoff is defined in the user input for the numeric variable. ' + // May consider deprecating this natural language description after units tests are implemented
-    checkField(dataset_json.DatasetPrompt) + checkField(DE_ds[0].SystemPrompt) + "The sqlite db in plain language is as follows:\n" + dataset_db_output.rag_docs.join(",") + " training data is as follows:" + training_data + " Question: {" + prompt + "} answer:";
-    let response;
-    if (llm_backend_type == "SJ") {
-      response = await call_sj_llm(system_prompt, comp_model_name, apilink);
-    } else if (llm_backend_type == "ollama") {
-      response = await call_ollama(system_prompt, comp_model_name, apilink);
+    const DE_ds = dataset_json.charts.find((chart) => chart.type == "DE");
+    if (!DE_ds) throw "DE information is not present in the dataset file.";
+    if (DE_ds.TrainingData.length == 0) throw "No training data is provided for the DE agent.";
+    const training_data = formatTrainingExamples(DE_ds.TrainingData);
+    const system_prompt = "I am an assistant that extracts the groups from the user prompt to carry out differential gene expression. The final output must be in the following JSON with NO extra comments. The schema is as follows: " + JSON.stringify(Schema) + ' . "group1" and "group2" fields are compulsory. Both "group1" and "group2" consist of an array of filter variables. ' + FILTER_DESCRIPTION + checkField(dataset_json.DatasetPrompt) + checkField(DE_ds.SystemPrompt) + "The sqlite db in plain language is as follows:\n" + dataset_db_output.rag_docs.join(",") + " training data is as follows:" + training_data + " Question: {" + prompt + "} answer:";
+    const response = await route_to_appropriate_llm_provider(system_prompt, llm);
+    if (testing) {
+      return { action: "dge", response: JSON.parse(response) };
     } else {
-      throw "Unknown LLM backend";
+      return await validate_DE_response(response, ds, dataset_db_output.db_rows);
     }
-    return await validate_DE_response(response, ds, dataset_db_output.db_rows);
   } else {
     return { type: "html", html: "Differential gene expression not supported for this dataset" };
   }
@@ -447,9 +504,7 @@ function find_label(filter, db_rows) {
   }
   return label;
 }
-async function extract_summary_terms(prompt, llm_backend_type, comp_model_name, apilink, dataset_db, dataset_json, genedb, ds) {
-  const dataset_db_output = await parse_dataset_db(dataset_db);
-  const genes_list = await parse_geneset_db(genedb);
+async function extract_summary_terms(prompt, llm, dataset_db_output, dataset_json, genes_list, ds, testing) {
   const Schema = {
     $schema: "http://json-schema.org/draft-07/schema#",
     $ref: "#/definitions/SummaryType",
@@ -463,73 +518,37 @@ async function extract_summary_terms(prompt, llm_backend_type, comp_model_name,
             type: "array",
             items: { $ref: "#/definitions/FilterTerm" },
             description: "Optional simple filter terms"
-          }
-        },
-        required: ["term", "simpleFilter"],
-        additionalProperties: false
-      },
-      FilterTerm: {
-        anyOf: [{ $ref: "#/definitions/CategoricalFilterTerm" }, { $ref: "#/definitions/NumericFilterTerm" }]
-      },
-      CategoricalFilterTerm: {
-        type: "object",
-        properties: {
-          term: { type: "string", description: "Name of categorical term" },
-          category: { type: "string", description: "The category of the term" },
-          join: {
+          },
+          childType: {
             type: "string",
-            enum: ["and", "or"],
-            description: "join term to be used only when there there is more than one filter term and should be placed in the 2nd filter term describing how it connects to the 1st term"
+            enum: ["violin", "boxplot", "sampleScatter", "barchart"],
+            description: "Optional explicit child type requested by the user. If omitted, the logic of the data types picks the child type."
           }
         },
-        required: ["term", "category"],
+        required: ["term", "simpleFilter"],
         additionalProperties: false
       },
-      NumericFilterTerm: {
-        type: "object",
-        properties: {
-          term: { type: "string", description: "Name of numeric term" },
-          start: { type: "number", description: "start position (or lower limit) of numeric term" },
-          stop: { type: "number", description: "stop position (or upper limit) of numeric term" },
-          join: {
-            type: "string",
-            enum: ["and", "or"],
-            description: "join term to be used only when there there is more than one filter term and should be placed in the 2nd filter term describing how it connects to the 1st term"
-          }
-        },
-        required: ["term"],
-        additionalProperties: false
-      }
+      ...FILTER_TERM_DEFINITIONS
     }
   };
-  const words = prompt.replace(/[^a-zA-Z0-9\s]/g, "").split(/\s+/).map((str) => str.toLowerCase());
-  const common_genes = words.filter((item) => genes_list.includes(item));
-  const summary_ds = dataset_json.charts.filter((chart) => chart.type == "Summary");
-  if (summary_ds.length == 0) throw "Summary information is not present in the dataset file.";
-  if (summary_ds[0].TrainingData.length == 0) throw "No training data is provided for the summary agent.";
-  let train_iter = 0;
-  let training_data = "";
-  for (const train_data of summary_ds[0].TrainingData) {
-    train_iter += 1;
-    training_data += "Example question" + train_iter.toString() + ": " + train_data.question + " Example answer" + train_iter.toString() + ":" + JSON.stringify(train_data.answer) + " ";
-  }
-  let system_prompt = "I am an assistant that extracts the summary terms from user query. The final output must be in the following JSON format with NO extra comments. The JSON schema is as follows: " + JSON.stringify(Schema) + ' term and term2 (if present) should ONLY contain names of the fields from the sqlite db. The "simpleFilter" field is optional and should contain an array of JSON terms with which the dataset will be filtered. A variable simultaneously CANNOT be part of both "term"/"term2" and "simpleFilter". There are two kinds of filter variables: "Categorical" and "Numeric". "Categorical" variables are those variables which can have a fixed set of values e.g. gender, race. They are defined by the "CategoricalFilterTerm" which consists of "term" (a field from the sqlite3 db)  and "category" (a value of the field from the sqlite db).  "Numeric" variables are those which can have any numeric value. They are defined by "NumericFilterTerm" and contain  the subfields "term" (a field from the sqlite3 db), "start" an optional filter which is defined when a lower cutoff is defined in the user input for the numeric variable and "stop" an optional filter which is defined when a higher cutoff is defined in the user input for the numeric variable. ' + // May consider deprecating this natural language description after unit tests are implemented
-  checkField(dataset_json.DatasetPrompt) + checkField(summary_ds[0].SystemPrompt) + "\n The DB content is as follows: " + dataset_db_output.rag_docs.join(",") + " training data is as follows:" + training_data;
+  const common_genes = extractGenesFromPrompt(prompt, genes_list);
+  const summary_ds = dataset_json.charts.find((chart) => chart.type == "Summary");
+  if (!summary_ds) throw "Summary information is not present in the dataset file.";
+  if (summary_ds.TrainingData.length == 0) throw "No training data is provided for the summary agent.";
+  const training_data = formatTrainingExamples(summary_ds.TrainingData);
+  let system_prompt = "I am an assistant that extracts the summary terms from user query. The final output must be in the following JSON format with NO extra comments. The JSON schema is as follows: " + JSON.stringify(Schema) + ' term and term2 (if present) should ONLY contain names of the fields from the sqlite db. The "simpleFilter" field is optional and should contain an array of JSON terms with which the dataset will be filtered. ' + FILTER_DESCRIPTION + checkField(dataset_json.DatasetPrompt) + checkField(summary_ds.SystemPrompt) + "\n The DB content is as follows: " + dataset_db_output.rag_docs.join(",") + " training data is as follows:" + training_data;
   if (dataset_json.hasGeneExpression) {
     if (common_genes.length > 0) {
       system_prompt += "\n List of relevant genes are as follows (separated by comma(,)):" + common_genes.join(",");
     }
   }
   system_prompt += " Question: {" + prompt + "} answer:";
-  let response;
-  if (llm_backend_type == "SJ") {
-    response = await call_sj_llm(system_prompt, comp_model_name, apilink);
-  } else if (llm_backend_type == "ollama") {
-    response = await call_ollama(system_prompt, comp_model_name, apilink);
+  const response = await route_to_appropriate_llm_provider(system_prompt, llm);
+  if (testing) {
+    return { action: "summary", response: JSON.parse(response) };
   } else {
-    throw "Unknown LLM backend";
+    return validate_summary_response(response, common_genes, dataset_json, ds);
   }
-  return validate_summary_response(response, common_genes, dataset_json, ds);
 }
 function validate_summary_response(response, common_genes, dataset_json, ds) {
   const response_type = JSON.parse(response);
@@ -542,6 +561,10 @@ function validate_summary_response(response, common_genes, dataset_json, ds) {
     html += term1_validation.html;
   } else {
     pp_plot_json.term = term1_validation.term_type;
+    if (term1_validation.category == "float" || term1_validation.category == "integer") {
+      pp_plot_json.term.q = { mode: "continuous" };
+    }
+    pp_plot_json.category = term1_validation.category;
   }
   if (response_type.term2) {
     const term2_validation = validate_term(response_type.term2, common_genes, dataset_json, ds);
@@ -549,6 +572,118 @@ function validate_summary_response(response, common_genes, dataset_json, ds) {
       html += term2_validation.html;
     } else {
       pp_plot_json.term2 = term2_validation.term_type;
+      if (term2_validation.category == "float" || term2_validation.category == "integer") {
+        pp_plot_json.term2.q = { mode: "continuous" };
+      }
+      pp_plot_json.category2 = term2_validation.category;
+    }
+  }
+  const llmChildType = response_type.childType && ["violin", "boxplot", "sampleScatter", "barchart"].includes(response_type.childType) ? response_type.childType : void 0;
+  const resolved = resolveChildType(pp_plot_json.category, pp_plot_json.category2, llmChildType);
+  if (resolved.error) {
+    html += resolved.error;
+  } else {
+    pp_plot_json.childType = resolved.childType;
+    if (resolved.bothNumeric && (resolved.childType == "violin" || resolved.childType == "boxplot")) {
+      pp_plot_json.term2.q = { mode: "discrete" };
+    }
+  }
+  delete pp_plot_json.category;
+  if (pp_plot_json.category2) delete pp_plot_json.category2;
+  if (response_type.simpleFilter && response_type.simpleFilter.length > 0) {
+    const validated_filters = validate_filter(response_type.simpleFilter, ds, "");
+    if (validated_filters.html.length > 0) {
+      html += validated_filters.html;
+    } else {
+      pp_plot_json.filter = validated_filters.simplefilter;
+    }
+  }
+  if (html.length > 0) {
+    return { type: "html", html };
+  } else {
+    return { type: "plot", plot: pp_plot_json };
+  }
+}
+async function extract_matrix_search_terms_from_query(prompt, llm, dataset_db_output, dataset_json, genes_list, ds, testing) {
+  const Schema = {
+    $schema: "http://json-schema.org/draft-07/schema#",
+    $ref: "#/definitions/MatrixType",
+    definitions: {
+      MatrixType: {
+        type: "object",
+        properties: {
+          terms: {
+            type: "array",
+            items: { type: "string" },
+            description: "Names of dictionary/clinical terms to include as rows in the matrix"
+          },
+          geneNames: {
+            type: "array",
+            items: { type: "string" },
+            description: "Names of genes to include as gene variant rows in the matrix"
+          },
+          simpleFilter: {
+            type: "array",
+            items: { $ref: "#/definitions/FilterTerm" },
+            description: "Optional simple filter terms to restrict the sample set"
+          }
+        },
+        additionalProperties: false
+      },
+      ...FILTER_TERM_DEFINITIONS
+    }
+  };
+  const common_genes = extractGenesFromPrompt(prompt, genes_list);
+  const matrix_ds = dataset_json.charts.filter((chart) => chart.type == "Matrix");
+  console.log("matrix_ds", matrix_ds);
+  console.log("dataset_json.charts", dataset_json.charts);
+  if (matrix_ds.length == 0) throw "Matrix information is not present in the dataset file.";
+  if (matrix_ds[0].TrainingData.length == 0) throw "No training data is provided for the matrix agent.";
+  const training_data = formatTrainingExamples(matrix_ds[0].TrainingData);
+  let system_prompt = "I am an assistant that extracts terms and gene names from the user query to create a matrix plot. A matrix plot displays multiple genes and/or clinical variables across samples in a grid layout. The final output must be in the following JSON format with NO extra comments. The JSON schema is as follows: " + JSON.stringify(Schema) + ' The "terms" field should ONLY contain names of clinical/dictionary fields from the sqlite db. The "geneNames" field should ONLY contain gene names. At least one of "terms" or "geneNames" must be provided. The "simpleFilter" field is optional and should contain an array of JSON terms with which the dataset will be filtered. ' + FILTER_DESCRIPTION + checkField(dataset_json.DatasetPrompt) + checkField(matrix_ds[0].SystemPrompt) + "\n The DB content is as follows: " + dataset_db_output.rag_docs.join(",") + " training data is as follows:" + training_data;
+  if (dataset_json.hasGeneExpression && common_genes.length > 0) {
+    system_prompt += "\n List of relevant genes are as follows (separated by comma(,)):" + common_genes.join(",");
+  }
+  system_prompt += " Question: {" + prompt + "} answer:";
+  const response = await route_to_appropriate_llm_provider(system_prompt, llm);
+  if (testing) {
+    return { action: "matrix", response: JSON.parse(response) };
+  } else {
+    return validate_matrix_response(response, common_genes, dataset_json, ds);
+  }
+}
+function validate_matrix_response(response, common_genes, dataset_json, ds) {
+  const response_type = JSON.parse(response);
+  const pp_plot_json = { chartType: "matrix" };
+  let html = "";
+  if (response_type.html) html = response_type.html;
+  if ((!response_type.terms || response_type.terms.length == 0) && (!response_type.geneNames || response_type.geneNames.length == 0)) {
+    html += "At least one clinical term or gene name is required for a matrix plot";
+  }
+  const twLst = [];
+  if (response_type.terms && Array.isArray(response_type.terms)) {
+    for (const t of response_type.terms) {
+      const term = ds.cohort.termdb.q.termjsonByOneid(t);
+      if (!term) {
+        html += "invalid term id:" + t + " ";
+      } else {
+        twLst.push({ id: term.id });
+      }
+    }
+  }
+  if (response_type.geneNames && Array.isArray(response_type.geneNames)) {
+    for (const g of response_type.geneNames) {
+      const gene_hits = common_genes.filter((gene) => gene == g.toLowerCase());
+      if (gene_hits.length == 0) {
+        html += "invalid gene name:" + g + " ";
+      } else {
+        const geneName = g.toUpperCase();
+        if (dataset_json.hasGeneExpression) {
+          twLst.push({ term: { gene: geneName, type: "geneExpression" } });
+        } else {
+          twLst.push({ term: { gene: geneName, name: geneName, type: "geneVariant" } });
+        }
+      }
     }
   }
   if (response_type.simpleFilter && response_type.simpleFilter.length > 0) {
@@ -562,12 +697,14 @@ function validate_summary_response(response, common_genes, dataset_json, ds) {
   if (html.length > 0) {
     return { type: "html", html };
   } else {
+    pp_plot_json.termgroups = [{ name: "", lst: twLst }];
     return { type: "plot", plot: pp_plot_json };
   }
 }
 function validate_term(response_term, common_genes, dataset_json, ds) {
   let html = "";
   let term_type;
+  let category = "";
   const term = ds.cohort.termdb.q.termjsonByOneid(response_term);
   if (!term) {
     const gene_hits = common_genes.filter((gene) => gene == response_term.toLowerCase());
@@ -576,14 +713,16 @@ function validate_term(response_term, common_genes, dataset_json, ds) {
     } else {
       if (dataset_json.hasGeneExpression) {
         term_type = { term: { gene: response_term.toUpperCase(), type: "geneExpression" } };
+        category = "float";
       } else {
         html += "Dataset does not support gene expression";
       }
     }
   } else {
     term_type = { id: term.id };
+    category = term.type;
   }
-  return { term_type, html };
+  return { term_type, html, category };
 }
 function countOccurrences(str, word) {
   if (word === "") return 0;
@@ -768,5 +907,7 @@ function parse_db_rows(db_row) {
   return output_string;
 }
 export {
-  api
+  api,
+  readJSONFile,
+  run_chat_pipeline
 };