Update main.py

main.py

@@ -1,5 +1,8 @@
-from fastapi import FastAPI
-from typing import List, Dict, Any
+from fastapi import FastAPI, HTTPException
+from rdkit import Chem
+from rdkit.Chem import Descriptors
+from rdkit.Chem import rdMolDescriptors
+import requests
 from fastapi.middleware.cors import CORSMiddleware
 
 app = FastAPI()
@@ -12,41 +15,50 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-class Database:
-    def __init__(self):
-        self.store = {}
+def name_to_smiles(name: str) -> str:
+    url = f"https://cactus.nci.nih.gov/chemical/structure/{name}/smiles"
+    response = requests.get(url)
+    if response.status_code == 200:
+        return response.text
+    return None
 
-    def insert(self, collection_name: str, document: Dict[str, Any]):
-        if collection_name not in self.store:
-            self.store[collection_name] = []
-        
-        self.store[collection_name].append(document)
-        return document
+def get_molecule_info(mol):
+    mol_weight = Descriptors.MolWt(mol)
+    num_atoms = mol.GetNumAtoms()
+    num_bonds = mol.GetNumBonds()
+    mol_formula = rdMolDescriptors.CalcMolFormula(mol)
+    tpsa = Descriptors.TPSA(mol)
+    mol_logp = Descriptors.MolLogP(mol)
+    num_rotatable_bonds = Descriptors.NumRotatableBonds(mol)
 
-    def find(self, collection_name: str, query: Dict[str, Any] = None):
-        if collection_name not in self.store:
-            return []
+    return {
+        'molecular_weight': mol_weight,
+        'number_of_atoms': num_atoms,
+        'number_of_bonds': num_bonds,
+        'molecular_formula': mol_formula,
+        'tpsa': tpsa,
+        'logP': mol_logp,
+        'number_of_rotatable_bonds': num_rotatable_bonds
+    }
 
-        if not query:
-            return self.store[collection_name]
-        
-        return [doc for doc in self.store[collection_name] if all(item in doc.items() for item in query.items())]
+@app.get("/molecule_info/")
+async def read_molecule_info(name: str):
+    if not name:
+        raise HTTPException(status_code=400, detail="No molecule name provided")
 
-db = Database()
+    smiles = name_to_smiles(name)
+    if not smiles:
+        raise HTTPException(status_code=400, detail="Could not fetch SMILES string for provided name")
 
-@app.get("/")
-def home():
-    home = "Hello, Welcome to LinDB!"
-    return home
+    mol = Chem.MolFromSmiles(smiles)
 
-@app.get("/items", response_model=List[Dict[str, Any]])
-async def get_items():
-    return db.find('items')
+    if not mol:
+        raise HTTPException(status_code=400, detail="Molecule not recognized")
 
-@app.post("/items", response_model=Dict[str, Any])
-async def add_item(item: Dict[str, Any]):
-    return db.insert('items', item)
+    info = get_molecule_info(mol)
+
+    return info
 
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=3000)
+    uvicorn.run(app, host="0.0.0.0", port=8000)