diff --git a/backend/main.py b/backend/main.py index 86a11021ab0b..c986d2726d63 100644 --- a/backend/main.py +++ b/backend/main.py @@ -16,6 +16,7 @@ from modules.chat.controller import chat_router from modules.contact_support.controller import contact_router from modules.knowledge.controller import knowledge_router +from modules.message.controller import messages_router from modules.misc.controller import misc_router from modules.notification.controller import notification_router from modules.onboarding.controller import onboarding_router @@ -49,6 +50,7 @@ ], ) + app = FastAPI() add_cors_middleware(app) @@ -58,9 +60,9 @@ app.include_router(crawl_router) app.include_router(onboarding_router) app.include_router(misc_router) - app.include_router(upload_router) app.include_router(user_router) +app.include_router(messages_router) app.include_router(api_key_router) app.include_router(subscription_router) app.include_router(prompt_router) diff --git a/backend/modules/message/__init__.py b/backend/modules/message/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/backend/modules/message/controller/__init__.py b/backend/modules/message/controller/__init__.py new file mode 100644 index 000000000000..08858bd92775 --- /dev/null +++ b/backend/modules/message/controller/__init__.py @@ -0,0 +1 @@ +from .message_routes import messages_router diff --git a/backend/modules/message/controller/message_routes.py b/backend/modules/message/controller/message_routes.py new file mode 100644 index 000000000000..40d6fd4d8689 --- /dev/null +++ b/backend/modules/message/controller/message_routes.py @@ -0,0 +1,79 @@ +from typing import List + +from fastapi import APIRouter, Depends +from middlewares.auth import ( # Assuming you have a get_current_user function + AuthBearer, + get_current_user, +) +from modules.message.dto.inputs import CreateMessageProperties, UpdateMessageProperties +from modules.message.entity.message import Message +from modules.message.service.messages_service import MessagesService +from modules.user.entity.user_identity import UserIdentity + +messages_router = APIRouter() + +messages_services = MessagesService() + + +@messages_router.get( + "/messages", + dependencies=[Depends(AuthBearer())], + tags=["Messages"], +) +async def get_messages_user_brain( + brain_id: str, + current_user: UserIdentity = Depends(get_current_user), +) -> List[Message] | List[None]: + """ + Get users messages information for the current user and brain + """ + + return messages_services.get_messages_brain(current_user.id, brain_id) + + +@messages_router.put( + "/messages", + dependencies=[Depends(AuthBearer())], + tags=["Messages"], +) +async def update_messages_brain( + message: UpdateMessageProperties, + current_user: UserIdentity = Depends(get_current_user), +) -> Message | None: + """ + Update user message information for the current user + """ + + return messages_services.update_message(current_user.id, message) + + +@messages_router.delete( + "/messages", + dependencies=[Depends(AuthBearer())], + tags=["Messages"], +) +async def delete_messages_brain( + message_id: str, + current_user: UserIdentity = Depends(get_current_user), +) -> Message | None: + """ + Delete user message for the current user with message_id + """ + + return messages_services.remove_message(current_user.id, message_id) + + +@messages_router.post( + "/messages", + dependencies=[Depends(AuthBearer())], + tags=["Messages"], +) +async def create_messages_brain( + message: CreateMessageProperties, + current_user: UserIdentity = Depends(get_current_user), +) -> Message | None: + """ + Create user message for the current user + """ + + return messages_services.create_message(current_user.id, message) diff --git a/backend/modules/message/dto/__init__.py b/backend/modules/message/dto/__init__.py new file mode 100644 index 000000000000..9fe9f8b1a635 --- /dev/null +++ b/backend/modules/message/dto/__init__.py @@ -0,0 +1 @@ +from .inputs import CreateMessageProperties, UpdateMessageProperties diff --git a/backend/modules/message/dto/inputs.py b/backend/modules/message/dto/inputs.py new file mode 100644 index 000000000000..c71fb51d51a1 --- /dev/null +++ b/backend/modules/message/dto/inputs.py @@ -0,0 +1,25 @@ +from uuid import UUID + +from pydantic import BaseModel + + +class CreateMessageProperties(BaseModel): + + """Properties that can be received on message creation""" + + brain_id: UUID + content: str + + class Config: + extra = "forbid" + + +class UpdateMessageProperties(BaseModel): + + """Properties that can be received on message update""" + + message_id: UUID + content: str + + class Config: + extra = "forbid" diff --git a/backend/modules/message/entity/__init__.py b/backend/modules/message/entity/__init__.py new file mode 100644 index 000000000000..76a3dce92ab6 --- /dev/null +++ b/backend/modules/message/entity/__init__.py @@ -0,0 +1 @@ +from .message import Message diff --git a/backend/modules/message/entity/message.py b/backend/modules/message/entity/message.py new file mode 100644 index 000000000000..3104372aee20 --- /dev/null +++ b/backend/modules/message/entity/message.py @@ -0,0 +1,13 @@ +from uuid import UUID + +from pydantic import BaseModel + + +class Message(BaseModel): + """Response when getting messages""" + + message_id: UUID + brain_id: UUID + user_id: UUID + content: str + created_at: str diff --git a/backend/modules/message/repository/__init__.py b/backend/modules/message/repository/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/backend/modules/message/repository/messages.py b/backend/modules/message/repository/messages.py new file mode 100644 index 000000000000..a760c1fb722c --- /dev/null +++ b/backend/modules/message/repository/messages.py @@ -0,0 +1,79 @@ +from models.settings import get_supabase_client +from modules.message.entity.message import Message + +from .messages_interface import MessagesInterface + + +class Messages(MessagesInterface): + def __init__(self): + supabase_client = get_supabase_client() + self.db = supabase_client + + def get_messages_brain(self, user_id, brain_id): + """ + Get user messages information by user_id and brain_id + """ + messages_data = ( + self.db.from_("messages") + .select("*") + .match({"user_id": str(user_id), "brain_id": str(brain_id)}) + .execute() + ).data + + if messages_data == []: + return [] + + messages = [] + for message in messages_data: + messages.append(Message(**message)) + + return messages + + def update_message(self, user_id, update): + """ + Update user messages information by user_id and message_id + """ + response = ( + self.db.table("messages") + .update({"content": update.content}) + .match({"user_id": str(user_id), "message_id": str(update.message_id)}) + .execute() + ) + + if len(response.data) == 0: + return None + return Message(**response.data[0]) + + def remove_message(self, user_id, message_id): + """ + Remove message by user_id and message_id + """ + response = ( + self.db.table("messages") + .delete() + .match({"user_id": str(user_id), "message_id": str(message_id)}) + .execute() + ) + + if len(response.data) == 0: + return None + return Message(**response.data[0]) + + def create_message(self, user_id, message_create): + """ + Create user messages information by user_id + """ + response = ( + self.db.table("messages") + .insert( + [ + { + "brain_id": str(message_create.brain_id), + "user_id": str(user_id), + "content": message_create.content, + } + ] + ) + .execute() + ) + return Message(**response.data[0]) diff --git a/backend/modules/message/repository/messages_interface.py b/backend/modules/message/repository/messages_interface.py new file mode 100644 index 000000000000..8801ccc0418b --- /dev/null +++ b/backend/modules/message/repository/messages_interface.py @@ -0,0 +1,40 @@ +from abc import ABC, abstractmethod +from typing import List +from uuid import UUID + +from modules.message.dto.inputs import CreateMessageProperties, UpdateMessageProperties +from modules.message.entity.message import Message + + +class MessagesInterface(ABC): + @abstractmethod + def get_messages_brain( + self, user_id: UUID, brain: UUID + ) -> List[Message] | List[None]: + """ + Get messages by user_id and brain_id + """ + pass + + @abstractmethod + def update_message( + self, user_id: UUID, update: UpdateMessageProperties + ) -> Message | None: + """Update user onboarding information by user_id""" + pass + + @abstractmethod + def remove_message(self, user_id: UUID, message_id: UUID) -> Message | None: + """ + Remove message by user_id and message_id + """ + pass + + @abstractmethod + def create_message( + self, user_id: UUID, message_create: CreateMessageProperties + ) -> Message: + """ + Create user onboarding information by user_id + """ + pass diff --git a/backend/modules/message/service/__init__.py b/backend/modules/message/service/__init__.py new file mode 100644 index 000000000000..4e05db98879b --- /dev/null +++ b/backend/modules/message/service/__init__.py @@ -0,0 +1 @@ +from .messages_service import MessagesService \ No newline at end of file diff --git a/backend/modules/message/service/messages_service.py b/backend/modules/message/service/messages_service.py new file mode 100644 index 000000000000..89b31b57265e --- /dev/null +++ b/backend/modules/message/service/messages_service.py @@ -0,0 +1,40 @@ +from typing import List +from uuid import UUID + +from modules.message.dto.inputs import CreateMessageProperties, UpdateMessageProperties +from modules.message.entity.message import Message +from modules.message.repository.messages import Messages +from modules.message.repository.messages_interface import MessagesInterface + + +class MessagesService: + repository: MessagesInterface + + def __init__(self): + self.repository = Messages() + + def get_messages_brain( + self, user_id: UUID, brain_id: UUID + ) -> List[Message] | List[None]: + """Update user onboarding information by user_id""" + + return self.repository.get_messages_brain(user_id, brain_id) + + def update_message( + self, user_id: UUID, message: UpdateMessageProperties + ) -> Message | None: + """Update user onboarding information by user_id""" + + return self.repository.update_message(user_id, message) + + def remove_message(self, user_id: UUID, message_id: UUID) -> Message | None: + """Update user onboarding information by user_id""" + + return self.repository.remove_message(user_id, message_id) + + def create_message( + self, user_id: UUID, message_create: CreateMessageProperties + ) -> Message: + """Update user onboarding information by user_id""" + + return self.repository.create_message(user_id, message_create) diff --git a/scripts/20231209003200_new_message_table.sql b/scripts/20231209003200_new_message_table.sql new file mode 100644 index 000000000000..ee8541b3edec --- /dev/null +++ b/scripts/20231209003200_new_message_table.sql @@ -0,0 +1,33 @@ +DO $$ +BEGIN + -- Check if the messages table does not exist + IF NOT EXISTS (SELECT FROM pg_catalog.pg_tables + WHERE schemaname = 'public' AND tablename = 'messages') THEN + + -- Create the table if it does not exist + CREATE TABLE public.messages ( + message_id uuid NOT NULL DEFAULT gen_random_uuid(), + brain_id uuid NOT NULL, + user_id uuid NOT NULL, + content text NOT NULL, + created_at timestamp without time zone NOT NULL DEFAULT current_timestamp, + CONSTRAINT messages_pkey PRIMARY KEY (message_id), + CONSTRAINT messages_brain_id_fkey FOREIGN KEY (brain_id) REFERENCES public.brains (brain_id), + CONSTRAINT messages_user_id_fkey FOREIGN KEY (user_id) REFERENCES auth.users (id) + ) TABLESPACE pg_default; + + RAISE NOTICE 'Created table public.messages'; + ELSE + RAISE NOTICE 'Table public.messages already exists'; + END IF; +END $$; + + +-- Update migrations table +INSERT INTO migrations (name) +SELECT '20231209003200_new_message_table' +WHERE NOT EXISTS ( + SELECT 1 FROM migrations WHERE name = '20231209003200_new_message_table' +); + +COMMIT; diff --git a/scripts/tables-ollama.sql b/scripts/tables-ollama.sql index c999df860b6a..e6c2d297ee3e 100644 --- a/scripts/tables-ollama.sql +++ b/scripts/tables-ollama.sql @@ -133,7 +133,7 @@ DO $$ BEGIN IF NOT EXISTS (SELECT 1 FROM pg_type WHERE typname = 'brain_type_enum') THEN -- Create the ENUM type 'brain_type' if it doesn't exist - CREATE TYPE brain_type_enum AS ENUM ('doc', 'api'); + CREATE TYPE brain_type_enum AS ENUM ('doc', 'api', 'composite'); END IF; END $$; @@ -206,6 +206,14 @@ CREATE TABLE IF NOT EXISTS brain_subscription_invitations ( FOREIGN KEY (brain_id) REFERENCES brains (brain_id) ); +-- Table for storing the relationship between brains for composite brains +CREATE TABLE IF NOT EXISTS composite_brain_connections ( + composite_brain_id UUID NOT NULL REFERENCES brains(brain_id), + connected_brain_id UUID NOT NULL REFERENCES brains(brain_id), + PRIMARY KEY (composite_brain_id, connected_brain_id), + CHECK (composite_brain_id != connected_brain_id) +); + --- Create user_identity table CREATE TABLE IF NOT EXISTS user_identity ( user_id UUID PRIMARY KEY, @@ -444,9 +452,28 @@ begin end; $$; +create schema if not exists extensions; + +create table if not exists + extensions.wrappers_fdw_stats (); + +grant all on extensions.wrappers_fdw_stats to service_role; + + +CREATE TABLE public.messages ( + message_id uuid NOT NULL DEFAULT gen_random_uuid(), + brain_id uuid NOT NULL, + user_id uuid NOT NULL, + content text NOT NULL, + created_at timestamp without time zone NOT NULL DEFAULT current_timestamp, + CONSTRAINT messages_pkey PRIMARY KEY (message_id), + CONSTRAINT messages_brain_id_fkey FOREIGN KEY (brain_id) REFERENCES public.brains (brain_id), + CONSTRAINT messages_user_id_fkey FOREIGN KEY (user_id) REFERENCES auth.users (id) +) TABLESPACE pg_default; + INSERT INTO migrations (name) -SELECT '20231203173900_new_api_key_format' +SELECT '20231205163000_new_table_composite_brain_connections' WHERE NOT EXISTS ( - SELECT 1 FROM migrations WHERE name = '20231203173900_new_api_key_format' + SELECT 1 FROM migrations WHERE name = '20231205163000_new_table_composite_brain_connections' ); diff --git a/scripts/tables.sql b/scripts/tables.sql index 435fb04ff1a9..49aeb6c0e2fb 100644 --- a/scripts/tables.sql +++ b/scripts/tables.sql @@ -460,6 +460,18 @@ create table if not exists grant all on extensions.wrappers_fdw_stats to service_role; +CREATE TABLE public.messages ( + message_id uuid NOT NULL DEFAULT gen_random_uuid(), + brain_id uuid NOT NULL, + user_id uuid NOT NULL, + content text NOT NULL, + created_at timestamp without time zone NOT NULL DEFAULT current_timestamp, + CONSTRAINT messages_pkey PRIMARY KEY (message_id), + CONSTRAINT messages_brain_id_fkey FOREIGN KEY (brain_id) REFERENCES public.brains (brain_id), + CONSTRAINT messages_user_id_fkey FOREIGN KEY (user_id) REFERENCES auth.users (id) +) TABLESPACE pg_default; + + INSERT INTO migrations (name) SELECT '20231205163000_new_table_composite_brain_connections' WHERE NOT EXISTS (