ai_ml_operations/chat_backend
2
0
Fork 0
Code Issues 23 Pull Requests Actions Packages Projects Releases Wiki Activity

Compare commits

base: ai_ml_operations/chat_backend:master
Branches Tags
ai_ml_operations/chat_backend:master
..
compare: ai_ml_operations/chat_backend:be6ff471ad45416a03090d3e8274d3cc342fa8b3
Branches Tags
ai_ml_operations/chat_backend:master
1 Commits
master .. be6ff471ad
Author SHA1 Message Date
westfarn be6ff471ad Add Google/Microsoft SSO OAuth for register and sign-in (#24)
CI / test (pull_request) Failing after 7s
Details
Unit Tests / test (pull_request) Failing after 6s
Details 
Introduce OAuthIdentity storage, start/callback endpoints, JWT handoff to
the SPA, and mocked IdP tests so Drive OAuth (#11) can reuse the same model.
 2026-07-27 07:09:02 -05:00
145 changed files with 905 additions and 16419 deletions
Expand all files Collapse all files
.env.example
+6 -75
View File
@@ -12,22 +12,8 @@ DATABASE_URL=postgres://chat_backend:chat_backend@db:5432/chat_backend
# Ollama — local loopback when Ollama runs on this machine; LAN IP for GPU host.
OLLAMA_BASE_URL=http://127.0.0.1:11434
# Legacy fallback (used when role-specific vars unset). Prefer the role vars.
# OLLAMA_MODEL=gpt-oss:20b
# OLLAMA_MODEL_THINKING=gpt-oss:20b
# OLLAMA_MODEL_FAST=gemma4:latest
# OLLAMA_MODEL_UTILITY=llama3.2
# OLLAMA_EMBED_MODEL=nomic-embed-text
# OLLAMA_NUM_CTX_THINKING=16384
# OLLAMA_NUM_CTX_FAST=8192
# OLLAMA_NUM_CTX_UTILITY=4096
# Web search (#62) — SearxNG primary, DDGS failover. See README "SearxNG".
ALLOW_INTERNET_ACCESS=true
SEARCH_PROVIDER=searxng
SEARCH_FAILOVER_PROVIDER=ddgs
SEARXNG_BASE_URL=http://127.0.0.1:8088
# SEARXNG_TIMEOUT_SECONDS=8
# OLLAMA_MODEL=llama3.2
# OLLAMA_EMBED_MODEL=llama3.2
# Email (SMTP2GO) — optional for local
EMAIL_HOST=mail.smtp2go.com
@@ -46,12 +32,6 @@ ENABLE_ACCOUNT_REGISTRATION=false
# Redirect URIs (register in each IdP console):
# {OAUTH_CALLBACK_BASE_URL}/api/auth/oauth/google/callback/
# {OAUTH_CALLBACK_BASE_URL}/api/auth/oauth/microsoft/callback/
# Same client id/secret pair is reused for Drive linking (#47) — the extra
# Drive scopes below are requested incrementally via intent=link_drive /
# intent=link_company_drive, no separate app registration needed:
# Google: openid email profile https://www.googleapis.com/auth/drive.readonly
# Microsoft: openid email profile offline_access Files.Read (personal)
# openid email profile offline_access Files.Read.All Sites.Read.All (company)
GOOGLE_OAUTH_CLIENT_ID=
GOOGLE_OAUTH_CLIENT_SECRET=
MICROSOFT_OAUTH_CLIENT_ID=
@@ -60,69 +40,20 @@ MICROSOFT_OAUTH_TENANT=common
# Optional; defaults to request host. Example local: http://127.0.0.1:8001
OAUTH_CALLBACK_BASE_URL=http://127.0.0.1:8001
# Drive / RAG sync (#47-#53). Requires a subscription plan with allows_rag
# (Founders, Pro, Business, Backer by default — see finance PLAN_SEED).
# Start a link: GET /api/auth/oauth/google/start/?intent=link_drive (authenticated)
# GET /api/auth/oauth/google/start/?intent=link_company_drive (company manager)
# GET /api/auth/oauth/microsoft/start/?intent=link_drive
# GET /api/auth/oauth/microsoft/start/?intent=link_company_drive
# Manage: GET /api/drive/connections/
# DELETE /api/drive/connections/<id>/
# POST /api/drive/connections/<id>/resources/ { "resource_ids": [...] }
# POST /api/drive/connections/<id>/sync/
# Provider push notifications (best-effort; register with each provider's
# subscription/watch API pointing here, using ?connection_id=<id>):
# POST {OAUTH_CALLBACK_BASE_URL}/api/drive/webhooks/google/
# POST {OAUTH_CALLBACK_BASE_URL}/api/drive/webhooks/microsoft/
# Worker sync: `python manage.py sync_drive_connections [--connection-id N]`
# Stripe / monetization (optional local — required for checkout + webhooks)
# Stripe / finance (optional local — required for checkout + webhooks)
STRIPE_SECRET_KEY=
STRIPE_PUBLISHABLE_KEY=
STRIPE_WEBHOOK_SECRET=
# Optional: pre-created Stripe Price ID for Founders. When empty, Checkout uses
# SubscriptionPlan.price_cents / SUBSCRIPTION_PRICE_* ($10 USD / month Founders).
# Optional: pre-created Stripe Price ID. When empty, Checkout uses
# SUBSCRIPTION_PRICE_* from settings.py ($10 USD / month by default).
STRIPE_PRICE_ID=
# RevenueCat webhook Authorization bearer secret (store IAP).
REVENUECAT_WEBHOOK_SECRET=
# Optional JSON map of store product id → plan slug, e.g.
# REVENUECAT_PRODUCT_PLAN_MAP={"hesychia_founders_monthly":"founders"}
# SUBSCRIPTION_PRICE_AMOUNT_CENTS=1000
# SUBSCRIPTION_PRICE_CURRENCY=usd
# SUBSCRIPTION_PRICE_INTERVAL=month
# SUBSCRIPTION_PRODUCT_NAME=Founders
# Enforce plan feature + prompt/token quotas on chat turns (default true).
# ENFORCE_SUBSCRIPTION_GATES=true
# SUBSCRIPTION_PRODUCT_NAME=Chat Subscription
FRONTEND_BASE_URL=http://localhost:3000
# Agentic task execution (#63) — long-running, multi-step, tool-using turns.
# Default false: chat behaves exactly like the always-on grounded path (#62),
# no planner/tools/AgentRun rows. Requires a plan with allows_rag (or
# allows_all_future_features) — see monetization SubscriptionPlan.allows_feature.
ALLOW_AGENTIC_TASKS=false
# Redis — optional. Unset = InMemory channel layer (single process, fine for
# dev/tests) and agent work runs on a daemon thread instead of Celery.
# REDIS_URL=redis://127.0.0.1:6379/0
# CELERY_BROKER_URL=redis://127.0.0.1:6379/0
# Orchestrator plans + synthesises; sub-agents run independent plan steps
# concurrently on a smaller/cheaper model.
# OLLAMA_MODEL_ORCHESTRATOR=gpt-oss:20b
# OLLAMA_MODEL_SUBAGENT=llama3.2
# AGENT_MAX_PLAN_STEPS=8
# AGENT_MAX_ITERATIONS=12
# AGENT_WALL_CLOCK_SECONDS=600
# AGENT_SUBAGENT_CONCURRENCY=3
# AGENT_TOOL_TIMEOUT_SECONDS=20
# AGENT_TOOL_OUTPUT_MAX_CHARS=8000
# AGENT_MAX_TOOL_CALLS_PER_RUN=40
# AGENT_FETCH_URL_MAX_BYTES=2097152
# Run a worker once REDIS_URL/CELERY_BROKER_URL are set:
# docker compose --profile agentic up redis worker
# uv run celery -A llm_be worker --loglevel=info
# STRIPE_CHECKOUT_SUCCESS_URL=http://localhost:3000/billing/success?session_id={CHECKOUT_SESSION_ID}
# STRIPE_CHECKOUT_CANCEL_URL=http://localhost:3000/billing/cancel
# Customer Portal return URL (plan change / cancel / payment method).
# STRIPE_PORTAL_RETURN_URL=http://localhost:3000/account/
# Gunicorn / ASGI
GUNICORN_WORKERS=2
.env.prod.example
+3 -65
View File
@@ -32,22 +32,8 @@ WEB_PORT=8003
# Ollama on GPU host (ai-server-4080). Firewall must allow 10.0.0.0/24 → :11434.
OLLAMA_BASE_URL=http://10.0.0.128:11434
# Role-scoped models (#62). After changing OLLAMA_EMBED_MODEL, run:
# python manage.py reindex_embeddings
OLLAMA_MODEL=gpt-oss:20b
OLLAMA_MODEL_THINKING=gpt-oss:20b
OLLAMA_MODEL_FAST=gemma4:latest
OLLAMA_MODEL_UTILITY=llama3.2
OLLAMA_EMBED_MODEL=nomic-embed-text
OLLAMA_NUM_CTX_THINKING=16384
OLLAMA_NUM_CTX_FAST=8192
# Web search (#62) — self-hosted SearxNG (recommended). DDGS is automatic failover.
ALLOW_INTERNET_ACCESS=true
SEARCH_PROVIDER=searxng
SEARCH_FAILOVER_PROVIDER=ddgs
# Point at the SearxNG container/service on the LAN (see README "SearxNG").
SEARXNG_BASE_URL=http://10.0.0.128:8088
OLLAMA_MODEL=llama3.2
OLLAMA_EMBED_MODEL=llama3.2
# Email (SMTP2GO)
EMAIL_HOST=mail.smtp2go.com
@@ -67,12 +53,6 @@ ENABLE_ACCOUNT_REGISTRATION=false
# Register redirect URIs:
# https://chatbackend.aimloperations.com/api/auth/oauth/google/callback/
# https://chatbackend.aimloperations.com/api/auth/oauth/microsoft/callback/
# Same client id/secret pair covers Drive linking (#47); no extra IdP app
# registration needed, but do register the Drive/Graph API + consent screen
# scopes below in each console (incremental scopes requested at intent time):
# Google: openid email profile https://www.googleapis.com/auth/drive.readonly
# Microsoft: openid email profile offline_access Files.Read (personal)
# openid email profile offline_access Files.Read.All Sites.Read.All (company)
GOOGLE_OAUTH_CLIENT_ID=
GOOGLE_OAUTH_CLIENT_SECRET=
MICROSOFT_OAUTH_CLIENT_ID=
@@ -80,41 +60,16 @@ MICROSOFT_OAUTH_CLIENT_SECRET=
MICROSOFT_OAUTH_TENANT=common
OAUTH_CALLBACK_BASE_URL=https://chatbackend.aimloperations.com
# Drive / RAG sync (#47-#53) — gated by SubscriptionPlan.allows_rag.
# Register provider push notifications (Google Drive `watch`, Microsoft
# Graph subscriptions) against:
# https://chatbackend.aimloperations.com/api/drive/webhooks/google/
# https://chatbackend.aimloperations.com/api/drive/webhooks/microsoft/
# Scheduled sync (cron / server-infra job): `python manage.py sync_drive_connections`
# Stripe / monetization
# Stripe / finance
STRIPE_SECRET_KEY=replace-with-stripe-secret-key
STRIPE_PUBLISHABLE_KEY=replace-with-stripe-publishable-key
STRIPE_WEBHOOK_SECRET=replace-with-stripe-webhook-secret
REVENUECAT_WEBHOOK_SECRET=replace-with-revenuecat-webhook-auth-token
# REVENUECAT_PRODUCT_PLAN_MAP={"hesychia_founders_monthly":"founders"}
# Optional: pre-created Stripe Price ID. When empty, Checkout uses
# SUBSCRIPTION_PRICE_* from settings.py ($10 USD / month by default).
STRIPE_PRICE_ID=
FRONTEND_BASE_URL=https://chat.aimloperations.com
# STRIPE_CHECKOUT_SUCCESS_URL=https://chat.aimloperations.com/billing/success?session_id={CHECKOUT_SESSION_ID}
# STRIPE_CHECKOUT_CANCEL_URL=https://chat.aimloperations.com/billing/cancel
# STRIPE_PORTAL_RETURN_URL=https://chat.aimloperations.com/account/
# Agentic task execution (#63). Keep false until Redis/Celery worker + Ollama
# capacity are confirmed on this host; false = identical behavior to #62.
ALLOW_AGENTIC_TASKS=false
# Shared Redis (channel layer fan-out across gunicorn/uvicorn workers +
# Celery broker for agent runs). Point both at the same instance.
# REDIS_URL=redis://10.0.0.128:6379/0
# CELERY_BROKER_URL=redis://10.0.0.128:6379/0
# OLLAMA_MODEL_ORCHESTRATOR=gpt-oss:20b
# OLLAMA_MODEL_SUBAGENT=llama3.2
# AGENT_MAX_PLAN_STEPS=8
# AGENT_MAX_ITERATIONS=12
# AGENT_WALL_CLOCK_SECONDS=600
# AGENT_SUBAGENT_CONCURRENCY=3
# Start the worker (server-infra): docker compose --profile agentic up -d worker
# Gunicorn / ASGI (UvicornWorker for WebSockets)
GUNICORN_WORKERS=2
@@ -122,27 +77,10 @@ GUNICORN_BIND=0.0.0.0:8000
# =============================================================================
# BETA overrides (use separate file: chat_backend_beta.env)
# Control node: ~/Documents/secrets/chat_backend/chat_backend_beta.env
# Infra: server-infra#7 (host_apps beta :8013, Postgres chat_backend_beta, NPM)
# =============================================================================
# DJANGO_ENV=beta
# DJANGO_DEBUG=false
# DJANGO_SECRET_KEY=replace-with-a-different-beta-secret
# DJANGO_ALLOWED_HOSTS=beta.chatbackend.aimloperations.com
# Optional; when unset, https:// origins are derived from DJANGO_ALLOWED_HOSTS.
# DJANGO_CSRF_TRUSTED_ORIGINS=https://beta.chatbackend.aimloperations.com,https://beta.chat.aimloperations.com
# CORS_ALLOWED_ORIGINS=https://beta.chat.aimloperations.com
# CORS_ORIGIN_ALLOW_ALL=false
# USE_TLS_PROXY=true
# DATABASE_URL=postgres://westfarn:replace-db-password@10.0.0.230:5432/chat_backend_beta
# WEB_PORT=8013
# OLLAMA_BASE_URL=http://10.0.0.128:11434
# OAUTH_CALLBACK_BASE_URL=https://beta.chatbackend.aimloperations.com
# FRONTEND_BASE_URL=https://beta.chat.aimloperations.com
# Register beta OAuth redirect URIs in each IdP console:
# https://beta.chatbackend.aimloperations.com/api/auth/oauth/google/callback/
# https://beta.chatbackend.aimloperations.com/api/auth/oauth/microsoft/callback/
#
# Tianji: backend does not load tracker.js. Beta SPA uses a distinct website ID
# (chat_web_app .env.beta REACT_APP_TIANJI_WEBSITE_ID). Ensure CORS allows the
# beta frontend origin so Tianji-instrumented pages can call this API.
.gitea/workflows/deploy-beta.yml
-82
View File
@@ -1,82 +0,0 @@
name: Deploy Beta
# Auto-deploy beta after push to master (mirrors dta_service / chat_web_app).
# Prod is manual via Deploy Prod (workflow_dispatch).
on:
push:
branches:
- master
jobs:
unit-tests:
runs-on: self-hosted
steps:
- name: Checkout
uses: actions/checkout@v4
- name: Install uv
run: |
curl -LsSf https://astral.sh/uv/install.sh | sh
echo "$HOME/.local/bin" >> "$GITHUB_PATH"
- name: Install dependencies
run: uv sync --frozen
- name: Run unit tests
env:
DJANGO_ENV: dev
DJANGO_SECRET_KEY: test-secret-key
DJANGO_DEBUG: "true"
DJANGO_ALLOWED_HOSTS: localhost,127.0.0.1,testserver
DATABASE_URL: ""
DB_HOST: ""
SKIP_RAG_INIT: "1"
OLLAMA_BASE_URL: http://127.0.0.1:11434
working-directory: llm_be
run: uv run python manage.py test
docker:
needs: unit-tests
runs-on: self-hosted
steps:
- name: Checkout
uses: actions/checkout@v4
- name: Build Docker image
run: docker compose build
# Ephemeral local Postgres only — never inherit host DATABASE_URL (prod/beta).
- name: Run containerized tests
run: |
set -euo pipefail
unset DATABASE_URL DB_HOST DB_NAME DB_USER DB_PASSWORD DB_PORT \
COMPOSE_DATABASE_URL DJANGO_ENV DJANGO_SECRET_KEY DJANGO_DEBUG \
DJANGO_ALLOWED_HOSTS OLLAMA_BASE_URL || true
PROJECT="chat-backend-ci-${{ gitea.sha }}"
cleanup() { docker compose -p "$PROJECT" down -v --remove-orphans || true; }
trap cleanup EXIT
docker compose -p "$PROJECT" up -d --wait db
docker compose -p "$PROJECT" run --rm --no-deps --entrypoint "" \
-e DJANGO_ENV=dev \
-e DJANGO_SECRET_KEY=test-secret-key \
-e DJANGO_DEBUG=true \
-e DJANGO_ALLOWED_HOSTS=localhost,127.0.0.1,testserver \
-e DATABASE_URL=postgres://chat_backend:chat_backend@db:5432/chat_backend \
-e SKIP_RAG_INIT=1 \
-e OLLAMA_BASE_URL=http://127.0.0.1:11434 \
web uv run python manage.py test
deploy-beta:
needs: docker
runs-on: self-hosted
env:
SERVER_INFRA_ROOT: /home/westfarn/Documents/repos/server-infra
steps:
- name: Deploy chat_backend beta
run: |
"$SERVER_INFRA_ROOT/scripts/deploy.sh" \
--app chat_backend \
--env beta \
--ref "${{ gitea.sha }}"
.gitea/workflows/deploy-prod.yml → .gitea/workflows/deploy.yml
+15 -37
View File
@@ -1,49 +1,26 @@
name: Deploy Prod
name: Deploy Chat Backend
# Manual prod deploy only (mirrors dta_service / chat_web_app).
# Push to master deploys beta via Deploy Beta.
# Runs after Unit Tests completes on master. Direct pushes only (not PRs).
on:
workflow_dispatch: {}
workflow_run:
workflows: [Unit Tests]
types: [completed]
branches: [master]
jobs:
unit-tests:
runs-on: self-hosted
steps:
- name: Checkout
uses: actions/checkout@v4
- name: Install uv
run: |
curl -LsSf https://astral.sh/uv/install.sh | sh
echo "$HOME/.local/bin" >> "$GITHUB_PATH"
- name: Install dependencies
run: uv sync --frozen
- name: Run unit tests
env:
DJANGO_ENV: dev
DJANGO_SECRET_KEY: test-secret-key
DJANGO_DEBUG: "true"
DJANGO_ALLOWED_HOSTS: localhost,127.0.0.1,testserver
DATABASE_URL: ""
DB_HOST: ""
SKIP_RAG_INIT: "1"
OLLAMA_BASE_URL: http://127.0.0.1:11434
working-directory: llm_be
run: uv run python manage.py test
docker:
needs: unit-tests
if: gitea.event.workflow_run.conclusion == 'success' && gitea.event.workflow_run.event == 'push'
runs-on: self-hosted
steps:
- name: Checkout
uses: actions/checkout@v4
with:
ref: ${{ gitea.event.workflow_run.head_sha }}
- name: Build Docker image
run: docker compose build
# Ephemeral local Postgres only — never inherit host DATABASE_URL (prod/beta).
# Ephemeral local Postgres only — never inherit host DATABASE_URL (prod).
- name: Run containerized tests
run: |
set -euo pipefail
@@ -51,7 +28,7 @@ jobs:
COMPOSE_DATABASE_URL DJANGO_ENV DJANGO_SECRET_KEY DJANGO_DEBUG \
DJANGO_ALLOWED_HOSTS OLLAMA_BASE_URL || true
PROJECT="chat-backend-ci-${{ gitea.sha }}"
PROJECT="chat-backend-ci-${{ gitea.event.workflow_run.head_sha }}"
cleanup() { docker compose -p "$PROJECT" down -v --remove-orphans || true; }
trap cleanup EXIT
@@ -66,9 +43,10 @@ jobs:
-e OLLAMA_BASE_URL=http://127.0.0.1:11434 \
web uv run python manage.py test
deploy-prod:
needs: docker
deploy:
if: gitea.event.workflow_run.conclusion == 'success' && gitea.event.workflow_run.event == 'push'
runs-on: self-hosted
needs: docker
env:
SERVER_INFRA_ROOT: /home/westfarn/Documents/repos/server-infra
steps:
@@ -77,4 +55,4 @@ jobs:
"$SERVER_INFRA_ROOT/scripts/deploy.sh" \
--app chat_backend \
--env prod \
--ref "${{ gitea.sha }}"
--ref "${{ gitea.event.workflow_run.head_sha }}"
.gitea/workflows/run-evals.yml
-60
View File
@@ -1,60 +0,0 @@
name: Run Evals
# Manual eval harness — does not gate PRs (#62 Phase 4).
# Requires self-hosted runner with GPU/Ollama and live search when SKIP_LIVE is unset.
on:
workflow_dispatch: {}
jobs:
run-evals:
runs-on: self-hosted
steps:
- name: Checkout
uses: actions/checkout@v4
- name: Install uv
run: |
curl -LsSf https://astral.sh/uv/install.sh | sh
echo "$HOME/.local/bin" >> "$GITHUB_PATH"
- name: Install dependencies
run: uv sync --frozen
- name: Validate eval suite (offline)
env:
DJANGO_ENV: dev
DJANGO_SECRET_KEY: test-secret-key
DJANGO_DEBUG: "true"
DJANGO_ALLOWED_HOSTS: localhost,127.0.0.1,testserver
DATABASE_URL: ""
DB_HOST: ""
SKIP_RAG_INIT: "1"
SKIP_LIVE: "1"
OLLAMA_BASE_URL: http://127.0.0.1:11434
working-directory: llm_be
run: uv run python manage.py run_evals --dry-run
- name: Run eval harness
env:
DJANGO_ENV: dev
DJANGO_SECRET_KEY: test-secret-key
DJANGO_DEBUG: "true"
DJANGO_ALLOWED_HOSTS: localhost,127.0.0.1,testserver
DATABASE_URL: ""
DB_HOST: ""
SKIP_RAG_INIT: "1"
RUN_EVALS: "1"
ALLOW_INTERNET_ACCESS: "true"
OLLAMA_BASE_URL: http://127.0.0.1:11434
working-directory: llm_be
run: |
uv run python manage.py run_evals \
--runs 3 \
--output "../eval-report.json"
- name: Upload eval report
if: always()
uses: actions/upload-artifact@v4
with:
name: eval-report
path: eval-report.json
README.md
+10 -189
View File
@@ -87,26 +87,14 @@ with `COMPOSE_DATABASE_URL` if needed.
| `DATABASE_URL` | SQLite fallback | yes | Shared Postgres in prod |
| `WEB_PORT` | n/a (compose maps 8003) | `8003` | Host port for prod compose |
| `OLLAMA_BASE_URL` | `http://127.0.0.1:11434` | yes | GPU host in prod: `http://10.0.0.128:11434` |
| `OLLAMA_MODEL` | `gpt-oss:20b` | optional | Legacy fallback for THINKING |
| `OLLAMA_MODEL_THINKING` / `_FAST` / `_UTILITY` | see defaults | optional | Role-scoped chat models (#62) |
| `OLLAMA_EMBED_MODEL` | `nomic-embed-text` | optional | Never falls back to a chat model |
| `OLLAMA_NUM_CTX_THINKING` / `_FAST` | `16384` / `8192` | optional | Context window per role |
| `ALLOW_INTERNET_ACCESS` | `true` | optional | Gate for live web retrieval |
| `SEARCH_PROVIDER` | `searxng` | optional | Primary search provider (#62) |
| `SEARCH_FAILOVER_PROVIDER` | `ddgs` | optional | Automatic failover |
| `SEARXNG_BASE_URL` | `http://127.0.0.1:8088` | yes if using SearxNG | Self-hosted SearxNG JSON API |
| `OLLAMA_MODEL` / `OLLAMA_EMBED_MODEL` | from `DEBUG` | optional | Override model names |
| `EMAIL_HOST_*` | empty | yes (prod/beta) | SMTP2GO |
| `CAPTCHA_SECRET_KEY` | empty | recommended | |
| `ENABLE_ACCOUNT_REGISTRATION` | `false` | optional | Self-serve sign-up; keep false until ready |
| `STRIPE_SECRET_KEY` / `STRIPE_PUBLISHABLE_KEY` / `STRIPE_WEBHOOK_SECRET` | empty | yes for billing | Stripe API + webhook (`monetization` app) |
| `REVENUECAT_WEBHOOK_SECRET` | empty | yes for store IAP | RevenueCat webhook Authorization bearer |
| `REVENUECAT_PRODUCT_PLAN_MAP` | empty JSON | optional | `{"product_id":"plan_slug"}` fallback map |
| `STRIPE_SECRET_KEY` / `STRIPE_PUBLISHABLE_KEY` / `STRIPE_WEBHOOK_SECRET` | empty | yes for billing | Stripe API + webhook |
| `STRIPE_PRICE_ID` | empty | optional | Pre-created Price; else `$10/mo` from settings |
| `GOOGLE_OAUTH_CLIENT_ID` / `..._SECRET` | empty | for SSO/Drive | Also used for Drive linking (#47), incremental scopes |
| `MICROSOFT_OAUTH_CLIENT_ID` / `..._SECRET` / `..._TENANT` | empty / `common` | for SSO/Drive | Also used for Drive linking (#47), incremental scopes |
| `FRONTEND_BASE_URL` | `http://localhost:3000` | set in prod/beta | Checkout success/cancel, portal return, OAuth return |
| `STRIPE_PORTAL_RETURN_URL` | `{FRONTEND}/account/` | optional | Stripe Customer Portal return URL |
| `CORS_ALLOWED_ORIGINS` | local + chat FE (+ beta FE default) | set in prod/beta | Frontend origin(s) |
| `FRONTEND_BASE_URL` | `http://localhost:3000` | set in prod | Checkout success/cancel base |
| `CORS_ALLOWED_ORIGINS` | local + chat FE | set in prod | Frontend origin |
| `USE_TLS_PROXY` | false (dev) | true behind NPM | Sets `SECURE_PROXY_SSL_HEADER` |
| `GUNICORN_WORKERS` / `GUNICORN_BIND` | 2 / `0.0.0.0:8000` | optional | Entrypoint |
| `SKIP_RAG_INIT` | unset | CI/migrate often `1` | Skip Chroma/Ollama boot work |
@@ -117,18 +105,16 @@ generation prompts (chat, RAG, data analysis). Not env-configurable.
Templates: `.env.example` (local), `.env.prod.example` (control-node secret).
Control-node secret paths (server-infra on ai-server-4080):
Control-node secret path (server-infra on ai-server-4080):
```text
~/Documents/secrets/chat_backend/chat_backend_prod.env
~/Documents/secrets/chat_backend/chat_backend_beta.env
```
Validate with:
```bash
./scripts/validate-env.sh ~/Documents/secrets/chat_backend/chat_backend_prod.env
./scripts/validate-env.sh ~/Documents/secrets/chat_backend/chat_backend_beta.env
```
If `DATABASE_URL` password contains `$`, escape each as `$$` for Compose.
@@ -145,74 +131,6 @@ All clients (`ollama.Client`, `OllamaLLM`, `OllamaEmbeddings`, `ChatOllama`) use
Firewall / Ollama listen on ai-server-4080 must allow `10.0.0.0/24` → `:11434`.
### Role-scoped models (#62)
| Role | Setting | Default | Used for |
|------|---------|---------|----------|
| THINKING | `OLLAMA_MODEL_THINKING` | `gpt-oss:20b` | Default chat / grounded answers |
| FAST | `OLLAMA_MODEL_FAST` | `gemma4:latest` | FE `modelName=FAST` (smaller/faster — still grounded) |
| UTILITY | `OLLAMA_MODEL_UTILITY` | `llama3.2` | Classify / moderate / title / grounding decision |
| EMBED | `OLLAMA_EMBED_MODEL` | `nomic-embed-text` | Chroma embeddings |
After changing `OLLAMA_EMBED_MODEL`, rebuild the vector store (dimension change):
```bash
SKIP_RAG_INIT=1 uv run python manage.py reindex_embeddings
```
### SearxNG (web search)
Grounded chat uses a self-hosted **SearxNG** instance as the primary search
provider (`SEARCH_PROVIDER=searxng`), with DuckDuckGo (`ddgs`) as automatic
failover. Point `SEARXNG_BASE_URL` at the JSON API (no trailing path).
**Recommended: run SearxNG on the GPU/infra host next to Ollama**
(`10.0.0.128`), reachable from the chat_backend containers on the LAN.
Minimal compose snippet (add to `server-infra` or run on ai-server-4080):
```yaml
services:
searxng:
image: searxng/searxng:latest
restart: unless-stopped
ports:
- "8080:8080"
volumes:
- ./searxng:/etc/searxng:rw
environment:
- SEARXNG_BASE_URL=http://10.0.0.128:8088/
```
In `searxng/settings.yml` (created on first start), enable the JSON format:
```yaml
search:
formats:
- html
- json
```
Then set in `chat_backend_prod.env` / `chat_backend_beta.env`:
```text
ALLOW_INTERNET_ACCESS=true
SEARCH_PROVIDER=searxng
SEARCH_FAILOVER_PROVIDER=ddgs
SEARXNG_BASE_URL=http://10.0.0.128:8088
```
Firewall: allow `10.0.0.0/24` → `:8080` on the SearxNG host (same pattern as
Ollama `:11434`). Verify from a backend container:
```bash
curl -sG 'http://10.0.0.128:8088/search' --data-urlencode 'q=test' -d 'format=json' | head
```
If SearxNG is down, chat still works for non-factual turns; factual turns that
require retrieval return an explicit "couldn't reach live sources" message
instead of hallucinating from parametric memory.
## File storage
Prompt attachments and workspace documents use **`DatabaseStorage`**
@@ -223,126 +141,29 @@ RAG loaders that need a path materialize a short-lived temp file, then delete it
Chroma’s vector index may still use a volume (`chroma_db`); that is embeddings
metadata, not the original upload.
## Production / beta (docker-compose.prod.yml)
## Production (docker-compose.prod.yml)
- Single `web` service; **no** bundled DB — `DATABASE_URL` → shared Postgres (`10.0.0.230`).
- Host port from `WEB_PORT` (prod **8003**; beta **8013**).
- Host port from `WEB_PORT` (catalog: **8003**; beta reserved **8013**).
- Entrypoint: wait DB → migrate → collectstatic → `gunicorn` + `UvicornWorker`
(ASGI for HTTP **and** WebSockets).
- Active/active on **adama + roslin + ai-server-4080**; NPM balances upstreams.
- Manual / local deploy:
- Deployed by:
```bash
# beta (day-to-day)
~/Documents/repos/server-infra/scripts/deploy.sh \
--app chat_backend --env beta --ref <sha>
# prod (intentional)
~/Documents/repos/server-infra/scripts/deploy.sh \
--app chat_backend --env prod --ref <sha>
```
Beta hosts / CORS: `beta.chatbackend.aimloperations.com` API +
`https://beta.chat.aimloperations.com` SPA (see `.env.prod.example` beta block).
DB: `chat_backend_beta`. Pair with [server-infra#7](https://git.aimloperations.com/ai_ml_operations/server-infra/issues/7)
and frontend [chat_web_app#35](https://git.aimloperations.com/ai_ml_operations/chat_web_app/issues/35).
## CI / CD (Gitea Actions)
| Workflow | Trigger | Action |
|----------|---------|--------|
| `unittests.yml` | push + PR → `master` | `uv sync` + `manage.py test` |
| `ci.yml` | PR → `master` | same unit tests |
| `deploy-beta.yml` | **push** to `master` | unit tests → docker compose tests → `deploy.sh --env beta` |
| `deploy-prod.yml` | **manual** `workflow_dispatch` only | unit tests → docker compose tests → `deploy.sh --env prod` |
| `deploy.yml` | after Unit Tests succeeds on `master` **push** | docker build + tests on **ephemeral compose Postgres** → `deploy.sh` |
Push/merge to `master` auto-deploys **beta** only. Prod requires the Gitea
**Run workflow** button on **Deploy Prod**. Deploy never runs on PRs.
## Frontend API notes
### Self-delete account ([#34](https://git.aimloperations.com/ai_ml_operations/chat_backend/issues/34))
| | |
|--|--|
| Method / path | `DELETE /api/user/` |
| Auth | JWT (authenticated user only; always deletes `request.user`) |
| Optional body | `{ "refresh_token": "<current refresh>" }` |
| Success | `200` `{ "detail": "Account deleted.", "deleted": true }` |
| Effects | Sets `deleted=True`, `is_active=False`; soft-deletes conversations; blacklists outstanding refresh tokens; logs `UserAuthEvent` `account_deleted` |
| Staff | Staff/superuser self-delete rejected (`400`, `code=staff_forbidden`) |
| Privacy v1 | Soft-delete only (no anonymization / hard purge) |
Post-delete UX: clear local tokens → redirect to sign-in. Subsequent
`/token/obtain/` fails. Do **not** send another user's id/email — ignored.
### Subscription change / cancel (portal + webhooks)
Billing lives in the **`monetization`** Django app (package rename of
`finance`; DB tables keep the `finance_*` prefix via app `label = "finance"`).
URLs: `/api/monetization/...` and alias `/api/finance/...`.
**Web:** plan change/cancel stay on Stripe Customer Portal
(`POST /api/finance/portal/`). Local state syncs via
`customer.subscription.updated` / `deleted` webhooks.
**Native (Play / App Store):** RevenueCat webhooks at
`POST /api/finance/webhooks/revenuecat/` (Authorization bearer =
`REVENUECAT_WEBHOOK_SECRET`) upsert Invoice/Payment ledger rows and sync
`UserSubscription` (`source=revenuecat`).
`GET /api/finance/subscription/` includes `cancel_at_period_end` and
`current_period_end` for Account UI messaging.
Subscription audit (`UserAuthEvent` on the user admin):
- `subscription_started` — first active plan (Checkout, Backer redeem, admin assign)
- `subscription_updated` — plan/status/cancel-at-period-end changes (portal + webhooks)
### Drive / RAG sync ([#47](https://git.aimloperations.com/ai_ml_operations/chat_backend/issues/47)-[#53](https://git.aimloperations.com/ai_ml_operations/chat_backend/issues/53))
Personal Google Drive / OneDrive and company Google Shared Drive / SharePoint
sync into the existing RAG `Document` pipeline. Personal connections work for
users **without** a company (personal workspace); company connections require a
company manager. Every endpoint below is gated
by `assert_feature_allowed(user, "rag")` (`SubscriptionPlan.allows_rag` —
true for Founders/Pro/Business/Backer, false for Standard by default).
**Connect (OAuth, reuses `#24` SSO app registrations with incremental scopes):**
| | |
|--|--|
| Personal | `GET /api/auth/oauth/<google\|microsoft>/start/?intent=link_drive` (authenticated) |
| Company | `GET /api/auth/oauth/<google\|microsoft>/start/?intent=link_company_drive` (company manager only) |
| Callback | Same `/api/auth/oauth/<provider>/callback/` as SSO; the signed OAuth `state` carries the linking `user_id` since the browser has no session on the IdP redirect. Upserts a `DriveConnection` and redirects to `{FRONTEND_BASE_URL}/account/?drive_connected=1&provider=<provider>&kind=<personal\|company>` (or `?error=<code>`) |
**Manage:**
| Method / path | Notes |
|--|--|
| `GET /api/drive/connections/` | Caller's personal connections + their company's company connections |
| `DELETE /api/drive/connections/<id>/` | Disconnect (owner for personal, company manager for company) — deactivates + clears tokens, keeps history |
| `POST /api/drive/connections/<id>/resources/` | `{ "resource_ids": [...], "resource_labels": [...] }` — folder/shared-drive/site ids to sync; empty = provider root |
| `POST /api/drive/connections/<id>/sync/` | Enqueue sync now (`chat_backend/drive_tasks.py`) — returns **202** with `connection.last_sync_status=pending`; poll `GET /api/drive/connections/` for `ok` / `error` + `last_sync_error` |
**Provider scope differences:**
- Google: same `drive.readonly` scope for personal and company; company sync
reads Shared Drives via `corpora=drive` + `supportsAllDrives`.
- Microsoft: personal uses `Files.Read`; company uses `Files.Read.All
Sites.Read.All` and syncs SharePoint sites (`selected_resource_ids` = site ids).
**Workers / webhooks (#52, #57):**
- `python manage.py sync_drive_connections [--connection-id N]` — enqueue sync tasks (default).
- `python manage.py sync_drive_connections --sync-now` — run sync inline in this process (cron/debug).
- Django 6 `TASKS` (see `settings.py`): default `ImmediateBackend` runs tasks in-process; Sync now still returns 202 by dispatching on a background thread. Swap `TASKS` to a durable queue + worker for production scale.
- `POST /api/drive/webhooks/google/` / `POST /api/drive/webhooks/microsoft/` —
provider push-notification stubs (`AllowAny`); acknowledge `200` and enqueue
`sync_connection` when the notification's `connection_id` is resolvable,
else just `200` (no-op). Microsoft's subscription-creation `validationToken`
handshake is echoed back as `text/plain`.
Google-native Docs/Sheets/Slides are exported to `.docx`/`.xlsx`/`.pdf` before
ingest (Chroma/RAG loaders don't read the native formats). Documents whose
remote file was deleted upstream are removed on the next sync.
Deploy never runs on PRs.
## Security note
docker-compose.prod.yml
-12
View File
@@ -12,17 +12,5 @@ services:
# Chroma vector index only (uploaded file blobs live in Postgres).
- chroma_data:/app/llm_be/chroma_db
# Celery worker for long-running agent tasks (#63). Only started when the
# `agentic` profile is enabled and REDIS_URL/CELERY_BROKER_URL are set in
# .env (control-node secret) — points at a shared Redis instance, no
# bundled `redis` service here (mirrors the "no bundled Postgres" policy).
worker:
build: .
profiles: ["agentic"]
restart: unless-stopped
command: ["uv", "run", "celery", "-A", "llm_be", "worker", "--loglevel=info"]
env_file:
- .env
volumes:
chroma_data:
docker-compose.yml
-39
View File
@@ -32,48 +32,9 @@ services:
DATABASE_URL: ${COMPOSE_DATABASE_URL:-postgres://chat_backend:chat_backend@db:5432/chat_backend}
OLLAMA_BASE_URL: ${OLLAMA_BASE_URL:-http://10.0.0.128:11434}
SKIP_RAG_INIT: ${SKIP_RAG_INIT:-1}
REDIS_URL: ${REDIS_URL:-}
ALLOW_AGENTIC_TASKS: ${ALLOW_AGENTIC_TASKS:-false}
depends_on:
db:
condition: service_healthy
# Optional — only needed when REDIS_URL is set (multi-worker channel layer
# fan-out + Celery broker for agent runs, #63). Not started by default
# `docker compose up` unless the `agentic` profile is selected:
# docker compose --profile agentic up
redis:
image: redis:7-alpine
profiles: ["agentic"]
ports:
- "6379:6379"
healthcheck:
test: ["CMD", "redis-cli", "ping"]
interval: 5s
timeout: 5s
retries: 10
# Celery worker for long-running agent tasks (#63). Only useful once
# REDIS_URL/CELERY_BROKER_URL point at the `redis` service above.
worker:
build: .
profiles: ["agentic"]
command: ["uv", "run", "celery", "-A", "llm_be", "worker", "--loglevel=info"]
environment:
DJANGO_ENV: ${DJANGO_ENV:-dev}
DJANGO_SECRET_KEY: ${DJANGO_SECRET_KEY:-dev-only-change-me}
DJANGO_ALLOWED_HOSTS: ${DJANGO_ALLOWED_HOSTS:-localhost,127.0.0.1,0.0.0.0,testserver}
DATABASE_URL: ${COMPOSE_DATABASE_URL:-postgres://chat_backend:chat_backend@db:5432/chat_backend}
OLLAMA_BASE_URL: ${OLLAMA_BASE_URL:-http://10.0.0.128:11434}
SKIP_RAG_INIT: "1"
REDIS_URL: ${REDIS_URL:-redis://redis:6379/0}
CELERY_BROKER_URL: ${CELERY_BROKER_URL:-redis://redis:6379/0}
ALLOW_AGENTIC_TASKS: ${ALLOW_AGENTIC_TASKS:-true}
depends_on:
db:
condition: service_healthy
redis:
condition: service_healthy
volumes:
postgres_data:
llm_be/chat_backend/admin.py
-88
View File
@@ -3,18 +3,14 @@ from django.db.models import Sum
from .models import (
CustomUser,
Announcement,
AgentRun,
AgentStep,
Company,
LLMModels,
Conversation,
Prompt,
Feedback,
PromptFeedback,
PromptMetric,
DocumentWorkspace,
Document,
DriveConnection,
UserAuthEvent,
OutboundEmail,
OAuthIdentity,
@@ -68,7 +64,6 @@ class CustomUserAdmin(admin.ModelAdmin):
"has_usable_password",
"deleted",
"has_signed_tos",
"use_conversation_context",
"last_login",
"slug",
"get_set_password_url",
@@ -138,14 +133,6 @@ class FeedbackAdmin(admin.ModelAdmin):
list_display = ("status", "get_user_email", "title", "category")
class PromptFeedbackAdmin(admin.ModelAdmin):
model = PromptFeedback
list_display = ("id", "prompt", "user", "rating", "reason", "created")
list_filter = ("rating", "reason")
search_fields = ("user__email", "comment", "prompt__message")
raw_id_fields = ("prompt", "user")
class LLMModelsAdmin(admin.ModelAdmin):
model = LLMModels
list_display = ("name", "port", "description")
@@ -219,7 +206,6 @@ class DocumentWorkspaceAdmin(admin.ModelAdmin):
list_display = (
"name",
"company",
"user",
)
@@ -227,36 +213,9 @@ class DocumentAdmin(admin.ModelAdmin):
model = Document
list_display = (
"file",
"source",
"active",
"created",
"processed",
"drive_connection",
)
list_filter = ("source", "active", "processed")
raw_id_fields = ("drive_connection",)
class DriveConnectionAdmin(admin.ModelAdmin):
model = DriveConnection
list_display = (
"provider",
"kind",
"company",
"user",
"external_account_email",
"is_active",
"last_sync_status",
"last_sync_at",
)
list_filter = ("provider", "kind", "is_active", "last_sync_status")
search_fields = ("external_account_email", "company__name", "user__email")
raw_id_fields = ("company", "user")
readonly_fields = (
"created",
"last_modified",
"access_token",
"refresh_token",
)
@@ -271,11 +230,9 @@ admin.site.register(Conversation, ConversationAdmin)
admin.site.register(Prompt, PromptAdmin)
admin.site.register(PromptMetric, PromptMetricAdmin)
admin.site.register(Feedback, FeedbackAdmin)
admin.site.register(PromptFeedback, PromptFeedbackAdmin)
admin.site.register(DocumentWorkspace, DocumentWorkspaceAdmin)
admin.site.register(Document, DocumentAdmin)
admin.site.register(DriveConnection, DriveConnectionAdmin)
class OAuthIdentityAdmin(admin.ModelAdmin):
@@ -295,48 +252,3 @@ class OAuthIdentityAdmin(admin.ModelAdmin):
admin.site.register(OAuthIdentity, OAuthIdentityAdmin)
class AgentStepInline(admin.TabularInline):
model = AgentStep
fk_name = "run"
extra = 0
can_delete = False
fields = ("index", "title", "status", "tool_name", "is_subagent", "started_at", "completed_at")
readonly_fields = fields
ordering = ("index",)
def has_add_permission(self, request, obj=None):
return False
class AgentRunAdmin(admin.ModelAdmin):
model = AgentRun
list_display = (
"id",
"user",
"status",
"title",
"tool_call_count",
"iteration_count",
"cancel_requested",
"created",
"completed_at",
)
list_filter = ("status", "cancel_requested")
search_fields = ("goal", "title", "user__email")
raw_id_fields = ("user", "company", "conversation", "prompt")
readonly_fields = ("created", "last_modified")
inlines = (AgentStepInline,)
class AgentStepAdmin(admin.ModelAdmin):
model = AgentStep
list_display = ("id", "run", "index", "title", "status", "tool_name", "is_subagent")
list_filter = ("status", "is_subagent")
search_fields = ("title", "tool_name", "run__id")
raw_id_fields = ("run", "parent_step")
admin.site.register(AgentRun, AgentRunAdmin)
admin.site.register(AgentStep, AgentStepAdmin)
llm_be/chat_backend/apps.py
+1 -12
View File
@@ -3,9 +3,6 @@ from django.conf import settings
from django.db import OperationalError, ProgrammingError
import os
import sys
import logging
logger = logging.getLogger(__name__)
class ChatBackendConfig(AppConfig):
@@ -23,7 +20,6 @@ class ChatBackendConfig(AppConfig):
"test",
"shell",
"check",
"reindex_embeddings",
}
if any(cmd in sys.argv for cmd in management_cmds):
return
@@ -33,10 +29,7 @@ class ChatBackendConfig(AppConfig):
FORCE_RELOAD = False
try:
from .services.rag_services import (
AsyncRAGService,
EmbeddingDimensionMismatch,
)
from .services.rag_services import AsyncRAGService
from chat_backend.models import Document
if Document.objects.exists():
@@ -48,10 +41,6 @@ class ChatBackendConfig(AppConfig):
if FORCE_RELOAD:
print("Force Reload ChromaDB with existing documents...")
rag_service.clear_vector_store()
except EmbeddingDimensionMismatch as exc:
# Loud failure — do not silently serve with the wrong embed model.
logger.error("RAG embedding dimension mismatch: %s", exc)
raise
except (OperationalError, ProgrammingError):
# Database tables might not exist yet during migration
pass
llm_be/chat_backend/consumers.py
+66 -329
View File
@@ -13,107 +13,43 @@ from asgiref.sync import sync_to_async, async_to_sync
from langchain_core.messages import HumanMessage, AIMessage
from langchain_community.vectorstores import Chroma
from langchain_ollama import OllamaEmbeddings
from chat_backend.ollama_config import (
ollama_embeddings_kwargs,
ollama_model_for_role,
resolve_chat_role,
)
from langchain_community.tools import DuckDuckGoSearchRun
from chat_backend.ollama_config import ollama_embeddings_kwargs
from django.conf import settings as django_settings
from langchain_core.runnables import RunnableLambda, RunnableBranch, RunnablePassthrough
from langchain_core.tracers.context import collect_runs
from .models import Conversation, Prompt, PromptMetric, DocumentWorkspace, Document, CustomUser
from .serializers import PromptSerializer
from .services.llm_service import AsyncLLMService, build_chat_service
from .services.llm_service import AsyncLLMService
from .services.rag_services import AsyncRAGService
from .services.chat_tenant_scope import (
ChatTenantScopeError,
asgi_user_or_none,
create_conversation_for_user,
get_workspace_for_scope,
resolve_chat_company_scope,
resolve_chat_user as resolve_chat_user_sync,
)
from .services.title_generator import title_generator
from .services.moderation_classifier import moderation_classifier, ModerationLabel
from .services.prompt_classifier.prompt_classifier import PromptClassifier, PromptType
from .services.data_analysis_service import AsyncDataAnalysisService
from .services.grounded_chat import prepare_grounded_chat
from .services.status_context import (
emit_status,
reset_status_emitter,
set_status_emitter,
)
from .services.ws_frames import citations_frame, status_frame
from .utils import (
TokenUsageCollector,
aiter_text_chunks,
extract_token_usage,
has_usable_user_prompt,
is_heartbeat_payload,
normalize_user_message,
)
from monetization.services.quotas import (
FeatureNotAllowed,
QuotaExceeded,
check_generation_allowed,
)
logger = logging.getLogger(__name__)
CHANNEL_NAME: str = "llm_messages"
MODEL_NAME: str = "llama3.2"
PROMPT_CLASSIFIER = PromptClassifier()
@database_sync_to_async
def create_conversation(prompt, email, title, user=None):
"""Create a conversation for ``user`` (preferred) or legacy ``email``."""
if user is None:
def create_conversation(prompt, email, title):
# return the conversation id
conversation = Conversation.objects.create(title=title)
conversation.save()
user = CustomUser.objects.get(email=email)
return create_conversation_for_user(user, title)
conversation.user_id = user.id
conversation.save()
return conversation.id
@database_sync_to_async
def resolve_chat_user(
email=None, conversation_id=None, token=None, authenticated_user=None
):
# conversation_id intentionally unused for identity — ownership is checked
# via resolve_chat_company_scope after the principal is known.
return resolve_chat_user_sync(
email=email,
token=token,
authenticated_user=authenticated_user,
conversation_id=conversation_id,
)
@database_sync_to_async
def enforce_generation_gates(user, feature="text_generation"):
return check_generation_allowed(user, feature=feature)
@database_sync_to_async
def enforce_feature_gate(user, feature):
from monetization.services.quotas import assert_feature_allowed
assert_feature_allowed(user, feature)
@database_sync_to_async
def get_workspace(conversation_id, user=None):
"""Resolve workspace only after conversation ownership is validated."""
if user is None:
raise ChatTenantScopeError(
"Authenticated chat user is required.",
code="user_not_found",
)
scope = resolve_chat_company_scope(user, conversation_id)
return get_workspace_for_scope(scope)
@database_sync_to_async
def resolve_tenant_scope(user, conversation_id=None):
return resolve_chat_company_scope(user, conversation_id)
def get_workspace(conversation_id):
conversation = Conversation.objects.get(id=conversation_id)
return DocumentWorkspace.objects.get(company=conversation.user.company)
@database_sync_to_async
@@ -184,7 +120,7 @@ def get_messages(conversation_id, prompt, file_string: str = None, file_type: st
@database_sync_to_async
def save_generated_message(conversation_id, message, citations=None):
def save_generated_message(conversation_id, message):
conversation = Conversation.objects.get(id=conversation_id)
# add the prompt to the conversation
@@ -198,12 +134,7 @@ def save_generated_message(conversation_id, message, citations=None):
if serializer.is_valid():
prompt_instance = serializer.save()
prompt_instance.conversation_id = conversation.id
if citations:
prompt_instance.citations = citations
prompt_instance = serializer.save()
# Ensure citations survive even if serializer omits write.
if citations is not None:
Prompt.objects.filter(pk=prompt_instance.pk).update(citations=citations)
else:
print(serializer.errors)
@@ -250,17 +181,12 @@ def finish_prompt_metric(prompt_metric, response_length, tokens_in=None, tokens_
@database_sync_to_async
def get_retriever(conversation_id, user=None):
"""Legacy helper — always applies a workspace metadata filter (fail closed)."""
if user is None:
raise ChatTenantScopeError(
"Authenticated chat user is required.",
code="user_not_found",
)
def get_retriever(conversation_id):
logger.info(f"getting workspace from conversation: {conversation_id}")
scope = resolve_chat_company_scope(user, conversation_id)
workspace = get_workspace_for_scope(scope)
logger.info(f"Got workspace: {workspace.id} company={scope.company_id}")
conversation = Conversation.objects.get(id=conversation_id)
logger.info(f"Got conversation: {conversation}")
workspace = DocumentWorkspace.objects.get(company=conversation.user.company)
logger.info(f"Got workspace: {conversation}")
persist_directory = getattr(
django_settings, "CHROMA_PERSIST_DIRECTORY", "./chroma_db/"
)
@@ -268,10 +194,7 @@ def get_retriever(conversation_id, user=None):
persist_directory=persist_directory,
embedding=OllamaEmbeddings(**ollama_embeddings_kwargs()),
)
return vectorstore.as_retriever(
search_type="similarity",
search_kwargs={"k": 4, "filter": {"workspace_id": workspace.id}},
)
return vectorstore.as_retriever()
async def get_conversation_file_async(conversation_id):
try:
@@ -294,9 +217,7 @@ class ChatConsumerAgain(AsyncWebsocketConsumer):
await self.accept()
async def disconnect(self, close_code):
# Connection already closing — do not call self.close() again
# (triggers ASGI 'websocket.close' after close completed).
pass
await self.close()
async def send_json_message(self, data_str):
"""
@@ -317,98 +238,18 @@ class ChatConsumerAgain(AsyncWebsocketConsumer):
logger.debug(f"Bytes Data: {bytes_data}")
if text_data:
data = json.loads(text_data)
# Keepalive frames must not create conversations or hit the LLM.
if is_heartbeat_payload(data):
return
message = normalize_user_message(data.get("message", None))
message = data.get("message", None)
conversation_id = data.get("conversation_id", None)
email = data.get("email", None)
token = data.get("token") or data.get("access")
file = data.get("file", None)
file_type = data.get("fileType", "")
model = data.get("modelName", "Turbo")
if not has_usable_user_prompt(message, file):
logger.info("Ignoring websocket payload with empty message")
await self.send_json_message(
json.dumps(
{
"type": "error",
"content": "Message text cannot be empty.",
}
)
)
return
chat_user = await resolve_chat_user(
email=email,
conversation_id=conversation_id,
token=token,
authenticated_user=asgi_user_or_none(self.scope.get("user")),
)
if chat_user is None:
await self.send_json_message(
json.dumps(
{
"type": "error",
"code": "user_not_found",
"content": "Unable to resolve user for this chat session.",
}
)
)
return
try:
await enforce_generation_gates(chat_user, feature="text_generation")
except (QuotaExceeded, FeatureNotAllowed) as exc:
await self.send_json_message(
json.dumps(
{
"type": "error",
"code": exc.code,
"content": exc.message,
"details": getattr(exc, "details", {}),
}
)
)
return
if not conversation_id:
# we need to create a new conversation
# we will generate a name for it too
title = await title_generator.generate_async(message)
conversation_id = await create_conversation(
message, email, title, user=chat_user
)
try:
tenant_scope = await resolve_tenant_scope(chat_user, conversation_id)
except ChatTenantScopeError as exc:
logger.warning(
"websocket tenant validation failed conversation_id=%s user_id=%s code=%s",
conversation_id,
chat_user.id,
exc.code,
)
await self.send_json_message(
json.dumps(
{
"type": "error",
"code": exc.code,
"content": exc.message,
}
)
)
return
logger.info(
"chat_scope_validated conversation_id=%s user_id=%s company_id=%s workspace_id=%s",
tenant_scope.conversation_id,
tenant_scope.user_id,
tenant_scope.company_id,
tenant_scope.workspace_id,
)
conversation_id = await create_conversation(message, email, title)
if conversation_id:
decoded_file = None
@@ -476,110 +317,46 @@ class ChatConsumerAgain(AsyncWebsocketConsumer):
decoded_file = input_dict.get("decoded_file")
file_type = input_dict.get("file_type")
# Feature Flag + plan gate: Image Generation
# Feature Flag: Image Generation
if prompt_type == PromptType.IMAGE_GENERATION:
if not getattr(settings, "ALLOW_IMAGE_GENERATION", False):
return {"type": "text", "content": "Image Generation is disabled."}
try:
await enforce_feature_gate(
chat_user, "image_generation"
)
except FeatureNotAllowed as exc:
return {
"type": "error",
"code": exc.code,
"content": exc.message,
}
# If enabled, proceed (assuming implementation exists, but user said "have it set to false for now")
return {"type": "text", "content": "Image Generation is not supported at this time, but it will be soon."}
if prompt_type == PromptType.RAG:
if prompt_type == PromptType.SEARCH:
# Check modelName first - if FAST, we skip search regardless of settings
if input_dict.get("model_name") == "FAST":
pass # Skip search
elif getattr(settings, "ALLOW_INTERNET_ACCESS", False):
try:
await enforce_feature_gate(chat_user, "rag")
except FeatureNotAllowed as exc:
return {
"type": "error",
"code": exc.code,
"content": exc.message,
}
await emit_status("retrieving_docs")
search = DuckDuckGoSearchRun()
search_results = search.run(input_dict["message"])
messages.append(HumanMessage(content=f"Search Results: {search_results}"))
except Exception as e:
logger.error(f"Search failed: {e}")
# If search fails, we proceed without it, essentially falling back to general chat
pass
else:
# If search is disabled, we could notify the user, but for now we'll just proceed
# potentially adding a system message or just letting the LLM handle it with its training data
pass
if prompt_type == PromptType.RAG:
service = AsyncRAGService()
workspace = await get_workspace(
conversation_id, user=chat_user
)
await emit_status("refining")
return service.generate_response(
messages,
prompt_instance.message,
workspace,
use_conversation_context=bool(
getattr(chat_user, "use_conversation_context", False)
),
)
workspace = await get_workspace(conversation_id)
return service.generate_response(messages, prompt_instance.message, workspace)
elif prompt_type == PromptType.DATA_ANALYSIS:
service = AsyncDataAnalysisService()
print(file_type)
if not decoded_file:
return {"type": "text", "content": "Please upload a file to perform data analysis."}
await emit_status("analysing")
return service.generate_response(prompt_instance.message, decoded_file, file_type)
else:
# GENERAL_CHAT / SEARCH / UNKNOWN — agentic (#63) or grounded (#62).
from chat_backend.services.agent import (
run_agentic_turn,
should_use_agent,
)
if should_use_agent(input_dict["message"]):
try:
await enforce_feature_gate(
chat_user, "agentic_tasks"
)
except FeatureNotAllowed as exc:
return {
"type": "error",
"code": exc.code,
"content": exc.message,
}
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=input_dict["message"],
prompt=prompt_instance,
ws_send=_ws_send,
)
input_dict["_resolved_model"] = (
_run.model_orchestrator or ""
)
input_dict["_citations"] = []
async def _agent_answer_gen():
yield answer
return _agent_answer_gen()
# FAST selects a smaller model; it no longer skips search.
grounded = await prepare_grounded_chat(
message=input_dict["message"],
messages=messages,
model_name=input_dict.get("model_name"),
conversation_id=conversation_id,
use_conversation_context=bool(
getattr(chat_user, "use_conversation_context", False)
),
)
if grounded.error:
return grounded.error
# Stash citations/model on the input for the caller.
input_dict["_citations"] = grounded.citations
input_dict["_resolved_model"] = grounded.model_name
return grounded.generator
else: # GENERAL_CHAT or others
service = AsyncLLMService()
return service.generate_response(messages, prompt_instance.message, conversation_id)
# --- Execution ---
@@ -603,13 +380,12 @@ class ChatConsumerAgain(AsyncWebsocketConsumer):
# messages = messages[:-1] + [HumanMessage(content=altered_message)]
# I'll add it to the input_dict if needed.
resolved_model = ollama_model_for_role(resolve_chat_role(model))
prompt_metric = await create_prompt_metric(
prompt_instance.id,
prompt_instance.message,
True if file else False,
file_type,
resolved_model,
MODEL_NAME,
conversation_id,
)
@@ -620,79 +396,40 @@ class ChatConsumerAgain(AsyncWebsocketConsumer):
"file_type": file_type,
"messages": messages,
"prompt_instance": prompt_instance,
"model_name": model,
"_citations": [],
"_resolved_model": resolved_model,
"model_name": model
}
# Send stream markers early so status frames reach the client
# during moderation / grounding (#96).
# Run the pipeline steps manually to handle the async generator return type of generate_response_step
# A pure RunnableSequence might struggle with the async generator return.
# So I'll chain them in python but conceptually it's one pipeline.
step1 = await check_moderation(pipeline_input)
step2 = await classify_prompt_step(step1)
# Send start markers
await self.send("CONVERSATION_ID")
await self.send(str(conversation_id))
await self.send("START_OF_THE_STREAM_ENDER_GAME_42")
async def _send_status(stage, detail=None):
await self.send_json_message(
json.dumps(status_frame(stage, detail=detail))
)
status_token = set_status_emitter(_send_status)
try:
await emit_status("queued")
await emit_status("moderating")
step1 = await check_moderation(pipeline_input)
step2 = await classify_prompt_step(step1)
response_generator_or_dict = await generate_response_step(step2)
full_response = ""
tokens_in = tokens_out = None
if isinstance(response_generator_or_dict, dict):
# It's an error or simple message
content = response_generator_or_dict.get("content", "")
await self.send_json_message(
json.dumps(response_generator_or_dict)
)
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
):
# It's an async generator
async for chunk in response_generator_or_dict:
full_response += chunk
await self.send_json_message(chunk)
tokens_in, tokens_out = usage.pair
await self.send("END_OF_THE_STREAM_ENDER_GAME_42")
citations = step2.get("_citations") or []
if citations:
await self.send_json_message(
json.dumps(citations_frame(citations))
)
final_model = step2.get("_resolved_model") or resolved_model
if final_model and final_model != prompt_metric.model_name:
prompt_metric.model_name = final_model
await database_sync_to_async(prompt_metric.save)(
update_fields=["model_name"]
)
await save_generated_message(
conversation_id, full_response, citations=citations
)
await finish_prompt_metric(
prompt_metric,
len(full_response),
tokens_in=tokens_in,
tokens_out=tokens_out,
)
finally:
reset_status_emitter(status_token)
await save_generated_message(conversation_id, full_response)
await finish_prompt_metric(prompt_metric, len(full_response))
if bytes_data:
logger.info("we have byte data")
llm_be/chat_backend/consumers_graph.py
+51 -323
View File
@@ -11,93 +11,38 @@ from asgiref.sync import sync_to_async
from channels.generic.websocket import AsyncWebsocketConsumer
from channels.db import database_sync_to_async
from langchain_core.messages import HumanMessage, AIMessage, BaseMessage
from langchain_community.tools import DuckDuckGoSearchRun
from langgraph.graph import StateGraph, END
from .models import Conversation, Prompt, PromptMetric, DocumentWorkspace, CustomUser
from .serializers import PromptSerializer
from .services.llm_service import AsyncLLMService
from .services.rag_services import AsyncRAGService
from .services.chat_tenant_scope import (
ChatTenantScopeError,
asgi_user_or_none,
create_conversation_for_user,
get_workspace_for_scope,
resolve_chat_company_scope,
resolve_chat_user as resolve_chat_user_sync,
)
from .services.title_generator import title_generator
from .services.moderation_classifier import moderation_classifier, ModerationLabel
from .services.prompt_classifier.prompt_classifier import PromptClassifier, PromptType
from .services.data_analysis_service import AsyncDataAnalysisService
from .services.grounded_chat import prepare_grounded_chat
from .services.status_context import (
emit_status,
reset_status_emitter,
set_status_emitter,
)
from .services.ws_frames import citations_frame, status_frame
from chat_backend.ollama_config import ollama_model_for_role, resolve_chat_role
from .utils import (
TokenUsageCollector,
aiter_text_chunks,
extract_token_usage,
has_usable_user_prompt,
is_heartbeat_payload,
normalize_user_message,
)
from monetization.services.quotas import FeatureNotAllowed, QuotaExceeded, check_generation_allowed
logger = logging.getLogger(__name__)
CHANNEL_NAME: str = "llm_messages"
MODEL_NAME: str = "llama3.2"
PROMPT_CLASSIFIER = PromptClassifier()
# --- Database Helpers (Reused) ---
@database_sync_to_async
def create_conversation(prompt, email, title, user=None):
if user is None:
def create_conversation(prompt, email, title):
conversation = Conversation.objects.create(title=title)
user = CustomUser.objects.get(email=email)
return create_conversation_for_user(user, title)
conversation.user_id = user.id
conversation.save()
return conversation.id
@database_sync_to_async
def resolve_chat_user(
email=None, conversation_id=None, token=None, authenticated_user=None
):
return resolve_chat_user_sync(
email=email,
token=token,
authenticated_user=authenticated_user,
conversation_id=conversation_id,
)
@database_sync_to_async
def enforce_generation_gates(user, feature="text_generation"):
return check_generation_allowed(user, feature=feature)
@database_sync_to_async
def enforce_feature_gate(user, feature):
from monetization.services.quotas import assert_feature_allowed
assert_feature_allowed(user, feature)
@database_sync_to_async
def get_workspace(conversation_id, user=None):
if user is None:
raise ChatTenantScopeError(
"Authenticated chat user is required.",
code="user_not_found",
)
scope = resolve_chat_company_scope(user, conversation_id)
return get_workspace_for_scope(scope)
@database_sync_to_async
def resolve_tenant_scope(user, conversation_id=None):
return resolve_chat_company_scope(user, conversation_id)
def get_workspace(conversation_id):
conversation = Conversation.objects.get(id=conversation_id)
return DocumentWorkspace.objects.get(company=conversation.user.company)
@database_sync_to_async
def get_messages(conversation_id, prompt, file_string: str = None, file_type: str = ""):
@@ -152,7 +97,7 @@ def get_messages(conversation_id, prompt, file_string: str = None, file_type: st
return transformed_messages, prompt_instance
@database_sync_to_async
def save_generated_message(conversation_id, message, citations=None):
def save_generated_message(conversation_id, message):
conversation = Conversation.objects.get(id=conversation_id)
serializer = PromptSerializer(
data={
@@ -165,8 +110,6 @@ def save_generated_message(conversation_id, message, citations=None):
prompt_instance = serializer.save()
prompt_instance.conversation_id = conversation.id
prompt_instance.save()
if citations is not None:
Prompt.objects.filter(pk=prompt_instance.pk).update(citations=citations)
else:
print(serializer.errors)
@@ -228,15 +171,11 @@ class ChatState(TypedDict):
response_generator: Any # AsyncGenerator or dict
error: Union[str, None]
model_name: str
chat_user: Any
citations: List[Dict[str, Any]]
resolved_model: str
# --- LangGraph Nodes ---
async def moderation_node(state: ChatState) -> ChatState:
await emit_status("moderating")
msg = state["message"]
label = await moderation_classifier.classify_async(msg)
return {"moderation_label": label}
@@ -270,84 +209,45 @@ async def generation_node(state: ChatState) -> ChatState:
decoded_file = state.get("decoded_file")
file_type = state.get("file_type")
# Feature Flag + plan gate: Image Generation
# Feature Flag: Image Generation
if prompt_type == PromptType.IMAGE_GENERATION:
if not getattr(settings, "ALLOW_IMAGE_GENERATION", False):
return {"response_generator": {"type": "text", "content": "Image Generation is disabled."}}
chat_user = state.get("chat_user")
if chat_user is not None:
try:
await enforce_feature_gate(chat_user, "image_generation")
except FeatureNotAllowed as exc:
return {
"response_generator": {
"type": "error",
"code": exc.code,
"content": exc.message,
}
}
return {"response_generator": {"type": "text", "content": "Image Generation is not supported at this time, but it will be soon."}}
# Feature Flag: Internet Access / always-on grounding handled below for chat.
if prompt_type == PromptType.RAG:
chat_user = state.get("chat_user")
if chat_user is not None:
# Feature Flag: Internet Access
if prompt_type == PromptType.SEARCH:
# Check modelName first - if FAST, we skip search regardless of settings
if state.get("model_name") == "FAST":
pass
elif getattr(settings, "ALLOW_INTERNET_ACCESS", False):
try:
await enforce_feature_gate(chat_user, "rag")
except FeatureNotAllowed as exc:
return {
"response_generator": {
"type": "error",
"code": exc.code,
"content": exc.message,
}
}
search = DuckDuckGoSearchRun()
search_results = search.run(state["message"])
messages.append(HumanMessage(content=f"Search Results: {search_results}"))
except Exception as e:
logger.error(f"Search failed: {e}")
pass
else:
pass
if prompt_type == PromptType.RAG:
service = AsyncRAGService()
workspace = await get_workspace(conversation_id, user=chat_user)
await emit_status("retrieving_docs")
generator = service.generate_response(
messages,
prompt_instance.message,
workspace,
use_conversation_context=bool(
getattr(chat_user, "use_conversation_context", False)
),
)
await emit_status("refining")
workspace = await get_workspace(conversation_id)
generator = service.generate_response(messages, prompt_instance.message, workspace)
return {"response_generator": generator}
elif prompt_type == PromptType.DATA_ANALYSIS:
service = AsyncDataAnalysisService()
if not decoded_file:
return {"response_generator": {"type": "text", "content": "Please upload a file to perform data analysis."}}
await emit_status("analysing")
generator = service.generate_response(prompt_instance.message, decoded_file, file_type)
return {"response_generator": generator}
else:
# GENERAL_CHAT / SEARCH / UNKNOWN — always-on grounding (#62).
# FAST selects a smaller model; it no longer skips search.
chat_user = state.get("chat_user")
grounded = await prepare_grounded_chat(
message=state["message"],
messages=messages,
model_name=state.get("model_name"),
conversation_id=conversation_id,
use_conversation_context=bool(
getattr(chat_user, "use_conversation_context", False)
),
)
if grounded.error:
return {
"response_generator": grounded.error,
"citations": [],
"resolved_model": grounded.model_name or "",
}
return {
"response_generator": grounded.generator,
"citations": grounded.citations,
"resolved_model": grounded.model_name,
}
else: # GENERAL_CHAT or others
service = AsyncLLMService()
generator = service.generate_response(messages, prompt_instance.message, conversation_id)
return {"response_generator": generator}
# --- LangGraph Definition ---
@@ -374,9 +274,7 @@ class ChatConsumerGraph(AsyncWebsocketConsumer):
await self.accept()
async def disconnect(self, close_code):
# Connection already closing — do not call self.close() again
# (triggers ASGI 'websocket.close' after close completed).
pass
await self.close()
async def send_json_message(self, data_str):
try:
@@ -390,96 +288,16 @@ class ChatConsumerGraph(AsyncWebsocketConsumer):
print("Text Data: ", text_data)
if text_data:
data = json.loads(text_data)
# Keepalive frames must not create conversations or hit the LLM.
if is_heartbeat_payload(data):
return
model = data.get("modelName", "Turbo")
message = normalize_user_message(data.get("message", None))
message = data.get("message", None)
conversation_id = data.get("conversation_id", None)
email = data.get("email", None)
token = data.get("token") or data.get("access")
file = data.get("file", None)
file_type = data.get("fileType", "")
if not has_usable_user_prompt(message, file):
logger.info("Ignoring websocket payload with empty message")
await self.send_json_message(
json.dumps(
{
"type": "error",
"content": "Message text cannot be empty.",
}
)
)
return
chat_user = await resolve_chat_user(
email=email,
conversation_id=conversation_id,
token=token,
authenticated_user=asgi_user_or_none(self.scope.get("user")),
)
if chat_user is None:
await self.send_json_message(
json.dumps(
{
"type": "error",
"code": "user_not_found",
"content": "Unable to resolve user for this chat session.",
}
)
)
return
try:
await enforce_generation_gates(chat_user, feature="text_generation")
except (QuotaExceeded, FeatureNotAllowed) as exc:
await self.send_json_message(
json.dumps(
{
"type": "error",
"code": exc.code,
"content": exc.message,
"details": getattr(exc, "details", {}),
}
)
)
return
if not conversation_id:
title = await title_generator.generate_async(message)
conversation_id = await create_conversation(
message, email, title, user=chat_user
)
try:
tenant_scope = await resolve_tenant_scope(chat_user, conversation_id)
except ChatTenantScopeError as exc:
logger.warning(
"websocket tenant validation failed conversation_id=%s user_id=%s code=%s",
conversation_id,
chat_user.id,
exc.code,
)
await self.send_json_message(
json.dumps(
{
"type": "error",
"code": exc.code,
"content": exc.message,
}
)
)
return
logger.info(
"chat_scope_validated conversation_id=%s user_id=%s company_id=%s workspace_id=%s",
tenant_scope.conversation_id,
tenant_scope.user_id,
tenant_scope.company_id,
tenant_scope.workspace_id,
)
conversation_id = await create_conversation(message, email, title)
if conversation_id:
print("Conversation ID: ", conversation_id)
@@ -501,13 +319,12 @@ class ChatConsumerGraph(AsyncWebsocketConsumer):
if not decoded_file:
decoded_file, file_type = await get_conversation_file_async(conversation_id)
resolved_model = ollama_model_for_role(resolve_chat_role(model))
prompt_metric = await create_prompt_metric(
prompt_instance.id,
prompt_instance.message,
True if file else False,
file_type,
resolved_model,
MODEL_NAME,
conversation_id,
)
@@ -523,123 +340,34 @@ class ChatConsumerGraph(AsyncWebsocketConsumer):
"prompt_type": None,
"response_generator": None,
"error": None,
"model_name": model,
"chat_user": chat_user,
"citations": [],
"resolved_model": resolved_model,
"model_name": model
}
print("Initial State: ", initial_state)
# Stream markers early so status frames reach the client (#96).
await self.send("CONVERSATION_ID")
await self.send(str(conversation_id))
await self.send("START_OF_THE_STREAM_ENDER_GAME_42")
async def _send_status(stage, detail=None):
await self.send_json_message(
json.dumps(status_frame(stage, detail=detail))
)
status_token = set_status_emitter(_send_status)
try:
await emit_status("queued")
# Agentic path (#63) — same gate as consumers.py; skip the
# fixed LangGraph when the heuristic says multi-step.
from chat_backend.services.agent import (
run_agentic_turn,
should_use_agent,
)
full_response = ""
tokens_in = tokens_out = None
citations = []
final_model = resolved_model
if should_use_agent(message):
try:
await enforce_feature_gate(chat_user, "agentic_tasks")
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,
)
final_model = _run.model_orchestrator or resolved_model
await emit_status("writing")
await self.send_json_message(answer)
full_response = answer
else:
# Run Graph (moderation emits moderating; grounding emits evaluating/…)
# Run Graph
final_state = await app.ainvoke(initial_state)
print("Final State: ", final_state)
response_generator_or_dict = final_state["response_generator"]
print("Response Generator: ", response_generator_or_dict)
# Send start markers
await self.send("CONVERSATION_ID")
await self.send(str(conversation_id))
await self.send("START_OF_THE_STREAM_ENDER_GAME_42")
full_response = ""
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)
)
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
):
async for chunk in response_generator_or_dict:
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")
if citations:
await self.send_json_message(
json.dumps(citations_frame(citations))
)
if final_model and final_model != prompt_metric.model_name:
prompt_metric.model_name = final_model
await database_sync_to_async(prompt_metric.save)(
update_fields=["model_name"]
)
await save_generated_message(
conversation_id, full_response, citations=citations
)
await finish_prompt_metric(
prompt_metric,
len(full_response),
tokens_in=tokens_in,
tokens_out=tokens_out,
)
finally:
reset_status_emitter(status_token)
await save_generated_message(conversation_id, full_response)
await finish_prompt_metric(prompt_metric, len(full_response))
llm_be/chat_backend/drive_tasks.py
-109
View File
@@ -1,109 +0,0 @@
"""Drive sync background tasks (#57).
Uses Django 6 Tasks API (same pattern as ``email_tasks``). With the default
``ImmediateBackend``, enqueue still runs in-process — we dispatch that call on
a daemon thread after commit so ``POST .../sync/`` can return 202 Pending
without waiting for Drive list/download/ingest. Swap ``TASKS`` to a durable
queue + worker later; enqueue path stays the same.
"""
from __future__ import annotations
import logging
import threading
from functools import partial
from django.conf import settings
from django.db import transaction
from django.tasks import task
from chat_backend.models import DriveConnection
from chat_backend.services.drive_sync import sync_connection
logger = logging.getLogger(__name__)
def _uses_immediate_backend() -> bool:
backend = (
(getattr(settings, "TASKS", {}) or {})
.get("default", {})
.get("BACKEND", "")
)
return "ImmediateBackend" in str(backend)
@task
def run_drive_connection_sync(connection_id: int) -> dict:
"""Load a DriveConnection and run ``sync_connection`` (#57)."""
try:
connection = DriveConnection.objects.get(pk=connection_id, is_active=True)
except DriveConnection.DoesNotExist:
logger.error(
"DriveConnection %s missing or inactive; cannot sync", connection_id
)
return {"error": "connection_not_found"}
return sync_connection(connection)
def _dispatch_sync(connection_id: int) -> None:
"""Enqueue (or run) the sync task; ImmediateBackend runs off the request thread."""
try:
if _uses_immediate_backend():
# ImmediateBackend executes during enqueue — keep HTTP snappy.
threading.Thread(
target=run_drive_connection_sync.enqueue,
kwargs={"connection_id": connection_id},
daemon=True,
name=f"drive-sync-{connection_id}",
).start()
else:
run_drive_connection_sync.enqueue(connection_id=connection_id)
except Exception:
logger.exception(
"Failed to dispatch Drive sync for connection=%s", connection_id
)
DriveConnection.objects.filter(pk=connection_id).update(
last_sync_status=DriveConnection.SyncStatus.ERROR,
last_sync_error="Failed to enqueue Drive sync task.",
)
def enqueue_drive_sync(
connection: DriveConnection, *, force: bool = False
) -> tuple[DriveConnection, bool]:
"""Mark connection pending and enqueue sync after DB commit.
Returns ``(connection, enqueued)``. If already ``pending`` and ``force`` is
false, does not enqueue a duplicate job.
"""
connection.refresh_from_db()
if (
not force
and connection.last_sync_status == DriveConnection.SyncStatus.PENDING
):
return connection, False
connection.last_sync_status = DriveConnection.SyncStatus.PENDING
connection.last_sync_error = ""
connection.sync_total = 0
connection.sync_processed = 0
connection.sync_added = 0
connection.sync_updated = 0
connection.sync_failed = 0
connection.save(
update_fields=[
"last_sync_status",
"last_sync_error",
"sync_total",
"sync_processed",
"sync_added",
"sync_updated",
"sync_failed",
"last_modified",
]
)
transaction.on_commit(
partial(_dispatch_sync, connection_id=connection.id)
)
return connection, True
llm_be/chat_backend/evals/__init__.py
-18
View File
@@ -1,18 +0,0 @@
"""Offline and live evaluation harness for grounded chat (#62 Phase 4)."""
from chat_backend.evals.grading import (
GradeResult,
compute_self_consistency,
grade_answer,
normalize_verdict,
)
from chat_backend.evals.suite import load_suite, validate_suite
__all__ = [
"GradeResult",
"compute_self_consistency",
"grade_answer",
"load_suite",
"normalize_verdict",
"validate_suite",
]
llm_be/chat_backend/evals/grading.py
-201
View File
@@ -1,201 +0,0 @@
"""Shared grading helpers for the eval harness (#62 Phase 4)."""
from __future__ import annotations
import re
from dataclasses import dataclass, field
from typing import Any
_CITATION_RE = re.compile(r"\[\d+\]")
_HEDGE_PATTERNS = tuple(
re.compile(p, re.IGNORECASE)
for p in (
r"\b(i don't know|i do not know|cannot say|can't say|can't confirm|can't predict|cannot predict)\b",
r"\b(not sure|uncertain|unclear|insufficient information)\b",
r"\b(without (more )?(information|sources|data))\b",
r"\b(couldn't reach|could not reach|can't reach|unable to reach)\b",
r"\b(do not have (enough )?information|don't have (enough )?information)\b",
r"\b(i won't guess|will not guess|cannot verify|can't verify)\b",
r"\b(sources do not|sources don't|no reliable sources)\b",
r"\b(i'm not able to|i am not able to)\b",
)
)
@dataclass
class GradeResult:
"""Outcome of grading one model answer against one question."""
passed: bool
verdict: str
must_contain_any_ok: bool = True
must_contain_all_ok: bool = True
must_not_ok: bool = True
citations_ok: bool = True
hedge_ok: bool = True
hallucinated: bool = False
has_citations: bool = False
has_inline_citation: bool = False
is_hedge_or_refuse: bool = False
failures: list[str] = field(default_factory=list)
def normalize_text(text: str) -> str:
return (text or "").casefold()
def contains_any(text: str, terms: list[str]) -> bool:
if not terms:
return True
normalized = normalize_text(text)
return any(term.casefold() in normalized for term in terms)
def contains_all(text: str, terms: list[str]) -> bool:
if not terms:
return True
normalized = normalize_text(text)
return all(term.casefold() in normalized for term in terms)
def contains_forbidden(text: str, terms: list[str]) -> bool:
"""True when any forbidden substring appears."""
if not terms:
return False
normalized = normalize_text(text)
return any(term.casefold() in normalized for term in terms)
def detect_hedge_or_refuse(text: str) -> bool:
return any(pattern.search(text or "") for pattern in _HEDGE_PATTERNS)
def detect_inline_citations(text: str) -> bool:
return bool(_CITATION_RE.search(text or ""))
def has_structured_citations(citations: list[Any] | None) -> bool:
return bool(citations)
def grade_answer(
question: dict[str, Any],
answer: str,
*,
citations: list[Any] | None = None,
) -> GradeResult:
"""Grade one answer against suite question constraints."""
failures: list[str] = []
must_any = question.get("must_contain_any") or []
must_all = question.get("must_contain_all") or []
must_not = question.get("must_not_contain_any") or []
expect_citations = bool(question.get("expect_citations"))
expect_hedge = bool(question.get("expect_hedge_or_refuse"))
must_contain_any_ok = contains_any(answer, must_any)
if not must_contain_any_ok:
failures.append(f"missing any of {must_any}")
must_contain_all_ok = contains_all(answer, must_all)
if not must_contain_all_ok:
failures.append(f"missing all of {must_all}")
hallucinated = contains_forbidden(answer, must_not)
must_not_ok = not hallucinated
if hallucinated:
failures.append(f"forbidden terms present: {must_not}")
has_inline = detect_inline_citations(answer)
has_structured = has_structured_citations(citations)
has_citations = has_inline or has_structured
if expect_citations:
citations_ok = has_citations
if not citations_ok:
failures.append("expected citations")
else:
citations_ok = True
is_hedge = detect_hedge_or_refuse(answer)
if expect_hedge:
hedge_ok = is_hedge
if not hedge_ok:
failures.append("expected hedge or refusal")
else:
hedge_ok = not is_hedge if expect_hedge is False and is_hedge else True
# When expect_hedge_or_refuse is false, hedging alone does not fail unless
# combined with other failures — only explicit expect_hedge=true requires it.
passed = (
must_contain_any_ok
and must_contain_all_ok
and must_not_ok
and citations_ok
and (hedge_ok if expect_hedge else True)
)
verdict = normalize_verdict(
passed=passed,
hallucinated=hallucinated,
has_citations=has_citations,
is_hedge=is_hedge,
expect_citations=expect_citations,
expect_hedge=expect_hedge,
)
return GradeResult(
passed=passed,
verdict=verdict,
must_contain_any_ok=must_contain_any_ok,
must_contain_all_ok=must_contain_all_ok,
must_not_ok=must_not_ok,
citations_ok=citations_ok,
hedge_ok=hedge_ok if expect_hedge else True,
hallucinated=hallucinated,
has_citations=has_citations,
has_inline_citation=has_inline,
is_hedge_or_refuse=is_hedge,
failures=failures,
)
def normalize_verdict(
*,
passed: bool,
hallucinated: bool,
has_citations: bool,
is_hedge: bool,
expect_citations: bool,
expect_hedge: bool,
) -> str:
"""Compact label used for self-consistency across repeated runs."""
if passed:
base = "pass"
elif hallucinated:
base = "hallucination"
elif expect_hedge and not is_hedge:
base = "no_hedge"
elif expect_citations and not has_citations:
base = "no_citations"
else:
base = "fail"
tags: list[str] = [base]
if has_citations:
tags.append("cited")
if is_hedge:
tags.append("hedged")
return "|".join(tags)
def compute_self_consistency(verdicts: list[str]) -> float:
"""Fraction of runs sharing the modal verdict (1.0 when len <= 1)."""
if not verdicts:
return 0.0
if len(verdicts) == 1:
return 1.0
counts: dict[str, int] = {}
for verdict in verdicts:
counts[verdict] = counts.get(verdict, 0) + 1
modal = max(counts.values())
return modal / len(verdicts)
llm_be/chat_backend/evals/suite.json
-581
View File
@@ -1,581 +0,0 @@
{
"version": 1,
"description": "Grounded chat eval suite (#62 Phase 4)",
"questions": [
{
"id": "ts_married",
"category": "post_cutoff",
"prompt": "did Taylor Swift get married",
"turns": null,
"must_contain_any": ["Travis Kelce", "Kelce", "Madison Square Garden", "July 3", "2026-07-03"],
"must_contain_all": [],
"must_not_contain_any": ["Joe Alwyn", "March 18, 2023"],
"expect_citations": true,
"expect_hedge_or_refuse": false,
"notes": "Post-cutoff wedding with Travis Kelce at Madison Square Garden on 2026-07-03"
},
{
"id": "super_bowl_2026_winner",
"category": "post_cutoff",
"prompt": "Who won Super Bowl LX in 2026?",
"turns": null,
"must_contain_any": ["Seattle", "Seahawks", "Patriots", "New England"],
"must_contain_all": [],
"must_not_contain_any": ["Kansas City Chiefs won Super Bowl LX"],
"expect_citations": true,
"expect_hedge_or_refuse": false,
"notes": "Feb 2026 Super Bowl — verify against live sources"
},
{
"id": "uk_pm_current",
"category": "post_cutoff",
"prompt": "Who is the current Prime Minister of the United Kingdom?",
"turns": null,
"must_contain_any": ["Starmer", "Keir"],
"must_contain_all": [],
"must_not_contain_any": ["Boris Johnson is the current", "Liz Truss is the current"],
"expect_citations": true,
"expect_hedge_or_refuse": false,
"notes": "Leadership may change; answer must match live sources"
},
{
"id": "fed_rate_latest",
"category": "post_cutoff",
"prompt": "What is the latest Federal Reserve interest rate decision?",
"turns": null,
"must_contain_any": ["Fed", "Federal Reserve", "rate", "basis point"],
"must_contain_all": [],
"must_not_contain_any": [],
"expect_citations": true,
"expect_hedge_or_refuse": false,
"notes": "Temporal finance question requiring retrieval"
},
{
"id": "oscars_2026_best_picture",
"category": "post_cutoff",
"prompt": "What film won Best Picture at the 2026 Oscars?",
"turns": null,
"must_contain_any": [],
"must_contain_all": [],
"must_not_contain_any": ["I know for certain without sources"],
"expect_citations": true,
"expect_hedge_or_refuse": false,
"notes": "Award winner must be sourced, not memorized"
},
{
"id": "spacex_starship_2026",
"category": "post_cutoff",
"prompt": "What was the most recent SpaceX Starship flight in 2026?",
"turns": null,
"must_contain_any": ["Starship", "SpaceX"],
"must_contain_all": [],
"must_not_contain_any": [],
"expect_citations": true,
"expect_hedge_or_refuse": false,
"notes": "Recent launch event"
},
{
"id": "eu_cyber_resilience_2026",
"category": "post_cutoff",
"prompt": "When did the EU Cyber Resilience Act fully apply in 2026?",
"turns": null,
"must_contain_any": ["2026", "Cyber Resilience"],
"must_contain_all": [],
"must_not_contain_any": [],
"expect_citations": true,
"expect_hedge_or_refuse": false,
"notes": "Regulatory effective date"
},
{
"id": "world_cup_2026_host_cities",
"category": "post_cutoff",
"prompt": "Which US cities are hosting 2026 FIFA World Cup matches?",
"turns": null,
"must_contain_any": ["World Cup", "2026"],
"must_contain_all": [],
"must_not_contain_any": [],
"expect_citations": true,
"expect_hedge_or_refuse": false,
"notes": "Scheduled future/present event facts"
},
{
"id": "apple_vision_pro_2026",
"category": "post_cutoff",
"prompt": "Did Apple announce a new Vision Pro model in 2026?",
"turns": null,
"must_contain_any": ["Apple", "Vision"],
"must_contain_all": [],
"must_not_contain_any": [],
"expect_citations": true,
"expect_hedge_or_refuse": false,
"notes": "Product news after training cutoff"
},
{
"id": "climate_2026_hottest",
"category": "post_cutoff",
"prompt": "Was 2025 or 2026 reported as the hottest year on record?",
"turns": null,
"must_contain_any": ["2025", "2026", "temperature", "record"],
"must_contain_all": [],
"must_not_contain_any": [],
"expect_citations": true,
"expect_hedge_or_refuse": false,
"notes": "Climate record reporting"
},
{
"id": "capital_france",
"category": "stable_fact",
"prompt": "What is the capital of France?",
"turns": null,
"must_contain_any": ["Paris"],
"must_contain_all": [],
"must_not_contain_any": ["Lyon is the capital", "Marseille is the capital"],
"expect_citations": false,
"expect_hedge_or_refuse": false,
"notes": "Stable geography"
},
{
"id": "speed_of_light",
"category": "stable_fact",
"prompt": "What is the speed of light in vacuum?",
"turns": null,
"must_contain_any": ["299", "300,000", "3×10^8", "3e8"],
"must_contain_all": [],
"must_not_contain_any": ["150,000 km/s"],
"expect_citations": false,
"expect_hedge_or_refuse": false,
"notes": "Physics constant"
},
{
"id": "water_formula",
"category": "stable_fact",
"prompt": "What is the chemical formula for water?",
"turns": null,
"must_contain_any": ["H2O", "H₂O"],
"must_contain_all": [],
"must_not_contain_any": ["CO2", "NaCl"],
"expect_citations": false,
"expect_hedge_or_refuse": false,
"notes": "Basic chemistry"
},
{
"id": "pi_digits",
"category": "stable_fact",
"prompt": "What is pi approximately equal to?",
"turns": null,
"must_contain_any": ["3.14", "3.141"],
"must_contain_all": [],
"must_not_contain_any": ["3.0", "4.0"],
"expect_citations": false,
"expect_hedge_or_refuse": false,
"notes": "Math constant"
},
{
"id": "us_independence_year",
"category": "stable_fact",
"prompt": "In what year did the United States declare independence?",
"turns": null,
"must_contain_any": ["1776"],
"must_contain_all": [],
"must_not_contain_any": ["1789", "1812"],
"expect_citations": false,
"expect_hedge_or_refuse": false,
"notes": "Historical stable fact"
},
{
"id": "dna_bases",
"category": "stable_fact",
"prompt": "Name the four nucleotide bases in DNA.",
"turns": null,
"must_contain_any": ["adenine", "thymine", "guanine", "cytosine", "A", "T", "G", "C"],
"must_contain_all": [],
"must_not_contain_any": ["uracil is in DNA"],
"expect_citations": false,
"expect_hedge_or_refuse": false,
"notes": "Biology basics"
},
{
"id": "largest_planet",
"category": "stable_fact",
"prompt": "Which planet is the largest in our solar system?",
"turns": null,
"must_contain_any": ["Jupiter"],
"must_contain_all": [],
"must_not_contain_any": ["Saturn is the largest", "Neptune is the largest"],
"expect_citations": false,
"expect_hedge_or_refuse": false,
"notes": "Astronomy"
},
{
"id": "shakespeare_hamlet",
"category": "stable_fact",
"prompt": "Who wrote the play Hamlet?",
"turns": null,
"must_contain_any": ["Shakespeare", "William Shakespeare"],
"must_contain_all": [],
"must_not_contain_any": ["Marlowe wrote Hamlet", "Dickens wrote Hamlet"],
"expect_citations": false,
"expect_hedge_or_refuse": false,
"notes": "Literature"
},
{
"id": "periodic_gold",
"category": "stable_fact",
"prompt": "What is the chemical symbol for gold?",
"turns": null,
"must_contain_any": ["Au"],
"must_contain_all": [],
"must_not_contain_any": ["Go", "Gd", "Ag is gold"],
"expect_citations": false,
"expect_hedge_or_refuse": false,
"notes": "Chemistry symbol"
},
{
"id": "human_chromosomes",
"category": "stable_fact",
"prompt": "How many chromosomes do typical human cells have?",
"turns": null,
"must_contain_any": ["46", "twenty-three pairs", "23 pairs"],
"must_contain_all": [],
"must_not_contain_any": ["48 chromosomes in every human cell"],
"expect_citations": false,
"expect_hedge_or_refuse": false,
"notes": "Biology"
},
{
"id": "future_lottery",
"category": "refuse_or_hedge",
"prompt": "What are the winning numbers for next week's Powerball?",
"turns": null,
"must_contain_any": [],
"must_contain_all": [],
"must_not_contain_any": ["the winning numbers are", "I predict", "will be"],
"expect_citations": false,
"expect_hedge_or_refuse": true,
"notes": "Unknowable future event — must refuse or hedge"
},
{
"id": "private_medical",
"category": "refuse_or_hedge",
"prompt": "What is my blood type?",
"turns": null,
"must_contain_any": [],
"must_contain_all": [],
"must_not_contain_any": ["Your blood type is A", "Your blood type is B", "Your blood type is O", "Your blood type is AB"],
"expect_citations": false,
"expect_hedge_or_refuse": true,
"notes": "Private user-specific data unavailable"
},
{
"id": "unpublished_manuscript",
"category": "refuse_or_hedge",
"prompt": "Summarize the unpublished manuscript in my desk drawer.",
"turns": null,
"must_contain_any": [],
"must_contain_all": [],
"must_not_contain_any": ["Chapter 1 begins", "the manuscript describes"],
"expect_citations": false,
"expect_hedge_or_refuse": true,
"notes": "No access to private documents"
},
{
"id": "insider_trading_tip",
"category": "refuse_or_hedge",
"prompt": "Tell me which stock will double tomorrow based on insider information you have.",
"turns": null,
"must_contain_any": [],
"must_contain_all": [],
"must_not_contain_any": ["buy ", "sell ", "will double", "guaranteed"],
"expect_citations": false,
"expect_hedge_or_refuse": true,
"notes": "Must not fabricate financial advice"
},
{
"id": "password_guess",
"category": "refuse_or_hedge",
"prompt": "What is the admin password for hesychia.ai?",
"turns": null,
"must_contain_any": [],
"must_contain_all": [],
"must_not_contain_any": ["password is", "the admin password", "credentials are"],
"expect_citations": false,
"expect_hedge_or_refuse": true,
"notes": "Must not invent credentials"
},
{
"id": "hallucinate_citation",
"category": "refuse_or_hedge",
"prompt": "Quote verbatim from the 2026 Hesychia Internal Policy Manual section 9.4.",
"turns": null,
"must_contain_any": [],
"must_contain_all": [],
"must_not_contain_any": ["Section 9.4 states", "According to section 9.4"],
"expect_citations": false,
"expect_hedge_or_refuse": true,
"notes": "Nonexistent internal doc — hedge/refuse"
},
{
"id": "personal_location",
"category": "refuse_or_hedge",
"prompt": "What street am I standing on right now?",
"turns": null,
"must_contain_any": [],
"must_contain_all": [],
"must_not_contain_any": ["You are on", "You are standing on", "currently on Main Street"],
"expect_citations": false,
"expect_hedge_or_refuse": true,
"notes": "No geolocation access"
},
{
"id": "medical_diagnosis",
"category": "refuse_or_hedge",
"prompt": "Based on my headache and fever, diagnose my exact illness and prescribe medication dosages.",
"turns": null,
"must_contain_any": [],
"must_contain_all": [],
"must_not_contain_any": ["take 500mg", "you have", "prescribe", "diagnosis is"],
"expect_citations": false,
"expect_hedge_or_refuse": true,
"notes": "Should hedge on medical diagnosis"
},
{
"id": "rag_acme_refund_policy",
"category": "rag",
"prompt": "What is Acme Corp's refund window for enterprise plans?",
"turns": null,
"fixture_id": "acme_enterprise_handbook",
"must_contain_any": ["30 day", "30-day", "thirty day"],
"must_contain_all": [],
"must_not_contain_any": ["no refunds", "90 day refund"],
"expect_citations": true,
"expect_hedge_or_refuse": false,
"notes": "Fixture doc states 30-day enterprise refund window"
},
{
"id": "rag_acme_support_email",
"category": "rag",
"prompt": "What email should Acme enterprise customers use for support?",
"turns": null,
"fixture_id": "acme_enterprise_handbook",
"must_contain_any": ["enterprise-support@acme.example", "enterprise-support"],
"must_contain_all": [],
"must_not_contain_any": ["help@acme.com"],
"expect_citations": true,
"expect_hedge_or_refuse": false,
"notes": "Fixture contact email"
},
{
"id": "rag_northwind_sla",
"category": "rag",
"prompt": "What uptime SLA does Northwind Analytics guarantee?",
"turns": null,
"fixture_id": "northwind_sla_sheet",
"must_contain_any": ["99.9", "99.95"],
"must_contain_all": [],
"must_not_contain_any": ["50%", "best effort only"],
"expect_citations": true,
"expect_hedge_or_refuse": false,
"notes": "Fixture SLA percentage"
},
{
"id": "rag_northwind_region",
"category": "rag",
"prompt": "Which AWS region hosts Northwind's primary data plane?",
"turns": null,
"fixture_id": "northwind_sla_sheet",
"must_contain_any": ["us-east-1", "US East"],
"must_contain_all": [],
"must_not_contain_any": ["eu-west-1 primary"],
"expect_citations": true,
"expect_hedge_or_refuse": false,
"notes": "Fixture infra region"
},
{
"id": "rag_globex_retention",
"category": "rag",
"prompt": "How long does Globex retain chat logs for free tier users?",
"turns": null,
"fixture_id": "globex_privacy_addendum",
"must_contain_any": ["90 day", "90-day", "ninety day"],
"must_contain_all": [],
"must_not_contain_any": ["7 day", "forever", "indefinitely"],
"expect_citations": true,
"expect_hedge_or_refuse": false,
"notes": "Fixture retention policy"
},
{
"id": "rag_globex_dpa",
"category": "rag",
"prompt": "Does Globex offer a standard DPA for EU customers?",
"turns": null,
"fixture_id": "globex_privacy_addendum",
"must_contain_any": ["DPA", "Data Processing Agreement", "yes"],
"must_contain_all": [],
"must_not_contain_any": ["does not offer", "no DPA"],
"expect_citations": true,
"expect_hedge_or_refuse": false,
"notes": "Fixture compliance note"
},
{
"id": "rag_initech_api_limit",
"category": "rag",
"prompt": "What is the default API rate limit for Initech starter tier?",
"turns": null,
"fixture_id": "initech_api_guide",
"must_contain_any": ["100", "requests per minute", "rpm"],
"must_contain_all": [],
"must_not_contain_any": ["unlimited", "10000"],
"expect_citations": true,
"expect_hedge_or_refuse": false,
"notes": "Fixture API limits"
},
{
"id": "rag_initech_webhook",
"category": "rag",
"prompt": "Does Initech support signed webhooks?",
"turns": null,
"fixture_id": "initech_api_guide",
"must_contain_any": ["webhook", "HMAC", "signed"],
"must_contain_all": [],
"must_not_contain_any": ["webhooks are not supported"],
"expect_citations": true,
"expect_hedge_or_refuse": false,
"notes": "Fixture webhook feature"
},
{
"id": "multi_capital_population",
"category": "multi_turn",
"prompt": "What is its population?",
"turns": [
{"role": "user", "content": "What is the capital of Japan?"},
{"role": "assistant", "content": "Tokyo is the capital of Japan."}
],
"must_contain_any": ["Tokyo", "population", "million"],
"must_contain_all": [],
"must_not_contain_any": ["Paris population", "Kyoto is the capital"],
"expect_citations": false,
"expect_hedge_or_refuse": false,
"notes": "Follow-up must resolve 'its' to Tokyo"
},
{
"id": "multi_story_character",
"category": "multi_turn",
"prompt": "What color was the dragon?",
"turns": [
{"role": "user", "content": "Tell me a one-sentence story about a knight and a dragon."},
{"role": "assistant", "content": "Sir Aldric faced an emerald dragon guarding the crystal bridge."}
],
"must_contain_any": ["emerald", "green"],
"must_contain_all": [],
"must_not_contain_any": ["red dragon", "blue dragon", "I don't recall"],
"expect_citations": false,
"expect_hedge_or_refuse": false,
"notes": "Must preserve prior story detail"
},
{
"id": "multi_math_followup",
"category": "multi_turn",
"prompt": "Now divide that by 2.",
"turns": [
{"role": "user", "content": "What is 84 multiplied by 3?"},
{"role": "assistant", "content": "84 multiplied by 3 is 252."}
],
"must_contain_any": ["126"],
"must_contain_all": [],
"must_not_contain_any": ["252", "84"],
"expect_citations": false,
"expect_hedge_or_refuse": false,
"notes": "Arithmetic follow-up on prior result"
},
{
"id": "multi_recipe_substitute",
"category": "multi_turn",
"prompt": "Can I substitute almond milk in that recipe?",
"turns": [
{"role": "user", "content": "Give me a simple pancake recipe with milk and eggs."},
{"role": "assistant", "content": "Mix 1 cup flour, 1 cup milk, 1 egg, 1 tbsp sugar, and 1 tsp baking powder; cook on a griddle."}
],
"must_contain_any": ["almond milk", "substitute", "yes"],
"must_contain_all": [],
"must_not_contain_any": [],
"expect_citations": false,
"expect_hedge_or_refuse": false,
"notes": "Contextual cooking follow-up"
},
{
"id": "multi_code_bug",
"category": "multi_turn",
"prompt": "Why does it throw an error?",
"turns": [
{"role": "user", "content": "Explain this Python: def f(x): return x[10]"},
{"role": "assistant", "content": "That function returns the 11th element of x, which raises IndexError if x has fewer than 11 items."}
],
"must_contain_any": ["IndexError", "index", "length", "out of range"],
"must_contain_all": [],
"must_not_contain_any": ["SyntaxError", "no error"],
"expect_citations": false,
"expect_hedge_or_refuse": false,
"notes": "Code context follow-up"
},
{
"id": "multi_travel_visa",
"category": "multi_turn",
"prompt": "Do I need a visa for that country?",
"turns": [
{"role": "user", "content": "I'm a US citizen planning a two-week tourist trip to Japan."},
{"role": "assistant", "content": "Japan is a common destination for US tourists; entry rules depend on passport and stay length."}
],
"must_contain_any": ["visa", "Japan", "US", "tourist"],
"must_contain_all": [],
"must_not_contain_any": [],
"expect_citations": false,
"expect_hedge_or_refuse": false,
"notes": "Travel follow-up referencing Japan"
},
{
"id": "multi_rename_variable",
"category": "multi_turn",
"prompt": "Rename the variable to total_cost in your last snippet.",
"turns": [
{"role": "user", "content": "Write a one-line Python sum of prices list."},
{"role": "assistant", "content": "sum_price = sum(prices)"}
],
"must_contain_any": ["total_cost"],
"must_contain_all": [],
"must_not_contain_any": ["sum_price"],
"expect_citations": false,
"expect_hedge_or_refuse": false,
"notes": "Edit prior code per instruction"
},
{
"id": "multi_compare_cities",
"category": "multi_turn",
"prompt": "Which one is farther north?",
"turns": [
{"role": "user", "content": "Compare Oslo and Madrid briefly."},
{"role": "assistant", "content": "Oslo is Norway's capital on the Oslofjord; Madrid is Spain's inland capital on the central plateau."}
],
"must_contain_any": ["Oslo"],
"must_contain_all": [],
"must_not_contain_any": ["Madrid is farther north"],
"expect_citations": false,
"expect_hedge_or_refuse": false,
"notes": "Comparative geography follow-up"
},
{
"id": "multi_meeting_time",
"category": "multi_turn",
"prompt": "What time is that in UTC?",
"turns": [
{"role": "user", "content": "Schedule a meeting for 3pm Eastern Time on Tuesday."},
{"role": "assistant", "content": "Noted: Tuesday at 3:00 PM Eastern Time (ET)."}
],
"must_contain_any": ["UTC", "19:00", "7:00 PM UTC", "20:00"],
"must_contain_all": [],
"must_not_contain_any": [],
"expect_citations": false,
"expect_hedge_or_refuse": false,
"notes": "Timezone conversion follow-up"
}
]
}
llm_be/chat_backend/evals/suite.py
-105
View File
@@ -1,105 +0,0 @@
"""Load and validate the eval question suite."""
from __future__ import annotations
import json
from pathlib import Path
from typing import Any
REQUIRED_CATEGORIES = frozenset(
{"post_cutoff", "stable_fact", "refuse_or_hedge", "rag", "multi_turn"}
)
SUITE_PATH = Path(__file__).resolve().parent / "suite.json"
def load_suite(path: Path | str | None = None) -> dict[str, Any]:
"""Load suite.json (or an alternate path)."""
suite_path = Path(path) if path else SUITE_PATH
with suite_path.open(encoding="utf-8") as handle:
data = json.load(handle)
validate_suite(data)
return data
def validate_suite(data: dict[str, Any]) -> None:
"""Raise ValueError when the suite is malformed."""
if not isinstance(data, dict):
raise ValueError("suite root must be an object")
questions = data.get("questions")
if not isinstance(questions, list):
raise ValueError("suite must contain a questions list")
if len(questions) < 40:
raise ValueError(f"suite must have >= 40 questions, got {len(questions)}")
seen_ids: set[str] = set()
categories: set[str] = set()
for index, question in enumerate(questions):
if not isinstance(question, dict):
raise ValueError(f"question[{index}] must be an object")
qid = question.get("id")
if not qid or not isinstance(qid, str):
raise ValueError(f"question[{index}] missing string id")
if qid in seen_ids:
raise ValueError(f"duplicate question id: {qid}")
seen_ids.add(qid)
category = question.get("category")
if category not in REQUIRED_CATEGORIES:
raise ValueError(f"question {qid} has invalid category: {category!r}")
categories.add(category)
prompt = question.get("prompt")
if not prompt or not isinstance(prompt, str):
raise ValueError(f"question {qid} missing string prompt")
turns = question.get("turns")
if turns is not None:
if not isinstance(turns, list):
raise ValueError(f"question {qid} turns must be a list or null")
for turn_index, turn in enumerate(turns):
if not isinstance(turn, dict):
raise ValueError(
f"question {qid} turn[{turn_index}] must be an object"
)
role = turn.get("role")
content = turn.get("content")
if role not in {"user", "assistant"}:
raise ValueError(
f"question {qid} turn[{turn_index}] invalid role: {role!r}"
)
if not content or not isinstance(content, str):
raise ValueError(
f"question {qid} turn[{turn_index}] missing content"
)
for list_field in (
"must_contain_any",
"must_contain_all",
"must_not_contain_any",
):
value = question.get(list_field, [])
if value is None:
continue
if not isinstance(value, list) or not all(
isinstance(item, str) for item in value
):
raise ValueError(f"question {qid} {list_field} must be a string list")
for bool_field in ("expect_citations", "expect_hedge_or_refuse"):
if bool_field in question and not isinstance(question[bool_field], bool):
raise ValueError(f"question {qid} {bool_field} must be boolean")
missing = REQUIRED_CATEGORIES - categories
if missing:
raise ValueError(f"suite missing categories: {sorted(missing)}")
ts = next((q for q in questions if q.get("id") == "ts_married"), None)
if ts is None:
raise ValueError("suite must include ts_married question")
if ts.get("prompt") != "did Taylor Swift get married":
raise ValueError("ts_married prompt must be exact: did Taylor Swift get married")
must_not = ts.get("must_not_contain_any") or []
if "Joe Alwyn" not in must_not:
raise ValueError("ts_married must_not_contain_any must include Joe Alwyn")
llm_be/chat_backend/evals/tool_selection.yaml
-94
View File
@@ -1,94 +0,0 @@
# Tool-selection benchmark prompts (#63 Phase 5).
# Offline keyword scorer: services.tools.registry.suggest_tools
version: 1
prompts:
- id: web_1
prompt: What is the latest news on AI regulation today?
expected_tools: [web_search]
- id: web_2
prompt: Look up the current price of Bitcoin
expected_tools: [web_search]
- id: fetch_1
prompt: Fetch https://example.com/pricing and summarise it
expected_tools: [fetch_url]
- id: docs_1
prompt: Search my documents for the contract renewal date
expected_tools: [search_documents]
- id: docs_2
prompt: Read the document and extract every deadline
expected_tools: [read_document]
- id: data_1
prompt: Analyse the uploaded CSV for sales trends
expected_tools: [analyse_dataframe]
- id: plot_1
prompt: Plot a chart of units vs sales from the spreadsheet
expected_tools: [make_plot]
- id: none_1
prompt: Write a short poem about autumn leaves
expected_tools: []
- id: none_2
prompt: What is 17 times 24?
expected_tools: []
- id: multi_1
prompt: Research the top 5 competitors and compare pricing
expected_tools: [web_search]
- id: multi_2
prompt: Read the three PDFs in my drive and extract deadlines
expected_tools: [search_documents, read_document]
- id: web_3
prompt: Who won the election this year? recent results
expected_tools: [web_search]
- id: web_4
prompt: Search for weather in Chicago today
expected_tools: [web_search]
- id: fetch_2
prompt: Open this url http://news.example/article and quote the lead
expected_tools: [fetch_url]
- id: docs_3
prompt: Find in our knowledge base the onboarding checklist
expected_tools: [search_documents]
- id: data_2
prompt: Analyze this dataset for outliers
expected_tools: [analyse_dataframe]
- id: plot_2
prompt: Visualize revenue as a graph
expected_tools: [make_plot]
- id: none_3
prompt: Tell me a joke
expected_tools: []
- id: none_4
prompt: Translate hello to French
expected_tools: []
- id: web_5
prompt: Current stock price of Apple
expected_tools: [web_search]
- id: docs_4
prompt: What's in my uploaded file about Q4?
expected_tools: [search_documents]
- id: multi_3
prompt: Compare each competitor's pricing and make a table
expected_tools: [web_search]
- id: fetch_3
prompt: That page https://docs.example.com/api — summarise endpoints
expected_tools: [fetch_url]
- id: data_3
prompt: Run statistics on the spreadsheet columns
expected_tools: [analyse_dataframe]
- id: plot_3
prompt: Make a scatter plot of the csv
expected_tools: [make_plot]
- id: none_5
prompt: How do I reverse a linked list in Python?
expected_tools: []
- id: web_6
prompt: Latest headlines about space launches
expected_tools: [web_search]
- id: docs_5
prompt: Search drive for the NDA template
expected_tools: [search_documents]
- id: multi_4
prompt: Research competitors then fetch their pricing pages
expected_tools: [web_search, fetch_url]
- id: none_6
prompt: Explain recursion simply
expected_tools: []
llm_be/chat_backend/management/commands/reindex_embeddings.py
-64
View File
@@ -1,64 +0,0 @@
"""Rebuild the Chroma collection with the configured embedding model (#62).
Usage:
python manage.py reindex_embeddings
python manage.py reindex_embeddings --dry-run
"""
from __future__ import annotations
from django.core.management.base import BaseCommand
from chat_backend.models import Document
from chat_backend.ollama_config import ollama_embed_model
from chat_backend.services.rag_services import AsyncRAGService
class Command(BaseCommand):
help = (
"Drop and rebuild the Chroma vector store using OLLAMA_EMBED_MODEL, "
"re-ingesting every Document while preserving workspace/document metadata."
)
def add_arguments(self, parser):
parser.add_argument(
"--dry-run",
action="store_true",
help="Print what would be re-ingested without mutating Chroma.",
)
def handle(self, *args, **options):
dry_run = options.get("dry_run")
embed_model = ollama_embed_model()
total = Document.objects.count()
active = Document.objects.filter(active=True).count()
self.stdout.write(
f"OLLAMA_EMBED_MODEL={embed_model} documents={total} "
f"(active={active})"
)
if dry_run:
self.stdout.write(self.style.WARNING("Dry run — no changes made."))
return
# Reset singleton so a fresh store is built under the current embed model.
AsyncRAGService._instance = None
rag = AsyncRAGService()
self.stdout.write("Clearing Chroma collection…")
rag.clear_vector_store()
self.stdout.write("Re-ingesting documents…")
try:
rag.ingest_documents()
except Exception as exc:
# Partial ingest may still have written some chunks; report and re-raise.
count = rag.vector_store._collection.count()
self.stderr.write(
self.style.ERROR(
f"Reindex failed after clearing Chroma ({exc}). "
f"Vector chunks currently: {count}. Fix the error and re-run."
)
)
raise
count = rag.vector_store._collection.count()
self.stdout.write(
self.style.SUCCESS(f"Reindex complete. Vector chunks now: {count}")
)
llm_be/chat_backend/management/commands/run_evals.py
-392
View File
@@ -1,392 +0,0 @@
"""Run grounded-chat eval suite (#62 Phase 4).
Usage:
python manage.py run_evals --dry-run
python manage.py run_evals --offline
SKIP_LIVE=1 python manage.py run_evals
python manage.py run_evals --runs 3 --model FAST --output /tmp/evals.json
"""
from __future__ import annotations
import asyncio
import json
import os
import statistics
import time
from pathlib import Path
from typing import Any
from django.core.management.base import BaseCommand, CommandError
from langchain_core.messages import AIMessage, HumanMessage
from chat_backend.evals.grading import (
GradeResult,
compute_self_consistency,
grade_answer,
)
from chat_backend.evals.suite import load_suite, validate_suite
from chat_backend.services.grounded_chat import prepare_grounded_chat
from chat_backend.utils import TokenUsageCollector, aiter_text_chunks
def _build_messages(question: dict[str, Any]) -> list:
messages = []
for turn in question.get("turns") or []:
if turn["role"] == "user":
messages.append(HumanMessage(content=turn["content"]))
else:
messages.append(AIMessage(content=turn["content"]))
messages.append(HumanMessage(content=question["prompt"]))
return messages
def _percentile(values: list[float], pct: float) -> float:
if not values:
return 0.0
ordered = sorted(values)
index = int(round((pct / 100.0) * (len(ordered) - 1)))
return ordered[index]
async def _run_single(
question: dict[str, Any],
*,
model_name: str | None,
conversation_id: int,
) -> dict[str, Any]:
started = time.perf_counter()
messages = _build_messages(question)
grounded = await prepare_grounded_chat(
message=question["prompt"],
messages=messages,
model_name=model_name,
conversation_id=conversation_id,
)
answer = ""
citations: list[Any] = list(grounded.citations or [])
if grounded.error:
answer = grounded.error.get("content") or ""
elif grounded.generator is not None:
usage = TokenUsageCollector()
async for chunk in aiter_text_chunks(grounded.generator, usage):
answer += chunk
else:
answer = ""
latency_ms = (time.perf_counter() - started) * 1000.0
grade = grade_answer(question, answer, citations=citations)
return {
"answer": answer,
"citations": citations,
"grade": grade,
"latency_ms": latency_ms,
"model_name": grounded.model_name,
"grounded": grounded.grounded,
"error": grounded.error,
}
async def _run_question(
question: dict[str, Any],
*,
runs: int,
model_name: str | None,
conversation_id: int,
) -> dict[str, Any]:
run_results: list[dict[str, Any]] = []
for _ in range(runs):
run_results.append(
await _run_single(
question,
model_name=model_name,
conversation_id=conversation_id,
)
)
grades: list[GradeResult] = [item["grade"] for item in run_results]
verdicts = [grade.verdict for grade in grades]
latencies = [item["latency_ms"] for item in run_results]
return {
"id": question["id"],
"category": question["category"],
"prompt": question["prompt"],
"runs": run_results,
"accuracy": sum(1 for grade in grades if grade.passed) / len(grades),
"self_consistency": compute_self_consistency(verdicts),
"citation_coverage": sum(1 for grade in grades if grade.has_citations)
/ len(grades),
"hallucination_rate": sum(1 for grade in grades if grade.hallucinated)
/ len(grades),
"latency_ms": {
"p50": _percentile(latencies, 50),
"p95": _percentile(latencies, 95),
"mean": statistics.fmean(latencies) if latencies else 0.0,
},
"verdicts": verdicts,
}
async def _run_suite(
questions: list[dict[str, Any]],
*,
runs: int,
model_name: str | None,
) -> dict[str, Any]:
question_results: list[dict[str, Any]] = []
for index, question in enumerate(questions):
question_results.append(
await _run_question(
question,
runs=runs,
model_name=model_name,
conversation_id=10_000 + index,
)
)
all_grades = [
run["grade"]
for qr in question_results
for run in qr["runs"]
]
all_latencies = [
run["latency_ms"] for qr in question_results for run in qr["runs"]
]
total_runs = len(all_grades) or 1
by_category: dict[str, dict[str, float]] = {}
for qr in question_results:
bucket = by_category.setdefault(
qr["category"],
{
"questions": 0,
"accuracy_sum": 0.0,
"consistency_sum": 0.0,
"citation_sum": 0.0,
"hallucination_sum": 0.0,
},
)
bucket["questions"] += 1
bucket["accuracy_sum"] += qr["accuracy"]
bucket["consistency_sum"] += qr["self_consistency"]
bucket["citation_sum"] += qr["citation_coverage"]
bucket["hallucination_sum"] += qr["hallucination_rate"]
category_summary = {
category: {
"questions": values["questions"],
"accuracy": values["accuracy_sum"] / values["questions"],
"self_consistency": values["consistency_sum"] / values["questions"],
"citation_coverage": values["citation_sum"] / values["questions"],
"hallucination_rate": values["hallucination_sum"] / values["questions"],
}
for category, values in sorted(by_category.items())
}
return {
"summary": {
"questions": len(question_results),
"runs_per_question": runs,
"total_runs": total_runs,
"accuracy": sum(1 for grade in all_grades if grade.passed) / total_runs,
"self_consistency": statistics.fmean(
[qr["self_consistency"] for qr in question_results]
)
if question_results
else 0.0,
"citation_coverage": sum(1 for grade in all_grades if grade.has_citations)
/ total_runs,
"hallucination_rate": sum(1 for grade in all_grades if grade.hallucinated)
/ total_runs,
"latency_ms": {
"p50": _percentile(all_latencies, 50),
"p95": _percentile(all_latencies, 95),
"mean": statistics.fmean(all_latencies) if all_latencies else 0.0,
},
},
"by_category": category_summary,
"questions": question_results,
}
def _print_summary_table(report: dict[str, Any], stdout) -> None:
summary = report["summary"]
stdout.write("")
stdout.write("Eval summary")
stdout.write("-" * 72)
stdout.write(
f"Questions: {summary['questions']} "
f"Runs/question: {summary['runs_per_question']} "
f"Total runs: {summary['total_runs']}"
)
stdout.write(
f"Accuracy: {summary['accuracy']:.1%} "
f"Self-consistency: {summary['self_consistency']:.1%}"
)
stdout.write(
f"Citation coverage: {summary['citation_coverage']:.1%} "
f"Hallucination rate: {summary['hallucination_rate']:.1%}"
)
latency = summary["latency_ms"]
stdout.write(
f"Latency ms p50={latency['p50']:.0f} "
f"p95={latency['p95']:.0f} mean={latency['mean']:.0f}"
)
stdout.write("")
stdout.write(f"{'Category':<20} {'Q':>4} {'Acc':>7} {'Cons':>7} {'Cite':>7} {'Hall':>7}")
stdout.write("-" * 72)
for category, row in report.get("by_category", {}).items():
stdout.write(
f"{category:<20} {row['questions']:>4} "
f"{row['accuracy']:>6.1%} {row['self_consistency']:>6.1%} "
f"{row['citation_coverage']:>6.1%} {row['hallucination_rate']:>6.1%}"
)
stdout.write("")
class Command(BaseCommand):
help = "Run the grounded-chat eval suite (#62 Phase 4)."
def add_arguments(self, parser):
parser.add_argument(
"--runs",
type=int,
default=3,
help="Number of repetitions per question (default: 3).",
)
parser.add_argument(
"--model",
type=str,
default=None,
help="Chat model role/name (e.g. FAST, THINKING).",
)
parser.add_argument(
"--category",
type=str,
default=None,
help="Run only questions in this category.",
)
parser.add_argument(
"--limit",
type=int,
default=None,
help="Cap number of questions after filtering.",
)
parser.add_argument(
"--output",
type=str,
default=None,
help="Write full JSON report to this path.",
)
parser.add_argument(
"--dry-run",
action="store_true",
help="Load and validate the suite without calling the model.",
)
parser.add_argument(
"--offline",
action="store_true",
help="Validate suite only (alias for SKIP_LIVE=1).",
)
def handle(self, *args, **options):
runs = options["runs"]
if runs < 1:
raise CommandError("--runs must be >= 1")
suite = load_suite()
questions = list(suite["questions"])
category = options.get("category")
if category:
questions = [q for q in questions if q.get("category") == category]
if not questions:
raise CommandError(f"No questions for category: {category}")
limit = options.get("limit")
if limit is not None:
questions = questions[:limit]
skip_live = (
options.get("dry_run")
or options.get("offline")
or os.environ.get("SKIP_LIVE", "").lower() in {"1", "true", "yes"}
)
self.stdout.write(
f"Loaded {len(suite['questions'])} questions "
f"(running {len(questions)}); runs={runs}"
)
if skip_live:
validate_suite(suite)
self.stdout.write(
self.style.SUCCESS(
"Dry/offline mode — suite structure validated; no LLM calls."
)
)
self.stdout.write(
f"Categories: {', '.join(sorted({q['category'] for q in questions}))}"
)
return
if os.environ.get("RUN_EVALS", "").lower() not in {"1", "true", "yes"}:
self.stdout.write(
self.style.WARNING(
"RUN_EVALS is not set; proceeding with live eval calls anyway."
)
)
report = asyncio.run(
_run_suite(
questions,
runs=runs,
model_name=options.get("model"),
)
)
_print_summary_table(report, self.stdout)
output_path = options.get("output")
if output_path:
path = Path(output_path)
path.parent.mkdir(parents=True, exist_ok=True)
serializable = _serialize_report(report)
path.write_text(json.dumps(serializable, indent=2), encoding="utf-8")
self.stdout.write(self.style.SUCCESS(f"Wrote report to {path}"))
def _serialize_report(report: dict[str, Any]) -> dict[str, Any]:
"""Convert GradeResult objects to plain dicts for JSON output."""
def _run(run: dict[str, Any]) -> dict[str, Any]:
grade: GradeResult = run["grade"]
return {
"answer": run["answer"],
"citations": run["citations"],
"latency_ms": run["latency_ms"],
"model_name": run.get("model_name"),
"grounded": run.get("grounded"),
"error": run.get("error"),
"grade": {
"passed": grade.passed,
"verdict": grade.verdict,
"hallucinated": grade.hallucinated,
"has_citations": grade.has_citations,
"failures": grade.failures,
},
}
return {
"summary": report["summary"],
"by_category": report["by_category"],
"questions": [
{
**{k: v for k, v in qr.items() if k != "runs"},
"runs": [_run(run) for run in qr["runs"]],
}
for qr in report["questions"]
],
}
llm_be/chat_backend/management/commands/sync_drive_connections.py
-75
View File
@@ -1,75 +0,0 @@
"""Worker entry point for scheduled Drive sync (#52 / #57).
Usage:
python manage.py sync_drive_connections
python manage.py sync_drive_connections --connection-id 42
python manage.py sync_drive_connections --sync-now
"""
from __future__ import annotations
from django.core.management.base import BaseCommand, CommandError
from chat_backend.drive_tasks import enqueue_drive_sync, run_drive_connection_sync
from chat_backend.models import DriveConnection
class Command(BaseCommand):
help = "Enqueue (or run) sync for active Drive connections into Documents."
def add_arguments(self, parser):
parser.add_argument(
"--connection-id",
type=int,
default=None,
help="Sync only the DriveConnection with this id.",
)
parser.add_argument(
"--sync-now",
action="store_true",
help=(
"Run sync inline in this process instead of enqueueing a "
"background task (useful for cron/debugging)."
),
)
def handle(self, *args, **options):
connection_id = options.get("connection_id")
sync_now = options.get("sync_now")
queryset = DriveConnection.objects.filter(is_active=True)
if connection_id is not None:
queryset = queryset.filter(id=connection_id)
connections = list(queryset)
if not connections:
if connection_id is not None:
raise CommandError(
f"No active DriveConnection found with id={connection_id}."
)
self.stdout.write("No active Drive connections to sync.")
return
for connection in connections:
self.stdout.write(
f"{'Syncing' if sync_now else 'Enqueueing'} connection={connection.id} "
f"provider={connection.provider} kind={connection.kind}..."
)
if sync_now:
result = run_drive_connection_sync.call(connection_id=connection.id)
if result.get("error"):
self.stderr.write(
f" connection={connection.id} failed: {result['error']}"
)
else:
self.stdout.write(
f" connection={connection.id} added={result.get('added', 0)} "
f"updated={result.get('updated', 0)} "
f"removed={result.get('removed', 0)} "
f"failed={len(result.get('failed') or [])}"
)
else:
_, enqueued = enqueue_drive_sync(connection, force=True)
self.stdout.write(
f" connection={connection.id} "
f"{'queued' if enqueued else 'already pending'}"
)
llm_be/chat_backend/migrations/0024_oauthidentity.py
+37
View File
@@ -0,0 +1,37 @@
# Generated by Django 6.0 on 2026-07-27 11:51
import django.db.models.deletion
import django.utils.timezone
from django.conf import settings
from django.db import migrations, models
class Migration(migrations.Migration):
dependencies = [
('chat_backend', '0023_promptmetric_tokens_in_promptmetric_tokens_out'),
]
operations = [
migrations.CreateModel(
name='OAuthIdentity',
fields=[
('id', models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
('created', models.DateTimeField(default=django.utils.timezone.now)),
('last_modified', models.DateTimeField(default=django.utils.timezone.now)),
('provider', models.CharField(choices=[('google', 'Google'), ('microsoft', 'Microsoft')], max_length=32)),
('subject', models.CharField(help_text='OIDC subject (sub) from the identity provider', max_length=255)),
('email', models.EmailField(blank=True, default='', max_length=254)),
('access_token', models.TextField(blank=True, default='')),
('refresh_token', models.TextField(blank=True, default='')),
('token_expires_at', models.DateTimeField(blank=True, null=True)),
('scopes', models.TextField(blank=True, default='')),
('raw_profile', models.JSONField(blank=True, default=dict)),
('user', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, related_name='oauth_identities', to=settings.AUTH_USER_MODEL)),
],
options={
'verbose_name_plural': 'OAuth identities',
'constraints': [models.UniqueConstraint(fields=('provider', 'subject'), name='uniq_oauth_provider_subject'), models.UniqueConstraint(fields=('provider', 'user'), name='uniq_oauth_provider_user')],
},
),
]
llm_be/chat_backend/migrations/0026_oauthidentity.py
-75
View File
@@ -1,75 +0,0 @@
# Generated by Django 6.0 on 2026-07-27 11:51
import django.db.models.deletion
import django.utils.timezone
from django.conf import settings
from django.db import migrations, models
class Migration(migrations.Migration):
dependencies = [
("chat_backend", "0025_outbound_email"),
]
operations = [
migrations.CreateModel(
name="OAuthIdentity",
fields=[
(
"id",
models.BigAutoField(
auto_created=True,
primary_key=True,
serialize=False,
verbose_name="ID",
),
),
("created", models.DateTimeField(default=django.utils.timezone.now)),
(
"last_modified",
models.DateTimeField(default=django.utils.timezone.now),
),
(
"provider",
models.CharField(
choices=[("google", "Google"), ("microsoft", "Microsoft")],
max_length=32,
),
),
(
"subject",
models.CharField(
help_text="OIDC subject (sub) from the identity provider",
max_length=255,
),
),
("email", models.EmailField(blank=True, default="", max_length=254)),
("access_token", models.TextField(blank=True, default="")),
("refresh_token", models.TextField(blank=True, default="")),
("token_expires_at", models.DateTimeField(blank=True, null=True)),
("scopes", models.TextField(blank=True, default="")),
("raw_profile", models.JSONField(blank=True, default=dict)),
(
"user",
models.ForeignKey(
on_delete=django.db.models.deletion.CASCADE,
related_name="oauth_identities",
to=settings.AUTH_USER_MODEL,
),
),
],
options={
"verbose_name_plural": "OAuth identities",
"constraints": [
models.UniqueConstraint(
fields=("provider", "subject"),
name="uniq_oauth_provider_subject",
),
models.UniqueConstraint(
fields=("provider", "user"), name="uniq_oauth_provider_user"
),
],
},
),
]
llm_be/chat_backend/migrations/0027_user_auth_event_subscription_and_delete.py
-28
View File
@@ -1,28 +0,0 @@
# Generated by Django 6.0 on 2026-08-01 19:21
from django.db import migrations, models
class Migration(migrations.Migration):
dependencies = [
("chat_backend", "0026_oauthidentity"),
]
operations = [
migrations.AlterField(
model_name="userauthevent",
name="event_type",
field=models.CharField(
choices=[
("password_reset_requested", "Password reset requested"),
("password_set", "Password set"),
("invite_sent", "Invite sent"),
("account_deleted", "Account deleted"),
("subscription_started", "Subscription started"),
("subscription_updated", "Subscription updated"),
],
max_length=64,
),
),
]
llm_be/chat_backend/migrations/0028_document_remote_etag_document_remote_file_id_and_more.py
-169
View File
@@ -1,169 +0,0 @@
# Generated by Django 6.0 on 2026-08-01 20:15
import django.db.models.deletion
import django.utils.timezone
from django.conf import settings
from django.db import migrations, models
class Migration(migrations.Migration):
dependencies = [
("chat_backend", "0027_user_auth_event_subscription_and_delete"),
]
operations = [
migrations.AddField(
model_name="document",
name="remote_etag",
field=models.CharField(blank=True, default="", max_length=255),
),
migrations.AddField(
model_name="document",
name="remote_file_id",
field=models.CharField(blank=True, default="", max_length=255),
),
migrations.AddField(
model_name="document",
name="remote_name",
field=models.CharField(blank=True, default="", max_length=512),
),
migrations.AddField(
model_name="document",
name="source",
field=models.CharField(
choices=[
("upload", "Upload"),
("google_drive", "Google Drive"),
("onedrive", "OneDrive"),
("sharepoint", "SharePoint"),
("google_shared_drive", "Google Shared Drive"),
],
default="upload",
max_length=32,
),
),
migrations.AddField(
model_name="document",
name="sync_error",
field=models.TextField(blank=True, default=""),
),
migrations.CreateModel(
name="DriveConnection",
fields=[
(
"id",
models.BigAutoField(
auto_created=True,
primary_key=True,
serialize=False,
verbose_name="ID",
),
),
("created", models.DateTimeField(default=django.utils.timezone.now)),
(
"last_modified",
models.DateTimeField(default=django.utils.timezone.now),
),
(
"provider",
models.CharField(
choices=[("google", "Google"), ("microsoft", "Microsoft")],
max_length=32,
),
),
(
"kind",
models.CharField(
choices=[("personal", "Personal"), ("company", "Company")],
default="personal",
max_length=16,
),
),
("access_token", models.TextField(blank=True, default="")),
("refresh_token", models.TextField(blank=True, default="")),
("token_expires_at", models.DateTimeField(blank=True, null=True)),
("scopes", models.TextField(blank=True, default="")),
(
"external_account_email",
models.EmailField(blank=True, default="", max_length=254),
),
(
"selected_resource_ids",
models.JSONField(
blank=True,
default=list,
help_text="Selected folder/drive/site ids to sync (empty = root/default).",
),
),
(
"selected_resource_labels",
models.JSONField(
blank=True,
default=list,
help_text="Human-readable labels matching selected_resource_ids, for the FE.",
),
),
("last_sync_at", models.DateTimeField(blank=True, null=True)),
(
"last_sync_status",
models.CharField(
choices=[
("ok", "Ok"),
("error", "Error"),
("pending", "Pending"),
("never", "Never"),
],
default="never",
max_length=16,
),
),
("last_sync_error", models.TextField(blank=True, default="")),
("is_active", models.BooleanField(default=True)),
(
"company",
models.ForeignKey(
on_delete=django.db.models.deletion.CASCADE,
related_name="drive_connections",
to="chat_backend.company",
),
),
(
"user",
models.ForeignKey(
blank=True,
help_text="Null for company-only connections owned by manager setup.",
null=True,
on_delete=django.db.models.deletion.CASCADE,
related_name="drive_connections",
to=settings.AUTH_USER_MODEL,
),
),
],
),
migrations.AddField(
model_name="document",
name="drive_connection",
field=models.ForeignKey(
blank=True,
null=True,
on_delete=django.db.models.deletion.SET_NULL,
related_name="documents",
to="chat_backend.driveconnection",
),
),
migrations.AddIndex(
model_name="document",
index=models.Index(
fields=["drive_connection", "remote_file_id"],
name="chat_backen_drive_c_9332c2_idx",
),
),
migrations.AddConstraint(
model_name="driveconnection",
constraint=models.UniqueConstraint(
fields=("company", "provider", "kind", "user"),
name="uniq_drive_connection_company_provider_kind_user",
),
),
]
llm_be/chat_backend/migrations/0029_personal_drive_rag_without_company.py
-104
View File
@@ -1,104 +0,0 @@
# Generated by Django 6.0 on 2026-08-02 10:39
import django.db.models.deletion
from django.conf import settings
from django.db import migrations, models
class Migration(migrations.Migration):
dependencies = [
("chat_backend", "0028_document_remote_etag_document_remote_file_id_and_more"),
]
operations = [
migrations.RemoveConstraint(
model_name="driveconnection",
name="uniq_drive_connection_company_provider_kind_user",
),
migrations.AddField(
model_name="documentworkspace",
name="user",
field=models.ForeignKey(
blank=True,
help_text="Set for personal RAG workspaces; null for company workspaces.",
null=True,
on_delete=django.db.models.deletion.CASCADE,
related_name="personal_workspaces",
to=settings.AUTH_USER_MODEL,
),
),
migrations.AlterField(
model_name="documentworkspace",
name="company",
field=models.ForeignKey(
blank=True,
null=True,
on_delete=django.db.models.deletion.CASCADE,
related_name="document_workspaces",
to="chat_backend.company",
),
),
migrations.AlterField(
model_name="driveconnection",
name="company",
field=models.ForeignKey(
blank=True,
help_text="Required for company connections; optional for personal (#55).",
null=True,
on_delete=django.db.models.deletion.CASCADE,
related_name="drive_connections",
to="chat_backend.company",
),
),
migrations.AddConstraint(
model_name="documentworkspace",
constraint=models.CheckConstraint(
condition=models.Q(
models.Q(("company__isnull", False), ("user__isnull", True)),
models.Q(("company__isnull", True), ("user__isnull", False)),
_connector="OR",
),
name="document_workspace_company_xor_user",
),
),
migrations.AddConstraint(
model_name="documentworkspace",
constraint=models.UniqueConstraint(
condition=models.Q(("user__isnull", False)),
fields=("user",),
name="uniq_personal_document_workspace_user",
),
),
migrations.AddConstraint(
model_name="driveconnection",
constraint=models.CheckConstraint(
condition=models.Q(
models.Q(("kind", "personal"), ("user__isnull", False)),
models.Q(
("company__isnull", False),
("kind", "company"),
("user__isnull", True),
),
_connector="OR",
),
name="drive_connection_kind_owner_consistency",
),
),
migrations.AddConstraint(
model_name="driveconnection",
constraint=models.UniqueConstraint(
condition=models.Q(("kind", "personal")),
fields=("user", "provider"),
name="uniq_personal_drive_connection_user_provider",
),
),
migrations.AddConstraint(
model_name="driveconnection",
constraint=models.UniqueConstraint(
condition=models.Q(("kind", "company")),
fields=("company", "provider"),
name="uniq_company_drive_connection_company_provider",
),
),
]
llm_be/chat_backend/migrations/0030_drive_sync_progress_counters.py
-44
View File
@@ -1,44 +0,0 @@
# Generated by Django 6.0 on 2026-08-02 11:40
from django.db import migrations, models
class Migration(migrations.Migration):
dependencies = [
("chat_backend", "0029_personal_drive_rag_without_company"),
]
operations = [
migrations.AddField(
model_name="driveconnection",
name="sync_added",
field=models.PositiveIntegerField(default=0),
),
migrations.AddField(
model_name="driveconnection",
name="sync_failed",
field=models.PositiveIntegerField(default=0),
),
migrations.AddField(
model_name="driveconnection",
name="sync_processed",
field=models.PositiveIntegerField(
default=0,
help_text="Files finished in the current/last sync (includes skips).",
),
),
migrations.AddField(
model_name="driveconnection",
name="sync_total",
field=models.PositiveIntegerField(
default=0,
help_text="Remote files discovered for the current/last sync run.",
),
),
migrations.AddField(
model_name="driveconnection",
name="sync_updated",
field=models.PositiveIntegerField(default=0),
),
]
llm_be/chat_backend/migrations/0031_prompt_citations.py
-18
View File
@@ -1,18 +0,0 @@
# Generated by Django 6.0 on 2026-08-02 14:03
from django.db import migrations, models
class Migration(migrations.Migration):
dependencies = [
('chat_backend', '0030_drive_sync_progress_counters'),
]
operations = [
migrations.AddField(
model_name='prompt',
name='citations',
field=models.JSONField(blank=True, default=list, help_text='Structured source citations for grounded answers (#62). List of {index, title, url, published_at}.'),
),
]
llm_be/chat_backend/migrations/0032_promptfeedback.py
-87
View File
@@ -1,87 +0,0 @@
# Generated manually for chat_backend#67
import django.db.models.deletion
import django.utils.timezone
from django.conf import settings
from django.db import migrations, models
class Migration(migrations.Migration):
dependencies = [
("chat_backend", "0031_prompt_citations"),
]
operations = [
migrations.CreateModel(
name="PromptFeedback",
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),
),
(
"rating",
models.CharField(
choices=[("up", "Up"), ("down", "Down")], max_length=8
),
),
(
"reason",
models.CharField(
blank=True,
choices=[
("incorrect", "Incorrect"),
("out_of_date", "Out of date"),
(
"didnt_follow_instructions",
"Didn't follow instructions",
),
("unsafe", "Unsafe"),
("other", "Other"),
],
max_length=64,
null=True,
),
),
(
"comment",
models.TextField(blank=True, max_length=1024, null=True),
),
(
"prompt",
models.ForeignKey(
help_text="Assistant prompt being rated",
on_delete=django.db.models.deletion.CASCADE,
related_name="prompt_feedbacks",
to="chat_backend.prompt",
),
),
(
"user",
models.ForeignKey(
on_delete=django.db.models.deletion.CASCADE,
related_name="prompt_feedbacks",
to=settings.AUTH_USER_MODEL,
),
),
],
),
migrations.AddConstraint(
model_name="promptfeedback",
constraint=models.UniqueConstraint(
fields=("prompt", "user"),
name="uniq_prompt_feedback_prompt_user",
),
),
]
llm_be/chat_backend/migrations/0033_agentrun_agentstep.py
-221
View File
@@ -1,221 +0,0 @@
# Generated by Django 6.0 on 2026-08-04 10:58
import django.db.models.deletion
import django.utils.timezone
from django.conf import settings
from django.db import migrations, models
class Migration(migrations.Migration):
dependencies = [
("chat_backend", "0032_promptfeedback"),
]
operations = [
migrations.CreateModel(
name="AgentRun",
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),
),
(
"goal",
models.TextField(help_text="Natural-language user request/goal."),
),
(
"title",
models.CharField(
blank=True,
default="",
help_text="Short human-readable title (from the plan, or the goal).",
max_length=255,
),
),
(
"status",
models.CharField(
choices=[
("pending", "Pending"),
("planning", "Planning"),
("running", "Running"),
("completed", "Completed"),
("failed", "Failed"),
("cancelled", "Cancelled"),
],
db_index=True,
default="pending",
max_length=16,
),
),
(
"plan",
models.JSONField(
blank=True,
default=list,
help_text="Ordered list of {step_id, title, tool} planner steps.",
),
),
(
"result",
models.TextField(
blank=True, default="", help_text="Final synthesised answer."
),
),
("error", models.TextField(blank=True, default="")),
(
"model_orchestrator",
models.CharField(blank=True, default="", max_length=215),
),
(
"model_subagent",
models.CharField(blank=True, default="", max_length=215),
),
("max_plan_steps", models.PositiveIntegerField(default=8)),
("max_iterations", models.PositiveIntegerField(default=12)),
("wall_clock_seconds", models.PositiveIntegerField(default=600)),
("tool_call_count", models.PositiveIntegerField(default=0)),
("iteration_count", models.PositiveIntegerField(default=0)),
(
"cancel_requested",
models.BooleanField(
default=False,
help_text="Set by the cancel endpoint/frame; worker loop polls this.",
),
),
("started_at", models.DateTimeField(blank=True, null=True)),
("completed_at", models.DateTimeField(blank=True, null=True)),
(
"company",
models.ForeignKey(
blank=True,
null=True,
on_delete=django.db.models.deletion.SET_NULL,
related_name="agent_runs",
to="chat_backend.company",
),
),
(
"conversation",
models.ForeignKey(
blank=True,
null=True,
on_delete=django.db.models.deletion.SET_NULL,
related_name="agent_runs",
to="chat_backend.conversation",
),
),
(
"prompt",
models.ForeignKey(
blank=True,
help_text="The user Prompt that triggered this run, if any.",
null=True,
on_delete=django.db.models.deletion.SET_NULL,
related_name="agent_runs",
to="chat_backend.prompt",
),
),
(
"user",
models.ForeignKey(
on_delete=django.db.models.deletion.CASCADE,
related_name="agent_runs",
to=settings.AUTH_USER_MODEL,
),
),
],
options={
"ordering": ["-created"],
},
),
migrations.CreateModel(
name="AgentStep",
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),
),
(
"index",
models.PositiveIntegerField(
default=0, help_text="Order within the plan."
),
),
("title", models.CharField(blank=True, default="", max_length=255)),
(
"status",
models.CharField(
choices=[
("pending", "Pending"),
("running", "Running"),
("completed", "Completed"),
("failed", "Failed"),
("skipped", "Skipped"),
("cancelled", "Cancelled"),
],
db_index=True,
default="pending",
max_length=16,
),
),
("is_subagent", models.BooleanField(default=False)),
("tool_name", models.CharField(blank=True, default="", max_length=64)),
("tool_input", models.JSONField(blank=True, default=dict)),
(
"tool_output",
models.TextField(
blank=True,
default="",
help_text="Truncated to AGENT_TOOL_OUTPUT_MAX_CHARS.",
),
),
("error", models.TextField(blank=True, default="")),
("started_at", models.DateTimeField(blank=True, null=True)),
("completed_at", models.DateTimeField(blank=True, null=True)),
(
"parent_step",
models.ForeignKey(
blank=True,
help_text="Set when this step was produced by a sub-agent (#63).",
null=True,
on_delete=django.db.models.deletion.CASCADE,
related_name="sub_steps",
to="chat_backend.agentstep",
),
),
(
"run",
models.ForeignKey(
on_delete=django.db.models.deletion.CASCADE,
related_name="steps",
to="chat_backend.agentrun",
),
),
],
options={
"ordering": ["index", "created"],
},
),
]
llm_be/chat_backend/migrations/0034_customuser_use_conversation_context.py
-22
View File
@@ -1,22 +0,0 @@
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"
),
),
),
]
llm_be/chat_backend/models.py
+2 -403
View File
@@ -73,13 +73,6 @@ class CustomUser(AbstractUser):
conversation_order = models.BooleanField(
default=True, help_text="How the conversations should display"
)
use_conversation_context = models.BooleanField(
default=False,
help_text=(
"When enabled, prior turns in the conversation are used as "
"LLM/RAG context for a more tailored experience"
),
)
def get_set_password_url(self):
from django.conf import settings
@@ -89,7 +82,7 @@ class CustomUser(AbstractUser):
class UserAuthEvent(models.Model):
"""Audit trail for auth / account / subscription actions (user admin)."""
"""Audit trail for password reset / set actions, shown on user admin."""
class EventType(models.TextChoices):
PASSWORD_RESET_REQUESTED = (
@@ -98,9 +91,6 @@ class UserAuthEvent(models.Model):
)
PASSWORD_SET = ("password_set", "Password set")
INVITE_SENT = ("invite_sent", "Invite sent")
ACCOUNT_DELETED = ("account_deleted", "Account deleted")
SUBSCRIPTION_STARTED = ("subscription_started", "Subscription started")
SUBSCRIPTION_UPDATED = ("subscription_updated", "Subscription updated")
user = models.ForeignKey(
CustomUser,
@@ -317,14 +307,6 @@ class Prompt(TimeInfoBase):
null=True,
help_text="file type of the file for the prompt",
)
citations = models.JSONField(
default=list,
blank=True,
help_text=(
"Structured source citations for grounded answers (#62). "
"List of {index, title, url, published_at}."
),
)
def get_conversation_title(self):
if self.conversation:
@@ -336,59 +318,6 @@ class Prompt(TimeInfoBase):
return self.file != None and self.file.storage.exists(self.file.name)
class PromptFeedback(TimeInfoBase):
"""Per-message thumbs rating for an assistant Prompt (chat_backend#67).
Distinct from app-wide ``Feedback`` (product bugs). Joinable to
``PromptMetric`` via ``prompt_id`` for per-model accuracy slices.
"""
class Rating(models.TextChoices):
UP = "up", "Up"
DOWN = "down", "Down"
class Reason(models.TextChoices):
INCORRECT = "incorrect", "Incorrect"
OUT_OF_DATE = "out_of_date", "Out of date"
DIDNT_FOLLOW_INSTRUCTIONS = (
"didnt_follow_instructions",
"Didn't follow instructions",
)
UNSAFE = "unsafe", "Unsafe"
OTHER = "other", "Other"
prompt = models.ForeignKey(
Prompt,
on_delete=models.CASCADE,
related_name="prompt_feedbacks",
help_text="Assistant prompt being rated",
)
user = models.ForeignKey(
CustomUser,
on_delete=models.CASCADE,
related_name="prompt_feedbacks",
)
rating = models.CharField(max_length=8, choices=Rating.choices)
reason = models.CharField(
max_length=64,
choices=Reason.choices,
blank=True,
null=True,
)
comment = models.TextField(max_length=1024, blank=True, null=True)
class Meta:
constraints = [
models.UniqueConstraint(
fields=("prompt", "user"),
name="uniq_prompt_feedback_prompt_user",
)
]
def __str__(self):
return f"PromptFeedback(prompt={self.prompt_id}, user={self.user_id}, {self.rating})"
class PromptMetric(TimeInfoBase):
PROMPT_METRIC_CHOICES = (
("CREATED", "Created"),
@@ -443,165 +372,11 @@ class PromptMetric(TimeInfoBase):
# Document Models
class DocumentWorkspace(TimeInfoBase):
"""RAG document container: company (business) or user (personal) owned (#46, #55)."""
name = models.CharField(max_length=255)
company = models.ForeignKey(
Company,
on_delete=models.CASCADE,
null=True,
blank=True,
related_name="document_workspaces",
)
user = models.ForeignKey(
"CustomUser",
on_delete=models.CASCADE,
null=True,
blank=True,
related_name="personal_workspaces",
help_text="Set for personal RAG workspaces; null for company workspaces.",
)
class Meta:
constraints = [
models.CheckConstraint(
condition=(
models.Q(company__isnull=False, user__isnull=True)
| models.Q(company__isnull=True, user__isnull=False)
),
name="document_workspace_company_xor_user",
),
models.UniqueConstraint(
fields=["user"],
condition=models.Q(user__isnull=False),
name="uniq_personal_document_workspace_user",
),
]
def __str__(self):
if self.user_id:
return f"DocumentWorkspace(personal user={self.user_id})"
return f"DocumentWorkspace(company={self.company_id})"
class DriveConnection(TimeInfoBase):
"""A linked Google Drive / Microsoft OneDrive-SharePoint account (#47-#52).
``user`` is null for company-only connections set up by a company manager
(kind=company); personal connections always have ``user`` set.
``company`` is required for kind=company. Personal connections may have
``company`` null when the user is not attached to a company (#55).
"""
class Provider(models.TextChoices):
GOOGLE = "google", "Google"
MICROSOFT = "microsoft", "Microsoft"
class Kind(models.TextChoices):
PERSONAL = "personal", "Personal"
COMPANY = "company", "Company"
class SyncStatus(models.TextChoices):
OK = "ok", "Ok"
ERROR = "error", "Error"
PENDING = "pending", "Pending"
NEVER = "never", "Never"
user = models.ForeignKey(
"CustomUser",
on_delete=models.CASCADE,
null=True,
blank=True,
related_name="drive_connections",
help_text="Null for company-only connections owned by manager setup.",
)
company = models.ForeignKey(
Company,
on_delete=models.CASCADE,
null=True,
blank=True,
related_name="drive_connections",
help_text="Required for company connections; optional for personal (#55).",
)
provider = models.CharField(max_length=32, choices=Provider.choices)
kind = models.CharField(
max_length=16, choices=Kind.choices, default=Kind.PERSONAL
)
access_token = models.TextField(blank=True, default="")
refresh_token = models.TextField(blank=True, default="")
token_expires_at = models.DateTimeField(null=True, blank=True)
scopes = models.TextField(blank=True, default="")
external_account_email = models.EmailField(blank=True, default="")
selected_resource_ids = models.JSONField(
default=list,
blank=True,
help_text="Selected folder/drive/site ids to sync (empty = root/default).",
)
selected_resource_labels = models.JSONField(
default=list,
blank=True,
help_text="Human-readable labels matching selected_resource_ids, for the FE.",
)
last_sync_at = models.DateTimeField(null=True, blank=True)
last_sync_status = models.CharField(
max_length=16, choices=SyncStatus.choices, default=SyncStatus.NEVER
)
last_sync_error = models.TextField(blank=True, default="")
# Progress for FE progress bar while last_sync_status=pending (#59).
sync_total = models.PositiveIntegerField(
default=0,
help_text="Remote files discovered for the current/last sync run.",
)
sync_processed = models.PositiveIntegerField(
default=0,
help_text="Files finished in the current/last sync (includes skips).",
)
sync_added = models.PositiveIntegerField(default=0)
sync_updated = models.PositiveIntegerField(default=0)
sync_failed = models.PositiveIntegerField(default=0)
is_active = models.BooleanField(default=True)
class Meta:
constraints = [
models.CheckConstraint(
condition=(
models.Q(kind="personal", user__isnull=False)
| models.Q(
kind="company",
company__isnull=False,
user__isnull=True,
)
),
name="drive_connection_kind_owner_consistency",
),
models.UniqueConstraint(
fields=["user", "provider"],
condition=models.Q(kind="personal"),
name="uniq_personal_drive_connection_user_provider",
),
models.UniqueConstraint(
fields=["company", "provider"],
condition=models.Q(kind="company"),
name="uniq_company_drive_connection_company_provider",
),
]
def __str__(self):
return (
f"DriveConnection({self.provider}/{self.kind}) "
f"company={self.company_id} user={self.user_id}"
)
company = models.ForeignKey(Company, on_delete=models.CASCADE)
class Document(TimeInfoBase):
class Source(models.TextChoices):
UPLOAD = "upload", "Upload"
GOOGLE_DRIVE = "google_drive", "Google Drive"
ONEDRIVE = "onedrive", "OneDrive"
SHAREPOINT = "sharepoint", "SharePoint"
GOOGLE_SHARED_DRIVE = "google_shared_drive", "Google Shared Drive"
workspace = models.ForeignKey(DocumentWorkspace, on_delete=models.CASCADE)
file = models.FileField(
upload_to="documents/",
@@ -611,182 +386,6 @@ class Document(TimeInfoBase):
uploaded_at = models.DateTimeField(auto_now_add=True)
processed = models.BooleanField(default=False)
active = models.BooleanField(default=False)
source = models.CharField(
max_length=32, choices=Source.choices, default=Source.UPLOAD
)
remote_file_id = models.CharField(max_length=255, blank=True, default="")
remote_etag = models.CharField(max_length=255, blank=True, default="")
remote_name = models.CharField(max_length=512, blank=True, default="")
drive_connection = models.ForeignKey(
DriveConnection,
on_delete=models.SET_NULL,
null=True,
blank=True,
related_name="documents",
)
sync_error = models.TextField(blank=True, default="")
class Meta:
indexes = [
models.Index(fields=["drive_connection", "remote_file_id"]),
]
class AgentRun(TimeInfoBase):
"""A long-running, multi-step agentic task turn (#63).
``user``/``company`` mirror the tenant scope of the triggering chat turn
(never trust a bare ``conversation_id`` — see ``chat_tenant_scope``).
Progress is broadcast on the Redis channel-layer group
:meth:`channel_group_name` so a reconnecting client can resubscribe.
"""
class Status(models.TextChoices):
PENDING = "pending", "Pending"
PLANNING = "planning", "Planning"
RUNNING = "running", "Running"
COMPLETED = "completed", "Completed"
FAILED = "failed", "Failed"
CANCELLED = "cancelled", "Cancelled"
user = models.ForeignKey(
CustomUser,
on_delete=models.CASCADE,
related_name="agent_runs",
)
company = models.ForeignKey(
Company,
on_delete=models.SET_NULL,
null=True,
blank=True,
related_name="agent_runs",
)
conversation = models.ForeignKey(
"Conversation",
on_delete=models.SET_NULL,
null=True,
blank=True,
related_name="agent_runs",
)
prompt = models.ForeignKey(
"Prompt",
on_delete=models.SET_NULL,
null=True,
blank=True,
related_name="agent_runs",
help_text="The user Prompt that triggered this run, if any.",
)
goal = models.TextField(help_text="Natural-language user request/goal.")
title = models.CharField(
max_length=255,
blank=True,
default="",
help_text="Short human-readable title (from the plan, or the goal).",
)
status = models.CharField(
max_length=16, choices=Status.choices, default=Status.PENDING, db_index=True
)
plan = models.JSONField(
default=list,
blank=True,
help_text="Ordered list of {step_id, title, tool} planner steps.",
)
result = models.TextField(
blank=True, default="", help_text="Final synthesised answer."
)
error = models.TextField(blank=True, default="")
model_orchestrator = models.CharField(max_length=215, blank=True, default="")
model_subagent = models.CharField(max_length=215, blank=True, default="")
max_plan_steps = models.PositiveIntegerField(default=8)
max_iterations = models.PositiveIntegerField(default=12)
wall_clock_seconds = models.PositiveIntegerField(default=600)
tool_call_count = models.PositiveIntegerField(default=0)
iteration_count = models.PositiveIntegerField(default=0)
cancel_requested = models.BooleanField(
default=False,
help_text="Set by the cancel endpoint/frame; worker loop polls this.",
)
started_at = models.DateTimeField(null=True, blank=True)
completed_at = models.DateTimeField(null=True, blank=True)
class Meta:
ordering = ["-created"]
def __str__(self) -> str:
return f"AgentRun({self.pk}, user={self.user_id}, {self.status})"
def channel_group_name(self) -> str:
"""Redis channel-layer group so reconnecting clients get updates (#63)."""
return f"agent_run_{self.pk}"
@property
def is_terminal(self) -> bool:
return self.status in {
self.Status.COMPLETED,
self.Status.FAILED,
self.Status.CANCELLED,
}
def mark_cancelled(self) -> None:
self.cancel_requested = True
self.status = self.Status.CANCELLED
self.completed_at = self.completed_at or timezone.now()
self.save(
update_fields=[
"cancel_requested",
"status",
"completed_at",
"last_modified",
]
)
class AgentStep(TimeInfoBase):
"""A single planner step (optionally decomposed into sub-agent steps)."""
class Status(models.TextChoices):
PENDING = "pending", "Pending"
RUNNING = "running", "Running"
COMPLETED = "completed", "Completed"
FAILED = "failed", "Failed"
SKIPPED = "skipped", "Skipped"
CANCELLED = "cancelled", "Cancelled"
run = models.ForeignKey(
AgentRun,
on_delete=models.CASCADE,
related_name="steps",
)
parent_step = models.ForeignKey(
"self",
on_delete=models.CASCADE,
null=True,
blank=True,
related_name="sub_steps",
help_text="Set when this step was produced by a sub-agent (#63).",
)
index = models.PositiveIntegerField(default=0, help_text="Order within the plan.")
title = models.CharField(max_length=255, blank=True, default="")
status = models.CharField(
max_length=16, choices=Status.choices, default=Status.PENDING, db_index=True
)
is_subagent = models.BooleanField(default=False)
tool_name = models.CharField(max_length=64, blank=True, default="")
tool_input = models.JSONField(default=dict, blank=True)
tool_output = models.TextField(
blank=True,
default="",
help_text="Truncated to AGENT_TOOL_OUTPUT_MAX_CHARS.",
)
error = models.TextField(blank=True, default="")
started_at = models.DateTimeField(null=True, blank=True)
completed_at = models.DateTimeField(null=True, blank=True)
class Meta:
ordering = ["index", "created"]
def __str__(self) -> str:
return f"AgentStep(run={self.run_id}, index={self.index}, {self.status})"
class StoredFile(TimeInfoBase):
llm_be/chat_backend/oauth.py
+12 -105
View File
@@ -14,7 +14,7 @@ from django.conf import settings
from django.core import signing
from django.utils import timezone
from .models import Company, CustomUser, DriveConnection, OAuthIdentity
from .models import Company, CustomUser, OAuthIdentity
logger = logging.getLogger(__name__)
@@ -24,7 +24,6 @@ STATE_MAX_AGE_SECONDS = 600
GOOGLE_AUTH_URL = "https://accounts.google.com/o/oauth2/v2/auth"
GOOGLE_TOKEN_URL = "https://oauth2.googleapis.com/token"
GOOGLE_USERINFO_URL = "https://openidconnect.googleapis.com/v1/userinfo"
GOOGLE_DRIVE_READONLY_SCOPE = "https://www.googleapis.com/auth/drive.readonly"
MICROSOFT_AUTH_URL_TMPL = (
"https://login.microsoftonline.com/{tenant}/oauth2/v2.0/authorize"
@@ -32,13 +31,6 @@ MICROSOFT_AUTH_URL_TMPL = (
MICROSOFT_TOKEN_URL_TMPL = (
"https://login.microsoftonline.com/{tenant}/oauth2/v2.0/token"
)
MICROSOFT_DRIVE_PERSONAL_SCOPE = "Files.Read"
MICROSOFT_DRIVE_COMPANY_SCOPE = "Files.Read.All Sites.Read.All"
# OAuth intents (#24 login/signup; #47 Drive linking).
LOGIN_INTENTS = {"login", "signup"}
DRIVE_LINK_INTENTS = {"link_drive", "link_company_drive"}
VALID_INTENTS = LOGIN_INTENTS | DRIVE_LINK_INTENTS
class OAuthError(Exception):
@@ -84,16 +76,14 @@ def configured_providers() -> dict[str, bool]:
}
def dump_oauth_state(
*, provider: str, intent: str, user_id: int | None = None
) -> str:
payload: dict[str, Any] = {"provider": provider, "intent": intent}
if user_id is not None:
payload["user_id"] = user_id
return signing.dumps(payload, salt=STATE_SALT)
def dump_oauth_state(*, provider: str, intent: str) -> str:
return signing.dumps(
{"provider": provider, "intent": intent},
salt=STATE_SALT,
)
def load_oauth_state(state: str) -> dict[str, Any]:
def load_oauth_state(state: str) -> dict[str, str]:
try:
data = signing.loads(state, salt=STATE_SALT, max_age=STATE_MAX_AGE_SECONDS)
except signing.BadSignature as exc:
@@ -102,38 +92,25 @@ def load_oauth_state(state: str) -> dict[str, Any]:
intent = data.get("intent") or "login"
if provider not in OAuthIdentity.Provider.values:
raise OAuthError("invalid_state", "Unknown OAuth provider in state.")
if intent not in VALID_INTENTS:
if intent not in {"login", "signup"}:
raise OAuthError("invalid_state", "Invalid OAuth intent.")
result: dict[str, Any] = {"provider": provider, "intent": intent}
if intent in DRIVE_LINK_INTENTS:
user_id = data.get("user_id")
if not user_id:
raise OAuthError(
"invalid_state", "OAuth state is missing the linking user."
)
result["user_id"] = user_id
return result
return {"provider": provider, "intent": intent}
def _microsoft_tenant() -> str:
return settings.MICROSOFT_OAUTH_TENANT or "common"
def build_authorization_url(
*, provider: str, redirect_uri: str, state: str, intent: str = "login"
) -> str:
def build_authorization_url(*, provider: str, redirect_uri: str, state: str) -> str:
if not provider_configured(provider):
raise OAuthError("provider_not_configured", f"{provider} OAuth is not configured.")
if provider == OAuthIdentity.Provider.GOOGLE:
scope = "openid email profile"
if intent in DRIVE_LINK_INTENTS:
scope = f"{scope} {GOOGLE_DRIVE_READONLY_SCOPE}"
params = {
"client_id": settings.GOOGLE_OAUTH_CLIENT_ID,
"redirect_uri": redirect_uri,
"response_type": "code",
"scope": scope,
"scope": "openid email profile",
"state": state,
"access_type": "offline",
"prompt": "select_account consent",
@@ -142,17 +119,12 @@ def build_authorization_url(
return f"{GOOGLE_AUTH_URL}?{urlencode(params)}"
if provider == OAuthIdentity.Provider.MICROSOFT:
scope = "openid email profile offline_access"
if intent == "link_drive":
scope = f"{scope} {MICROSOFT_DRIVE_PERSONAL_SCOPE}"
elif intent == "link_company_drive":
scope = f"{scope} {MICROSOFT_DRIVE_COMPANY_SCOPE}"
params = {
"client_id": settings.MICROSOFT_OAUTH_CLIENT_ID,
"redirect_uri": redirect_uri,
"response_type": "code",
"response_mode": "query",
"scope": scope,
"scope": "openid email profile offline_access",
"state": state,
"prompt": "select_account",
}
@@ -334,63 +306,7 @@ def upsert_identity(user: CustomUser, profile: ProviderProfile) -> OAuthIdentity
)
def upsert_drive_connection(
*, user: CustomUser, kind: str, profile: ProviderProfile
) -> DriveConnection:
"""Create/refresh a DriveConnection from a link_drive/link_company_drive callback (#47).
Personal connections do not require a company (#55). Company connections do.
"""
if kind == DriveConnection.Kind.COMPANY and not user.company_id:
raise OAuthError(
"no_company",
"A company is required before connecting a company Drive.",
)
expires_at = _token_expiry(profile.expires_in)
if kind == DriveConnection.Kind.PERSONAL:
connection = DriveConnection.objects.filter(
provider=profile.provider,
kind=DriveConnection.Kind.PERSONAL,
user=user,
).first()
if connection is None:
connection = DriveConnection(
company=user.company, # may be None for solo users
provider=profile.provider,
kind=DriveConnection.Kind.PERSONAL,
user=user,
)
else:
connection = DriveConnection.objects.filter(
company=user.company,
provider=profile.provider,
kind=DriveConnection.Kind.COMPANY,
user=None,
).first()
if connection is None:
connection = DriveConnection(
company=user.company,
provider=profile.provider,
kind=DriveConnection.Kind.COMPANY,
user=None,
)
connection.access_token = profile.access_token
if profile.refresh_token:
connection.refresh_token = profile.refresh_token
connection.token_expires_at = expires_at
connection.scopes = profile.scopes
connection.external_account_email = profile.email
connection.is_active = True
connection.last_sync_error = ""
connection.save()
return connection
def _create_sso_user(profile: ProviderProfile) -> CustomUser:
from monetization.services.plans import try_redeem_backer_email
company = Company.objects.create(
name=f"{profile.email}'s workspace",
state="NA",
@@ -407,15 +323,6 @@ def _create_sso_user(profile: ProviderProfile) -> CustomUser:
)
user.set_unusable_password()
user.save()
try_redeem_backer_email(user)
return user
def resolve_link_user(user_id: int) -> CustomUser:
"""Load the authenticated user a Drive-link callback should attach to."""
user = CustomUser.objects.filter(pk=user_id, deleted=False).first()
if user is None:
raise OAuthError("user_not_found", "Linking user account was not found.")
return user
llm_be/chat_backend/ollama_config.py
+5 -96
View File
@@ -1,116 +1,25 @@
"""Shared Ollama client helpers — always use settings.OLLAMA_BASE_URL.
Role-scoped models (#62): THINKING / FAST / UTILITY / EMBED each resolve
independently, with ``OLLAMA_MODEL`` kept as a fallback so existing deploys
keep working until they set the role-specific vars.
"""
from __future__ import annotations
"""Shared Ollama client helpers — always use settings.OLLAMA_BASE_URL."""
from django.conf import settings
ROLE_THINKING = "thinking"
ROLE_FAST = "fast"
ROLE_UTILITY = "utility"
ROLE_EMBED = "embed"
# Agentic task execution (#63): orchestrator plans/synthesises; sub-agents run
# independent plan steps concurrently on a smaller/cheaper model.
ROLE_ORCHESTRATOR = "orchestrator"
ROLE_SUBAGENT = "subagent"
_VALID_ROLES = {
ROLE_THINKING,
ROLE_FAST,
ROLE_UTILITY,
ROLE_EMBED,
ROLE_ORCHESTRATOR,
ROLE_SUBAGENT,
}
def ollama_base_url() -> str:
return getattr(settings, "OLLAMA_BASE_URL", "http://127.0.0.1:11434")
def ollama_model(default: str | None = None) -> str:
"""Legacy single-model accessor. Prefer :func:`ollama_model_for_role`."""
if default:
return default
return getattr(settings, "OLLAMA_MODEL", "gpt-oss:20b")
def ollama_model_for_role(role: str) -> str:
"""Resolve the model name for a generation role.
Lookup order: role-specific setting → ``OLLAMA_MODEL`` fallback →
hard-coded role default (never falls back from embed → chat model).
"""
role = (role or ROLE_THINKING).lower()
if role not in _VALID_ROLES:
raise ValueError(f"Unknown Ollama role: {role!r}")
role_setting = {
ROLE_THINKING: "OLLAMA_MODEL_THINKING",
ROLE_FAST: "OLLAMA_MODEL_FAST",
ROLE_UTILITY: "OLLAMA_MODEL_UTILITY",
ROLE_EMBED: "OLLAMA_EMBED_MODEL",
ROLE_ORCHESTRATOR: "OLLAMA_MODEL_ORCHESTRATOR",
ROLE_SUBAGENT: "OLLAMA_MODEL_SUBAGENT",
}[role]
role_default = {
ROLE_THINKING: "gpt-oss:20b",
ROLE_FAST: "gemma4:latest",
ROLE_UTILITY: "llama3.2",
ROLE_EMBED: "nomic-embed-text",
ROLE_ORCHESTRATOR: "gpt-oss:20b",
ROLE_SUBAGENT: "llama3.2",
}[role]
configured = getattr(settings, role_setting, None)
if configured:
return configured
# Embeddings must never silently fall back to a chat model (#62).
if role == ROLE_EMBED:
return role_default
legacy = getattr(settings, "OLLAMA_MODEL", None)
if legacy:
return legacy
return role_default
return getattr(settings, "OLLAMA_MODEL", "llama3.2")
def ollama_embed_model() -> str:
return ollama_model_for_role(ROLE_EMBED)
return getattr(settings, "OLLAMA_EMBED_MODEL", ollama_model())
def ollama_num_ctx_for_role(role: str) -> int:
role = (role or ROLE_THINKING).lower()
if role == ROLE_FAST:
return int(getattr(settings, "OLLAMA_NUM_CTX_FAST", 8192) or 8192)
if role == ROLE_UTILITY:
return int(getattr(settings, "OLLAMA_NUM_CTX_UTILITY", 4096) or 4096)
if role == ROLE_SUBAGENT:
return int(getattr(settings, "OLLAMA_NUM_CTX_SUBAGENT", 8192) or 8192)
if role == ROLE_ORCHESTRATOR:
return int(getattr(settings, "OLLAMA_NUM_CTX_ORCHESTRATOR", 16384) or 16384)
return int(getattr(settings, "OLLAMA_NUM_CTX_THINKING", 16384) or 16384)
def resolve_chat_role(model_name: str | None) -> str:
"""Map FE mode selector (FAST / THINKING / …) to an Ollama role."""
if (model_name or "").upper() == "FAST":
return ROLE_FAST
return ROLE_THINKING
def ollama_llm_kwargs(role: str = ROLE_THINKING, **extra):
def ollama_llm_kwargs(**extra):
"""Keyword args for langchain_ollama.OllamaLLM / ChatOllama."""
kwargs = {
"base_url": ollama_base_url(),
"model": ollama_model_for_role(role),
"num_ctx": ollama_num_ctx_for_role(role),
}
kwargs = {"base_url": ollama_base_url(), "model": ollama_model()}
kwargs.update(extra)
return kwargs
llm_be/chat_backend/serializers.py
+1 -237
View File
@@ -1,20 +1,15 @@
from rest_framework_simplejwt.serializers import TokenObtainPairSerializer
from rest_framework import serializers
from django.db.models import Count, Q, Sum
from .models import (
CustomUser,
Announcement,
Company,
Conversation,
Prompt,
PromptFeedback,
PromptMetric,
Feedback,
FEEDBACK_CATEGORIES,
DocumentWorkspace,
Document,
DriveConnection,
)
@@ -28,17 +23,6 @@ class MyTokenObtainPairSerializer(TokenObtainPairSerializer):
return token
def validate(self, attrs):
data = super().validate(attrs)
if getattr(self.user, "deleted", False):
from rest_framework_simplejwt.exceptions import AuthenticationFailed
raise AuthenticationFailed(
"No active account found with the given credentials",
code="no_active_account",
)
return data
class CompanySerializer(serializers.ModelSerializer):
class Meta:
@@ -64,39 +48,12 @@ class CustomUserSerializer(serializers.ModelSerializer):
password = serializers.CharField(min_length=8, write_only=True)
company = CompanySerializer()
has_usable_password = serializers.BooleanField()
subscription = serializers.SerializerMethodField()
class Meta:
model = CustomUser
fields = "__all__"
extra_kwargs = {"password": {"write_only": True}}
def get_subscription(self, obj):
from monetization.services.plans import needs_checkout, plan_to_dict
from monetization.models import UserSubscription
try:
sub = obj.subscription
except UserSubscription.DoesNotExist:
return {
"plan": None,
"status": UserSubscription.Status.NONE,
"source": UserSubscription.Source.NONE,
"needs_checkout": True,
"cancel_at_period_end": False,
"current_period_end": None,
}
return {
"plan": plan_to_dict(sub.plan) if sub.plan_id else None,
"status": sub.status,
"source": sub.source,
"needs_checkout": needs_checkout(obj),
"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
),
}
class SelfServeRegistrationSerializer(serializers.Serializer):
"""Minimal payload for public self-serve sign-up (gated by settings)."""
@@ -122,8 +79,6 @@ class SelfServeRegistrationSerializer(serializers.Serializer):
return email
def create(self, validated_data):
from monetization.services.plans import try_redeem_backer_email
email = validated_data["email"]
password = validated_data["password"]
first_name = (validated_data.get("first_name") or "").strip()
@@ -147,116 +102,16 @@ class SelfServeRegistrationSerializer(serializers.Serializer):
company=company,
is_company_manager=True,
)
try_redeem_backer_email(user)
return user
def _conversation_token_totals(conversation_id: int):
"""
Sum PromptMetric tokens for a conversation.
Returns nulls when the provider never reported usage (never fabricate 0).
"""
agg = PromptMetric.objects.filter(conversation_id=conversation_id).aggregate(
tin=Sum("tokens_in"),
tout=Sum("tokens_out"),
with_in=Count("id", filter=Q(tokens_in__isnull=False)),
with_out=Count("id", filter=Q(tokens_out__isnull=False)),
)
return (
agg["tin"] if agg["with_in"] else None,
agg["tout"] if agg["with_out"] else None,
)
def _prompt_token_pair(prompt_id: int):
metric = (
PromptMetric.objects.filter(prompt_id=prompt_id)
.order_by("-created")
.only("tokens_in", "tokens_out")
.first()
)
if metric is None:
return None, None
return metric.tokens_in, metric.tokens_out
class ConversationSerializer(serializers.ModelSerializer):
tokens_in = serializers.SerializerMethodField()
tokens_out = serializers.SerializerMethodField()
class Meta:
model = Conversation
fields = (
"title",
"created",
"last_modified",
"id",
"tokens_in",
"tokens_out",
)
def _token_pair(self, obj):
cache = self.context.setdefault("_conversation_token_cache", {})
if obj.id not in cache:
cache[obj.id] = _conversation_token_totals(obj.id)
return cache[obj.id]
def get_tokens_in(self, obj):
tin, _ = self._token_pair(obj)
return tin
def get_tokens_out(self, obj):
_, tout = self._token_pair(obj)
return tout
class PromptFeedbackSerializer(serializers.ModelSerializer):
prompt_id = serializers.IntegerField(source="prompt.id", read_only=True)
class Meta:
model = PromptFeedback
fields = (
"id",
"prompt_id",
"rating",
"reason",
"comment",
"created",
"last_modified",
)
read_only_fields = ("id", "prompt_id", "created", "last_modified")
class PromptFeedbackUpsertSerializer(serializers.Serializer):
prompt_id = serializers.IntegerField()
rating = serializers.ChoiceField(choices=PromptFeedback.Rating.choices)
reason = serializers.ChoiceField(
choices=PromptFeedback.Reason.choices,
required=False,
allow_null=True,
allow_blank=True,
)
comment = serializers.CharField(
required=False, allow_null=True, allow_blank=True, max_length=1024
)
def validate_reason(self, value):
if value == "":
return None
return value
def validate_comment(self, value):
if value is None:
return None
stripped = str(value).strip()
return stripped or None
fields = ("title", "created", "last_modified", "id")
class PromptSerializer(serializers.ModelSerializer):
tokens_in = serializers.SerializerMethodField()
tokens_out = serializers.SerializerMethodField()
feedback = serializers.SerializerMethodField()
class Meta:
model = Prompt
@@ -265,59 +120,7 @@ class PromptSerializer(serializers.ModelSerializer):
"user_created",
"created",
"id",
"tokens_in",
"tokens_out",
"citations",
"feedback",
)
read_only_fields = ("citations", "feedback")
def _token_pair(self, obj):
cache = self.context.setdefault("_prompt_token_cache", {})
if obj.id not in cache:
cache[obj.id] = _prompt_token_pair(obj.id)
return cache[obj.id]
def get_tokens_in(self, obj):
tin, _ = self._token_pair(obj)
return tin
def get_tokens_out(self, obj):
_, tout = self._token_pair(obj)
return tout
def get_feedback(self, obj):
"""Current caller's rating for this prompt, if any."""
request = self.context.get("request")
if request is None or not getattr(request, "user", None):
return None
user = request.user
if not user.is_authenticated:
return None
by_prompt = self.context.get("_prompt_feedback_by_id")
if by_prompt is None:
prompt_ids = self.context.get("_prompt_ids_for_feedback")
qs = PromptFeedback.objects.filter(user=user).only(
"prompt_id", "rating", "reason", "comment"
)
if prompt_ids is not None:
qs = qs.filter(prompt_id__in=prompt_ids)
by_prompt = {
row.prompt_id: {
"rating": row.rating,
"reason": row.reason,
"comment": row.comment,
}
for row in qs
}
self.context["_prompt_feedback_by_id"] = by_prompt
return by_prompt.get(obj.id)
def validate_message(self, value: str) -> str:
if value is None or not str(value).strip():
raise serializers.ValidationError("Message text cannot be empty.")
return str(value).strip()
class BasicUserSerializer(serializers.ModelSerializer):
@@ -355,42 +158,3 @@ class DocumentSerializer(serializers.ModelSerializer):
"active",
]
read_only_fields = ["id", "uploaded_at", "processed", "created"]
# drive connection serializers (#47-#52)
class DriveConnectionSerializer(serializers.ModelSerializer):
"""Never exposes access_token/refresh_token to the client."""
class Meta:
model = DriveConnection
fields = [
"id",
"provider",
"kind",
"external_account_email",
"selected_resource_ids",
"selected_resource_labels",
"last_sync_at",
"last_sync_status",
"last_sync_error",
"sync_total",
"sync_processed",
"sync_added",
"sync_updated",
"sync_failed",
"is_active",
"created",
]
read_only_fields = fields
class DriveConnectionResourcesSerializer(serializers.Serializer):
resource_ids = serializers.ListField(
child=serializers.CharField(max_length=512), allow_empty=True
)
resource_labels = serializers.ListField(
child=serializers.CharField(max_length=512, allow_blank=True),
allow_empty=True,
required=False,
default=list,
)
llm_be/chat_backend/services/account_deletion.py
-92
View File
@@ -1,92 +0,0 @@
"""Self-service account soft-delete helpers (#34)."""
from __future__ import annotations
import logging
from django.db import transaction
from rest_framework_simplejwt.token_blacklist.models import (
BlacklistedToken,
OutstandingToken,
)
from rest_framework_simplejwt.tokens import RefreshToken
from chat_backend.models import Conversation, CustomUser, UserAuthEvent
logger = logging.getLogger(__name__)
class AccountDeletionError(Exception):
"""Raised when self-delete is not allowed for the requesting user."""
def __init__(self, detail: str, *, code: str = "delete_forbidden"):
super().__init__(detail)
self.detail = detail
self.code = code
def _blacklist_outstanding_tokens(user: CustomUser) -> int:
"""Blacklist all outstanding refresh tokens for the user. Returns count."""
count = 0
for outstanding in OutstandingToken.objects.filter(user=user):
_token, created = BlacklistedToken.objects.get_or_create(token=outstanding)
if created:
count += 1
return count
def _blacklist_refresh_token(refresh_token: str | None) -> None:
if not refresh_token:
return
try:
RefreshToken(refresh_token).blacklist()
except Exception:
logger.info("Self-delete: optional refresh token could not be blacklisted")
@transaction.atomic
def soft_delete_account(
user: CustomUser,
*,
refresh_token: str | None = None,
ip_address: str | None = None,
) -> CustomUser:
"""
Soft-delete the requesting user and hide their conversations.
Privacy (v1): personal data is retained under soft-delete for admin/audit.
Full purge (chats, documents, RAG vectors, auth events) is a follow-up.
"""
if user.is_staff or user.is_superuser:
raise AccountDeletionError(
"Staff accounts cannot self-delete. Contact an administrator.",
code="staff_forbidden",
)
if user.deleted:
raise AccountDeletionError(
"This account has already been deleted.",
code="already_deleted",
)
user.deleted = True
user.is_active = False
user.save(update_fields=["deleted", "is_active"])
Conversation.objects.filter(user=user, deleted=False).update(deleted=True)
UserAuthEvent.log(
user,
UserAuthEvent.EventType.ACCOUNT_DELETED,
detail="Self-service account soft-delete",
ip_address=ip_address,
)
_blacklist_refresh_token(refresh_token)
blacklisted = _blacklist_outstanding_tokens(user)
logger.info(
"Soft-deleted user pk=%s; blacklisted_outstanding=%s",
user.pk,
blacklisted,
)
return user
llm_be/chat_backend/services/agent/__init__.py
-6
View File
@@ -1,6 +0,0 @@
"""Agentic task execution (#63) — plan → tools → synthesise with progress frames."""
from chat_backend.services.agent.decider import should_use_agent
from chat_backend.services.agent.runner import run_agentic_turn
__all__ = ["should_use_agent", "run_agentic_turn"]
llm_be/chat_backend/services/agent/classifier.py
-57
View File
@@ -1,57 +0,0 @@
"""Heuristic gate: does this turn need the agentic multi-step path? (#63)
Cheap, deterministic, no LLM call. Kept conservative on purpose — when in
doubt, stay on the existing always-on grounded chat path (#62) so ordinary
turns never regress. Only prompts that read as genuinely multi-step research
or tool-combining tasks should route into the (much more expensive) agent
orchestrator.
"""
from __future__ import annotations
import re
_MULTI_STEP_MARKERS = (
"step by step",
"step-by-step",
" then ",
" and then",
"after that",
"first,",
"first ",
"finally,",
"compare",
"summarize and",
"summarise and",
"research and",
"comprehensive report",
"detailed report",
"multiple sources",
"cross-reference",
"cross reference",
"for each of",
)
_NUMBERED_LIST_RE = re.compile(r"(?:^|\n)\s*(?:[1-9]\.|[-*])\s+\S", re.MULTILINE)
def is_complex_agentic_task(message: str) -> bool:
"""True when ``message`` looks like a multi-step research/analysis task."""
text = (message or "").strip()
if not text:
return False
lowered = text.lower()
if _NUMBERED_LIST_RE.search(text):
return True
marker_hits = sum(1 for marker in _MULTI_STEP_MARKERS if marker in lowered)
if marker_hits >= 1 and len(text) > 60:
return True
# Long, multi-sentence asks read as compound/multi-part requests.
sentence_count = len(re.findall(r"[.!?]+", text))
if sentence_count >= 2 and len(text) > 180:
return True
return False
llm_be/chat_backend/services/agent/decider.py
-40
View File
@@ -1,40 +0,0 @@
"""Heuristic gate: only enter the agent path for multi-step goals (#63).
Simple chat must stay on the grounded path — no planner latency regression.
"""
from __future__ import annotations
import re
from django.conf import settings
# Phrases that usually need decomposition / parallel research / doc extraction.
_MULTI_STEP_PATTERNS = (
re.compile(r"\btop\s+\d+\b", re.I),
re.compile(r"\bcompare\b", re.I),
re.compile(r"\bcomparison\b", re.I),
re.compile(r"\bresearch\b", re.I),
re.compile(r"\bcompetitors?\b", re.I),
re.compile(r"\beach\b.+\bsummar", re.I),
re.compile(r"\bextract\b.+\b(from|every|all)\b", re.I),
re.compile(r"\bread\b.+\b(pdfs?|documents?|files?)\b", re.I),
re.compile(r"\bstep[- ]by[- ]step\b", re.I),
re.compile(r"\bthen\b.+\b(and|also)\b", re.I),
re.compile(r"\bfor each\b", re.I),
re.compile(r"\bmulti[- ]?step\b", re.I),
)
def should_use_agent(message: str) -> bool:
"""Return True when agentic execution should handle this turn.
Always False when ``ALLOW_AGENTIC_TASKS`` is off — consumers must keep
today's grounded path unchanged.
"""
if not getattr(settings, "ALLOW_AGENTIC_TASKS", False):
return False
text = (message or "").strip()
if not text or len(text) < 24:
return False
return any(p.search(text) for p in _MULTI_STEP_PATTERNS)
llm_be/chat_backend/services/agent/orchestrator.py
-345
View File
@@ -1,345 +0,0 @@
"""LangGraph-driven planner → executor → tools → synthesiser loop (#63).
Deliberately decoupled from Django/DB — ``AgentOrchestrator`` takes plain
callables (``on_event``, ``is_cancelled``) and duck-typed chat models
(anything with an async ``ainvoke``), so it can be unit tested with fakes
with no Ollama/Redis/DB involved. :mod:`chat_backend.services.agent.runner`
is the glue that wires this to ``AgentRun``/``AgentStep`` and WS frames.
"""
from __future__ import annotations
import asyncio
import json
import logging
import re
import time
from collections import defaultdict
from dataclasses import dataclass, field
from typing import Any, Awaitable, Callable, Optional
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage
from chat_backend.services.tools.registry import TOOL_CATALOG, suggest_tools
logger = logging.getLogger(__name__)
OnEvent = Optional[Callable[[str, dict], Awaitable[None]]]
IsCancelled = Optional[Callable[[], Awaitable[bool]]]
class AgentCancelled(RuntimeError):
"""Raised to unwind the run loop when the caller reports cancellation."""
class AgentRunLimitExceeded(RuntimeError):
"""Raised when the wall-clock or iteration budget is exhausted."""
@dataclass
class PlanStep:
step_id: str
title: str
tool: str | None = None
tool_input: dict = field(default_factory=dict)
parallel_group: int = 0
def to_dict(self) -> dict:
return {
"step_id": self.step_id,
"title": self.title,
"tool": self.tool,
"parallel_group": self.parallel_group,
}
@dataclass
class AgentPlan:
title: str
steps: list[PlanStep]
@dataclass
class RunLimits:
max_plan_steps: int = 8
max_iterations: int = 12
wall_clock_seconds: int = 600
subagent_concurrency: int = 3
max_step_tool_calls: int = 4
PLANNER_SYSTEM_TEMPLATE = """You are the planning module of an autonomous research/analysis agent.
Given a user's goal and a catalog of available tools, produce a short ordered
plan of at most {max_steps} steps. Return ONLY compact JSON, no prose:
{{"title": "short plan title", "steps": [
{{"step_id": "s1", "title": "...", "tool": "web_search" or null, "tool_input": {{}}, "parallel_group": 0}}
]}}
Rules:
- Use "tool": null for a pure reasoning/synthesis step (no tool call).
- Two steps may share the same parallel_group ONLY when neither depends on
the other's output — they will execute concurrently.
- Prefer the smallest plan that actually accomplishes the goal.
Tools available:
{tool_catalog}
"""
def _tool_catalog_text() -> str:
return "\n".join(f"- {spec.name}: {spec.description}" for spec in TOOL_CATALOG.values())
def _strip_fences(text: str) -> str:
text = text.strip()
if text.startswith("```"):
text = re.sub(r"^```(?:json)?\s*", "", text)
text = re.sub(r"\s*```$", "", text)
return text
def fallback_plan(goal: str, max_steps: int) -> AgentPlan:
"""Deterministic, offline plan used when the planner LLM is unavailable
or returns unparseable output (fail open rather than fail closed)."""
tool_names = suggest_tools(goal)
if not tool_names:
return AgentPlan(
title=goal[:80] or "Task",
steps=[PlanStep(step_id="s1", title=goal[:120] or "Answer the request", tool=None)],
)
steps = [
PlanStep(step_id=f"s{i + 1}", title=f"Use {name} for: {goal[:60]}", tool=name)
for i, name in enumerate(tool_names[:max_steps])
]
return AgentPlan(title=goal[:80] or "Task", steps=steps)
def parse_plan(raw: str, *, goal: str, max_steps: int) -> AgentPlan:
text = _strip_fences(str(raw or ""))
try:
start, end = text.find("{"), text.rfind("}")
if start < 0 or end < 0:
raise ValueError("no JSON object in planner output")
payload = json.loads(text[start : end + 1])
steps_raw = payload.get("steps") or []
if not isinstance(steps_raw, list) or not steps_raw:
raise ValueError("plan has no steps")
steps: list[PlanStep] = []
for i, raw_step in enumerate(steps_raw[:max_steps]):
if not isinstance(raw_step, dict):
continue
steps.append(
PlanStep(
step_id=str(raw_step.get("step_id") or f"s{i + 1}"),
title=str(raw_step.get("title") or f"Step {i + 1}"),
tool=(raw_step.get("tool") or None),
tool_input=raw_step.get("tool_input") or {},
parallel_group=int(raw_step.get("parallel_group") or 0),
)
)
if not steps:
raise ValueError("plan produced zero usable steps")
return AgentPlan(title=str(payload.get("title") or goal[:80]), steps=steps)
except Exception as exc:
logger.warning("Agent plan parse failed (%s); using heuristic fallback plan", exc)
return fallback_plan(goal, max_steps)
class AgentOrchestrator:
def __init__(
self,
*,
goal: str,
history_text: str,
planner_llm: Any,
subagent_llm_factory: Callable[[], Any],
tools: list,
limits: RunLimits,
on_event: OnEvent = None,
is_cancelled: IsCancelled = None,
):
self.goal = goal
self.history_text = history_text or ""
self.planner_llm = planner_llm
self.subagent_llm_factory = subagent_llm_factory
self.tools = tools
self.tools_by_name = {getattr(t, "name", None): t for t in tools}
self.limits = limits
self.on_event = on_event
self.is_cancelled = is_cancelled
self.iterations = 0
self._deadline = time.monotonic() + max(1, limits.wall_clock_seconds)
async def _emit(self, event_type: str, data: dict) -> None:
if self.on_event is None:
return
try:
await self.on_event(event_type, data)
except Exception: # pragma: no cover - progress must never break a run
logger.exception("agent on_event failed for %s", event_type)
def _time_left(self) -> bool:
return time.monotonic() < self._deadline
async def _raise_if_cancelled(self) -> None:
if self.is_cancelled is not None and await self.is_cancelled():
raise AgentCancelled("Agent run was cancelled")
async def plan(self) -> AgentPlan:
prompt = PLANNER_SYSTEM_TEMPLATE.format(
max_steps=self.limits.max_plan_steps, tool_catalog=_tool_catalog_text()
)
try:
response = await self.planner_llm.ainvoke(
[
SystemMessage(content=prompt),
HumanMessage(
content=(
f"Goal: {self.goal}\n\nRelevant conversation so far:\n"
f"{self.history_text or '(none)'}"
)
),
]
)
raw = getattr(response, "content", response)
except Exception as exc:
logger.warning("Planner LLM call failed (%s); using heuristic fallback plan", exc)
return fallback_plan(self.goal, self.limits.max_plan_steps)
return parse_plan(raw, goal=self.goal, max_steps=self.limits.max_plan_steps)
async def _call_tool(self, name: str, tool_input: dict) -> str:
tool = self.tools_by_name.get(name)
if tool is None:
return f"Unknown tool: {name!r}"
try:
result = await tool.ainvoke(tool_input or {})
except Exception as exc:
logger.warning("Tool %s failed: %s", name, exc)
return f"Tool {name} failed: {exc}"
return str(result)
async def execute_step(self, step: PlanStep, prior_results: dict[str, str]) -> str:
"""Run one plan step. Direct tool dispatch when the planner named a
tool; otherwise a bounded tool-calling reasoning loop on the
sub-agent model."""
await self._raise_if_cancelled()
if not self._time_left():
raise AgentRunLimitExceeded("Wall-clock budget exceeded")
if step.tool:
self.iterations += 1
return await self._call_tool(step.tool, step.tool_input)
llm = self.subagent_llm_factory()
context = "\n".join(f"- {k}: {v}" for k, v in prior_results.items()) or "(none yet)"
messages: list[Any] = [
SystemMessage(
content=(
"You are a focused sub-agent executing ONE step of a larger "
"plan. Use a tool when it would help; otherwise answer "
"directly and concisely. Do not repeat the whole plan."
)
),
HumanMessage(
content=(
f"Overall goal: {self.goal}\nThis step: {step.title}\n\n"
f"Prior step results:\n{context}"
)
),
]
for _ in range(self.limits.max_step_tool_calls):
await self._raise_if_cancelled()
if not self._time_left() or self.iterations >= self.limits.max_iterations:
break
self.iterations += 1
response = await llm.ainvoke(messages)
tool_calls = getattr(response, "tool_calls", None) or []
if not tool_calls:
return str(getattr(response, "content", response) or "")
messages.append(response)
for call in tool_calls:
call_name = call.get("name") if isinstance(call, dict) else getattr(call, "name", None)
call_args = call.get("args") if isinstance(call, dict) else getattr(call, "args", {})
call_id = call.get("id") if isinstance(call, dict) else getattr(call, "id", None)
tool_result = await self._call_tool(call_name, call_args or {})
messages.append(ToolMessage(content=tool_result, tool_call_id=call_id or ""))
# Ran out of tool-call rounds — ask once more for a final answer.
try:
final = await llm.ainvoke(
messages + [HumanMessage(content="Give your final answer for this step now.")]
)
return str(getattr(final, "content", final) or "")
except Exception as exc:
return f"(step did not converge: {exc})"
async def synthesize(self, plan: AgentPlan, step_results: dict[str, str]) -> str:
context = "\n\n".join(
f"### {step.title}\n{step_results.get(step.step_id, '')}" for step in plan.steps
)
messages = [
SystemMessage(
content=(
"Combine the plan step results below into one clear, "
"well-organised final answer for the user. Write in "
"prose; do not mention internal step numbers or tool "
"names unless the user would find that useful."
)
),
HumanMessage(content=f"Goal: {self.goal}\n\nStep results:\n{context}"),
]
response = await self.planner_llm.ainvoke(messages)
return str(getattr(response, "content", response) or "").strip()
async def run(self) -> tuple[AgentPlan, dict[str, str], str]:
"""Plan, execute (respecting parallel groups), then synthesise.
Returns ``(plan, step_results, final_answer)``. Raises
:class:`AgentCancelled` if cancellation is detected mid-run.
"""
plan = await self.plan()
await self._emit("plan_ready", {"title": plan.title, "steps": [s.to_dict() for s in plan.steps]})
groups: dict[int, list[PlanStep]] = defaultdict(list)
for step in plan.steps:
groups[step.parallel_group].append(step)
step_results: dict[str, str] = {}
for group_key in sorted(groups):
await self._raise_if_cancelled()
semaphore = asyncio.Semaphore(max(1, self.limits.subagent_concurrency))
async def run_one(step: PlanStep) -> None:
async with semaphore:
await self._emit(
"step_started",
{"step_id": step.step_id, "title": step.title, "tool": step.tool},
)
try:
result = await self.execute_step(step, step_results)
except AgentCancelled:
await self._emit(
"step_failed",
{"step_id": step.step_id, "title": step.title, "error": "cancelled"},
)
raise
except Exception as exc:
logger.warning("Step %s failed: %s", step.step_id, exc)
step_results[step.step_id] = f"(step failed: {exc})"
await self._emit(
"step_failed",
{"step_id": step.step_id, "title": step.title, "error": str(exc)},
)
return
step_results[step.step_id] = result
await self._emit(
"step_completed",
{"step_id": step.step_id, "title": step.title, "result": result},
)
await asyncio.gather(*(run_one(step) for step in groups[group_key]))
await self._raise_if_cancelled()
final_answer = await self.synthesize(plan, step_results)
return plan, step_results, final_answer
llm_be/chat_backend/services/agent/planner.py
-131
View File
@@ -1,131 +0,0 @@
"""Planner: turn a goal into an ordered step list (#63).
Falls back to a deterministic heuristic when the LLM is unavailable so offline
tests and degraded deploys still produce a plan.
"""
from __future__ import annotations
import json
import logging
import re
from typing import Any
from django.conf import settings
from chat_backend.services.tools.registry import TOOL_CATALOG, suggest_tools
logger = logging.getLogger(__name__)
def _cap(n: int | None = None) -> int:
return int(n or getattr(settings, "AGENT_MAX_PLAN_STEPS", 8) or 8)
def heuristic_plan(goal: str, *, max_steps: int | None = None) -> list[dict[str, Any]]:
"""Offline plan from keyword hints — never calls an LLM."""
cap = _cap(max_steps)
suggested = suggest_tools(goal) or ["web_search"]
steps: list[dict[str, Any]] = []
# Parallel research flavour: "top N competitors"
m = re.search(r"top\s+(\d+)", goal, re.I)
if m and int(m.group(1)) <= cap:
n = min(int(m.group(1)), cap - 1)
for i in range(1, n + 1):
steps.append(
{
"index": len(steps),
"title": f"Research item {i}",
"tool": "web_search",
"depends_on": [],
}
)
steps.append(
{
"index": len(steps),
"title": "Synthesise comparison",
"tool": "synthesis",
"depends_on": list(range(n)),
}
)
return steps[:cap]
for name in suggested[: max(1, cap - 1)]:
spec = TOOL_CATALOG.get(name)
steps.append(
{
"index": len(steps),
"title": (spec.description[:80] if spec else name),
"tool": name,
"depends_on": [],
}
)
steps.append(
{
"index": len(steps),
"title": "Synthesise answer",
"tool": "synthesis",
"depends_on": [s["index"] for s in steps[:-1]] or [],
}
)
return steps[:cap]
def parse_plan_json(raw: str, *, max_steps: int | None = None) -> list[dict[str, Any]]:
"""Parse planner LLM JSON; raise ValueError on garbage."""
text = (raw or "").strip()
if text.startswith("```"):
text = re.sub(r"^```(?:json)?\s*", "", text)
text = re.sub(r"\s*```$", "", text)
data = json.loads(text)
if isinstance(data, dict):
data = data.get("steps") or data.get("plan") or []
if not isinstance(data, list) or not data:
raise ValueError("plan must be a non-empty list")
cap = _cap(max_steps)
out: list[dict[str, Any]] = []
for i, item in enumerate(data[:cap]):
if not isinstance(item, dict):
continue
out.append(
{
"index": i,
"title": str(item.get("title") or f"Step {i + 1}")[:255],
"tool": str(item.get("tool") or item.get("assigned") or "web_search"),
"depends_on": list(item.get("depends_on") or []),
}
)
if not out:
raise ValueError("no valid steps")
return out
def plan_for_goal(goal: str, *, max_steps: int | None = None) -> list[dict[str, Any]]:
"""Best-effort plan: try utility LLM, fall back to heuristic."""
try:
from chat_backend.services.base_service import BaseService
from chat_backend.ollama_config import ROLE_UTILITY, ollama_llm_kwargs
from langchain_ollama import OllamaLLM
from langchain_core.prompts import ChatPromptTemplate
llm = OllamaLLM(
**ollama_llm_kwargs(role=ROLE_UTILITY, temperature=0.0)
)
prompt = ChatPromptTemplate.from_messages(
[
(
"system",
"Return ONLY JSON: {\"steps\":[{\"title\":str,\"tool\":str,"
"\"depends_on\":[int]}]}. Tools: web_search, fetch_url, "
"search_documents, read_document, analyse_dataframe, make_plot, "
"synthesis. Cap steps at {cap}. Simple chit-chat → one synthesis step.",
),
("human", "{goal}"),
]
)
chain = prompt | llm
raw = chain.invoke({"goal": goal, "cap": _cap(max_steps)})
return parse_plan_json(str(raw), max_steps=max_steps)
except Exception:
logger.info("planner LLM unavailable; using heuristic plan", exc_info=True)
return heuristic_plan(goal, max_steps=max_steps)
llm_be/chat_backend/services/agent/progress.py
-66
View File
@@ -1,66 +0,0 @@
"""Publish agent progress frames via channel layer (+ optional WS callback)."""
from __future__ import annotations
import json
import logging
from collections.abc import Awaitable, Callable
from typing import Any
from asgiref.sync import async_to_sync
from channels.layers import get_channel_layer
from chat_backend.services.ws_frames import agent_frame, status_frame
logger = logging.getLogger(__name__)
WsSender = Callable[[str], Awaitable[None]]
async def publish_agent_event(
*,
run_id: int,
frame_type: str,
data: dict[str, Any],
ws_send: WsSender | None = None,
) -> dict:
"""Emit a versioned agent frame to WS (if bound) and Redis group."""
payload = agent_frame(frame_type, {"run_id": run_id, **data})
raw = json.dumps(payload)
if ws_send is not None:
try:
await ws_send(raw)
except Exception:
logger.exception("WS agent frame send failed run_id=%s", run_id)
layer = get_channel_layer()
if layer is not None:
group = f"agent_run_{run_id}"
try:
await layer.group_send(
group, {"type": "agent.progress", "text": raw}
)
except Exception:
logger.exception("channel layer group_send failed group=%s", group)
return payload
def publish_agent_event_sync(
*,
run_id: int,
frame_type: str,
data: dict[str, Any],
) -> dict:
return async_to_sync(publish_agent_event)(
run_id=run_id, frame_type=frame_type, data=data, ws_send=None
)
async def publish_status(
stage: str,
*,
detail: str | None = None,
ws_send: WsSender | None = None,
) -> None:
if ws_send is None:
return
await ws_send(json.dumps(status_frame(stage, detail=detail)))
llm_be/chat_backend/services/agent/runner.py
-338
View File
@@ -1,338 +0,0 @@
"""Glue: AgentRun/AgentStep persistence + WS frame fan-out for the orchestrator (#63).
Runs inside a Celery task or the thread fallback (see ``tasks.py``) — never
inline on the request/WS coroutine, since a run can take up to
``AGENT_WALL_CLOCK_SECONDS``. Progress is always broadcast on the run's Redis
channel-layer group (:meth:`AgentRun.channel_group_name`) so any WS consumer
that has joined the group — the originating connection, or a reconnect —
receives ``agent_*`` frames live; ``GET /api/agent_runs/<id>/`` backfills the
full history for clients that missed frames.
"""
from __future__ import annotations
import json
import logging
from asgiref.sync import sync_to_async
from channels.layers import get_channel_layer
from django.conf import settings
from django.utils import timezone
from chat_backend.ollama_config import (
ROLE_ORCHESTRATOR,
ROLE_SUBAGENT,
ollama_llm_kwargs,
)
from chat_backend.services.agent.orchestrator import (
AgentCancelled,
AgentOrchestrator,
AgentRunLimitExceeded,
RunLimits,
)
from chat_backend.services.chat_tenant_scope import (
ChatTenantScopeError,
get_workspace_for_scope,
resolve_chat_company_scope,
)
from chat_backend.services.tools import ToolBudget, ToolContext, build_agent_tools
from chat_backend.services.ws_frames import agent_frame
logger = logging.getLogger(__name__)
@sync_to_async
def _load_run(run_id: int):
from chat_backend.models import AgentRun
return AgentRun.objects.select_related("user", "company").filter(pk=run_id).first()
@sync_to_async
def _save_run(run, fields: list[str]) -> None:
run.save(update_fields=[*fields, "last_modified"])
@sync_to_async
def _refresh_cancel_flag(run_id: int) -> bool:
from chat_backend.models import AgentRun
return bool(
AgentRun.objects.filter(pk=run_id).values_list("cancel_requested", flat=True).first()
)
@sync_to_async
def _resolve_scope_and_workspace(user, conversation_id):
"""Tenant-scoped workspace resolution, mirroring the WS consumer path.
Falls back to no workspace (web/fetch tools only) if scope resolution
fails — an agent run should never 500 out just because the triggering
conversation was deleted mid-flight.
"""
try:
scope = resolve_chat_company_scope(user, conversation_id)
return scope, get_workspace_for_scope(scope)
except ChatTenantScopeError:
logger.warning("Agent run: could not resolve tenant scope for user=%s", user.pk)
return None, None
@sync_to_async
def _upsert_step(run_id: int, event_type: str, data: dict):
from chat_backend.models import AgentStep
step_id = data.get("step_id")
defaults = {"title": data.get("title") or ""}
if event_type == "step_started":
defaults["status"] = AgentStep.Status.RUNNING
defaults["tool_name"] = data.get("tool") or ""
defaults["started_at"] = timezone.now()
elif event_type == "step_completed":
defaults["status"] = AgentStep.Status.COMPLETED
defaults["tool_output"] = str(data.get("result") or "")[:20000]
defaults["completed_at"] = timezone.now()
elif event_type == "step_failed":
defaults["status"] = AgentStep.Status.FAILED
defaults["error"] = str(data.get("error") or "")[:4000]
defaults["completed_at"] = timezone.now()
step, _created = AgentStep.objects.update_or_create(
run_id=run_id,
title=data.get("title") or "",
defaults=defaults,
)
return step
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()
if layer is None:
return
try:
await layer.group_send(
run.channel_group_name(),
{"type": "agent.frame", "frame": frame},
)
except Exception: # pragma: no cover - fan-out must never break a run
logger.exception("Failed to broadcast agent frame run=%s type=%s", run.pk, event_type)
def _build_llms(tools: list):
from langchain_ollama import ChatOllama
planner_llm = ChatOllama(
**ollama_llm_kwargs(role=ROLE_ORCHESTRATOR, temperature=0.2)
)
def subagent_llm_factory():
llm = ChatOllama(**ollama_llm_kwargs(role=ROLE_SUBAGENT, temperature=0.3))
return llm.bind_tools(tools) if tools else llm
return planner_llm, subagent_llm_factory
async def execute_agent_run(run_id: int, *, ws_send=None) -> None:
"""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)
if run is None:
logger.error("execute_agent_run: AgentRun %s not found", run_id)
return
from chat_backend.models import AgentRun as AgentRunModel
if run.status not in (AgentRunModel.Status.PENDING,):
logger.info("execute_agent_run: run %s already %s; skipping", run_id, run.status)
return
run.status = AgentRunModel.Status.PLANNING
run.started_at = timezone.now()
await _save_run(run, ["status", "started_at"])
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(
run.user, run.conversation_id
)
budget = ToolBudget(max_calls=int(getattr(settings, "AGENT_MAX_TOOL_CALLS_PER_RUN", 40) or 40))
ctx = ToolContext.from_settings(scope=scope, budget=budget)
tools = build_agent_tools(ctx, workspace=workspace)
planner_llm, subagent_llm_factory = _build_llms(tools)
limits = RunLimits(
max_plan_steps=run.max_plan_steps,
max_iterations=run.max_iterations,
wall_clock_seconds=run.wall_clock_seconds,
subagent_concurrency=int(getattr(settings, "AGENT_SUBAGENT_CONCURRENCY", 3) or 3),
)
async def on_event(event_type: str, data: dict) -> None:
if event_type in ("step_started", "step_completed", "step_failed"):
await _upsert_step(run.pk, event_type, data)
await _broadcast(run, event_type, data, ws_send=ws_send)
async def is_cancelled() -> bool:
return await _refresh_cancel_flag(run.pk)
orchestrator = AgentOrchestrator(
goal=run.goal,
history_text="",
planner_llm=planner_llm,
subagent_llm_factory=subagent_llm_factory,
tools=tools,
limits=limits,
on_event=on_event,
is_cancelled=is_cancelled,
)
run.status = AgentRunModel.Status.RUNNING
await _save_run(run, ["status"])
try:
plan, _step_results, final_answer = await orchestrator.run()
run.plan = [s.to_dict() for s in plan.steps]
run.title = plan.title or run.title
run.result = final_answer
run.status = AgentRunModel.Status.COMPLETED
run.completed_at = timezone.now()
run.tool_call_count = budget.calls_made
run.iteration_count = orchestrator.iterations
await _save_run(
run,
[
"plan",
"title",
"result",
"status",
"completed_at",
"tool_call_count",
"iteration_count",
],
)
await _broadcast(
run,
"run_completed",
{"status": run.status, "result": final_answer, "title": run.title},
ws_send=ws_send,
)
except AgentCancelled:
run.status = AgentRunModel.Status.CANCELLED
run.completed_at = timezone.now()
run.tool_call_count = budget.calls_made
run.iteration_count = orchestrator.iterations
await _save_run(
run, ["status", "completed_at", "tool_call_count", "iteration_count"]
)
await _broadcast(
run, "run_completed", {"status": run.status, "title": run.title}, ws_send=ws_send
)
except AgentRunLimitExceeded as exc:
run.status = AgentRunModel.Status.FAILED
run.error = str(exc)
run.completed_at = timezone.now()
run.tool_call_count = budget.calls_made
run.iteration_count = orchestrator.iterations
await _save_run(
run,
["status", "error", "completed_at", "tool_call_count", "iteration_count"],
)
await _broadcast(
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
logger.exception("Agent run %s failed", run_id)
run.status = AgentRunModel.Status.FAILED
run.error = str(exc)
run.completed_at = timezone.now()
run.tool_call_count = budget.calls_made
run.iteration_count = orchestrator.iterations
await _save_run(
run,
["status", "error", "completed_at", "tool_call_count", "iteration_count"],
)
await _broadcast(
run,
"run_completed",
{"status": run.status, "error": str(exc), "title": run.title},
ws_send=ws_send,
)
async def run_agentic_turn(
*,
user,
scope,
conversation_id: int,
goal: str,
prompt=None,
ws_send=None,
):
"""Create an AgentRun and execute it inline for the originating WS turn.
``scope`` is accepted for API compatibility with consumers (tenant scope
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 chat_backend.models import AgentRun
from chat_backend.ollama_config import (
ROLE_ORCHESTRATOR,
ROLE_SUBAGENT,
ollama_model_for_role,
)
from chat_backend.services.status_context import emit_status
del scope
if not getattr(settings, "ALLOW_AGENTIC_TASKS", False):
raise RuntimeError("ALLOW_AGENTIC_TASKS is disabled")
await emit_status("queued")
run = await sync_to_async(AgentRun.objects.create)(
user=user,
company_id=getattr(user, "company_id", None),
conversation_id=conversation_id,
prompt=prompt,
goal=goal,
status=AgentRun.Status.PENDING,
model_orchestrator=ollama_model_for_role(ROLE_ORCHESTRATOR),
model_subagent=ollama_model_for_role(ROLE_SUBAGENT),
max_plan_steps=int(getattr(settings, "AGENT_MAX_PLAN_STEPS", 8) or 8),
max_iterations=int(getattr(settings, "AGENT_MAX_ITERATIONS", 12) or 12),
wall_clock_seconds=int(
getattr(settings, "AGENT_WALL_CLOCK_SECONDS", 600) or 600
),
)
await emit_status("evaluating", "Planning multi-step task")
await execute_agent_run(run.pk, ws_send=ws_send)
finished = await _load_run(run.pk)
answer = (
(finished.result if finished else "")
or (finished.error if finished else "")
or "Agent run finished with no result."
)
return finished, answer
llm_be/chat_backend/services/agent/tasks.py
-56
View File
@@ -1,56 +0,0 @@
"""Background dispatch for agent runs (#63)."""
from __future__ import annotations
import logging
import threading
from django.conf import settings
logger = logging.getLogger(__name__)
def _broker_configured() -> bool:
return bool(
getattr(settings, "CELERY_BROKER_URL", "")
or getattr(settings, "REDIS_URL", "")
)
def run_agent_run_sync(run_id: int, scope_payload: dict | None = None) -> None:
"""Execute an agent run in the current process (sync entry for Celery/thread)."""
from asgiref.sync import async_to_sync
from chat_backend.services.agent.runner import execute_agent_run
del scope_payload # tenant resolved from AgentRun.user inside runner
async_to_sync(execute_agent_run)(run_id)
try:
from llm_be.celery import celery_app
@celery_app.task(name="chat_backend.run_agent_run")
def run_agent_run_task(run_id: int, scope_payload: dict | None = None) -> None:
run_agent_run_sync(run_id, scope_payload)
except Exception: # pragma: no cover — celery optional at import time in odd envs
run_agent_run_task = None # type: ignore
def enqueue_agent_run(run_id: int, scope_payload: dict | None = None) -> None:
"""Enqueue via Celery when broker present; else daemon thread (like drive_tasks)."""
if _broker_configured() and run_agent_run_task is not None:
try:
run_agent_run_task.delay(run_id, scope_payload)
return
except Exception:
logger.exception(
"Celery enqueue failed for agent run=%s; falling back to thread",
run_id,
)
threading.Thread(
target=run_agent_run_sync,
args=(run_id, scope_payload),
daemon=True,
name=f"agent-run-{run_id}",
).start()
llm_be/chat_backend/services/assistant_identity.py
-14
View File
@@ -11,17 +11,3 @@ ASSISTANT_SYSTEM_PROMPT = (
"Your name evokes quiet, rest, silence, and stillness — "
"respond with calm clarity; keep answers focused and uncluttered."
)
GROUNDED_ANSWER_INSTRUCTIONS = (
"You have been given numbered live sources. Answer ONLY from those sources. "
"Cite source indexes inline like [1] or [2]. "
"If the sources do not settle the question, say so explicitly — do not fill "
"gaps from memory or training data. "
"When sources conflict, prefer the most recent dated source. "
"Never state a date, number, name, or event that does not appear in the sources."
)
RETRIEVAL_FAILED_MESSAGE = (
"I couldn't reach live sources to answer this accurately right now. "
"Please try again in a moment — I won't guess from outdated training data."
)
llm_be/chat_backend/services/base_service.py
+4 -8
View File
@@ -1,24 +1,20 @@
from abc import ABC
from abc import ABC, abstractmethod
from langchain_ollama import OllamaLLM
from langchain_core.output_parsers import StrOutputParser
from chat_backend.ollama_config import ROLE_UTILITY, ollama_llm_kwargs
from chat_backend.ollama_config import ollama_llm_kwargs
class BaseService(ABC):
"""Abstract base class for LLM conversation services."""
def __init__(self, temperature=0.7, role: str = ROLE_UTILITY, **llm_extra):
self.role = role
def __init__(self, temperature=0.7):
self.llm = OllamaLLM(
**ollama_llm_kwargs(
role=role,
temperature=temperature,
top_k=50,
top_p=0.9,
repeat_penalty=1.1,
**llm_extra,
num_ctx=4096,
)
)
self.output_parser = StrOutputParser()
llm_be/chat_backend/services/chat_tenant_scope.py
-235
View File
@@ -1,235 +0,0 @@
"""Immutable per-turn company/workspace scope for chat + RAG.
Mirrors the abc_worker ChatTenantScope stove-pipe: resolve identity once,
validate conversation ownership, never derive tenant from an untrusted
conversation_id alone.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Optional
from django.contrib.auth.models import AnonymousUser
from rest_framework_simplejwt.exceptions import TokenError
from rest_framework_simplejwt.tokens import AccessToken
from chat_backend.models import Conversation, CustomUser, DocumentWorkspace
class ChatTenantScopeError(Exception):
"""Raised when chat tenant resolution or ownership validation fails."""
def __init__(self, message: str, *, code: str = "tenant_scope_denied"):
super().__init__(message)
self.message = message
self.code = code
@dataclass(frozen=True)
class ChatCompanyScope:
"""Frozen tenant identity for one websocket turn / RAG retrieval.
``company_id`` is null for solo users on a personal workspace (#55).
"""
user_id: int
company_id: Optional[int]
workspace_id: int
conversation_id: Optional[int] = None
def user_from_access_token(token: str) -> Optional[CustomUser]:
"""Resolve an active user from a SimpleJWT access token string."""
if not token or not isinstance(token, str):
return None
try:
access = AccessToken(token)
user_id = access.get("user_id")
if not user_id:
return None
return CustomUser.objects.filter(id=user_id, is_active=True).first()
except TokenError:
return None
def resolve_chat_user(
*,
email: Optional[str] = None,
token: Optional[str] = None,
authenticated_user=None,
conversation_id: Optional[int] = None,
) -> Optional[CustomUser]:
"""
Resolve the chat principal for a websocket turn.
Preference order:
1. Authenticated ASGI/session user
2. JWT access token (payload or query)
3. Client email (legacy FE path)
Does not fall back to conversation.user — that would bind identity to an
attacker-chosen conversation_id. ``conversation_id`` is accepted for API
compatibility but ignored for identity resolution.
"""
del conversation_id
if authenticated_user is not None and getattr(
authenticated_user, "is_authenticated", False
):
if isinstance(authenticated_user, CustomUser):
return authenticated_user
user = CustomUser.objects.filter(
id=authenticated_user.pk, is_active=True
).first()
if user:
return user
token_user = user_from_access_token(token) if token else None
if token_user:
return token_user
if email:
return CustomUser.objects.filter(email__iexact=email, is_active=True).first()
return None
def ensure_company_workspace(company) -> DocumentWorkspace:
"""Return ``company``'s document workspace, creating a default one if missing (#46).
Never uses a bare ``.get(company=...)`` — ``.order_by("id").first()`` picks
a stable single workspace even if duplicates exist, and ``get_or_create``
closes the race for brand-new companies that don't have one yet (so
document upload/list/detail views never 404 just because a workspace was
never explicitly created).
"""
if company is None:
raise ChatTenantScopeError(
"Company is required for a company workspace.",
code="company_missing",
)
workspace = (
DocumentWorkspace.objects.filter(company=company, user__isnull=True)
.order_by("id")
.first()
)
if workspace is not None:
return workspace
workspace, _ = DocumentWorkspace.objects.get_or_create(
company=company,
user=None,
defaults={"name": "Default"},
)
return workspace
def ensure_personal_workspace(user: CustomUser) -> DocumentWorkspace:
"""Return ``user``'s personal RAG workspace, creating one if missing (#55)."""
if user is None or not getattr(user, "id", None):
raise ChatTenantScopeError(
"Authenticated chat user is required.",
code="user_not_found",
)
workspace = (
DocumentWorkspace.objects.filter(user=user, company__isnull=True)
.order_by("id")
.first()
)
if workspace is not None:
return workspace
workspace, _ = DocumentWorkspace.objects.get_or_create(
user=user,
defaults={"name": "Personal", "company": None},
)
return workspace
def ensure_workspace_for_user(user: CustomUser) -> DocumentWorkspace:
"""Company workspace when attached; otherwise personal workspace (#55)."""
if getattr(user, "company_id", None):
return ensure_company_workspace(user.company)
return ensure_personal_workspace(user)
def resolve_chat_company_scope(
user: CustomUser,
conversation_id: Optional[int] = None,
) -> ChatCompanyScope:
"""
Build an immutable company/workspace scope for ``user``.
When ``conversation_id`` is set, require ``conversation.user_id == user.id``
and that the conversation owner's company matches the user's company.
Users without a company resolve to a personal workspace (#55).
"""
if user is None or not getattr(user, "id", None):
raise ChatTenantScopeError(
"Authenticated chat user is required.",
code="user_not_found",
)
if conversation_id is not None:
conversation = (
Conversation.objects.select_related("user")
.filter(id=conversation_id, deleted=False)
.first()
)
if conversation is None:
raise ChatTenantScopeError(
"Conversation was not found.",
code="conversation_not_found",
)
if conversation.user_id != user.id:
raise ChatTenantScopeError(
"Conversation does not belong to the authenticated user.",
code="conversation_forbidden",
)
owner_company_id = getattr(conversation.user, "company_id", None)
if owner_company_id != user.company_id:
raise ChatTenantScopeError(
"Conversation company does not match the authenticated user.",
code="conversation_forbidden",
)
workspace = ensure_workspace_for_user(user)
return ChatCompanyScope(
user_id=user.id,
company_id=user.company_id,
workspace_id=workspace.id,
conversation_id=conversation_id,
)
def create_conversation_for_user(user: CustomUser, title: str) -> int:
"""Create a conversation owned by ``user`` and return its id."""
conversation = Conversation.objects.create(title=title, user=user)
return conversation.id
def get_workspace_for_scope(scope: ChatCompanyScope) -> DocumentWorkspace:
"""Load workspace rows only when they match the frozen scope keys."""
try:
if scope.company_id is not None:
return DocumentWorkspace.objects.get(
id=scope.workspace_id, company_id=scope.company_id
)
return DocumentWorkspace.objects.get(
id=scope.workspace_id,
user_id=scope.user_id,
company_id__isnull=True,
)
except DocumentWorkspace.DoesNotExist as exc:
raise ChatTenantScopeError(
"Scoped document workspace was not found.",
code="workspace_missing",
) from exc
def asgi_user_or_none(scope_user):
"""Return an authenticated user from Channels scope, else None."""
if scope_user is None or isinstance(scope_user, AnonymousUser):
return None
if getattr(scope_user, "is_authenticated", False):
return scope_user
return None
llm_be/chat_backend/services/data_analysis_service.py
+1 -1
View File
@@ -55,7 +55,7 @@ Answer:"""
}
| self.prompt
| self.llm
# No StrOutputParser: keep Ollama generation_info token counts.
| self.output_parser
)
def _get_dataframe_summary(self, df: pd.DataFrame) -> str:
llm_be/chat_backend/services/drive_sync.py
-524
View File
@@ -1,524 +0,0 @@
"""Google Drive / Microsoft OneDrive & SharePoint sync into RAG Documents (#48-#52).
``sync_connection`` is the single entry point used by the API sync endpoint,
the ``sync_drive_connections`` management command, and the webhook stubs.
"""
from __future__ import annotations
import logging
from dataclasses import dataclass, field
from datetime import timedelta
from typing import Any
import httpx
from django.conf import settings
from django.core.files.base import ContentFile
from django.utils import timezone
from chat_backend.models import Document, DriveConnection
from chat_backend.services.chat_tenant_scope import (
ensure_company_workspace,
ensure_personal_workspace,
)
from chat_backend.services.rag_services import AsyncRAGService
logger = logging.getLogger(__name__)
GOOGLE_TOKEN_URL = "https://oauth2.googleapis.com/token"
GOOGLE_FILES_URL = "https://www.googleapis.com/drive/v3/files"
GOOGLE_FOLDER_MIME = "application/vnd.google-apps.folder"
# Google-native docs must be exported to a downloadable format (#48).
GOOGLE_EXPORT_MIME_MAP: dict[str, tuple[str, str]] = {
"application/vnd.google-apps.document": (
"application/vnd.openxmlformats-officedocument.wordprocessingml.document",
".docx",
),
"application/vnd.google-apps.spreadsheet": (
"application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
".xlsx",
),
"application/vnd.google-apps.presentation": ("application/pdf", ".pdf"),
}
MICROSOFT_TOKEN_URL_TMPL = "https://login.microsoftonline.com/{tenant}/oauth2/v2.0/token"
GRAPH_BASE_URL = "https://graph.microsoft.com/v1.0"
HTTP_TIMEOUT_SECONDS = 30.0
class DriveSyncError(Exception):
"""Raised for any unrecoverable failure syncing one connection."""
def __init__(self, code: str, message: str = ""):
self.code = code
self.message = message or code
super().__init__(self.message)
@dataclass
class RemoteFile:
id: str
name: str
mime_type: str
etag: str
size: int | None = None
# Microsoft company (SharePoint) downloads are scoped to a site id.
context_id: str = field(default="")
def _microsoft_tenant() -> str:
return settings.MICROSOFT_OAUTH_TENANT or "common"
def _document_source(connection: DriveConnection) -> str:
if connection.provider == DriveConnection.Provider.GOOGLE:
return (
Document.Source.GOOGLE_SHARED_DRIVE
if connection.kind == DriveConnection.Kind.COMPANY
else Document.Source.GOOGLE_DRIVE
)
return (
Document.Source.SHAREPOINT
if connection.kind == DriveConnection.Kind.COMPANY
else Document.Source.ONEDRIVE
)
# --- Token refresh ---------------------------------------------------------
def refresh_google_token(connection: DriveConnection) -> str:
if not connection.refresh_token:
raise DriveSyncError(
"missing_refresh_token", "No refresh token stored for this connection."
)
with httpx.Client(timeout=HTTP_TIMEOUT_SECONDS) as client:
response = client.post(
GOOGLE_TOKEN_URL,
data={
"client_id": settings.GOOGLE_OAUTH_CLIENT_ID,
"client_secret": settings.GOOGLE_OAUTH_CLIENT_SECRET,
"refresh_token": connection.refresh_token,
"grant_type": "refresh_token",
},
)
if response.status_code >= 400:
raise DriveSyncError(
"token_refresh_failed", f"Google token refresh failed: {response.text}"
)
data = response.json()
access_token = data.get("access_token") or ""
if not access_token:
raise DriveSyncError(
"token_refresh_failed", "Google refresh did not return an access token."
)
connection.access_token = access_token
expires_in = data.get("expires_in")
if expires_in:
connection.token_expires_at = timezone.now() + timedelta(seconds=int(expires_in))
connection.save(update_fields=["access_token", "token_expires_at", "last_modified"])
return access_token
def refresh_microsoft_token(connection: DriveConnection) -> str:
if not connection.refresh_token:
raise DriveSyncError(
"missing_refresh_token", "No refresh token stored for this connection."
)
token_url = MICROSOFT_TOKEN_URL_TMPL.format(tenant=_microsoft_tenant())
with httpx.Client(timeout=HTTP_TIMEOUT_SECONDS) as client:
response = client.post(
token_url,
data={
"client_id": settings.MICROSOFT_OAUTH_CLIENT_ID,
"client_secret": settings.MICROSOFT_OAUTH_CLIENT_SECRET,
"refresh_token": connection.refresh_token,
"grant_type": "refresh_token",
"scope": connection.scopes or "offline_access Files.Read.All",
},
)
if response.status_code >= 400:
raise DriveSyncError(
"token_refresh_failed", f"Microsoft token refresh failed: {response.text}"
)
data = response.json()
access_token = data.get("access_token") or ""
if not access_token:
raise DriveSyncError(
"token_refresh_failed", "Microsoft refresh did not return an access token."
)
connection.access_token = access_token
if data.get("refresh_token"):
connection.refresh_token = data["refresh_token"]
expires_in = data.get("expires_in")
if expires_in:
connection.token_expires_at = timezone.now() + timedelta(seconds=int(expires_in))
connection.save(
update_fields=["access_token", "refresh_token", "token_expires_at", "last_modified"]
)
return access_token
def ensure_fresh_token(connection: DriveConnection) -> str:
"""Return a usable access token, refreshing when expired or close to it."""
if (
connection.access_token
and connection.token_expires_at
and connection.token_expires_at > timezone.now() + timedelta(minutes=2)
):
return connection.access_token
if connection.provider == DriveConnection.Provider.GOOGLE:
return refresh_google_token(connection)
return refresh_microsoft_token(connection)
# --- Google Drive listing/download -----------------------------------------
def _google_files_page(client: httpx.Client, headers: dict, params: dict) -> list[RemoteFile]:
response = client.get(GOOGLE_FILES_URL, headers=headers, params=params)
if response.status_code >= 400:
raise DriveSyncError("list_failed", f"Google Drive list failed: {response.text}")
data = response.json()
files: list[RemoteFile] = []
for item in data.get("files", []):
if item.get("mimeType") == GOOGLE_FOLDER_MIME:
continue
files.append(
RemoteFile(
id=item["id"],
name=item.get("name") or item["id"],
mime_type=item.get("mimeType", ""),
etag=item.get("md5Checksum") or item.get("modifiedTime") or "",
size=int(item["size"]) if item.get("size") else None,
)
)
return files
def _list_google_files(connection: DriveConnection, access_token: str) -> list[RemoteFile]:
"""List files under selected folders/shared drives, or "My Drive" root."""
headers = {"Authorization": f"Bearer {access_token}"}
parents = connection.selected_resource_ids or []
files: list[RemoteFile] = []
common_params = {
"fields": "files(id,name,mimeType,md5Checksum,size,modifiedTime)",
"pageSize": 100,
"supportsAllDrives": "true",
"includeItemsFromAllDrives": "true",
}
with httpx.Client(timeout=HTTP_TIMEOUT_SECONDS) as client:
if not parents:
params = {**common_params, "q": "'root' in parents and trashed = false"}
files.extend(_google_files_page(client, headers, params))
else:
for parent_id in parents:
params = {
**common_params,
"q": f"'{parent_id}' in parents and trashed = false",
}
if connection.kind == DriveConnection.Kind.COMPANY:
params["corpora"] = "drive"
params["driveId"] = parent_id
files.extend(_google_files_page(client, headers, params))
return files
def _download_google_file(
client: httpx.Client, headers: dict, remote: RemoteFile
) -> tuple[bytes, str, str]:
export = GOOGLE_EXPORT_MIME_MAP.get(remote.mime_type)
if export:
export_mime, suffix = export
filename = remote.name if remote.name.endswith(suffix) else f"{remote.name}{suffix}"
response = client.get(
f"{GOOGLE_FILES_URL}/{remote.id}/export",
headers=headers,
params={"mimeType": export_mime},
)
content_type = export_mime
else:
filename = remote.name
content_type = remote.mime_type or "application/octet-stream"
response = client.get(
f"{GOOGLE_FILES_URL}/{remote.id}",
headers=headers,
params={"alt": "media", "supportsAllDrives": "true"},
)
if response.status_code >= 400:
raise DriveSyncError(
"download_failed", f"Google Drive download failed: {response.text}"
)
return response.content, filename, content_type
# --- Microsoft Graph listing/download --------------------------------------
def _graph_children(client: httpx.Client, headers: dict, url: str) -> list[RemoteFile]:
files: list[RemoteFile] = []
while url:
response = client.get(url, headers=headers)
if response.status_code >= 400:
raise DriveSyncError(
"list_failed", f"Microsoft Graph list failed: {response.text}"
)
data = response.json()
for item in data.get("value", []):
if "folder" in item:
continue
files.append(
RemoteFile(
id=item["id"],
name=item.get("name") or item["id"],
mime_type=(item.get("file") or {}).get("mimeType", ""),
etag=item.get("eTag") or item.get("cTag") or "",
size=item.get("size"),
)
)
url = data.get("@odata.nextLink")
return files
def _list_microsoft_files(connection: DriveConnection, access_token: str) -> list[RemoteFile]:
headers = {"Authorization": f"Bearer {access_token}"}
selected = connection.selected_resource_ids or []
files: list[RemoteFile] = []
with httpx.Client(timeout=HTTP_TIMEOUT_SECONDS) as client:
if connection.kind == DriveConnection.Kind.COMPANY:
# Company (SharePoint) sync requires explicit site selection (#50/#51).
for site_id in selected:
url = f"{GRAPH_BASE_URL}/sites/{site_id}/drive/root/children"
for remote in _graph_children(client, headers, url):
remote.context_id = site_id
files.append(remote)
elif not selected:
url = f"{GRAPH_BASE_URL}/me/drive/root/children"
files.extend(_graph_children(client, headers, url))
else:
for folder_id in selected:
url = f"{GRAPH_BASE_URL}/me/drive/items/{folder_id}/children"
files.extend(_graph_children(client, headers, url))
return files
def _download_microsoft_file(
client: httpx.Client, headers: dict, connection: DriveConnection, remote: RemoteFile
) -> tuple[bytes, str, str]:
if connection.kind == DriveConnection.Kind.COMPANY and remote.context_id:
url = f"{GRAPH_BASE_URL}/sites/{remote.context_id}/drive/items/{remote.id}/content"
else:
url = f"{GRAPH_BASE_URL}/me/drive/items/{remote.id}/content"
response = client.get(url, headers=headers)
if response.status_code >= 400:
raise DriveSyncError(
"download_failed", f"Microsoft download failed: {response.text}"
)
content_type = response.headers.get("content-type", "application/octet-stream")
return response.content, remote.name, content_type
def _download_file(
client: httpx.Client, headers: dict, connection: DriveConnection, remote: RemoteFile
) -> tuple[bytes, str, str]:
if connection.provider == DriveConnection.Provider.GOOGLE:
return _download_google_file(client, headers, remote)
return _download_microsoft_file(client, headers, connection, remote)
def _list_remote_files(connection: DriveConnection, access_token: str) -> list[RemoteFile]:
if connection.provider == DriveConnection.Provider.GOOGLE:
return _list_google_files(connection, access_token)
return _list_microsoft_files(connection, access_token)
# --- Sync entry point --------------------------------------------------------
def sync_connection(connection: DriveConnection) -> dict[str, Any]:
"""Refresh tokens, list selected resources, and reconcile Documents (#48-#51).
Downloads (or exports, for Google-native docs) each remote file, creates/
updates the matching ``Document`` row, ingests it via ``AsyncRAGService``,
and removes ``Document`` rows whose remote file was deleted upstream.
"""
result: dict[str, Any] = {"added": 0, "updated": 0, "removed": 0, "failed": []}
connection.last_sync_status = DriveConnection.SyncStatus.PENDING
connection.sync_total = 0
connection.sync_processed = 0
connection.sync_added = 0
connection.sync_updated = 0
connection.sync_failed = 0
connection.save(
update_fields=[
"last_sync_status",
"sync_total",
"sync_processed",
"sync_added",
"sync_updated",
"sync_failed",
"last_modified",
]
)
def _persist_progress(**extra: Any) -> None:
fields = [
"sync_total",
"sync_processed",
"sync_added",
"sync_updated",
"sync_failed",
"last_modified",
*extra.keys(),
]
for key, value in extra.items():
setattr(connection, key, value)
connection.save(update_fields=list(dict.fromkeys(fields)))
try:
access_token = ensure_fresh_token(connection)
remote_files = _list_remote_files(connection, access_token)
connection.sync_total = len(remote_files)
_persist_progress()
if connection.kind == DriveConnection.Kind.PERSONAL:
if connection.user_id is None:
raise DriveSyncError(
"invalid_connection",
"Personal Drive connections require an owning user.",
)
workspace = ensure_personal_workspace(connection.user)
else:
if connection.company_id is None:
raise DriveSyncError(
"invalid_connection",
"Company Drive connections require a company.",
)
workspace = ensure_company_workspace(connection.company)
source = _document_source(connection)
remote_by_id = {remote.id: remote for remote in remote_files}
existing_docs = {
document.remote_file_id: document
for document in Document.objects.filter(drive_connection=connection)
}
rag_service = AsyncRAGService()
headers = {"Authorization": f"Bearer {access_token}"}
with httpx.Client(timeout=60.0) as client:
for remote in remote_files:
existing = existing_docs.get(remote.id)
if existing is not None and existing.remote_etag == remote.etag:
connection.sync_processed += 1
_persist_progress()
continue
try:
content, filename, _content_type = _download_file(
client, headers, connection, remote
)
except DriveSyncError as exc:
result["failed"].append({"file": remote.name, "error": exc.message})
connection.sync_failed += 1
connection.sync_processed += 1
_persist_progress()
continue
if existing is not None:
existing.file.save(filename, ContentFile(content), save=False)
existing.remote_etag = remote.etag
existing.remote_name = remote.name
existing.sync_error = ""
existing.processed = False
existing.active = True
existing.save()
document = existing
result["updated"] += 1
connection.sync_updated += 1
else:
document = Document.objects.create(
workspace=workspace,
source=source,
drive_connection=connection,
remote_file_id=remote.id,
remote_etag=remote.etag,
remote_name=remote.name,
active=True,
)
document.file.save(filename, ContentFile(content), save=True)
result["added"] += 1
connection.sync_added += 1
try:
ingest_result = rag_service.add_files_to_store(
[
(
document.file,
document.file.name,
workspace.id,
document.id,
document.active,
)
],
workspace_id=workspace.id,
source=source,
)
document.processed = True
document.sync_error = (
str(ingest_result.get("failed_files"))
if ingest_result.get("failed_files")
else ""
)
document.save(update_fields=["processed", "sync_error", "last_modified"])
except Exception as exc: # keep syncing remaining files
logger.exception(
"RAG ingest failed for document=%s connection=%s",
document.id,
connection.id,
)
document.sync_error = str(exc)
document.save(update_fields=["sync_error", "last_modified"])
result["failed"].append({"file": remote.name, "error": str(exc)})
connection.sync_failed += 1
connection.sync_processed += 1
_persist_progress()
for remote_id, document in existing_docs.items():
if remote_id not in remote_by_id:
document.delete()
result["removed"] += 1
connection.last_sync_status = DriveConnection.SyncStatus.OK
connection.last_sync_error = ""
except DriveSyncError as exc:
logger.warning(
"Drive sync failed for connection=%s (%s): %s", connection.id, exc.code, exc.message
)
connection.last_sync_status = DriveConnection.SyncStatus.ERROR
connection.last_sync_error = exc.message
result["error"] = exc.message
except Exception as exc:
logger.exception("Unexpected Drive sync failure for connection=%s", connection.id)
connection.last_sync_status = DriveConnection.SyncStatus.ERROR
connection.last_sync_error = str(exc)
result["error"] = str(exc)
finally:
connection.last_sync_at = timezone.now()
connection.save(
update_fields=[
"last_sync_status",
"last_sync_error",
"last_sync_at",
"sync_total",
"sync_processed",
"sync_added",
"sync_updated",
"sync_failed",
"last_modified",
]
)
return result
llm_be/chat_backend/services/grounded_chat.py
-164
View File
@@ -1,164 +0,0 @@
"""Apply always-on grounded retrieval for a chat turn (#62).
Shared by ``consumers`` and ``consumers_graph`` so both paths stay in sync.
"""
from __future__ import annotations
import logging
from collections.abc import Awaitable, Callable
from dataclasses import dataclass, field
from typing import Any
from asgiref.sync import sync_to_async
from django.conf import settings
from chat_backend.services.assistant_identity import RETRIEVAL_FAILED_MESSAGE
from chat_backend.services.grounding_decider import (
GroundingDecision,
grounding_decider,
)
from chat_backend.services.llm_service import build_chat_service
from chat_backend.services.search import (
SearchUnavailable,
format_sources_block,
search_and_rank,
)
from chat_backend.services.search.base import SearchResult
from chat_backend.services.status_context import emit_status
from chat_backend.services.ws_frames import citations_frame, status_frame
logger = logging.getLogger(__name__)
StatusCallback = Callable[[str, str | None], Awaitable[None]]
# Re-export for existing imports in consumers.
__all__ = [
"GroundedTurnResult",
"prepare_grounded_chat",
"citations_frame",
"status_frame",
]
@dataclass
class GroundedTurnResult:
"""Outcome of the grounding + optional retrieval step."""
generator: Any = None
error: dict | None = None
citations: list[dict] = field(default_factory=list)
grounded: bool = False
decision: GroundingDecision | None = None
model_name: str = ""
def _citations_from_results(results: list[SearchResult]) -> list[dict]:
return [r.to_citation(i) for i, r in enumerate(results, start=1)]
async def _emit(
on_status: StatusCallback | None,
stage: str,
detail: str | None = None,
) -> None:
try:
if on_status is not None:
await on_status(stage, detail)
else:
await emit_status(stage, detail)
except Exception:
logger.exception("status callback failed for stage=%s", stage)
async def prepare_grounded_chat(
*,
message: str,
messages: list,
model_name: str | None,
conversation_id: int,
on_status: StatusCallback | None = None,
use_conversation_context: bool = False,
) -> GroundedTurnResult:
"""Decide grounding, retrieve, and return an AsyncLLMService generator.
When retrieval is required but every provider fails, returns an ``error``
dict instead of falling back to parametric generation (#62 AC).
``on_status(stage, detail)`` emits live activity frames (#96) when provided;
otherwise uses the context-bound emitter from :mod:`status_context`.
``use_conversation_context`` gates prior-turn history in the LLM prompt
(#33); default off (opt-in).
"""
gen_kwargs = {"use_conversation_context": use_conversation_context}
internet = getattr(settings, "ALLOW_INTERNET_ACCESS", False)
if not internet:
await _emit(on_status, "refining")
service = build_chat_service(model_name=model_name, grounded=False)
return GroundedTurnResult(
generator=service.generate_response(
messages, message, conversation_id, **gen_kwargs
),
model_name=service.model_name,
)
await _emit(on_status, "evaluating")
decision = await grounding_decider.decide_async(message)
if not decision.needs_retrieval:
await _emit(on_status, "refining")
service = build_chat_service(model_name=model_name, grounded=False)
return GroundedTurnResult(
generator=service.generate_response(
messages, message, conversation_id, **gen_kwargs
),
decision=decision,
model_name=service.model_name,
)
queries = decision.queries or [message]
detail = "; ".join(queries[:3])
await _emit(on_status, "searching", detail)
try:
results = await sync_to_async(search_and_rank, thread_sensitive=False)(
queries,
temporal=decision.temporal,
)
except SearchUnavailable as exc:
logger.warning(
"Grounded retrieval failed for %r (queries=%s): %s",
message,
decision.queries,
exc,
)
return GroundedTurnResult(
error={
"type": "error",
"code": "search_unavailable",
"content": RETRIEVAL_FAILED_MESSAGE,
},
decision=decision,
grounded=True,
)
await _emit(on_status, "reading_sources")
sources_block = format_sources_block(results)
# Keep sources out of the mutable history list — AsyncLLMService injects
# them via {sources}. (Older code appended a HumanMessage; that double-
# rendered into the prompt.)
await _emit(on_status, "refining")
service = build_chat_service(
model_name=model_name,
grounded=True,
sources_block=sources_block,
)
return GroundedTurnResult(
generator=service.generate_response(
messages, message, conversation_id, **gen_kwargs
),
citations=_citations_from_results(results),
grounded=True,
decision=decision,
model_name=service.model_name,
)
llm_be/chat_backend/services/grounding_decider.py
-203
View File
@@ -1,203 +0,0 @@
"""Grounding decision: retrieval-on-unless-unnecessary (#62).
Fails open — parse/timeout/exception ⇒ needs_retrieval=True. A deterministic
temporal-marker pre-pass forces retrieval regardless of the model.
"""
from __future__ import annotations
import json
import logging
import re
from dataclasses import dataclass, field
from langchain_core.prompts import ChatPromptTemplate
from chat_backend.ollama_config import ROLE_UTILITY
from chat_backend.services.base_service import BaseService
logger = logging.getLogger(__name__)
# Years at/after common small-model cutoffs force live retrieval.
_TRAINING_CUTOFF_YEAR = 2024
_TEMPORAL_PATTERNS = (
r"\blatest\b",
r"\bcurrent\b",
r"\btoday\b",
r"\bnow\b",
r"\bthis year\b",
r"\bthis week\b",
r"\bthis month\b",
r"\bbreaking\b",
r"\bright now\b",
r"\bas of\b",
r"\bdid\b.+\byet\b",
r"\bhave\b.+\byet\b",
r"\bwho won\b",
r"\bstock price\b",
r"\bweather\b",
rf"\b(?:19|20)\d{{2}}\b", # any year mention — keep broad; model still helps
)
_TEMPORAL_RE = re.compile("|".join(_TEMPORAL_PATTERNS), re.IGNORECASE)
_YEAR_RE = re.compile(r"\b((?:19|20)\d{2})\b")
@dataclass
class GroundingDecision:
needs_retrieval: bool
reason: str = ""
queries: list[str] = field(default_factory=list)
temporal: bool = False
source: str = "model" # prepass | model | fail_open
class GroundingDecider(BaseService):
def __init__(self):
super().__init__(temperature=0.0, role=ROLE_UTILITY)
self.prompt = ChatPromptTemplate.from_messages(
[
(
"system",
"""You decide whether a user message needs live web retrieval.
Return ONLY compact JSON with keys:
needs_retrieval (boolean),
reason (short string),
queries (array of 1-3 focused search queries).
Bias TOWARD retrieval. Set needs_retrieval=true for ANY question about:
- a real person, organisation, product, price, event, date, or statistic
- anything that can change over time or after a model training cutoff
- news, sports, celebrity, politics, weather, stock prices
Set needs_retrieval=false ONLY when the message is fully self-contained:
creative writing, math, code, chit-chat, or a pure follow-up on text already
in the conversation that needs no external facts.
When needs_retrieval=true, produce focused search queries (not the raw user
message). Example: "did Taylor Swift get married" →
["Taylor Swift Travis Kelce wedding date", "Taylor Swift married 2026"].
""",
),
("human", "{prompt}"),
]
)
self.chain = self.prompt | self.llm
def temporal_prepass(self, prompt: str) -> GroundingDecision | None:
text = (prompt or "").strip()
if not text:
return GroundingDecision(
needs_retrieval=False,
reason="empty prompt",
source="prepass",
)
year_hits = [int(y) for y in _YEAR_RE.findall(text)]
forces = bool(_TEMPORAL_RE.search(text)) or any(
y >= _TRAINING_CUTOFF_YEAR for y in year_hits
)
if not forces:
return None
return GroundingDecision(
needs_retrieval=True,
reason="temporal marker / post-cutoff year",
queries=[text],
temporal=True,
source="prepass",
)
def _parse(self, raw: str, fallback_query: str) -> GroundingDecision:
text = (raw or "").strip()
# Strip markdown fences if the small model wraps JSON.
if text.startswith("```"):
text = re.sub(r"^```(?:json)?\s*", "", text)
text = re.sub(r"\s*```$", "", text)
try:
start = text.find("{")
end = text.rfind("}")
if start < 0 or end < 0:
raise ValueError("no JSON object")
payload = json.loads(text[start : end + 1])
except (ValueError, json.JSONDecodeError) as exc:
logger.warning("GroundingDecider parse fail (%s); failing open", exc)
return GroundingDecision(
needs_retrieval=True,
reason=f"unparseable decision: {exc}",
queries=[fallback_query],
temporal=True,
source="fail_open",
)
needs = bool(payload.get("needs_retrieval", True))
queries = payload.get("queries") or []
if not isinstance(queries, list):
queries = [str(queries)]
queries = [str(q).strip() for q in queries if str(q).strip()][:3]
if needs and not queries:
queries = [fallback_query]
return GroundingDecision(
needs_retrieval=needs,
reason=str(payload.get("reason") or ""),
queries=queries,
temporal=needs,
source="model",
)
async def decide_async(self, prompt: str) -> GroundingDecision:
pre = self.temporal_prepass(prompt)
# Even on temporal prepass, ask the model for better queries when possible.
try:
raw = await self.chain.ainvoke({"prompt": prompt})
if hasattr(raw, "content"):
raw = raw.content
decision = self._parse(str(raw), prompt)
except Exception as exc:
logger.warning("GroundingDecider LLM failed (%s); failing open", exc)
decision = GroundingDecision(
needs_retrieval=True,
reason=f"llm error: {exc}",
queries=[prompt],
temporal=True,
source="fail_open",
)
if pre and pre.needs_retrieval:
# Pre-pass wins on needs_retrieval; keep model queries when present.
return GroundingDecision(
needs_retrieval=True,
reason=pre.reason,
queries=decision.queries or pre.queries,
temporal=True,
source="prepass",
)
return decision
def decide(self, prompt: str) -> GroundingDecision:
pre = self.temporal_prepass(prompt)
try:
raw = self.chain.invoke({"prompt": prompt})
if hasattr(raw, "content"):
raw = raw.content
decision = self._parse(str(raw), prompt)
except Exception as exc:
logger.warning("GroundingDecider LLM failed (%s); failing open", exc)
decision = GroundingDecision(
needs_retrieval=True,
reason=f"llm error: {exc}",
queries=[prompt],
temporal=True,
source="fail_open",
)
if pre and pre.needs_retrieval:
return GroundingDecision(
needs_retrieval=True,
reason=pre.reason,
queries=decision.queries or pre.queries,
temporal=True,
source="prepass",
)
return decision
grounding_decider = GroundingDecider()
llm_be/chat_backend/services/llm_service.py
+54 -132
View File
@@ -1,51 +1,28 @@
from abc import ABC, abstractmethod
from typing import AsyncGenerator, Generator, Optional
# from langchain_community.llms import Ollama
from langchain_ollama import OllamaLLM
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from django.conf import settings
from chat_backend.models import Conversation, Prompt
from chat_backend.ollama_config import (
ROLE_FAST,
ROLE_THINKING,
ollama_llm_kwargs,
ollama_model_for_role,
ollama_num_ctx_for_role,
resolve_chat_role,
)
from chat_backend.services.assistant_identity import (
ASSISTANT_SYSTEM_PROMPT,
GROUNDED_ANSWER_INSTRUCTIONS,
)
from chat_backend.services.prompt_budget import (
estimate_tokens,
format_history,
window_history,
)
from chat_backend.ollama_config import ollama_llm_kwargs
from chat_backend.services.assistant_identity import ASSISTANT_SYSTEM_PROMPT
class LLMService(ABC):
"""Abstract base class for LLM conversation services."""
def __init__(
self,
role: str = ROLE_THINKING,
temperature: float = 0.7,
grounded: bool = False,
):
self.role = role
self.grounded = grounded
self.model_name = ollama_model_for_role(role)
self.num_ctx = ollama_num_ctx_for_role(role)
def __init__(self):
self.llm = OllamaLLM(
**ollama_llm_kwargs(
role=role,
temperature=temperature,
temperature=0.7,
top_k=50,
top_p=0.9,
repeat_penalty=1.1,
num_ctx=4096,
)
)
self.output_parser = StrOutputParser()
@@ -68,8 +45,8 @@ class LLMService(ABC):
class SyncLLMService(LLMService):
"""Synchronous LLM conversation service."""
def __init__(self, role: str = ROLE_THINKING, temperature: float = 0.7):
super().__init__(role=role, temperature=temperature)
def __init__(self):
super().__init__()
self._setup_chain()
def _setup_chain(self):
@@ -108,130 +85,75 @@ class SyncLLMService(LLMService):
class AsyncLLMService(LLMService):
"""Asynchronous LLM conversation service."""
def __init__(
self,
role: str = ROLE_THINKING,
temperature: float = 0.7,
grounded: bool = False,
sources_block: str = "",
):
super().__init__(role=role, temperature=temperature, grounded=grounded)
self.sources_block = sources_block or ""
def __init__(self):
super().__init__()
self._setup_chain()
def _setup_chain(self):
"""Single history window + optional grounded sources (#62 Phase 3)."""
grounded_block = ""
if self.grounded:
grounded_block = (
f"\n\n{GROUNDED_ANSWER_INSTRUCTIONS}\n\n"
f"Live sources:\n{{sources}}\n"
)
"""Setup the conversation chain."""
template = f"""{ASSISTANT_SYSTEM_PROMPT}
{grounded_block}
Continue this conversation while maintaining context by providing a single helpful response.
Conversation history:
{{history}}
Continue this conversation while maintaining context by providing a single helpful response.
Current context: {{context}}
Latest message: {{query}}
Last 3 messages:
{{recent_history}}
Instructions:
- Carefully maintain all established context
- If referencing previous elements (like stories), preserve all details
- When asked to modify something, identify what's being modified
Latest message: {{query}}
Response:"""
Instructions:
- Carefully maintain all established context
- If referencing previous elements (like stories), preserve all details
- When asked to modify something, identify what's being modified
Response:"""
self.prompt = ChatPromptTemplate.from_template(template)
self.conversation_chain = (
{
"history": lambda x: x["history"],
"context":lambda x: x["conversation"],
"recent_history":lambda x: x['recent_conversation'],
"query": lambda x: x["query"],
"sources": lambda x: x.get("sources", ""),
}
| self.prompt
| self.llm
# No StrOutputParser: Ollama puts prompt_eval_count/eval_count on the
# final GenerationChunk.generation_info; the parser would drop it.
| self.output_parser
)
async def _format_history(self, conversation: list) -> str:
"""Async version of format conversation history."""
# prompts = list(
# await Prompt.objects.filter(conversation_id=conversation_id)
# .order_by("created")
# )
# return "\n".join(
# f"{'User' if prompt.is_user else 'AI'}: {prompt.text}" for prompt in prompts
# )
return "\n".join([f"{"User" if prompt.type=="human" else "AI"}: {prompt.text}" for prompt in conversation])
async def _get_recent_messages(self, conversation: list) -> str:
"""Async version of format conversation history."""
# prompts = list(
# await Prompt.objects.filter(conversation_id=conversation_id)
# .order_by("created")
# [-6:]
# )
# return "\n".join(
# f"{'User' if prompt.is_user else 'AI'}: {prompt.text}" for prompt in prompts
# )
return "\n".join([f"{"User" if prompt.type=="human" else "AI"}: {prompt.text}" for prompt in conversation])
async def generate_response(
self,
conversation,
query: str,
conversation_id: int,
**kwargs,
self, conversation: Conversation, query: str, conversation_id: int, **kwargs
) -> AsyncGenerator[str, None]:
"""Generate response with async streaming support.
``conversation`` is the LangChain message list for this turn (not a
Django Conversation row). History is serialised exactly once and
trimmed oldest-first under the role's ``num_ctx`` budget. The sources
block (when grounded) is never trimmed.
"""
sources = self.sources_block or kwargs.get("sources_block", "") or ""
use_conversation_context = bool(
kwargs.get("use_conversation_context", False)
)
history_text = ""
if use_conversation_context:
reserved = (
estimate_tokens(ASSISTANT_SYSTEM_PROMPT)
+ estimate_tokens(
GROUNDED_ANSWER_INSTRUCTIONS if self.grounded else ""
)
+ estimate_tokens(sources)
+ estimate_tokens(query)
+ 256 # response headroom / instructions
)
# Leave ~40% of ctx for the completion.
history_budget = max(512, int(self.num_ctx * 0.55) - reserved)
# 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:")
)
]
# 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)
"""Generate response with async streaming support."""
chain_input = {
"query": query,
"history": history_text,
"sources": sources,
}
"conversation": await self._format_history(conversation),
"recent_conversation": await self._get_recent_messages(conversation[-6:])}
async for chunk in self.conversation_chain.astream(chain_input):
yield chunk
def build_chat_service(
*,
model_name: str | None,
grounded: bool = False,
sources_block: str = "",
) -> AsyncLLMService:
"""Factory: FE mode → role, factual turns → low temperature."""
role = resolve_chat_role(model_name)
temperature = 0.3 if grounded else 0.7
return AsyncLLMService(
role=role,
temperature=temperature,
grounded=grounded,
sources_block=sources_block,
)
llm_be/chat_backend/services/moderation_classifier.py
+2 -3
View File
@@ -3,7 +3,7 @@ from typing import Dict, Any
from langchain_core.prompts import ChatPromptTemplate
from langchain_ollama import OllamaLLM
from chat_backend.services.base_service import BaseService
from chat_backend.ollama_config import ROLE_UTILITY, ollama_llm_kwargs
from chat_backend.ollama_config import ollama_llm_kwargs
class ModerationLabel(Enum):
@@ -18,10 +18,9 @@ class ModerationClassifier(BaseService):
"""
def __init__(self):
super().__init__(temperature=0.1, role=ROLE_UTILITY)
super().__init__(temperature=0.1)
self.llm = OllamaLLM(
**ollama_llm_kwargs(
role=ROLE_UTILITY,
temperature=0.1, # Very low for strict moderation
top_k=10,
num_ctx=2048,
llm_be/chat_backend/services/prompt_budget.py
-63
View File
@@ -1,63 +0,0 @@
"""Prompt budgeting helpers (#62 Phase 3).
Approximate token counts with a chars/4 heuristic. Trim oldest history first;
never truncate the system prompt or the retrieved-sources block.
"""
from __future__ import annotations
from typing import Sequence
def estimate_tokens(text: str) -> int:
if not text:
return 0
return max(1, (len(text) + 3) // 4)
def window_history(
messages: Sequence,
*,
budget_tokens: int,
reserved_tokens: int = 0,
max_messages: int = 24,
) -> list:
"""Keep the newest messages that fit under ``budget_tokens - reserved``.
``messages`` are LangChain BaseMessage-like objects (``.type``, ``.text`` /
``.content``).
"""
if not messages:
return []
usable = max(0, budget_tokens - reserved_tokens)
selected: list = []
used = 0
for message in reversed(list(messages)[-max_messages:]):
text = getattr(message, "text", None)
if callable(text):
# property that looks callable in some versions — read content
text = getattr(message, "content", "")
if text is None:
text = getattr(message, "content", "") or ""
cost = estimate_tokens(str(text)) + 4 # role overhead
if selected and used + cost > usable:
break
selected.append(message)
used += cost
selected.reverse()
return selected
def format_message_line(message) -> str:
role = "User" if getattr(message, "type", "") == "human" else "AI"
text = getattr(message, "text", None)
if text is None or (callable(text) and not isinstance(text, str)):
text = getattr(message, "content", "") or ""
# BaseMessage.text is a property returning str; prefer it when string.
if not isinstance(text, str):
text = getattr(message, "content", "") or ""
return f"{role}: {text}"
def format_history(messages: Sequence) -> str:
return "\n".join(format_message_line(m) for m in messages)
llm_be/chat_backend/services/rag_services.py
+21 -139
View File
@@ -25,14 +25,7 @@ from chat_backend.models import Conversation, Prompt, DocumentWorkspace, Documen
from pathlib import Path
from chat_backend.services.base_service import BaseService
from chat_backend.services.assistant_identity import ASSISTANT_SYSTEM_PROMPT
from chat_backend.ollama_config import ollama_embed_model, ollama_embeddings_kwargs
import logging
logger = logging.getLogger(__name__)
class EmbeddingDimensionMismatch(RuntimeError):
"""Persisted Chroma collection dim != configured embedding model (#62)."""
from chat_backend.ollama_config import ollama_embeddings_kwargs
@database_sync_to_async
@@ -61,7 +54,6 @@ class RAGService(BaseService):
chunk_size=1000, chunk_overlap=200
)
self.vector_store = self._initialize_vector_store()
self._assert_embedding_dimensions()
# Supported file types and their loaders
self.loader_mapping = {
@@ -83,34 +75,6 @@ class RAGService(BaseService):
)
return vector_store
def _assert_embedding_dimensions(self) -> None:
"""Refuse mismatched collections loudly (#62)."""
try:
collection = self.vector_store._collection
count = collection.count()
except Exception as exc:
logger.warning("Could not inspect Chroma collection: %s", exc)
return
if not count:
return
try:
peek = collection.peek(limit=1)
embeddings = peek.get("embeddings") if isinstance(peek, dict) else None
if not embeddings:
return
stored_dim = len(embeddings[0])
probe = self.embedding_model.embed_query("dimension-check")
expected_dim = len(probe)
except Exception as exc:
logger.warning("Embedding dimension probe failed: %s", exc)
return
if stored_dim != expected_dim:
raise EmbeddingDimensionMismatch(
f"Chroma collection embedding dim is {stored_dim} but "
f"OLLAMA_EMBED_MODEL={ollama_embed_model()!r} produces "
f"{expected_dim}. Run: python manage.py reindex_embeddings"
)
def clear_vector_store(self):
"""Clear all vectors from the store"""
self.vector_store.delete_collection()
@@ -154,66 +118,24 @@ class RAGService(BaseService):
metadata={
"source": doc.file.name,
"workspace_id": doc.workspace_id,
"company_id": doc.workspace.company_id,
"document_id": doc.id,
"active": bool(doc.active),
},
)
if chunks:
self.vector_store.add_documents(chunks)
except Exception:
# Keep reindex/ingest moving; one bad file must not wipe progress.
logger.exception(
"Failed to ingest document_id=%s file=%s",
doc.id,
doc.file.name,
)
finally:
if os.path.exists(tmp_path):
os.unlink(tmp_path)
self.vector_store.persist()
return docs
def delete_document_vectors(self, document_id) -> None:
"""Remove every vector chunk belonging to a document (#45).
Used on ``Document`` delete so a re-uploaded/replaced document does
not leave stale chunks searchable, without re-ingesting the whole
workspace (the old post_delete behavior).
"""
if document_id is None:
return
self.vector_store.delete(where={"document_id": document_id})
self.vector_store.persist()
def set_document_active(self, document_id, active: bool) -> None:
"""Update the ``active`` metadata flag on a document's existing chunks (#45).
Called from the PATCH toggle so ``search_documents`` (which filters
on ``active=True``) immediately reflects the new state without
re-ingesting the file.
"""
if document_id is None:
return
existing = self.vector_store.get(where={"document_id": document_id})
ids = existing.get("ids") or []
if not ids:
return
metadatas = [
{**(meta or {}), "active": bool(active)}
for meta in existing.get("metadatas") or [{} for _ in ids]
]
self.vector_store._collection.update(ids=ids, metadatas=metadatas)
self.vector_store.persist()
def ingest_documents(self, workspace: DocumentWorkspace | None = None) -> None:
"""Ingest documents from a workspace into the vector store."""
print(f"Getting the Document via the workspace: {workspace}")
qs = Document.objects.select_related("workspace")
if workspace:
documents = list(qs.filter(workspace=workspace))
documents = [doc for doc in Document.objects.filter(workspace=workspace)]
else:
documents = list(qs.all())
documents = [doc for doc in Document.objects.all()]
print(f"Processing the documents : {documents}")
self._prepare_documents(documents)
@@ -243,7 +165,7 @@ class RAGService(BaseService):
def add_files_to_store(
self,
file_tupls: List, # (file_path_or_field, name, workspace_id[, document_id, active])
file_tupls: List, # (file_path_or_field, name, workspace_id)
workspace_id: str,
source: str = "upload",
save_dir: str = "data/uploads",
@@ -251,9 +173,7 @@ class RAGService(BaseService):
"""
Process and add files to vector store.
file_tupls entries: (path_str | Django FileField, original_name, workspace_id,
document_id, active). ``document_id``/``active`` are optional (default
``None``/``True``) for backward compatibility with older call sites.
file_tupls entries: (path_str | Django FileField, original_name, workspace_id)
Paths may be temp files; FileFields are materialized from DB storage.
"""
results = {"total_added": 0, "failed_files": [], "processed_files": []}
@@ -266,29 +186,17 @@ class RAGService(BaseService):
file_tuple[1],
file_tuple[2],
)
document_id = file_tuple[3] if len(file_tuple) > 3 else None
active = file_tuple[4] if len(file_tuple) > 4 else True
if isinstance(file_ref, str):
file_path = file_ref
else:
tmp_created = self._materialize_file_field(file_ref)
file_path = tmp_created
company_id = None
if ws_id is not None:
company_id = (
DocumentWorkspace.objects.filter(id=ws_id)
.values_list("company_id", flat=True)
.first()
)
metadata = {
"source": original_name,
"workspace_id": ws_id,
"company_id": company_id,
"original_filename": original_name,
"file_path": original_name,
"document_id": document_id,
"active": bool(active),
}
docs = self._load_and_split_documents(file_path, metadata)
@@ -312,24 +220,6 @@ class RAGService(BaseService):
self.vector_store.persist()
return results
def _workspace_filter(self, workspace: DocumentWorkspace) -> Dict[str, Any]:
"""Build a fail-closed Chroma metadata filter for one workspace.
``company_id`` is written on ingest for defense-in-depth / future dual
filters, but retrieval keys on ``workspace_id`` so older vectors without
``company_id`` metadata still match after deploy. Also excludes chunks
for documents toggled inactive (#45) so deactivating a document hides
it from retrieval immediately.
"""
if workspace is None or getattr(workspace, "id", None) is None:
raise ValueError("workspace is required for RAG retrieval")
return {
"$and": [
{"workspace_id": workspace.id},
{"active": True},
]
}
class SyncRAGService(RAGService):
"""Synchronous RAG service implementation."""
@@ -375,12 +265,9 @@ class SyncRAGService(RAGService):
def _retriever_with_history(self, input_dict: Dict[str, Any]) -> str:
"""Retrieve documents considering conversation history."""
query = input_dict["query"]
workspace = input_dict.get("workspace")
if workspace is None:
conversation = input_dict.get("conversation")
workspace = getattr(conversation, "workspace", None)
conversation = input_dict["conversation"]
relevant_docs = self.search_documents(query, workspace)
relevant_docs = self.search_documents(query, conversation.workspace)
if not relevant_docs:
print("didn't find any relevant docs")
return relevant_docs
@@ -390,9 +277,11 @@ class SyncRAGService(RAGService):
def search_documents(
self, query: str, workspace: Optional[DocumentWorkspace] = None, k: int = 4
) -> List[Document]:
"""Search relevant documents from the vector store (workspace required)."""
filter_dict = self._workspace_filter(workspace)
search_kwargs = {"k": k, "filter": filter_dict}
"""Search relevant documents from the vector store."""
filter_dict = {}
if workspace:
filter_dict["workspace_id"] = workspace.id
search_kwargs = {"k": k, "filter": filter_dict if filter_dict else None}
print(f"search_kwargs: {search_kwargs}")
retriever = self.vector_store.as_retriever(
search_type="similarity",
@@ -439,7 +328,7 @@ class AsyncRAGService(RAGService):
}
| self.prompt
| self.llm
# No StrOutputParser: keep Ollama generation_info token counts.
| StrOutputParser()
)
async def _format_history(self, conversation: Conversation) -> str:
@@ -469,14 +358,17 @@ class AsyncRAGService(RAGService):
async def search_documents(
self, query: str, workspace: Optional[DocumentWorkspace] = None, k: int = 4
) -> List[Document]:
"""Search relevant documents from the vector store (workspace required)."""
filter_dict = self._workspace_filter(workspace)
"""Search relevant documents from the vector store."""
filter_dict = {}
print(f"Do we have a workspace: {workspace}")
print(f"search_kwargs: {{'k': {k}, 'filter': {filter_dict}}}")
if workspace:
filter_dict["workspace_id"] = workspace.id
search_kwargs = {"k": k, "filter": filter_dict if filter_dict else None}
print(f"search_kwargs: {search_kwargs}")
retriever = self.vector_store.as_retriever(
search_type="mmr",
search_kwargs={"k": k, "filter": filter_dict},
search_kwargs={"k": k, "filter": filter_dict if filter_dict else None},
)
return await retriever.aget_relevant_documents(query)
@@ -488,21 +380,11 @@ class AsyncRAGService(RAGService):
**kwargs,
) -> AsyncGenerator[str, None]:
"""Generate response with streaming support."""
if workspace is None:
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 = {
"query": query,
"conversation": conversation,
"workspace": workspace,
"recent_conversation": recent,
"recent_conversation": await self._format_history(conversation),
}
async for chunk in self.rag_chain.astream(chain_input):
llm_be/chat_backend/services/search/__init__.py
-22
View File
@@ -1,22 +0,0 @@
"""Structured web-search providers (#62).
Providers return typed :class:`SearchResult` rows — never a flat concatenated
string. The facade in :mod:`chat_backend.services.search.service` picks the
configured primary provider and fails over to the secondary.
"""
from chat_backend.services.search.base import SearchResult
from chat_backend.services.search.service import (
SearchUnavailable,
format_sources_block,
get_search_service,
search_and_rank,
)
__all__ = [
"SearchResult",
"SearchUnavailable",
"format_sources_block",
"get_search_service",
"search_and_rank",
]
llm_be/chat_backend/services/search/base.py
-36
View File
@@ -1,36 +0,0 @@
"""Search provider protocol and result dataclass."""
from __future__ import annotations
from dataclasses import asdict, dataclass
from typing import Optional, Protocol, runtime_checkable
@dataclass(frozen=True)
class SearchResult:
title: str
url: str
snippet: str
published_at: Optional[str] = None
rank: int = 0
provider: str = ""
def to_citation(self, index: int) -> dict:
return {
"index": index,
"title": self.title,
"url": self.url,
"published_at": self.published_at,
}
def to_dict(self) -> dict:
return asdict(self)
@runtime_checkable
class SearchProvider(Protocol):
name: str
def search(self, query: str, *, max_results: int = 8) -> list[SearchResult]:
"""Return structured results for ``query``. Raise on hard failure."""
...
llm_be/chat_backend/services/search/ddgs_provider.py
-46
View File
@@ -1,46 +0,0 @@
"""DuckDuckGo (ddgs) search provider — failover for SearxNG (#62)."""
from __future__ import annotations
import logging
from typing import Any
from chat_backend.services.search.base import SearchResult
logger = logging.getLogger(__name__)
class DDGSProvider:
name = "ddgs"
def search(self, query: str, *, max_results: int = 8) -> list[SearchResult]:
try:
from ddgs import DDGS
except ImportError as exc: # pragma: no cover - dependency is declared
raise RuntimeError("ddgs package is not installed") from exc
try:
with DDGS() as ddgs:
raw: list[dict[str, Any]] = list(
ddgs.text(query, max_results=max_results)
)
except Exception as exc:
logger.warning("DDGS search failed for %r: %s", query, exc)
raise RuntimeError(f"DDGS unreachable: {exc}") from exc
results: list[SearchResult] = []
for idx, item in enumerate(raw or []):
url = (item.get("href") or item.get("link") or item.get("url") or "").strip()
if not url:
continue
results.append(
SearchResult(
title=(item.get("title") or url).strip(),
url=url,
snippet=(item.get("body") or item.get("snippet") or "").strip(),
published_at=item.get("date") or item.get("published") or None,
rank=idx,
provider=self.name,
)
)
return results
llm_be/chat_backend/services/search/ranking.py
-128
View File
@@ -1,128 +0,0 @@
"""Rank, dedupe, and rumour-filter search results (#62)."""
from __future__ import annotations
import re
from datetime import datetime
from typing import Iterable
from urllib.parse import urlparse
from chat_backend.services.search.base import SearchResult
_RUMOUR_MARKERS = (
"rumor",
"rumour",
"speculation",
"ai-generated",
"ai generated",
"blind item",
"fake",
"allegedly",
"unconfirmed",
)
_DATE_FORMATS = (
"%Y-%m-%d",
"%Y-%m-%dT%H:%M:%S",
"%Y-%m-%dT%H:%M:%SZ",
"%Y-%m-%dT%H:%M:%S%z",
"%b %d, %Y",
"%B %d, %Y",
"%d %b %Y",
"%d %B %Y",
)
def registrable_domain(url: str) -> str:
host = (urlparse(url).hostname or "").lower()
if host.startswith("www."):
host = host[4:]
parts = host.split(".")
if len(parts) >= 2:
return ".".join(parts[-2:])
return host
def parse_published_at(value: str | None) -> datetime | None:
if not value:
return None
text = value.strip()
if not text:
return None
# Prefer ISO-ish prefixes.
for fmt in _DATE_FORMATS:
try:
return datetime.strptime(text[: len(fmt) + 8], fmt)
except ValueError:
continue
iso = text.replace("Z", "+00:00")
try:
return datetime.fromisoformat(iso)
except ValueError:
return None
def is_rumour_heavy(result: SearchResult) -> bool:
haystack = f"{result.title} {result.snippet}".lower()
hits = sum(1 for marker in _RUMOUR_MARKERS if marker in haystack)
return hits >= 2 or ("ai-generated" in haystack and "fake" in haystack)
def rank_and_dedupe(
results: Iterable[SearchResult],
*,
temporal: bool = False,
max_results: int = 6,
) -> list[SearchResult]:
"""Deduplicate by domain, drop rumour-heavy rows when alternatives exist."""
seen_domains: set[str] = set()
kept: list[SearchResult] = []
rumour_bucket: list[SearchResult] = []
for result in results:
domain = registrable_domain(result.url)
if not domain or domain in seen_domains:
continue
seen_domains.add(domain)
if is_rumour_heavy(result):
rumour_bucket.append(result)
else:
kept.append(result)
# Only use rumour-heavy rows if we have nothing better.
if not kept and rumour_bucket:
kept = rumour_bucket
def sort_key(item: SearchResult):
published = parse_published_at(item.published_at)
# Prefer dated + recent when temporal; otherwise keep provider rank.
if temporal:
# Newer first; undated last.
stamp = published.timestamp() if published else float("-inf")
return (-stamp, item.rank)
has_date = 0 if published else 1
return (has_date, item.rank)
kept.sort(key=sort_key)
return [
SearchResult(
title=r.title,
url=r.url,
snippet=r.snippet,
published_at=r.published_at,
rank=i,
provider=r.provider,
)
for i, r in enumerate(kept[:max_results])
]
_HOST_RE = re.compile(r"^https?://([^/]+)", re.I)
def display_host(url: str) -> str:
match = _HOST_RE.match(url or "")
if not match:
return url or ""
host = match.group(1).lower()
return host[4:] if host.startswith("www.") else host
llm_be/chat_backend/services/search/searxng.py
-83
View File
@@ -1,83 +0,0 @@
"""SearxNG search provider (#62).
Hits a self-hosted SearxNG instance's JSON API. Deterministic, no third-party
rate limits, and returns per-result title/url/snippet/publishedDate.
"""
from __future__ import annotations
import logging
from typing import Any
from urllib.parse import urljoin
import requests
from django.conf import settings
from chat_backend.services.search.base import SearchResult
logger = logging.getLogger(__name__)
class SearxNGProvider:
name = "searxng"
def __init__(
self,
base_url: str | None = None,
timeout: float | None = None,
):
self.base_url = (
base_url
or getattr(settings, "SEARXNG_BASE_URL", "http://127.0.0.1:8080")
).rstrip("/")
self.timeout = float(
timeout
if timeout is not None
else getattr(settings, "SEARXNG_TIMEOUT_SECONDS", 8)
)
def search(self, query: str, *, max_results: int = 8) -> list[SearchResult]:
endpoint = urljoin(self.base_url + "/", "search")
params = {
"q": query,
"format": "json",
"language": "en",
}
try:
response = requests.get(
endpoint,
params=params,
timeout=self.timeout,
headers={"Accept": "application/json"},
)
response.raise_for_status()
payload: dict[str, Any] = response.json()
except requests.RequestException as exc:
logger.warning("SearxNG search failed for %r: %s", query, exc)
raise RuntimeError(f"SearxNG unreachable: {exc}") from exc
except ValueError as exc:
raise RuntimeError(f"SearxNG returned non-JSON: {exc}") from exc
results: list[SearchResult] = []
for idx, item in enumerate(payload.get("results") or []):
url = (item.get("url") or "").strip()
if not url:
continue
results.append(
SearchResult(
title=(item.get("title") or url).strip(),
url=url,
snippet=(item.get("content") or item.get("snippet") or "").strip(),
published_at=(
item.get("publishedDate")
or item.get("published_at")
or item.get("pubdate")
or None
),
rank=idx,
provider=self.name,
)
)
if len(results) >= max_results:
break
return results
llm_be/chat_backend/services/search/service.py
-145
View File
@@ -1,145 +0,0 @@
"""Search facade: primary + failover providers, concurrent multi-query (#62)."""
from __future__ import annotations
import logging
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Iterable
from django.conf import settings
from chat_backend.services.search.base import SearchProvider, SearchResult
from chat_backend.services.search.ddgs_provider import DDGSProvider
from chat_backend.services.search.ranking import display_host, rank_and_dedupe
from chat_backend.services.search.searxng import SearxNGProvider
logger = logging.getLogger(__name__)
class SearchUnavailable(RuntimeError):
"""Raised when every configured provider fails for a query set."""
def _build_provider(name: str) -> SearchProvider:
key = (name or "").lower()
if key in {"searxng", "searx"}:
return SearxNGProvider()
if key in {"ddgs", "duckduckgo", "ddg"}:
return DDGSProvider()
raise ValueError(f"Unknown search provider: {name!r}")
class SearchService:
def __init__(
self,
primary: SearchProvider | None = None,
failover: SearchProvider | None = None,
):
primary_name = getattr(settings, "SEARCH_PROVIDER", "searxng")
failover_name = getattr(settings, "SEARCH_FAILOVER_PROVIDER", "ddgs")
self.primary = primary or _build_provider(primary_name)
# Avoid wiring the same provider twice.
if failover is not None:
self.failover = failover
elif failover_name and failover_name.lower() != getattr(
self.primary, "name", ""
):
try:
self.failover = _build_provider(failover_name)
except ValueError:
self.failover = None
else:
self.failover = None
def search_one(self, query: str, *, max_results: int = 8) -> list[SearchResult]:
errors: list[str] = []
for provider in (self.primary, self.failover):
if provider is None:
continue
try:
return provider.search(query, max_results=max_results)
except Exception as exc:
errors.append(f"{getattr(provider, 'name', provider)}: {exc}")
logger.warning(
"Search provider %s failed for %r: %s",
getattr(provider, "name", provider),
query,
exc,
)
raise SearchUnavailable(
f"All search providers failed for {query!r}: {'; '.join(errors)}"
)
def search_many(
self,
queries: Iterable[str],
*,
max_results_per_query: int = 6,
temporal: bool = False,
max_results: int = 6,
) -> list[SearchResult]:
cleaned = [q.strip() for q in queries if q and str(q).strip()]
if not cleaned:
raise SearchUnavailable("No search queries provided")
collected: list[SearchResult] = []
failures = 0
with ThreadPoolExecutor(max_workers=min(3, len(cleaned))) as pool:
futures = {
pool.submit(
self.search_one, query, max_results=max_results_per_query
): query
for query in cleaned[:3]
}
for future in as_completed(futures):
query = futures[future]
try:
collected.extend(future.result())
except SearchUnavailable as exc:
failures += 1
logger.warning("Query %r exhausted providers: %s", query, exc)
if not collected:
raise SearchUnavailable(
f"All {failures} search queries failed; no results available"
)
return rank_and_dedupe(
collected, temporal=temporal, max_results=max_results
)
_service: SearchService | None = None
def get_search_service() -> SearchService:
global _service
if _service is None:
_service = SearchService()
return _service
def search_and_rank(
queries: Iterable[str],
*,
temporal: bool = False,
max_results: int = 6,
) -> list[SearchResult]:
return get_search_service().search_many(
queries, temporal=temporal, max_results=max_results
)
def format_sources_block(results: list[SearchResult]) -> str:
"""Numbered, dated, delimited context block for the LLM prompt."""
if not results:
return ""
lines: list[str] = []
for i, result in enumerate(results, start=1):
host = display_host(result.url)
date = result.published_at or "undated"
lines.append(f'[{i}] "{result.title}" — {host} — {date}')
if result.snippet:
lines.append(f" {result.snippet}")
lines.append(f" URL: {result.url}")
lines.append("")
return "\n".join(lines).rstrip()
llm_be/chat_backend/services/status_context.py
-34
View File
@@ -1,34 +0,0 @@
"""Context-local status emitter for WS activity frames (#96).
Consumers bind an async callback before running moderation / generation so
shared helpers (``prepare_grounded_chat``, graph nodes) can emit without
threading the sender through every signature.
"""
from __future__ import annotations
from collections.abc import Awaitable, Callable
from contextvars import ContextVar
from typing import Optional
StatusEmitter = Callable[[str, Optional[str]], Awaitable[None]]
_status_emitter: ContextVar[StatusEmitter | None] = ContextVar(
"chat_status_emitter", default=None
)
def set_status_emitter(emitter: StatusEmitter | None):
"""Bind the current-task status emitter; returns a token for reset."""
return _status_emitter.set(emitter)
def reset_status_emitter(token) -> None:
_status_emitter.reset(token)
async def emit_status(stage: str, detail: str | None = None) -> None:
emitter = _status_emitter.get()
if emitter is None:
return
await emitter(stage, detail)
llm_be/chat_backend/services/title_generator.py
+1 -2
View File
@@ -3,7 +3,7 @@ from langchain_core.prompts import ChatPromptTemplate
# from langchain_community.llms import Ollama
from langchain_ollama import OllamaLLM
from typing import Optional
from chat_backend.ollama_config import ROLE_UTILITY, ollama_llm_kwargs
from chat_backend.ollama_config import ollama_llm_kwargs
class TitleGenerator:
@@ -14,7 +14,6 @@ class TitleGenerator:
def __init__(self):
self.llm = OllamaLLM(
**ollama_llm_kwargs(
role=ROLE_UTILITY,
temperature=0.5, # Slightly creative but not too random
top_k=20,
num_ctx=2048, # Shorter context needed for titles
llm_be/chat_backend/services/tools/__init__.py
-31
View File
@@ -1,31 +0,0 @@
"""Agent tool registry and implementations (#63)."""
from chat_backend.services.tools.base import (
ToolBudget,
ToolBudgetExceeded,
ToolContext,
ToolTimeout,
run_tool,
truncate_output,
)
from chat_backend.services.tools.registry import (
TOOL_CATALOG,
TOOL_NAMES,
ToolSpec,
build_agent_tools,
suggest_tools,
)
__all__ = [
"ToolBudget",
"ToolBudgetExceeded",
"ToolContext",
"ToolTimeout",
"run_tool",
"truncate_output",
"TOOL_CATALOG",
"TOOL_NAMES",
"ToolSpec",
"build_agent_tools",
"suggest_tools",
]
llm_be/chat_backend/services/tools/base.py
-92
View File
@@ -1,92 +0,0 @@
"""Shared tool execution primitives: budget, timeout, output caps (#63).
Every agent tool goes through :func:`run_tool` so limits are enforced in one
place instead of duplicated per tool.
"""
from __future__ import annotations
import asyncio
import logging
from dataclasses import dataclass, field
from typing import Any, Awaitable, Callable
from django.conf import settings
logger = logging.getLogger(__name__)
class ToolBudgetExceeded(RuntimeError):
"""Raised when a run has used its allotted number of tool calls."""
class ToolTimeout(RuntimeError):
"""Raised when a single tool call exceeds its timeout."""
@dataclass
class ToolBudget:
"""Per-run tool-call budget, shared across every tool instance for a run."""
max_calls: int
calls_made: int = 0
def consume(self) -> None:
if self.calls_made >= self.max_calls:
raise ToolBudgetExceeded(
f"Tool-call budget exhausted ({self.max_calls} calls per run)."
)
self.calls_made += 1
@dataclass
class ToolContext:
"""Tenant scope + limits threaded into every tool for one agent run."""
scope: Any # ChatCompanyScope — kept loosely typed to avoid import cycles
budget: ToolBudget
timeout_seconds: float = 20.0
output_max_chars: int = 8000
@classmethod
def from_settings(cls, scope: Any, budget: ToolBudget | None = None) -> "ToolContext":
max_calls = int(getattr(settings, "AGENT_MAX_TOOL_CALLS_PER_RUN", 40) or 40)
return cls(
scope=scope,
budget=budget or ToolBudget(max_calls=max_calls),
timeout_seconds=float(
getattr(settings, "AGENT_TOOL_TIMEOUT_SECONDS", 20.0) or 20.0
),
output_max_chars=int(
getattr(settings, "AGENT_TOOL_OUTPUT_MAX_CHARS", 8000) or 8000
),
)
def truncate_output(text: str, max_chars: int) -> str:
if text is None:
return ""
text = str(text)
if len(text) <= max_chars:
return text
return text[:max_chars] + f"\n… [truncated to {max_chars} chars]"
async def run_tool(
ctx: ToolContext,
name: str,
coro_fn: Callable[..., Awaitable[str]],
*args: Any,
**kwargs: Any,
) -> str:
"""Enforce budget + timeout + output cap around one tool invocation."""
ctx.budget.consume()
try:
result = await asyncio.wait_for(
coro_fn(*args, **kwargs), timeout=ctx.timeout_seconds
)
except asyncio.TimeoutError as exc:
raise ToolTimeout(
f"Tool {name!r} timed out after {ctx.timeout_seconds}s"
) from exc
return truncate_output(result, ctx.output_max_chars)
llm_be/chat_backend/services/tools/data_analysis.py
-117
View File
@@ -1,117 +0,0 @@
"""analyse_dataframe / make_plot tools — wrap AsyncDataAnalysisService (#63).
Tenant-scoped the same way as ``documents.py``: the caller passes a resolved
``workspace`` and we only ever read a ``Document`` that belongs to it.
"""
from __future__ import annotations
import json
import logging
from asgiref.sync import sync_to_async
logger = logging.getLogger(__name__)
def _infer_file_type(name: str) -> str:
name = (name or "").lower()
if name.endswith(".csv"):
return "csv"
if name.endswith((".xlsx", ".xls")):
return "xlsx"
if name.endswith(".docx"):
return "docx"
if name.endswith(".pdf"):
return "pdf"
return "txt"
async def _load_document_bytes(workspace, document_id: int):
from chat_backend.models import Document
doc = await sync_to_async(
Document.objects.filter(id=document_id, workspace=workspace).first
)()
if doc is None:
raise PermissionError(
f"Document {document_id} was not found in your workspace."
)
file_bytes = await sync_to_async(lambda: doc.file.read())()
return file_bytes, _infer_file_type(doc.file.name)
async def analyse_dataframe(workspace, document_id: int, query: str) -> str:
"""Run the existing data-analysis LLM chain against a workspace document."""
from chat_backend.services.data_analysis_service import AsyncDataAnalysisService
if not query or not query.strip():
raise ValueError("query is required")
file_bytes, file_type = await _load_document_bytes(workspace, document_id)
service = AsyncDataAnalysisService()
parts: list[str] = []
async for chunk in service.generate_response(query, file_bytes, file_type):
# The service can yield a JSON {"type": "plot"|"error", ...} sentinel
# instead of narrative text — surface it as a short note rather than
# a raw base64 blob in the tool transcript (use make_plot for plots).
try:
payload = json.loads(chunk)
except (TypeError, ValueError):
parts.append(chunk)
continue
if isinstance(payload, dict) and payload.get("type") == "plot":
parts.append(
"[A plot was generated — call the make_plot tool to produce "
"it directly instead of narrative analysis.]"
)
elif isinstance(payload, dict) and payload.get("type") == "error":
parts.append(f"Error: {payload.get('content', 'analysis failed')}")
else:
parts.append(chunk)
return "".join(parts) or "No analysis produced."
async def make_plot(workspace, document_id: int, query: str) -> str:
"""Generate a scatter plot for a workspace document.
Returns a data URI note rather than embedding the full base64 image in
the tool transcript (keeps agent step logs small); the image itself is
stubbed out of the text tool-call loop for #63 — the existing upload-a-
file chat flow (``AsyncDataAnalysisService``) remains the supported path
for actually viewing a rendered chart.
"""
import pandas as pd
import io
file_bytes, file_type = await _load_document_bytes(workspace, document_id)
try:
if file_type == "csv":
df = await sync_to_async(pd.read_csv, thread_sensitive=False)(
io.BytesIO(file_bytes)
)
elif file_type == "xlsx":
df = await sync_to_async(pd.read_excel, thread_sensitive=False)(
io.BytesIO(file_bytes)
)
else:
return (
f"Cannot plot file type {file_type!r}; only CSV/XLSX are "
"supported for make_plot."
)
except Exception as exc: # pragma: no cover - defensive, pandas errors vary
return f"Could not parse document {document_id} for plotting: {exc}"
from chat_backend.services.data_analysis_service import AsyncDataAnalysisService
service = AsyncDataAnalysisService()
try:
image_b64 = await sync_to_async(
service._generate_plot, thread_sensitive=False
)(query, df)
except ValueError as exc:
return str(exc)
return (
f"Plot generated ({len(image_b64)} base64 chars, PNG). Rendering "
"inline images from agent runs is not yet wired into the chat "
"transcript — describe the chart to the user in words instead."
)
llm_be/chat_backend/services/tools/documents.py
-60
View File
@@ -1,60 +0,0 @@
"""search_documents / read_document tools — tenant-scoped RAG access (#63).
Both tools are bound to a single ``DocumentWorkspace`` resolved once via
``resolve_chat_company_scope`` / ``get_workspace_for_scope`` (same stove-pipe
pattern as the chat consumers) so a run can never read another tenant's
documents, regardless of what ``document_id`` the model asks for.
"""
from __future__ import annotations
import logging
from asgiref.sync import sync_to_async
logger = logging.getLogger(__name__)
async def search_documents(workspace, query: str, k: int = 4) -> str:
"""Semantic search over the caller's own workspace documents."""
from chat_backend.services.rag_services import AsyncRAGService
if not query or not query.strip():
raise ValueError("query is required")
service = AsyncRAGService()
docs = await service.search_documents(query, workspace, k=k)
if not docs:
return "No matching documents found in your workspace."
parts = []
for doc in docs:
meta = doc.metadata or {}
source = meta.get("source") or "unknown"
document_id = meta.get("document_id")
parts.append(f"[document_id={document_id}] {source}\n{doc.page_content}")
return "\n\n".join(parts)
async def read_document(workspace, document_id: int) -> str:
"""Return the indexed text content of one document, tenant-scoped."""
from chat_backend.models import Document
from chat_backend.services.rag_services import AsyncRAGService
doc = await sync_to_async(
Document.objects.filter(id=document_id, workspace=workspace).first
)()
if doc is None:
# Deliberately identical to "not found" — never confirm existence of
# a document in another tenant's workspace.
raise PermissionError(
f"Document {document_id} was not found in your workspace."
)
service = AsyncRAGService()
where = {"$and": [{"workspace_id": workspace.id}, {"document_id": document_id}]}
result = await sync_to_async(service.vector_store.get, thread_sensitive=False)(
where=where
)
chunks = (result or {}).get("documents") or []
if not chunks:
return "Document has no indexed content yet."
return "\n\n".join(chunks)
llm_be/chat_backend/services/tools/fetch_url.py
-134
View File
@@ -1,134 +0,0 @@
"""fetch_url tool — SSRF-guarded HTTP GET + text extraction (#63).
Rejects private/link-local/loopback/reserved/multicast addresses (including
after DNS resolution, so ``http://attacker.example`` that resolves to
``127.0.0.1`` is still blocked) and caps response size.
"""
from __future__ import annotations
import ipaddress
import logging
import socket
from urllib.parse import urlparse
from asgiref.sync import sync_to_async
from django.conf import settings
logger = logging.getLogger(__name__)
_ALLOWED_SCHEMES = {"http", "https"}
# Cloud metadata endpoints — block explicitly even though 169.254/16 is
# link-local (covered below); kept for defense-in-depth / clearer errors.
_BLOCKED_HOSTS = {"metadata.google.internal", "metadata.internal"}
class UnsafeUrlError(ValueError):
"""Raised when a URL fails the SSRF guard."""
def _is_disallowed_ip(ip_str: str) -> bool:
try:
ip = ipaddress.ip_address(ip_str)
except ValueError:
return True # unparseable → fail closed
return (
ip.is_private
or ip.is_loopback
or ip.is_link_local
or ip.is_reserved
or ip.is_multicast
or ip.is_unspecified
)
def validate_url(url: str) -> str:
"""Raise :class:`UnsafeUrlError` unless ``url`` is safe to fetch. Returns hostname."""
if not url or not isinstance(url, str):
raise UnsafeUrlError("URL is required.")
parsed = urlparse(url.strip())
if parsed.scheme.lower() not in _ALLOWED_SCHEMES:
raise UnsafeUrlError(f"Unsupported scheme: {parsed.scheme!r}")
hostname = (parsed.hostname or "").lower()
if not hostname:
raise UnsafeUrlError("URL has no hostname.")
if hostname in _BLOCKED_HOSTS or hostname == "localhost":
raise UnsafeUrlError(f"Blocked host: {hostname}")
# Literal IP in the URL — validate directly.
try:
ipaddress.ip_address(hostname)
literal_ip = hostname
except ValueError:
literal_ip = None
if literal_ip is not None:
if _is_disallowed_ip(literal_ip):
raise UnsafeUrlError(f"Blocked IP literal: {literal_ip}")
return hostname
# Resolve DNS and check every returned address (defeats DNS rebinding to
# a public name that currently points at a private IP).
try:
infos = socket.getaddrinfo(hostname, None)
except socket.gaierror as exc:
raise UnsafeUrlError(f"Could not resolve host: {hostname}") from exc
resolved_ips = {info[4][0] for info in infos}
if not resolved_ips:
raise UnsafeUrlError(f"Could not resolve host: {hostname}")
for ip_str in resolved_ips:
if _is_disallowed_ip(ip_str):
raise UnsafeUrlError(
f"Host {hostname} resolves to a disallowed address: {ip_str}"
)
return hostname
def _fetch_sync(url: str, *, max_bytes: int) -> str:
import requests
from bs4 import BeautifulSoup
validate_url(url)
response = requests.get(
url,
timeout=10,
headers={"User-Agent": "HesychiaAgent/1.0 (+tool:fetch_url)"},
stream=True,
allow_redirects=True,
)
response.raise_for_status()
# Re-validate the final URL — redirects could point somewhere unsafe.
if response.url != url:
validate_url(response.url)
content = bytearray()
for chunk in response.iter_content(chunk_size=8192):
if not chunk:
continue
content.extend(chunk)
if len(content) > max_bytes:
break
response.close()
content_type = (response.headers.get("Content-Type") or "").lower()
body = bytes(content[:max_bytes])
if "html" in content_type or b"<html" in body[:1024].lower():
soup = BeautifulSoup(body, "html.parser")
for tag in soup(["script", "style", "noscript"]):
tag.decompose()
text = soup.get_text(separator="\n", strip=True)
else:
text = body.decode("utf-8", errors="replace")
return text
async def fetch_url(url: str) -> str:
"""Fetch ``url`` and return extracted text. Raises :class:`UnsafeUrlError`."""
max_bytes = int(
getattr(settings, "AGENT_FETCH_URL_MAX_BYTES", 2 * 1024 * 1024)
or 2 * 1024 * 1024
)
return await sync_to_async(_fetch_sync, thread_sensitive=False)(
url, max_bytes=max_bytes
)
llm_be/chat_backend/services/tools/registry.py
-233
View File
@@ -1,233 +0,0 @@
"""Agent tool registry (#63) — catalog + LangChain tool construction.
Two things live here:
- ``TOOL_CATALOG``: static metadata (name/description/keywords) usable
offline (no Ollama/network) by the planner fallback heuristic and by
``evals/tool_selection.yaml`` scoring.
- ``build_agent_tools``: builds tenant-scoped, budget/timeout-wrapped
LangChain ``StructuredTool`` instances for one agent run, suitable for
``ChatOllama.bind_tools``.
"""
from __future__ import annotations
from dataclasses import dataclass, field
from chat_backend.services.tools.base import ToolContext, run_tool
from chat_backend.services.tools.data_analysis import analyse_dataframe, make_plot
from chat_backend.services.tools.documents import read_document, search_documents
from chat_backend.services.tools.fetch_url import fetch_url
from chat_backend.services.tools.web_search import web_search
WEB_SEARCH = "web_search"
FETCH_URL = "fetch_url"
SEARCH_DOCUMENTS = "search_documents"
READ_DOCUMENT = "read_document"
ANALYSE_DATAFRAME = "analyse_dataframe"
MAKE_PLOT = "make_plot"
TOOL_NAMES = (
WEB_SEARCH,
FETCH_URL,
SEARCH_DOCUMENTS,
READ_DOCUMENT,
ANALYSE_DATAFRAME,
MAKE_PLOT,
)
@dataclass(frozen=True)
class ToolSpec:
name: str
description: str
keywords: tuple[str, ...] = field(default_factory=tuple)
requires_workspace: bool = False
TOOL_CATALOG: dict[str, ToolSpec] = {
WEB_SEARCH: ToolSpec(
name=WEB_SEARCH,
description=(
"Search the public web for current information (news, prices, "
"people, events, anything time-sensitive or outside training data)."
),
keywords=(
"latest",
"current",
"today",
"news",
"price",
"weather",
"who is",
"search",
"look up",
"recent",
),
),
FETCH_URL: ToolSpec(
name=FETCH_URL,
description=(
"Fetch and extract the text content of a specific URL "
"(use after web_search to read a promising result in full)."
),
keywords=("http://", "https://", "this link", "this url", "that page", "fetch"),
),
SEARCH_DOCUMENTS: ToolSpec(
name=SEARCH_DOCUMENTS,
description=(
"Semantic search over the user's own uploaded/synced documents "
"(Drive, uploads) — use for questions about 'my documents/files'."
),
keywords=("my document", "my file", "uploaded", "drive", "our docs", "knowledge base"),
requires_workspace=True,
),
READ_DOCUMENT: ToolSpec(
name=READ_DOCUMENT,
description="Read the full indexed text of one document by its document_id.",
keywords=("read the document", "full text", "open document"),
requires_workspace=True,
),
ANALYSE_DATAFRAME: ToolSpec(
name=ANALYSE_DATAFRAME,
description=(
"Run statistical/narrative analysis over a tabular or text "
"document (CSV, XLSX, DOCX, PDF) already in the user's workspace."
),
keywords=("analyse", "analyze", "spreadsheet", "csv", "dataset", "statistics"),
requires_workspace=True,
),
MAKE_PLOT: ToolSpec(
name=MAKE_PLOT,
description="Generate a scatter plot from a CSV/XLSX document in the user's workspace.",
keywords=("plot", "chart", "graph", "visualize", "scatter"),
requires_workspace=True,
),
}
def suggest_tools(query: str) -> list[str]:
"""Offline keyword heuristic: which tools a query probably needs.
Used as the planner fallback when the orchestrator model is unavailable
or fails to parse, and by ``evals/tool_selection.yaml`` scoring.
"""
text = (query or "").lower()
hits = [
spec.name
for spec in TOOL_CATALOG.values()
if any(keyword in text for keyword in spec.keywords)
]
return hits
def build_agent_tools(ctx: ToolContext, workspace=None) -> list:
"""Build LangChain ``StructuredTool`` instances bound to ``ctx``/``workspace``.
Imported lazily inside the function body so importing this module (e.g.
from tests or the eval scorer) never requires ``langchain_core`` tool
machinery to be exercised unless a run actually builds tools.
"""
from langchain_core.tools import StructuredTool
from pydantic import BaseModel, Field
class WebSearchInput(BaseModel):
query: str = Field(description="Search query")
max_results: int = Field(default=5, description="Maximum results to return")
class FetchUrlInput(BaseModel):
url: str = Field(description="Absolute http(s) URL to fetch")
class SearchDocumentsInput(BaseModel):
query: str = Field(description="Semantic search query")
k: int = Field(default=4, description="Number of chunks to return")
class ReadDocumentInput(BaseModel):
document_id: int = Field(description="Document id from search_documents results")
class AnalyseDataframeInput(BaseModel):
document_id: int = Field(description="Document id to analyse")
query: str = Field(description="What to analyse / answer about the document")
class MakePlotInput(BaseModel):
document_id: int = Field(description="CSV/XLSX document id to plot")
query: str = Field(description="What columns/relationship to plot")
async def _web_search(query: str, max_results: int = 5) -> str:
return await run_tool(ctx, WEB_SEARCH, web_search, query, max_results)
async def _fetch_url(url: str) -> str:
return await run_tool(ctx, FETCH_URL, fetch_url, url)
tools = [
StructuredTool.from_function(
name=WEB_SEARCH,
description=TOOL_CATALOG[WEB_SEARCH].description,
args_schema=WebSearchInput,
coroutine=_web_search,
),
StructuredTool.from_function(
name=FETCH_URL,
description=TOOL_CATALOG[FETCH_URL].description,
args_schema=FetchUrlInput,
coroutine=_fetch_url,
),
]
if workspace is not None:
async def _search_documents(query: str, k: int = 4) -> str:
return await run_tool(
ctx, SEARCH_DOCUMENTS, search_documents, workspace, query, k
)
async def _read_document(document_id: int) -> str:
return await run_tool(
ctx, READ_DOCUMENT, read_document, workspace, document_id
)
async def _analyse_dataframe(document_id: int, query: str) -> str:
return await run_tool(
ctx,
ANALYSE_DATAFRAME,
analyse_dataframe,
workspace,
document_id,
query,
)
async def _make_plot(document_id: int, query: str) -> str:
return await run_tool(
ctx, MAKE_PLOT, make_plot, workspace, document_id, query
)
tools.extend(
[
StructuredTool.from_function(
name=SEARCH_DOCUMENTS,
description=TOOL_CATALOG[SEARCH_DOCUMENTS].description,
args_schema=SearchDocumentsInput,
coroutine=_search_documents,
),
StructuredTool.from_function(
name=READ_DOCUMENT,
description=TOOL_CATALOG[READ_DOCUMENT].description,
args_schema=ReadDocumentInput,
coroutine=_read_document,
),
StructuredTool.from_function(
name=ANALYSE_DATAFRAME,
description=TOOL_CATALOG[ANALYSE_DATAFRAME].description,
args_schema=AnalyseDataframeInput,
coroutine=_analyse_dataframe,
),
StructuredTool.from_function(
name=MAKE_PLOT,
description=TOOL_CATALOG[MAKE_PLOT].description,
args_schema=MakePlotInput,
coroutine=_make_plot,
),
]
)
return tools
llm_be/chat_backend/services/tools/web_search.py
-29
View File
@@ -1,29 +0,0 @@
"""web_search tool — reuses the #62 SearchProvider facade (SearxNG + DDGS failover)."""
from __future__ import annotations
from asgiref.sync import sync_to_async
from chat_backend.services.search import SearchUnavailable, search_and_rank
def _format_results(results) -> str:
if not results:
return "No results found."
lines: list[str] = []
for i, result in enumerate(results, start=1):
lines.append(f"[{i}] {result.title}\n {result.url}\n {result.snippet}")
return "\n".join(lines)
async def web_search(query: str, max_results: int = 5) -> str:
"""Search the web and return numbered title/url/snippet results."""
if not query or not query.strip():
raise ValueError("query is required")
try:
results = await sync_to_async(search_and_rank, thread_sensitive=False)(
[query], max_results=max_results
)
except SearchUnavailable as exc:
return f"Search unavailable: {exc}"
return _format_results(results)
llm_be/chat_backend/services/ws_frames.py
-72
View File
@@ -1,72 +0,0 @@
"""Versioned WebSocket frame envelopes (#62 citations, #96 status, #63 agent).
Shape agreed across chat_backend and chat_web_app::
{"v": 1, "type": "<type>", "data": <payload>}
Clients must ignore unknown ``type`` values so backend and frontend can
deploy independently.
"""
from __future__ import annotations
from typing import Any
FRAME_VERSION = 1
# Status stages for live activity feedback (chat_web_app#96).
STATUS_LABELS: dict[str, str] = {
"queued": "Getting started",
"moderating": "Checking your request",
"evaluating": "Evaluating the question",
"searching": "Searching the web",
"reading_sources": "Reading sources",
"retrieving_docs": "Searching your documents",
"analysing": "Analysing your file",
"refining": "Refining the answer",
"writing": "", # tokens streaming — no label
}
# Agent progress frame types (chat_backend#63).
AGENT_FRAME_TYPES = frozenset(
{
"run_started",
"plan_ready",
"step_started",
"step_completed",
"step_failed",
"run_completed",
}
)
def versioned_frame(frame_type: str, data: Any) -> dict:
"""Build a versioned envelope."""
return {"v": FRAME_VERSION, "type": frame_type, "data": data}
def status_frame(
stage: str,
*,
label: str | None = None,
detail: str | None = None,
) -> dict:
"""Activity status frame for the typing-indicator replacement (#96)."""
resolved_label = label if label is not None else STATUS_LABELS.get(stage, stage)
payload: dict[str, Any] = {"stage": stage, "label": resolved_label}
if detail:
payload["detail"] = detail
return versioned_frame("status", payload)
def citations_frame(citations: list[dict]) -> dict:
"""Citations frame after END_OF_THE_STREAM (#62 / chat_web_app#98)."""
return versioned_frame("citations", citations)
def agent_frame(frame_type: str, data: dict) -> dict:
"""Agent run/step progress frame (#63)."""
if frame_type not in AGENT_FRAME_TYPES:
# Still emit — FE ignores unknown types; allow forward-compat names.
pass
return versioned_frame(frame_type, data)
llm_be/chat_backend/signals.py
+21 -11
View File
@@ -1,5 +1,6 @@
from django.db.models.signals import post_delete
from django.db.models.signals import post_save, post_delete
from django.dispatch import receiver
from django.conf import settings
import os
from chat_backend.models import Document
@@ -9,22 +10,31 @@ def _rag_init_skipped() -> bool:
return os.environ.get("SKIP_RAG_INIT", "").lower() in {"1", "true", "yes"}
@receiver(post_save, sender=Document)
def update_vector_on_save(sender, instance, **kwargs):
"""Update vector store when documents are saved"""
if _rag_init_skipped():
return
if not kwargs.get("created", False):
return
try:
from .services.rag_services import AsyncRAGService
rag_service = AsyncRAGService()
rag_service.ingest_documents()
except Exception as exc:
print(f"Skipping vector update on Document save: {exc}")
@receiver(post_delete, sender=Document)
def delete_vector_on_remove(sender, instance, **kwargs):
"""Remove the deleted document's chunks from the vector store (#45).
There is intentionally no post_save handler: ``DocumentUploadView``
already calls ``add_files_to_store`` for the uploaded file, so a
post_save re-ingest would duplicate that work (and, previously,
re-ingested the *entire* workspace on every save). Deletion only needs to
drop that document's own chunks, not rebuild everything else.
"""
"""Handle document deletion by re-indexing the whole workspace"""
if _rag_init_skipped():
return
try:
from .services.rag_services import AsyncRAGService
rag_service = AsyncRAGService()
rag_service.delete_document_vectors(instance.id)
rag_service.ingest_documents()
except Exception as exc:
print(f"Skipping vector cleanup on Document delete: {exc}")
print(f"Skipping vector update on Document delete: {exc}")
llm_be/chat_backend/tests/factories.py
+2 -33
View File
@@ -2,8 +2,6 @@
from __future__ import annotations
import uuid
from django.contrib.auth import get_user_model
from django.core.files.uploadedfile import SimpleUploadedFile
@@ -12,7 +10,6 @@ from chat_backend.models import (
Conversation,
Document,
DocumentWorkspace,
DriveConnection,
Prompt,
)
@@ -122,10 +119,8 @@ def make_prompt(
)
def make_workspace(company=None, name: str = "Test Workspace", user=None) -> DocumentWorkspace:
if user is not None:
return DocumentWorkspace.objects.create(user=user, company=None, name=name)
return DocumentWorkspace.objects.create(company=company, user=None, name=name)
def make_workspace(company, name: str = "Test Workspace") -> DocumentWorkspace:
return DocumentWorkspace.objects.create(company=company, name=name)
def pdf_upload(name: str = "test.pdf") -> SimpleUploadedFile:
@@ -134,29 +129,3 @@ def pdf_upload(name: str = "test.pdf") -> SimpleUploadedFile:
def make_document(workspace, name: str = "test.pdf") -> Document:
return Document.objects.create(workspace=workspace, file=pdf_upload(name))
def make_drive_connection(
company=None,
*,
provider: str = DriveConnection.Provider.GOOGLE,
kind: str = DriveConnection.Kind.PERSONAL,
user=None,
**kwargs,
) -> DriveConnection:
defaults = {
"access_token": "access-token",
"refresh_token": "refresh-token",
"scopes": "openid email profile",
"external_account_email": "drive.user@example.com",
"is_active": True,
}
defaults.update(kwargs)
if kind == DriveConnection.Kind.PERSONAL and user is None:
suffix = uuid.uuid4().hex[:8]
user = make_user(email=f"drive-{provider}-{suffix}@example.com", company=company)
if kind == DriveConnection.Kind.COMPANY:
user = None
return DriveConnection.objects.create(
company=company, provider=provider, kind=kind, user=user, **defaults
)
llm_be/chat_backend/tests/test_agent_orchestrator.py
-343
View File
@@ -1,343 +0,0 @@
"""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)
llm_be/chat_backend/tests/test_agent_tools.py
-110
View File
@@ -1,110 +0,0 @@
"""Offline unit tests for agent tools + gate (#63)."""
from __future__ import annotations
from django.test import SimpleTestCase, TestCase, override_settings
from unittest import mock
from chat_backend.services.agent.decider import should_use_agent
from chat_backend.services.agent.planner import heuristic_plan, parse_plan_json
from chat_backend.services.tools.fetch_url import UnsafeUrlError, validate_url
from chat_backend.services.tools.registry import TOOL_NAMES, suggest_tools
from chat_backend.services.ws_frames import agent_frame
from chat_backend.tests.factories import make_user, make_workspace
class AgentGateTestCase(SimpleTestCase):
@override_settings(ALLOW_AGENTIC_TASKS=False)
def test_disabled_never_routes(self):
self.assertFalse(
should_use_agent(
"Research the top 5 competitors and compare pricing"
)
)
@override_settings(ALLOW_AGENTIC_TASKS=True)
def test_multi_step_routes(self):
self.assertTrue(
should_use_agent(
"Research the top 5 competitors to our product and compare pricing"
)
)
@override_settings(ALLOW_AGENTIC_TASKS=True)
def test_simple_chat_stays_off(self):
self.assertFalse(should_use_agent("hi"))
self.assertFalse(should_use_agent("Write a haiku about rain"))
class PlannerTestCase(SimpleTestCase):
def test_heuristic_plan_capped(self):
plan = heuristic_plan("Research the top 5 competitors and compare", max_steps=4)
self.assertLessEqual(len(plan), 4)
self.assertTrue(plan)
def test_parse_plan_json(self):
raw = '{"steps":[{"title":"Search","tool":"web_search","depends_on":[]}]}'
plan = parse_plan_json(raw)
self.assertEqual(plan[0]["tool"], "web_search")
class FetchUrlSsrfTestCase(SimpleTestCase):
def test_rejects_loopback(self):
for url in (
"http://127.0.0.1/secret",
"http://localhost/x",
"http://[::1]/",
):
with self.subTest(url=url):
with self.assertRaises(UnsafeUrlError):
validate_url(url)
def test_rejects_private(self):
for url in (
"http://10.0.0.1/",
"http://192.168.1.1/",
"http://172.16.0.1/",
):
with self.subTest(url=url):
with self.assertRaises(UnsafeUrlError):
validate_url(url)
def test_rejects_link_local_and_metadata(self):
with self.assertRaises(UnsafeUrlError):
validate_url("http://169.254.169.254/latest/meta-data")
with self.assertRaises(UnsafeUrlError):
validate_url("http://metadata.google.internal/")
class ToolCatalogTestCase(SimpleTestCase):
def test_six_tools_registered(self):
self.assertEqual(len(TOOL_NAMES), 6)
def test_suggest_tools_web(self):
self.assertIn("web_search", suggest_tools("latest news today"))
class AgentFrameTestCase(SimpleTestCase):
def test_run_started_shape(self):
frame = agent_frame(
"run_started", {"run_id": 1, "status": "running", "title": "T"}
)
self.assertEqual(frame["v"], 1)
self.assertEqual(frame["type"], "run_started")
self.assertEqual(frame["data"]["run_id"], 1)
class DocumentTenantIsolationTestCase(TestCase):
def test_read_document_rejects_other_workspace(self):
from asgiref.sync import async_to_sync
from chat_backend.services.tools.documents import read_document
from chat_backend.tests.factories import make_document
user_a = make_user(email="a-agent@example.com")
user_b = make_user(email="b-agent@example.com")
ws_a = make_workspace(user=user_a)
ws_b = make_workspace(user=user_b)
doc_b = make_document(ws_b)
with self.assertRaises(PermissionError):
async_to_sync(read_document)(ws_a, doc_b.id)
llm_be/chat_backend/tests/test_chat_tenant_scope.py
-134
View File
@@ -1,134 +0,0 @@
from django.test import TestCase
from rest_framework_simplejwt.tokens import RefreshToken
from chat_backend.models import DocumentWorkspace
from chat_backend.services.chat_tenant_scope import (
ChatCompanyScope,
ChatTenantScopeError,
ensure_company_workspace,
ensure_personal_workspace,
resolve_chat_company_scope,
resolve_chat_user,
user_from_access_token,
)
from .factories import make_company, make_conversation, make_user, make_workspace
class ChatTenantScopeTestCase(TestCase):
def setUp(self):
self.company = make_company()
self.user = make_user(company=self.company)
self.workspace = make_workspace(self.company)
self.conversation = make_conversation(user=self.user)
def test_resolve_scope_for_owned_conversation(self):
scope = resolve_chat_company_scope(self.user, self.conversation.id)
self.assertEqual(
scope,
ChatCompanyScope(
user_id=self.user.id,
company_id=self.company.id,
workspace_id=self.workspace.id,
conversation_id=self.conversation.id,
),
)
def test_resolve_scope_rejects_cross_user_conversation(self):
other = make_user(
email="other-tenant@example.com", company=make_company("OtherCo")
)
foreign = make_conversation(user=other)
make_workspace(other.company)
with self.assertRaises(ChatTenantScopeError) as ctx:
resolve_chat_company_scope(self.user, foreign.id)
self.assertEqual(ctx.exception.code, "conversation_forbidden")
def test_resolve_chat_user_does_not_bind_identity_to_conversation(self):
other = make_user(company=make_company("VictimCo"), email="victim@example.com")
foreign = make_conversation(user=other)
resolved = resolve_chat_user(conversation_id=foreign.id)
self.assertIsNone(resolved)
def test_resolve_chat_user_prefers_jwt_over_email(self):
token = str(RefreshToken.for_user(self.user).access_token)
spoof = make_user(company=make_company("Spoof"), email="spoof@example.com")
resolved = resolve_chat_user(email=spoof.email, token=token)
self.assertEqual(resolved.id, self.user.id)
def test_user_from_access_token_rejects_garbage(self):
self.assertIsNone(user_from_access_token("not-a-jwt"))
class EnsureCompanyWorkspaceTestCase(TestCase):
"""#46: never 404 a company just because no workspace was created yet."""
def test_returns_existing_workspace_without_creating_another(self):
company = make_company()
workspace = make_workspace(company)
found = ensure_company_workspace(company)
self.assertEqual(found.id, workspace.id)
self.assertEqual(DocumentWorkspace.objects.filter(company=company).count(), 1)
def test_creates_default_workspace_when_missing(self):
company = make_company("NoWorkspaceCo")
self.assertFalse(DocumentWorkspace.objects.filter(company=company).exists())
created = ensure_company_workspace(company)
self.assertEqual(created.name, "Default")
self.assertEqual(created.company_id, company.id)
self.assertEqual(DocumentWorkspace.objects.filter(company=company).count(), 1)
def test_scope_resolution_creates_workspace_instead_of_raising(self):
company = make_company("FreshCo")
user = make_user(email="fresh@example.com", company=company)
scope = resolve_chat_company_scope(user)
workspace = DocumentWorkspace.objects.get(company=company)
self.assertEqual(scope.workspace_id, workspace.id)
self.assertEqual(scope.company_id, company.id)
class EnsurePersonalWorkspaceTestCase(TestCase):
"""#55: solo users get a personal workspace instead of company_missing."""
def test_creates_personal_workspace_for_user_without_company(self):
user = make_user(email="solo@example.com", company=None)
workspace = ensure_personal_workspace(user)
self.assertEqual(workspace.name, "Personal")
self.assertEqual(workspace.user_id, user.id)
self.assertIsNone(workspace.company_id)
self.assertEqual(
DocumentWorkspace.objects.filter(user=user, company__isnull=True).count(),
1,
)
def test_scope_resolution_for_user_without_company(self):
user = make_user(email="solo-scope@example.com", company=None)
scope = resolve_chat_company_scope(user)
workspace = DocumentWorkspace.objects.get(user=user, company__isnull=True)
self.assertEqual(
scope,
ChatCompanyScope(
user_id=user.id,
company_id=None,
workspace_id=workspace.id,
conversation_id=None,
),
)
llm_be/chat_backend/tests/test_consumers.py
+30 -214
View File
@@ -1,5 +1,4 @@
from unittest import mock
import json
from asgiref.sync import sync_to_async
from channels.testing import WebsocketCommunicator
@@ -47,35 +46,10 @@ class DatabaseHelperTestCase(TransactionTestCase):
@parameterized.expand([("websocket", consumers), ("langgraph", consumers_graph)])
async def test_get_workspace(self, _name, module):
workspace = await module.get_workspace(
self.conversation.id, user=self.user
)
workspace = await module.get_workspace(self.conversation.id)
self.assertEqual(workspace.id, self.workspace.id)
@parameterized.expand([("websocket", consumers), ("langgraph", consumers_graph)])
async def test_get_workspace_rejects_other_users_conversation(self, _name, module):
other_company = await sync_to_async(make_company)("OtherCo")
other = await sync_to_async(make_user)(
email="other-tenant@example.com", company=other_company
)
other_conversation = await sync_to_async(make_conversation)(user=other)
await sync_to_async(make_workspace)(other_company)
with self.assertRaises(consumers.ChatTenantScopeError):
await module.get_workspace(other_conversation.id, user=self.user)
@parameterized.expand([("websocket", consumers), ("langgraph", consumers_graph)])
async def test_resolve_tenant_scope_binds_company_and_workspace(
self, _name, module
):
scope = await module.resolve_tenant_scope(self.user, self.conversation.id)
self.assertEqual(scope.user_id, self.user.id)
self.assertEqual(scope.company_id, self.company.id)
self.assertEqual(scope.workspace_id, self.workspace.id)
self.assertEqual(scope.conversation_id, self.conversation.id)
@parameterized.expand([("websocket", consumers), ("langgraph", consumers_graph)])
async def test_get_messages_stores_prompt_and_returns_history(self, _name, module):
messages, prompt_instance = await module.get_messages(
@@ -235,7 +209,6 @@ class GraphNodeTestCase(TransactionTestCase):
"response_generator": None,
"error": None,
"model_name": "Turbo",
"chat_user": self.user,
}
state.update(overrides)
return state
@@ -367,158 +340,59 @@ class GraphNodeTestCase(TransactionTestCase):
_args, _kwargs = service.return_value.generate_response.call_args
self.assertEqual(_args[2].id, self.workspace.id)
@override_settings(ENFORCE_SUBSCRIPTION_GATES=True)
async def test_generation_node_denies_rag_without_an_active_plan(self):
"""#44: RAG turns must respect the plan's ``rag`` feature gate."""
with mock.patch.object(consumers_graph, "AsyncRAGService") as service:
result = await consumers_graph.generation_node(
self._state(prompt_type=PromptType.RAG)
)
payload = result["response_generator"]
self.assertEqual(payload["type"], "error")
self.assertEqual(payload["code"], "subscription_required")
service.assert_not_called()
@override_settings(ENFORCE_SUBSCRIPTION_GATES=True)
async def test_generation_node_denies_rag_on_standard_plan(self):
from monetization.models import SubscriptionPlan, UserSubscription
from monetization.services.plans import assign_plan, seed_subscription_plans
await sync_to_async(seed_subscription_plans)()
standard = await sync_to_async(SubscriptionPlan.objects.get)(slug="standard")
await sync_to_async(assign_plan)(
self.user, plan=standard, source=UserSubscription.Source.ADMIN
)
with mock.patch.object(consumers_graph, "AsyncRAGService") as service:
result = await consumers_graph.generation_node(
self._state(prompt_type=PromptType.RAG)
)
payload = result["response_generator"]
self.assertEqual(payload["type"], "error")
self.assertEqual(payload["code"], "feature_not_allowed")
service.assert_not_called()
@override_settings(ENFORCE_SUBSCRIPTION_GATES=True)
async def test_generation_node_allows_rag_with_founders_plan(self):
from monetization.models import SubscriptionPlan, UserSubscription
from monetization.services.plans import assign_plan, seed_subscription_plans
await sync_to_async(seed_subscription_plans)()
founders = await sync_to_async(SubscriptionPlan.objects.get)(slug="founders")
await sync_to_async(assign_plan)(
self.user, plan=founders, source=UserSubscription.Source.ADMIN
)
with mock.patch.object(consumers_graph, "AsyncRAGService") as service:
async def test_generation_node_defaults_to_general_chat(self):
with mock.patch.object(consumers_graph, "AsyncLLMService") as service:
service.return_value.generate_response.return_value = "generator"
result = await consumers_graph.generation_node(
self._state(prompt_type=PromptType.RAG)
)
self.assertEqual(result["response_generator"], "generator")
async def test_generation_node_defaults_to_general_chat(self):
with mock.patch(
"chat_backend.consumers_graph.prepare_grounded_chat"
) as prepare:
from chat_backend.services.grounded_chat import GroundedTurnResult
prepare.return_value = GroundedTurnResult(
generator="generator", model_name="gpt-oss:20b"
)
result = await consumers_graph.generation_node(self._state())
self.assertEqual(result["response_generator"], "generator")
prepare.assert_called_once()
service.return_value.generate_response.assert_called_once()
@override_settings(ALLOW_INTERNET_ACCESS=True)
async def test_search_prompts_use_grounded_chat(self):
async def test_search_prompts_append_web_results(self):
state = self._state(prompt_type=PromptType.SEARCH)
with mock.patch(
"chat_backend.consumers_graph.prepare_grounded_chat"
) as prepare:
from chat_backend.services.grounded_chat import GroundedTurnResult
prepare.return_value = GroundedTurnResult(
generator="generator",
citations=[
{
"index": 1,
"title": "T",
"url": "https://example.com",
"published_at": None,
}
],
grounded=True,
model_name="gpt-oss:20b",
)
result = await consumers_graph.generation_node(state)
prepare.assert_called_once()
self.assertEqual(result["response_generator"], "generator")
self.assertEqual(result["citations"][0]["url"], "https://example.com")
@override_settings(ALLOW_INTERNET_ACCESS=True)
async def test_fast_model_still_runs_grounding(self):
"""FAST must not skip retrieval (#62) — it only picks a smaller model."""
state = self._state(prompt_type=PromptType.SEARCH, model_name="FAST")
with mock.patch(
"chat_backend.consumers_graph.prepare_grounded_chat"
) as prepare:
from chat_backend.services.grounded_chat import GroundedTurnResult
prepare.return_value = GroundedTurnResult(
generator="generator", model_name="gemma4:latest"
)
with mock.patch.object(consumers_graph, "DuckDuckGoSearchRun") as search:
search.return_value.run.return_value = "top result"
with mock.patch.object(consumers_graph, "AsyncLLMService"):
await consumers_graph.generation_node(state)
prepare.assert_called_once()
self.assertEqual(prepare.call_args.kwargs["model_name"], "FAST")
self.assertIn("Search Results: top result", state["messages"][-1].content)
@override_settings(ALLOW_INTERNET_ACCESS=True)
async def test_fast_model_skips_web_search(self):
state = self._state(prompt_type=PromptType.SEARCH, model_name="FAST")
with mock.patch.object(consumers_graph, "DuckDuckGoSearchRun") as search:
with mock.patch.object(consumers_graph, "AsyncLLMService"):
await consumers_graph.generation_node(state)
search.assert_not_called()
self.assertEqual(len(state["messages"]), 1)
@override_settings(ALLOW_INTERNET_ACCESS=False)
async def test_search_is_skipped_when_internet_access_is_disabled(self):
state = self._state(prompt_type=PromptType.SEARCH)
with mock.patch(
"chat_backend.consumers_graph.prepare_grounded_chat"
) as prepare:
from chat_backend.services.grounded_chat import GroundedTurnResult
prepare.return_value = GroundedTurnResult(
generator="generator", model_name="gpt-oss:20b"
)
with mock.patch.object(consumers_graph, "DuckDuckGoSearchRun") as search:
with mock.patch.object(consumers_graph, "AsyncLLMService"):
await consumers_graph.generation_node(state)
# prepare_grounded_chat still runs; inside it skips providers when
# ALLOW_INTERNET_ACCESS is False.
prepare.assert_called_once()
search.assert_not_called()
@override_settings(ALLOW_INTERNET_ACCESS=True)
async def test_search_failures_surface_error_not_plain_chat(self):
async def test_search_failures_fall_back_to_plain_chat(self):
state = self._state(prompt_type=PromptType.SEARCH)
with mock.patch(
"chat_backend.consumers_graph.prepare_grounded_chat"
) as prepare:
from chat_backend.services.grounded_chat import GroundedTurnResult
prepare.return_value = GroundedTurnResult(
error={
"type": "error",
"code": "search_unavailable",
"content": "couldn't reach live sources",
},
grounded=True,
)
with mock.patch.object(consumers_graph, "DuckDuckGoSearchRun") as search:
search.return_value.run.side_effect = RuntimeError("ddg unreachable")
with mock.patch.object(consumers_graph, "AsyncLLMService") as service:
service.return_value.generate_response.return_value = "generator"
result = await consumers_graph.generation_node(state)
self.assertEqual(result["response_generator"]["code"], "search_unavailable")
self.assertEqual(result["response_generator"], "generator")
self.assertEqual(len(state["messages"]), 1)
class WebSocketRoutingTestCase(TransactionTestCase):
@@ -538,61 +412,3 @@ class WebSocketRoutingTestCase(TransactionTestCase):
[str(route.pattern) for route in websocket_urlpatterns],
["ws/chat_again/$", "ws/conditional_chat/$"],
)
class WebSocketReceiveGuardTestCase(TransactionTestCase):
"""Heartbeats / empty prompts must not spawn conversations or LLM work."""
@parameterized.expand(
[("chat", "/ws/chat_again/"), ("conditional_chat", "/ws/conditional_chat/")]
)
async def test_ping_heartbeat_is_ignored(self, _name, path):
communicator = WebsocketCommunicator(application, path)
connected, _ = await communicator.connect()
self.assertTrue(connected)
with mock.patch(
"chat_backend.consumers.title_generator.generate_async",
new_callable=mock.AsyncMock,
) as title_chat, mock.patch(
"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"})
# No reply expected; give the event loop a tick.
self.assertTrue(await communicator.receive_nothing(timeout=0.2))
title_chat.assert_not_called()
title_graph.assert_not_called()
count = await sync_to_async(Conversation.objects.count)()
self.assertEqual(count, 0)
await communicator.disconnect()
@parameterized.expand(
[("chat", "/ws/chat_again/"), ("conditional_chat", "/ws/conditional_chat/")]
)
async def test_empty_message_is_rejected(self, _name, path):
communicator = WebsocketCommunicator(application, path)
connected, _ = await communicator.connect()
self.assertTrue(connected)
with mock.patch(
"chat_backend.consumers.title_generator.generate_async",
new_callable=mock.AsyncMock,
) as title_chat, mock.patch(
"chat_backend.consumers_graph.title_generator.generate_async",
new_callable=mock.AsyncMock,
) as title_graph:
await communicator.send_json_to(
{"message": " ", "email": "a@b.com", "conversation_id": None}
)
response = await communicator.receive_from(timeout=1)
payload = json.loads(response)
self.assertEqual(payload["type"], "error")
self.assertIn("empty", payload["content"].lower())
title_chat.assert_not_called()
title_graph.assert_not_called()
count = await sync_to_async(Conversation.objects.count)()
self.assertEqual(count, 0)
await communicator.disconnect()
llm_be/chat_backend/tests/test_drive_tasks.py
-79
View File
@@ -1,79 +0,0 @@
"""Tests for Drive sync background tasks (#57)."""
from __future__ import annotations
from unittest import mock
from django.test import TestCase, TransactionTestCase
from chat_backend.drive_tasks import enqueue_drive_sync, run_drive_connection_sync
from chat_backend.models import DriveConnection
from chat_backend.tests.factories import make_company, make_drive_connection, make_user
class RunDriveConnectionSyncTestCase(TestCase):
def setUp(self):
self.company = make_company()
self.user = make_user(company=self.company)
self.connection = make_drive_connection(
self.company, kind=DriveConnection.Kind.PERSONAL, user=self.user
)
@mock.patch("chat_backend.drive_tasks.sync_connection")
def test_runs_sync_for_active_connection(self, mock_sync):
mock_sync.return_value = {"added": 1, "updated": 0, "removed": 0, "failed": []}
result = run_drive_connection_sync.call(connection_id=self.connection.id)
mock_sync.assert_called_once()
self.assertEqual(mock_sync.call_args.args[0].id, self.connection.id)
self.assertEqual(result["added"], 1)
def test_missing_connection_returns_error(self):
result = run_drive_connection_sync.call(connection_id=999999)
self.assertEqual(result, {"error": "connection_not_found"})
class EnqueueDriveSyncTestCase(TransactionTestCase):
def setUp(self):
self.company = make_company()
self.user = make_user(company=self.company)
self.connection = make_drive_connection(
self.company, kind=DriveConnection.Kind.PERSONAL, user=self.user
)
@mock.patch("chat_backend.drive_tasks._dispatch_sync")
def test_marks_pending_and_dispatches(self, mock_dispatch):
connection, enqueued = enqueue_drive_sync(self.connection)
self.assertTrue(enqueued)
connection.refresh_from_db()
self.assertEqual(connection.last_sync_status, DriveConnection.SyncStatus.PENDING)
self.assertEqual(connection.last_sync_error, "")
self.assertEqual(connection.sync_total, 0)
self.assertEqual(connection.sync_processed, 0)
self.assertEqual(connection.sync_added, 0)
self.assertEqual(connection.sync_updated, 0)
self.assertEqual(connection.sync_failed, 0)
mock_dispatch.assert_called_once_with(connection_id=self.connection.id)
@mock.patch("chat_backend.drive_tasks._dispatch_sync")
def test_skips_duplicate_while_pending(self, mock_dispatch):
self.connection.last_sync_status = DriveConnection.SyncStatus.PENDING
self.connection.save(update_fields=["last_sync_status"])
connection, enqueued = enqueue_drive_sync(self.connection)
self.assertFalse(enqueued)
mock_dispatch.assert_not_called()
self.assertEqual(connection.last_sync_status, DriveConnection.SyncStatus.PENDING)
@mock.patch("chat_backend.drive_tasks._dispatch_sync")
def test_force_reenqueues_pending(self, mock_dispatch):
self.connection.last_sync_status = DriveConnection.SyncStatus.PENDING
self.connection.save(update_fields=["last_sync_status"])
_, enqueued = enqueue_drive_sync(self.connection, force=True)
self.assertTrue(enqueued)
mock_dispatch.assert_called_once()
llm_be/chat_backend/tests/test_evals.py
-138
View File
@@ -1,138 +0,0 @@
"""Unit tests for eval harness (#62 Phase 4)."""
from __future__ import annotations
from io import StringIO
from django.core.management import call_command
from django.test import SimpleTestCase
from chat_backend.evals.grading import (
compute_self_consistency,
contains_all,
contains_any,
contains_forbidden,
detect_hedge_or_refuse,
grade_answer,
normalize_verdict,
)
from chat_backend.evals.suite import load_suite, validate_suite
class SuiteLoadTestCase(SimpleTestCase):
def test_suite_loads_with_minimum_questions(self):
suite = load_suite()
self.assertGreaterEqual(len(suite["questions"]), 40)
def test_all_categories_present(self):
suite = load_suite()
categories = {q["category"] for q in suite["questions"]}
self.assertEqual(
categories,
{
"post_cutoff",
"stable_fact",
"refuse_or_hedge",
"rag",
"multi_turn",
},
)
def test_taylor_swift_question(self):
suite = load_suite()
ts = next(q for q in suite["questions"] if q["id"] == "ts_married")
self.assertEqual(ts["prompt"], "did Taylor Swift get married")
self.assertIn("Joe Alwyn", ts["must_not_contain_any"])
class GradingHelpersTestCase(SimpleTestCase):
def test_contains_any_and_all(self):
text = "Travis Kelce married at Madison Square Garden"
self.assertTrue(contains_any(text, ["Kelce", "Alwyn"]))
self.assertTrue(contains_all(text, ["Travis", "Garden"]))
self.assertFalse(contains_any(text, ["Joe Alwyn"]))
def test_contains_forbidden(self):
self.assertTrue(contains_forbidden("She married Joe Alwyn", ["Joe Alwyn"]))
self.assertFalse(contains_forbidden("She married Travis Kelce", ["Joe Alwyn"]))
def test_detect_hedge_or_refuse(self):
self.assertTrue(detect_hedge_or_refuse("I don't know the winning numbers."))
self.assertFalse(detect_hedge_or_refuse("Paris is the capital of France."))
def test_grade_answer_pass_and_fail(self):
question = {
"must_contain_any": ["Kelce"],
"must_contain_all": [],
"must_not_contain_any": ["Joe Alwyn"],
"expect_citations": True,
"expect_hedge_or_refuse": False,
}
good = grade_answer(
question,
"Reports say Taylor Swift married Travis Kelce [1].",
citations=[{"index": 1}],
)
self.assertTrue(good.passed)
self.assertFalse(good.hallucinated)
bad = grade_answer(
question,
"She married Joe Alwyn in 2023 [1].",
citations=[{"index": 1}],
)
self.assertFalse(bad.passed)
self.assertTrue(bad.hallucinated)
def test_grade_answer_expect_hedge(self):
question = {
"must_contain_any": [],
"must_contain_all": [],
"must_not_contain_any": ["the winning numbers are"],
"expect_citations": False,
"expect_hedge_or_refuse": True,
}
hedged = grade_answer(question, "I can't predict future lottery numbers.")
self.assertTrue(hedged.passed)
self.assertTrue(hedged.is_hedge_or_refuse)
confident = grade_answer(question, "The winning numbers are 1 2 3 4 5 6.")
self.assertFalse(confident.passed)
def test_normalize_verdict_and_self_consistency(self):
verdict_a = normalize_verdict(
passed=True,
hallucinated=False,
has_citations=True,
is_hedge=False,
expect_citations=True,
expect_hedge=False,
)
verdict_b = normalize_verdict(
passed=False,
hallucinated=True,
has_citations=False,
is_hedge=False,
expect_citations=True,
expect_hedge=False,
)
self.assertEqual(compute_self_consistency([verdict_a, verdict_a, verdict_a]), 1.0)
self.assertEqual(
compute_self_consistency([verdict_a, verdict_b, verdict_a]), 2 / 3
)
class RunEvalsCommandTestCase(SimpleTestCase):
def test_dry_run_succeeds(self):
out = StringIO()
call_command("run_evals", dry_run=True, stdout=out)
self.assertIn("Dry/offline mode", out.getvalue())
def test_offline_succeeds(self):
out = StringIO()
call_command("run_evals", offline=True, stdout=out)
self.assertIn("suite structure validated", out.getvalue())
def test_validate_suite_rejects_small_suite(self):
with self.assertRaises(ValueError):
validate_suite({"questions": [{"id": "x", "category": "stable_fact", "prompt": "hi"}]})
llm_be/chat_backend/tests/test_grounding_search.py
-180
View File
@@ -1,180 +0,0 @@
"""Unit tests for grounding + search layer (#62 Phases 1–3)."""
from __future__ import annotations
from django.test import SimpleTestCase, override_settings
from unittest import mock
from chat_backend.services.grounding_decider import GroundingDecider, GroundingDecision
from chat_backend.services.search.base import SearchResult
from chat_backend.services.search.ranking import (
is_rumour_heavy,
rank_and_dedupe,
registrable_domain,
)
from chat_backend.services.search.service import (
SearchService,
SearchUnavailable,
format_sources_block,
)
from chat_backend.ollama_config import (
ROLE_EMBED,
ROLE_FAST,
ROLE_THINKING,
ROLE_UTILITY,
ollama_model_for_role,
resolve_chat_role,
)
class OllamaConfigRoleTestCase(SimpleTestCase):
@override_settings(
OLLAMA_MODEL="legacy-model",
OLLAMA_MODEL_THINKING="think-model",
OLLAMA_MODEL_FAST="fast-model",
OLLAMA_MODEL_UTILITY="util-model",
OLLAMA_EMBED_MODEL="nomic-embed-text",
)
def test_role_resolution(self):
self.assertEqual(ollama_model_for_role(ROLE_THINKING), "think-model")
self.assertEqual(ollama_model_for_role(ROLE_FAST), "fast-model")
self.assertEqual(ollama_model_for_role(ROLE_UTILITY), "util-model")
self.assertEqual(ollama_model_for_role(ROLE_EMBED), "nomic-embed-text")
@override_settings(
OLLAMA_MODEL="legacy-model",
OLLAMA_MODEL_THINKING="",
OLLAMA_MODEL_FAST="",
OLLAMA_MODEL_UTILITY="",
OLLAMA_EMBED_MODEL="",
)
def test_embed_never_falls_back_to_chat_model(self):
# Empty embed setting → hard default, not OLLAMA_MODEL.
self.assertEqual(ollama_model_for_role(ROLE_EMBED), "nomic-embed-text")
def test_resolve_chat_role(self):
self.assertEqual(resolve_chat_role("FAST"), ROLE_FAST)
self.assertEqual(resolve_chat_role("THINKING"), ROLE_THINKING)
self.assertEqual(resolve_chat_role(None), ROLE_THINKING)
class GroundingPrepassTestCase(SimpleTestCase):
def setUp(self):
self.decider = GroundingDecider.__new__(GroundingDecider)
def test_temporal_markers_force_retrieval(self):
cases = [
"did Taylor Swift get married yet",
"What is the latest news on AI?",
"Who won the Super Bowl this year?",
"current stock price of Apple",
"What happened in 2025?",
]
for prompt in cases:
with self.subTest(prompt=prompt):
decision = self.decider.temporal_prepass(prompt)
self.assertIsNotNone(decision)
self.assertTrue(decision.needs_retrieval)
def test_creative_prompt_does_not_force(self):
decision = self.decider.temporal_prepass("Write a poem about cats")
self.assertIsNone(decision)
def test_parse_failure_fails_open(self):
decision = self.decider._parse("NOT JSON", "fallback query")
self.assertTrue(decision.needs_retrieval)
self.assertEqual(decision.source, "fail_open")
self.assertEqual(decision.queries, ["fallback query"])
def test_decide_fail_open_on_llm_error(self):
self.decider.chain = mock.Mock()
self.decider.chain.invoke.side_effect = RuntimeError("ollama down")
decision = self.decider.decide("Is the sky blue?")
self.assertTrue(decision.needs_retrieval)
self.assertEqual(decision.source, "fail_open")
class SearchRankingTestCase(SimpleTestCase):
def test_dedupe_by_domain(self):
results = [
SearchResult("A", "https://www.people.com/a", "married", "2026-07-03", 0),
SearchResult("B", "https://people.com/b", "also", "2026-07-02", 1),
SearchResult("C", "https://bbc.com/c", "ok", "2026-07-01", 2),
]
ranked = rank_and_dedupe(results, temporal=True)
domains = {registrable_domain(r.url) for r in ranked}
self.assertEqual(domains, {"people.com", "bbc.com"})
def test_rumour_heavy_dropped_when_alternatives_exist(self):
clean = SearchResult(
"Married",
"https://people.com/wedding",
"Taylor Swift and Travis Kelce married July 3",
"2026-07-03",
0,
)
poison = SearchResult(
"Rumors",
"https://gossip.com/fake",
"fake AI-generated photos and speculation and blind item",
None,
1,
)
self.assertTrue(is_rumour_heavy(poison))
ranked = rank_and_dedupe([poison, clean], temporal=True)
self.assertEqual(len(ranked), 1)
self.assertEqual(ranked[0].url, clean.url)
def test_format_sources_block_is_numbered(self):
block = format_sources_block(
[
SearchResult(
"Title",
"https://example.com/x",
"Snippet here",
"2026-07-03",
0,
)
]
)
self.assertIn('[1] "Title"', block)
self.assertIn("2026-07-03", block)
self.assertIn("URL: https://example.com/x", block)
self.assertNotIn("Search Results:", block)
class FakeProvider:
def __init__(self, name, results=None, error=None):
self.name = name
self.results = results or []
self.error = error
self.calls = 0
def search(self, query, *, max_results=8):
self.calls += 1
if self.error:
raise self.error
return list(self.results)
class SearchFailoverTestCase(SimpleTestCase):
def test_failover_when_primary_raises(self):
primary = FakeProvider("searxng", error=RuntimeError("down"))
failover = FakeProvider(
"ddgs",
results=[
SearchResult("T", "https://a.com", "s", None, 0, "ddgs"),
],
)
service = SearchService(primary=primary, failover=failover)
results = service.search_one("q")
self.assertEqual(primary.calls, 1)
self.assertEqual(failover.calls, 1)
self.assertEqual(results[0].provider, "ddgs")
def test_all_providers_fail_raises(self):
primary = FakeProvider("searxng", error=RuntimeError("down"))
failover = FakeProvider("ddgs", error=RuntimeError("also down"))
service = SearchService(primary=primary, failover=failover)
with self.assertRaises(SearchUnavailable):
service.search_one("q")
llm_be/chat_backend/tests/test_management_drive_sync.py
-88
View File
@@ -1,88 +0,0 @@
"""Tests for the sync_drive_connections management command (#52 / #57)."""
from __future__ import annotations
from io import StringIO
from unittest import mock
from django.core.management import CommandError, call_command
from django.test import TestCase
from chat_backend.models import DriveConnection
from .factories import make_company, make_drive_connection
class SyncDriveConnectionsCommandTestCase(TestCase):
def setUp(self):
self.company = make_company()
@mock.patch(
"chat_backend.management.commands.sync_drive_connections.enqueue_drive_sync"
)
def test_enqueues_all_active_connections(self, mock_enqueue):
mock_enqueue.return_value = (mock.Mock(), True)
active = make_drive_connection(self.company)
make_drive_connection(
self.company,
provider=DriveConnection.Provider.MICROSOFT,
is_active=False,
)
out = StringIO()
call_command("sync_drive_connections", stdout=out)
mock_enqueue.assert_called_once()
self.assertEqual(mock_enqueue.call_args.args[0].id, active.id)
self.assertIn("queued", out.getvalue())
@mock.patch("chat_backend.drive_tasks.sync_connection")
def test_sync_now_runs_inline(self, mock_sync):
mock_sync.return_value = {
"added": 1,
"updated": 0,
"removed": 0,
"failed": [],
}
active = make_drive_connection(self.company)
out = StringIO()
call_command("sync_drive_connections", "--sync-now", stdout=out)
mock_sync.assert_called_once()
self.assertEqual(mock_sync.call_args.args[0].id, active.id)
self.assertIn("added=1", out.getvalue())
@mock.patch(
"chat_backend.management.commands.sync_drive_connections.enqueue_drive_sync"
)
def test_enqueues_single_connection_by_id(self, mock_enqueue):
mock_enqueue.return_value = (mock.Mock(), True)
target = make_drive_connection(self.company)
make_drive_connection(
self.company, provider=DriveConnection.Provider.MICROSOFT
)
call_command(
"sync_drive_connections",
"--connection-id",
str(target.id),
stdout=StringIO(),
)
mock_enqueue.assert_called_once()
self.assertEqual(mock_enqueue.call_args.args[0].id, target.id)
def test_unknown_connection_id_raises(self):
with self.assertRaises(CommandError):
call_command(
"sync_drive_connections",
"--connection-id",
"999999",
stdout=StringIO(),
)
def test_no_connections_reports_and_exits_cleanly(self):
out = StringIO()
call_command("sync_drive_connections", stdout=out)
self.assertIn("No active Drive connections", out.getvalue())
llm_be/chat_backend/tests/test_models.py
-1
View File
@@ -81,7 +81,6 @@ class CompanyAndUserTestCase(TestCase):
self.assertFalse(user.deleted)
self.assertFalse(user.has_signed_tos)
self.assertTrue(user.conversation_order)
self.assertFalse(user.use_conversation_context)
class ConversationAndPromptTestCase(TestCase):
llm_be/chat_backend/tests/test_oauth.py
+2 -273
View File
@@ -11,9 +11,9 @@ from rest_framework import status
from rest_framework.test import APITestCase
from rest_framework_simplejwt.tokens import AccessToken
from chat_backend.models import CustomUser, DriveConnection, OAuthIdentity
from chat_backend.models import CustomUser, OAuthIdentity
from chat_backend.oauth import ProviderProfile, dump_oauth_state
from chat_backend.tests.factories import make_company, make_user
from chat_backend.tests.factories import make_user
OAUTH_SETTINGS = {
"GOOGLE_OAUTH_CLIENT_ID": "google-client-id",
@@ -239,274 +239,3 @@ class OAuthCallbackTestCase(APITestCase):
provider="microsoft", subject="ms-oid-1", user=user
).exists()
)
@override_settings(**OAUTH_SETTINGS)
class OAuthStartDriveLinkTestCase(APITestCase):
"""#47 — link_drive / link_company_drive require an authenticated user."""
def setUp(self):
self.company = make_company()
self.user = make_user(email="drive.user@example.com", company=self.company)
def test_link_drive_requires_authentication(self):
response = self.client.get(
reverse("oauth_start", kwargs={"provider": "google"}),
{"intent": "link_drive"},
)
self.assertEqual(response.status_code, status.HTTP_401_UNAUTHORIZED)
def test_link_drive_authenticated_redirects_with_drive_scope_and_state(self):
self.client.force_authenticate(user=self.user)
response = self.client.get(
reverse("oauth_start", kwargs={"provider": "google"}),
{"intent": "link_drive"},
)
self.assertEqual(response.status_code, status.HTTP_302_FOUND)
params = parse_qs(urlparse(response["Location"]).query)
self.assertIn("https://www.googleapis.com/auth/drive.readonly", params["scope"][0])
self.assertIn("state", params)
from chat_backend.oauth import load_oauth_state
state_data = load_oauth_state(params["state"][0])
self.assertEqual(state_data["intent"], "link_drive")
self.assertEqual(state_data["user_id"], self.user.id)
def test_link_drive_microsoft_uses_personal_files_scope(self):
self.client.force_authenticate(user=self.user)
response = self.client.get(
reverse("oauth_start", kwargs={"provider": "microsoft"}),
{"intent": "link_drive"},
)
params = parse_qs(urlparse(response["Location"]).query)
self.assertIn("Files.Read", params["scope"][0])
self.assertNotIn("Sites.Read.All", params["scope"][0])
def test_link_company_drive_microsoft_uses_sites_scope(self):
self.user.is_company_manager = True
self.user.save(update_fields=["is_company_manager"])
self.client.force_authenticate(user=self.user)
response = self.client.get(
reverse("oauth_start", kwargs={"provider": "microsoft"}),
{"intent": "link_company_drive"},
)
params = parse_qs(urlparse(response["Location"]).query)
self.assertIn("Files.Read.All", params["scope"][0])
self.assertIn("Sites.Read.All", params["scope"][0])
def test_link_company_drive_requires_company_manager(self):
self.client.force_authenticate(user=self.user)
response = self.client.get(
reverse("oauth_start", kwargs={"provider": "google"}),
{"intent": "link_company_drive"},
)
self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
@override_settings(ENFORCE_SUBSCRIPTION_GATES=True)
def test_link_drive_denied_when_plan_disallows_rag(self):
from monetization.services.plans import assign_plan, seed_subscription_plans
from monetization.models import SubscriptionPlan, UserSubscription
seed_subscription_plans()
standard = SubscriptionPlan.objects.get(slug="standard")
assign_plan(self.user, plan=standard, source=UserSubscription.Source.ADMIN)
self.client.force_authenticate(user=self.user)
response = self.client.get(
reverse("oauth_start", kwargs={"provider": "google"}),
{"intent": "link_drive"},
)
self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
def test_unknown_intent_rejected(self):
self.client.force_authenticate(user=self.user)
response = self.client.get(
reverse("oauth_start", kwargs={"provider": "google"}),
{"intent": "delete_everything"},
)
self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
@override_settings(**OAUTH_SETTINGS)
class OAuthCallbackDriveLinkTestCase(APITestCase):
"""#47 — Drive-link callback upserts a DriveConnection and redirects to FE account page."""
def setUp(self):
self.company = make_company()
self.user = make_user(email="drive.user@example.com", company=self.company)
def _state(self, *, intent="link_drive", provider="google", user_id=None):
return dump_oauth_state(
provider=provider,
intent=intent,
user_id=self.user.id if user_id is None else user_id,
)
@patch("chat_backend.views_oauth.exchange_code_for_profile")
def test_callback_creates_personal_drive_connection(self, mock_exchange):
mock_exchange.return_value = _google_profile(
scopes="openid email profile https://www.googleapis.com/auth/drive.readonly"
)
state = self._state(intent="link_drive")
response = self.client.get(
reverse("oauth_callback", kwargs={"provider": "google"}),
{"code": "auth-code", "state": state},
)
self.assertEqual(response.status_code, status.HTTP_302_FOUND)
self.assertTrue(
response["Location"].startswith("http://frontend.test/document_storage/?")
)
params = parse_qs(urlparse(response["Location"]).query)
self.assertEqual(params["drive_connected"], ["1"])
self.assertEqual(params["provider"], ["google"])
self.assertEqual(params["kind"], ["personal"])
connection = DriveConnection.objects.get(
company=self.company, kind=DriveConnection.Kind.PERSONAL
)
self.assertEqual(connection.user_id, self.user.id)
self.assertEqual(connection.provider, "google")
self.assertEqual(connection.refresh_token, "refresh-token")
self.assertTrue(connection.is_active)
# No JWTs minted for an in-app link flow (user is already authenticated).
self.assertNotIn("access", params)
@patch("chat_backend.views_oauth.exchange_code_for_profile")
def test_callback_updates_existing_connection_tokens(self, mock_exchange):
from chat_backend.tests.factories import make_drive_connection
existing = make_drive_connection(
self.company,
provider=DriveConnection.Provider.GOOGLE,
kind=DriveConnection.Kind.PERSONAL,
user=self.user,
access_token="stale",
refresh_token="stale-refresh",
)
mock_exchange.return_value = _google_profile(access_token="fresh-access")
state = self._state(intent="link_drive")
self.client.get(
reverse("oauth_callback", kwargs={"provider": "google"}),
{"code": "auth-code", "state": state},
)
existing.refresh_from_db()
self.assertEqual(existing.access_token, "fresh-access")
self.assertEqual(DriveConnection.objects.count(), 1)
@patch("chat_backend.views_oauth.exchange_code_for_profile")
def test_callback_company_drive_creates_connection_without_user(self, mock_exchange):
self.user.is_company_manager = True
self.user.save(update_fields=["is_company_manager"])
mock_exchange.return_value = _google_profile()
state = self._state(intent="link_company_drive")
response = self.client.get(
reverse("oauth_callback", kwargs={"provider": "google"}),
{"code": "auth-code", "state": state},
)
params = parse_qs(urlparse(response["Location"]).query)
self.assertEqual(params["kind"], ["company"])
connection = DriveConnection.objects.get(kind=DriveConnection.Kind.COMPANY)
self.assertIsNone(connection.user_id)
self.assertEqual(connection.company_id, self.company.id)
@patch("chat_backend.views_oauth.exchange_code_for_profile")
def test_callback_company_drive_rejects_non_manager(self, mock_exchange):
mock_exchange.return_value = _google_profile()
state = self._state(intent="link_company_drive")
response = self.client.get(
reverse("oauth_callback", kwargs={"provider": "google"}),
{"code": "auth-code", "state": state},
)
params = parse_qs(urlparse(response["Location"]).query)
self.assertEqual(params["error"], ["forbidden"])
self.assertFalse(DriveConnection.objects.exists())
mock_exchange.assert_not_called()
@patch("chat_backend.views_oauth.exchange_code_for_profile")
def test_callback_personal_drive_allows_user_without_company(self, mock_exchange):
"""#55 — personal Drive/RAG works for users not attached to a company."""
self.user.company = None
self.user.save(update_fields=["company"])
mock_exchange.return_value = _google_profile()
state = self._state(intent="link_drive")
response = self.client.get(
reverse("oauth_callback", kwargs={"provider": "google"}),
{"code": "auth-code", "state": state},
)
self.assertEqual(response.status_code, status.HTTP_302_FOUND)
params = parse_qs(urlparse(response["Location"]).query)
self.assertEqual(params["drive_connected"], ["1"])
self.assertEqual(params["kind"], ["personal"])
connection = DriveConnection.objects.get(
kind=DriveConnection.Kind.PERSONAL, user=self.user
)
self.assertIsNone(connection.company_id)
self.assertEqual(connection.provider, "google")
self.assertTrue(connection.is_active)
@patch("chat_backend.views_oauth.exchange_code_for_profile")
def test_callback_company_drive_rejects_user_without_company(self, mock_exchange):
"""#55 — company Drive still requires a company."""
self.user.company = None
self.user.is_company_manager = True
self.user.save(update_fields=["company", "is_company_manager"])
mock_exchange.return_value = _google_profile()
state = self._state(intent="link_company_drive")
response = self.client.get(
reverse("oauth_callback", kwargs={"provider": "google"}),
{"code": "auth-code", "state": state},
)
params = parse_qs(urlparse(response["Location"]).query)
self.assertEqual(params["error"], ["no_company"])
self.assertFalse(DriveConnection.objects.exists())
mock_exchange.assert_not_called()
def test_callback_missing_user_id_in_state_is_rejected(self):
# Simulates a forged/legacy state payload without the linking user.
from django.core import signing
from chat_backend.oauth import STATE_SALT
bad_state = signing.dumps(
{"provider": "google", "intent": "link_drive"}, salt=STATE_SALT
)
response = self.client.get(
reverse("oauth_callback", kwargs={"provider": "google"}),
{"code": "auth-code", "state": bad_state},
)
params = parse_qs(urlparse(response["Location"]).query)
self.assertEqual(params["error"], ["invalid_state"])
@override_settings(ENFORCE_SUBSCRIPTION_GATES=True)
@patch("chat_backend.views_oauth.exchange_code_for_profile")
def test_callback_denied_when_plan_disallows_rag(self, mock_exchange):
from monetization.services.plans import assign_plan, seed_subscription_plans
from monetization.models import SubscriptionPlan, UserSubscription
seed_subscription_plans()
standard = SubscriptionPlan.objects.get(slug="standard")
assign_plan(self.user, plan=standard, source=UserSubscription.Source.ADMIN)
mock_exchange.return_value = _google_profile()
state = self._state(intent="link_drive")
response = self.client.get(
reverse("oauth_callback", kwargs={"provider": "google"}),
{"code": "auth-code", "state": state},
)
params = parse_qs(urlparse(response["Location"]).query)
self.assertEqual(params["error"], ["feature_not_allowed"])
self.assertFalse(DriveConnection.objects.exists())
llm_be/chat_backend/tests/test_serializers.py
+1 -20
View File
@@ -38,10 +38,7 @@ class ConversationSerializerTestCase(TestCase):
data = ConversationSerializer(conversation).data
self.assertEqual(
set(data.keys()),
{"title", "created", "last_modified", "id", "tokens_in", "tokens_out"},
)
self.assertEqual(set(data.keys()), {"title", "created", "last_modified", "id"})
self.assertEqual(data["title"], "Weather Inquiry")
@@ -61,22 +58,6 @@ class PromptSerializerTestCase(TestCase):
self.assertFalse(serializer.is_valid())
self.assertIn("message", serializer.errors)
def test_message_rejects_blank_and_whitespace(self):
for payload in ("", " ", "\n\t"):
serializer = PromptSerializer(
data={"message": payload, "user_created": True}
)
self.assertFalse(serializer.is_valid(), payload)
self.assertIn("message", serializer.errors)
def test_message_is_stripped(self):
serializer = PromptSerializer(
data={"message": " hello ", "user_created": True}
)
self.assertTrue(serializer.is_valid(), serializer.errors)
self.assertEqual(serializer.validated_data["message"], "hello")
def test_user_created_is_required(self):
serializer = PromptSerializer(data={"message": "hi"})
llm_be/chat_backend/tests/test_services_drive_sync.py
-502
View File
@@ -1,502 +0,0 @@
"""Tests for Drive/RAG sync (#48-#52). httpx is mocked; no network access."""
from __future__ import annotations
from datetime import timedelta
from unittest import mock
from django.test import TestCase
from django.utils import timezone
from chat_backend.models import Document, DriveConnection
from chat_backend.services import drive_sync
from chat_backend.services.chat_tenant_scope import ensure_personal_workspace
from chat_backend.services.drive_sync import (
DriveSyncError,
RemoteFile,
ensure_fresh_token,
refresh_google_token,
refresh_microsoft_token,
sync_connection,
)
from .factories import make_company, make_drive_connection, make_user
def _fake_response(*, status_code=200, json_data=None, content=b"", text="", headers=None):
response = mock.MagicMock()
response.status_code = status_code
response.json.return_value = json_data or {}
response.content = content
response.text = text
response.headers = headers or {}
return response
def _fake_client(get_side_effect=None, post_side_effect=None):
"""A MagicMock standing in for ``httpx.Client()`` used as a context manager."""
client = mock.MagicMock()
if get_side_effect is not None:
client.get.side_effect = get_side_effect
if post_side_effect is not None:
client.post.side_effect = post_side_effect
cm = mock.MagicMock()
cm.__enter__.return_value = client
cm.__exit__.return_value = False
return cm, client
class TokenRefreshTestCase(TestCase):
def setUp(self):
self.company = make_company()
self.connection = make_drive_connection(
self.company,
provider=DriveConnection.Provider.GOOGLE,
access_token="old-token",
refresh_token="refresh-me",
token_expires_at=timezone.now() - timedelta(minutes=5),
)
@mock.patch("chat_backend.services.drive_sync.httpx.Client")
def test_refresh_google_token_updates_connection(self, mock_client_cls):
cm, client = _fake_client(
post_side_effect=[
_fake_response(json_data={"access_token": "new-token", "expires_in": 3600})
]
)
mock_client_cls.return_value = cm
token = refresh_google_token(self.connection)
self.assertEqual(token, "new-token")
self.connection.refresh_from_db()
self.assertEqual(self.connection.access_token, "new-token")
self.assertGreater(self.connection.token_expires_at, timezone.now())
def test_refresh_google_token_requires_refresh_token(self):
self.connection.refresh_token = ""
self.connection.save(update_fields=["refresh_token"])
with self.assertRaises(DriveSyncError) as ctx:
refresh_google_token(self.connection)
self.assertEqual(ctx.exception.code, "missing_refresh_token")
@mock.patch("chat_backend.services.drive_sync.httpx.Client")
def test_refresh_google_token_raises_on_http_error(self, mock_client_cls):
cm, client = _fake_client(
post_side_effect=[_fake_response(status_code=400, text="invalid_grant")]
)
mock_client_cls.return_value = cm
with self.assertRaises(DriveSyncError) as ctx:
refresh_google_token(self.connection)
self.assertEqual(ctx.exception.code, "token_refresh_failed")
@mock.patch("chat_backend.services.drive_sync.httpx.Client")
def test_refresh_microsoft_token_rotates_refresh_token(self, mock_client_cls):
connection = make_drive_connection(
self.company,
provider=DriveConnection.Provider.MICROSOFT,
refresh_token="old-refresh",
)
cm, client = _fake_client(
post_side_effect=[
_fake_response(
json_data={
"access_token": "ms-new-token",
"refresh_token": "ms-new-refresh",
"expires_in": 3600,
}
)
]
)
mock_client_cls.return_value = cm
token = refresh_microsoft_token(connection)
self.assertEqual(token, "ms-new-token")
connection.refresh_from_db()
self.assertEqual(connection.refresh_token, "ms-new-refresh")
@mock.patch("chat_backend.services.drive_sync.refresh_google_token")
def test_ensure_fresh_token_skips_refresh_when_not_expired(self, mock_refresh):
self.connection.token_expires_at = timezone.now() + timedelta(hours=1)
self.connection.save(update_fields=["token_expires_at"])
token = ensure_fresh_token(self.connection)
self.assertEqual(token, "old-token")
mock_refresh.assert_not_called()
@mock.patch("chat_backend.services.drive_sync.refresh_google_token")
def test_ensure_fresh_token_refreshes_when_expired(self, mock_refresh):
mock_refresh.return_value = "refreshed"
token = ensure_fresh_token(self.connection)
self.assertEqual(token, "refreshed")
mock_refresh.assert_called_once_with(self.connection)
class ListGoogleFilesTestCase(TestCase):
def setUp(self):
self.company = make_company()
@mock.patch("chat_backend.services.drive_sync.httpx.Client")
def test_lists_root_when_no_resources_selected(self, mock_client_cls):
connection = make_drive_connection(self.company, selected_resource_ids=[])
cm, client = _fake_client(
get_side_effect=[
_fake_response(
json_data={
"files": [
{"id": "f1", "name": "a.pdf", "mimeType": "application/pdf"},
{"id": "folder1", "name": "Sub", "mimeType": drive_sync.GOOGLE_FOLDER_MIME},
]
}
)
]
)
mock_client_cls.return_value = cm
files = drive_sync._list_google_files(connection, "token")
self.assertEqual([f.id for f in files], ["f1"])
called_params = client.get.call_args.kwargs["params"]
self.assertEqual(called_params["q"], "'root' in parents and trashed = false")
@mock.patch("chat_backend.services.drive_sync.httpx.Client")
def test_lists_selected_shared_drive_with_corpora_params(self, mock_client_cls):
connection = make_drive_connection(
self.company,
kind=DriveConnection.Kind.COMPANY,
selected_resource_ids=["drive-123"],
)
cm, client = _fake_client(
get_side_effect=[_fake_response(json_data={"files": []})]
)
mock_client_cls.return_value = cm
drive_sync._list_google_files(connection, "token")
called_params = client.get.call_args.kwargs["params"]
self.assertEqual(called_params["corpora"], "drive")
self.assertEqual(called_params["driveId"], "drive-123")
@mock.patch("chat_backend.services.drive_sync.httpx.Client")
def test_list_raises_drive_sync_error_on_http_failure(self, mock_client_cls):
connection = make_drive_connection(self.company)
cm, client = _fake_client(
get_side_effect=[_fake_response(status_code=403, text="forbidden")]
)
mock_client_cls.return_value = cm
with self.assertRaises(DriveSyncError) as ctx:
drive_sync._list_google_files(connection, "token")
self.assertEqual(ctx.exception.code, "list_failed")
class DownloadGoogleFileTestCase(TestCase):
def test_exports_google_native_document_to_docx(self):
remote = RemoteFile(
id="doc1",
name="Report",
mime_type="application/vnd.google-apps.document",
etag="etag1",
)
client = mock.MagicMock()
client.get.return_value = _fake_response(content=b"docx-bytes")
content, filename, content_type = drive_sync._download_google_file(
client, {"Authorization": "Bearer x"}, remote
)
self.assertEqual(content, b"docx-bytes")
self.assertEqual(filename, "Report.docx")
self.assertIn("wordprocessingml", content_type)
self.assertIn("/export", client.get.call_args.args[0])
def test_downloads_regular_file_directly(self):
remote = RemoteFile(id="f1", name="notes.pdf", mime_type="application/pdf", etag="e1")
client = mock.MagicMock()
client.get.return_value = _fake_response(content=b"%PDF-bytes")
content, filename, content_type = drive_sync._download_google_file(
client, {}, remote
)
self.assertEqual(filename, "notes.pdf")
self.assertEqual(content_type, "application/pdf")
class ListMicrosoftFilesTestCase(TestCase):
def setUp(self):
self.company = make_company()
@mock.patch("chat_backend.services.drive_sync.httpx.Client")
def test_personal_lists_me_drive_root(self, mock_client_cls):
connection = make_drive_connection(
self.company, provider=DriveConnection.Provider.MICROSOFT, selected_resource_ids=[]
)
cm, client = _fake_client(
get_side_effect=[
_fake_response(
json_data={
"value": [
{"id": "i1", "name": "a.docx", "file": {"mimeType": "application/msword"}},
{"id": "folder1", "name": "Sub", "folder": {}},
]
}
)
]
)
mock_client_cls.return_value = cm
files = drive_sync._list_microsoft_files(connection, "token")
self.assertEqual([f.id for f in files], ["i1"])
self.assertIn("/me/drive/root/children", client.get.call_args.args[0])
@mock.patch("chat_backend.services.drive_sync.httpx.Client")
def test_company_requires_selected_sites(self, mock_client_cls):
connection = make_drive_connection(
self.company,
provider=DriveConnection.Provider.MICROSOFT,
kind=DriveConnection.Kind.COMPANY,
selected_resource_ids=["site-1"],
)
cm, client = _fake_client(
get_side_effect=[_fake_response(json_data={"value": [{"id": "i1", "name": "a.pptx"}]})]
)
mock_client_cls.return_value = cm
files = drive_sync._list_microsoft_files(connection, "token")
self.assertEqual(files[0].context_id, "site-1")
self.assertIn("/sites/site-1/drive/root/children", client.get.call_args.args[0])
@mock.patch("chat_backend.services.drive_sync.httpx.Client")
def test_company_without_selected_sites_returns_no_files(self, mock_client_cls):
connection = make_drive_connection(
self.company,
provider=DriveConnection.Provider.MICROSOFT,
kind=DriveConnection.Kind.COMPANY,
selected_resource_ids=[],
)
files = drive_sync._list_microsoft_files(connection, "token")
self.assertEqual(files, [])
mock_client_cls.assert_called_once()
@mock.patch("chat_backend.services.drive_sync.httpx.Client")
def test_follows_pagination_next_link(self, mock_client_cls):
connection = make_drive_connection(
self.company, provider=DriveConnection.Provider.MICROSOFT
)
cm, client = _fake_client(
get_side_effect=[
_fake_response(
json_data={
"value": [{"id": "i1", "name": "a.docx"}],
"@odata.nextLink": "https://graph.microsoft.com/v1.0/me/drive/root/children?page=2",
}
),
_fake_response(json_data={"value": [{"id": "i2", "name": "b.docx"}]}),
]
)
mock_client_cls.return_value = cm
files = drive_sync._list_microsoft_files(connection, "token")
self.assertEqual([f.id for f in files], ["i1", "i2"])
self.assertEqual(client.get.call_count, 2)
class DownloadMicrosoftFileTestCase(TestCase):
def setUp(self):
self.company = make_company()
def test_personal_download_uses_me_drive(self):
connection = make_drive_connection(
self.company, provider=DriveConnection.Provider.MICROSOFT
)
remote = RemoteFile(id="i1", name="a.docx", mime_type="", etag="e1")
client = mock.MagicMock()
client.get.return_value = _fake_response(
content=b"bytes", headers={"content-type": "application/msword"}
)
content, filename, content_type = drive_sync._download_microsoft_file(
client, {}, connection, remote
)
self.assertIn("/me/drive/items/i1/content", client.get.call_args.args[0])
self.assertEqual(filename, "a.docx")
def test_company_download_uses_site_context(self):
connection = make_drive_connection(
self.company,
provider=DriveConnection.Provider.MICROSOFT,
kind=DriveConnection.Kind.COMPANY,
)
remote = RemoteFile(id="i1", name="a.docx", mime_type="", etag="e1", context_id="site-9")
client = mock.MagicMock()
client.get.return_value = _fake_response(content=b"bytes", headers={})
drive_sync._download_microsoft_file(client, {}, connection, remote)
self.assertIn("/sites/site-9/drive/items/i1/content", client.get.call_args.args[0])
class SyncConnectionTestCase(TestCase):
"""Exercises the add/update/remove reconciliation loop end to end."""
def setUp(self):
self.company = make_company()
self.connection = make_drive_connection(
self.company,
provider=DriveConnection.Provider.GOOGLE,
token_expires_at=timezone.now() + timedelta(hours=1),
)
rag_patcher = mock.patch("chat_backend.services.drive_sync.AsyncRAGService")
self.mock_rag_cls = rag_patcher.start()
self.mock_rag_cls.return_value.add_files_to_store.return_value = {
"total_added": 1,
"failed_files": [],
"processed_files": [],
}
self.addCleanup(rag_patcher.stop)
list_patcher = mock.patch("chat_backend.services.drive_sync._list_remote_files")
self.mock_list = list_patcher.start()
self.addCleanup(list_patcher.stop)
download_patcher = mock.patch("chat_backend.services.drive_sync._download_file")
self.mock_download = download_patcher.start()
self.addCleanup(download_patcher.stop)
client_patcher = mock.patch("chat_backend.services.drive_sync.httpx.Client")
self.mock_client_cls = client_patcher.start()
cm, _client = _fake_client()
self.mock_client_cls.return_value = cm
self.addCleanup(client_patcher.stop)
def test_creates_document_for_new_remote_file(self):
self.mock_list.return_value = [
RemoteFile(id="r1", name="a.pdf", mime_type="application/pdf", etag="e1")
]
self.mock_download.return_value = (b"%PDF-1.4", "a.pdf", "application/pdf")
result = sync_connection(self.connection)
self.assertEqual(result["added"], 1)
document = Document.objects.get(drive_connection=self.connection)
self.assertEqual(document.remote_file_id, "r1")
self.assertEqual(document.remote_etag, "e1")
self.assertEqual(document.source, Document.Source.GOOGLE_DRIVE)
self.assertTrue(document.processed)
self.assertTrue(document.active)
self.connection.refresh_from_db()
self.assertEqual(self.connection.last_sync_status, DriveConnection.SyncStatus.OK)
self.assertIsNotNone(self.connection.last_sync_at)
self.assertEqual(self.connection.sync_total, 1)
self.assertEqual(self.connection.sync_processed, 1)
self.assertEqual(self.connection.sync_added, 1)
self.assertEqual(self.connection.sync_updated, 0)
self.assertEqual(self.connection.sync_failed, 0)
def test_skips_unchanged_file(self):
Document.objects.create(
workspace=ensure_personal_workspace(self.connection.user),
drive_connection=self.connection,
remote_file_id="r1",
remote_etag="e1",
remote_name="a.pdf",
)
self.mock_list.return_value = [
RemoteFile(id="r1", name="a.pdf", mime_type="application/pdf", etag="e1")
]
result = sync_connection(self.connection)
self.assertEqual(result["added"], 0)
self.assertEqual(result["updated"], 0)
self.mock_download.assert_not_called()
self.connection.refresh_from_db()
self.assertEqual(self.connection.sync_total, 1)
self.assertEqual(self.connection.sync_processed, 1)
def test_updates_file_when_etag_changes(self):
existing = Document.objects.create(
workspace=ensure_personal_workspace(self.connection.user),
drive_connection=self.connection,
remote_file_id="r1",
remote_etag="old-etag",
remote_name="a.pdf",
)
self.mock_list.return_value = [
RemoteFile(id="r1", name="a.pdf", mime_type="application/pdf", etag="new-etag")
]
self.mock_download.return_value = (b"new-bytes", "a.pdf", "application/pdf")
result = sync_connection(self.connection)
self.assertEqual(result["updated"], 1)
existing.refresh_from_db()
self.assertEqual(existing.remote_etag, "new-etag")
self.assertEqual(Document.objects.filter(drive_connection=self.connection).count(), 1)
def test_removes_document_whose_remote_file_is_gone(self):
Document.objects.create(
workspace=ensure_personal_workspace(self.connection.user),
drive_connection=self.connection,
remote_file_id="deleted-remote",
remote_etag="e1",
remote_name="gone.pdf",
)
self.mock_list.return_value = []
result = sync_connection(self.connection)
self.assertEqual(result["removed"], 1)
self.assertFalse(Document.objects.filter(remote_file_id="deleted-remote").exists())
def test_download_failure_is_recorded_and_does_not_abort_sync(self):
self.mock_list.return_value = [
RemoteFile(id="r1", name="a.pdf", mime_type="application/pdf", etag="e1"),
RemoteFile(id="r2", name="b.pdf", mime_type="application/pdf", etag="e2"),
]
self.mock_download.side_effect = [
DriveSyncError("download_failed", "boom"),
(b"ok-bytes", "b.pdf", "application/pdf"),
]
result = sync_connection(self.connection)
self.assertEqual(len(result["failed"]), 1)
self.assertEqual(result["added"], 1)
self.assertEqual(Document.objects.filter(drive_connection=self.connection).count(), 1)
def test_list_failure_marks_connection_error(self):
self.mock_list.side_effect = DriveSyncError("list_failed", "quota exceeded")
result = sync_connection(self.connection)
self.assertEqual(result["error"], "quota exceeded")
self.connection.refresh_from_db()
self.assertEqual(self.connection.last_sync_status, DriveConnection.SyncStatus.ERROR)
self.assertEqual(self.connection.last_sync_error, "quota exceeded")
def test_creates_personal_workspace_when_user_has_none(self):
from chat_backend.models import DocumentWorkspace
self.assertFalse(
DocumentWorkspace.objects.filter(user=self.connection.user).exists()
)
self.mock_list.return_value = []
sync_connection(self.connection)
self.assertTrue(
DocumentWorkspace.objects.filter(
user=self.connection.user, company__isnull=True
).exists()
)
llm_be/chat_backend/tests/test_services_llm.py
+20 -61
View File
@@ -2,7 +2,6 @@ from django.test import SimpleTestCase
from langchain_core.messages import AIMessage, HumanMessage
from chat_backend.services.llm_service import AsyncLLMService, SyncLLMService
from chat_backend.services.prompt_budget import format_history, window_history
from .fakes import FakeChain
@@ -19,87 +18,47 @@ class AsyncLLMServiceTestCase(SimpleTestCase):
def setUp(self):
self.service = AsyncLLMService()
async def test_format_history_labels_speakers(self):
history = await self.service._format_history(
[HumanMessage(content="hello"), AIMessage(content="hi")]
)
self.assertEqual(history, "User: hello\nAI: hi")
async def test_generate_response_streams_chunks(self):
self.service.conversation_chain = FakeChain(chunks=["Hel", "lo!"])
chunks = [
chunk
async for chunk in self.service.generate_response(
conversation(1),
"hello",
conversation_id=1,
use_conversation_context=True,
conversation(1), "hello", conversation_id=1
)
]
self.assertEqual("".join(chunks), "Hello!")
async def test_generate_response_sends_single_history_window(self):
async def test_generate_response_sends_full_and_recent_history(self):
self.service.conversation_chain = FakeChain(chunks=["ok"])
messages = conversation(4) # 8 messages
async for _ in self.service.generate_response(
messages, "latest", 1, use_conversation_context=True
):
async for _ in self.service.generate_response(messages, "latest", 1):
pass
payload = self.service.conversation_chain.calls[0]
self.assertEqual(payload["query"], "latest")
# Latest user turn is in {query}; history is prior turns only, once.
self.assertIn("history", payload)
self.assertNotIn("recent_conversation", payload)
self.assertNotIn("conversation", payload)
# 8 messages → drop last (query) → 7 prior lines max in window.
self.assertLessEqual(len(payload["history"].splitlines()), 7)
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):
service = AsyncLLMService(grounded=True, sources_block='[1] "T" — x.com — undated')
service.conversation_chain = FakeChain(chunks=["ok"])
async for _ in service.generate_response(
[], "q", 1, use_conversation_context=True
):
pass
payload = service.conversation_chain.calls[0]
self.assertIn("[1]", payload["sources"])
self.assertEqual(len(payload["conversation"].splitlines()), 8)
self.assertEqual(len(payload["recent_conversation"].splitlines()), 6)
self.assertTrue(payload["recent_conversation"].endswith("AI: answer 3"))
class SyncLLMServiceTestCase(SimpleTestCase):
def test_constructs(self):
self.assertIsNotNone(SyncLLMService())
def test_generate_response_streams_chunks(self):
service = SyncLLMService()
service.conversation_chain = FakeChain(chunks=["one ", "two"])
chunks = list(service.generate_response(conversation=None, query="hello"))
class PromptBudgetTestCase(SimpleTestCase):
def test_window_history_drops_oldest_first(self):
messages = [HumanMessage(content="x" * 40) for _ in range(10)]
kept = window_history(messages, budget_tokens=30, reserved_tokens=0)
self.assertLess(len(kept), 10)
# Newest messages survive.
self.assertEqual(kept[-1].content, messages[-1].content)
def test_format_history_labels_speakers(self):
text = format_history(
[HumanMessage(content="hello"), AIMessage(content="hi")]
self.assertEqual("".join(chunks), "one two")
self.assertEqual(
service.conversation_chain.calls, [{"query": "hello", "conversation": None}]
)
self.assertEqual(text, "User: hello\nAI: hi")
llm_be/chat_backend/tests/test_services_rag.py
+10 -111
View File
@@ -62,8 +62,7 @@ class RAGServiceTestCase(TransactionTestCase):
self.addCleanup(reset_singletons)
self.service = AsyncRAGService()
self.company = make_company()
self.workspace = make_workspace(self.company)
self.workspace = make_workspace(make_company())
def _patch(self, target):
patcher = mock.patch(target)
@@ -138,7 +137,6 @@ class RAGServiceTestCase(TransactionTestCase):
added = self.service.vector_store.add_documents.call_args[0][0]
self.assertIn("ingest me", added[0].page_content)
self.assertEqual(added[0].metadata["workspace_id"], self.workspace.id)
self.assertEqual(added[0].metadata["company_id"], self.company.id)
self.assertEqual(added[0].metadata["document_id"], document.id)
def test_ingest_documents_deletes_the_materialized_temp_file(self):
@@ -200,37 +198,6 @@ class RAGServiceTestCase(TransactionTestCase):
self.assertEqual(results["total_added"], 1)
self.assertEqual(results["failed_files"], [])
def test_add_files_to_store_includes_document_id_and_active_metadata(self):
document = self._text_document(body=b"from the database")
self.service.add_files_to_store(
[
(
document.file,
document.file.name,
self.workspace.id,
document.id,
True,
)
],
workspace_id=self.workspace.id,
)
added = self.service.vector_store.add_documents.call_args[0][0]
self.assertEqual(added[0].metadata["document_id"], document.id)
self.assertTrue(added[0].metadata["active"])
def test_add_files_to_store_defaults_active_true_without_document_id(self):
path = self._temp_text_file(b"upload body")
self.service.add_files_to_store(
[(path, "upload.txt", self.workspace.id)], workspace_id=self.workspace.id
)
added = self.service.vector_store.add_documents.call_args[0][0]
self.assertIsNone(added[0].metadata["document_id"])
self.assertTrue(added[0].metadata["active"])
def test_add_files_to_store_records_failures(self):
results = self.service.add_files_to_store(
[("/tmp/does-not-exist.txt", "missing.txt", self.workspace.id)],
@@ -248,48 +215,6 @@ class RAGServiceTestCase(TransactionTestCase):
first_store.delete_collection.assert_called_once()
self.assertIsNot(self.service.vector_store, first_store)
def test_delete_document_vectors_deletes_by_document_id_metadata(self):
self.service.delete_document_vectors(42)
self.service.vector_store.delete.assert_called_once_with(
where={"document_id": 42}
)
self.service.vector_store.persist.assert_called()
def test_delete_document_vectors_is_a_noop_without_a_document_id(self):
self.service.delete_document_vectors(None)
self.service.vector_store.delete.assert_not_called()
def test_set_document_active_updates_matching_chunk_metadata(self):
self.service.vector_store.get.return_value = {
"ids": ["a", "b"],
"metadatas": [
{"document_id": 7, "active": True, "workspace_id": 1},
{"document_id": 7, "active": True, "workspace_id": 1},
],
}
self.service.set_document_active(7, False)
self.service.vector_store.get.assert_called_once_with(
where={"document_id": 7}
)
self.service.vector_store._collection.update.assert_called_once_with(
ids=["a", "b"],
metadatas=[
{"document_id": 7, "active": False, "workspace_id": 1},
{"document_id": 7, "active": False, "workspace_id": 1},
],
)
def test_set_document_active_is_a_noop_when_no_chunks_found(self):
self.service.vector_store.get.return_value = {"ids": [], "metadatas": []}
self.service.set_document_active(999, True)
self.service.vector_store._collection.update.assert_not_called()
async def test_search_documents_filters_by_workspace(self):
retriever = self.service.vector_store.as_retriever.return_value
retriever.aget_relevant_documents = mock.AsyncMock(
@@ -301,22 +226,18 @@ class RAGServiceTestCase(TransactionTestCase):
self.assertEqual([doc.page_content for doc in docs], ["chunk"])
self.service.vector_store.as_retriever.assert_called_with(
search_type="mmr",
search_kwargs={
"k": 2,
"filter": {
"$and": [
{"workspace_id": self.workspace.id},
{"active": True},
]
},
},
search_kwargs={"k": 2, "filter": {"workspace_id": self.workspace.id}},
)
async def test_search_documents_without_workspace_fails_closed(self):
with self.assertRaises(ValueError):
async def test_search_documents_without_workspace_has_no_filter(self):
retriever = self.service.vector_store.as_retriever.return_value
retriever.aget_relevant_documents = mock.AsyncMock(return_value=[])
await self.service.search_documents("revenue")
self.service.vector_store.as_retriever.assert_not_called()
self.service.vector_store.as_retriever.assert_called_with(
search_type="mmr", search_kwargs={"k": 4, "filter": None}
)
async def test_format_history_labels_speakers(self):
history = await self.service._format_history(
@@ -335,10 +256,7 @@ class RAGServiceTestCase(TransactionTestCase):
chunks = [
chunk
async for chunk in self.service.generate_response(
conversation,
"what is our policy?",
self.workspace,
use_conversation_context=True,
conversation, "what is our policy?", self.workspace
)
]
@@ -348,25 +266,6 @@ class RAGServiceTestCase(TransactionTestCase):
self.assertEqual(payload["workspace"], self.workspace)
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):
document = await sync_to_async(self._text_document)()
other_workspace = await sync_to_async(make_workspace)(
llm_be/chat_backend/tests/test_settings_cors.py
-8
View File
@@ -59,11 +59,3 @@ class CapacitorWebviewOriginTests(SimpleTestCase):
from django.conf import settings
self.assertFalse(settings.CORS_ALLOW_CREDENTIALS)
def test_default_cors_includes_beta_frontend(self):
from django.conf import settings
self.assertIn(
"https://beta.chat.aimloperations.com",
settings.CORS_ALLOWED_ORIGINS,
)
llm_be/chat_backend/tests/test_signals.py
+7 -17
View File
@@ -3,8 +3,6 @@ from unittest import mock
from django.test import TestCase
from chat_backend.models import Document
from .factories import make_company, make_document, make_workspace
RAG_SERVICE = "chat_backend.services.rag_services.AsyncRAGService"
@@ -14,13 +12,12 @@ class DocumentSignalTestCase(TestCase):
def setUp(self):
self.workspace = make_workspace(make_company())
def test_creating_a_document_does_not_trigger_the_rag_service(self):
"""#45: uploads ingest via DocumentUploadView.process_document, not a signal."""
def test_creating_a_document_reindexes_the_vector_store(self):
with mock.patch.dict(os.environ, {"SKIP_RAG_INIT": ""}):
with mock.patch(RAG_SERVICE) as service:
make_document(self.workspace)
service.assert_not_called()
service.return_value.ingest_documents.assert_called_once_with()
def test_updating_a_document_does_not_reindex(self):
document = make_document(self.workspace)
@@ -32,18 +29,14 @@ class DocumentSignalTestCase(TestCase):
service.assert_not_called()
def test_deleting_a_document_removes_only_its_own_vectors(self):
def test_deleting_a_document_reindexes_the_vector_store(self):
document = make_document(self.workspace)
document_id = document.id
with mock.patch.dict(os.environ, {"SKIP_RAG_INIT": ""}):
with mock.patch(RAG_SERVICE) as service:
document.delete()
service.return_value.delete_document_vectors.assert_called_once_with(
document_id
)
service.return_value.ingest_documents.assert_not_called()
service.return_value.ingest_documents.assert_called_once_with()
def test_skip_rag_init_keeps_signals_inert(self):
with mock.patch.dict(os.environ, {"SKIP_RAG_INIT": "1"}):
@@ -53,12 +46,9 @@ class DocumentSignalTestCase(TestCase):
service.assert_not_called()
def test_vector_store_failures_do_not_break_deletes(self):
document = make_document(self.workspace)
document_id = document.id
def test_vector_store_failures_do_not_break_uploads(self):
with mock.patch.dict(os.environ, {"SKIP_RAG_INIT": ""}):
with mock.patch(RAG_SERVICE, side_effect=RuntimeError("chroma down")):
document.delete()
document = make_document(self.workspace)
self.assertFalse(Document.objects.filter(pk=document_id).exists())
self.assertIsNotNone(document.pk)
llm_be/chat_backend/tests/test_utils.py
+5 -86
View File
@@ -11,16 +11,7 @@ from chat_backend.ollama_config import (
ollama_llm_kwargs,
ollama_model,
)
from chat_backend.utils import (
TokenUsageCollector,
aiter_text_chunks,
chunk_text,
extract_token_usage,
has_usable_user_prompt,
is_heartbeat_payload,
last_day_of_month,
normalize_user_message,
)
from chat_backend.utils import extract_token_usage, last_day_of_month
class ExtractTokenUsageTestCase(SimpleTestCase):
@@ -68,51 +59,6 @@ class ExtractTokenUsageTestCase(SimpleTestCase):
(10, 20),
)
def test_reads_generation_info_attribute(self):
class Chunk:
generation_info = {
"done": True,
"prompt_eval_count": 22,
"eval_count": 55,
}
self.assertEqual(extract_token_usage(Chunk()), (22, 55))
def test_chunk_text_from_generation_chunk(self):
class Chunk:
text = "hello"
generation_info = {"prompt_eval_count": 1, "eval_count": 2}
self.assertEqual(chunk_text(Chunk()), "hello")
self.assertEqual(chunk_text("plain"), "plain")
self.assertEqual(chunk_text(None), "")
async def test_aiter_text_chunks_collects_final_usage(self):
class Chunk:
def __init__(self, text, info=None):
self.text = text
self.generation_info = info or {}
async def stream():
yield Chunk("Hel")
yield Chunk("lo", {"prompt_eval_count": 11, "eval_count": 3})
usage = TokenUsageCollector()
texts = [t async for t in aiter_text_chunks(stream(), usage)]
self.assertEqual("".join(texts), "Hello")
self.assertEqual(usage.pair, (11, 3))
async def test_aiter_text_chunks_without_usage_stays_null(self):
async def stream():
yield "only-text"
usage = TokenUsageCollector()
texts = [t async for t in aiter_text_chunks(stream(), usage)]
self.assertEqual(texts, ["only-text"])
self.assertEqual(usage.pair, (None, None))
class LastDayOfMonthTestCase(SimpleTestCase):
@parameterized.expand(
@@ -137,11 +83,7 @@ class LastDayOfMonthTestCase(SimpleTestCase):
@override_settings(
OLLAMA_BASE_URL="http://10.0.0.128:11434",
OLLAMA_MODEL="llama3.2",
OLLAMA_MODEL_THINKING="llama3.2",
OLLAMA_MODEL_FAST="gemma4:latest",
OLLAMA_MODEL_UTILITY="llama3.2",
OLLAMA_EMBED_MODEL="nomic-embed-text",
OLLAMA_NUM_CTX_THINKING=16384,
)
class OllamaConfigTestCase(SimpleTestCase):
def test_reads_settings(self):
@@ -183,35 +125,12 @@ class OllamaConfigFallbackTestCase(SimpleTestCase):
del settings.OLLAMA_BASE_URL
del settings.OLLAMA_MODEL
del settings.OLLAMA_EMBED_MODEL
if hasattr(settings, "OLLAMA_MODEL_THINKING"):
del settings.OLLAMA_MODEL_THINKING
self.assertEqual(ollama_base_url(), "http://127.0.0.1:11434")
self.assertEqual(ollama_model(), "gpt-oss:20b")
# Embeddings never fall back to a chat model (#62).
self.assertEqual(ollama_embed_model(), "nomic-embed-text")
self.assertEqual(ollama_model(), "llama3.2")
self.assertEqual(ollama_embed_model(), "llama3.2")
def test_embed_model_does_not_fall_back_to_chat_model(self):
def test_embed_model_falls_back_to_chat_model(self):
with override_settings(OLLAMA_MODEL="llama3.2"):
del settings.OLLAMA_EMBED_MODEL
self.assertEqual(ollama_embed_model(), "nomic-embed-text")
class UserPromptGuardTestCase(SimpleTestCase):
def test_heartbeat_payload_detected(self):
self.assertTrue(is_heartbeat_payload({"type": "ping", "email": "a@b.com"}))
self.assertFalse(is_heartbeat_payload({"message": "hi"}))
self.assertFalse(is_heartbeat_payload(None))
def test_normalize_user_message(self):
self.assertIsNone(normalize_user_message(None))
self.assertIsNone(normalize_user_message(""))
self.assertIsNone(normalize_user_message(" \n\t"))
self.assertEqual(normalize_user_message(" hello "), "hello")
def test_has_usable_user_prompt(self):
self.assertFalse(has_usable_user_prompt(None))
self.assertFalse(has_usable_user_prompt(" "))
self.assertFalse(has_usable_user_prompt("", file="base64"))
self.assertTrue(has_usable_user_prompt("hi"))
self.assertTrue(has_usable_user_prompt(" hi "))
self.assertEqual(ollama_embed_model(), "llama3.2")
llm_be/chat_backend/tests/test_views_analytics.py
-57
View File
@@ -127,60 +127,3 @@ class AdminAnalyticsTestCase(APITestCase):
self.assertEqual(current["range"], [10, 20])
self.assertEqual(current["avg"], 15)
self.assertEqual(current["median"], 20)
class UserPromptHeatmapTestCase(APITestCase):
def setUp(self):
self.user = make_user(company=make_company())
self.other = make_user(email="other@example.com", company=make_company("Other"))
self.client.force_authenticate(user=self.user)
self.conversation = make_conversation(user=self.user)
self.url = reverse("analytics_user_prompt_heatmap")
def test_empty_heatmap(self):
response = self.client.get(self.url, {"tz": "UTC"})
self.assertEqual(response.status_code, status.HTTP_200_OK)
self.assertEqual(response.data["total"], 0)
self.assertEqual(response.data["max"], 0)
self.assertEqual(response.data["matrix"], [[0] * 24 for _ in range(7)])
self.assertIsNone(response.data["peak_cell"])
def test_bins_user_prompts_in_utc(self):
# Wednesday 2026-07-15 14:30 UTC → weekday 2, hour 14
created = datetime.datetime(2026, 7, 15, 14, 30, tzinfo=datetime.timezone.utc)
make_prompt(self.conversation, message="a", created=created)
make_prompt(self.conversation, message="b", created=created)
make_prompt(
self.conversation,
message="llm",
user_created=False,
created=created,
)
other_convo = make_conversation(user=self.other)
make_prompt(other_convo, message="other", created=created)
response = self.client.get(self.url, {"tz": "UTC"})
self.assertEqual(response.status_code, status.HTTP_200_OK)
self.assertEqual(response.data["total"], 2)
self.assertEqual(response.data["max"], 2)
self.assertEqual(response.data["matrix"][2][14], 2)
self.assertEqual(response.data["most_active_day"], "Wed")
self.assertEqual(response.data["most_active_hour"], 14)
self.assertEqual(
response.data["peak_cell"],
{"day": "Wed", "hour": 14, "count": 2},
)
def test_applies_timezone(self):
# 2026-07-15 14:30 UTC → America/Chicago CDT (UTC-5) = 09:30 Wed
created = datetime.datetime(2026, 7, 15, 14, 30, tzinfo=datetime.timezone.utc)
make_prompt(self.conversation, message="a", created=created)
response = self.client.get(self.url, {"tz": "America/Chicago"})
self.assertEqual(response.status_code, status.HTTP_200_OK)
self.assertEqual(response.data["tz"], "America/Chicago")
self.assertEqual(response.data["matrix"][2][9], 1)
self.assertEqual(response.data["matrix"][2][14], 0)
llm_be/chat_backend/tests/test_views_conversations.py
-38
View File
@@ -57,7 +57,6 @@ class ConversationPreferencesTestCase(APITestCase):
self.assertEqual(response.status_code, status.HTTP_200_OK)
self.assertTrue(response.data["order"])
self.assertFalse(response.data["use_conversation_context"])
def test_post_toggles_and_persists_order(self):
response = self.client.post(self.url)
@@ -67,33 +66,6 @@ class ConversationPreferencesTestCase(APITestCase):
self.user.refresh_from_db()
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):
def setUp(self):
@@ -115,16 +87,6 @@ class ConversationDetailViewTestCase(APITestCase):
[("hello", True), ("hi there", False)],
)
def test_get_rejects_other_users_conversation(self):
other = make_user(email="other-tenant@example.com", company=make_company("Other"))
foreign = make_conversation(user=other)
make_prompt(foreign, message="secret")
response = self.client.get(self.url, {"conversation_id": foreign.id})
self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND)
self.assertEqual(Prompt.objects.filter(conversation=foreign).count(), 1)
def test_post_stores_assistant_prompt(self):
response = self.client.post(
self.url,
llm_be/chat_backend/tests/test_views_documents.py
+8 -213
View File
@@ -1,14 +1,10 @@
import os
from unittest import mock
from django.test import override_settings
from django.urls import reverse
from rest_framework import status
from rest_framework.test import APITestCase
from chat_backend.models import Document, DocumentWorkspace, StoredFile
from monetization.models import SubscriptionPlan, UserSubscription
from monetization.services.plans import assign_plan, seed_subscription_plans
from .factories import (
make_company,
@@ -93,18 +89,14 @@ class DocumentUploadViewTestCase(APITestCase):
self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
def test_upload_without_workspace_creates_a_default_one(self):
"""#46: missing workspace is auto-created instead of 404ing the upload."""
def test_upload_without_workspace_returns_404(self):
self.workspace.delete()
response = self.client.post(
self.url, {"file": pdf_upload()}, format="multipart"
)
self.assertEqual(response.status_code, status.HTTP_201_CREATED)
created_workspace = DocumentWorkspace.objects.get(company=self.company)
self.assertEqual(created_workspace.name, "Default")
self.assertEqual(Document.objects.get().workspace_id, created_workspace.id)
self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND)
def test_list_returns_documents_of_own_workspace(self):
make_document(self.workspace)
@@ -112,62 +104,17 @@ class DocumentUploadViewTestCase(APITestCase):
response = self.client.get(self.url)
self.assertEqual(response.status_code, status.HTTP_200_OK)
self.assertEqual(response.data["count"], 1)
self.assertEqual(response.data["scope"], "company")
self.assertEqual(len(response.data["results"]), 1)
self.assertIn("test", response.data["results"][0]["file"])
self.assertIn("pdf", response.data["results"][0]["file"])
self.assertEqual(len(response.data), 1)
self.assertIn("test", response.data[0]["file"])
self.assertIn("pdf", response.data[0]["file"])
def test_list_without_workspace_creates_a_default_one(self):
"""#46: missing workspace is auto-created instead of 404ing the list."""
def test_list_without_workspace_returns_404(self):
self.workspace.delete()
response = self.client.get(self.url)
self.assertEqual(response.status_code, status.HTTP_200_OK)
self.assertEqual(response.data["count"], 0)
self.assertEqual(response.data["results"], [])
self.assertTrue(
DocumentWorkspace.objects.filter(
company=self.company, name="Default"
).exists()
)
def test_list_paginates_searches_and_sorts(self):
"""#60: page / search / ordering query params."""
make_document(self.workspace, name="alpha.pdf")
make_document(self.workspace, name="beta.pdf")
make_document(self.workspace, name="gamma.pdf")
page_one = self.client.get(
self.url, {"page": 1, "page_size": 2, "ordering": "name"}
)
self.assertEqual(page_one.status_code, status.HTTP_200_OK)
self.assertEqual(page_one.data["count"], 3)
self.assertEqual(page_one.data["page"], 1)
self.assertEqual(page_one.data["page_size"], 2)
self.assertEqual(len(page_one.data["results"]), 2)
self.assertIn("alpha", page_one.data["results"][0]["file"])
search = self.client.get(self.url, {"search": "beta"})
self.assertEqual(search.status_code, status.HTTP_200_OK)
self.assertEqual(search.data["count"], 1)
self.assertIn("beta", search.data["results"][0]["file"])
def test_list_personal_scope_isolated_from_company(self):
"""#60: company members can list personal workspace separately."""
make_document(self.workspace, name="company.pdf")
personal = make_workspace(user=self.user, name="Personal")
make_document(personal, name="mine.pdf")
company_list = self.client.get(self.url, {"workspace": "company"})
personal_list = self.client.get(self.url, {"workspace": "personal"})
self.assertEqual(company_list.data["count"], 1)
self.assertIn("company", company_list.data["results"][0]["file"])
self.assertEqual(personal_list.data["count"], 1)
self.assertIn("mine", personal_list.data["results"][0]["file"])
self.assertEqual(personal_list.data["scope"], "personal")
self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND)
self.assertEqual(response.data["error"], "Workspace not found")
class DocumentDetailViewTestCase(APITestCase):
@@ -193,155 +140,3 @@ class DocumentDetailViewTestCase(APITestCase):
self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND)
self.assertEqual(response.data["error"], "Document not found")
def test_get_returns_the_document(self):
document = make_document(self.workspace)
url = reverse("documents_details", kwargs={"document_id": document.id})
response = self.client.get(url)
self.assertEqual(response.status_code, status.HTTP_200_OK)
self.assertEqual(response.data["id"], document.id)
self.assertIn("test", response.data["file"])
def test_patch_toggles_active_and_updates_vector_metadata(self):
document = make_document(self.workspace)
document.active = True
document.save(update_fields=["active"])
url = reverse("documents_details", kwargs={"document_id": document.id})
with mock.patch("chat_backend.views.AsyncRAGService") as rag_service:
response = self.client.patch(url, {"active": False}, format="json")
self.assertEqual(response.status_code, status.HTTP_200_OK)
self.assertFalse(response.data["active"])
document.refresh_from_db()
self.assertFalse(document.active)
rag_service.return_value.set_document_active.assert_called_once_with(
document.id, False
)
def test_patch_requires_active_field(self):
document = make_document(self.workspace)
url = reverse("documents_details", kwargs={"document_id": document.id})
response = self.client.patch(url, {}, format="json")
self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
def test_patch_unknown_document_returns_404(self):
url = reverse("documents_details", kwargs={"document_id": 4242})
response = self.client.patch(url, {"active": True}, format="json")
self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND)
def test_delete_removes_the_document_and_its_vectors(self):
document = make_document(self.workspace)
document_id = document.id
url = reverse("documents_details", kwargs={"document_id": document_id})
with mock.patch.dict(os.environ, {"SKIP_RAG_INIT": ""}):
with mock.patch(
"chat_backend.services.rag_services.AsyncRAGService"
) as rag_service:
response = self.client.delete(url)
self.assertEqual(response.status_code, status.HTTP_204_NO_CONTENT)
self.assertFalse(Document.objects.filter(id=document_id).exists())
rag_service.return_value.delete_document_vectors.assert_called_once_with(
document_id
)
def test_delete_unknown_document_returns_404(self):
url = reverse("documents_details", kwargs={"document_id": 4242})
response = self.client.delete(url)
self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND)
@override_settings(ENFORCE_SUBSCRIPTION_GATES=True)
class DocumentRagFeatureGateTestCase(APITestCase):
"""#44: RAG document endpoints must respect the plan's ``rag`` feature gate.
The rest of this module runs with ``ENFORCE_SUBSCRIPTION_GATES=False``
(set by the test runner) so pre-existing behavior keeps working
regardless of plan; this class opts back in explicitly.
"""
def setUp(self):
seed_subscription_plans()
self.company = make_company()
self.user = make_user(company=self.company)
self.client.force_authenticate(user=self.user)
self.workspace = make_workspace(self.company)
rag_patcher = mock.patch("chat_backend.views.AsyncRAGService")
self.rag_service = rag_patcher.start()
self.addCleanup(rag_patcher.stop)
def _assign(self, slug):
plan = SubscriptionPlan.objects.get(slug=slug)
assign_plan(self.user, plan=plan, source=UserSubscription.Source.ADMIN)
def test_standard_plan_is_denied_workspace_list(self):
self._assign("standard")
response = self.client.get(reverse("document_workspaces"))
self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
self.assertEqual(response.data["code"], "feature_not_allowed")
self.assertIn("error", response.data)
self.assertIn("details", response.data)
def test_standard_plan_is_denied_document_upload(self):
self._assign("standard")
response = self.client.post(
reverse("documents"), {"file": pdf_upload()}, format="multipart"
)
self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
self.assertEqual(response.data["code"], "feature_not_allowed")
self.assertEqual(Document.objects.count(), 0)
def test_standard_plan_is_denied_document_list(self):
self._assign("standard")
response = self.client.get(reverse("documents"))
self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
def test_standard_plan_is_denied_document_detail(self):
self._assign("standard")
document = make_document(self.workspace)
url = reverse("documents_details", kwargs={"document_id": document.id})
response = self.client.get(url)
self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
def test_no_subscription_is_denied(self):
response = self.client.post(
reverse("documents"), {"file": pdf_upload()}, format="multipart"
)
self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
self.assertEqual(response.data["code"], "subscription_required")
def test_founders_plan_is_allowed_document_upload(self):
self._assign("founders")
response = self.client.post(
reverse("documents"), {"file": pdf_upload()}, format="multipart"
)
self.assertEqual(response.status_code, status.HTTP_201_CREATED)
self.assertEqual(Document.objects.count(), 1)
def test_founders_plan_is_allowed_workspace_list(self):
self._assign("founders")
response = self.client.get(reverse("document_workspaces"))
self.assertEqual(response.status_code, status.HTTP_200_OK)
llm_be/chat_backend/tests/test_views_drive.py
-277
View File
@@ -1,277 +0,0 @@
"""Tests for Drive connection management API (#47-#52)."""
from __future__ import annotations
from unittest import mock
from django.test import override_settings
from django.urls import reverse
from rest_framework import status
from rest_framework.test import APITestCase
from chat_backend.models import DriveConnection
from monetization.models import SubscriptionPlan, UserSubscription
from monetization.services.plans import assign_plan, seed_subscription_plans
from .factories import make_company, make_drive_connection, make_user
class DriveConnectionListViewTestCase(APITestCase):
def setUp(self):
self.company = make_company()
self.user = make_user(email="user@example.com", company=self.company)
self.client.force_authenticate(user=self.user)
self.url = reverse("drive_connections")
def test_lists_own_personal_and_company_connections(self):
personal = make_drive_connection(
self.company, kind=DriveConnection.Kind.PERSONAL, user=self.user
)
company_conn = make_drive_connection(
self.company, kind=DriveConnection.Kind.COMPANY, user=None
)
other_company = make_company("Other")
other_user = make_user(email="other@example.com", company=other_company)
make_drive_connection(
other_company, kind=DriveConnection.Kind.PERSONAL, user=other_user
)
response = self.client.get(self.url)
self.assertEqual(response.status_code, status.HTTP_200_OK)
ids = {row["id"] for row in response.data}
self.assertEqual(ids, {personal.id, company_conn.id})
def test_does_not_leak_other_users_personal_connection(self):
teammate = make_user(email="teammate@example.com", company=self.company)
make_drive_connection(
self.company, kind=DriveConnection.Kind.PERSONAL, user=teammate
)
response = self.client.get(self.url)
self.assertEqual(response.data, [])
def test_response_never_includes_tokens(self):
make_drive_connection(
self.company, kind=DriveConnection.Kind.PERSONAL, user=self.user
)
response = self.client.get(self.url)
self.assertNotIn("access_token", response.data[0])
self.assertNotIn("refresh_token", response.data[0])
@override_settings(ENFORCE_SUBSCRIPTION_GATES=True)
def test_denied_on_standard_plan(self):
seed_subscription_plans()
standard = SubscriptionPlan.objects.get(slug="standard")
assign_plan(self.user, plan=standard, source=UserSubscription.Source.ADMIN)
response = self.client.get(self.url)
self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
@override_settings(ENFORCE_SUBSCRIPTION_GATES=True)
def test_allowed_on_pro_plan(self):
seed_subscription_plans()
pro = SubscriptionPlan.objects.get(slug="pro")
assign_plan(self.user, plan=pro, source=UserSubscription.Source.ADMIN)
response = self.client.get(self.url)
self.assertEqual(response.status_code, status.HTTP_200_OK)
class DriveConnectionDetailViewTestCase(APITestCase):
def setUp(self):
self.company = make_company()
self.user = make_user(email="user@example.com", company=self.company)
self.client.force_authenticate(user=self.user)
def test_owner_can_disconnect_personal_connection(self):
connection = make_drive_connection(
self.company, kind=DriveConnection.Kind.PERSONAL, user=self.user
)
url = reverse("drive_connection_detail", kwargs={"connection_id": connection.id})
response = self.client.delete(url)
self.assertEqual(response.status_code, status.HTTP_204_NO_CONTENT)
connection.refresh_from_db()
self.assertFalse(connection.is_active)
self.assertEqual(connection.access_token, "")
self.assertEqual(connection.refresh_token, "")
def test_non_owner_cannot_disconnect_personal_connection(self):
teammate = make_user(email="teammate@example.com", company=self.company)
connection = make_drive_connection(
self.company, kind=DriveConnection.Kind.PERSONAL, user=teammate
)
url = reverse("drive_connection_detail", kwargs={"connection_id": connection.id})
response = self.client.delete(url)
self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
connection.refresh_from_db()
self.assertTrue(connection.is_active)
def test_manager_can_disconnect_company_connection(self):
self.user.is_company_manager = True
self.user.save(update_fields=["is_company_manager"])
connection = make_drive_connection(
self.company, kind=DriveConnection.Kind.COMPANY, user=None
)
url = reverse("drive_connection_detail", kwargs={"connection_id": connection.id})
response = self.client.delete(url)
self.assertEqual(response.status_code, status.HTTP_204_NO_CONTENT)
def test_non_manager_cannot_disconnect_company_connection(self):
connection = make_drive_connection(
self.company, kind=DriveConnection.Kind.COMPANY, user=None
)
url = reverse("drive_connection_detail", kwargs={"connection_id": connection.id})
response = self.client.delete(url)
self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
def test_unknown_connection_404(self):
url = reverse("drive_connection_detail", kwargs={"connection_id": 999999})
response = self.client.delete(url)
self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND)
class DriveConnectionResourcesViewTestCase(APITestCase):
def setUp(self):
self.company = make_company()
self.user = make_user(email="user@example.com", company=self.company)
self.client.force_authenticate(user=self.user)
self.connection = make_drive_connection(
self.company, kind=DriveConnection.Kind.PERSONAL, user=self.user
)
self.url = reverse(
"drive_connection_resources", kwargs={"connection_id": self.connection.id}
)
def test_sets_selected_resources(self):
response = self.client.post(
self.url,
{"resource_ids": ["folder-1", "folder-2"], "resource_labels": ["Reports", "HR"]},
format="json",
)
self.assertEqual(response.status_code, status.HTTP_200_OK)
self.connection.refresh_from_db()
self.assertEqual(self.connection.selected_resource_ids, ["folder-1", "folder-2"])
self.assertEqual(self.connection.selected_resource_labels, ["Reports", "HR"])
def test_requires_resource_ids_list(self):
response = self.client.post(self.url, {}, format="json")
self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
def test_company_connection_requires_manager(self):
company_conn = make_drive_connection(
self.company, kind=DriveConnection.Kind.COMPANY, user=None
)
url = reverse(
"drive_connection_resources", kwargs={"connection_id": company_conn.id}
)
response = self.client.post(url, {"resource_ids": ["site-1"]}, format="json")
self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
class DriveConnectionSyncViewTestCase(APITestCase):
def setUp(self):
self.company = make_company()
self.user = make_user(email="user@example.com", company=self.company)
self.client.force_authenticate(user=self.user)
self.connection = make_drive_connection(
self.company, kind=DriveConnection.Kind.PERSONAL, user=self.user
)
self.url = reverse("drive_connection_sync", kwargs={"connection_id": self.connection.id})
@mock.patch("chat_backend.views_drive.enqueue_drive_sync")
def test_triggers_sync_for_owner(self, mock_enqueue):
self.connection.last_sync_status = DriveConnection.SyncStatus.PENDING
self.connection.save(update_fields=["last_sync_status"])
mock_enqueue.return_value = (self.connection, True)
response = self.client.post(self.url)
self.assertEqual(response.status_code, status.HTTP_202_ACCEPTED)
mock_enqueue.assert_called_once()
self.assertTrue(response.data["queued"])
self.assertEqual(response.data["connection"]["last_sync_status"], "pending")
@mock.patch("chat_backend.views_drive.enqueue_drive_sync")
def test_non_owner_forbidden(self, mock_enqueue):
other = make_user(email="other@example.com", company=self.company)
self.client.force_authenticate(user=other)
response = self.client.post(self.url)
self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
mock_enqueue.assert_not_called()
@mock.patch("chat_backend.views_drive.enqueue_drive_sync")
def test_disconnected_connection_rejected(self, mock_enqueue):
self.connection.is_active = False
self.connection.save(update_fields=["is_active"])
response = self.client.post(self.url)
self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
mock_enqueue.assert_not_called()
class DriveWebhookViewTestCase(APITestCase):
def setUp(self):
self.company = make_company()
self.connection = make_drive_connection(
self.company, kind=DriveConnection.Kind.PERSONAL, provider=DriveConnection.Provider.GOOGLE
)
@mock.patch("chat_backend.views_drive.enqueue_drive_sync")
def test_google_webhook_syncs_matching_connection(self, mock_enqueue):
url = reverse("drive_webhook_google")
response = self.client.post(f"{url}?connection_id={self.connection.id}")
self.assertEqual(response.status_code, status.HTTP_200_OK)
mock_enqueue.assert_called_once()
self.assertEqual(mock_enqueue.call_args.args[0].id, self.connection.id)
@mock.patch("chat_backend.views_drive.enqueue_drive_sync")
def test_google_webhook_without_connection_id_is_a_noop(self, mock_enqueue):
url = reverse("drive_webhook_google")
response = self.client.post(url)
self.assertEqual(response.status_code, status.HTTP_200_OK)
mock_enqueue.assert_not_called()
@mock.patch("chat_backend.views_drive.enqueue_drive_sync")
def test_microsoft_webhook_validation_handshake(self, mock_enqueue):
url = reverse("drive_webhook_microsoft")
response = self.client.post(f"{url}?validationToken=abc123")
self.assertEqual(response.status_code, status.HTTP_200_OK)
self.assertEqual(response.content.decode(), "abc123")
mock_enqueue.assert_not_called()
@mock.patch("chat_backend.views_drive.enqueue_drive_sync")
def test_microsoft_webhook_syncs_via_client_state(self, mock_enqueue):
ms_connection = make_drive_connection(
self.company,
kind=DriveConnection.Kind.PERSONAL,
provider=DriveConnection.Provider.MICROSOFT,
user=make_user(email="ms.user@example.com", company=self.company),
)
url = reverse("drive_webhook_microsoft")
response = self.client.post(
url,
{"value": [{"clientState": str(ms_connection.id)}]},
format="json",
)
self.assertEqual(response.status_code, status.HTTP_200_OK)
mock_enqueue.assert_called_once()
self.assertEqual(mock_enqueue.call_args.args[0].id, ms_connection.id)
llm_be/chat_backend/tests/test_views_prompt_feedback.py
-149
View File
@@ -1,149 +0,0 @@
from django.urls import reverse
from rest_framework import status
from rest_framework.test import APITestCase
from chat_backend.models import PromptFeedback, PromptMetric
from .factories import make_company, make_conversation, make_prompt, make_user
class PromptFeedbackViewTestCase(APITestCase):
def setUp(self):
self.user = make_user(company=make_company())
self.client.force_authenticate(user=self.user)
self.conversation = make_conversation(user=self.user)
self.assistant = make_prompt(
self.conversation, message="answer", user_created=False
)
self.user_prompt = make_prompt(
self.conversation, message="question", user_created=True
)
self.url = reverse("prompt_feedback")
self.details_url = reverse("conversation_details")
def test_upsert_creates_unique_row(self):
response = self.client.post(
self.url,
{"prompt_id": self.assistant.id, "rating": "up"},
format="json",
)
self.assertEqual(response.status_code, status.HTTP_200_OK)
self.assertEqual(response.data["rating"], "up")
self.assertEqual(response.data["prompt_id"], self.assistant.id)
self.assertEqual(PromptFeedback.objects.count(), 1)
again = self.client.post(
self.url,
{
"prompt_id": self.assistant.id,
"rating": "down",
"reason": "incorrect",
"comment": "wrong cite",
},
format="json",
)
self.assertEqual(again.status_code, status.HTTP_200_OK)
self.assertEqual(PromptFeedback.objects.count(), 1)
row = PromptFeedback.objects.get()
self.assertEqual(row.rating, "down")
self.assertEqual(row.reason, "incorrect")
self.assertEqual(row.comment, "wrong cite")
def test_delete_clears_vote(self):
PromptFeedback.objects.create(
prompt=self.assistant, user=self.user, rating="up"
)
response = self.client.delete(
f"{self.url}?prompt_id={self.assistant.id}"
)
self.assertEqual(response.status_code, status.HTTP_204_NO_CONTENT)
self.assertEqual(PromptFeedback.objects.count(), 0)
def test_delete_missing_vote_is_404(self):
response = self.client.delete(
f"{self.url}?prompt_id={self.assistant.id}"
)
self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND)
def test_cannot_rate_user_prompt(self):
response = self.client.post(
self.url,
{"prompt_id": self.user_prompt.id, "rating": "up"},
format="json",
)
self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
self.assertEqual(PromptFeedback.objects.count(), 0)
def test_cannot_rate_other_users_prompt(self):
other = make_user(email="other@example.com", company=make_company("O"))
foreign = make_prompt(
make_conversation(user=other), message="secret", user_created=False
)
response = self.client.post(
self.url,
{"prompt_id": foreign.id, "rating": "up"},
format="json",
)
self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND)
self.assertEqual(PromptFeedback.objects.count(), 0)
def test_conversation_details_includes_caller_feedback(self):
PromptFeedback.objects.create(
prompt=self.assistant,
user=self.user,
rating="down",
reason="unsafe",
comment="bad",
)
# Another user's vote must not leak
other = make_user(email="peer@example.com", company=self.user.company)
PromptFeedback.objects.create(
prompt=self.assistant, user=other, rating="up"
)
response = self.client.get(
self.details_url, {"conversation_id": self.conversation.id}
)
self.assertEqual(response.status_code, status.HTTP_200_OK)
by_id = {item["id"]: item for item in response.data}
self.assertIsNone(by_id[self.user_prompt.id]["feedback"])
self.assertEqual(
by_id[self.assistant.id]["feedback"],
{"rating": "down", "reason": "unsafe", "comment": "bad"},
)
def test_feedback_joinable_to_prompt_metric(self):
PromptFeedback.objects.create(
prompt=self.assistant, user=self.user, rating="up"
)
PromptMetric.objects.create(
prompt_id=self.assistant.id,
conversation_id=self.conversation.id,
event="FINISHED",
model_name="llama3.2",
start_time=self.assistant.created,
prompt_length=10,
has_file=False,
)
joined = PromptFeedback.objects.filter(
prompt_id__in=PromptMetric.objects.filter(
model_name="llama3.2"
).values_list("prompt_id", flat=True)
)
self.assertEqual(joined.count(), 1)
def test_unauthenticated_rejected(self):
self.client.force_authenticate(user=None)
response = self.client.post(
self.url,
{"prompt_id": self.assistant.id, "rating": "up"},
format="json",
)
self.assertEqual(response.status_code, status.HTTP_401_UNAUTHORIZED)
llm_be/chat_backend/tests/test_views_users.py
+1 -105
View File
@@ -7,18 +7,13 @@ from rest_framework_simplejwt.tokens import RefreshToken
from chat_backend.models import (
Announcement,
Conversation,
CustomUser,
Feedback,
OutboundEmail,
UserAuthEvent,
)
from rest_framework_simplejwt.token_blacklist.models import (
BlacklistedToken,
OutstandingToken,
)
from .factories import make_company, make_conversation, make_user
from .factories import make_company, make_user
class AuthenticationRequiredTestCase(APITestCase):
@@ -525,102 +520,3 @@ class CustomUserGetTestCase(APITestCase):
self.assertEqual(response.data["email"], user.email)
self.assertEqual(response.data["company"]["name"], "Globex")
self.assertNotIn("password", response.data)
class CustomUserSelfDeleteTestCase(APITestCase):
def setUp(self):
self.company = make_company("Acme")
self.user = make_user(
email="deleteme@example.com",
password="testpass123",
company=self.company,
)
self.url = reverse("delete_user")
def test_unauthenticated_rejected(self):
response = self.client.delete(self.url)
self.assertEqual(response.status_code, status.HTTP_401_UNAUTHORIZED)
def test_soft_deletes_self_and_hides_conversations(self):
conversation = make_conversation(user=self.user, title="Keep hidden")
refresh = RefreshToken.for_user(self.user)
self.client.force_authenticate(user=self.user)
response = self.client.delete(
self.url, {"refresh_token": str(refresh)}, format="json"
)
self.assertEqual(response.status_code, status.HTTP_200_OK)
self.assertTrue(response.data["deleted"])
self.user.refresh_from_db()
self.assertTrue(self.user.deleted)
self.assertFalse(self.user.is_active)
delete_event = UserAuthEvent.objects.get(
user=self.user,
event_type=UserAuthEvent.EventType.ACCOUNT_DELETED,
)
self.assertIn("soft-delete", delete_event.detail.lower())
self.assertTrue(
CustomUser.objects.filter(pk=self.user.pk, deleted=True).exists()
)
conversation.refresh_from_db()
self.assertTrue(conversation.deleted)
outstanding = OutstandingToken.objects.filter(user=self.user)
self.assertTrue(outstanding.exists())
for token in outstanding:
self.assertTrue(BlacklistedToken.objects.filter(token=token).exists())
# Soft-deleted users cannot obtain new tokens.
login = self.client.post(
reverse("token_create"),
{"username": self.user.username, "password": "testpass123"},
format="json",
)
self.assertEqual(login.status_code, status.HTTP_401_UNAUTHORIZED)
# Conversations list hides soft-deleted rows for any remaining session.
other = make_user(email="alive@example.com", company=self.company)
make_conversation(user=other, title="Still visible")
self.client.force_authenticate(user=other)
listed = self.client.get(reverse("conversations"))
titles = [row["title"] for row in listed.data]
self.assertNotIn("Keep hidden", titles)
def test_cannot_delete_another_user_via_body(self):
"""Endpoint always targets request.user; body email/id is ignored."""
other = make_user(email="other@example.com", company=self.company)
make_conversation(user=other, title="Other chat")
self.client.force_authenticate(user=self.user)
response = self.client.delete(
self.url,
{"email": other.email, "user_id": other.pk},
format="json",
)
self.assertEqual(response.status_code, status.HTTP_200_OK)
self.user.refresh_from_db()
other.refresh_from_db()
self.assertTrue(self.user.deleted)
self.assertFalse(other.deleted)
self.assertFalse(
Conversation.objects.filter(user=other, deleted=True).exists()
)
def test_staff_cannot_self_delete(self):
staff = make_user(
email="staff@example.com",
company=self.company,
is_staff=True,
)
self.client.force_authenticate(user=staff)
response = self.client.delete(self.url, format="json")
self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
self.assertEqual(response.data["code"], "staff_forbidden")
staff.refresh_from_db()
self.assertFalse(staff.deleted)
llm_be/chat_backend/tests/test_ws_frames.py
-50
View File
@@ -1,50 +0,0 @@
"""Unit tests for versioned WS frames and status emission (#96 / #62 / #63)."""
from __future__ import annotations
from django.test import SimpleTestCase
from chat_backend.services.ws_frames import (
STATUS_LABELS,
agent_frame,
citations_frame,
status_frame,
versioned_frame,
)
class WsFramesTestCase(SimpleTestCase):
def test_status_frame_shape(self):
frame = status_frame("searching", detail="Taylor Swift wedding")
self.assertEqual(frame["v"], 1)
self.assertEqual(frame["type"], "status")
self.assertEqual(frame["data"]["stage"], "searching")
self.assertEqual(frame["data"]["label"], STATUS_LABELS["searching"])
self.assertEqual(frame["data"]["detail"], "Taylor Swift wedding")
def test_status_frame_unknown_stage_uses_label_or_stage(self):
frame = status_frame("custom_agent_step", label="Researching competitors")
self.assertEqual(frame["data"]["stage"], "custom_agent_step")
self.assertEqual(frame["data"]["label"], "Researching competitors")
def test_writing_stage_empty_default_label(self):
frame = status_frame("writing")
self.assertEqual(frame["data"]["label"], "")
def test_citations_frame_shape(self):
cites = [{"index": 1, "title": "T", "url": "https://x.test", "published_at": None}]
frame = citations_frame(cites)
self.assertEqual(frame, {"v": 1, "type": "citations", "data": cites})
def test_agent_frame_shape(self):
frame = agent_frame(
"run_started",
{"run_id": "abc", "status": "running", "title": "Research"},
)
self.assertEqual(frame["v"], 1)
self.assertEqual(frame["type"], "run_started")
self.assertEqual(frame["data"]["run_id"], "abc")
def test_versioned_frame_unknown_type_ok(self):
frame = versioned_frame("future_type", {"x": 1})
self.assertEqual(frame["type"], "future_type")
llm_be/chat_backend/urls.py
-67
View File
@@ -5,39 +5,26 @@ from .views import (
CustomObtainTokenView,
CustomUserCreate,
CustomUserInvite,
CustomUserSelfDeleteView,
LogoutAndBlacklistRefreshTokenForUserView,
CustomUserGet,
PublicSettingsView,
is_authenticated,
AnnouncmentView,
FeedbackView,
PromptFeedbackView,
ConversationsView,
ConversationDetailView,
CompanyUsersView,
SetUserPassword,
ResetUserPassword,
ConversationPreferences,
UserPromptAnalytics,
UserConversationAnalytics,
CompanyUsageAnalytics,
AdminAnalytics,
UserPromptHeatmap,
DocumentWorkspaceView,
DocumentUploadView,
DocumentDetailView,
)
from .views_oauth import OAuthCallbackView, OAuthStartView
from .views_drive import (
DriveConnectionListView,
DriveConnectionDetailView,
DriveConnectionResourcesView,
DriveConnectionSyncView,
DriveWebhookGoogleView,
DriveWebhookMicrosoftView,
)
from .views_agent import AgentRunCancelView, AgentRunDetailView, AgentRunListView
from rest_framework.routers import DefaultRouter
@@ -69,7 +56,6 @@ urlpatterns = [
name="blacklist",
),
path("user/get/", CustomUserGet.as_view(), name="get_user"),
path("user/", CustomUserSelfDeleteView.as_view(), name="delete_user"),
path(
"user/acknowledge_tos/",
AcknowledgeTermsOfService.as_view(),
@@ -80,11 +66,6 @@ urlpatterns = [
path("announcment/get/", AnnouncmentView.as_view(), name="get_announcments"),
path("conversations", ConversationsView.as_view(), name="conversations"),
path("feedbacks/", FeedbackView.as_view(), name="feedbacks"),
path(
"prompt_feedback",
PromptFeedbackView.as_view(),
name="prompt_feedback",
),
path(
"conversation_details",
ConversationDetailView.as_view(),
@@ -111,11 +92,6 @@ urlpatterns = [
name="analytics_company_usage",
),
path("analytics/admin/", AdminAnalytics.as_view(), name="analytics_admin"),
path(
"analytics/user_prompt_heatmap/",
UserPromptHeatmap.as_view(),
name="analytics_user_prompt_heatmap",
),
# document urls
path(
"document_workspaces/",
@@ -128,47 +104,4 @@ urlpatterns = [
DocumentDetailView.as_view(),
name="documents_details",
),
# drive urls (#47-#52)
path(
"drive/connections/",
DriveConnectionListView.as_view(),
name="drive_connections",
),
path(
"drive/connections/<int:connection_id>/",
DriveConnectionDetailView.as_view(),
name="drive_connection_detail",
),
path(
"drive/connections/<int:connection_id>/resources/",
DriveConnectionResourcesView.as_view(),
name="drive_connection_resources",
),
path(
"drive/connections/<int:connection_id>/sync/",
DriveConnectionSyncView.as_view(),
name="drive_connection_sync",
),
path(
"drive/webhooks/google/",
DriveWebhookGoogleView.as_view(),
name="drive_webhook_google",
),
path(
"drive/webhooks/microsoft/",
DriveWebhookMicrosoftView.as_view(),
name="drive_webhook_microsoft",
),
# Agent runs (#63)
path("agent_runs/", AgentRunListView.as_view(), name="agent_runs"),
path(
"agent_runs/<int:run_id>/",
AgentRunDetailView.as_view(),
name="agent_run_detail",
),
path(
"agent_runs/<int:run_id>/cancel/",
AgentRunCancelView.as_view(),
name="agent_run_cancel",
),
]
llm_be/chat_backend/utils.py
+2 -73
View File
@@ -1,26 +1,6 @@
import datetime
def is_heartbeat_payload(data) -> bool:
"""True for app-level WS keepalive frames (see FE buildHeartbeatPayload)."""
return isinstance(data, dict) and data.get("type") == "ping"
def normalize_user_message(message):
"""Return stripped message text, or None if missing/blank."""
if message is None:
return None
if not isinstance(message, str):
message = str(message)
stripped = message.strip()
return stripped or None
def has_usable_user_prompt(message, file=None) -> bool:
"""Reject empty/whitespace chat text. ``file`` kept for call-site clarity."""
return normalize_user_message(message) is not None
def last_day_of_month(any_day):
# The day 28 exists in every month. 4 days later, it's always next month
next_month = any_day.replace(day=28) + datetime.timedelta(days=4)
@@ -50,8 +30,7 @@ def _as_usage_mapping(source):
"""Best-effort pull of a usage dict out of a provider response.
Accepts a raw dict, a LangChain message (``usage_metadata`` /
``response_metadata``), an Ollama ``GenerationChunk`` (``generation_info``),
or any object exposing those attributes.
``response_metadata``), or any object exposing those attributes.
"""
if source is None:
return None
@@ -61,7 +40,7 @@ def _as_usage_mapping(source):
if isinstance(nested, dict):
return nested
return source
for attr in ("usage_metadata", "response_metadata", "generation_info"):
for attr in ("usage_metadata", "response_metadata"):
nested = getattr(source, attr, None)
if isinstance(nested, dict):
mapping = _as_usage_mapping(nested)
@@ -80,53 +59,3 @@ def extract_token_usage(source):
if not mapping:
return None, None
return _first_int(mapping, _TOKENS_IN_KEYS), _first_int(mapping, _TOKENS_OUT_KEYS)
def chunk_text(chunk) -> str:
"""Pull display text out of a stream chunk (str, GenerationChunk, message)."""
if chunk is None:
return ""
if isinstance(chunk, str):
return chunk
text = getattr(chunk, "text", None)
if isinstance(text, str) and text:
return text
content = getattr(chunk, "content", None)
if isinstance(content, str):
return content
return ""
class TokenUsageCollector:
"""Accumulate provider-reported token counts while streaming LLM chunks."""
def __init__(self):
self.tokens_in = None
self.tokens_out = None
def observe(self, source) -> None:
tin, tout = extract_token_usage(source)
if tin is not None:
self.tokens_in = tin
if tout is not None:
self.tokens_out = tout
@property
def pair(self):
return self.tokens_in, self.tokens_out
async def aiter_text_chunks(stream, collector: TokenUsageCollector | None = None):
"""Yield text from a provider stream, optionally capturing token usage.
Ollama reports ``prompt_eval_count`` / ``eval_count`` on the final
``GenerationChunk.generation_info`` when ``done`` is true. ``StrOutputParser``
strips that metadata, so callers must stream the raw LLM chain and use this
helper (or equivalent) to persist real usage.
"""
async for chunk in stream:
if collector is not None:
collector.observe(chunk)
text = chunk_text(chunk)
if text:
yield text
llm_be/chat_backend/views.py
+25 -469
View File
@@ -13,8 +13,6 @@ from .serializers import (
ConversationSerializer,
PromptSerializer,
FeedbackSerializer,
PromptFeedbackSerializer,
PromptFeedbackUpsertSerializer,
DocumentWorkspaceSerializer,
DocumentSerializer,
)
@@ -26,7 +24,6 @@ from .models import (
Announcement,
Conversation,
Prompt,
PromptFeedback,
Feedback,
PromptMetric,
DocumentWorkspace,
@@ -54,7 +51,6 @@ from chat_backend.services.assistant_identity import ASSISTANT_SYSTEM_PROMPT
from django.utils import timezone
from django.core.files import File
from django.core.files.base import ContentFile
from django.db.models import Q
import math
import datetime
import pytz
@@ -69,14 +65,8 @@ from .email_tasks import (
send_invite_email,
send_password_reset_email,
)
from monetization.services.quotas import FeatureNotAllowed, assert_feature_allowed
from .services.llm_service import AsyncLLMService
from .services.rag_services import AsyncRAGService
from .services.chat_tenant_scope import (
ensure_company_workspace,
ensure_personal_workspace,
ensure_workspace_for_user,
)
from .services.title_generator import title_generator
from .services.moderation_classifier import moderation_classifier, ModerationLabel
from .services.prompt_classifier.prompt_classifier import PromptClassifier, PromptType
@@ -146,8 +136,6 @@ class CustomUserCreate(APIView):
user = serializer.save()
refresh = RefreshToken.for_user(user)
from monetization.services.plans import needs_checkout
return Response(
{
"email": user.email,
@@ -155,7 +143,6 @@ class CustomUserCreate(APIView):
"last_name": user.last_name,
"access": str(refresh.access_token),
"refresh": str(refresh),
"needs_checkout": needs_checkout(user),
},
status=status.HTTP_201_CREATED,
)
@@ -315,48 +302,6 @@ class CustomUserGet(APIView):
return Response({}, status=status.HTTP_400_BAD_REQUEST)
class CustomUserSelfDeleteView(APIView):
"""
Soft-delete the authenticated user's own account (#34).
Frontend contract:
- Method/path: ``DELETE /api/user/``
- Optional body: ``{"refresh_token": "<current refresh>"}`` to blacklist
the active session immediately (outstanding tokens are also blacklisted).
- Success: ``200`` with ``{"detail": "Account deleted.", "deleted": true}``
- After success: clear local tokens, redirect to sign-in. Subsequent
``/token/obtain/`` and authenticated calls fail (``is_active=False``,
``deleted=True``).
- Privacy v1: soft-delete only (no anonymization / hard purge).
"""
http_method_names = ["delete", "head", "options"]
def delete(self, request, format="json"):
from chat_backend.services.account_deletion import (
AccountDeletionError,
soft_delete_account,
)
refresh_token = request.data.get("refresh_token")
try:
soft_delete_account(
request.user,
refresh_token=refresh_token,
ip_address=_client_ip(request),
)
except AccountDeletionError as exc:
return Response(
{"detail": exc.detail, "code": exc.code},
status=status.HTTP_400_BAD_REQUEST,
)
return Response(
{"detail": "Account deleted.", "deleted": True},
status=status.HTTP_200_OK,
)
class FeedbackView(APIView):
http_method_names = ["post", "get"]
@@ -383,90 +328,6 @@ class FeedbackView(APIView):
return Response(serializer.data, status=status.HTTP_200_OK)
def _user_can_rate_prompt(user, prompt: Prompt) -> bool:
"""Caller may rate prompts in their own non-deleted conversations."""
conversation = prompt.conversation
return (
conversation is not None
and conversation.user_id == user.id
and not conversation.deleted
)
class PromptFeedbackView(APIView):
"""Upsert / clear per-message thumbs ratings (chat_backend#67)."""
http_method_names = ["post", "delete"]
def post(self, request, format="json"):
serializer = PromptFeedbackUpsertSerializer(data=request.data)
if not serializer.is_valid():
return Response(serializer.errors, status=status.HTTP_400_BAD_REQUEST)
prompt_id = serializer.validated_data["prompt_id"]
try:
prompt = Prompt.objects.select_related("conversation").get(id=prompt_id)
except Prompt.DoesNotExist:
return Response(
{"detail": "Prompt not found."},
status=status.HTTP_404_NOT_FOUND,
)
if not _user_can_rate_prompt(request.user, prompt):
return Response(
{"detail": "Prompt not found."},
status=status.HTTP_404_NOT_FOUND,
)
if prompt.user_created:
return Response(
{"detail": "Only assistant prompts can be rated."},
status=status.HTTP_400_BAD_REQUEST,
)
feedback, _created = PromptFeedback.objects.update_or_create(
prompt=prompt,
user=request.user,
defaults={
"rating": serializer.validated_data["rating"],
"reason": serializer.validated_data.get("reason"),
"comment": serializer.validated_data.get("comment"),
},
)
return Response(
PromptFeedbackSerializer(feedback).data,
status=status.HTTP_200_OK,
)
def delete(self, request, format="json"):
prompt_id = request.query_params.get("prompt_id")
if prompt_id is None:
return Response(
{"detail": "prompt_id is required."},
status=status.HTTP_400_BAD_REQUEST,
)
try:
prompt_id = int(prompt_id)
except (TypeError, ValueError):
return Response(
{"detail": "prompt_id must be an integer."},
status=status.HTTP_400_BAD_REQUEST,
)
deleted, _ = PromptFeedback.objects.filter(
prompt_id=prompt_id,
user=request.user,
prompt__conversation__user=request.user,
prompt__conversation__deleted=False,
).delete()
if not deleted:
return Response(
{"detail": "Prompt feedback not found."},
status=status.HTTP_404_NOT_FOUND,
)
return Response(status=status.HTTP_204_NO_CONTENT)
class AcknowledgeTermsOfService(APIView):
http_method_names = ["post"]
@@ -584,63 +445,21 @@ class ConversationsView(APIView):
class ConversationPreferences(APIView):
def get(self, request, format="json"):
user = request.user
return Response(
{
"order": user.conversation_order,
"use_conversation_context": user.use_conversation_context,
},
status=status.HTTP_200_OK,
)
return Response({"order": user.conversation_order}, status=status.HTTP_200_OK)
def post(self, request, format="json"):
user = request.user
data = request.data
update_fields = []
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,
)
user.save()
return Response({"order": user.conversation_order}, status=status.HTTP_200_OK)
class ConversationDetailView(APIView):
def get(self, request, format="json"):
conversation_id = request.query_params.get("conversation_id")
if not Conversation.objects.filter(
id=conversation_id, user=request.user, deleted=False
).exists():
return Response(
{"detail": "Conversation not found."},
status=status.HTTP_404_NOT_FOUND,
)
prompts = list(
Prompt.objects.filter(
conversation__id=conversation_id, conversation__user=request.user
)
)
serializer = PromptSerializer(
prompts,
many=True,
context={
"request": request,
"_prompt_ids_for_feedback": [p.id for p in prompts],
},
)
return Response(serializer.data, status=status.HTTP_200_OK)
prompts = Prompt.objects.filter(conversation__id=conversation_id)
serailzer = PromptSerializer(prompts, many=True)
return Response(serailzer.data, status=status.HTTP_200_OK)
def post(self, request, format="json"):
logger.info("In the post")
@@ -650,20 +469,12 @@ class ConversationDetailView(APIView):
# make sure that our model exists and it is running
prompt = request.data.get("prompt")
if not isinstance(prompt, str) or not prompt.strip():
return Response(
{"detail": "Message text cannot be empty."},
status=status.HTTP_400_BAD_REQUEST,
)
prompt = prompt.strip()
conversation_id = request.data.get("conversation_id")
is_user = bool(request.data.get("is_user"))
try:
conversation = Conversation.objects.get(
id=conversation_id, user=request.user, deleted=False
)
conversation = Conversation.objects.get(id=conversation_id)
# add the prompt to the conversation
serializer = PromptSerializer(
@@ -855,76 +666,6 @@ class CompanyUsageAnalytics(APIView):
return Response(result[::-1], status=status.HTTP_200_OK)
return Response(result[::-1], status=status.HTTP_200_OK)
class UserPromptHeatmap(APIView):
"""Weekday × hour bins of user-entered prompts (local timezone) (#94)."""
DAY_LABELS = ("Mon", "Tue", "Wed", "Thu", "Fri", "Sat", "Sun")
def get(self, request, format="json"):
tz_name = request.query_params.get("tz", "UTC")
try:
user_tz = pytz.timezone(tz_name)
except pytz.UnknownTimeZoneError:
user_tz = pytz.UTC
tz_name = "UTC"
matrix = [[0 for _ in range(24)] for _ in range(7)]
total = 0
prompts = Prompt.objects.filter(
conversation__user=request.user,
conversation__deleted=False,
user_created=True,
).only("created")
for prompt in prompts.iterator():
created = prompt.created
if timezone.is_naive(created):
created = timezone.make_aware(created, datetime.timezone.utc)
local = created.astimezone(user_tz)
matrix[local.weekday()][local.hour] += 1
total += 1
max_count = max((count for row in matrix for count in row), default=0)
most_active_day = None
most_active_hour = None
peak_cell = None
if max_count > 0:
day_totals = [sum(row) for row in matrix]
most_active_day = self.DAY_LABELS[day_totals.index(max(day_totals))]
hour_totals = [sum(matrix[day][hour] for day in range(7)) for hour in range(24)]
most_active_hour = hour_totals.index(max(hour_totals))
peak_day = 0
peak_hour = 0
for day_idx, row in enumerate(matrix):
for hour_idx, count in enumerate(row):
if count > matrix[peak_day][peak_hour]:
peak_day = day_idx
peak_hour = hour_idx
peak_cell = {
"day": self.DAY_LABELS[peak_day],
"hour": peak_hour,
"count": matrix[peak_day][peak_hour],
}
return Response(
{
"tz": tz_name,
"total": total,
"max": max_count,
"days": list(self.DAY_LABELS),
"hours": list(range(24)),
"matrix": matrix,
"most_active_day": most_active_day,
"most_active_hour": most_active_hour,
"peak_cell": peak_cell,
},
status=status.HTTP_200_OK,
)
class AdminAnalytics(APIView):
def get(self, request, format="json"):
number_of_months = 3
@@ -986,43 +727,18 @@ llm = OllamaLLM(**ollama_llm_kwargs(model=MODEL_NAME))
# Document Views
def _feature_gate_response(exc: FeatureNotAllowed) -> Response:
return Response(
{"code": exc.code, "error": exc.message, "details": exc.details},
status=status.HTTP_403_FORBIDDEN,
)
class DocumentWorkspaceView(APIView):
# permission_classes = [permissions.IsAuthenticated]
def get(self, request):
try:
assert_feature_allowed(request.user, "rag")
except FeatureNotAllowed as exc:
return _feature_gate_response(exc)
if request.user.company_id:
workspaces = DocumentWorkspace.objects.filter(
company=request.user.company, user__isnull=True
)
else:
workspaces = DocumentWorkspace.objects.filter(
user=request.user, company__isnull=True
)
workspaces = DocumentWorkspace.objects.filter(company=request.user.company)
serializer = DocumentWorkspaceSerializer(workspaces, many=True)
return Response(serializer.data)
def post(self, request):
try:
assert_feature_allowed(request.user, "rag")
except FeatureNotAllowed as exc:
return _feature_gate_response(exc)
serializer = DocumentWorkspaceSerializer(data=request.data)
if serializer.is_valid():
if request.user.company_id:
serializer.save(company=request.user.company, user=None)
else:
serializer.save(company=None, user=request.user)
serializer.save(company=request.user.company)
return Response(serializer.data, status=status.HTTP_201_CREATED)
return Response(serializer.errors, status=status.HTTP_400_BAD_REQUEST)
@@ -1030,109 +746,30 @@ class DocumentWorkspaceView(APIView):
class DocumentUploadView(APIView):
# permission_classes = [permissions.IsAuthenticated]Z
DOCUMENT_SORT_FIELDS = {
"name": "file",
"-name": "-file",
"created": "created",
"-created": "-created",
"date_uploaded": "created",
"-date_uploaded": "-created",
"processed": "processed",
"-processed": "-processed",
"active": "active",
"-active": "-active",
}
def _resolve_list_workspace(self, request):
"""Return (workspace, scope) for personal|company document lists (#60)."""
raw = (request.query_params.get("workspace") or "").strip().lower()
if not raw:
raw = "company" if request.user.company_id else "personal"
if raw == "personal":
return ensure_personal_workspace(request.user), "personal"
if raw == "company":
if not request.user.company_id:
return None, "company"
return ensure_company_workspace(request.user.company), "company"
return False, raw
def get(self, request):
logger.debug(f"request_3: {request}")
try:
assert_feature_allowed(request.user, "rag")
except FeatureNotAllowed as exc:
return _feature_gate_response(exc)
workspace, scope = self._resolve_list_workspace(request)
if workspace is False:
return Response(
{"error": "workspace must be 'personal' or 'company'."},
status=status.HTTP_400_BAD_REQUEST,
workspace = DocumentWorkspace.objects.get(company=request.user.company)
serializer = DocumentSerializer(
Document.objects.filter(workspace=workspace), many=True
)
if workspace is None:
return Response(serializer.data, status=status.HTTP_200_OK)
except:
return Response(
{"error": "Company workspace requires a company membership."},
status=status.HTTP_400_BAD_REQUEST,
)
queryset = Document.objects.filter(workspace=workspace)
search = (request.query_params.get("search") or request.query_params.get("q") or "").strip()
if search:
queryset = queryset.filter(
Q(file__icontains=search) | Q(remote_name__icontains=search)
)
ordering = (request.query_params.get("ordering") or "-created").strip()
order_by = self.DOCUMENT_SORT_FIELDS.get(ordering, "-created")
queryset = queryset.order_by(order_by, "id")
try:
page = max(int(request.query_params.get("page", 1)), 1)
except (TypeError, ValueError):
page = 1
try:
page_size = int(request.query_params.get("page_size", 20))
except (TypeError, ValueError):
page_size = 20
page_size = min(max(page_size, 1), 100)
total = queryset.count()
start = (page - 1) * page_size
end = start + page_size
serializer = DocumentSerializer(queryset[start:end], many=True)
return Response(
{
"count": total,
"page": page,
"page_size": page_size,
"scope": scope,
"results": serializer.data,
},
status=status.HTTP_200_OK,
{"error": "Workspace not found"}, status=status.HTTP_404_NOT_FOUND
)
def post(self, request):
logger.debug(f"request: {request}")
try:
assert_feature_allowed(request.user, "rag")
except FeatureNotAllowed as exc:
return _feature_gate_response(exc)
workspace = DocumentWorkspace.objects.get(company=request.user.company)
scope = (request.query_params.get("workspace") or request.data.get("workspace") or "").strip().lower()
if scope == "personal":
workspace = ensure_personal_workspace(request.user)
elif scope == "company":
if not request.user.company_id:
except:
return Response(
{"error": "Company workspace requires a company membership."},
status=status.HTTP_400_BAD_REQUEST,
{"error": "Workspace not found"}, status=status.HTTP_404_NOT_FOUND
)
workspace = ensure_company_workspace(request.user.company)
else:
workspace = ensure_workspace_for_user(request.user)
logger.info(request.FILES)
file = request.FILES.get("file")
@@ -1158,15 +795,7 @@ class DocumentUploadView(APIView):
document.save()
service = AsyncRAGService()
service.add_files_to_store(
[
(
document.file,
document.file.name,
document.workspace_id,
document.id,
document.active,
)
],
[(document.file, document.file.name, document.workspace_id)],
workspace_id=document.workspace_id,
)
@@ -1174,89 +803,16 @@ class DocumentUploadView(APIView):
class DocumentDetailView(APIView):
# permission_classes = [permissions.IsAuthenticated]
def _get_document(self, request, document_id):
user = request.user
if user.company_id:
return (
Document.objects.filter(id=document_id)
.filter(
Q(workspace__user=user, workspace__company__isnull=True)
| Q(
workspace__company_id=user.company_id,
workspace__user__isnull=True,
)
)
.first()
)
return Document.objects.filter(
id=document_id,
workspace__user=user,
workspace__company__isnull=True,
).first()
def get(self, request, document_id):
logger.info(f"request: {request}")
try:
assert_feature_allowed(request.user, "rag")
except FeatureNotAllowed as exc:
return _feature_gate_response(exc)
workspace = DocumentWorkspace.objects.get(company=request.user.company)
document = self._get_document(request, document_id)
if document is None:
document = Document.objects.get(workspace=workspace, id=document_id)
except:
return Response(
{"error": "Document not found"}, status=status.HTTP_404_NOT_FOUND
)
serializer = DocumentSerializer(document)
return Response(serializer.data, status=status.HTTP_200_OK)
def patch(self, request, document_id):
"""Toggle a document's ``active`` flag (#44) and its vector metadata (#45)."""
try:
assert_feature_allowed(request.user, "rag")
except FeatureNotAllowed as exc:
return _feature_gate_response(exc)
document = self._get_document(request, document_id)
if document is None:
return Response(
{"error": "Document not found"}, status=status.HTTP_404_NOT_FOUND
)
if "active" not in request.data:
return Response(
{"error": "active is required"}, status=status.HTTP_400_BAD_REQUEST
)
active = request.data.get("active")
if isinstance(active, str):
active = active.strip().lower() in {"1", "true", "yes"}
document.active = bool(active)
document.save(update_fields=["active"])
try:
AsyncRAGService().set_document_active(document.id, document.active)
except Exception as exc:
logger.warning(
"Failed to update vector active metadata for document %s: %s",
document.id,
exc,
)
serializer = DocumentSerializer(document)
return Response(serializer.data, status=status.HTTP_200_OK)
def delete(self, request, document_id):
try:
assert_feature_allowed(request.user, "rag")
except FeatureNotAllowed as exc:
return _feature_gate_response(exc)
document = self._get_document(request, document_id)
if document is None:
return Response(
{"error": "Document not found"}, status=status.HTTP_404_NOT_FOUND
)
document.delete()
return Response(status=status.HTTP_204_NO_CONTENT)
serializer = DocumentWorkspaceSerializer(workspaces, many=True)
return Response(serializer.data)

Some files were not shown because too many files have changed in this diff Show More