## Summary Implements epic [#42](#42) (children #43–#53) and advances [#11](#11). - **Entitlement:** `allows_rag` on plans (founders / backer / pro / business; not standard); exposed as `features.rag` - **Gates:** document REST + WS `PromptType.RAG` use `assert_feature_allowed(..., "rag")` - **Lifecycle:** dedupe ingest, delete vectors by `document_id`, honor `active`, fix document detail PATCH/DELETE - **Workspaces:** auto-create default company workspace; fail-closed scoping - **Drive:** personal + company Google/Microsoft connect (`link_drive` / `link_company_drive`), resource selection, sync, webhooks stubs, `sync_drive_connections` management command - **Docs/env:** README + `.env*.example` updated Companion FE: `chat_web_app` branch `feature/rag-epic-42-ui` (#81–#85). ## Test plan - [x] `SKIP_RAG_INIT=1 uv run python manage.py test` (457 OK) - [ ] Migrate finance `0004` + chat_backend `0028` on beta - [ ] Verify Standard user: Documents API 403 + no RAG retrieval - [ ] Verify Founders/Pro: upload + list + active toggle - [ ] Connect Google/Microsoft Drive (incremental scopes) and Sync - [ ] Company manager: `link_company_drive`; non-manager 403 - [ ] Run `manage.py sync_drive_connections`Reviewed-on: #54
This commit was merged in pull request #54.
This commit is contained in:
@@ -32,6 +32,12 @@ ENABLE_ACCOUNT_REGISTRATION=false
|
|||||||
# Redirect URIs (register in each IdP console):
|
# Redirect URIs (register in each IdP console):
|
||||||
# {OAUTH_CALLBACK_BASE_URL}/api/auth/oauth/google/callback/
|
# {OAUTH_CALLBACK_BASE_URL}/api/auth/oauth/google/callback/
|
||||||
# {OAUTH_CALLBACK_BASE_URL}/api/auth/oauth/microsoft/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_ID=
|
||||||
GOOGLE_OAUTH_CLIENT_SECRET=
|
GOOGLE_OAUTH_CLIENT_SECRET=
|
||||||
MICROSOFT_OAUTH_CLIENT_ID=
|
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
|
# Optional; defaults to request host. Example local: http://127.0.0.1:8001
|
||||||
OAUTH_CALLBACK_BASE_URL=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/<id>/
|
||||||
|
# POST /api/drive/connections/<id>/resources/ { "resource_ids": [...] }
|
||||||
|
# POST /api/drive/connections/<id>/sync/
|
||||||
|
# Provider push notifications (best-effort; register with each provider's
|
||||||
|
# subscription/watch API pointing here, using ?connection_id=<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 / finance (optional local — required for checkout + webhooks)
|
||||||
STRIPE_SECRET_KEY=
|
STRIPE_SECRET_KEY=
|
||||||
STRIPE_PUBLISHABLE_KEY=
|
STRIPE_PUBLISHABLE_KEY=
|
||||||
|
|||||||
@@ -53,6 +53,12 @@ ENABLE_ACCOUNT_REGISTRATION=false
|
|||||||
# Register redirect URIs:
|
# Register redirect URIs:
|
||||||
# https://chatbackend.aimloperations.com/api/auth/oauth/google/callback/
|
# https://chatbackend.aimloperations.com/api/auth/oauth/google/callback/
|
||||||
# https://chatbackend.aimloperations.com/api/auth/oauth/microsoft/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_ID=
|
||||||
GOOGLE_OAUTH_CLIENT_SECRET=
|
GOOGLE_OAUTH_CLIENT_SECRET=
|
||||||
MICROSOFT_OAUTH_CLIENT_ID=
|
MICROSOFT_OAUTH_CLIENT_ID=
|
||||||
@@ -60,6 +66,13 @@ MICROSOFT_OAUTH_CLIENT_SECRET=
|
|||||||
MICROSOFT_OAUTH_TENANT=common
|
MICROSOFT_OAUTH_TENANT=common
|
||||||
OAUTH_CALLBACK_BASE_URL=https://chatbackend.aimloperations.com
|
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 / finance
|
||||||
STRIPE_SECRET_KEY=replace-with-stripe-secret-key
|
STRIPE_SECRET_KEY=replace-with-stripe-secret-key
|
||||||
STRIPE_PUBLISHABLE_KEY=replace-with-stripe-publishable-key
|
STRIPE_PUBLISHABLE_KEY=replace-with-stripe-publishable-key
|
||||||
|
|||||||
@@ -93,6 +93,8 @@ with `COMPOSE_DATABASE_URL` if needed.
|
|||||||
| `ENABLE_ACCOUNT_REGISTRATION` | `false` | optional | Self-serve sign-up; keep false until ready |
|
| `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_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 |
|
| `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 |
|
| `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 |
|
| `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) |
|
| `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_started` — first active plan (Checkout, Backer redeem, admin assign)
|
||||||
- `subscription_updated` — plan/status/cancel-at-period-end changes (portal + webhooks)
|
- `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/<google\|microsoft>/start/?intent=link_drive` (authenticated) |
|
||||||
|
| Company | `GET /api/auth/oauth/<google\|microsoft>/start/?intent=link_company_drive` (company manager only) |
|
||||||
|
| Callback | Same `/api/auth/oauth/<provider>/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=<provider>&kind=<personal\|company>` (or `?error=<code>`) |
|
||||||
|
|
||||||
|
**Manage:**
|
||||||
|
|
||||||
|
| Method / path | Notes |
|
||||||
|
|--|--|
|
||||||
|
| `GET /api/drive/connections/` | Caller's personal connections + their company's company connections |
|
||||||
|
| `DELETE /api/drive/connections/<id>/` | Disconnect (owner for personal, company manager for company) — deactivates + clears tokens, keeps history |
|
||||||
|
| `POST /api/drive/connections/<id>/resources/` | `{ "resource_ids": [...], "resource_labels": [...] }` — folder/shared-drive/site ids to sync; empty = provider root |
|
||||||
|
| `POST /api/drive/connections/<id>/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
|
## Security note
|
||||||
|
|
||||||
Secrets previously hardcoded in `settings.py` (email password, captcha, Django
|
Secrets previously hardcoded in `settings.py` (email password, captcha, Django
|
||||||
|
|||||||
@@ -11,6 +11,7 @@ from .models import (
|
|||||||
PromptMetric,
|
PromptMetric,
|
||||||
DocumentWorkspace,
|
DocumentWorkspace,
|
||||||
Document,
|
Document,
|
||||||
|
DriveConnection,
|
||||||
UserAuthEvent,
|
UserAuthEvent,
|
||||||
OutboundEmail,
|
OutboundEmail,
|
||||||
OAuthIdentity,
|
OAuthIdentity,
|
||||||
@@ -213,9 +214,36 @@ class DocumentAdmin(admin.ModelAdmin):
|
|||||||
model = Document
|
model = Document
|
||||||
list_display = (
|
list_display = (
|
||||||
"file",
|
"file",
|
||||||
|
"source",
|
||||||
"active",
|
"active",
|
||||||
"created",
|
"created",
|
||||||
"processed",
|
"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(DocumentWorkspace, DocumentWorkspaceAdmin)
|
||||||
admin.site.register(Document, DocumentAdmin)
|
admin.site.register(Document, DocumentAdmin)
|
||||||
|
admin.site.register(DriveConnection, DriveConnectionAdmin)
|
||||||
|
|
||||||
|
|
||||||
class OAuthIdentityAdmin(admin.ModelAdmin):
|
class OAuthIdentityAdmin(admin.ModelAdmin):
|
||||||
|
|||||||
@@ -497,6 +497,14 @@ class ChatConsumerAgain(AsyncWebsocketConsumer):
|
|||||||
pass
|
pass
|
||||||
|
|
||||||
if prompt_type == PromptType.RAG:
|
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()
|
service = AsyncRAGService()
|
||||||
workspace = await get_workspace(
|
workspace = await get_workspace(
|
||||||
conversation_id, user=chat_user
|
conversation_id, user=chat_user
|
||||||
|
|||||||
@@ -294,8 +294,19 @@ async def generation_node(state: ChatState) -> ChatState:
|
|||||||
pass
|
pass
|
||||||
|
|
||||||
if prompt_type == PromptType.RAG:
|
if prompt_type == PromptType.RAG:
|
||||||
service = AsyncRAGService()
|
|
||||||
chat_user = state.get("chat_user")
|
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)
|
workspace = await get_workspace(conversation_id, user=chat_user)
|
||||||
generator = service.generate_response(messages, prompt_instance.message, workspace)
|
generator = service.generate_response(messages, prompt_instance.message, workspace)
|
||||||
return {"response_generator": generator}
|
return {"response_generator": generator}
|
||||||
|
|||||||
@@ -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'])}"
|
||||||
|
)
|
||||||
+169
@@ -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",
|
||||||
|
),
|
||||||
|
),
|
||||||
|
]
|
||||||
@@ -379,7 +379,87 @@ class DocumentWorkspace(TimeInfoBase):
|
|||||||
company = models.ForeignKey(Company, on_delete=models.CASCADE)
|
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 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)
|
workspace = models.ForeignKey(DocumentWorkspace, on_delete=models.CASCADE)
|
||||||
file = models.FileField(
|
file = models.FileField(
|
||||||
upload_to="documents/",
|
upload_to="documents/",
|
||||||
@@ -389,6 +469,25 @@ class Document(TimeInfoBase):
|
|||||||
uploaded_at = models.DateTimeField(auto_now_add=True)
|
uploaded_at = models.DateTimeField(auto_now_add=True)
|
||||||
processed = models.BooleanField(default=False)
|
processed = models.BooleanField(default=False)
|
||||||
active = 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):
|
class StoredFile(TimeInfoBase):
|
||||||
|
|||||||
@@ -14,7 +14,7 @@ from django.conf import settings
|
|||||||
from django.core import signing
|
from django.core import signing
|
||||||
from django.utils import timezone
|
from django.utils import timezone
|
||||||
|
|
||||||
from .models import Company, CustomUser, OAuthIdentity
|
from .models import Company, CustomUser, DriveConnection, OAuthIdentity
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
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_AUTH_URL = "https://accounts.google.com/o/oauth2/v2/auth"
|
||||||
GOOGLE_TOKEN_URL = "https://oauth2.googleapis.com/token"
|
GOOGLE_TOKEN_URL = "https://oauth2.googleapis.com/token"
|
||||||
GOOGLE_USERINFO_URL = "https://openidconnect.googleapis.com/v1/userinfo"
|
GOOGLE_USERINFO_URL = "https://openidconnect.googleapis.com/v1/userinfo"
|
||||||
|
GOOGLE_DRIVE_READONLY_SCOPE = "https://www.googleapis.com/auth/drive.readonly"
|
||||||
|
|
||||||
MICROSOFT_AUTH_URL_TMPL = (
|
MICROSOFT_AUTH_URL_TMPL = (
|
||||||
"https://login.microsoftonline.com/{tenant}/oauth2/v2.0/authorize"
|
"https://login.microsoftonline.com/{tenant}/oauth2/v2.0/authorize"
|
||||||
@@ -31,6 +32,13 @@ MICROSOFT_AUTH_URL_TMPL = (
|
|||||||
MICROSOFT_TOKEN_URL_TMPL = (
|
MICROSOFT_TOKEN_URL_TMPL = (
|
||||||
"https://login.microsoftonline.com/{tenant}/oauth2/v2.0/token"
|
"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):
|
class OAuthError(Exception):
|
||||||
@@ -76,14 +84,16 @@ def configured_providers() -> dict[str, bool]:
|
|||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
def dump_oauth_state(*, provider: str, intent: str) -> str:
|
def dump_oauth_state(
|
||||||
return signing.dumps(
|
*, provider: str, intent: str, user_id: int | None = None
|
||||||
{"provider": provider, "intent": intent},
|
) -> str:
|
||||||
salt=STATE_SALT,
|
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:
|
try:
|
||||||
data = signing.loads(state, salt=STATE_SALT, max_age=STATE_MAX_AGE_SECONDS)
|
data = signing.loads(state, salt=STATE_SALT, max_age=STATE_MAX_AGE_SECONDS)
|
||||||
except signing.BadSignature as exc:
|
except signing.BadSignature as exc:
|
||||||
@@ -92,25 +102,38 @@ def load_oauth_state(state: str) -> dict[str, str]:
|
|||||||
intent = data.get("intent") or "login"
|
intent = data.get("intent") or "login"
|
||||||
if provider not in OAuthIdentity.Provider.values:
|
if provider not in OAuthIdentity.Provider.values:
|
||||||
raise OAuthError("invalid_state", "Unknown OAuth provider in state.")
|
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.")
|
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:
|
def _microsoft_tenant() -> str:
|
||||||
return settings.MICROSOFT_OAUTH_TENANT or "common"
|
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):
|
if not provider_configured(provider):
|
||||||
raise OAuthError("provider_not_configured", f"{provider} OAuth is not configured.")
|
raise OAuthError("provider_not_configured", f"{provider} OAuth is not configured.")
|
||||||
|
|
||||||
if provider == OAuthIdentity.Provider.GOOGLE:
|
if provider == OAuthIdentity.Provider.GOOGLE:
|
||||||
|
scope = "openid email profile"
|
||||||
|
if intent in DRIVE_LINK_INTENTS:
|
||||||
|
scope = f"{scope} {GOOGLE_DRIVE_READONLY_SCOPE}"
|
||||||
params = {
|
params = {
|
||||||
"client_id": settings.GOOGLE_OAUTH_CLIENT_ID,
|
"client_id": settings.GOOGLE_OAUTH_CLIENT_ID,
|
||||||
"redirect_uri": redirect_uri,
|
"redirect_uri": redirect_uri,
|
||||||
"response_type": "code",
|
"response_type": "code",
|
||||||
"scope": "openid email profile",
|
"scope": scope,
|
||||||
"state": state,
|
"state": state,
|
||||||
"access_type": "offline",
|
"access_type": "offline",
|
||||||
"prompt": "select_account consent",
|
"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)}"
|
return f"{GOOGLE_AUTH_URL}?{urlencode(params)}"
|
||||||
|
|
||||||
if provider == OAuthIdentity.Provider.MICROSOFT:
|
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 = {
|
params = {
|
||||||
"client_id": settings.MICROSOFT_OAUTH_CLIENT_ID,
|
"client_id": settings.MICROSOFT_OAUTH_CLIENT_ID,
|
||||||
"redirect_uri": redirect_uri,
|
"redirect_uri": redirect_uri,
|
||||||
"response_type": "code",
|
"response_type": "code",
|
||||||
"response_mode": "query",
|
"response_mode": "query",
|
||||||
"scope": "openid email profile offline_access",
|
"scope": scope,
|
||||||
"state": state,
|
"state": state,
|
||||||
"prompt": "select_account",
|
"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:
|
def _create_sso_user(profile: ProviderProfile) -> CustomUser:
|
||||||
from finance.services.plans import try_redeem_backer_email
|
from finance.services.plans import try_redeem_backer_email
|
||||||
|
|
||||||
@@ -329,6 +389,14 @@ def _create_sso_user(profile: ProviderProfile) -> CustomUser:
|
|||||||
return user
|
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(
|
def resolve_user_from_profile(
|
||||||
*, profile: ProviderProfile, intent: str
|
*, profile: ProviderProfile, intent: str
|
||||||
) -> tuple[CustomUser, bool]:
|
) -> tuple[CustomUser, bool]:
|
||||||
|
|||||||
@@ -13,6 +13,7 @@ from .models import (
|
|||||||
FEEDBACK_CATEGORIES,
|
FEEDBACK_CATEGORIES,
|
||||||
DocumentWorkspace,
|
DocumentWorkspace,
|
||||||
Document,
|
Document,
|
||||||
|
DriveConnection,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
@@ -279,3 +280,37 @@ class DocumentSerializer(serializers.ModelSerializer):
|
|||||||
"active",
|
"active",
|
||||||
]
|
]
|
||||||
read_only_fields = ["id", "uploaded_at", "processed", "created"]
|
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,
|
||||||
|
)
|
||||||
|
|||||||
@@ -91,6 +91,27 @@ def resolve_chat_user(
|
|||||||
return None
|
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(
|
def resolve_chat_company_scope(
|
||||||
user: CustomUser,
|
user: CustomUser,
|
||||||
conversation_id: Optional[int] = None,
|
conversation_id: Optional[int] = None,
|
||||||
@@ -135,16 +156,7 @@ def resolve_chat_company_scope(
|
|||||||
code="conversation_forbidden",
|
code="conversation_forbidden",
|
||||||
)
|
)
|
||||||
|
|
||||||
workspace = (
|
workspace = ensure_company_workspace(user.company)
|
||||||
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",
|
|
||||||
)
|
|
||||||
|
|
||||||
return ChatCompanyScope(
|
return ChatCompanyScope(
|
||||||
user_id=user.id,
|
user_id=user.id,
|
||||||
|
|||||||
@@ -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
|
||||||
@@ -120,6 +120,7 @@ class RAGService(BaseService):
|
|||||||
"workspace_id": doc.workspace_id,
|
"workspace_id": doc.workspace_id,
|
||||||
"company_id": doc.workspace.company_id,
|
"company_id": doc.workspace.company_id,
|
||||||
"document_id": doc.id,
|
"document_id": doc.id,
|
||||||
|
"active": bool(doc.active),
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
if chunks:
|
if chunks:
|
||||||
@@ -130,6 +131,38 @@ class RAGService(BaseService):
|
|||||||
self.vector_store.persist()
|
self.vector_store.persist()
|
||||||
return docs
|
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:
|
def ingest_documents(self, workspace: DocumentWorkspace | None = None) -> None:
|
||||||
"""Ingest documents from a workspace into the vector store."""
|
"""Ingest documents from a workspace into the vector store."""
|
||||||
print(f"Getting the Document via the workspace: {workspace}")
|
print(f"Getting the Document via the workspace: {workspace}")
|
||||||
@@ -167,7 +200,7 @@ class RAGService(BaseService):
|
|||||||
|
|
||||||
def add_files_to_store(
|
def add_files_to_store(
|
||||||
self,
|
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,
|
workspace_id: str,
|
||||||
source: str = "upload",
|
source: str = "upload",
|
||||||
save_dir: str = "data/uploads",
|
save_dir: str = "data/uploads",
|
||||||
@@ -175,7 +208,9 @@ class RAGService(BaseService):
|
|||||||
"""
|
"""
|
||||||
Process and add files to vector store.
|
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.
|
Paths may be temp files; FileFields are materialized from DB storage.
|
||||||
"""
|
"""
|
||||||
results = {"total_added": 0, "failed_files": [], "processed_files": []}
|
results = {"total_added": 0, "failed_files": [], "processed_files": []}
|
||||||
@@ -188,6 +223,8 @@ class RAGService(BaseService):
|
|||||||
file_tuple[1],
|
file_tuple[1],
|
||||||
file_tuple[2],
|
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):
|
if isinstance(file_ref, str):
|
||||||
file_path = file_ref
|
file_path = file_ref
|
||||||
else:
|
else:
|
||||||
@@ -207,6 +244,8 @@ class RAGService(BaseService):
|
|||||||
"company_id": company_id,
|
"company_id": company_id,
|
||||||
"original_filename": original_name,
|
"original_filename": original_name,
|
||||||
"file_path": original_name,
|
"file_path": original_name,
|
||||||
|
"document_id": document_id,
|
||||||
|
"active": bool(active),
|
||||||
}
|
}
|
||||||
|
|
||||||
docs = self._load_and_split_documents(file_path, metadata)
|
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
|
``company_id`` is written on ingest for defense-in-depth / future dual
|
||||||
filters, but retrieval keys on ``workspace_id`` so older vectors without
|
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:
|
if workspace is None or getattr(workspace, "id", None) is None:
|
||||||
raise ValueError("workspace is required for RAG retrieval")
|
raise ValueError("workspace is required for RAG retrieval")
|
||||||
return {"workspace_id": workspace.id}
|
return {
|
||||||
|
"$and": [
|
||||||
|
{"workspace_id": workspace.id},
|
||||||
|
{"active": True},
|
||||||
|
]
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
class SyncRAGService(RAGService):
|
class SyncRAGService(RAGService):
|
||||||
|
|||||||
@@ -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.dispatch import receiver
|
||||||
from django.conf import settings
|
|
||||||
import os
|
import os
|
||||||
|
|
||||||
from chat_backend.models import Document
|
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"}
|
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)
|
@receiver(post_delete, sender=Document)
|
||||||
def delete_vector_on_remove(sender, instance, **kwargs):
|
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():
|
if _rag_init_skipped():
|
||||||
return
|
return
|
||||||
try:
|
try:
|
||||||
from .services.rag_services import AsyncRAGService
|
from .services.rag_services import AsyncRAGService
|
||||||
|
|
||||||
rag_service = AsyncRAGService()
|
rag_service = AsyncRAGService()
|
||||||
rag_service.ingest_documents()
|
rag_service.delete_document_vectors(instance.id)
|
||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
print(f"Skipping vector update on Document delete: {exc}")
|
print(f"Skipping vector cleanup on Document delete: {exc}")
|
||||||
|
|||||||
@@ -10,6 +10,7 @@ from chat_backend.models import (
|
|||||||
Conversation,
|
Conversation,
|
||||||
Document,
|
Document,
|
||||||
DocumentWorkspace,
|
DocumentWorkspace,
|
||||||
|
DriveConnection,
|
||||||
Prompt,
|
Prompt,
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -129,3 +130,24 @@ def pdf_upload(name: str = "test.pdf") -> SimpleUploadedFile:
|
|||||||
|
|
||||||
def make_document(workspace, name: str = "test.pdf") -> Document:
|
def make_document(workspace, name: str = "test.pdf") -> Document:
|
||||||
return Document.objects.create(workspace=workspace, file=pdf_upload(name))
|
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
|
||||||
|
)
|
||||||
|
|||||||
@@ -1,9 +1,11 @@
|
|||||||
from django.test import TestCase
|
from django.test import TestCase
|
||||||
from rest_framework_simplejwt.tokens import RefreshToken
|
from rest_framework_simplejwt.tokens import RefreshToken
|
||||||
|
|
||||||
|
from chat_backend.models import DocumentWorkspace
|
||||||
from chat_backend.services.chat_tenant_scope import (
|
from chat_backend.services.chat_tenant_scope import (
|
||||||
ChatCompanyScope,
|
ChatCompanyScope,
|
||||||
ChatTenantScopeError,
|
ChatTenantScopeError,
|
||||||
|
ensure_company_workspace,
|
||||||
resolve_chat_company_scope,
|
resolve_chat_company_scope,
|
||||||
resolve_chat_user,
|
resolve_chat_user,
|
||||||
user_from_access_token,
|
user_from_access_token,
|
||||||
@@ -62,3 +64,36 @@ class ChatTenantScopeTestCase(TestCase):
|
|||||||
|
|
||||||
def test_user_from_access_token_rejects_garbage(self):
|
def test_user_from_access_token_rejects_garbage(self):
|
||||||
self.assertIsNone(user_from_access_token("not-a-jwt"))
|
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)
|
||||||
|
|||||||
@@ -367,6 +367,60 @@ class GraphNodeTestCase(TransactionTestCase):
|
|||||||
_args, _kwargs = service.return_value.generate_response.call_args
|
_args, _kwargs = service.return_value.generate_response.call_args
|
||||||
self.assertEqual(_args[2].id, self.workspace.id)
|
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):
|
async def test_generation_node_defaults_to_general_chat(self):
|
||||||
with mock.patch.object(consumers_graph, "AsyncLLMService") as service:
|
with mock.patch.object(consumers_graph, "AsyncLLMService") as service:
|
||||||
service.return_value.generate_response.return_value = "generator"
|
service.return_value.generate_response.return_value = "generator"
|
||||||
|
|||||||
@@ -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())
|
||||||
@@ -11,9 +11,9 @@ from rest_framework import status
|
|||||||
from rest_framework.test import APITestCase
|
from rest_framework.test import APITestCase
|
||||||
from rest_framework_simplejwt.tokens import AccessToken
|
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.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 = {
|
OAUTH_SETTINGS = {
|
||||||
"GOOGLE_OAUTH_CLIENT_ID": "google-client-id",
|
"GOOGLE_OAUTH_CLIENT_ID": "google-client-id",
|
||||||
@@ -239,3 +239,231 @@ class OAuthCallbackTestCase(APITestCase):
|
|||||||
provider="microsoft", subject="ms-oid-1", user=user
|
provider="microsoft", subject="ms-oid-1", user=user
|
||||||
).exists()
|
).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())
|
||||||
|
|||||||
@@ -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())
|
||||||
@@ -200,6 +200,37 @@ class RAGServiceTestCase(TransactionTestCase):
|
|||||||
self.assertEqual(results["total_added"], 1)
|
self.assertEqual(results["total_added"], 1)
|
||||||
self.assertEqual(results["failed_files"], [])
|
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):
|
def test_add_files_to_store_records_failures(self):
|
||||||
results = self.service.add_files_to_store(
|
results = self.service.add_files_to_store(
|
||||||
[("/tmp/does-not-exist.txt", "missing.txt", self.workspace.id)],
|
[("/tmp/does-not-exist.txt", "missing.txt", self.workspace.id)],
|
||||||
@@ -217,6 +248,48 @@ class RAGServiceTestCase(TransactionTestCase):
|
|||||||
first_store.delete_collection.assert_called_once()
|
first_store.delete_collection.assert_called_once()
|
||||||
self.assertIsNot(self.service.vector_store, first_store)
|
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):
|
async def test_search_documents_filters_by_workspace(self):
|
||||||
retriever = self.service.vector_store.as_retriever.return_value
|
retriever = self.service.vector_store.as_retriever.return_value
|
||||||
retriever.aget_relevant_documents = mock.AsyncMock(
|
retriever.aget_relevant_documents = mock.AsyncMock(
|
||||||
@@ -228,7 +301,15 @@ class RAGServiceTestCase(TransactionTestCase):
|
|||||||
self.assertEqual([doc.page_content for doc in docs], ["chunk"])
|
self.assertEqual([doc.page_content for doc in docs], ["chunk"])
|
||||||
self.service.vector_store.as_retriever.assert_called_with(
|
self.service.vector_store.as_retriever.assert_called_with(
|
||||||
search_type="mmr",
|
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):
|
async def test_search_documents_without_workspace_fails_closed(self):
|
||||||
|
|||||||
@@ -3,6 +3,8 @@ from unittest import mock
|
|||||||
|
|
||||||
from django.test import TestCase
|
from django.test import TestCase
|
||||||
|
|
||||||
|
from chat_backend.models import Document
|
||||||
|
|
||||||
from .factories import make_company, make_document, make_workspace
|
from .factories import make_company, make_document, make_workspace
|
||||||
|
|
||||||
RAG_SERVICE = "chat_backend.services.rag_services.AsyncRAGService"
|
RAG_SERVICE = "chat_backend.services.rag_services.AsyncRAGService"
|
||||||
@@ -12,12 +14,13 @@ class DocumentSignalTestCase(TestCase):
|
|||||||
def setUp(self):
|
def setUp(self):
|
||||||
self.workspace = make_workspace(make_company())
|
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.dict(os.environ, {"SKIP_RAG_INIT": ""}):
|
||||||
with mock.patch(RAG_SERVICE) as service:
|
with mock.patch(RAG_SERVICE) as service:
|
||||||
make_document(self.workspace)
|
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):
|
def test_updating_a_document_does_not_reindex(self):
|
||||||
document = make_document(self.workspace)
|
document = make_document(self.workspace)
|
||||||
@@ -29,14 +32,18 @@ class DocumentSignalTestCase(TestCase):
|
|||||||
|
|
||||||
service.assert_not_called()
|
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 = make_document(self.workspace)
|
||||||
|
document_id = document.id
|
||||||
|
|
||||||
with mock.patch.dict(os.environ, {"SKIP_RAG_INIT": ""}):
|
with mock.patch.dict(os.environ, {"SKIP_RAG_INIT": ""}):
|
||||||
with mock.patch(RAG_SERVICE) as service:
|
with mock.patch(RAG_SERVICE) as service:
|
||||||
document.delete()
|
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):
|
def test_skip_rag_init_keeps_signals_inert(self):
|
||||||
with mock.patch.dict(os.environ, {"SKIP_RAG_INIT": "1"}):
|
with mock.patch.dict(os.environ, {"SKIP_RAG_INIT": "1"}):
|
||||||
@@ -46,9 +53,12 @@ class DocumentSignalTestCase(TestCase):
|
|||||||
|
|
||||||
service.assert_not_called()
|
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.dict(os.environ, {"SKIP_RAG_INIT": ""}):
|
||||||
with mock.patch(RAG_SERVICE, side_effect=RuntimeError("chroma down")):
|
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())
|
||||||
|
|||||||
@@ -1,10 +1,14 @@
|
|||||||
|
import os
|
||||||
from unittest import mock
|
from unittest import mock
|
||||||
|
|
||||||
|
from django.test import override_settings
|
||||||
from django.urls import reverse
|
from django.urls import reverse
|
||||||
from rest_framework import status
|
from rest_framework import status
|
||||||
from rest_framework.test import APITestCase
|
from rest_framework.test import APITestCase
|
||||||
|
|
||||||
from chat_backend.models import Document, DocumentWorkspace, StoredFile
|
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 (
|
from .factories import (
|
||||||
make_company,
|
make_company,
|
||||||
@@ -89,14 +93,18 @@ class DocumentUploadViewTestCase(APITestCase):
|
|||||||
|
|
||||||
self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
|
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()
|
self.workspace.delete()
|
||||||
|
|
||||||
response = self.client.post(
|
response = self.client.post(
|
||||||
self.url, {"file": pdf_upload()}, format="multipart"
|
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):
|
def test_list_returns_documents_of_own_workspace(self):
|
||||||
make_document(self.workspace)
|
make_document(self.workspace)
|
||||||
@@ -108,13 +116,19 @@ class DocumentUploadViewTestCase(APITestCase):
|
|||||||
self.assertIn("test", response.data[0]["file"])
|
self.assertIn("test", response.data[0]["file"])
|
||||||
self.assertIn("pdf", 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()
|
self.workspace.delete()
|
||||||
|
|
||||||
response = self.client.get(self.url)
|
response = self.client.get(self.url)
|
||||||
|
|
||||||
self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND)
|
self.assertEqual(response.status_code, status.HTTP_200_OK)
|
||||||
self.assertEqual(response.data["error"], "Workspace not found")
|
self.assertEqual(response.data, [])
|
||||||
|
self.assertTrue(
|
||||||
|
DocumentWorkspace.objects.filter(
|
||||||
|
company=self.company, name="Default"
|
||||||
|
).exists()
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
class DocumentDetailViewTestCase(APITestCase):
|
class DocumentDetailViewTestCase(APITestCase):
|
||||||
@@ -140,3 +154,155 @@ class DocumentDetailViewTestCase(APITestCase):
|
|||||||
|
|
||||||
self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND)
|
self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND)
|
||||||
self.assertEqual(response.data["error"], "Document 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)
|
||||||
|
|||||||
@@ -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)
|
||||||
@@ -27,6 +27,14 @@ from .views import (
|
|||||||
DocumentDetailView,
|
DocumentDetailView,
|
||||||
)
|
)
|
||||||
from .views_oauth import OAuthCallbackView, OAuthStartView
|
from .views_oauth import OAuthCallbackView, OAuthStartView
|
||||||
|
from .views_drive import (
|
||||||
|
DriveConnectionListView,
|
||||||
|
DriveConnectionDetailView,
|
||||||
|
DriveConnectionResourcesView,
|
||||||
|
DriveConnectionSyncView,
|
||||||
|
DriveWebhookGoogleView,
|
||||||
|
DriveWebhookMicrosoftView,
|
||||||
|
)
|
||||||
from rest_framework.routers import DefaultRouter
|
from rest_framework.routers import DefaultRouter
|
||||||
|
|
||||||
|
|
||||||
@@ -107,4 +115,35 @@ urlpatterns = [
|
|||||||
DocumentDetailView.as_view(),
|
DocumentDetailView.as_view(),
|
||||||
name="documents_details",
|
name="documents_details",
|
||||||
),
|
),
|
||||||
|
# drive urls (#47-#52)
|
||||||
|
path(
|
||||||
|
"drive/connections/",
|
||||||
|
DriveConnectionListView.as_view(),
|
||||||
|
name="drive_connections",
|
||||||
|
),
|
||||||
|
path(
|
||||||
|
"drive/connections/<int:connection_id>/",
|
||||||
|
DriveConnectionDetailView.as_view(),
|
||||||
|
name="drive_connection_detail",
|
||||||
|
),
|
||||||
|
path(
|
||||||
|
"drive/connections/<int:connection_id>/resources/",
|
||||||
|
DriveConnectionResourcesView.as_view(),
|
||||||
|
name="drive_connection_resources",
|
||||||
|
),
|
||||||
|
path(
|
||||||
|
"drive/connections/<int:connection_id>/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",
|
||||||
|
),
|
||||||
]
|
]
|
||||||
|
|||||||
@@ -65,8 +65,10 @@ from .email_tasks import (
|
|||||||
send_invite_email,
|
send_invite_email,
|
||||||
send_password_reset_email,
|
send_password_reset_email,
|
||||||
)
|
)
|
||||||
|
from finance.services.quotas import FeatureNotAllowed, assert_feature_allowed
|
||||||
from .services.llm_service import AsyncLLMService
|
from .services.llm_service import AsyncLLMService
|
||||||
from .services.rag_services import AsyncRAGService
|
from .services.rag_services import AsyncRAGService
|
||||||
|
from .services.chat_tenant_scope import ensure_company_workspace
|
||||||
from .services.title_generator import title_generator
|
from .services.title_generator import title_generator
|
||||||
from .services.moderation_classifier import moderation_classifier, ModerationLabel
|
from .services.moderation_classifier import moderation_classifier, ModerationLabel
|
||||||
from .services.prompt_classifier.prompt_classifier import PromptClassifier, PromptType
|
from .services.prompt_classifier.prompt_classifier import PromptClassifier, PromptType
|
||||||
@@ -789,15 +791,30 @@ llm = OllamaLLM(**ollama_llm_kwargs(model=MODEL_NAME))
|
|||||||
|
|
||||||
|
|
||||||
# Document Views
|
# 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):
|
class DocumentWorkspaceView(APIView):
|
||||||
# permission_classes = [permissions.IsAuthenticated]
|
# permission_classes = [permissions.IsAuthenticated]
|
||||||
|
|
||||||
def get(self, request):
|
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)
|
workspaces = DocumentWorkspace.objects.filter(company=request.user.company)
|
||||||
serializer = DocumentWorkspaceSerializer(workspaces, many=True)
|
serializer = DocumentWorkspaceSerializer(workspaces, many=True)
|
||||||
return Response(serializer.data)
|
return Response(serializer.data)
|
||||||
|
|
||||||
def post(self, request):
|
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)
|
serializer = DocumentWorkspaceSerializer(data=request.data)
|
||||||
if serializer.is_valid():
|
if serializer.is_valid():
|
||||||
serializer.save(company=request.user.company)
|
serializer.save(company=request.user.company)
|
||||||
@@ -811,27 +828,25 @@ class DocumentUploadView(APIView):
|
|||||||
def get(self, request):
|
def get(self, request):
|
||||||
logger.debug(f"request_3: {request}")
|
logger.debug(f"request_3: {request}")
|
||||||
try:
|
try:
|
||||||
workspace = DocumentWorkspace.objects.get(company=request.user.company)
|
assert_feature_allowed(request.user, "rag")
|
||||||
|
except FeatureNotAllowed as exc:
|
||||||
|
return _feature_gate_response(exc)
|
||||||
|
|
||||||
|
workspace = ensure_company_workspace(request.user.company)
|
||||||
serializer = DocumentSerializer(
|
serializer = DocumentSerializer(
|
||||||
Document.objects.filter(workspace=workspace), many=True
|
Document.objects.filter(workspace=workspace), many=True
|
||||||
)
|
)
|
||||||
return Response(serializer.data, status=status.HTTP_200_OK)
|
return Response(serializer.data, status=status.HTTP_200_OK)
|
||||||
|
|
||||||
except:
|
|
||||||
return Response(
|
|
||||||
{"error": "Workspace not found"}, status=status.HTTP_404_NOT_FOUND
|
|
||||||
)
|
|
||||||
|
|
||||||
def post(self, request):
|
def post(self, request):
|
||||||
logger.debug(f"request: {request}")
|
logger.debug(f"request: {request}")
|
||||||
|
|
||||||
try:
|
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:
|
workspace = ensure_company_workspace(request.user.company)
|
||||||
return Response(
|
|
||||||
{"error": "Workspace not found"}, status=status.HTTP_404_NOT_FOUND
|
|
||||||
)
|
|
||||||
|
|
||||||
logger.info(request.FILES)
|
logger.info(request.FILES)
|
||||||
file = request.FILES.get("file")
|
file = request.FILES.get("file")
|
||||||
@@ -857,7 +872,15 @@ class DocumentUploadView(APIView):
|
|||||||
document.save()
|
document.save()
|
||||||
service = AsyncRAGService()
|
service = AsyncRAGService()
|
||||||
service.add_files_to_store(
|
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,
|
workspace_id=document.workspace_id,
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -865,16 +888,73 @@ class DocumentUploadView(APIView):
|
|||||||
class DocumentDetailView(APIView):
|
class DocumentDetailView(APIView):
|
||||||
# permission_classes = [permissions.IsAuthenticated]
|
# 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):
|
def get(self, request, document_id):
|
||||||
logger.info(f"request: {request}")
|
logger.info(f"request: {request}")
|
||||||
try:
|
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)
|
document = self._get_document(request, document_id)
|
||||||
except:
|
if document is None:
|
||||||
return Response(
|
return Response(
|
||||||
{"error": "Document not found"}, status=status.HTTP_404_NOT_FOUND
|
{"error": "Document not found"}, status=status.HTTP_404_NOT_FOUND
|
||||||
)
|
)
|
||||||
|
|
||||||
serializer = DocumentWorkspaceSerializer(workspaces, many=True)
|
serializer = DocumentSerializer(document)
|
||||||
return Response(serializer.data)
|
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)
|
||||||
|
|||||||
@@ -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/<id>/ — 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/<id>/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/<id>/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)
|
||||||
@@ -11,10 +11,14 @@ from django.urls import reverse
|
|||||||
from rest_framework import permissions, status
|
from rest_framework import permissions, status
|
||||||
from rest_framework.response import Response
|
from rest_framework.response import Response
|
||||||
from rest_framework.views import APIView
|
from rest_framework.views import APIView
|
||||||
|
from rest_framework_simplejwt.authentication import JWTAuthentication
|
||||||
from rest_framework_simplejwt.tokens import RefreshToken
|
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 (
|
from .oauth import (
|
||||||
|
DRIVE_LINK_INTENTS,
|
||||||
OAuthError,
|
OAuthError,
|
||||||
build_authorization_url,
|
build_authorization_url,
|
||||||
configured_providers,
|
configured_providers,
|
||||||
@@ -22,7 +26,9 @@ from .oauth import (
|
|||||||
exchange_code_for_profile,
|
exchange_code_for_profile,
|
||||||
load_oauth_state,
|
load_oauth_state,
|
||||||
provider_configured,
|
provider_configured,
|
||||||
|
resolve_link_user,
|
||||||
resolve_user_from_profile,
|
resolve_user_from_profile,
|
||||||
|
upsert_drive_connection,
|
||||||
upsert_identity,
|
upsert_identity,
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -44,17 +50,32 @@ def _frontend_callback_url(**params: str) -> str:
|
|||||||
return f"{base}/auth/callback/?{query}"
|
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:
|
def _redirect_error(code: str, message: str = "") -> HttpResponseRedirect:
|
||||||
return HttpResponseRedirect(
|
return HttpResponseRedirect(
|
||||||
_frontend_callback_url(error=code, error_description=message or code)
|
_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):
|
class OAuthStartView(APIView):
|
||||||
"""Redirect the browser to Google / Microsoft authorize URL."""
|
"""Redirect the browser to Google / Microsoft authorize URL."""
|
||||||
|
|
||||||
permission_classes = (permissions.AllowAny,)
|
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"]
|
http_method_names = ["get"]
|
||||||
|
|
||||||
def get(self, request, provider: str):
|
def get(self, request, provider: str):
|
||||||
@@ -71,9 +92,14 @@ class OAuthStartView(APIView):
|
|||||||
)
|
)
|
||||||
|
|
||||||
intent = (request.query_params.get("intent") or "login").lower()
|
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(
|
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,
|
status=status.HTTP_400_BAD_REQUEST,
|
||||||
)
|
)
|
||||||
if intent == "signup" and not settings.ENABLE_ACCOUNT_REGISTRATION:
|
if intent == "signup" and not settings.ENABLE_ACCOUNT_REGISTRATION:
|
||||||
@@ -82,15 +108,43 @@ class OAuthStartView(APIView):
|
|||||||
status=status.HTTP_403_FORBIDDEN,
|
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)
|
redirect_uri = _callback_redirect_uri(request, provider)
|
||||||
try:
|
try:
|
||||||
auth_url = build_authorization_url(
|
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:
|
except OAuthError as exc:
|
||||||
return Response({"detail": exc.message}, status=status.HTTP_400_BAD_REQUEST)
|
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)
|
return HttpResponseRedirect(auth_url)
|
||||||
|
|
||||||
|
|
||||||
@@ -121,14 +175,22 @@ class OAuthCallbackView(APIView):
|
|||||||
state_data = load_oauth_state(state)
|
state_data = load_oauth_state(state)
|
||||||
if state_data["provider"] != provider:
|
if state_data["provider"] != provider:
|
||||||
raise OAuthError("invalid_state", "OAuth provider mismatch.")
|
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)
|
redirect_uri = _callback_redirect_uri(request, provider)
|
||||||
profile = exchange_code_for_profile(
|
profile = exchange_code_for_profile(
|
||||||
provider=provider, code=code, redirect_uri=redirect_uri
|
provider=provider, code=code, redirect_uri=redirect_uri
|
||||||
)
|
)
|
||||||
user, created = resolve_user_from_profile(
|
user, created = resolve_user_from_profile(profile=profile, intent=intent)
|
||||||
profile=profile, intent=state_data["intent"]
|
|
||||||
)
|
|
||||||
# Refresh stored tokens on every successful login.
|
# Refresh stored tokens on every successful login.
|
||||||
upsert_identity(user, profile)
|
upsert_identity(user, profile)
|
||||||
except OAuthError as exc:
|
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:
|
def oauth_public_flags() -> dict:
|
||||||
"""Feature flags for /public/settings/."""
|
"""Feature flags for /public/settings/."""
|
||||||
|
|||||||
@@ -12,13 +12,14 @@ class SubscriptionPlanAdmin(admin.ModelAdmin):
|
|||||||
"is_public",
|
"is_public",
|
||||||
"is_selectable",
|
"is_selectable",
|
||||||
"allows_image_generation",
|
"allows_image_generation",
|
||||||
|
"allows_rag",
|
||||||
"allows_all_future_features",
|
"allows_all_future_features",
|
||||||
"prompt_quota_per_window",
|
"prompt_quota_per_window",
|
||||||
"prompt_window_hours",
|
"prompt_window_hours",
|
||||||
"monthly_token_quota",
|
"monthly_token_quota",
|
||||||
"sort_order",
|
"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")
|
search_fields = ("slug", "name", "stripe_price_id")
|
||||||
readonly_fields = ("created", "last_modified")
|
readonly_fields = ("created", "last_modified")
|
||||||
prepopulated_fields = {"slug": ("name",)}
|
prepopulated_fields = {"slug": ("name",)}
|
||||||
|
|||||||
@@ -7,9 +7,113 @@ from django.db import migrations, models
|
|||||||
|
|
||||||
|
|
||||||
def seed_plans(apps, schema_editor):
|
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):
|
class Migration(migrations.Migration):
|
||||||
|
|||||||
@@ -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).",
|
||||||
|
),
|
||||||
|
),
|
||||||
|
]
|
||||||
@@ -40,6 +40,10 @@ class SubscriptionPlan(TimeInfoBase):
|
|||||||
)
|
)
|
||||||
allows_text_generation = models.BooleanField(default=True)
|
allows_text_generation = models.BooleanField(default=True)
|
||||||
allows_image_generation = models.BooleanField(default=False)
|
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(
|
allows_all_future_features = models.BooleanField(
|
||||||
default=False,
|
default=False,
|
||||||
help_text="Founders/Backer: unlock new capabilities as they ship.",
|
help_text="Founders/Backer: unlock new capabilities as they ship.",
|
||||||
@@ -74,6 +78,8 @@ class SubscriptionPlan(TimeInfoBase):
|
|||||||
return self.allows_text_generation
|
return self.allows_text_generation
|
||||||
if feature in ("image", "image_generation"):
|
if feature in ("image", "image_generation"):
|
||||||
return self.allows_image_generation
|
return self.allows_image_generation
|
||||||
|
if feature in ("rag", "document_rag"):
|
||||||
|
return self.allows_rag
|
||||||
return False
|
return False
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -82,5 +82,6 @@ class SubscriptionPlanSerializer(serializers.ModelSerializer):
|
|||||||
return {
|
return {
|
||||||
"text_generation": obj.allows_feature("text_generation"),
|
"text_generation": obj.allows_feature("text_generation"),
|
||||||
"image_generation": obj.allows_feature("image_generation"),
|
"image_generation": obj.allows_feature("image_generation"),
|
||||||
|
"rag": obj.allows_feature("rag"),
|
||||||
"all_future_features": obj.allows_all_future_features,
|
"all_future_features": obj.allows_all_future_features,
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -41,6 +41,7 @@ PLAN_SEED: list[dict[str, Any]] = [
|
|||||||
"is_selectable": True,
|
"is_selectable": True,
|
||||||
"allows_text_generation": True,
|
"allows_text_generation": True,
|
||||||
"allows_image_generation": True,
|
"allows_image_generation": True,
|
||||||
|
"allows_rag": True,
|
||||||
"allows_all_future_features": True,
|
"allows_all_future_features": True,
|
||||||
"prompt_quota_per_window": 300,
|
"prompt_quota_per_window": 300,
|
||||||
"prompt_window_hours": 6,
|
"prompt_window_hours": 6,
|
||||||
@@ -59,6 +60,7 @@ PLAN_SEED: list[dict[str, Any]] = [
|
|||||||
"is_selectable": False,
|
"is_selectable": False,
|
||||||
"allows_text_generation": True,
|
"allows_text_generation": True,
|
||||||
"allows_image_generation": False,
|
"allows_image_generation": False,
|
||||||
|
"allows_rag": False,
|
||||||
"allows_all_future_features": False,
|
"allows_all_future_features": False,
|
||||||
"prompt_quota_per_window": 100,
|
"prompt_quota_per_window": 100,
|
||||||
"prompt_window_hours": 6,
|
"prompt_window_hours": 6,
|
||||||
@@ -70,13 +72,14 @@ PLAN_SEED: list[dict[str, Any]] = [
|
|||||||
"name": "Pro / Creator",
|
"name": "Pro / Creator",
|
||||||
"description": (
|
"description": (
|
||||||
"Higher message caps and multi-modal workflows for heavy users, "
|
"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,
|
"price_cents": 4000,
|
||||||
"is_public": False,
|
"is_public": False,
|
||||||
"is_selectable": False,
|
"is_selectable": False,
|
||||||
"allows_text_generation": True,
|
"allows_text_generation": True,
|
||||||
"allows_image_generation": True,
|
"allows_image_generation": True,
|
||||||
|
"allows_rag": True,
|
||||||
"allows_all_future_features": False,
|
"allows_all_future_features": False,
|
||||||
"prompt_quota_per_window": 200,
|
"prompt_quota_per_window": 200,
|
||||||
"prompt_window_hours": 6,
|
"prompt_window_hours": 6,
|
||||||
@@ -87,14 +90,16 @@ PLAN_SEED: list[dict[str, Any]] = [
|
|||||||
"slug": SubscriptionPlan.Slug.BUSINESS,
|
"slug": SubscriptionPlan.Slug.BUSINESS,
|
||||||
"name": "Business Team",
|
"name": "Business Team",
|
||||||
"description": (
|
"description": (
|
||||||
"Team seats, centralized auth, priority support, and absolute data "
|
"Team seats, centralized auth, priority support, company Drive/RAG "
|
||||||
"privacy for local companies handling sensitive data."
|
"sync, and absolute data privacy for local companies handling "
|
||||||
|
"sensitive data."
|
||||||
),
|
),
|
||||||
"price_cents": 9900,
|
"price_cents": 9900,
|
||||||
"is_public": False,
|
"is_public": False,
|
||||||
"is_selectable": False,
|
"is_selectable": False,
|
||||||
"allows_text_generation": True,
|
"allows_text_generation": True,
|
||||||
"allows_image_generation": True,
|
"allows_image_generation": True,
|
||||||
|
"allows_rag": True,
|
||||||
"allows_all_future_features": False,
|
"allows_all_future_features": False,
|
||||||
"prompt_quota_per_window": 300,
|
"prompt_quota_per_window": 300,
|
||||||
"prompt_window_hours": 6,
|
"prompt_window_hours": 6,
|
||||||
@@ -113,6 +118,7 @@ PLAN_SEED: list[dict[str, Any]] = [
|
|||||||
"is_selectable": False,
|
"is_selectable": False,
|
||||||
"allows_text_generation": True,
|
"allows_text_generation": True,
|
||||||
"allows_image_generation": True,
|
"allows_image_generation": True,
|
||||||
|
"allows_rag": True,
|
||||||
"allows_all_future_features": True,
|
"allows_all_future_features": True,
|
||||||
"prompt_quota_per_window": 300,
|
"prompt_quota_per_window": 300,
|
||||||
"prompt_window_hours": 6,
|
"prompt_window_hours": 6,
|
||||||
@@ -415,6 +421,7 @@ def plan_to_dict(plan: SubscriptionPlan | None) -> dict[str, Any] | None:
|
|||||||
"features": {
|
"features": {
|
||||||
"text_generation": plan.allows_feature("text_generation"),
|
"text_generation": plan.allows_feature("text_generation"),
|
||||||
"image_generation": plan.allows_feature("image_generation"),
|
"image_generation": plan.allows_feature("image_generation"),
|
||||||
|
"rag": plan.allows_feature("rag"),
|
||||||
"all_future_features": plan.allows_all_future_features,
|
"all_future_features": plan.allows_all_future_features,
|
||||||
},
|
},
|
||||||
"prompt_quota_per_window": plan.prompt_quota_per_window,
|
"prompt_quota_per_window": plan.prompt_quota_per_window,
|
||||||
|
|||||||
@@ -55,6 +55,24 @@ class PlanCatalogTestCase(TestCase):
|
|||||||
self.assertFalse(plans["backer"].is_selectable)
|
self.assertFalse(plans["backer"].is_selectable)
|
||||||
self.assertTrue(plans["backer"].allows_all_future_features)
|
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):
|
class BackerRedeemTestCase(TestCase):
|
||||||
def setUp(self):
|
def setUp(self):
|
||||||
@@ -126,6 +144,23 @@ class QuotaGateTestCase(TestCase):
|
|||||||
)
|
)
|
||||||
assert_feature_allowed(self.user, "image_generation")
|
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):
|
def test_token_quota_blocks_when_reported(self):
|
||||||
self.plan.monthly_token_quota = 50
|
self.plan.monthly_token_quota = 50
|
||||||
self.plan.prompt_quota_per_window = 1000
|
self.plan.prompt_quota_per_window = 1000
|
||||||
|
|||||||
Reference in New Issue
Block a user