maykcaldas/MAPI_LLM

1from langchain.agents import Tool, tool2import requests3from langchain import OpenAI4from langchain import LLMMathChain, SerpAPIWrapper5import os6from rdkit import Chem7 8@tool9def query2smiles(text):10  '''This function queries the one given molecule name and returns a SMILES string from the record'''11  try:#query the PubChem database12    r = requests.get('https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/name/' + text + '/property/IsomericSMILES/JSON')13    #convert the response to a json object14    data = r.json() 15    #return the SMILES string16    smi = data['PropertyTable']['Properties'][0]['IsomericSMILES']17    # remove salts18    return smi19  except:20    f"Could not find the IUPAC name for {text}"21 22@tool23def smiles2IUPAC(text):24  '''This function queries the one given smiles name and returns a IUPAC name from the record'''25  #query the PubChem database26  try:27    r = requests.get('https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/smiles/' + text + '/property/IUPACName/JSON')28    data = r.json()29    smi = data["PropertyTable"]["Properties"][0]["IUPACName"]30    return smi31  except:32    return f"Could not find the IUPAC name for {text}"33 34@tool35def formula2IUPAC(text):36  '''This function queries the one given chemical formula and returns a material name from the record.'''37  try:38    r = requests.get('https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/formula/' + text + '/property/IUPACName/JSON')39    data = r.json()40    print(data)41    smi = data["PropertyTable"]["Properties"][0]["IUPACName"]42    return smi43  except:44    return f"Could not find the IUPAC name for {text}"45 46@tool47def name2formula(text):48  '''This function queries the one given material name and returns a chemical formula from the record.'''49  try:50    r = requests.get('https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/name/' + text + '/property/MolecularFormula/JSON')51    data = r.json()52    print(data)53    smi = data["PropertyTable"]["Properties"][0]["MolecularFormula"]54    return smi55  except:56    return f"Could not find the molecular formula for {text}"57 58@tool59def canonicalizeSMILES(smiles):60  '''Given a smiles representation, this function returns a canonicalized version of the same smiles.61  It's better to search for molecules in its canonicalized form'''62  return Chem.MolToSmiles(Chem.MolFromSmiles(smiles))63 64@tool65def web_search(keywords, search_engine="google"):66  '''Useful to do a simple google search. 67      Use this tool to find general information from websites.68      Use keywords for your search. 69  '''70  return SerpAPIWrapper(71    serpapi_api_key=os.getenv("SERP_API_KEY"),72    search_engine=search_engine73  ).run(keywords)74 75@tool76def LLM_predict(prompt):77  ''' This function receives a prompt generate with context by the create_context_prompt tool and request a completion to a language model. Then returns the completion'''78  llm = OpenAI(79        model_name='text-ada-001',  #TODO: Maybe change to gpt-4 when ready80        temperature=0.7,81        n=1,82        best_of=5,83        top_p=1.0,84        stop=["\n\n", "###", "#", "##"],85        # model_kwargs=kwargs,86    )87  return llm.generate([prompt]).generations[0][0].text88 89common_tools = [90    query2smiles,91    smiles2IUPAC,92    # formula2IUPAC,93    # name2formula,94    canonicalizeSMILES,95    web_search,96    LLM_predict97]