enviPy-bayer/bb4g/__init__.py

import enum
import json
import logging
import math
from datetime import datetime
from typing import List

import requests
from django.conf import settings as s
from envipy_additional_information import register, EnviPyModel, UIConfig, WidgetType

from bridge.contracts import Classifier  # noqa: I001
from bridge.dto import (
    BuildResult,
    EnviPyDTO,
    EvaluationResult,
    RunResult,
    TransformationProductPrediction,
)  # noqa: I001

logger = logging.getLogger("epdb")


class SamplingAlgorithm(enum.Enum):
    EXACT = "exact"


@register("bb4gconfig")
class BB4GConfig(EnviPyModel):
    sampling_algorithm: SamplingAlgorithm = SamplingAlgorithm.EXACT
    cutoff: int = -5

    class UI:
        title = "BB4G Configuration"
        sampling_algorithm = UIConfig(
            widget=WidgetType.SELECT,
            label="BB4G Sampling Algorithm",
            order=1,
            placeholder="If unset defaults to 'exact'"
        )
        cutoff = UIConfig(
            widget=WidgetType.NUMBER,
            label="BB4G Cutoff",
            order=2,
            placeholder="If unset defaults to -5"
        )


# Once stable these will be exposed by enviPy-plugins lib
class BB4G(Classifier):
    Config = BB4GConfig

    def __init__(self, config: BB4GConfig | None = None):
        super().__init__(config)
        self.url = f"{s.BB4G_URL}"

        self.token = self.acquire_token()
        self.header = {
            "Authorization": f"Bearer {self.token}",
            "Content-Type": "application/json",
        }

    def acquire_token(self):
        BB4G_TENANT_ID = s.BB4G_TENANT_ID
        BB4G_CLIENT_ID = s.BB4G_CLIENT_ID
        BB4G_CLIENT_SECRET = s.BB4G_CLIENT_SECRET
        BB4G_SCOPE = s.BB4G_SCOPE

        BB4G_TOKEN_URL = f"https://login.microsoftonline.com/{BB4G_TENANT_ID}/oauth2/v2.0/token"

        payload = {
            "client_id": BB4G_CLIENT_ID,
            "client_secret": BB4G_CLIENT_SECRET,
            "scope": BB4G_SCOPE,
            "grant_type": "client_credentials"
        }

        # No Proxy required, URL is whitelisted
        res = requests.post(BB4G_TOKEN_URL, data=payload)

        res.raise_for_status()

        return res.json()["access_token"]

    def start(self):
        header = {
            "Authorization": f"Bearer {self.token}",
            "Content-Type": "application/json",
        }

        started = False

        while not started:
            res = requests.post(f"{self.url}/start", headers=header, data={}, proxies=s.PROXIES or None)
            logger.info(f"Starting BB4G: {res.status_code}")
            if res.status_code == 200:
                started = True
            elif res.status_code in [500, 502]:
                import time
                time.sleep(5)
            else:
                raise ValueError(f"Unexpected status code: {res.status_code}")

    @classmethod
    def requires_rule_packages(cls) -> bool:
        return False

    @classmethod
    def requires_data_packages(cls) -> bool:
        return False

    @classmethod
    def identifier(cls) -> str:
        return "bb4g"

    @classmethod
    def name(cls) -> str:
        return "BB4G Template Free Model"

    @classmethod
    def display(cls) -> str:
        return "BB4G Template Free Model"

    def build(self, eP: EnviPyDTO, *args, **kwargs) -> BuildResult | None:
        return

    def run(self, eP: EnviPyDTO, *args, **kwargs) -> RunResult:

        # Ensure Service is running
        self.start()

        smiles = [c.smiles for c in eP.get_compounds()]
        preds = self._post(smiles)

        results = []

        for substrate in preds.keys():
            results.append(
                TransformationProductPrediction(
                    substrate=substrate,
                    products=preds[substrate],
                )
            )

        return RunResult(
            producer=eP.get_context().url,
            description=f"Generated at {datetime.now()}",
            result=results,
        )

    def evaluate(self, eP: EnviPyDTO, *args, **kwargs) -> EvaluationResult:
        pass

    def build_and_evaluate(self, eP: EnviPyDTO, *args, **kwargs) -> EvaluationResult:
        pass

    def _post(self, smiles: List[str]) -> dict[str, dict[str, float]]:
        header = {
            "Authorization": f"Bearer {self.token}",
            "Content-Type": "application/json",
        }

        result = {}

        for smi in smiles:
            data = {
                "smiles": smi,
                "sampling_alg": self.config.sampling_algorithm.value,
                "cutoff": self.config.cutoff,
            }

            retries = 0
            while retries < 100:
                resp = requests.post(f"{self.url}/compute", headers=header, data=json.dumps(data),
                                     proxies=s.PROXIES or None)

                if resp.status_code == 418:
                    retries += 1
                    logger.info(f"BB4G predict hit a 418, retrying in 60 seconds")
                    import time
                    time.sleep(3)
                    continue

                resp.raise_for_status()

                for substrate, predictions in resp.json().items():
                    preds = {}

                    for pred in predictions:
                        prod = pred["prediction"]
                        prob = math.exp(pred["log_likelihood"])
                        preds[prod] = prob

                    result[substrate] = preds

                break

        return result