add organization

api/core/app/apps/message_based_app_generator.py

@@ -4,8 +4,6 @@ from collections.abc import Generator
 from datetime import UTC, datetime
 from typing import Optional, Union, cast
 
-from sqlalchemy import and_
-
 from core.app.app_config.entities import EasyUIBasedAppConfig, EasyUIBasedAppModelConfigFrom
 from core.app.apps.base_app_generator import BaseAppGenerator
 from core.app.apps.base_app_queue_manager import AppQueueManager, GenerateTaskStoppedError
@@ -26,11 +24,12 @@ from core.app.entities.task_entities import (
 from core.app.task_pipeline.easy_ui_based_generate_task_pipeline import EasyUIBasedGenerateTaskPipeline
 from core.prompt.utils.prompt_template_parser import PromptTemplateParser
 from extensions.ext_database import db
-from models import Account
+from models import Account, EndUser
 from models.enums import CreatedByRole
-from models.model import App, AppMode, AppModelConfig, Conversation, EndUser, Message, MessageFile
+from models.model import App, AppMode, AppModelConfig, Conversation, Message, MessageFile
 from services.errors.app_model_config import AppModelConfigBrokenError
 from services.errors.conversation import ConversationCompletedError, ConversationNotExistsError
+from sqlalchemy import and_
 
 logger = logging.getLogger(__name__)
 
@@ -139,23 +138,33 @@ class MessageBasedAppGenerator(BaseAppGenerator):
         conversation: Optional[Conversation] = None,
     ) -> tuple[Conversation, Message]:
         """
-        Initialize generate records
-        :param application_generate_entity: application generate entity
-        :conversation conversation
-        :return:
+        Initialize generation records, including conversation and message
         """
-        app_config: EasyUIBasedAppConfig = cast(EasyUIBasedAppConfig, application_generate_entity.app_config)
+        app_config = cast(EasyUIBasedAppConfig, application_generate_entity.app_config)
 
-        # get from source
         end_user_id = None
         account_id = None
+        organization_id = None
+
         if application_generate_entity.invoke_from in {InvokeFrom.WEB_APP, InvokeFrom.SERVICE_API}:
             from_source = "api"
             end_user_id = application_generate_entity.user_id
+
+            # Get organization_id from end_user if available
+            if end_user_id:
+                end_user = db.session.query(EndUser).filter(EndUser.id == end_user_id).first()
+                if end_user and end_user.organization_id:
+                    organization_id = end_user.organization_id
         else:
             from_source = "console"
             account_id = application_generate_entity.user_id
 
+            # Get organization_id from account if available
+            if account_id:
+                account = db.session.query(Account).filter(Account.id == account_id).first()
+                if account and account.current_organization_id:
+                    organization_id = account.current_organization_id
+
         if isinstance(application_generate_entity, AdvancedChatAppGenerateEntity):
             app_model_config_id = None
             override_model_configs = None
@@ -179,6 +188,7 @@ class MessageBasedAppGenerator(BaseAppGenerator):
         if not conversation:
             conversation = Conversation(
                 app_id=app_config.app_id,
+                organization_id=organization_id,
                 app_model_config_id=app_model_config_id,
                 model_provider=model_provider,
                 model_id=model_id,
@@ -205,6 +215,7 @@ class MessageBasedAppGenerator(BaseAppGenerator):
 
         message = Message(
             app_id=app_config.app_id,
+            organization_id=organization_id,
             model_provider=model_provider,
             model_id=model_id,
             override_model_configs=json.dumps(override_model_configs) if override_model_configs else None,

api/migrations/versions/2025_03_24_2002-18dd49e03533_add_organization_support.py

@@ -0,0 +1,57 @@
+"""add organization support
+
+Revision ID: 18dd49e03533
+Revises: ceaf4dfed584
+Create Date: 2025-03-24 20:02:56.847845
+
+"""
+from alembic import op
+import models as models
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = '18dd49e03533'
+down_revision = 'ceaf4dfed584'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('accounts', schema=None) as batch_op:
+        batch_op.add_column(sa.Column('current_organization_id', models.types.StringUUID(), nullable=True))
+
+    with op.batch_alter_table('conversations', schema=None) as batch_op:
+        batch_op.add_column(sa.Column('organization_id', models.types.StringUUID(), nullable=True))
+        batch_op.create_index('conversation_organization_idx', ['organization_id'], unique=False)
+
+    with op.batch_alter_table('end_users', schema=None) as batch_op:
+        batch_op.add_column(sa.Column('organization_id', models.types.StringUUID(), nullable=True))
+        batch_op.create_index('end_user_organization_id_idx', ['organization_id'], unique=False)
+
+    with op.batch_alter_table('messages', schema=None) as batch_op:
+        batch_op.add_column(sa.Column('organization_id', models.types.StringUUID(), nullable=True))
+        batch_op.create_index('message_organization_id_idx', ['organization_id'], unique=False)
+
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('messages', schema=None) as batch_op:
+        batch_op.drop_index('message_organization_id_idx')
+        batch_op.drop_column('organization_id')
+
+    with op.batch_alter_table('end_users', schema=None) as batch_op:
+        batch_op.drop_index('end_user_organization_id_idx')
+        batch_op.drop_column('organization_id')
+
+    with op.batch_alter_table('conversations', schema=None) as batch_op:
+        batch_op.drop_index('conversation_organization_idx')
+        batch_op.drop_column('organization_id')
+
+    with op.batch_alter_table('accounts', schema=None) as batch_op:
+        batch_op.drop_column('current_organization_id')
+
+    # ### end Alembic commands ###

api/models/account.py

@@ -31,6 +31,7 @@ class Account(UserMixin, db.Model):  # type: ignore[name-defined]
     interface_language = db.Column(db.String(255))
     interface_theme = db.Column(db.String(255))
     timezone = db.Column(db.String(255))
+    current_organization_id = db.Column(StringUUID, nullable=True)  # Added for organization support
     last_login_at = db.Column(db.DateTime)
     last_login_ip = db.Column(db.String(255))
     last_active_at = db.Column(db.DateTime, nullable=False, server_default=func.current_timestamp())
@@ -38,6 +39,7 @@ class Account(UserMixin, db.Model):  # type: ignore[name-defined]
     initialized_at = db.Column(db.DateTime)
     created_at = db.Column(db.DateTime, nullable=False, server_default=func.current_timestamp())
     updated_at = db.Column(db.DateTime, nullable=False, server_default=func.current_timestamp())
+    _current_tenant = None  # Initialize to avoid AttributeError
 
     @property
     def is_password_set(self):
@@ -123,6 +125,70 @@ class Account(UserMixin, db.Model):  # type: ignore[name-defined]
     def is_dataset_operator(self):
         return self.current_role == TenantAccountRole.DATASET_OPERATOR
 
+    @property
+    def organizations(self):
+        """Get all organizations the account is a member of"""
+        from .organization import Organization, OrganizationMember
+
+        org_members = db.session.query(OrganizationMember).filter(OrganizationMember.account_id == self.id).all()
+
+        organization_ids = [om.organization_id for om in org_members]
+        if not organization_ids:
+            return []
+
+        return db.session.query(Organization).filter(Organization.id.in_(organization_ids)).all()
+
+    @property
+    def current_organization(self):
+        """Get the current organization for this account"""
+        if not self.current_organization_id:
+            return None
+
+        from .organization import Organization
+
+        return db.session.query(Organization).filter(Organization.id == self.current_organization_id).first()
+
+    @property
+    def current_org_role(self):
+        """Get the role in the current organization"""
+        if not self.current_organization_id:
+            return None
+
+        from .organization import OrganizationMember
+
+        member = (
+            db.session.query(OrganizationMember)
+            .filter(
+                OrganizationMember.organization_id == self.current_organization_id,
+                OrganizationMember.account_id == self.id,
+            )
+            .first()
+        )
+
+        if not member:
+            return None
+
+        return member.role
+
+    def is_org_admin(self, organization_id=None):
+        """Check if user is admin in the specified or current organization"""
+        from .organization import OrganizationMember, OrganizationRole
+
+        org_id = organization_id or self.current_organization_id
+        if not org_id:
+            return False
+
+        member = (
+            db.session.query(OrganizationMember)
+            .filter(OrganizationMember.organization_id == org_id, OrganizationMember.account_id == self.id)
+            .first()
+        )
+
+        if not member:
+            return False
+
+        return member.role == OrganizationRole.ADMIN
+
 
 class TenantStatus(enum.StrEnum):
     NORMAL = "normal"
@@ -197,6 +263,12 @@ class Tenant(db.Model):  # type: ignore[name-defined]
             .all()
         )
 
+    def get_organizations(self) -> list:
+        """Get all organizations under this tenant"""
+        from .organization import Organization
+
+        return db.session.query(Organization).filter(Organization.tenant_id == self.id).all()
+
     @property
     def custom_config_dict(self) -> dict:
         return json.loads(self.custom_config) if self.custom_config else {}

api/models/model.py

@@ -217,6 +217,19 @@ class App(db.Model):  # type: ignore[name-defined]
 
         return tags or []
 
+    @property
+    def organization_access(self):
+        """Get list of organizations with access to this app"""
+        from .organization import AppOrganizationAccess, Organization
+
+        access_records = db.session.query(AppOrganizationAccess).filter(AppOrganizationAccess.app_id == self.id).all()
+
+        if not access_records:
+            return []
+
+        organization_ids = [record.organization_id for record in access_records]
+        return db.session.query(Organization).filter(Organization.id.in_(organization_ids)).all()
+
 
 class AppModelConfig(db.Model):  # type: ignore[name-defined]
     __tablename__ = "app_model_configs"
@@ -527,10 +540,12 @@ class Conversation(db.Model):  # type: ignore[name-defined]
     __table_args__ = (
         db.PrimaryKeyConstraint("id", name="conversation_pkey"),
         db.Index("conversation_app_from_user_idx", "app_id", "from_source", "from_end_user_id"),
+        db.Index("conversation_organization_idx", "organization_id"),
     )
 
     id: Mapped[str] = mapped_column(StringUUID, server_default=db.text("uuid_generate_v4()"))
     app_id = db.Column(StringUUID, nullable=False)
+    organization_id = db.Column(StringUUID, nullable=True)
     app_model_config_id = db.Column(StringUUID, nullable=True)
     model_provider = db.Column(db.String(255), nullable=True)
     override_model_configs = db.Column(db.Text)
@@ -560,6 +575,16 @@ class Conversation(db.Model):  # type: ignore[name-defined]
 
     is_deleted = db.Column(db.Boolean, nullable=False, server_default=db.text("false"))
 
+    @property
+    def organization(self):
+        """Get the organization this conversation belongs to"""
+        if not self.organization_id:
+            return None
+
+        from .organization import Organization
+
+        return db.session.query(Organization).filter(Organization.id == self.organization_id).first()
+
     @property
     def inputs(self):
         inputs = self._inputs.copy()
@@ -767,10 +792,12 @@ class Message(db.Model):  # type: ignore[name-defined]
         db.Index("message_account_idx", "app_id", "from_source", "from_account_id"),
         db.Index("message_workflow_run_id_idx", "conversation_id", "workflow_run_id"),
         db.Index("message_created_at_idx", "created_at"),
+        db.Index("message_organization_id_idx", "organization_id"),
     )
 
     id: Mapped[str] = mapped_column(StringUUID, server_default=db.text("uuid_generate_v4()"))
     app_id = db.Column(StringUUID, nullable=False)
+    organization_id = db.Column(StringUUID, nullable=True)
     model_provider = db.Column(db.String(255), nullable=True)
     model_id = db.Column(db.String(255), nullable=True)
     override_model_configs = db.Column(db.Text)
@@ -801,6 +828,16 @@ class Message(db.Model):  # type: ignore[name-defined]
     agent_based = db.Column(db.Boolean, nullable=False, server_default=db.text("false"))
     workflow_run_id = db.Column(StringUUID)
 
+    @property
+    def organization(self):
+        """Get the organization this message belongs to"""
+        if not self.organization_id:
+            return None
+
+        from .organization import Organization
+
+        return db.session.query(Organization).filter(Organization.id == self.organization_id).first()
+
     @property
     def inputs(self):
         inputs = self._inputs.copy()
@@ -1312,10 +1349,12 @@ class EndUser(UserMixin, db.Model):  # type: ignore[name-defined]
         db.PrimaryKeyConstraint("id", name="end_user_pkey"),
         db.Index("end_user_session_id_idx", "session_id", "type"),
         db.Index("end_user_tenant_session_id_idx", "tenant_id", "session_id", "type"),
+        db.Index("end_user_organization_id_idx", "organization_id"),
     )
 
     id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()"))
     tenant_id = db.Column(StringUUID, nullable=False)
+    organization_id = db.Column(StringUUID, nullable=True)
     app_id = db.Column(StringUUID, nullable=True)
     type = db.Column(db.String(255), nullable=False)
     external_user_id = db.Column(db.String(255), nullable=True)

api/models/organization.py

@@ -0,0 +1,186 @@
+import enum
+import json
+from datetime import datetime
+
+from sqlalchemy import func
+from sqlalchemy.orm import Mapped, mapped_column
+
+from .engine import db
+from .types import StringUUID
+
+
+class OrganizationType(enum.StrEnum):
+    SCHOOL = "school"
+    UNIVERSITY = "university"
+    COMPANY = "company"
+    ORGANIZATION = "organization"
+
+
+class Organization(db.Model):  # type: ignore[name-defined]
+    """
+    Organization model to represent schools or companies under a single tenant.
+    This allows a single app provider (tenant) to serve multiple organizations
+    with separate data and configurations.
+    """
+
+    __tablename__ = "organizations"
+    __table_args__ = (
+        db.PrimaryKeyConstraint("id", name="organization_pkey"),
+        db.Index("organization_tenant_id_idx", "tenant_id"),
+        db.Index("organization_code_idx", "code"),
+    )
+
+    id: Mapped[str] = mapped_column(StringUUID, server_default=db.text("uuid_generate_v4()"))
+    tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)  # The owning tenant (app provider)
+    name: Mapped[str] = mapped_column(db.String(255), nullable=False)
+    code: Mapped[str] = mapped_column(db.String(64), nullable=False, unique=True)  # Unique code for the organization
+    description: Mapped[str] = mapped_column(db.Text, nullable=True)
+    type: Mapped[str] = mapped_column(db.String(64), nullable=False, default="school")
+    logo: Mapped[str] = mapped_column(db.String(255), nullable=True)
+    settings: Mapped[str] = mapped_column(db.Text, nullable=True)  # JSON settings
+    status: Mapped[str] = mapped_column(
+        db.String(16), nullable=False, server_default=db.text("'active'::character varying")
+    )
+    created_by: Mapped[str] = mapped_column(StringUUID, nullable=False)
+    created_at: Mapped[datetime] = mapped_column(db.DateTime, nullable=False, server_default=func.current_timestamp())
+    updated_at: Mapped[datetime] = mapped_column(db.DateTime, nullable=False, server_default=func.current_timestamp())
+
+    @property
+    def settings_dict(self) -> dict:
+        """Get organization settings as a dictionary"""
+        return json.loads(self.settings) if self.settings else {}
+
+    @settings_dict.setter
+    def settings_dict(self, value: dict):
+        """Set organization settings from a dictionary"""
+        self.settings = json.dumps(value)
+
+    @property
+    def allowed_email_domains(self) -> list[str]:
+        """Get list of allowed email domains for this organization"""
+        settings = self.settings_dict
+        return settings.get('allowed_email_domains', [])
+
+    @allowed_email_domains.setter
+    def allowed_email_domains(self, domains: list[str]):
+        """Set allowed email domains for this organization"""
+        settings = self.settings_dict
+        settings['allowed_email_domains'] = domains
+        self.settings_dict = settings
+
+    @property
+    def is_email_restricted(self) -> bool:
+        """Check if organization restricts registration by email domain"""
+        return len(self.allowed_email_domains) > 0
+
+    def validate_email(self, email: str) -> bool:
+        """Validate if an email is allowed for this organization"""
+        if not self.is_email_restricted:
+            return True
+
+        email_domain = email.split('@')[-1].lower()
+        return email_domain in self.allowed_email_domains
+
+    @property
+    def available_apps(self):
+        """Get apps available for this organization"""
+        app_access = (
+            db.session.query(AppOrganizationAccess).filter(AppOrganizationAccess.organization_id == self.id).all()
+        )
+
+        if not app_access:
+            return []
+
+        from .model import App
+
+        app_ids = [access.app_id for access in app_access]
+        return db.session.query(App).filter(App.id.in_(app_ids)).all()
+
+
+class OrganizationRole(enum.StrEnum):
+    """Roles within an organization (school/company)"""
+
+    ADMIN = "admin"  # Can manage the organization
+    TEACHER = "teacher"  # For educational orgs
+    STUDENT = "student"  # For educational orgs
+    STAFF = "staff"  # General staff
+    MANAGER = "manager"  # Department manager
+    EMPLOYEE = "employee"  # Regular employee
+    GUEST = "guest"  # Guest access
+
+    @property
+    def is_admin(self) -> bool:
+        return self == OrganizationRole.ADMIN
+
+    @property
+    def is_staff(self) -> bool:
+        return self in {
+            OrganizationRole.ADMIN,
+            OrganizationRole.TEACHER,
+            OrganizationRole.STAFF,
+            OrganizationRole.MANAGER,
+        }
+
+
+class OrganizationMember(db.Model):  # type: ignore[name-defined]
+    """Represents membership of an account in an organization"""
+
+    __tablename__ = "organization_members"
+    __table_args__ = (
+        db.PrimaryKeyConstraint("id", name="organization_member_pkey"),
+        db.Index("org_member_org_idx", "organization_id"),
+        db.Index("org_member_account_idx", "account_id"),
+        db.UniqueConstraint("organization_id", "account_id", name="unique_org_account"),
+    )
+
+    id: Mapped[str] = mapped_column(StringUUID, server_default=db.text("uuid_generate_v4()"))
+    organization_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
+    account_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
+    role: Mapped[str] = mapped_column(db.String(64), nullable=False)
+    department: Mapped[str] = mapped_column(db.String(255), nullable=True)
+    title: Mapped[str] = mapped_column(db.String(255), nullable=True)
+    is_default: Mapped[bool] = mapped_column(db.Boolean, nullable=False, server_default=db.text("false"))
+    metadata: Mapped[str] = mapped_column(db.Text, nullable=True)  # Additional metadata as JSON
+    created_by: Mapped[str] = mapped_column(StringUUID, nullable=False)
+    created_at: Mapped[datetime] = mapped_column(db.DateTime, nullable=False, server_default=func.current_timestamp())
+    updated_at: Mapped[datetime] = mapped_column(db.DateTime, nullable=False, server_default=func.current_timestamp())
+
+    @property
+    def metadata_dict(self) -> dict:
+        """Get member metadata as a dictionary"""
+        return json.loads(self.metadata) if self.metadata else {}
+
+    @metadata_dict.setter
+    def metadata_dict(self, value: dict):
+        """Set member metadata from a dictionary"""
+        self.metadata = json.dumps(value)
+
+
+class AppOrganizationAccess(db.Model):  # type: ignore[name-defined]
+    """Controls which apps are accessible to which organizations"""
+
+    __tablename__ = "app_organization_access"
+    __table_args__ = (
+        db.PrimaryKeyConstraint("id", name="app_organization_access_pkey"),
+        db.Index("app_org_access_app_idx", "app_id"),
+        db.Index("app_org_access_org_idx", "organization_id"),
+        db.UniqueConstraint("app_id", "organization_id", name="unique_app_organization"),
+    )
+
+    id: Mapped[str] = mapped_column(StringUUID, server_default=db.text("uuid_generate_v4()"))
+    app_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
+    organization_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
+    permissions: Mapped[str] = mapped_column(db.Text, nullable=True)  # JSON permissions
+    created_by: Mapped[str] = mapped_column(StringUUID, nullable=False)
+    created_at: Mapped[datetime] = mapped_column(db.DateTime, nullable=False, server_default=func.current_timestamp())
+    updated_at: Mapped[datetime] = mapped_column(db.DateTime, nullable=False, server_default=func.current_timestamp())
+
+    @property
+    def permissions_dict(self) -> dict:
+        """Get permissions as a dictionary"""
+        return json.loads(self.permissions) if self.permissions else {}
+
+    @permissions_dict.setter
+    def permissions_dict(self, value: dict):
+        """Set permissions from a dictionary"""
+        self.permissions = json.dumps(value)

api/services/conversation_service.py

@@ -2,9 +2,6 @@ from collections.abc import Callable, Sequence
 from datetime import UTC, datetime
 from typing import Optional, Union
 
-from sqlalchemy import asc, desc, func, or_, select
-from sqlalchemy.orm import Session
-
 from core.app.entities.app_invoke_entities import InvokeFrom
 from core.llm_generator.llm_generator import LLMGenerator
 from extensions.ext_database import db
@@ -13,6 +10,8 @@ from models.account import Account
 from models.model import App, Conversation, EndUser, Message
 from services.errors.conversation import ConversationNotExistsError, LastConversationNotExistsError
 from services.errors.message import MessageNotExistsError
+from sqlalchemy import asc, desc, func, or_, select
+from sqlalchemy.orm import Session
 
 
 class ConversationService:
@@ -33,6 +32,13 @@ class ConversationService:
         if not user:
             return InfiniteScrollPagination(data=[], limit=limit, has_more=False)
 
+        # Get organization_id if available
+        organization_id = None
+        if isinstance(user, EndUser) and user.organization_id:
+            organization_id = user.organization_id
+        elif isinstance(user, Account) and user.current_organization_id:
+            organization_id = user.current_organization_id
+
         stmt = select(Conversation).where(
             Conversation.is_deleted == False,
             Conversation.app_id == app_model.id,
@@ -41,6 +47,11 @@ class ConversationService:
             Conversation.from_account_id == (user.id if isinstance(user, Account) else None),
             or_(Conversation.invoke_from.is_(None), Conversation.invoke_from == invoke_from.value),
         )
+
+        # Add organization filter if available
+        if organization_id:
+            stmt = stmt.where(Conversation.organization_id == organization_id)
+
         if include_ids is not None:
             stmt = stmt.where(Conversation.id.in_(include_ids))
         if exclude_ids is not None:
@@ -141,9 +152,15 @@ class ConversationService:
 
     @classmethod
     def get_conversation(cls, app_model: App, conversation_id: str, user: Optional[Union[Account, EndUser]]):
-        conversation = (
-            db.session.query(Conversation)
-            .filter(
+        # Get organization_id if available
+        organization_id = None
+        if user:
+            if isinstance(user, EndUser) and user.organization_id:
+                organization_id = user.organization_id
+            elif isinstance(user, Account) and user.current_organization_id:
+                organization_id = user.current_organization_id
+
+        query = db.session.query(Conversation).filter(
             Conversation.id == conversation_id,
             Conversation.app_id == app_model.id,
             Conversation.from_source == ("api" if isinstance(user, EndUser) else "console"),
@@ -151,8 +168,12 @@ class ConversationService:
             Conversation.from_account_id == (user.id if isinstance(user, Account) else None),
             Conversation.is_deleted == False,
         )
-            .first()
-        )
+
+        # Add organization filter if available
+        if organization_id:
+            query = query.filter(Conversation.organization_id == organization_id)
+
+        conversation = query.first()
 
         if not conversation:
             raise ConversationNotExistsError()

api/services/message_service.py

@@ -192,17 +192,27 @@ class MessageService:
 
     @classmethod
     def get_message(cls, app_model: App, user: Optional[Union[Account, EndUser]], message_id: str):
-        message = (
-            db.session.query(Message)
-            .filter(
+        # Get organization_id if available
+        organization_id = None
+        if user:
+            if isinstance(user, EndUser) and user.organization_id:
+                organization_id = user.organization_id
+            elif isinstance(user, Account) and user.current_organization_id:
+                organization_id = user.current_organization_id
+
+        query = db.session.query(Message).filter(
             Message.id == message_id,
             Message.app_id == app_model.id,
             Message.from_source == ("api" if isinstance(user, EndUser) else "console"),
             Message.from_end_user_id == (user.id if isinstance(user, EndUser) else None),
             Message.from_account_id == (user.id if isinstance(user, Account) else None),
         )
-            .first()
-        )
+
+        # Add organization filter if available
+        if organization_id:
+            query = query.filter(Message.organization_id == organization_id)
+
+        message = query.first()
 
         if not message:
             raise MessageNotExistsError()