quivr/backend/routes/brain_routes.py

import os
import time
from typing import Optional
from uuid import UUID

from auth.auth_bearer import AuthBearer, get_current_user
from fastapi import APIRouter, Depends, Request
from logger import get_logger
from models.brains import Brain
from models.settings import CommonsDep, common_dependencies
from models.users import User
from pydantic import BaseModel
from utils.users import fetch_user_id_from_credentials

logger = get_logger(__name__)

brain_router = APIRouter()


class BrainToUpdate(BaseModel):
    brain_id: UUID
    name: Optional[str] = "New Brain"
    status: Optional[str] = "public"
    model: Optional[str] = "gpt-3.5-turbo-0613"
    temperature: Optional[float] = 0.0
    max_tokens: Optional[int] = 256
    file_sha1: Optional[str] = ""


# get all brains
@brain_router.get("/brains", dependencies=[Depends(AuthBearer())], tags=["Brain"])
async def brain_endpoint(current_user: User = 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.
    """
    commons = common_dependencies()
    brain = Brain()
    user_id = fetch_user_id_from_credentials(commons, {"email": current_user.email})
    brains = brain.get_user_brains(user_id)
    return {"brains": brains}


# get one brain
@brain_router.get(
    "/brains/{brain_id}", dependencies=[Depends(AuthBearer())], tags=["Brain"]
)
async def 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 = Brain(brain_id=brain_id)
    brains = brain.get_brain_details()
    if len(brains) > 0:
        return {
            "brainId": brain_id,
            "brainName": brains[0]["name"],
            "status": brains[0]["status"],
        }
    else:
        return {"error": f"No brain found with brain_id {brain_id}"}


# delete one brain
@brain_router.delete(
    "/brains/{brain_id}", dependencies=[Depends(AuthBearer())], tags=["Brain"]
)
async def brain_endpoint(brain_id: UUID):
    """
    Delete a specific brain by brain ID.
    """
    brain = Brain(brain_id=brain_id)
    brain.delete_brain()
    return {"message": f"{brain_id}  has been deleted."}


class BrainObject(BaseModel):
    brain_id: Optional[UUID]
    name: Optional[str] = "New Brain"
    status: Optional[str] = "public"
    model: Optional[str] = "gpt-3.5-turbo-0613"
    temperature: Optional[float] = 0.0
    max_tokens: Optional[int] = 256
    file_sha1: Optional[str] = ""


# create new brain
@brain_router.post("/brains", dependencies=[Depends(AuthBearer())], tags=["Brain"])
async def brain_endpoint(
    request: Request,
    brain: BrainObject,
    current_user: User = 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'
    """
    commons = common_dependencies()
    brain = Brain(name=brain.name)
    user_id = fetch_user_id_from_credentials(commons, {"email": current_user.email})
    created_brain = brain.create_brain(brain.name)[0]
    # create a brain X user entry
    brain.create_brain_user(created_brain["brain_id"], user_id, rights="Owner")

    return {"id": created_brain["brain_id"], "name": created_brain["name"]}


# update existing brain
@brain_router.put(
    "/brains/{brain_id}", dependencies=[Depends(AuthBearer())], tags=["Brain"]
)
async def brain_endpoint(
    request: Request,
    brain_id: UUID,
    input_brain: Brain,
    fileName: Optional[str],
    current_user: User = Depends(get_current_user),
):
    """
    Update an existing brain with new brain parameters/files.
    If the file is contained in Add file to brain :
        if given a fileName/ file sha1 / -> add all the vector Ids to the brains_vectors
    Modify other brain fields:
        name, status, model, max_tokens, temperature
    Return modified brain ? No need -> do an optimistic update
    """
    commons = common_dependencies()
    brain = Brain(brain_id=brain_id)

    # Add new file to brain , il file_sha1 already exists in brains_vectors -> out (not now)
    if brain.file_sha1:
        # add all the vector Ids to the brains_vectors  with the given brain.brain_id
        brain.update_brain_with_file(file_sha1=input_brain.file_sha1)
        print("brain:", brain)

    brain.update_brain_fields(commons, brain)
    return {"message": f"Brain {brain_id} has been updated."}