quivr/backend/routes/brain_routes.py

from uuid import UUID

from auth import AuthBearer, get_current_user
from fastapi import APIRouter, Depends, HTTPException
from logger import get_logger
from models import UserIdentity, UserUsage
from models.databases.supabase.brains import (
    BrainQuestionRequest,
    BrainUpdatableProperties,
    CreateBrainProperties,
)
from repository.brain import (
    create_brain,
    create_brain_user,
    get_brain_details,
    get_default_user_brain_or_create_new,
    get_question_context_from_brain,
    get_user_brains,
    get_user_default_brain,
    set_as_default_brain_for_user,
    update_brain_by_id,
)
from repository.prompt import delete_prompt_by_id, get_prompt_by_id
from routes.authorizations.brain_authorization import has_brain_authorization
from routes.authorizations.types import RoleEnum

logger = get_logger(__name__)

brain_router = APIRouter()


# get all brains
@brain_router.get("/brains/", dependencies=[Depends(AuthBearer())], tags=["Brain"])
async def brain_endpoint(current_user: UserIdentity = Depends(get_current_user)):
    """
    Retrieve all brains for the current user.

    - `current_user`: The current authenticated user.
    - Returns a list of all brains registered for the user.

    This endpoint retrieves all the brains associated with the current authenticated user. It returns a list of brains objects
    containing the brain ID and brain name for each brain.
    """
    brains = get_user_brains(current_user.id)
    return {"brains": brains}


# get default brain
@brain_router.get(
    "/brains/default/", dependencies=[Depends(AuthBearer())], tags=["Brain"]
)
async def get_default_brain_endpoint(
    current_user: UserIdentity = Depends(get_current_user),
):
    """
    Retrieve the default brain for the current user. If the user doesnt have one, it creates one.

    - `current_user`: The current authenticated user.
    - Returns the default brain for the user.

    This endpoint retrieves the default brain associated with the current authenticated user.
    The default brain is defined as the brain marked as default in the brains_users table.
    """

    brain = get_default_user_brain_or_create_new(current_user)
    return {"id": brain.brain_id, "name": brain.name, "rights": "Owner"}


@brain_router.get(
    "/brains/{brain_id}/",
    dependencies=[Depends(AuthBearer()), Depends(has_brain_authorization())],
    tags=["Brain"],
)
async def get_brain_endpoint(
    brain_id: UUID,
):
    """
    Retrieve details of a specific brain by brain ID.

    - `brain_id`: The ID of the brain to retrieve details for.
    - Returns the brain ID and its history.

    This endpoint retrieves the details of a specific brain identified by the provided brain ID. It returns the brain ID and its
    history, which includes the brain messages exchanged in the brain.
    """

    brain_details = get_brain_details(brain_id)
    if brain_details is None:
        raise HTTPException(
            status_code=404,
            detail="Brain details not found",
        )

    return brain_details


# create new brain
@brain_router.post("/brains/", dependencies=[Depends(AuthBearer())], tags=["Brain"])
async def create_brain_endpoint(
    brain: CreateBrainProperties,
    current_user: UserIdentity = Depends(get_current_user),
):
    """
    Create a new brain with given
        name
        status
        model
        max_tokens
        temperature
    In the brains table & in the brains_users table and put the creator user as 'Owner'
    """

    user_brains = get_user_brains(current_user.id)
    userDailyUsage = UserUsage(
        id=current_user.id,
        email=current_user.email,
        openai_api_key=current_user.openai_api_key,
    )
    userSettings = userDailyUsage.get_user_settings()

    if len(user_brains) >= userSettings.get("max_brains", 5):
        raise HTTPException(
            status_code=429,
            detail=f"Maximum number of brains reached ({userSettings.get('max_brains', 5)}).",
        )

    new_brain = create_brain(
        brain,
    )
    default_brain = get_user_default_brain(current_user.id)
    if default_brain:
        logger.info(f"Default brain already exists for user {current_user.id}")
        create_brain_user(
            user_id=current_user.id,
            brain_id=new_brain.brain_id,
            rights=RoleEnum.Owner,
            is_default_brain=False,
        )
    else:
        logger.info(
            f"Default brain does not exist for user {current_user.id}. It will be created."
        )
        create_brain_user(
            user_id=current_user.id,
            brain_id=new_brain.brain_id,
            rights=RoleEnum.Owner,
            is_default_brain=True,
        )

    return {
        "id": new_brain.brain_id,
        "name": brain.name,
        "rights": "Owner",
    }


# update existing brain
@brain_router.put(
    "/brains/{brain_id}/",
    dependencies=[
        Depends(
            AuthBearer(),
        ),
        Depends(has_brain_authorization([RoleEnum.Editor, RoleEnum.Owner])),
    ],
    tags=["Brain"],
)
async def update_brain_endpoint(
    brain_id: UUID,
    input_brain: BrainUpdatableProperties,
):
    """
    Update an existing brain with new brain configuration
    """

    # Remove prompt if it is private and no longer used by brain
    if input_brain.prompt_id is None:
        existing_brain = get_brain_details(brain_id)
        if existing_brain is None:
            raise HTTPException(
                status_code=404,
                detail="Brain not found",
            )
        prompt_id = existing_brain.prompt_id
        if prompt_id is not None:
            prompt = get_prompt_by_id(prompt_id)
            if prompt is not None and prompt.status == "private":
                delete_prompt_by_id(prompt_id)

    update_brain_by_id(brain_id, input_brain)

    return {"message": f"Brain {brain_id} has been updated."}


# set as default brain
@brain_router.post(
    "/brains/{brain_id}/default",
    dependencies=[
        Depends(
            AuthBearer(),
        ),
        Depends(has_brain_authorization()),
    ],
    tags=["Brain"],
)
async def set_as_default_brain_endpoint(
    brain_id: UUID,
    user: UserIdentity = Depends(get_current_user),
):
    """
    Set a brain as default for the current user.
    """

    set_as_default_brain_for_user(user.id, brain_id)

    return {"message": f"Brain {brain_id} has been set as default brain."}


@brain_router.post(
    "/brains/{brain_id}/question_context",
    dependencies=[
        Depends(
            AuthBearer(),
        ),
        Depends(has_brain_authorization()),
    ],
    tags=["Brain"],
)
async def get_question_context_from_brain_endpoint(
    brain_id: UUID,
    request: BrainQuestionRequest,
):
    """
    Get question context from brain
    """

    context = get_question_context_from_brain(brain_id, request.question)

    return {"context": context}