diff --git a/backend/modules/brain/controller/brain_routes.py b/backend/modules/brain/controller/brain_routes.py index 1627c3e54..1d4927203 100644 --- a/backend/modules/brain/controller/brain_routes.py +++ b/backend/modules/brain/controller/brain_routes.py @@ -81,9 +81,12 @@ async def retrieve_default_brain( ], tags=["Brain"], ) -async def retrieve_brain_by_id(brain_id: UUID): +async def retrieve_brain_by_id( + brain_id: UUID, + current_user: UserIdentity = Depends(get_current_user), +): """Retrieve details of a specific brain by its ID.""" - brain_details = brain_service.get_brain_details(brain_id) + brain_details = brain_service.get_brain_details(brain_id, current_user.id) if brain_details is None: raise HTTPException(status_code=404, detail="Brain details not found") return brain_details @@ -131,10 +134,12 @@ async def create_new_brain( tags=["Brain"], ) async def update_existing_brain( - brain_id: UUID, brain_update_data: BrainUpdatableProperties + brain_id: UUID, + brain_update_data: BrainUpdatableProperties, + current_user: UserIdentity = Depends(get_current_user), ): """Update an existing brain's configuration.""" - existing_brain = brain_service.get_brain_details(brain_id) + existing_brain = brain_service.get_brain_details(brain_id, current_user.id) if existing_brain is None: raise HTTPException(status_code=404, detail="Brain not found") @@ -169,7 +174,7 @@ async def update_existing_brain_secrets( ): """Update an existing brain's secrets.""" - existing_brain = brain_service.get_brain_details(brain_id) + existing_brain = brain_service.get_brain_details(brain_id, None) if existing_brain is None: raise HTTPException(status_code=404, detail="Brain not found") diff --git a/backend/modules/brain/dto/inputs.py b/backend/modules/brain/dto/inputs.py index efc7908fc..ccfd6b96b 100644 --- a/backend/modules/brain/dto/inputs.py +++ b/backend/modules/brain/dto/inputs.py @@ -1,4 +1,3 @@ -from enum import Enum from typing import Optional from uuid import UUID @@ -10,17 +9,13 @@ from modules.brain.entity.api_brain_definition_entity import ( ApiBrainDefinitionSecret, ) from modules.brain.entity.brain_entity import BrainType +from modules.brain.entity.integration_brain import IntegrationType from pydantic import BaseModel, Extra logger = get_logger(__name__) -class IntegrationType(str, Enum): - CUSTOM = "custom" - SYNC = "sync" - - -class CreateApiBrainDefinition(BaseModel, extra=Extra.forbid): +class CreateApiBrainDefinition(BaseModel, extra="ignore"): method: ApiBrainAllowedMethods url: str params: Optional[ApiBrainDefinitionSchema] = ApiBrainDefinitionSchema() @@ -30,7 +25,7 @@ class CreateApiBrainDefinition(BaseModel, extra=Extra.forbid): jq_instructions: Optional[str] = None -class CreateIntegrationBrain(BaseModel, extra=Extra.forbid): +class CreateIntegrationBrain(BaseModel, extra="ignore"): integration_name: str integration_logo_url: str connection_settings: dict @@ -39,16 +34,16 @@ class CreateIntegrationBrain(BaseModel, extra=Extra.forbid): max_files: int -class BrainIntegrationSettings(BaseModel, extra=Extra.forbid): +class BrainIntegrationSettings(BaseModel, extra="ignore"): integration_id: str settings: dict -class BrainIntegrationUpdateSettings(BaseModel, extra=Extra.forbid): +class BrainIntegrationUpdateSettings(BaseModel, extra="ignore"): settings: dict -class CreateBrainProperties(BaseModel, extra=Extra.forbid): +class CreateBrainProperties(BaseModel, extra="ignore"): name: Optional[str] = "Default brain" description: str = "This is a description" status: Optional[str] = "private" @@ -69,7 +64,8 @@ class CreateBrainProperties(BaseModel, extra=Extra.forbid): return brain_dict -class BrainUpdatableProperties(BaseModel): +class BrainUpdatableProperties(BaseModel, extra="ignore"): + name: Optional[str] = None description: Optional[str] = None temperature: Optional[float] = None diff --git a/backend/modules/brain/entity/brain_entity.py b/backend/modules/brain/entity/brain_entity.py index be3aecd09..e7b55ac70 100644 --- a/backend/modules/brain/entity/brain_entity.py +++ b/backend/modules/brain/entity/brain_entity.py @@ -3,6 +3,10 @@ from typing import List, Optional from uuid import UUID from modules.brain.entity.api_brain_definition_entity import ApiBrainDefinitionEntity +from modules.brain.entity.integration_brain import ( + IntegrationDescriptionEntity, + IntegrationEntity, +) from pydantic import BaseModel @@ -28,6 +32,8 @@ class BrainEntity(BaseModel): connected_brains_ids: Optional[List[UUID]] = None raw: Optional[bool] = None jq_instructions: Optional[str] = None + integration: Optional[IntegrationEntity] = None + integration_description: Optional[IntegrationDescriptionEntity] = None @property def id(self) -> UUID: diff --git a/backend/modules/brain/entity/integration_brain.py b/backend/modules/brain/entity/integration_brain.py index 5b2d46f7b..3a9a364e0 100644 --- a/backend/modules/brain/entity/integration_brain.py +++ b/backend/modules/brain/entity/integration_brain.py @@ -1,10 +1,15 @@ +from enum import Enum from typing import Optional from uuid import UUID -from modules.brain.dto.inputs import IntegrationType from pydantic import BaseModel +class IntegrationType(str, Enum): + CUSTOM = "custom" + SYNC = "sync" + + class IntegrationDescriptionEntity(BaseModel): id: UUID integration_name: str diff --git a/backend/modules/brain/repository/interfaces/integration_brains_interface.py b/backend/modules/brain/repository/interfaces/integration_brains_interface.py index 346400a83..0d8d7140e 100644 --- a/backend/modules/brain/repository/interfaces/integration_brains_interface.py +++ b/backend/modules/brain/repository/interfaces/integration_brains_interface.py @@ -9,7 +9,7 @@ from modules.brain.entity.integration_brain import ( class IntegrationBrainInterface(ABC): @abstractmethod - def get_integration_brain(self, brain_id: UUID) -> IntegrationEntity: + def get_integration_brain(self, brain_id: UUID, user_id: UUID) -> IntegrationEntity: """Get the integration brain entity Args: diff --git a/backend/modules/brain/service/brain_authorization_service.py b/backend/modules/brain/service/brain_authorization_service.py index e0431dabd..c0c2c5e2c 100644 --- a/backend/modules/brain/service/brain_authorization_service.py +++ b/backend/modules/brain/service/brain_authorization_service.py @@ -47,7 +47,7 @@ def validate_brain_authorization( return: None """ - brain = brain_service.get_brain_details(brain_id) + brain = brain_service.get_brain_details(brain_id, user_id) if brain and brain.status == "public": return diff --git a/backend/modules/brain/service/brain_service.py b/backend/modules/brain/service/brain_service.py index 876f53631..e152039ae 100644 --- a/backend/modules/brain/service/brain_service.py +++ b/backend/modules/brain/service/brain_service.py @@ -22,6 +22,7 @@ from modules.brain.repository.interfaces import ( BrainsVectorsInterface, CompositeBrainsConnectionsInterface, ExternalApiSecretsInterface, + IntegrationBrainInterface, IntegrationDescriptionInterface, ) from modules.brain.service.api_brain_definition_service import ApiBrainDefinitionService @@ -42,7 +43,7 @@ class BrainService: brain_vector_repository: BrainsVectorsInterface external_api_secrets_repository: ExternalApiSecretsInterface composite_brains_connections_repository: CompositeBrainsConnectionsInterface - integration_brains_repository: IntegrationDescriptionInterface + integration_brains_repository: IntegrationBrainInterface integration_description_repository: IntegrationDescriptionInterface def __init__(self): @@ -324,7 +325,7 @@ class BrainService: 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: + def get_brain_details(self, brain_id: UUID, user_id: UUID) -> BrainEntity | None: brain = self.brain_repository.get_brain_details(brain_id) if brain == None: return None @@ -341,6 +342,17 @@ class BrainService: brain_id ) ) + if brain.brain_type == BrainType.INTEGRATION: + brain.integration = ( + self.integration_brains_repository.get_integration_brain( + brain_id, user_id + ) + ) + brain.integration_description = ( + self.integration_description_repository.get_integration_description( + brain.integration.integration_id + ) + ) return brain def get_connected_brains(self, brain_id: UUID) -> list[BrainEntity]: