diff --git a/fedn/fedn/network/api/interface.py b/fedn/fedn/network/api/interface.py index 2c2588149..5fa1cc1dd 100644 --- 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 index d2e454a12..980801569 100644 --- 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)