Create utils.py

utils.py

@@ -0,0 +1,70 @@
+from typing import List
+import requests
+import os
+
+query_templates = {
+    "SMILES": """query BySMILES($smiles: String!, $limit: Int = {}) {{
+        molecules(SMILES: $smiles, limit: $limit) {{
+            {}
+        }}
+    }}""",
+    "IUPAC": """query ByIUPAC($iupac: String!, $limit: Int = {}) {{
+        molecules(IUPAC: $iupac, limit: $limit) {{
+            {}
+        }}
+    }}""",
+    "NAME": """query ByName($name: String!, $limit: Int = {}) {{
+        molecules(name: $name, limit: $limit) {{
+            {}
+        }}
+    }}""",
+    "InChI": """query ByInChI($inchi: String!, $limit: Int = {}) {{
+        molecules(InChI: $inchi, limit: $limit) {{
+            {}
+        }}
+    }}""",
+    "InChIKey": """query ByInChIKey($inchikey: String!, $limit: Int = {}) {{
+        molecules(InChIKey: $inchikey, limit: $limit) {{
+            {}
+        }}
+    }}""",
+    "CID": """query ByCID($cid: Int!) {{
+        molecules(CID: $cid) {{
+            {}
+        }}
+    }}"""
+}
+
+SUPPORTED_RETURN_FIELDS = ["CID", "IUPAC", "SMILES", "InChI", "InChIKey", "synonyms"]
+API_URL = os.getenv('API_URL')
+
+class DbProcessor:
+    def __init__(self, url: str = API_URL) -> None:
+        self.url = url
+
+    def request_data(self, text_input: str, query_type: str, return_fields: List[str], limit: int = 1):
+        if query_type not in query_templates:
+            raise ValueError(f"Query type '{query_type}' is not supported.")
+
+        for field in return_fields:
+            if field not in SUPPORTED_RETURN_FIELDS:
+                raise RuntimeError(f"Field '{field}' is not supported, try to use one from {SUPPORTED_RETURN_FIELDS}")
+
+        return_fields_str = "\n".join(return_fields)
+
+        query = query_templates[query_type].format(limit, return_fields_str)
+        variables = {query_type.lower(): text_input}
+
+        payload = {
+            "query": query,
+            "variables": variables
+        }
+
+        response = requests.post(
+            url=f"{self.url}/{query_type.lower()}",
+            json=payload
+        )
+
+        response.raise_for_status()
+
+        return response.json()['molecules'][0]