💡 Add monitors to SDK

Author: momegas

README.md

@@ -87,7 +87,7 @@ ENV=dev uvicorn whitebox.main:app --reload
 
 ### Tests:
 
-- Run: `ENV=test pytest -s`
+- Run: `ENV=test pytest` or `ENV=test pytest -s` to preserve logs.
 - Watch: `ENV=test ptw`
 - Run test coverage `ENV=test coverage run -m pytest`
 - Look at coverage report: `coverage report` or `coverage html` to generate an html. To view it in your browser open the `htmlcov/index.html` file.

whitebox/cron_tasks/shared.py

@@ -1,4 +1,4 @@
-from typing import List, Union
+from typing import List, Tuple, Union
 
 import pandas as pd
 from sqlalchemy.orm import Session
@@ -24,7 +24,7 @@ async def get_model_dataset_rows_df(db: Session, model_id: str) -> pd.DataFrame:
 
 async def get_model_inference_rows_df(
     db: Session, model_id: str
-) -> tuple[pd.DataFrame, pd.DataFrame, pd.Series]:
+) -> Tuple[pd.DataFrame, pd.DataFrame, pd.Series]:
     inference_rows_in_db = crud.inference_rows.get_inference_rows_by_model(
         db=db, model_id=model_id
     )

whitebox/sdk/whitebox.py

@@ -2,11 +2,18 @@
 import numpy as np
 import pandas as pd
 from whitebox.schemas.model import FeatureTypes, ModelCreateDto, ModelType
-from typing import Dict
+from typing import Dict, Optional
 import requests
 import logging
 from fastapi import status
 
+from whitebox.schemas.modelMonitor import (
+    AlertSeverity,
+    ModelMonitorCreateDto,
+    MonitorMetrics,
+    MonitorStatus,
+)
+
 
 class APiVersion(str, Enum):
     v1 = "v1"
@@ -177,3 +184,38 @@ def _check_processed_and_non_processed_length(
                 "Processed and non processed dataframes must have the same length."
             )
         return True
+
+    def create_model_monitor(
+        self,
+        model_id: str,
+        name: str,
+        status: MonitorStatus,
+        metric: MonitorMetrics,
+        feature: Optional[str],
+        lower_threshold: Optional[float],
+        severity: AlertSeverity,
+        email: str,
+    ):
+        """
+        Creates a monitor for a model.
+        """
+
+        model_monitor = ModelMonitorCreateDto(
+            model_id=model_id,
+            name=name,
+            status=status,
+            metric=metric,
+            feature=feature,
+            lower_threshold=lower_threshold,
+            severity=severity,
+            email=email,
+        )
+
+        result = requests.post(
+            url=f"{self.host}/{self.api_version}/model-monitors",
+            json=model_monitor.dict(),
+            headers={"api-key": self.api_key},
+        )
+
+        logger.info(result.json())
+        return result.json()

whitebox/tests/utils/maps.py

@@ -33,4 +33,5 @@
     "sdk_delete_model",
     "sdk_log_training_dataset",
     "sdk_log_inferences",
+    "sdk_create_model_monitor",
 ]

whitebox/tests/v1/test_sdk.py

@@ -1,5 +1,6 @@
 import pandas as pd
 import pytest
+from whitebox.schemas.modelMonitor import AlertSeverity, MonitorMetrics, MonitorStatus
 from whitebox.sdk import Whitebox
 from whitebox.tests.v1.conftest import get_order_number, state, state_sdk
 from whitebox.tests.v1.mock_data import (
@@ -171,3 +172,26 @@ def test_sdk_log_inferences(client):
             timestamps=timestamps,
             actuals=mixed_actuals,
         )
+
+
+@pytest.mark.order(get_order_number("sdk_create_model_monitor"))
+def test_sdk_create_model_monitor(client):
+    with requests_mock.Mocker() as m:
+        m.post(
+            url=f"{state_sdk.wb.host}/v1/model-monitors",
+            json=model_multi_create_payload,
+            headers={"api-key": state_sdk.wb.api_key},
+        )
+
+        model_monitor = state_sdk.wb.create_model_monitor(
+            model_id="mock_model_id",
+            name="test",
+            status=MonitorStatus.active,
+            metric=MonitorMetrics.accuracy,
+            feature="feature1",
+            lower_threshold=0.7,
+            severity=AlertSeverity.high,
+            email="jaclie.chan@chinamail.io",
+        )
+
+        assert model_monitor == model_monitor