MemGPT/memgpt/server/rest_api/admin/users.py

import uuid
from typing import List, Optional

from fastapi import APIRouter, Body, HTTPException, Query
from pydantic import BaseModel, Field

from memgpt.data_types import User
from memgpt.server.rest_api.interface import QueuingInterface
from memgpt.server.server import SyncServer

router = APIRouter()


class GetAllUsersResponse(BaseModel):
    user_list: List[dict] = Field(..., description="A list of users.")


class CreateUserRequest(BaseModel):
    user_id: Optional[uuid.UUID] = Field(None, description="Identifier of the user (optional, generated automatically if null).")
    api_key_name: Optional[str] = Field(None, description="Name for API key autogenerated on user creation (optional).")


class CreateUserResponse(BaseModel):
    user_id: uuid.UUID = Field(..., description="Identifier of the user (UUID).")
    api_key: str = Field(..., description="New API key generated for user.")


class CreateAPIKeyRequest(BaseModel):
    user_id: uuid.UUID = Field(..., description="Identifier of the user (UUID).")
    name: Optional[str] = Field(None, description="Name for the API key (optional).")


class CreateAPIKeyResponse(BaseModel):
    api_key: str = Field(..., description="New API key generated.")


class GetAPIKeysRequest(BaseModel):
    user_id: uuid.UUID = Field(..., description="Identifier of the user (UUID).")


class GetAPIKeysResponse(BaseModel):
    api_key_list: List[str] = Field(..., description="Identifier of the user (UUID).")


class DeleteAPIKeyResponse(BaseModel):
    message: str
    api_key_deleted: str


class DeleteUserResponse(BaseModel):
    message: str
    user_id_deleted: uuid.UUID


def setup_admin_router(server: SyncServer, interface: QueuingInterface):
    @router.get("/users", tags=["admin"], response_model=GetAllUsersResponse)
    def get_all_users():
        """
        Get a list of all users in the database
        """
        try:
            users = server.ms.get_all_users()
            processed_users = [{"user_id": user.id} for user in users]
        except HTTPException:
            raise
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"{e}")
        return GetAllUsersResponse(user_list=processed_users)

    @router.post("/users", tags=["admin"], response_model=CreateUserResponse)
    def create_user(request: Optional[CreateUserRequest] = Body(None)):
        """
        Create a new user in the database
        """
        if request is None:
            request = CreateUserRequest()

        new_user = User(
            id=None if not request.user_id else request.user_id,
            # TODO can add more fields (name? metadata?)
        )

        try:
            server.ms.create_user(new_user)

            # initialize default presets automatically for user
            server.initialize_default_presets(new_user.id)

            # make sure we can retrieve the user from the DB too
            new_user_ret = server.ms.get_user(new_user.id)
            if new_user_ret is None:
                raise HTTPException(status_code=500, detail=f"Failed to verify user creation")

            # create an API key for the user
            token = server.ms.create_api_key(user_id=new_user.id, name=request.api_key_name)

        except HTTPException:
            raise
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"{e}")
        return CreateUserResponse(user_id=new_user_ret.id, api_key=token.token)

    @router.delete("/users", tags=["admin"], response_model=DeleteUserResponse)
    def delete_user(
        user_id: uuid.UUID = Query(..., description="The user_id key to be deleted."),
    ):
        # TODO make a soft deletion, instead of a hard deletion
        try:
            user = server.ms.get_user(user_id=user_id)
            if user is None:
                raise HTTPException(status_code=404, detail=f"User does not exist")
            server.ms.delete_user(user_id=user_id)
        except HTTPException:
            raise
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"{e}")
        return DeleteUserResponse(message="User successfully deleted.", user_id_deleted=user_id)

    @router.post("/users/keys", tags=["admin"], response_model=CreateAPIKeyResponse)
    def create_new_api_key(request: CreateAPIKeyRequest = Body(...)):
        """
        Create a new API key for a user
        """
        try:
            token = server.ms.create_api_key(user_id=request.user_id, name=request.name)
        except HTTPException:
            raise
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"{e}")
        return CreateAPIKeyResponse(api_key=token.token)

    @router.get("/users/keys", tags=["admin"], response_model=GetAPIKeysResponse)
    def get_api_keys(
        user_id: uuid.UUID = Query(..., description="The unique identifier of the user."),
    ):
        """
        Get a list of all API keys for a user
        """
        print("GET USERS", user_id)
        try:
            tokens = server.ms.get_all_api_keys_for_user(user_id=user_id)
            processed_tokens = [t.token for t in tokens]
        except HTTPException:
            raise
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"{e}")
        print("TOKENS", processed_tokens)
        return GetAPIKeysResponse(api_key_list=processed_tokens)

    @router.delete("/users/keys", tags=["admin"], response_model=DeleteAPIKeyResponse)
    def delete_api_key(
        api_key: str = Query(..., description="The API key to be deleted."),
    ):
        try:
            token = server.ms.get_api_key(api_key=api_key)
            if token is None:
                raise HTTPException(status_code=404, detail=f"API key does not exist")
            server.ms.delete_api_key(api_key=api_key)
        except HTTPException:
            raise
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"{e}")
        return DeleteAPIKeyResponse(message="API key successfully deleted.", api_key_deleted=api_key)

    return router