Minor bug fixes

app/db.py

@@ -6,10 +6,9 @@ from time import time
 
 import gradio as gr
 import pytz
-import requests
 from tinydb import TinyDB, where
 
-from app.fn import send_email
+from app.fn import send_email, get_timezone_by_ip
 
 SERVER_DATA_DIR = os.getenv('DATA', 'results')
 DB_EXPIRY = timedelta(hours=48).total_seconds()
@@ -33,12 +32,6 @@ def ts_to_str(timestamp, timezone):
     # Format the datetime object to the specified string format
     return localized_dt.strftime('%Y-%m-%d %H:%M:%S (%Z%z)')
 
-def get_timezone_by_ip(ip):
-    try:
-        data = requests.get(f'https://worldtimeapi.org/api/ip/{ip}').json()
-        return data['timezone']
-    except Exception:
-        return 'UTC'
 
 class JobDB(TinyDB):
     def remove_job_record(self,job_id):

app/fn.py

@@ -182,9 +182,9 @@ FILTER_MAP = {
 }
 
 
-def get_timezone_by_ip(ip, session):
+def get_timezone_by_ip(ip):
     try:
-        data = session.get(f'https://worldtimeapi.org/api/ip/{ip}').json()
+        data = requests.get(f'https://worldtimeapi.org/api/ip/{ip}').json()
         return data['timezone']
     except Exception:
         return 'UTC'

app/main.py

@@ -1,4 +1,7 @@
+import spaces
+
 import os
+import tempfile
 import uuid
 import zipfile
 from datetime import datetime
@@ -21,7 +24,8 @@ from inference import (read_fragment_library, process_fragment_library, extract_
 from app import static, fn, db
 
 
-gr.set_static_paths(paths=["./results/", "./app/"])
+Path(tempfile.gettempdir(), 'gradio').mkdir(exist_ok=True)
+gr.set_static_paths(paths=["data/", "results/", "app/"])
 job_db = db.init_job_db()
 os.chmod('./fpocket', 0o755)
 
@@ -192,6 +196,7 @@ def job_validate(
 
     return job_info
 
+@spaces.GPU(duration=600)
 def dock_link(
     frag_lib, prot,
     dock_n_steps, dock_n_poses, dock_confidence_threshold,
@@ -233,7 +238,7 @@ def dock_link(
             temp_sigma_data_tr=config.temp_sigma_data_tr,
             temp_sigma_data_rot=config.temp_sigma_data_rot,
             temp_sigma_data_tor=config.temp_sigma_data_tor,
-            save_docking=(pocket_extract_method == 'clustering'),
+            save_docking=True,
             device=device,
         )
 
@@ -650,7 +655,7 @@ with gr.Blocks(theme=THEME, title='GenFBDD', css=static.CSS, delete_cache=(3600,
                 raise ValueError(f"Unsupported method: {method}")
             return {input_prot_file: gr.File(str(file), visible=True)}
         except Exception as e:
-            gr.Error(f"Query error: {str(e)}")
+            raise gr.Error(f"Query error: {str(e)}")
 
     prot_query_btn.click(
         fn=pdb_query,
@@ -658,7 +663,7 @@ with gr.Blocks(theme=THEME, title='GenFBDD', css=static.CSS, delete_cache=(3600,
         outputs=[input_prot_file],
     )
 
-    input_prot_file_change = input_prot_file.change(
+    input_prot_file.change(
         fn=lambda: gr.Info('Rendering 3DMol view...'),
     ).then(
         fn=lambda x, y: gr.Info('3DMol view rendered.'),

utils/ui.py

@@ -1,65 +0,0 @@
-import os
-import smtplib
-from datetime import datetime
-from email.mime.multipart import MIMEMultipart
-from email.mime.text import MIMEText
-from email.utils import formatdate, make_msgid
-
-import pytz
-from markdown import markdown
-
-
-def get_timezone_by_ip(ip, session):
-    try:
-        data = session.get(f'https://worldtimeapi.org/api/ip/{ip}').json()
-        return data['timezone']
-    except Exception:
-        return 'UTC'
-
-
-def ts_to_str(timestamp, timezone):
-    # Create a timezone-aware datetime object from the UNIX timestamp
-    dt = datetime.fromtimestamp(timestamp, pytz.utc)
-
-    # Convert the timezone-aware datetime object to the target timezone
-    target_timezone = pytz.timezone(timezone)
-    localized_dt = dt.astimezone(target_timezone)
-
-    # Format the datetime object to the specified string format
-    return localized_dt.strftime('%Y-%m-%d %H:%M:%S (%Z%z)')
-
-
-def send_email(job_info):
-    if job_info.get('email'):
-        try:
-            email_info = job_info.copy()
-            email_serv = os.getenv('EMAIL_SERV')
-            email_port = os.getenv('EMAIL_PORT')
-            email_addr = os.getenv('EMAIL_ADDR')
-            email_pass = os.getenv('EMAIL_PASS')
-            email_form = os.getenv('EMAIL_FORM')
-            email_subj = os.getenv('EMAIL_SUBJ')
-
-            for key, value in email_info.items():
-                if key.endswith("time") and value:
-                    email_info[key] = ts_to_str(value, get_timezone_by_ip(email_info['ip']))
-
-            server = smtplib.SMTP(email_serv, int(email_port))
-            # server.starttls()
-
-            server.login(email_addr, email_pass)
-            msg = MIMEMultipart("alternative")
-            msg["From"] = email_addr
-            msg["To"] = email_info['email']
-            msg["Subject"] = email_subj.format(**email_info)
-            msg["Date"] = formatdate(localtime=True)
-            msg["Message-ID"] = make_msgid()
-
-            msg.attach(MIMEText(markdown(email_form.format(**email_info)), 'html'))
-            msg.attach(MIMEText(email_form.format(**email_info), 'plain'))
-
-            server.sendmail(email_addr, email_info['email'], msg.as_string())
-            server.quit()
-            gr.Info('Email notification sent.')
-        except Exception as e:
-            gr.Warning('Failed to send email notification due to error: ' + str(e))