refactor: to modules (#1754)

# Description

Please include a summary of the changes and the related issue. Please
also include relevant motivation and context.

## Checklist before requesting a review

Please delete options that are not relevant.

- [ ] My code follows the style guidelines of this project
- [ ] I have performed a self-review of my code
- [ ] I have commented hard-to-understand areas
- [ ] I have ideally added tests that prove my fix is effective or that
my feature works
- [ ] New and existing unit tests pass locally with my changes
- [ ] Any dependent changes have been merged

## Screenshots (if appropriate):
This commit is contained in:
Zineb El Bachiri 2023-11-30 22:29:28 +01:00 committed by GitHub
parent 68da22aa1d
commit f48dab4a7d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
125 changed files with 1665 additions and 1439 deletions

View File

@ -1,6 +1,6 @@
from celery import shared_task
from models.brains import Brain
from models.settings import get_supabase_db
from modules.brain.service.brain_vector_service import BrainVectorService
from packages.embeddings.vectors import Neurons
from repository.files.upload_file import DocumentSerializable
@ -15,5 +15,5 @@ def create_embedding_for_document(brain_id, doc_with_metadata, file_sha1):
created_vector_id = created_vector[0] # pyright: ignore reportPrivateUsage=none
brain = Brain(id=brain_id) # pyright: ignore
brain.create_brain_vector(created_vector_id, file_sha1)
brain_vector_service = BrainVectorService(brain_id)
brain_vector_service.create_brain_vector(created_vector_id, file_sha1)

View File

@ -7,6 +7,7 @@ from celery.schedules import crontab
from fastapi import UploadFile
from models.files import File
from models.settings import get_supabase_client
from modules.brain.service.brain_service import BrainService
from modules.notification.dto.inputs import NotificationUpdatableProperties
from modules.notification.entity.notification import NotificationsStatusEnum
from modules.notification.service.notification_service import NotificationService
@ -14,13 +15,13 @@ from modules.onboarding.service.onboarding_service import OnboardingService
from packages.files.crawl.crawler import CrawlWebsite
from packages.files.parsers.github import process_github
from packages.files.processors import filter_file
from repository.brain.update_brain_last_update_time import update_brain_last_update_time
CELERY_BROKER_URL = os.getenv("CELERY_BROKER_URL", "")
CELEBRY_BROKER_QUEUE_NAME = os.getenv("CELEBRY_BROKER_QUEUE_NAME", "quivr")
onboardingService = OnboardingService()
notification_service = NotificationService()
brain_service = BrainService()
if CELERY_BROKER_URL.startswith("sqs"):
broker_transport_options = {
@ -100,7 +101,7 @@ def process_file_and_notify(
message=str(notification_message),
),
)
update_brain_last_update_time(brain_id)
brain_service.update_brain_last_update_time(brain_id)
return True
except Exception as e:

View File

@ -1,5 +1,3 @@
import os
from packages.utils.handle_request_validation_error import (
handle_request_validation_error,
)
@ -10,13 +8,12 @@ if __name__ == "__main__":
from dotenv import load_dotenv # type: ignore
load_dotenv()
import sentry_sdk
from fastapi import FastAPI, HTTPException
from fastapi.responses import JSONResponse
from logger import get_logger
from middlewares.cors import add_cors_middleware
from modules.misc.controller import misc_router
from routes.chat_routes import chat_router
from routes.misc_routes import misc_router
logger = get_logger(__name__)

View File

@ -1,4 +1,3 @@
from packages.utils import handle_request_validation_error
if __name__ == "__main__":
@ -11,8 +10,8 @@ from fastapi import FastAPI, HTTPException
from fastapi.responses import JSONResponse
from logger import get_logger
from middlewares.cors import add_cors_middleware
from modules.misc.controller import misc_router
from routes.crawl_routes import crawl_router
from routes.misc_routes import misc_router
logger = get_logger(__name__)
app = FastAPI()

View File

@ -5,18 +5,19 @@ from uuid import UUID
from fastapi import HTTPException
from logger import get_logger
from litellm import completion
from models.chats import ChatQuestion
from models.databases.supabase.chats import CreateChatHistory
from repository.brain.get_brain_by_id import get_brain_by_id
from repository.chat.get_chat_history import GetChatHistoryOutput, get_chat_history
from repository.chat.update_chat_history import update_chat_history
from repository.chat.update_message_by_id import update_message_by_id
from llm.qa_base import QABaseBrainPicking
from llm.utils.call_brain_api import call_brain_api
from llm.utils.get_api_brain_definition_as_json_schema import (
get_api_brain_definition_as_json_schema,
)
from models.chats import ChatQuestion
from models.databases.supabase.chats import CreateChatHistory
from modules.brain.service.brain_service import BrainService
from repository.chat.get_chat_history import GetChatHistoryOutput, get_chat_history
from repository.chat.update_chat_history import update_chat_history
from repository.chat.update_message_by_id import update_message_by_id
brain_service = BrainService()
logger = get_logger(__name__)
@ -141,7 +142,7 @@ class APIBrainQA(
status_code=400, detail="No brain id provided in the question"
)
brain = get_brain_by_id(question.brain_id)
brain = brain_service.get_brain_by_id(question.brain_id)
if not brain:
raise HTTPException(status_code=404, detail="Brain not found")

View File

@ -15,12 +15,14 @@ from langchain.prompts.chat import (
HumanMessagePromptTemplate,
SystemMessagePromptTemplate,
)
from llm.utils.get_prompt_to_use import get_prompt_to_use
from llm.utils.get_prompt_to_use_id import get_prompt_to_use_id
from logger import get_logger
from models import BrainSettings # Importing settings related to the 'brain'
from models.chats import ChatQuestion
from models.databases.supabase.chats import CreateChatHistory
from modules.brain.service.brain_service import BrainService
from pydantic import BaseModel
from repository.brain import get_brain_by_id
from repository.chat import (
GetChatHistoryOutput,
format_chat_history,
@ -31,15 +33,15 @@ from repository.chat import (
from supabase.client import Client, create_client
from vectorstore.supabase import CustomSupabaseVectorStore
from llm.utils.get_prompt_to_use import get_prompt_to_use
from llm.utils.get_prompt_to_use_id import get_prompt_to_use_id
from .prompts.CONDENSE_PROMPT import CONDENSE_QUESTION_PROMPT
logger = get_logger(__name__)
QUIVR_DEFAULT_PROMPT = "Your name is Quivr. You're a helpful assistant. If you don't know the answer, just say that you don't know, don't try to make up an answer."
brain_service = BrainService()
class QABaseBrainPicking(BaseModel):
"""
Main class for the Brain Picking functionality.
@ -243,7 +245,7 @@ class QABaseBrainPicking(BaseModel):
brain = None
if question.brain_id:
brain = get_brain_by_id(question.brain_id)
brain = brain_service.get_brain_by_id(question.brain_id)
return GetChatHistoryOutput(
**{
@ -319,7 +321,7 @@ class QABaseBrainPicking(BaseModel):
brain = None
if question.brain_id:
brain = get_brain_by_id(question.brain_id)
brain = brain_service.get_brain_by_id(question.brain_id)
streamed_chat_history = update_chat_history(
CreateChatHistory(

View File

@ -1,5 +1,5 @@
from fastapi import HTTPException
from models.ApiBrainDefinition import ApiBrainDefinitionSchema
from modules.brain.entity.api_brain_definition_entity import ApiBrainDefinitionSchema
def extract_api_brain_definition_values_from_llm_output(

View File

@ -1,4 +1,6 @@
from models.ApiBrainDefinition import ApiBrainDefinitionSchemaProperty
from modules.brain.entity.api_brain_definition_entity import (
ApiBrainDefinitionSchemaProperty,
)
def format_api_brain_property(property: ApiBrainDefinitionSchemaProperty):

View File

@ -1,9 +1,7 @@
from fastapi import HTTPException
from llm.utils.extract_api_definition import (
format_api_brain_property,
)
from llm.utils.extract_api_definition import format_api_brain_property
from llm.utils.sanitize_function_name import sanitize_function_name
from models.brain_entity import BrainEntity
from modules.brain.entity.brain_entity import BrainEntity
from repository.api_brain_definition.get_api_brain_definition import (
get_api_brain_definition,
)

View File

@ -1,7 +1,9 @@
from typing import Optional
from uuid import UUID
from repository.brain import get_brain_prompt_id
from modules.brain.service.brain_service import BrainService
brain_service = BrainService()
def get_prompt_to_use_id(
@ -11,5 +13,9 @@ def get_prompt_to_use_id(
return None
return (
prompt_id if prompt_id else get_brain_prompt_id(brain_id) if brain_id else None
prompt_id
if prompt_id
else brain_service.get_brain_prompt_id(brain_id)
if brain_id
else None
)

View File

@ -13,19 +13,19 @@ from fastapi import FastAPI, HTTPException
from fastapi.responses import JSONResponse
from logger import get_logger
from middlewares.cors import add_cors_middleware
from modules.knowledge.controller.knowledge_routes import knowledge_router
from modules.notification.controller.notification_routes import notification_router
from modules.onboarding.controller.onboarding_routes import onboarding_router
from modules.prompt.controller.prompt_routes import prompt_router
from modules.user.controller.user_controller import user_router
from routes.api_key_routes import api_key_router
from modules.api_key.controller import api_key_router
from modules.contact_support.controller import contact_router
from modules.knowledge.controller import knowledge_router
from modules.misc.controller import misc_router
from modules.notification.controller import notification_router
from modules.onboarding.controller import onboarding_router
from modules.prompt.controller import prompt_router
from modules.upload.controller import upload_router
from modules.user.controller import user_router
from routes.brain_routes import brain_router
from routes.chat_routes import chat_router
from routes.contact_routes import router as contact_router
from routes.crawl_routes import crawl_router
from routes.misc_routes import misc_router
from routes.subscription_routes import subscription_router
from routes.upload_routes import upload_router
from sentry_sdk.integrations.fastapi import FastApiIntegration
from sentry_sdk.integrations.starlette import StarletteIntegration

View File

@ -1,46 +0,0 @@
from datetime import datetime
from fastapi import HTTPException
from models.settings import get_supabase_db
from modules.user.entity.user_identity import UserIdentity
from pydantic import DateError
async def verify_api_key(
api_key: str,
) -> bool:
try:
# Use UTC time to avoid timezone issues
current_date = datetime.utcnow().date()
supabase_db = get_supabase_db()
result = supabase_db.get_active_api_key(api_key)
if result.data is not None and len(result.data) > 0:
api_key_creation_date = datetime.strptime(
result.data[0]["creation_time"], "%Y-%m-%dT%H:%M:%S"
).date()
if (api_key_creation_date.month == current_date.month) and (
api_key_creation_date.year == current_date.year
):
return True
return False
except DateError:
return False
async def get_user_from_api_key(
api_key: str,
) -> UserIdentity:
supabase_db = get_supabase_db()
user_id_data = supabase_db.get_user_id_by_api_key(api_key)
if not user_id_data.data:
raise HTTPException(status_code=400, detail="Invalid API key.")
user_id = user_id_data.data[0]["user_id"]
email = supabase_db.get_user_email(user_id)
return UserIdentity(email=email, id=user_id)

View File

@ -3,10 +3,12 @@ from typing import Optional
from fastapi import Depends, HTTPException, Request
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from middlewares.auth.api_key_handler import get_user_from_api_key, verify_api_key
from middlewares.auth.jwt_token_handler import decode_access_token, verify_token
from modules.api_key.service.api_key_service import ApiKeyService
from modules.user.entity.user_identity import UserIdentity
api_key_service = ApiKeyService()
class AuthBearer(HTTPBearer):
def __init__(self, auto_error: bool = True):
@ -41,10 +43,10 @@ class AuthBearer(HTTPBearer):
return self.get_test_user()
elif verify_token(token):
return decode_access_token(token)
elif await verify_api_key(
elif await api_key_service.verify_api_key(
token,
):
return await get_user_from_api_key(
return await api_key_service.get_user_from_api_key(
token,
)
else:

View File

@ -1,12 +1,10 @@
from .brain_entity import BrainEntity, MinimalBrainEntity
from .brains import Brain
from .brains_subscription_invitations import BrainSubscription
from .chat import Chat, ChatHistory
from .chats import ChatMessage, ChatQuestion
from .files import File
from .settings import (BrainRateLimiting, BrainSettings, ContactsSettings,
ResendSettings, get_documents_vector_store,
get_embeddings, get_supabase_client, get_supabase_db)
from .settings import (BrainRateLimiting, BrainSettings, ResendSettings,
get_documents_vector_store, get_embeddings,
get_supabase_client, get_supabase_db)
from .user_usage import UserUsage
# TODO uncomment the below import when start using SQLalchemy

View File

@ -1,85 +0,0 @@
from typing import Any, List, Optional
from uuid import UUID
from logger import get_logger
from models.databases.supabase.supabase import SupabaseDB
from models.settings import get_supabase_client, get_supabase_db
from packages.embeddings.vectors import get_unique_files_from_vector_ids
from pydantic import BaseModel
from supabase.client import Client
logger = get_logger(__name__)
class Brain(BaseModel):
id: Optional[UUID] = None
name: Optional[str] = "Default brain"
description: Optional[str] = "This is a description"
status: Optional[str] = "private"
model: Optional[str] = None
temperature: Optional[float] = 0.0
max_tokens: Optional[int] = 256
files: List[Any] = []
prompt_id: Optional[UUID] = None
class Config:
arbitrary_types_allowed = True
@property
def supabase_client(self) -> Client:
return get_supabase_client()
@property
def supabase_db(self) -> SupabaseDB:
return get_supabase_db()
@property
def brain_size(self):
self.get_unique_brain_files()
current_brain_size = sum(float(doc["size"]) for doc in self.files)
return current_brain_size
@classmethod
def create(cls, *args, **kwargs):
commons = {"supabase": get_supabase_client()}
return cls(
commons=commons, *args, **kwargs # pyright: ignore reportPrivateUsage=none
) # pyright: ignore reportPrivateUsage=none
def create_brain_vector(self, vector_id, file_sha1):
return self.supabase_db.create_brain_vector(self.id, vector_id, file_sha1) # type: ignore
def get_vector_ids_from_file_sha1(self, file_sha1: str):
return self.supabase_db.get_vector_ids_from_file_sha1(file_sha1)
def update_brain_with_file(self, file_sha1: str):
# not used
vector_ids = self.get_vector_ids_from_file_sha1(file_sha1)
for vector_id in vector_ids:
self.create_brain_vector(vector_id, file_sha1)
def get_unique_brain_files(self):
"""
Retrieve unique brain data (i.e. uploaded files and crawled websites).
"""
vector_ids = self.supabase_db.get_brain_vector_ids(self.id) # type: ignore
self.files = get_unique_files_from_vector_ids(vector_ids)
return self.files
def delete_file_from_brain(self, file_name: str):
file_name_with_brain_id = f"{self.id}/{file_name}"
self.supabase_client.storage.from_("quivr").remove([file_name_with_brain_id])
return self.supabase_db.delete_file_from_brain(self.id, file_name) # type: ignore
def get_all_knowledge_in_brain(self):
"""
Retrieve unique brain data (i.e. uploaded files and crawled websites).
"""
vector_ids = self.supabase_db.get_brain_vector_ids(self.id) # type: ignore
self.files = get_unique_files_from_vector_ids(vector_ids)
return self.files

View File

@ -2,68 +2,8 @@ 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) -> list[BrainEntity]:
pass
@abstractmethod
def get_brain_for_user(self, user_id: str):
pass
@abstractmethod
def delete_brain_user_by_id(self, user_id: str, brain_id: str):
pass
@abstractmethod
def delete_brain_vector(self, brain_id: str):
pass
@abstractmethod
def delete_brain_users(self, brain_id: str):
pass
@abstractmethod
def delete_brain(self, brain_id: str):
pass
@abstractmethod
def create_brain(self, brain: str):
pass
@abstractmethod
def create_brain_user(
self, user_id: UUID, brain_id: UUID, rights: str, default_brain: bool
):
pass
@abstractmethod
def create_brain_vector(self, brain_id: UUID, vector_id: UUID, file_sha1: str):
pass
@abstractmethod
def get_vector_ids_from_file_sha1(self, file_sha1: str):
pass
@abstractmethod
def get_brain_vector_ids(self, brain_id: UUID):
pass
@abstractmethod
def delete_file_from_brain(self, brain_id: UUID, file_name: str):
pass
@abstractmethod
def get_default_user_brain_id(self, user_id: UUID) -> UUID:
pass
@abstractmethod
def get_brain_by_id(self, brain_id: UUID):
pass
@abstractmethod
def create_user_daily_usage(self, user_id: UUID, user_email: str, date: datetime):
pass
@ -86,10 +26,6 @@ class Repository(ABC):
):
pass
@abstractmethod
def get_user_email(self, user_id: UUID):
pass
@abstractmethod
def set_file_vectors_ids(self, file_sha1: str):
pass
@ -118,22 +54,6 @@ class Repository(ABC):
):
pass
@abstractmethod
def create_api_key(self, new_key_id: UUID, new_api_key: str, user_id: UUID):
pass
@abstractmethod
def delete_api_key(self, key_id: UUID, user_id: UUID):
pass
@abstractmethod
def get_active_api_key(self, api_key: UUID):
pass
@abstractmethod
def get_user_id_by_api_key(self, api_key: UUID):
pass
@abstractmethod
def create_chat(self, new_chat):
pass

View File

@ -1,6 +1,4 @@
from models.databases.supabase.api_brain_definition import ApiBrainDefinitions
from models.databases.supabase.api_key_handler import ApiKeyHandler
from models.databases.supabase.brains import Brain
from models.databases.supabase.brains_subscription_invitations import \
BrainSubscription
from models.databases.supabase.chats import Chats

View File

@ -2,12 +2,12 @@ from enum import Enum
from typing import Optional
from uuid import UUID
from models.ApiBrainDefinition import (
from models.databases.repository import Repository
from modules.brain.entity.api_brain_definition_entity import (
ApiBrainDefinition,
ApiBrainDefinitionSchema,
ApiBrainDefinitionSecret,
)
from models.databases.repository import Repository
from pydantic import BaseModel, Extra

View File

@ -1,365 +0,0 @@
from typing import Optional
from uuid import UUID
from logger import get_logger
from models.ApiBrainDefinition import ApiBrainDefinition
from models.brain_entity import (
BrainEntity,
BrainType,
BrainUser,
MinimalBrainEntity,
PublicBrain,
)
from models.databases.repository import Repository
from models.databases.supabase.api_brain_definition import (
CreateApiBrainDefinition,
)
from pydantic import BaseModel, Extra
logger = get_logger(__name__)
class CreateBrainProperties(BaseModel, extra=Extra.forbid):
name: Optional[str] = "Default brain"
description: Optional[str] = "This is a description"
status: Optional[str] = "private"
model: Optional[str]
temperature: Optional[float] = 0.0
max_tokens: Optional[int] = 256
prompt_id: Optional[UUID] = None
brain_type: Optional[BrainType] = BrainType.DOC
brain_definition: Optional[CreateApiBrainDefinition]
brain_secrets_values: dict = {}
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]
status: Optional[str]
prompt_id: Optional[UUID]
brain_definition: Optional[ApiBrainDefinition]
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 BrainQuestionRequest(BaseModel):
question: str
class Brain(Repository):
def __init__(self, supabase_client):
self.db = supabase_client
def create_brain(self, brain: CreateBrainProperties):
response = (
self.db.table("brains").insert(
brain.dict(exclude={"brain_definition", "brain_secrets_values"})
)
).execute()
return BrainEntity(**response.data[0])
def get_user_brains(self, user_id) -> list[MinimalBrainEntity]:
response = (
self.db.from_("brains_users")
.select("id:brain_id, rights, brains (brain_id, name, status, brain_type)")
.filter("user_id", "eq", user_id)
.execute()
)
user_brains: list[MinimalBrainEntity] = []
for item in response.data:
user_brains.append(
MinimalBrainEntity(
id=item["brains"]["brain_id"],
name=item["brains"]["name"],
rights=item["rights"],
status=item["brains"]["status"],
brain_type=item["brains"]["brain_type"],
)
)
user_brains[-1].rights = item["rights"]
return user_brains
def get_public_brains(self) -> list[PublicBrain]:
response = (
self.db.from_("brains")
.select(
"id:brain_id, name, description, last_update, brain_type, brain_definition: api_brain_definition(*), number_of_subscribers:brains_users(count)"
)
.filter("status", "eq", "public")
.execute()
)
public_brains: list[PublicBrain] = []
for item in response.data:
item["number_of_subscribers"] = item["number_of_subscribers"][0]["count"]
if not item["brain_definition"]:
del item["brain_definition"]
else:
item["brain_definition"] = item["brain_definition"][0]
item["brain_definition"]["secrets"] = []
public_brains.append(PublicBrain(**item))
return public_brains
def update_brain_last_update_time(self, brain_id: UUID) -> None:
self.db.table("brains").update({"last_update": "now()"}).match(
{"brain_id": brain_id}
).execute()
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, status, name, brain_type)"
)
.filter("user_id", "eq", user_id)
.filter("brain_id", "eq", brain_id)
.execute()
)
if len(response.data) == 0:
return None
brain_data = response.data[0]
return MinimalBrainEntity(
id=brain_data["brains"]["id"],
name=brain_data["brains"]["name"],
rights=brain_data["rights"],
status=brain_data["brains"]["status"],
brain_type=brain_data["brains"]["brain_type"],
)
def get_brain_details(self, brain_id):
response = (
self.db.from_("brains")
.select("id:brain_id, name, *")
.filter("brain_id", "eq", brain_id)
.execute()
)
return response.data
def delete_brain_user_by_id(
self,
user_id: UUID,
brain_id: UUID,
):
results = (
self.db.table("brains_users")
.delete()
.match({"brain_id": str(brain_id), "user_id": str(user_id)})
.execute()
)
return results.data
def delete_brain_vector(self, brain_id: str):
results = (
self.db.table("brains_vectors")
.delete()
.match({"brain_id": brain_id})
.execute()
)
return results
def delete_brain_users(self, brain_id: str):
results = (
self.db.table("brains_users")
.delete()
.match({"brain_id": brain_id})
.execute()
)
return results
def delete_brain_subscribers(self, brain_id: UUID):
results = (
self.db.table("brains_users")
.delete()
.match({"brain_id": str(brain_id)})
.match({"rights": "Viewer"})
.execute()
).data
return results
def delete_brain(self, brain_id: str):
results = (
self.db.table("brains").delete().match({"brain_id": brain_id}).execute()
)
return results
def create_brain_user(self, user_id: UUID, brain_id, rights, default_brain: bool):
response = (
self.db.table("brains_users")
.insert(
{
"brain_id": str(brain_id),
"user_id": str(user_id),
"rights": rights,
"default_brain": default_brain,
}
)
.execute()
)
return response
def create_brain_vector(self, brain_id, vector_id, file_sha1):
response = (
self.db.table("brains_vectors")
.insert(
{
"brain_id": str(brain_id),
"vector_id": str(vector_id),
"file_sha1": file_sha1,
}
)
.execute()
)
return response.data
def get_vector_ids_from_file_sha1(self, file_sha1: str):
# move to vectors class
vectorsResponse = (
self.db.table("vectors")
.select("id")
.filter("file_sha1", "eq", file_sha1)
.execute()
)
return vectorsResponse.data
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):
"""
Retrieve unique brain data (i.e. uploaded files and crawled websites).
"""
response = (
self.db.from_("brains_vectors")
.select("vector_id")
.filter("brain_id", "eq", brain_id)
.execute()
)
vector_ids = [item["vector_id"] for item in response.data]
if len(vector_ids) == 0:
return []
return vector_ids
def delete_file_from_brain(self, brain_id, file_name: str):
# First, get the vector_ids associated with the file_name
file_vectors = (
self.db.table("vectors")
.select("id")
.filter("metadata->>file_name", "eq", file_name)
.execute()
)
file_vectors_ids = [item["id"] for item in file_vectors.data]
# remove current file vectors from brain vectors
self.db.table("brains_vectors").delete().filter(
"vector_id", "in", file_vectors_ids
).filter("brain_id", "eq", brain_id).execute()
vectors_used_by_another_brain = (
self.db.table("brains_vectors")
.select("vector_id")
.filter("vector_id", "in", file_vectors_ids)
.filter("brain_id", "neq", brain_id)
.execute()
)
vectors_used_by_another_brain_ids = [
item["vector_id"] for item in vectors_used_by_another_brain.data
]
vectors_no_longer_used_ids = [
id for id in file_vectors_ids if id not in vectors_used_by_another_brain_ids
]
self.db.table("vectors").delete().filter(
"id", "in", vectors_no_longer_used_ids
).execute()
return {"message": f"File {file_name} in brain {brain_id} has been deleted."}
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"))
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
if len(response) == 0:
return None
return BrainEntity(**response[0])
def get_brain_subscribers_count(self, brain_id: UUID) -> int:
response = (
self.db.from_("brains_users")
.select(
"count",
)
.filter("brain_id", "eq", str(brain_id))
.execute()
).data
if len(response) == 0:
raise ValueError(f"Brain with id {brain_id} does not exist.")
return response[0]["count"]
def get_brain_users(self, brain_id: UUID) -> list[BrainUser]:
response = (
self.db.table("brains_users")
.select("id:brain_id, *")
.filter("brain_id", "eq", str(brain_id))
.execute()
)
return [BrainUser(**item) for item in response.data]

View File

@ -1,8 +1,6 @@
from logger import get_logger
from models.databases.supabase import (
ApiBrainDefinitions,
ApiKeyHandler,
Brain,
BrainSubscription,
Chats,
File,
@ -14,22 +12,18 @@ logger = get_logger(__name__)
class SupabaseDB(
Brain,
UserUsage,
File,
BrainSubscription,
ApiKeyHandler,
Chats,
Vector,
ApiBrainDefinitions,
):
def __init__(self, supabase_client):
self.db = supabase_client
Brain.__init__(self, supabase_client)
UserUsage.__init__(self, supabase_client)
File.__init__(self, supabase_client)
BrainSubscription.__init__(self, supabase_client)
ApiKeyHandler.__init__(self, supabase_client)
Chats.__init__(self, supabase_client)
Vector.__init__(self, supabase_client)
ApiBrainDefinitions.__init__(self, supabase_client)

View File

@ -141,18 +141,3 @@ class UserUsage(Repository):
)
return response
def get_user_email(self, user_id):
"""
Fetch the user email from the database
"""
response = (
self.db.from_("user_daily_usage")
.select("email")
.filter("user_id", "eq", user_id)
.execute()
).data
if response and len(response) > 0:
return response[0]["email"]
return None

View File

@ -6,9 +6,9 @@ from uuid import UUID
from fastapi import UploadFile
from langchain.text_splitter import RecursiveCharacterTextSplitter
from logger import get_logger
from models.brains import Brain
from models.databases.supabase.supabase import SupabaseDB
from models.settings import get_supabase_db
from modules.brain.service.brain_vector_service import BrainVectorService
from packages.files.file import compute_sha1_from_file
from pydantic import BaseModel
@ -130,11 +130,13 @@ class File(BaseModel):
"""
return self.file.size < 1 # pyright: ignore reportPrivateUsage=none
def link_file_to_brain(self, brain: Brain):
def link_file_to_brain(self, brain_id):
self.set_file_vectors_ids()
if self.vectors_ids is None:
return
brain_vector_service = BrainVectorService(brain_id)
for vector_id in self.vectors_ids: # pyright: ignore reportPrivateUsage=none
brain.create_brain_vector(vector_id["id"], self.file_sha1)
brain_vector_service.create_brain_vector(vector_id["id"], self.file_sha1)

View File

@ -22,11 +22,6 @@ class BrainSettings(BaseSettings):
ollama_api_base_url: str = None
class ContactsSettings(BaseSettings):
resend_contact_sales_from: str = "null"
resend_contact_sales_to: str = "null"
class ResendSettings(BaseSettings):
resend_api_key: str = "null"

View File

@ -0,0 +1 @@
from .api_key_routes import api_key_router

View File

@ -6,25 +6,18 @@ from asyncpg.exceptions import UniqueViolationError
from fastapi import APIRouter, Depends
from logger import get_logger
from middlewares.auth import AuthBearer, get_current_user
from models import get_supabase_db
from modules.api_key.dto.outputs import ApiKeyInfo
from modules.api_key.entity.api_key import ApiKey
from modules.api_key.repository.api_keys import ApiKeys
from modules.user.entity.user_identity import UserIdentity
from pydantic import BaseModel
logger = get_logger(__name__)
class ApiKeyInfo(BaseModel):
key_id: str
creation_time: str
class ApiKey(BaseModel):
api_key: str
key_id: str
api_key_router = APIRouter()
api_keys_repository = ApiKeys()
@api_key_router.post(
"/api-key",
@ -46,12 +39,11 @@ async def create_api_key(current_user: UserIdentity = Depends(get_current_user))
new_key_id = uuid4()
new_api_key = token_hex(16)
api_key_inserted = False
supabase_db = get_supabase_db()
while not api_key_inserted:
try:
# Attempt to insert new API key into database
supabase_db.create_api_key(new_key_id, new_api_key, current_user.id)
api_keys_repository.create_api_key(new_key_id, new_api_key, current_user.id)
api_key_inserted = True
except UniqueViolationError:
@ -80,8 +72,7 @@ async def delete_api_key(
as inactive in the database.
"""
supabase_db = get_supabase_db()
supabase_db.delete_api_key(key_id, current_user.id)
api_keys_repository.delete_api_key(key_id, current_user.id)
return {"message": "API key deleted."}
@ -102,6 +93,5 @@ async def get_api_keys(current_user: UserIdentity = Depends(get_current_user)):
This endpoint retrieves all the active API keys associated with the current user. It returns a list of API key objects
containing the key ID and creation time for each API key.
"""
supabase_db = get_supabase_db()
response = supabase_db.get_user_api_keys(current_user.id)
response = api_keys_repository.get_user_api_keys(current_user.id)
return response.data

View File

@ -0,0 +1 @@
from .outputs import ApiKeyInfo

View File

@ -0,0 +1,6 @@
from pydantic import BaseModel
class ApiKeyInfo(BaseModel):
key_id: str
creation_time: str

View File

@ -0,0 +1,6 @@
from pydantic import BaseModel
class ApiKey(BaseModel):
api_key: str
key_id: str

View File

@ -0,0 +1,27 @@
from abc import ABC, abstractmethod
from typing import List
from uuid import UUID
from modules.api_key.entity.api_key import ApiKey
class ApiKeysInterface(ABC):
@abstractmethod
def create_api_key(self, new_key_id: UUID, new_api_key: str, user_id: UUID):
pass
@abstractmethod
def delete_api_key(self, key_id: UUID, user_id: UUID):
pass
@abstractmethod
def get_active_api_key(self, api_key: UUID):
pass
@abstractmethod
def get_user_id_by_api_key(self, api_key: UUID):
pass
@abstractmethod
def get_user_api_keys(self, user_id: UUID) -> List[ApiKey]:
pass

View File

@ -1,11 +1,13 @@
from datetime import datetime
from uuid import UUID
from models.databases.repository import Repository
from models.settings import get_supabase_client
from modules.api_key.repository.api_key_interface import ApiKeysInterface
class ApiKeyHandler(Repository):
def __init__(self, supabase_client):
class ApiKeys(ApiKeysInterface):
def __init__(self):
supabase_client = get_supabase_client()
self.db = supabase_client # type: ignore
def create_api_key(self, new_key_id, new_api_key, user_id):
@ -60,7 +62,7 @@ class ApiKeyHandler(Repository):
)
return response
def get_user_api_keys(self, user_id: UUID):
def get_user_api_keys(self, user_id):
response = (
self.db.table("api_keys")
.select("key_id, creation_time")

View File

@ -0,0 +1,62 @@
from datetime import datetime
from fastapi import HTTPException
from logger import get_logger
from modules.api_key.repository.api_key_interface import ApiKeysInterface
from modules.api_key.repository.api_keys import ApiKeys
from modules.user.entity.user_identity import UserIdentity
from modules.user.service.user_service import UserService
from pydantic import DateError
logger = get_logger(__name__)
user_service = UserService()
class ApiKeyService:
repository: ApiKeysInterface
def __init__(self):
self.repository = ApiKeys()
async def verify_api_key(
self,
api_key: str,
) -> bool:
try:
# Use UTC time to avoid timezone issues
current_date = datetime.utcnow().date()
result = self.repository.get_active_api_key(api_key)
if result.data is not None and len(result.data) > 0:
api_key_creation_date = datetime.strptime(
result.data[0]["creation_time"], "%Y-%m-%dT%H:%M:%S"
).date()
if (api_key_creation_date.month == current_date.month) and (
api_key_creation_date.year == current_date.year
):
return True
return False
except DateError:
return False
async def get_user_from_api_key(
self,
api_key: str,
) -> UserIdentity:
user_id_data = self.repository.get_user_id_by_api_key(api_key)
if not user_id_data.data:
raise HTTPException(status_code=400, detail="Invalid API key.")
user_id = user_id_data.data[0]["user_id"]
# TODO: directly UserService instead
email = user_service.get_user_email_by_user_id(user_id)
if email is None:
raise HTTPException(status_code=400, detail="Invalid API key.")
return UserIdentity(email=email, id=user_id)

View File

View File

@ -0,0 +1,50 @@
from typing import Optional
from uuid import UUID
from logger import get_logger
from models.databases.supabase.api_brain_definition import CreateApiBrainDefinition
from modules.brain.entity.api_brain_definition_entity import ApiBrainDefinition
from modules.brain.entity.brain_entity import BrainType
from pydantic import BaseModel, Extra
logger = get_logger(__name__)
class CreateBrainProperties(BaseModel, extra=Extra.forbid):
name: Optional[str] = "Default brain"
description: Optional[str] = "This is a description"
status: Optional[str] = "private"
model: Optional[str]
temperature: Optional[float] = 0.0
max_tokens: Optional[int] = 256
prompt_id: Optional[UUID] = None
brain_type: Optional[BrainType] = BrainType.DOC
brain_definition: Optional[CreateApiBrainDefinition]
brain_secrets_values: dict = {}
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]
status: Optional[str]
prompt_id: Optional[UUID]
brain_definition: Optional[ApiBrainDefinition]
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 BrainQuestionRequest(BaseModel):
question: str

View File

View File

@ -2,10 +2,8 @@ from enum import Enum
from typing import Optional
from uuid import UUID
from modules.brain.entity.api_brain_definition_entity import ApiBrainDefinition
from pydantic import BaseModel
from routes.authorizations.types import RoleEnum
from models.ApiBrainDefinition import ApiBrainDefinition
class BrainType(str, Enum):
@ -38,14 +36,6 @@ class BrainEntity(BaseModel):
return data
class MinimalBrainEntity(BaseModel):
id: UUID
name: str
rights: RoleEnum
status: str
brain_type: BrainType
class PublicBrain(BaseModel):
id: UUID
name: str
@ -56,8 +46,22 @@ class PublicBrain(BaseModel):
brain_definition: Optional[ApiBrainDefinition]
class RoleEnum(str, Enum):
Viewer = "Viewer"
Editor = "Editor"
Owner = "Owner"
class BrainUser(BaseModel):
id: UUID
user_id: UUID
rights: RoleEnum
default_brain: bool = False
class MinimalUserBrainEntity(BaseModel):
id: UUID
name: str
rights: RoleEnum
status: str
brain_type: BrainType

View File

@ -0,0 +1,111 @@
from uuid import UUID
from logger import get_logger
from models.settings import get_supabase_client
from modules.brain.dto.inputs import BrainUpdatableProperties
from modules.brain.entity.brain_entity import BrainEntity, PublicBrain
from modules.brain.repository.interfaces.brains_interface import BrainsInterface
from repository.external_api_secret.utils import build_secret_unique_name
logger = get_logger(__name__)
class Brains(BrainsInterface):
def __init__(self):
supabase_client = get_supabase_client()
self.db = supabase_client
def create_brain(self, brain):
response = (
self.db.table("brains").insert(
brain.dict(exclude={"brain_definition", "brain_secrets_values"})
)
).execute()
return BrainEntity(**response.data[0])
def get_public_brains(self):
response = (
self.db.from_("brains")
.select(
"id:brain_id, name, description, last_update, brain_type, brain_definition: api_brain_definition(*), number_of_subscribers:brains_users(count)"
)
.filter("status", "eq", "public")
.execute()
)
public_brains: list[PublicBrain] = []
for item in response.data:
item["number_of_subscribers"] = item["number_of_subscribers"][0]["count"]
if not item["brain_definition"]:
del item["brain_definition"]
else:
item["brain_definition"] = item["brain_definition"][0]
item["brain_definition"]["secrets"] = []
public_brains.append(PublicBrain(**item))
return public_brains
def update_brain_last_update_time(self, brain_id):
self.db.table("brains").update({"last_update": "now()"}).match(
{"brain_id": brain_id}
).execute()
def get_brain_details(self, brain_id):
response = (
self.db.table("brains")
.select("*")
.filter("brain_id", "eq", str(brain_id))
.execute()
)
if response.data == []:
return None
return BrainEntity(**response.data[0])
def delete_brain(self, brain_id: str):
results = (
self.db.table("brains").delete().match({"brain_id": brain_id}).execute()
)
return results
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_by_id(self, brain_id: UUID) -> BrainEntity | None:
# TODO: merge this method with get_brain_details
response = (
self.db.from_("brains")
.select("id:brain_id, name, *")
.filter("brain_id", "eq", brain_id)
.execute()
).data
if len(response) == 0:
return None
return BrainEntity(**response[0])
def delete_secret(self, user_id: UUID, brain_id: UUID, secret_name: str) -> bool:
response = self.db.rpc(
"delete_secret",
{
"secret_name": build_secret_unique_name(
user_id=user_id, brain_id=brain_id, secret_name=secret_name
),
},
).execute()
return response.data

View File

@ -0,0 +1,159 @@
from uuid import UUID
from logger import get_logger
from models.settings import get_supabase_client
from modules.brain.entity.brain_entity import BrainUser, MinimalUserBrainEntity
from modules.brain.repository.interfaces.brains_users_interface import (
BrainsUsersInterface,
)
logger = get_logger(__name__)
class BrainsUsers(BrainsUsersInterface):
def __init__(self):
supabase_client = get_supabase_client()
self.db = supabase_client
def get_user_brains(self, user_id) -> list[MinimalUserBrainEntity]:
response = (
self.db.from_("brains_users")
.select("id:brain_id, rights, brains (brain_id, name, status, brain_type)")
.filter("user_id", "eq", user_id)
.execute()
)
user_brains: list[MinimalUserBrainEntity] = []
for item in response.data:
user_brains.append(
MinimalUserBrainEntity(
id=item["brains"]["brain_id"],
name=item["brains"]["name"],
rights=item["rights"],
status=item["brains"]["status"],
brain_type=item["brains"]["brain_type"],
)
)
user_brains[-1].rights = item["rights"]
return user_brains
def get_brain_for_user(self, user_id, brain_id):
response = (
self.db.from_("brains_users")
.select(
"id:brain_id, rights, brains (id: brain_id, status, name, brain_type)"
)
.filter("user_id", "eq", user_id)
.filter("brain_id", "eq", brain_id)
.execute()
)
if len(response.data) == 0:
return None
brain_data = response.data[0]
return MinimalUserBrainEntity(
id=brain_data["brains"]["id"],
name=brain_data["brains"]["name"],
rights=brain_data["rights"],
status=brain_data["brains"]["status"],
brain_type=brain_data["brains"]["brain_type"],
)
def delete_brain_user_by_id(
self,
user_id: UUID,
brain_id: UUID,
):
results = (
self.db.table("brains_users")
.delete()
.match({"brain_id": str(brain_id), "user_id": str(user_id)})
.execute()
)
return results.data
def delete_brain_users(self, brain_id: str):
results = (
self.db.table("brains_users")
.delete()
.match({"brain_id": brain_id})
.execute()
)
return results
def create_brain_user(self, user_id: UUID, brain_id, rights, default_brain: bool):
response = (
self.db.table("brains_users")
.insert(
{
"brain_id": str(brain_id),
"user_id": str(user_id),
"rights": rights,
"default_brain": default_brain,
}
)
.execute()
)
return response
def get_user_default_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"))
def get_brain_users(self, brain_id: UUID) -> list[BrainUser]:
response = (
self.db.table("brains_users")
.select("id:brain_id, *")
.filter("brain_id", "eq", str(brain_id))
.execute()
)
return [BrainUser(**item) for item in response.data]
def delete_brain_subscribers(self, brain_id: UUID):
results = (
self.db.table("brains_users")
.delete()
.match({"brain_id": str(brain_id)})
.match({"rights": "Viewer"})
.execute()
).data
return results
def get_brain_subscribers_count(self, brain_id: UUID) -> int:
response = (
self.db.from_("brains_users")
.select(
"count",
)
.filter("brain_id", "eq", str(brain_id))
.execute()
).data
if len(response) == 0:
raise ValueError(f"Brain with id {brain_id} does not exist.")
return response[0]["count"]
def update_brain_user_default_status(
self, user_id: UUID, brain_id: UUID, default_brain: bool
):
self.db.table("brains_users").update({"default_brain": default_brain}).match(
{"brain_id": brain_id, "user_id": user_id}
).execute()
def update_brain_user_rights(
self, user_id: UUID, brain_id: UUID, rights: str
) -> None:
self.db.table("brains_users").update({"rights": rights}).match(
{"brain_id": brain_id, "user_id": user_id}
).execute()

View File

@ -0,0 +1,104 @@
from logger import get_logger
from models.settings import get_supabase_client
from modules.brain.repository.interfaces.brains_vectors_interface import (
BrainsVectorsInterface,
)
logger = get_logger(__name__)
class BrainsVectors(BrainsVectorsInterface):
def __init__(self):
supabase_client = get_supabase_client()
self.db = supabase_client
def create_brain_vector(self, brain_id, vector_id, file_sha1):
response = (
self.db.table("brains_vectors")
.insert(
{
"brain_id": str(brain_id),
"vector_id": str(vector_id),
"file_sha1": file_sha1,
}
)
.execute()
)
return response.data
def get_vector_ids_from_file_sha1(self, file_sha1: str):
# move to vectors class
vectorsResponse = (
self.db.table("vectors")
.select("id")
.filter("file_sha1", "eq", file_sha1)
.execute()
)
return vectorsResponse.data
def get_brain_vector_ids(self, brain_id):
"""
Retrieve unique brain data (i.e. uploaded files and crawled websites).
"""
response = (
self.db.from_("brains_vectors")
.select("vector_id")
.filter("brain_id", "eq", brain_id)
.execute()
)
vector_ids = [item["vector_id"] for item in response.data]
if len(vector_ids) == 0:
return []
return vector_ids
def delete_file_from_brain(self, brain_id, file_name: str):
# First, get the vector_ids associated with the file_name
file_vectors = (
self.db.table("vectors")
.select("id")
.filter("metadata->>file_name", "eq", file_name)
.execute()
)
file_vectors_ids = [item["id"] for item in file_vectors.data]
# remove current file vectors from brain vectors
self.db.table("brains_vectors").delete().filter(
"vector_id", "in", f"({','.join(map(str, file_vectors_ids))})"
).filter("brain_id", "eq", brain_id).execute()
vectors_used_by_another_brain = (
self.db.table("brains_vectors")
.select("vector_id")
.filter("vector_id", "in", f"({','.join(map(str, file_vectors_ids))})")
.filter("brain_id", "neq", brain_id)
.execute()
)
vectors_used_by_another_brain_ids = [
item["vector_id"] for item in vectors_used_by_another_brain.data
]
vectors_no_longer_used_ids = [
id for id in file_vectors_ids if id not in vectors_used_by_another_brain_ids
]
self.db.table("vectors").delete().filter(
"id", "in", f"({','.join(map(str, vectors_no_longer_used_ids))})"
).execute()
return {"message": f"File {file_name} in brain {brain_id} has been deleted."}
def delete_brain_vector(self, brain_id: str):
results = (
self.db.table("brains_vectors")
.delete()
.match({"brain_id": brain_id})
.execute()
)
return results

View File

@ -0,0 +1,65 @@
from abc import ABC, abstractmethod
from uuid import UUID
from modules.brain.dto.inputs import BrainUpdatableProperties, CreateBrainProperties
from modules.brain.entity.brain_entity import BrainEntity, PublicBrain
class BrainsInterface(ABC):
@abstractmethod
def create_brain(self, brain: CreateBrainProperties) -> BrainEntity:
"""
Create a brain in the brains table
"""
pass
@abstractmethod
def get_public_brains(self) -> list[PublicBrain]:
"""
Get all public brains
"""
pass
@abstractmethod
def get_brain_details(self, brain_id: UUID) -> BrainEntity | None:
"""
Get all public brains
"""
pass
@abstractmethod
def update_brain_last_update_time(self, brain_id: UUID) -> None:
"""
Update the last update time of the brain
"""
pass
@abstractmethod
def delete_brain(self, brain_id: UUID):
"""
Delete a brain
"""
pass
@abstractmethod
def update_brain_by_id(
self, brain_id: UUID, brain: BrainUpdatableProperties
) -> BrainEntity | None:
"""
Update a brain by id
"""
pass
@abstractmethod
def get_brain_by_id(self, brain_id: UUID) -> BrainEntity | None:
"""
Get a brain by id
"""
pass
@abstractmethod
def delete_secret(self, user_id: UUID, brain_id: UUID, secret_name: str) -> bool:
"""
Delete a secret from a brain
"""
pass

View File

@ -0,0 +1,92 @@
from abc import ABC, abstractmethod
from typing import List
from uuid import UUID
from modules.brain.entity.brain_entity import BrainUser, MinimalUserBrainEntity
class BrainsUsersInterface(ABC):
@abstractmethod
def get_user_brains(self, user_id) -> list[MinimalUserBrainEntity]:
"""
Create a brain in the brains table
"""
pass
@abstractmethod
def get_brain_for_user(self, user_id, brain_id) -> MinimalUserBrainEntity | None:
"""
Get a brain for a user
"""
pass
@abstractmethod
def delete_brain_user_by_id(
self,
user_id: UUID,
brain_id: UUID,
):
"""
Delete a user in a brain
"""
pass
@abstractmethod
def delete_brain_users(self, brain_id: str):
"""
Delete all users for a brain
"""
pass
@abstractmethod
def create_brain_user(self, user_id: UUID, brain_id, rights, default_brain: bool):
"""
Create a brain user
"""
pass
@abstractmethod
def get_user_default_brain_id(self, user_id: UUID) -> UUID | None:
"""
Get the default brain id for a user
"""
pass
@abstractmethod
def get_brain_users(self, brain_id: UUID) -> List[BrainUser]:
"""
Get all users for a brain
"""
pass
@abstractmethod
def delete_brain_subscribers(self, brain_id: UUID):
"""
Delete all subscribers for a brain with Viewer rights
"""
pass
@abstractmethod
def get_brain_subscribers_count(self, brain_id: UUID) -> int:
"""
Get the number of subscribers for a brain
"""
pass
@abstractmethod
def update_brain_user_default_status(
self, user_id: UUID, brain_id: UUID, default_brain: bool
):
"""
Update the default brain status for a user
"""
pass
@abstractmethod
def update_brain_user_rights(
self, brain_id: UUID, user_id: UUID, rights: str
) -> BrainUser:
"""
Update the rights for a user in a brain
"""
pass

View File

@ -0,0 +1,41 @@
from abc import ABC, abstractmethod
from typing import List
from uuid import UUID
# TODO: Replace BrainsVectors with KnowledgeVectors interface instead
class BrainsVectorsInterface(ABC):
@abstractmethod
def create_brain_vector(self, brain_id, vector_id, file_sha1):
"""
Create a brain vector
"""
pass
@abstractmethod
def get_vector_ids_from_file_sha1(self, file_sha1: str):
"""
Get vector ids from file sha1
"""
pass
@abstractmethod
def get_brain_vector_ids(self, brain_id) -> List[UUID]:
"""
Get brain vector ids
"""
pass
@abstractmethod
def delete_file_from_brain(self, brain_id, file_name: str):
"""
Delete file from brain
"""
pass
@abstractmethod
def delete_brain_vector(self, brain_id: str):
"""
Delete brain vector
"""
pass

View File

@ -0,0 +1,41 @@
from abc import ABC, abstractmethod
from typing import List
from uuid import UUID
# TODO: Replace BrainsVectors with KnowledgeVectors interface instead
class BrainsVectorsInterface(ABC):
@abstractmethod
def create_brain_vector(self, brain_id, vector_id, file_sha1):
"""
Create a brain vector
"""
pass
@abstractmethod
def get_vector_ids_from_file_sha1(self, file_sha1: str):
"""
Get vector ids from file sha1
"""
pass
@abstractmethod
def get_brain_vector_ids(self, brain_id) -> List[UUID]:
"""
Get brain vector ids
"""
pass
@abstractmethod
def delete_file_from_brain(self, brain_id, file_name: str):
"""
Delete file from brain
"""
pass
@abstractmethod
def delete_brain_vector(self, brain_id: str):
"""
Delete brain vector
"""
pass

View File

@ -3,10 +3,13 @@ from uuid import UUID
from fastapi import Depends, HTTPException, status
from middlewares.auth.auth_bearer import get_current_user
from modules.brain.entity.brain_entity import RoleEnum
from modules.brain.service.brain_service import BrainService
from modules.brain.service.brain_user_service import BrainUserService
from modules.user.entity.user_identity import UserIdentity
from repository.brain import get_brain_for_user
from repository.brain.get_brain_details import get_brain_details
from routes.authorizations.types import RoleEnum
brain_user_service = BrainUserService()
brain_service = BrainService()
def has_brain_authorization(
@ -44,7 +47,7 @@ def validate_brain_authorization(
return: None
"""
brain = get_brain_details(brain_id)
brain = brain_service.get_brain_details(brain_id)
if brain and brain.status == "public":
return
@ -55,7 +58,7 @@ def validate_brain_authorization(
detail="Missing required role",
)
user_brain = get_brain_for_user(user_id, brain_id)
user_brain = brain_user_service.get_brain_for_user(user_id, brain_id)
if user_brain is None:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,

View File

@ -0,0 +1,230 @@
from typing import Optional
from uuid import UUID
from fastapi import HTTPException
from modules.brain.dto.inputs import BrainUpdatableProperties, CreateBrainProperties
from modules.brain.entity.brain_entity import BrainEntity, BrainType, PublicBrain
from modules.brain.repository.brains import Brains
from modules.brain.repository.brains_users import BrainsUsers
from modules.brain.repository.brains_vectors import BrainsVectors
from modules.brain.repository.interfaces.brains_interface import BrainsInterface
from modules.brain.repository.interfaces.brains_users_interface import (
BrainsUsersInterface,
)
from modules.brain.repository.interfaces.brains_vectors_interface import (
BrainsVectorsInterface,
)
from modules.knowledge.service.knowledge_service import KnowledgeService
from repository.api_brain_definition.add_api_brain_definition import (
add_api_brain_definition,
)
from repository.api_brain_definition.delete_api_brain_definition import (
delete_api_brain_definition,
)
from repository.api_brain_definition.get_api_brain_definition import (
get_api_brain_definition,
)
from repository.api_brain_definition.update_api_brain_definition import (
update_api_brain_definition,
)
from repository.external_api_secret.create_secret import create_secret
knowledge_service = KnowledgeService()
class BrainService:
brain_repository: BrainsInterface
brain_user_repository: BrainsUsersInterface
brain_vector_repository: BrainsVectorsInterface
def __init__(self):
self.brain_repository = Brains()
self.brain_user_repository = BrainsUsers()
self.brain_vector = BrainsVectors()
def get_brain_by_id(self, brain_id: UUID):
return self.brain_repository.get_brain_by_id(brain_id)
def create_brain(
self,
user_id: UUID,
brain: Optional[CreateBrainProperties],
) -> BrainEntity:
if brain == None:
brain = CreateBrainProperties() # type: ignore model and brain_definition
if brain.brain_type == BrainType.API:
if brain.brain_definition is None:
raise HTTPException(
status_code=404, detail="Brain definition not found"
)
if brain.brain_definition.url is None:
raise HTTPException(status_code=404, detail="Brain url not found")
if brain.brain_definition.method is None:
raise HTTPException(status_code=404, detail="Brain method not found")
created_brain = self.brain_repository.create_brain(brain)
if brain.brain_type == BrainType.API and brain.brain_definition is not None:
add_api_brain_definition(
brain_id=created_brain.brain_id,
api_brain_definition=brain.brain_definition,
)
secrets_values = brain.brain_secrets_values
for secret_name in secrets_values:
create_secret(
user_id=user_id,
brain_id=created_brain.brain_id,
secret_name=secret_name,
secret_value=secrets_values[secret_name],
)
return created_brain
def delete_brain_secrets_values(self, brain_id: UUID) -> None:
brain_definition = get_api_brain_definition(brain_id=brain_id)
if brain_definition is None:
raise HTTPException(status_code=404, detail="Brain definition not found.")
secrets = brain_definition.secrets
if len(secrets) > 0:
brain_users = self.brain_user_repository.get_brain_users(brain_id=brain_id)
for user in brain_users:
for secret in secrets:
self.brain_repository.delete_secret(
user_id=user.user_id,
brain_id=brain_id,
secret_name=secret.name,
)
def delete_brain(self, brain_id: UUID) -> dict[str, str]:
brain_to_delete = self.get_brain_by_id(brain_id=brain_id)
if brain_to_delete is None:
raise HTTPException(status_code=404, detail="Brain not found.")
if brain_to_delete.brain_type == BrainType.API:
self.delete_brain_secrets_values(
brain_id=brain_id,
)
delete_api_brain_definition(brain_id=brain_id)
else:
knowledge_service.remove_brain_all_knowledge(brain_id)
self.brain_vector.delete_brain_vector(str(brain_id))
self.brain_user_repository.delete_brain_users(str(brain_id))
self.brain_repository.delete_brain(str(brain_id)) # type: ignore
return {"message": "Brain deleted."}
def get_brain_prompt_id(self, brain_id: UUID) -> UUID | None:
brain = self.get_brain_by_id(brain_id)
prompt_id = brain.prompt_id if brain else None
return prompt_id
def update_brain_by_id(
self, brain_id: UUID, brain_new_values: BrainUpdatableProperties
) -> BrainEntity:
"""Update a prompt by id"""
existing_brain = self.brain_repository.get_brain_by_id(brain_id)
if existing_brain is None:
raise HTTPException(
status_code=404,
detail=f"Brain with id {brain_id} not found",
)
brain_update_answer = self.brain_repository.update_brain_by_id(
brain_id,
brain=BrainUpdatableProperties(
**brain_new_values.dict(exclude={"brain_definition"})
),
)
if brain_update_answer is None:
raise HTTPException(
status_code=404,
detail=f"Brain with id {brain_id} not found",
)
if (
brain_update_answer.brain_type == BrainType.API
and brain_new_values.brain_definition
):
existing_brain_secrets_definition = (
existing_brain.brain_definition.secrets
if existing_brain.brain_definition
else None
)
brain_new_values_secrets_definition = (
brain_new_values.brain_definition.secrets
if brain_new_values.brain_definition
else None
)
should_remove_existing_secrets_values = (
existing_brain_secrets_definition
and brain_new_values_secrets_definition
and existing_brain_secrets_definition
!= brain_new_values_secrets_definition
)
if should_remove_existing_secrets_values:
self.delete_brain_secrets_values(brain_id=brain_id)
update_api_brain_definition(
brain_id,
api_brain_definition=brain_new_values.brain_definition,
)
if brain_update_answer is None:
raise HTTPException(
status_code=404,
detail=f"Brain with id {brain_id} not found",
)
self.brain_repository.update_brain_last_update_time(brain_id)
return brain_update_answer
def update_brain_last_update_time(self, brain_id: UUID):
self.brain_repository.update_brain_last_update_time(brain_id)
def get_brain_details(self, brain_id: UUID) -> BrainEntity | None:
brain = self.brain_repository.get_brain_details(brain_id)
# id ?
if brain == None:
return None
if brain.brain_type == BrainType.API:
brain_definition = get_api_brain_definition(brain_id)
brain.brain_definition = brain_definition
return brain
def get_public_brains(self) -> list[PublicBrain]:
return self.brain_repository.get_public_brains()
def update_secret_value(
self,
user_id: UUID,
brain_id: UUID,
secret_name: str,
secret_value: str,
) -> None:
"""Update an existing secret."""
self.brain_repository.delete_secret(
user_id=user_id,
brain_id=brain_id,
secret_name=secret_name,
)
create_secret(
user_id=user_id,
brain_id=brain_id,
secret_name=secret_name,
secret_value=secret_value,
)

View File

@ -0,0 +1,134 @@
from typing import List
from uuid import UUID
from fastapi import HTTPException
from logger import get_logger
from modules.brain.entity.brain_entity import (
BrainEntity,
BrainType,
BrainUser,
MinimalUserBrainEntity,
RoleEnum,
)
from modules.brain.repository.brains import Brains
from modules.brain.repository.brains_users import BrainsUsers
from modules.brain.repository.interfaces.brains_interface import BrainsInterface
from modules.brain.repository.interfaces.brains_users_interface import (
BrainsUsersInterface,
)
from modules.brain.service.brain_service import BrainService
from modules.user.entity.user_identity import UserIdentity
from repository.api_brain_definition.get_api_brain_definition import (
get_api_brain_definition,
)
logger = get_logger(__name__)
brain_service = BrainService()
class BrainUserService:
brain_repository: BrainsInterface
brain_user_repository: BrainsUsersInterface
def __init__(self):
self.brain_repository = Brains()
self.brain_user_repository = BrainsUsers()
def get_user_default_brain(self, user_id: UUID) -> BrainEntity | None:
brain_id = self.brain_user_repository.get_user_default_brain_id(user_id)
logger.info(f"Default brain response: {brain_id}")
if brain_id is None:
return None
logger.info(f"Default brain id: {brain_id}")
return brain_service.get_brain_by_id(brain_id)
def delete_brain_user(self, user_id: UUID, brain_id: UUID) -> None:
brain_to_delete_user_from = brain_service.get_brain_by_id(brain_id=brain_id)
if brain_to_delete_user_from is None:
raise HTTPException(status_code=404, detail="Brain not found.")
if brain_to_delete_user_from.brain_type == BrainType.API:
brain_definition = get_api_brain_definition(brain_id=brain_id)
if brain_definition is None:
raise HTTPException(
status_code=404, detail="Brain definition not found."
)
secrets = brain_definition.secrets
for secret in secrets:
self.brain_repository.delete_secret(
user_id=user_id,
brain_id=brain_id,
secret_name=secret.name,
)
self.brain_user_repository.delete_brain_user_by_id(
user_id=user_id,
brain_id=brain_id,
)
def get_default_user_brain_or_create_new(self, user: UserIdentity):
default_brain = self.get_user_default_brain(user.id)
if not default_brain:
default_brain = brain_service.create_brain(brain=None, user_id=user.id)
self.brain_user_repository.create_brain_user(
user.id, default_brain.brain_id, RoleEnum.Owner, True
)
return default_brain
def set_as_default_brain_for_user(self, user_id: UUID, brain_id: UUID):
old_default_brain = self.get_user_default_brain(user_id)
if old_default_brain is not None:
self.brain_user_repository.update_brain_user_default_status(
user_id=user_id,
brain_id=old_default_brain.brain_id,
default_brain=False,
)
self.brain_user_repository.update_brain_user_default_status(
user_id=user_id,
brain_id=brain_id,
default_brain=True,
)
def delete_brain_users(self, brain_id: UUID) -> None:
self.brain_user_repository.delete_brain_subscribers(
brain_id=brain_id,
)
def create_brain_user(
self, user_id: UUID, brain_id: UUID, rights: RoleEnum, is_default_brain: bool
):
self.brain_user_repository.create_brain_user(
user_id=user_id,
brain_id=brain_id,
rights=rights,
default_brain=is_default_brain,
)
def get_brain_for_user(self, user_id: UUID, brain_id: UUID):
return self.brain_user_repository.get_brain_for_user(user_id, brain_id) # type: ignore
def get_user_brains(self, user_id: UUID) -> list[MinimalUserBrainEntity]:
results = self.brain_user_repository.get_user_brains(user_id) # type: ignore
return results # type: ignore
def get_brain_users(self, brain_id: UUID) -> List[BrainUser]:
return self.brain_user_repository.get_brain_users(brain_id)
def update_brain_user_rights(
self, brain_id: UUID, user_id: UUID, rights: str
) -> None:
self.brain_user_repository.update_brain_user_rights(
brain_id=brain_id,
user_id=user_id,
rights=rights,
)

View File

@ -0,0 +1,62 @@
from typing import Any, List
from uuid import UUID
from logger import get_logger
from modules.brain.repository.brains_vectors import BrainsVectors
from modules.brain.repository.interfaces.brains_vectors_interface import (
BrainsVectorsInterface,
)
from modules.knowledge.repository.storage import Storage
from packages.embeddings.vectors import get_unique_files_from_vector_ids
logger = get_logger(__name__)
class BrainVectorService:
repository: BrainsVectorsInterface
id: UUID
files: List[Any] = []
def __init__(self, brain_id: UUID):
self.repository = BrainsVectors()
self.id = brain_id
def create_brain_vector(self, vector_id, file_sha1):
return self.repository.create_brain_vector(self.id, vector_id, file_sha1) # type: ignore
def update_brain_with_file(self, file_sha1: str):
# not used
vector_ids = self.repository.get_vector_ids_from_file_sha1(file_sha1)
if vector_ids == None or len(vector_ids) == 0:
logger.info(f"No vector ids found for file {file_sha1}")
return
for vector_id in vector_ids:
self.create_brain_vector(vector_id, file_sha1)
def get_unique_brain_files(self):
"""
Retrieve unique brain data (i.e. uploaded files and crawled websites).
"""
vector_ids = self.repository.get_brain_vector_ids(self.id) # type: ignore
self.files = get_unique_files_from_vector_ids(vector_ids)
return self.files
def delete_file_from_brain(self, file_name: str):
file_name_with_brain_id = f"{self.id}/{file_name}"
storage = Storage()
storage.remove_file(file_name_with_brain_id)
return self.repository.delete_file_from_brain(self.id, file_name) # type: ignore
def delete_file_url_from_brain(self, file_name: str):
return self.repository.delete_file_from_brain(self.id, file_name) # type: ignore
@property
def brain_size(self):
# TODO: change the calculation of the brain size, calculate the size stored for the embeddings + what's in the storage
self.get_unique_brain_files()
current_brain_size = sum(float(doc["size"]) for doc in self.files)
return current_brain_size

View File

@ -0,0 +1 @@
from .contact_routes import contact_router

View File

@ -1,6 +1,6 @@
from fastapi import APIRouter
from logger import get_logger
from models import ContactsSettings
from modules.contact_support.controller.settings import ContactsSettings
from packages.emails.send_email import send_email
from pydantic import BaseModel
@ -10,7 +10,7 @@ class ContactMessage(BaseModel):
content: str
router = APIRouter()
contact_router = APIRouter()
logger = get_logger(__name__)
@ -33,7 +33,7 @@ def resend_contact_sales_email(customer_email: str, content: str):
return send_email(params)
@router.post("/contact")
@contact_router.post("/contact")
def post_contact(message: ContactMessage):
try:
resend_contact_sales_email(message.customer_email, message.content)

View File

@ -0,0 +1,6 @@
from pydantic import BaseSettings
class ContactsSettings(BaseSettings):
resend_contact_sales_from: str = "null"
resend_contact_sales_to: str = "null"

View File

@ -3,16 +3,15 @@ from uuid import UUID
from fastapi import APIRouter, Depends, HTTPException, Query
from logger import get_logger
from middlewares.auth import AuthBearer, get_current_user
from models import Brain
from modules.knowledge.service.knowledge_service import KnowledgeService
from modules.user.entity.user_identity import UserIdentity
from repository.files.delete_file import delete_file_from_storage
from repository.files.generate_file_signed_url import generate_file_signed_url
from routes.authorizations.brain_authorization import (
RoleEnum,
from modules.brain.entity.brain_entity import RoleEnum
from modules.brain.service.brain_authorization_service import (
has_brain_authorization,
validate_brain_authorization,
)
from modules.brain.service.brain_vector_service import BrainVectorService
from modules.knowledge.service.knowledge_service import KnowledgeService
from modules.user.entity.user_identity import UserIdentity
from repository.files.generate_file_signed_url import generate_file_signed_url
knowledge_router = APIRouter()
logger = get_logger(__name__)
@ -56,17 +55,15 @@ async def delete_endpoint(
Delete a specific knowledge from a brain.
"""
brain = Brain(id=brain_id)
knowledge = knowledge_service.get_knowledge(knowledge_id)
file_name = knowledge.file_name if knowledge.file_name else knowledge.url
knowledge_service.remove_knowledge(knowledge_id)
brain_vector_service = BrainVectorService(brain_id)
if knowledge.file_name:
delete_file_from_storage(f"{brain_id}/{knowledge.file_name}")
brain.delete_file_from_brain(knowledge.file_name)
brain_vector_service.delete_file_from_brain(knowledge.file_name)
elif knowledge.url:
brain.delete_file_from_brain(knowledge.url)
brain_vector_service.delete_file_url_from_brain(knowledge.url)
return {
"message": f"{file_name} of brain {brain_id} has been deleted by user {current_user.email}."

View File

@ -0,0 +1,28 @@
from logger import get_logger
from models.settings import get_supabase_client
from modules.knowledge.repository.storage_interface import StorageInterface
logger = get_logger(__name__)
class Storage(StorageInterface):
def __init__(self):
supabase_client = get_supabase_client()
self.db = supabase_client
def upload_file(self, file_name: str):
"""
Upload file to storage
"""
self.db.storage.from_("quivr").download(file_name)
def remove_file(self, file_name: str):
"""
Remove file from storage
"""
try:
response = self.db.storage.from_("quivr").remove([file_name])
return response
except Exception as e:
logger.error(e)
# raise e

View File

@ -0,0 +1,10 @@
from abc import ABC, abstractmethod
class StorageInterface(ABC):
@abstractmethod
def remove_file(self, file_name: str):
"""
Remove file from storage
"""
pass

View File

@ -0,0 +1 @@
from .misc_routes import misc_router

View File

@ -0,0 +1 @@
from .notification_routes import notification_router

View File

@ -0,0 +1 @@
from .onboarding_routes import onboarding_router

View File

@ -0,0 +1 @@
from .prompt_routes import prompt_router

View File

View File

@ -0,0 +1 @@
from .upload_routes import upload_router

View File

@ -7,6 +7,10 @@ from fastapi import APIRouter, Depends, HTTPException, Query, Request, UploadFil
from logger import get_logger
from middlewares.auth import AuthBearer, get_current_user
from models import UserUsage
from modules.brain.entity.brain_entity import RoleEnum
from modules.brain.service.brain_authorization_service import (
validate_brain_authorization,
)
from modules.knowledge.dto.inputs import CreateKnowledgeProperties
from modules.knowledge.service.knowledge_service import KnowledgeService
from modules.notification.dto.inputs import (
@ -18,10 +22,6 @@ from modules.notification.service.notification_service import NotificationServic
from modules.user.entity.user_identity import UserIdentity
from packages.files.file import convert_bytes, get_file_size
from repository.files.upload_file import upload_file_storage
from routes.authorizations.brain_authorization import (
RoleEnum,
validate_brain_authorization,
)
logger = get_logger(__name__)
upload_router = APIRouter()

View File

@ -0,0 +1 @@
from .user_controller import user_router

View File

@ -2,16 +2,16 @@ import time
from fastapi import APIRouter, Depends, Request
from middlewares.auth import AuthBearer, get_current_user
from models import Brain, UserUsage
from models import UserUsage
from modules.brain.service.brain_user_service import BrainUserService
from modules.brain.service.brain_vector_service import BrainVectorService
from modules.user.dto.inputs import UserUpdatableProperties
from modules.user.entity.user_identity import UserIdentity
from modules.user.repository import (
UserUpdatableProperties,
get_user_identity,
update_user_properties,
)
from repository.brain import get_user_default_brain
from modules.user.repository.users import Users
user_router = APIRouter()
brain_user_service = BrainUserService()
user_repository = Users()
@user_router.get("/user", dependencies=[Depends(AuthBearer())], tags=["User"])
@ -42,10 +42,10 @@ async def get_user_endpoint(
user_daily_usage = UserUsage(id=current_user.id)
requests_stats = user_daily_usage.get_user_usage()
default_brain = get_user_default_brain(current_user.id)
default_brain = brain_user_service.get_user_default_brain(current_user.id)
if default_brain:
defaul_brain_size = Brain(id=default_brain.brain_id).brain_size
defaul_brain_size = BrainVectorService(default_brain.brain_id).brain_size
else:
defaul_brain_size = 0
@ -74,7 +74,9 @@ def update_user_identity_route(
"""
Update user identity.
"""
return update_user_properties(current_user.id, user_identity_updatable_properties)
return user_repository.update_user_properties(
current_user.id, user_identity_updatable_properties
)
@user_router.get(
@ -88,4 +90,4 @@ def get_user_identity_route(
"""
Get user identity.
"""
return get_user_identity(current_user.id)
return user_repository.get_user_identity(current_user.id)

View File

@ -0,0 +1,6 @@
from pydantic import BaseModel
class UserUpdatableProperties(BaseModel):
# Nothing for now
empty: bool = True

View File

@ -1,5 +1 @@
from .get_user_identity import get_user_identity
from .create_user_identity import create_user_identity
from .update_user_properties import update_user_properties, UserUpdatableProperties
from .get_user_id_by_user_email import get_user_id_by_user_email
from .get_user_email_by_user_id import get_user_email_by_user_id
from .users import Users

View File

@ -1,22 +0,0 @@
from uuid import UUID
from models import get_supabase_client
from modules.user.entity.user_identity import UserIdentity
def create_user_identity(id: UUID) -> UserIdentity:
supabase_client = get_supabase_client()
response = (
supabase_client.from_("user_identity")
.insert(
{
"user_id": str(id),
}
)
.execute()
)
user_identity = response.data[0]
return UserIdentity(
id=user_identity.get("user_id")
)

View File

@ -1,11 +0,0 @@
from uuid import UUID
from models import get_supabase_client
def get_user_email_by_user_id(user_id: UUID) -> str:
supabase_client = get_supabase_client()
response = supabase_client.rpc(
"get_user_email_by_user_id", {"user_id": str(user_id)}
).execute()
return response.data[0]["email"]

View File

@ -1,15 +0,0 @@
from uuid import UUID
from models import get_supabase_client
def get_user_id_by_user_email(email: str) -> UUID | None:
supabase_client = get_supabase_client()
response = (
supabase_client.rpc("get_user_id_by_user_email", {"user_email": email})
.execute()
.data
)
if len(response) > 0:
return response[0]["user_id"]
return None

View File

@ -1,26 +0,0 @@
from multiprocessing import get_logger
from uuid import UUID
from models import get_supabase_client
from modules.user.entity.user_identity import UserIdentity
from .create_user_identity import create_user_identity
logger = get_logger()
def get_user_identity(user_id: UUID) -> UserIdentity:
supabase_client = get_supabase_client()
response = (
supabase_client.from_("user_identity")
.select("*")
.filter("user_id", "eq", str(user_id))
.execute()
)
if len(response.data) == 0:
return create_user_identity(user_id)
user_identity = response.data[0]
return UserIdentity(id=user_id)

View File

@ -1,34 +0,0 @@
from uuid import UUID
from models.settings import get_supabase_client
from modules.user.entity.user_identity import UserIdentity
from modules.user.repository import create_user_identity
from pydantic import BaseModel
class UserUpdatableProperties(BaseModel):
# Nothing for now
empty: bool = True
def update_user_properties(
user_id: UUID,
user_identity_updatable_properties: UserUpdatableProperties,
) -> UserIdentity:
supabase_client = get_supabase_client()
response = (
supabase_client.from_("user_identity")
.update(user_identity_updatable_properties.__dict__)
.filter("user_id", "eq", user_id) # type: ignore
.execute()
)
if len(response.data) == 0:
return create_user_identity(
user_id
)
user_identity = response.data[0]
return UserIdentity(id=user_id)

View File

@ -0,0 +1,73 @@
from models.settings import get_supabase_client
from modules.user.entity.user_identity import UserIdentity
from modules.user.repository.users_interface import UsersInterface
class Users(UsersInterface):
def __init__(self):
supabase_client = get_supabase_client()
self.db = supabase_client
def create_user_identity(self, id):
response = (
self.db.from_("user_identity")
.insert(
{
"user_id": str(id),
}
)
.execute()
)
user_identity = response.data[0]
return UserIdentity(id=user_identity.get("user_id"))
def update_user_properties(
self,
user_id,
user_identity_updatable_properties,
):
response = (
self.db.from_("user_identity")
.update(user_identity_updatable_properties.__dict__)
.filter("user_id", "eq", user_id) # type: ignore
.execute()
)
if len(response.data) == 0:
return self.create_user_identity(user_id)
user_identity = response.data[0]
print("USER_IDENTITY", user_identity)
return UserIdentity(id=user_id)
def get_user_identity(self, user_id):
response = (
self.db.from_("user_identity")
.select("*")
.filter("user_id", "eq", str(user_id))
.execute()
)
if len(response.data) == 0:
return self.create_user_identity(user_id)
user_identity = response.data[0]
print("USER_IDENTITY", user_identity)
return UserIdentity(id=user_id)
def get_user_id_by_user_email(self, email):
response = (
self.db.rpc("get_user_id_by_user_email", {"user_email": email})
.execute()
.data
)
if len(response) > 0:
return response[0]["user_id"]
return None
def get_user_email_by_user_id(self, user_id):
response = self.db.rpc(
"get_user_email_by_user_id", {"user_id": str(user_id)}
).execute()
return response.data[0]["email"]

View File

@ -0,0 +1,46 @@
from abc import ABC, abstractmethod
from uuid import UUID
from modules.user.dto.inputs import UserUpdatableProperties
from modules.user.entity.user_identity import UserIdentity
class UsersInterface(ABC):
@abstractmethod
def create_user_identity(self, id: UUID) -> UserIdentity:
"""
Create a user identity
"""
pass
@abstractmethod
def update_user_properties(
self,
user_id: UUID,
user_identity_updatable_properties: UserUpdatableProperties,
) -> UserIdentity:
"""
Update the user properties
"""
pass
@abstractmethod
def get_user_identity(self, user_id: UUID) -> UserIdentity:
"""
Get the user identity
"""
pass
@abstractmethod
def get_user_id_by_user_email(self, email: str) -> UUID | None:
"""
Get the user id by user email
"""
pass
@abstractmethod
def get_user_email_by_user_id(self, user_id: UUID) -> str:
"""
Get the user email by user id
"""
pass

View File

@ -1,2 +1 @@
from .get_user_id_by_email import get_user_id_by_email
from modules.user.repository import get_user_email_by_user_id
from .user_service import UserService

View File

@ -1,7 +0,0 @@
from uuid import UUID
from modules.user.repository import get_user_id_by_user_email
def get_user_id_by_email(email: str) -> UUID | None:
return get_user_id_by_user_email(email)

View File

@ -0,0 +1,17 @@
from uuid import UUID
from modules.user.repository.users import Users
from modules.user.repository.users_interface import UsersInterface
class UserService:
repository: UsersInterface
def __init__(self):
self.repository = Users()
def get_user_id_by_email(self, email: str) -> UUID | None:
return self.repository.get_user_id_by_user_email(email)
def get_user_email_by_user_id(self, user_id: UUID) -> str | None:
return self.repository.get_user_email_by_user_id(user_id)

View File

@ -46,7 +46,7 @@ def process_batch(batch_ids: List[str]):
# TODO: move to Knowledge class
def get_unique_files_from_vector_ids(vectors_ids: List[str]):
def get_unique_files_from_vector_ids(vectors_ids):
# Move into Vectors class
"""
Retrieve unique user data vectors.

View File

@ -4,9 +4,8 @@ import time
from langchain.document_loaders import GitLoader
from langchain.schema import Document
from langchain.text_splitter import RecursiveCharacterTextSplitter
# from models import Brain, File
# from packages.embeddings.vectors import Neurons
from models.files import File
from packages.embeddings.vectors import Neurons
from packages.files.file import compute_sha1_from_content
@ -58,24 +57,21 @@ async def process_github(
print(doc_with_metadata.metadata["file_name"])
# TO FIX: Import of file and brain creates a circular dependency
# file = File(
# file_sha1=compute_sha1_from_content(doc.page_content.encode("utf-8"))
# )
file = File(
file_sha1=compute_sha1_from_content(doc.page_content.encode("utf-8"))
)
# file_exists = file.file_already_exists()
file_exists = file.file_already_exists()
# if not file_exists:
# neurons = Neurons()
# created_vector = neurons.create_vector(doc_with_metadata)
if not file_exists:
neurons = Neurons()
created_vector = neurons.create_vector(doc_with_metadata)
# file_exists_in_brain = file.file_already_exists_in_brain(brain_id)
file_exists_in_brain = file.file_already_exists_in_brain(brain_id)
# if not file_exists_in_brain:
# brain = Brain(id=brain_id)
# file.link_file_to_brain(brain)
if not file_exists_in_brain:
file.link_file_to_brain(brain_id)
return {
# "message": f"✅ Github with {len(documents)} files has been uploaded.",
"message": "Github processor is currently unavailable.",
"message": f"✅ Github with {len(documents)} files has been uploaded.",
"type": "success",
}

View File

@ -1,5 +1,4 @@
from models.brains import Brain
from repository.brain.get_brain_by_id import get_brain_by_id
from modules.brain.service.brain_service import BrainService
from .parsers.audio import process_audio
from .parsers.code_python import process_python
@ -46,6 +45,9 @@ def create_response(message, type):
return {"message": message, "type": type}
brain_service = BrainService()
# TODO: Move filter_file to a file service to avoid circular imports from models/files.py for File class
async def filter_file(
file,
@ -58,7 +60,7 @@ async def filter_file(
file_exists_in_brain = file.file_already_exists_in_brain(brain_id)
using_file_name = original_file_name or file.file.filename if file.file else ""
brain = get_brain_by_id(brain_id)
brain = brain_service.get_brain_by_id(brain_id)
if brain is None:
raise Exception("It seems like you're uploading knowledge to an unknown brain.")
@ -73,7 +75,7 @@ async def filter_file(
"error", # pyright: ignore reportPrivateUsage=none
)
elif file_exists:
file.link_file_to_brain(brain=Brain(id=brain_id))
file.link_file_to_brain(brain_id)
return create_response(
f"{using_file_name} has been uploaded to brain {brain.name}.", # pyright: ignore reportPrivateUsage=none
"success",

View File

@ -2,7 +2,7 @@ from typing import Optional
from uuid import UUID
from models import get_supabase_db
from models.ApiBrainDefinition import ApiBrainDefinition
from modules.brain.entity.api_brain_definition_entity import ApiBrainDefinition
def get_api_brain_definition(brain_id: UUID) -> Optional[ApiBrainDefinition]:

View File

@ -1,8 +1,8 @@
from typing import Optional
from uuid import UUID
from models.ApiBrainDefinition import ApiBrainDefinition
from models.settings import get_supabase_db
from modules.brain.entity.api_brain_definition_entity import ApiBrainDefinition
def update_api_brain_definition(

View File

@ -1,16 +1 @@
from .create_brain import create_brain
from .create_brain_user import create_brain_user
from .delete_brain_users import delete_brain_users
from .get_brain_by_id import get_brain_by_id
from .get_public_brains import get_public_brains
from .get_brain_details import get_brain_details
from .get_brain_for_user import get_brain_for_user
from .get_brain_prompt_id import get_brain_prompt_id
from .get_default_user_brain import get_user_default_brain
from .get_default_user_brain_or_create_new import \
get_default_user_brain_or_create_new
from .get_question_context_from_brain import get_question_context_from_brain
from .get_user_brains import get_user_brains
from .set_as_default_brain_for_user import set_as_default_brain_for_user
from .update_brain import update_brain_by_id
from .update_user_rights import update_brain_user_rights

View File

@ -1,45 +0,0 @@
from uuid import UUID
from fastapi import HTTPException
from models import BrainEntity, get_supabase_db
from models.brain_entity import BrainType
from models.databases.supabase.brains import CreateBrainProperties
from repository.api_brain_definition.add_api_brain_definition import (
add_api_brain_definition,
)
from repository.external_api_secret.create_secret import create_secret
def create_brain(brain: CreateBrainProperties, user_id: UUID) -> BrainEntity:
if brain.brain_type == BrainType.API:
if brain.brain_definition is None:
raise HTTPException(status_code=404, detail="Brain definition not found")
if brain.brain_definition.url is None:
raise HTTPException(status_code=404, detail="Brain url not found")
if brain.brain_definition.method is None:
raise HTTPException(status_code=404, detail="Brain method not found")
supabase_db = get_supabase_db()
created_brain = supabase_db.create_brain(brain)
if brain.brain_type == BrainType.API and brain.brain_definition is not None:
add_api_brain_definition(
brain_id=created_brain.brain_id,
api_brain_definition=brain.brain_definition,
)
secrets_values = brain.brain_secrets_values
for secret_name in secrets_values:
create_secret(
user_id=user_id,
brain_id=created_brain.brain_id,
secret_name=secret_name,
secret_value=secrets_values[secret_name],
)
return created_brain

View File

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

View File

@ -1,35 +0,0 @@
from uuid import UUID
from fastapi import HTTPException
from models.brain_entity import BrainType
from models.settings import get_supabase_db
from modules.knowledge.service.knowledge_service import KnowledgeService
from repository.api_brain_definition.delete_api_brain_definition import (
delete_api_brain_definition,
)
from repository.brain import get_brain_by_id
from repository.brain.delete_brain_secrets import delete_brain_secrets_values
knowledge_service = KnowledgeService()
def delete_brain(brain_id: UUID) -> dict[str, str]:
supabase_db = get_supabase_db()
brain_to_delete = get_brain_by_id(brain_id=brain_id)
if brain_to_delete is None:
raise HTTPException(status_code=404, detail="Brain not found.")
if brain_to_delete.brain_type == BrainType.API:
delete_brain_secrets_values(
brain_id=brain_id,
)
delete_api_brain_definition(brain_id=brain_id)
else:
knowledge_service.remove_brain_all_knowledge(brain_id)
supabase_db.delete_brain_vector(str(brain_id))
supabase_db.delete_brain_users(str(brain_id))
supabase_db.delete_brain(str(brain_id))
return {"message": "Brain deleted."}

View File

@ -1,30 +1,3 @@
from uuid import UUID
from modules.brain.service.brain_user_service import BrainUserService
from fastapi import HTTPException
from models.settings import get_supabase_db
from repository.api_brain_definition.get_api_brain_definition import (
get_api_brain_definition,
)
from repository.external_api_secret import delete_secret
def delete_brain_secrets_values(brain_id: UUID) -> None:
supabase_db = get_supabase_db()
brain_definition = get_api_brain_definition(brain_id=brain_id)
if brain_definition is None:
raise HTTPException(status_code=404, detail="Brain definition not found.")
secrets = brain_definition.secrets
if len(secrets) > 0:
brain_users = supabase_db.get_brain_users(brain_id=brain_id)
for user in brain_users:
for secret in secrets:
delete_secret(
user_id=user.user_id,
brain_id=brain_id,
secret_name=secret.name,
)
brain_user_service = BrainUserService()

View File

@ -1,35 +0,0 @@
from uuid import UUID
from fastapi import HTTPException
from models.brain_entity import BrainType
from models.settings import get_supabase_db
from repository.api_brain_definition.get_api_brain_definition import (
get_api_brain_definition,
)
from repository.brain.get_brain_by_id import get_brain_by_id
from repository.external_api_secret.delete_secret import delete_secret
def delete_brain_user(user_id: UUID, brain_id: UUID) -> None:
supabase_db = get_supabase_db()
brain_to_delete_user_from = get_brain_by_id(brain_id=brain_id)
if brain_to_delete_user_from is None:
raise HTTPException(status_code=404, detail="Brain not found.")
if brain_to_delete_user_from.brain_type == BrainType.API:
brain_definition = get_api_brain_definition(brain_id=brain_id)
if brain_definition is None:
raise HTTPException(status_code=404, detail="Brain definition not found.")
secrets = brain_definition.secrets
for secret in secrets:
delete_secret(
user_id=user_id,
brain_id=brain_id,
secret_name=secret.name,
)
supabase_db.delete_brain_user_by_id(
user_id=user_id,
brain_id=brain_id,
)

View File

@ -1,10 +0,0 @@
from uuid import UUID
from models.settings import get_supabase_db
def delete_brain_users(brain_id: UUID) -> None:
supabase_db = get_supabase_db()
supabase_db.delete_brain_subscribers(
brain_id=brain_id,
)

View File

@ -1,9 +0,0 @@
from uuid import UUID
from models import BrainEntity, get_supabase_db
def get_brain_by_id(brain_id: UUID) -> BrainEntity | None:
supabase_db = get_supabase_db()
return supabase_db.get_brain_by_id(brain_id)

View File

@ -1,27 +0,0 @@
from uuid import UUID
from models import BrainEntity, get_supabase_client
from models.brain_entity import BrainType
from repository.api_brain_definition.get_api_brain_definition import (
get_api_brain_definition,
)
def get_brain_details(brain_id: UUID) -> BrainEntity | None:
supabase_client = get_supabase_client()
response = (
supabase_client.table("brains")
.select("*")
.filter("brain_id", "eq", str(brain_id))
.execute()
)
if response.data == []:
return None
brain = BrainEntity(**response.data[0])
if brain.brain_type == BrainType.API:
brain_definition = get_api_brain_definition(brain_id)
brain.brain_definition = brain_definition
return brain

View File

@ -1,8 +0,0 @@
from uuid import UUID
from models import MinimalBrainEntity, get_supabase_db
def get_brain_for_user(user_id: UUID, brain_id: UUID) -> MinimalBrainEntity | None:
supabase_db = get_supabase_db()
return supabase_db.get_brain_for_user(user_id, brain_id) # type: ignore

View File

@ -1,10 +0,0 @@
from uuid import UUID
from repository.brain import get_brain_by_id
def get_brain_prompt_id(brain_id: UUID) -> UUID | None:
brain = get_brain_by_id(brain_id)
prompt_id = brain.prompt_id if brain else None
return prompt_id

View File

@ -1,10 +0,0 @@
from uuid import UUID
from models.brain_entity import BrainUser
from models.settings import get_supabase_db
def get_brain_users(brain_id: UUID) -> list[BrainUser]:
supabase_db = get_supabase_db()
return supabase_db.get_brain_users(brain_id)

View File

@ -1,21 +0,0 @@
from uuid import UUID
from logger import get_logger
from models import BrainEntity, get_supabase_db
from repository.brain import get_brain_by_id
logger = get_logger(__name__)
def get_user_default_brain(user_id: UUID) -> BrainEntity | None:
supabase_db = get_supabase_db()
brain_id = supabase_db.get_default_user_brain_id(user_id)
logger.info(f"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

@ -1,15 +0,0 @@
from models import BrainEntity
from models.databases.supabase.brains import CreateBrainProperties
from modules.user.entity.user_identity import UserIdentity
from repository.brain import create_brain, create_brain_user, get_user_default_brain
from routes.authorizations.types import RoleEnum
def get_default_user_brain_or_create_new(user: UserIdentity) -> BrainEntity:
default_brain = get_user_default_brain(user.id)
if not default_brain:
default_brain = create_brain(brain=CreateBrainProperties(), user_id=user.id)
create_brain_user(user.id, default_brain.brain_id, RoleEnum.Owner, True)
return default_brain

Some files were not shown because too many files have changed in this diff Show More