Update app.py

app.py

@@ -12,7 +12,6 @@ from langchain_huggingface import ChatHuggingFace
 from langchain_core.prompts import PromptTemplate, ChatPromptTemplate
 from langchain_core.runnables import chain
 from uuid import uuid4
-import re
 import matplotlib.pyplot as plt
 
 from rdkit import Chem
@@ -32,6 +31,9 @@ from chembl_webresource_client.new_client import new_client
 from tqdm.auto import tqdm
 import requests 
 import spaces
+from rcsbapi.search import TextQuery
+import requests
+import itertools
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -57,6 +59,7 @@ class State(TypedDict):
   which_tool: int
   props_string: str
   loop_again: str
+  which_pdbs: int
   #(Literal["lipinski_tool", "substitution_tool", "pharm_feature_tool"],
   #                   Literal["lipinski_tool", "substitution_tool", "pharm_feature_tool"])
 
@@ -394,6 +397,47 @@ def pdb_node(state: State) -> State:
   state["which_tool"] += 1
   return state
 
+  def find_node(state: State) -> State:
+  '''
+    Accepts a protein name and searches the protein databack for PDB IDs that match along with the entry titles. 
+
+      Args:
+        protein_name: the protein to query
+      Returns:
+        props_string: a string of the 
+  '''
+  test_protein = state["query_protein"].strip()
+  which_pdbs = state["which_pdbs"]
+  current_props_string = state["props_string"]
+
+  print(f"find tool using {test_protein}")
+  print('===================================================')
+
+  try:
+    query = TextQuery(value=test_protein)
+    results = query()
+
+    def pdb_gen():
+      for rid in results:
+        yield(rid)
+
+    take10 = itertools.islice(pdb_gen(), which_pdbs, which_pdbs+10, 1)
+
+    pdb_string = f'10 PDBs that match the protein {test_protein} are: \n'
+    for pdb in take10:
+      data = requests.get(f"https://data.rcsb.org/rest/v1/core/entry/{pdb}").json()
+      title = data['struct']['title']
+      pdb_string += f'PDB ID: {pdb}, with title: {title} \n'
+    state["which_pdbs"] = which_pdbs+10
+  except:
+    pdb_string = ''
+  
+
+  current_props_string += pdb_string
+  state["props_string"] = current_props_string
+  state["which_tool"] += 1
+  return state
+
 def first_node(state: State) -> State:
   '''
     The first node of the agent. This node receives the input and asks the LLM
@@ -427,6 +471,7 @@ def first_node(state: State) -> State:
   props_string = ""
   state["props_string"] = props_string
   state["loop_again"] = None
+  state['which_pdbs'] = 0
 
   raw_input = state["messages"][-1].content
   parts = raw_input.split(',')
@@ -471,7 +516,8 @@ It returns a string containing the protein ID, gene name, organism, and protein
 list_bioactives_tool: Accepts a given UNIPROT ID and searches for bioactive molecules \n \
 get_bioactives_tool: Accepts a Chembl ID and get all bioactives molecule SMILES and IC50s for that ID\n \
 pdb_tool: Accepts a PDB ID and queires the protein databank for the number of chains in and sequence of the \n \
-protein, as well as other information such as ligands in the structure.\
+protein, as well as other information such as ligands in the structure. \
+find_tool: Accepts a protein name and seaches for PDB IDs that match, returning the PDB ID and the title. \
 '
   res = chat_model.invoke(prompt)
 
@@ -538,7 +584,8 @@ It returns a string containing the protein ID, gene name, organism, and protein
 list_bioactives_tool: Accepts a given UNIPROT ID and searches for bioactive molecules \n \
 get_bioactives_tool: Accepts a Chembl ID and get all bioactives molecule SMILES and IC50s for that ID\n \
 pdb_tool: Accepts a PDB ID and queires the protein databank for the number of chains in and sequence of the \
-protein, as well as other information such as ligands in the structure. \n'
+protein, as well as other information such as ligands in the structure. \
+find_tool: Accepts a protein name and seaches for PDB IDs that match, returning the PDB ID and the title.'
 
   res = chat_model.invoke(prompt)
 
@@ -712,6 +759,7 @@ builder.add_node("uniprot_node", uniprot_node)
 builder.add_node("listbioactives_node", listbioactives_node)
 builder.add_node("getbioactives_node", getbioactives_node)
 builder.add_node("pdb_node", pdb_node)
+builder.add_node("find_node", find_node)
 
 builder.add_node("loop_node", loop_node)
 builder.add_node("parser_node", parser_node)
@@ -724,6 +772,7 @@ builder.add_conditional_edges("first_node", get_chemtool, {
     "list_bioactives_tool": "listbioactives_node",
     "get_bioactives_tool": "getbioactives_node",
     "pdb_tool": "pdb_node",
+    "find_tool": "find_node",
     None: "parser_node"})
 
 builder.add_conditional_edges("retry_node", get_chemtool, {
@@ -731,18 +780,21 @@ builder.add_conditional_edges("retry_node", get_chemtool, {
     "list_bioactives_tool": "listbioactives_node",
     "get_bioactives_tool": "getbioactives_node",
     "pdb_tool": "pdb_node",
+    "find_tool": "find_node",
     None: "parser_node"})
 
 builder.add_edge("uniprot_node", "loop_node")
 builder.add_edge("listbioactives_node", "loop_node")
 builder.add_edge("getbioactives_node", "loop_node")
 builder.add_edge("pdb_node", "loop_node")
+builder.add_edge("find_node": "loop_node")
 
 builder.add_conditional_edges("loop_node", get_chemtool, {
     "uniprot_tool": "uniprot_node",
     "list_bioactives_tool": "listbioactives_node",
     "get_bioactives_tool": "getbioactives_node",
     "pdb_tool": "pdb_node",
+    "find_tool": "find_node",
     None: "parser_node"})
 
 builder.add_conditional_edges("parser_node", loop_or_not, {
@@ -793,6 +845,7 @@ with gr.Blocks(fill_height=True) as forest:
               - calls Chembl to find hits for a given uniprot id and reports number of bioactive molecules in the hit
               - calls Chembl to find a list bioactive molecules for a given chembl id and their IC50 values
               - calls PDB to find the number of chains in a protein, proteins sequences and small molecules in the structure
+              - calls PDB to find PDB IDs that match a protein name.
               ''')
 
   with gr.Row():