feat(backend): implement brain-prompt link (#831)

* feat: add prompt_id field to brain

* feat(Prompt controller): update prompt routes

* feat: remove unused private prompts

* refactor: add BrainEntity and repo and service

* tests: partially type main Repository

* feat: add PromptStatusEnum enum

* feat: change delete prompt repository return type
This commit is contained in:
Mamadou DICKO 2023-08-03 10:37:13 +02:00 committed by GitHub
parent e3b6114248
commit 4ca6c667da
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
30 changed files with 473 additions and 192 deletions

View File

@ -0,0 +1,23 @@
from typing import Optional
from uuid import UUID
from pydantic import BaseModel
from routes.authorizations.types import RoleEnum
class BrainEntity(BaseModel):
brain_id: UUID
name: str
description: Optional[str]
temperature: Optional[float]
model: Optional[str]
max_tokens: Optional[int]
openai_api_key: Optional[str]
status: Optional[str]
prompt_id: Optional[UUID]
class MinimalBrainEntity(BaseModel):
id: UUID
name: str
rights: RoleEnum

View File

@ -6,7 +6,6 @@ from pydantic import BaseModel
from utils.vectors import get_unique_files_from_vector_ids from utils.vectors import get_unique_files_from_vector_ids
from models.settings import BrainRateLimiting, CommonsDep, common_dependencies from models.settings import BrainRateLimiting, CommonsDep, common_dependencies
from models.users import User
logger = get_logger(__name__) logger = get_logger(__name__)
@ -22,6 +21,7 @@ class Brain(BaseModel):
openai_api_key: Optional[str] = None openai_api_key: Optional[str] = None
files: List[Any] = [] files: List[Any] = []
max_brain_size = BrainRateLimiting().max_brain_size max_brain_size = BrainRateLimiting().max_brain_size
prompt_id: Optional[UUID] = None
class Config: class Config:
arbitrary_types_allowed = True arbitrary_types_allowed = True
@ -77,24 +77,6 @@ class Brain(BaseModel):
{"brain_id": self.id, "user_id": user_id} {"brain_id": self.id, "user_id": user_id}
).execute() ).execute()
def get_user_brains(self, user_id):
return self.commons["db"].get_user_brains(user_id)
def get_brain_for_user(self, user_id):
return self.commons["db"].get_brain_for_user(user_id, self.id)
def get_brain_details(self):
response = (
self.commons["supabase"]
.from_("brains")
.select("id:brain_id, name, *")
.filter("brain_id", "eq", self.id)
.execute()
)
if response.data == []:
return None
return response.data[0]
def delete_brain(self, user_id): def delete_brain(self, user_id):
results = self.commons["db"].delete_brain_user_by_id(user_id, self.id) results = self.commons["db"].delete_brain_user_by_id(user_id, self.id)
@ -105,37 +87,12 @@ class Brain(BaseModel):
self.commons["db"].delete_brain_user(self.id) self.commons["db"].delete_brain_user(self.id)
self.commons["db"].delete_brain(self.id) self.commons["db"].delete_brain(self.id)
def create_brain(self):
response = self.commons["db"].create_brain(self.name)
self.id = response.data[0]["brain_id"]
return response.data
def create_brain_user(self, user_id: UUID, rights, default_brain):
response = self.commons["db"].create_brain_user(user_id=user_id, brain_id=self.id, rights=rights, default_brain=default_brain)
self.id = response.data[0]["brain_id"]
return response.data
def set_as_default_brain_for_user(self, user: User):
old_default_brain = get_default_user_brain(user)
if old_default_brain is not None:
self.commons["supabase"].table("brains_users").update(
{"default_brain": False}
).match({"brain_id": old_default_brain["id"], "user_id": user.id}).execute()
self.commons["supabase"].table("brains_users").update(
{"default_brain": True}
).match({"brain_id": self.id, "user_id": user.id}).execute()
def create_brain_vector(self, vector_id, file_sha1): def create_brain_vector(self, vector_id, file_sha1):
return self.commons["db"].create_brain_vector(self.id, vector_id, file_sha1) return self.commons["db"].create_brain_vector(self.id, vector_id, file_sha1)
def get_vector_ids_from_file_sha1(self, file_sha1: str): def get_vector_ids_from_file_sha1(self, file_sha1: str):
return self.commons["db"].get_vector_ids_from_file_sha1(file_sha1) return self.commons["db"].get_vector_ids_from_file_sha1(file_sha1)
def update_brain_fields(self):
return self.commons["db"].update_brain_fields(brain_id=self.id, brain_name=self.name)
def update_brain_with_file(self, file_sha1: str): def update_brain_with_file(self, file_sha1: str):
# not used # not used
vector_ids = self.get_vector_ids_from_file_sha1(file_sha1) vector_ids = self.get_vector_ids_from_file_sha1(file_sha1)
@ -154,29 +111,3 @@ class Brain(BaseModel):
def delete_file_from_brain(self, file_name: str): def delete_file_from_brain(self, file_name: str):
return self.commons["db"].delete_file_from_brain(self.id, file_name) return self.commons["db"].delete_file_from_brain(self.id, file_name)
def get_default_user_brain(user: User):
commons = common_dependencies()
response = commons["db"].get_default_user_brain_id(user.id)
logger.info("Default brain response:", response.data)
default_brain_id = response.data[0]["brain_id"] if response.data else None
logger.info(f"Default brain id: {default_brain_id}")
if default_brain_id:
brain_response = commons["db"].get_brain_by_id(default_brain_id)
return brain_response.data[0] if brain_response.data else None
def get_default_user_brain_or_create_new(user: User) -> Brain:
default_brain = get_default_user_brain(user)
if default_brain:
return Brain.create(**default_brain)
else:
brain = Brain.create()
brain.create_brain()
brain.create_brain_user(user.id, "Owner", True)
return brain

View File

@ -2,10 +2,12 @@ from abc import ABC, abstractmethod
from datetime import datetime from datetime import datetime
from uuid import UUID from uuid import UUID
from models.brain_entity import BrainEntity
class Repository(ABC): class Repository(ABC):
@abstractmethod @abstractmethod
def get_user_brains(self, user_id: str): def get_user_brains(self, user_id: str) -> list[BrainEntity]:
pass pass
@abstractmethod @abstractmethod
@ -29,7 +31,7 @@ class Repository(ABC):
pass pass
@abstractmethod @abstractmethod
def create_brain(self, name: str): def create_brain(self, brain: str):
pass pass
@abstractmethod @abstractmethod
@ -46,10 +48,6 @@ class Repository(ABC):
def get_vector_ids_from_file_sha1(self, file_sha1: str): def get_vector_ids_from_file_sha1(self, file_sha1: str):
pass pass
@abstractmethod
def update_brain_fields(self, brain_id: UUID, brain_name: str):
pass
@abstractmethod @abstractmethod
def get_brain_vector_ids(self, brain_id: UUID): def get_brain_vector_ids(self, brain_id: UUID):
pass pass
@ -59,7 +57,7 @@ class Repository(ABC):
pass pass
@abstractmethod @abstractmethod
def get_default_user_brain_id(self, user_id: UUID): def get_default_user_brain_id(self, user_id: UUID) -> UUID:
pass pass
@abstractmethod @abstractmethod

View File

@ -1,30 +1,77 @@
from typing import Optional
from uuid import UUID from uuid import UUID
from models.databases.repository import Repository
from logger import get_logger from logger import get_logger
from models.brain_entity import BrainEntity, MinimalBrainEntity
from models.databases.repository import Repository
from pydantic import BaseModel
logger = get_logger(__name__) logger = get_logger(__name__)
class CreateBrainProperties(BaseModel):
name: Optional[str] = "Default brain"
description: Optional[str] = "This is a description"
status: Optional[str] = "private"
model: Optional[str] = "gpt-3.5-turbo-0613"
temperature: Optional[float] = 0.0
max_tokens: Optional[int] = 256
openai_api_key: Optional[str] = None
prompt_id: Optional[UUID] = None
def dict(self, *args, **kwargs):
brain_dict = super().dict(*args, **kwargs)
if brain_dict.get("prompt_id"):
brain_dict["prompt_id"] = str(brain_dict.get("prompt_id"))
return brain_dict
class BrainUpdatableProperties(BaseModel):
name: Optional[str]
description: Optional[str]
temperature: Optional[float]
model: Optional[str]
max_tokens: Optional[int]
openai_api_key: Optional[str]
status: Optional[str]
prompt_id: Optional[UUID]
def dict(self, *args, **kwargs):
brain_dict = super().dict(*args, **kwargs)
if brain_dict.get("prompt_id"):
brain_dict["prompt_id"] = str(brain_dict.get("prompt_id"))
return brain_dict
class Brain(Repository): class Brain(Repository):
def __init__(self, supabase_client): def __init__(self, supabase_client):
self.db = supabase_client self.db = supabase_client
def get_user_brains(self, user_id): def create_brain(self, brain: CreateBrainProperties):
return BrainEntity(
**((self.db.table("brains").insert(brain)).execute()).data[0]
)
def get_user_brains(self, user_id) -> list[MinimalBrainEntity]:
response = ( response = (
self.db self.db.from_("brains_users")
.from_("brains_users")
.select("id:brain_id, rights, brains (id: brain_id, name)") .select("id:brain_id, rights, brains (id: brain_id, name)")
.filter("user_id", "eq", user_id) .filter("user_id", "eq", user_id)
.execute() .execute()
) )
user_brains = [] user_brains: list[MinimalBrainEntity] = []
for item in response.data: for item in response.data:
user_brains.append(item["brains"]) user_brains.append(
user_brains[-1]["rights"] = item["rights"] MinimalBrainEntity(
id=item["brains"]["id"],
name=item["brains"]["name"],
rights=item["rights"],
)
)
user_brains[-1].rights = item["rights"]
return user_brains return user_brains
def get_brain_for_user(self, user_id, brain_id): def get_brain_for_user(self, user_id, brain_id) -> MinimalBrainEntity | None:
response = ( response = (
self.db.from_("brains_users") self.db.from_("brains_users")
.select("id:brain_id, rights, brains (id: brain_id, name)") .select("id:brain_id, rights, brains (id: brain_id, name)")
@ -34,7 +81,13 @@ class Brain(Repository):
) )
if len(response.data) == 0: if len(response.data) == 0:
return None return None
return response.data[0] brain_data = response.data[0]
return MinimalBrainEntity(
id=brain_data["brains"]["id"],
name=brain_data["brains"]["name"],
rights=brain_data["rights"],
)
def get_brain_details(self, brain_id): def get_brain_details(self, brain_id):
response = ( response = (
@ -81,10 +134,7 @@ class Brain(Repository):
return results return results
def create_brain(self, name): def create_brain_user(self, user_id: UUID, brain_id, rights, default_brain: bool):
return self.db.table("brains").insert({"name": name}).execute()
def create_brain_user(self, user_id: UUID, brain_id, rights, default_brain):
response = ( response = (
self.db.table("brains_users") self.db.table("brains_users")
.insert( .insert(
@ -124,10 +174,20 @@ class Brain(Repository):
) )
return vectorsResponse.data return vectorsResponse.data
def update_brain_fields(self, brain_id, brain_name): def update_brain_by_id(
self.db.table("brains").update({"name": brain_name}).match( self, brain_id: UUID, brain: BrainUpdatableProperties
{"brain_id": brain_id} ) -> BrainEntity | None:
).execute() update_brain_response = (
self.db.table("brains")
.update(brain.dict(exclude_unset=True))
.match({"brain_id": brain_id})
.execute()
).data
if len(update_brain_response) == 0:
return None
return BrainEntity(**update_brain_response[0])
def get_brain_vector_ids(self, brain_id): def get_brain_vector_ids(self, brain_id):
""" """
@ -183,23 +243,29 @@ class Brain(Repository):
return {"message": f"File {file_name} in brain {brain_id} has been deleted."} return {"message": f"File {file_name} in brain {brain_id} has been deleted."}
def get_default_user_brain_id(self, user_id: UUID): def get_default_user_brain_id(self, user_id: UUID) -> UUID | None:
response = ( response = (
(
self.db.from_("brains_users") self.db.from_("brains_users")
.select("brain_id") .select("brain_id")
.filter("user_id", "eq", user_id) .filter("user_id", "eq", user_id)
.filter("default_brain", "eq", True) .filter("default_brain", "eq", True)
.execute() .execute()
) )
).data
if len(response) == 0:
return None
return UUID(response[0].get("brain_id"))
return response def get_brain_by_id(self, brain_id: UUID) -> BrainEntity | None:
def get_brain_by_id(self, brain_id: UUID):
response = ( response = (
self.db.from_("brains") self.db.from_("brains")
.select("id:brain_id, name, *") .select("id:brain_id, name, *")
.filter("brain_id", "eq", brain_id) .filter("brain_id", "eq", brain_id)
.execute() .execute()
) ).data
return response if len(response) == 0:
return None
return BrainEntity(**response[0])

View File

@ -3,7 +3,7 @@ from uuid import UUID
from fastapi import HTTPException from fastapi import HTTPException
from models.databases.repository import Repository from models.databases.repository import Repository
from models.prompt import Prompt from models.prompt import Prompt, PromptStatusEnum
from pydantic import BaseModel from pydantic import BaseModel
@ -12,7 +12,7 @@ class CreatePromptProperties(BaseModel):
title: str title: str
content: str content: str
status: str = "private" status: PromptStatusEnum = PromptStatusEnum.private
class PromptUpdatableProperties(BaseModel): class PromptUpdatableProperties(BaseModel):
@ -20,7 +20,14 @@ class PromptUpdatableProperties(BaseModel):
title: Optional[str] title: Optional[str]
content: Optional[str] content: Optional[str]
status: Optional[str] status: Optional[PromptStatusEnum]
class DeletePromptResponse(BaseModel):
"""Response when deleting a prompt"""
status: str = "delete"
prompt_id: UUID
class Prompts(Repository): class Prompts(Repository):
@ -28,20 +35,22 @@ class Prompts(Repository):
self.db = supabase_client self.db = supabase_client
def create_prompt(self, prompt: CreatePromptProperties) -> Prompt: def create_prompt(self, prompt: CreatePromptProperties) -> Prompt:
"""Create a prompt by id""" """
Create a prompt
"""
response = (self.db.from_("prompts").insert(prompt.dict()).execute()).data response = (self.db.from_("prompts").insert(prompt.dict()).execute()).data
return Prompt(**response[0]) return Prompt(**response[0])
def delete_prompt_by_id(self, prompt_id: UUID) -> Prompt | None: def delete_prompt_by_id(self, prompt_id: UUID) -> DeletePromptResponse:
""" """
Delete a prompt by id Delete a prompt by id
Args: Args:
prompt_id (UUID): The id of the prompt prompt_id (UUID): The id of the prompt
Returns: Returns:
Prompt: The prompt A dictionary containing the status of the delete and prompt_id of the deleted prompt
""" """
response = ( response = (
self.db.from_("prompts") self.db.from_("prompts")
@ -50,9 +59,11 @@ class Prompts(Repository):
.execute() .execute()
.data .data
) )
if response == []: if response == []:
raise HTTPException(404, "Prompt not found") raise HTTPException(404, "Prompt not found")
return Prompt(**response[0])
return DeletePromptResponse(status="deleted", prompt_id=prompt_id)
def get_prompt_by_id(self, prompt_id: UUID) -> Prompt | None: def get_prompt_by_id(self, prompt_id: UUID) -> Prompt | None:
""" """

View File

@ -1,10 +1,16 @@
from enum import Enum
from uuid import UUID from uuid import UUID
from pydantic import BaseModel from pydantic import BaseModel
class PromptStatusEnum(str, Enum):
private = "private"
public = "public"
class Prompt(BaseModel): class Prompt(BaseModel):
title: str title: str
content: str content: str
status: str = "private" status: PromptStatusEnum = PromptStatusEnum.private
id: UUID id: UUID

View File

@ -0,0 +1,9 @@
from models.brain_entity import BrainEntity
from models.databases.supabase.brains import CreateBrainProperties
from models.settings import common_dependencies
def create_brain(brain: CreateBrainProperties) -> BrainEntity:
commons = common_dependencies()
return commons["db"].create_brain(brain.dict(exclude_unset=True))

View File

@ -0,0 +1,16 @@
from uuid import UUID
from models.settings import common_dependencies
from routes.authorizations.types import RoleEnum
def create_brain_user(
user_id: UUID, brain_id: UUID, rights: RoleEnum, is_default_brain: bool
) -> None:
commons = common_dependencies()
commons["db"].create_brain_user(
user_id=user_id,
brain_id=brain_id,
rights=rights,
default_brain=is_default_brain,
).data[0]

View File

@ -0,0 +1,10 @@
from uuid import UUID
from models.brain_entity import BrainEntity
from models.settings import common_dependencies
def get_brain_by_id(brain_id: UUID) -> BrainEntity | None:
commons = common_dependencies()
return commons["db"].get_brain_by_id(brain_id)

View File

@ -0,0 +1,19 @@
from uuid import UUID
from models.brain_entity import BrainEntity
from models.settings import common_dependencies
def get_brain_details(brain_id: UUID) -> BrainEntity | None:
commons = common_dependencies()
response = (
commons["supabase"]
.from_("brains")
.select("*")
.filter("brain_id", "eq", brain_id)
.execute()
)
if response.data == []:
return None
return BrainEntity(**response.data[0])

View File

@ -0,0 +1,9 @@
from uuid import UUID
from models.brain_entity import MinimalBrainEntity
from models.settings import common_dependencies
def get_brain_for_user(user_id: UUID, brain_id: UUID) -> MinimalBrainEntity:
commons = common_dependencies()
return commons["db"].get_brain_for_user(user_id, brain_id)

View File

@ -0,0 +1,22 @@
from uuid import UUID
from logger import get_logger
from models.brain_entity import BrainEntity
from models.settings import common_dependencies
from repository.brain.get_brain_by_id import get_brain_by_id
logger = get_logger(__name__)
def get_user_default_brain(user_id: UUID) -> BrainEntity | None:
commons = common_dependencies()
brain_id = commons["db"].get_default_user_brain_id(user_id)
logger.info("Default brain response:", brain_id)
if brain_id is None:
return None
logger.info(f"Default brain id: {brain_id}")
return get_brain_by_id(brain_id)

View File

@ -0,0 +1,17 @@
from models.brain_entity import BrainEntity
from models.databases.supabase.brains import CreateBrainProperties
from models.users import User
from repository.brain.create_brain import create_brain
from repository.brain.create_brain_user import create_brain_user
from repository.brain.get_default_user_brain import get_user_default_brain
from routes.authorizations.types import RoleEnum
def get_default_user_brain_or_create_new(user: User) -> BrainEntity:
default_brain = get_user_default_brain(user.id)
if not default_brain:
default_brain = create_brain(CreateBrainProperties())
create_brain_user(user.id, default_brain.brain_id, RoleEnum.Owner, True)
return default_brain

View File

@ -0,0 +1,11 @@
from uuid import UUID
from models.brain_entity import BrainEntity
from models.settings import common_dependencies
def get_user_brains(user_id: UUID) -> list[BrainEntity]:
commons = common_dependencies()
results = commons["db"].get_user_brains(user_id)
return results

View File

@ -0,0 +1,19 @@
from uuid import UUID
from models.settings import common_dependencies
from repository.brain.get_default_user_brain import get_user_default_brain
def set_as_default_brain_for_user(user_id: UUID, brain_id: UUID):
commons = common_dependencies()
old_default_brain = get_user_default_brain(user_id)
if old_default_brain is not None:
commons["supabase"].table("brains_users").update(
{"default_brain": False}
).match({"brain_id": old_default_brain.brain_id, "user_id": user_id}).execute()
commons["supabase"].table("brains_users").update({"default_brain": True}).match(
{"brain_id": brain_id, "user_id": user_id}
).execute()

View File

@ -0,0 +1,12 @@
from uuid import UUID
from models.brain_entity import BrainEntity
from models.databases.supabase.brains import BrainUpdatableProperties
from models.settings import common_dependencies
def update_brain_by_id(brain_id: UUID, brain: BrainUpdatableProperties) -> BrainEntity:
"""Update a prompt by id"""
commons = common_dependencies()
return commons["db"].update_brain_by_id(brain_id, brain)

View File

@ -1,9 +1,9 @@
import resend import resend
from logger import get_logger from logger import get_logger
from models.brains import Brain
from models.brains_subscription_invitations import BrainSubscription from models.brains_subscription_invitations import BrainSubscription
from models.settings import BrainSettings from models.settings import BrainSettings
from repository.brain.get_brain_details import get_brain_details
from repository.brain_subscription.get_brain_url import get_brain_url from repository.brain_subscription.get_brain_url import get_brain_url
logger = get_logger(__name__) logger = get_logger(__name__)
@ -19,11 +19,10 @@ def resend_invitation_email(
brain_url = get_brain_url(origin, brain_subscription.brain_id) brain_url = get_brain_url(origin, brain_subscription.brain_id)
invitation_brain_client = Brain(id=brain_subscription.brain_id) invitation_brain = get_brain_details(brain_subscription.brain_id)
invitation_brain = invitation_brain_client.get_brain_details()
if invitation_brain is None: if invitation_brain is None:
raise Exception("Brain not found") raise Exception("Brain not found")
brain_name = invitation_brain["name"] brain_name = invitation_brain.name
html_body = f""" html_body = f"""
<p>Brain {brain_name} has been shared with you by {inviter_email}.</p> <p>Brain {brain_name} has been shared with you by {inviter_email}.</p>

View File

@ -4,7 +4,6 @@ from models.settings import common_dependencies
def create_prompt(prompt: CreatePromptProperties) -> Prompt: def create_prompt(prompt: CreatePromptProperties) -> Prompt:
"""Create a prompt by id"""
commons = common_dependencies() commons = common_dependencies()
return commons["db"].create_prompt(prompt) return commons["db"].create_prompt(prompt)

View File

@ -1,10 +1,10 @@
from uuid import UUID from uuid import UUID
from models.prompt import Prompt from models.databases.supabase.prompts import DeletePromptResponse
from models.settings import common_dependencies from models.settings import common_dependencies
def delete_prompt_by_id(prompt_id: UUID) -> Prompt | None: def delete_prompt_by_id(prompt_id: UUID) -> DeletePromptResponse:
""" """
Delete a prompt by id Delete a prompt by id
Args: Args:

View File

@ -1,17 +1,12 @@
from enum import Enum
from typing import List, Optional, Union from typing import List, Optional, Union
from uuid import UUID from uuid import UUID
from auth.auth_bearer import get_current_user from auth.auth_bearer import get_current_user
from fastapi import Depends, HTTPException, status from fastapi import Depends, HTTPException, status
from models.brains import Brain
from models.users import User from models.users import User
from repository.brain.get_brain_for_user import get_brain_for_user
from routes.authorizations.types import RoleEnum
class RoleEnum(str, Enum):
Viewer = "Viewer"
Editor = "Editor"
Owner = "Owner"
def has_brain_authorization( def has_brain_authorization(
@ -53,8 +48,7 @@ def validate_brain_authorization(
detail="Missing required role", detail="Missing required role",
) )
brain = Brain(id=brain_id) user_brain = get_brain_for_user(user_id, brain_id)
user_brain = brain.get_brain_for_user(user_id)
if user_brain is None: if user_brain is None:
raise HTTPException( raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND, status_code=status.HTTP_404_NOT_FOUND,
@ -66,7 +60,7 @@ def validate_brain_authorization(
required_roles = [required_roles] required_roles = [required_roles]
# Check if the user has at least one of the required roles # Check if the user has at least one of the required roles
if user_brain.get("rights") not in required_roles: if user_brain.rights not in required_roles:
raise HTTPException( raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN, status_code=status.HTTP_403_FORBIDDEN,
detail="You don't have the required role(s) for this brain", detail="You don't have the required role(s) for this brain",

View File

@ -0,0 +1,7 @@
from enum import Enum
class RoleEnum(str, Enum):
Viewer = "Viewer"
Editor = "Editor"
Owner = "Owner"

View File

@ -3,15 +3,31 @@ from uuid import UUID
from auth import AuthBearer, get_current_user from auth import AuthBearer, get_current_user
from fastapi import APIRouter, Depends, HTTPException from fastapi import APIRouter, Depends, HTTPException
from logger import get_logger from logger import get_logger
from models.brains import ( from models.databases.supabase.brains import (
Brain, BrainUpdatableProperties,
get_default_user_brain, CreateBrainProperties,
get_default_user_brain_or_create_new,
) )
from models.settings import BrainRateLimiting from models.settings import BrainRateLimiting
from models.users import User from models.users import User
from repository.brain.create_brain import create_brain
from repository.brain.create_brain_user import create_brain_user
from repository.brain.get_brain_details import get_brain_details
from repository.brain.get_default_user_brain import get_user_default_brain
from repository.brain.get_default_user_brain_or_create_new import (
get_default_user_brain_or_create_new,
)
from repository.brain.get_user_brains import get_user_brains
from repository.brain.set_as_default_brain_for_user import (
set_as_default_brain_for_user,
)
from repository.brain.update_brain import update_brain_by_id
from repository.prompt.delete_prompt_py_id import delete_prompt_by_id
from repository.prompt.get_prompt_by_id import get_prompt_by_id
from routes.authorizations.brain_authorization import RoleEnum, has_brain_authorization from routes.authorizations.brain_authorization import (
has_brain_authorization,
)
from routes.authorizations.types import RoleEnum
logger = get_logger(__name__) logger = get_logger(__name__)
@ -30,8 +46,7 @@ async def brain_endpoint(current_user: User = Depends(get_current_user)):
This endpoint retrieves all the brains associated with the current authenticated user. It returns a list of brains objects 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. containing the brain ID and brain name for each brain.
""" """
brain = Brain() brains = get_user_brains(current_user.id)
brains = brain.get_user_brains(current_user.id)
return {"brains": brains} return {"brains": brains}
@ -51,10 +66,9 @@ async def get_default_brain_endpoint(current_user: User = Depends(get_current_us
""" """
brain = get_default_user_brain_or_create_new(current_user) brain = get_default_user_brain_or_create_new(current_user)
return {"id": brain.id, "name": brain.name, "rights": "Owner"} return {"id": brain.brain_id, "name": brain.name, "rights": "Owner"}
# get one brain - Currently not used in FE
@brain_router.get( @brain_router.get(
"/brains/{brain_id}/", "/brains/{brain_id}/",
dependencies=[Depends(AuthBearer()), Depends(has_brain_authorization())], dependencies=[Depends(AuthBearer()), Depends(has_brain_authorization())],
@ -72,9 +86,8 @@ async def get_brain_endpoint(
This endpoint retrieves the details of a specific brain identified by the provided brain ID. It returns the brain ID and its 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. history, which includes the brain messages exchanged in the brain.
""" """
brain = Brain(id=brain_id)
brain_details = brain.get_brain_details() brain_details = get_brain_details(brain_id)
if brain_details is None: if brain_details is None:
raise HTTPException( raise HTTPException(
status_code=404, status_code=404,
@ -87,7 +100,7 @@ async def get_brain_endpoint(
# create new brain # create new brain
@brain_router.post("/brains/", dependencies=[Depends(AuthBearer())], tags=["Brain"]) @brain_router.post("/brains/", dependencies=[Depends(AuthBearer())], tags=["Brain"])
async def create_brain_endpoint( async def create_brain_endpoint(
brain: Brain, brain: CreateBrainProperties,
current_user: User = Depends(get_current_user), current_user: User = Depends(get_current_user),
): ):
""" """
@ -100,7 +113,7 @@ async def create_brain_endpoint(
In the brains table & in the brains_users table and put the creator user as 'Owner' In the brains table & in the brains_users table and put the creator user as 'Owner'
""" """
user_brains = brain.get_user_brains(current_user.id) user_brains = get_user_brains(current_user.id)
max_brain_per_user = BrainRateLimiting().max_brain_per_user max_brain_per_user = BrainRateLimiting().max_brain_per_user
if len(user_brains) >= max_brain_per_user: if len(user_brains) >= max_brain_per_user:
@ -109,23 +122,31 @@ async def create_brain_endpoint(
detail=f"Maximum number of brains reached ({max_brain_per_user}).", detail=f"Maximum number of brains reached ({max_brain_per_user}).",
) )
brain.create_brain() # pyright: ignore reportPrivateUsage=none new_brain = create_brain(
default_brain = get_default_user_brain(current_user) brain,
)
default_brain = get_user_default_brain(current_user.id)
if default_brain: if default_brain:
logger.info(f"Default brain already exists for user {current_user.id}") logger.info(f"Default brain already exists for user {current_user.id}")
brain.create_brain_user( # pyright: ignore reportPrivateUsage=none create_brain_user(
user_id=current_user.id, rights="Owner", default_brain=False user_id=current_user.id,
brain_id=new_brain.brain_id,
rights=RoleEnum.Owner,
is_default_brain=False,
) )
else: else:
logger.info( logger.info(
f"Default brain does not exist for user {current_user.id}. It will be created." f"Default brain does not exist for user {current_user.id}. It will be created."
) )
brain.create_brain_user( # pyright: ignore reportPrivateUsage=none create_brain_user(
user_id=current_user.id, rights="Owner", default_brain=True user_id=current_user.id,
brain_id=new_brain.brain_id,
rights=RoleEnum.Owner,
is_default_brain=True,
) )
return { return {
"id": brain.id, # pyright: ignore reportPrivateUsage=none "id": new_brain.brain_id,
"name": brain.name, "name": brain.name,
"rights": "Owner", "rights": "Owner",
} }
@ -144,15 +165,28 @@ async def create_brain_endpoint(
) )
async def update_brain_endpoint( async def update_brain_endpoint(
brain_id: UUID, brain_id: UUID,
input_brain: Brain, input_brain: BrainUpdatableProperties,
): ):
""" """
Update an existing brain with new brain configuration Update an existing brain with new brain configuration
""" """
input_brain.id = brain_id
print("brain", input_brain)
input_brain.update_brain_fields() # 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."} return {"message": f"Brain {brain_id} has been updated."}
@ -174,8 +208,7 @@ async def set_as_default_brain_endpoint(
""" """
Set a brain as default for the current user. Set a brain as default for the current user.
""" """
brain = Brain(id=brain_id)
brain.set_as_default_brain_for_user(user) set_as_default_brain_for_user(user.id, brain_id)
return {"message": f"Brain {brain_id} has been set as default brain."} return {"message": f"Brain {brain_id} has been set as default brain."}

View File

@ -9,11 +9,15 @@ from auth import AuthBearer, get_current_user
from fastapi import APIRouter, Depends, Query, Request from fastapi import APIRouter, Depends, Query, Request
from fastapi.responses import StreamingResponse from fastapi.responses import StreamingResponse
from llm.openai import OpenAIBrainPicking from llm.openai import OpenAIBrainPicking
from models.brains import Brain, get_default_user_brain_or_create_new from models.brains import Brain
from models.chat import Chat, ChatHistory from models.chat import Chat, ChatHistory
from models.chats import ChatQuestion from models.chats import ChatQuestion
from models.settings import LLMSettings, common_dependencies from models.settings import LLMSettings, common_dependencies
from models.users import User from models.users import User
from repository.brain.get_brain_details import get_brain_details
from repository.brain.get_default_user_brain_or_create_new import (
get_default_user_brain_or_create_new,
)
from repository.chat.create_chat import CreateChatProperties, create_chat from repository.chat.create_chat import CreateChatProperties, create_chat
from repository.chat.get_chat_by_id import get_chat_by_id from repository.chat.get_chat_by_id import get_chat_by_id
from repository.chat.get_chat_history import get_chat_history from repository.chat.get_chat_history import get_chat_history
@ -57,10 +61,7 @@ def delete_chat_from_db(commons, chat_id):
def fetch_user_stats(commons, user, date): def fetch_user_stats(commons, user, date):
response = ( response = commons["db"].get_user_stats(user.email, date)
commons["db"]
.get_user_stats(user.email, date)
)
userItem = next(iter(response.data or []), {"requests_count": 0}) userItem = next(iter(response.data or []), {"requests_count": 0})
return userItem return userItem
@ -173,9 +174,10 @@ async def create_question_handler(
brain = Brain(id=brain_id) brain = Brain(id=brain_id)
if not current_user.user_openai_api_key: if not current_user.user_openai_api_key:
brain_details = brain.get_brain_details() if brain_id:
brain_details = get_brain_details(brain_id)
if brain_details: if brain_details:
current_user.user_openai_api_key = brain_details["openai_api_key"] current_user.user_openai_api_key = brain_details.openai_api_key
if not current_user.user_openai_api_key: if not current_user.user_openai_api_key:
user_identity = get_user_identity(current_user.id) user_identity = get_user_identity(current_user.id)
@ -199,7 +201,7 @@ async def create_question_handler(
LLMSettings() LLMSettings()
if not brain_id: if not brain_id:
brain_id = get_default_user_brain_or_create_new(current_user).id brain_id = get_default_user_brain_or_create_new(current_user).brain_id
gpt_answer_generator = OpenAIBrainPicking( gpt_answer_generator = OpenAIBrainPicking(
chat_id=str(chat_id), chat_id=str(chat_id),
@ -244,10 +246,10 @@ async def create_stream_question_handler(
current_user.user_openai_api_key = request.headers.get("Openai-Api-Key") current_user.user_openai_api_key = request.headers.get("Openai-Api-Key")
brain = Brain(id=brain_id) brain = Brain(id=brain_id)
if not current_user.user_openai_api_key: if not current_user.user_openai_api_key and brain_id:
brain_details = brain.get_brain_details() brain_details = get_brain_details(brain_id)
if brain_details: if brain_details:
current_user.user_openai_api_key = brain_details["openai_api_key"] current_user.user_openai_api_key = brain_details.openai_api_key
if not current_user.user_openai_api_key: if not current_user.user_openai_api_key:
user_identity = get_user_identity(current_user.id) user_identity = get_user_identity(current_user.id)
@ -270,7 +272,7 @@ async def create_stream_question_handler(
logger.info(f"Streaming request for {chat_question.model}") logger.info(f"Streaming request for {chat_question.model}")
check_user_limit(current_user) check_user_limit(current_user)
if not brain_id: if not brain_id:
brain_id = get_default_user_brain_or_create_new(current_user).id brain_id = get_default_user_brain_or_create_new(current_user).brain_id
gpt_answer_generator = OpenAIBrainPicking( gpt_answer_generator = OpenAIBrainPicking(
chat_id=str(chat_id), chat_id=str(chat_id),

View File

@ -1,7 +1,16 @@
from uuid import UUID
from auth import AuthBearer from auth import AuthBearer
from fastapi import APIRouter, Depends from fastapi import APIRouter, Depends
from models.databases.supabase.prompts import (
CreatePromptProperties,
PromptUpdatableProperties,
)
from models.prompt import Prompt from models.prompt import Prompt
from repository.prompt.create_prompt import create_prompt
from repository.prompt.get_prompt_by_id import get_prompt_by_id
from repository.prompt.get_public_prompts import get_public_prompts from repository.prompt.get_public_prompts import get_public_prompts
from repository.prompt.update_prompt_by_id import update_prompt_by_id
prompt_router = APIRouter() prompt_router = APIRouter()
@ -13,3 +22,36 @@ async def get_prompts() -> list[Prompt]:
""" """
return get_public_prompts() return get_public_prompts()
@prompt_router.get(
"/prompts/{prompt_id}", dependencies=[Depends(AuthBearer())], tags=["Prompt"]
)
async def get_prompt(prompt_id: UUID) -> Prompt | None:
"""
Retrieve a prompt by its id
"""
return get_prompt_by_id(prompt_id)
@prompt_router.put(
"/prompts/{prompt_id}", dependencies=[Depends(AuthBearer())], tags=["Prompt"]
)
async def update_prompt(
prompt_id: UUID, prompt: PromptUpdatableProperties
) -> Prompt | None:
"""
Update a prompt by its id
"""
return update_prompt_by_id(prompt_id, prompt)
@prompt_router.post("/prompts", dependencies=[Depends(AuthBearer())], tags=["Prompt"])
async def create_prompt_route(prompt: CreatePromptProperties) -> Prompt | None:
"""
Create a prompt by its id
"""
return create_prompt(prompt)

View File

@ -7,6 +7,9 @@ from models.brains import Brain
from models.brains_subscription_invitations import BrainSubscription from models.brains_subscription_invitations import BrainSubscription
from models.users import User from models.users import User
from pydantic import BaseModel from pydantic import BaseModel
from repository.brain.create_brain_user import create_brain_user
from repository.brain.get_brain_details import get_brain_details
from repository.brain.get_brain_for_user import get_brain_for_user
from repository.brain.update_user_rights import update_brain_user_rights from repository.brain.update_user_rights import update_brain_user_rights
from repository.brain_subscription.resend_invitation_email import ( from repository.brain_subscription.resend_invitation_email import (
resend_invitation_email, resend_invitation_email,
@ -121,14 +124,14 @@ async def remove_user_subscription(
brain = Brain( brain = Brain(
id=brain_id, id=brain_id,
) )
user_brain = brain.get_brain_for_user(current_user.id) user_brain = get_brain_for_user(current_user.id, brain_id)
if user_brain is None: if user_brain is None:
raise HTTPException( raise HTTPException(
status_code=403, status_code=403,
detail="You don't have permission for this brain", detail="You don't have permission for this brain",
) )
if user_brain.get("rights") != "Owner": if user_brain.rights != "Owner":
brain.delete_user_from_brain(current_user.id) brain.delete_user_from_brain(current_user.id)
else: else:
brain_users = brain.get_brain_users() brain_users = brain.get_brain_users()
@ -170,8 +173,7 @@ def get_user_invitation(brain_id: UUID, current_user: User = Depends(get_current
detail="You have not been invited to this brain", detail="You have not been invited to this brain",
) )
brain = Brain(id=brain_id) brain_details = get_brain_details(brain_id)
brain_details = brain.get_brain_details()
if brain_details is None: if brain_details is None:
raise HTTPException( raise HTTPException(
@ -179,7 +181,7 @@ def get_user_invitation(brain_id: UUID, current_user: User = Depends(get_current
detail="Brain not found while trying to get invitation", detail="Brain not found while trying to get invitation",
) )
return {"name": brain_details["name"], "rights": invitation["rights"]} return {"name": brain_details.name, "rights": invitation["rights"]}
@subscription_router.post( @subscription_router.post(
@ -208,9 +210,11 @@ async def accept_invitation(
raise HTTPException(status_code=404, detail="Invitation not found") raise HTTPException(status_code=404, detail="Invitation not found")
try: try:
brain = Brain(id=brain_id) create_brain_user(
brain.create_brain_user( user_id=current_user.id,
user_id=current_user.id, rights=invitation["rights"], default_brain=False brain_id=brain_id,
rights=invitation["rights"],
is_default_brain=False,
) )
except Exception as e: except Exception as e:
raise HTTPException(status_code=400, detail=f"Error adding user to brain: {e}") raise HTTPException(status_code=400, detail=f"Error adding user to brain: {e}")
@ -299,8 +303,8 @@ def update_brain_subscription(
) )
# check if user is not an editor trying to update an owner right which is not allowed # check if user is not an editor trying to update an owner right which is not allowed
current_invitation = brain.get_brain_for_user(user_id) current_invitation = get_brain_for_user(user_id, brain_id)
if current_invitation is not None and current_invitation.get("rights") == "Owner": if current_invitation is not None and current_invitation.rights == "Owner":
try: try:
validate_brain_authorization( validate_brain_authorization(
brain_id, brain_id,

View File

@ -7,6 +7,7 @@ from models.brains import Brain
from models.files import File from models.files import File
from models.settings import common_dependencies from models.settings import common_dependencies
from models.users import User from models.users import User
from repository.brain.get_brain_details import get_brain_details
from repository.user_identity.get_user_identity import get_user_identity from repository.user_identity.get_user_identity import get_user_identity
from utils.file import convert_bytes, get_file_size from utils.file import convert_bytes, get_file_size
from utils.processors import filter_file from utils.processors import filter_file
@ -62,9 +63,9 @@ async def upload_file(
else: else:
openai_api_key = request.headers.get("Openai-Api-Key", None) openai_api_key = request.headers.get("Openai-Api-Key", None)
if openai_api_key is None: if openai_api_key is None:
brain_details = brain.get_brain_details() brain_details = get_brain_details(brain_id)
if brain_details: if brain_details:
openai_api_key = brain_details["openai_api_key"] openai_api_key = brain_details.openai_api_key
if openai_api_key is None: if openai_api_key is None:
openai_api_key = get_user_identity(current_user.id).openai_api_key openai_api_key = get_user_identity(current_user.id).openai_api_key

View File

@ -3,10 +3,11 @@ import time
from auth import AuthBearer, get_current_user from auth import AuthBearer, get_current_user
from fastapi import APIRouter, Depends, Request from fastapi import APIRouter, Depends, Request
from models.brains import Brain, get_default_user_brain from models.brains import Brain
from models.settings import BrainRateLimiting from models.settings import BrainRateLimiting
from models.user_identity import UserIdentity from models.user_identity import UserIdentity
from models.users import User from models.users import User
from repository.brain.get_default_user_brain import get_user_default_brain
from repository.user_identity.get_user_identity import get_user_identity from repository.user_identity.get_user_identity import get_user_identity
from repository.user_identity.update_user_identity import ( from repository.user_identity.update_user_identity import (
UserIdentityUpdatableProperties, UserIdentityUpdatableProperties,
@ -47,10 +48,10 @@ async def get_user_endpoint(
date = time.strftime("%Y%m%d") date = time.strftime("%Y%m%d")
max_requests_number = os.getenv("MAX_REQUESTS_NUMBER") max_requests_number = os.getenv("MAX_REQUESTS_NUMBER")
requests_stats = current_user.get_user_request_stats() requests_stats = current_user.get_user_request_stats()
default_brain = get_default_user_brain(current_user) default_brain = get_user_default_brain(current_user.id)
if default_brain: if default_brain:
defaul_brain_size = Brain(id=default_brain["id"]).brain_size defaul_brain_size = Brain(id=default_brain.brain_id).brain_size
else: else:
defaul_brain_size = 0 defaul_brain_size = 0

View File

@ -1,7 +1,7 @@
import random import random
import string import string
from models.brains import get_default_user_brain from repository.brain.get_default_user_brain import get_user_default_brain
def test_retrieve_default_brain(client, api_key): def test_retrieve_default_brain(client, api_key):
@ -211,7 +211,6 @@ def test_set_as_default_brain_endpoint(client, api_key):
user_info = response.json() user_info = response.json()
user_id = user_info["id"] user_id = user_info["id"]
default_brain = get_default_user_brain(user_id) default_brain = get_user_default_brain(user_id)
assert default_brain is not None assert default_brain is not None
assert default_brain["id"] == brain_id assert default_brain.brain_id == brain_id
assert default_brain["default_brain"] is True

View File

@ -0,0 +1,19 @@
BEGIN;
-- Check if prompt_id column exists
DO $$
BEGIN
IF NOT EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name = 'brains' AND column_name = 'prompt_id') THEN
-- Add prompt_id column and reference the table prompts' id column
ALTER TABLE brains ADD COLUMN prompt_id UUID REFERENCES prompts(id);
END IF;
END $$;
-- Update migrations table
INSERT INTO migrations (name)
SELECT '20230802120700_add_prompt_id_to_brain'
WHERE NOT EXISTS (
SELECT 1 FROM migrations WHERE name = '20230802120700_add_prompt_id_to_brain'
);
COMMIT;

View File

@ -134,9 +134,11 @@ CREATE TABLE IF NOT EXISTS brains (
model TEXT, model TEXT,
max_tokens INT, max_tokens INT,
temperature FLOAT, temperature FLOAT,
openai_api_key TEXT openai_api_key TEXT,
prompt_id UUID REFERENCES prompts(id)
); );
-- Create brains X users table -- Create brains X users table
CREATE TABLE IF NOT EXISTS brains_users ( CREATE TABLE IF NOT EXISTS brains_users (
brain_id UUID, brain_id UUID,
@ -210,7 +212,7 @@ CREATE TABLE IF NOT EXISTS migrations (
); );
INSERT INTO migrations (name) INSERT INTO migrations (name)
SELECT '20230701180101_add_prompts_table' SELECT '20230802120700_add_prompt_id_to_brain'
WHERE NOT EXISTS ( WHERE NOT EXISTS (
SELECT 1 FROM migrations WHERE name = '20230701180101_add_prompts_table' SELECT 1 FROM migrations WHERE name = '20230802120700_add_prompt_id_to_brain'
); );