feat: migrate /api/_status endpoint to /api/v1/status

CHANGELOG.md

@@ -25,7 +25,8 @@ These are the section headers that we use:
 - Added `GET /api/v1/workspaces/:workspace_id/users` endpoint to get the users of a workspace. ([#153](https://github.com/argilla-io/argilla-server/pull/153))
 - Added `POST /api/v1/workspaces/:workspace_id/users` endpoind to add a user to a workspace. ([#156](https://github.com/argilla-io/argilla-server/pull/156))
 - Added `DELETE /api/v1/workspaces/:workspace_id/users/:user_id` endpoint to remove a user from a workspace. ([#158](https://github.com/argilla-io/argilla-server/pull/158))
-- Added `GET /api/v1/version` endpoin to get the current Argilla version. ([#162](https://github.com/argilla-io/argilla-server/pull/162))
+- Added `GET /api/v1/version` endpoint to get the current Argilla version. ([#162](https://github.com/argilla-io/argilla-server/pull/162))
+- Added `GET /api/v1/status` endpoint to get Argilla service status. ([#165](https://github.com/argilla-io/argilla-server/pull/165))
 
 ## [Unreleased]()
 

src/argilla_server/apis/v1/handlers/info.py

@@ -12,14 +12,24 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-from fastapi import APIRouter
+from fastapi import APIRouter, Depends
 
-from argilla_server._version import __version__
-from argilla_server.schemas.v1.info import Version
+from argilla_server.contexts import info
+from argilla_server.schemas.v1.info import Status, Version
+from argilla_server.search_engine import SearchEngine, get_search_engine
 
 router = APIRouter(tags=["info"])
 
 
 @router.get("/version", response_model=Version)
 async def get_version():
-    return Version(version=__version__)
+    return Version(version=info.argilla_version())
+
+
+@router.get("/status", response_model=Status)
+async def get_status(search_engine: SearchEngine = Depends(get_search_engine)):
+    return Status(
+        version=info.argilla_version(),
+        search_engine=await search_engine.info(),
+        memory=info.memory_status(),
+    )

src/argilla_server/contexts/info.py

@@ -0,0 +1,55 @@
+#  Copyright 2021-present, the Recognai S.L. team.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+import os
+
+import psutil
+
+from argilla_server._version import __version__
+
+
+def argilla_version() -> str:
+    return __version__
+
+
+def memory_status() -> dict:
+    process = psutil.Process(os.getpid())
+
+    return {k: _memory_size(v) for k, v in process.memory_info()._asdict().items()}
+
+
+def _memory_size(bytes) -> str:
+    system = [
+        (1024**5, "P"),
+        (1024**4, "T"),
+        (1024**3, "G"),
+        (1024**2, "M"),
+        (1024**1, "K"),
+        (1024**0, "B"),
+    ]
+
+    factor, suffix = None, None
+    for factor, suffix in system:
+        if bytes >= factor:
+            break
+
+    amount = int(bytes / factor)
+    if isinstance(suffix, tuple):
+        singular, multiple = suffix
+        if amount == 1:
+            suffix = singular
+        else:
+            suffix = multiple
+
+    return str(amount) + suffix

src/argilla_server/schemas/v1/info.py

@@ -17,3 +17,9 @@
 
 class Version(BaseModel):
     version: str
+
+
+class Status(BaseModel):
+    version: str
+    search_engine: dict
+    memory: dict

src/argilla_server/search_engine/base.py

@@ -264,6 +264,10 @@ async def new_instance(cls) -> "SearchEngine":
     async def close(self):
         pass
 
+    @abstractmethod
+    async def info(self) -> dict:
+        pass
+
     @classmethod
     def register(cls, engine_name: str):
         def decorator(engine_class):

src/argilla_server/search_engine/elasticsearch.py

@@ -64,6 +64,9 @@ async def new_instance(cls) -> "ElasticSearchEngine":
     async def close(self):
         await self.client.close()
 
+    async def info(self) -> dict:
+        return await self.client.info()
+
     def _mapping_for_vector_settings(self, vector_settings: VectorSettings) -> dict:
         return {
             es_field_for_vector_settings(vector_settings): {

src/argilla_server/search_engine/opensearch.py

@@ -56,6 +56,9 @@ async def new_instance(cls) -> "OpenSearchEngine":
     async def close(self):
         await self.client.close()
 
+    async def info(self) -> dict:
+        return await self.client.info()
+
     def _configure_index_settings(self):
         base_settings = super()._configure_index_settings()
         return {**base_settings, "index.knn": False}

src/argilla_server/services/info.py

@@ -13,6 +13,8 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
+# TODO: Once we delete all API v0 endpoints, we can remove this file.
+
 import os
 from typing import Any, Dict
 

tests/unit/api/v1/info/test_get_status.py

@@ -0,0 +1,36 @@
+#  Copyright 2021-present, the Recognai S.L. team.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+import pytest
+from argilla_server._version import __version__
+from argilla_server.search_engine import SearchEngine
+from httpx import AsyncClient
+
+
+@pytest.mark.asyncio
+class TestGetStatus:
+    def url(self) -> str:
+        return "/api/v1/status"
+
+    async def test_get_status(self, async_client: AsyncClient, mock_search_engine: SearchEngine):
+        mock_search_engine.info.return_value = {}
+
+        response = await async_client.get(self.url())
+
+        assert response.status_code == 200
+
+        response_json = response.json()
+        assert response_json["version"] == __version__
+        assert "search_engine" in response_json
+        assert "memory" in response_json