Upload app.py

app.py

@@ -12,6 +12,10 @@ import plotly.graph_objects as go
 from shinywidgets import output_widget, render_widget
 import requests
 import asyncio
+from Bio import PDB
+from Bio.PDB.PDBList import PDBList
+from Bio.PDB.Polypeptide import protein_letters_3to1
+import shutil
 warnings.filterwarnings('ignore')
 
 # Load environment variables from .env file
@@ -52,56 +56,70 @@ class PDBSearchAssistant:
             Examples:
             Query: "Find human insulin structures with X-ray better than 2.5Å resolution"
             Protein: insulin
-            Organism: human
+            Organism: Homo sapiens
             Resolution: 2.5
             Sequence: none
             PDB_ID: none
             Method: X-RAY
 
-            Query: "Get sequence of PDB ID 8ET6"
+            Query: "Find structures containing sequence MNIFEMLRIDEGLRLKIYKDTEGYYTIGIGHLLTKSPSLNAAKSELDKAIGRNTNGVITKDEAEKLFNQDVDAAVRGILRNAKLKPVYDSLDAVRRAALINMVFQMGETGVAGFTNSLRMLQQKRWDEAAVNLAKSRWYNQTPNRAKRVITTFRTGTWDAYKNL"
             Protein: none
             Organism: none
             Resolution: none
-            Sequence: none
-            PDB_ID: 8ET6
+            Sequence: MNIFEMLRIDEGLRLKIYKDTEGYYTIGIGHLLTKSPSLNAAKSELDKAIGRNTNGVITKDEAEKLFNQDVDAAVRGILRNAKLKPVYDSLDAVRRAALINMVFQMGETGVAGFTNSLRMLQQKRWDEAAVNLAKSRWYNQTPNRAKRVITTFRTGTWDAYKNL
+            PDB_ID: none
             Method: none
 
-            Query: "Sequence of 7BZ5"
+            Query: "Get sequence of PDB ID 8ET6"
             Protein: none
             Organism: none
             Resolution: none
             Sequence: none
-            PDB_ID: 7BZ5
+            PDB_ID: 8ET6
             Method: none
 
-            Query: "7BZ5"
-            Protein: none
-            Organism: none
+            Query: "Find mouse lysozyme structures"
+            Protein: lysozyme
+            Organism: Mus musculus
             Resolution: none
             Sequence: none
-            PDB_ID: 7BZ5
+            PDB_ID: none
             Method: none
 
-            Query: "6KAO"
+            Query: "Show me E. coli protein structures solved by Cryo-EM"
             Protein: none
-            Organism: none
+            Organism: Escherichia coli
             Resolution: none
             Sequence: none
-            PDB_ID: 6KAO
+            PDB_ID: none
+            Method: EM
+
+            Query: "Find S. cerevisiae structures with resolution better than 1.8Å"
+            Protein: none
+            Organism: Saccharomyces cerevisiae
+            Resolution: 1.8
+            Sequence: none
+            PDB_ID: none
             Method: none
-            
-            Query: "Find structures containing sequence MNIFEMLRIDEGLRLKIYKDTEGYYTIGIGHLLTKSPSLNAAKSELDKAIGRNTNGVITKDEAEKLFNQDVDAAVRGILRNAKLKPVYDSLDAVRRAALINMVFQMGETGVAGFTNSLRMLQQKRWDEAAVNLAKSRWYNQTPNRAKRVITTFRTGTWDAYKNL"
+
+            Query: "Sequence of 7BZ5"
             Protein: none
             Organism: none
             Resolution: none
-            Sequence: MNIFEMLRIDEGLRLKIYKDTEGYYTIGIGHLLTKSPSLNAAKSELDKAIGRNTNGVITKDEAEKLFNQDVDAAVRGILRNAKLKPVYDSLDAVRRAALINMVFQMGETGVAGFTNSLRMLQQKRWDEAAVNLAKSRWYNQTPNRAKRVITTFRTGTWDAYKNL
-            PDB_ID: none
+            Sequence: none
+            PDB_ID: 7BZ5
             Method: none
 
+
+
             Now analyze:
             Query: {query}
             """
 
+        self.pdb_dir = "pdb_tmp"  # 임시 PDB 파일 저장 디렉토리
+        os.makedirs(self.pdb_dir, exist_ok=True)
+        self.pdbl = PDBList()
+    
     def search_pdb(self, query):
         try:
             # Get search parameters from LLM
@@ -114,6 +132,7 @@ class PDBSearchAssistant:
             pdb_id = None
             sequence = None
             method = None
+            organism = None
             has_resolution_query = False
             resolution_direction = "less"
             
@@ -175,6 +194,10 @@ class PDBSearchAssistant:
                     value = line.split('PDB_ID:')[1].strip()
                     if value.lower() not in ['none', 'n/a']:
                         pdb_id = value
+                elif 'Organism:' in line:
+                    value = line.split('Organism:')[1].strip()
+                    if value.lower() not in ['none', 'n/a']:
+                        organism = value
             
             # Build search query
             queries = []
@@ -260,6 +283,16 @@ class PDBSearchAssistant:
                 )
                 queries.append(method_query)
             
+            # Add organism filter if specified
+            if organism:
+                print(f"Adding organism filter: {organism}")
+                organism_query = AttributeQuery(
+                    attribute="rcsb_entity_source_organism.taxonomy_lineage.name",
+                    operator="exact_match",
+                    value=organism
+                )
+                queries.append(organism_query)
+            
             # Combine queries with AND operator
             if queries:
                 final_query = queries[0]
@@ -334,91 +367,70 @@ class PDBSearchAssistant:
             return []
 
     def get_sequences_by_pdb_id(self, pdb_id):
-        """Get sequences for all chains in a PDB structure"""
+        """Get sequences for all chains in a PDB structure using Biopython"""
         try:
-            # ProteinSearchEngine 인스턴스 생성
-            search_engine = ProteinSearchEngine()
-            
-            # ProteinQuery 객체 생성 (resolution limit은 높게 설정하여 모든 결과 포함)
-            query = ProteinQuery(
-                name=pdb_id,
-                max_resolution=100.0  # 높은 값으로 설정하여 모든 구조 포함
+            # Download PDB file
+            pdb_path = self.pdbl.retrieve_pdb_file(
+                pdb_id, 
+                pdir=self.pdb_dir, 
+                file_format="pdb"
             )
             
-            # 검색 실행
-            results = search_engine.search(query)
-            
-            if not results:
+            if not pdb_path or not os.path.exists(pdb_path):
+                print(f"Failed to download PDB file for {pdb_id}")
                 return []
-                
-            sequences = []
-            # 결과에서 sequence 정보 추출
-            for structure in results:
-                if structure.pdb_id.upper() == pdb_id.upper():
-                    chain_info = {
-                        'chain_id': 'ALL',  # 체인 정보는 통합
-                        'entity_id': '1',
-                        'description': structure.title,
-                        'sequence': structure.sequence,
-                        'length': len(structure.sequence),
-                        'resolution': structure.resolution,
-                        'method': structure.method,
-                        'release_date': structure.release_date
-                    }
-                    sequences.append(chain_info)
-                    break  # 정확한 PDB ID 매치를 찾으면 중단
             
-            # 결과가 없으면 직접 API 호출 시도
-            if not sequences:
-                print(f"No results found using ProteinSearchEngine, trying direct API call...")
-                return self._get_sequences_by_direct_api(pdb_id)
-                
-            return sequences
+            # Parse structure
+            parser = PDB.PDBParser(QUIET=True)
+            structure = parser.get_structure(pdb_id, pdb_path)
             
-        except Exception as e:
-            print(f"Error in ProteinSearchEngine search for PDB ID {pdb_id}: {str(e)}")
-            # 에러 발생 시 직접 API 호출로 폴백
-            return self._get_sequences_by_direct_api(pdb_id)
-
-    def _get_sequences_by_direct_api(self, pdb_id):
-        """Fallback method using direct API calls"""
-        # 기존의 get_sequences_by_pdb_id 메소드 내용을 여기로 이동
-        try:
-            url = f"https://data.rcsb.org/rest/v1/core/polymer_entity_instances/{pdb_id}"
-            response = requests.get(url)
+            # Get structure info from RCSB API for additional details
+            structure_url = f"https://data.rcsb.org/rest/v1/core/entry/{pdb_id}"
+            response = requests.get(structure_url)
+            structure_data = response.json() if response.status_code == 200 else {}
             
-            if response.status_code != 200:
-                return []
-                
-            chains_data = response.json()
             sequences = []
-            
-            for chain_id in chains_data.keys():
-                entity_id = chains_data[chain_id].get('rcsb_polymer_entity_instance_container_identifiers', {}).get('entity_id')
-                
-                if entity_id:
-                    entity_url = f"https://data.rcsb.org/rest/v1/core/polymer_entity/{pdb_id}/{entity_id}"
-                    entity_response = requests.get(entity_url)
+            # Extract sequences from each chain
+            for model in structure:
+                for chain in model:
+                    sequence = ""
+                    for residue in chain:
+                        if PDB.is_aa(residue, standard=True):
+                            try:
+                                # 3글자 아미노산 코드를 1글자로 변환
+                                resname = residue.get_resname()
+                                if resname in protein_letters_3to1:
+                                    sequence += protein_letters_3to1[resname]
+                            except:
+                                continue
                     
-                    if entity_response.status_code == 200:
-                        entity_data = entity_response.json()
-                        sequence = entity_data.get('entity_poly', {}).get('pdbx_seq_one_letter_code', '')
-                        description = entity_data.get('rcsb_polymer_entity', {}).get('pdbx_description', 'N/A')
-                        
+                    if sequence:  # Only add if sequence is not empty
                         chain_info = {
-                            'chain_id': chain_id,
-                            'entity_id': entity_id,
-                            'description': description,
+                            'chain_id': chain.id,
+                            'entity_id': '1',  # Default entity ID
+                            'description': structure_data.get('struct', {}).get('title', 'N/A'),
                             'sequence': sequence,
-                            'length': len(sequence)
+                            'length': len(sequence),
+                            'resolution': structure_data.get('rcsb_entry_info', {}).get('resolution_combined', [0.0])[0],
+                            'method': structure_data.get('exptl', [{}])[0].get('method', 'Unknown'),
+                            'release_date': structure_data.get('rcsb_accession_info', {}).get('initial_release_date', 'N/A')
                         }
                         sequences.append(chain_info)
             
+            # Cleanup downloaded file
+            if os.path.exists(pdb_path):
+                os.remove(pdb_path)
+            
             return sequences
             
         except Exception as e:
-            print(f"Error in direct API call for PDB ID {pdb_id}: {str(e)}")
+            print(f"Error getting sequences for PDB ID {pdb_id}: {str(e)}")
             return []
+    
+    def __del__(self):
+        """Cleanup temporary directory on object destruction"""
+        if hasattr(self, 'pdb_dir') and os.path.exists(self.pdb_dir):
+            shutil.rmtree(self.pdb_dir)
 
     def process_query(self, query):
         """Process query and return results"""
@@ -480,37 +492,6 @@ def pdbsummary(name):
 
     return answer
 
-
-def render_html(pdb_id, chain="A"):
-    if pdb_id is None or chain is None:
-        return ""
-    html_content = f"""
-    <html>
-    <header>
-        <script src="https://3Dmol.org/build/3Dmol-min.js"></script>
-        <script src="https://3Dmol.org/build/3Dmol.ui-min.js"></script>
-    </header>
-    <body>
-        <div style="height: 400px; position: relative;" class="viewer_3Dmoljs"
-             data-pdb="{pdb_id}"
-             data-backgroundalpha="0.0"
-             data-style="cartoon:color=white"
-             data-select1="chain:{chain}"
-             data-zoomto="chain:{chain}"
-             data-style1="cartoon:color=spectrum"
-             data-spin="axis:y;speed:0.2">
-        </div>
-    </body>
-    </html>
-    """
-    iframe = f"""
-    <iframe style="width: 100%; height: 480px; border: none;" 
-            srcdoc='{html_content}'>
-    </iframe>
-    """
-    return iframe
-
-
 def create_interactive_table(df):
     if df.empty:
         return go.Figure()
@@ -578,12 +559,16 @@ app_ui = ui.page_fluid(
                 margin: 0 auto;
             }
             #query {
-                height: 100px;
+                height: 150px;
                 font-size: 16px;
                 padding: 15px;
                 width: 80%;
                 margin: 0 auto;
                 display: block;
+                white-space: pre-wrap;
+                word-wrap: break-word;
+                resize: vertical;
+                overflow-y: auto;
             }
             .content-wrapper {
                 text-align: center;
@@ -678,11 +663,12 @@ app_ui = ui.page_fluid(
                     "Search Query", 
                     {"class": "query-label", "for": "query"}
                 ),
-                ui.input_text(
+                ui.input_text_area(
                     "query", 
                     "", 
                     value="Human insulin",
-                    width="100%"
+                    width="100%",
+                    resize="vertical"
                 ),
             )
         ),
@@ -738,15 +724,6 @@ app_ui = ui.page_fluid(
                     ui.output_text("sequence_output")
                 )
             )
-        ),
-        ui.row(
-            ui.column(12,
-                ui.div(
-                    {"class": "3d-iframe", "id": "3d-iframe"}, # css 미설정 
-                    ui.h4("3D Rendering"),
-                    ui.output_ui("output_iframe")
-                )    
-            )
         )
     )
 )
@@ -759,10 +736,8 @@ def server(input, output, session):
     @reactive.Effect
     @reactive.event(input.search)
     def _():
-        # 검색 시작 시 상태 변경
         status_store.set("Searching...")
         
-        # 프롬프트 처리
         query_results = assistant.process_query(input.query())
         results_store.set(query_results)
         
@@ -770,16 +745,18 @@ def server(input, output, session):
             if not query_results["results"]:
                 status_store.set("No sequences found")
             else:
-                status_store.set("Ready")  # 검색 완료 시 Ready로 변경
+                status_store.set("Ready")
         else:
             df = pd.DataFrame(query_results["results"])
-            status_store.set("Ready")  # 검색 완료 시 Ready로 변경
-            
-            @output
-            @render_widget
-            def results_table():
-                return create_interactive_table(df)
-    
+            if df.empty:
+                status_store.set("No structures found")
+            else:
+                status_store.set("Ready")
+                @output
+                @render_widget
+                def results_table():
+                    return create_interactive_table(df)
+
     @output
     @render.text
     def search_status():
@@ -819,23 +796,10 @@ def server(input, output, session):
             
             return "\n".join(output_text)
         return ""
-    
-    @output
-    @render.text
-    def output_iframe():
-        current_results = results_store.get()
-        if current_results["type"] == "structure":
-            pdb_id = current_results["results"][0]['PDB ID']
-            # chain 가져오는 건 아직
-            return render_html(pdb_id, "A")
-        else:
-            return ""
 
 app = App(app_ui, server)
 
 if __name__ == "__main__":
     import nest_asyncio
     nest_asyncio.apply()
-    app.run(host="0.0.0.0", port=7862)
-
-
+    app.run(host="0.0.0.0", port=7862)