@sjcrh/proteinpaint-rust 2.169.0 → 2.170.0

package/Cargo.toml

@@ -123,10 +123,14 @@ path="src/cerno.rs"
 name="readH5"
 path="src/readH5.rs"
 
-[[bin]]
-name="aichatbot"
-path="src/aichatbot.rs"
-
 [[bin]]
 name="manhattan_plot"
 path="src/manhattan_plot.rs"
+
+[[bin]]
+name="query_classification"
+path="src/query_classification.rs"
+
+[[bin]]
+name="summary_agent"
+path="src/summary_agent.rs"

package/package.json

@@ -1,5 +1,5 @@
 {
-	"version": "2.169.0",
+	"version": "2.170.0",
 	"name": "@sjcrh/proteinpaint-rust",
 	"type": "module",
 	"description": "Rust-based utilities for proteinpaint",

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,69 +11,52 @@ 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>,
-}
-
-#[derive(PartialEq, Debug, Clone, schemars::JsonSchema, serde::Serialize, serde::Deserialize)]
-enum Charts {
-    // More chart types will be added here later
-    Summary(TrainTestDataSummary),
-    DE(TrainTestDataDE),
+    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)]
-struct TrainTestDataSummary {
-    SystemPrompt: String,
-    TrainingData: Vec<QuestionAnswerSummary>,
-    TestData: Vec<QuestionAnswerSummary>,
+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)]
-struct QuestionAnswerSummary {
-    question: String,
-    answer: SummaryType,
+pub struct QuestionAnswer {
+    pub question: String,
+    pub answer: AnswerFormat,
 }
 
 #[derive(PartialEq, Debug, Clone, schemars::JsonSchema, serde::Serialize, serde::Deserialize)]
-struct TrainTestDataDE {
-    SystemPrompt: String,
-    TrainingData: Vec<QuestionAnswerDE>,
-    TestData: Vec<QuestionAnswerDE>,
-}
-
-#[derive(PartialEq, Debug, Clone, schemars::JsonSchema, serde::Serialize, serde::Deserialize)]
-struct QuestionAnswerDE {
-    question: String,
-    answer: DEType,
+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 DEType {
+pub struct DEType {
     action: String,
     DE_output: DETerms,
 }
 
 #[derive(PartialEq, Debug, Clone, schemars::JsonSchema, serde::Serialize, serde::Deserialize)]
-struct DETerms {
+pub struct DETerms {
     group1: GroupType,
     group2: GroupType,
 }
 
 #[derive(PartialEq, Debug, Clone, schemars::JsonSchema, serde::Serialize, serde::Deserialize)]
-struct GroupType {
+pub struct GroupType {
     name: String,
 }
 
@@ -86,204 +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"),
-                    }
-
-                    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(())
-}
+//#[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,
@@ -429,7 +396,8 @@ 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,
@@ -449,9 +417,12 @@ async fn classify_query_by_dataset_type(
     let mut contents = String::from("");
 
     if let Some(object) = ai_json.as_object() {
-        for (_key, value) in object {
-            contents += &value.as_str().unwrap();
-            contents += "---"; // Adding delimiter
+        contents = object["general"].to_string();
+        for (key, value) in object {
+            if key != "general" {
+                contents += &value.as_str().unwrap();
+                contents += "---"; // Adding delimiter
+            }
         }
     }
 
@@ -568,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,
@@ -841,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,
@@ -897,10 +869,10 @@ async fn extract_summary_information(
         .filter(|x| user_words2.contains(&x.to_lowercase()))
         .collect();
 
-    let mut summary_data_check: Option<TrainTestDataSummary> = None;
+    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;
         }
     }
@@ -910,18 +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 {
-                let summary_answer: SummaryType = ques_ans.answer;
-                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";
+                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(
@@ -990,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.