Compare commits

...
9 Commits
Author SHA1 Message Date
westfarn 331cff4730 Keep follow-ups on the open chat and send prompt history in order (#85)
Deploy Beta / unit-tests (push) Successful in 9s
Unit Tests / test (push) Successful in 10s
Deploy Beta / docker (push) Successful in 24s
Deploy Beta / deploy-beta (push) Successful in 1m19s
## Summary
- Closes #84.
- Same-chat history now defaults on, and existing accounts are switched on. The account toggle stays an opt-out. Follow-ups such as "I don't know, why?" and "put that list in a table" can see the previous turn.
- Grounded answers may use the conversation for follow-ups. Citations stay on facts that come from live sources. A pure follow-up skips web search; a follow-up that asks for a price or rating still searches.
- Conversation prompts are loaded oldest-first (`created`, then `id`) for the model and for `conversation_details`. The current user turn is dropped from history by matching its text.
- Display order of the bubbles is chat_web_app#136.

## Test plan
- [ ] `manage.py test chat_backend.tests.test_services_llm chat_backend.tests.test_grounding_search.GroundingPrepassTestCase chat_backend.tests.test_views_conversations chat_backend.tests.test_models.CompanyAndUserTestCase chat_backend.tests.test_migration_graph chat_backend.tests.test_consumers.DatabaseHelperTestCase chat_backend.tests.test_services_rag`
- [ ] Apply migration `0036_enable_conversation_context`.
- [ ] Ask for a joke, then "I don't know, why?". The reply explains the joke and does not cite an unrelated page.
- [ ] Ask for stroller options, then ask for that list in a table with prices and ratings. The table uses the products already named. Missing prices are called out as missing, not the list.
- [ ] Turn "Remember this chat" off. The next reply ignores earlier turns. Turn it back on and a follow-up uses them again.
- [ ] Reload a multi-turn chat. Bubbles stay oldest at the top.

Reviewed-on: #85
2026-10-25 06:05:07 -07:00
westfarn 10f02c7cc4 Drop compromised eMudhra roots from outbound TLS trust (#83)
Deploy Beta / unit-tests (push) Successful in 26s
Unit Tests / test (push) Successful in 14s
Deploy Beta / docker (push) Successful in 30s
Deploy Beta / deploy-beta (push) Successful in 12m45s
Outbound TLS in this image still trusted the compromised eMudhra roots (`emSign ECC Root CA - G3` and `Comodo AAA Services`).

- Image install now upgrades `ca-certificates` and regenerates the trust bundle.
- `certifi` is locked at 2026.10.23, which drops those roots for `requests` and `httpx`.

Merge this, then redeploy beta and prod.

Reviewed-on: #83
2026-10-23 09:02:11 -07:00
westfarn fc06202a1f Read attached chat images with a vision model (#82)
Deploy Beta / unit-tests (push) Successful in 10s
Unit Tests / test (push) Successful in 9s
Deploy Beta / docker (push) Successful in 23s
Deploy Beta / deploy-beta (push) Successful in 1m22s
## Summary
- Image attachments (PNG, JPEG, WebP, GIF) skip the document parser and go to a vision-capable Ollama model (`OLLAMA_MODEL_VISION`, default `gemma4:latest`).
- The image is stored on the prompt. `PromptSerializer` returns `file_type` and base64 `image` so history can redraw it.
- Unsupported image types and files over 8 MB return a websocket error and are not saved.
- A later text turn does not reuse a previous image as a spreadsheet.

Pairs with `chat_web_app` `feat/128-vision-image-prompt` ([#128](ai_ml_operations/chat_web_app#128)).

## Test plan
- [ ] `manage.py test chat_backend.tests.test_vision chat_backend.tests.test_consumers.GraphNodeTestCase`
- [ ] Send a prompt plus PNG over `ws/chat_again/` and confirm the vision model is called.
- [ ] `GET conversation_details` includes `image` and `file_type` for that turn.
- [ ] CSV upload still follows the document path.

Reviewed-on: #82
2026-10-20 06:13:31 -07:00
westfarn 32f57c2376 Enable password self-registration and keep SSO off (#81)
Deploy Beta / unit-tests (push) Successful in 10s
Unit Tests / test (push) Successful in 9s
Deploy Beta / docker (push) Successful in 31s
Deploy Beta / deploy-beta (push) Successful in 11m48s
## Summary
- Closes #80.
- Password self-registration is on by default (`ENABLE_ACCOUNT_REGISTRATION`). `/user/create/` issues tokens and `needs_checkout`, and the web sign-up page already redirects into Stripe Checkout.
- Browser SSO stays off (`ENABLE_SSO`, default false). Live Google and Microsoft client ids stay available for Drive linking, but public settings no longer advertise them, and login/signup OAuth start and callback are rejected.

Stripe path checked before this change: Checkout builds a subscription session (Founders at $10/month, or `STRIPE_PRICE_ID`), stores a draft invoice, and signed `checkout.session.completed` / `invoice.paid` webhooks grant the plan. Checkout, webhook, portal, registration, and OAuth tests passed.

## Deploy note
If `chat_backend_prod.env` or `chat_backend_beta.env` on the control node still sets `ENABLE_ACCOUNT_REGISTRATION=false`, that line overrides this default. Set it to `true` or remove it, and leave `ENABLE_SSO` unset or `false`.

## Test plan
- [ ] `GET /api/public/settings/` returns `enable_account_registration: true` and `oauth.google` / `oauth.microsoft` false.
- [ ] Sign-up form creates an account and redirects to Stripe Checkout. No Google or Microsoft button.
- [ ] Checkout success grants the Founders plan after the webhook.
- [ ] `/api/auth/oauth/google/start/?intent=signup` returns 403. Drive connect still starts.
- [ ] Setting `ENABLE_ACCOUNT_REGISTRATION=false` closes `/user/create/` again.

Reviewed-on: #81
2026-10-20 04:44:47 -07:00
westfarn 4fc2b87bb4 Grant mobile plans from verified App Store and Play purchases (#78) (#79)
Deploy Beta / unit-tests (push) Successful in 10s
Unit Tests / test (push) Successful in 9s
Deploy Beta / docker (push) Successful in 23s
Deploy Beta / deploy-beta (push) Successful in 1m19s
## Summary
- Mobile access is granted by `POST /api/finance/store/confirm/` after the app finishes a StoreKit or Play Billing purchase.
- App Store signed transactions are checked against Apple Root CA G3. Play purchase tokens are checked with the Play Developer API.
- The RevenueCat webhook no longer grants a plan. Stripe Checkout and the Stripe webhook stay for the website.
- A purchase account token must match the signed-in user, and a transaction already linked to another account is rejected. Access ends when the verified period end passes unless the app confirms again.

Client: ai_ml_operations/chat_web_app#122

## Config
- `APPLE_IAP_BUNDLE_ID` (default `ai.hesychia.chat`)
- `GOOGLE_PLAY_PACKAGE_NAME`
- `GOOGLE_PLAY_SERVICE_ACCOUNT_JSON` (JSON or a file path)
- Plan `revenuecat_product_id` or `REVENUECAT_PRODUCT_PLAN_MAP` maps the store product to a plan

## Test plan
- [x] `python manage.py test monetization.tests.test_store_purchases monetization.tests.test_revenuecat monetization.tests.test_plans_quotas`
- [ ] Confirm an App Store sandbox transaction and a Play license-test purchase against a configured backend
- [ ] Confirm a token already linked to another user is rejected
- [ ] Confirm `GET /subscription/` returns `store_account_token` and flips a lapsed store sub to canceled

Reviewed-on: #79
2026-10-19 04:37:49 -07:00
westfarn bb5c43a4e7 Make chat_backend migrations a single chain (#77)
Unit Tests / test (push) Successful in 9s
Deploy Beta / unit-tests (push) Successful in 9s
Deploy Beta / deploy-beta (push) Successful in 1m8s
Deploy Beta / docker (push) Successful in 22s
## Summary
- Closes #76.
- `0034_waitlistentry` and `0034_customuser_use_conversation_context` both depended on `0033_agentrun_agentstep`, so the graph had two leaves.
- Waitlist is now `0035_waitlistentry` and depends on `0034_customuser_use_conversation_context`. One leaf: `0035_waitlistentry`.

## Deploy note
If `django_migrations` already has a row named `0034_waitlistentry`, rename it before migrate:

```sql
UPDATE django_migrations
SET name = '0035_waitlistentry'
WHERE app = 'chat_backend' AND name = '0034_waitlistentry';
```

Then run migrate. A database that never applied `0034_waitlistentry` can migrate as usual.

## Test plan
- [x] `chat_backend.tests.test_migration_graph` asserts a single leaf and the new parent.
- [x] `chat_backend.tests.test_public_marketing` still passes.
- [ ] On a database that already applied both 0034 siblings, run the rename SQL, then `migrate`, and confirm no `InconsistentMigrationHistory`.

Reviewed-on: #77
2026-10-19 01:08:54 -07:00
westfarn 257890e600 Public waitlist and usage-count endpoints (#75)
Deploy Beta / unit-tests (push) Failing after 5s
Deploy Beta / docker (push) Skipped
Deploy Beta / deploy-beta (push) Skipped
Unit Tests / test (push) Failing after 5s
## Summary
- `POST /api/public/waitlist/` stores an email. The same response is returned for a new or existing address. No account is created.
- `GET /api/public/usage/` returns the active, non-deleted account count, the document count, and `as_of`. No emails, names, or avatars.
- Both views are `AllowAny` with empty authentication classes, same pattern as public settings.
- Closes #74.

Companion frontend: chat_web_app public homepage PR.

`0034_waitlistentry` depends on `0033_agentrun_agentstep`. `feat/33-conversation-context-rag` also adds a `0034_*` migration, so one of those branches will need a renumber after the other merges.

## Test plan
- [x] `chat_backend.tests.test_public_marketing` — 5 passed
- [ ] `POST /api/public/waitlist/` with a valid email, then the same email again
- [ ] Invalid email returns 400 and stores nothing
- [ ] `GET /api/public/usage/` without a token

Reviewed-on: #75
2026-10-17 05:39:34 -07:00
westfarn ab3dfffa6b Add opt-in use_conversation_context flag for chat history (#33) (#73)
Unit Tests / test (push) Successful in 11s
Deploy Beta / unit-tests (push) Successful in 11s
Deploy Beta / docker (push) Successful in 29s
Deploy Beta / deploy-beta (push) Successful in 7m24s
## Summary
- Closes [#33](#33) — per-user `use_conversation_context` boolean on `CustomUser` (default `false`, opt-in).
- Migration `0034_customuser_use_conversation_context`.
- `GET/POST /api/conversation_preferences` reads/writes the flag; POST sets absolute boolean without toggling `conversation_order` when only this field is sent.
- Chat + document RAG paths skip prior-turn history when the flag is off (`AsyncLLMService`, `AsyncRAGService`, both WS consumers).

## Frontend contract
- Field: `use_conversation_context` (boolean, default false)
- Read: `GET /api/conversation_preferences` → `{ order, use_conversation_context }`
- Write: `POST /api/conversation_preferences` with `{ "use_conversation_context": true|false }`
- Tooltip copy: use previous conversations to better customize the experience
- Counterpart: [chat_web_app#66](ai_ml_operations/chat_web_app#66)

## Test plan
- [ ] `manage.py test chat_backend.tests.test_models chat_backend.tests.test_views_conversations.ConversationPreferencesTestCase chat_backend.tests.test_services_llm chat_backend.tests.test_services_rag`
- [ ] New user / unset flag → `use_conversation_context` is false; history empty in LLM/RAG prompts
- [ ] Enable via preferences API → prior turns appear in history
- [ ] Disable again → history skipped; order preference unchanged when posting only the context flagReviewed-on: #73
2026-08-04 06:27:44 -07:00
westfarn bef151c92c Fix agent progress frames on live WebSocket (#63) (#72)
Deploy Beta / unit-tests (push) Successful in 11s
Unit Tests / test (push) Successful in 11s
Deploy Beta / docker (push) Successful in 23s
Deploy Beta / deploy-beta (push) Successful in 51s
## Summary
- Follow-up to [#71](#71) / [#63](#63).
- **Bugfix:** `run_agentic_turn` discarded `ws_send`, so `run_started` / `plan_ready` / `step_*` / `run_completed` frames only hit an empty Redis group and never the connected client during inline agent runs.
- Thread `ws_send` through `execute_agent_run` / `_broadcast` (WS + channel-layer group).
- Route agentic turns in `consumers_graph` the same way as `consumers.py`.
- Add offline `test_agent_orchestrator.py` coverage.

## Test plan
- [x] `SKIP_RAG_INIT=1 uv run python manage.py test chat_backend.tests.test_agent_orchestrator chat_backend.tests.test_agent_tools chat_backend.tests.test_consumers`
- [ ] Manual: `ALLOW_AGENTIC_TASKS=true`, multi-step prompt — confirm agent frames arrive on the live WS
- [ ] Manual: graph consumer path (`ws/conditional_chat/`) also routes agentic turnsReviewed-on: #72
2026-08-04 06:26:26 -07:00
49 changed files with 2762 additions and 296 deletions

No files matched your search

+13 -4
View File
@@ -16,6 +16,8 @@ OLLAMA_BASE_URL=http://127.0.0.1:11434
# OLLAMA_MODEL=gpt-oss:20b # OLLAMA_MODEL=gpt-oss:20b
# OLLAMA_MODEL_THINKING=gpt-oss:20b # OLLAMA_MODEL_THINKING=gpt-oss:20b
# OLLAMA_MODEL_FAST=gemma4:latest # OLLAMA_MODEL_FAST=gemma4:latest
# Vision (image + prompt). Must be a vision-capable tag. Default gemma4:latest.
# OLLAMA_MODEL_VISION=gemma4:latest
# OLLAMA_MODEL_UTILITY=llama3.2 # OLLAMA_MODEL_UTILITY=llama3.2
# OLLAMA_EMBED_MODEL=nomic-embed-text # OLLAMA_EMBED_MODEL=nomic-embed-text
# OLLAMA_NUM_CTX_THINKING=16384 # OLLAMA_NUM_CTX_THINKING=16384
@@ -39,8 +41,11 @@ EMAIL_USE_TLS=true
# Captcha (optional local) # Captcha (optional local)
CAPTCHA_SECRET_KEY= CAPTCHA_SECRET_KEY=
# Self-serve sign-up (default false — set true to allow /user/create/) # Password self-serve sign-up (default true). Set false to close /user/create/.
ENABLE_ACCOUNT_REGISTRATION=false ENABLE_ACCOUNT_REGISTRATION=true
# Browser SSO. Leave false until Google / Microsoft sign-in should be public.
# Client ids below can still be set for Drive linking.
ENABLE_SSO=false
# OAuth SSO — Google / Microsoft (#24). Leave blank to hide SSO buttons. # OAuth SSO — Google / Microsoft (#24). Leave blank to hide SSO buttons.
# Redirect URIs (register in each IdP console): # Redirect URIs (register in each IdP console):
@@ -83,8 +88,12 @@ STRIPE_WEBHOOK_SECRET=
# Optional: pre-created Stripe Price ID for Founders. When empty, Checkout uses # Optional: pre-created Stripe Price ID for Founders. When empty, Checkout uses
# SubscriptionPlan.price_cents / SUBSCRIPTION_PRICE_* ($10 USD / month Founders). # SubscriptionPlan.price_cents / SUBSCRIPTION_PRICE_* ($10 USD / month Founders).
STRIPE_PRICE_ID= STRIPE_PRICE_ID=
# RevenueCat webhook Authorization bearer secret (store IAP). # Native store purchases (#78). iOS verifies the StoreKit signed transaction
REVENUECAT_WEBHOOK_SECRET= # against Apple Root CA G3. Android calls the Play Developer API.
APPLE_IAP_BUNDLE_ID=ai.hesychia.chat
GOOGLE_PLAY_PACKAGE_NAME=ai.hesychia.chat
# Service-account JSON (one line) or a path to the JSON file. Empty disables Play confirm.
GOOGLE_PLAY_SERVICE_ACCOUNT_JSON=
# Optional JSON map of store product id → plan slug, e.g. # Optional JSON map of store product id → plan slug, e.g.
# REVENUECAT_PRODUCT_PLAN_MAP={"hesychia_founders_monthly":"founders"} # REVENUECAT_PRODUCT_PLAN_MAP={"hesychia_founders_monthly":"founders"}
# SUBSCRIPTION_PRICE_AMOUNT_CENTS=1000 # SUBSCRIPTION_PRICE_AMOUNT_CENTS=1000
+12 -4
View File
@@ -37,6 +37,8 @@ OLLAMA_BASE_URL=http://10.0.0.128:11434
OLLAMA_MODEL=gpt-oss:20b OLLAMA_MODEL=gpt-oss:20b
OLLAMA_MODEL_THINKING=gpt-oss:20b OLLAMA_MODEL_THINKING=gpt-oss:20b
OLLAMA_MODEL_FAST=gemma4:latest OLLAMA_MODEL_FAST=gemma4:latest
# Vision (image + prompt). Must be a vision-capable tag.
OLLAMA_MODEL_VISION=gemma4:latest
OLLAMA_MODEL_UTILITY=llama3.2 OLLAMA_MODEL_UTILITY=llama3.2
OLLAMA_EMBED_MODEL=nomic-embed-text OLLAMA_EMBED_MODEL=nomic-embed-text
OLLAMA_NUM_CTX_THINKING=16384 OLLAMA_NUM_CTX_THINKING=16384
@@ -59,9 +61,13 @@ EMAIL_USE_TLS=true
# Captcha # Captcha
CAPTCHA_SECRET_KEY=replace-with-captcha-secret CAPTCHA_SECRET_KEY=replace-with-captcha-secret
# Self-serve account registration (sign-up). Keep false until ready to open # Password self-serve sign-up. Default true. Set false in
# public registration; set true in chat_backend_prod.env / chat_backend_beta.env. # chat_backend_prod.env / chat_backend_beta.env to close public registration.
ENABLE_ACCOUNT_REGISTRATION=false # An explicit false in that secret overrides the code default.
ENABLE_ACCOUNT_REGISTRATION=true
# Browser SSO stays off. OAuth client ids below are for Drive linking until
# this is set true.
ENABLE_SSO=false
# OAuth SSO — Google / Microsoft (#24). Never commit real secrets. # OAuth SSO — Google / Microsoft (#24). Never commit real secrets.
# Register redirect URIs: # Register redirect URIs:
@@ -91,7 +97,9 @@ OAUTH_CALLBACK_BASE_URL=https://chatbackend.aimloperations.com
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
STRIPE_WEBHOOK_SECRET=replace-with-stripe-webhook-secret STRIPE_WEBHOOK_SECRET=replace-with-stripe-webhook-secret
REVENUECAT_WEBHOOK_SECRET=replace-with-revenuecat-webhook-auth-token APPLE_IAP_BUNDLE_ID=ai.hesychia.chat
GOOGLE_PLAY_PACKAGE_NAME=ai.hesychia.chat
GOOGLE_PLAY_SERVICE_ACCOUNT_JSON=replace-with-play-console-service-account-json
# REVENUECAT_PRODUCT_PLAN_MAP={"hesychia_founders_monthly":"founders"} # REVENUECAT_PRODUCT_PLAN_MAP={"hesychia_founders_monthly":"founders"}
# Optional: pre-created Stripe Price ID. When empty, Checkout uses # Optional: pre-created Stripe Price ID. When empty, Checkout uses
# SUBSCRIPTION_PRICE_* from settings.py ($10 USD / month by default). # SUBSCRIPTION_PRICE_* from settings.py ($10 USD / month by default).
+2 -1
View File
@@ -9,7 +9,8 @@ ENV PYTHONDONTWRITEBYTECODE=1 \
WORKDIR /app WORKDIR /app
RUN apt-get update \ RUN apt-get update \
&& apt-get install -y --no-install-recommends libpq5 \ && apt-get install -y --no-install-recommends ca-certificates libpq5 \
&& update-ca-certificates \
&& rm -rf /var/lib/apt/lists/* && rm -rf /var/lib/apt/lists/*
COPY --from=ghcr.io/astral-sh/uv:latest /uv /usr/local/bin/uv COPY --from=ghcr.io/astral-sh/uv:latest /uv /usr/local/bin/uv
+13 -6
View File
@@ -97,9 +97,12 @@ with `COMPOSE_DATABASE_URL` if needed.
| `SEARXNG_BASE_URL` | `http://127.0.0.1:8088` | yes if using SearxNG | Self-hosted SearxNG JSON API | | `SEARXNG_BASE_URL` | `http://127.0.0.1:8088` | yes if using SearxNG | Self-hosted SearxNG JSON API |
| `EMAIL_HOST_*` | empty | yes (prod/beta) | SMTP2GO | | `EMAIL_HOST_*` | empty | yes (prod/beta) | SMTP2GO |
| `CAPTCHA_SECRET_KEY` | empty | recommended | | | `CAPTCHA_SECRET_KEY` | empty | recommended | |
| `ENABLE_ACCOUNT_REGISTRATION` | `false` | optional | Self-serve sign-up; keep false until ready | | `ENABLE_ACCOUNT_REGISTRATION` | `true` | optional | Password self-serve sign-up. Set `false` to close it |
| `ENABLE_SSO` | `false` | optional | Google / Microsoft sign-in and sign-up. Drive linking does not need this |
| `STRIPE_SECRET_KEY` / `STRIPE_PUBLISHABLE_KEY` / `STRIPE_WEBHOOK_SECRET` | empty | yes for billing | Stripe API + webhook (`monetization` app) | | `STRIPE_SECRET_KEY` / `STRIPE_PUBLISHABLE_KEY` / `STRIPE_WEBHOOK_SECRET` | empty | yes for billing | Stripe API + webhook (`monetization` app) |
| `REVENUECAT_WEBHOOK_SECRET` | empty | yes for store IAP | RevenueCat webhook Authorization bearer | | `APPLE_IAP_BUNDLE_ID` | `ai.hesychia.chat` | yes for iOS IAP | Must match the signed StoreKit transaction |
| `GOOGLE_PLAY_PACKAGE_NAME` | `ai.hesychia.chat` | yes for Play IAP | Android application id |
| `GOOGLE_PLAY_SERVICE_ACCOUNT_JSON` | empty | yes for Play IAP | Play Console service-account JSON (or a file path) |
| `REVENUECAT_PRODUCT_PLAN_MAP` | empty JSON | optional | `{"product_id":"plan_slug"}` fallback map | | `REVENUECAT_PRODUCT_PLAN_MAP` | empty JSON | optional | `{"product_id":"plan_slug"}` fallback map |
| `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 | | `GOOGLE_OAUTH_CLIENT_ID` / `..._SECRET` | empty | for SSO/Drive | Also used for Drive linking (#47), incremental scopes |
@@ -286,10 +289,14 @@ URLs: `/api/monetization/...` and alias `/api/finance/...`.
(`POST /api/finance/portal/`). Local state syncs via (`POST /api/finance/portal/`). Local state syncs via
`customer.subscription.updated` / `deleted` webhooks. `customer.subscription.updated` / `deleted` webhooks.
**Native (Play / App Store):** RevenueCat webhooks at **Native (Play / App Store):** the app buys through StoreKit or Play Billing,
`POST /api/finance/webhooks/revenuecat/` (Authorization bearer = then `POST /api/finance/store/confirm/` with the signed App Store transaction
`REVENUECAT_WEBHOOK_SECRET`) upsert Invoice/Payment ledger rows and sync or the Play purchase token. The backend verifies it (Apple certificate chain,
`UserSubscription` (`source=revenuecat`). or Play Developer API `purchases.subscriptionsv2`) and upserts the ledger.
`source` is `app_store` or `play_store`. There is no store webhook. Renewals
are another confirm from the app; access ends when `current_period_end` passes.
`GET /api/finance/subscription/` includes `store_account_token` to pass into
the store purchase as `appAccountToken` / `obfuscatedAccountId`.
`GET /api/finance/subscription/` includes `cancel_at_period_end` and `GET /api/finance/subscription/` includes `cancel_at_period_end` and
`current_period_end` for Account UI messaging. `current_period_end` for Account UI messaging.
+1
View File
@@ -68,6 +68,7 @@ class CustomUserAdmin(admin.ModelAdmin):
"has_usable_password", "has_usable_password",
"deleted", "deleted",
"has_signed_tos", "has_signed_tos",
"use_conversation_context",
"last_login", "last_login",
"slug", "slug",
"get_set_password_url", "get_set_password_url",
+43 -28
View File
@@ -1,5 +1,4 @@
import json import json
import base64
import logging import logging
import pandas as pd import pandas as pd
from datetime import datetime from datetime import datetime
@@ -14,6 +13,7 @@ from langchain_core.messages import HumanMessage, AIMessage
from langchain_community.vectorstores import Chroma from langchain_community.vectorstores import Chroma
from langchain_ollama import OllamaEmbeddings from langchain_ollama import OllamaEmbeddings
from chat_backend.ollama_config import ( from chat_backend.ollama_config import (
ROLE_VISION,
ollama_embeddings_kwargs, ollama_embeddings_kwargs,
ollama_model_for_role, ollama_model_for_role,
resolve_chat_role, resolve_chat_role,
@@ -38,6 +38,7 @@ 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
from .services.data_analysis_service import AsyncDataAnalysisService from .services.data_analysis_service import AsyncDataAnalysisService
from .services.vision import AsyncVisionService, is_image_type, prepare_attachment
from .services.grounded_chat import prepare_grounded_chat from .services.grounded_chat import prepare_grounded_chat
from .services.status_context import ( from .services.status_context import (
emit_status, emit_status,
@@ -142,7 +143,10 @@ def get_messages(conversation_id, prompt, file_string: str = None, file_type: st
prompt_instance.file_type = file_type prompt_instance.file_type = file_type
prompt_instance.save() prompt_instance.save()
for prompt_obj in Prompt.objects.filter(conversation__id=conversation_id): for prompt_obj in (
Prompt.objects.filter(conversation__id=conversation_id)
.order_by("created", "id")
):
messages.append( messages.append(
{ {
"content": prompt_obj.message, "content": prompt_obj.message,
@@ -278,7 +282,9 @@ async def get_conversation_file_async(conversation_id):
# Get the very first prompt in the conversation that has a file # Get the very first prompt in the conversation that has a file
prompt_with_file = await Prompt.objects.filter( prompt_with_file = await Prompt.objects.filter(
conversation_id=conversation_id conversation_id=conversation_id
).exclude(file='').order_by('created').afirst() ).exclude(file='').exclude(
file_type__in=["png", "jpeg", "jpg", "webp", "gif"]
).order_by('created').afirst()
if prompt_with_file and prompt_with_file.file: if prompt_with_file and prompt_with_file.file:
# Opening a DatabaseStorage file hits the DB, so read inside the thread. # Opening a DatabaseStorage file hits the DB, so read inside the thread.
@@ -414,29 +420,19 @@ class ChatConsumerAgain(AsyncWebsocketConsumer):
decoded_file = None decoded_file = None
if file: if file:
decoded_file = base64.b64decode(file) decoded_file, file_type, attachment_error = prepare_attachment(
logger.debug(decoded_file) file, file_type
# The `altered_message` should only be created if a file exists )
# and you want to pass its content directly to the classifier. if attachment_error:
# Here, we'll let the classifier decide based on the user's prompt await self.send_json_message(
# and then handle the file content separately. json.dumps(
altered_message = message {
if "csv" in file_type: "type": "error",
file_type = "csv" "content": attachment_error,
#altered_message = f"{message}\n The file type is csv and the file contents are: {decoded_file}" }
elif "xmlformats-officedocument" in file_type: )
file_type = "xlsx" )
#df = pd.read_excel(decoded_file) return
#altered_message = f"{message}\n The file type is xlsx and the file contents are: {df}"
elif "word" in file_type:
file_type = "docx"
elif "pdf" in file_type:
file_type = "pdf"
elif "text" in file_type:
file_type = "txt"
#altered_message = f"{message}\n The file type is txt and the file contents are: {decoded_file}"
else:
file_type = "Not Sure"
logger.info(f'received: "{message}" for conversation {conversation_id}') logger.info(f'received: "{message}" for conversation {conversation_id}')
@@ -476,6 +472,15 @@ class ChatConsumerAgain(AsyncWebsocketConsumer):
decoded_file = input_dict.get("decoded_file") decoded_file = input_dict.get("decoded_file")
file_type = input_dict.get("file_type") file_type = input_dict.get("file_type")
if is_image_type(file_type) and decoded_file:
service = AsyncVisionService()
input_dict["_resolved_model"] = service.model_name
input_dict["_citations"] = []
await emit_status("reading_image")
return service.generate_response(
prompt_instance.message, decoded_file, file_type
)
# Feature Flag + plan gate: Image Generation # Feature Flag + plan gate: Image Generation
if prompt_type == PromptType.IMAGE_GENERATION: if prompt_type == PromptType.IMAGE_GENERATION:
if not getattr(settings, "ALLOW_IMAGE_GENERATION", False): if not getattr(settings, "ALLOW_IMAGE_GENERATION", False):
@@ -508,7 +513,12 @@ class ChatConsumerAgain(AsyncWebsocketConsumer):
) )
await emit_status("refining") await emit_status("refining")
return service.generate_response( return service.generate_response(
messages, prompt_instance.message, workspace messages,
prompt_instance.message,
workspace,
use_conversation_context=bool(
getattr(chat_user, "use_conversation_context", False)
),
) )
elif prompt_type == PromptType.DATA_ANALYSIS: elif prompt_type == PromptType.DATA_ANALYSIS:
@@ -565,6 +575,9 @@ class ChatConsumerAgain(AsyncWebsocketConsumer):
messages=messages, messages=messages,
model_name=input_dict.get("model_name"), model_name=input_dict.get("model_name"),
conversation_id=conversation_id, conversation_id=conversation_id,
use_conversation_context=bool(
getattr(chat_user, "use_conversation_context", False)
),
) )
if grounded.error: if grounded.error:
return grounded.error return grounded.error
@@ -595,7 +608,9 @@ class ChatConsumerAgain(AsyncWebsocketConsumer):
# messages = messages[:-1] + [HumanMessage(content=altered_message)] # messages = messages[:-1] + [HumanMessage(content=altered_message)]
# I'll add it to the input_dict if needed. # I'll add it to the input_dict if needed.
resolved_model = ollama_model_for_role(resolve_chat_role(model)) resolved_model = ollama_model_for_role(
ROLE_VISION if is_image_type(file_type) else resolve_chat_role(model)
)
prompt_metric = await create_prompt_metric( prompt_metric = await create_prompt_metric(
prompt_instance.id, prompt_instance.id,
prompt_instance.message, prompt_instance.message,
+124 -36
View File
@@ -1,5 +1,4 @@
import json import json
import base64
import logging import logging
import pandas as pd import pandas as pd
from datetime import datetime from datetime import datetime
@@ -29,6 +28,7 @@ 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
from .services.data_analysis_service import AsyncDataAnalysisService from .services.data_analysis_service import AsyncDataAnalysisService
from .services.vision import AsyncVisionService, is_image_type, prepare_attachment
from .services.grounded_chat import prepare_grounded_chat from .services.grounded_chat import prepare_grounded_chat
from .services.status_context import ( from .services.status_context import (
emit_status, emit_status,
@@ -36,7 +36,11 @@ from .services.status_context import (
set_status_emitter, set_status_emitter,
) )
from .services.ws_frames import citations_frame, status_frame from .services.ws_frames import citations_frame, status_frame
from chat_backend.ollama_config import ollama_model_for_role, resolve_chat_role from chat_backend.ollama_config import (
ROLE_VISION,
ollama_model_for_role,
resolve_chat_role,
)
from .utils import ( from .utils import (
TokenUsageCollector, TokenUsageCollector,
aiter_text_chunks, aiter_text_chunks,
@@ -122,7 +126,10 @@ def get_messages(conversation_id, prompt, file_string: str = None, file_type: st
prompt_instance.file_type = file_type prompt_instance.file_type = file_type
prompt_instance.save() prompt_instance.save()
for prompt_obj in Prompt.objects.filter(conversation__id=conversation_id): for prompt_obj in (
Prompt.objects.filter(conversation__id=conversation_id)
.order_by("created", "id")
):
messages.append( messages.append(
{ {
"content": prompt_obj.message, "content": prompt_obj.message,
@@ -202,7 +209,9 @@ async def get_conversation_file_async(conversation_id):
try: try:
prompt_with_file = await Prompt.objects.filter( prompt_with_file = await Prompt.objects.filter(
conversation_id=conversation_id conversation_id=conversation_id
).exclude(file='').order_by('created').afirst() ).exclude(file='').exclude(
file_type__in=["png", "jpeg", "jpg", "webp", "gif"]
).order_by('created').afirst()
if prompt_with_file and prompt_with_file.file: if prompt_with_file and prompt_with_file.file:
# Opening a DatabaseStorage file hits the DB, so read inside the thread. # Opening a DatabaseStorage file hits the DB, so read inside the thread.
@@ -250,6 +259,10 @@ async def classification_node(state: ChatState) -> ChatState:
prompt_type = await PROMPT_CLASSIFIER.classify_async(msg) prompt_type = await PROMPT_CLASSIFIER.classify_async(msg)
# Images are read by the vision model, not the spreadsheet parser.
if is_image_type(state.get("file_type")):
return {"prompt_type": PromptType.GENERAL_CHAT}
# Override logic # Override logic
if decoded_file and (prompt_type == PromptType.DATA_ANALYSIS or 'analyze' in msg.lower() or 'data' in msg.lower()): if decoded_file and (prompt_type == PromptType.DATA_ANALYSIS or 'analyze' in msg.lower() or 'data' in msg.lower()):
prompt_type = PromptType.DATA_ANALYSIS prompt_type = PromptType.DATA_ANALYSIS
@@ -270,6 +283,17 @@ async def generation_node(state: ChatState) -> ChatState:
decoded_file = state.get("decoded_file") decoded_file = state.get("decoded_file")
file_type = state.get("file_type") file_type = state.get("file_type")
if is_image_type(file_type) and decoded_file:
service = AsyncVisionService()
await emit_status("reading_image")
generator = service.generate_response(
prompt_instance.message, decoded_file, file_type
)
return {
"response_generator": generator,
"resolved_model": service.model_name,
}
# Feature Flag + plan gate: Image Generation # Feature Flag + plan gate: Image Generation
if prompt_type == PromptType.IMAGE_GENERATION: if prompt_type == PromptType.IMAGE_GENERATION:
if not getattr(settings, "ALLOW_IMAGE_GENERATION", False): if not getattr(settings, "ALLOW_IMAGE_GENERATION", False):
@@ -305,7 +329,14 @@ async def generation_node(state: ChatState) -> ChatState:
service = AsyncRAGService() service = AsyncRAGService()
workspace = await get_workspace(conversation_id, user=chat_user) workspace = await get_workspace(conversation_id, user=chat_user)
await emit_status("retrieving_docs") await emit_status("retrieving_docs")
generator = service.generate_response(messages, prompt_instance.message, workspace) generator = service.generate_response(
messages,
prompt_instance.message,
workspace,
use_conversation_context=bool(
getattr(chat_user, "use_conversation_context", False)
),
)
await emit_status("refining") await emit_status("refining")
return {"response_generator": generator} return {"response_generator": generator}
@@ -320,11 +351,15 @@ async def generation_node(state: ChatState) -> ChatState:
else: else:
# GENERAL_CHAT / SEARCH / UNKNOWN — always-on grounding (#62). # GENERAL_CHAT / SEARCH / UNKNOWN — always-on grounding (#62).
# FAST selects a smaller model; it no longer skips search. # FAST selects a smaller model; it no longer skips search.
chat_user = state.get("chat_user")
grounded = await prepare_grounded_chat( grounded = await prepare_grounded_chat(
message=state["message"], message=state["message"],
messages=messages, messages=messages,
model_name=state.get("model_name"), model_name=state.get("model_name"),
conversation_id=conversation_id, conversation_id=conversation_id,
use_conversation_context=bool(
getattr(chat_user, "use_conversation_context", False)
),
) )
if grounded.error: if grounded.error:
return { return {
@@ -474,13 +509,19 @@ class ChatConsumerGraph(AsyncWebsocketConsumer):
print("Conversation ID: ", conversation_id) print("Conversation ID: ", conversation_id)
decoded_file = None decoded_file = None
if file: if file:
decoded_file = base64.b64decode(file) decoded_file, file_type, attachment_error = prepare_attachment(
if "csv" in file_type: file_type = "csv" file, file_type
elif "xmlformats-officedocument" in file_type: file_type = "xlsx" )
elif "word" in file_type: file_type = "docx" if attachment_error:
elif "pdf" in file_type: file_type = "pdf" await self.send_json_message(
elif "text" in file_type: file_type = "txt" json.dumps(
else: file_type = "Not Sure" {
"type": "error",
"content": attachment_error,
}
)
)
return
# Pre-fetch messages and file # Pre-fetch messages and file
messages, prompt_instance = await get_messages( messages, prompt_instance = await get_messages(
@@ -490,7 +531,9 @@ class ChatConsumerGraph(AsyncWebsocketConsumer):
if not decoded_file: if not decoded_file:
decoded_file, file_type = await get_conversation_file_async(conversation_id) decoded_file, file_type = await get_conversation_file_async(conversation_id)
resolved_model = ollama_model_for_role(resolve_chat_role(model)) resolved_model = ollama_model_for_role(
ROLE_VISION if is_image_type(file_type) else resolve_chat_role(model)
)
prompt_metric = await create_prompt_metric( prompt_metric = await create_prompt_metric(
prompt_instance.id, prompt_instance.id,
prompt_instance.message, prompt_instance.message,
@@ -532,44 +575,89 @@ class ChatConsumerGraph(AsyncWebsocketConsumer):
status_token = set_status_emitter(_send_status) status_token = set_status_emitter(_send_status)
try: try:
await emit_status("queued") await emit_status("queued")
# Run Graph (moderation emits moderating; grounding emits evaluating/…)
final_state = await app.ainvoke(initial_state)
print("Final State: ", final_state)
response_generator_or_dict = final_state["response_generator"] # Agentic path (#63) — same gate as consumers.py; skip the
print("Response Generator: ", response_generator_or_dict) # fixed LangGraph when the heuristic says multi-step.
from chat_backend.services.agent import (
run_agentic_turn,
should_use_agent,
)
full_response = "" full_response = ""
tokens_in = tokens_out = None tokens_in = tokens_out = None
citations = []
final_model = resolved_model
if isinstance(response_generator_or_dict, dict): if should_use_agent(message) and not is_image_type(file_type):
content = response_generator_or_dict.get("content", "") try:
await self.send_json_message( await enforce_feature_gate(chat_user, "agentic_tasks")
json.dumps(response_generator_or_dict) except FeatureNotAllowed as exc:
await self.send_json_message(
json.dumps(
{
"type": "error",
"code": exc.code,
"content": exc.message,
}
)
)
await self.send("END_OF_THE_STREAM_ENDER_GAME_42")
return
async def _ws_send(raw: str):
await self.send_json_message(raw)
_run, answer = await run_agentic_turn(
user=chat_user,
scope=tenant_scope,
conversation_id=conversation_id,
goal=message,
prompt=prompt_instance,
ws_send=_ws_send,
) )
full_response = content final_model = _run.model_orchestrator or resolved_model
tokens_in, tokens_out = extract_token_usage(
response_generator_or_dict
)
else:
await emit_status("writing") await emit_status("writing")
usage = TokenUsageCollector() await self.send_json_message(answer)
async for chunk in aiter_text_chunks( full_response = answer
response_generator_or_dict, usage else:
): # Run Graph (moderation emits moderating; grounding emits evaluating/…)
full_response += chunk final_state = await app.ainvoke(initial_state)
await self.send_json_message(chunk) print("Final State: ", final_state)
tokens_in, tokens_out = usage.pair
response_generator_or_dict = final_state["response_generator"]
print("Response Generator: ", response_generator_or_dict)
if isinstance(response_generator_or_dict, dict):
content = response_generator_or_dict.get("content", "")
await self.send_json_message(
json.dumps(response_generator_or_dict)
)
full_response = content
tokens_in, tokens_out = extract_token_usage(
response_generator_or_dict
)
else:
await emit_status("writing")
usage = TokenUsageCollector()
async for chunk in aiter_text_chunks(
response_generator_or_dict, usage
):
full_response += chunk
await self.send_json_message(chunk)
tokens_in, tokens_out = usage.pair
citations = final_state.get("citations") or []
final_model = (
final_state.get("resolved_model") or resolved_model
)
await self.send("END_OF_THE_STREAM_ENDER_GAME_42") await self.send("END_OF_THE_STREAM_ENDER_GAME_42")
citations = final_state.get("citations") or []
if citations: if citations:
await self.send_json_message( await self.send_json_message(
json.dumps(citations_frame(citations)) json.dumps(citations_frame(citations))
) )
final_model = final_state.get("resolved_model") or resolved_model
if final_model and final_model != prompt_metric.model_name: if final_model and final_model != prompt_metric.model_name:
prompt_metric.model_name = final_model prompt_metric.model_name = final_model
await database_sync_to_async(prompt_metric.save)( await database_sync_to_async(prompt_metric.save)(
@@ -0,0 +1,22 @@
from django.db import migrations, models
class Migration(migrations.Migration):
dependencies = [
("chat_backend", "0033_agentrun_agentstep"),
]
operations = [
migrations.AddField(
model_name="customuser",
name="use_conversation_context",
field=models.BooleanField(
default=False,
help_text=(
"When enabled, prior turns in the conversation are used as "
"LLM/RAG context for a more tailored experience"
),
),
),
]
@@ -0,0 +1,37 @@
# Generated by Django 6.0 on 2026-10-17
import django.utils.timezone
from django.db import migrations, models
class Migration(migrations.Migration):
dependencies = [
("chat_backend", "0034_customuser_use_conversation_context"),
]
operations = [
migrations.CreateModel(
name="WaitlistEntry",
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),
),
("email", models.EmailField(max_length=254, unique=True)),
],
options={
"ordering": ["-created"],
},
),
]
@@ -0,0 +1,36 @@
from django.db import migrations, models
def enable_same_chat_history(apps, schema_editor):
"""Turn on thread history for accounts still on the old default (#84).
The column defaulted to false, so an explicit opt-out cannot be told apart
from "never touched". Follow-ups were ignoring the open chat. Accounts that
still want one-shot answers can turn the setting off again.
"""
CustomUser = apps.get_model("chat_backend", "CustomUser")
CustomUser.objects.filter(use_conversation_context=False).update(
use_conversation_context=True
)
class Migration(migrations.Migration):
dependencies = [
("chat_backend", "0035_waitlistentry"),
]
operations = [
migrations.AlterField(
model_name="customuser",
name="use_conversation_context",
field=models.BooleanField(
default=True,
help_text=(
"When enabled, earlier turns in the open chat are included "
"when answering a follow-up. Turn off to answer each message alone."
),
),
),
migrations.RunPython(enable_same_chat_history, migrations.RunPython.noop),
]
+20 -3
View File
@@ -73,6 +73,13 @@ class CustomUser(AbstractUser):
conversation_order = models.BooleanField( conversation_order = models.BooleanField(
default=True, help_text="How the conversations should display" default=True, help_text="How the conversations should display"
) )
use_conversation_context = models.BooleanField(
default=True,
help_text=(
"When enabled, earlier turns in the open chat are included "
"when answering a follow-up. Turn off to answer each message alone."
),
)
def get_set_password_url(self): def get_set_password_url(self):
from django.conf import settings from django.conf import settings
@@ -518,9 +525,7 @@ class DriveConnection(TimeInfoBase):
help_text="Required for company connections; optional for personal (#55).", help_text="Required for company connections; optional for personal (#55).",
) )
provider = models.CharField(max_length=32, choices=Provider.choices) provider = models.CharField(max_length=32, choices=Provider.choices)
kind = models.CharField( kind = models.CharField(max_length=16, choices=Kind.choices, default=Kind.PERSONAL)
max_length=16, choices=Kind.choices, default=Kind.PERSONAL
)
access_token = models.TextField(blank=True, default="") access_token = models.TextField(blank=True, default="")
refresh_token = models.TextField(blank=True, default="") refresh_token = models.TextField(blank=True, default="")
token_expires_at = models.DateTimeField(null=True, blank=True) token_expires_at = models.DateTimeField(null=True, blank=True)
@@ -792,3 +797,15 @@ class StoredFile(TimeInfoBase):
def __str__(self): def __str__(self):
return self.name return self.name
class WaitlistEntry(TimeInfoBase):
"""Email captured from the public site when self-serve registration is off."""
email = models.EmailField(unique=True)
class Meta:
ordering = ["-created"]
def __str__(self):
return self.email
+2 -2
View File
@@ -459,8 +459,8 @@ def resolve_user_from_profile(
upsert_identity(email_user, profile) upsert_identity(email_user, profile)
return email_user, False return email_user, False
# New account path # New account path. Password registration does not enable SSO sign-up.
allow_create = settings.ENABLE_ACCOUNT_REGISTRATION allow_create = settings.ENABLE_ACCOUNT_REGISTRATION and settings.ENABLE_SSO
if intent == "signup" and not allow_create: if intent == "signup" and not allow_create:
raise OAuthError("registration_disabled", "Account registration is disabled.") raise OAuthError("registration_disabled", "Account registration is disabled.")
if intent == "login" and not allow_create: if intent == "login" and not allow_create:
+9 -2
View File
@@ -17,6 +17,8 @@ ROLE_EMBED = "embed"
# independent plan steps concurrently on a smaller/cheaper model. # independent plan steps concurrently on a smaller/cheaper model.
ROLE_ORCHESTRATOR = "orchestrator" ROLE_ORCHESTRATOR = "orchestrator"
ROLE_SUBAGENT = "subagent" ROLE_SUBAGENT = "subagent"
# Image + prompt turns (#128). Never falls back to the text chat model.
ROLE_VISION = "vision"
_VALID_ROLES = { _VALID_ROLES = {
ROLE_THINKING, ROLE_THINKING,
@@ -25,6 +27,7 @@ _VALID_ROLES = {
ROLE_EMBED, ROLE_EMBED,
ROLE_ORCHESTRATOR, ROLE_ORCHESTRATOR,
ROLE_SUBAGENT, ROLE_SUBAGENT,
ROLE_VISION,
} }
@@ -56,6 +59,7 @@ def ollama_model_for_role(role: str) -> str:
ROLE_EMBED: "OLLAMA_EMBED_MODEL", ROLE_EMBED: "OLLAMA_EMBED_MODEL",
ROLE_ORCHESTRATOR: "OLLAMA_MODEL_ORCHESTRATOR", ROLE_ORCHESTRATOR: "OLLAMA_MODEL_ORCHESTRATOR",
ROLE_SUBAGENT: "OLLAMA_MODEL_SUBAGENT", ROLE_SUBAGENT: "OLLAMA_MODEL_SUBAGENT",
ROLE_VISION: "OLLAMA_MODEL_VISION",
}[role] }[role]
role_default = { role_default = {
ROLE_THINKING: "gpt-oss:20b", ROLE_THINKING: "gpt-oss:20b",
@@ -64,14 +68,15 @@ def ollama_model_for_role(role: str) -> str:
ROLE_EMBED: "nomic-embed-text", ROLE_EMBED: "nomic-embed-text",
ROLE_ORCHESTRATOR: "gpt-oss:20b", ROLE_ORCHESTRATOR: "gpt-oss:20b",
ROLE_SUBAGENT: "llama3.2", ROLE_SUBAGENT: "llama3.2",
ROLE_VISION: "gemma4:latest",
}[role] }[role]
configured = getattr(settings, role_setting, None) configured = getattr(settings, role_setting, None)
if configured: if configured:
return configured return configured
# Embeddings must never silently fall back to a chat model (#62). # Embeddings and vision must never silently fall back to a text chat model.
if role == ROLE_EMBED: if role in (ROLE_EMBED, ROLE_VISION):
return role_default return role_default
legacy = getattr(settings, "OLLAMA_MODEL", None) legacy = getattr(settings, "OLLAMA_MODEL", None)
@@ -94,6 +99,8 @@ def ollama_num_ctx_for_role(role: str) -> int:
return int(getattr(settings, "OLLAMA_NUM_CTX_SUBAGENT", 8192) or 8192) return int(getattr(settings, "OLLAMA_NUM_CTX_SUBAGENT", 8192) or 8192)
if role == ROLE_ORCHESTRATOR: if role == ROLE_ORCHESTRATOR:
return int(getattr(settings, "OLLAMA_NUM_CTX_ORCHESTRATOR", 16384) or 16384) return int(getattr(settings, "OLLAMA_NUM_CTX_ORCHESTRATOR", 16384) or 16384)
if role == ROLE_VISION:
return int(getattr(settings, "OLLAMA_NUM_CTX_VISION", 8192) or 8192)
return int(getattr(settings, "OLLAMA_NUM_CTX_THINKING", 16384) or 16384) return int(getattr(settings, "OLLAMA_NUM_CTX_THINKING", 16384) or 16384)
+9 -1
View File
@@ -257,6 +257,7 @@ class PromptSerializer(serializers.ModelSerializer):
tokens_in = serializers.SerializerMethodField() tokens_in = serializers.SerializerMethodField()
tokens_out = serializers.SerializerMethodField() tokens_out = serializers.SerializerMethodField()
feedback = serializers.SerializerMethodField() feedback = serializers.SerializerMethodField()
image = serializers.SerializerMethodField()
class Meta: class Meta:
model = Prompt model = Prompt
@@ -269,8 +270,10 @@ class PromptSerializer(serializers.ModelSerializer):
"tokens_out", "tokens_out",
"citations", "citations",
"feedback", "feedback",
"file_type",
"image",
) )
read_only_fields = ("citations", "feedback") read_only_fields = ("citations", "feedback", "file_type", "image")
def _token_pair(self, obj): def _token_pair(self, obj):
cache = self.context.setdefault("_prompt_token_cache", {}) cache = self.context.setdefault("_prompt_token_cache", {})
@@ -314,6 +317,11 @@ class PromptSerializer(serializers.ModelSerializer):
self.context["_prompt_feedback_by_id"] = by_prompt self.context["_prompt_feedback_by_id"] = by_prompt
return by_prompt.get(obj.id) return by_prompt.get(obj.id)
def get_image(self, obj):
from chat_backend.services.vision import image_base64
return image_base64(obj)
def validate_message(self, value: str) -> str: def validate_message(self, value: str) -> str:
if value is None or not str(value).strip(): if value is None or not str(value).strip():
raise serializers.ValidationError("Message text cannot be empty.") raise serializers.ValidationError("Message text cannot be empty.")
+46 -13
View File
@@ -11,6 +11,7 @@ full history for clients that missed frames.
from __future__ import annotations from __future__ import annotations
import json
import logging import logging
from asgiref.sync import sync_to_async from asgiref.sync import sync_to_async
@@ -103,11 +104,19 @@ def _upsert_step(run_id: int, event_type: str, data: dict):
return step return step
async def _broadcast(run, event_type: str, data: dict) -> None: async def _broadcast(run, event_type: str, data: dict, *, ws_send=None) -> None:
"""Emit one agent frame both on the originating WS (if bound) and the
run's channel-layer group (so a reconnect / ``GET /api/agent_runs/<id>/``
still sees it)."""
frame = agent_frame(event_type, {"run_id": str(run.pk), **data})
if ws_send is not None:
try:
await ws_send(json.dumps(frame))
except Exception: # pragma: no cover - WS send must never break a run
logger.exception("Failed to send agent frame on WS run=%s type=%s", run.pk, event_type)
layer = get_channel_layer() layer = get_channel_layer()
if layer is None: if layer is None:
return return
frame = agent_frame(event_type, {"run_id": str(run.pk), **data})
try: try:
await layer.group_send( await layer.group_send(
run.channel_group_name(), run.channel_group_name(),
@@ -131,8 +140,15 @@ def _build_llms(tools: list):
return planner_llm, subagent_llm_factory return planner_llm, subagent_llm_factory
async def execute_agent_run(run_id: int) -> None: async def execute_agent_run(run_id: int, *, ws_send=None) -> None:
"""Load, execute, and persist the outcome of one :class:`AgentRun`.""" """Load, execute, and persist the outcome of one :class:`AgentRun`.
``ws_send`` is optional — set by :func:`run_agentic_turn` when a run is
started inline on a live WS connection so frames land there immediately,
in addition to the channel-layer group broadcast every run always gets
(for reconnects / ``GET /api/agent_runs/<id>/`` backfill). Background
dispatch via :mod:`chat_backend.services.agent.tasks` leaves it unset.
"""
run = await _load_run(run_id) run = await _load_run(run_id)
if run is None: if run is None:
logger.error("execute_agent_run: AgentRun %s not found", run_id) logger.error("execute_agent_run: AgentRun %s not found", run_id)
@@ -146,7 +162,12 @@ async def execute_agent_run(run_id: int) -> None:
run.status = AgentRunModel.Status.PLANNING run.status = AgentRunModel.Status.PLANNING
run.started_at = timezone.now() run.started_at = timezone.now()
await _save_run(run, ["status", "started_at"]) await _save_run(run, ["status", "started_at"])
await _broadcast(run, "run_started", {"title": run.title or run.goal[:80], "status": run.status}) await _broadcast(
run,
"run_started",
{"title": run.title or run.goal[:80], "status": run.status},
ws_send=ws_send,
)
scope, workspace = await _resolve_scope_and_workspace( scope, workspace = await _resolve_scope_and_workspace(
run.user, run.conversation_id run.user, run.conversation_id
@@ -166,7 +187,7 @@ async def execute_agent_run(run_id: int) -> None:
async def on_event(event_type: str, data: dict) -> None: async def on_event(event_type: str, data: dict) -> None:
if event_type in ("step_started", "step_completed", "step_failed"): if event_type in ("step_started", "step_completed", "step_failed"):
await _upsert_step(run.pk, event_type, data) await _upsert_step(run.pk, event_type, data)
await _broadcast(run, event_type, data) await _broadcast(run, event_type, data, ws_send=ws_send)
async def is_cancelled() -> bool: async def is_cancelled() -> bool:
return await _refresh_cancel_flag(run.pk) return await _refresh_cancel_flag(run.pk)
@@ -210,6 +231,7 @@ async def execute_agent_run(run_id: int) -> None:
run, run,
"run_completed", "run_completed",
{"status": run.status, "result": final_answer, "title": run.title}, {"status": run.status, "result": final_answer, "title": run.title},
ws_send=ws_send,
) )
except AgentCancelled: except AgentCancelled:
run.status = AgentRunModel.Status.CANCELLED run.status = AgentRunModel.Status.CANCELLED
@@ -219,7 +241,9 @@ async def execute_agent_run(run_id: int) -> None:
await _save_run( await _save_run(
run, ["status", "completed_at", "tool_call_count", "iteration_count"] run, ["status", "completed_at", "tool_call_count", "iteration_count"]
) )
await _broadcast(run, "run_completed", {"status": run.status, "title": run.title}) await _broadcast(
run, "run_completed", {"status": run.status, "title": run.title}, ws_send=ws_send
)
except AgentRunLimitExceeded as exc: except AgentRunLimitExceeded as exc:
run.status = AgentRunModel.Status.FAILED run.status = AgentRunModel.Status.FAILED
run.error = str(exc) run.error = str(exc)
@@ -231,7 +255,10 @@ async def execute_agent_run(run_id: int) -> None:
["status", "error", "completed_at", "tool_call_count", "iteration_count"], ["status", "error", "completed_at", "tool_call_count", "iteration_count"],
) )
await _broadcast( await _broadcast(
run, "run_completed", {"status": run.status, "error": str(exc), "title": run.title} run,
"run_completed",
{"status": run.status, "error": str(exc), "title": run.title},
ws_send=ws_send,
) )
except Exception as exc: # pragma: no cover - defensive top-level guard except Exception as exc: # pragma: no cover - defensive top-level guard
logger.exception("Agent run %s failed", run_id) logger.exception("Agent run %s failed", run_id)
@@ -245,7 +272,10 @@ async def execute_agent_run(run_id: int) -> None:
["status", "error", "completed_at", "tool_call_count", "iteration_count"], ["status", "error", "completed_at", "tool_call_count", "iteration_count"],
) )
await _broadcast( await _broadcast(
run, "run_completed", {"status": run.status, "error": str(exc), "title": run.title} run,
"run_completed",
{"status": run.status, "error": str(exc), "title": run.title},
ws_send=ws_send,
) )
@@ -260,8 +290,12 @@ async def run_agentic_turn(
): ):
"""Create an AgentRun and execute it inline for the originating WS turn. """Create an AgentRun and execute it inline for the originating WS turn.
``ws_send`` / ``scope`` accepted for API compatibility with consumers. ``scope`` is accepted for API compatibility with consumers (tenant scope
Progress broadcasts via the channel-layer group. Returns ``(run, answer)``. was already validated there) but re-derived from ``user``/
``conversation_id`` inside :func:`execute_agent_run` to keep a single
resolution path. ``ws_send`` — when provided — receives every
``agent_*`` frame live, in addition to the channel-layer group broadcast
every run always gets. Returns ``(run, answer)``.
""" """
from django.conf import settings from django.conf import settings
from chat_backend.models import AgentRun from chat_backend.models import AgentRun
@@ -273,7 +307,6 @@ async def run_agentic_turn(
from chat_backend.services.status_context import emit_status from chat_backend.services.status_context import emit_status
del scope del scope
del ws_send
if not getattr(settings, "ALLOW_AGENTIC_TASKS", False): if not getattr(settings, "ALLOW_AGENTIC_TASKS", False):
raise RuntimeError("ALLOW_AGENTIC_TASKS is disabled") raise RuntimeError("ALLOW_AGENTIC_TASKS is disabled")
@@ -295,7 +328,7 @@ async def run_agentic_turn(
), ),
) )
await emit_status("evaluating", "Planning multi-step task") await emit_status("evaluating", "Planning multi-step task")
await execute_agent_run(run.pk) await execute_agent_run(run.pk, ws_send=ws_send)
finished = await _load_run(run.pk) finished = await _load_run(run.pk)
answer = ( answer = (
(finished.result if finished else "") (finished.result if finished else "")
@@ -13,12 +13,18 @@ ASSISTANT_SYSTEM_PROMPT = (
) )
GROUNDED_ANSWER_INSTRUCTIONS = ( GROUNDED_ANSWER_INSTRUCTIONS = (
"You have been given numbered live sources. Answer ONLY from those sources. " "You may have numbered live sources and the conversation so far. "
"Cite source indexes inline like [1] or [2]. " "Use the conversation for follow-ups: explain a previous answer, reformat it "
"If the sources do not settle the question, say so explicitly — do not fill " "(including a table of items already named), and reply to lines such as "
"gaps from memory or training data. " '"I don\'t know, why?". Do not say that list or answer is missing when it '
"is already in the conversation. "
"Cite source indexes inline like [1] or [2] only for facts taken from live sources. "
"If a new external fact (price, rating, date, news) is in neither the conversation "
"nor the sources, say that specific fact is unavailable, and still complete the "
"part the conversation already answers. "
"When sources conflict, prefer the most recent dated source. " "When sources conflict, prefer the most recent dated source. "
"Never state a date, number, name, or event that does not appear in the sources." "Never invent a date, number, name, or event that is in neither the conversation "
"nor the sources."
) )
RETRIEVAL_FAILED_MESSAGE = ( RETRIEVAL_FAILED_MESSAGE = (
+13 -4
View File
@@ -78,6 +78,7 @@ async def prepare_grounded_chat(
model_name: str | None, model_name: str | None,
conversation_id: int, conversation_id: int,
on_status: StatusCallback | None = None, on_status: StatusCallback | None = None,
use_conversation_context: bool = False,
) -> GroundedTurnResult: ) -> GroundedTurnResult:
"""Decide grounding, retrieve, and return an AsyncLLMService generator. """Decide grounding, retrieve, and return an AsyncLLMService generator.
@@ -86,27 +87,33 @@ async def prepare_grounded_chat(
``on_status(stage, detail)`` emits live activity frames (#96) when provided; ``on_status(stage, detail)`` emits live activity frames (#96) when provided;
otherwise uses the context-bound emitter from :mod:`status_context`. otherwise uses the context-bound emitter from :mod:`status_context`.
``use_conversation_context`` gates prior-turn history in the LLM prompt.
The account setting defaults on so follow-ups see this chat (#84).
""" """
gen_kwargs = {"use_conversation_context": use_conversation_context}
internet = getattr(settings, "ALLOW_INTERNET_ACCESS", False) internet = getattr(settings, "ALLOW_INTERNET_ACCESS", False)
if not internet: if not internet:
await _emit(on_status, "refining") await _emit(on_status, "refining")
service = build_chat_service(model_name=model_name, grounded=False) service = build_chat_service(model_name=model_name, grounded=False)
return GroundedTurnResult( return GroundedTurnResult(
generator=service.generate_response( generator=service.generate_response(
messages, message, conversation_id messages, message, conversation_id, **gen_kwargs
), ),
model_name=service.model_name, model_name=service.model_name,
) )
await _emit(on_status, "evaluating") await _emit(on_status, "evaluating")
decision = await grounding_decider.decide_async(message) # Prior turns exist once the current user message is not the only row.
has_history = bool(use_conversation_context) and len(messages) > 1
decision = await grounding_decider.decide_async(message, has_history=has_history)
if not decision.needs_retrieval: if not decision.needs_retrieval:
await _emit(on_status, "refining") await _emit(on_status, "refining")
service = build_chat_service(model_name=model_name, grounded=False) service = build_chat_service(model_name=model_name, grounded=False)
return GroundedTurnResult( return GroundedTurnResult(
generator=service.generate_response( generator=service.generate_response(
messages, message, conversation_id messages, message, conversation_id, **gen_kwargs
), ),
decision=decision, decision=decision,
model_name=service.model_name, model_name=service.model_name,
@@ -149,7 +156,9 @@ async def prepare_grounded_chat(
sources_block=sources_block, sources_block=sources_block,
) )
return GroundedTurnResult( return GroundedTurnResult(
generator=service.generate_response(messages, message, conversation_id), generator=service.generate_response(
messages, message, conversation_id, **gen_kwargs
),
citations=_citations_from_results(results), citations=_citations_from_results(results),
grounded=True, grounded=True,
decision=decision, decision=decision,
@@ -43,6 +43,28 @@ _TEMPORAL_PATTERNS = (
_TEMPORAL_RE = re.compile("|".join(_TEMPORAL_PATTERNS), re.IGNORECASE) _TEMPORAL_RE = re.compile("|".join(_TEMPORAL_PATTERNS), re.IGNORECASE)
_YEAR_RE = re.compile(r"\b((?:19|20)\d{2})\b") _YEAR_RE = re.compile(r"\b((?:19|20)\d{2})\b")
# Pure follow-ups should reuse the open chat, not a fresh web search (#84).
_PURE_FOLLOWUP_RE = re.compile(
r"("
r"\bi don'?t know\b"
r"|^\s*why\b"
r"|\bwhy\s*\??\s*$"
r"|\bwhat do you mean\b"
r"|\bexplain\b"
r"|\b(that|this|the|your) (list|joke|answer|setup|punchline)\b"
r"|\bmake (that|this|it|them)\b"
r"|\bput (that|this|it|them)\b"
r"|\bin a table\b"
r"|\binto a table\b"
r")",
re.IGNORECASE,
)
# These still need live sources even when the sentence points at earlier turns.
_NEW_EXTERNAL_FACT_RE = re.compile(
r"\b(price|prices|rating|ratings|cost|costs|review|reviews|cheapest)\b",
re.IGNORECASE,
)
@dataclass @dataclass
class GroundingDecision: class GroundingDecision:
@@ -107,6 +129,28 @@ message). Example: "did Taylor Swift get married" →
source="prepass", source="prepass",
) )
def followup_prepass(
self, prompt: str, *, has_history: bool
) -> GroundingDecision | None:
"""Skip web search when the line only continues the open chat (#84)."""
if not has_history:
return None
text = (prompt or "").strip()
if not text or not _PURE_FOLLOWUP_RE.search(text):
return None
if _NEW_EXTERNAL_FACT_RE.search(text):
return None
temporal = self.temporal_prepass(text)
if temporal is not None and temporal.needs_retrieval:
return None
return GroundingDecision(
needs_retrieval=False,
reason="follow-up on the current conversation",
queries=[],
temporal=False,
source="prepass",
)
def _parse(self, raw: str, fallback_query: str) -> GroundingDecision: def _parse(self, raw: str, fallback_query: str) -> GroundingDecision:
text = (raw or "").strip() text = (raw or "").strip()
# Strip markdown fences if the small model wraps JSON. # Strip markdown fences if the small model wraps JSON.
@@ -144,7 +188,12 @@ message). Example: "did Taylor Swift get married" →
source="model", source="model",
) )
async def decide_async(self, prompt: str) -> GroundingDecision: async def decide_async(
self, prompt: str, *, has_history: bool = False
) -> GroundingDecision:
follow = self.followup_prepass(prompt, has_history=has_history)
if follow is not None:
return follow
pre = self.temporal_prepass(prompt) pre = self.temporal_prepass(prompt)
# Even on temporal prepass, ask the model for better queries when possible. # Even on temporal prepass, ask the model for better queries when possible.
try: try:
@@ -173,7 +222,10 @@ message). Example: "did Taylor Swift get married" →
) )
return decision return decision
def decide(self, prompt: str) -> GroundingDecision: def decide(self, prompt: str, *, has_history: bool = False) -> GroundingDecision:
follow = self.followup_prepass(prompt, has_history=has_history)
if follow is not None:
return follow
pre = self.temporal_prepass(prompt) pre = self.temporal_prepass(prompt)
try: try:
raw = self.chain.invoke({"prompt": prompt}) raw = self.chain.invoke({"prompt": prompt})
+58 -27
View File
@@ -173,35 +173,39 @@ Response:"""
block (when grounded) is never trimmed. block (when grounded) is never trimmed.
""" """
sources = self.sources_block or kwargs.get("sources_block", "") or "" sources = self.sources_block or kwargs.get("sources_block", "") or ""
reserved = ( use_conversation_context = bool(kwargs.get("use_conversation_context", False))
estimate_tokens(ASSISTANT_SYSTEM_PROMPT) history_text = ""
+ estimate_tokens(GROUNDED_ANSWER_INSTRUCTIONS if self.grounded else "") if use_conversation_context:
+ estimate_tokens(sources) reserved = (
+ estimate_tokens(query) estimate_tokens(ASSISTANT_SYSTEM_PROMPT)
+ 256 # response headroom / instructions + estimate_tokens(GROUNDED_ANSWER_INSTRUCTIONS if self.grounded else "")
) + estimate_tokens(sources)
# Leave ~40% of ctx for the completion. + estimate_tokens(query)
history_budget = max(512, int(self.num_ctx * 0.55) - reserved) + 256 # response headroom / instructions
# Drop any prior "Search Results:" blobs from history — sources are
# passed separately now so we don't double-inject.
clean = [
m
for m in conversation
if not (
getattr(m, "type", "") == "human"
and str(getattr(m, "content", "")).startswith("Search Results:")
) )
and not ( # Leave ~40% of ctx for the completion.
getattr(m, "type", "") == "human" history_budget = max(512, int(self.num_ctx * 0.55) - reserved)
and str(getattr(m, "content", "")).startswith("Live sources:") # Drop any prior "Search Results:" blobs from history — sources are
# passed separately now so we don't double-inject.
clean = [
m
for m in conversation
if not (
getattr(m, "type", "") == "human"
and str(getattr(m, "content", "")).startswith("Search Results:")
)
and not (
getattr(m, "type", "") == "human"
and str(getattr(m, "content", "")).startswith("Live sources:")
)
]
# Current user turn is passed as {query}. Drop that text, not
# whichever row happens to be last.
prior = prior_turns(clean, query)
windowed = window_history(
prior, budget_tokens=history_budget, reserved_tokens=0
) )
] history_text = format_history(windowed)
# Exclude the latest user turn from history (it's in {query}).
prior = clean[:-1] if clean else []
windowed = window_history(
prior, budget_tokens=history_budget, reserved_tokens=0
)
history_text = format_history(windowed)
chain_input = { chain_input = {
"query": query, "query": query,
@@ -213,6 +217,33 @@ Response:"""
yield chunk yield chunk
def message_text(message) -> str:
text = getattr(message, "content", "") or ""
if not isinstance(text, str):
text = str(text)
return text
def prior_turns(messages, query: str) -> list:
"""Drop the current user turn from history.
Match the query text, newest hit first, so an unordered transcript cannot
drop a different turn just because it sits at the end of the list (#84).
If the query is not in the list, keep every turn.
"""
if not messages:
return []
needle = (query or "").strip()
for index in range(len(messages) - 1, -1, -1):
message = messages[index]
if (
getattr(message, "type", "") == "human"
and message_text(message).strip() == needle
):
return list(messages[:index]) + list(messages[index + 1 :])
return list(messages)
def build_chat_service( def build_chat_service(
*, *,
model_name: str | None, model_name: str | None,
+9 -1
View File
@@ -490,11 +490,19 @@ class AsyncRAGService(RAGService):
"""Generate response with streaming support.""" """Generate response with streaming support."""
if workspace is None: if workspace is None:
raise ValueError("workspace is required for RAG generation") raise ValueError("workspace is required for RAG generation")
use_conversation_context = bool(
kwargs.get("use_conversation_context", False)
)
recent = (
await self._format_history(conversation)
if use_conversation_context
else ""
)
chain_input = { chain_input = {
"query": query, "query": query,
"conversation": conversation, "conversation": conversation,
"workspace": workspace, "workspace": workspace,
"recent_conversation": await self._format_history(conversation), "recent_conversation": recent,
} }
async for chunk in self.rag_chain.astream(chain_input): async for chunk in self.rag_chain.astream(chain_input):
+186
View File
@@ -0,0 +1,186 @@
"""Image attachments on a chat turn (chat_web_app#128).
Document uploads stay on the csv/xlsx/pdf path. An image is stored on the
prompt and sent to a vision-capable Ollama model with the user's text.
"""
from __future__ import annotations
import base64
import json
import logging
import httpx
from chat_backend.ollama_config import (
ROLE_VISION,
ollama_base_url,
ollama_model_for_role,
ollama_num_ctx_for_role,
)
logger = logging.getLogger(__name__)
MAX_IMAGE_BYTES = 8 * 1024 * 1024
IMAGE_MIME = {
"png": "image/png",
"jpeg": "image/jpeg",
"jpg": "image/jpeg",
"webp": "image/webp",
"gif": "image/gif",
}
_IMAGE_EXTENSIONS = frozenset(IMAGE_MIME)
class VisionChunk:
"""Stream piece that ``chunk_text`` and ``extract_token_usage`` both read."""
def __init__(self, text: str = "", info: dict | None = None):
self.content = text
self.generation_info = info or None
def canonical_file_type(file_type: str) -> str:
"""Map a browser MIME (or short extension) to the stored file_type.
Order matches the historical websocket normalizer so spreadsheets that
share the openxml MIME family still land on ``xlsx``.
"""
lowered = (file_type or "").lower()
if "csv" in lowered:
return "csv"
if "xmlformats-officedocument" in lowered:
return "xlsx"
if "word" in lowered:
return "docx"
if "pdf" in lowered:
return "pdf"
if "text" in lowered:
return "txt"
if "png" in lowered:
return "png"
if "jpeg" in lowered or "jpg" in lowered:
return "jpeg"
if "webp" in lowered:
return "webp"
if "gif" in lowered:
return "gif"
if lowered.startswith("image/"):
return "unsupported_image"
return "Not Sure"
def is_image_type(file_type: str | None) -> bool:
return (file_type or "").lower() in _IMAGE_EXTENSIONS
def image_mime(file_type: str | None) -> str:
return IMAGE_MIME.get((file_type or "").lower(), "application/octet-stream")
def prepare_attachment(file_b64: str, file_type: str) -> tuple[bytes, str, str | None]:
"""Decode an upload. ``error`` is a user-facing string when the image is rejected."""
try:
decoded = base64.b64decode(file_b64, validate=False)
except Exception:
return b"", "Not Sure", "Could not read the attached file."
canonical = canonical_file_type(file_type)
if canonical == "unsupported_image":
return (
decoded,
canonical,
"Unsupported image. Use PNG, JPEG, WebP, or GIF.",
)
if is_image_type(canonical) and len(decoded) > MAX_IMAGE_BYTES:
return (
decoded,
canonical,
"Image is too large. Maximum size is 8 MB.",
)
return decoded, canonical, None
def image_base64(prompt) -> str | None:
"""Base64 of a stored prompt image, or None for non-images."""
if not is_image_type(getattr(prompt, "file_type", None)):
return None
if not prompt.file_exists():
return None
try:
data = prompt.file.read()
except Exception:
logger.exception(
"failed to read prompt image prompt_id=%s", getattr(prompt, "id", None)
)
return None
if not data:
return None
if isinstance(data, str):
data = data.encode()
return base64.b64encode(data).decode("ascii")
class AsyncVisionService:
"""Stream an answer from Ollama's chat API with one image."""
def __init__(self):
self.model_name = ollama_model_for_role(ROLE_VISION)
self.num_ctx = ollama_num_ctx_for_role(ROLE_VISION)
async def generate_response(self, query: str, image_bytes: bytes, file_type: str):
payload = {
"model": self.model_name,
"stream": True,
"options": {"num_ctx": self.num_ctx},
"messages": [
{
"role": "user",
"content": (
"The user attached an image. Answer their question using what "
"is visible in the image, including any text in it.\n\n"
f"{query}"
),
"images": [base64.b64encode(image_bytes).decode("ascii")],
}
],
}
url = f"{ollama_base_url().rstrip('/')}/api/chat"
try:
async with httpx.AsyncClient(timeout=httpx.Timeout(180.0)) as client:
async with client.stream("POST", url, json=payload) as response:
if response.status_code >= 400:
body = await response.aread()
logger.error(
"vision model rejected request status=%s body=%s",
response.status_code,
body[:500],
)
yield VisionChunk(
text="The vision model could not read that image."
)
return
async for line in response.aiter_lines():
if not line:
continue
try:
event = json.loads(line)
except json.JSONDecodeError:
continue
message = event.get("message") or {}
text = message.get("content") or ""
if text:
yield VisionChunk(text=text)
if event.get("done"):
info = {
key: event[key]
for key in ("prompt_eval_count", "eval_count")
if key in event
}
if info:
yield VisionChunk(info=info)
except httpx.HTTPError:
logger.exception("vision request failed model=%s", self.model_name)
yield VisionChunk(text="The vision model could not read that image.")
@@ -23,6 +23,7 @@ STATUS_LABELS: dict[str, str] = {
"reading_sources": "Reading sources", "reading_sources": "Reading sources",
"retrieving_docs": "Searching your documents", "retrieving_docs": "Searching your documents",
"analysing": "Analysing your file", "analysing": "Analysing your file",
"reading_image": "Reading your image",
"refining": "Refining the answer", "refining": "Refining the answer",
"writing": "", # tokens streaming — no label "writing": "", # tokens streaming — no label
} }
@@ -0,0 +1,343 @@
"""Offline unit tests for the agent orchestrator + runner (#63).
No live Ollama/Redis/network — planner/subagent LLMs and tools are fakes,
following the ``async_to_sync`` convention already used in
``test_agent_tools.py`` for exercising async code from Django's sync
``TestCase``.
"""
from __future__ import annotations
from types import SimpleNamespace
from unittest import mock
from asgiref.sync import async_to_sync
from django.test import SimpleTestCase, TestCase, override_settings
from chat_backend.services.agent.orchestrator import (
AgentCancelled,
AgentOrchestrator,
AgentRunLimitExceeded,
RunLimits,
fallback_plan,
parse_plan,
)
from chat_backend.tests.factories import make_conversation, make_prompt, make_user
class FakeMessage:
def __init__(self, content="", tool_calls=None):
self.content = content
self.tool_calls = tool_calls or []
class FakePlannerLLM:
"""Returns a fixed JSON plan, then a fixed synthesis string on the 2nd call."""
def __init__(self, plan_json: str, synthesis_text: str = "Final answer."):
self.plan_json = plan_json
self.synthesis_text = synthesis_text
self.calls = 0
async def ainvoke(self, messages):
self.calls += 1
if self.calls == 1:
return FakeMessage(content=self.plan_json)
return FakeMessage(content=self.synthesis_text)
class FailingLLM:
async def ainvoke(self, messages):
raise RuntimeError("ollama down")
def make_fake_tool(name: str, output: str = "tool result"):
async def _ainvoke(_input):
return output
return SimpleNamespace(name=name, ainvoke=_ainvoke)
class PlanParsingTestCase(SimpleTestCase):
def test_parse_plan_valid_json(self):
raw = (
'{"title": "T", "steps": ['
'{"step_id": "s1", "title": "Search", "tool": "web_search"},'
'{"step_id": "s2", "title": "Write", "tool": null}]}'
)
plan = parse_plan(raw, goal="goal", max_steps=8)
self.assertEqual(plan.title, "T")
self.assertEqual(len(plan.steps), 2)
self.assertEqual(plan.steps[0].tool, "web_search")
self.assertIsNone(plan.steps[1].tool)
def test_parse_plan_strips_code_fences(self):
raw = '```json\n{"title": "T", "steps": [{"step_id": "s1", "title": "X"}]}\n```'
plan = parse_plan(raw, goal="goal", max_steps=8)
self.assertEqual(len(plan.steps), 1)
def test_parse_plan_falls_back_on_garbage(self):
plan = parse_plan("not json at all", goal="latest news today", max_steps=8)
self.assertTrue(plan.steps)
self.assertEqual(plan.steps[0].tool, "web_search")
def test_fallback_plan_no_tool_hints(self):
plan = fallback_plan("write a haiku about rain", max_steps=8)
self.assertEqual(len(plan.steps), 1)
self.assertIsNone(plan.steps[0].tool)
def test_parse_plan_caps_at_max_steps(self):
steps = ",".join(
f'{{"step_id": "s{i}", "title": "Step {i}", "tool": null}}' for i in range(10)
)
raw = f'{{"title": "T", "steps": [{steps}]}}'
plan = parse_plan(raw, goal="goal", max_steps=3)
self.assertEqual(len(plan.steps), 3)
class AgentOrchestratorRunTestCase(SimpleTestCase):
def _run(self, orchestrator):
return async_to_sync(orchestrator.run)()
def test_full_run_with_tool_and_synthesis(self):
plan_json = (
'{"title": "Research", "steps": ['
'{"step_id": "s1", "title": "Search web", "tool": "web_search", '
'"tool_input": {"query": "x"}}]}'
)
planner_llm = FakePlannerLLM(plan_json, synthesis_text="Here is the synthesis.")
tool = make_fake_tool("web_search", output="search results about x")
events: list[tuple[str, dict]] = []
async def on_event(event_type, data):
events.append((event_type, data))
orchestrator = AgentOrchestrator(
goal="Research x",
history_text="",
planner_llm=planner_llm,
subagent_llm_factory=lambda: FailingLLM(),
tools=[tool],
limits=RunLimits(),
on_event=on_event,
)
plan, results, answer = self._run(orchestrator)
self.assertEqual(plan.title, "Research")
self.assertEqual(results["s1"], "search results about x")
self.assertEqual(answer, "Here is the synthesis.")
event_types = [e[0] for e in events]
self.assertEqual(
event_types, ["plan_ready", "step_started", "step_completed"]
)
def test_planner_failure_uses_fallback_plan(self):
# plan() itself is exercised directly here (rather than the full
# run(), which also calls the same LLM for synthesis) so a fully
# unavailable planner model is isolated to the planning step.
orchestrator = AgentOrchestrator(
goal="latest news today",
history_text="",
planner_llm=FailingLLM(),
subagent_llm_factory=lambda: FailingLLM(),
tools=[make_fake_tool("web_search")],
limits=RunLimits(),
)
plan = async_to_sync(orchestrator.plan)()
self.assertTrue(plan.steps)
self.assertEqual(plan.steps[0].tool, "web_search")
def test_step_failure_is_recorded_but_run_continues(self):
plan_json = (
'{"title": "T", "steps": ['
'{"step_id": "s1", "title": "Broken", "tool": "broken_tool"}]}'
)
planner_llm = FakePlannerLLM(plan_json, synthesis_text="Done despite failure.")
orchestrator = AgentOrchestrator(
goal="do a thing",
history_text="",
planner_llm=planner_llm,
subagent_llm_factory=lambda: FailingLLM(),
tools=[], # "broken_tool" is not registered
limits=RunLimits(),
)
_plan, results, answer = self._run(orchestrator)
self.assertIn("Unknown tool", results["s1"])
self.assertEqual(answer, "Done despite failure.")
def test_cancellation_raises_agent_cancelled(self):
plan_json = (
'{"title": "T", "steps": ['
'{"step_id": "s1", "title": "Search", "tool": "web_search"}]}'
)
planner_llm = FakePlannerLLM(plan_json)
tool = make_fake_tool("web_search")
async def is_cancelled():
return True
orchestrator = AgentOrchestrator(
goal="goal",
history_text="",
planner_llm=planner_llm,
subagent_llm_factory=lambda: FailingLLM(),
tools=[tool],
limits=RunLimits(),
is_cancelled=is_cancelled,
)
with self.assertRaises(AgentCancelled):
self._run(orchestrator)
def test_wall_clock_exceeded_fails_step_but_run_still_completes(self):
# A blown wall-clock budget fails the in-flight step (caught by the
# generic per-step exception handler in run()) rather than aborting
# the whole run — synthesis still runs over whatever is available.
plan_json = (
'{"title": "T", "steps": ['
'{"step_id": "s1", "title": "Search", "tool": "web_search"}]}'
)
planner_llm = FakePlannerLLM(plan_json, synthesis_text="Best effort answer.")
tool = make_fake_tool("web_search")
orchestrator = AgentOrchestrator(
goal="goal",
history_text="",
planner_llm=planner_llm,
subagent_llm_factory=lambda: FailingLLM(),
tools=[tool],
limits=RunLimits(wall_clock_seconds=1),
)
# Force the deadline into the past without sleeping in the test.
orchestrator._deadline = 0.0
_plan, results, answer = self._run(orchestrator)
self.assertIn("Wall-clock budget exceeded", results["s1"])
self.assertEqual(answer, "Best effort answer.")
def test_execute_step_raises_when_deadline_passed(self):
orchestrator = AgentOrchestrator(
goal="goal",
history_text="",
planner_llm=FailingLLM(),
subagent_llm_factory=lambda: FailingLLM(),
tools=[],
limits=RunLimits(),
)
orchestrator._deadline = 0.0
from chat_backend.services.agent.orchestrator import PlanStep
with self.assertRaises(AgentRunLimitExceeded):
async_to_sync(orchestrator.execute_step)(
PlanStep(step_id="s1", title="x", tool="web_search"), {}
)
@override_settings(ALLOW_AGENTIC_TASKS=True)
class RunAgenticTurnTestCase(TestCase):
def setUp(self):
self.user = make_user(email="agent-runner@example.com")
self.conversation = make_conversation(user=self.user)
self.prompt = make_prompt(self.conversation, message="Research the top 5 things")
def _patch_llms(self, plan_json="", synthesis_text="All done."):
planner_llm = FakePlannerLLM(
plan_json
or '{"title": "T", "steps": [{"step_id": "s1", "title": "Answer", "tool": null}]}',
synthesis_text=synthesis_text,
)
return mock.patch(
"chat_backend.services.agent.runner._build_llms",
return_value=(planner_llm, lambda: FailingLLM()),
)
def test_run_agentic_turn_completes_and_persists(self):
from chat_backend.models import AgentRun
sent_frames = []
async def ws_send(raw):
sent_frames.append(raw)
with self._patch_llms(synthesis_text="The answer is 42."):
run, answer = async_to_sync(self._call_run_agentic_turn)(ws_send)
run.refresh_from_db()
self.assertEqual(run.status, AgentRun.Status.COMPLETED)
self.assertEqual(run.result, "The answer is 42.")
self.assertEqual(answer, "The answer is 42.")
self.assertTrue(sent_frames) # progress frames reached the WS callback
async def _call_run_agentic_turn(self, ws_send):
from chat_backend.services.agent.runner import run_agentic_turn
from chat_backend.services.chat_tenant_scope import resolve_chat_company_scope
from asgiref.sync import sync_to_async
scope = await sync_to_async(resolve_chat_company_scope)(
self.user, self.conversation.id
)
return await run_agentic_turn(
user=self.user,
scope=scope,
conversation_id=self.conversation.id,
goal=self.prompt.message,
prompt=self.prompt,
ws_send=ws_send,
)
def test_disabled_flag_short_circuits(self):
from chat_backend.services.chat_tenant_scope import resolve_chat_company_scope
with override_settings(ALLOW_AGENTIC_TASKS=False):
with self.assertRaises(RuntimeError):
async_to_sync(self._call_run_agentic_turn)(None)
@override_settings(ALLOW_AGENTIC_TASKS=True)
class ExecuteAgentRunTestCase(TestCase):
def setUp(self):
self.user = make_user(email="agent-bg@example.com")
self.conversation = make_conversation(user=self.user)
def test_execute_agent_run_marks_cancelled(self):
from chat_backend.models import AgentRun, AgentStep
run = AgentRun.objects.create(
user=self.user,
conversation=self.conversation,
goal="Research the top 5 things and compare",
status=AgentRun.Status.PENDING,
)
run.cancel_requested = True
run.save(update_fields=["cancel_requested"])
planner_llm = FakePlannerLLM(
'{"title": "T", "steps": [{"step_id": "s1", "title": "Search", '
'"tool": "web_search"}]}'
)
with mock.patch(
"chat_backend.services.agent.runner._build_llms",
return_value=(planner_llm, lambda: FailingLLM()),
), mock.patch(
"chat_backend.services.agent.runner.build_agent_tools",
return_value=[make_fake_tool("web_search")],
):
from chat_backend.services.agent.runner import execute_agent_run
async_to_sync(execute_agent_run)(run.pk)
run.refresh_from_db()
self.assertEqual(run.status, AgentRun.Status.CANCELLED)
self.assertEqual(AgentStep.objects.filter(run=run).count(), 0)
def test_execute_agent_run_skips_non_pending(self):
from chat_backend.models import AgentRun
from chat_backend.services.agent.runner import execute_agent_run
run = AgentRun.objects.create(
user=self.user,
conversation=self.conversation,
goal="already running",
status=AgentRun.Status.RUNNING,
)
# Should return immediately without touching the run.
async_to_sync(execute_agent_run)(run.pk)
run.refresh_from_db()
self.assertEqual(run.status, AgentRun.Status.RUNNING)
+83 -18
View File
@@ -1,9 +1,11 @@
from unittest import mock
import json import json
from datetime import timedelta
from unittest import mock
from asgiref.sync import sync_to_async from asgiref.sync import sync_to_async
from channels.testing import WebsocketCommunicator from channels.testing import WebsocketCommunicator
from django.test import TestCase, TransactionTestCase, override_settings from django.test import TestCase, TransactionTestCase, override_settings
from django.utils import timezone
from langchain_core.messages import AIMessage, HumanMessage from langchain_core.messages import AIMessage, HumanMessage
from parameterized import parameterized from parameterized import parameterized
@@ -47,9 +49,7 @@ class DatabaseHelperTestCase(TransactionTestCase):
@parameterized.expand([("websocket", consumers), ("langgraph", consumers_graph)]) @parameterized.expand([("websocket", consumers), ("langgraph", consumers_graph)])
async def test_get_workspace(self, _name, module): async def test_get_workspace(self, _name, module):
workspace = await module.get_workspace( workspace = await module.get_workspace(self.conversation.id, user=self.user)
self.conversation.id, user=self.user
)
self.assertEqual(workspace.id, self.workspace.id) self.assertEqual(workspace.id, self.workspace.id)
@@ -99,6 +99,29 @@ class DatabaseHelperTestCase(TransactionTestCase):
[type(message) for message in messages], [AIMessage, HumanMessage] [type(message) for message in messages], [AIMessage, HumanMessage]
) )
@parameterized.expand([("websocket", consumers), ("langgraph", consumers_graph)])
async def test_get_messages_orders_by_created_not_insert_order(self, _name, module):
now = timezone.now()
await sync_to_async(Prompt.objects.create)(
conversation=self.conversation,
message="newer answer",
user_created=False,
created=now,
)
await sync_to_async(Prompt.objects.create)(
conversation=self.conversation,
message="older question",
user_created=True,
created=now - timedelta(minutes=5),
)
messages, _prompt = await module.get_messages(self.conversation.id, "follow up")
self.assertEqual(
[message.content for message in messages],
["older question", "newer answer", "follow up"],
)
async def test_get_messages_stores_attachment_and_inlines_csv(self): async def test_get_messages_stores_attachment_and_inlines_csv(self):
messages, prompt_instance = await consumers.get_messages( messages, prompt_instance = await consumers.get_messages(
self.conversation.id, "analyze this", b"name,sales\nalice,10\n", "csv" self.conversation.id, "analyze this", b"name,sales\nalice,10\n", "csv"
@@ -355,6 +378,42 @@ class GraphNodeTestCase(TransactionTestCase):
self.prompt.message, b"a,b\n1,2\n", "csv" self.prompt.message, b"a,b\n1,2\n", "csv"
) )
async def test_generation_node_sends_images_to_vision(self):
with mock.patch.object(consumers_graph, "AsyncVisionService") as service:
service.return_value.model_name = "gemma4:latest"
service.return_value.generate_response.return_value = "generator"
result = await consumers_graph.generation_node(
self._state(
prompt_type=PromptType.DATA_ANALYSIS,
message="what does this say?",
decoded_file=b"\x89PNG",
file_type="png",
)
)
self.assertEqual(result["response_generator"], "generator")
self.assertEqual(result["resolved_model"], "gemma4:latest")
service.return_value.generate_response.assert_called_once_with(
self.prompt.message, b"\x89PNG", "png"
)
async def test_attached_image_is_not_data_analysis(self):
with mock.patch.object(
consumers_graph.PROMPT_CLASSIFIER,
"classify_async",
mock.AsyncMock(return_value=PromptType.DATA_ANALYSIS),
):
result = await consumers_graph.classification_node(
self._state(
message="analyze this chart",
decoded_file=b"\x89PNG",
file_type="png",
)
)
self.assertEqual(result["prompt_type"], PromptType.GENERAL_CHAT)
async def test_generation_node_routes_rag_prompts_with_workspace(self): async def test_generation_node_routes_rag_prompts_with_workspace(self):
with mock.patch.object(consumers_graph, "AsyncRAGService") as service: with mock.patch.object(consumers_graph, "AsyncRAGService") as service:
service.return_value.generate_response.return_value = "generator" service.return_value.generate_response.return_value = "generator"
@@ -551,13 +610,16 @@ class WebSocketReceiveGuardTestCase(TransactionTestCase):
connected, _ = await communicator.connect() connected, _ = await communicator.connect()
self.assertTrue(connected) self.assertTrue(connected)
with mock.patch( with (
"chat_backend.consumers.title_generator.generate_async", mock.patch(
new_callable=mock.AsyncMock, "chat_backend.consumers.title_generator.generate_async",
) as title_chat, mock.patch( new_callable=mock.AsyncMock,
"chat_backend.consumers_graph.title_generator.generate_async", ) as title_chat,
new_callable=mock.AsyncMock, mock.patch(
) as title_graph: "chat_backend.consumers_graph.title_generator.generate_async",
new_callable=mock.AsyncMock,
) as title_graph,
):
await communicator.send_json_to({"type": "ping", "email": "a@b.com"}) await communicator.send_json_to({"type": "ping", "email": "a@b.com"})
# No reply expected; give the event loop a tick. # No reply expected; give the event loop a tick.
self.assertTrue(await communicator.receive_nothing(timeout=0.2)) self.assertTrue(await communicator.receive_nothing(timeout=0.2))
@@ -576,13 +638,16 @@ class WebSocketReceiveGuardTestCase(TransactionTestCase):
connected, _ = await communicator.connect() connected, _ = await communicator.connect()
self.assertTrue(connected) self.assertTrue(connected)
with mock.patch( with (
"chat_backend.consumers.title_generator.generate_async", mock.patch(
new_callable=mock.AsyncMock, "chat_backend.consumers.title_generator.generate_async",
) as title_chat, mock.patch( new_callable=mock.AsyncMock,
"chat_backend.consumers_graph.title_generator.generate_async", ) as title_chat,
new_callable=mock.AsyncMock, mock.patch(
) as title_graph: "chat_backend.consumers_graph.title_generator.generate_async",
new_callable=mock.AsyncMock,
) as title_graph,
):
await communicator.send_json_to( await communicator.send_json_to(
{"message": " ", "email": "a@b.com", "conversation_id": None} {"message": " ", "email": "a@b.com", "conversation_id": None}
) )
@@ -86,6 +86,25 @@ class GroundingPrepassTestCase(SimpleTestCase):
self.assertEqual(decision.source, "fail_open") self.assertEqual(decision.source, "fail_open")
self.assertEqual(decision.queries, ["fallback query"]) self.assertEqual(decision.queries, ["fallback query"])
def test_pure_followup_skips_retrieval_without_calling_model(self):
self.decider.chain = mock.Mock()
decision = self.decider.decide("I don't know. Why?", has_history=True)
self.assertFalse(decision.needs_retrieval)
self.assertEqual(decision.reason, "follow-up on the current conversation")
self.decider.chain.invoke.assert_not_called()
def test_followup_without_history_does_not_short_circuit(self):
self.assertIsNone(
self.decider.followup_prepass("I don't know. Why?", has_history=False)
)
def test_price_followup_still_allows_retrieval(self):
decision = self.decider.followup_prepass(
"Can you make that list for me and add the prices and ratings in a table",
has_history=True,
)
self.assertIsNone(decision)
def test_decide_fail_open_on_llm_error(self): def test_decide_fail_open_on_llm_error(self):
self.decider.chain = mock.Mock() self.decider.chain = mock.Mock()
self.decider.chain.invoke.side_effect = RuntimeError("ollama down") self.decider.chain.invoke.side_effect = RuntimeError("ollama down")
@@ -0,0 +1,16 @@
from django.db import connection
from django.db.migrations.loader import MigrationLoader
from django.test import TestCase
class MigrationGraphTests(TestCase):
def test_chat_backend_has_a_single_leaf(self):
loader = MigrationLoader(connection, ignore_no_migrations=True)
leaves = [
name for app, name in loader.graph.leaf_nodes() if app == "chat_backend"
]
self.assertEqual(leaves, ["0036_enable_conversation_context"])
node = loader.graph.nodes[("chat_backend", "0036_enable_conversation_context")]
parents = [name for app, name in node.dependencies if app == "chat_backend"]
self.assertEqual(parents, ["0035_waitlistentry"])
+1
View File
@@ -81,6 +81,7 @@ class CompanyAndUserTestCase(TestCase):
self.assertFalse(user.deleted) self.assertFalse(user.deleted)
self.assertFalse(user.has_signed_tos) self.assertFalse(user.has_signed_tos)
self.assertTrue(user.conversation_order) self.assertTrue(user.conversation_order)
self.assertTrue(user.use_conversation_context)
class ConversationAndPromptTestCase(TestCase): class ConversationAndPromptTestCase(TestCase):
+38
View File
@@ -24,6 +24,7 @@ OAUTH_SETTINGS = {
"FRONTEND_BASE_URL": "http://frontend.test", "FRONTEND_BASE_URL": "http://frontend.test",
"OAUTH_CALLBACK_BASE_URL": "http://backend.test", "OAUTH_CALLBACK_BASE_URL": "http://backend.test",
"ENABLE_ACCOUNT_REGISTRATION": True, "ENABLE_ACCOUNT_REGISTRATION": True,
"ENABLE_SSO": True,
} }
@@ -53,6 +54,14 @@ class PublicSettingsOAuthTestCase(APITestCase):
self.assertTrue(response.data["oauth"]["google"]) self.assertTrue(response.data["oauth"]["google"])
self.assertTrue(response.data["oauth"]["microsoft"]) self.assertTrue(response.data["oauth"]["microsoft"])
@override_settings(ENABLE_SSO=False)
def test_hides_oauth_when_sso_disabled(self):
response = self.client.get(reverse("public_settings"))
self.assertEqual(response.status_code, status.HTTP_200_OK)
self.assertTrue(response.data["enable_account_registration"])
self.assertFalse(response.data["oauth"]["google"])
self.assertFalse(response.data["oauth"]["microsoft"])
@override_settings(GOOGLE_OAUTH_CLIENT_ID="", GOOGLE_OAUTH_CLIENT_SECRET="") @override_settings(GOOGLE_OAUTH_CLIENT_ID="", GOOGLE_OAUTH_CLIENT_SECRET="")
def test_hides_unconfigured_provider(self): def test_hides_unconfigured_provider(self):
response = self.client.get(reverse("public_settings")) response = self.client.get(reverse("public_settings"))
@@ -83,6 +92,15 @@ class OAuthStartTestCase(APITestCase):
) )
self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND) self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND)
@override_settings(ENABLE_SSO=False)
def test_login_and_signup_blocked_when_sso_disabled(self):
for intent in ("login", "signup"):
response = self.client.get(
reverse("oauth_start", kwargs={"provider": "google"}),
{"intent": intent},
)
self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
@override_settings(ENABLE_ACCOUNT_REGISTRATION=False) @override_settings(ENABLE_ACCOUNT_REGISTRATION=False)
def test_signup_intent_blocked_when_registration_disabled(self): def test_signup_intent_blocked_when_registration_disabled(self):
response = self.client.get( response = self.client.get(
@@ -180,6 +198,17 @@ class OAuthCallbackTestCase(APITestCase):
params = parse_qs(urlparse(response["Location"]).query) params = parse_qs(urlparse(response["Location"]).query)
self.assertEqual(params["error"], ["email_missing"]) self.assertEqual(params["error"], ["email_missing"])
@patch("chat_backend.views_oauth.exchange_code_for_profile")
@override_settings(ENABLE_SSO=False)
def test_callback_signup_rejected_when_sso_disabled(self, mock_exchange):
mock_exchange.return_value = _google_profile()
response = self._callback(intent="signup")
params = parse_qs(urlparse(response["Location"]).query)
self.assertEqual(params["error"], ["sso_disabled"])
self.assertEqual(CustomUser.objects.count(), 0)
mock_exchange.assert_not_called()
@patch("chat_backend.views_oauth.exchange_code_for_profile") @patch("chat_backend.views_oauth.exchange_code_for_profile")
@override_settings(ENABLE_ACCOUNT_REGISTRATION=False) @override_settings(ENABLE_ACCOUNT_REGISTRATION=False)
def test_callback_login_without_account_when_registration_off(self, mock_exchange): def test_callback_login_without_account_when_registration_off(self, mock_exchange):
@@ -256,6 +285,15 @@ class OAuthStartDriveLinkTestCase(APITestCase):
) )
self.assertEqual(response.status_code, status.HTTP_401_UNAUTHORIZED) self.assertEqual(response.status_code, status.HTTP_401_UNAUTHORIZED)
@override_settings(ENABLE_SSO=False)
def test_link_drive_still_starts_when_sso_disabled(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)
def test_link_drive_authenticated_redirects_with_drive_scope_and_state(self): def test_link_drive_authenticated_redirects_with_drive_scope_and_state(self):
self.client.force_authenticate(user=self.user) self.client.force_authenticate(user=self.user)
response = self.client.get( response = self.client.get(
@@ -0,0 +1,65 @@
from django.contrib.auth import get_user_model
from django.urls import reverse
from rest_framework.test import APITestCase
from chat_backend.models import Document, WaitlistEntry
from chat_backend.tests.factories import make_company, make_document, make_user, make_workspace
class PublicUsageViewTests(APITestCase):
def test_counts_active_accounts_and_documents_without_personal_data(self):
company = make_company()
make_user(email="active@example.com", company=company)
inactive = make_user(email="inactive@example.com")
inactive.is_active = False
inactive.save(update_fields=["is_active"])
deleted = make_user(email="deleted@example.com")
deleted.deleted = True
deleted.save(update_fields=["deleted"])
make_document(make_workspace(company=company))
response = self.client.get(reverse("public_usage"))
self.assertEqual(response.status_code, 200)
self.assertEqual(
set(response.data.keys()),
{"account_count", "document_count", "as_of"},
)
self.assertEqual(response.data["account_count"], 1)
self.assertEqual(response.data["document_count"], Document.objects.count())
self.assertIn("T", response.data["as_of"])
body = response.content.decode()
self.assertNotIn("active@example.com", body)
self.assertNotIn("avatar", body)
def test_does_not_require_authentication(self):
response = self.client.get(reverse("public_usage"))
self.assertEqual(response.status_code, 200)
self.assertEqual(response.data["account_count"], 0)
class PublicWaitlistViewTests(APITestCase):
def test_stores_an_email_once(self):
url = reverse("public_waitlist")
first = self.client.post(url, {"email": " Person@Example.com "}, format="json")
second = self.client.post(url, {"email": "person@example.com"}, format="json")
self.assertEqual(first.status_code, 200)
self.assertEqual(second.status_code, 200)
self.assertEqual(first.data["detail"], "You're on the list.")
self.assertEqual(second.data["detail"], "You're on the list.")
self.assertEqual(WaitlistEntry.objects.count(), 1)
self.assertEqual(WaitlistEntry.objects.get().email, "person@example.com")
self.assertEqual(get_user_model().objects.count(), 0)
def test_rejects_a_bad_email(self):
response = self.client.post(
reverse("public_waitlist"), {"email": "not-an-email"}, format="json"
)
self.assertEqual(response.status_code, 400)
self.assertEqual(WaitlistEntry.objects.count(), 0)
def test_rejects_a_missing_email(self):
response = self.client.post(reverse("public_waitlist"), {}, format="json")
self.assertEqual(response.status_code, 400)
self.assertEqual(WaitlistEntry.objects.count(), 0)
+84 -11
View File
@@ -1,7 +1,12 @@
from django.test import SimpleTestCase from django.test import SimpleTestCase
from langchain_core.messages import AIMessage, HumanMessage from langchain_core.messages import AIMessage, HumanMessage
from chat_backend.services.llm_service import AsyncLLMService, SyncLLMService from chat_backend.services.assistant_identity import GROUNDED_ANSWER_INSTRUCTIONS
from chat_backend.services.llm_service import (
AsyncLLMService,
SyncLLMService,
prior_turns,
)
from chat_backend.services.prompt_budget import format_history, window_history from chat_backend.services.prompt_budget import format_history, window_history
from .fakes import FakeChain from .fakes import FakeChain
@@ -25,7 +30,10 @@ class AsyncLLMServiceTestCase(SimpleTestCase):
chunks = [ chunks = [
chunk chunk
async for chunk in self.service.generate_response( async for chunk in self.service.generate_response(
conversation(1), "hello", conversation_id=1 conversation(1),
"hello",
conversation_id=1,
use_conversation_context=True,
) )
] ]
@@ -33,25 +41,74 @@ class AsyncLLMServiceTestCase(SimpleTestCase):
async def test_generate_response_sends_single_history_window(self): async def test_generate_response_sends_single_history_window(self):
self.service.conversation_chain = FakeChain(chunks=["ok"]) self.service.conversation_chain = FakeChain(chunks=["ok"])
messages = conversation(4) # 8 messages messages = conversation(4) + [HumanMessage(content="latest")]
async for _ in self.service.generate_response(messages, "latest", 1): async for _ in self.service.generate_response(
messages, "latest", 1, use_conversation_context=True
):
pass pass
payload = self.service.conversation_chain.calls[0] payload = self.service.conversation_chain.calls[0]
self.assertEqual(payload["query"], "latest") self.assertEqual(payload["query"], "latest")
# Latest user turn is in {query}; history is prior turns only, once. # Latest user turn is in {query}; history is prior turns only, once.
self.assertIn("history", payload) self.assertIn("history", payload)
self.assertNotIn("User: latest", payload["history"])
self.assertIn("AI: answer 3", payload["history"])
self.assertNotIn("recent_conversation", payload) self.assertNotIn("recent_conversation", payload)
self.assertNotIn("conversation", payload) self.assertNotIn("conversation", payload)
# 8 messages → drop last (query) → 7 prior lines max in window. self.assertLessEqual(len(payload["history"].splitlines()), 8)
self.assertLessEqual(len(payload["history"].splitlines()), 7)
async def test_history_drops_matching_turn_even_if_it_is_not_last(self):
self.service.conversation_chain = FakeChain(chunks=["ok"])
# Unordered transcript: current question is first, older turns follow.
messages = [
HumanMessage(content="I don't know. Why?"),
AIMessage(content="Why did the meditation instructor bring a ladder?"),
HumanMessage(content="Tell me a joke"),
]
async for _ in self.service.generate_response(
messages,
"I don't know. Why?",
1,
use_conversation_context=True,
):
pass
history = self.service.conversation_chain.calls[0]["history"]
self.assertNotIn("I don't know", history)
self.assertIn("Tell me a joke", history)
self.assertIn("meditation instructor", history)
async def test_generate_response_skips_history_when_context_disabled(self):
self.service.conversation_chain = FakeChain(chunks=["ok"])
messages = conversation(4)
async for _ in self.service.generate_response(
messages, "latest", 1, use_conversation_context=False
):
pass
payload = self.service.conversation_chain.calls[0]
self.assertEqual(payload["history"], "")
async def test_generate_response_defaults_to_no_history(self):
self.service.conversation_chain = FakeChain(chunks=["ok"])
async for _ in self.service.generate_response(conversation(2), "q", 1):
pass
self.assertEqual(self.service.conversation_chain.calls[0]["history"], "")
async def test_grounded_service_includes_sources(self): async def test_grounded_service_includes_sources(self):
service = AsyncLLMService(grounded=True, sources_block='[1] "T" — x.com — undated') service = AsyncLLMService(
grounded=True, sources_block='[1] "T" — x.com — undated'
)
service.conversation_chain = FakeChain(chunks=["ok"]) service.conversation_chain = FakeChain(chunks=["ok"])
async for _ in service.generate_response([], "q", 1): async for _ in service.generate_response(
[], "q", 1, use_conversation_context=True
):
pass pass
payload = service.conversation_chain.calls[0] payload = service.conversation_chain.calls[0]
@@ -63,6 +120,24 @@ class SyncLLMServiceTestCase(SimpleTestCase):
self.assertIsNotNone(SyncLLMService()) self.assertIsNotNone(SyncLLMService())
class PriorTurnsTestCase(SimpleTestCase):
def test_grounded_instructions_keep_the_conversation(self):
self.assertIn("conversation", GROUNDED_ANSWER_INSTRUCTIONS.lower())
self.assertNotIn("Answer ONLY from those sources", GROUNDED_ANSWER_INSTRUCTIONS)
def test_prior_turns_drops_newest_matching_user_text(self):
messages = [
HumanMessage(content="I don't know. Why?"),
AIMessage(content="the joke"),
HumanMessage(content="Tell me a joke"),
]
kept = prior_turns(messages, "I don't know. Why?")
self.assertEqual(
[message.content for message in kept],
["the joke", "Tell me a joke"],
)
class PromptBudgetTestCase(SimpleTestCase): class PromptBudgetTestCase(SimpleTestCase):
def test_window_history_drops_oldest_first(self): def test_window_history_drops_oldest_first(self):
messages = [HumanMessage(content="x" * 40) for _ in range(10)] messages = [HumanMessage(content="x" * 40) for _ in range(10)]
@@ -72,7 +147,5 @@ class PromptBudgetTestCase(SimpleTestCase):
self.assertEqual(kept[-1].content, messages[-1].content) self.assertEqual(kept[-1].content, messages[-1].content)
def test_format_history_labels_speakers(self): def test_format_history_labels_speakers(self):
text = format_history( text = format_history([HumanMessage(content="hello"), AIMessage(content="hi")])
[HumanMessage(content="hello"), AIMessage(content="hi")]
)
self.assertEqual(text, "User: hello\nAI: hi") self.assertEqual(text, "User: hello\nAI: hi")
+23 -1
View File
@@ -335,7 +335,10 @@ class RAGServiceTestCase(TransactionTestCase):
chunks = [ chunks = [
chunk chunk
async for chunk in self.service.generate_response( async for chunk in self.service.generate_response(
conversation, "what is our policy?", self.workspace conversation,
"what is our policy?",
self.workspace,
use_conversation_context=True,
) )
] ]
@@ -345,6 +348,25 @@ class RAGServiceTestCase(TransactionTestCase):
self.assertEqual(payload["workspace"], self.workspace) self.assertEqual(payload["workspace"], self.workspace)
self.assertEqual(payload["recent_conversation"], "User: what is our policy?") self.assertEqual(payload["recent_conversation"], "User: what is our policy?")
async def test_generate_response_skips_history_when_context_disabled(self):
self.service.rag_chain = FakeChain(chunks=["ok"])
conversation = [
HumanMessage(content="prior"),
AIMessage(content="answer"),
HumanMessage(content="what is our policy?"),
]
async for _ in self.service.generate_response(
conversation,
"what is our policy?",
self.workspace,
use_conversation_context=False,
):
pass
payload = self.service.rag_chain.calls[0]
self.assertEqual(payload["recent_conversation"], "")
async def test_get_documents_helper_scopes_by_workspace(self): async def test_get_documents_helper_scopes_by_workspace(self):
document = await sync_to_async(self._text_document)() document = await sync_to_async(self._text_document)()
other_workspace = await sync_to_async(make_workspace)( other_workspace = await sync_to_async(make_workspace)(
@@ -1,4 +1,7 @@
from datetime import timedelta
from django.urls import reverse from django.urls import reverse
from django.utils import timezone
from rest_framework import status from rest_framework import status
from rest_framework.test import APITestCase from rest_framework.test import APITestCase
@@ -57,6 +60,7 @@ class ConversationPreferencesTestCase(APITestCase):
self.assertEqual(response.status_code, status.HTTP_200_OK) self.assertEqual(response.status_code, status.HTTP_200_OK)
self.assertTrue(response.data["order"]) self.assertTrue(response.data["order"])
self.assertTrue(response.data["use_conversation_context"])
def test_post_toggles_and_persists_order(self): def test_post_toggles_and_persists_order(self):
response = self.client.post(self.url) response = self.client.post(self.url)
@@ -66,6 +70,33 @@ class ConversationPreferencesTestCase(APITestCase):
self.user.refresh_from_db() self.user.refresh_from_db()
self.assertFalse(self.user.conversation_order) self.assertFalse(self.user.conversation_order)
def test_post_sets_use_conversation_context_without_toggling_order(self):
before_order = self.user.conversation_order
response = self.client.post(
self.url, {"use_conversation_context": True}, format="json"
)
self.assertEqual(response.status_code, status.HTTP_200_OK)
self.assertTrue(response.data["use_conversation_context"])
self.assertEqual(response.data["order"], before_order)
self.user.refresh_from_db()
self.assertTrue(self.user.use_conversation_context)
self.assertEqual(self.user.conversation_order, before_order)
def test_post_can_disable_use_conversation_context(self):
self.user.use_conversation_context = True
self.user.save(update_fields=["use_conversation_context"])
response = self.client.post(
self.url, {"use_conversation_context": False}, format="json"
)
self.assertEqual(response.status_code, status.HTTP_200_OK)
self.assertFalse(response.data["use_conversation_context"])
self.user.refresh_from_db()
self.assertFalse(self.user.use_conversation_context)
class ConversationDetailViewTestCase(APITestCase): class ConversationDetailViewTestCase(APITestCase):
def setUp(self): def setUp(self):
@@ -87,8 +118,33 @@ class ConversationDetailViewTestCase(APITestCase):
[("hello", True), ("hi there", False)], [("hello", True), ("hi there", False)],
) )
def test_get_returns_prompts_oldest_first_even_if_inserted_newest_first(self):
now = timezone.now()
make_prompt(
self.conversation,
message="later",
user_created=False,
created=now,
)
make_prompt(
self.conversation,
message="earlier",
user_created=True,
created=now - timedelta(minutes=5),
)
response = self.client.get(self.url, {"conversation_id": self.conversation.id})
self.assertEqual(response.status_code, status.HTTP_200_OK)
self.assertEqual(
[item["message"] for item in response.data],
["earlier", "later"],
)
def test_get_rejects_other_users_conversation(self): def test_get_rejects_other_users_conversation(self):
other = make_user(email="other-tenant@example.com", company=make_company("Other")) other = make_user(
email="other-tenant@example.com", company=make_company("Other")
)
foreign = make_conversation(user=other) foreign = make_conversation(user=other)
make_prompt(foreign, message="secret") make_prompt(foreign, message="secret")
+19 -1
View File
@@ -115,15 +115,33 @@ class PublicSettingsTestCase(APITestCase):
self.assertEqual(response.status_code, status.HTTP_200_OK) self.assertEqual(response.status_code, status.HTTP_200_OK)
self.assertEqual(response.data["enable_account_registration"], False) self.assertEqual(response.data["enable_account_registration"], False)
def test_registration_on_and_sso_hidden_by_default(self):
response = self.client.get(reverse("public_settings"))
self.assertEqual(response.status_code, status.HTTP_200_OK)
self.assertEqual(response.data["enable_account_registration"], True)
self.assertEqual(response.data["oauth"], {"google": False, "microsoft": False})
class UserCreateTestCase(APITestCase): class UserCreateTestCase(APITestCase):
def test_registration_disabled_by_default(self): def test_registration_enabled_by_default(self):
response = self.client.post( response = self.client.post(
reverse("create_user"), reverse("create_user"),
{"email": "new@example.com", "password": "securepass1"}, {"email": "new@example.com", "password": "securepass1"},
format="json", format="json",
) )
self.assertEqual(response.status_code, status.HTTP_201_CREATED)
self.assertTrue(CustomUser.objects.filter(email="new@example.com").exists())
def test_registration_disabled_when_setting_off(self):
with self.settings(ENABLE_ACCOUNT_REGISTRATION=False):
response = self.client.post(
reverse("create_user"),
{"email": "new@example.com", "password": "securepass1"},
format="json",
)
self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN) self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
self.assertEqual(CustomUser.objects.count(), 0) self.assertEqual(CustomUser.objects.count(), 0)
+64
View File
@@ -0,0 +1,64 @@
import base64
from django.core.files.uploadedfile import SimpleUploadedFile
from django.test import TestCase
from chat_backend.services.vision import (
canonical_file_type,
image_base64,
is_image_type,
prepare_attachment,
)
from chat_backend.tests.factories import make_conversation, make_prompt
class VisionAttachmentTestCase(TestCase):
def test_canonical_types(self):
self.assertEqual(canonical_file_type("image/png"), "png")
self.assertEqual(canonical_file_type("image/jpeg"), "jpeg")
self.assertEqual(canonical_file_type("image/jpg"), "jpeg")
self.assertEqual(canonical_file_type("image/webp"), "webp")
self.assertEqual(canonical_file_type("image/gif"), "gif")
self.assertEqual(canonical_file_type("text/plain"), "txt")
self.assertEqual(canonical_file_type("application/pdf"), "pdf")
self.assertEqual(canonical_file_type("image/bmp"), "unsupported_image")
self.assertEqual(canonical_file_type("image/svg+xml"), "unsupported_image")
def test_prepare_attachment_rejects_oversize_and_unknown_images(self):
png = base64.b64encode(b"\x89PNG").decode()
decoded, kind, error = prepare_attachment(png, "image/png")
self.assertEqual(decoded, b"\x89PNG")
self.assertEqual(kind, "png")
self.assertIsNone(error)
self.assertTrue(is_image_type(kind))
huge = base64.b64encode(b"x" * (8 * 1024 * 1024 + 1)).decode()
_, _, error = prepare_attachment(huge, "image/png")
self.assertIn("too large", error)
_, _, error = prepare_attachment(png, "image/bmp")
self.assertIn("Unsupported image", error)
def test_document_types_unchanged(self):
raw = base64.b64encode(b"a,b\n1,2\n").decode()
decoded, kind, error = prepare_attachment(raw, "text/csv")
self.assertEqual(kind, "csv")
self.assertIsNone(error)
self.assertEqual(decoded, b"a,b\n1,2\n")
self.assertFalse(is_image_type(kind))
def test_image_base64_reads_stored_png(self):
prompt = make_prompt(make_conversation())
prompt.file_type = "png"
prompt.file.save("shot.png", SimpleUploadedFile("shot.png", b"\x89PNG"))
prompt.save()
self.assertEqual(image_base64(prompt), base64.b64encode(b"\x89PNG").decode())
def test_image_base64_skips_documents(self):
prompt = make_prompt(make_conversation())
prompt.file_type = "csv"
prompt.file.save("data.csv", SimpleUploadedFile("data.csv", b"a,b"))
prompt.save()
self.assertIsNone(image_base64(prompt))
+4
View File
@@ -9,6 +9,8 @@ from .views import (
LogoutAndBlacklistRefreshTokenForUserView, LogoutAndBlacklistRefreshTokenForUserView,
CustomUserGet, CustomUserGet,
PublicSettingsView, PublicSettingsView,
PublicUsageView,
PublicWaitlistView,
is_authenticated, is_authenticated,
AnnouncmentView, AnnouncmentView,
FeedbackView, FeedbackView,
@@ -46,6 +48,8 @@ urlpatterns = [
path("token/refresh/", jwt_views.TokenRefreshView.as_view(), name="token_refresh"), path("token/refresh/", jwt_views.TokenRefreshView.as_view(), name="token_refresh"),
path("user/create/", CustomUserCreate.as_view(), name="create_user"), path("user/create/", CustomUserCreate.as_view(), name="create_user"),
path("public/settings/", PublicSettingsView.as_view(), name="public_settings"), path("public/settings/", PublicSettingsView.as_view(), name="public_settings"),
path("public/usage/", PublicUsageView.as_view(), name="public_usage"),
path("public/waitlist/", PublicWaitlistView.as_view(), name="public_waitlist"),
path( path(
"auth/oauth/<str:provider>/start/", "auth/oauth/<str:provider>/start/",
OAuthStartView.as_view(), OAuthStartView.as_view(),
+81 -7
View File
@@ -32,6 +32,7 @@ from .models import (
DocumentWorkspace, DocumentWorkspace,
Document, Document,
UserAuthEvent, UserAuthEvent,
WaitlistEntry,
) )
from django.views.decorators.cache import never_cache from django.views.decorators.cache import never_cache
from django.http import JsonResponse from django.http import JsonResponse
@@ -127,8 +128,57 @@ class PublicSettingsView(APIView):
) )
class PublicUsageView(APIView):
"""Aggregate counts for the public site. No emails, names, or avatars."""
permission_classes = (permissions.AllowAny,)
authentication_classes = ()
def get(self, request):
from django.utils import timezone
return Response(
{
"account_count": CustomUser.objects.filter(
deleted=False, is_active=True
).count(),
"document_count": Document.objects.count(),
"as_of": timezone.now().isoformat(),
}
)
class PublicWaitlistView(APIView):
"""Store a waitlist email. Same response for new and existing addresses."""
permission_classes = (permissions.AllowAny,)
authentication_classes = ()
def post(self, request):
from django.core.exceptions import ValidationError
from django.core.validators import validate_email
raw = request.data.get("email") if isinstance(request.data, dict) else None
email = str(raw or "").strip().lower()
if len(email) > 254:
return Response(
{"detail": "Enter a valid email."},
status=status.HTTP_400_BAD_REQUEST,
)
try:
validate_email(email)
except ValidationError:
return Response(
{"detail": "Enter a valid email."},
status=status.HTTP_400_BAD_REQUEST,
)
WaitlistEntry.objects.get_or_create(email=email)
return Response({"detail": "You're on the list."}, status=status.HTTP_200_OK)
class CustomUserCreate(APIView): class CustomUserCreate(APIView):
"""Self-serve registration. Gated by ENABLE_ACCOUNT_REGISTRATION (default off).""" """Password self-serve registration. Gated by ENABLE_ACCOUNT_REGISTRATION."""
permission_classes = (permissions.AllowAny,) permission_classes = (permissions.AllowAny,)
authentication_classes = () authentication_classes = ()
@@ -584,13 +634,37 @@ class ConversationsView(APIView):
class ConversationPreferences(APIView): class ConversationPreferences(APIView):
def get(self, request, format="json"): def get(self, request, format="json"):
user = request.user user = request.user
return Response({"order": user.conversation_order}, status=status.HTTP_200_OK) return Response(
{
"order": user.conversation_order,
"use_conversation_context": user.use_conversation_context,
},
status=status.HTTP_200_OK,
)
def post(self, request, format="json"): def post(self, request, format="json"):
user = request.user user = request.user
user.conversation_order = not user.conversation_order data = request.data
user.save() update_fields = []
return Response({"order": user.conversation_order}, status=status.HTTP_200_OK)
if "use_conversation_context" in data:
user.use_conversation_context = bool(data.get("use_conversation_context"))
update_fields.append("use_conversation_context")
# Legacy: bare POST or POST with ``order`` toggles conversation_order.
# POST that only sets use_conversation_context leaves order unchanged.
if "order" in data or "use_conversation_context" not in data:
user.conversation_order = not user.conversation_order
update_fields.append("conversation_order")
user.save(update_fields=update_fields)
return Response(
{
"order": user.conversation_order,
"use_conversation_context": user.use_conversation_context,
},
status=status.HTTP_200_OK,
)
class ConversationDetailView(APIView): class ConversationDetailView(APIView):
@@ -606,7 +680,7 @@ class ConversationDetailView(APIView):
prompts = list( prompts = list(
Prompt.objects.filter( Prompt.objects.filter(
conversation__id=conversation_id, conversation__user=request.user conversation__id=conversation_id, conversation__user=request.user
) ).order_by("created", "id")
) )
serializer = PromptSerializer( serializer = PromptSerializer(
prompts, prompts,
@@ -660,7 +734,7 @@ class ConversationDetailView(APIView):
messages = [] messages = []
for prompt_obj in Prompt.objects.filter( for prompt_obj in Prompt.objects.filter(
conversation__id=conversation_id conversation__id=conversation_id
): ).order_by("created", "id"):
messages.append( messages.append(
{ {
"content": prompt_obj.message, "content": prompt_obj.message,
+10 -1
View File
@@ -102,6 +102,11 @@ class OAuthStartView(APIView):
}, },
status=status.HTTP_400_BAD_REQUEST, status=status.HTTP_400_BAD_REQUEST,
) )
if intent in {"login", "signup"} and not settings.ENABLE_SSO:
return Response(
{"detail": "Single sign-on is not enabled."},
status=status.HTTP_403_FORBIDDEN,
)
if intent == "signup" and not settings.ENABLE_ACCOUNT_REGISTRATION: if intent == "signup" and not settings.ENABLE_ACCOUNT_REGISTRATION:
return Response( return Response(
{"detail": "Account registration is disabled."}, {"detail": "Account registration is disabled."},
@@ -184,6 +189,8 @@ class OAuthCallbackView(APIView):
return self._handle_drive_link_callback( return self._handle_drive_link_callback(
request, provider=provider, code=code, state_data=state_data request, provider=provider, code=code, state_data=state_data
) )
if not settings.ENABLE_SSO:
return _redirect_error("sso_disabled", "Single sign-on is not enabled.")
try: try:
redirect_uri = _callback_redirect_uri(request, provider) redirect_uri = _callback_redirect_uri(request, provider)
@@ -261,5 +268,7 @@ class OAuthCallbackView(APIView):
def oauth_public_flags() -> dict: def oauth_public_flags() -> dict:
"""Feature flags for /public/settings/.""" """Feature flags for /public/settings/. SSO buttons stay hidden until ENABLE_SSO."""
if not settings.ENABLE_SSO:
return {"google": False, "microsoft": False}
return configured_providers() return configured_providers()
+23 -4
View File
@@ -169,6 +169,10 @@ OLLAMA_NUM_CTX_UTILITY = int(env("OLLAMA_NUM_CTX_UTILITY", "4096") or "4096")
# sub-agents run independent plan steps on a smaller/cheaper model. # sub-agents run independent plan steps on a smaller/cheaper model.
OLLAMA_MODEL_ORCHESTRATOR = env("OLLAMA_MODEL_ORCHESTRATOR", "") or OLLAMA_MODEL OLLAMA_MODEL_ORCHESTRATOR = env("OLLAMA_MODEL_ORCHESTRATOR", "") or OLLAMA_MODEL
OLLAMA_MODEL_SUBAGENT = env("OLLAMA_MODEL_SUBAGENT", "") or OLLAMA_MODEL_UTILITY OLLAMA_MODEL_SUBAGENT = env("OLLAMA_MODEL_SUBAGENT", "") or OLLAMA_MODEL_UTILITY
# Vision turns (image + prompt). Must be a vision-capable Ollama tag.
# Empty falls back to gemma4:latest inside ollama_model_for_role, not OLLAMA_MODEL.
OLLAMA_MODEL_VISION = env("OLLAMA_MODEL_VISION", "") or ""
OLLAMA_NUM_CTX_VISION = int(env("OLLAMA_NUM_CTX_VISION", "8192") or "8192")
OLLAMA_NUM_CTX_ORCHESTRATOR = int( OLLAMA_NUM_CTX_ORCHESTRATOR = int(
env("OLLAMA_NUM_CTX_ORCHESTRATOR", "16384") or "16384" env("OLLAMA_NUM_CTX_ORCHESTRATOR", "16384") or "16384"
) )
@@ -381,9 +385,12 @@ AGENT_FETCH_URL_MAX_BYTES = int(
env("AGENT_FETCH_URL_MAX_BYTES", str(2 * 1024 * 1024)) or str(2 * 1024 * 1024) env("AGENT_FETCH_URL_MAX_BYTES", str(2 * 1024 * 1024)) or str(2 * 1024 * 1024)
) )
# Self-serve account registration (sign-up page). Default off — enable via # Password self-serve sign-up. On by default (#80). Set false in the
# control-node secret (chat_backend_<env>.env) when ready for public sign-up. # control-node secret (chat_backend_<env>.env) to close public registration.
ENABLE_ACCOUNT_REGISTRATION = env_bool("ENABLE_ACCOUNT_REGISTRATION", False) ENABLE_ACCOUNT_REGISTRATION = env_bool("ENABLE_ACCOUNT_REGISTRATION", True)
# Browser SSO (Google / Microsoft sign-in and sign-up). Off until explicitly
# enabled. OAuth client ids may still be set for Drive linking (#47).
ENABLE_SSO = env_bool("ENABLE_SSO", False)
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# OAuth SSO (Google / Microsoft) — #24 # OAuth SSO (Google / Microsoft) — #24
@@ -408,8 +415,20 @@ STRIPE_WEBHOOK_SECRET = env("STRIPE_WEBHOOK_SECRET", "") or ""
# price_data built from SUBSCRIPTION_PRICE_* below. # price_data built from SUBSCRIPTION_PRICE_* below.
STRIPE_PRICE_ID = env("STRIPE_PRICE_ID", "") or "" STRIPE_PRICE_ID = env("STRIPE_PRICE_ID", "") or ""
# RevenueCat webhook Authorization bearer secret + optional product→plan map. # Legacy RevenueCat webhook secret. Store grants no longer use that webhook.
REVENUECAT_WEBHOOK_SECRET = env("REVENUECAT_WEBHOOK_SECRET", "") or "" REVENUECAT_WEBHOOK_SECRET = env("REVENUECAT_WEBHOOK_SECRET", "") or ""
# Native store verification. Apple signed transactions chain to the bundled
# Apple Root CA G3 unless APPLE_IAP_ROOT_PEM adds another PEM.
APPLE_IAP_BUNDLE_ID = (
env("APPLE_IAP_BUNDLE_ID", "ai.hesychia.chat") or "ai.hesychia.chat"
)
APPLE_IAP_ROOT_PEM = env("APPLE_IAP_ROOT_PEM", "") or ""
GOOGLE_PLAY_PACKAGE_NAME = (
env("GOOGLE_PLAY_PACKAGE_NAME", "ai.hesychia.chat") or "ai.hesychia.chat"
)
GOOGLE_PLAY_SERVICE_ACCOUNT_JSON = (
env("GOOGLE_PLAY_SERVICE_ACCOUNT_JSON", "") or ""
)
_rc_product_map_raw = env("REVENUECAT_PRODUCT_PLAN_MAP", "") or "" _rc_product_map_raw = env("REVENUECAT_PRODUCT_PLAN_MAP", "") or ""
REVENUECAT_PRODUCT_PLAN_MAP: dict = {} REVENUECAT_PRODUCT_PLAN_MAP: dict = {}
if _rc_product_map_raw.strip(): if _rc_product_map_raw.strip():
@@ -0,0 +1,15 @@
-----BEGIN CERTIFICATE-----
MIICQzCCAcmgAwIBAgIILcX8iNLFS5UwCgYIKoZIzj0EAwMwZzEbMBkGA1UEAwwS
QXBwbGUgUm9vdCBDQSAtIEczMSYwJAYDVQQLDB1BcHBsZSBDZXJ0aWZpY2F0aW9u
IEF1dGhvcml0eTETMBEGA1UECgwKQXBwbGUgSW5jLjELMAkGA1UEBhMCVVMwHhcN
MTQwNDMwMTgxOTA2WhcNMzkwNDMwMTgxOTA2WjBnMRswGQYDVQQDDBJBcHBsZSBS
b290IENBIC0gRzMxJjAkBgNVBAsMHUFwcGxlIENlcnRpZmljYXRpb24gQXV0aG9y
aXR5MRMwEQYDVQQKDApBcHBsZSBJbmMuMQswCQYDVQQGEwJVUzB2MBAGByqGSM49
AgEGBSuBBAAiA2IABJjpLz1AcqTtkyJygRMc3RCV8cWjTnHcFBbZDuWmBSp3ZHtf
TjjTuxxEtX/1H7YyYl3J6YRbTzBPEVoA/VhYDKX1DyxNB0cTddqXl5dvMVztK517
IDvYuVTZXpmkOlEKMaNCMEAwHQYDVR0OBBYEFLuw3qFYM4iapIqZ3r6966/ayySr
MA8GA1UdEwEB/wQFMAMBAf8wDgYDVR0PAQH/BAQDAgEGMAoGCCqGSM49BAMDA2gA
MGUCMQCD6cHEFl4aXTQY2e3v9GwOAEZLuN+yRhHFD/3meoyhpmvOwgPUnPWTxnS4
at+qIxUCMG1mihDK1A3UT82NQz60imOlM27jbdoXt2QfyFMm+YhidDkLF1vLUagM
6BgD56KyKA==
-----END CERTIFICATE-----
@@ -0,0 +1,56 @@
from django.db import migrations, models
class Migration(migrations.Migration):
dependencies = [
("finance", "0005_revenuecat_fields"),
]
operations = [
migrations.AlterField(
model_name="invoice",
name="provider",
field=models.CharField(
choices=[
("stripe", "Stripe"),
("revenuecat", "RevenueCat"),
("app_store", "App Store"),
("play_store", "Play Store"),
],
default="stripe",
max_length=32,
),
),
migrations.AlterField(
model_name="payment",
name="provider",
field=models.CharField(
choices=[
("stripe", "Stripe"),
("revenuecat", "RevenueCat"),
("app_store", "App Store"),
("play_store", "Play Store"),
],
default="stripe",
max_length=32,
),
),
migrations.AlterField(
model_name="usersubscription",
name="source",
field=models.CharField(
choices=[
("none", "None"),
("stripe", "Stripe"),
("revenuecat", "RevenueCat"),
("app_store", "App Store"),
("play_store", "Play Store"),
("backer", "Backer"),
("admin", "Admin"),
],
default="none",
max_length=32,
),
),
]
+6
View File
@@ -134,6 +134,8 @@ class UserSubscription(TimeInfoBase):
NONE = "none", "None" NONE = "none", "None"
STRIPE = "stripe", "Stripe" STRIPE = "stripe", "Stripe"
REVENUECAT = "revenuecat", "RevenueCat" REVENUECAT = "revenuecat", "RevenueCat"
APP_STORE = "app_store", "App Store"
PLAY_STORE = "play_store", "Play Store"
BACKER = "backer", "Backer" BACKER = "backer", "Backer"
ADMIN = "admin", "Admin" ADMIN = "admin", "Admin"
@@ -214,6 +216,8 @@ class Invoice(TimeInfoBase):
class Provider(models.TextChoices): class Provider(models.TextChoices):
STRIPE = "stripe", "Stripe" STRIPE = "stripe", "Stripe"
REVENUECAT = "revenuecat", "RevenueCat" REVENUECAT = "revenuecat", "RevenueCat"
APP_STORE = "app_store", "App Store"
PLAY_STORE = "play_store", "Play Store"
class Status(models.TextChoices): class Status(models.TextChoices):
DRAFT = "draft", "Draft" DRAFT = "draft", "Draft"
@@ -308,6 +312,8 @@ class Payment(TimeInfoBase):
class Provider(models.TextChoices): class Provider(models.TextChoices):
STRIPE = "stripe", "Stripe" STRIPE = "stripe", "Stripe"
REVENUECAT = "revenuecat", "RevenueCat" REVENUECAT = "revenuecat", "RevenueCat"
APP_STORE = "app_store", "App Store"
PLAY_STORE = "play_store", "Play Store"
class Status(models.TextChoices): class Status(models.TextChoices):
PENDING = "pending", "Pending" PENDING = "pending", "Pending"
+16
View File
@@ -59,8 +59,23 @@ class PortalSessionSerializer(serializers.Serializer):
return_url = serializers.URLField(required=False, allow_blank=False) return_url = serializers.URLField(required=False, allow_blank=False)
class StorePurchaseConfirmSerializer(serializers.Serializer):
store = serializers.ChoiceField(choices=("app_store", "play_store"))
product_id = serializers.CharField(max_length=255)
transaction_id = serializers.CharField(
max_length=255, required=False, allow_blank=True
)
purchase_token = serializers.CharField(
max_length=4096, required=False, allow_blank=True
)
signed_transaction = serializers.CharField(required=False, allow_blank=True)
class SubscriptionPlanSerializer(serializers.ModelSerializer): class SubscriptionPlanSerializer(serializers.ModelSerializer):
features = serializers.SerializerMethodField() features = serializers.SerializerMethodField()
store_product_id = serializers.CharField(
source="revenuecat_product_id", read_only=True
)
class Meta: class Meta:
model = SubscriptionPlan model = SubscriptionPlan
@@ -78,6 +93,7 @@ class SubscriptionPlanSerializer(serializers.ModelSerializer):
"prompt_window_hours", "prompt_window_hours",
"monthly_token_quota", "monthly_token_quota",
"sort_order", "sort_order",
"store_product_id",
] ]
read_only_fields = fields read_only_fields = fields
+5 -2
View File
@@ -247,6 +247,8 @@ def needs_checkout(user) -> bool:
UserSubscription.Source.ADMIN, UserSubscription.Source.ADMIN,
UserSubscription.Source.STRIPE, UserSubscription.Source.STRIPE,
UserSubscription.Source.REVENUECAT, UserSubscription.Source.REVENUECAT,
UserSubscription.Source.APP_STORE,
UserSubscription.Source.PLAY_STORE,
): ):
return False return False
return True return True
@@ -457,8 +459,9 @@ def assign_plan_from_revenuecat(
cancel_at_period_end: bool | None = None, cancel_at_period_end: bool | None = None,
current_period_end=None, current_period_end=None,
keep_existing_plan_if_unknown: bool = False, keep_existing_plan_if_unknown: bool = False,
source: str | None = None,
) -> UserSubscription: ) -> UserSubscription:
"""Assign a catalog plan from a RevenueCat store purchase event.""" """Assign a catalog plan from a verified store purchase."""
seed_subscription_plans(update_existing=False) seed_subscription_plans(update_existing=False)
existing = ( existing = (
UserSubscription.objects.filter(user=user).select_related("plan").first() UserSubscription.objects.filter(user=user).select_related("plan").first()
@@ -493,7 +496,7 @@ def assign_plan_from_revenuecat(
sub = assign_plan( sub = assign_plan(
user, user,
plan=plan, plan=plan,
source=UserSubscription.Source.REVENUECAT, source=source or UserSubscription.Source.REVENUECAT,
status=status, status=status,
revenuecat_original_transaction_id=revenuecat_original_transaction_id or "", revenuecat_original_transaction_id=revenuecat_original_transaction_id or "",
log_auth_event=False, log_auth_event=False,
@@ -0,0 +1,617 @@
"""Verify App Store and Play Store purchases and grant entitlements.
The mobile app completes the purchase with StoreKit or Play Billing, then
POSTs the signed transaction or purchase token. This module checks that
payload with Apple or Google and writes the subscription ledger. Store
webhooks are not accepted.
"""
from __future__ import annotations
import base64
import json
import logging
import time
import uuid
from dataclasses import dataclass
from datetime import datetime, timezone as dt_timezone
from pathlib import Path
from typing import Any
from urllib.parse import quote
import jwt
import requests
from cryptography import x509
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import ec, padding
from django.conf import settings
from django.db import transaction
from django.utils import timezone
from monetization.models import Invoice, Payment, UserSubscription
from monetization.services.plans import (
assign_plan_from_revenuecat,
log_subscription_auth_event,
resolve_plan_from_revenuecat_product,
)
logger = logging.getLogger(__name__)
# uuid5(NAMESPACE_URL, this name). Do not change; clients receive the result.
_STORE_ACCOUNT_NAMESPACE = uuid.uuid5(
uuid.NAMESPACE_URL,
"https://chat.aimloperations.com/store-account",
)
_APP_STORE_CERT_OID = x509.ObjectIdentifier("1.2.840.113635.100.6.11.1")
_PLAY_SCOPE = "https://www.googleapis.com/auth/androidpublisher"
_BUNDLED_APPLE_ROOT = (
Path(__file__).resolve().parent.parent / "certs" / "apple_root_g3.pem"
)
_STORE_SOURCES = frozenset(
{
UserSubscription.Source.APP_STORE,
UserSubscription.Source.PLAY_STORE,
UserSubscription.Source.REVENUECAT,
}
)
class StoreNotConfiguredError(Exception):
"""Apple or Google credentials required to verify a purchase are missing."""
class StorePurchaseRejected(Exception):
"""The store payload is not an active purchase for this user."""
class StoreVerificationError(Exception):
"""The store API could not be reached."""
@dataclass(frozen=True)
class VerifiedStorePurchase:
source: str
store_label: str
product_id: str
transaction_id: str
original_transaction_id: str
expires_at: datetime | None
purchased_at: datetime | None
cancel_at_period_end: bool
status: str
price_cents: int
currency: str
account_token: str
def store_account_token(user_id: int) -> str:
"""Stable store account id for this user. Pass it to StoreKit and Play."""
return str(uuid.uuid5(_STORE_ACCOUNT_NAMESPACE, str(user_id)))
def expire_lapsed_store_subscription(sub: UserSubscription | None):
"""End store access once the verified period has passed.
Renewals arrive when the app confirms the store entitlement again.
"""
if sub is None or sub.source not in _STORE_SOURCES:
return sub
if sub.status != UserSubscription.Status.ACTIVE or not sub.current_period_end:
return sub
if sub.current_period_end > timezone.now():
return sub
prev_status = sub.status
sub.status = UserSubscription.Status.CANCELED
sub.cancel_at_period_end = False
sub.save(update_fields=["status", "cancel_at_period_end", "last_modified"])
if prev_status != UserSubscription.Status.CANCELED:
log_subscription_auth_event(
sub.user,
started=False,
detail=(
f"plan={sub.plan.slug if sub.plan_id else 'none'} "
f"source={sub.source} status={sub.status} "
"reason=store_period_ended"
),
)
return sub
def confirm_store_purchase(user, payload: dict[str, Any]) -> UserSubscription:
"""Verify one store purchase for the authenticated user and grant it."""
store = (payload.get("store") or "").strip().lower()
product_id = (payload.get("product_id") or "").strip()
if store == UserSubscription.Source.APP_STORE:
verified = verify_apple_signed_transaction(
signed_transaction=(payload.get("signed_transaction") or "").strip(),
product_id=product_id,
transaction_id=(payload.get("transaction_id") or "").strip(),
)
elif store == UserSubscription.Source.PLAY_STORE:
verified = verify_google_subscription(
product_id=product_id,
purchase_token=(payload.get("purchase_token") or "").strip(),
transaction_id=(payload.get("transaction_id") or "").strip(),
)
else:
raise StorePurchaseRejected("store must be app_store or play_store.")
if verified.product_id != product_id:
raise StorePurchaseRejected("Store product does not match this purchase.")
return grant_verified_store_purchase(user, verified)
def verify_apple_signed_transaction(
*,
signed_transaction: str,
product_id: str,
transaction_id: str = "",
) -> VerifiedStorePurchase:
"""Verify a StoreKit 2 signed transaction against Apple's root."""
if not signed_transaction:
raise StorePurchaseRejected("App Store purchases require signed_transaction.")
bundle_id = (getattr(settings, "APPLE_IAP_BUNDLE_ID", "") or "").strip()
if not bundle_id:
raise StoreNotConfiguredError("APPLE_IAP_BUNDLE_ID is not configured.")
roots = _load_apple_roots()
if not roots:
raise StoreNotConfiguredError("No Apple root certificate is configured.")
try:
header = jwt.get_unverified_header(signed_transaction)
except jwt.PyJWTError as exc:
raise StorePurchaseRejected("Invalid App Store signed transaction.") from exc
x5c = header.get("x5c") if isinstance(header, dict) else None
if not isinstance(x5c, list) or not x5c:
raise StorePurchaseRejected("App Store transaction is missing its certificate.")
try:
certs = [_load_der_cert(item) for item in x5c]
leaf = _trusted_leaf(certs, roots)
except StorePurchaseRejected:
raise
except Exception as exc:
raise StorePurchaseRejected("App Store certificate chain is not trusted.") from exc
_require_app_store_oid(leaf)
try:
claims = jwt.decode(
signed_transaction,
leaf.public_key(),
algorithms=["ES256"],
options={"verify_aud": False, "verify_exp": False},
)
except jwt.PyJWTError as exc:
raise StorePurchaseRejected("App Store transaction signature is invalid.") from exc
if not isinstance(claims, dict):
raise StorePurchaseRejected("App Store transaction payload is invalid.")
claim_bundle = str(claims.get("bundleId") or "")
if claim_bundle != bundle_id:
raise StorePurchaseRejected("App Store bundle id does not match this app.")
claim_product = str(claims.get("productId") or "")
if product_id and claim_product != product_id:
raise StorePurchaseRejected("App Store product does not match this purchase.")
claim_txn = str(claims.get("transactionId") or "")
if transaction_id and claim_txn and transaction_id != claim_txn:
raise StorePurchaseRejected("App Store transaction id does not match.")
if not claim_txn:
raise StorePurchaseRejected("App Store transaction id is missing.")
expires_at = _ms_to_dt(claims.get("expiresDate"))
purchased_at = _ms_to_dt(claims.get("purchaseDate"))
revoked = _ms_to_dt(claims.get("revocationDate")) is not None
now = timezone.now()
if revoked or (expires_at is not None and expires_at <= now):
status = UserSubscription.Status.CANCELED
else:
status = UserSubscription.Status.ACTIVE
price_cents = 0
raw_price = claims.get("price")
try:
if raw_price is not None:
price_cents = max(0, int(raw_price) // 10)
except (TypeError, ValueError):
price_cents = 0
currency = str(claims.get("currency") or "usd").lower()
original = str(claims.get("originalTransactionId") or claim_txn)
return VerifiedStorePurchase(
source=UserSubscription.Source.APP_STORE,
store_label="APP_STORE",
product_id=claim_product,
transaction_id=claim_txn,
original_transaction_id=original,
expires_at=expires_at,
purchased_at=purchased_at,
cancel_at_period_end=False,
status=status,
price_cents=price_cents,
currency=currency,
account_token=str(claims.get("appAccountToken") or ""),
)
def verify_google_subscription(
*,
product_id: str,
purchase_token: str,
transaction_id: str = "",
) -> VerifiedStorePurchase:
"""Verify a Play Billing purchase token with the Android Publisher API."""
if not purchase_token:
raise StorePurchaseRejected("Play Store purchases require purchase_token.")
if not product_id:
raise StorePurchaseRejected("Play Store purchases require product_id.")
body = fetch_google_subscription(purchase_token)
return _verified_from_google(
body,
product_id=product_id,
purchase_token=purchase_token,
transaction_id=transaction_id,
)
def fetch_google_subscription(purchase_token: str) -> dict[str, Any]:
"""GET purchases.subscriptionsv2 for this token. Tests patch this."""
package_name = (getattr(settings, "GOOGLE_PLAY_PACKAGE_NAME", "") or "").strip()
raw = (getattr(settings, "GOOGLE_PLAY_SERVICE_ACCOUNT_JSON", "") or "").strip()
if not package_name or not raw:
raise StoreNotConfiguredError(
"Google Play verification is not configured "
"(GOOGLE_PLAY_PACKAGE_NAME and GOOGLE_PLAY_SERVICE_ACCOUNT_JSON)."
)
info = _load_service_account_info(raw)
try:
access_token = _google_access_token(info)
url = (
"https://androidpublisher.googleapis.com/androidpublisher/v3/"
f"applications/{quote(package_name, safe='')}/purchases/"
f"subscriptionsv2/tokens/{quote(purchase_token, safe='')}"
)
response = requests.get(
url,
headers={"Authorization": f"Bearer {access_token}"},
timeout=15,
)
except (requests.RequestException, ValueError, OSError) as exc:
logger.exception("Play Store verification request failed")
raise StoreVerificationError("Play Store verification failed.") from exc
if response.status_code == 404:
raise StorePurchaseRejected("Play Store purchase was not found.")
if response.status_code >= 400:
logger.error(
"Play Store verification HTTP %s",
response.status_code,
)
if response.status_code >= 500:
raise StoreVerificationError("Play Store verification failed.")
raise StorePurchaseRejected("Play Store rejected this purchase token.")
try:
body = response.json()
except ValueError as exc:
raise StoreVerificationError("Play Store returned an invalid response.") from exc
if not isinstance(body, dict):
raise StoreVerificationError("Play Store returned an invalid response.")
return body
def _verified_from_google(
body: dict[str, Any],
*,
product_id: str,
purchase_token: str,
transaction_id: str,
) -> VerifiedStorePurchase:
line_items = body.get("lineItems") or []
match = None
if isinstance(line_items, list):
for item in line_items:
if isinstance(item, dict) and str(item.get("productId") or "") == product_id:
match = item
break
if match is None:
raise StorePurchaseRejected("Play Store product does not match this purchase.")
state = str(body.get("subscriptionState") or "")
expires_at = _parse_rfc3339(match.get("expiryTime"))
now = timezone.now()
expired = expires_at is not None and expires_at <= now
cancel_at_period_end = False
if state in {"SUBSCRIPTION_STATE_PENDING", "SUBSCRIPTION_STATE_PENDING_PURCHASE_CANCELED"}:
raise StorePurchaseRejected("Play Store purchase is still pending.")
if state in {"SUBSCRIPTION_STATE_ACTIVE", "SUBSCRIPTION_STATE_IN_GRACE_PERIOD"} and not expired:
status = UserSubscription.Status.ACTIVE
elif state == "SUBSCRIPTION_STATE_CANCELED" and not expired:
status = UserSubscription.Status.ACTIVE
cancel_at_period_end = True
elif state == "SUBSCRIPTION_STATE_ON_HOLD":
status = UserSubscription.Status.PAST_DUE
else:
status = UserSubscription.Status.CANCELED
auto_renew = (match.get("autoRenewingPlan") or {}) if isinstance(match, dict) else {}
if isinstance(auto_renew, dict) and auto_renew.get("autoRenewEnabled") is False and not expired:
if status == UserSubscription.Status.ACTIVE:
cancel_at_period_end = True
order_id = str(match.get("latestSuccessfulOrderId") or transaction_id or purchase_token)
identifiers = body.get("externalAccountIdentifiers") or {}
account_token = ""
if isinstance(identifiers, dict):
account_token = str(identifiers.get("obfuscatedExternalAccountId") or "")
price_cents, currency = _google_price(auto_renew if isinstance(auto_renew, dict) else {})
return VerifiedStorePurchase(
source=UserSubscription.Source.PLAY_STORE,
store_label="PLAY_STORE",
product_id=product_id,
transaction_id=order_id,
original_transaction_id=purchase_token,
expires_at=expires_at,
purchased_at=None,
cancel_at_period_end=cancel_at_period_end,
status=status,
price_cents=price_cents,
currency=currency,
account_token=account_token,
)
@transaction.atomic
def grant_verified_store_purchase(user, verified: VerifiedStorePurchase) -> UserSubscription:
plan = resolve_plan_from_revenuecat_product(verified.product_id)
if plan is None:
raise StorePurchaseRejected("This store product is not mapped to a plan.")
expected = store_account_token(user.pk)
if verified.account_token and verified.account_token.lower() != expected.lower():
raise StorePurchaseRejected("Purchase account token does not match this user.")
original = verified.original_transaction_id
if original:
taken = (
UserSubscription.objects.select_for_update()
.filter(revenuecat_original_transaction_id=original)
.exclude(user=user)
.exists()
)
if taken:
raise StorePurchaseRejected(
"This purchase is already linked to another account."
)
sub = assign_plan_from_revenuecat(
user,
plan_slug=plan.slug,
product_id=verified.product_id,
revenuecat_original_transaction_id=original,
status=verified.status,
cancel_at_period_end=verified.cancel_at_period_end,
current_period_end=verified.expires_at,
keep_existing_plan_if_unknown=False,
source=verified.source,
)
if verified.status != UserSubscription.Status.ACTIVE:
return sub
amount = verified.price_cents or plan.price_cents
currency = (verified.currency or plan.currency or "usd").lower()
event_id = f"{verified.source}:{verified.transaction_id}"
invoice, _created = Invoice.objects.update_or_create(
revenuecat_event_id=event_id,
defaults={
"user": user,
"company": getattr(user, "company", None),
"provider": verified.source,
"status": Invoice.Status.PAID,
"currency": currency,
"amount_due": amount,
"amount_paid": amount,
"period_start": verified.purchased_at,
"period_end": verified.expires_at,
"revenuecat_store": verified.store_label,
"description": f"{plan.name} — {verified.store_label}",
"hosted_invoice_url": "",
},
)
Payment.objects.update_or_create(
revenuecat_transaction_id=verified.transaction_id,
defaults={
"user": user,
"company": getattr(user, "company", None),
"invoice": invoice,
"provider": verified.source,
"status": Payment.Status.SUCCEEDED,
"currency": currency,
"amount": amount,
"paid_at": verified.purchased_at or timezone.now(),
"failure_message": "",
},
)
return sub
def _load_apple_roots() -> list[x509.Certificate]:
pems: list[str] = []
extra = (getattr(settings, "APPLE_IAP_ROOT_PEM", "") or "").strip()
if extra:
pems.append(extra)
if _BUNDLED_APPLE_ROOT.is_file():
pems.append(_BUNDLED_APPLE_ROOT.read_text(encoding="utf-8"))
roots: list[x509.Certificate] = []
for pem in pems:
try:
roots.append(x509.load_pem_x509_certificate(pem.encode("utf-8")))
except ValueError:
logger.error("Ignoring an unreadable Apple root certificate")
return roots
def _load_der_cert(value: Any) -> x509.Certificate:
if not isinstance(value, str) or not value.strip():
raise StorePurchaseRejected("App Store certificate is missing.")
try:
der = base64.b64decode(value)
return x509.load_der_x509_certificate(der)
except (ValueError, TypeError) as exc:
raise StorePurchaseRejected("App Store certificate is invalid.") from exc
def _trusted_leaf(
certs: list[x509.Certificate],
roots: list[x509.Certificate],
) -> x509.Certificate:
if not certs:
raise StorePurchaseRejected("App Store certificate chain is empty.")
now = datetime.now(dt_timezone.utc)
for cert in certs:
if cert.not_valid_before_utc > now or cert.not_valid_after_utc < now:
raise StorePurchaseRejected("App Store certificate is not currently valid.")
for index in range(len(certs) - 1):
_verify_signed_by(certs[index], certs[index + 1])
if certs[index].issuer != certs[index + 1].subject:
raise StorePurchaseRejected("App Store certificate issuer mismatch.")
last = certs[-1]
root_fps = {root.fingerprint(hashes.SHA256()) for root in roots}
if last.fingerprint(hashes.SHA256()) in root_fps:
return certs[0]
for root in roots:
try:
_verify_signed_by(last, root)
except Exception:
continue
if last.issuer == root.subject:
return certs[0]
raise StorePurchaseRejected("App Store certificate chain is not trusted.")
def _verify_signed_by(child: x509.Certificate, issuer: x509.Certificate) -> None:
public_key = issuer.public_key()
signature_hash = child.signature_hash_algorithm
if signature_hash is None:
raise StorePurchaseRejected("App Store certificate has no signature hash.")
if isinstance(public_key, ec.EllipticCurvePublicKey):
public_key.verify(
child.signature,
child.tbs_certificate_bytes,
ec.ECDSA(signature_hash),
)
return
public_key.verify(
child.signature,
child.tbs_certificate_bytes,
padding.PKCS1v15(),
signature_hash,
)
def _require_app_store_oid(leaf: x509.Certificate) -> None:
try:
leaf.extensions.get_extension_for_oid(_APP_STORE_CERT_OID)
except x509.ExtensionNotFound as exc:
raise StorePurchaseRejected(
"Signed transaction certificate is not an App Store certificate."
) from exc
def _google_access_token(info: dict[str, Any]) -> str:
"""OAuth JWT bearer grant for a Play Console service account."""
client_email = info.get("client_email")
private_key = info.get("private_key")
token_uri = info.get("token_uri") or "https://oauth2.googleapis.com/token"
if not client_email or not private_key:
raise StoreNotConfiguredError(
"GOOGLE_PLAY_SERVICE_ACCOUNT_JSON is missing client_email or private_key."
)
now = int(time.time())
assertion = jwt.encode(
{
"iss": client_email,
"scope": _PLAY_SCOPE,
"aud": token_uri,
"iat": now,
"exp": now + 3600,
},
private_key,
algorithm="RS256",
)
try:
response = requests.post(
token_uri,
data={
"grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer",
"assertion": assertion,
},
timeout=15,
)
except requests.RequestException as exc:
raise StoreVerificationError("Play Store verification failed.") from exc
if response.status_code >= 400:
raise StoreVerificationError("Play Store authentication failed.")
try:
token = response.json().get("access_token")
except ValueError as exc:
raise StoreVerificationError("Play Store authentication failed.") from exc
if not token:
raise StoreVerificationError("Play Store authentication failed.")
return str(token)
def _load_service_account_info(raw: str) -> dict[str, Any]:
text = raw
if not raw.lstrip().startswith("{"):
path = Path(raw)
if path.is_file():
text = path.read_text(encoding="utf-8")
try:
info = json.loads(text)
except json.JSONDecodeError as exc:
raise StoreNotConfiguredError(
"GOOGLE_PLAY_SERVICE_ACCOUNT_JSON is not valid JSON."
) from exc
if not isinstance(info, dict):
raise StoreNotConfiguredError(
"GOOGLE_PLAY_SERVICE_ACCOUNT_JSON must be a JSON object."
)
return info
def _ms_to_dt(value: Any) -> datetime | None:
if value is None or value == "":
return None
try:
ms = int(value)
except (TypeError, ValueError):
return None
if ms <= 0:
return None
return datetime.fromtimestamp(ms / 1000.0, tz=dt_timezone.utc)
def _parse_rfc3339(value: Any) -> datetime | None:
if not value or not isinstance(value, str):
return None
try:
parsed = datetime.fromisoformat(value.replace("Z", "+00:00"))
except ValueError:
return None
if parsed.tzinfo is None:
return parsed.replace(tzinfo=dt_timezone.utc)
return parsed
def _google_price(auto_renew: dict[str, Any]) -> tuple[int, str]:
price = auto_renew.get("recurringPrice") if isinstance(auto_renew, dict) else None
if not isinstance(price, dict):
return 0, "usd"
currency = str(price.get("currencyCode") or "usd").lower()
try:
units = int(price.get("units") or 0)
except (TypeError, ValueError):
units = 0
try:
nanos = int(price.get("nanos") or 0)
except (TypeError, ValueError):
nanos = 0
cents = units * 100 + nanos // 10_000_000
return max(0, cents), currency
+1 -51
View File
@@ -2,10 +2,7 @@
from __future__ import annotations from __future__ import annotations
from django.test import TestCase, override_settings from django.test import TestCase
from django.urls import reverse
from rest_framework import status
from rest_framework.test import APIClient
from chat_backend.tests.factories import make_user from chat_backend.tests.factories import make_user
from monetization.models import Invoice, Payment, SubscriptionPlan, UserSubscription from monetization.models import Invoice, Payment, SubscriptionPlan, UserSubscription
@@ -129,50 +126,3 @@ class RevenueCatDispatchTests(TestCase):
self.assertEqual(sub.status, UserSubscription.Status.CANCELED) self.assertEqual(sub.status, UserSubscription.Status.CANCELED)
@override_settings(REVENUECAT_WEBHOOK_SECRET="test-rc-secret")
class RevenueCatWebhookViewTests(TestCase):
def setUp(self):
seed_subscription_plans(update_existing=False)
self.client = APIClient()
self.url = reverse("finance_revenuecat_webhook")
self.user = make_user(email="rcview@example.com")
founders = SubscriptionPlan.objects.get(slug="founders")
founders.revenuecat_product_id = "hesychia_founders_monthly"
founders.save(update_fields=["revenuecat_product_id", "last_modified"])
def test_missing_secret_config(self):
with override_settings(REVENUECAT_WEBHOOK_SECRET=""):
response = self.client.post(
self.url, {"event": _rc_event()}, format="json"
)
self.assertEqual(response.status_code, status.HTTP_503_SERVICE_UNAVAILABLE)
def test_unauthorized(self):
response = self.client.post(
self.url,
{"event": _rc_event(app_user_id=str(self.user.pk))},
format="json",
HTTP_AUTHORIZATION="Bearer wrong",
)
self.assertEqual(response.status_code, status.HTTP_401_UNAUTHORIZED)
def test_success(self):
response = self.client.post(
self.url,
{"api_version": "1.0", "event": _rc_event(app_user_id=str(self.user.pk))},
format="json",
HTTP_AUTHORIZATION="Bearer test-rc-secret",
)
self.assertEqual(response.status_code, status.HTTP_200_OK)
self.assertTrue(
UserSubscription.objects.filter(
user=self.user,
source=UserSubscription.Source.REVENUECAT,
status=UserSubscription.Status.ACTIVE,
).exists()
)
self.assertTrue(
Invoice.objects.filter(
user=self.user, provider=Invoice.Provider.REVENUECAT
).exists()
)
@@ -0,0 +1,336 @@
"""StoreKit and Play Billing confirmation. No store webhooks."""
from __future__ import annotations
import base64
from datetime import datetime, timedelta, timezone as dt_timezone
from unittest.mock import patch
import jwt
from cryptography import x509
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.x509.oid import NameOID
from django.test import TestCase, override_settings
from django.urls import reverse
from django.utils import timezone
from rest_framework import status
from rest_framework.test import APITestCase
from chat_backend.tests.factories import make_user
from monetization.models import Invoice, Payment, SubscriptionPlan, UserSubscription
from monetization.services.plans import seed_subscription_plans
from monetization.services.store_purchases import (
StoreNotConfiguredError,
StorePurchaseRejected,
confirm_store_purchase,
store_account_token,
)
_APP_STORE_OID = x509.ObjectIdentifier("1.2.840.113635.100.6.11.1")
_PRODUCT = "hesychia_founders_monthly"
def _cert(*, subject_cn, subject_key, issuer_cert, issuer_key, is_ca, app_store=False):
now = datetime.now(dt_timezone.utc)
name = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, subject_cn)])
issuer_name = issuer_cert.subject if issuer_cert is not None else name
builder = (
x509.CertificateBuilder()
.subject_name(name)
.issuer_name(issuer_name)
.public_key(subject_key.public_key())
.serial_number(x509.random_serial_number())
.not_valid_before(now - timedelta(days=1))
.not_valid_after(now + timedelta(days=30))
)
if is_ca:
builder = builder.add_extension(
x509.BasicConstraints(ca=True, path_length=None),
critical=True,
)
if app_store:
builder = builder.add_extension(
x509.UnrecognizedExtension(_APP_STORE_OID, b"\x05\x00"),
critical=False,
)
signer = issuer_key if issuer_key is not None else subject_key
return builder.sign(signer, hashes.SHA256())
def _chain():
root_key = ec.generate_private_key(ec.SECP256R1())
leaf_key = ec.generate_private_key(ec.SECP256R1())
root = _cert(
subject_cn="Test Apple Root",
subject_key=root_key,
issuer_cert=None,
issuer_key=None,
is_ca=True,
)
leaf = _cert(
subject_cn="Test StoreKit",
subject_key=leaf_key,
issuer_cert=root,
issuer_key=root_key,
is_ca=False,
app_store=True,
)
root_pem = root.public_bytes(serialization.Encoding.PEM).decode("ascii")
return leaf_key, leaf, root, root_pem
def _b64_der(cert) -> str:
der = cert.public_bytes(serialization.Encoding.DER)
return base64.b64encode(der).decode("ascii")
def _signed_transaction(leaf_key, leaf, root, claims):
return jwt.encode(
claims,
leaf_key,
algorithm="ES256",
headers={"x5c": [_b64_der(leaf), _b64_der(root)]},
)
def _future_ms(days=30) -> int:
moment = datetime.now(dt_timezone.utc) + timedelta(days=days)
return int(moment.timestamp() * 1000)
class AppleConfirmTests(TestCase):
def setUp(self):
seed_subscription_plans(update_existing=False)
self.user = make_user(email="apple@example.com")
founders = SubscriptionPlan.objects.get(slug="founders")
founders.revenuecat_product_id = _PRODUCT
founders.save(update_fields=["revenuecat_product_id", "last_modified"])
self.leaf_key, self.leaf, self.root, self.root_pem = _chain()
def _claims(self, **overrides):
base = {
"bundleId": "ai.hesychia.chat",
"productId": _PRODUCT,
"transactionId": "2000000000001",
"originalTransactionId": "2000000000001",
"appAccountToken": store_account_token(self.user.pk),
"purchaseDate": _future_ms(days=-1),
"expiresDate": _future_ms(days=30),
"price": 10000,
"currency": "USD",
}
base.update(overrides)
return base
def _confirm(self, claims):
token = _signed_transaction(self.leaf_key, self.leaf, self.root, claims)
with override_settings(
APPLE_IAP_BUNDLE_ID="ai.hesychia.chat",
APPLE_IAP_ROOT_PEM=self.root_pem,
):
return confirm_store_purchase(
self.user,
{
"store": "app_store",
"product_id": _PRODUCT,
"transaction_id": claims["transactionId"],
"signed_transaction": token,
},
)
def test_signed_transaction_grants_plan(self):
sub = self._confirm(self._claims())
self.assertEqual(sub.source, UserSubscription.Source.APP_STORE)
self.assertEqual(sub.status, UserSubscription.Status.ACTIVE)
self.assertEqual(sub.plan.slug, "founders")
self.assertEqual(sub.revenuecat_original_transaction_id, "2000000000001")
invoice = Invoice.objects.get(user=self.user)
self.assertEqual(invoice.provider, Invoice.Provider.APP_STORE)
self.assertEqual(invoice.amount_paid, 1000)
self.assertEqual(invoice.revenuecat_store, "APP_STORE")
payment = Payment.objects.get(user=self.user)
self.assertEqual(payment.provider, Payment.Provider.APP_STORE)
self.assertEqual(payment.status, Payment.Status.SUCCEEDED)
def test_repeat_confirm_is_idempotent(self):
claims = self._claims()
self._confirm(claims)
self._confirm(claims)
self.assertEqual(Invoice.objects.filter(user=self.user).count(), 1)
self.assertEqual(Payment.objects.filter(user=self.user).count(), 1)
def test_wrong_account_token_rejected(self):
with self.assertRaises(StorePurchaseRejected):
self._confirm(self._claims(appAccountToken="00000000-0000-0000-0000-000000000000"))
def test_untrusted_chain_rejected(self):
token = _signed_transaction(
self.leaf_key, self.leaf, self.root, self._claims()
)
with override_settings(APPLE_IAP_BUNDLE_ID="ai.hesychia.chat", APPLE_IAP_ROOT_PEM=""):
with self.assertRaises(StorePurchaseRejected):
confirm_store_purchase(
self.user,
{
"store": "app_store",
"product_id": _PRODUCT,
"signed_transaction": token,
},
)
def test_revoked_transaction_cancels(self):
sub = self._confirm(self._claims(revocationDate=_future_ms(days=-1)))
self.assertEqual(sub.status, UserSubscription.Status.CANCELED)
self.assertFalse(Payment.objects.filter(user=self.user).exists())
class PlayConfirmTests(TestCase):
def setUp(self):
seed_subscription_plans(update_existing=False)
self.user = make_user(email="play@example.com")
founders = SubscriptionPlan.objects.get(slug="founders")
founders.revenuecat_product_id = _PRODUCT
founders.save(update_fields=["revenuecat_product_id", "last_modified"])
def _body(self, **overrides):
expiry = (datetime.now(dt_timezone.utc) + timedelta(days=30)).strftime(
"%Y-%m-%dT%H:%M:%SZ"
)
body = {
"subscriptionState": "SUBSCRIPTION_STATE_ACTIVE",
"lineItems": [
{
"productId": _PRODUCT,
"expiryTime": expiry,
"latestSuccessfulOrderId": "GPA.1234-5678",
"autoRenewingPlan": {
"autoRenewEnabled": True,
"recurringPrice": {
"currencyCode": "USD",
"units": "10",
"nanos": 0,
},
},
}
],
"externalAccountIdentifiers": {
"obfuscatedExternalAccountId": store_account_token(self.user.pk),
},
}
body.update(overrides)
return body
@patch("monetization.services.store_purchases.fetch_google_subscription")
def test_play_token_grants_plan(self, mock_fetch):
mock_fetch.return_value = self._body()
sub = confirm_store_purchase(
self.user,
{
"store": "play_store",
"product_id": _PRODUCT,
"purchase_token": "token-abc",
"transaction_id": "GPA.1234-5678",
},
)
self.assertEqual(sub.source, UserSubscription.Source.PLAY_STORE)
self.assertEqual(sub.status, UserSubscription.Status.ACTIVE)
self.assertEqual(sub.revenuecat_original_transaction_id, "token-abc")
self.assertFalse(sub.cancel_at_period_end)
invoice = Invoice.objects.get(user=self.user)
self.assertEqual(invoice.provider, Invoice.Provider.PLAY_STORE)
self.assertEqual(invoice.amount_paid, 1000)
@patch("monetization.services.store_purchases.fetch_google_subscription")
def test_canceled_but_unexpired_keeps_access(self, mock_fetch):
mock_fetch.return_value = self._body(
subscriptionState="SUBSCRIPTION_STATE_CANCELED"
)
sub = confirm_store_purchase(
self.user,
{
"store": "play_store",
"product_id": _PRODUCT,
"purchase_token": "token-cancel",
},
)
self.assertEqual(sub.status, UserSubscription.Status.ACTIVE)
self.assertTrue(sub.cancel_at_period_end)
@patch("monetization.services.store_purchases.fetch_google_subscription")
def test_token_already_linked_to_someone_else(self, mock_fetch):
other = make_user(email="other@example.com")
UserSubscription.objects.create(
user=other,
plan=SubscriptionPlan.objects.get(slug="founders"),
status=UserSubscription.Status.ACTIVE,
source=UserSubscription.Source.PLAY_STORE,
revenuecat_original_transaction_id="token-taken",
)
mock_fetch.return_value = self._body()
with self.assertRaises(StorePurchaseRejected):
confirm_store_purchase(
self.user,
{
"store": "play_store",
"product_id": _PRODUCT,
"purchase_token": "token-taken",
},
)
def test_missing_play_credentials(self):
with override_settings(GOOGLE_PLAY_SERVICE_ACCOUNT_JSON=""):
with self.assertRaises(StoreNotConfiguredError):
confirm_store_purchase(
self.user,
{
"store": "play_store",
"product_id": _PRODUCT,
"purchase_token": "token-abc",
},
)
class StoreConfirmApiTests(APITestCase):
def setUp(self):
seed_subscription_plans(update_existing=False)
self.user = make_user(email="api@example.com")
self.client.force_authenticate(user=self.user)
self.url = reverse("finance_store_confirm")
def test_requires_auth(self):
self.client.force_authenticate(user=None)
response = self.client.post(self.url, {}, format="json")
self.assertEqual(response.status_code, status.HTTP_401_UNAUTHORIZED)
@patch("monetization.views.confirm_store_purchase")
def test_rejected_purchase_is_400(self, mock_confirm):
mock_confirm.side_effect = StorePurchaseRejected("nope")
response = self.client.post(
self.url,
{
"store": "play_store",
"product_id": _PRODUCT,
"purchase_token": "token",
},
format="json",
)
self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
def test_lapsed_store_period_ends_on_read(self):
plan = SubscriptionPlan.objects.get(slug="founders")
UserSubscription.objects.create(
user=self.user,
plan=plan,
status=UserSubscription.Status.ACTIVE,
source=UserSubscription.Source.APP_STORE,
current_period_end=timezone.now() - timedelta(days=1),
)
response = self.client.get(reverse("finance_subscription_me"))
self.assertEqual(response.status_code, status.HTTP_200_OK)
self.assertEqual(response.data["status"], "canceled")
self.assertTrue(response.data["needs_checkout"])
self.assertEqual(
response.data["store_account_token"],
store_account_token(self.user.pk),
)
+6 -6
View File
@@ -1,12 +1,12 @@
from django.urls import path from django.urls import path
from monetization.views import ( from monetization.views import (
ConfirmStorePurchaseView,
CreateBillingPortalSessionView, CreateBillingPortalSessionView,
CreateCheckoutSessionView, CreateCheckoutSessionView,
InvoiceListView, InvoiceListView,
PaymentListView, PaymentListView,
PlanListView, PlanListView,
RevenueCatWebhookView,
StripeWebhookView, StripeWebhookView,
SubscriptionMeView, SubscriptionMeView,
) )
@@ -42,14 +42,14 @@ urlpatterns = [
SubscriptionMeView.as_view(), SubscriptionMeView.as_view(),
name="finance_subscription_me", name="finance_subscription_me",
), ),
path(
"store/confirm/",
ConfirmStorePurchaseView.as_view(),
name="finance_store_confirm",
),
path( path(
"webhooks/stripe/", "webhooks/stripe/",
StripeWebhookView.as_view(), StripeWebhookView.as_view(),
name="finance_stripe_webhook", name="finance_stripe_webhook",
), ),
path(
"webhooks/revenuecat/",
RevenueCatWebhookView.as_view(),
name="finance_revenuecat_webhook",
),
] ]
+50 -51
View File
@@ -12,6 +12,7 @@ from monetization.serializers import (
InvoiceSerializer, InvoiceSerializer,
PaymentSerializer, PaymentSerializer,
PortalSessionSerializer, PortalSessionSerializer,
StorePurchaseConfirmSerializer,
SubscriptionPlanSerializer, SubscriptionPlanSerializer,
) )
from monetization.services.plans import ( from monetization.services.plans import (
@@ -20,10 +21,13 @@ from monetization.services.plans import (
seed_subscription_plans, seed_subscription_plans,
) )
from monetization.services.quotas import get_usage_snapshot from monetization.services.quotas import get_usage_snapshot
from monetization.services.revenuecat import ( from monetization.services.store_purchases import (
RevenueCatWebhookAuthError, StoreNotConfiguredError,
dispatch_revenuecat_event, StorePurchaseRejected,
verify_revenuecat_authorization, StoreVerificationError,
confirm_store_purchase,
expire_lapsed_store_subscription,
store_account_token,
) )
from monetization.services.stripe import ( from monetization.services.stripe import (
StripeNotConfiguredError, StripeNotConfiguredError,
@@ -194,6 +198,7 @@ class SubscriptionMeView(APIView):
sub = request.user.subscription sub = request.user.subscription
except UserSubscription.DoesNotExist: except UserSubscription.DoesNotExist:
sub = None sub = None
sub = expire_lapsed_store_subscription(sub)
usage = get_usage_snapshot(request.user) usage = get_usage_snapshot(request.user)
payload = { payload = {
@@ -213,11 +218,52 @@ class SubscriptionMeView(APIView):
if sub and sub.current_period_end if sub and sub.current_period_end
else None else None
), ),
"store_account_token": store_account_token(request.user.pk),
"usage": usage.to_dict(), "usage": usage.to_dict(),
} }
return Response(payload) return Response(payload)
class ConfirmStorePurchaseView(APIView):
"""Verify a StoreKit or Play Billing purchase and grant the plan.
The app calls this after the native store sheet succeeds. There is no
store webhook.
"""
def post(self, request):
serializer = StorePurchaseConfirmSerializer(data=request.data)
serializer.is_valid(raise_exception=True)
try:
sub = confirm_store_purchase(request.user, serializer.validated_data)
except StorePurchaseRejected as exc:
return Response({"detail": str(exc)}, status=status.HTTP_400_BAD_REQUEST)
except StoreNotConfiguredError as exc:
return Response(
{"detail": str(exc)},
status=status.HTTP_503_SERVICE_UNAVAILABLE,
)
except StoreVerificationError as exc:
logger.exception("Store purchase verification failed")
return Response({"detail": str(exc)}, status=status.HTTP_502_BAD_GATEWAY)
return Response(
{
"source": sub.source,
"status": sub.status,
"plan_slug": sub.plan.slug if sub.plan_id else None,
"cancel_at_period_end": bool(sub.cancel_at_period_end),
"current_period_end": (
sub.current_period_end.isoformat()
if sub.current_period_end
else None
),
"store_account_token": store_account_token(request.user.pk),
},
status=status.HTTP_200_OK,
)
class StripeWebhookView(APIView): class StripeWebhookView(APIView):
"""Verify Stripe signatures and upsert local invoice/payment rows.""" """Verify Stripe signatures and upsert local invoice/payment rows."""
@@ -266,50 +312,3 @@ class StripeWebhookView(APIView):
) )
return Response({"received": True}, status=status.HTTP_200_OK) return Response({"received": True}, status=status.HTTP_200_OK)
class RevenueCatWebhookView(APIView):
"""Verify RevenueCat Authorization and upsert subscription + ledger rows."""
permission_classes = (permissions.AllowAny,)
authentication_classes = ()
def post(self, request):
webhook_secret = settings.REVENUECAT_WEBHOOK_SECRET
if not webhook_secret:
logger.error("REVENUECAT_WEBHOOK_SECRET is not configured")
return Response(
{"detail": "Webhook secret not configured"},
status=status.HTTP_503_SERVICE_UNAVAILABLE,
)
try:
verify_revenuecat_authorization(
authorization_header=request.META.get("HTTP_AUTHORIZATION"),
expected_secret=webhook_secret,
)
except RevenueCatWebhookAuthError as exc:
return Response(
{"detail": str(exc)},
status=status.HTTP_401_UNAUTHORIZED,
)
payload = request.data
if not isinstance(payload, dict):
return Response(
{"detail": "Invalid payload"},
status=status.HTTP_400_BAD_REQUEST,
)
try:
dispatch_revenuecat_event(payload)
except Exception:
event = payload.get("event") if isinstance(payload, dict) else {}
event_id = event.get("id") if isinstance(event, dict) else None
logger.exception("Error handling RevenueCat event %s", event_id)
return Response(
{"detail": "Webhook handler error"},
status=status.HTTP_500_INTERNAL_SERVER_ERROR,
)
return Response({"received": True}, status=status.HTTP_200_OK)
Generated
+3 -3
View File
@@ -497,11 +497,11 @@ wheels = [
[[package]] [[package]]
name = "certifi" name = "certifi"
version = "2026.7.22" version = "2026.10.23"
source = { registry = "https://pypi.org/simple" } source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/a3/c2/24167ea9858356b47a87a50d39908bfdb72ceeefe0041586e704e5376b3a/certifi-2026.7.22.tar.gz", hash = "sha256:741e2c3b351ddf169a738da9f2c048608ff7f2c5cc02f1ebc6b118bb090d5d55", size = 138112, upload-time = "2026-07-22T03:35:12.644Z" } sdist = { url = "https://files.pythonhosted.org/packages/fd/11/b524cbd9bbdaca58608e4cf4053fd463417549e1959fa8d3aa18319f9555/certifi-2026.10.23.tar.gz", hash = "sha256:a9895e6e2758b371e0545e321a3a8fd1fcfd9829ada8824a936dfb16aaaf36cd", size = 136225, upload-time = "2026-10-23T11:36:16.478Z" }
wheels = [ wheels = [
{ url = "https://files.pythonhosted.org/packages/0b/a7/71ac2cff56fec219ed242bb11b8efb69fcc4bec75db06fb7bfe35de520e6/certifi-2026.7.22-py3-none-any.whl", hash = "sha256:62f22742b58a1a33014a2b6b706588a8d7e2a88ae7bd1a6ebe8c992928483775", size = 136983, upload-time = "2026-07-22T03:35:11.276Z" }, { url = "https://files.pythonhosted.org/packages/32/d1/a3652fe2c19134d0d15922177821d863cbd7e842d53a21bbe7e8e3a1b5ac/certifi-2026.10.23-py3-none-any.whl", hash = "sha256:e9b54de76e530bb6cf0ccd036787103acb58fd42ba930ea8bf948ddbaa2f61d1", size = 135082, upload-time = "2026-10-23T11:36:14.763Z" },
] ]
[[package]] [[package]]