list models, sort on committed_at

scaleoutsystems · Oct 20, 2023 · a8929b1 · a8929b1
1 parent e10594b
commit a8929b1
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 21 deletions.
diff --git a/fedn/fedn/network/api/interface.py b/fedn/fedn/network/api/interface.py
@@ -632,28 +632,20 @@ def get_latest_model(self):
             )
 
     def get_models(self, session_id=None, limit=None, skip=None):
-        models_object = self.statestore.list_models(session_id, limit, skip)
+        result = self.statestore.list_models(session_id, limit, skip)
 
-        if models_object is None:
+        if result is None:
             return (
                 jsonify({"success": False, "message": "No models found."}),
                 404,
             )
 
-        payload = {}
-        for object in models_object["result"]:
-            model = object["model"]
-            session_id = object["session_id"]
-            commited_at = object["commited_at"]
-
-            obj = {
-                "session_id": session_id,
-                "commited_at": commited_at,
-                "model": model,
-            }
-            payload[model] = obj
+        arr = []
+
+        for model in result["result"]:
+            arr.append(model)
 
-        result = {"result": payload, "count": models_object["count"]}
+        result = {"result": arr, "count": result["count"]}
 
         return jsonify(result)
 

diff --git a/fedn/fedn/network/statestore/mongostatestore.py b/fedn/fedn/network/statestore/mongostatestore.py
@@ -153,14 +153,14 @@ def set_latest_model(self, model_id, session_id=None):
         :return:
         """
 
-        commited_at = str(datetime.now())
+        committed_at = datetime.now()
 
         self.model.insert_one(
             {
                 "key": "models",
                 "model": model_id,
                 "session_id": session_id,
-                "committed_at": commited_at,
+                "committed_at": committed_at,
             }
         )
 
@@ -172,7 +172,7 @@ def set_latest_model(self, model_id, session_id=None):
             {
                 "$push": {
                     "model": model_id,
-                    "committed_at": commited_at,
+                    "committed_at": str(committed_at),
                 }
             },
             True,
@@ -337,7 +337,14 @@ def get_helper(self):
         except (KeyError, IndexError):
             return None
 
-    def list_models(self, session_id=None, limit=None, skip=None):
+    def list_models(
+        self,
+        session_id=None,
+        limit=None,
+        skip=None,
+        sort_key="committed_at",
+        sort_order=pymongo.DESCENDING,
+    ):
         """List all models in the statestore.
 
         :param session_id: The session id.
@@ -357,14 +364,23 @@ def list_models(self, session_id=None, limit=None, skip=None):
             else {"key": "models", "session_id": session_id}
         )
 
+        projection = {"_id": False, "key": False}
+
         if limit is not None and skip is not None:
             limit = int(limit)
             skip = int(skip)
 
-            result = self.model.find(find_option).limit(limit).skip(skip)
+            result = (
+                self.model.find(find_option, projection)
+                .limit(limit)
+                .skip(skip)
+                .sort(sort_key, sort_order)
+            )
 
         else:
-            result = self.model.find(find_option)
+            result = self.model.find(find_option, projection).sort(
+                sort_key, sort_order
+            )
 
         count = self.model.count_documents(find_option)