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 models.settings import BrainRateLimiting, CommonsDep, common_dependencies
from models.users import User
logger = get_logger(__name__)
@ -22,6 +21,7 @@ class Brain(BaseModel):
openai_api_key: Optional[str] = None
files: List[Any] = []
max_brain_size = BrainRateLimiting().max_brain_size
prompt_id: Optional[UUID] = None
class Config:
arbitrary_types_allowed = True
@ -77,24 +77,6 @@ class Brain(BaseModel):
{"brain_id": self.id, "user_id": user_id}
).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):
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(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):
return self.commons["db"].create_brain_vector(self.id, vector_id, file_sha1)
def get_vector_ids_from_file_sha1(self, file_sha1: str):
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):
# not used
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):
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 uuid import UUID
from models.brain_entity import BrainEntity
class Repository(ABC):
@abstractmethod
def get_user_brains(self, user_id: str):
def get_user_brains(self, user_id: str) -> list[BrainEntity]:
pass
@abstractmethod
@ -29,7 +31,7 @@ class Repository(ABC):
pass
@abstractmethod
def create_brain(self, name: str):
def create_brain(self, brain: str):
pass
@abstractmethod
@ -46,10 +48,6 @@ class Repository(ABC):
def get_vector_ids_from_file_sha1(self, file_sha1: str):
pass
@abstractmethod
def update_brain_fields(self, brain_id: UUID, brain_name: str):
pass
@abstractmethod
def get_brain_vector_ids(self, brain_id: UUID):
pass
@ -59,7 +57,7 @@ class Repository(ABC):
pass
@abstractmethod
def get_default_user_brain_id(self, user_id: UUID):
def get_default_user_brain_id(self, user_id: UUID) -> UUID:
pass
@abstractmethod

View File

@ -1,30 +1,77 @@
from typing import Optional
from uuid import UUID
from models.databases.repository import Repository
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__)
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):
def __init__(self, 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 = (
self.db
.from_("brains_users")
self.db.from_("brains_users")
.select("id:brain_id, rights, brains (id: brain_id, name)")
.filter("user_id", "eq", user_id)
.execute()
)
user_brains = []
user_brains: list[MinimalBrainEntity] = []
for item in response.data:
user_brains.append(item["brains"])
user_brains[-1]["rights"] = item["rights"]
user_brains.append(
MinimalBrainEntity(
id=item["brains"]["id"],
name=item["brains"]["name"],
rights=item["rights"],
)
)
user_brains[-1].rights = item["rights"]
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 = (
self.db.from_("brains_users")
.select("id:brain_id, rights, brains (id: brain_id, name)")
@ -34,7 +81,13 @@ class Brain(Repository):
)
if len(response.data) == 0:
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):
response = (
@ -81,10 +134,7 @@ class Brain(Repository):
return results
def create_brain(self, name):
return self.db.table("brains").insert({"name": name}).execute()
def create_brain_user(self, user_id: UUID, brain_id, rights, default_brain):
def create_brain_user(self, user_id: UUID, brain_id, rights, default_brain: bool):
response = (
self.db.table("brains_users")
.insert(
@ -124,10 +174,20 @@ class Brain(Repository):
)
return vectorsResponse.data
def update_brain_fields(self, brain_id, brain_name):
self.db.table("brains").update({"name": brain_name}).match(
{"brain_id": brain_id}
).execute()
def update_brain_by_id(
self, brain_id: UUID, brain: BrainUpdatableProperties
) -> BrainEntity | None:
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):
"""
@ -183,23 +243,29 @@ class Brain(Repository):
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 = (
(
self.db.from_("brains_users")
.select("brain_id")
.filter("user_id", "eq", user_id)
.filter("default_brain", "eq", True)
.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):
def get_brain_by_id(self, brain_id: UUID) -> BrainEntity | None:
response = (
self.db.from_("brains")
.select("id:brain_id, name, *")
.filter("brain_id", "eq", brain_id)
.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 models.databases.repository import Repository
from models.prompt import Prompt
from models.prompt import Prompt, PromptStatusEnum
from pydantic import BaseModel
@ -12,7 +12,7 @@ class CreatePromptProperties(BaseModel):
title: str
content: str
status: str = "private"
status: PromptStatusEnum = PromptStatusEnum.private
class PromptUpdatableProperties(BaseModel):
@ -20,7 +20,14 @@ class PromptUpdatableProperties(BaseModel):
title: 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):
@ -28,20 +35,22 @@ class Prompts(Repository):
self.db = supabase_client
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
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
Args:
prompt_id (UUID): The id of the prompt
Returns:
Prompt: The prompt
A dictionary containing the status of the delete and prompt_id of the deleted prompt
"""
response = (
self.db.from_("prompts")
@ -50,9 +59,11 @@ class Prompts(Repository):
.execute()
.data
)
if response == []:
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:
"""

View File

@ -1,10 +1,16 @@
from enum import Enum
from uuid import UUID
from pydantic import BaseModel
class PromptStatusEnum(str, Enum):
private = "private"
public = "public"
class Prompt(BaseModel):
title: str
content: str
status: str = "private"
status: PromptStatusEnum = PromptStatusEnum.private
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
from logger import get_logger
from models.brains import Brain
from models.brains_subscription_invitations import BrainSubscription
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
logger = get_logger(__name__)
@ -19,11 +19,10 @@ def resend_invitation_email(
brain_url = get_brain_url(origin, brain_subscription.brain_id)
invitation_brain_client = Brain(id=brain_subscription.brain_id)
invitation_brain = invitation_brain_client.get_brain_details()
invitation_brain = get_brain_details(brain_subscription.brain_id)
if invitation_brain is None:
raise Exception("Brain not found")
brain_name = invitation_brain["name"]
brain_name = invitation_brain.name
html_body = f"""
<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:
"""Create a prompt by id"""
commons = common_dependencies()
return commons["db"].create_prompt(prompt)

View File

@ -1,10 +1,10 @@
from uuid import UUID
from models.prompt import Prompt
from models.databases.supabase.prompts import DeletePromptResponse
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
Args:

View File

@ -1,17 +1,12 @@
from enum import Enum
from typing import List, Optional, Union
from uuid import UUID
from auth.auth_bearer import get_current_user
from fastapi import Depends, HTTPException, status
from models.brains import Brain
from models.users import User
from repository.brain.get_brain_for_user import get_brain_for_user
class RoleEnum(str, Enum):
Viewer = "Viewer"
Editor = "Editor"
Owner = "Owner"
from routes.authorizations.types import RoleEnum
def has_brain_authorization(
@ -53,8 +48,7 @@ def validate_brain_authorization(
detail="Missing required role",
)
brain = Brain(id=brain_id)
user_brain = brain.get_brain_for_user(user_id)
user_brain = get_brain_for_user(user_id, brain_id)
if user_brain is None:
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
@ -66,7 +60,7 @@ def validate_brain_authorization(
required_roles = [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(
status_code=status.HTTP_403_FORBIDDEN,
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 fastapi import APIRouter, Depends, HTTPException
from logger import get_logger
from models.brains import (
Brain,
get_default_user_brain,
get_default_user_brain_or_create_new,
from models.databases.supabase.brains import (
BrainUpdatableProperties,
CreateBrainProperties,
)
from models.settings import BrainRateLimiting
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__)
@ -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
containing the brain ID and brain name for each brain.
"""
brain = Brain()
brains = brain.get_user_brains(current_user.id)
brains = get_user_brains(current_user.id)
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)
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(
"/brains/{brain_id}/",
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
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:
raise HTTPException(
status_code=404,
@ -87,7 +100,7 @@ async def get_brain_endpoint(
# create new brain
@brain_router.post("/brains/", dependencies=[Depends(AuthBearer())], tags=["Brain"])
async def create_brain_endpoint(
brain: Brain,
brain: CreateBrainProperties,
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'
"""
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
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}).",
)
brain.create_brain() # pyright: ignore reportPrivateUsage=none
default_brain = get_default_user_brain(current_user)
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}")
brain.create_brain_user( # pyright: ignore reportPrivateUsage=none
user_id=current_user.id, rights="Owner", default_brain=False
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."
)
brain.create_brain_user( # pyright: ignore reportPrivateUsage=none
user_id=current_user.id, rights="Owner", default_brain=True
create_brain_user(
user_id=current_user.id,
brain_id=new_brain.brain_id,
rights=RoleEnum.Owner,
is_default_brain=True,
)
return {
"id": brain.id, # pyright: ignore reportPrivateUsage=none
"id": new_brain.brain_id,
"name": brain.name,
"rights": "Owner",
}
@ -144,15 +165,28 @@ async def create_brain_endpoint(
)
async def update_brain_endpoint(
brain_id: UUID,
input_brain: Brain,
input_brain: BrainUpdatableProperties,
):
"""
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."}
@ -174,8 +208,7 @@ async def set_as_default_brain_endpoint(
"""
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."}

View File

@ -9,11 +9,15 @@ from auth import AuthBearer, get_current_user
from fastapi import APIRouter, Depends, Query, Request
from fastapi.responses import StreamingResponse
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.chats import ChatQuestion
from models.settings import LLMSettings, common_dependencies
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.get_chat_by_id import get_chat_by_id
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):
response = (
commons["db"]
.get_user_stats(user.email, date)
)
response = commons["db"].get_user_stats(user.email, date)
userItem = next(iter(response.data or []), {"requests_count": 0})
return userItem
@ -173,9 +174,10 @@ async def create_question_handler(
brain = Brain(id=brain_id)
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:
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:
user_identity = get_user_identity(current_user.id)
@ -199,7 +201,7 @@ async def create_question_handler(
LLMSettings()
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(
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")
brain = Brain(id=brain_id)
if not current_user.user_openai_api_key:
brain_details = brain.get_brain_details()
if not current_user.user_openai_api_key and brain_id:
brain_details = get_brain_details(brain_id)
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:
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}")
check_user_limit(current_user)
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(
chat_id=str(chat_id),

View File

@ -1,7 +1,16 @@
from uuid import UUID
from auth import AuthBearer
from fastapi import APIRouter, Depends
from models.databases.supabase.prompts import (
CreatePromptProperties,
PromptUpdatableProperties,
)
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.update_prompt_by_id import update_prompt_by_id
prompt_router = APIRouter()
@ -13,3 +22,36 @@ async def get_prompts() -> list[Prompt]:
"""
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.users import User
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_subscription.resend_invitation_email import (
resend_invitation_email,
@ -121,14 +124,14 @@ async def remove_user_subscription(
brain = Brain(
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:
raise HTTPException(
status_code=403,
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)
else:
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",
)
brain = Brain(id=brain_id)
brain_details = brain.get_brain_details()
brain_details = get_brain_details(brain_id)
if brain_details is None:
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",
)
return {"name": brain_details["name"], "rights": invitation["rights"]}
return {"name": brain_details.name, "rights": invitation["rights"]}
@subscription_router.post(
@ -208,9 +210,11 @@ async def accept_invitation(
raise HTTPException(status_code=404, detail="Invitation not found")
try:
brain = Brain(id=brain_id)
brain.create_brain_user(
user_id=current_user.id, rights=invitation["rights"], default_brain=False
create_brain_user(
user_id=current_user.id,
brain_id=brain_id,
rights=invitation["rights"],
is_default_brain=False,
)
except Exception as 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
current_invitation = brain.get_brain_for_user(user_id)
if current_invitation is not None and current_invitation.get("rights") == "Owner":
current_invitation = get_brain_for_user(user_id, brain_id)
if current_invitation is not None and current_invitation.rights == "Owner":
try:
validate_brain_authorization(
brain_id,

View File

@ -7,6 +7,7 @@ from models.brains import Brain
from models.files import File
from models.settings import common_dependencies
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 utils.file import convert_bytes, get_file_size
from utils.processors import filter_file
@ -62,9 +63,9 @@ async def upload_file(
else:
openai_api_key = request.headers.get("Openai-Api-Key", None)
if openai_api_key is None:
brain_details = brain.get_brain_details()
brain_details = get_brain_details(brain_id)
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:
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 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.user_identity import UserIdentity
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.update_user_identity import (
UserIdentityUpdatableProperties,
@ -47,10 +48,10 @@ async def get_user_endpoint(
date = time.strftime("%Y%m%d")
max_requests_number = os.getenv("MAX_REQUESTS_NUMBER")
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:
defaul_brain_size = Brain(id=default_brain["id"]).brain_size
defaul_brain_size = Brain(id=default_brain.brain_id).brain_size
else:
defaul_brain_size = 0

View File

@ -1,7 +1,7 @@
import random
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):
@ -211,7 +211,6 @@ def test_set_as_default_brain_endpoint(client, api_key):
user_info = response.json()
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["id"] == brain_id
assert default_brain["default_brain"] is True
assert default_brain.brain_id == brain_id

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,
max_tokens INT,
temperature FLOAT,
openai_api_key TEXT
openai_api_key TEXT,
prompt_id UUID REFERENCES prompts(id)
);
-- Create brains X users table
CREATE TABLE IF NOT EXISTS brains_users (
brain_id UUID,
@ -210,7 +212,7 @@ CREATE TABLE IF NOT EXISTS migrations (
);
INSERT INTO migrations (name)
SELECT '20230701180101_add_prompts_table'
SELECT '20230802120700_add_prompt_id_to_brain'
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'
);