Update app.py

app.py

@@ -1,7 +1,7 @@
 import concurrent.futures
 import glob
 import smtplib
-from datetime import datetime
+from datetime import datetime, timedelta
 import itertools
 import textwrap
 from email.mime.multipart import MIMEMultipart
@@ -52,9 +52,9 @@ from deepscreen.predict import predict
 sys.path.append(os.path.join(RDConfig.RDContribDir, 'SA_Score'))
 import sascorer
 
-UNIPROT_ENDPOINT = 'https://rest.uniprot.org/uniprotkb/{query}'
 DATASET_MAX_LEN = 10_000
 SERVER_DATA_DIR = os.getenv('DATA')  # '/data'
+DB_EXPIRY = timedelta(hours=48).total_seconds()
 
 CSS = """
 .help-tip {
@@ -244,20 +244,32 @@ for job in jobs:
 scheduler = BackgroundScheduler()
 
 
+def remove_job_record(job_id):
+    # Delete the job from the database
+    db.remove(Job.id == job_id)
+    # Delete the corresponding files
+    files = glob.glob(f"/data/{job_id}*")
+    for file_path in files:
+        if os.path.exists(file_path):
+            os.remove(file_path)
+
+
 def check_expiry():
     Job = Query()
     jobs = db.all()
 
     for job in jobs:
         # Check if the job has expired
-        if job['expiry_time'] < time():
-            # Delete the job from the database
-            db.remove(Job.id == job['id'])
-            # Delete the corresponding file
-            files = glob.glob(f"/data/{job['id']}*")
-            for file_path in files:
-                if os.path.exists(file_path):
-                    os.remove(file_path)
+        if job['status'] != 'RUNNING':
+            expiry_time = job['expiry_time'] if job['expiry_time'] is not None else job['start_time'] + DB_EXPIRY
+            if expiry_time < time():
+                # Delete the job from the database
+                db.remove(Job.id == job['id'])
+                # Delete the corresponding file
+                files = glob.glob(f"/data/{job['id']}*")
+                for file_path in files:
+                    if os.path.exists(file_path):
+                        os.remove(file_path)
         elif job['status'] == 'RUNNING' and time() - job['start_time'] > 4 * 60 * 60:  # 4 hours
             # Mark the job as failed
             db.update({'status': 'FAILED',
@@ -699,7 +711,7 @@ def submit_predict(predict_filepath, task, preset, target_family, state):
         job_query = (Job.id == job_id)
 
         end_time = time()
-        expiry_time = end_time + 48 * 60 * 60  # Add 48 hours
+        expiry_time = end_time + DB_EXPIRY
 
         db.update({'end_time': end_time,
                    'expiry_time': expiry_time,
@@ -758,7 +770,7 @@ def update_df(file, progress=gr.Progress(track_tqdm=True)):
                 raw_df: df,
                 report_df: df.copy(),
                 analyze_btn: gr.Button(interactive=True),
-                report_task: gr.Dropdown(value=task)}  # pie_chart
+                report_task: task}  # pie_chart
     else:
         return {analyze_btn: gr.Button(interactive=False)}
 
@@ -1567,7 +1579,7 @@ with gr.Blocks(theme=theme, title='DeepSEQreen', css=CSS, delete_cache=(3600, 48
                 with gr.Column():
                     file_for_report = gr.File(interactive=True, type='filepath')
                     report_task = gr.Dropdown(list(TASK_MAP.keys()), visible=False, value=None,
-                                              label='Specify the Task for the Labels in the Upload Dataset')
+                                              label='Specify the Task Labels in the Upload Dataset')
                 raw_df = gr.State(value=pd.DataFrame())
                 report_df = gr.State(value=pd.DataFrame())
                 scores = gr.CheckboxGroup(list(SCORE_MAP.keys()), label='Scores')
@@ -1660,6 +1672,7 @@ with gr.Blocks(theme=theme, title='DeepSEQreen', css=CSS, delete_cache=(3600, 48
 
 
     def uniprot_query(input_type, uid, gene, organism='Human'):
+        uniprot_endpoint = 'https://rest.uniprot.org/uniprotkb/{query}'
         fasta_rec = ''
 
         match input_type:
@@ -1670,11 +1683,11 @@ with gr.Blocks(theme=theme, title='DeepSEQreen', css=CSS, delete_cache=(3600, 48
                 query = f'search?query=organism_name:{organism.strip()}+AND+gene:{gene.strip()}&format=fasta'
 
         try:
-            fasta = session.get(UNIPROT_ENDPOINT.format(query=query))
+            fasta = session.get(uniprot_endpoint.format(query=query))
             fasta.raise_for_status()
             if fasta.text:
                 fasta_rec = next(SeqIO.parse(io.StringIO(fasta.text), format='fasta'))
-                fasta_rec = f"{fasta_rec.description}\n{fasta_rec.seq}"
+                fasta_rec = f">{fasta_rec.description}\n{fasta_rec.seq}"
 
         except Exception as e:
             raise gr.Warning(f"Failed to query FASTA from UniProt database due to {str(e)}")
@@ -1698,18 +1711,21 @@ with gr.Blocks(theme=theme, title='DeepSEQreen', css=CSS, delete_cache=(3600, 48
 
 
     def target_family_detect(fasta, progress=gr.Progress(track_tqdm=True)):
-        aligner = PairwiseAligner(scoring='blastp', mode='local')
-        alignment_df = pd.read_csv('data/target_libraries/ChEMBL33_all_spe_single_prot_info.csv')
-
-        def align_score(query):
-            return aligner.align(process_target_fasta(fasta), query).score
-
-        alignment_df['score'] = alignment_df['X2'].swifter.progress_bar(
-            desc="Detecting protein family of the target...").apply(align_score)
-        row = alignment_df.loc[alignment_df['score'].idxmax()]
-        return gr.Dropdown(value=row['protein_family'].capitalize(),
-                           info=f"Reason: Best BLASTP score ({row['score']}) "
-                                f"with {row['ID2']} from family {row['protein_family']}")
+        try:
+            aligner = PairwiseAligner(scoring='blastp', mode='local')
+            alignment_df = pd.read_csv('data/target_libraries/ChEMBL33_all_spe_single_prot_info.csv')
+
+            def align_score(query):
+                return aligner.align(process_target_fasta(fasta), query).score
+
+            alignment_df['score'] = alignment_df['X2'].swifter.progress_bar(
+                desc="Detecting protein family of the target...").apply(align_score)
+            row = alignment_df.loc[alignment_df['score'].idxmax()]
+            return gr.Dropdown(value=row['protein_family'].capitalize(),
+                               info=f"Reason: Best BLASTP score ({row['score']}) "
+                                    f"with {row['ID2']} from family {row['protein_family']}")
+        except Exception as e:
+            gr.Warning("Failed to detect the protein family due to error: " + str(e))
 
 
     target_family_detect_btn.click(fn=target_family_detect, inputs=target_fasta, outputs=drug_screen_target_family)
@@ -2243,11 +2259,11 @@ QALAHAYFAQYHDPDDEPVADPYDQSFESRDLLIDEWKSLTYDEVISFVPPPLDQEEMES
             return None
 
 
-    def create_html_report_file(df, file_report, progress=gr.Progress(track_tqdm=True)):
+    def create_html_report_file(df, file_report, task, progress=gr.Progress(track_tqdm=True)):
         try:
             now = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
             filename = f"/data/{Path(file_report.name).stem}_DeepSEQreen_report_{now}.html"
-            create_html_report(df, filename)
+            create_html_report(df, filename, task)
             return gr.File(filename, visible=True)
         except Exception as e:
             gr.Warning(f"Failed to generate HTML due to error: {str(e)}")
@@ -2261,7 +2277,7 @@ QALAHAYFAQYHDPDDEPVADPYDQSFESRDLLIDEWKSLTYDEVISFVPPPLDQEEMES
            outputs=csv_download_file, show_progress='full')
     html_generate.click(
         lambda: [gr.Button(visible=False), gr.File(visible=True)], outputs=[html_generate, html_download_file],
-    ).then(fn=create_html_report_file, inputs=[report_df, file_for_report],
+    ).then(fn=create_html_report_file, inputs=[report_df, file_for_report, report_task],
            outputs=html_download_file, show_progress='full')
 
 if __name__ == "__main__":