Upload app.py

app.py

@@ -36,7 +36,7 @@ class PDBSearchAssistant:
             Extract specific search parameters from the protein-related query:
             1. Protein name or type
             2. Resolution cutoff (in Å)
-            3. Sequence information
+            3. Protein sequence information
             4. Specific PDB ID
             5. Experimental method (X-RAY, EM, NMR)
             6. Organism/Species information
@@ -58,6 +58,46 @@ class PDBSearchAssistant:
             PDB_ID: none
             Method: X-RAY
 
+            Query: "Get sequence of PDB ID 8ET6"
+            Protein: none
+            Organism: none
+            Resolution: none
+            Sequence: none
+            PDB_ID: 8ET6
+            Method: none
+
+            Query: "Sequence of 7BZ5"
+            Protein: none
+            Organism: none
+            Resolution: none
+            Sequence: none
+            PDB_ID: 7BZ5
+            Method: none
+
+            Query: "7BZ5"
+            Protein: none
+            Organism: none
+            Resolution: none
+            Sequence: none
+            PDB_ID: 7BZ5
+            Method: none
+
+            Query: "6KAO"
+            Protein: none
+            Organism: none
+            Resolution: none
+            Sequence: none
+            PDB_ID: 6KAO
+            Method: none
+            
+            Query: "Find structures containing sequence MNIFEMLRIDEGLRLKIYKDTEGYYTIGIGHLLTKSPSLNAAKSELDKAIGRNTNGVITKDEAEKLFNQDVDAAVRGILRNAKLKPVYDSLDAVRRAALINMVFQMGETGVAGFTNSLRMLQQKRWDEAAVNLAKSRWYNQTPNRAKRVITTFRTGTWDAYKNL"
+            Protein: none
+            Organism: none
+            Resolution: none
+            Sequence: MNIFEMLRIDEGLRLKIYKDTEGYYTIGIGHLLTKSPSLNAAKSELDKAIGRNTNGVITKDEAEKLFNQDVDAAVRGILRNAKLKPVYDSLDAVRRAALINMVFQMGETGVAGFTNSLRMLQQKRWDEAAVNLAKSRWYNQTPNRAKRVITTFRTGTWDAYKNL
+            PDB_ID: none
+            Method: none
+
             Now analyze:
             Query: {query}
             """
@@ -380,52 +420,65 @@ class PDBSearchAssistant:
             print(f"Error in direct API call for PDB ID {pdb_id}: {str(e)}")
             return []
 
-    def analyze_query_type(self, query):
-        """Analyze query type and extract relevant information"""
-        print(f"\nAnalyzing query: '{query}'")  # 입력된 쿼리 출력
-        query = query.lower().strip()
-        print(f"Lowercase query: '{query}'")  # 소문자로 변환된 쿼리 출력
-        
-        # Check for sequence query pattern
-        sequence_patterns = [
-            r"sequence\s+of\s+pdb\s+id\s+([a-zA-Z0-9]{4})",
-            r"sequence\s+for\s+pdb\s+id\s+([a-zA-Z0-9]{4})",
-            r"get\s+sequence\s+([a-zA-Z0-9]{4})",
-            r"([a-zA-Z0-9]{4})\s+sequence"
-        ]
-        
-        for i, pattern in enumerate(sequence_patterns):
-            print(f"Trying pattern {i+1}: {pattern}")  # 각 패턴 시도 출력
-            match = re.search(pattern, query)
-            if match:
-                pdb_id = match.group(1).upper()
-                print(f"Match found! PDB ID: {pdb_id}")  # 매치된 PDB ID 출력
+    def process_query(self, query):
+        """Process query and return results"""
+        try:
+            # Get search parameters from LLM
+            formatted_prompt = self.prompt_template.format(query=query)
+            response = self.pipe(formatted_prompt)[0]['generated_text']
+            print("Generated parameters:", response)
+            
+            # Parse LLM response for PDB ID
+            pdb_id = None
+            for line in response.split('\n'):
+                if 'PDB_ID:' in line:
+                    value = line.split('PDB_ID:')[1].strip()
+                    if value.lower() not in ['none', 'n/a']:
+                        pdb_id = value.upper()
+                        break
+            
+            # Check if query is asking for sequence
+            sequence_keywords = ['sequence', 'seq']
+            is_sequence_query = any(keyword in query.lower() for keyword in sequence_keywords)
+            
+            if is_sequence_query and pdb_id:
+                # Get sequences for the PDB ID
+                sequences = self.get_sequences_by_pdb_id(pdb_id)
                 return {
                     "type": "sequence",
-                    "pdb_id": pdb_id
+                    "results": sequences
                 }
-        
-        print("No sequence pattern matched, treating as structure search")  # 구조 검색으로 처리
-        return {
-            "type": "structure",
-            "query": query
-        }
-
-    def process_query(self, query):
-        """Process query and return appropriate results"""
-        query_info = self.analyze_query_type(query)
-        
-        if query_info["type"] == "sequence":
-            return {
-                "type": "sequence",
-                "results": self.get_sequences_by_pdb_id(query_info["pdb_id"])
-            }
-        else:
+            
+            # If not a sequence query or no PDB ID found, proceed with normal structure search
             return {
                 "type": "structure",
-                "results": self.search_pdb(query_info["query"])
+                "results": self.search_pdb(query)
             }
+            
+        except Exception as e:
+            print(f"Error processing query: {str(e)}")
+            return {"type": "structure", "results": []}
+
+def pdbsummary(name):
+
+    search_engine = ProteinSearchEngine()
+
+    query = ProteinQuery(
+        name,
+        max_resolution= 5.0
+    )
+
+    results = search_engine.search(query)
+
+    answer = ""
+    for i, structure in enumerate(results, 1):
+        answer += f"\n{i}. PDB ID : {structure.pdb_id}\n"
+        answer += f"\nResolution : {structure.resolution:.2f} A \n"
+        answer += f"Method : {structure.method}\n Title : {structure.title}\n"
+        answer += f"Release Date : {structure.release_date}\n Sequence length: {len(structure.sequence)} aa\n"
+        answer += f"    Sequence:\n {structure.sequence}\n"
 
+    return answer
 
 def create_interactive_table(df):
     if df.empty:
@@ -610,8 +663,7 @@ app_ui = ui.page_fluid(
                     ui.tags.ul(
                         ui.tags.li("Human hemoglobin C resolution better than 2.5Å"),
                         ui.tags.li("Find structures containing sequence MNIFEMLRIDEGLRLKIYKDTEGYYTIGIGHLLTKSPSLNAAKSELDKAIGRNTNGVITKDEAEKLFNQDVDAAVRGILRNAKLKPVYDSLDAVRRAALINMVFQMGETGVAGFTNSLRMLQQKRWDEAAVNLAKSRWYNQTPNRAKRVITTFRTGTWDAYKNL"),
-                        ui.tags.li("Sequence of PDB ID 8ET6"),
-                        ui.tags.li("Get sequence 7BZ5")
+                        ui.tags.li("Sequence of PDB ID 8ET6")
                     )
                 )
             )
@@ -707,8 +759,6 @@ def server(input, output, session):
     @render.text
     def sequence_output():
         current_results = results_store.get()
-        print(current_results["type"])
-        print(current_results["results"])
         if current_results["type"] == "sequence":
             sequences = current_results["results"]
             if not sequences: