"""FLorist server FastAPI endpoints and routes."""
from contextlib import asynccontextmanager
from typing import Any, AsyncGenerator
from fastapi import Depends, FastAPI
from fastapi.responses import JSONResponse
from motor.motor_asyncio import AsyncIOMotorClient
from florist.api.auth.token import DEFAULT_USERNAME, Token, make_default_server_user
from florist.api.clients.clients import Client
from florist.api.clients.optimizers import Optimizer
from florist.api.db.config import DatabaseConfig
from florist.api.db.server_entities import User
from florist.api.models.models import Model
from florist.api.routes.server.auth import check_default_user_token
from florist.api.routes.server.auth import router as auth_router
from florist.api.routes.server.job import router as job_router
from florist.api.routes.server.status import router as status_router
from florist.api.routes.server.training import router as training_router
from florist.api.servers.strategies import Strategy
[docs]
@asynccontextmanager
async def lifespan(app: FastAPI) -> AsyncGenerator[Any, Any]:
"""Set up function for app startup and shutdown."""
# Set up mongodb
app.db_client = AsyncIOMotorClient(DatabaseConfig.get_mongodb_uri()) # type: ignore[attr-defined]
app.database = app.db_client[DatabaseConfig.get_mongodb_db_name()] # type: ignore[attr-defined]
# Create default user if it does not exist
user = await User.find_by_username(DEFAULT_USERNAME, app.database) # type: ignore[attr-defined]
if user is None:
await make_default_server_user(app.database) # type: ignore[attr-defined]
# Making a dictionary to hold the authentication tokens of the clients
# this server is connected to
app.clients_auth_tokens: dict[str, Token] = {} # type: ignore[attr-defined, misc]
yield
# Shut down mongodb
app.db_client.close() # type: ignore[attr-defined]
app = FastAPI(lifespan=lifespan)
app.include_router(training_router, tags=["training"], prefix="/api/server/training")
app.include_router(job_router, tags=["job"], prefix="/api/server/job")
app.include_router(status_router, tags=["status"], prefix="/api/server/check_status")
app.include_router(auth_router, tags=["auth"], prefix="/api/server/auth")
[docs]
@app.get(
path="/api/server/models",
response_description="Returns a list of all available models",
dependencies=[Depends(check_default_user_token)],
)
def list_models() -> JSONResponse:
"""
Return a list of all available models.
:return: (JSONResponse) A JSON response with a list of all elements in the `api.servers.common.Model` enum.
"""
return JSONResponse(Model.list())
[docs]
@app.get(
path="/api/server/clients/{strategy}",
response_description="Returns a list of all available clients by strategy",
dependencies=[Depends(check_default_user_token)],
)
def list_clients(strategy: Strategy) -> JSONResponse:
"""
Return a list of all available clients by strategy.
:param strategy: (Strategy) The strategy to find the compatible clients.
:return: (JSONResponse) A JSON response with a list of all elements in the `api.clients.common.Client` enum
that are compatible with the given strategy.
"""
return JSONResponse(Client.list_by_strategy(strategy))
[docs]
@app.get(
path="/api/server/strategies",
response_description="Returns a list of all available strategies",
dependencies=[Depends(check_default_user_token)],
)
def list_strategies() -> JSONResponse:
"""
Return a list of all available strategies.
:return: (JSONResponse) A JSON response with a list of all elements in the `api.servers.strategy.Strategies` enum.
"""
return JSONResponse(Strategy.list())
[docs]
@app.get(
path="/api/server/optimizers",
response_description="Returns a list of all available optimizers",
dependencies=[Depends(check_default_user_token)],
)
def list_optimizers() -> JSONResponse:
"""
Return a list of all available optimizers.
:return: (JSONResponse) A JSON response with a list of all elements in the `api.clients.optimizers.Optimizer` enum.
"""
return JSONResponse(Optimizer.list())