minor

bb4g/__init__.py

@@ -2,9 +2,11 @@ import json
 import math
 from datetime import datetime
 from typing import List
-
+import enum
 import requests
 from django.conf import settings as s
+from envipy_additional_information import EnviPyModel, UIConfig, WidgetType
+from envipy_additional_information import register
 
 from bridge.contracts import Classifier  # noqa: I001
 from bridge.dto import (
@@ -15,12 +17,36 @@ from bridge.dto import (
     TransformationProductPrediction,
 )  # noqa: I001
 
+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 = None
+    Config = BB4GConfig
 
-    def __init__(self, config=None):
+    def __init__(self, config: BB4GConfig | None = None):
         super().__init__(config)
         self.url = f"{s.BB4G_URL}"
 
@@ -29,10 +55,6 @@ class BB4G(Classifier):
             "Authorization": f"Bearer {self.token}",
             "Content-Type": "application/json",
         }
-        self.proxies = {
-            "http": s.HTTP_PROXY,
-            "https": s.HTTPS_PROXY,
-        }
 
     def acquire_token(self):
         BB4G_TENANT_ID = s.BB4G_TENANT_ID
@@ -65,7 +87,7 @@ class BB4G(Classifier):
         started = False
         retries = 0
         while not started and retries < 5:
-            res = requests.post(f"{self.url}/start", headers=header, data={}, proxies=self.proxies)
+            res = requests.post(f"{self.url}/start", headers=header, data={}, proxies=s.PROXIES or None)
 
             if res.status_code == 200:
                 started = True
@@ -140,11 +162,11 @@ class BB4G(Classifier):
         for smi in smiles:
             data = {
                 "smiles": smi,
-                "sampling_alg": "exact",
-                "cutoff": -5,
+                "sampling_alg": self.config.sampling_algorithm.value,
+                "cutoff": self.config.cutoff,
             }
 
-            resp = requests.post(f"{self.url}/compute", headers=header, data=json.dumps(data), proxies=self.proxies)
+            resp = requests.post(f"{self.url}/compute", headers=header, data=json.dumps(data), proxies=s.PROXIES or None)
 
             resp.raise_for_status()