@sjcrh/proteinpaint-rust 2.169.0 → 2.170.1

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (7) hide show
  1. package/Cargo.toml +8 -4
  2. package/package.json +1 -1
  3. package/src/aichatbot.rs +231 -254
  4. package/src/manhattan_plot.rs +225 -23
  5. package/src/query_classification.rs +152 -0
  6. package/src/summary_agent.rs +201 -0
  7. package/src/test_ai.rs +79 -72
package/src/summary_agent.rs ADDED
@@ -0,0 +1,201 @@
1
+ // Syntax: cd .. && cargo build --release && time cat ~/sjpp/test.txt | target/release/aichatbot
2
+ #![allow(non_snake_case)]
3
+ use crate::aichatbot::AiJsonFormat;
4
+ use anyhow::Result;
5
+ use json::JsonValue;
6
+ use std::fs;
7
+ use std::io;
8
+ use std::path::Path;
9
+ mod aichatbot; // Get summary agent
10
+
11
+ mod ollama; // Importing custom rig module for invoking ollama server
12
+ mod sjprovider; // Importing custom rig module for invoking SJ GPU server
13
+ mod test_ai; // Test examples for AI chatbot
14
+
15
+ #[tokio::main]
16
+ async fn main() -> Result<()> {
17
+ let mut input = String::new();
18
+ match io::stdin().read_line(&mut input) {
19
+ // Accepting the piped input from nodejs (or command line from testing)
20
+ Ok(_n) => {
21
+ let input_json = json::parse(&input);
22
+ match input_json {
23
+ Ok(json_string) => {
24
+ //println!("json_string:{}", json_string);
25
+ let user_input_json: &JsonValue = &json_string["user_input"];
26
+ let user_input: &str;
27
+ match user_input_json.as_str() {
28
+ Some(inp) => user_input = inp,
29
+ None => panic!("user_input field is missing in input json"),
30
+ }
31
+
32
+ let dataset_db_json: &JsonValue = &json_string["dataset_db"];
33
+ let dataset_db_str: &str;
34
+ match dataset_db_json.as_str() {
35
+ Some(inp) => dataset_db_str = inp,
36
+ None => panic!("dataset_db field is missing in input json"),
37
+ }
38
+
39
+ let genedb_json: &JsonValue = &json_string["genedb"];
40
+ let genedb_str: &str;
41
+ match genedb_json.as_str() {
42
+ Some(inp) => genedb_str = inp,
43
+ None => panic!("genedb field is missing in input json"),
44
+ }
45
+
46
+ if user_input.len() == 0 {
47
+ panic!("The user input is empty");
48
+ }
49
+
50
+ let tpmasterdir_json: &JsonValue = &json_string["tpmasterdir"];
51
+ let tpmasterdir: &str;
52
+ match tpmasterdir_json.as_str() {
53
+ Some(inp) => tpmasterdir = inp,
54
+ None => panic!("tpmasterdir not found"),
55
+ }
56
+
57
+ let binpath_json: &JsonValue = &json_string["binpath"];
58
+ let binpath: &str;
59
+ match binpath_json.as_str() {
60
+ Some(inp) => binpath = inp,
61
+ None => panic!("binpath not found"),
62
+ }
63
+
64
+ let ai_json_file_json: &JsonValue = &json_string["aifiles"];
65
+ let ai_json_file: String;
66
+ match ai_json_file_json.as_str() {
67
+ Some(inp) => ai_json_file = String::from(binpath) + &"/../../" + &inp,
68
+ None => {
69
+ panic!("ai json file not found")
70
+ }
71
+ }
72
+
73
+ let ai_json_file = Path::new(&ai_json_file);
74
+ let ai_json_file_path;
75
+ let current_dir = std::env::current_dir().unwrap();
76
+ match ai_json_file.canonicalize() {
77
+ Ok(p) => ai_json_file_path = p,
78
+ Err(_) => {
79
+ panic!(
80
+ "AI JSON file path not found:{:?}, current directory:{:?}",
81
+ ai_json_file, current_dir
82
+ )
83
+ }
84
+ }
85
+
86
+ // Read the file
87
+ let ai_data = fs::read_to_string(ai_json_file_path).unwrap();
88
+
89
+ // Parse the JSON data
90
+ let ai_json: AiJsonFormat =
91
+ serde_json::from_str(&ai_data).expect("AI JSON file does not have the correct format");
92
+
93
+ let genedb = String::from(tpmasterdir) + &"/" + &genedb_str;
94
+ let dataset_db = String::from(tpmasterdir) + &"/" + &dataset_db_str;
95
+
96
+ let apilink_json: &JsonValue = &json_string["apilink"];
97
+ let apilink: &str;
98
+ match apilink_json.as_str() {
99
+ Some(inp) => apilink = inp,
100
+ None => panic!("apilink field is missing in input json"),
101
+ }
102
+
103
+ let comp_model_name_json: &JsonValue = &json_string["comp_model_name"];
104
+ let comp_model_name: &str;
105
+ match comp_model_name_json.as_str() {
106
+ Some(inp) => comp_model_name = inp,
107
+ None => panic!("comp_model_name field is missing in input json"),
108
+ }
109
+
110
+ let embedding_model_name_json: &JsonValue = &json_string["embedding_model_name"];
111
+ let embedding_model_name: &str;
112
+ match embedding_model_name_json.as_str() {
113
+ Some(inp) => embedding_model_name = inp,
114
+ None => panic!("embedding_model_name field is missing in input json"),
115
+ }
116
+
117
+ let llm_backend_name_json: &JsonValue = &json_string["llm_backend_name"];
118
+ let llm_backend_name: &str;
119
+ match llm_backend_name_json.as_str() {
120
+ Some(inp) => llm_backend_name = inp,
121
+ None => panic!("llm_backend_name field is missing in input json"),
122
+ }
123
+
124
+ let llm_backend_type: aichatbot::llm_backend;
125
+ let mut final_output: Option<String> = None;
126
+ let temperature: f64 = 0.01;
127
+ let max_new_tokens: usize = 512;
128
+ let top_p: f32 = 0.95;
129
+ let testing = false; // This variable is always false in production, this is true in test_ai.rs for testing code
130
+ if llm_backend_name != "ollama" && llm_backend_name != "SJ" {
131
+ panic!(
132
+ "This code currently supports only Ollama and SJ provider. llm_backend_name must be \"ollama\" or \"SJ\""
133
+ );
134
+ } else if llm_backend_name == "ollama".to_string() {
135
+ llm_backend_type = aichatbot::llm_backend::Ollama();
136
+ // Initialize Ollama client
137
+ let ollama_client = ollama::Client::builder()
138
+ .base_url(apilink)
139
+ .build()
140
+ .expect("Ollama server not found");
141
+ let embedding_model = ollama_client.embedding_model(embedding_model_name);
142
+ let comp_model = ollama_client.completion_model(comp_model_name);
143
+ final_output = Some(
144
+ aichatbot::extract_summary_information(
145
+ user_input,
146
+ comp_model,
147
+ embedding_model,
148
+ &llm_backend_type,
149
+ temperature,
150
+ max_new_tokens,
151
+ top_p,
152
+ &dataset_db,
153
+ &genedb,
154
+ &ai_json,
155
+ testing,
156
+ )
157
+ .await,
158
+ );
159
+ } else if llm_backend_name == "SJ".to_string() {
160
+ llm_backend_type = aichatbot::llm_backend::Sj();
161
+ // Initialize Sj provider client
162
+ let sj_client = sjprovider::Client::builder()
163
+ .base_url(apilink)
164
+ .build()
165
+ .expect("SJ server not found");
166
+ let embedding_model = sj_client.embedding_model(embedding_model_name);
167
+ let comp_model = sj_client.completion_model(comp_model_name);
168
+ final_output = Some(
169
+ aichatbot::extract_summary_information(
170
+ user_input,
171
+ comp_model,
172
+ embedding_model,
173
+ &llm_backend_type,
174
+ temperature,
175
+ max_new_tokens,
176
+ top_p,
177
+ &dataset_db,
178
+ &genedb,
179
+ &ai_json,
180
+ testing,
181
+ )
182
+ .await,
183
+ );
184
+ }
185
+
186
+ match final_output {
187
+ Some(fin_out) => {
188
+ println!("final_output:{:?}", fin_out.replace("\\", ""));
189
+ }
190
+ None => {
191
+ println!("final_output:{{\"{}\":\"{}\"}}", "action", "unknown");
192
+ }
193
+ }
194
+ }
195
+ Err(error) => println!("Incorrect json:{}", error),
196
+ }
197
+ }
198
+ Err(error) => println!("Piping error: {}", error),
199
+ }
200
+ Ok(())
201
+ }
package/src/test_ai.rs CHANGED
@@ -5,6 +5,7 @@ fn main() {}
5
5
 
6
6
  #[cfg(test)]
7
7
  mod tests {
8
+ use crate::aichatbot::{AiJsonFormat, AnswerFormat, SummaryType, llm_backend, run_pipeline};
8
9
  use serde_json;
9
10
  use std::fs::{self};
10
11
  use std::path::Path;
@@ -62,13 +63,13 @@ mod tests {
62
63
  // Read the file
63
64
  let ai_data = fs::read_to_string(ai_json_file).unwrap();
64
65
  // Parse the JSON data
65
- let ai_json: super::super::AiJsonFormat =
66
+ let ai_json: AiJsonFormat =
66
67
  serde_json::from_str(&ai_data).expect("AI JSON file does not have the correct format");
67
68
  //println!("ai_json:{:?}", ai_json);
68
69
  let genedb = String::from(&serverconfig.tpmasterdir) + &"/" + &ai_json.genedb;
69
70
  let dataset_db = String::from(&serverconfig.tpmasterdir) + &"/" + &ai_json.db;
70
71
  let llm_backend_name = &serverconfig.llm_backend;
71
- let llm_backend_type: super::super::llm_backend;
72
+ let llm_backend_type: llm_backend;
72
73
 
73
74
  if llm_backend_name != "ollama" && llm_backend_name != "SJ" {
74
75
  panic!(
@@ -78,7 +79,7 @@ mod tests {
78
79
  let ollama_host = &serverconfig.ollama_apilink;
79
80
  let ollama_embedding_model_name = &serverconfig.ollama_embedding_model_name;
80
81
  let ollama_comp_model_name = &serverconfig.ollama_comp_model_name;
81
- llm_backend_type = super::super::llm_backend::Ollama();
82
+ llm_backend_type = llm_backend::Ollama();
82
83
  let ollama_client = super::super::ollama::Client::builder()
83
84
  .base_url(ollama_host)
84
85
  .build()
@@ -86,42 +87,45 @@ mod tests {
86
87
  let embedding_model = ollama_client.embedding_model(ollama_embedding_model_name);
87
88
  let comp_model = ollama_client.completion_model(ollama_comp_model_name);
88
89
  for chart in ai_json.charts.clone() {
89
- match chart {
90
- super::super::Charts::Summary(testdata) => {
91
- for ques_ans in testdata.TestData {
92
- let user_input = ques_ans.question;
93
- let llm_output = super::super::run_pipeline(
94
- &user_input,
95
- comp_model.clone(),
96
- embedding_model.clone(),
97
- llm_backend_type.clone(),
98
- temperature,
99
- max_new_tokens,
100
- top_p,
101
- &dataset_db,
102
- &genedb,
103
- &ai_json,
104
- &airoute,
105
- testing,
106
- )
107
- .await;
108
- let llm_json_value: super::super::SummaryType = serde_json::from_str(&llm_output.unwrap()).expect("Did not get a valid JSON of type {action: summary, summaryterms:[{clinical: term1}, {geneExpression: gene}], filter:[{term: term1, value: value1}]} from the LLM");
109
- let sum: super::super::SummaryType = ques_ans.answer;
110
- //println!("expected answer:{:?}", &sum);
111
- assert_eq!(
112
- llm_json_value.sort_summarytype_struct(),
113
- sum.sort_summarytype_struct()
114
- );
90
+ if chart.r#type == "Summary" {
91
+ for ques_ans in chart.TestData {
92
+ let user_input = ques_ans.question;
93
+ let llm_output = run_pipeline(
94
+ &user_input,
95
+ comp_model.clone(),
96
+ embedding_model.clone(),
97
+ llm_backend_type.clone(),
98
+ temperature,
99
+ max_new_tokens,
100
+ top_p,
101
+ &dataset_db,
102
+ &genedb,
103
+ &ai_json,
104
+ &airoute,
105
+ testing,
106
+ )
107
+ .await;
108
+ let llm_json_value: SummaryType = serde_json::from_str(&llm_output.unwrap()).expect("Did not get a valid JSON of type {action: summary, summaryterms:[{clinical: term1}, {geneExpression: gene}], filter:[{term: term1, value: value1}]} from the LLM");
109
+ match ques_ans.answer {
110
+ AnswerFormat::summary_type(sum) => {
111
+ //println!("expected answer:{:?}", &sum);
112
+ assert_eq!(
113
+ llm_json_value.sort_summarytype_struct(),
114
+ sum.sort_summarytype_struct()
115
+ );
116
+ }
117
+ AnswerFormat::DE_type(_) => {
118
+ panic!("DE type not valid for summary")
119
+ }
115
120
  }
116
121
  }
117
- super::super::Charts::DE(_testdata) => {} // To do
118
122
  }
119
123
  }
120
124
  } else if *llm_backend_name == "SJ".to_string() {
121
125
  let sjprovider_host = &serverconfig.sj_apilink;
122
126
  let sj_embedding_model_name = &serverconfig.sj_embedding_model_name;
123
127
  let sj_comp_model_name = &serverconfig.sj_comp_model_name;
124
- llm_backend_type = super::super::llm_backend::Sj();
128
+ llm_backend_type = llm_backend::Sj();
125
129
  let sj_client = super::super::sjprovider::Client::builder()
126
130
  .base_url(sjprovider_host)
127
131
  .build()
@@ -130,50 +134,53 @@ mod tests {
130
134
  let comp_model = sj_client.completion_model(sj_comp_model_name);
131
135
 
132
136
  for chart in ai_json.charts.clone() {
133
- match chart {
134
- super::super::Charts::Summary(testdata) => {
135
- for ques_ans in testdata.TestData {
136
- let user_input = ques_ans.question;
137
- if user_input.len() > 0 {
138
- let llm_output = super::super::run_pipeline(
139
- &user_input,
140
- comp_model.clone(),
141
- embedding_model.clone(),
142
- llm_backend_type.clone(),
143
- temperature,
144
- max_new_tokens,
145
- top_p,
146
- &dataset_db,
147
- &genedb,
148
- &ai_json,
149
- &airoute,
150
- testing,
151
- )
152
- .await;
153
- //println!("user_input:{}", user_input);
154
- //println!("llm_answer:{:?}", llm_output);
155
- //println!("expected answer:{:?}", &ques_ans.answer);
156
- let llm_json_value: super::super::SummaryType = serde_json::from_str(&llm_output.unwrap()).expect("Did not get a valid JSON of type {action: summary, summaryterms:[{clinical: term1}, {geneExpression: gene}], filter:[{term: term1, value: value1}]} from the LLM");
157
- //println!(
158
- // "llm_answer:{:?}",
159
- // llm_json_value.clone().sort_summarytype_struct()
160
- //);
161
- //println!(
162
- // "expected answer:{:?}",
163
- // &expected_json_value.clone().sort_summarytype_struct()
164
- //);
165
- let sum: super::super::SummaryType = ques_ans.answer;
166
- //println!("expected answer:{:?}", &sum);
167
- assert_eq!(
168
- llm_json_value.sort_summarytype_struct(),
169
- sum.sort_summarytype_struct()
170
- );
171
- } else {
172
- panic!("The user input is empty");
137
+ if chart.r#type == "Summary" {
138
+ for ques_ans in chart.TestData {
139
+ let user_input = ques_ans.question;
140
+ if user_input.len() > 0 {
141
+ let llm_output = run_pipeline(
142
+ &user_input,
143
+ comp_model.clone(),
144
+ embedding_model.clone(),
145
+ llm_backend_type.clone(),
146
+ temperature,
147
+ max_new_tokens,
148
+ top_p,
149
+ &dataset_db,
150
+ &genedb,
151
+ &ai_json,
152
+ &airoute,
153
+ testing,
154
+ )
155
+ .await;
156
+ //println!("user_input:{}", user_input);
157
+ //println!("llm_answer:{:?}", llm_output);
158
+ //println!("expected answer:{:?}", &ques_ans.answer);
159
+ let llm_json_value: SummaryType = serde_json::from_str(&llm_output.unwrap()).expect("Did not get a valid JSON of type {action: summary, summaryterms:[{clinical: term1}, {geneExpression: gene}], filter:[{term: term1, value: value1}]} from the LLM");
160
+ //println!(
161
+ // "llm_answer:{:?}",
162
+ // llm_json_value.clone().sort_summarytype_struct()
163
+ //);
164
+ //println!(
165
+ // "expected answer:{:?}",
166
+ // &expected_json_value.clone().sort_summarytype_struct()
167
+ //);
168
+ match ques_ans.answer {
169
+ AnswerFormat::summary_type(sum) => {
170
+ //println!("expected answer:{:?}", &sum);
171
+ assert_eq!(
172
+ llm_json_value.sort_summarytype_struct(),
173
+ sum.sort_summarytype_struct()
174
+ );
175
+ }
176
+ AnswerFormat::DE_type(_) => {
177
+ panic!("DE type not valid for summary")
178
+ }
173
179
  }
180
+ } else {
181
+ panic!("The user input is empty");
174
182
  }
175
183
  }
176
- super::super::Charts::DE(_testdata) => {} // To do
177
184
  }
178
185
  }
179
186
  }