diff --git a/README.md b/README.md index fcc627c..a72cd9e 100644 --- a/README.md +++ b/README.md @@ -214,7 +214,9 @@ Subscription audit (`UserAuthEvent` on the user admin): ### Drive / RAG sync ([#47](https://git.aimloperations.com/ai_ml_operations/chat_backend/issues/47)-[#53](https://git.aimloperations.com/ai_ml_operations/chat_backend/issues/53)) Personal Google Drive / OneDrive and company Google Shared Drive / SharePoint -sync into the existing RAG `Document` pipeline. Every endpoint below is gated +sync into the existing RAG `Document` pipeline. Personal connections work for +users **without** a company (personal workspace); company connections require a +company manager. Every endpoint below is gated by `assert_feature_allowed(user, "rag")` (`SubscriptionPlan.allows_rag` — true for Founders/Pro/Business/Backer, false for Standard by default). diff --git a/llm_be/chat_backend/admin.py b/llm_be/chat_backend/admin.py index fd73c25..3a8852c 100644 --- a/llm_be/chat_backend/admin.py +++ b/llm_be/chat_backend/admin.py @@ -207,6 +207,7 @@ class DocumentWorkspaceAdmin(admin.ModelAdmin): list_display = ( "name", "company", + "user", ) diff --git a/llm_be/chat_backend/migrations/0029_personal_drive_rag_without_company.py b/llm_be/chat_backend/migrations/0029_personal_drive_rag_without_company.py new file mode 100644 index 0000000..fa00f13 --- /dev/null +++ b/llm_be/chat_backend/migrations/0029_personal_drive_rag_without_company.py @@ -0,0 +1,104 @@ +# Generated by Django 6.0 on 2026-08-02 10:39 + +import django.db.models.deletion +from django.conf import settings +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ("chat_backend", "0028_document_remote_etag_document_remote_file_id_and_more"), + ] + + operations = [ + migrations.RemoveConstraint( + model_name="driveconnection", + name="uniq_drive_connection_company_provider_kind_user", + ), + migrations.AddField( + model_name="documentworkspace", + name="user", + field=models.ForeignKey( + blank=True, + help_text="Set for personal RAG workspaces; null for company workspaces.", + null=True, + on_delete=django.db.models.deletion.CASCADE, + related_name="personal_workspaces", + to=settings.AUTH_USER_MODEL, + ), + ), + migrations.AlterField( + model_name="documentworkspace", + name="company", + field=models.ForeignKey( + blank=True, + null=True, + on_delete=django.db.models.deletion.CASCADE, + related_name="document_workspaces", + to="chat_backend.company", + ), + ), + migrations.AlterField( + model_name="driveconnection", + name="company", + field=models.ForeignKey( + blank=True, + help_text="Required for company connections; optional for personal (#55).", + null=True, + on_delete=django.db.models.deletion.CASCADE, + related_name="drive_connections", + to="chat_backend.company", + ), + ), + migrations.AddConstraint( + model_name="documentworkspace", + constraint=models.CheckConstraint( + condition=models.Q( + models.Q(("company__isnull", False), ("user__isnull", True)), + models.Q(("company__isnull", True), ("user__isnull", False)), + _connector="OR", + ), + name="document_workspace_company_xor_user", + ), + ), + migrations.AddConstraint( + model_name="documentworkspace", + constraint=models.UniqueConstraint( + condition=models.Q(("user__isnull", False)), + fields=("user",), + name="uniq_personal_document_workspace_user", + ), + ), + migrations.AddConstraint( + model_name="driveconnection", + constraint=models.CheckConstraint( + condition=models.Q( + models.Q(("kind", "personal"), ("user__isnull", False)), + models.Q( + ("company__isnull", False), + ("kind", "company"), + ("user__isnull", True), + ), + _connector="OR", + ), + name="drive_connection_kind_owner_consistency", + ), + ), + migrations.AddConstraint( + model_name="driveconnection", + constraint=models.UniqueConstraint( + condition=models.Q(("kind", "personal")), + fields=("user", "provider"), + name="uniq_personal_drive_connection_user_provider", + ), + ), + migrations.AddConstraint( + model_name="driveconnection", + constraint=models.UniqueConstraint( + condition=models.Q(("kind", "company")), + fields=("company", "provider"), + name="uniq_company_drive_connection_company_provider", + ), + ), + ] diff --git a/llm_be/chat_backend/models.py b/llm_be/chat_backend/models.py index d708054..08b704c 100644 --- a/llm_be/chat_backend/models.py +++ b/llm_be/chat_backend/models.py @@ -375,8 +375,45 @@ class PromptMetric(TimeInfoBase): # Document Models class DocumentWorkspace(TimeInfoBase): + """RAG document container: company (business) or user (personal) owned (#46, #55).""" + name = models.CharField(max_length=255) - company = models.ForeignKey(Company, on_delete=models.CASCADE) + company = models.ForeignKey( + Company, + on_delete=models.CASCADE, + null=True, + blank=True, + related_name="document_workspaces", + ) + user = models.ForeignKey( + "CustomUser", + on_delete=models.CASCADE, + null=True, + blank=True, + related_name="personal_workspaces", + help_text="Set for personal RAG workspaces; null for company workspaces.", + ) + + class Meta: + constraints = [ + models.CheckConstraint( + condition=( + models.Q(company__isnull=False, user__isnull=True) + | models.Q(company__isnull=True, user__isnull=False) + ), + name="document_workspace_company_xor_user", + ), + models.UniqueConstraint( + fields=["user"], + condition=models.Q(user__isnull=False), + name="uniq_personal_document_workspace_user", + ), + ] + + def __str__(self): + if self.user_id: + return f"DocumentWorkspace(personal user={self.user_id})" + return f"DocumentWorkspace(company={self.company_id})" class DriveConnection(TimeInfoBase): @@ -384,6 +421,9 @@ class DriveConnection(TimeInfoBase): ``user`` is null for company-only connections set up by a company manager (kind=company); personal connections always have ``user`` set. + + ``company`` is required for kind=company. Personal connections may have + ``company`` null when the user is not attached to a company (#55). """ class Provider(models.TextChoices): @@ -409,7 +449,12 @@ class DriveConnection(TimeInfoBase): help_text="Null for company-only connections owned by manager setup.", ) company = models.ForeignKey( - Company, on_delete=models.CASCADE, related_name="drive_connections" + Company, + on_delete=models.CASCADE, + null=True, + blank=True, + related_name="drive_connections", + help_text="Required for company connections; optional for personal (#55).", ) provider = models.CharField(max_length=32, choices=Provider.choices) kind = models.CharField( @@ -439,9 +484,26 @@ class DriveConnection(TimeInfoBase): class Meta: constraints = [ + models.CheckConstraint( + condition=( + models.Q(kind="personal", user__isnull=False) + | models.Q( + kind="company", + company__isnull=False, + user__isnull=True, + ) + ), + name="drive_connection_kind_owner_consistency", + ), models.UniqueConstraint( - fields=["company", "provider", "kind", "user"], - name="uniq_drive_connection_company_provider_kind_user", + fields=["user", "provider"], + condition=models.Q(kind="personal"), + name="uniq_personal_drive_connection_user_provider", + ), + models.UniqueConstraint( + fields=["company", "provider"], + condition=models.Q(kind="company"), + name="uniq_company_drive_connection_company_provider", ), ] diff --git a/llm_be/chat_backend/oauth.py b/llm_be/chat_backend/oauth.py index 92a3c77..c3ab900 100644 --- a/llm_be/chat_backend/oauth.py +++ b/llm_be/chat_backend/oauth.py @@ -337,22 +337,44 @@ def upsert_identity(user: CustomUser, profile: ProviderProfile) -> OAuthIdentity def upsert_drive_connection( *, user: CustomUser, kind: str, profile: ProviderProfile ) -> DriveConnection: - """Create/refresh a DriveConnection from a link_drive/link_company_drive callback (#47).""" + """Create/refresh a DriveConnection from a link_drive/link_company_drive callback (#47). + + Personal connections do not require a company (#55). Company connections do. + """ + if kind == DriveConnection.Kind.COMPANY and not user.company_id: + raise OAuthError( + "no_company", + "A company is required before connecting a company Drive.", + ) + expires_at = _token_expiry(profile.expires_in) - lookup_user = user if kind == DriveConnection.Kind.PERSONAL else None - connection = DriveConnection.objects.filter( - company=user.company, - provider=profile.provider, - kind=kind, - user=lookup_user, - ).first() - if connection is None: - connection = DriveConnection( + if kind == DriveConnection.Kind.PERSONAL: + connection = DriveConnection.objects.filter( + provider=profile.provider, + kind=DriveConnection.Kind.PERSONAL, + user=user, + ).first() + if connection is None: + connection = DriveConnection( + company=user.company, # may be None for solo users + provider=profile.provider, + kind=DriveConnection.Kind.PERSONAL, + user=user, + ) + else: + connection = DriveConnection.objects.filter( company=user.company, provider=profile.provider, - kind=kind, - user=lookup_user, - ) + kind=DriveConnection.Kind.COMPANY, + user=None, + ).first() + if connection is None: + connection = DriveConnection( + company=user.company, + provider=profile.provider, + kind=DriveConnection.Kind.COMPANY, + user=None, + ) connection.access_token = profile.access_token if profile.refresh_token: diff --git a/llm_be/chat_backend/services/chat_tenant_scope.py b/llm_be/chat_backend/services/chat_tenant_scope.py index 5bb0ecf..d086123 100644 --- a/llm_be/chat_backend/services/chat_tenant_scope.py +++ b/llm_be/chat_backend/services/chat_tenant_scope.py @@ -28,10 +28,13 @@ class ChatTenantScopeError(Exception): @dataclass(frozen=True) class ChatCompanyScope: - """Frozen tenant identity for one websocket turn / RAG retrieval.""" + """Frozen tenant identity for one websocket turn / RAG retrieval. + + ``company_id`` is null for solo users on a personal workspace (#55). + """ user_id: int - company_id: int + company_id: Optional[int] workspace_id: int conversation_id: Optional[int] = None @@ -100,18 +103,54 @@ def ensure_company_workspace(company) -> DocumentWorkspace: document upload/list/detail views never 404 just because a workspace was never explicitly created). """ + if company is None: + raise ChatTenantScopeError( + "Company is required for a company workspace.", + code="company_missing", + ) workspace = ( - DocumentWorkspace.objects.filter(company=company).order_by("id").first() + DocumentWorkspace.objects.filter(company=company, user__isnull=True) + .order_by("id") + .first() ) if workspace is not None: return workspace workspace, _ = DocumentWorkspace.objects.get_or_create( company=company, + user=None, defaults={"name": "Default"}, ) return workspace +def ensure_personal_workspace(user: CustomUser) -> DocumentWorkspace: + """Return ``user``'s personal RAG workspace, creating one if missing (#55).""" + if user is None or not getattr(user, "id", None): + raise ChatTenantScopeError( + "Authenticated chat user is required.", + code="user_not_found", + ) + workspace = ( + DocumentWorkspace.objects.filter(user=user, company__isnull=True) + .order_by("id") + .first() + ) + if workspace is not None: + return workspace + workspace, _ = DocumentWorkspace.objects.get_or_create( + user=user, + defaults={"name": "Personal", "company": None}, + ) + return workspace + + +def ensure_workspace_for_user(user: CustomUser) -> DocumentWorkspace: + """Company workspace when attached; otherwise personal workspace (#55).""" + if getattr(user, "company_id", None): + return ensure_company_workspace(user.company) + return ensure_personal_workspace(user) + + def resolve_chat_company_scope( user: CustomUser, conversation_id: Optional[int] = None, @@ -121,17 +160,13 @@ def resolve_chat_company_scope( When ``conversation_id`` is set, require ``conversation.user_id == user.id`` and that the conversation owner's company matches the user's company. + Users without a company resolve to a personal workspace (#55). """ if user is None or not getattr(user, "id", None): raise ChatTenantScopeError( "Authenticated chat user is required.", code="user_not_found", ) - if not getattr(user, "company_id", None): - raise ChatTenantScopeError( - "User is not attached to a company workspace.", - code="company_missing", - ) if conversation_id is not None: conversation = ( @@ -156,7 +191,7 @@ def resolve_chat_company_scope( code="conversation_forbidden", ) - workspace = ensure_company_workspace(user.company) + workspace = ensure_workspace_for_user(user) return ChatCompanyScope( user_id=user.id, @@ -175,8 +210,14 @@ def create_conversation_for_user(user: CustomUser, title: str) -> int: def get_workspace_for_scope(scope: ChatCompanyScope) -> DocumentWorkspace: """Load workspace rows only when they match the frozen scope keys.""" try: + if scope.company_id is not None: + return DocumentWorkspace.objects.get( + id=scope.workspace_id, company_id=scope.company_id + ) return DocumentWorkspace.objects.get( - id=scope.workspace_id, company_id=scope.company_id + id=scope.workspace_id, + user_id=scope.user_id, + company_id__isnull=True, ) except DocumentWorkspace.DoesNotExist as exc: raise ChatTenantScopeError( diff --git a/llm_be/chat_backend/services/drive_sync.py b/llm_be/chat_backend/services/drive_sync.py index 1b2f854..f2ee549 100644 --- a/llm_be/chat_backend/services/drive_sync.py +++ b/llm_be/chat_backend/services/drive_sync.py @@ -17,7 +17,10 @@ from django.core.files.base import ContentFile from django.utils import timezone from chat_backend.models import Document, DriveConnection -from chat_backend.services.chat_tenant_scope import ensure_company_workspace +from chat_backend.services.chat_tenant_scope import ( + ensure_company_workspace, + ensure_personal_workspace, +) from chat_backend.services.rag_services import AsyncRAGService logger = logging.getLogger(__name__) @@ -348,7 +351,20 @@ def sync_connection(connection: DriveConnection) -> dict[str, Any]: try: access_token = ensure_fresh_token(connection) remote_files = _list_remote_files(connection, access_token) - workspace = ensure_company_workspace(connection.company) + if connection.kind == DriveConnection.Kind.PERSONAL: + if connection.user_id is None: + raise DriveSyncError( + "invalid_connection", + "Personal Drive connections require an owning user.", + ) + workspace = ensure_personal_workspace(connection.user) + else: + if connection.company_id is None: + raise DriveSyncError( + "invalid_connection", + "Company Drive connections require a company.", + ) + workspace = ensure_company_workspace(connection.company) source = _document_source(connection) remote_by_id = {remote.id: remote for remote in remote_files} diff --git a/llm_be/chat_backend/tests/factories.py b/llm_be/chat_backend/tests/factories.py index 9f2a0cb..2a4f147 100644 --- a/llm_be/chat_backend/tests/factories.py +++ b/llm_be/chat_backend/tests/factories.py @@ -2,6 +2,8 @@ from __future__ import annotations +import uuid + from django.contrib.auth import get_user_model from django.core.files.uploadedfile import SimpleUploadedFile @@ -120,8 +122,10 @@ def make_prompt( ) -def make_workspace(company, name: str = "Test Workspace") -> DocumentWorkspace: - return DocumentWorkspace.objects.create(company=company, name=name) +def make_workspace(company=None, name: str = "Test Workspace", user=None) -> DocumentWorkspace: + if user is not None: + return DocumentWorkspace.objects.create(user=user, company=None, name=name) + return DocumentWorkspace.objects.create(company=company, user=None, name=name) def pdf_upload(name: str = "test.pdf") -> SimpleUploadedFile: @@ -133,7 +137,7 @@ def make_document(workspace, name: str = "test.pdf") -> Document: def make_drive_connection( - company, + company=None, *, provider: str = DriveConnection.Provider.GOOGLE, kind: str = DriveConnection.Kind.PERSONAL, @@ -148,6 +152,11 @@ def make_drive_connection( "is_active": True, } defaults.update(kwargs) + if kind == DriveConnection.Kind.PERSONAL and user is None: + suffix = uuid.uuid4().hex[:8] + user = make_user(email=f"drive-{provider}-{suffix}@example.com", company=company) + if kind == DriveConnection.Kind.COMPANY: + user = None return DriveConnection.objects.create( company=company, provider=provider, kind=kind, user=user, **defaults ) diff --git a/llm_be/chat_backend/tests/test_chat_tenant_scope.py b/llm_be/chat_backend/tests/test_chat_tenant_scope.py index b3274e9..0a8687a 100644 --- a/llm_be/chat_backend/tests/test_chat_tenant_scope.py +++ b/llm_be/chat_backend/tests/test_chat_tenant_scope.py @@ -6,6 +6,7 @@ from chat_backend.services.chat_tenant_scope import ( ChatCompanyScope, ChatTenantScopeError, ensure_company_workspace, + ensure_personal_workspace, resolve_chat_company_scope, resolve_chat_user, user_from_access_token, @@ -97,3 +98,37 @@ class EnsureCompanyWorkspaceTestCase(TestCase): workspace = DocumentWorkspace.objects.get(company=company) self.assertEqual(scope.workspace_id, workspace.id) + self.assertEqual(scope.company_id, company.id) + + +class EnsurePersonalWorkspaceTestCase(TestCase): + """#55: solo users get a personal workspace instead of company_missing.""" + + def test_creates_personal_workspace_for_user_without_company(self): + user = make_user(email="solo@example.com", company=None) + + workspace = ensure_personal_workspace(user) + + self.assertEqual(workspace.name, "Personal") + self.assertEqual(workspace.user_id, user.id) + self.assertIsNone(workspace.company_id) + self.assertEqual( + DocumentWorkspace.objects.filter(user=user, company__isnull=True).count(), + 1, + ) + + def test_scope_resolution_for_user_without_company(self): + user = make_user(email="solo-scope@example.com", company=None) + + scope = resolve_chat_company_scope(user) + + workspace = DocumentWorkspace.objects.get(user=user, company__isnull=True) + self.assertEqual( + scope, + ChatCompanyScope( + user_id=user.id, + company_id=None, + workspace_id=workspace.id, + conversation_id=None, + ), + ) diff --git a/llm_be/chat_backend/tests/test_oauth.py b/llm_be/chat_backend/tests/test_oauth.py index 3962a6d..902bb92 100644 --- a/llm_be/chat_backend/tests/test_oauth.py +++ b/llm_be/chat_backend/tests/test_oauth.py @@ -431,6 +431,49 @@ class OAuthCallbackDriveLinkTestCase(APITestCase): self.assertFalse(DriveConnection.objects.exists()) mock_exchange.assert_not_called() + @patch("chat_backend.views_oauth.exchange_code_for_profile") + def test_callback_personal_drive_allows_user_without_company(self, mock_exchange): + """#55 — personal Drive/RAG works for users not attached to a company.""" + self.user.company = None + self.user.save(update_fields=["company"]) + mock_exchange.return_value = _google_profile() + state = self._state(intent="link_drive") + + response = self.client.get( + reverse("oauth_callback", kwargs={"provider": "google"}), + {"code": "auth-code", "state": state}, + ) + + self.assertEqual(response.status_code, status.HTTP_302_FOUND) + params = parse_qs(urlparse(response["Location"]).query) + self.assertEqual(params["drive_connected"], ["1"]) + self.assertEqual(params["kind"], ["personal"]) + connection = DriveConnection.objects.get( + kind=DriveConnection.Kind.PERSONAL, user=self.user + ) + self.assertIsNone(connection.company_id) + self.assertEqual(connection.provider, "google") + self.assertTrue(connection.is_active) + + @patch("chat_backend.views_oauth.exchange_code_for_profile") + def test_callback_company_drive_rejects_user_without_company(self, mock_exchange): + """#55 — company Drive still requires a company.""" + self.user.company = None + self.user.is_company_manager = True + self.user.save(update_fields=["company", "is_company_manager"]) + mock_exchange.return_value = _google_profile() + state = self._state(intent="link_company_drive") + + response = self.client.get( + reverse("oauth_callback", kwargs={"provider": "google"}), + {"code": "auth-code", "state": state}, + ) + + params = parse_qs(urlparse(response["Location"]).query) + self.assertEqual(params["error"], ["no_company"]) + self.assertFalse(DriveConnection.objects.exists()) + mock_exchange.assert_not_called() + def test_callback_missing_user_id_in_state_is_rejected(self): # Simulates a forged/legacy state payload without the linking user. from django.core import signing diff --git a/llm_be/chat_backend/tests/test_services_drive_sync.py b/llm_be/chat_backend/tests/test_services_drive_sync.py index a71e6c6..3fab3b9 100644 --- a/llm_be/chat_backend/tests/test_services_drive_sync.py +++ b/llm_be/chat_backend/tests/test_services_drive_sync.py @@ -10,7 +10,7 @@ from django.utils import timezone from chat_backend.models import Document, DriveConnection from chat_backend.services import drive_sync -from chat_backend.services.chat_tenant_scope import ensure_company_workspace +from chat_backend.services.chat_tenant_scope import ensure_personal_workspace from chat_backend.services.drive_sync import ( DriveSyncError, RemoteFile, @@ -400,7 +400,7 @@ class SyncConnectionTestCase(TestCase): def test_skips_unchanged_file(self): Document.objects.create( - workspace=ensure_company_workspace(self.company), + workspace=ensure_personal_workspace(self.connection.user), drive_connection=self.connection, remote_file_id="r1", remote_etag="e1", @@ -418,7 +418,7 @@ class SyncConnectionTestCase(TestCase): def test_updates_file_when_etag_changes(self): existing = Document.objects.create( - workspace=ensure_company_workspace(self.company), + workspace=ensure_personal_workspace(self.connection.user), drive_connection=self.connection, remote_file_id="r1", remote_etag="old-etag", @@ -438,7 +438,7 @@ class SyncConnectionTestCase(TestCase): def test_removes_document_whose_remote_file_is_gone(self): Document.objects.create( - workspace=ensure_company_workspace(self.company), + workspace=ensure_personal_workspace(self.connection.user), drive_connection=self.connection, remote_file_id="deleted-remote", remote_etag="e1", @@ -477,12 +477,18 @@ class SyncConnectionTestCase(TestCase): self.assertEqual(self.connection.last_sync_status, DriveConnection.SyncStatus.ERROR) self.assertEqual(self.connection.last_sync_error, "quota exceeded") - def test_creates_workspace_when_company_has_none(self): + def test_creates_personal_workspace_when_user_has_none(self): from chat_backend.models import DocumentWorkspace - self.assertFalse(DocumentWorkspace.objects.filter(company=self.company).exists()) + self.assertFalse( + DocumentWorkspace.objects.filter(user=self.connection.user).exists() + ) self.mock_list.return_value = [] sync_connection(self.connection) - self.assertTrue(DocumentWorkspace.objects.filter(company=self.company).exists()) + self.assertTrue( + DocumentWorkspace.objects.filter( + user=self.connection.user, company__isnull=True + ).exists() + ) diff --git a/llm_be/chat_backend/views.py b/llm_be/chat_backend/views.py index e123591..a24b2d2 100644 --- a/llm_be/chat_backend/views.py +++ b/llm_be/chat_backend/views.py @@ -68,7 +68,7 @@ from .email_tasks import ( from finance.services.quotas import FeatureNotAllowed, assert_feature_allowed from .services.llm_service import AsyncLLMService from .services.rag_services import AsyncRAGService -from .services.chat_tenant_scope import ensure_company_workspace +from .services.chat_tenant_scope import ensure_workspace_for_user from .services.title_generator import title_generator from .services.moderation_classifier import moderation_classifier, ModerationLabel from .services.prompt_classifier.prompt_classifier import PromptClassifier, PromptType @@ -806,7 +806,14 @@ class DocumentWorkspaceView(APIView): assert_feature_allowed(request.user, "rag") except FeatureNotAllowed as exc: return _feature_gate_response(exc) - workspaces = DocumentWorkspace.objects.filter(company=request.user.company) + if request.user.company_id: + workspaces = DocumentWorkspace.objects.filter( + company=request.user.company, user__isnull=True + ) + else: + workspaces = DocumentWorkspace.objects.filter( + user=request.user, company__isnull=True + ) serializer = DocumentWorkspaceSerializer(workspaces, many=True) return Response(serializer.data) @@ -817,7 +824,10 @@ class DocumentWorkspaceView(APIView): return _feature_gate_response(exc) serializer = DocumentWorkspaceSerializer(data=request.data) if serializer.is_valid(): - serializer.save(company=request.user.company) + if request.user.company_id: + serializer.save(company=request.user.company, user=None) + else: + serializer.save(company=None, user=request.user) return Response(serializer.data, status=status.HTTP_201_CREATED) return Response(serializer.errors, status=status.HTTP_400_BAD_REQUEST) @@ -832,7 +842,7 @@ class DocumentUploadView(APIView): except FeatureNotAllowed as exc: return _feature_gate_response(exc) - workspace = ensure_company_workspace(request.user.company) + workspace = ensure_workspace_for_user(request.user) serializer = DocumentSerializer( Document.objects.filter(workspace=workspace), many=True ) @@ -846,7 +856,7 @@ class DocumentUploadView(APIView): except FeatureNotAllowed as exc: return _feature_gate_response(exc) - workspace = ensure_company_workspace(request.user.company) + workspace = ensure_workspace_for_user(request.user) logger.info(request.FILES) file = request.FILES.get("file") @@ -889,7 +899,7 @@ class DocumentDetailView(APIView): # permission_classes = [permissions.IsAuthenticated] def _get_document(self, request, document_id): - workspace = ensure_company_workspace(request.user.company) + workspace = ensure_workspace_for_user(request.user) return Document.objects.filter(workspace=workspace, id=document_id).first() def get(self, request, document_id): diff --git a/llm_be/chat_backend/views_oauth.py b/llm_be/chat_backend/views_oauth.py index de4104e..334e645 100644 --- a/llm_be/chat_backend/views_oauth.py +++ b/llm_be/chat_backend/views_oauth.py @@ -229,6 +229,11 @@ class OAuthCallbackView(APIView): "forbidden", "Only a company manager can connect a company Drive.", ) + if kind == DriveConnection.Kind.COMPANY and not user.company_id: + raise OAuthError( + "no_company", + "A company is required before connecting a company Drive.", + ) assert_feature_allowed(user, "rag") redirect_uri = _callback_redirect_uri(request, provider)