utils.py

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]