Refactoring parse_message()

lib/http.py

@@ -36,7 +36,7 @@ def process_item(process_name: str, message: Dict[str, Any]):
     queue_prefix = Queue.get_queue_prefix()
     queue_suffix = Queue.get_queue_suffix()
     queue = QueueWorker.create(process_name)
-    queue.push_message(f"{queue_prefix}{process_name}{queue_suffix}", schemas.parse_message({"body": message, "model_name": process_name}))
+    queue.push_message(f"{queue_prefix}{process_name}{queue_suffix}", schemas.parse_input_message({"body": message, "model_name": process_name}))
     return {"message": "Message pushed successfully", "queue": process_name, "body": message}
 
 @app.post("/trigger_callback")

lib/model/audio.py

@@ -1,4 +1,4 @@
-from typing import Union, List, Dict
+from typing import Union, List, Dict, Any
 import os
 import tempfile
 
@@ -27,3 +27,18 @@ def process(self, audio: schemas.Message) -> Dict[str, Union[str, List[int]]]:
         finally:
             os.remove(temp_file_name)
         return {"hash_value": hash_value}
+
+    @classmethod
+    def validate_input(cls, data: Dict) -> None:
+        """
+        Validate input data. Must be implemented by all child "Model" classes.
+        """
+        pass
+
+
+    @classmethod
+    def parse_input_message(cls, data: Dict) -> Any:
+        """
+        Validate input data. Must be implemented by all child "Model" classes.
+        """
+        return None

lib/model/classycat.py

@@ -1,7 +1,7 @@
-from typing import Union
+from typing import Union, Dict, Any
 from lib.logger import logger
 from lib.model.model import Model
-from lib.schemas import Message, ClassyCatSchemaResponse, ClassyCatBatchClassificationResponse
+from lib.schemas import Message, ClassyCatSchemaResponse, ClassyCatBatchClassificationResponse, ClassyCatResponse
 from lib.model.classycat_classify import Model as ClassifyModel
 from lib.model.classycat_schema_create import Model as ClassyCatSchemaCreateModel
 from lib.model.classycat_schema_lookup import Model as ClassyCatSchemaLookupModel
@@ -23,3 +23,40 @@ def process(self, message: Message) -> Union[ClassyCatSchemaResponse, ClassyCatB
         else:
             logger.error(f"Unknown event type {event_type}")
             raise PrestoBaseException(f"Unknown event type {event_type}", 422)
+
+    @classmethod
+    def validate_input(cls, data: Dict) -> None:
+        """
+        Validate input data. Must be implemented by all child "Model" classes.
+        """
+        event_type = data['parameters']['event_type']
+
+        if event_type == 'classify':
+            ClassifyModel.validate_input(data)
+        elif event_type == 'schema_lookup':
+            ClassyCatSchemaLookupModel.validate_input(data)
+        elif event_type == 'schema_create':
+            ClassyCatSchemaCreateModel.validate_input(data)
+        else:
+            logger.error(f"Unknown event type {event_type}")
+            raise PrestoBaseException(f"Unknown event type {event_type}", 422)
+
+    @classmethod
+    def parse_input_message(cls, data: Dict) -> Any:
+        """
+        Parse input into appropriate response instances.
+        """
+        event_type = data['parameters']['event_type']
+
+        if event_type == 'classify':
+            result_instance_class = ClassifyModel
+        elif event_type == 'schema_lookup':
+            result_instance_class = ClassyCatSchemaLookupModel
+        elif event_type == 'schema_create':
+            result_instance_class = ClassyCatSchemaCreateModel
+
+        else:
+            logger.error(f"Unknown event type {event_type}")
+            raise PrestoBaseException(f"Unknown event type {event_type}", 422)
+
+        return result_instance_class.parse_input_message(data)

lib/model/classycat_classify.py

@@ -1,3 +1,4 @@
+from typing import Dict, Any
 import os
 import json
 import uuid
@@ -230,3 +231,35 @@ def process(self, message: Message) -> ClassyCatBatchClassificationResponse:
                 raise e
             else:
                 raise PrestoBaseException(f"Error classifying items: {e}", 500) from e
+
+
+    @classmethod
+    def validate_input(cls, data: Dict) -> None:
+        """
+        Validate input data. Must be implemented by all child "Model" classes.
+        """
+        if "schema_id" not in data["parameters"] or data["parameters"]["schema_id"] == "":
+            raise PrestoBaseException("schema_id is required as input to classify", 422)
+
+        if "items" not in data["parameters"] or len(data["parameters"]["items"]) == 0:
+            raise PrestoBaseException("items are required as input to classify", 422)
+
+        for item in data["parameters"]["items"]:
+            if "id" not in item or item["id"] == "":
+                raise PrestoBaseException("id is required for each item", 422)
+            if "text" not in item or item["text"] == "":
+                raise PrestoBaseException("text is required for each item", 422)
+
+    @classmethod
+    def parse_input_message(cls, data: Dict) -> Any:
+        """
+        Parse input into appropriate response instances.
+        """
+        event_type = data['parameters']['event_type']
+        result_data = data.get('result', {})
+
+        if event_type == 'classify':
+            return ClassyCatBatchClassificationResponse(**result_data)
+        else:
+            logger.error(f"Unknown event type {event_type}")
+            raise PrestoBaseException(f"Unknown event type {event_type}", 422)

lib/model/classycat_schema_create.py

@@ -1,3 +1,4 @@
+from typing import Dict, Any
 import os
 import json
 import uuid
@@ -210,7 +211,8 @@ def process(self, message: Message) -> ClassyCatSchemaResponse:
             raise PrestoBaseException(f"Error creating schema: {e}", 500) from e
 
 
-    def verify_schema_parameters(self, schema_name, topics, examples, languages): #todo
+    @classmethod
+    def verify_schema_parameters(cls, schema_name, topics, examples, languages):
 
         if not schema_name or not isinstance(schema_name, str) or len(schema_name) == 0:
             raise ValueError("schema_name is invalid. It must be a non-empty string")
@@ -243,3 +245,36 @@ def verify_schema_parameters(self, schema_name, topics, examples, languages): #t
 
     def schema_name_exists(self, schema_name):
         return file_exists_in_s3(self.output_bucket, f"{schema_name}.json")
+
+
+    @classmethod
+    def validate_input(cls, data: Dict) -> None:
+        """
+        Validate input data. Must be implemented by all child "Model" classes.
+        """
+        schema_specs = data['parameters']
+
+        schema_name = schema_specs["schema_name"]
+        topics = schema_specs["topics"]
+        examples = schema_specs["examples"]
+        languages = schema_specs["languages"]  # ['English', 'Spanish']
+
+        try:
+            cls.verify_schema_parameters(schema_name, topics, examples, languages)
+        except Exception as e:
+            logger.exception(f"Error verifying schema parameters: {e}")
+            raise PrestoBaseException(f"Error verifying schema parameters: {e}", 422) from e
+
+    @classmethod
+    def parse_input_message(cls, data: Dict) -> Any:
+        """
+        Parse input into appropriate response instances.
+        """
+        event_type = data['parameters']['event_type']
+        result_data = data.get('result', {})
+
+        if event_type == 'schema_create':
+            return ClassyCatSchemaResponse(**result_data)
+        else:
+            logger.error(f"Unknown event type {event_type}")
+            raise PrestoBaseException(f"Unknown event type {event_type}", 422)

lib/model/classycat_schema_lookup.py

@@ -1,3 +1,4 @@
+from typing import Dict, Any
 import os
 import json
 from lib.logger import logger
@@ -74,4 +75,27 @@ def process(self, message: Message) -> ClassyCatSchemaResponse:
             return result
         except Exception as e:
             logger.error(f"Error looking up schema name {schema_name}: {e}")
-            raise PrestoBaseException(f"Error looking up schema name {schema_name}", 500) from e
+            raise PrestoBaseException(f"Error looking up schema name {schema_name}", 500) from e
+
+
+    @classmethod
+    def validate_input(cls, data: Dict) -> None:
+        """
+        Validate input data. Must be implemented by all child "Model" classes.
+        """
+        if "schema_name" not in data["parameters"] or data["parameters"]["schema_name"] == "":
+            raise PrestoBaseException("schema_name is required as input to schema look up", 422)
+
+    @classmethod
+    def parse_input_message(cls, data: Dict) -> Any:
+        """
+        Parse input into appropriate response instances.
+        """
+        event_type = data['parameters']['event_type']
+        result_data = data.get('result', {})
+
+        if event_type == 'schema_lookup':
+            return ClassyCatSchemaResponse(**result_data)
+        else:
+            logger.error(f"Unknown event type {event_type}")
+            raise PrestoBaseException(f"Unknown event type {event_type}", 422)

lib/model/fasttext.py

@@ -1,4 +1,4 @@
-from typing import Union, Dict, List
+from typing import Union, Dict, List, Any
 
 import fasttext
 from huggingface_hub import hf_hub_download
@@ -44,3 +44,17 @@ def respond(self, docs: Union[List[schemas.Message], schemas.Message]) -> List[s
         for doc, detected_lang in zip(docs, detected_langs):
             doc.body.result = detected_lang
         return docs
+
+    @classmethod
+    def validate_input(cls, data: Dict) -> None:
+        """
+        Validate input data. Must be implemented by all child "Model" classes.
+        """
+        pass
+
+    @classmethod
+    def parse_input_message(cls, data: Dict) -> Any:
+        """
+        Validate input data. Must be implemented by all child "Model" classes.
+        """
+        return None

lib/model/fptg.py

@@ -1,7 +1,11 @@
 from lib.model.generic_transformer import GenericTransformerModel
+
 MODEL_NAME = 'meedan/paraphrase-filipino-mpnet-base-v2'
+
+
 class Model(GenericTransformerModel):
     BATCH_SIZE = 100
+
     def __init__(self):
         """
         Init FPTG model. Fairly standard for all vectorizers.

lib/model/generic_transformer.py

@@ -1,5 +1,5 @@
 import os
-from typing import Union, Dict, List
+from typing import Union, Dict, List, Any
 from sentence_transformers import SentenceTransformer
 from lib.logger import logger
 from lib.model.model import Model
@@ -34,3 +34,17 @@ def vectorize(self, texts: List[str]) -> List[List[float]]:
         Vectorize the text! Run as batch.
         """
         return {"hash_value": self.model.encode(texts).tolist()}
+
+    @classmethod
+    def validate_input(cls, data: Dict) -> None:
+        """
+        Validate input data. Must be implemented by all child "Model" classes.
+        """
+        pass
+
+    @classmethod
+    def parse_input_message(cls, data: Dict) -> Any:
+        """
+        Validate input data. Must be implemented by all child "Model" classes.
+        """
+        return None

lib/model/image.py

@@ -1,4 +1,4 @@
-from typing import Dict
+from typing import Dict, Any
 import io
 import urllib.request
 
@@ -35,3 +35,18 @@ def process(self, image: schemas.Message) -> schemas.GenericItem:
         Generic function for returning the actual response.
         """
         return {"hash_value": self.compute_pdq(self.get_iobytes_for_image(image))}
+
+
+    @classmethod
+    def validate_input(cls, data: Dict) -> None:
+        """
+        Validate input data. Must be implemented by all child "Model" classes.
+        """
+        pass
+
+    @classmethod
+    def parse_input_message(cls, data: Dict) -> Any:
+        """
+        Validate input data. Must be implemented by all child "Model" classes.
+        """
+        return None

lib/model/indian_sbert.py

@@ -1,7 +1,11 @@
 from lib.model.generic_transformer import GenericTransformerModel
+
 MODEL_NAME = 'meedan/indian-sbert'
+
+
 class Model(GenericTransformerModel):
     BATCH_SIZE = 100
+
     def __init__(self):
         """
         Init IndianSbert model. Fairly standard for all vectorizers.

lib/model/mean_tokens.py

@@ -1,7 +1,11 @@
 from lib.model.generic_transformer import GenericTransformerModel
+
 MODEL_NAME = 'xlm-r-bert-base-nli-stsb-mean-tokens'
+
+
 class Model(GenericTransformerModel):
     BATCH_SIZE = 100
+
     def __init__(self):
         """
         Init MeanTokens model. Fairly standard for all vectorizers.

lib/model/model.py

@@ -83,6 +83,23 @@ def respond(self, messages: Union[List[schemas.Message], schemas.Message]) -> Li
             message.body.result = self.get_response(message)
         return messages
 
+
+    @classmethod
+    def validate_input(cls, data: Dict) -> None:
+        """
+        Validate input data. Must be implemented by child classes.
+        """
+        raise NotImplementedError
+
+
+    @classmethod
+    def parse_input_message(cls, data: Dict) -> Any:
+        """
+        Parse input data. Must be implemented by child classes.
+        """
+        raise NotImplementedError
+
+
     @classmethod
     def create(cls):
         """

lib/model/video.py

@@ -1,4 +1,4 @@
-from typing import Dict
+from typing import Dict, Any
 import os
 import uuid
 import shutil
@@ -64,3 +64,17 @@ def process(self, video: schemas.Message) -> schemas.GenericItem:
                 if os.path.exists(file_path):
                     os.remove(file_path)
         return {"folder": self.tmk_bucket(), "filepath": self.tmk_file_path(video_filename), "hash_value": hash_value}
+
+    @classmethod
+    def validate_input(cls, data: Dict) -> None:
+        """
+        Validate input data. Must be implemented by all child "Model" classes.
+        """
+        pass
+
+    @classmethod
+    def parse_input_message(cls, data: Dict) -> Any:
+        """
+        Validate input data. Must be implemented by all child "Model" classes.
+        """
+        return None

lib/model/yake_keywords.py

@@ -1,4 +1,4 @@
-from typing import Dict
+from typing import Dict, Any
 import io
 import urllib.request
 
@@ -50,3 +50,17 @@ def process(self, message: schemas.Message) -> schemas.YakeKeywordsResponse:
         """
         keywords = self.run_yake(**self.get_params(message))
         return keywords
+
+    @classmethod
+    def validate_input(cls, data: Dict) -> None:
+        """
+        Validate input data. Must be implemented by all child "Model" classes.
+        """
+        pass
+
+    @classmethod
+    def parse_input_message(cls, data: Dict) -> Any:
+        """
+        Validate input data. Must be implemented by all child "Model" classes.
+        """
+        return None