Personal Drive/RAG without a company (#55) (#56)

## Summary
- Closes [#55](#55) (related [#46](#46))
- **Personal** Google Drive / OneDrive works for users **without** a company (personal RAG)
- **Company** Drive still requires company + manager
- Schema: nullable `DriveConnection.company`, personal `DocumentWorkspace.user`, ownership check constraints + conditional uniques
- Runtime: `ensure_personal_workspace`, personal sync → personal WS, chat/document APIs fall back to personal WS when `company_id` is null
- Supersedes the interim "reject with `no_company`" approach (wrong for personal connect)

## Test plan
- [x] OAuth: personal Drive callback with `user.company=NULL` succeeds (`company_id=NULL` on connection)
- [x] OAuth: company Drive still returns `no_company` / `forbidden` appropriately
- [x] Chat tenant scope creates personal workspace for solo users
- [x] Drive sync + document view suites (`91` related tests)
- [ ] Manual: solo entitled user connects Google Drive → success, sync lands in personal workspace
- [ ] Manual: company manager company Drive still works
- [ ] Migrate prod/staging with `0029_personal_drive_rag_without_company`Reviewed-on: #56

llm_be/chat_backend/admin.py

@@ -207,6 +207,7 @@ class DocumentWorkspaceAdmin(admin.ModelAdmin):
     list_display = (
         "name",
         "company",
+        "user",
     )
 
 

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",
+            ),
+        ),
+    ]

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",
             ),
         ]
 

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:

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(

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}

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
     )

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,
+            ),
+        )

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

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()
+        )

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):

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)