diff --git a/.env.example b/.env.example index 7670d01..0a3ac97 100644 --- a/.env.example +++ b/.env.example @@ -32,6 +32,12 @@ ENABLE_ACCOUNT_REGISTRATION=false # Redirect URIs (register in each IdP console): # {OAUTH_CALLBACK_BASE_URL}/api/auth/oauth/google/callback/ # {OAUTH_CALLBACK_BASE_URL}/api/auth/oauth/microsoft/callback/ +# Same client id/secret pair is reused for Drive linking (#47) — the extra +# Drive scopes below are requested incrementally via intent=link_drive / +# intent=link_company_drive, no separate app registration needed: +# Google: openid email profile https://www.googleapis.com/auth/drive.readonly +# Microsoft: openid email profile offline_access Files.Read (personal) +# openid email profile offline_access Files.Read.All Sites.Read.All (company) GOOGLE_OAUTH_CLIENT_ID= GOOGLE_OAUTH_CLIENT_SECRET= MICROSOFT_OAUTH_CLIENT_ID= @@ -40,6 +46,22 @@ MICROSOFT_OAUTH_TENANT=common # Optional; defaults to request host. Example local: http://127.0.0.1:8001 OAUTH_CALLBACK_BASE_URL=http://127.0.0.1:8001 +# Drive / RAG sync (#47-#53). Requires a subscription plan with allows_rag +# (Founders, Pro, Business, Backer by default — see finance PLAN_SEED). +# Start a link: GET /api/auth/oauth/google/start/?intent=link_drive (authenticated) +# GET /api/auth/oauth/google/start/?intent=link_company_drive (company manager) +# GET /api/auth/oauth/microsoft/start/?intent=link_drive +# GET /api/auth/oauth/microsoft/start/?intent=link_company_drive +# Manage: GET /api/drive/connections/ +# DELETE /api/drive/connections// +# POST /api/drive/connections//resources/ { "resource_ids": [...] } +# POST /api/drive/connections//sync/ +# Provider push notifications (best-effort; register with each provider's +# subscription/watch API pointing here, using ?connection_id=): +# POST {OAUTH_CALLBACK_BASE_URL}/api/drive/webhooks/google/ +# POST {OAUTH_CALLBACK_BASE_URL}/api/drive/webhooks/microsoft/ +# Worker sync: `python manage.py sync_drive_connections [--connection-id N]` + # Stripe / finance (optional local — required for checkout + webhooks) STRIPE_SECRET_KEY= STRIPE_PUBLISHABLE_KEY= diff --git a/.env.prod.example b/.env.prod.example index 1c6321e..63f6750 100644 --- a/.env.prod.example +++ b/.env.prod.example @@ -53,6 +53,12 @@ ENABLE_ACCOUNT_REGISTRATION=false # Register redirect URIs: # https://chatbackend.aimloperations.com/api/auth/oauth/google/callback/ # https://chatbackend.aimloperations.com/api/auth/oauth/microsoft/callback/ +# Same client id/secret pair covers Drive linking (#47); no extra IdP app +# registration needed, but do register the Drive/Graph API + consent screen +# scopes below in each console (incremental scopes requested at intent time): +# Google: openid email profile https://www.googleapis.com/auth/drive.readonly +# Microsoft: openid email profile offline_access Files.Read (personal) +# openid email profile offline_access Files.Read.All Sites.Read.All (company) GOOGLE_OAUTH_CLIENT_ID= GOOGLE_OAUTH_CLIENT_SECRET= MICROSOFT_OAUTH_CLIENT_ID= @@ -60,6 +66,13 @@ MICROSOFT_OAUTH_CLIENT_SECRET= MICROSOFT_OAUTH_TENANT=common OAUTH_CALLBACK_BASE_URL=https://chatbackend.aimloperations.com +# Drive / RAG sync (#47-#53) — gated by SubscriptionPlan.allows_rag. +# Register provider push notifications (Google Drive `watch`, Microsoft +# Graph subscriptions) against: +# https://chatbackend.aimloperations.com/api/drive/webhooks/google/ +# https://chatbackend.aimloperations.com/api/drive/webhooks/microsoft/ +# Scheduled sync (cron / server-infra job): `python manage.py sync_drive_connections` + # Stripe / finance STRIPE_SECRET_KEY=replace-with-stripe-secret-key STRIPE_PUBLISHABLE_KEY=replace-with-stripe-publishable-key diff --git a/README.md b/README.md index f644d9e..fcc627c 100644 --- a/README.md +++ b/README.md @@ -93,6 +93,8 @@ with `COMPOSE_DATABASE_URL` if needed. | `ENABLE_ACCOUNT_REGISTRATION` | `false` | optional | Self-serve sign-up; keep false until ready | | `STRIPE_SECRET_KEY` / `STRIPE_PUBLISHABLE_KEY` / `STRIPE_WEBHOOK_SECRET` | empty | yes for billing | Stripe API + webhook | | `STRIPE_PRICE_ID` | empty | optional | Pre-created Price; else `$10/mo` from settings | +| `GOOGLE_OAUTH_CLIENT_ID` / `..._SECRET` | empty | for SSO/Drive | Also used for Drive linking (#47), incremental scopes | +| `MICROSOFT_OAUTH_CLIENT_ID` / `..._SECRET` / `..._TENANT` | empty / `common` | for SSO/Drive | Also used for Drive linking (#47), incremental scopes | | `FRONTEND_BASE_URL` | `http://localhost:3000` | set in prod/beta | Checkout success/cancel, portal return, OAuth return | | `STRIPE_PORTAL_RETURN_URL` | `{FRONTEND}/account/` | optional | Stripe Customer Portal return URL | | `CORS_ALLOWED_ORIGINS` | local + chat FE (+ beta FE default) | set in prod/beta | Frontend origin(s) | @@ -209,6 +211,48 @@ Subscription audit (`UserAuthEvent` on the user admin): - `subscription_started` — first active plan (Checkout, Backer redeem, admin assign) - `subscription_updated` — plan/status/cancel-at-period-end changes (portal + webhooks) +### 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 +by `assert_feature_allowed(user, "rag")` (`SubscriptionPlan.allows_rag` — +true for Founders/Pro/Business/Backer, false for Standard by default). + +**Connect (OAuth, reuses `#24` SSO app registrations with incremental scopes):** + +| | | +|--|--| +| Personal | `GET /api/auth/oauth//start/?intent=link_drive` (authenticated) | +| Company | `GET /api/auth/oauth//start/?intent=link_company_drive` (company manager only) | +| Callback | Same `/api/auth/oauth//callback/` as SSO; the signed OAuth `state` carries the linking `user_id` since the browser has no session on the IdP redirect. Upserts a `DriveConnection` and redirects to `{FRONTEND_BASE_URL}/account/?drive_connected=1&provider=&kind=` (or `?error=`) | + +**Manage:** + +| Method / path | Notes | +|--|--| +| `GET /api/drive/connections/` | Caller's personal connections + their company's company connections | +| `DELETE /api/drive/connections//` | Disconnect (owner for personal, company manager for company) — deactivates + clears tokens, keeps history | +| `POST /api/drive/connections//resources/` | `{ "resource_ids": [...], "resource_labels": [...] }` — folder/shared-drive/site ids to sync; empty = provider root | +| `POST /api/drive/connections//sync/` | Sync now (`chat_backend/services/drive_sync.py::sync_connection`) | + +**Provider scope differences:** +- Google: same `drive.readonly` scope for personal and company; company sync + reads Shared Drives via `corpora=drive` + `supportsAllDrives`. +- Microsoft: personal uses `Files.Read`; company uses `Files.Read.All + Sites.Read.All` and syncs SharePoint sites (`selected_resource_ids` = site ids). + +**Workers / webhooks (#52):** +- `python manage.py sync_drive_connections [--connection-id N]` — cron/worker entry point. +- `POST /api/drive/webhooks/google/` / `POST /api/drive/webhooks/microsoft/` — + provider push-notification stubs (`AllowAny`); acknowledge `200` and call + `sync_connection` when the notification's `connection_id` is resolvable, + else just `200` (no-op). Microsoft's subscription-creation `validationToken` + handshake is echoed back as `text/plain`. + +Google-native Docs/Sheets/Slides are exported to `.docx`/`.xlsx`/`.pdf` before +ingest (Chroma/RAG loaders don't read the native formats). Documents whose +remote file was deleted upstream are removed on the next sync. + ## Security note Secrets previously hardcoded in `settings.py` (email password, captcha, Django diff --git a/llm_be/chat_backend/admin.py b/llm_be/chat_backend/admin.py index 9760a5f..fd73c25 100644 --- a/llm_be/chat_backend/admin.py +++ b/llm_be/chat_backend/admin.py @@ -11,6 +11,7 @@ from .models import ( PromptMetric, DocumentWorkspace, Document, + DriveConnection, UserAuthEvent, OutboundEmail, OAuthIdentity, @@ -213,9 +214,36 @@ class DocumentAdmin(admin.ModelAdmin): model = Document list_display = ( "file", + "source", "active", "created", "processed", + "drive_connection", + ) + list_filter = ("source", "active", "processed") + raw_id_fields = ("drive_connection",) + + +class DriveConnectionAdmin(admin.ModelAdmin): + model = DriveConnection + list_display = ( + "provider", + "kind", + "company", + "user", + "external_account_email", + "is_active", + "last_sync_status", + "last_sync_at", + ) + list_filter = ("provider", "kind", "is_active", "last_sync_status") + search_fields = ("external_account_email", "company__name", "user__email") + raw_id_fields = ("company", "user") + readonly_fields = ( + "created", + "last_modified", + "access_token", + "refresh_token", ) @@ -233,6 +261,7 @@ admin.site.register(Feedback, FeedbackAdmin) admin.site.register(DocumentWorkspace, DocumentWorkspaceAdmin) admin.site.register(Document, DocumentAdmin) +admin.site.register(DriveConnection, DriveConnectionAdmin) class OAuthIdentityAdmin(admin.ModelAdmin): diff --git a/llm_be/chat_backend/consumers.py b/llm_be/chat_backend/consumers.py index 2dc900a..86a947d 100644 --- a/llm_be/chat_backend/consumers.py +++ b/llm_be/chat_backend/consumers.py @@ -497,6 +497,14 @@ class ChatConsumerAgain(AsyncWebsocketConsumer): pass if prompt_type == PromptType.RAG: + try: + await enforce_feature_gate(chat_user, "rag") + except FeatureNotAllowed as exc: + return { + "type": "error", + "code": exc.code, + "content": exc.message, + } service = AsyncRAGService() workspace = await get_workspace( conversation_id, user=chat_user diff --git a/llm_be/chat_backend/consumers_graph.py b/llm_be/chat_backend/consumers_graph.py index 10329e7..fda32f7 100644 --- a/llm_be/chat_backend/consumers_graph.py +++ b/llm_be/chat_backend/consumers_graph.py @@ -294,8 +294,19 @@ async def generation_node(state: ChatState) -> ChatState: pass if prompt_type == PromptType.RAG: - service = AsyncRAGService() chat_user = state.get("chat_user") + if chat_user is not None: + try: + await enforce_feature_gate(chat_user, "rag") + except FeatureNotAllowed as exc: + return { + "response_generator": { + "type": "error", + "code": exc.code, + "content": exc.message, + } + } + service = AsyncRAGService() workspace = await get_workspace(conversation_id, user=chat_user) generator = service.generate_response(messages, prompt_instance.message, workspace) return {"response_generator": generator} diff --git a/llm_be/chat_backend/management/__init__.py b/llm_be/chat_backend/management/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/llm_be/chat_backend/management/commands/__init__.py b/llm_be/chat_backend/management/commands/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/llm_be/chat_backend/management/commands/sync_drive_connections.py b/llm_be/chat_backend/management/commands/sync_drive_connections.py new file mode 100644 index 0000000..d5a65d5 --- /dev/null +++ b/llm_be/chat_backend/management/commands/sync_drive_connections.py @@ -0,0 +1,55 @@ +"""Worker entry point for scheduled Drive sync (#52). + +Usage: + python manage.py sync_drive_connections + python manage.py sync_drive_connections --connection-id 42 +""" + +from __future__ import annotations + +from django.core.management.base import BaseCommand, CommandError + +from chat_backend.models import DriveConnection +from chat_backend.services.drive_sync import sync_connection + + +class Command(BaseCommand): + help = "Sync active Drive connections (Google Drive / OneDrive / SharePoint) into Documents." + + def add_arguments(self, parser): + parser.add_argument( + "--connection-id", + type=int, + default=None, + help="Sync only the DriveConnection with this id.", + ) + + def handle(self, *args, **options): + connection_id = options.get("connection_id") + queryset = DriveConnection.objects.filter(is_active=True) + if connection_id is not None: + queryset = queryset.filter(id=connection_id) + + connections = list(queryset) + if not connections: + if connection_id is not None: + raise CommandError( + f"No active DriveConnection found with id={connection_id}." + ) + self.stdout.write("No active Drive connections to sync.") + return + + for connection in connections: + self.stdout.write( + f"Syncing connection={connection.id} " + f"provider={connection.provider} kind={connection.kind}..." + ) + result = sync_connection(connection) + if result.get("error"): + self.stderr.write(f" connection={connection.id} failed: {result['error']}") + else: + self.stdout.write( + f" connection={connection.id} added={result['added']} " + f"updated={result['updated']} removed={result['removed']} " + f"failed={len(result['failed'])}" + ) diff --git a/llm_be/chat_backend/migrations/0028_document_remote_etag_document_remote_file_id_and_more.py b/llm_be/chat_backend/migrations/0028_document_remote_etag_document_remote_file_id_and_more.py new file mode 100644 index 0000000..1f2e46e --- /dev/null +++ b/llm_be/chat_backend/migrations/0028_document_remote_etag_document_remote_file_id_and_more.py @@ -0,0 +1,169 @@ +# Generated by Django 6.0 on 2026-08-01 20:15 + +import django.db.models.deletion +import django.utils.timezone +from django.conf import settings +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ("chat_backend", "0027_user_auth_event_subscription_and_delete"), + ] + + operations = [ + migrations.AddField( + model_name="document", + name="remote_etag", + field=models.CharField(blank=True, default="", max_length=255), + ), + migrations.AddField( + model_name="document", + name="remote_file_id", + field=models.CharField(blank=True, default="", max_length=255), + ), + migrations.AddField( + model_name="document", + name="remote_name", + field=models.CharField(blank=True, default="", max_length=512), + ), + migrations.AddField( + model_name="document", + name="source", + field=models.CharField( + choices=[ + ("upload", "Upload"), + ("google_drive", "Google Drive"), + ("onedrive", "OneDrive"), + ("sharepoint", "SharePoint"), + ("google_shared_drive", "Google Shared Drive"), + ], + default="upload", + max_length=32, + ), + ), + migrations.AddField( + model_name="document", + name="sync_error", + field=models.TextField(blank=True, default=""), + ), + migrations.CreateModel( + name="DriveConnection", + fields=[ + ( + "id", + models.BigAutoField( + auto_created=True, + primary_key=True, + serialize=False, + verbose_name="ID", + ), + ), + ("created", models.DateTimeField(default=django.utils.timezone.now)), + ( + "last_modified", + models.DateTimeField(default=django.utils.timezone.now), + ), + ( + "provider", + models.CharField( + choices=[("google", "Google"), ("microsoft", "Microsoft")], + max_length=32, + ), + ), + ( + "kind", + models.CharField( + choices=[("personal", "Personal"), ("company", "Company")], + default="personal", + max_length=16, + ), + ), + ("access_token", models.TextField(blank=True, default="")), + ("refresh_token", models.TextField(blank=True, default="")), + ("token_expires_at", models.DateTimeField(blank=True, null=True)), + ("scopes", models.TextField(blank=True, default="")), + ( + "external_account_email", + models.EmailField(blank=True, default="", max_length=254), + ), + ( + "selected_resource_ids", + models.JSONField( + blank=True, + default=list, + help_text="Selected folder/drive/site ids to sync (empty = root/default).", + ), + ), + ( + "selected_resource_labels", + models.JSONField( + blank=True, + default=list, + help_text="Human-readable labels matching selected_resource_ids, for the FE.", + ), + ), + ("last_sync_at", models.DateTimeField(blank=True, null=True)), + ( + "last_sync_status", + models.CharField( + choices=[ + ("ok", "Ok"), + ("error", "Error"), + ("pending", "Pending"), + ("never", "Never"), + ], + default="never", + max_length=16, + ), + ), + ("last_sync_error", models.TextField(blank=True, default="")), + ("is_active", models.BooleanField(default=True)), + ( + "company", + models.ForeignKey( + on_delete=django.db.models.deletion.CASCADE, + related_name="drive_connections", + to="chat_backend.company", + ), + ), + ( + "user", + models.ForeignKey( + blank=True, + help_text="Null for company-only connections owned by manager setup.", + null=True, + on_delete=django.db.models.deletion.CASCADE, + related_name="drive_connections", + to=settings.AUTH_USER_MODEL, + ), + ), + ], + ), + migrations.AddField( + model_name="document", + name="drive_connection", + field=models.ForeignKey( + blank=True, + null=True, + on_delete=django.db.models.deletion.SET_NULL, + related_name="documents", + to="chat_backend.driveconnection", + ), + ), + migrations.AddIndex( + model_name="document", + index=models.Index( + fields=["drive_connection", "remote_file_id"], + name="chat_backen_drive_c_9332c2_idx", + ), + ), + migrations.AddConstraint( + model_name="driveconnection", + constraint=models.UniqueConstraint( + fields=("company", "provider", "kind", "user"), + name="uniq_drive_connection_company_provider_kind_user", + ), + ), + ] diff --git a/llm_be/chat_backend/models.py b/llm_be/chat_backend/models.py index 8a692a8..d708054 100644 --- a/llm_be/chat_backend/models.py +++ b/llm_be/chat_backend/models.py @@ -379,7 +379,87 @@ class DocumentWorkspace(TimeInfoBase): company = models.ForeignKey(Company, on_delete=models.CASCADE) +class DriveConnection(TimeInfoBase): + """A linked Google Drive / Microsoft OneDrive-SharePoint account (#47-#52). + + ``user`` is null for company-only connections set up by a company manager + (kind=company); personal connections always have ``user`` set. + """ + + class Provider(models.TextChoices): + GOOGLE = "google", "Google" + MICROSOFT = "microsoft", "Microsoft" + + class Kind(models.TextChoices): + PERSONAL = "personal", "Personal" + COMPANY = "company", "Company" + + class SyncStatus(models.TextChoices): + OK = "ok", "Ok" + ERROR = "error", "Error" + PENDING = "pending", "Pending" + NEVER = "never", "Never" + + user = models.ForeignKey( + "CustomUser", + on_delete=models.CASCADE, + null=True, + blank=True, + related_name="drive_connections", + help_text="Null for company-only connections owned by manager setup.", + ) + company = models.ForeignKey( + Company, on_delete=models.CASCADE, related_name="drive_connections" + ) + provider = models.CharField(max_length=32, choices=Provider.choices) + kind = models.CharField( + max_length=16, choices=Kind.choices, default=Kind.PERSONAL + ) + access_token = models.TextField(blank=True, default="") + refresh_token = models.TextField(blank=True, default="") + token_expires_at = models.DateTimeField(null=True, blank=True) + scopes = models.TextField(blank=True, default="") + external_account_email = models.EmailField(blank=True, default="") + selected_resource_ids = models.JSONField( + default=list, + blank=True, + help_text="Selected folder/drive/site ids to sync (empty = root/default).", + ) + selected_resource_labels = models.JSONField( + default=list, + blank=True, + help_text="Human-readable labels matching selected_resource_ids, for the FE.", + ) + last_sync_at = models.DateTimeField(null=True, blank=True) + last_sync_status = models.CharField( + max_length=16, choices=SyncStatus.choices, default=SyncStatus.NEVER + ) + last_sync_error = models.TextField(blank=True, default="") + is_active = models.BooleanField(default=True) + + class Meta: + constraints = [ + models.UniqueConstraint( + fields=["company", "provider", "kind", "user"], + name="uniq_drive_connection_company_provider_kind_user", + ), + ] + + def __str__(self): + return ( + f"DriveConnection({self.provider}/{self.kind}) " + f"company={self.company_id} user={self.user_id}" + ) + + class Document(TimeInfoBase): + class Source(models.TextChoices): + UPLOAD = "upload", "Upload" + GOOGLE_DRIVE = "google_drive", "Google Drive" + ONEDRIVE = "onedrive", "OneDrive" + SHAREPOINT = "sharepoint", "SharePoint" + GOOGLE_SHARED_DRIVE = "google_shared_drive", "Google Shared Drive" + workspace = models.ForeignKey(DocumentWorkspace, on_delete=models.CASCADE) file = models.FileField( upload_to="documents/", @@ -389,6 +469,25 @@ class Document(TimeInfoBase): uploaded_at = models.DateTimeField(auto_now_add=True) processed = models.BooleanField(default=False) active = models.BooleanField(default=False) + source = models.CharField( + max_length=32, choices=Source.choices, default=Source.UPLOAD + ) + remote_file_id = models.CharField(max_length=255, blank=True, default="") + remote_etag = models.CharField(max_length=255, blank=True, default="") + remote_name = models.CharField(max_length=512, blank=True, default="") + drive_connection = models.ForeignKey( + DriveConnection, + on_delete=models.SET_NULL, + null=True, + blank=True, + related_name="documents", + ) + sync_error = models.TextField(blank=True, default="") + + class Meta: + indexes = [ + models.Index(fields=["drive_connection", "remote_file_id"]), + ] class StoredFile(TimeInfoBase): diff --git a/llm_be/chat_backend/oauth.py b/llm_be/chat_backend/oauth.py index 0186010..92a3c77 100644 --- a/llm_be/chat_backend/oauth.py +++ b/llm_be/chat_backend/oauth.py @@ -14,7 +14,7 @@ from django.conf import settings from django.core import signing from django.utils import timezone -from .models import Company, CustomUser, OAuthIdentity +from .models import Company, CustomUser, DriveConnection, OAuthIdentity logger = logging.getLogger(__name__) @@ -24,6 +24,7 @@ STATE_MAX_AGE_SECONDS = 600 GOOGLE_AUTH_URL = "https://accounts.google.com/o/oauth2/v2/auth" GOOGLE_TOKEN_URL = "https://oauth2.googleapis.com/token" GOOGLE_USERINFO_URL = "https://openidconnect.googleapis.com/v1/userinfo" +GOOGLE_DRIVE_READONLY_SCOPE = "https://www.googleapis.com/auth/drive.readonly" MICROSOFT_AUTH_URL_TMPL = ( "https://login.microsoftonline.com/{tenant}/oauth2/v2.0/authorize" @@ -31,6 +32,13 @@ MICROSOFT_AUTH_URL_TMPL = ( MICROSOFT_TOKEN_URL_TMPL = ( "https://login.microsoftonline.com/{tenant}/oauth2/v2.0/token" ) +MICROSOFT_DRIVE_PERSONAL_SCOPE = "Files.Read" +MICROSOFT_DRIVE_COMPANY_SCOPE = "Files.Read.All Sites.Read.All" + +# OAuth intents (#24 login/signup; #47 Drive linking). +LOGIN_INTENTS = {"login", "signup"} +DRIVE_LINK_INTENTS = {"link_drive", "link_company_drive"} +VALID_INTENTS = LOGIN_INTENTS | DRIVE_LINK_INTENTS class OAuthError(Exception): @@ -76,14 +84,16 @@ def configured_providers() -> dict[str, bool]: } -def dump_oauth_state(*, provider: str, intent: str) -> str: - return signing.dumps( - {"provider": provider, "intent": intent}, - salt=STATE_SALT, - ) +def dump_oauth_state( + *, provider: str, intent: str, user_id: int | None = None +) -> str: + payload: dict[str, Any] = {"provider": provider, "intent": intent} + if user_id is not None: + payload["user_id"] = user_id + return signing.dumps(payload, salt=STATE_SALT) -def load_oauth_state(state: str) -> dict[str, str]: +def load_oauth_state(state: str) -> dict[str, Any]: try: data = signing.loads(state, salt=STATE_SALT, max_age=STATE_MAX_AGE_SECONDS) except signing.BadSignature as exc: @@ -92,25 +102,38 @@ def load_oauth_state(state: str) -> dict[str, str]: intent = data.get("intent") or "login" if provider not in OAuthIdentity.Provider.values: raise OAuthError("invalid_state", "Unknown OAuth provider in state.") - if intent not in {"login", "signup"}: + if intent not in VALID_INTENTS: raise OAuthError("invalid_state", "Invalid OAuth intent.") - return {"provider": provider, "intent": intent} + result: dict[str, Any] = {"provider": provider, "intent": intent} + if intent in DRIVE_LINK_INTENTS: + user_id = data.get("user_id") + if not user_id: + raise OAuthError( + "invalid_state", "OAuth state is missing the linking user." + ) + result["user_id"] = user_id + return result def _microsoft_tenant() -> str: return settings.MICROSOFT_OAUTH_TENANT or "common" -def build_authorization_url(*, provider: str, redirect_uri: str, state: str) -> str: +def build_authorization_url( + *, provider: str, redirect_uri: str, state: str, intent: str = "login" +) -> str: if not provider_configured(provider): raise OAuthError("provider_not_configured", f"{provider} OAuth is not configured.") if provider == OAuthIdentity.Provider.GOOGLE: + scope = "openid email profile" + if intent in DRIVE_LINK_INTENTS: + scope = f"{scope} {GOOGLE_DRIVE_READONLY_SCOPE}" params = { "client_id": settings.GOOGLE_OAUTH_CLIENT_ID, "redirect_uri": redirect_uri, "response_type": "code", - "scope": "openid email profile", + "scope": scope, "state": state, "access_type": "offline", "prompt": "select_account consent", @@ -119,12 +142,17 @@ def build_authorization_url(*, provider: str, redirect_uri: str, state: str) -> return f"{GOOGLE_AUTH_URL}?{urlencode(params)}" if provider == OAuthIdentity.Provider.MICROSOFT: + scope = "openid email profile offline_access" + if intent == "link_drive": + scope = f"{scope} {MICROSOFT_DRIVE_PERSONAL_SCOPE}" + elif intent == "link_company_drive": + scope = f"{scope} {MICROSOFT_DRIVE_COMPANY_SCOPE}" params = { "client_id": settings.MICROSOFT_OAUTH_CLIENT_ID, "redirect_uri": redirect_uri, "response_type": "code", "response_mode": "query", - "scope": "openid email profile offline_access", + "scope": scope, "state": state, "prompt": "select_account", } @@ -306,6 +334,38 @@ 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).""" + 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( + company=user.company, + provider=profile.provider, + kind=kind, + user=lookup_user, + ) + + connection.access_token = profile.access_token + if profile.refresh_token: + connection.refresh_token = profile.refresh_token + connection.token_expires_at = expires_at + connection.scopes = profile.scopes + connection.external_account_email = profile.email + connection.is_active = True + connection.last_sync_error = "" + connection.save() + return connection + + def _create_sso_user(profile: ProviderProfile) -> CustomUser: from finance.services.plans import try_redeem_backer_email @@ -329,6 +389,14 @@ def _create_sso_user(profile: ProviderProfile) -> CustomUser: return user +def resolve_link_user(user_id: int) -> CustomUser: + """Load the authenticated user a Drive-link callback should attach to.""" + user = CustomUser.objects.filter(pk=user_id, deleted=False).first() + if user is None: + raise OAuthError("user_not_found", "Linking user account was not found.") + return user + + def resolve_user_from_profile( *, profile: ProviderProfile, intent: str ) -> tuple[CustomUser, bool]: diff --git a/llm_be/chat_backend/serializers.py b/llm_be/chat_backend/serializers.py index 66490a6..a9f1c24 100644 --- a/llm_be/chat_backend/serializers.py +++ b/llm_be/chat_backend/serializers.py @@ -13,6 +13,7 @@ from .models import ( FEEDBACK_CATEGORIES, DocumentWorkspace, Document, + DriveConnection, ) @@ -279,3 +280,37 @@ class DocumentSerializer(serializers.ModelSerializer): "active", ] read_only_fields = ["id", "uploaded_at", "processed", "created"] + + +# drive connection serializers (#47-#52) +class DriveConnectionSerializer(serializers.ModelSerializer): + """Never exposes access_token/refresh_token to the client.""" + + class Meta: + model = DriveConnection + fields = [ + "id", + "provider", + "kind", + "external_account_email", + "selected_resource_ids", + "selected_resource_labels", + "last_sync_at", + "last_sync_status", + "last_sync_error", + "is_active", + "created", + ] + read_only_fields = fields + + +class DriveConnectionResourcesSerializer(serializers.Serializer): + resource_ids = serializers.ListField( + child=serializers.CharField(max_length=512), allow_empty=True + ) + resource_labels = serializers.ListField( + child=serializers.CharField(max_length=512, allow_blank=True), + allow_empty=True, + required=False, + default=list, + ) diff --git a/llm_be/chat_backend/services/chat_tenant_scope.py b/llm_be/chat_backend/services/chat_tenant_scope.py index aa9d9e2..5bb0ecf 100644 --- a/llm_be/chat_backend/services/chat_tenant_scope.py +++ b/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, diff --git a/llm_be/chat_backend/services/drive_sync.py b/llm_be/chat_backend/services/drive_sync.py new file mode 100644 index 0000000..1b2f854 --- /dev/null +++ b/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 diff --git a/llm_be/chat_backend/services/rag_services.py b/llm_be/chat_backend/services/rag_services.py index 01d6dd4..0290a88 100644 --- a/llm_be/chat_backend/services/rag_services.py +++ b/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): diff --git a/llm_be/chat_backend/signals.py b/llm_be/chat_backend/signals.py index a954252..b624768 100644 --- a/llm_be/chat_backend/signals.py +++ b/llm_be/chat_backend/signals.py @@ -1,6 +1,5 @@ -from django.db.models.signals import post_save, post_delete +from django.db.models.signals import post_delete from django.dispatch import receiver -from django.conf import settings import os from chat_backend.models import Document @@ -10,31 +9,22 @@ def _rag_init_skipped() -> bool: return os.environ.get("SKIP_RAG_INIT", "").lower() in {"1", "true", "yes"} -@receiver(post_save, sender=Document) -def update_vector_on_save(sender, instance, **kwargs): - """Update vector store when documents are saved""" - if _rag_init_skipped(): - return - if not kwargs.get("created", False): - return - try: - from .services.rag_services import AsyncRAGService - - rag_service = AsyncRAGService() - rag_service.ingest_documents() - except Exception as exc: - print(f"Skipping vector update on Document save: {exc}") - - @receiver(post_delete, sender=Document) def delete_vector_on_remove(sender, instance, **kwargs): - """Handle document deletion by re-indexing the whole workspace""" + """Remove the deleted document's chunks from the vector store (#45). + + There is intentionally no post_save handler: ``DocumentUploadView`` + already calls ``add_files_to_store`` for the uploaded file, so a + post_save re-ingest would duplicate that work (and, previously, + re-ingested the *entire* workspace on every save). Deletion only needs to + drop that document's own chunks, not rebuild everything else. + """ if _rag_init_skipped(): return try: from .services.rag_services import AsyncRAGService rag_service = AsyncRAGService() - rag_service.ingest_documents() + rag_service.delete_document_vectors(instance.id) except Exception as exc: - print(f"Skipping vector update on Document delete: {exc}") + print(f"Skipping vector cleanup on Document delete: {exc}") diff --git a/llm_be/chat_backend/tests/factories.py b/llm_be/chat_backend/tests/factories.py index 350460a..9f2a0cb 100644 --- a/llm_be/chat_backend/tests/factories.py +++ b/llm_be/chat_backend/tests/factories.py @@ -10,6 +10,7 @@ from chat_backend.models import ( Conversation, Document, DocumentWorkspace, + DriveConnection, Prompt, ) @@ -129,3 +130,24 @@ def pdf_upload(name: str = "test.pdf") -> SimpleUploadedFile: def make_document(workspace, name: str = "test.pdf") -> Document: return Document.objects.create(workspace=workspace, file=pdf_upload(name)) + + +def make_drive_connection( + company, + *, + provider: str = DriveConnection.Provider.GOOGLE, + kind: str = DriveConnection.Kind.PERSONAL, + user=None, + **kwargs, +) -> DriveConnection: + defaults = { + "access_token": "access-token", + "refresh_token": "refresh-token", + "scopes": "openid email profile", + "external_account_email": "drive.user@example.com", + "is_active": True, + } + defaults.update(kwargs) + 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 848b1d2..b3274e9 100644 --- a/llm_be/chat_backend/tests/test_chat_tenant_scope.py +++ b/llm_be/chat_backend/tests/test_chat_tenant_scope.py @@ -1,9 +1,11 @@ from django.test import TestCase from rest_framework_simplejwt.tokens import RefreshToken +from chat_backend.models import DocumentWorkspace from chat_backend.services.chat_tenant_scope import ( ChatCompanyScope, ChatTenantScopeError, + ensure_company_workspace, resolve_chat_company_scope, resolve_chat_user, user_from_access_token, @@ -62,3 +64,36 @@ class ChatTenantScopeTestCase(TestCase): def test_user_from_access_token_rejects_garbage(self): self.assertIsNone(user_from_access_token("not-a-jwt")) + + +class EnsureCompanyWorkspaceTestCase(TestCase): + """#46: never 404 a company just because no workspace was created yet.""" + + def test_returns_existing_workspace_without_creating_another(self): + company = make_company() + workspace = make_workspace(company) + + found = ensure_company_workspace(company) + + self.assertEqual(found.id, workspace.id) + self.assertEqual(DocumentWorkspace.objects.filter(company=company).count(), 1) + + def test_creates_default_workspace_when_missing(self): + company = make_company("NoWorkspaceCo") + + self.assertFalse(DocumentWorkspace.objects.filter(company=company).exists()) + + created = ensure_company_workspace(company) + + self.assertEqual(created.name, "Default") + self.assertEqual(created.company_id, company.id) + self.assertEqual(DocumentWorkspace.objects.filter(company=company).count(), 1) + + def test_scope_resolution_creates_workspace_instead_of_raising(self): + company = make_company("FreshCo") + user = make_user(email="fresh@example.com", company=company) + + scope = resolve_chat_company_scope(user) + + workspace = DocumentWorkspace.objects.get(company=company) + self.assertEqual(scope.workspace_id, workspace.id) diff --git a/llm_be/chat_backend/tests/test_consumers.py b/llm_be/chat_backend/tests/test_consumers.py index fd03501..961af65 100644 --- a/llm_be/chat_backend/tests/test_consumers.py +++ b/llm_be/chat_backend/tests/test_consumers.py @@ -367,6 +367,60 @@ class GraphNodeTestCase(TransactionTestCase): _args, _kwargs = service.return_value.generate_response.call_args self.assertEqual(_args[2].id, self.workspace.id) + @override_settings(ENFORCE_SUBSCRIPTION_GATES=True) + async def test_generation_node_denies_rag_without_an_active_plan(self): + """#44: RAG turns must respect the plan's ``rag`` feature gate.""" + with mock.patch.object(consumers_graph, "AsyncRAGService") as service: + result = await consumers_graph.generation_node( + self._state(prompt_type=PromptType.RAG) + ) + + payload = result["response_generator"] + self.assertEqual(payload["type"], "error") + self.assertEqual(payload["code"], "subscription_required") + service.assert_not_called() + + @override_settings(ENFORCE_SUBSCRIPTION_GATES=True) + async def test_generation_node_denies_rag_on_standard_plan(self): + from finance.models import SubscriptionPlan, UserSubscription + from finance.services.plans import assign_plan, seed_subscription_plans + + await sync_to_async(seed_subscription_plans)() + standard = await sync_to_async(SubscriptionPlan.objects.get)(slug="standard") + await sync_to_async(assign_plan)( + self.user, plan=standard, source=UserSubscription.Source.ADMIN + ) + + with mock.patch.object(consumers_graph, "AsyncRAGService") as service: + result = await consumers_graph.generation_node( + self._state(prompt_type=PromptType.RAG) + ) + + payload = result["response_generator"] + self.assertEqual(payload["type"], "error") + self.assertEqual(payload["code"], "feature_not_allowed") + service.assert_not_called() + + @override_settings(ENFORCE_SUBSCRIPTION_GATES=True) + async def test_generation_node_allows_rag_with_founders_plan(self): + from finance.models import SubscriptionPlan, UserSubscription + from finance.services.plans import assign_plan, seed_subscription_plans + + await sync_to_async(seed_subscription_plans)() + founders = await sync_to_async(SubscriptionPlan.objects.get)(slug="founders") + await sync_to_async(assign_plan)( + self.user, plan=founders, source=UserSubscription.Source.ADMIN + ) + + with mock.patch.object(consumers_graph, "AsyncRAGService") as service: + service.return_value.generate_response.return_value = "generator" + + result = await consumers_graph.generation_node( + self._state(prompt_type=PromptType.RAG) + ) + + self.assertEqual(result["response_generator"], "generator") + async def test_generation_node_defaults_to_general_chat(self): with mock.patch.object(consumers_graph, "AsyncLLMService") as service: service.return_value.generate_response.return_value = "generator" diff --git a/llm_be/chat_backend/tests/test_management_drive_sync.py b/llm_be/chat_backend/tests/test_management_drive_sync.py new file mode 100644 index 0000000..5118975 --- /dev/null +++ b/llm_be/chat_backend/tests/test_management_drive_sync.py @@ -0,0 +1,55 @@ +"""Tests for the sync_drive_connections management command (#52).""" + +from __future__ import annotations + +from io import StringIO +from unittest import mock + +from django.core.management import CommandError, call_command +from django.test import TestCase + +from chat_backend.models import DriveConnection + +from .factories import make_company, make_drive_connection + + +class SyncDriveConnectionsCommandTestCase(TestCase): + def setUp(self): + self.company = make_company() + + @mock.patch("chat_backend.management.commands.sync_drive_connections.sync_connection") + def test_syncs_all_active_connections(self, mock_sync): + mock_sync.return_value = {"added": 1, "updated": 0, "removed": 0, "failed": []} + active = make_drive_connection(self.company) + make_drive_connection( + self.company, + provider=DriveConnection.Provider.MICROSOFT, + is_active=False, + ) + + out = StringIO() + call_command("sync_drive_connections", stdout=out) + + mock_sync.assert_called_once_with(active) + self.assertIn("added=1", out.getvalue()) + + @mock.patch("chat_backend.management.commands.sync_drive_connections.sync_connection") + def test_syncs_single_connection_by_id(self, mock_sync): + mock_sync.return_value = {"added": 0, "updated": 1, "removed": 0, "failed": []} + target = make_drive_connection(self.company) + other = make_drive_connection( + self.company, provider=DriveConnection.Provider.MICROSOFT + ) + + call_command("sync_drive_connections", "--connection-id", str(target.id), stdout=StringIO()) + + mock_sync.assert_called_once_with(target) + + def test_unknown_connection_id_raises(self): + with self.assertRaises(CommandError): + call_command("sync_drive_connections", "--connection-id", "999999", stdout=StringIO()) + + def test_no_connections_reports_and_exits_cleanly(self): + out = StringIO() + call_command("sync_drive_connections", stdout=out) + self.assertIn("No active Drive connections", out.getvalue()) diff --git a/llm_be/chat_backend/tests/test_oauth.py b/llm_be/chat_backend/tests/test_oauth.py index 662d13c..3962a6d 100644 --- a/llm_be/chat_backend/tests/test_oauth.py +++ b/llm_be/chat_backend/tests/test_oauth.py @@ -11,9 +11,9 @@ from rest_framework import status from rest_framework.test import APITestCase from rest_framework_simplejwt.tokens import AccessToken -from chat_backend.models import CustomUser, OAuthIdentity +from chat_backend.models import CustomUser, DriveConnection, OAuthIdentity from chat_backend.oauth import ProviderProfile, dump_oauth_state -from chat_backend.tests.factories import make_user +from chat_backend.tests.factories import make_company, make_user OAUTH_SETTINGS = { "GOOGLE_OAUTH_CLIENT_ID": "google-client-id", @@ -239,3 +239,231 @@ class OAuthCallbackTestCase(APITestCase): provider="microsoft", subject="ms-oid-1", user=user ).exists() ) + + +@override_settings(**OAUTH_SETTINGS) +class OAuthStartDriveLinkTestCase(APITestCase): + """#47 — link_drive / link_company_drive require an authenticated user.""" + + def setUp(self): + self.company = make_company() + self.user = make_user(email="drive.user@example.com", company=self.company) + + def test_link_drive_requires_authentication(self): + response = self.client.get( + reverse("oauth_start", kwargs={"provider": "google"}), + {"intent": "link_drive"}, + ) + self.assertEqual(response.status_code, status.HTTP_401_UNAUTHORIZED) + + def test_link_drive_authenticated_redirects_with_drive_scope_and_state(self): + self.client.force_authenticate(user=self.user) + response = self.client.get( + reverse("oauth_start", kwargs={"provider": "google"}), + {"intent": "link_drive"}, + ) + self.assertEqual(response.status_code, status.HTTP_302_FOUND) + params = parse_qs(urlparse(response["Location"]).query) + self.assertIn("https://www.googleapis.com/auth/drive.readonly", params["scope"][0]) + self.assertIn("state", params) + + from chat_backend.oauth import load_oauth_state + + state_data = load_oauth_state(params["state"][0]) + self.assertEqual(state_data["intent"], "link_drive") + self.assertEqual(state_data["user_id"], self.user.id) + + def test_link_drive_microsoft_uses_personal_files_scope(self): + self.client.force_authenticate(user=self.user) + response = self.client.get( + reverse("oauth_start", kwargs={"provider": "microsoft"}), + {"intent": "link_drive"}, + ) + params = parse_qs(urlparse(response["Location"]).query) + self.assertIn("Files.Read", params["scope"][0]) + self.assertNotIn("Sites.Read.All", params["scope"][0]) + + def test_link_company_drive_microsoft_uses_sites_scope(self): + self.user.is_company_manager = True + self.user.save(update_fields=["is_company_manager"]) + self.client.force_authenticate(user=self.user) + response = self.client.get( + reverse("oauth_start", kwargs={"provider": "microsoft"}), + {"intent": "link_company_drive"}, + ) + params = parse_qs(urlparse(response["Location"]).query) + self.assertIn("Files.Read.All", params["scope"][0]) + self.assertIn("Sites.Read.All", params["scope"][0]) + + def test_link_company_drive_requires_company_manager(self): + self.client.force_authenticate(user=self.user) + response = self.client.get( + reverse("oauth_start", kwargs={"provider": "google"}), + {"intent": "link_company_drive"}, + ) + self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN) + + @override_settings(ENFORCE_SUBSCRIPTION_GATES=True) + def test_link_drive_denied_when_plan_disallows_rag(self): + from finance.services.plans import assign_plan, seed_subscription_plans + from finance.models import SubscriptionPlan, UserSubscription + + seed_subscription_plans() + standard = SubscriptionPlan.objects.get(slug="standard") + assign_plan(self.user, plan=standard, source=UserSubscription.Source.ADMIN) + + self.client.force_authenticate(user=self.user) + response = self.client.get( + reverse("oauth_start", kwargs={"provider": "google"}), + {"intent": "link_drive"}, + ) + self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN) + + def test_unknown_intent_rejected(self): + self.client.force_authenticate(user=self.user) + response = self.client.get( + reverse("oauth_start", kwargs={"provider": "google"}), + {"intent": "delete_everything"}, + ) + self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST) + + +@override_settings(**OAUTH_SETTINGS) +class OAuthCallbackDriveLinkTestCase(APITestCase): + """#47 — Drive-link callback upserts a DriveConnection and redirects to FE account page.""" + + def setUp(self): + self.company = make_company() + self.user = make_user(email="drive.user@example.com", company=self.company) + + def _state(self, *, intent="link_drive", provider="google", user_id=None): + return dump_oauth_state( + provider=provider, + intent=intent, + user_id=self.user.id if user_id is None else user_id, + ) + + @patch("chat_backend.views_oauth.exchange_code_for_profile") + def test_callback_creates_personal_drive_connection(self, mock_exchange): + mock_exchange.return_value = _google_profile( + scopes="openid email profile https://www.googleapis.com/auth/drive.readonly" + ) + 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) + self.assertTrue( + response["Location"].startswith("http://frontend.test/document_storage/?") + ) + params = parse_qs(urlparse(response["Location"]).query) + self.assertEqual(params["drive_connected"], ["1"]) + self.assertEqual(params["provider"], ["google"]) + self.assertEqual(params["kind"], ["personal"]) + + connection = DriveConnection.objects.get( + company=self.company, kind=DriveConnection.Kind.PERSONAL + ) + self.assertEqual(connection.user_id, self.user.id) + self.assertEqual(connection.provider, "google") + self.assertEqual(connection.refresh_token, "refresh-token") + self.assertTrue(connection.is_active) + # No JWTs minted for an in-app link flow (user is already authenticated). + self.assertNotIn("access", params) + + @patch("chat_backend.views_oauth.exchange_code_for_profile") + def test_callback_updates_existing_connection_tokens(self, mock_exchange): + from chat_backend.tests.factories import make_drive_connection + + existing = make_drive_connection( + self.company, + provider=DriveConnection.Provider.GOOGLE, + kind=DriveConnection.Kind.PERSONAL, + user=self.user, + access_token="stale", + refresh_token="stale-refresh", + ) + mock_exchange.return_value = _google_profile(access_token="fresh-access") + state = self._state(intent="link_drive") + + self.client.get( + reverse("oauth_callback", kwargs={"provider": "google"}), + {"code": "auth-code", "state": state}, + ) + + existing.refresh_from_db() + self.assertEqual(existing.access_token, "fresh-access") + self.assertEqual(DriveConnection.objects.count(), 1) + + @patch("chat_backend.views_oauth.exchange_code_for_profile") + def test_callback_company_drive_creates_connection_without_user(self, mock_exchange): + self.user.is_company_manager = True + self.user.save(update_fields=["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["kind"], ["company"]) + connection = DriveConnection.objects.get(kind=DriveConnection.Kind.COMPANY) + self.assertIsNone(connection.user_id) + self.assertEqual(connection.company_id, self.company.id) + + @patch("chat_backend.views_oauth.exchange_code_for_profile") + def test_callback_company_drive_rejects_non_manager(self, mock_exchange): + 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"], ["forbidden"]) + 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 + + from chat_backend.oauth import STATE_SALT + + bad_state = signing.dumps( + {"provider": "google", "intent": "link_drive"}, salt=STATE_SALT + ) + response = self.client.get( + reverse("oauth_callback", kwargs={"provider": "google"}), + {"code": "auth-code", "state": bad_state}, + ) + params = parse_qs(urlparse(response["Location"]).query) + self.assertEqual(params["error"], ["invalid_state"]) + + @override_settings(ENFORCE_SUBSCRIPTION_GATES=True) + @patch("chat_backend.views_oauth.exchange_code_for_profile") + def test_callback_denied_when_plan_disallows_rag(self, mock_exchange): + from finance.services.plans import assign_plan, seed_subscription_plans + from finance.models import SubscriptionPlan, UserSubscription + + seed_subscription_plans() + standard = SubscriptionPlan.objects.get(slug="standard") + assign_plan(self.user, plan=standard, source=UserSubscription.Source.ADMIN) + 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}, + ) + + params = parse_qs(urlparse(response["Location"]).query) + self.assertEqual(params["error"], ["feature_not_allowed"]) + self.assertFalse(DriveConnection.objects.exists()) diff --git a/llm_be/chat_backend/tests/test_services_drive_sync.py b/llm_be/chat_backend/tests/test_services_drive_sync.py new file mode 100644 index 0000000..a71e6c6 --- /dev/null +++ b/llm_be/chat_backend/tests/test_services_drive_sync.py @@ -0,0 +1,488 @@ +"""Tests for Drive/RAG sync (#48-#52). httpx is mocked; no network access.""" + +from __future__ import annotations + +from datetime import timedelta +from unittest import mock + +from django.test import TestCase +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.drive_sync import ( + DriveSyncError, + RemoteFile, + ensure_fresh_token, + refresh_google_token, + refresh_microsoft_token, + sync_connection, +) + +from .factories import make_company, make_drive_connection, make_user + + +def _fake_response(*, status_code=200, json_data=None, content=b"", text="", headers=None): + response = mock.MagicMock() + response.status_code = status_code + response.json.return_value = json_data or {} + response.content = content + response.text = text + response.headers = headers or {} + return response + + +def _fake_client(get_side_effect=None, post_side_effect=None): + """A MagicMock standing in for ``httpx.Client()`` used as a context manager.""" + client = mock.MagicMock() + if get_side_effect is not None: + client.get.side_effect = get_side_effect + if post_side_effect is not None: + client.post.side_effect = post_side_effect + cm = mock.MagicMock() + cm.__enter__.return_value = client + cm.__exit__.return_value = False + return cm, client + + +class TokenRefreshTestCase(TestCase): + def setUp(self): + self.company = make_company() + self.connection = make_drive_connection( + self.company, + provider=DriveConnection.Provider.GOOGLE, + access_token="old-token", + refresh_token="refresh-me", + token_expires_at=timezone.now() - timedelta(minutes=5), + ) + + @mock.patch("chat_backend.services.drive_sync.httpx.Client") + def test_refresh_google_token_updates_connection(self, mock_client_cls): + cm, client = _fake_client( + post_side_effect=[ + _fake_response(json_data={"access_token": "new-token", "expires_in": 3600}) + ] + ) + mock_client_cls.return_value = cm + + token = refresh_google_token(self.connection) + + self.assertEqual(token, "new-token") + self.connection.refresh_from_db() + self.assertEqual(self.connection.access_token, "new-token") + self.assertGreater(self.connection.token_expires_at, timezone.now()) + + def test_refresh_google_token_requires_refresh_token(self): + self.connection.refresh_token = "" + self.connection.save(update_fields=["refresh_token"]) + + with self.assertRaises(DriveSyncError) as ctx: + refresh_google_token(self.connection) + self.assertEqual(ctx.exception.code, "missing_refresh_token") + + @mock.patch("chat_backend.services.drive_sync.httpx.Client") + def test_refresh_google_token_raises_on_http_error(self, mock_client_cls): + cm, client = _fake_client( + post_side_effect=[_fake_response(status_code=400, text="invalid_grant")] + ) + mock_client_cls.return_value = cm + + with self.assertRaises(DriveSyncError) as ctx: + refresh_google_token(self.connection) + self.assertEqual(ctx.exception.code, "token_refresh_failed") + + @mock.patch("chat_backend.services.drive_sync.httpx.Client") + def test_refresh_microsoft_token_rotates_refresh_token(self, mock_client_cls): + connection = make_drive_connection( + self.company, + provider=DriveConnection.Provider.MICROSOFT, + refresh_token="old-refresh", + ) + cm, client = _fake_client( + post_side_effect=[ + _fake_response( + json_data={ + "access_token": "ms-new-token", + "refresh_token": "ms-new-refresh", + "expires_in": 3600, + } + ) + ] + ) + mock_client_cls.return_value = cm + + token = refresh_microsoft_token(connection) + + self.assertEqual(token, "ms-new-token") + connection.refresh_from_db() + self.assertEqual(connection.refresh_token, "ms-new-refresh") + + @mock.patch("chat_backend.services.drive_sync.refresh_google_token") + def test_ensure_fresh_token_skips_refresh_when_not_expired(self, mock_refresh): + self.connection.token_expires_at = timezone.now() + timedelta(hours=1) + self.connection.save(update_fields=["token_expires_at"]) + + token = ensure_fresh_token(self.connection) + + self.assertEqual(token, "old-token") + mock_refresh.assert_not_called() + + @mock.patch("chat_backend.services.drive_sync.refresh_google_token") + def test_ensure_fresh_token_refreshes_when_expired(self, mock_refresh): + mock_refresh.return_value = "refreshed" + token = ensure_fresh_token(self.connection) + self.assertEqual(token, "refreshed") + mock_refresh.assert_called_once_with(self.connection) + + +class ListGoogleFilesTestCase(TestCase): + def setUp(self): + self.company = make_company() + + @mock.patch("chat_backend.services.drive_sync.httpx.Client") + def test_lists_root_when_no_resources_selected(self, mock_client_cls): + connection = make_drive_connection(self.company, selected_resource_ids=[]) + cm, client = _fake_client( + get_side_effect=[ + _fake_response( + json_data={ + "files": [ + {"id": "f1", "name": "a.pdf", "mimeType": "application/pdf"}, + {"id": "folder1", "name": "Sub", "mimeType": drive_sync.GOOGLE_FOLDER_MIME}, + ] + } + ) + ] + ) + mock_client_cls.return_value = cm + + files = drive_sync._list_google_files(connection, "token") + + self.assertEqual([f.id for f in files], ["f1"]) + called_params = client.get.call_args.kwargs["params"] + self.assertEqual(called_params["q"], "'root' in parents and trashed = false") + + @mock.patch("chat_backend.services.drive_sync.httpx.Client") + def test_lists_selected_shared_drive_with_corpora_params(self, mock_client_cls): + connection = make_drive_connection( + self.company, + kind=DriveConnection.Kind.COMPANY, + selected_resource_ids=["drive-123"], + ) + cm, client = _fake_client( + get_side_effect=[_fake_response(json_data={"files": []})] + ) + mock_client_cls.return_value = cm + + drive_sync._list_google_files(connection, "token") + + called_params = client.get.call_args.kwargs["params"] + self.assertEqual(called_params["corpora"], "drive") + self.assertEqual(called_params["driveId"], "drive-123") + + @mock.patch("chat_backend.services.drive_sync.httpx.Client") + def test_list_raises_drive_sync_error_on_http_failure(self, mock_client_cls): + connection = make_drive_connection(self.company) + cm, client = _fake_client( + get_side_effect=[_fake_response(status_code=403, text="forbidden")] + ) + mock_client_cls.return_value = cm + + with self.assertRaises(DriveSyncError) as ctx: + drive_sync._list_google_files(connection, "token") + self.assertEqual(ctx.exception.code, "list_failed") + + +class DownloadGoogleFileTestCase(TestCase): + def test_exports_google_native_document_to_docx(self): + remote = RemoteFile( + id="doc1", + name="Report", + mime_type="application/vnd.google-apps.document", + etag="etag1", + ) + client = mock.MagicMock() + client.get.return_value = _fake_response(content=b"docx-bytes") + + content, filename, content_type = drive_sync._download_google_file( + client, {"Authorization": "Bearer x"}, remote + ) + + self.assertEqual(content, b"docx-bytes") + self.assertEqual(filename, "Report.docx") + self.assertIn("wordprocessingml", content_type) + self.assertIn("/export", client.get.call_args.args[0]) + + def test_downloads_regular_file_directly(self): + remote = RemoteFile(id="f1", name="notes.pdf", mime_type="application/pdf", etag="e1") + client = mock.MagicMock() + client.get.return_value = _fake_response(content=b"%PDF-bytes") + + content, filename, content_type = drive_sync._download_google_file( + client, {}, remote + ) + + self.assertEqual(filename, "notes.pdf") + self.assertEqual(content_type, "application/pdf") + + +class ListMicrosoftFilesTestCase(TestCase): + def setUp(self): + self.company = make_company() + + @mock.patch("chat_backend.services.drive_sync.httpx.Client") + def test_personal_lists_me_drive_root(self, mock_client_cls): + connection = make_drive_connection( + self.company, provider=DriveConnection.Provider.MICROSOFT, selected_resource_ids=[] + ) + cm, client = _fake_client( + get_side_effect=[ + _fake_response( + json_data={ + "value": [ + {"id": "i1", "name": "a.docx", "file": {"mimeType": "application/msword"}}, + {"id": "folder1", "name": "Sub", "folder": {}}, + ] + } + ) + ] + ) + mock_client_cls.return_value = cm + + files = drive_sync._list_microsoft_files(connection, "token") + + self.assertEqual([f.id for f in files], ["i1"]) + self.assertIn("/me/drive/root/children", client.get.call_args.args[0]) + + @mock.patch("chat_backend.services.drive_sync.httpx.Client") + def test_company_requires_selected_sites(self, mock_client_cls): + connection = make_drive_connection( + self.company, + provider=DriveConnection.Provider.MICROSOFT, + kind=DriveConnection.Kind.COMPANY, + selected_resource_ids=["site-1"], + ) + cm, client = _fake_client( + get_side_effect=[_fake_response(json_data={"value": [{"id": "i1", "name": "a.pptx"}]})] + ) + mock_client_cls.return_value = cm + + files = drive_sync._list_microsoft_files(connection, "token") + + self.assertEqual(files[0].context_id, "site-1") + self.assertIn("/sites/site-1/drive/root/children", client.get.call_args.args[0]) + + @mock.patch("chat_backend.services.drive_sync.httpx.Client") + def test_company_without_selected_sites_returns_no_files(self, mock_client_cls): + connection = make_drive_connection( + self.company, + provider=DriveConnection.Provider.MICROSOFT, + kind=DriveConnection.Kind.COMPANY, + selected_resource_ids=[], + ) + files = drive_sync._list_microsoft_files(connection, "token") + self.assertEqual(files, []) + mock_client_cls.assert_called_once() + + @mock.patch("chat_backend.services.drive_sync.httpx.Client") + def test_follows_pagination_next_link(self, mock_client_cls): + connection = make_drive_connection( + self.company, provider=DriveConnection.Provider.MICROSOFT + ) + cm, client = _fake_client( + get_side_effect=[ + _fake_response( + json_data={ + "value": [{"id": "i1", "name": "a.docx"}], + "@odata.nextLink": "https://graph.microsoft.com/v1.0/me/drive/root/children?page=2", + } + ), + _fake_response(json_data={"value": [{"id": "i2", "name": "b.docx"}]}), + ] + ) + mock_client_cls.return_value = cm + + files = drive_sync._list_microsoft_files(connection, "token") + + self.assertEqual([f.id for f in files], ["i1", "i2"]) + self.assertEqual(client.get.call_count, 2) + + +class DownloadMicrosoftFileTestCase(TestCase): + def setUp(self): + self.company = make_company() + + def test_personal_download_uses_me_drive(self): + connection = make_drive_connection( + self.company, provider=DriveConnection.Provider.MICROSOFT + ) + remote = RemoteFile(id="i1", name="a.docx", mime_type="", etag="e1") + client = mock.MagicMock() + client.get.return_value = _fake_response( + content=b"bytes", headers={"content-type": "application/msword"} + ) + + content, filename, content_type = drive_sync._download_microsoft_file( + client, {}, connection, remote + ) + + self.assertIn("/me/drive/items/i1/content", client.get.call_args.args[0]) + self.assertEqual(filename, "a.docx") + + def test_company_download_uses_site_context(self): + connection = make_drive_connection( + self.company, + provider=DriveConnection.Provider.MICROSOFT, + kind=DriveConnection.Kind.COMPANY, + ) + remote = RemoteFile(id="i1", name="a.docx", mime_type="", etag="e1", context_id="site-9") + client = mock.MagicMock() + client.get.return_value = _fake_response(content=b"bytes", headers={}) + + drive_sync._download_microsoft_file(client, {}, connection, remote) + + self.assertIn("/sites/site-9/drive/items/i1/content", client.get.call_args.args[0]) + + +class SyncConnectionTestCase(TestCase): + """Exercises the add/update/remove reconciliation loop end to end.""" + + def setUp(self): + self.company = make_company() + self.connection = make_drive_connection( + self.company, + provider=DriveConnection.Provider.GOOGLE, + token_expires_at=timezone.now() + timedelta(hours=1), + ) + rag_patcher = mock.patch("chat_backend.services.drive_sync.AsyncRAGService") + self.mock_rag_cls = rag_patcher.start() + self.mock_rag_cls.return_value.add_files_to_store.return_value = { + "total_added": 1, + "failed_files": [], + "processed_files": [], + } + self.addCleanup(rag_patcher.stop) + + list_patcher = mock.patch("chat_backend.services.drive_sync._list_remote_files") + self.mock_list = list_patcher.start() + self.addCleanup(list_patcher.stop) + + download_patcher = mock.patch("chat_backend.services.drive_sync._download_file") + self.mock_download = download_patcher.start() + self.addCleanup(download_patcher.stop) + + client_patcher = mock.patch("chat_backend.services.drive_sync.httpx.Client") + self.mock_client_cls = client_patcher.start() + cm, _client = _fake_client() + self.mock_client_cls.return_value = cm + self.addCleanup(client_patcher.stop) + + def test_creates_document_for_new_remote_file(self): + self.mock_list.return_value = [ + RemoteFile(id="r1", name="a.pdf", mime_type="application/pdf", etag="e1") + ] + self.mock_download.return_value = (b"%PDF-1.4", "a.pdf", "application/pdf") + + result = sync_connection(self.connection) + + self.assertEqual(result["added"], 1) + document = Document.objects.get(drive_connection=self.connection) + self.assertEqual(document.remote_file_id, "r1") + self.assertEqual(document.remote_etag, "e1") + self.assertEqual(document.source, Document.Source.GOOGLE_DRIVE) + self.assertTrue(document.processed) + self.assertTrue(document.active) + + self.connection.refresh_from_db() + self.assertEqual(self.connection.last_sync_status, DriveConnection.SyncStatus.OK) + self.assertIsNotNone(self.connection.last_sync_at) + + def test_skips_unchanged_file(self): + Document.objects.create( + workspace=ensure_company_workspace(self.company), + drive_connection=self.connection, + remote_file_id="r1", + remote_etag="e1", + remote_name="a.pdf", + ) + self.mock_list.return_value = [ + RemoteFile(id="r1", name="a.pdf", mime_type="application/pdf", etag="e1") + ] + + result = sync_connection(self.connection) + + self.assertEqual(result["added"], 0) + self.assertEqual(result["updated"], 0) + self.mock_download.assert_not_called() + + def test_updates_file_when_etag_changes(self): + existing = Document.objects.create( + workspace=ensure_company_workspace(self.company), + drive_connection=self.connection, + remote_file_id="r1", + remote_etag="old-etag", + remote_name="a.pdf", + ) + self.mock_list.return_value = [ + RemoteFile(id="r1", name="a.pdf", mime_type="application/pdf", etag="new-etag") + ] + self.mock_download.return_value = (b"new-bytes", "a.pdf", "application/pdf") + + result = sync_connection(self.connection) + + self.assertEqual(result["updated"], 1) + existing.refresh_from_db() + self.assertEqual(existing.remote_etag, "new-etag") + self.assertEqual(Document.objects.filter(drive_connection=self.connection).count(), 1) + + def test_removes_document_whose_remote_file_is_gone(self): + Document.objects.create( + workspace=ensure_company_workspace(self.company), + drive_connection=self.connection, + remote_file_id="deleted-remote", + remote_etag="e1", + remote_name="gone.pdf", + ) + self.mock_list.return_value = [] + + result = sync_connection(self.connection) + + self.assertEqual(result["removed"], 1) + self.assertFalse(Document.objects.filter(remote_file_id="deleted-remote").exists()) + + def test_download_failure_is_recorded_and_does_not_abort_sync(self): + self.mock_list.return_value = [ + RemoteFile(id="r1", name="a.pdf", mime_type="application/pdf", etag="e1"), + RemoteFile(id="r2", name="b.pdf", mime_type="application/pdf", etag="e2"), + ] + self.mock_download.side_effect = [ + DriveSyncError("download_failed", "boom"), + (b"ok-bytes", "b.pdf", "application/pdf"), + ] + + result = sync_connection(self.connection) + + self.assertEqual(len(result["failed"]), 1) + self.assertEqual(result["added"], 1) + self.assertEqual(Document.objects.filter(drive_connection=self.connection).count(), 1) + + def test_list_failure_marks_connection_error(self): + self.mock_list.side_effect = DriveSyncError("list_failed", "quota exceeded") + + result = sync_connection(self.connection) + + self.assertEqual(result["error"], "quota exceeded") + self.connection.refresh_from_db() + 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): + from chat_backend.models import DocumentWorkspace + + self.assertFalse(DocumentWorkspace.objects.filter(company=self.company).exists()) + self.mock_list.return_value = [] + + sync_connection(self.connection) + + self.assertTrue(DocumentWorkspace.objects.filter(company=self.company).exists()) diff --git a/llm_be/chat_backend/tests/test_services_rag.py b/llm_be/chat_backend/tests/test_services_rag.py index 03839aa..0c94331 100644 --- a/llm_be/chat_backend/tests/test_services_rag.py +++ b/llm_be/chat_backend/tests/test_services_rag.py @@ -200,6 +200,37 @@ class RAGServiceTestCase(TransactionTestCase): self.assertEqual(results["total_added"], 1) self.assertEqual(results["failed_files"], []) + def test_add_files_to_store_includes_document_id_and_active_metadata(self): + document = self._text_document(body=b"from the database") + + self.service.add_files_to_store( + [ + ( + document.file, + document.file.name, + self.workspace.id, + document.id, + True, + ) + ], + workspace_id=self.workspace.id, + ) + + added = self.service.vector_store.add_documents.call_args[0][0] + self.assertEqual(added[0].metadata["document_id"], document.id) + self.assertTrue(added[0].metadata["active"]) + + def test_add_files_to_store_defaults_active_true_without_document_id(self): + path = self._temp_text_file(b"upload body") + + self.service.add_files_to_store( + [(path, "upload.txt", self.workspace.id)], workspace_id=self.workspace.id + ) + + added = self.service.vector_store.add_documents.call_args[0][0] + self.assertIsNone(added[0].metadata["document_id"]) + self.assertTrue(added[0].metadata["active"]) + def test_add_files_to_store_records_failures(self): results = self.service.add_files_to_store( [("/tmp/does-not-exist.txt", "missing.txt", self.workspace.id)], @@ -217,6 +248,48 @@ class RAGServiceTestCase(TransactionTestCase): first_store.delete_collection.assert_called_once() self.assertIsNot(self.service.vector_store, first_store) + def test_delete_document_vectors_deletes_by_document_id_metadata(self): + self.service.delete_document_vectors(42) + + self.service.vector_store.delete.assert_called_once_with( + where={"document_id": 42} + ) + self.service.vector_store.persist.assert_called() + + def test_delete_document_vectors_is_a_noop_without_a_document_id(self): + self.service.delete_document_vectors(None) + + self.service.vector_store.delete.assert_not_called() + + def test_set_document_active_updates_matching_chunk_metadata(self): + self.service.vector_store.get.return_value = { + "ids": ["a", "b"], + "metadatas": [ + {"document_id": 7, "active": True, "workspace_id": 1}, + {"document_id": 7, "active": True, "workspace_id": 1}, + ], + } + + self.service.set_document_active(7, False) + + self.service.vector_store.get.assert_called_once_with( + where={"document_id": 7} + ) + self.service.vector_store._collection.update.assert_called_once_with( + ids=["a", "b"], + metadatas=[ + {"document_id": 7, "active": False, "workspace_id": 1}, + {"document_id": 7, "active": False, "workspace_id": 1}, + ], + ) + + def test_set_document_active_is_a_noop_when_no_chunks_found(self): + self.service.vector_store.get.return_value = {"ids": [], "metadatas": []} + + self.service.set_document_active(999, True) + + self.service.vector_store._collection.update.assert_not_called() + async def test_search_documents_filters_by_workspace(self): retriever = self.service.vector_store.as_retriever.return_value retriever.aget_relevant_documents = mock.AsyncMock( @@ -228,7 +301,15 @@ class RAGServiceTestCase(TransactionTestCase): self.assertEqual([doc.page_content for doc in docs], ["chunk"]) self.service.vector_store.as_retriever.assert_called_with( search_type="mmr", - search_kwargs={"k": 2, "filter": {"workspace_id": self.workspace.id}}, + search_kwargs={ + "k": 2, + "filter": { + "$and": [ + {"workspace_id": self.workspace.id}, + {"active": True}, + ] + }, + }, ) async def test_search_documents_without_workspace_fails_closed(self): diff --git a/llm_be/chat_backend/tests/test_signals.py b/llm_be/chat_backend/tests/test_signals.py index d9294e6..c4a30e5 100644 --- a/llm_be/chat_backend/tests/test_signals.py +++ b/llm_be/chat_backend/tests/test_signals.py @@ -3,6 +3,8 @@ from unittest import mock from django.test import TestCase +from chat_backend.models import Document + from .factories import make_company, make_document, make_workspace RAG_SERVICE = "chat_backend.services.rag_services.AsyncRAGService" @@ -12,12 +14,13 @@ class DocumentSignalTestCase(TestCase): def setUp(self): self.workspace = make_workspace(make_company()) - def test_creating_a_document_reindexes_the_vector_store(self): + def test_creating_a_document_does_not_trigger_the_rag_service(self): + """#45: uploads ingest via DocumentUploadView.process_document, not a signal.""" with mock.patch.dict(os.environ, {"SKIP_RAG_INIT": ""}): with mock.patch(RAG_SERVICE) as service: make_document(self.workspace) - service.return_value.ingest_documents.assert_called_once_with() + service.assert_not_called() def test_updating_a_document_does_not_reindex(self): document = make_document(self.workspace) @@ -29,14 +32,18 @@ class DocumentSignalTestCase(TestCase): service.assert_not_called() - def test_deleting_a_document_reindexes_the_vector_store(self): + def test_deleting_a_document_removes_only_its_own_vectors(self): document = make_document(self.workspace) + document_id = document.id with mock.patch.dict(os.environ, {"SKIP_RAG_INIT": ""}): with mock.patch(RAG_SERVICE) as service: document.delete() - service.return_value.ingest_documents.assert_called_once_with() + service.return_value.delete_document_vectors.assert_called_once_with( + document_id + ) + service.return_value.ingest_documents.assert_not_called() def test_skip_rag_init_keeps_signals_inert(self): with mock.patch.dict(os.environ, {"SKIP_RAG_INIT": "1"}): @@ -46,9 +53,12 @@ class DocumentSignalTestCase(TestCase): service.assert_not_called() - def test_vector_store_failures_do_not_break_uploads(self): + def test_vector_store_failures_do_not_break_deletes(self): + document = make_document(self.workspace) + document_id = document.id + with mock.patch.dict(os.environ, {"SKIP_RAG_INIT": ""}): with mock.patch(RAG_SERVICE, side_effect=RuntimeError("chroma down")): - document = make_document(self.workspace) + document.delete() - self.assertIsNotNone(document.pk) + self.assertFalse(Document.objects.filter(pk=document_id).exists()) diff --git a/llm_be/chat_backend/tests/test_views_documents.py b/llm_be/chat_backend/tests/test_views_documents.py index 0f89e29..ee38382 100644 --- a/llm_be/chat_backend/tests/test_views_documents.py +++ b/llm_be/chat_backend/tests/test_views_documents.py @@ -1,10 +1,14 @@ +import os from unittest import mock +from django.test import override_settings from django.urls import reverse from rest_framework import status from rest_framework.test import APITestCase from chat_backend.models import Document, DocumentWorkspace, StoredFile +from finance.models import SubscriptionPlan, UserSubscription +from finance.services.plans import assign_plan, seed_subscription_plans from .factories import ( make_company, @@ -89,14 +93,18 @@ class DocumentUploadViewTestCase(APITestCase): self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST) - def test_upload_without_workspace_returns_404(self): + def test_upload_without_workspace_creates_a_default_one(self): + """#46: missing workspace is auto-created instead of 404ing the upload.""" self.workspace.delete() response = self.client.post( self.url, {"file": pdf_upload()}, format="multipart" ) - self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND) + self.assertEqual(response.status_code, status.HTTP_201_CREATED) + created_workspace = DocumentWorkspace.objects.get(company=self.company) + self.assertEqual(created_workspace.name, "Default") + self.assertEqual(Document.objects.get().workspace_id, created_workspace.id) def test_list_returns_documents_of_own_workspace(self): make_document(self.workspace) @@ -108,13 +116,19 @@ class DocumentUploadViewTestCase(APITestCase): self.assertIn("test", response.data[0]["file"]) self.assertIn("pdf", response.data[0]["file"]) - def test_list_without_workspace_returns_404(self): + def test_list_without_workspace_creates_a_default_one(self): + """#46: missing workspace is auto-created instead of 404ing the list.""" self.workspace.delete() response = self.client.get(self.url) - self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND) - self.assertEqual(response.data["error"], "Workspace not found") + self.assertEqual(response.status_code, status.HTTP_200_OK) + self.assertEqual(response.data, []) + self.assertTrue( + DocumentWorkspace.objects.filter( + company=self.company, name="Default" + ).exists() + ) class DocumentDetailViewTestCase(APITestCase): @@ -140,3 +154,155 @@ class DocumentDetailViewTestCase(APITestCase): self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND) self.assertEqual(response.data["error"], "Document not found") + + def test_get_returns_the_document(self): + document = make_document(self.workspace) + url = reverse("documents_details", kwargs={"document_id": document.id}) + + response = self.client.get(url) + + self.assertEqual(response.status_code, status.HTTP_200_OK) + self.assertEqual(response.data["id"], document.id) + self.assertIn("test", response.data["file"]) + + def test_patch_toggles_active_and_updates_vector_metadata(self): + document = make_document(self.workspace) + document.active = True + document.save(update_fields=["active"]) + url = reverse("documents_details", kwargs={"document_id": document.id}) + + with mock.patch("chat_backend.views.AsyncRAGService") as rag_service: + response = self.client.patch(url, {"active": False}, format="json") + + self.assertEqual(response.status_code, status.HTTP_200_OK) + self.assertFalse(response.data["active"]) + document.refresh_from_db() + self.assertFalse(document.active) + rag_service.return_value.set_document_active.assert_called_once_with( + document.id, False + ) + + def test_patch_requires_active_field(self): + document = make_document(self.workspace) + url = reverse("documents_details", kwargs={"document_id": document.id}) + + response = self.client.patch(url, {}, format="json") + + self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST) + + def test_patch_unknown_document_returns_404(self): + url = reverse("documents_details", kwargs={"document_id": 4242}) + + response = self.client.patch(url, {"active": True}, format="json") + + self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND) + + def test_delete_removes_the_document_and_its_vectors(self): + document = make_document(self.workspace) + document_id = document.id + url = reverse("documents_details", kwargs={"document_id": document_id}) + + with mock.patch.dict(os.environ, {"SKIP_RAG_INIT": ""}): + with mock.patch( + "chat_backend.services.rag_services.AsyncRAGService" + ) as rag_service: + response = self.client.delete(url) + + self.assertEqual(response.status_code, status.HTTP_204_NO_CONTENT) + self.assertFalse(Document.objects.filter(id=document_id).exists()) + rag_service.return_value.delete_document_vectors.assert_called_once_with( + document_id + ) + + def test_delete_unknown_document_returns_404(self): + url = reverse("documents_details", kwargs={"document_id": 4242}) + + response = self.client.delete(url) + + self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND) + + +@override_settings(ENFORCE_SUBSCRIPTION_GATES=True) +class DocumentRagFeatureGateTestCase(APITestCase): + """#44: RAG document endpoints must respect the plan's ``rag`` feature gate. + + The rest of this module runs with ``ENFORCE_SUBSCRIPTION_GATES=False`` + (set by the test runner) so pre-existing behavior keeps working + regardless of plan; this class opts back in explicitly. + """ + + def setUp(self): + seed_subscription_plans() + self.company = make_company() + self.user = make_user(company=self.company) + self.client.force_authenticate(user=self.user) + self.workspace = make_workspace(self.company) + rag_patcher = mock.patch("chat_backend.views.AsyncRAGService") + self.rag_service = rag_patcher.start() + self.addCleanup(rag_patcher.stop) + + def _assign(self, slug): + plan = SubscriptionPlan.objects.get(slug=slug) + assign_plan(self.user, plan=plan, source=UserSubscription.Source.ADMIN) + + def test_standard_plan_is_denied_workspace_list(self): + self._assign("standard") + + response = self.client.get(reverse("document_workspaces")) + + self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN) + self.assertEqual(response.data["code"], "feature_not_allowed") + self.assertIn("error", response.data) + self.assertIn("details", response.data) + + def test_standard_plan_is_denied_document_upload(self): + self._assign("standard") + + response = self.client.post( + reverse("documents"), {"file": pdf_upload()}, format="multipart" + ) + + self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN) + self.assertEqual(response.data["code"], "feature_not_allowed") + self.assertEqual(Document.objects.count(), 0) + + def test_standard_plan_is_denied_document_list(self): + self._assign("standard") + + response = self.client.get(reverse("documents")) + + self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN) + + def test_standard_plan_is_denied_document_detail(self): + self._assign("standard") + document = make_document(self.workspace) + + url = reverse("documents_details", kwargs={"document_id": document.id}) + response = self.client.get(url) + + self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN) + + def test_no_subscription_is_denied(self): + response = self.client.post( + reverse("documents"), {"file": pdf_upload()}, format="multipart" + ) + + self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN) + self.assertEqual(response.data["code"], "subscription_required") + + def test_founders_plan_is_allowed_document_upload(self): + self._assign("founders") + + response = self.client.post( + reverse("documents"), {"file": pdf_upload()}, format="multipart" + ) + + self.assertEqual(response.status_code, status.HTTP_201_CREATED) + self.assertEqual(Document.objects.count(), 1) + + def test_founders_plan_is_allowed_workspace_list(self): + self._assign("founders") + + response = self.client.get(reverse("document_workspaces")) + + self.assertEqual(response.status_code, status.HTTP_200_OK) diff --git a/llm_be/chat_backend/tests/test_views_drive.py b/llm_be/chat_backend/tests/test_views_drive.py new file mode 100644 index 0000000..d44a983 --- /dev/null +++ b/llm_be/chat_backend/tests/test_views_drive.py @@ -0,0 +1,272 @@ +"""Tests for Drive connection management API (#47-#52).""" + +from __future__ import annotations + +from unittest import mock + +from django.test import override_settings +from django.urls import reverse +from rest_framework import status +from rest_framework.test import APITestCase + +from chat_backend.models import DriveConnection +from finance.models import SubscriptionPlan, UserSubscription +from finance.services.plans import assign_plan, seed_subscription_plans + +from .factories import make_company, make_drive_connection, make_user + + +class DriveConnectionListViewTestCase(APITestCase): + def setUp(self): + self.company = make_company() + self.user = make_user(email="user@example.com", company=self.company) + self.client.force_authenticate(user=self.user) + self.url = reverse("drive_connections") + + def test_lists_own_personal_and_company_connections(self): + personal = make_drive_connection( + self.company, kind=DriveConnection.Kind.PERSONAL, user=self.user + ) + company_conn = make_drive_connection( + self.company, kind=DriveConnection.Kind.COMPANY, user=None + ) + other_company = make_company("Other") + other_user = make_user(email="other@example.com", company=other_company) + make_drive_connection( + other_company, kind=DriveConnection.Kind.PERSONAL, user=other_user + ) + + response = self.client.get(self.url) + + self.assertEqual(response.status_code, status.HTTP_200_OK) + ids = {row["id"] for row in response.data} + self.assertEqual(ids, {personal.id, company_conn.id}) + + def test_does_not_leak_other_users_personal_connection(self): + teammate = make_user(email="teammate@example.com", company=self.company) + make_drive_connection( + self.company, kind=DriveConnection.Kind.PERSONAL, user=teammate + ) + + response = self.client.get(self.url) + + self.assertEqual(response.data, []) + + def test_response_never_includes_tokens(self): + make_drive_connection( + self.company, kind=DriveConnection.Kind.PERSONAL, user=self.user + ) + response = self.client.get(self.url) + self.assertNotIn("access_token", response.data[0]) + self.assertNotIn("refresh_token", response.data[0]) + + @override_settings(ENFORCE_SUBSCRIPTION_GATES=True) + def test_denied_on_standard_plan(self): + seed_subscription_plans() + standard = SubscriptionPlan.objects.get(slug="standard") + assign_plan(self.user, plan=standard, source=UserSubscription.Source.ADMIN) + + response = self.client.get(self.url) + self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN) + + @override_settings(ENFORCE_SUBSCRIPTION_GATES=True) + def test_allowed_on_pro_plan(self): + seed_subscription_plans() + pro = SubscriptionPlan.objects.get(slug="pro") + assign_plan(self.user, plan=pro, source=UserSubscription.Source.ADMIN) + + response = self.client.get(self.url) + self.assertEqual(response.status_code, status.HTTP_200_OK) + + +class DriveConnectionDetailViewTestCase(APITestCase): + def setUp(self): + self.company = make_company() + self.user = make_user(email="user@example.com", company=self.company) + self.client.force_authenticate(user=self.user) + + def test_owner_can_disconnect_personal_connection(self): + connection = make_drive_connection( + self.company, kind=DriveConnection.Kind.PERSONAL, user=self.user + ) + url = reverse("drive_connection_detail", kwargs={"connection_id": connection.id}) + + response = self.client.delete(url) + + self.assertEqual(response.status_code, status.HTTP_204_NO_CONTENT) + connection.refresh_from_db() + self.assertFalse(connection.is_active) + self.assertEqual(connection.access_token, "") + self.assertEqual(connection.refresh_token, "") + + def test_non_owner_cannot_disconnect_personal_connection(self): + teammate = make_user(email="teammate@example.com", company=self.company) + connection = make_drive_connection( + self.company, kind=DriveConnection.Kind.PERSONAL, user=teammate + ) + url = reverse("drive_connection_detail", kwargs={"connection_id": connection.id}) + + response = self.client.delete(url) + + self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN) + connection.refresh_from_db() + self.assertTrue(connection.is_active) + + def test_manager_can_disconnect_company_connection(self): + self.user.is_company_manager = True + self.user.save(update_fields=["is_company_manager"]) + connection = make_drive_connection( + self.company, kind=DriveConnection.Kind.COMPANY, user=None + ) + url = reverse("drive_connection_detail", kwargs={"connection_id": connection.id}) + + response = self.client.delete(url) + + self.assertEqual(response.status_code, status.HTTP_204_NO_CONTENT) + + def test_non_manager_cannot_disconnect_company_connection(self): + connection = make_drive_connection( + self.company, kind=DriveConnection.Kind.COMPANY, user=None + ) + url = reverse("drive_connection_detail", kwargs={"connection_id": connection.id}) + + response = self.client.delete(url) + + self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN) + + def test_unknown_connection_404(self): + url = reverse("drive_connection_detail", kwargs={"connection_id": 999999}) + response = self.client.delete(url) + self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND) + + +class DriveConnectionResourcesViewTestCase(APITestCase): + def setUp(self): + self.company = make_company() + self.user = make_user(email="user@example.com", company=self.company) + self.client.force_authenticate(user=self.user) + self.connection = make_drive_connection( + self.company, kind=DriveConnection.Kind.PERSONAL, user=self.user + ) + self.url = reverse( + "drive_connection_resources", kwargs={"connection_id": self.connection.id} + ) + + def test_sets_selected_resources(self): + response = self.client.post( + self.url, + {"resource_ids": ["folder-1", "folder-2"], "resource_labels": ["Reports", "HR"]}, + format="json", + ) + + self.assertEqual(response.status_code, status.HTTP_200_OK) + self.connection.refresh_from_db() + self.assertEqual(self.connection.selected_resource_ids, ["folder-1", "folder-2"]) + self.assertEqual(self.connection.selected_resource_labels, ["Reports", "HR"]) + + def test_requires_resource_ids_list(self): + response = self.client.post(self.url, {}, format="json") + self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST) + + def test_company_connection_requires_manager(self): + company_conn = make_drive_connection( + self.company, kind=DriveConnection.Kind.COMPANY, user=None + ) + url = reverse( + "drive_connection_resources", kwargs={"connection_id": company_conn.id} + ) + response = self.client.post(url, {"resource_ids": ["site-1"]}, format="json") + self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN) + + +class DriveConnectionSyncViewTestCase(APITestCase): + def setUp(self): + self.company = make_company() + self.user = make_user(email="user@example.com", company=self.company) + self.client.force_authenticate(user=self.user) + self.connection = make_drive_connection( + self.company, kind=DriveConnection.Kind.PERSONAL, user=self.user + ) + self.url = reverse("drive_connection_sync", kwargs={"connection_id": self.connection.id}) + + @mock.patch("chat_backend.views_drive.sync_connection") + def test_triggers_sync_for_owner(self, mock_sync): + mock_sync.return_value = {"added": 2, "updated": 0, "removed": 0, "failed": []} + + response = self.client.post(self.url) + + self.assertEqual(response.status_code, status.HTTP_200_OK) + mock_sync.assert_called_once_with(self.connection) + self.assertEqual(response.data["result"]["added"], 2) + + @mock.patch("chat_backend.views_drive.sync_connection") + def test_non_owner_forbidden(self, mock_sync): + other = make_user(email="other@example.com", company=self.company) + self.client.force_authenticate(user=other) + + response = self.client.post(self.url) + + self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN) + mock_sync.assert_not_called() + + @mock.patch("chat_backend.views_drive.sync_connection") + def test_disconnected_connection_rejected(self, mock_sync): + self.connection.is_active = False + self.connection.save(update_fields=["is_active"]) + + response = self.client.post(self.url) + + self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST) + mock_sync.assert_not_called() + + +class DriveWebhookViewTestCase(APITestCase): + def setUp(self): + self.company = make_company() + self.connection = make_drive_connection( + self.company, kind=DriveConnection.Kind.PERSONAL, provider=DriveConnection.Provider.GOOGLE + ) + + @mock.patch("chat_backend.views_drive.sync_connection") + def test_google_webhook_syncs_matching_connection(self, mock_sync): + url = reverse("drive_webhook_google") + response = self.client.post(f"{url}?connection_id={self.connection.id}") + + self.assertEqual(response.status_code, status.HTTP_200_OK) + mock_sync.assert_called_once_with(self.connection) + + @mock.patch("chat_backend.views_drive.sync_connection") + def test_google_webhook_without_connection_id_is_a_noop(self, mock_sync): + url = reverse("drive_webhook_google") + response = self.client.post(url) + + self.assertEqual(response.status_code, status.HTTP_200_OK) + mock_sync.assert_not_called() + + @mock.patch("chat_backend.views_drive.sync_connection") + def test_microsoft_webhook_validation_handshake(self, mock_sync): + url = reverse("drive_webhook_microsoft") + response = self.client.post(f"{url}?validationToken=abc123") + + self.assertEqual(response.status_code, status.HTTP_200_OK) + self.assertEqual(response.content.decode(), "abc123") + mock_sync.assert_not_called() + + @mock.patch("chat_backend.views_drive.sync_connection") + def test_microsoft_webhook_syncs_via_client_state(self, mock_sync): + ms_connection = make_drive_connection( + self.company, + kind=DriveConnection.Kind.PERSONAL, + provider=DriveConnection.Provider.MICROSOFT, + user=make_user(email="ms.user@example.com", company=self.company), + ) + url = reverse("drive_webhook_microsoft") + + response = self.client.post( + url, + {"value": [{"clientState": str(ms_connection.id)}]}, + format="json", + ) + + self.assertEqual(response.status_code, status.HTTP_200_OK) + mock_sync.assert_called_once_with(ms_connection) diff --git a/llm_be/chat_backend/urls.py b/llm_be/chat_backend/urls.py index e9244ba..87537b8 100644 --- a/llm_be/chat_backend/urls.py +++ b/llm_be/chat_backend/urls.py @@ -27,6 +27,14 @@ from .views import ( DocumentDetailView, ) from .views_oauth import OAuthCallbackView, OAuthStartView +from .views_drive import ( + DriveConnectionListView, + DriveConnectionDetailView, + DriveConnectionResourcesView, + DriveConnectionSyncView, + DriveWebhookGoogleView, + DriveWebhookMicrosoftView, +) from rest_framework.routers import DefaultRouter @@ -107,4 +115,35 @@ urlpatterns = [ DocumentDetailView.as_view(), name="documents_details", ), + # drive urls (#47-#52) + path( + "drive/connections/", + DriveConnectionListView.as_view(), + name="drive_connections", + ), + path( + "drive/connections//", + DriveConnectionDetailView.as_view(), + name="drive_connection_detail", + ), + path( + "drive/connections//resources/", + DriveConnectionResourcesView.as_view(), + name="drive_connection_resources", + ), + path( + "drive/connections//sync/", + DriveConnectionSyncView.as_view(), + name="drive_connection_sync", + ), + path( + "drive/webhooks/google/", + DriveWebhookGoogleView.as_view(), + name="drive_webhook_google", + ), + path( + "drive/webhooks/microsoft/", + DriveWebhookMicrosoftView.as_view(), + name="drive_webhook_microsoft", + ), ] diff --git a/llm_be/chat_backend/views.py b/llm_be/chat_backend/views.py index 9edef51..e123591 100644 --- a/llm_be/chat_backend/views.py +++ b/llm_be/chat_backend/views.py @@ -65,8 +65,10 @@ from .email_tasks import ( send_invite_email, send_password_reset_email, ) +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.title_generator import title_generator from .services.moderation_classifier import moderation_classifier, ModerationLabel from .services.prompt_classifier.prompt_classifier import PromptClassifier, PromptType @@ -789,15 +791,30 @@ llm = OllamaLLM(**ollama_llm_kwargs(model=MODEL_NAME)) # Document Views +def _feature_gate_response(exc: FeatureNotAllowed) -> Response: + return Response( + {"code": exc.code, "error": exc.message, "details": exc.details}, + status=status.HTTP_403_FORBIDDEN, + ) + + class DocumentWorkspaceView(APIView): # permission_classes = [permissions.IsAuthenticated] def get(self, request): + try: + assert_feature_allowed(request.user, "rag") + except FeatureNotAllowed as exc: + return _feature_gate_response(exc) workspaces = DocumentWorkspace.objects.filter(company=request.user.company) serializer = DocumentWorkspaceSerializer(workspaces, many=True) return Response(serializer.data) def post(self, request): + try: + assert_feature_allowed(request.user, "rag") + except FeatureNotAllowed as exc: + return _feature_gate_response(exc) serializer = DocumentWorkspaceSerializer(data=request.data) if serializer.is_valid(): serializer.save(company=request.user.company) @@ -811,27 +828,25 @@ class DocumentUploadView(APIView): def get(self, request): logger.debug(f"request_3: {request}") try: - workspace = DocumentWorkspace.objects.get(company=request.user.company) - serializer = DocumentSerializer( - Document.objects.filter(workspace=workspace), many=True - ) - return Response(serializer.data, status=status.HTTP_200_OK) + assert_feature_allowed(request.user, "rag") + except FeatureNotAllowed as exc: + return _feature_gate_response(exc) - except: - return Response( - {"error": "Workspace not found"}, status=status.HTTP_404_NOT_FOUND - ) + workspace = ensure_company_workspace(request.user.company) + serializer = DocumentSerializer( + Document.objects.filter(workspace=workspace), many=True + ) + return Response(serializer.data, status=status.HTTP_200_OK) def post(self, request): logger.debug(f"request: {request}") try: - workspace = DocumentWorkspace.objects.get(company=request.user.company) + assert_feature_allowed(request.user, "rag") + except FeatureNotAllowed as exc: + return _feature_gate_response(exc) - except: - return Response( - {"error": "Workspace not found"}, status=status.HTTP_404_NOT_FOUND - ) + workspace = ensure_company_workspace(request.user.company) logger.info(request.FILES) file = request.FILES.get("file") @@ -857,7 +872,15 @@ class DocumentUploadView(APIView): document.save() service = AsyncRAGService() service.add_files_to_store( - [(document.file, document.file.name, document.workspace_id)], + [ + ( + document.file, + document.file.name, + document.workspace_id, + document.id, + document.active, + ) + ], workspace_id=document.workspace_id, ) @@ -865,16 +888,73 @@ class DocumentUploadView(APIView): class DocumentDetailView(APIView): # permission_classes = [permissions.IsAuthenticated] + def _get_document(self, request, document_id): + workspace = ensure_company_workspace(request.user.company) + return Document.objects.filter(workspace=workspace, id=document_id).first() + def get(self, request, document_id): logger.info(f"request: {request}") try: - workspace = DocumentWorkspace.objects.get(company=request.user.company) + assert_feature_allowed(request.user, "rag") + except FeatureNotAllowed as exc: + return _feature_gate_response(exc) - document = Document.objects.get(workspace=workspace, id=document_id) - except: + document = self._get_document(request, document_id) + if document is None: return Response( {"error": "Document not found"}, status=status.HTTP_404_NOT_FOUND ) - serializer = DocumentWorkspaceSerializer(workspaces, many=True) - return Response(serializer.data) + serializer = DocumentSerializer(document) + return Response(serializer.data, status=status.HTTP_200_OK) + + def patch(self, request, document_id): + """Toggle a document's ``active`` flag (#44) and its vector metadata (#45).""" + try: + assert_feature_allowed(request.user, "rag") + except FeatureNotAllowed as exc: + return _feature_gate_response(exc) + + document = self._get_document(request, document_id) + if document is None: + return Response( + {"error": "Document not found"}, status=status.HTTP_404_NOT_FOUND + ) + + if "active" not in request.data: + return Response( + {"error": "active is required"}, status=status.HTTP_400_BAD_REQUEST + ) + + active = request.data.get("active") + if isinstance(active, str): + active = active.strip().lower() in {"1", "true", "yes"} + document.active = bool(active) + document.save(update_fields=["active"]) + + try: + AsyncRAGService().set_document_active(document.id, document.active) + except Exception as exc: + logger.warning( + "Failed to update vector active metadata for document %s: %s", + document.id, + exc, + ) + + serializer = DocumentSerializer(document) + return Response(serializer.data, status=status.HTTP_200_OK) + + def delete(self, request, document_id): + try: + assert_feature_allowed(request.user, "rag") + except FeatureNotAllowed as exc: + return _feature_gate_response(exc) + + document = self._get_document(request, document_id) + if document is None: + return Response( + {"error": "Document not found"}, status=status.HTTP_404_NOT_FOUND + ) + + document.delete() + return Response(status=status.HTTP_204_NO_CONTENT) diff --git a/llm_be/chat_backend/views_drive.py b/llm_be/chat_backend/views_drive.py new file mode 100644 index 0000000..eb1721a --- /dev/null +++ b/llm_be/chat_backend/views_drive.py @@ -0,0 +1,204 @@ +"""Drive connection management, sync, and webhook API (#47-#52).""" + +from __future__ import annotations + +import logging + +from django.db.models import Q +from django.http import HttpResponse +from rest_framework import permissions, status +from rest_framework.response import Response +from rest_framework.views import APIView + +from finance.services.quotas import FeatureNotAllowed, assert_feature_allowed + +from .models import DriveConnection +from .serializers import DriveConnectionResourcesSerializer, DriveConnectionSerializer +from .services.drive_sync import sync_connection + +logger = logging.getLogger(__name__) + + +def _require_rag_feature(user) -> Response | None: + """Return a 403 Response when the plan doesn't allow Drive/RAG, else None.""" + try: + assert_feature_allowed(user, "rag") + except FeatureNotAllowed as exc: + return Response({"detail": exc.message, "code": exc.code}, status=status.HTTP_403_FORBIDDEN) + return None + + +def _user_can_manage(connection: DriveConnection, user) -> bool: + """Owner may manage personal connections; only managers manage company ones (#50/#51).""" + if connection.kind == DriveConnection.Kind.PERSONAL: + return connection.user_id == user.id + return connection.company_id == user.company_id and bool(user.is_company_manager) + + +class DriveConnectionListView(APIView): + """GET /api/drive/connections/ — user's personal + own-company connections.""" + + def get(self, request): + denied = _require_rag_feature(request.user) + if denied is not None: + return denied + + user = request.user + connections = DriveConnection.objects.filter( + Q(kind=DriveConnection.Kind.PERSONAL, user=user) + | Q(kind=DriveConnection.Kind.COMPANY, company_id=user.company_id) + ).order_by("-created") + serializer = DriveConnectionSerializer(connections, many=True) + return Response(serializer.data) + + +class DriveConnectionDetailView(APIView): + """DELETE /api/drive/connections// — disconnect (owner or company manager).""" + + def delete(self, request, connection_id): + denied = _require_rag_feature(request.user) + if denied is not None: + return denied + + connection = DriveConnection.objects.filter(id=connection_id).first() + if connection is None: + return Response({"detail": "Connection not found."}, status=status.HTTP_404_NOT_FOUND) + if not _user_can_manage(connection, request.user): + return Response( + {"detail": "You are not allowed to disconnect this connection."}, + status=status.HTTP_403_FORBIDDEN, + ) + + connection.is_active = False + connection.access_token = "" + connection.refresh_token = "" + connection.last_sync_status = DriveConnection.SyncStatus.NEVER + connection.save( + update_fields=[ + "is_active", + "access_token", + "refresh_token", + "last_sync_status", + "last_modified", + ] + ) + return Response(status=status.HTTP_204_NO_CONTENT) + + +class DriveConnectionResourcesView(APIView): + """POST /api/drive/connections//resources/ — set selected folder/drive/site ids.""" + + def post(self, request, connection_id): + denied = _require_rag_feature(request.user) + if denied is not None: + return denied + + connection = DriveConnection.objects.filter(id=connection_id).first() + if connection is None: + return Response({"detail": "Connection not found."}, status=status.HTTP_404_NOT_FOUND) + if not _user_can_manage(connection, request.user): + return Response( + {"detail": "You are not allowed to configure this connection."}, + status=status.HTTP_403_FORBIDDEN, + ) + + serializer = DriveConnectionResourcesSerializer(data=request.data) + if not serializer.is_valid(): + return Response(serializer.errors, status=status.HTTP_400_BAD_REQUEST) + + connection.selected_resource_ids = serializer.validated_data["resource_ids"] + connection.selected_resource_labels = serializer.validated_data.get( + "resource_labels", [] + ) + connection.save( + update_fields=["selected_resource_ids", "selected_resource_labels", "last_modified"] + ) + return Response(DriveConnectionSerializer(connection).data) + + +class DriveConnectionSyncView(APIView): + """POST /api/drive/connections//sync/ — trigger a sync now.""" + + def post(self, request, connection_id): + denied = _require_rag_feature(request.user) + if denied is not None: + return denied + + connection = DriveConnection.objects.filter(id=connection_id).first() + if connection is None: + return Response({"detail": "Connection not found."}, status=status.HTTP_404_NOT_FOUND) + if not _user_can_manage(connection, request.user): + return Response( + {"detail": "You are not allowed to sync this connection."}, + status=status.HTTP_403_FORBIDDEN, + ) + if not connection.is_active: + return Response( + {"detail": "This connection is disconnected."}, + status=status.HTTP_400_BAD_REQUEST, + ) + + result = sync_connection(connection) + connection.refresh_from_db() + return Response( + {"result": result, "connection": DriveConnectionSerializer(connection).data} + ) + + +def _maybe_sync_from_webhook(connection_id, provider: str) -> None: + """Best-effort sync trigger for a webhook payload (#52). + + Providers don't reliably identify the local connection without a prior + ``watch``/subscription setup that stores our id as the channel token / + ``clientState``. When we can't resolve a connection, ack with 200 anyway + per the provider contract (retrying would just repeat the same lookup). + """ + if not connection_id: + return + connection = DriveConnection.objects.filter( + id=connection_id, provider=provider, is_active=True + ).first() + if connection is None: + return + try: + sync_connection(connection) + except Exception: + logger.exception( + "Webhook-triggered Drive sync failed for connection=%s", connection_id + ) + + +class DriveWebhookGoogleView(APIView): + """POST /api/drive/webhooks/google/ — Google Drive push notification stub.""" + + permission_classes = (permissions.AllowAny,) + authentication_classes = () + + def post(self, request): + connection_id = request.query_params.get("connection_id") or request.META.get( + "HTTP_X_GOOG_CHANNEL_TOKEN" + ) + _maybe_sync_from_webhook(connection_id, DriveConnection.Provider.GOOGLE) + return Response(status=status.HTTP_200_OK) + + +class DriveWebhookMicrosoftView(APIView): + """POST /api/drive/webhooks/microsoft/ — Microsoft Graph change notification stub.""" + + permission_classes = (permissions.AllowAny,) + authentication_classes = () + + def post(self, request): + validation_token = request.query_params.get("validationToken") + if validation_token: + # Graph subscription-creation handshake: echo the token back as text/plain. + return HttpResponse(validation_token, content_type="text/plain") + + connection_id = request.query_params.get("connection_id") + if not connection_id: + for notification in (request.data or {}).get("value", []): + connection_id = notification.get("clientState") + if connection_id: + break + _maybe_sync_from_webhook(connection_id, DriveConnection.Provider.MICROSOFT) + return Response(status=status.HTTP_200_OK) diff --git a/llm_be/chat_backend/views_oauth.py b/llm_be/chat_backend/views_oauth.py index 11c87c8..de4104e 100644 --- a/llm_be/chat_backend/views_oauth.py +++ b/llm_be/chat_backend/views_oauth.py @@ -11,10 +11,14 @@ from django.urls import reverse from rest_framework import permissions, status from rest_framework.response import Response from rest_framework.views import APIView +from rest_framework_simplejwt.authentication import JWTAuthentication from rest_framework_simplejwt.tokens import RefreshToken -from .models import OAuthIdentity +from finance.services.quotas import FeatureNotAllowed, assert_feature_allowed + +from .models import DriveConnection, OAuthIdentity from .oauth import ( + DRIVE_LINK_INTENTS, OAuthError, build_authorization_url, configured_providers, @@ -22,7 +26,9 @@ from .oauth import ( exchange_code_for_profile, load_oauth_state, provider_configured, + resolve_link_user, resolve_user_from_profile, + upsert_drive_connection, upsert_identity, ) @@ -44,17 +50,32 @@ def _frontend_callback_url(**params: str) -> str: return f"{base}/auth/callback/?{query}" +def _frontend_drive_redirect_url(**params: str) -> str: + """Drive-link callbacks return to Documents storage (#47 / #83).""" + base = settings.FRONTEND_BASE_URL.rstrip("/") + query = urlencode({k: v for k, v in params.items() if v is not None and v != ""}) + return f"{base}/document_storage/?{query}" + + def _redirect_error(code: str, message: str = "") -> HttpResponseRedirect: return HttpResponseRedirect( _frontend_callback_url(error=code, error_description=message or code) ) +def _redirect_drive_error(code: str, message: str = "") -> HttpResponseRedirect: + return HttpResponseRedirect( + _frontend_drive_redirect_url(error=code, error_description=message or code) + ) + + class OAuthStartView(APIView): """Redirect the browser to Google / Microsoft authorize URL.""" permission_classes = (permissions.AllowAny,) - authentication_classes = () + # Login/signup are anonymous; link_drive/link_company_drive need + # request.user, so JWT auth runs but never blocks the anonymous flows. + authentication_classes = (JWTAuthentication,) http_method_names = ["get"] def get(self, request, provider: str): @@ -71,9 +92,14 @@ class OAuthStartView(APIView): ) intent = (request.query_params.get("intent") or "login").lower() - if intent not in {"login", "signup"}: + if intent not in {"login", "signup", *DRIVE_LINK_INTENTS}: return Response( - {"detail": "intent must be 'login' or 'signup'."}, + { + "detail": ( + "intent must be one of 'login', 'signup', " + "'link_drive', 'link_company_drive'." + ) + }, status=status.HTTP_400_BAD_REQUEST, ) if intent == "signup" and not settings.ENABLE_ACCOUNT_REGISTRATION: @@ -82,15 +108,43 @@ class OAuthStartView(APIView): status=status.HTTP_403_FORBIDDEN, ) - state = dump_oauth_state(provider=provider, intent=intent) + user_id = None + if intent in DRIVE_LINK_INTENTS: + user = request.user + if user is None or not user.is_authenticated: + return Response( + {"detail": "Authentication is required to link a Drive account."}, + status=status.HTTP_401_UNAUTHORIZED, + ) + try: + assert_feature_allowed(user, "rag") + except FeatureNotAllowed as exc: + return Response( + {"detail": exc.message, "code": exc.code}, + status=status.HTTP_403_FORBIDDEN, + ) + if intent == "link_company_drive" and not user.is_company_manager: + return Response( + {"detail": "Only a company manager can connect a company Drive."}, + status=status.HTTP_403_FORBIDDEN, + ) + user_id = user.id + + state = dump_oauth_state(provider=provider, intent=intent, user_id=user_id) redirect_uri = _callback_redirect_uri(request, provider) try: auth_url = build_authorization_url( - provider=provider, redirect_uri=redirect_uri, state=state + provider=provider, redirect_uri=redirect_uri, state=state, intent=intent ) except OAuthError as exc: return Response({"detail": exc.message}, status=status.HTTP_400_BAD_REQUEST) + # Authenticated Drive-link flows are started via XHR (Bearer JWT). A + # full-page redirect would drop the Authorization header, so return + # the IdP URL as JSON when the client asks for it (?response=json). + if intent in DRIVE_LINK_INTENTS and request.query_params.get("response") == "json": + return Response({"authorize_url": auth_url}) + return HttpResponseRedirect(auth_url) @@ -121,14 +175,22 @@ class OAuthCallbackView(APIView): state_data = load_oauth_state(state) if state_data["provider"] != provider: raise OAuthError("invalid_state", "OAuth provider mismatch.") + except OAuthError as exc: + logger.info("OAuth callback failed (%s): %s", exc.code, exc.message) + return _redirect_error(exc.code, exc.message) + intent = state_data["intent"] + if intent in DRIVE_LINK_INTENTS: + return self._handle_drive_link_callback( + request, provider=provider, code=code, state_data=state_data + ) + + try: redirect_uri = _callback_redirect_uri(request, provider) profile = exchange_code_for_profile( provider=provider, code=code, redirect_uri=redirect_uri ) - user, created = resolve_user_from_profile( - profile=profile, intent=state_data["intent"] - ) + user, created = resolve_user_from_profile(profile=profile, intent=intent) # Refresh stored tokens on every successful login. upsert_identity(user, profile) except OAuthError as exc: @@ -151,6 +213,47 @@ class OAuthCallbackView(APIView): ) ) + def _handle_drive_link_callback( + self, request, *, provider: str, code: str, state_data: dict + ) -> HttpResponseRedirect: + """Exchange code + upsert a DriveConnection for an already-authenticated user (#47).""" + try: + user = resolve_link_user(state_data["user_id"]) + kind = ( + DriveConnection.Kind.COMPANY + if state_data["intent"] == "link_company_drive" + else DriveConnection.Kind.PERSONAL + ) + if kind == DriveConnection.Kind.COMPANY and not user.is_company_manager: + raise OAuthError( + "forbidden", + "Only a company manager can connect a company Drive.", + ) + assert_feature_allowed(user, "rag") + + redirect_uri = _callback_redirect_uri(request, provider) + profile = exchange_code_for_profile( + provider=provider, code=code, redirect_uri=redirect_uri + ) + connection = upsert_drive_connection(user=user, kind=kind, profile=profile) + except FeatureNotAllowed as exc: + logger.info("Drive link denied by feature gate (%s): %s", exc.code, exc.message) + return _redirect_drive_error(exc.code, exc.message) + except OAuthError as exc: + logger.info("Drive link callback failed (%s): %s", exc.code, exc.message) + return _redirect_drive_error(exc.code, exc.message) + except Exception: + logger.exception("Unexpected Drive link callback failure") + return _redirect_drive_error("server_error", "Unexpected Drive link error.") + + return HttpResponseRedirect( + _frontend_drive_redirect_url( + drive_connected="1", + provider=provider, + kind=connection.kind, + ) + ) + def oauth_public_flags() -> dict: """Feature flags for /public/settings/.""" diff --git a/llm_be/finance/admin.py b/llm_be/finance/admin.py index eea703a..de93b58 100644 --- a/llm_be/finance/admin.py +++ b/llm_be/finance/admin.py @@ -12,13 +12,14 @@ class SubscriptionPlanAdmin(admin.ModelAdmin): "is_public", "is_selectable", "allows_image_generation", + "allows_rag", "allows_all_future_features", "prompt_quota_per_window", "prompt_window_hours", "monthly_token_quota", "sort_order", ) - list_filter = ("is_public", "is_selectable", "allows_image_generation") + list_filter = ("is_public", "is_selectable", "allows_image_generation", "allows_rag") search_fields = ("slug", "name", "stripe_price_id") readonly_fields = ("created", "last_modified") prepopulated_fields = {"slug": ("name",)} diff --git a/llm_be/finance/migrations/0002_subscription_plans_quotas.py b/llm_be/finance/migrations/0002_subscription_plans_quotas.py index 36229e8..700ddfa 100644 --- a/llm_be/finance/migrations/0002_subscription_plans_quotas.py +++ b/llm_be/finance/migrations/0002_subscription_plans_quotas.py @@ -7,9 +7,113 @@ from django.db import migrations, models def seed_plans(apps, schema_editor): - from finance.services.plans import seed_subscription_plans + """Seed the original 5-plan catalog frozen at this migration's schema. - seed_subscription_plans(update_existing=True) + Deliberately does NOT import ``finance.services.plans`` — that module's + ``PLAN_SEED``/model class reflect the *current* code, so a later required + field (e.g. ``allows_rag`` added in #43) would make this historical + RunPython try to write a column that doesn't exist yet when a fresh + database replays migrations in order. Live code re-seeds (and adds any + new fields) via ``seed_subscription_plans()`` calls elsewhere (app + startup, quota checks, test setUp), so this only needs to create the + original rows. + """ + SubscriptionPlan = apps.get_model("finance", "SubscriptionPlan") + seed = [ + { + "slug": "founders", + "name": "Founders", + "description": ( + "Unlimited product access for early supporters: text plus all future " + "capabilities as they ship. $10/mo." + ), + "price_cents": 1000, + "is_public": True, + "is_selectable": True, + "allows_text_generation": True, + "allows_image_generation": True, + "allows_all_future_features": True, + "prompt_quota_per_window": 300, + "prompt_window_hours": 6, + "monthly_token_quota": None, + "sort_order": 10, + }, + { + "slug": "standard", + "name": "Standard", + "description": ( + "Secure conversational chat and coding assistance for developers " + "and privacy-conscious individuals." + ), + "price_cents": 1500, + "is_public": False, + "is_selectable": False, + "allows_text_generation": True, + "allows_image_generation": False, + "allows_all_future_features": False, + "prompt_quota_per_window": 100, + "prompt_window_hours": 6, + "monthly_token_quota": 1_000_000, + "sort_order": 20, + }, + { + "slug": "pro", + "name": "Pro / Creator", + "description": ( + "Higher message caps and multi-modal workflows for heavy users, " + "including image generation when available." + ), + "price_cents": 4000, + "is_public": False, + "is_selectable": False, + "allows_text_generation": True, + "allows_image_generation": True, + "allows_all_future_features": False, + "prompt_quota_per_window": 200, + "prompt_window_hours": 6, + "monthly_token_quota": 3_000_000, + "sort_order": 30, + }, + { + "slug": "business", + "name": "Business Team", + "description": ( + "Team seats, centralized auth, priority support, and absolute data " + "privacy for local companies handling sensitive data." + ), + "price_cents": 9900, + "is_public": False, + "is_selectable": False, + "allows_text_generation": True, + "allows_image_generation": True, + "allows_all_future_features": False, + "prompt_quota_per_window": 300, + "prompt_window_hours": 6, + "monthly_token_quota": 5_000_000, + "sort_order": 40, + }, + { + "slug": "backer", + "name": "Backer", + "description": ( + "Complimentary Founders-level access for pre-approved emails. " + "Not shown at checkout." + ), + "price_cents": 0, + "is_public": False, + "is_selectable": False, + "allows_text_generation": True, + "allows_image_generation": True, + "allows_all_future_features": True, + "prompt_quota_per_window": 300, + "prompt_window_hours": 6, + "monthly_token_quota": None, + "sort_order": 5, + }, + ] + for row in seed: + slug = row.pop("slug") + SubscriptionPlan.objects.get_or_create(slug=slug, defaults=row) class Migration(migrations.Migration): diff --git a/llm_be/finance/migrations/0004_subscriptionplan_allows_rag.py b/llm_be/finance/migrations/0004_subscriptionplan_allows_rag.py new file mode 100644 index 0000000..d026c82 --- /dev/null +++ b/llm_be/finance/migrations/0004_subscriptionplan_allows_rag.py @@ -0,0 +1,21 @@ +# Generated by Django 6.0 on 2026-08-01 20:15 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ("finance", "0003_subscription_cancel_period_fields"), + ] + + operations = [ + migrations.AddField( + model_name="subscriptionplan", + name="allows_rag", + field=models.BooleanField( + default=False, + help_text="Drive/RAG document sync (Google Drive, OneDrive, SharePoint).", + ), + ), + ] diff --git a/llm_be/finance/models.py b/llm_be/finance/models.py index 7084e87..47596f3 100644 --- a/llm_be/finance/models.py +++ b/llm_be/finance/models.py @@ -40,6 +40,10 @@ class SubscriptionPlan(TimeInfoBase): ) allows_text_generation = models.BooleanField(default=True) allows_image_generation = models.BooleanField(default=False) + allows_rag = models.BooleanField( + default=False, + help_text="Drive/RAG document sync (Google Drive, OneDrive, SharePoint).", + ) allows_all_future_features = models.BooleanField( default=False, help_text="Founders/Backer: unlock new capabilities as they ship.", @@ -74,6 +78,8 @@ class SubscriptionPlan(TimeInfoBase): return self.allows_text_generation if feature in ("image", "image_generation"): return self.allows_image_generation + if feature in ("rag", "document_rag"): + return self.allows_rag return False diff --git a/llm_be/finance/serializers.py b/llm_be/finance/serializers.py index fe710cf..bed2fd5 100644 --- a/llm_be/finance/serializers.py +++ b/llm_be/finance/serializers.py @@ -82,5 +82,6 @@ class SubscriptionPlanSerializer(serializers.ModelSerializer): return { "text_generation": obj.allows_feature("text_generation"), "image_generation": obj.allows_feature("image_generation"), + "rag": obj.allows_feature("rag"), "all_future_features": obj.allows_all_future_features, } diff --git a/llm_be/finance/services/plans.py b/llm_be/finance/services/plans.py index e2cad08..9c792e2 100644 --- a/llm_be/finance/services/plans.py +++ b/llm_be/finance/services/plans.py @@ -41,6 +41,7 @@ PLAN_SEED: list[dict[str, Any]] = [ "is_selectable": True, "allows_text_generation": True, "allows_image_generation": True, + "allows_rag": True, "allows_all_future_features": True, "prompt_quota_per_window": 300, "prompt_window_hours": 6, @@ -59,6 +60,7 @@ PLAN_SEED: list[dict[str, Any]] = [ "is_selectable": False, "allows_text_generation": True, "allows_image_generation": False, + "allows_rag": False, "allows_all_future_features": False, "prompt_quota_per_window": 100, "prompt_window_hours": 6, @@ -70,13 +72,14 @@ PLAN_SEED: list[dict[str, Any]] = [ "name": "Pro / Creator", "description": ( "Higher message caps and multi-modal workflows for heavy users, " - "including image generation when available." + "including image generation and Drive/RAG sync when available." ), "price_cents": 4000, "is_public": False, "is_selectable": False, "allows_text_generation": True, "allows_image_generation": True, + "allows_rag": True, "allows_all_future_features": False, "prompt_quota_per_window": 200, "prompt_window_hours": 6, @@ -87,14 +90,16 @@ PLAN_SEED: list[dict[str, Any]] = [ "slug": SubscriptionPlan.Slug.BUSINESS, "name": "Business Team", "description": ( - "Team seats, centralized auth, priority support, and absolute data " - "privacy for local companies handling sensitive data." + "Team seats, centralized auth, priority support, company Drive/RAG " + "sync, and absolute data privacy for local companies handling " + "sensitive data." ), "price_cents": 9900, "is_public": False, "is_selectable": False, "allows_text_generation": True, "allows_image_generation": True, + "allows_rag": True, "allows_all_future_features": False, "prompt_quota_per_window": 300, "prompt_window_hours": 6, @@ -113,6 +118,7 @@ PLAN_SEED: list[dict[str, Any]] = [ "is_selectable": False, "allows_text_generation": True, "allows_image_generation": True, + "allows_rag": True, "allows_all_future_features": True, "prompt_quota_per_window": 300, "prompt_window_hours": 6, @@ -415,6 +421,7 @@ def plan_to_dict(plan: SubscriptionPlan | None) -> dict[str, Any] | None: "features": { "text_generation": plan.allows_feature("text_generation"), "image_generation": plan.allows_feature("image_generation"), + "rag": plan.allows_feature("rag"), "all_future_features": plan.allows_all_future_features, }, "prompt_quota_per_window": plan.prompt_quota_per_window, diff --git a/llm_be/finance/tests/test_plans_quotas.py b/llm_be/finance/tests/test_plans_quotas.py index 6cd1770..b25f9a0 100644 --- a/llm_be/finance/tests/test_plans_quotas.py +++ b/llm_be/finance/tests/test_plans_quotas.py @@ -55,6 +55,24 @@ class PlanCatalogTestCase(TestCase): self.assertFalse(plans["backer"].is_selectable) self.assertTrue(plans["backer"].allows_all_future_features) + def test_seed_allows_rag_matrix(self): + """#43: RAG is gated per-plan — standard is the only tier without it.""" + plans = {p.slug: p for p in seed_subscription_plans()} + + self.assertTrue(plans["founders"].allows_rag) + self.assertFalse(plans["standard"].allows_rag) + self.assertTrue(plans["pro"].allows_rag) + self.assertTrue(plans["business"].allows_rag) + self.assertTrue(plans["backer"].allows_rag) + + def test_allows_feature_recognizes_rag_aliases(self): + plans = {p.slug: p for p in seed_subscription_plans()} + + self.assertTrue(plans["pro"].allows_feature("rag")) + self.assertTrue(plans["pro"].allows_feature("document_rag")) + self.assertFalse(plans["standard"].allows_feature("rag")) + self.assertFalse(plans["standard"].allows_feature("document_rag")) + class BackerRedeemTestCase(TestCase): def setUp(self): @@ -126,6 +144,23 @@ class QuotaGateTestCase(TestCase): ) assert_feature_allowed(self.user, "image_generation") + def test_business_allows_rag(self): + business = SubscriptionPlan.objects.get(slug="business") + assign_plan(self.user, plan=business, source=UserSubscription.Source.ADMIN) + assert_feature_allowed(self.user, "rag") + + def test_feature_gate_blocks_rag_on_standard(self): + with self.assertRaises(FeatureNotAllowed) as ctx: + assert_feature_allowed(self.user, "rag") + self.assertEqual(ctx.exception.code, "feature_not_allowed") + + def test_pro_allows_rag(self): + pro = SubscriptionPlan.objects.get(slug="pro") + assign_plan( + self.user, plan=pro, source=UserSubscription.Source.ADMIN + ) + assert_feature_allowed(self.user, "rag") + def test_token_quota_blocks_when_reported(self): self.plan.monthly_token_quota = 50 self.plan.prompt_quota_per_window = 1000