@sjcrh/proteinpaint-rust 2.167.0 → 2.170.0

package/src/aichatbot.rs

@@ -1,7 +1,7 @@
 // Syntax: cd .. && cargo build --release && time cat ~/sjpp/test.txt | target/release/aichatbot
 #![allow(non_snake_case)]
-use anyhow::Result;
-use json::JsonValue;
+//use anyhow::Result;
+//use json::JsonValue;
 use r2d2_sqlite::SqliteConnectionManager;
 use rig::agent::AgentBuilder;
 use rig::completion::Prompt;
@@ -11,40 +11,53 @@ use schemars::JsonSchema;
 use serde_json::{Map, Value, json};
 use std::collections::HashMap;
 use std::fs;
-use std::io;
-use std::path::Path;
-mod ollama; // Importing custom rig module for invoking ollama server
-mod sjprovider; // Importing custom rig module for invoking SJ GPU server
-
-mod test_ai; // Test examples for AI chatbot
 
 // Struct for intaking data from dataset json
 #[derive(PartialEq, Debug, Clone, schemars::JsonSchema, serde::Serialize, serde::Deserialize)]
 pub struct AiJsonFormat {
-    hasGeneExpression: bool,
-    db: String,     // Dataset db
-    genedb: String, // Gene db
-    charts: Vec<Charts>,
+    pub hasGeneExpression: bool,
+    pub db: String,     // Dataset db
+    pub genedb: String, // Gene db
+    pub charts: Vec<TrainTestData>,
+}
+
+#[derive(PartialEq, Debug, Clone, schemars::JsonSchema, serde::Serialize, serde::Deserialize)]
+pub struct TrainTestData {
+    pub r#type: String,
+    pub SystemPrompt: String,
+    pub TrainingData: Vec<QuestionAnswer>,
+    pub TestData: Vec<QuestionAnswer>,
 }
 
 #[derive(PartialEq, Debug, Clone, schemars::JsonSchema, serde::Serialize, serde::Deserialize)]
-enum Charts {
-    // More chart types will be added here later
-    Summary(TrainTestData),
-    DE(TrainTestData),
+pub struct QuestionAnswer {
+    pub question: String,
+    pub answer: AnswerFormat,
 }
 
 #[derive(PartialEq, Debug, Clone, schemars::JsonSchema, serde::Serialize, serde::Deserialize)]
-struct TrainTestData {
-    SystemPrompt: String,
-    TrainingData: Vec<QuestionAnswer>,
-    TestData: Vec<QuestionAnswer>,
+pub enum AnswerFormat {
+    #[allow(non_camel_case_types)]
+    summary_type(SummaryType),
+    #[allow(non_camel_case_types)]
+    DE_type(DEType),
 }
 
 #[derive(PartialEq, Debug, Clone, schemars::JsonSchema, serde::Serialize, serde::Deserialize)]
-struct QuestionAnswer {
-    question: String,
-    answer: String,
+pub struct DEType {
+    action: String,
+    DE_output: DETerms,
+}
+
+#[derive(PartialEq, Debug, Clone, schemars::JsonSchema, serde::Serialize, serde::Deserialize)]
+pub struct DETerms {
+    group1: GroupType,
+    group2: GroupType,
+}
+
+#[derive(PartialEq, Debug, Clone, schemars::JsonSchema, serde::Serialize, serde::Deserialize)]
+pub struct GroupType {
+    name: String,
 }
 
 #[allow(non_camel_case_types)]
@@ -56,180 +69,188 @@ pub enum llm_backend {
 
 #[derive(Debug, JsonSchema)]
 #[allow(dead_code)]
-struct OutputJson {
+pub struct OutputJson {
     pub answer: String,
 }
 
-#[tokio::main]
-async fn main() -> Result<()> {
-    let mut input = String::new();
-    match io::stdin().read_line(&mut input) {
-        // Accepting the piped input from nodejs (or command line from testing)
-        Ok(_n) => {
-            let input_json = json::parse(&input);
-            match input_json {
-                Ok(json_string) => {
-                    //println!("json_string:{}", json_string);
-                    let user_input_json: &JsonValue = &json_string["user_input"];
-                    let user_input: &str;
-                    match user_input_json.as_str() {
-                        Some(inp) => user_input = inp,
-                        None => panic!("user_input field is missing in input json"),
-                    }
-
-                    if user_input.len() == 0 {
-                        panic!("The user input is empty");
-                    }
-
-                    let tpmasterdir_json: &JsonValue = &json_string["tpmasterdir"];
-                    let tpmasterdir: &str;
-                    match tpmasterdir_json.as_str() {
-                        Some(inp) => tpmasterdir = inp,
-                        None => panic!("tpmasterdir not found"),
-                    }
-
-                    let binpath_json: &JsonValue = &json_string["binpath"];
-                    let binpath: &str;
-                    match binpath_json.as_str() {
-                        Some(inp) => binpath = inp,
-                        None => panic!("binpath not found"),
-                    }
-
-                    let ai_json_file_json: &JsonValue = &json_string["aifiles"];
-                    let ai_json_file: String;
-                    match ai_json_file_json.as_str() {
-                        Some(inp) => ai_json_file = String::from(binpath) + &"/../../" + &inp,
-                        None => {
-                            panic!("ai json file not found")
-                        }
-                    }
-
-                    let ai_json_file = Path::new(&ai_json_file);
-                    let ai_json_file_path;
-                    let current_dir = std::env::current_dir().unwrap();
-                    match ai_json_file.canonicalize() {
-                        Ok(p) => ai_json_file_path = p,
-                        Err(_) => {
-                            panic!(
-                                "AI JSON file path not found:{:?}, current directory:{:?}",
-                                ai_json_file, current_dir
-                            )
-                        }
-                    }
-
-                    // Read the file
-                    let ai_data = fs::read_to_string(ai_json_file_path).unwrap();
-
-                    // Parse the JSON data
-                    let ai_json: AiJsonFormat =
-                        serde_json::from_str(&ai_data).expect("AI JSON file does not have the correct format");
-
-                    let genedb = String::from(tpmasterdir) + &"/" + &ai_json.genedb;
-                    let dataset_db = String::from(tpmasterdir) + &"/" + &ai_json.db;
-
-                    let apilink_json: &JsonValue = &json_string["apilink"];
-                    let apilink: &str;
-                    match apilink_json.as_str() {
-                        Some(inp) => apilink = inp,
-                        None => panic!("apilink field is missing in input json"),
-                    }
-
-                    let comp_model_name_json: &JsonValue = &json_string["comp_model_name"];
-                    let comp_model_name: &str;
-                    match comp_model_name_json.as_str() {
-                        Some(inp) => comp_model_name = inp,
-                        None => panic!("comp_model_name field is missing in input json"),
-                    }
-
-                    let embedding_model_name_json: &JsonValue = &json_string["embedding_model_name"];
-                    let embedding_model_name: &str;
-                    match embedding_model_name_json.as_str() {
-                        Some(inp) => embedding_model_name = inp,
-                        None => panic!("embedding_model_name field is missing in input json"),
-                    }
-
-                    let llm_backend_name_json: &JsonValue = &json_string["llm_backend_name"];
-                    let llm_backend_name: &str;
-                    match llm_backend_name_json.as_str() {
-                        Some(inp) => llm_backend_name = inp,
-                        None => panic!("llm_backend_name field is missing in input json"),
-                    }
-
-                    let llm_backend_type: llm_backend;
-                    let mut final_output: Option<String> = None;
-                    let temperature: f64 = 0.01;
-                    let max_new_tokens: usize = 512;
-                    let top_p: f32 = 0.95;
-                    let testing = false; // This variable is always false in production, this is true in test_ai.rs for testing code
-                    if llm_backend_name != "ollama" && llm_backend_name != "SJ" {
-                        panic!(
-                            "This code currently supports only Ollama and SJ provider. llm_backend_name must be \"ollama\" or \"SJ\""
-                        );
-                    } else if llm_backend_name == "ollama".to_string() {
-                        llm_backend_type = llm_backend::Ollama();
-                        // Initialize Ollama client
-                        let ollama_client = ollama::Client::builder()
-                            .base_url(apilink)
-                            .build()
-                            .expect("Ollama server not found");
-                        let embedding_model = ollama_client.embedding_model(embedding_model_name);
-                        let comp_model = ollama_client.completion_model(comp_model_name);
-                        final_output = run_pipeline(
-                            user_input,
-                            comp_model,
-                            embedding_model,
-                            llm_backend_type,
-                            temperature,
-                            max_new_tokens,
-                            top_p,
-                            &dataset_db,
-                            &genedb,
-                            &ai_json,
-                            testing,
-                        )
-                        .await;
-                    } else if llm_backend_name == "SJ".to_string() {
-                        llm_backend_type = llm_backend::Sj();
-                        // Initialize Sj provider client
-                        let sj_client = sjprovider::Client::builder()
-                            .base_url(apilink)
-                            .build()
-                            .expect("SJ server not found");
-                        let embedding_model = sj_client.embedding_model(embedding_model_name);
-                        let comp_model = sj_client.completion_model(comp_model_name);
-                        final_output = run_pipeline(
-                            user_input,
-                            comp_model,
-                            embedding_model,
-                            llm_backend_type,
-                            temperature,
-                            max_new_tokens,
-                            top_p,
-                            &dataset_db,
-                            &genedb,
-                            &ai_json,
-                            testing,
-                        )
-                        .await;
-                    }
-
-                    match final_output {
-                        Some(fin_out) => {
-                            println!("final_output:{:?}", fin_out.replace("\\", ""));
-                        }
-                        None => {
-                            println!("final_output:{{\"{}\":\"{}\"}}", "action", "unknown");
-                        }
-                    }
-                }
-                Err(error) => println!("Incorrect json:{}", error),
-            }
-        }
-        Err(error) => println!("Piping error: {}", error),
-    }
-    Ok(())
-}
+//#[tokio::main]
+//async fn main() -> Result<()> {
+//    let mut input = String::new();
+//    match io::stdin().read_line(&mut input) {
+//        // Accepting the piped input from nodejs (or command line from testing)
+//        Ok(_n) => {
+//            let input_json = json::parse(&input);
+//            match input_json {
+//                Ok(json_string) => {
+//                    //println!("json_string:{}", json_string);
+//                    let user_input_json: &JsonValue = &json_string["user_input"];
+//                    let user_input: &str;
+//                    match user_input_json.as_str() {
+//                        Some(inp) => user_input = inp,
+//                        None => panic!("user_input field is missing in input json"),
+//                    }
+//                    let dataset_db_json: &JsonValue = &json_string["dataset_db"];
+//                    let dataset_db_str: &str;
+//                    match dataset_db_json.as_str() {
+//                        Some(inp) => dataset_db_str = inp,
+//                        None => panic!("dataset_db field is missing in input json"),
+//                    }
+//                    let genedb_json: &JsonValue = &json_string["genedb"];
+//                    let genedb_str: &str;
+//                    match genedb_json.as_str() {
+//                        Some(inp) => genedb_str = inp,
+//                        None => panic!("genedb field is missing in input json"),
+//                    }
+//                    let aiRoute_json: &JsonValue = &json_string["aiRoute"];
+//                    let aiRoute_str: &str;
+//                    match aiRoute_json.as_str() {
+//                        Some(inp) => aiRoute_str = inp,
+//                        None => panic!("aiRoute field is missing in input json"),
+//                    }
+//                    if user_input.len() == 0 {
+//                        panic!("The user input is empty");
+//                    }
+//                    let tpmasterdir_json: &JsonValue = &json_string["tpmasterdir"];
+//                    let tpmasterdir: &str;
+//                    match tpmasterdir_json.as_str() {
+//                        Some(inp) => tpmasterdir = inp,
+//                        None => panic!("tpmasterdir not found"),
+//                    }
+//                    let binpath_json: &JsonValue = &json_string["binpath"];
+//                    let binpath: &str;
+//                    match binpath_json.as_str() {
+//                        Some(inp) => binpath = inp,
+//                        None => panic!("binpath not found"),
+//                    }
+//                    let ai_json_file_json: &JsonValue = &json_string["aifiles"];
+//                    let ai_json_file: String;
+//                    match ai_json_file_json.as_str() {
+//                        Some(inp) => ai_json_file = String::from(binpath) + &"/../../" + &inp,
+//                        None => {
+//                            panic!("ai json file not found")
+//                        }
+//                    }
+//                    let ai_json_file = Path::new(&ai_json_file);
+//                    let ai_json_file_path;
+//                    let current_dir = std::env::current_dir().unwrap();
+//                    match ai_json_file.canonicalize() {
+//                        Ok(p) => ai_json_file_path = p,
+//                        Err(_) => {
+//                            panic!(
+//                                "AI JSON file path not found:{:?}, current directory:{:?}",
+//                                ai_json_file, current_dir
+//                            )
+//                        }
+//                    }
+//                    // Read the file
+//                    let ai_data = fs::read_to_string(ai_json_file_path).unwrap();
+//                    // Parse the JSON data
+//                    let ai_json: AiJsonFormat =
+//                        serde_json::from_str(&ai_data).expect("AI JSON file does not have the correct format");
+//                    let genedb = String::from(tpmasterdir) + &"/" + &genedb_str;
+//                    let dataset_db = String::from(tpmasterdir) + &"/" + &dataset_db_str;
+//                    let airoute = String::from(binpath) + &"/../../" + &aiRoute_str;
+//                    let apilink_json: &JsonValue = &json_string["apilink"];
+//                    let apilink: &str;
+//                    match apilink_json.as_str() {
+//                        Some(inp) => apilink = inp,
+//                        None => panic!("apilink field is missing in input json"),
+//                    }
+//                    let comp_model_name_json: &JsonValue = &json_string["comp_model_name"];
+//                    let comp_model_name: &str;
+//                    match comp_model_name_json.as_str() {
+//                        Some(inp) => comp_model_name = inp,
+//                        None => panic!("comp_model_name field is missing in input json"),
+//                    }
+//                    let embedding_model_name_json: &JsonValue = &json_string["embedding_model_name"];
+//                    let embedding_model_name: &str;
+//                    match embedding_model_name_json.as_str() {
+//                        Some(inp) => embedding_model_name = inp,
+//                        None => panic!("embedding_model_name field is missing in input json"),
+//                    }
+//                    let llm_backend_name_json: &JsonValue = &json_string["llm_backend_name"];
+//                    let llm_backend_name: &str;
+//                    match llm_backend_name_json.as_str() {
+//                        Some(inp) => llm_backend_name = inp,
+//                        None => panic!("llm_backend_name field is missing in input json"),
+//                    }
+//                    let llm_backend_type: llm_backend;
+//                    let mut final_output: Option<String> = None;
+//                    let temperature: f64 = 0.01;
+//                    let max_new_tokens: usize = 512;
+//                    let top_p: f32 = 0.95;
+//                    let testing = false; // This variable is always false in production, this is true in test_ai.rs for testing code
+//                    if llm_backend_name != "ollama" && llm_backend_name != "SJ" {
+//                        panic!(
+//                           "This code currently supports only Ollama and SJ provider. llm_backend_name must be \"ollama\" or \"SJ\""
+//                       );
+//                    } else if llm_backend_name == "ollama".to_string() {
+//                        llm_backend_type = llm_backend::Ollama();
+//                        // Initialize Ollama client
+//                        let ollama_client = ollama::Client::builder()
+//                            .base_url(apilink)
+//                            .build()
+//                            .expect("Ollama server not found");
+//                        let embedding_model = ollama_client.embedding_model(embedding_model_name);
+//                        let comp_model = ollama_client.completion_model(comp_model_name);
+//                        final_output = run_pipeline(
+//                            user_input,
+//                            comp_model,
+//                            embedding_model,
+//                            llm_backend_type,
+//                            temperature,
+//                            max_new_tokens,
+//                            top_p,
+//                            &dataset_db,
+//                            &genedb,
+//                            &ai_json,
+//                            &airoute,
+//                            testing,
+//                        )
+//                        .await;
+//                    } else if llm_backend_name == "SJ".to_string() {
+//                        llm_backend_type = llm_backend::Sj();
+//                        // Initialize Sj provider client
+//                        let sj_client = sjprovider::Client::builder()
+//                            .base_url(apilink)
+//                            .build()
+//                            .expect("SJ server not found");
+//                        let embedding_model = sj_client.embedding_model(embedding_model_name);
+//                        let comp_model = sj_client.completion_model(comp_model_name);
+//                        final_output = run_pipeline(
+//                            user_input,
+//                            comp_model,
+//                            embedding_model,
+//                            llm_backend_type,
+//                            temperature,
+//                            max_new_tokens,
+//                            top_p,
+//                            &dataset_db,
+//                            &genedb,
+//                            &ai_json,
+//                            &airoute,
+//                            testing,
+//                        )
+//                        .await;
+//                    }
+//                    match final_output {
+//                        Some(fin_out) => {
+//                            println!("final_output:{:?}", fin_out.replace("\\", ""));
+//                        }
+//                        None => {
+//                            println!("final_output:{{\"{}\":\"{}\"}}", "action", "unknown");
+//                        }
+//                    }
+//                }
+//                Err(error) => println!("Incorrect json:{}", error),
+//            }
+//        }
+//        Err(error) => println!("Piping error: {}", error),
+//    }
+//    Ok(())
+//}
 
+#[allow(dead_code)]
 pub async fn run_pipeline(
     user_input: &str,
     comp_model: impl rig::completion::CompletionModel + 'static,
@@ -241,6 +262,7 @@ pub async fn run_pipeline(
     dataset_db: &str,
     genedb: &str,
     ai_json: &AiJsonFormat,
+    ai_route: &str,
     testing: bool,
 ) -> Option<String> {
     let mut classification: String = classify_query_by_dataset_type(
@@ -251,6 +273,7 @@ pub async fn run_pipeline(
         temperature,
         max_new_tokens,
         top_p,
+        ai_route,
     )
     .await;
     classification = classification.replace("\"", "");
@@ -373,104 +396,40 @@ pub async fn run_pipeline(
     Some(final_output)
 }
 
-async fn classify_query_by_dataset_type(
+#[allow(dead_code)]
+pub async fn classify_query_by_dataset_type(
     user_input: &str,
     comp_model: impl rig::completion::CompletionModel + 'static,
-    embedding_model: impl rig::embeddings::EmbeddingModel + 'static,
+    _embedding_model: impl rig::embeddings::EmbeddingModel + 'static,
     llm_backend_type: &llm_backend,
     temperature: f64,
     max_new_tokens: usize,
     top_p: f32,
+    ai_route: &str,
 ) -> String {
-    // Create a string to hold the file contents
-    let contents = String::from("SNV/SNP or point mutations nucleotide mutations are very common forms of mutations which can often give rise to genetic diseases such as cancer, Alzheimer's disease etc. They can be duw to substitution of nucleotide, or insertion or deletion of a nucleotide. Indels are multi-nucleotide insertion/deletion/substitutions. Complex indels are indels where insertion and deletion have happened in the same genomic locus. Every genomic sample from each patient has its own set of mutations therefore requiring personalized treatment. 
-
-If a ProteinPaint dataset contains SNV/Indel/SV data then return JSON with single key, 'snv_indel'.
-
----
-
-Copy number variation (CNV) is a phenomenon in which sections of the genome are repeated and the number of repeats in the genome varies between individuals.[1] Copy number variation is a special type of structural variation: specifically, it is a type of duplication or deletion event that affects a considerable number of base pairs.
+    // Read the file
+    let ai_route_data = fs::read_to_string(ai_route).unwrap();
 
-If a ProteinPaint dataset contains copy number variation data then return JSON with single key, 'cnv'.
-
----
-
-Structural variants/fusions (SV) are genomic mutations when eith a DNA region is translocated or copied to an entirely different genomic locus. In case of transcriptomic data, when RNA is fused from two different genes its called a gene fusion.
-
-If a ProteinPaint dataset contains structural variation or gene fusion data then return JSON with single key, 'sv_fusion'.
----
-
-Hierarchical clustering of gene expression is an unsupervised learning technique where several number of relevant genes and the samples are clustered so as to determine (previously unknown) cohorts of samples (or patients) or structure in data. It is very commonly used to determine subtypes of a particular disease based on RNA sequencing data. 
-
-If a ProteinPaint dataset contains hierarchical data then return JSON with single key, 'hierarchical'.
-
----
-
-Differential Gene Expression (DGE or DE) is a technique where the most upregulated (or highest) and downregulated (or lowest) genes between two cohorts of samples (or patients) are determined from a pool of THOUSANDS of genes. Differential gene expression CANNOT be computed for a SINGLE gene. A volcano plot is shown with fold-change in the x-axis and adjusted p-value on the y-axis. So, the upregulated and downregulared genes are on opposite sides of the graph and the most significant genes (based on adjusted p-value) is on the top of the graph. Following differential gene expression generally GeneSet Enrichment Analysis (GSEA) is carried out where based on the genes and their corresponding fold changes the upregulation/downregulation of genesets (or pathways) is determined.
-
-Sample Query1: \"Which gene has the highest expression between the two genders\"
-Sample Answer1: { \"answer\": \"dge\" }
-
-Sample Query2: \"Which gene has the lowest expression between the two races\"
-Sample Answer2: { \"answer\": \"dge\" }
-
-Sample Query1: \"Which genes are the most upregulated genes between group A and group B\"
-Sample Answer1: { \"answer\": \"dge\" }
-
-Sample Query3: \"Which gene are overexpressed between male and female\"
-Sample Answer3: { \"answer\": \"dge\" }
-
-Sample Query4: \"Which gene are housekeeping genes between male and female\"
-Sample Answer4: { \"answer\": \"dge\" } 
-
-
-If a ProteinPaint dataset contains differential gene expression data then return JSON with single key, 'dge'.
-
----
-
-Survival analysis (also called time-to-event analysis or duration analysis) is a branch of statistics aimed at analyzing the duration of time from a well-defined time origin until one or more events happen, called survival times or duration times. In other words, in survival analysis, we are interested in a certain event and want to analyze the time until the event happens. Generally in survival analysis survival rates between two (or more) cohorts of patients  is compared. 
-
-There are two main methods of survival analysis:
-
-1) Kaplan-Meier (HM) analysis is a univariate test that only takes into account a single categorical variable.
-2) Cox proportional hazards model (coxph) is a multivariate test that can take into account multiple variables.
-
-   The hazard ratio (HR) is an indicator of the effect of the stimulus (e.g. drug dose, treatment) between two cohorts of patients.
-   HR = 1: No effect
-   HR < 1: Reduction in the hazard
-   HR > 1: Increase in Hazard
-
-Sample Query1: \"Compare survival rates between group A and B\"
-Sample Answer1: { \"answer\": \"survival\" }
-
-
-If a ProteinPaint dataset contains survival data then return JSON with single key, 'survival'.
-
----
-
-Next generation sequencing reads (NGS) are mapped to a human genome using alignment algorithm such as burrows-wheelers alignment algorithm. Then these reads are called using variant calling algorithms such as GATK (Genome Analysis Toolkit). However this type of analysis is too compute intensive and beyond the scope of visualization software such as ProteinPaint.
-
-If a user query asks about variant calling or mapping reads then JSON with single key, 'variant_calling'.
-
----
-
-Summary plot in ProteinPaint shows the various facets of the datasets. Show expression of a SINGLE gene or compare the expression of a SINGLE gene across two different cohorts defined by the user. It may show all the samples according to their respective diagnosis or subtypes of cancer. It is also useful for comparing and correlating different clinical variables. It can show all possible distributions, frequency of a category, overlay, correlate or cross-tabulate with another variable on top of it. If a user query asks about a SINGLE gene expression or correlating clinical variables then return JSON with single key, 'summary'.
-
-Sample Query1: \"Show all fusions for patients with age less than 30\"
-Sample Answer1: { \"answer\": \"summary\" }
-
-Sample Query2: \"List all molecular subtypes of leukemia\"
-Sample Answer2: { \"answer\": \"summary\" } 
-
-Sample Query3: \"is tp53 expression higher in men than women ?\"
-Sample Answer3: { \"answer\": \"summary\" }
-
-Sample Query4: \"Compare ATM expression between races for women greater than 80yrs\"
-Sample Answer4: { \"answer\": \"summary\" }
+    // Parse the JSON data
+    let ai_json: Value = serde_json::from_str(&ai_route_data).expect("AI JSON file does not have the correct format");
 
+    // Create a string to hold the file contents
+    let mut contents = String::from("");
+
+    if let Some(object) = ai_json.as_object() {
+        contents = object["general"].to_string();
+        for (key, value) in object {
+            if key != "general" {
+                contents += &value.as_str().unwrap();
+                contents += "---"; // Adding delimiter
+            }
+        }
+    }
 
-If a query does not match any of the fields described above, then return JSON with single key, 'none'
-");
+    // Removing the last "---" characters
+    contents.pop();
+    contents.pop();
+    contents.pop();
 
     // Split the contents by the delimiter "---"
     let parts: Vec<&str> = contents.split("---").collect();
@@ -501,18 +460,18 @@ If a query does not match any of the fields described above, then return JSON wi
         rag_docs.push(part.trim().to_string())
     }
 
-    //let top_k: usize = 3;
+    //let top_k: usize = 3; // Embedding model not used currently
     // Create embeddings and add to vector store
-    let embeddings = EmbeddingsBuilder::new(embedding_model.clone())
-        .documents(rag_docs)
-        .expect("Reason1")
-        .build()
-        .await
-        .unwrap();
+    //let embeddings = EmbeddingsBuilder::new(embedding_model.clone())
+    //    .documents(rag_docs)
+    //    .expect("Reason1")
+    //    .build()
+    //    .await
+    //    .unwrap();
 
-    // Create vector store
-    let mut vector_store = InMemoryVectorStore::<String>::default();
-    InMemoryVectorStore::add_documents(&mut vector_store, embeddings);
+    //// Create vector store
+    //let mut vector_store = InMemoryVectorStore::<String>::default();
+    //InMemoryVectorStore::add_documents(&mut vector_store, embeddings);
 
     // Create RAG agent
     let agent = AgentBuilder::new(comp_model).preamble(&(String::from("Generate classification for the user query into summary, dge, hierarchical, snv_indel, cnv, variant_calling, sv_fusion and none categories. Return output in JSON with ALWAYS a single word answer { \"answer\": \"dge\" }, that is 'summary' for summary plot, 'dge' for differential gene expression, 'hierarchical' for hierarchical clustering, 'snv_indel' for SNV/Indel, 'cnv' for CNV and 'sv_fusion' for SV/fusion, 'variant_calling' for variant calling, 'surivial' for survival data, 'none' for none of the previously described categories. The summary plot list and summarizes the cohort of patients according to the user query. The answer should always be in lower case\n The options are as follows:\n") + &contents + "\nQuestion= {question} \nanswer")).temperature(temperature).additional_params(additional).build();
@@ -580,6 +539,7 @@ struct DEOutput {
     group2: Group,
 }
 
+#[allow(dead_code)]
 #[allow(non_snake_case)]
 async fn extract_DE_search_terms_from_query(
     user_input: &str,
@@ -853,7 +813,7 @@ async fn parse_dataset_db(db: &str) -> (Vec<String>, Vec<DbRows>) {
     (rag_docs, db_vec)
 }
 
-async fn extract_summary_information(
+pub async fn extract_summary_information(
     user_input: &str,
     comp_model: impl rig::completion::CompletionModel + 'static,
     _embedding_model: impl rig::embeddings::EmbeddingModel + 'static,
@@ -911,8 +871,8 @@ async fn extract_summary_information(
 
     let mut summary_data_check: Option<TrainTestData> = None;
     for chart in ai_json.charts.clone() {
-        if let Charts::Summary(traindata) = chart {
-            summary_data_check = Some(traindata);
+        if chart.r#type == "Summary" {
+            summary_data_check = Some(chart);
             break;
         }
     }
@@ -922,17 +882,23 @@ async fn extract_summary_information(
             let mut training_data: String = String::from("");
             let mut train_iter = 0;
             for ques_ans in summary_data.TrainingData {
-                train_iter += 1;
-                training_data += "Example question";
-                training_data += &train_iter.to_string();
-                training_data += &":";
-                training_data += &ques_ans.question;
-                training_data += &" ";
-                training_data += "Example answer";
-                training_data += &train_iter.to_string();
-                training_data += &":";
-                training_data += &ques_ans.answer;
-                training_data += &"\n";
+                match ques_ans.answer {
+                    AnswerFormat::summary_type(sum) => {
+                        let summary_answer: SummaryType = sum;
+                        train_iter += 1;
+                        training_data += "Example question";
+                        training_data += &train_iter.to_string();
+                        training_data += &":";
+                        training_data += &ques_ans.question;
+                        training_data += &" ";
+                        training_data += "Example answer";
+                        training_data += &train_iter.to_string();
+                        training_data += &":";
+                        training_data += &serde_json::to_string(&summary_answer).unwrap();
+                        training_data += &"\n";
+                    }
+                    AnswerFormat::DE_type(_) => panic!("DE type not valid for summary"),
+                }
             }
 
             let system_prompt: String = String::from(
@@ -1001,7 +967,7 @@ fn get_summary_string() -> String {
 //const geneExpression: &str = &"geneExpression";
 
 #[derive(PartialEq, Debug, Clone, schemars::JsonSchema, serde::Serialize, serde::Deserialize)]
-struct SummaryType {
+pub struct SummaryType {
     // Serde uses this for deserialization.
     #[serde(default = "get_summary_string")]
     // Schemars uses this for schema generation.
@@ -1014,7 +980,7 @@ struct SummaryType {
 
 impl SummaryType {
     #[allow(dead_code)]
-    pub fn sort_summarytype_struct(&mut self) {
+    pub fn sort_summarytype_struct(mut self) -> SummaryType {
         // This function is necessary for testing (test_ai.rs) to see if two variables of type "SummaryType" are equal or not. Without this a vector of two Summarytype holding the same values but in different order will be classified separately.
         self.summaryterms.sort();
 
@@ -1022,6 +988,7 @@ impl SummaryType {
             Some(ref mut filterterms) => filterterms.sort(),
             None => {}
         }
+        self.clone()
     }
 }
 
@@ -1039,7 +1006,7 @@ impl PartialOrd for SummaryTerms {
             (SummaryTerms::clinical(_), SummaryTerms::clinical(_)) => Some(std::cmp::Ordering::Equal),
             (SummaryTerms::geneExpression(_), SummaryTerms::geneExpression(_)) => Some(std::cmp::Ordering::Equal),
             (SummaryTerms::clinical(_), SummaryTerms::geneExpression(_)) => Some(std::cmp::Ordering::Greater),
-            (SummaryTerms::geneExpression(_), SummaryTerms::clinical(_)) => Some(std::cmp::Ordering::Greater),
+            (SummaryTerms::geneExpression(_), SummaryTerms::clinical(_)) => Some(std::cmp::Ordering::Less),
         }
     }
 }
@@ -1313,10 +1280,40 @@ fn validate_summary_output(
                                     + &categorical_filter.value
                                     + &"\"},";
                                 validated_filter_terms_PP += &string_json;
-                                filter_hits += 1; // Once numeric term is also implemented, this statement will go outside the match block
                             }
-                            FilterTerm::Numeric(_numeric_term) => {} // To be implemented later
+                            FilterTerm::Numeric(numeric_filter) => {
+                                let string_json;
+                                if numeric_filter.greaterThan.is_some() && numeric_filter.lessThan.is_none() {
+                                    string_json = "{\"term\":\"".to_string()
+                                        + &numeric_filter.term
+                                        + &"\", \"gt\":\""
+                                        + &numeric_filter.greaterThan.unwrap().to_string()
+                                        + &"\"},";
+                                } else if numeric_filter.greaterThan.is_none() && numeric_filter.lessThan.is_some() {
+                                    string_json = "{\"term\":\"".to_string()
+                                        + &numeric_filter.term
+                                        + &"\", \"lt\":\""
+                                        + &numeric_filter.lessThan.unwrap().to_string()
+                                        + &"\"},";
+                                } else if numeric_filter.greaterThan.is_some() && numeric_filter.lessThan.is_some() {
+                                    string_json = "{\"term\":\"".to_string()
+                                        + &numeric_filter.term
+                                        + &"\", \"lt\":\""
+                                        + &numeric_filter.lessThan.unwrap().to_string()
+                                        + &"\", \"gt\":\""
+                                        + &numeric_filter.greaterThan.unwrap().to_string()
+                                        + &"\"},";
+                                } else {
+                                    // When both greater and less than are none
+                                    panic!(
+                                        "Numeric filter term {} is missing both greater than and less than values. One of them must be defined",
+                                        &numeric_filter.term
+                                    );
+                                }
+                                validated_filter_terms_PP += &string_json;
+                            }
                         };
+                        filter_hits += 1;
                     }
                     println!("validated_filter_terms_PP:{}", validated_filter_terms_PP);
                     if filter_hits > 0 {