From 72ee27d859082c8ce3dd9b4795eda65b704f14d0 Mon Sep 17 00:00:00 2001
From: Srinivas Gorur-Shandilya <code@srinivas.gs>
Date: Tue, 8 Jul 2025 15:09:24 -0400
Subject: [PATCH] feat: support for new version of loop modelling

---
 src/drug_discovery/structures/protein.py |  4 +--
 src/functions/loop_modelling.py          | 46 ++++++++++++++++--------
 2 files changed, 34 insertions(+), 16 deletions(-)

diff --git a/src/drug_discovery/structures/protein.py b/src/drug_discovery/structures/protein.py
index 89bf1517..0fd3404e 100644
--- a/src/drug_discovery/structures/protein.py
+++ b/src/drug_discovery/structures/protein.py
@@ -247,7 +247,7 @@ def sequence(self) -> list[str]:
 
         return sequences
 
-    def model_loops(self) -> None:
+    def model_loops(self, use_cache: bool = True) -> None:
         """model loops in protein structure"""
 
         from deeporigin.functions.loop_modelling import model_loops
@@ -257,7 +257,7 @@ def model_loops(self) -> None:
         if pdb_id is None:
             raise ValueError("Currently, PDB ID is required to model loops.")
 
-        file_path = model_loops(pdb_id=pdb_id)
+        file_path = model_loops(pdb_id=pdb_id, use_cache=use_cache)
         protein = Protein.from_file(file_path)
         self.structure = protein.structure
 
diff --git a/src/functions/loop_modelling.py b/src/functions/loop_modelling.py
index 994a4c13..5e5c9c40 100644
--- a/src/functions/loop_modelling.py
+++ b/src/functions/loop_modelling.py
@@ -6,7 +6,8 @@
 from beartype import beartype
 import requests
 
-URL = "http://loop-modelling.default.jobs.edge.deeporigin.io/model_loops"
+# URL = "http://loop-modelling.default.jobs.edge.deeporigin.io/model_loops"
+URL = "http://localhost:8080/model_loops"
 CACHE_DIR = os.path.expanduser("~/.deeporigin/model_loops")
 
 
@@ -14,7 +15,8 @@
 def model_loops(
     *,
     pdb_id: str,
-) -> str:
+    use_cache: bool = True,
+) -> dict:
     """
     Run system preparation on a protein-ligand complex.
 
@@ -29,28 +31,44 @@ def model_loops(
     hash_input = f"{pdb_id}"
     cache_key = hashlib.md5(hash_input.encode()).hexdigest()
     cache_path = os.path.join(CACHE_DIR, cache_key)
-    output_pdb_path = os.path.join(cache_path, "loops_modelled.pdb")
+    output_json_path = os.path.join(cache_path, "data.json")
+    output_pdb_path = os.path.join(cache_path, "data.pdb")
+
+    # Create cache directory if it doesn't exist
+    os.makedirs(cache_path, exist_ok=True)
 
     # Check if cached results exist
-    if os.path.exists(output_pdb_path):
-        return output_pdb_path
+    if use_cache and os.path.exists(output_json_path):
+        import json
+
+        with open(output_json_path, "r") as f:
+            return json.load(f)
 
     # If no cached results, proceed with server call
     payload = {
         "pdb_id": pdb_id,
     }
 
-    response = requests.post(URL, json=payload, stream=True)
+    try:
+        response = requests.post(URL, json=payload, stream=True)
+    except Exception as e:
+        raise RuntimeError(f"Failed to connect to server: {e}") from None
 
     if response.status_code == 200:
-        # Create cache directory if it doesn't exist
-        os.makedirs(cache_path, exist_ok=True)
-        # Save the result to the cache
-        with open(output_pdb_path, "wb") as f:
-            for chunk in response.iter_content(chunk_size=8192):
-                if chunk:
-                    f.write(chunk)
-        return output_pdb_path
+        import json
+
+        response_json = response.json()
+        # Write pdb_contents to output_pdb_path
+        pdb_contents = response_json.pop("pdb_contents", None)
+        if pdb_contents is not None:
+            with open(output_pdb_path, "w") as pdb_file:
+                pdb_file.write(pdb_contents)
+
+        response_json["pdb_file"] = output_pdb_path
+        # Write the modified JSON (without pdb_contents) to output_json_path
+        with open(output_json_path, "w") as f:
+            json.dump(response_json, f, indent=2)
+        return response_json
 
     # If the server request fails, raise an error
     raise RuntimeError(