Add tier-gated RAG and Drive document sources (#42)

Ship allows_rag entitlement (founders/backer/pro/business), enforce it on
document APIs and RAG chat, harden ingest/delete/active lifecycle, and add
Google/Microsoft Drive connect + sync for personal and company knowledge bases.

Closes #43 #44 #45 #46 #47 #48 #49 #50 #51 #52 #53
Parent epic: #42
Related: #11

llm_be/chat_backend/services/chat_tenant_scope.py

@@ -91,6 +91,27 @@ def resolve_chat_user(
     return None
 
 
+def ensure_company_workspace(company) -> DocumentWorkspace:
+    """Return ``company``'s document workspace, creating a default one if missing (#46).
+
+    Never uses a bare ``.get(company=...)`` — ``.order_by("id").first()`` picks
+    a stable single workspace even if duplicates exist, and ``get_or_create``
+    closes the race for brand-new companies that don't have one yet (so
+    document upload/list/detail views never 404 just because a workspace was
+    never explicitly created).
+    """
+    workspace = (
+        DocumentWorkspace.objects.filter(company=company).order_by("id").first()
+    )
+    if workspace is not None:
+        return workspace
+    workspace, _ = DocumentWorkspace.objects.get_or_create(
+        company=company,
+        defaults={"name": "Default"},
+    )
+    return workspace
+
+
 def resolve_chat_company_scope(
     user: CustomUser,
     conversation_id: Optional[int] = None,
@@ -135,16 +156,7 @@ def resolve_chat_company_scope(
                 code="conversation_forbidden",
             )
 
-    workspace = (
-        DocumentWorkspace.objects.filter(company_id=user.company_id)
-        .order_by("id")
-        .first()
-    )
-    if workspace is None:
-        raise ChatTenantScopeError(
-            "No document workspace exists for this company.",
-            code="workspace_missing",
-        )
+    workspace = ensure_company_workspace(user.company)
 
     return ChatCompanyScope(
         user_id=user.id,

llm_be/chat_backend/services/drive_sync.py

@@ -0,0 +1,455 @@
+"""Google Drive / Microsoft OneDrive & SharePoint sync into RAG Documents (#48-#52).
+
+``sync_connection`` is the single entry point used by the API sync endpoint,
+the ``sync_drive_connections`` management command, and the webhook stubs.
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass, field
+from datetime import timedelta
+from typing import Any
+
+import httpx
+from django.conf import settings
+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.rag_services import AsyncRAGService
+
+logger = logging.getLogger(__name__)
+
+GOOGLE_TOKEN_URL = "https://oauth2.googleapis.com/token"
+GOOGLE_FILES_URL = "https://www.googleapis.com/drive/v3/files"
+GOOGLE_FOLDER_MIME = "application/vnd.google-apps.folder"
+# Google-native docs must be exported to a downloadable format (#48).
+GOOGLE_EXPORT_MIME_MAP: dict[str, tuple[str, str]] = {
+    "application/vnd.google-apps.document": (
+        "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
+        ".docx",
+    ),
+    "application/vnd.google-apps.spreadsheet": (
+        "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
+        ".xlsx",
+    ),
+    "application/vnd.google-apps.presentation": ("application/pdf", ".pdf"),
+}
+
+MICROSOFT_TOKEN_URL_TMPL = "https://login.microsoftonline.com/{tenant}/oauth2/v2.0/token"
+GRAPH_BASE_URL = "https://graph.microsoft.com/v1.0"
+
+HTTP_TIMEOUT_SECONDS = 30.0
+
+
+class DriveSyncError(Exception):
+    """Raised for any unrecoverable failure syncing one connection."""
+
+    def __init__(self, code: str, message: str = ""):
+        self.code = code
+        self.message = message or code
+        super().__init__(self.message)
+
+
+@dataclass
+class RemoteFile:
+    id: str
+    name: str
+    mime_type: str
+    etag: str
+    size: int | None = None
+    # Microsoft company (SharePoint) downloads are scoped to a site id.
+    context_id: str = field(default="")
+
+
+def _microsoft_tenant() -> str:
+    return settings.MICROSOFT_OAUTH_TENANT or "common"
+
+
+def _document_source(connection: DriveConnection) -> str:
+    if connection.provider == DriveConnection.Provider.GOOGLE:
+        return (
+            Document.Source.GOOGLE_SHARED_DRIVE
+            if connection.kind == DriveConnection.Kind.COMPANY
+            else Document.Source.GOOGLE_DRIVE
+        )
+    return (
+        Document.Source.SHAREPOINT
+        if connection.kind == DriveConnection.Kind.COMPANY
+        else Document.Source.ONEDRIVE
+    )
+
+
+# --- Token refresh ---------------------------------------------------------
+
+
+def refresh_google_token(connection: DriveConnection) -> str:
+    if not connection.refresh_token:
+        raise DriveSyncError(
+            "missing_refresh_token", "No refresh token stored for this connection."
+        )
+    with httpx.Client(timeout=HTTP_TIMEOUT_SECONDS) as client:
+        response = client.post(
+            GOOGLE_TOKEN_URL,
+            data={
+                "client_id": settings.GOOGLE_OAUTH_CLIENT_ID,
+                "client_secret": settings.GOOGLE_OAUTH_CLIENT_SECRET,
+                "refresh_token": connection.refresh_token,
+                "grant_type": "refresh_token",
+            },
+        )
+    if response.status_code >= 400:
+        raise DriveSyncError(
+            "token_refresh_failed", f"Google token refresh failed: {response.text}"
+        )
+    data = response.json()
+    access_token = data.get("access_token") or ""
+    if not access_token:
+        raise DriveSyncError(
+            "token_refresh_failed", "Google refresh did not return an access token."
+        )
+    connection.access_token = access_token
+    expires_in = data.get("expires_in")
+    if expires_in:
+        connection.token_expires_at = timezone.now() + timedelta(seconds=int(expires_in))
+    connection.save(update_fields=["access_token", "token_expires_at", "last_modified"])
+    return access_token
+
+
+def refresh_microsoft_token(connection: DriveConnection) -> str:
+    if not connection.refresh_token:
+        raise DriveSyncError(
+            "missing_refresh_token", "No refresh token stored for this connection."
+        )
+    token_url = MICROSOFT_TOKEN_URL_TMPL.format(tenant=_microsoft_tenant())
+    with httpx.Client(timeout=HTTP_TIMEOUT_SECONDS) as client:
+        response = client.post(
+            token_url,
+            data={
+                "client_id": settings.MICROSOFT_OAUTH_CLIENT_ID,
+                "client_secret": settings.MICROSOFT_OAUTH_CLIENT_SECRET,
+                "refresh_token": connection.refresh_token,
+                "grant_type": "refresh_token",
+                "scope": connection.scopes or "offline_access Files.Read.All",
+            },
+        )
+    if response.status_code >= 400:
+        raise DriveSyncError(
+            "token_refresh_failed", f"Microsoft token refresh failed: {response.text}"
+        )
+    data = response.json()
+    access_token = data.get("access_token") or ""
+    if not access_token:
+        raise DriveSyncError(
+            "token_refresh_failed", "Microsoft refresh did not return an access token."
+        )
+    connection.access_token = access_token
+    if data.get("refresh_token"):
+        connection.refresh_token = data["refresh_token"]
+    expires_in = data.get("expires_in")
+    if expires_in:
+        connection.token_expires_at = timezone.now() + timedelta(seconds=int(expires_in))
+    connection.save(
+        update_fields=["access_token", "refresh_token", "token_expires_at", "last_modified"]
+    )
+    return access_token
+
+
+def ensure_fresh_token(connection: DriveConnection) -> str:
+    """Return a usable access token, refreshing when expired or close to it."""
+    if (
+        connection.access_token
+        and connection.token_expires_at
+        and connection.token_expires_at > timezone.now() + timedelta(minutes=2)
+    ):
+        return connection.access_token
+    if connection.provider == DriveConnection.Provider.GOOGLE:
+        return refresh_google_token(connection)
+    return refresh_microsoft_token(connection)
+
+
+# --- Google Drive listing/download -----------------------------------------
+
+
+def _google_files_page(client: httpx.Client, headers: dict, params: dict) -> list[RemoteFile]:
+    response = client.get(GOOGLE_FILES_URL, headers=headers, params=params)
+    if response.status_code >= 400:
+        raise DriveSyncError("list_failed", f"Google Drive list failed: {response.text}")
+    data = response.json()
+    files: list[RemoteFile] = []
+    for item in data.get("files", []):
+        if item.get("mimeType") == GOOGLE_FOLDER_MIME:
+            continue
+        files.append(
+            RemoteFile(
+                id=item["id"],
+                name=item.get("name") or item["id"],
+                mime_type=item.get("mimeType", ""),
+                etag=item.get("md5Checksum") or item.get("modifiedTime") or "",
+                size=int(item["size"]) if item.get("size") else None,
+            )
+        )
+    return files
+
+
+def _list_google_files(connection: DriveConnection, access_token: str) -> list[RemoteFile]:
+    """List files under selected folders/shared drives, or "My Drive" root."""
+    headers = {"Authorization": f"Bearer {access_token}"}
+    parents = connection.selected_resource_ids or []
+    files: list[RemoteFile] = []
+    common_params = {
+        "fields": "files(id,name,mimeType,md5Checksum,size,modifiedTime)",
+        "pageSize": 100,
+        "supportsAllDrives": "true",
+        "includeItemsFromAllDrives": "true",
+    }
+    with httpx.Client(timeout=HTTP_TIMEOUT_SECONDS) as client:
+        if not parents:
+            params = {**common_params, "q": "'root' in parents and trashed = false"}
+            files.extend(_google_files_page(client, headers, params))
+        else:
+            for parent_id in parents:
+                params = {
+                    **common_params,
+                    "q": f"'{parent_id}' in parents and trashed = false",
+                }
+                if connection.kind == DriveConnection.Kind.COMPANY:
+                    params["corpora"] = "drive"
+                    params["driveId"] = parent_id
+                files.extend(_google_files_page(client, headers, params))
+    return files
+
+
+def _download_google_file(
+    client: httpx.Client, headers: dict, remote: RemoteFile
+) -> tuple[bytes, str, str]:
+    export = GOOGLE_EXPORT_MIME_MAP.get(remote.mime_type)
+    if export:
+        export_mime, suffix = export
+        filename = remote.name if remote.name.endswith(suffix) else f"{remote.name}{suffix}"
+        response = client.get(
+            f"{GOOGLE_FILES_URL}/{remote.id}/export",
+            headers=headers,
+            params={"mimeType": export_mime},
+        )
+        content_type = export_mime
+    else:
+        filename = remote.name
+        content_type = remote.mime_type or "application/octet-stream"
+        response = client.get(
+            f"{GOOGLE_FILES_URL}/{remote.id}",
+            headers=headers,
+            params={"alt": "media", "supportsAllDrives": "true"},
+        )
+    if response.status_code >= 400:
+        raise DriveSyncError(
+            "download_failed", f"Google Drive download failed: {response.text}"
+        )
+    return response.content, filename, content_type
+
+
+# --- Microsoft Graph listing/download --------------------------------------
+
+
+def _graph_children(client: httpx.Client, headers: dict, url: str) -> list[RemoteFile]:
+    files: list[RemoteFile] = []
+    while url:
+        response = client.get(url, headers=headers)
+        if response.status_code >= 400:
+            raise DriveSyncError(
+                "list_failed", f"Microsoft Graph list failed: {response.text}"
+            )
+        data = response.json()
+        for item in data.get("value", []):
+            if "folder" in item:
+                continue
+            files.append(
+                RemoteFile(
+                    id=item["id"],
+                    name=item.get("name") or item["id"],
+                    mime_type=(item.get("file") or {}).get("mimeType", ""),
+                    etag=item.get("eTag") or item.get("cTag") or "",
+                    size=item.get("size"),
+                )
+            )
+        url = data.get("@odata.nextLink")
+    return files
+
+
+def _list_microsoft_files(connection: DriveConnection, access_token: str) -> list[RemoteFile]:
+    headers = {"Authorization": f"Bearer {access_token}"}
+    selected = connection.selected_resource_ids or []
+    files: list[RemoteFile] = []
+    with httpx.Client(timeout=HTTP_TIMEOUT_SECONDS) as client:
+        if connection.kind == DriveConnection.Kind.COMPANY:
+            # Company (SharePoint) sync requires explicit site selection (#50/#51).
+            for site_id in selected:
+                url = f"{GRAPH_BASE_URL}/sites/{site_id}/drive/root/children"
+                for remote in _graph_children(client, headers, url):
+                    remote.context_id = site_id
+                    files.append(remote)
+        elif not selected:
+            url = f"{GRAPH_BASE_URL}/me/drive/root/children"
+            files.extend(_graph_children(client, headers, url))
+        else:
+            for folder_id in selected:
+                url = f"{GRAPH_BASE_URL}/me/drive/items/{folder_id}/children"
+                files.extend(_graph_children(client, headers, url))
+    return files
+
+
+def _download_microsoft_file(
+    client: httpx.Client, headers: dict, connection: DriveConnection, remote: RemoteFile
+) -> tuple[bytes, str, str]:
+    if connection.kind == DriveConnection.Kind.COMPANY and remote.context_id:
+        url = f"{GRAPH_BASE_URL}/sites/{remote.context_id}/drive/items/{remote.id}/content"
+    else:
+        url = f"{GRAPH_BASE_URL}/me/drive/items/{remote.id}/content"
+    response = client.get(url, headers=headers)
+    if response.status_code >= 400:
+        raise DriveSyncError(
+            "download_failed", f"Microsoft download failed: {response.text}"
+        )
+    content_type = response.headers.get("content-type", "application/octet-stream")
+    return response.content, remote.name, content_type
+
+
+def _download_file(
+    client: httpx.Client, headers: dict, connection: DriveConnection, remote: RemoteFile
+) -> tuple[bytes, str, str]:
+    if connection.provider == DriveConnection.Provider.GOOGLE:
+        return _download_google_file(client, headers, remote)
+    return _download_microsoft_file(client, headers, connection, remote)
+
+
+def _list_remote_files(connection: DriveConnection, access_token: str) -> list[RemoteFile]:
+    if connection.provider == DriveConnection.Provider.GOOGLE:
+        return _list_google_files(connection, access_token)
+    return _list_microsoft_files(connection, access_token)
+
+
+# --- Sync entry point --------------------------------------------------------
+
+
+def sync_connection(connection: DriveConnection) -> dict[str, Any]:
+    """Refresh tokens, list selected resources, and reconcile Documents (#48-#51).
+
+    Downloads (or exports, for Google-native docs) each remote file, creates/
+    updates the matching ``Document`` row, ingests it via ``AsyncRAGService``,
+    and removes ``Document`` rows whose remote file was deleted upstream.
+    """
+    result: dict[str, Any] = {"added": 0, "updated": 0, "removed": 0, "failed": []}
+
+    connection.last_sync_status = DriveConnection.SyncStatus.PENDING
+    connection.save(update_fields=["last_sync_status", "last_modified"])
+
+    try:
+        access_token = ensure_fresh_token(connection)
+        remote_files = _list_remote_files(connection, access_token)
+        workspace = ensure_company_workspace(connection.company)
+        source = _document_source(connection)
+
+        remote_by_id = {remote.id: remote for remote in remote_files}
+        existing_docs = {
+            document.remote_file_id: document
+            for document in Document.objects.filter(drive_connection=connection)
+        }
+
+        rag_service = AsyncRAGService()
+        headers = {"Authorization": f"Bearer {access_token}"}
+        with httpx.Client(timeout=60.0) as client:
+            for remote in remote_files:
+                existing = existing_docs.get(remote.id)
+                if existing is not None and existing.remote_etag == remote.etag:
+                    continue
+
+                try:
+                    content, filename, _content_type = _download_file(
+                        client, headers, connection, remote
+                    )
+                except DriveSyncError as exc:
+                    result["failed"].append({"file": remote.name, "error": exc.message})
+                    continue
+
+                if existing is not None:
+                    existing.file.save(filename, ContentFile(content), save=False)
+                    existing.remote_etag = remote.etag
+                    existing.remote_name = remote.name
+                    existing.sync_error = ""
+                    existing.processed = False
+                    existing.active = True
+                    existing.save()
+                    document = existing
+                    result["updated"] += 1
+                else:
+                    document = Document.objects.create(
+                        workspace=workspace,
+                        source=source,
+                        drive_connection=connection,
+                        remote_file_id=remote.id,
+                        remote_etag=remote.etag,
+                        remote_name=remote.name,
+                        active=True,
+                    )
+                    document.file.save(filename, ContentFile(content), save=True)
+                    result["added"] += 1
+
+                try:
+                    ingest_result = rag_service.add_files_to_store(
+                        [
+                            (
+                                document.file,
+                                document.file.name,
+                                workspace.id,
+                                document.id,
+                                document.active,
+                            )
+                        ],
+                        workspace_id=workspace.id,
+                        source=source,
+                    )
+                    document.processed = True
+                    document.sync_error = (
+                        str(ingest_result.get("failed_files"))
+                        if ingest_result.get("failed_files")
+                        else ""
+                    )
+                    document.save(update_fields=["processed", "sync_error", "last_modified"])
+                except Exception as exc:  # keep syncing remaining files
+                    logger.exception(
+                        "RAG ingest failed for document=%s connection=%s",
+                        document.id,
+                        connection.id,
+                    )
+                    document.sync_error = str(exc)
+                    document.save(update_fields=["sync_error", "last_modified"])
+                    result["failed"].append({"file": remote.name, "error": str(exc)})
+
+            for remote_id, document in existing_docs.items():
+                if remote_id not in remote_by_id:
+                    document.delete()
+                    result["removed"] += 1
+
+        connection.last_sync_status = DriveConnection.SyncStatus.OK
+        connection.last_sync_error = ""
+    except DriveSyncError as exc:
+        logger.warning(
+            "Drive sync failed for connection=%s (%s): %s", connection.id, exc.code, exc.message
+        )
+        connection.last_sync_status = DriveConnection.SyncStatus.ERROR
+        connection.last_sync_error = exc.message
+        result["error"] = exc.message
+    except Exception as exc:
+        logger.exception("Unexpected Drive sync failure for connection=%s", connection.id)
+        connection.last_sync_status = DriveConnection.SyncStatus.ERROR
+        connection.last_sync_error = str(exc)
+        result["error"] = str(exc)
+    finally:
+        connection.last_sync_at = timezone.now()
+        connection.save(
+            update_fields=["last_sync_status", "last_sync_error", "last_sync_at", "last_modified"]
+        )
+
+    return result

llm_be/chat_backend/services/rag_services.py

@@ -120,6 +120,7 @@ class RAGService(BaseService):
                         "workspace_id": doc.workspace_id,
                         "company_id": doc.workspace.company_id,
                         "document_id": doc.id,
+                        "active": bool(doc.active),
                     },
                 )
                 if chunks:
@@ -130,6 +131,38 @@ class RAGService(BaseService):
         self.vector_store.persist()
         return docs
 
+    def delete_document_vectors(self, document_id) -> None:
+        """Remove every vector chunk belonging to a document (#45).
+
+        Used on ``Document`` delete so a re-uploaded/replaced document does
+        not leave stale chunks searchable, without re-ingesting the whole
+        workspace (the old post_delete behavior).
+        """
+        if document_id is None:
+            return
+        self.vector_store.delete(where={"document_id": document_id})
+        self.vector_store.persist()
+
+    def set_document_active(self, document_id, active: bool) -> None:
+        """Update the ``active`` metadata flag on a document's existing chunks (#45).
+
+        Called from the PATCH toggle so ``search_documents`` (which filters
+        on ``active=True``) immediately reflects the new state without
+        re-ingesting the file.
+        """
+        if document_id is None:
+            return
+        existing = self.vector_store.get(where={"document_id": document_id})
+        ids = existing.get("ids") or []
+        if not ids:
+            return
+        metadatas = [
+            {**(meta or {}), "active": bool(active)}
+            for meta in existing.get("metadatas") or [{} for _ in ids]
+        ]
+        self.vector_store._collection.update(ids=ids, metadatas=metadatas)
+        self.vector_store.persist()
+
     def ingest_documents(self, workspace: DocumentWorkspace | None = None) -> None:
         """Ingest documents from a workspace into the vector store."""
         print(f"Getting the Document via the workspace: {workspace}")
@@ -167,7 +200,7 @@ class RAGService(BaseService):
 
     def add_files_to_store(
         self,
-        file_tupls: List,  # (file_path_or_field, name, workspace_id)
+        file_tupls: List,  # (file_path_or_field, name, workspace_id[, document_id, active])
         workspace_id: str,
         source: str = "upload",
         save_dir: str = "data/uploads",
@@ -175,7 +208,9 @@ class RAGService(BaseService):
         """
         Process and add files to vector store.
 
-        file_tupls entries: (path_str | Django FileField, original_name, workspace_id)
+        file_tupls entries: (path_str | Django FileField, original_name, workspace_id,
+        document_id, active). ``document_id``/``active`` are optional (default
+        ``None``/``True``) for backward compatibility with older call sites.
         Paths may be temp files; FileFields are materialized from DB storage.
         """
         results = {"total_added": 0, "failed_files": [], "processed_files": []}
@@ -188,6 +223,8 @@ class RAGService(BaseService):
                     file_tuple[1],
                     file_tuple[2],
                 )
+                document_id = file_tuple[3] if len(file_tuple) > 3 else None
+                active = file_tuple[4] if len(file_tuple) > 4 else True
                 if isinstance(file_ref, str):
                     file_path = file_ref
                 else:
@@ -207,6 +244,8 @@ class RAGService(BaseService):
                     "company_id": company_id,
                     "original_filename": original_name,
                     "file_path": original_name,
+                    "document_id": document_id,
+                    "active": bool(active),
                 }
 
                 docs = self._load_and_split_documents(file_path, metadata)
@@ -235,11 +274,18 @@ class RAGService(BaseService):
 
         ``company_id`` is written on ingest for defense-in-depth / future dual
         filters, but retrieval keys on ``workspace_id`` so older vectors without
-        ``company_id`` metadata still match after deploy.
+        ``company_id`` metadata still match after deploy. Also excludes chunks
+        for documents toggled inactive (#45) so deactivating a document hides
+        it from retrieval immediately.
         """
         if workspace is None or getattr(workspace, "id", None) is None:
             raise ValueError("workspace is required for RAG retrieval")
-        return {"workspace_id": workspace.id}
+        return {
+            "$and": [
+                {"workspace_id": workspace.id},
+                {"active": True},
+            ]
+        }
 
 
 class SyncRAGService(RAGService):