Skip to content
This repository has been archived by the owner on Jun 14, 2024. It is now read-only.

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

Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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]()

Expand Down
18 changes: 14 additions & 4 deletions src/argilla_server/apis/v1/handlers/info.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
)
55 changes: 55 additions & 0 deletions src/argilla_server/contexts/info.py
Original file line number Diff line number Diff line change
@@ -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
6 changes: 6 additions & 0 deletions src/argilla_server/schemas/v1/info.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,3 +17,9 @@

class Version(BaseModel):
version: str


class Status(BaseModel):
version: str
search_engine: dict
memory: dict
4 changes: 4 additions & 0 deletions src/argilla_server/search_engine/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
3 changes: 3 additions & 0 deletions src/argilla_server/search_engine/elasticsearch.py
Original file line number Diff line number Diff line change
Expand Up @@ -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): {
Expand Down
3 changes: 3 additions & 0 deletions src/argilla_server/search_engine/opensearch.py
Original file line number Diff line number Diff line change
Expand Up @@ -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}
Expand Down
2 changes: 2 additions & 0 deletions src/argilla_server/services/info.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
36 changes: 36 additions & 0 deletions tests/unit/api/v1/info/test_get_status.py
Original file line number Diff line number Diff line change
@@ -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
Loading