From c6d779f591887c2be28a6b6ca4668e7a187363e5 Mon Sep 17 00:00:00 2001
From: faraphel <rc60650@hotmail.com>
Date: Fri, 10 Jan 2025 11:08:10 +0100
Subject: [PATCH] added support for tags to categorise the models

---
 samples/models/dummy/config.json | 1 +
 source/model/PythonModel.py      | 7 ++++++-
 source/model/base/BaseModel.py   | 2 ++
 3 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/samples/models/dummy/config.json b/samples/models/dummy/config.json
index d5eabc8..09be48c 100644
--- a/samples/models/dummy/config.json
+++ b/samples/models/dummy/config.json
@@ -1,5 +1,6 @@
 {
   "type": "python",
+  "tags": ["dummy"],
   "file": "model.py",
 
   "inputs": {
diff --git a/source/model/PythonModel.py b/source/model/PythonModel.py
index bb2531f..0d47fc6 100644
--- a/source/model/PythonModel.py
+++ b/source/model/PythonModel.py
@@ -68,7 +68,12 @@ class PythonModel(base.BaseModel):
         infer_api.__signature__ = inspect.Signature(parameters=parameters)
 
         # add the inference endpoint on the API
-        self.manager.application.add_api_route(f"/models/{self.name}/infer", infer_api, methods=["POST"])
+        self.manager.application.add_api_route(
+            f"/models/{self.name}/infer",
+            infer_api,
+            methods=["POST"],
+            tags=self.tags,
+        )
 
     def _load(self) -> None:
         return self.module.load(self)
diff --git a/source/model/base/BaseModel.py b/source/model/base/BaseModel.py
index 4978843..d39d7d2 100644
--- a/source/model/base/BaseModel.py
+++ b/source/model/base/BaseModel.py
@@ -18,6 +18,8 @@ class BaseModel(abc.ABC):
         self.manager = manager
         # the mimetype of the model responses
         self.output_type: str = configuration.get("output_type", "application/json")
+        # get the tags of the model
+        self.tags = configuration.get("tags", [])
 
         self._loaded = False