WIP: add typechecking for allowed files

vitalii-bulyzhyn · vitalii-bulyzhyn · commit c816c47e1409 · 2022-10-05T11:56:34.000-03:00
diff --git a/sdk/diffgram/models/base_model.py b/sdk/diffgram/models/base_model.py
@@ -1,7 +1,22 @@
+from fastapi import HTTPException
+from typing import List, Literal
 from .model_interfaces import DiffgramFile, Prediction, Attribute, Instance
 class DiffgramBaseModel():
-    def init(self):
-        pass
+    diffgram_allowed_types = Literal['image', 'frame', 'video', 'text', 'audio', 'sensor_fusion', 'geospatial']
+
+    def __init__(
+        self, 
+        allowed_types: list = None
+    ):
+        if allowed_types is not None:
+            if not isinstance(allowed_types, list):
+                raise ValueError('allowed_types must be of type list')
+
+            for allowed_type in allowed_types:
+                if allowed_type not in self.diffgram_allowed_types:
+                    raise ValueError(f"{allowed_type} is not valid Diffgram file type")
+            
+            self.allowed_types = allowed_types
 
     def infere(self, file: DiffgramFile) -> Prediction:
         raise NotImplementedError
@@ -14,7 +29,11 @@ def ping(self):
 
     def serve(self, app):
         @app.post("/infere")
-        async def predict(file: DiffgramFile):
+        async def infere_route(file: DiffgramFile):
+            if self.allowed_types is not None:
+                if file.type not in self.allowed_types:
+                    raise HTTPException(status_code=404, detail=f"This model does not support {file.type} files")
+
             predictions = self.infere(file)
 
             if not isinstance(predictions, Prediction):
@@ -38,7 +57,7 @@ async def predict(file: DiffgramFile):
             }
 
         @app.get("/get_schema")
-        async def schema():
+        async def get_schema_route():
             return {
                 "message": "Get schema here"
             }
diff --git a/sdk/diffgram/models/model_interfaces.py b/sdk/diffgram/models/model_interfaces.py
@@ -1,10 +1,8 @@
 from pydantic import BaseModel
-from typing import Optional
-from typing import List
-
+from typing import Optional, List, Literal
 class DiffgramFile(BaseModel):
     id: int
-    type: str
+    type: Literal['image', 'frame', 'video', 'text', 'audio', 'sensor_fusion', 'geospatial']
 
 class Attribute(BaseModel):
     id: int
diff --git a/sdk/samples/model_server.py b/sdk/samples/model_server.py
@@ -11,4 +11,4 @@ def get_schema(self):
         return super().get_schema()
 
 app = FastAPI()
-MyTestModel().serve(app)
+MyTestModel(allowed_types=["image"]).serve(app)
