diff --git a/.env.example b/.env.example index 316a294..02556bb 100644 --- a/.env.example +++ b/.env.example @@ -30,6 +30,16 @@ LOG_EACH_QUERY=false RESERVATION_TTL_SECONDS=600 EMBEDDING_MODEL=text-embedding-3-small +# ===== Ingestion Limits ===== +MAX_FILE_SIZE_BYTES=10485760 +MAX_PDF_PAGE_COUNT=10 +MIN_EXTRACTED_TEXT_CHARS=1 +MAX_BULK_UPLOAD_FILES=50 +EMBEDDING_BATCH_SIZE=32 +OPENAI_EMBEDDING_TIMEOUT_SECONDS=300 +INGEST_EXTRACT_JOB_TIMEOUT_SECONDS=900 +INGEST_INDEX_JOB_TIMEOUT_SECONDS=1800 + # ===== Client (React) ===== VITE_API_URL=http://localhost:8000 VITE_SUPABASE_URL=https://your-project.supabase.co diff --git a/.github/workflows/deploy.yml b/.github/workflows/deploy.yml index 18e5cd9..6df6ce3 100644 --- a/.github/workflows/deploy.yml +++ b/.github/workflows/deploy.yml @@ -7,7 +7,7 @@ on: description: 'Pull request number' required: true push: - branches: [ main ] + branches: [ main, develop ] pull_request: types: [opened, synchronize, reopened, ready_for_review] @@ -37,10 +37,11 @@ jobs: run: mkdir -p trivy-reports - name: Run Trivy FS Scan - uses: aquasecurity/trivy-action@0.24.0 + uses: aquasecurity/trivy-action@v0.36.0 with: scan-type: 'fs' scan-ref: '.' + version: 'v0.70.0' scanners: 'vuln,misconfig,secret,license' ignore-unfixed: true format: 'table' @@ -95,4 +96,4 @@ jobs: with: name: bandit-report path: bandit-report.html - retention-days: 30 \ No newline at end of file + retention-days: 30 diff --git a/README.md b/README.md index a24dd45..7d2a0b6 100644 --- a/README.md +++ b/README.md @@ -174,10 +174,11 @@ sequenceDiagram What happens in practice: - `upload-prepare` validates file size, content type, workspace limits, and idempotency. +- `upload-prepare-batch` can prepare many small PDFs as one ingestion run and returns per-file results. - The API stores a placeholder document record and returns a signed storage URL. -- `upload-complete` confirms the object exists in storage and enqueues extraction. +- `upload-complete` and `upload-complete-batch` confirm uploaded objects and enqueue extraction jobs. - `ingest_extract` downloads the PDF and writes extracted page text into `document_pages`. -- `ingest_index` chunks page text, generates embeddings, stores vectors, and marks the document ready. +- `ingest_index` chunks page text, generates embeddings in batches, stores vectors, and marks the document ready. Primary files: - `server/app/api/documents.py` @@ -284,7 +285,11 @@ enterprise-rag/ - `GET /documents/{document_id}` - `GET /documents/{document_id}/pages/{page_number}` - `POST /documents/upload-prepare` +- `POST /documents/upload-prepare-batch` - `POST /documents/upload-complete` +- `POST /documents/upload-complete-batch` +- `GET /documents/ingestion-runs/{run_id}` +- `GET /documents/ingestion-queues` - `POST /documents/{document_id}/retry` - `POST /documents/{document_id}/reindex` - `DELETE /documents/{document_id}` @@ -314,6 +319,7 @@ enterprise-rag/ Core tables in the current implementation: - `workspaces`: tenant root for all user content +- `ingestion_runs`: batch upload and ingestion progress grouping - `documents`: uploaded PDF metadata and pipeline status - `document_pages`: extracted page text - `chunks`: page-bounded text chunks used for retrieval @@ -337,7 +343,9 @@ Current enforced limits from the application config and rate limiter: - `1` workspace per user - up to `100` documents per workspace -- maximum file size: `20 MB` +- maximum file size: `10 MB` +- maximum PDF page count: `10` +- maximum files per backend batch upload: `50` - supported upload type: `application/pdf` - maximum query length: `500` characters - retrieval depth: `top_k = 5` @@ -347,6 +355,10 @@ Current enforced limits from the application config and rate limiter: - upload complete rate limit: `20` requests per minute per workspace - query rate limit: `100` requests per minute per workspace +Bulk ingestion endpoints have separate rate limits from single-file endpoints: +- batch upload prepare: `5` requests per minute per workspace +- batch upload complete: `10` requests per minute per workspace + ### Token Budget Model The token budget is managed with reservation semantics so concurrent requests do not overspend the daily allowance. @@ -487,6 +499,14 @@ DAILY_TOKEN_LIMIT=100000 RESERVATION_TTL_SECONDS=600 LOG_EACH_QUERY=false EMBEDDING_MODEL=text-embedding-3-small +MAX_FILE_SIZE_BYTES=10485760 +MAX_PDF_PAGE_COUNT=10 +MIN_EXTRACTED_TEXT_CHARS=1 +MAX_BULK_UPLOAD_FILES=50 +EMBEDDING_BATCH_SIZE=32 +OPENAI_EMBEDDING_TIMEOUT_SECONDS=300 +INGEST_EXTRACT_JOB_TIMEOUT_SECONDS=900 +INGEST_INDEX_JOB_TIMEOUT_SECONDS=1800 VITE_API_URL=http://localhost:8000 VITE_SUPABASE_URL=https://your-project.supabase.co VITE_SUPABASE_ANON_KEY=your-anon-key @@ -511,6 +531,7 @@ Embeddings are stored in `chunk_embeddings.embedding` using `pgvector`. Retrieva - failed documents can be retried with `POST /documents/{document_id}/retry` - already processed documents can be reindexed with `POST /documents/{document_id}/reindex` - stale token reservations can be cleared by the maintenance job +- RQ extraction/indexing jobs use explicit timeouts and a failure callback so timed-out jobs mark documents `failed` instead of leaving them stuck in processing - document deletion removes metadata first, then attempts storage cleanup ### Frontend Application Areas diff --git a/client/README.md b/client/README.md index edb8e19..b221f8b 100644 --- a/client/README.md +++ b/client/README.md @@ -90,14 +90,19 @@ Primary files: Upload UX is centered in `UploadPage` and `components/upload/UploadPanel`. Flow: -1. request `POST /documents/upload-prepare` -2. upload the PDF directly to the signed URL -3. notify the backend with `POST /documents/upload-complete` -4. poll the document list while processing is active -5. let the user delete documents from the table +1. select and review up to 50 PDFs as one upload run +2. request `POST /documents/upload-prepare-batch` +3. upload accepted PDFs directly to their signed URLs +4. notify the backend with `POST /documents/upload-complete-batch` +5. poll documents, queue status, and the active ingestion run while processing is active +6. separate active, failed, and ready documents for recovery and review Current behavior: -- shows live document status in a table +- shows selected files before upload starts +- shows accepted/rejected/processing/ready/failed counts for the current run +- shows live document status grouped by workflow state +- supports search, status filtering, and sorting for larger document sets +- supports retrying failed documents individually or in a retryable batch - refreshes every 4 seconds while documents are processing - treats `ready` and `indexed` as successful ingest states @@ -145,7 +150,12 @@ The client currently uses these backend APIs: - `GET /documents/{document_id}/pages/{page_number}` - `POST /documents/upload-prepare` - `POST /documents/upload-complete` +- `POST /documents/upload-prepare-batch` +- `POST /documents/upload-complete-batch` +- `GET /documents/ingestion-runs/{run_id}` +- `GET /documents/ingestion-queues` - `DELETE /documents/{document_id}` +- `POST /documents/{document_id}/retry` - `POST /query/stream` - `GET /citations/{chunk_id}` - `GET /queries` diff --git a/client/src/components/layout/AppShell.tsx b/client/src/components/layout/AppShell.tsx index 575282d..f6a5d7d 100644 --- a/client/src/components/layout/AppShell.tsx +++ b/client/src/components/layout/AppShell.tsx @@ -69,7 +69,7 @@ export default function AppShell() { } try { - const nextDocuments = await apiGetDocuments(accessToken); + const nextDocuments = await apiGetDocuments(accessToken, { limit: 100 }); setDocuments(nextDocuments); } catch { setDocuments([]); diff --git a/client/src/components/upload/UploadPanel.tsx b/client/src/components/upload/UploadPanel.tsx index a238d48..85a02f8 100644 --- a/client/src/components/upload/UploadPanel.tsx +++ b/client/src/components/upload/UploadPanel.tsx @@ -1,27 +1,39 @@ -import { File, UploadCloud, X } from "lucide-react"; +import { + AlertTriangle, + CheckCircle2, + File, + Loader2, + Play, + UploadCloud, + X, +} from "lucide-react"; import { useCallback, useEffect, useMemo, useRef, useState } from "react"; import type { DragEventHandler } from "react"; import { - apiUploadComplete, - apiUploadPrepare, + apiUploadCompleteBatch, + apiUploadPrepareBatch, apiUploadToSignedUrl, type DocumentRecord, type DocumentStatus, + type UploadCompleteBatchFile, } from "../../lib/api"; import { cn } from "../../lib/utils"; -const MAX_FILES = 100; +const MAX_FILES = 50; +const MAX_FILE_SIZE_BYTES = 10 * 1024 * 1024; const CONCURRENCY = 4; type UploadTaskState = - | "queued" - | "preparing" + | "selected" + | "rejected" + | "accepted" | "uploading" | "completing" + | "queued" | "extracting" | "indexing" - | "indexed" + | "ready" | "failed"; type UploadTask = { @@ -31,166 +43,350 @@ type UploadTask = { progress: number; message?: string; documentId?: string; + bucket?: string; + storagePath?: string; + uploadUrl?: string; }; type UploadPanelProps = { token: string; onAfterUpload: () => Promise; documents: DocumentRecord[]; + onRunCreated?: (runId: string) => void; }; +type PreparedTask = { + taskId: string; + documentId: string; + bucket: string; + storagePath: string; + uploadUrl: string; +}; + +function clientFileId(file: File, index: number): string { + const safeName = file.name.replace(/[^A-Za-z0-9_-]/g, "_").slice(0, 36); + const randomPart = Math.random().toString(36).slice(2, 10); + return `${Date.now()}-${index}-${safeName}-${randomPart}`; +} + +function fileSizeLabel(bytes: number): string { + if (bytes < 1024 * 1024) { + return `${Math.max(bytes / 1024, 1).toFixed(0)} KB`; + } + return `${(bytes / (1024 * 1024)).toFixed(1)} MB`; +} + +function validateFile(file: File): string | null { + const isPdf = file.type === "application/pdf" || file.name.toLowerCase().endsWith(".pdf"); + if (!isPdf) { + return "Only PDF files are supported."; + } + + if (file.size > MAX_FILE_SIZE_BYTES) { + return `PDF must be ${fileSizeLabel(MAX_FILE_SIZE_BYTES)} or smaller.`; + } + + return null; +} + function statusLabel(status: UploadTaskState): string { - if (status === "preparing") return "Preparing"; + if (status === "selected") return "Ready to upload"; + if (status === "rejected") return "Rejected"; + if (status === "accepted") return "Prepared"; if (status === "uploading") return "Uploading"; - if (status === "completing") return "Completing"; - if (status === "extracting") return "Extracting"; + if (status === "completing") return "Finalizing upload"; + if (status === "extracting") return "Extracting text"; if (status === "indexing") return "Indexing"; - if (status === "indexed") return "Indexed"; + if (status === "ready") return "Ready"; if (status === "failed") return "Failed"; return "Queued"; } -function mapDocumentStatus(status: string): UploadTaskState { - const docStatus = status.toLowerCase() as DocumentStatus; +function mapDocumentStatus(status: DocumentStatus): UploadTaskState { + if (status === "indexed" || status === "ready") return "ready"; + if (status === "extracting") return "extracting"; + if (status === "indexing") return "indexing"; + if (status === "failed") return "failed"; + if (status === "uploaded" || status === "queued") return "queued"; + return "accepted"; +} + +function progressForState(state: UploadTaskState): number { + if (state === "selected" || state === "rejected" || state === "failed") return 0; + if (state === "accepted") return 15; + if (state === "uploading") return 45; + if (state === "completing") return 75; + if (state === "queued") return 82; + if (state === "extracting") return 90; + if (state === "indexing") return 96; + if (state === "ready") return 100; + return 0; +} + +function taskTone(state: UploadTaskState): string { + if (state === "ready") { + return "border-green-200 bg-green-50"; + } + if (state === "failed" || state === "rejected") { + return "border-red-200 bg-red-50"; + } + if (state === "selected") { + return "border-app-border bg-app-surface"; + } + return "border-amber-200 bg-amber-50"; +} + +async function runWithConcurrency( + items: T[], + limit: number, + worker: (item: T) => Promise, +): Promise { + let index = 0; + const runners = Array.from({ length: Math.min(limit, items.length) }, async () => { + while (index < items.length) { + const current = items[index]; + index += 1; + await worker(current); + } + }); - if (docStatus === "indexed" || docStatus === "ready") return "indexed"; - if (docStatus === "extracting") return "extracting"; - if (docStatus === "indexing") return "indexing"; - if (docStatus === "uploaded") return "indexing"; - if (docStatus === "failed") return "failed"; - return "queued"; + await Promise.all(runners); } -export default function UploadPanel({ token, onAfterUpload, documents }: UploadPanelProps) { +export default function UploadPanel({ + token, + onAfterUpload, + documents, + onRunCreated, +}: UploadPanelProps) { const [tasks, setTasks] = useState([]); + const [runId, setRunId] = useState(null); + const [runState, setRunState] = useState<"idle" | "preparing" | "uploading" | "processing">("idle"); const [error, setError] = useState(null); const inputRef = useRef(null); - const startedTaskIdsRef = useRef>(new Set()); - const activeCount = useMemo( - () => tasks.filter((task) => ["preparing", "uploading", "completing"].includes(task.state)).length, + const selectedCount = useMemo(() => tasks.filter((task) => task.state === "selected").length, [tasks]); + const rejectedCount = useMemo(() => tasks.filter((task) => task.state === "rejected").length, [tasks]); + const activeUploadCount = useMemo( + () => tasks.filter((task) => ["accepted", "uploading", "completing", "queued", "extracting", "indexing"].includes(task.state)).length, [tasks], ); - - const queuedTasks = useMemo(() => tasks.filter((task) => task.state === "queued"), [tasks]); + const readyCount = useMemo(() => tasks.filter((task) => task.state === "ready").length, [tasks]); + const failedCount = useMemo(() => tasks.filter((task) => task.state === "failed").length, [tasks]); + const isRunning = runState === "preparing" || runState === "uploading"; const updateTask = useCallback((id: string, patch: Partial) => { setTasks((current) => current.map((task) => (task.id === id ? { ...task, ...patch } : task))); }, []); - const runTask = useCallback(async (task: UploadTask) => { - try { - const contentType = task.file.type || "application/pdf"; - - const prepared = await apiUploadPrepare(token, { - filename: task.file.name, - content_type: contentType, - file_size_bytes: task.file.size, - idempotency_key: task.id, - }); - - updateTask(task.id, { state: "uploading", progress: 35, documentId: prepared.document_id }); - await apiUploadToSignedUrl(prepared.upload_url, task.file); + useEffect(() => { + if (documents.length === 0) { + return; + } - updateTask(task.id, { state: "completing", progress: 80 }); - const completed = await apiUploadComplete(token, { - document_id: prepared.document_id, - bucket: prepared.bucket, - storage_path: prepared.storage_path, - }); + const documentById = new Map(documents.map((doc) => [doc.id, doc])); + setTasks((current) => + current.map((task) => { + if (!task.documentId) { + return task; + } - updateTask(task.id, { - state: mapDocumentStatus(completed.status), - progress: completed.status === "failed" ? 0 : 100, - }); + const document = documentById.get(task.documentId); + if (!document) { + return task; + } - await onAfterUpload(); - } catch (err) { - updateTask(task.id, { - state: "failed", - progress: 0, - message: err instanceof Error ? err.message : "Upload failed", - }); - } finally { - startedTaskIdsRef.current.delete(task.id); - } - }, [onAfterUpload, token, updateTask]); + const nextState = mapDocumentStatus(document.status); + return { + ...task, + state: nextState, + progress: progressForState(nextState), + message: document.error_message ?? task.message, + }; + }), + ); + }, [documents]); - useEffect(() => { - if (queuedTasks.length === 0) { + const addFiles = (incoming: FileList | File[]) => { + const incomingFiles = Array.from(incoming); + if (incomingFiles.length === 0) { return; } - const slots = Math.max(CONCURRENCY - activeCount, 0); - if (slots === 0) { - return; - } + setError(null); + setTasks((current) => { + const activeSelected = current.filter((task) => task.state !== "rejected").length; + let remainingSlots = Math.max(MAX_FILES - activeSelected, 0); + const additions: UploadTask[] = incomingFiles.map((file, index) => { + const id = clientFileId(file, index); + const validationError = validateFile(file); + if (validationError) { + return { + id, + file, + state: "rejected", + progress: 0, + message: validationError, + }; + } + + if (remainingSlots <= 0) { + return { + id, + file, + state: "rejected", + progress: 0, + message: `This upload run supports up to ${MAX_FILES} files.`, + }; + } - const toStart = queuedTasks - .filter((task) => !startedTaskIdsRef.current.has(task.id)) - .slice(0, slots); + remainingSlots -= 1; + return { + id, + file, + state: "selected", + progress: 0, + }; + }); - toStart.forEach((task) => { - startedTaskIdsRef.current.add(task.id); - updateTask(task.id, { state: "preparing", progress: 10 }); - void runTask(task); + return [...current, ...additions]; }); - }, [activeCount, queuedTasks, runTask, updateTask]); + }; - useEffect(() => { - if (documents.length === 0) { + const startBatchUpload = async () => { + const filesToPrepare = tasks.filter((task) => task.state === "selected"); + if (filesToPrepare.length === 0 || isRunning) { return; } - const indexedById = new Set( - documents - .filter((doc) => doc.status === "indexed" || doc.status === "ready") - .map((doc) => doc.id), - ); - const indexedByFilename = new Set( - documents - .filter((doc) => doc.status === "indexed" || doc.status === "ready") - .map((doc) => doc.filename), - ); + setError(null); + setRunState("preparing"); - setTasks((current) => - current.filter((task) => { - if (task.documentId && indexedById.has(task.documentId)) { - return false; + try { + const prepared = await apiUploadPrepareBatch(token, { + name: `Upload run ${new Date().toLocaleString()}`, + files: filesToPrepare.map((task) => ({ + filename: task.file.name, + content_type: task.file.type || "application/pdf", + file_size_bytes: task.file.size, + idempotency_key: task.id, + client_file_id: task.id, + })), + }); + + setRunId(prepared.ingestion_run_id); + onRunCreated?.(prepared.ingestion_run_id); + + const preparedTasks: PreparedTask[] = []; + const taskByClientId = new Map(filesToPrepare.map((task) => [task.id, task])); + + prepared.items.forEach((item) => { + const task = item.client_file_id ? taskByClientId.get(item.client_file_id) : filesToPrepare[item.index]; + if (!task) { + return; } - if (!task.documentId && indexedByFilename.has(task.file.name)) { - return false; + + if ((item.status === "prepared" || item.status === "accepted") && item.document_id && item.bucket && item.storage_path && item.upload_url) { + preparedTasks.push({ + taskId: task.id, + documentId: item.document_id, + bucket: item.bucket, + storagePath: item.storage_path, + uploadUrl: item.upload_url, + }); + updateTask(task.id, { + state: "accepted", + progress: progressForState("accepted"), + documentId: item.document_id, + bucket: item.bucket, + storagePath: item.storage_path, + uploadUrl: item.upload_url, + message: undefined, + }); + return; } - return true; - }), - ); - }, [documents]); - const addFiles = (incoming: FileList | File[]) => { - const files = Array.from(incoming).filter((file) => file.type === "application/pdf" || file.name.toLowerCase().endsWith(".pdf")); + updateTask(task.id, { + state: "rejected", + progress: 0, + message: item.error ?? "The server rejected this file.", + }); + }); - if (files.length === 0) { - setError("Select at least one PDF file."); - return; - } + if (preparedTasks.length === 0) { + setRunState("idle"); + await onAfterUpload(); + return; + } - setError(null); + setRunState("uploading"); + const completedFiles: UploadCompleteBatchFile[] = []; + const documentIdToTaskId = new Map(); - setTasks((current) => { - if (current.length + files.length > MAX_FILES) { - setError(`You can upload up to ${MAX_FILES} files at once.`); - return current; + await runWithConcurrency(preparedTasks, CONCURRENCY, async (preparedTask) => { + const task = filesToPrepare.find((item) => item.id === preparedTask.taskId); + if (!task) { + return; + } + + try { + updateTask(preparedTask.taskId, { state: "uploading", progress: progressForState("uploading") }); + await apiUploadToSignedUrl(preparedTask.uploadUrl, task.file); + updateTask(preparedTask.taskId, { state: "completing", progress: progressForState("completing") }); + completedFiles.push({ + document_id: preparedTask.documentId, + bucket: preparedTask.bucket, + storage_path: preparedTask.storagePath, + }); + documentIdToTaskId.set(preparedTask.documentId, preparedTask.taskId); + } catch (uploadError) { + updateTask(preparedTask.taskId, { + state: "failed", + progress: 0, + message: uploadError instanceof Error ? uploadError.message : "Storage upload failed.", + }); + } + }); + + if (completedFiles.length > 0) { + const completed = await apiUploadCompleteBatch(token, { + ingestion_run_id: prepared.ingestion_run_id, + files: completedFiles, + }); + + completed.items.forEach((item) => { + const taskId = documentIdToTaskId.get(item.document_id); + if (!taskId) { + return; + } + + if (item.error || item.status === "failed") { + updateTask(taskId, { + state: "failed", + progress: 0, + message: item.error ?? "Upload completion failed.", + }); + return; + } + + updateTask(taskId, { + state: "queued", + progress: progressForState("queued"), + message: "Queued for extraction", + }); + }); } - const additions: UploadTask[] = files.map((file, idx) => ({ - id: `${file.name}-${file.size}-${Date.now()}-${idx}`, - file, - state: "queued", - progress: 0, - })); + setRunState("processing"); + await onAfterUpload(); + } catch (err) { + setRunState("idle"); + setError(err instanceof Error ? err.message : "Batch upload failed."); + } + }; - return [...current, ...additions]; - }); + const clearFinished = () => { + setTasks((current) => current.filter((task) => !["ready", "rejected", "failed"].includes(task.state))); }; const onDrop: DragEventHandler = (event) => { @@ -200,19 +396,33 @@ export default function UploadPanel({ token, onAfterUpload, documents }: UploadP return (
-
+
-

Upload documents

-

Drop PDFs or browse to upload. Maximum 100 files.

+

Bulk document upload

+

Select up to {MAX_FILES} PDFs, review the batch, then start one ingestion run.

+
+
+ {tasks.length > 0 ? ( + + ) : null} +
-
+
+ + 0 ? "danger" : "neutral"} /> + 0 ? "warning" : "neutral"} /> + 0 ? "success" : "neutral"} /> + 0 ? "danger" : "neutral"} /> + {runId ? : null} +
+ {error ?

{error}

: null} +
+

+ {runState === "preparing" + ? "Preparing signed upload URLs..." + : runState === "uploading" + ? "Uploading accepted files..." + : runState === "processing" + ? "Backend ingestion is running. Status will refresh automatically." + : selectedCount > 0 + ? `${selectedCount} files are ready to start.` + : "No files selected yet."} +

+ +
+
{tasks.map((task) => ( -
+

{task.file.name}

-

{statusLabel(task.state)}

+

+ {statusLabel(task.state)} - {fileSizeLabel(task.file.size)} +

- {task.state === "failed" ? ( - + {task.state === "ready" ? : null} + {task.state === "failed" || task.state === "rejected" ? ( + ) : null}
@@ -267,13 +507,21 @@ export default function UploadPanel({ token, onAfterUpload, documents }: UploadP
- {task.message ?

{task.message}

: null} + {task.message ? ( +

+ {task.message} +

+ ) : null}
))} @@ -282,3 +530,28 @@ export default function UploadPanel({ token, onAfterUpload, documents }: UploadP
); } + +function Metric({ + label, + value, + tone = "neutral", +}: { + label: string; + value: number | string; + tone?: "neutral" | "success" | "warning" | "danger"; +}) { + return ( + + {label} + {value} + + ); +} diff --git a/client/src/components/usage/TokenMeter.tsx b/client/src/components/usage/TokenMeter.tsx index e698c90..45465a1 100644 --- a/client/src/components/usage/TokenMeter.tsx +++ b/client/src/components/usage/TokenMeter.tsx @@ -19,17 +19,19 @@ export default function TokenMeter({ usage, loading = false, error = null }: Tok return
No usage data yet.
; } - const consumed = usage.used + usage.reserved; + const used = usage.used ?? usage.tokens_used ?? 0; + const reserved = usage.reserved ?? usage.tokens_reserved ?? 0; + const consumed = used + reserved; const percent = usage.limit > 0 ? Math.min(100, Math.round((consumed / usage.limit) * 100)) : 0; return (

Today's Token Budget

- Used: {usage.used} + Used: {used}

- Reserved: {usage.reserved} + Reserved: {reserved}

Remaining: {usage.remaining} / {usage.limit} diff --git a/client/src/lib/api.ts b/client/src/lib/api.ts index 0f72949..4475cc5 100644 --- a/client/src/lib/api.ts +++ b/client/src/lib/api.ts @@ -41,6 +41,7 @@ export type DocumentRecord = { status: DocumentStatus; page_count?: number | null; file_size_bytes?: number; + ingestion_run_id?: string | null; created_at?: string; updated_at?: string; error_message?: string | null; @@ -60,12 +61,101 @@ export type UploadCompleteResponse = { job_id?: string; }; +export type UploadPrepareBatchFile = { + filename: string; + content_type: string; + file_size_bytes: number; + idempotency_key?: string; + client_file_id?: string; +}; + +export type UploadPrepareBatchItem = { + index: number; + filename: string; + client_file_id?: string | null; + status: string; + document_id?: string | null; + bucket?: string | null; + storage_path?: string | null; + upload_url?: string | null; + expires_in?: number | null; + error?: string | null; +}; + +export type UploadPrepareBatchResponse = { + ingestion_run_id: string; + bucket: string; + expires_in: number; + accepted_count: number; + rejected_count: number; + items: UploadPrepareBatchItem[]; +}; + +export type UploadCompleteBatchFile = { + document_id: string; + bucket: string; + storage_path: string; +}; + +export type UploadCompleteBatchItem = { + index: number; + document_id: string; + status: string; + job_id?: string | null; + error?: string | null; +}; + +export type UploadCompleteBatchResponse = { + ingestion_run_id?: string | null; + accepted_count: number; + failed_count: number; + items: UploadCompleteBatchItem[]; +}; + export type DocumentJobResponse = { document_id: string; status: string; job_id: string; }; +export type IngestionRunStatusCounts = { + pending_upload: number; + uploading: number; + queued: number; + uploaded: number; + extracting: number; + indexing: number; + ready: number; + indexed: number; + failed: number; + total: number; +}; + +export type IngestionRunResponse = { + id: string; + name?: string | null; + status: string; + total_documents: number; + accepted_documents: number; + rejected_documents: number; + document_statuses: IngestionRunStatusCounts; + created_at: string; + updated_at: string; +}; + +export type IngestionQueueStatusItem = { + name: string; + queued_count: number; + started_count: number; + deferred_count: number; + scheduled_count: number; + failed_count: number; +}; + +export type IngestionQueueStatusResponse = { + queues: IngestionQueueStatusItem[]; +}; + export type QueryCitation = { document_id: string; page_number: number; @@ -336,16 +426,13 @@ function ensureUsage(value: unknown): UsageToday { function normalizeDocumentStatus(input: unknown): DocumentStatus { const status = String(input ?? "uploaded").toLowerCase(); - if (status === "ready") { - return "indexed"; - } - if ( status === "pending_upload" || status === "uploaded" || status === "queued" || status === "extracting" || status === "indexing" || + status === "ready" || status === "indexed" || status === "failed" ) { @@ -371,6 +458,7 @@ function normalizeDocument(rawDoc: unknown): DocumentRecord | null { status: normalizeDocumentStatus(raw.status), page_count: raw.page_count == null ? null : Number(raw.page_count), file_size_bytes: raw.file_size_bytes == null ? undefined : Number(raw.file_size_bytes), + ingestion_run_id: raw.ingestion_run_id == null ? null : String(raw.ingestion_run_id), created_at: raw.created_at ? String(raw.created_at) : undefined, updated_at: raw.updated_at ? String(raw.updated_at) : undefined, error_message: raw.error_message ? String(raw.error_message) : null, @@ -424,8 +512,19 @@ export function apiAuthMe(token: string): Promise { return apiGetAuthMe(token); } -export async function apiGetDocuments(token: string): Promise { - const payload = await apiRequest("/documents", token, { method: "GET" }); +export async function apiGetDocuments(token: string, params?: { limit?: number; offset?: number; status?: string }): Promise { + const search = new URLSearchParams(); + if (typeof params?.limit === "number") { + search.set("limit", String(params.limit)); + } + if (typeof params?.offset === "number") { + search.set("offset", String(params.offset)); + } + if (params?.status) { + search.set("status", params.status); + } + const suffix = search.toString() ? `?${search.toString()}` : ""; + const payload = await apiRequest(`/documents${suffix}`, token, { method: "GET" }); return normalizeDocuments(payload); } @@ -492,6 +591,34 @@ export function apiUploadComplete( }); } +export function apiUploadPrepareBatch( + token: string, + payload: { files: UploadPrepareBatchFile[]; name?: string }, +): Promise { + return apiRequest("/documents/upload-prepare-batch", token, { + method: "POST", + body: JSON.stringify(payload), + }); +} + +export function apiUploadCompleteBatch( + token: string, + payload: { files: UploadCompleteBatchFile[]; ingestion_run_id?: string | null }, +): Promise { + return apiRequest("/documents/upload-complete-batch", token, { + method: "POST", + body: JSON.stringify(payload), + }); +} + +export function apiGetIngestionRun(token: string, runId: string): Promise { + return apiRequest(`/documents/ingestion-runs/${runId}`, token, { method: "GET" }); +} + +export function apiGetIngestionQueues(token: string): Promise { + return apiRequest("/documents/ingestion-queues", token, { method: "GET" }); +} + export async function apiUploadDocument(token: string, file: File): Promise { const prepare = await apiUploadPrepare(token, { filename: file.name, diff --git a/client/src/pages/UploadPage.tsx b/client/src/pages/UploadPage.tsx index a52c306..47d5c5b 100644 --- a/client/src/pages/UploadPage.tsx +++ b/client/src/pages/UploadPage.tsx @@ -1,62 +1,263 @@ -import { useEffect, useMemo, useState } from "react"; +import { + AlertTriangle, + CheckCircle2, + Clock3, + FileText, + RefreshCw, + RotateCcw, + Search, + Trash2, +} from "lucide-react"; +import { useCallback, useEffect, useMemo, useState } from "react"; +import type { ReactNode } from "react"; import UploadPanel from "../components/upload/UploadPanel"; import { useAuth } from "../context/AuthContext"; import { useAppShellContext } from "../components/layout/AppShell"; -import { apiDeleteDocument } from "../lib/api"; +import { + apiDeleteDocument, + apiGetIngestionQueues, + apiGetIngestionRun, + apiRetryDocument, + type DocumentRecord, + type DocumentStatus, + type IngestionQueueStatusResponse, + type IngestionRunResponse, + type IngestionRunStatusCounts, +} from "../lib/api"; +import { cn } from "../lib/utils"; -const processingStatuses = new Set(["pending_upload", "uploaded", "queued", "extracting", "indexing"]); +const processingStatuses = new Set(["pending_upload", "uploaded", "queued", "extracting", "indexing"]); +const readyStatuses = new Set(["ready", "indexed"]); +const terminalFailurePrefixes = [ + "Validation failure:", + "Upload/storage failure:", + "Page-limit failure:", + "Unsupported-content failure:", +]; -function statusBadge(status: string): string { - if (status === "indexed" || status === "ready") { - return "border-app-success/30 bg-green-50 text-green-700"; +type StatusFilter = "all" | "processing" | "failed" | "ready"; +type SortMode = "newest" | "oldest" | "status" | "updated"; + +function isProcessing(status: DocumentStatus): boolean { + return processingStatuses.has(status); +} + +function isReady(status: DocumentStatus): boolean { + return readyStatuses.has(status); +} + +function isRetryableFailure(errorMessage?: string | null): boolean { + if (!errorMessage) { + return true; + } + return !terminalFailurePrefixes.some((prefix) => errorMessage.startsWith(prefix)); +} + +function statusLabel(status: DocumentStatus): string { + if (status === "indexed" || status === "ready") return "Ready"; + if (status === "pending_upload") return "Waiting to upload"; + if (status === "uploaded" || status === "queued") return "Queued"; + if (status === "extracting") return "Extracting text"; + if (status === "indexing") return "Indexing"; + if (status === "failed") return "Failed"; + return String(status); +} + +function statusHelp(status: DocumentStatus): string { + if (status === "pending_upload") return "The document has a placeholder but no uploaded object yet."; + if (status === "uploaded" || status === "queued") return "The file is uploaded and waiting for a worker."; + if (status === "extracting") return "Text is being extracted from the PDF."; + if (status === "indexing") return "Chunks and embeddings are being created."; + if (status === "ready" || status === "indexed") return "The document is available for chat."; + if (status === "failed") return "Ingestion stopped for this document."; + return "Document status"; +} + +function statusBadge(status: DocumentStatus): string { + if (isReady(status)) { + return "border-green-200 bg-green-50 text-green-700"; } if (status === "failed") { - return "border-app-danger/30 bg-red-50 text-red-700"; + return "border-red-200 bg-red-50 text-red-700"; } return "border-amber-200 bg-amber-50 text-amber-700"; } -function statusLabel(status: string): string { - if (status === "indexed" || status === "ready") { - return "Indexed"; - } +function statusRank(status: DocumentStatus): number { + if (isProcessing(status)) return 0; + if (status === "failed") return 1; + if (isReady(status)) return 2; + return 3; +} - if (status === "pending_upload") { - return "Pending upload"; - } +function formatDate(value?: string): string { + return value ? new Date(value).toLocaleString() : "--"; +} + +function countDocuments(documents: DocumentRecord[]): IngestionRunStatusCounts { + const counts: IngestionRunStatusCounts = { + pending_upload: 0, + uploading: 0, + queued: 0, + uploaded: 0, + extracting: 0, + indexing: 0, + ready: 0, + indexed: 0, + failed: 0, + total: 0, + }; + + documents.forEach((doc) => { + counts[doc.status] += 1; + counts.total += 1; + }); + + return counts; +} - return status.charAt(0).toUpperCase() + status.slice(1); +function normalizeCounts(counts: Partial): IngestionRunStatusCounts { + return { + pending_upload: counts.pending_upload ?? 0, + uploading: counts.uploading ?? 0, + queued: counts.queued ?? 0, + uploaded: counts.uploaded ?? 0, + extracting: counts.extracting ?? 0, + indexing: counts.indexing ?? 0, + ready: counts.ready ?? 0, + indexed: counts.indexed ?? 0, + failed: counts.failed ?? 0, + total: counts.total ?? 0, + }; +} + +function completedCount(counts: IngestionRunStatusCounts): number { + return counts.ready + counts.indexed + counts.failed; +} + +function queueTotals(queueStatus: IngestionQueueStatusResponse | null) { + return (queueStatus?.queues ?? []).reduce( + (acc, queue) => ({ + queued: acc.queued + queue.queued_count, + running: acc.running + queue.started_count, + failed: acc.failed + queue.failed_count, + }), + { queued: 0, running: 0, failed: 0 }, + ); } export default function UploadPage() { const { accessToken } = useAuth(); - const { documents, loading, refreshDocuments, activeDocument, setActiveDocumentId } = useAppShellContext(); + const { + documents, + loading, + refreshDocuments, + refreshWorkspace, + activeDocument, + setActiveDocumentId, + } = useAppShellContext(); const [deletingId, setDeletingId] = useState(null); + const [retryingIds, setRetryingIds] = useState>(new Set()); + const [activeRunId, setActiveRunId] = useState(null); + const [activeRun, setActiveRun] = useState(null); + const [queueStatus, setQueueStatus] = useState(null); + const [statusError, setStatusError] = useState(null); + const [query, setQuery] = useState(""); + const [statusFilter, setStatusFilter] = useState("all"); + const [sortMode, setSortMode] = useState("newest"); - const hasProcessing = useMemo( - () => documents.some((doc) => processingStatuses.has(doc.status)), - [documents], - ); + const hasProcessing = useMemo(() => documents.some((doc) => isProcessing(doc.status)), [documents]); useEffect(() => { - if (!hasProcessing) { + if (activeRunId) { + return; + } + + const latestRunDocument = documents.find((doc) => doc.ingestion_run_id && isProcessing(doc.status)) + ?? documents.find((doc) => doc.ingestion_run_id); + + if (latestRunDocument?.ingestion_run_id) { + setActiveRunId(latestRunDocument.ingestion_run_id); + } + }, [activeRunId, documents]); + + const refreshIngestionState = useCallback(async () => { + if (!accessToken) { + return; + } + + try { + const [queues, run] = await Promise.all([ + apiGetIngestionQueues(accessToken), + activeRunId ? apiGetIngestionRun(accessToken, activeRunId) : Promise.resolve(null), + ]); + setQueueStatus(queues); + setActiveRun(run); + setStatusError(null); + } catch (error) { + setStatusError(error instanceof Error ? error.message : "Unable to refresh ingestion state."); + } + }, [accessToken, activeRunId]); + + useEffect(() => { + if (!accessToken) { + return; + } + + void refreshIngestionState(); + + const runStillProcessing = activeRun?.status === "processing"; + if (!hasProcessing && !runStillProcessing) { return; } const interval = window.setInterval(() => { void refreshDocuments(); + void refreshWorkspace(); + void refreshIngestionState(); }, 4000); return () => { window.clearInterval(interval); }; - }, [hasProcessing, refreshDocuments]); + }, [accessToken, activeRun?.status, hasProcessing, refreshDocuments, refreshIngestionState, refreshWorkspace]); if (!accessToken) { return null; } + const counts = countDocuments(documents); + const processingCount = counts.pending_upload + counts.uploaded + counts.queued + counts.extracting + counts.indexing; + const readyCount = counts.ready + counts.indexed; + const failedCount = counts.failed; + + const visibleDocuments = documents + .filter((doc) => doc.filename.toLowerCase().includes(query.toLowerCase())) + .filter((doc) => { + if (statusFilter === "processing") return isProcessing(doc.status); + if (statusFilter === "failed") return doc.status === "failed"; + if (statusFilter === "ready") return isReady(doc.status); + return true; + }) + .sort((a, b) => { + if (sortMode === "oldest") { + return new Date(a.created_at ?? "").getTime() - new Date(b.created_at ?? "").getTime(); + } + if (sortMode === "updated") { + return new Date(b.updated_at ?? "").getTime() - new Date(a.updated_at ?? "").getTime(); + } + if (sortMode === "status") { + return statusRank(a.status) - statusRank(b.status); + } + return new Date(b.created_at ?? "").getTime() - new Date(a.created_at ?? "").getTime(); + }); + + const activeDocuments = visibleDocuments.filter((doc) => isProcessing(doc.status)); + const failedDocuments = visibleDocuments.filter((doc) => doc.status === "failed"); + const readyDocuments = visibleDocuments.filter((doc) => isReady(doc.status)); + const retryableFailedDocuments = documents.filter((doc) => doc.status === "failed" && isRetryableFailure(doc.error_message)); + const handleDelete = async (documentId: string, filename: string) => { const confirmed = window.confirm(`Delete document "${filename}"? This cannot be undone.`); if (!confirmed) { @@ -69,7 +270,7 @@ export default function UploadPage() { if (activeDocument?.id === documentId) { setActiveDocumentId(null); } - await refreshDocuments(); + await Promise.all([refreshDocuments(), refreshWorkspace(), refreshIngestionState()]); } catch (error) { const message = error instanceof Error ? error.message : "Delete failed"; window.alert(message); @@ -78,62 +279,410 @@ export default function UploadPage() { } }; + const retryDocuments = async (targets: DocumentRecord[]) => { + const retryableTargets = targets.filter((doc) => isRetryableFailure(doc.error_message)); + if (retryableTargets.length === 0) { + window.alert("No retryable failed documents are selected."); + return; + } + + setRetryingIds((current) => { + const next = new Set(current); + retryableTargets.forEach((doc) => next.add(doc.id)); + return next; + }); + + const failures: string[] = []; + for (const document of retryableTargets) { + try { + await apiRetryDocument(accessToken, document.id); + } catch (error) { + failures.push(`${document.filename}: ${error instanceof Error ? error.message : "Retry failed"}`); + } + } + + setRetryingIds((current) => { + const next = new Set(current); + retryableTargets.forEach((doc) => next.delete(doc.id)); + return next; + }); + + await Promise.all([refreshDocuments(), refreshWorkspace(), refreshIngestionState()]); + + if (failures.length > 0) { + window.alert(`Some retries failed:\n${failures.join("\n")}`); + } + }; + return (

- + { + await Promise.all([refreshDocuments(), refreshWorkspace(), refreshIngestionState()]); + }} + documents={documents} + onRunCreated={setActiveRunId} + /> + + doc.ingestion_run_id === activeRunId) : []} + queueStatus={queueStatus} + statusError={statusError} + onRefresh={() => { + void Promise.all([refreshDocuments(), refreshWorkspace(), refreshIngestionState()]); + }} + /> -
-
-

Documents

-

Ingestion status updates automatically while processing is active.

+
+
+
+

Ingestion workspace

+

Monitor processing, recover failures, and confirm which documents are ready.

+
+ +
+ } /> + } tone="warning" /> + } tone="success" /> + } tone="danger" /> +
-
- - - - - - - - - - - - {documents.map((doc) => ( - - - - - - - +
+ + +
+ {(["all", "processing", "failed", "ready"] as StatusFilter[]).map((filter) => ( + + ))} +
+ + +
+ +
+ {(statusFilter === "all" || statusFilter === "processing") ? ( + + {activeDocuments.map((doc) => ( + void handleDelete(doc.id, doc.filename)} + /> + ))} + + ) : null} + + {(statusFilter === "all" || statusFilter === "failed") ? ( + void retryDocuments(retryableFailedDocuments)} + disabled={retryableFailedDocuments.length === 0} + > + + Retry all retryable + + )} + > + {failedDocuments.map((doc) => ( + void retryDocuments([doc]) : undefined} + onDelete={() => void handleDelete(doc.id, doc.filename)} + /> ))} + + ) : null} - {!loading && documents.length === 0 ? ( -
- - - ) : null} - -
FilenameStatusPagesUpdatedActions
{doc.filename} - - {statusLabel(doc.status)} - - {doc.page_count ?? "--"}{doc.updated_at ? new Date(doc.updated_at).toLocaleString() : "--"} - -
- No documents yet. -
+ {(statusFilter === "all" || statusFilter === "ready") ? ( + + {readyDocuments.map((doc) => ( + void handleDelete(doc.id, doc.filename)} + /> + ))} + + ) : null}
); } + +function IngestionSummary({ + run, + runDocuments, + queueStatus, + statusError, + onRefresh, +}: { + run: IngestionRunResponse | null; + runDocuments: DocumentRecord[]; + queueStatus: IngestionQueueStatusResponse | null; + statusError: string | null; + onRefresh: () => void; +}) { + const counts = normalizeCounts(run?.document_statuses ?? countDocuments(runDocuments)); + const acceptedDocuments = run?.accepted_documents ?? runDocuments.length; + const rejectedDocuments = run?.rejected_documents ?? 0; + const doneCount = completedCount(counts); + const progress = acceptedDocuments > 0 ? Math.round((doneCount / acceptedDocuments) * 100) : 0; + const queues = queueTotals(queueStatus); + const hasRun = Boolean(run) || runDocuments.length > 0; + + return ( +
+
+
+

Current ingestion run

+

+ {hasRun + ? `${doneCount} of ${acceptedDocuments} accepted documents have reached a final state.` + : "Start a batch upload to track one ingestion run here."} +

+
+ +
+ +
+
+
+ +
+ + 0 ? "danger" : "neutral"} /> + + + + +
+ +
+ {run ? ( + + Run status: {run.status} + + ) : null} + + Queue: {queues.queued} waiting,{" "} + {queues.running} running + + {queues.failed > 0 ? ( + + Queue failures: {queues.failed} + + ) : null} + {statusError ? ( + + {statusError} + + ) : null} +
+
+ ); +} + +function SummaryMetric({ + label, + value, + icon, + tone = "neutral", +}: { + label: string; + value: number; + icon: ReactNode; + tone?: "neutral" | "success" | "warning" | "danger"; +}) { + return ( +
+
+ {icon} + {label} +
+

{value}

+
+ ); +} + +function RunMetric({ + label, + value, + tone = "neutral", +}: { + label: string; + value: number; + tone?: "neutral" | "success" | "warning" | "danger"; +}) { + return ( +
+

{label}

+

{value}

+
+ ); +} + +function DocumentSection({ + title, + count, + emptyText, + action, + children, +}: { + title: string; + count: number; + emptyText: string; + action?: ReactNode; + children: ReactNode; +}) { + return ( +
+
+
+

{title}

+ {count} +
+ {action} +
+
+ {count > 0 ? children :

{emptyText}

} +
+
+ ); +} + +function DocumentRow({ + document, + deleting, + retrying, + onRetry, + onDelete, +}: { + document: DocumentRecord; + deleting: boolean; + retrying: boolean; + onRetry?: () => void; + onDelete: () => void; +}) { + return ( +
+
+
+

{document.filename}

+

{statusHelp(document.status)}

+ {document.error_message ?

{document.error_message}

: null} +
+ +
+ + {statusLabel(document.status)} + +
+ +
+

{document.page_count == null ? "Pages --" : `${document.page_count} pages`}

+

Updated {formatDate(document.updated_at)}

+
+ +
+ {onRetry ? ( + + ) : null} + +
+
+
+ ); +} diff --git a/docker-compose.yml b/docker-compose.yml index b91f191..55356af 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -1,25 +1,4 @@ -version: '3.8' - services: - # PostgreSQL (local dev only - production uses Supabase) - postgres: - image: pgvector/pgvector:pg16 - container_name: rag-postgres - environment: - POSTGRES_DB: enterprise_rag - POSTGRES_USER: postgres - POSTGRES_PASSWORD: postgres - volumes: - - postgres_data:/var/lib/postgresql/data - - ./scripts/init-db.sql:/docker-entrypoint-initdb.d/init.sql - ports: - - "5432:5432" - healthcheck: - test: ["CMD-SHELL", "pg_isready -U postgres"] - interval: 5s - timeout: 5s - retries: 5 - # Redis (queue backend + cache) redis: image: redis:7-alpine @@ -49,6 +28,10 @@ services: - SUPABASE_KEY=${SUPABASE_KEY} - SUPABASE_JWT_SECRET=${SUPABASE_JWT_SECRET} - OPENAI_API_KEY=${OPENAI_API_KEY} + - OPENAI_EMBEDDING_TIMEOUT_SECONDS=${OPENAI_EMBEDDING_TIMEOUT_SECONDS:-300} + - INGEST_EXTRACT_JOB_TIMEOUT_SECONDS=${INGEST_EXTRACT_JOB_TIMEOUT_SECONDS:-900} + - INGEST_INDEX_JOB_TIMEOUT_SECONDS=${INGEST_INDEX_JOB_TIMEOUT_SECONDS:-1800} + - EMBEDDING_BATCH_SIZE=${EMBEDDING_BATCH_SIZE:-32} - ENVIRONMENT=development volumes: - ./server:/app @@ -56,8 +39,6 @@ services: ports: - "8000:8000" depends_on: - postgres: - condition: service_healthy redis: condition: service_healthy command: uvicorn app.main:app --host 0.0.0.0 --port 8000 --reload @@ -74,6 +55,10 @@ services: - SUPABASE_SERVICE_ROLE_KEY=${SUPABASE_SERVICE_ROLE_KEY} - SUPABASE_KEY=${SUPABASE_KEY} - OPENAI_API_KEY=${OPENAI_API_KEY} + - OPENAI_EMBEDDING_TIMEOUT_SECONDS=${OPENAI_EMBEDDING_TIMEOUT_SECONDS:-300} + - INGEST_EXTRACT_JOB_TIMEOUT_SECONDS=${INGEST_EXTRACT_JOB_TIMEOUT_SECONDS:-900} + - INGEST_INDEX_JOB_TIMEOUT_SECONDS=${INGEST_INDEX_JOB_TIMEOUT_SECONDS:-1800} + - EMBEDDING_BATCH_SIZE=${EMBEDDING_BATCH_SIZE:-32} - QUEUE_NAME=ingest_extract - WORKER_COUNT=5 volumes: @@ -81,7 +66,6 @@ services: - ./server/app:/app/shared # Share code depends_on: - redis - - postgres deploy: replicas: 5 command: python worker.py ingest_extract @@ -98,6 +82,10 @@ services: - SUPABASE_SERVICE_ROLE_KEY=${SUPABASE_SERVICE_ROLE_KEY} - SUPABASE_KEY=${SUPABASE_KEY} - OPENAI_API_KEY=${OPENAI_API_KEY} + - OPENAI_EMBEDDING_TIMEOUT_SECONDS=${OPENAI_EMBEDDING_TIMEOUT_SECONDS:-300} + - INGEST_EXTRACT_JOB_TIMEOUT_SECONDS=${INGEST_EXTRACT_JOB_TIMEOUT_SECONDS:-900} + - INGEST_INDEX_JOB_TIMEOUT_SECONDS=${INGEST_INDEX_JOB_TIMEOUT_SECONDS:-1800} + - EMBEDDING_BATCH_SIZE=${EMBEDDING_BATCH_SIZE:-32} - QUEUE_NAME=ingest_index - WORKER_COUNT=3 volumes: @@ -105,7 +93,6 @@ services: - ./server/app:/app/shared depends_on: - redis - - postgres deploy: replicas: 3 command: python worker.py ingest_index @@ -140,6 +127,5 @@ services: command: npm run dev -- --host 0.0.0.0 volumes: - postgres_data: redis_data: client_node_modules: diff --git a/docs/ingestion-workflow-evaluation.md b/docs/ingestion-workflow-evaluation.md new file mode 100644 index 0000000..1eb0414 --- /dev/null +++ b/docs/ingestion-workflow-evaluation.md @@ -0,0 +1,322 @@ +# Ingestion Workflow Evaluation + +Date: 2026-06-05 + +## Executive Summary + +The ingestion workflow is working for the current small local test set. Batch upload, extraction, indexing, expected validation failures, retry, queue visibility, and final document states were all exercised through the running Docker Compose stack. + +The current system is ready for small PDF ingestion and manual validation. It is not yet proven for larger-scale ingestion such as 25, 50, or 100 PDFs in one run. The biggest remaining weaknesses are observability and run lifecycle accuracy, not the basic document pipeline. + +Main strengths: + +- Valid PDFs move to `ready`. +- Page-limit failures and scanned/image-only PDFs are rejected with clear document-level messages. +- Extraction and indexing are handled asynchronously by separate RQ queues. +- Current queue depth was visible and drained to zero after the batch. +- Embedding calls are batched through `EMBEDDING_BATCH_SIZE`. +- Timeout/failure callback handling now exists so killed indexing jobs should not leave documents stuck forever. + +Main weaknesses: + +- Historical ingestion runs can remain `processing` or `preparing` even after their documents are no longer active. +- Expected business-rule rejections still appear in RQ failed registries, which makes operational failure counts noisy. +- There is no durable per-stage timing for queue wait, extraction time, embedding time, and DB write time. +- Large-batch behavior was not tested in this pass. +- Token-budget exhaustion behavior during concurrent ingestion was not tested in this pass. + +## Test Setup + +Environment: + +- Docker Compose local stack on Windows. +- Client: `http://localhost:5173`. +- API: `http://localhost:8000`. +- RQ dashboard: `http://localhost:9181`. +- Postgres with pgvector: `rag-postgres`. +- Redis: `rag-redis`. +- Extraction workers: 5 `worker-extract` containers. +- Index workers: 3 `worker-index` containers. +- OpenAI calls: real embedding calls, not mocked. +- Supabase: used for auth and raw PDF storage. + +Current configured limits observed from docs/config: + +- Maximum PDF page count: `10`. +- Maximum file size: `10 MB`. +- Maximum backend batch upload size: `50`. +- Embedding batch size: `32`. +- Extract job timeout: `900` seconds. +- Index job timeout: `1800` seconds. +- OpenAI embedding request timeout: `300` seconds. + +Documents tested in the latest real run: + +| File | Expected Result | Actual Result | +| --- | --- | --- | +| `terms-and-conditions-template.pdf` | valid, under 10 pages | `ready`, 2 pages | +| `About_ChatGPT_Pro_tiers_OpenAI_Help_Center.pdf` | valid, under 10 pages | `ready`, 4 pages | +| `Terms-and-Conditions-Free-Template_PDF-2.pdf` | too many pages | `failed`, page-limit failure | +| `Screenshot_22_.pdf` | image-only/scanned | `failed`, unsupported-content failure | + +## Workflow Validation + +### Upload Preparation + +Batch upload preparation worked. The frontend sent `POST /documents/upload-prepare-batch` and the server returned `201 Created`. The batch run was created in `ingestion_runs`. + +The frontend/backend status mismatch found earlier was fixed: the backend returns prepared batch items, and the frontend now treats prepared items as uploadable instead of leaving documents at `pending_upload`. + +### Upload Completion + +Upload completion worked for the latest real run. Documents moved past `pending_upload`; no document remained stuck in upload state after the validated batch. + +### Extraction + +Extraction behaved correctly: + +- Two valid PDFs produced page rows. +- One 28-page PDF failed because it exceeded the 10-page limit. +- One scanned/image-only PDF failed because no extractable text was found. +- The latest database state had 6 extracted `document_pages` rows, matching the two valid PDFs. + +### Indexing + +Indexing behaved correctly after the timeout/failure handling work and retry: + +- The two valid PDFs became `ready`. +- The latest database state had 6 `chunks` rows and 6 `chunk_embeddings` rows. +- Index workers logged successful index jobs for the two valid documents. +- No latest document remained stuck in `indexing`. + +### Final Document State + +Latest document summary from Postgres: + +| Metric | Value | +| --- | ---: | +| Total latest documents | 4 | +| Ready | 2 | +| Failed | 2 | +| In progress | 0 | +| Extracted pages | 6 | +| Chunks | 6 | +| Embeddings | 6 | +| Extract queue depth after run | 0 | +| Index queue depth after run | 0 | + +The latest document-level states are correct. + +### Retry and Recovery + +Retry/recovery was exercised after a previous indexing worker timeout left a document in `indexing`. The stuck document was marked failed and retried, then reached `ready`. + +The code now also wires explicit RQ job timeouts and an `on_failure` callback that marks active documents failed when an RQ job is killed or times out. A forced worker-kill test after this callback was added was not performed in this pass. + +### Reindex + +The reindex endpoint exists and the indexing code rebuilds chunks/embeddings idempotently. A full user-facing reindex scenario was not executed in this evaluation pass. + +## Performance Findings + +This was a small functional/performance check, not a high-scale benchmark. + +Measured latest run: + +| File | Final Status | Time From Row Creation To Final State | +| --- | --- | ---: | +| `Screenshot_22_.pdf` | failed | 7.88 seconds | +| `Terms-and-Conditions-Free-Template_PDF-2.pdf` | failed | 8.17 seconds | +| `terms-and-conditions-template.pdf` | ready | 13.88 seconds | +| `About_ChatGPT_Pro_tiers_OpenAI_Help_Center.pdf` | ready | 14.25 seconds | + +Observed worker timing from logs: + +- Valid extraction jobs started around `19:29:00` to `19:29:01` UTC. +- Valid extraction jobs finished around `19:29:04` to `19:29:05` UTC. +- Index jobs started around `19:29:04` to `19:29:05` UTC. +- Index jobs finished around `19:29:10` UTC. +- Queue wait was small for this run because workers were already idle. + +Current bottleneck for small valid PDFs appears to be indexing/embedding, not extraction. That conclusion is limited because only two valid PDFs were indexed in the latest run. + +Not measured: + +- 10, 25, 50, or 100 PDF batch throughput. +- Average queue wait under sustained backlog. +- Worker utilization. +- DB query/write count. +- OpenAI request latency by batch. +- Throughput while queries are being served at the same time. + +## Reliability Findings + +What passed: + +- Valid documents reached `ready`. +- Invalid documents reached `failed`. +- Failure messages were specific and visible in the UI. +- Expected failures were isolated to the individual documents. +- Queues drained to zero after the run. +- The latest run had no documents stuck in `pending_upload`, `uploaded`, `extracting`, or `indexing`. + +Issues found: + +- Earlier, an index worker was killed after the default RQ timeout and the document remained `indexing`. This exposed a real failure-state bug. The code now has explicit job timeouts and a failure callback, but a deliberate post-fix kill test is still needed. +- RQ failed registries contain expected validation failures. At inspection time Redis had `11` failed extract jobs and `7` failed index jobs from current/prior runs. These counts are noisy because they mix expected document rejections with infrastructure failures. +- Historical ingestion runs are not always terminal in the database. At inspection time `ingestion_runs` had statuses: `partial = 1`, `processing = 5`, `preparing = 4`. Some of these are older runs from earlier failed attempts, but this shows run lifecycle cleanup/status derivation still needs hardening. + +## Queue and Worker Behavior + +Current worker topology is stronger than the original baseline: + +- 5 extraction workers listen on `ingest_extract`. +- 3 indexing workers listen on `ingest_index`. +- Each document still uses one extract job and one index job. +- Valid documents can extract and index in parallel when workers are available. + +Observed state after the latest run: + +- `rq:queue:ingest_extract` length: `0`. +- `rq:queue:ingest_index` length: `0`. +- Workers remained up and continued cleaning registries. + +Remaining gaps: + +- There is no backpressure policy beyond queueing. +- There is no persistent queue-depth history. +- There is no clear separation between expected document rejection and worker/system failure in RQ failure monitoring. +- There is no automatic stale-run reconciliation job. + +## Database Efficiency + +Improvements now present: + +- Page rows are inserted in a batch-style execution. +- Embedding vectors are inserted in batches. +- Indexing clears and rebuilds chunks/embeddings for idempotent retry/reindex behavior. + +Remaining concerns: + +- Work is still organized one document per extract job and one document per index job. +- There is no recorded DB write count per ingestion run. +- There is no stage timing stored alongside rows. +- There is no evidence yet for lock contention or write amplification under load because larger runs were not tested. + +## Embedding Efficiency + +The implementation now uses `EMBEDDING_BATCH_SIZE`, configured as `32`, so it no longer embeds one chunk per OpenAI call by design. + +In the latest run there were only 6 total chunks, so each document fit within a small number of embedding calls. This validates correctness but does not prove high-throughput embedding behavior. + +Embedding is still the most likely throughput bottleneck for larger text-heavy PDFs because it depends on external OpenAI latency, quota, and rate limits. + +## Token-Budget Safety + +What exists: + +- Indexing reserves token budget before embedding batches. +- Successful embedding batches commit usage. +- Failure paths attempt to release outstanding reservations. +- A maintenance job exists for stale reservation cleanup. + +What was not tested: + +- Low-budget ingestion failure. +- Concurrent uploads competing for the same workspace budget. +- Reservation cleanup after forced worker termination. +- Stale reservation detection after long-running or killed jobs. + +This area needs a dedicated failure test before the ingestion system can be considered robust under concurrency. + +## Observability Findings + +What is working: + +- The UI shows ready, failed, and in-progress document groups. +- The API exposes `GET /documents/ingestion-queues`. +- The API exposes `GET /documents/ingestion-runs/{run_id}`. +- Document-level error messages are clear. +- Docker logs are useful for worker-level diagnosis. + +What is weak: + +- Run status history is not fully reliable for older runs. +- Expected validation failures create RQ failure noise. +- There is no per-stage timing model. +- There is no queue-depth trend, worker utilization view, or OpenAI latency breakdown. +- There is no simple operator command/API that says "show stuck ingestion work". + +## Multi-Document Readiness + +The ingestion side now supports multiple uploaded documents and preserves workspace/document metadata in pages, chunks, and embeddings. + +This helps future multi-document retrieval because the data model already carries: + +- `workspace_id` +- `document_id` +- page rows +- chunk rows +- embedding rows + +The main remaining multi-document gap is on the query side, not ingestion. Retrieval and chat are still effectively scoped to one selected document at a time. + +## What Is Missing + +Important missing pieces: + +- Larger load tests with 10, 25, and 50 small PDFs. +- Forced worker-kill test after the new failure callback. +- Low token-budget ingestion test. +- Reindex test from the UI/API. +- Automated run-status reconciliation. +- Cleaner RQ failure semantics for expected document rejection. +- Durable stage timing fields or metrics. +- Queue-depth metrics over time. +- Integration tests for batch upload, partial success, retry, and stuck-job recovery. + +## Recommended Next Improvements + +Highest impact for reliability: + +1. Add a reconciliation job that marks stale `uploaded`, `extracting`, and `indexing` documents failed after a configured timeout. +2. Update ingestion-run status from document states consistently, or derive it dynamically instead of trusting stale stored statuses. +3. Separate expected validation failures from infrastructure failures so RQ failed registries are actionable. +4. Add an automated test for RQ worker timeout/kill recovery. + +Highest impact for throughput: + +1. Add per-stage timing metrics for upload completion, queue wait, extraction, indexing, embedding, and DB writes. +2. Run repeatable 10/25/50 PDF load tests with small text PDFs. +3. Record embedding batch counts and OpenAI latency per document. +4. Consider bulk DB insert optimization only after measuring real DB bottlenecks. + +Highest impact for observability: + +1. Add a `GET /documents/ingestion-health` endpoint with stuck counts, queue depths, failed registry counts, and oldest active document age. +2. Add run-level timing and final status summary fields. +3. Show infrastructure failures separately from validation failures in the UI. +4. Add an operator command or script to inspect stuck documents and stale reservations. + +Future multi-document retrieval work: + +1. Add selected-document or workspace-wide retrieval APIs. +2. Update retrieval filters to search across multiple document IDs. +3. Update the chat UI so a user can select multiple ready documents. +4. Preserve citation quality by returning document title, page, and chunk metadata for every source. + +## Recommended Next Tasks + +Concrete next task files to create: + +1. `docs/task/ingestion-observability-and-stuck-job-recovery.md` +2. `docs/task/ingestion-load-test-suite.md` +3. `docs/task/token-budget-ingestion-failure-tests.md` +4. `docs/task/multi-document-retrieval-api.md` +5. `docs/task/multi-document-chat-ui.md` + +## Final Assessment + +Task 4 validates that the improved ingestion workflow works for the current real local test set. The latest batch ended in correct document states, queues drained, and valid PDFs became queryable. + +The pipeline should not yet be called large-scale ready. The next engineering focus should be observability, stale-state recovery, and repeatable load testing before adding more ingestion features. diff --git a/docs/task/ingestion-follow-up-hardening-and-validation.md b/docs/task/ingestion-follow-up-hardening-and-validation.md new file mode 100644 index 0000000..bcfa46f --- /dev/null +++ b/docs/task/ingestion-follow-up-hardening-and-validation.md @@ -0,0 +1,294 @@ +# Task: Ingestion Follow-Up Hardening and Validation + +## Status + +Derived from the current ingestion workflow audit and the existing evaluation report. + +## Goal + +Capture the important ingestion gaps that are still not implemented or still weak, then define: + +1. how each gap should be implemented +2. what should be tested after implementation + +This file is meant to turn the audit into concrete follow-up work. + +## Why This Task Exists + +The repository already has a working ingestion pipeline for small PDFs, including: + +- batch upload prepare and complete +- ingestion runs +- extraction and indexing workers +- structured failure messages +- retry and reindex flows +- queue visibility +- token-budget reservation and cleanup + +But the system is still missing several important pieces before the ingestion path can be considered robust, measurable, and ready for higher-volume usage. + +## Scope + +This task covers the remaining weak points found in the audit: + +1. no repeatable load and throughput validation +2. weak ingestion observability +3. incomplete stale-state and run lifecycle recovery +4. schema and runtime status mismatch +5. query path still single-document +6. thin automated ingestion test coverage + +## Recommended Order + +1. Fix schema and status-model mismatch +2. Add stale document and run reconciliation +3. Add ingestion observability and operator health visibility +4. Add automated ingestion integration tests +5. Add repeatable load-test tooling +6. Extend retrieval and querying toward multi-document support + +The first three items are backend safety and operability work. They should happen before calling the ingestion system stable at larger scale. + +## Follow-Up Items + +### 1. Repeatable Load and Throughput Validation Is Missing + +#### Current Weakness + +The code supports batch ingestion, but there is no repeatable load-test harness for: + +- 10 PDFs +- 25 PDFs +- 50 PDFs +- 100 PDFs + +There is also no durable way to compare one ingestion revision against another. + +#### How To Implement + +- Add a dedicated load-test task doc and script set under `scripts/` for ingestion benchmarking. +- Create a repeatable test dataset of small text PDFs with known page counts and known expected outcomes. +- Add a driver that can: + - create or reuse a workspace + - prepare and complete uploads in bulk + - poll ingestion runs until terminal + - capture final document states + - measure total wall-clock run time + - measure per-document completion time +- Persist results to a simple JSON or markdown artifact so runs can be compared over time. +- Include scenarios with: + - all-valid batches + - mixed valid and invalid documents + - queue backlog conditions + +#### What To Test After Implementation + +- 10 valid small PDFs complete successfully. +- 25 valid small PDFs complete successfully. +- 50 valid small PDFs complete successfully. +- Mixed valid and invalid batches still complete with isolated failures. +- The load-test artifact records: + - total duration + - per-document duration + - success count + - failure count + - queue depth snapshots if available +- Re-running the same scenario produces comparable output structure. + +### 2. Ingestion Observability Is Still Weak + +#### Current Weakness + +The system exposes queue counts and document/run status, but it still lacks: + +- stage timing +- queue-depth history +- worker utilization visibility +- operator-facing ingestion health summary +- clear separation between expected document rejection and infrastructure failure + +#### How To Implement + +- Add stage timestamps or timing fields for key transitions: + - upload completed + - extract job enqueued + - extraction started + - extraction finished + - index job enqueued + - indexing started + - indexing finished +- Add an operator endpoint such as `GET /documents/ingestion-health` that returns: + - current queue counts + - failed registry counts + - oldest active document age + - stale active document count + - active ingestion run count + - stale reservation count if practical +- Add structured worker logs for: + - queue wait start/end + - extraction duration + - indexing duration + - embedding batch counts + - embedding API latency if practical +- Track expected validation failures separately from infrastructure failures in logs or API summaries. + +#### What To Test After Implementation + +- A successful document exposes complete stage timing. +- A failed validation document is visible as a document failure but not mislabeled as infrastructure failure. +- The health endpoint reports queue and stale-state summaries correctly during active ingestion. +- Worker logs contain enough information to distinguish: + - queue delay + - extraction time + - indexing time + - embedding time +- The UI or operator flow can identify whether a slowdown is caused by queue wait or processing. + +### 3. Stale-State and Run Lifecycle Recovery Is Incomplete + +#### Current Weakness + +There is failure callback coverage for killed jobs, but there is still no general reconciliation process for: + +- documents stuck in `uploaded` +- documents stuck in `extracting` +- documents stuck in `indexing` +- ingestion runs left in `preparing` or `processing` after work is effectively done + +#### How To Implement + +- Add a scheduled reconciliation job for ingestion state. +- Define configurable stale thresholds for: + - `uploaded` + - `extracting` + - `indexing` +- Reconciliation should: + - mark stale active documents as `failed` + - attach a structured transient-infrastructure error message + - refresh or recompute affected ingestion run statuses +- Consider deriving ingestion run status dynamically from current document states instead of trusting persisted run status alone. +- Ensure reconciliation is idempotent and safe to run frequently. + +#### What To Test After Implementation + +- A document intentionally left in `uploaded` beyond threshold is marked `failed`. +- A document intentionally left in `extracting` beyond threshold is marked `failed`. +- A document intentionally left in `indexing` beyond threshold is marked `failed`. +- A run with only terminal document states is reconciled to `completed`, `partial`, or `failed` correctly. +- Re-running reconciliation does not corrupt already-terminal rows. +- Retry still works correctly after reconciliation marks a document failed. + +### 4. Schema and Runtime Status Model Do Not Fully Match + +#### Current Weakness + +The code uses a richer runtime model including statuses like: + +- `uploading` +- `extracting` +- `indexed` + +But the SQL schema files still define a narrower `documents.status` check constraint. The code currently works around this mismatch defensively, which is fragile. + +#### How To Implement + +- Update `scripts/schema.local.sql` and `scripts/schema.supabase.sql` so the document status constraint matches the actual runtime statuses used by the application and workers. +- Review migrations or schema upgrade steps so existing environments can be brought to the same status model safely. +- Remove compatibility branches only after the schema is consistently aligned across environments. +- Ensure frontend status types, backend schemas, workers, and DB constraints all use the same status vocabulary. + +#### What To Test After Implementation + +- Fresh schema setup accepts every runtime status the code uses. +- Existing environments can migrate without data loss. +- Worker transitions no longer rely on fallback behavior for missing statuses. +- Upload, extract, index, retry, and reindex transitions all succeed against the aligned schema. +- Status values shown in the UI match the persisted DB states exactly. + +### 5. Retrieval and Querying Are Still Single-Document + +#### Current Weakness + +The ingestion side now supports many documents per workspace, but the query side still accepts one `document_id` and retrieves chunks from only one document at a time. + +This is not an ingestion breakage, but it is a real readiness gap for the intended multi-document system. + +#### How To Implement + +- Add a new query contract that accepts multiple selected document IDs. +- Enforce explicit query-side limits for: + - max selected documents + - max total pages across selected documents +- Update retrieval SQL to search across a selected document set instead of one document. +- Preserve citation quality by returning document metadata together with page and chunk references. +- Update query logs and observability so they record all searched documents, not one. + +#### What To Test After Implementation + +- A query across two ready documents returns citations from both when relevant. +- Query limits reject oversized document selections safely. +- Workspace isolation still holds for multi-document queries. +- Query logs record the selected document set accurately. +- Existing single-document flows continue to work. + +### 6. Automated Ingestion Test Coverage Is Too Thin + +#### Current Weakness + +The repository has helper-level tests for ingestion policy, ingestion run status derivation, and token-budget primitives, but it does not have strong automated coverage for: + +- batch upload orchestration +- partial success batches +- retry flows +- reindex flows +- stuck-job recovery +- low-budget ingestion failure +- queue failure callback behavior + +#### How To Implement + +- Add integration-style tests around the documents API and worker jobs. +- Use fixtures or mocks for: + - Supabase storage interactions + - OpenAI embeddings + - Redis/RQ enqueue behavior where full worker execution is not needed +- Add job-level tests for: + - extraction success + - page-limit failure + - unsupported-content failure + - indexing success + - budget failure + - cleanup after failure +- Add API tests for: + - batch prepare + - batch complete + - retry allowed vs retry denied + - reindex with and without existing pages + +#### What To Test After Implementation + +- Batch prepare returns mixed accepted and rejected results correctly. +- Batch complete enqueues extraction correctly and fails safely on missing objects. +- Retry only works for retryable failures. +- Reindex chooses extract or index path correctly based on existing page rows. +- Index budget exhaustion leaves no stranded reservations. +- Failure callback marks timed-out or killed jobs as failed. +- Cleanup removes stale chunks, embeddings, and pages when expected. + +## Acceptance Criteria + +This task should be considered complete only when: + +- the remaining ingestion weak points are each turned into explicit implementation work +- every item has a clear post-implementation validation plan +- the schema matches the runtime status model +- stale-state recovery exists for documents and runs +- ingestion observability is materially better +- automated ingestion coverage is broader than helper-only unit tests +- repeatable load validation exists for larger batches + +## Summary + +The ingestion pipeline is no longer missing its basic architecture. The important remaining work is hardening, observability, lifecycle recovery, automated validation, and true multi-document readiness. + +This task captures that follow-up work in a way that is concrete enough to implement and test. diff --git a/scripts/schema.local.sql b/scripts/schema.local.sql index e4eb313..9d63331 100644 --- a/scripts/schema.local.sql +++ b/scripts/schema.local.sql @@ -11,6 +11,30 @@ CREATE TABLE IF NOT EXISTS workspaces ( CREATE INDEX IF NOT EXISTS idx_workspaces_owner ON workspaces(owner_id); +CREATE TABLE IF NOT EXISTS ingestion_runs ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + workspace_id UUID NOT NULL REFERENCES workspaces(id) ON DELETE CASCADE, + name TEXT, + status TEXT NOT NULL DEFAULT 'preparing', + total_documents INT NOT NULL DEFAULT 0, + accepted_documents INT NOT NULL DEFAULT 0, + rejected_documents INT NOT NULL DEFAULT 0, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + + CONSTRAINT chk_ingestion_runs_status CHECK (status IN ('preparing', 'processing', 'completed', 'partial', 'failed')), + CONSTRAINT chk_ingestion_runs_counts CHECK ( + total_documents >= 0 + AND accepted_documents >= 0 + AND rejected_documents >= 0 + ) +); + +CREATE INDEX IF NOT EXISTS idx_ingestion_runs_workspace_created + ON ingestion_runs(workspace_id, created_at DESC); +CREATE INDEX IF NOT EXISTS idx_ingestion_runs_workspace_status + ON ingestion_runs(workspace_id, status); + CREATE TABLE IF NOT EXISTS documents ( id UUID PRIMARY KEY DEFAULT gen_random_uuid(), workspace_id UUID NOT NULL REFERENCES workspaces(id) ON DELETE CASCADE, @@ -19,12 +43,13 @@ CREATE TABLE IF NOT EXISTS documents ( page_count INT, file_hash_sha256 TEXT NOT NULL, storage_path TEXT NOT NULL, + ingestion_run_id UUID, status TEXT NOT NULL DEFAULT 'pending_upload', error_message TEXT, created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), - CONSTRAINT chk_file_size CHECK (file_size_bytes > 0 AND file_size_bytes <= 20971520), + CONSTRAINT chk_file_size CHECK (file_size_bytes > 0 AND file_size_bytes <= 10485760), CONSTRAINT chk_page_count CHECK (page_count IS NULL OR (page_count > 0 AND page_count <= 10)), CONSTRAINT chk_status CHECK (status IN ('pending_upload', 'uploaded', 'indexing', 'ready', 'failed')) ); @@ -34,6 +59,23 @@ CREATE UNIQUE INDEX IF NOT EXISTS idx_documents_workspace_hash CREATE INDEX IF NOT EXISTS idx_documents_workspace ON documents(workspace_id); CREATE INDEX IF NOT EXISTS idx_documents_workspace_status ON documents(workspace_id, status); +ALTER TABLE documents + ADD COLUMN IF NOT EXISTS ingestion_run_id UUID; + +DO $$ +BEGIN + ALTER TABLE documents + ADD CONSTRAINT fk_documents_ingestion_run + FOREIGN KEY (ingestion_run_id) + REFERENCES ingestion_runs(id) + ON DELETE SET NULL; +EXCEPTION + WHEN duplicate_object THEN NULL; +END $$; + +CREATE INDEX IF NOT EXISTS idx_documents_ingestion_run + ON documents(workspace_id, ingestion_run_id); + CREATE TABLE IF NOT EXISTS document_pages ( id BIGSERIAL PRIMARY KEY, workspace_id UUID NOT NULL REFERENCES workspaces(id) ON DELETE CASCADE, diff --git a/scripts/schema.supabase.sql b/scripts/schema.supabase.sql index e4eb313..9d63331 100644 --- a/scripts/schema.supabase.sql +++ b/scripts/schema.supabase.sql @@ -11,6 +11,30 @@ CREATE TABLE IF NOT EXISTS workspaces ( CREATE INDEX IF NOT EXISTS idx_workspaces_owner ON workspaces(owner_id); +CREATE TABLE IF NOT EXISTS ingestion_runs ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + workspace_id UUID NOT NULL REFERENCES workspaces(id) ON DELETE CASCADE, + name TEXT, + status TEXT NOT NULL DEFAULT 'preparing', + total_documents INT NOT NULL DEFAULT 0, + accepted_documents INT NOT NULL DEFAULT 0, + rejected_documents INT NOT NULL DEFAULT 0, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + + CONSTRAINT chk_ingestion_runs_status CHECK (status IN ('preparing', 'processing', 'completed', 'partial', 'failed')), + CONSTRAINT chk_ingestion_runs_counts CHECK ( + total_documents >= 0 + AND accepted_documents >= 0 + AND rejected_documents >= 0 + ) +); + +CREATE INDEX IF NOT EXISTS idx_ingestion_runs_workspace_created + ON ingestion_runs(workspace_id, created_at DESC); +CREATE INDEX IF NOT EXISTS idx_ingestion_runs_workspace_status + ON ingestion_runs(workspace_id, status); + CREATE TABLE IF NOT EXISTS documents ( id UUID PRIMARY KEY DEFAULT gen_random_uuid(), workspace_id UUID NOT NULL REFERENCES workspaces(id) ON DELETE CASCADE, @@ -19,12 +43,13 @@ CREATE TABLE IF NOT EXISTS documents ( page_count INT, file_hash_sha256 TEXT NOT NULL, storage_path TEXT NOT NULL, + ingestion_run_id UUID, status TEXT NOT NULL DEFAULT 'pending_upload', error_message TEXT, created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), - CONSTRAINT chk_file_size CHECK (file_size_bytes > 0 AND file_size_bytes <= 20971520), + CONSTRAINT chk_file_size CHECK (file_size_bytes > 0 AND file_size_bytes <= 10485760), CONSTRAINT chk_page_count CHECK (page_count IS NULL OR (page_count > 0 AND page_count <= 10)), CONSTRAINT chk_status CHECK (status IN ('pending_upload', 'uploaded', 'indexing', 'ready', 'failed')) ); @@ -34,6 +59,23 @@ CREATE UNIQUE INDEX IF NOT EXISTS idx_documents_workspace_hash CREATE INDEX IF NOT EXISTS idx_documents_workspace ON documents(workspace_id); CREATE INDEX IF NOT EXISTS idx_documents_workspace_status ON documents(workspace_id, status); +ALTER TABLE documents + ADD COLUMN IF NOT EXISTS ingestion_run_id UUID; + +DO $$ +BEGIN + ALTER TABLE documents + ADD CONSTRAINT fk_documents_ingestion_run + FOREIGN KEY (ingestion_run_id) + REFERENCES ingestion_runs(id) + ON DELETE SET NULL; +EXCEPTION + WHEN duplicate_object THEN NULL; +END $$; + +CREATE INDEX IF NOT EXISTS idx_documents_ingestion_run + ON documents(workspace_id, ingestion_run_id); + CREATE TABLE IF NOT EXISTS document_pages ( id BIGSERIAL PRIMARY KEY, workspace_id UUID NOT NULL REFERENCES workspaces(id) ON DELETE CASCADE, diff --git a/server/app/api/documents.py b/server/app/api/documents.py index 335fd11..d088006 100644 --- a/server/app/api/documents.py +++ b/server/app/api/documents.py @@ -9,11 +9,28 @@ from fastapi import APIRouter, Depends, HTTPException, Query, Response, status from redis import Redis from rq import Queue +from rq.job import Callback +from rq.registry import ( + DeferredJobRegistry, + FailedJobRegistry, + ScheduledJobRegistry, + StartedJobRegistry, +) from sqlalchemy import func, inspect, select, text from sqlalchemy.orm import Session from app.api.deps import get_workspace_id from app.config import settings +from app.core.ingestion_runs import ( + derive_ingestion_run_status, + empty_document_status_counts, +) +from app.core.ingestion_policy import ( + IngestionFailureCategory, + format_bytes, + ingestion_failure_message, + is_retryable_failure, +) from app.core.storage import delete_object, generate_signed_upload_url, object_exists from app.db.models import Document from app.db.session import get_db @@ -24,15 +41,33 @@ DocumentListItem, DocumentListResponse, DocumentProgress, + IngestionQueueStatusItem, + IngestionQueueStatusResponse, + IngestionRunResponse, + IngestionRunStatusCounts, + UploadCompleteBatchItem, + UploadCompleteBatchRequest, + UploadCompleteBatchResponse, + UploadPrepareBatchItem, + UploadPrepareBatchRequest, + UploadPrepareBatchResponse, UploadCompleteRequest, UploadCompleteResponse, UploadPrepareRequest, UploadPrepareResponse, ) -from app.utils.rate_limit import QUERY_RATE_LIMIT, UPLOAD_COMPLETE_RATE_LIMIT, UPLOAD_PREPARE_RATE_LIMIT, enforce_workspace_rate_limit +from app.utils.rate_limit import ( + BULK_UPLOAD_COMPLETE_RATE_LIMIT, + BULK_UPLOAD_PREPARE_RATE_LIMIT, + QUERY_RATE_LIMIT, + UPLOAD_COMPLETE_RATE_LIMIT, + UPLOAD_PREPARE_RATE_LIMIT, + enforce_workspace_rate_limit, +) router = APIRouter() logger = logging.getLogger(__name__) +INGESTION_FAILURE_CALLBACK = "jobs.ingest_callbacks.mark_ingestion_job_failed" ALLOWED_STATUS_FILTERS = { "pending_upload", "uploading", @@ -56,18 +91,22 @@ def _find_existing_prepare_by_idempotency( workspace_id: uuid.UUID, idempotency_hash: str, ) -> dict | None: - row = db.execute( - text( - """ + row = ( + db.execute( + text( + """ SELECT id, filename, storage_path, status FROM documents WHERE workspace_id = :workspace_id AND file_hash_sha256 = :idempotency_hash LIMIT 1 """ - ), - {"workspace_id": workspace_id, "idempotency_hash": idempotency_hash}, - ).mappings().first() + ), + {"workspace_id": workspace_id, "idempotency_hash": idempotency_hash}, + ) + .mappings() + .first() + ) if row is None: return None return { @@ -81,12 +120,16 @@ def _find_existing_prepare_by_idempotency( def _sanitize_filename(filename: str) -> str: name = Path(filename).name.strip() if not name: - raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid filename") + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, detail=_validation_error("Invalid filename.") + ) sanitized = re.sub(r"[^A-Za-z0-9._-]", "_", name) sanitized = re.sub(r"_+", "_", sanitized).strip("_") if not sanitized: - raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid filename") + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, detail=_validation_error("Invalid filename.") + ) return sanitized @@ -95,7 +138,21 @@ def _document_columns(db: Session) -> set[str]: return {col["name"] for col in inspect(bind).get_columns("documents")} -def _count_for_document(db: Session, sql: str, workspace_id: uuid.UUID, document_id: uuid.UUID) -> int: +def _table_exists(db: Session, table_name: str) -> bool: + return inspect(db.get_bind()).has_table(table_name) + + +def _require_ingestion_run_schema(db: Session) -> None: + if not _table_exists(db, "ingestion_runs") or "ingestion_run_id" not in _document_columns(db): + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail="Ingestion-run schema is not initialized. Apply scripts/schema.local.sql or scripts/schema.supabase.sql.", + ) + + +def _count_for_document( + db: Session, sql: str, workspace_id: uuid.UUID, document_id: uuid.UUID +) -> int: value = db.execute( text(sql), { @@ -114,7 +171,199 @@ def _trim_text(value: str, max_chars: int) -> str: return f"{value[: max_chars - 3].rstrip()}..." -def _enqueue_extract(*, workspace_id: uuid.UUID, document_id: uuid.UUID, bucket: str, storage_path: str): +def _validation_error(detail: str) -> str: + return ingestion_failure_message(IngestionFailureCategory.VALIDATION, detail) + + +def _set_document_failed( + *, + db: Session, + workspace_id: uuid.UUID, + document_id: uuid.UUID, + category: IngestionFailureCategory, + detail: str, +) -> str: + error_message = ingestion_failure_message(category, detail) + db.execute( + text( + """ + UPDATE documents + SET status = 'failed', + error_message = :error_message, + updated_at = :updated_at + WHERE id = :document_id + AND workspace_id = :workspace_id + """ + ), + { + "workspace_id": workspace_id, + "document_id": document_id, + "error_message": error_message, + "updated_at": datetime.now(UTC), + }, + ) + return error_message + + +def _validate_upload_file( + *, filename: str, content_type: str, file_size_bytes: int +) -> tuple[str | None, str | None]: + if file_size_bytes > settings.MAX_FILE_SIZE_BYTES: + return ( + None, + _validation_error( + f"PDF file size must be {format_bytes(settings.MAX_FILE_SIZE_BYTES)} or smaller." + ), + ) + + allowed_content_types = { + allowed_content_type.lower() for allowed_content_type in settings.ALLOWED_CONTENT_TYPES + } + if content_type.lower() not in allowed_content_types: + return ( + None, + _validation_error( + f"Only PDF uploads are supported. Allowed content types: {sorted(allowed_content_types)}." + ), + ) + + try: + sanitized_filename = _sanitize_filename(filename) + except HTTPException as exc: + return None, str(exc.detail) + + if not sanitized_filename.lower().endswith(".pdf"): + return None, _validation_error("Only files with a .pdf extension are supported.") + + return sanitized_filename, None + + +def _create_ingestion_run( + *, + db: Session, + workspace_id: uuid.UUID, + run_id: uuid.UUID, + name: str | None, + total_documents: int, + accepted_documents: int, + rejected_documents: int, +) -> None: + run_status = "processing" if accepted_documents else "failed" + db.execute( + text( + """ + INSERT INTO ingestion_runs ( + id, workspace_id, name, status, total_documents, + accepted_documents, rejected_documents, created_at, updated_at + ) + VALUES ( + :id, :workspace_id, :name, :status, :total_documents, + :accepted_documents, :rejected_documents, :created_at, :updated_at + ) + """ + ), + { + "id": run_id, + "workspace_id": workspace_id, + "name": name, + "status": run_status, + "total_documents": total_documents, + "accepted_documents": accepted_documents, + "rejected_documents": rejected_documents, + "created_at": datetime.now(UTC), + "updated_at": datetime.now(UTC), + }, + ) + + +def _document_status_counts_for_run( + *, db: Session, workspace_id: uuid.UUID, run_id: uuid.UUID +) -> dict[str, int]: + counts = empty_document_status_counts() + rows = ( + db.execute( + text( + """ + SELECT status, COUNT(*) AS count + FROM documents + WHERE workspace_id = :workspace_id + AND ingestion_run_id = :run_id + GROUP BY status + """ + ), + {"workspace_id": workspace_id, "run_id": run_id}, + ) + .mappings() + .all() + ) + for row in rows: + status_name = str(row["status"]) + counts[status_name] = int(row["count"] or 0) + counts["total"] = sum(counts.values()) + return counts + + +def _refresh_ingestion_run_status( + *, db: Session, workspace_id: uuid.UUID, run_id: uuid.UUID +) -> str: + run_row = ( + db.execute( + text( + """ + SELECT accepted_documents, rejected_documents + FROM ingestion_runs + WHERE id = :run_id + AND workspace_id = :workspace_id + LIMIT 1 + """ + ), + {"workspace_id": workspace_id, "run_id": run_id}, + ) + .mappings() + .first() + ) + if run_row is None: + raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Ingestion run not found") + + status_counts = _document_status_counts_for_run(db=db, workspace_id=workspace_id, run_id=run_id) + run_status = derive_ingestion_run_status( + status_counts=status_counts, + accepted_documents=int(run_row["accepted_documents"] or 0), + rejected_documents=int(run_row["rejected_documents"] or 0), + ) + db.execute( + text( + """ + UPDATE ingestion_runs + SET status = :status, + updated_at = :updated_at + WHERE id = :run_id + AND workspace_id = :workspace_id + """ + ), + { + "status": run_status, + "updated_at": datetime.now(UTC), + "workspace_id": workspace_id, + "run_id": run_id, + }, + ) + return run_status + + +def _registry_count(registry) -> int: + count_value = registry.count + return int(count_value() if callable(count_value) else count_value) + + +def _enqueue_extract( + *, + workspace_id: uuid.UUID, + document_id: uuid.UUID, + bucket: str, + storage_path: str, + ingestion_run_id: uuid.UUID | None = None, +): redis_conn = Redis.from_url(settings.REDIS_URL) queue = Queue("ingest_extract", connection=redis_conn) return queue.enqueue( @@ -123,16 +372,34 @@ def _enqueue_extract(*, workspace_id: uuid.UUID, document_id: uuid.UUID, bucket: document_id=str(document_id), bucket=bucket, storage_path=storage_path, + ingestion_run_id=str(ingestion_run_id) if ingestion_run_id else None, + job_timeout=settings.INGEST_EXTRACT_JOB_TIMEOUT_SECONDS, + on_failure=Callback(INGESTION_FAILURE_CALLBACK, timeout=60), + meta={ + "workspace_id": str(workspace_id), + "document_id": str(document_id), + "ingestion_run_id": str(ingestion_run_id) if ingestion_run_id else None, + }, ) -def _enqueue_index(*, workspace_id: uuid.UUID, document_id: uuid.UUID): +def _enqueue_index( + *, workspace_id: uuid.UUID, document_id: uuid.UUID, ingestion_run_id: uuid.UUID | None = None +): redis_conn = Redis.from_url(settings.REDIS_URL) queue = Queue("ingest_index", connection=redis_conn) return queue.enqueue( "jobs.ingest_index.ingest_index", workspace_id=str(workspace_id), document_id=str(document_id), + ingestion_run_id=str(ingestion_run_id) if ingestion_run_id else None, + job_timeout=settings.INGEST_INDEX_JOB_TIMEOUT_SECONDS, + on_failure=Callback(INGESTION_FAILURE_CALLBACK, timeout=60), + meta={ + "workspace_id": str(workspace_id), + "document_id": str(document_id), + "ingestion_run_id": str(ingestion_run_id) if ingestion_run_id else None, + }, ) @@ -145,7 +412,9 @@ def list_documents( db: Session = Depends(get_db), ) -> DocumentListResponse: if limit < 1 or limit > 100: - raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="limit must be between 1 and 100") + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, detail="limit must be between 1 and 100" + ) if offset < 0: raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="offset must be >= 0") @@ -157,6 +426,13 @@ def list_documents( columns = _document_columns(db) has_content_type = "content_type" in columns + page_count_select = "page_count" if "page_count" in columns else "NULL AS page_count" + ingestion_run_id_select = ( + "ingestion_run_id" if "ingestion_run_id" in columns else "NULL AS ingestion_run_id" + ) + error_message_select = ( + "error_message" if "error_message" in columns else "NULL AS error_message" + ) where_clauses = ["workspace_id = :workspace_id"] params: dict[str, object] = {"workspace_id": workspace_id, "limit": limit, "offset": offset} @@ -168,10 +444,13 @@ def list_documents( total_sql = text(f"SELECT COUNT(*) FROM documents WHERE {where_sql}") total = int(db.execute(total_sql, params).scalar_one() or 0) - content_type_select = "content_type" if has_content_type else "'application/pdf' AS content_type" + content_type_select = ( + "content_type" if has_content_type else "'application/pdf' AS content_type" + ) list_sql = text( f""" - SELECT id, filename, {content_type_select}, file_size_bytes, status, created_at, updated_at + SELECT id, filename, {content_type_select}, file_size_bytes, {page_count_select}, + {ingestion_run_id_select}, status, {error_message_select}, created_at, updated_at FROM documents WHERE {where_sql} ORDER BY created_at DESC @@ -186,7 +465,10 @@ def list_documents( filename=row["filename"], content_type=row["content_type"], file_size_bytes=int(row["file_size_bytes"]), + page_count=row["page_count"], + ingestion_run_id=row["ingestion_run_id"], status=row["status"], + error_message=row["error_message"], created_at=row["created_at"], updated_at=row["updated_at"], ) @@ -196,6 +478,79 @@ def list_documents( return DocumentListResponse(items=items, limit=limit, offset=offset, total=total) +@router.get("/ingestion-queues", response_model=IngestionQueueStatusResponse) +def get_ingestion_queues( + workspace_id: uuid.UUID = Depends(get_workspace_id), +) -> IngestionQueueStatusResponse: + redis_conn = Redis.from_url(settings.REDIS_URL) + queue_items: list[IngestionQueueStatusItem] = [] + for queue_name in ("ingest_extract", "ingest_index"): + queue = Queue(queue_name, connection=redis_conn) + queue_items.append( + IngestionQueueStatusItem( + name=queue_name, + queued_count=len(queue), + started_count=_registry_count( + StartedJobRegistry(queue_name, connection=redis_conn) + ), + deferred_count=_registry_count( + DeferredJobRegistry(queue_name, connection=redis_conn) + ), + scheduled_count=_registry_count( + ScheduledJobRegistry(queue_name, connection=redis_conn) + ), + failed_count=_registry_count(FailedJobRegistry(queue_name, connection=redis_conn)), + ) + ) + + # Keep the endpoint workspace-authenticated even though queue depth is global. + _ = workspace_id + return IngestionQueueStatusResponse(queues=queue_items) + + +@router.get("/ingestion-runs/{run_id}", response_model=IngestionRunResponse) +def get_ingestion_run( + run_id: uuid.UUID, + workspace_id: uuid.UUID = Depends(get_workspace_id), + db: Session = Depends(get_db), +) -> IngestionRunResponse: + _require_ingestion_run_schema(db) + run_status = _refresh_ingestion_run_status(db=db, workspace_id=workspace_id, run_id=run_id) + run_row = ( + db.execute( + text( + """ + SELECT id, name, total_documents, accepted_documents, rejected_documents, + created_at, updated_at + FROM ingestion_runs + WHERE id = :run_id + AND workspace_id = :workspace_id + LIMIT 1 + """ + ), + {"workspace_id": workspace_id, "run_id": run_id}, + ) + .mappings() + .first() + ) + if run_row is None: + raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Ingestion run not found") + status_counts = _document_status_counts_for_run(db=db, workspace_id=workspace_id, run_id=run_id) + db.commit() + + return IngestionRunResponse( + id=run_row["id"], + name=run_row["name"], + status=run_status, + total_documents=int(run_row["total_documents"] or 0), + accepted_documents=int(run_row["accepted_documents"] or 0), + rejected_documents=int(run_row["rejected_documents"] or 0), + document_statuses=IngestionRunStatusCounts(**status_counts), + created_at=run_row["created_at"], + updated_at=run_row["updated_at"], + ) + + @router.get("/{document_id}", response_model=DocumentDetailResponse) def get_document( document_id: uuid.UUID, @@ -212,6 +567,10 @@ def get_document( "created_at", "updated_at", ] + if "error_message" in columns: + select_fields.append("error_message") + else: + select_fields.append("NULL AS error_message") if "content_type" in columns: select_fields.append("content_type") else: @@ -220,6 +579,10 @@ def get_document( select_fields.append("storage_bucket") else: select_fields.append("NULL AS storage_bucket") + if "ingestion_run_id" in columns: + select_fields.append("ingestion_run_id") + else: + select_fields.append("NULL AS ingestion_run_id") if "pages_total" in columns: select_fields.append("pages_total") else: @@ -234,10 +597,14 @@ def get_document( LIMIT 1 """ ) - row = db.execute( - detail_sql, - {"document_id": document_id, "workspace_id": workspace_id}, - ).mappings().first() + row = ( + db.execute( + detail_sql, + {"document_id": document_id, "workspace_id": workspace_id}, + ) + .mappings() + .first() + ) if row is None: raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Document not found") @@ -282,7 +649,9 @@ def get_document( filename=row["filename"], content_type=row["content_type"], file_size_bytes=int(row["file_size_bytes"]), + ingestion_run_id=row["ingestion_run_id"], status=row["status"], + error_message=row["error_message"], bucket=row["storage_bucket"] or settings.SUPABASE_STORAGE_BUCKET, storage_path=row["storage_path"], created_at=row["created_at"], @@ -305,7 +674,9 @@ def get_document_page( db: Session = Depends(get_db), ) -> DocumentPageSourceResponse: if page_number < 1: - raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="page_number must be >= 1") + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, detail="page_number must be >= 1" + ) enforce_workspace_rate_limit( workspace_id=workspace_id, @@ -313,9 +684,10 @@ def get_document_page( limit=QUERY_RATE_LIMIT, ) - row = db.execute( - text( - """ + row = ( + db.execute( + text( + """ SELECT content FROM document_pages WHERE workspace_id = :workspace_id @@ -323,13 +695,16 @@ def get_document_page( AND page_number = :page_number LIMIT 1 """ - ), - { - "workspace_id": workspace_id, - "document_id": document_id, - "page_number": page_number, - }, - ).mappings().first() + ), + { + "workspace_id": workspace_id, + "document_id": document_id, + "page_number": page_number, + }, + ) + .mappings() + .first() + ) if row is None: raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Page not found") @@ -340,7 +715,9 @@ def get_document_page( ) -@router.post("/upload-prepare", response_model=UploadPrepareResponse, status_code=status.HTTP_201_CREATED) +@router.post( + "/upload-prepare", response_model=UploadPrepareResponse, status_code=status.HTTP_201_CREATED +) def upload_prepare( payload: UploadPrepareRequest, workspace_id: uuid.UUID = Depends(get_workspace_id), @@ -352,30 +729,26 @@ def upload_prepare( limit=UPLOAD_PREPARE_RATE_LIMIT, ) - if payload.file_size_bytes > settings.MAX_FILE_SIZE_BYTES: - raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, - detail=f"File size exceeds limit ({settings.MAX_FILE_SIZE_BYTES} bytes)", - ) - - allowed_content_types = {content_type.lower() for content_type in settings.ALLOWED_CONTENT_TYPES} - if payload.content_type.lower() not in allowed_content_types: - raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, - detail=f"Unsupported content_type. Allowed: {sorted(allowed_content_types)}", - ) + sanitized_filename, validation_error = _validate_upload_file( + filename=payload.filename, + content_type=payload.content_type, + file_size_bytes=payload.file_size_bytes, + ) + if validation_error is not None or sanitized_filename is None: + raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=validation_error) count_stmt = select(func.count(Document.id)).where(Document.workspace_id == workspace_id) document_count = int(db.execute(count_stmt).scalar_one() or 0) if document_count >= settings.MAX_DOCUMENTS_PER_WORKSPACE: raise HTTPException( status_code=status.HTTP_409_CONFLICT, - detail="Workspace document limit reached", + detail=_validation_error( + f"Workspace document limit reached. Maximum documents per workspace: {settings.MAX_DOCUMENTS_PER_WORKSPACE}." + ), ) now = datetime.now(UTC) bucket = settings.SUPABASE_STORAGE_BUCKET - sanitized_filename = _sanitize_filename(payload.filename) columns = _document_columns(db) idempotency_hash: str | None = None @@ -481,41 +854,284 @@ def upload_prepare( ) -@router.post("/upload-complete", response_model=UploadCompleteResponse) -def upload_complete( - payload: UploadCompleteRequest, +@router.post( + "/upload-prepare-batch", + response_model=UploadPrepareBatchResponse, + status_code=status.HTTP_201_CREATED, +) +def upload_prepare_batch( + payload: UploadPrepareBatchRequest, workspace_id: uuid.UUID = Depends(get_workspace_id), db: Session = Depends(get_db), -) -> UploadCompleteResponse: +) -> UploadPrepareBatchResponse: enforce_workspace_rate_limit( workspace_id=workspace_id, - operation="documents_upload_complete", - limit=UPLOAD_COMPLETE_RATE_LIMIT, + operation="documents_upload_prepare_batch", + limit=BULK_UPLOAD_PREPARE_RATE_LIMIT, ) + _require_ingestion_run_schema(db) + + if not payload.files: + raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="files is required") + if len(payload.files) > settings.MAX_BULK_UPLOAD_FILES: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=f"Batch upload supports up to {settings.MAX_BULK_UPLOAD_FILES} files.", + ) + columns = _document_columns(db) + bucket = settings.SUPABASE_STORAGE_BUCKET + run_id = uuid.uuid4() + now = datetime.now(UTC) + document_count = int( + db.execute( + select(func.count(Document.id)).where(Document.workspace_id == workspace_id) + ).scalar_one() + or 0 + ) + remaining_new_document_slots = max(0, settings.MAX_DOCUMENTS_PER_WORKSPACE - document_count) + + items: list[UploadPrepareBatchItem] = [] + insert_rows: list[dict[str, object]] = [] + existing_document_ids: list[uuid.UUID] = [] + new_document_count = 0 + + for index, file_payload in enumerate(payload.files): + sanitized_filename, validation_error = _validate_upload_file( + filename=file_payload.filename, + content_type=file_payload.content_type, + file_size_bytes=file_payload.file_size_bytes, + ) + if validation_error is not None or sanitized_filename is None: + items.append( + UploadPrepareBatchItem( + index=index, + filename=file_payload.filename, + client_file_id=file_payload.client_file_id, + status="rejected", + error=validation_error, + ) + ) + continue + + idempotency_hash: str | None = None + if file_payload.idempotency_key and "file_hash_sha256" in columns: + idempotency_hash = _idempotency_hash(file_payload.idempotency_key) + existing = _find_existing_prepare_by_idempotency( + db=db, + workspace_id=workspace_id, + idempotency_hash=idempotency_hash, + ) + if existing is not None: + existing_status = str(existing["status"]) + if existing_status not in {"pending_upload", "uploading"}: + items.append( + UploadPrepareBatchItem( + index=index, + filename=sanitized_filename, + client_file_id=file_payload.client_file_id, + status="rejected", + error=f"Upload already prepared for this key (current status: {existing_status})", + ) + ) + continue + + upload_url = generate_signed_upload_url( + bucket=bucket, + path=str(existing["storage_path"]), + expires=settings.UPLOAD_URL_EXPIRES_SECONDS, + ) + existing_document_ids.append(existing["id"]) + items.append( + UploadPrepareBatchItem( + index=index, + filename=sanitized_filename, + client_file_id=file_payload.client_file_id, + status="prepared", + document_id=existing["id"], + bucket=bucket, + storage_path=str(existing["storage_path"]), + upload_url=upload_url, + expires_in=settings.UPLOAD_URL_EXPIRES_SECONDS, + ) + ) + continue + + if new_document_count >= remaining_new_document_slots: + items.append( + UploadPrepareBatchItem( + index=index, + filename=sanitized_filename, + client_file_id=file_payload.client_file_id, + status="rejected", + error=_validation_error( + "Workspace document limit reached. " + f"Maximum documents per workspace: {settings.MAX_DOCUMENTS_PER_WORKSPACE}." + ), + ) + ) + continue + + document_id = uuid.uuid4() + storage_path = f"{workspace_id}/{document_id}/{sanitized_filename}" + try: + upload_url = generate_signed_upload_url( + bucket=bucket, + path=storage_path, + expires=settings.UPLOAD_URL_EXPIRES_SECONDS, + ) + except ValueError as exc: + items.append( + UploadPrepareBatchItem( + index=index, + filename=sanitized_filename, + client_file_id=file_payload.client_file_id, + status="rejected", + error=str(exc), + ) + ) + continue + + insert_fields: dict[str, object] = { + "id": document_id, + "workspace_id": workspace_id, + "filename": sanitized_filename, + "file_size_bytes": file_payload.file_size_bytes, + "storage_path": storage_path, + "status": "pending_upload", + "ingestion_run_id": run_id, + "created_at": now, + "updated_at": now, + } + if "content_type" in columns: + insert_fields["content_type"] = file_payload.content_type + if "storage_bucket" in columns: + insert_fields["storage_bucket"] = bucket + if "file_hash_sha256" in columns: + insert_fields["file_hash_sha256"] = idempotency_hash or f"uploading:{document_id}" + insert_rows.append(insert_fields) + new_document_count += 1 + + items.append( + UploadPrepareBatchItem( + index=index, + filename=sanitized_filename, + client_file_id=file_payload.client_file_id, + status="prepared", + document_id=document_id, + bucket=bucket, + storage_path=storage_path, + upload_url=upload_url, + expires_in=settings.UPLOAD_URL_EXPIRES_SECONDS, + ) + ) + + accepted_count = sum(1 for item in items if item.status == "prepared") + rejected_count = len(items) - accepted_count + + try: + _create_ingestion_run( + db=db, + workspace_id=workspace_id, + run_id=run_id, + name=payload.name, + total_documents=len(payload.files), + accepted_documents=accepted_count, + rejected_documents=rejected_count, + ) + if insert_rows: + keys = list(insert_rows[0].keys()) + columns_sql = ", ".join(keys) + values_sql = ", ".join(f":{key}" for key in keys) + db.execute( + text(f"INSERT INTO documents ({columns_sql}) VALUES ({values_sql})"), + insert_rows, + ) + if existing_document_ids: + for existing_document_id in existing_document_ids: + db.execute( + text( + """ + UPDATE documents + SET ingestion_run_id = :run_id, + updated_at = :updated_at + WHERE workspace_id = :workspace_id + AND id = :document_id + """ + ), + { + "workspace_id": workspace_id, + "document_id": existing_document_id, + "run_id": run_id, + "updated_at": datetime.now(UTC), + }, + ) + db.commit() + except Exception as exc: # noqa: BLE001 + db.rollback() + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail=f"Failed to create batch upload placeholders: {exc}", + ) from exc + + return UploadPrepareBatchResponse( + ingestion_run_id=run_id, + bucket=bucket, + expires_in=settings.UPLOAD_URL_EXPIRES_SECONDS, + accepted_count=accepted_count, + rejected_count=rejected_count, + items=items, + ) + + +def _complete_upload_document( + *, + payload: UploadCompleteRequest, + workspace_id: uuid.UUID, + db: Session, + expected_run_id: uuid.UUID | None = None, +) -> UploadCompleteResponse: columns = _document_columns(db) select_fields = ["id", "status", "storage_path"] if "storage_bucket" in columns: select_fields.append("storage_bucket") else: select_fields.append("NULL AS storage_bucket") - - select_sql = text( - f""" - SELECT {", ".join(select_fields)} - FROM documents - WHERE id = :document_id - AND workspace_id = :workspace_id - LIMIT 1 - """ + if "error_message" in columns: + select_fields.append("error_message") + else: + select_fields.append("NULL AS error_message") + if "ingestion_run_id" in columns: + select_fields.append("ingestion_run_id") + else: + select_fields.append("NULL AS ingestion_run_id") + + row = ( + db.execute( + text( + f""" + SELECT {", ".join(select_fields)} + FROM documents + WHERE id = :document_id + AND workspace_id = :workspace_id + LIMIT 1 + """ + ), + {"document_id": payload.document_id, "workspace_id": workspace_id}, + ) + .mappings() + .first() ) - row = db.execute( - select_sql, - {"document_id": payload.document_id, "workspace_id": workspace_id}, - ).mappings().first() if row is None: raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Document not found") + row_run_id = row["ingestion_run_id"] + if expected_run_id is not None and str(row_run_id) != str(expected_run_id): + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="Document does not belong to the requested ingestion run", + ) + stored_bucket = row["storage_bucket"] or settings.SUPABASE_STORAGE_BUCKET if payload.bucket != stored_bucket or payload.storage_path != row["storage_path"]: raise HTTPException( @@ -536,23 +1152,32 @@ def upload_complete( exists = object_exists(bucket=payload.bucket, path=payload.storage_path) if exists is False: + error_message = _set_document_failed( + db=db, + workspace_id=workspace_id, + document_id=payload.document_id, + category=IngestionFailureCategory.UPLOAD_STORAGE, + detail="Uploaded object was not found in storage. Upload the PDF again before starting ingestion.", + ) + db.commit() raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, - detail="Uploaded object not found in storage", + detail=error_message, ) - update_sql = text( - """ - UPDATE documents - SET status = 'uploaded', - updated_at = :updated_at - WHERE id = :document_id - AND workspace_id = :workspace_id - AND status IN ('uploading', 'pending_upload') - """ - ) + update_fields = ["status = 'uploaded'", "updated_at = :updated_at"] + if "error_message" in columns: + update_fields.append("error_message = NULL") update_result = db.execute( - update_sql, + text( + f""" + UPDATE documents + SET {", ".join(update_fields)} + WHERE id = :document_id + AND workspace_id = :workspace_id + AND status IN ('uploading', 'pending_upload') + """ + ), { "updated_at": datetime.now(UTC), "document_id": payload.document_id, @@ -565,6 +1190,7 @@ def upload_complete( status_code=status.HTTP_409_CONFLICT, detail="Upload completion conflict", ) + db.commit() try: job = _enqueue_extract( @@ -572,16 +1198,22 @@ def upload_complete( document_id=payload.document_id, bucket=payload.bucket, storage_path=payload.storage_path, + ingestion_run_id=row_run_id, ) except Exception as exc: # noqa: BLE001 - db.rollback() + _set_document_failed( + db=db, + workspace_id=workspace_id, + document_id=payload.document_id, + category=IngestionFailureCategory.TRANSIENT_INFRASTRUCTURE, + detail=f"Failed to enqueue ingestion job: {exc}", + ) + db.commit() raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Failed to enqueue ingestion job: {exc}", ) from exc - db.commit() - return UploadCompleteResponse( document_id=payload.document_id, status="uploaded", @@ -589,6 +1221,87 @@ def upload_complete( ) +@router.post("/upload-complete", response_model=UploadCompleteResponse) +def upload_complete( + payload: UploadCompleteRequest, + workspace_id: uuid.UUID = Depends(get_workspace_id), + db: Session = Depends(get_db), +) -> UploadCompleteResponse: + enforce_workspace_rate_limit( + workspace_id=workspace_id, + operation="documents_upload_complete", + limit=UPLOAD_COMPLETE_RATE_LIMIT, + ) + + return _complete_upload_document(payload=payload, workspace_id=workspace_id, db=db) + + +@router.post("/upload-complete-batch", response_model=UploadCompleteBatchResponse) +def upload_complete_batch( + payload: UploadCompleteBatchRequest, + workspace_id: uuid.UUID = Depends(get_workspace_id), + db: Session = Depends(get_db), +) -> UploadCompleteBatchResponse: + enforce_workspace_rate_limit( + workspace_id=workspace_id, + operation="documents_upload_complete_batch", + limit=BULK_UPLOAD_COMPLETE_RATE_LIMIT, + ) + if not payload.files: + raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="files is required") + if len(payload.files) > settings.MAX_BULK_UPLOAD_FILES: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=f"Batch completion supports up to {settings.MAX_BULK_UPLOAD_FILES} files.", + ) + if payload.ingestion_run_id is not None: + _require_ingestion_run_schema(db) + + items: list[UploadCompleteBatchItem] = [] + accepted_count = 0 + failed_count = 0 + for index, file_payload in enumerate(payload.files): + try: + completed = _complete_upload_document( + payload=file_payload, + workspace_id=workspace_id, + db=db, + expected_run_id=payload.ingestion_run_id, + ) + accepted_count += 1 + items.append( + UploadCompleteBatchItem( + index=index, + document_id=completed.document_id, + status=completed.status, + job_id=completed.job_id, + ) + ) + except HTTPException as exc: + failed_count += 1 + items.append( + UploadCompleteBatchItem( + index=index, + document_id=file_payload.document_id, + status="failed", + error=str(exc.detail), + ) + ) + + if payload.ingestion_run_id is not None: + _refresh_ingestion_run_status( + db=db, workspace_id=workspace_id, run_id=payload.ingestion_run_id + ) + db.commit() + + return UploadCompleteBatchResponse( + ingestion_run_id=payload.ingestion_run_id, + accepted_count=accepted_count, + failed_count=failed_count, + items=items, + ) + + @router.delete("/{document_id}", status_code=status.HTTP_204_NO_CONTENT) def delete_document( document_id: uuid.UUID, @@ -596,7 +1309,9 @@ def delete_document( db: Session = Depends(get_db), ) -> Response: columns = _document_columns(db) - storage_bucket_select = "storage_bucket" if "storage_bucket" in columns else "NULL::text AS storage_bucket" + storage_bucket_select = ( + "storage_bucket" if "storage_bucket" in columns else "NULL::text AS storage_bucket" + ) delete_sql = text( f""" DELETE FROM documents @@ -606,10 +1321,14 @@ def delete_document( """ ) try: - doc_row = db.execute( - delete_sql, - {"workspace_id": workspace_id, "document_id": document_id}, - ).mappings().first() + doc_row = ( + db.execute( + delete_sql, + {"workspace_id": workspace_id, "document_id": document_id}, + ) + .mappings() + .first() + ) if doc_row is None: db.rollback() raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Document not found") @@ -664,19 +1383,31 @@ def retry_document( select_fields.append("storage_bucket") else: select_fields.append("NULL AS storage_bucket") + if "error_message" in columns: + select_fields.append("error_message") + else: + select_fields.append("NULL AS error_message") + if "ingestion_run_id" in columns: + select_fields.append("ingestion_run_id") + else: + select_fields.append("NULL AS ingestion_run_id") - row = db.execute( - text( - f""" + row = ( + db.execute( + text( + f""" SELECT {", ".join(select_fields)} FROM documents WHERE id = :document_id AND workspace_id = :workspace_id LIMIT 1 """ - ), - {"document_id": document_id, "workspace_id": workspace_id}, - ).mappings().first() + ), + {"document_id": document_id, "workspace_id": workspace_id}, + ) + .mappings() + .first() + ) if row is None: raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Document not found") if row["status"] != "failed": @@ -684,6 +1415,42 @@ def retry_document( status_code=status.HTTP_409_CONFLICT, detail="Retry is only allowed for failed documents", ) + if not is_retryable_failure(row["error_message"]): + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=f"Retry is not allowed for this failure. {row['error_message']}", + ) + + db.execute( + text( + """ + DELETE FROM chunk_embeddings + WHERE workspace_id = :workspace_id + AND document_id = :document_id + """ + ), + {"workspace_id": workspace_id, "document_id": document_id}, + ) + db.execute( + text( + """ + DELETE FROM chunks + WHERE workspace_id = :workspace_id + AND document_id = :document_id + """ + ), + {"workspace_id": workspace_id, "document_id": document_id}, + ) + db.execute( + text( + """ + DELETE FROM document_pages + WHERE workspace_id = :workspace_id + AND document_id = :document_id + """ + ), + {"workspace_id": workspace_id, "document_id": document_id}, + ) update_fields = ["status = 'uploaded'", "updated_at = :updated_at"] params: dict[str, object] = { @@ -693,6 +1460,8 @@ def retry_document( } if "error_message" in columns: update_fields.append("error_message = NULL") + if "page_count" in columns: + update_fields.append("page_count = NULL") db.execute( text( @@ -712,6 +1481,7 @@ def retry_document( document_id=document_id, bucket=row["storage_bucket"] or settings.SUPABASE_STORAGE_BUCKET, storage_path=row["storage_path"], + ingestion_run_id=row["ingestion_run_id"], ) except Exception as exc: # noqa: BLE001 db.rollback() @@ -742,19 +1512,31 @@ def reindex_document( select_fields.append("storage_bucket") else: select_fields.append("NULL AS storage_bucket") + if "error_message" in columns: + select_fields.append("error_message") + else: + select_fields.append("NULL AS error_message") + if "ingestion_run_id" in columns: + select_fields.append("ingestion_run_id") + else: + select_fields.append("NULL AS ingestion_run_id") - row = db.execute( - text( - f""" + row = ( + db.execute( + text( + f""" SELECT {", ".join(select_fields)} FROM documents WHERE id = :document_id AND workspace_id = :workspace_id LIMIT 1 """ - ), - {"document_id": document_id, "workspace_id": workspace_id}, - ).mappings().first() + ), + {"document_id": document_id, "workspace_id": workspace_id}, + ) + .mappings() + .first() + ) if row is None: raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Document not found") if row["status"] not in {"ready", "indexed", "failed"}: @@ -762,6 +1544,11 @@ def reindex_document( status_code=status.HTTP_409_CONFLICT, detail="Reindex is only allowed for ready/indexed/failed documents", ) + if row["status"] == "failed" and not is_retryable_failure(row["error_message"]): + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=f"Reindex is not allowed for this failure. {row['error_message']}", + ) pages_count = _count_for_document( db, @@ -817,13 +1604,18 @@ def reindex_document( try: if has_pages: - job = _enqueue_index(workspace_id=workspace_id, document_id=document_id) + job = _enqueue_index( + workspace_id=workspace_id, + document_id=document_id, + ingestion_run_id=row["ingestion_run_id"], + ) else: job = _enqueue_extract( workspace_id=workspace_id, document_id=document_id, bucket=row["storage_bucket"] or settings.SUPABASE_STORAGE_BUCKET, storage_path=row["storage_path"], + ingestion_run_id=row["ingestion_run_id"], ) except Exception as exc: # noqa: BLE001 db.rollback() diff --git a/server/app/config.py b/server/app/config.py index f3df499..73c30d2 100644 --- a/server/app/config.py +++ b/server/app/config.py @@ -28,9 +28,16 @@ class Settings(BaseSettings): MAX_QUESTION_CHARS: int = 500 EMBEDDING_MODEL: str = "text-embedding-3-small" EMBEDDING_DIM: int = 1536 + OPENAI_EMBEDDING_TIMEOUT_SECONDS: int = 300 LOG_EACH_QUERY: bool = False - MAX_FILE_SIZE_BYTES: int = 20 * 1024 * 1024 + MAX_FILE_SIZE_BYTES: int = 10 * 1024 * 1024 + MAX_PDF_PAGE_COUNT: int = 10 + MIN_EXTRACTED_TEXT_CHARS: int = 1 MAX_DOCUMENTS_PER_WORKSPACE: int = 100 + MAX_BULK_UPLOAD_FILES: int = 50 + EMBEDDING_BATCH_SIZE: int = 32 + INGEST_EXTRACT_JOB_TIMEOUT_SECONDS: int = 900 + INGEST_INDEX_JOB_TIMEOUT_SECONDS: int = 1800 ALLOWED_CONTENT_TYPES: list[str] = Field(default_factory=lambda: ["application/pdf"]) model_config = SettingsConfigDict(env_file=".env", extra="ignore") diff --git a/server/app/core/ingestion_policy.py b/server/app/core/ingestion_policy.py new file mode 100644 index 0000000..3676b94 --- /dev/null +++ b/server/app/core/ingestion_policy.py @@ -0,0 +1,76 @@ +from __future__ import annotations + +from enum import Enum + + +class IngestionFailureCategory(str, Enum): + VALIDATION = "validation" + UPLOAD_STORAGE = "upload_storage" + EXTRACTION = "extraction" + PAGE_LIMIT = "page_limit" + UNSUPPORTED_CONTENT = "unsupported_content" + INDEXING = "indexing" + BUDGET = "budget" + TRANSIENT_INFRASTRUCTURE = "transient_infrastructure" + + +FAILURE_LABELS: dict[IngestionFailureCategory, str] = { + IngestionFailureCategory.VALIDATION: "Validation failure", + IngestionFailureCategory.UPLOAD_STORAGE: "Upload/storage failure", + IngestionFailureCategory.EXTRACTION: "Extraction failure", + IngestionFailureCategory.PAGE_LIMIT: "Page-limit failure", + IngestionFailureCategory.UNSUPPORTED_CONTENT: "Unsupported-content failure", + IngestionFailureCategory.INDEXING: "Indexing failure", + IngestionFailureCategory.BUDGET: "Budget failure", + IngestionFailureCategory.TRANSIENT_INFRASTRUCTURE: "Transient infrastructure failure", +} + +RETRYABLE_FAILURE_CATEGORIES = { + IngestionFailureCategory.EXTRACTION, + IngestionFailureCategory.INDEXING, + IngestionFailureCategory.BUDGET, + IngestionFailureCategory.TRANSIENT_INFRASTRUCTURE, +} + +TERMINAL_FAILURE_CATEGORIES = { + IngestionFailureCategory.VALIDATION, + IngestionFailureCategory.UPLOAD_STORAGE, + IngestionFailureCategory.PAGE_LIMIT, + IngestionFailureCategory.UNSUPPORTED_CONTENT, +} + + +def format_bytes(size_bytes: int) -> str: + mib = 1024 * 1024 + if size_bytes % mib == 0: + return f"{size_bytes // mib} MB" + return f"{size_bytes} bytes" + + +def ingestion_failure_message(category: IngestionFailureCategory, detail: str) -> str: + cleaned_detail = " ".join(str(detail).split()) + return f"{FAILURE_LABELS[category]}: {cleaned_detail}" + + +class IngestionFailure(Exception): + def __init__(self, category: IngestionFailureCategory, detail: str) -> None: + self.category = category + self.detail = detail + super().__init__(ingestion_failure_message(category, detail)) + + +def failure_category_from_message(error_message: str | None) -> IngestionFailureCategory | None: + if not error_message: + return None + for category, label in FAILURE_LABELS.items(): + if error_message.startswith(f"{label}:"): + return category + return None + + +def is_retryable_failure(error_message: str | None) -> bool: + category = failure_category_from_message(error_message) + if category is None: + # Legacy failures did not have structured categories; keep them retryable. + return True + return category in RETRYABLE_FAILURE_CATEGORIES diff --git a/server/app/core/ingestion_runs.py b/server/app/core/ingestion_runs.py new file mode 100644 index 0000000..e82ada9 --- /dev/null +++ b/server/app/core/ingestion_runs.py @@ -0,0 +1,55 @@ +from __future__ import annotations + +PROCESSING_DOCUMENT_STATUSES = {"uploaded", "extracting", "indexing"} +SUCCESS_DOCUMENT_STATUSES = {"ready", "indexed"} +FAILED_DOCUMENT_STATUSES = {"failed"} +PENDING_DOCUMENT_STATUSES = {"pending_upload", "uploading"} + + +def empty_document_status_counts() -> dict[str, int]: + return { + "pending_upload": 0, + "uploading": 0, + "uploaded": 0, + "extracting": 0, + "indexing": 0, + "ready": 0, + "indexed": 0, + "failed": 0, + } + + +def derive_ingestion_run_status( + *, + status_counts: dict[str, int], + accepted_documents: int, + rejected_documents: int, +) -> str: + if accepted_documents <= 0: + return "failed" if rejected_documents > 0 else "preparing" + + success_count = sum( + status_counts.get(status_name, 0) for status_name in SUCCESS_DOCUMENT_STATUSES + ) + failed_count = sum( + status_counts.get(status_name, 0) for status_name in FAILED_DOCUMENT_STATUSES + ) + terminal_count = success_count + failed_count + + if terminal_count >= accepted_documents: + if failed_count and success_count: + return "partial" + if failed_count: + return "failed" + return "completed" + + processing_count = sum( + status_counts.get(status_name, 0) for status_name in PROCESSING_DOCUMENT_STATUSES + ) + pending_count = sum( + status_counts.get(status_name, 0) for status_name in PENDING_DOCUMENT_STATUSES + ) + + if processing_count > 0 or pending_count < accepted_documents: + return "processing" + return "preparing" diff --git a/server/app/core/token_budget.py b/server/app/core/token_budget.py index e75d4b9..8822b6c 100644 --- a/server/app/core/token_budget.py +++ b/server/app/core/token_budget.py @@ -2,17 +2,20 @@ from contextlib import nullcontext from datetime import UTC, date, datetime, time, timedelta +from typing import Callable, TypeVar import uuid from sqlalchemy import select from sqlalchemy.dialects.postgresql import insert as pg_insert from sqlalchemy.dialects.sqlite import insert as sqlite_insert -from sqlalchemy.orm import Session +from sqlalchemy.orm import Session, sessionmaker from app.config import settings from app.core.errors import BudgetExceededError, InvalidReservationError from app.db.models import WorkspaceDailyUsage +T = TypeVar("T") + def _as_utc_date(usage_date_utc: date | datetime) -> date: if isinstance(usage_date_utc, datetime): @@ -71,6 +74,20 @@ def _transaction_context(db: Session): return db.begin() +def _run_in_isolated_session(db: Session, operation: Callable[[Session], T]) -> T: + isolated_factory = sessionmaker(bind=db.get_bind(), autocommit=False, autoflush=False) + isolated_db = isolated_factory() + try: + result = operation(isolated_db) + isolated_db.commit() + return result + except Exception: + isolated_db.rollback() + raise + finally: + isolated_db.close() + + def get_or_create_daily_row(db: Session, workspace_id: uuid.UUID, usage_date_utc: date | datetime) -> WorkspaceDailyUsage: usage_date = _as_utc_date(usage_date_utc) _insert_usage_row_if_missing(db, workspace_id, usage_date) @@ -85,15 +102,12 @@ def get_or_create_daily_row(db: Session, workspace_id: uuid.UUID, usage_date_utc return db.execute(stmt).scalar_one() -def reserve_tokens( +def _reserve_tokens_in_session( db: Session, workspace_id: uuid.UUID, amount: int, usage_date_utc: date | datetime, - reservation_ttl_seconds: int = 600, ) -> dict[str, int]: - reservation_amount = _ensure_non_negative_amount(amount) - _ensure_non_negative_amount(reservation_ttl_seconds, "reservation_ttl_seconds") token_limit = int(settings.DAILY_TOKEN_LIMIT) with _transaction_context(db): @@ -101,7 +115,7 @@ def reserve_tokens( used = int(row.tokens_used) reserved = int(row.tokens_reserved) remaining = token_limit - (used + reserved) - if reservation_amount > remaining: + if amount > remaining: raise BudgetExceededError( "Daily token limit reached for this workspace", used=used, @@ -109,7 +123,7 @@ def reserve_tokens( limit=token_limit, ) - row.tokens_reserved = reserved + reservation_amount + row.tokens_reserved = reserved + amount db.flush() reserved_now = int(row.tokens_reserved) remaining_now = max(0, token_limit - (used + reserved_now)) @@ -121,30 +135,38 @@ def reserve_tokens( } -def release_tokens(db: Session, workspace_id: uuid.UUID, amount: int, usage_date_utc: date | datetime) -> dict[str, int]: - release_amount = _ensure_non_negative_amount(amount) +def _release_tokens_in_session( + db: Session, + workspace_id: uuid.UUID, + amount: int, + usage_date_utc: date | datetime, +) -> dict[str, int]: with _transaction_context(db): row = get_or_create_daily_row(db, workspace_id, usage_date_utc) reserved = int(row.tokens_reserved) - if release_amount > reserved: + if amount > reserved: raise InvalidReservationError("Cannot release more tokens than currently reserved") - row.tokens_reserved = reserved - release_amount + row.tokens_reserved = reserved - amount db.flush() reserved_now = int(row.tokens_reserved) return {"reserved_now": reserved_now} -def commit_usage(db: Session, workspace_id: uuid.UUID, amount: int, usage_date_utc: date | datetime) -> dict[str, int]: - usage_amount = _ensure_non_negative_amount(amount) +def _commit_usage_in_session( + db: Session, + workspace_id: uuid.UUID, + amount: int, + usage_date_utc: date | datetime, +) -> dict[str, int]: with _transaction_context(db): row = get_or_create_daily_row(db, workspace_id, usage_date_utc) reserved = int(row.tokens_reserved) used = int(row.tokens_used) - if usage_amount > reserved: + if amount > reserved: raise InvalidReservationError("Cannot commit more tokens than currently reserved") - row.tokens_reserved = reserved - usage_amount - row.tokens_used = used + usage_amount + row.tokens_reserved = reserved - amount + row.tokens_used = used + amount db.flush() used_now = int(row.tokens_used) reserved_now = int(row.tokens_reserved) @@ -152,6 +174,73 @@ def commit_usage(db: Session, workspace_id: uuid.UUID, amount: int, usage_date_u return {"used_now": used_now, "reserved_now": reserved_now} +def reserve_tokens( + db: Session, + workspace_id: uuid.UUID, + amount: int, + usage_date_utc: date | datetime, + reservation_ttl_seconds: int = 600, +) -> dict[str, int]: + reservation_amount = _ensure_non_negative_amount(amount) + _ensure_non_negative_amount(reservation_ttl_seconds, "reservation_ttl_seconds") + if db.in_transaction(): + return _run_in_isolated_session( + db, + lambda isolated_db: _reserve_tokens_in_session( + isolated_db, + workspace_id=workspace_id, + amount=reservation_amount, + usage_date_utc=usage_date_utc, + ), + ) + return _reserve_tokens_in_session( + db, + workspace_id=workspace_id, + amount=reservation_amount, + usage_date_utc=usage_date_utc, + ) + + +def release_tokens(db: Session, workspace_id: uuid.UUID, amount: int, usage_date_utc: date | datetime) -> dict[str, int]: + release_amount = _ensure_non_negative_amount(amount) + if db.in_transaction(): + return _run_in_isolated_session( + db, + lambda isolated_db: _release_tokens_in_session( + isolated_db, + workspace_id=workspace_id, + amount=release_amount, + usage_date_utc=usage_date_utc, + ), + ) + return _release_tokens_in_session( + db, + workspace_id=workspace_id, + amount=release_amount, + usage_date_utc=usage_date_utc, + ) + + +def commit_usage(db: Session, workspace_id: uuid.UUID, amount: int, usage_date_utc: date | datetime) -> dict[str, int]: + usage_amount = _ensure_non_negative_amount(amount) + if db.in_transaction(): + return _run_in_isolated_session( + db, + lambda isolated_db: _commit_usage_in_session( + isolated_db, + workspace_id=workspace_id, + amount=usage_amount, + usage_date_utc=usage_date_utc, + ), + ) + return _commit_usage_in_session( + db, + workspace_id=workspace_id, + amount=usage_amount, + usage_date_utc=usage_date_utc, + ) + + def get_budget_status( db: Session, workspace_id: uuid.UUID, diff --git a/server/app/db/models.py b/server/app/db/models.py index 562a811..fa5db69 100644 --- a/server/app/db/models.py +++ b/server/app/db/models.py @@ -1,7 +1,17 @@ import uuid from datetime import datetime -from sqlalchemy import BigInteger, CheckConstraint, Date, DateTime, ForeignKey, Integer, String, Text, func +from sqlalchemy import ( + BigInteger, + CheckConstraint, + Date, + DateTime, + ForeignKey, + Integer, + String, + Text, + func, +) from sqlalchemy.dialects.postgresql import JSONB, UUID from sqlalchemy.orm import Mapped, mapped_column @@ -14,7 +24,29 @@ class Workspace(Base): id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4) name: Mapped[str] = mapped_column(Text, nullable=False) owner_id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), nullable=False) - created_at: Mapped[DateTime] = mapped_column(DateTime(timezone=True), nullable=False, server_default=func.now()) + created_at: Mapped[DateTime] = mapped_column( + DateTime(timezone=True), nullable=False, server_default=func.now() + ) + + +class IngestionRun(Base): + __tablename__ = "ingestion_runs" + + id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4) + workspace_id: Mapped[uuid.UUID] = mapped_column( + UUID(as_uuid=True), ForeignKey("workspaces.id", ondelete="CASCADE"), nullable=False + ) + name: Mapped[str | None] = mapped_column(Text, nullable=True) + status: Mapped[str] = mapped_column(String(32), nullable=False, default="preparing") + total_documents: Mapped[int] = mapped_column(Integer, nullable=False, default=0) + accepted_documents: Mapped[int] = mapped_column(Integer, nullable=False, default=0) + rejected_documents: Mapped[int] = mapped_column(Integer, nullable=False, default=0) + created_at: Mapped[DateTime] = mapped_column( + DateTime(timezone=True), nullable=False, server_default=func.now() + ) + updated_at: Mapped[DateTime] = mapped_column( + DateTime(timezone=True), nullable=False, server_default=func.now(), onupdate=func.now() + ) class Document(Base): @@ -29,9 +61,14 @@ class Document(Base): page_count: Mapped[int | None] = mapped_column(Integer, nullable=True) file_hash_sha256: Mapped[str] = mapped_column(Text, nullable=False) storage_path: Mapped[str] = mapped_column(Text, nullable=False) + ingestion_run_id: Mapped[uuid.UUID | None] = mapped_column( + UUID(as_uuid=True), ForeignKey("ingestion_runs.id", ondelete="SET NULL"), nullable=True + ) status: Mapped[str] = mapped_column(String(32), nullable=False, default="pending_upload") error_message: Mapped[str | None] = mapped_column(Text, nullable=True) - created_at: Mapped[DateTime] = mapped_column(DateTime(timezone=True), nullable=False, server_default=func.now()) + created_at: Mapped[DateTime] = mapped_column( + DateTime(timezone=True), nullable=False, server_default=func.now() + ) updated_at: Mapped[DateTime] = mapped_column( DateTime(timezone=True), nullable=False, server_default=func.now(), onupdate=func.now() ) @@ -50,7 +87,9 @@ class WorkspaceDailyUsage(Base): date: Mapped[Date] = mapped_column(Date, primary_key=True) tokens_used: Mapped[int] = mapped_column(BigInteger, nullable=False, default=0) tokens_reserved: Mapped[int] = mapped_column(BigInteger, nullable=False, default=0) - created_at: Mapped[DateTime] = mapped_column(DateTime(timezone=True), nullable=False, server_default=func.now()) + created_at: Mapped[DateTime] = mapped_column( + DateTime(timezone=True), nullable=False, server_default=func.now() + ) updated_at: Mapped[DateTime] = mapped_column( DateTime(timezone=True), nullable=False, server_default=func.now(), onupdate=func.now() ) @@ -68,9 +107,13 @@ class ChatSession(Base): ) title: Mapped[str] = mapped_column(Text, nullable=False, default="") messages: Mapped[list[dict[str, object]]] = mapped_column(JSONB, nullable=False, default=list) - started_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False, server_default=func.now()) + started_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), nullable=False, server_default=func.now() + ) ended_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True) - created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False, server_default=func.now()) + created_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), nullable=False, server_default=func.now() + ) updated_at: Mapped[datetime] = mapped_column( DateTime(timezone=True), nullable=False, server_default=func.now(), onupdate=func.now() ) diff --git a/server/app/schemas/documents.py b/server/app/schemas/documents.py index bce86d5..7bf4f4e 100644 --- a/server/app/schemas/documents.py +++ b/server/app/schemas/documents.py @@ -19,6 +19,41 @@ class UploadPrepareResponse(BaseModel): expires_in: int +class UploadPrepareBatchFile(BaseModel): + filename: str = Field(min_length=1, max_length=255) + content_type: str = Field(min_length=1, max_length=100) + file_size_bytes: int = Field(gt=0) + idempotency_key: str | None = Field(default=None, min_length=1, max_length=120) + client_file_id: str | None = Field(default=None, min_length=1, max_length=120) + + +class UploadPrepareBatchRequest(BaseModel): + files: list[UploadPrepareBatchFile] + name: str | None = Field(default=None, min_length=1, max_length=255) + + +class UploadPrepareBatchItem(BaseModel): + index: int + filename: str + client_file_id: str | None = None + status: str + document_id: uuid.UUID | None = None + bucket: str | None = None + storage_path: str | None = None + upload_url: str | None = None + expires_in: int | None = None + error: str | None = None + + +class UploadPrepareBatchResponse(BaseModel): + ingestion_run_id: uuid.UUID + bucket: str + expires_in: int + accepted_count: int + rejected_count: int + items: list[UploadPrepareBatchItem] + + class UploadCompleteRequest(BaseModel): document_id: uuid.UUID bucket: str = Field(min_length=1, max_length=255) @@ -31,6 +66,26 @@ class UploadCompleteResponse(BaseModel): job_id: str +class UploadCompleteBatchRequest(BaseModel): + files: list[UploadCompleteRequest] + ingestion_run_id: uuid.UUID | None = None + + +class UploadCompleteBatchItem(BaseModel): + index: int + document_id: uuid.UUID + status: str + job_id: str | None = None + error: str | None = None + + +class UploadCompleteBatchResponse(BaseModel): + ingestion_run_id: uuid.UUID | None = None + accepted_count: int + failed_count: int + items: list[UploadCompleteBatchItem] + + class DocumentJobResponse(BaseModel): document_id: uuid.UUID status: str @@ -42,7 +97,10 @@ class DocumentListItem(BaseModel): filename: str content_type: str file_size_bytes: int + page_count: int | None = None + ingestion_run_id: uuid.UUID | None = None status: str + error_message: str | None = None created_at: datetime updated_at: datetime @@ -66,9 +124,48 @@ class DocumentDetailResponse(BaseModel): filename: str content_type: str file_size_bytes: int + ingestion_run_id: uuid.UUID | None = None status: str + error_message: str | None = None bucket: str storage_path: str created_at: datetime updated_at: datetime progress: DocumentProgress + + +class IngestionRunStatusCounts(BaseModel): + pending_upload: int = 0 + uploading: int = 0 + uploaded: int = 0 + extracting: int = 0 + indexing: int = 0 + ready: int = 0 + indexed: int = 0 + failed: int = 0 + total: int = 0 + + +class IngestionRunResponse(BaseModel): + id: uuid.UUID + name: str | None = None + status: str + total_documents: int + accepted_documents: int + rejected_documents: int + document_statuses: IngestionRunStatusCounts + created_at: datetime + updated_at: datetime + + +class IngestionQueueStatusItem(BaseModel): + name: str + queued_count: int + started_count: int + deferred_count: int + scheduled_count: int + failed_count: int + + +class IngestionQueueStatusResponse(BaseModel): + queues: list[IngestionQueueStatusItem] diff --git a/server/app/utils/rate_limit.py b/server/app/utils/rate_limit.py index 46ca86a..05d3004 100644 --- a/server/app/utils/rate_limit.py +++ b/server/app/utils/rate_limit.py @@ -10,6 +10,8 @@ UPLOAD_PREPARE_RATE_LIMIT = 10 UPLOAD_COMPLETE_RATE_LIMIT = 20 +BULK_UPLOAD_PREPARE_RATE_LIMIT = 5 +BULK_UPLOAD_COMPLETE_RATE_LIMIT = 10 QUERY_RATE_LIMIT = 100 RATE_LIMIT_WINDOW_SECONDS = 60 diff --git a/server/scripts/bootstrap_schema.py b/server/scripts/bootstrap_schema.py new file mode 100644 index 0000000..56d2409 --- /dev/null +++ b/server/scripts/bootstrap_schema.py @@ -0,0 +1,113 @@ +from __future__ import annotations + +from sqlalchemy import text + +from app.db import models # noqa: F401 +from app.db.session import Base, engine + + +DDL_STATEMENTS = [ + """ + CREATE EXTENSION IF NOT EXISTS vector + """, + """ + CREATE EXTENSION IF NOT EXISTS pgcrypto + """, + """ + CREATE TABLE IF NOT EXISTS document_pages ( + id BIGSERIAL PRIMARY KEY, + workspace_id UUID NOT NULL REFERENCES workspaces(id) ON DELETE CASCADE, + document_id UUID NOT NULL REFERENCES documents(id) ON DELETE CASCADE, + page_number INTEGER NOT NULL CHECK (page_number > 0), + content TEXT NOT NULL, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + UNIQUE (document_id, page_number) + ) + """, + """ + CREATE INDEX IF NOT EXISTS idx_document_pages_workspace + ON document_pages(workspace_id) + """, + """ + CREATE INDEX IF NOT EXISTS idx_document_pages_document + ON document_pages(document_id) + """, + """ + CREATE TABLE IF NOT EXISTS chunks ( + id UUID PRIMARY KEY, + workspace_id UUID NOT NULL REFERENCES workspaces(id) ON DELETE CASCADE, + document_id UUID NOT NULL REFERENCES documents(id) ON DELETE CASCADE, + page_start INTEGER NOT NULL CHECK (page_start > 0), + page_end INTEGER NOT NULL CHECK (page_end >= page_start), + chunk_index INTEGER NOT NULL CHECK (chunk_index >= 0), + content TEXT NOT NULL, + content_hash TEXT NOT NULL, + token_count INTEGER NOT NULL DEFAULT 0 CHECK (token_count >= 0), + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW() + ) + """, + """ + CREATE UNIQUE INDEX IF NOT EXISTS idx_chunks_document_chunk_index + ON chunks(document_id, chunk_index) + """, + """ + CREATE INDEX IF NOT EXISTS idx_chunks_workspace_document + ON chunks(workspace_id, document_id) + """, + """ + CREATE TABLE IF NOT EXISTS chunk_embeddings ( + chunk_id UUID PRIMARY KEY REFERENCES chunks(id) ON DELETE CASCADE, + workspace_id UUID NOT NULL REFERENCES workspaces(id) ON DELETE CASCADE, + document_id UUID NOT NULL REFERENCES documents(id) ON DELETE CASCADE, + embedding vector(1536) NOT NULL, + embedding_model TEXT NOT NULL, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW() + ) + """, + """ + CREATE INDEX IF NOT EXISTS idx_chunk_embeddings_workspace_document + ON chunk_embeddings(workspace_id, document_id) + """, + """ + CREATE INDEX IF NOT EXISTS idx_chunk_embeddings_vector + ON chunk_embeddings + USING hnsw (embedding vector_cosine_ops) + """, + """ + CREATE TABLE IF NOT EXISTS query_logs ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + workspace_id UUID NOT NULL REFERENCES workspaces(id) ON DELETE CASCADE, + user_id TEXT NOT NULL, + query_text TEXT NOT NULL, + documents_searched UUID[] NOT NULL DEFAULT '{}', + retrieved_chunk_ids UUID[] NOT NULL DEFAULT '{}', + chunk_scores DOUBLE PRECISION[] NOT NULL DEFAULT '{}', + answer_text TEXT NULL, + error_message TEXT NULL, + retrieval_latency_ms INTEGER NOT NULL DEFAULT 0, + llm_latency_ms INTEGER NULL, + total_latency_ms INTEGER NOT NULL DEFAULT 0, + embedding_tokens_used BIGINT NOT NULL DEFAULT 0, + llm_input_tokens INTEGER NULL, + llm_output_tokens INTEGER NULL, + total_tokens_used BIGINT NOT NULL DEFAULT 0, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW() + ) + """, + """ + CREATE INDEX IF NOT EXISTS idx_query_logs_workspace_created_at + ON query_logs(workspace_id, created_at DESC) + """, +] + + +def main() -> None: + Base.metadata.create_all(bind=engine) + with engine.begin() as conn: + for statement in DDL_STATEMENTS: + conn.execute(text(statement)) + print("Schema bootstrap completed.") + + +if __name__ == "__main__": + main() diff --git a/server/scripts/inspect_schema.py b/server/scripts/inspect_schema.py new file mode 100644 index 0000000..59073b2 --- /dev/null +++ b/server/scripts/inspect_schema.py @@ -0,0 +1,41 @@ +from __future__ import annotations + +from sqlalchemy import text + +from app.db.session import engine + + +def main() -> None: + with engine.connect() as conn: + table_rows = conn.execute( + text( + """ + SELECT table_schema, table_name + FROM information_schema.tables + WHERE table_schema NOT IN ('pg_catalog', 'information_schema') + ORDER BY table_schema, table_name + """ + ) + ) + extension_rows = conn.execute( + text( + """ + SELECT extname + FROM pg_extension + WHERE extname IN ('pgcrypto', 'vector') + ORDER BY extname + """ + ) + ) + + print("Tables:") + for row in table_rows: + print(f"{row.table_schema}.{row.table_name}") + + print("Extensions:") + for row in extension_rows: + print(row.extname) + + +if __name__ == "__main__": + main() diff --git a/server/tests/test_core/test_ingestion_policy.py b/server/tests/test_core/test_ingestion_policy.py new file mode 100644 index 0000000..126b444 --- /dev/null +++ b/server/tests/test_core/test_ingestion_policy.py @@ -0,0 +1,31 @@ +from app.core.ingestion_policy import ( + IngestionFailureCategory, + failure_category_from_message, + ingestion_failure_message, + is_retryable_failure, +) + + +def test_ingestion_failure_messages_are_categorized() -> None: + message = ingestion_failure_message( + IngestionFailureCategory.PAGE_LIMIT, + "PDF has 12 pages. Maximum supported page count is 10.", + ) + + assert message.startswith("Page-limit failure:") + assert failure_category_from_message(message) == IngestionFailureCategory.PAGE_LIMIT + + +def test_retryability_uses_failure_category() -> None: + terminal_message = ingestion_failure_message( + IngestionFailureCategory.UNSUPPORTED_CONTENT, + "No extractable text was found.", + ) + retryable_message = ingestion_failure_message( + IngestionFailureCategory.INDEXING, + "Embedding API failed.", + ) + + assert not is_retryable_failure(terminal_message) + assert is_retryable_failure(retryable_message) + assert is_retryable_failure("legacy raw exception") diff --git a/server/tests/test_core/test_ingestion_runs.py b/server/tests/test_core/test_ingestion_runs.py new file mode 100644 index 0000000..440477f --- /dev/null +++ b/server/tests/test_core/test_ingestion_runs.py @@ -0,0 +1,47 @@ +from app.core.ingestion_runs import ( + derive_ingestion_run_status, + empty_document_status_counts, +) + + +def test_ingestion_run_status_is_processing_while_documents_are_active() -> None: + counts = empty_document_status_counts() + counts["uploaded"] = 2 + counts["ready"] = 1 + + assert ( + derive_ingestion_run_status( + status_counts=counts, + accepted_documents=3, + rejected_documents=0, + ) + == "processing" + ) + + +def test_ingestion_run_status_is_partial_when_success_and_failure_finish() -> None: + counts = empty_document_status_counts() + counts["ready"] = 2 + counts["failed"] = 1 + + assert ( + derive_ingestion_run_status( + status_counts=counts, + accepted_documents=3, + rejected_documents=0, + ) + == "partial" + ) + + +def test_ingestion_run_status_handles_all_rejected_prepare() -> None: + counts = empty_document_status_counts() + + assert ( + derive_ingestion_run_status( + status_counts=counts, + accepted_documents=0, + rejected_documents=2, + ) + == "failed" + ) diff --git a/server/tests/test_token_budget.py b/server/tests/test_token_budget.py index 921c16a..44e44fa 100644 --- a/server/tests/test_token_budget.py +++ b/server/tests/test_token_budget.py @@ -71,6 +71,47 @@ def test_release_reduces_reserved(db_session: Session, workspace_id: uuid.UUID, assert status["used"] == 0 +def test_budget_operations_use_short_transaction_when_caller_session_is_active( + sqlite_session_factory: sessionmaker, + limit_guard, +) -> None: + settings.DAILY_TOKEN_LIMIT = 1_000 + today = date(2026, 2, 15) + + setup_session = sqlite_session_factory() + workspace = Workspace(name="Active Transaction", owner_id=uuid.uuid4()) + setup_session.add(workspace) + setup_session.commit() + setup_session.close() + + caller_session = sqlite_session_factory() + observer_session = sqlite_session_factory() + try: + # Trigger SQLAlchemy autobegin on the caller session before reserving tokens. + caller_session.get(Workspace, workspace.id) + + reserve_result = reserve_tokens(caller_session, workspace.id, 200, today) + status_after_reserve = get_budget_status(observer_session, workspace.id, today) + + commit_result = commit_usage(caller_session, workspace.id, 125, today) + status_after_commit = get_budget_status(observer_session, workspace.id, today) + + release_result = release_tokens(caller_session, workspace.id, 75, today) + status_after_release = get_budget_status(observer_session, workspace.id, today) + + assert reserve_result["reserved"] == 200 + assert status_after_reserve["reserved"] == 200 + assert commit_result["used_now"] == 125 + assert status_after_commit["used"] == 125 + assert status_after_commit["reserved"] == 75 + assert release_result["reserved_now"] == 0 + assert status_after_release["used"] == 125 + assert status_after_release["reserved"] == 0 + finally: + caller_session.close() + observer_session.close() + + def test_resets_at_is_next_midnight_utc(db_session: Session, workspace_id: uuid.UUID, limit_guard) -> None: settings.DAILY_TOKEN_LIMIT = 1_000 usage_day = date(2026, 2, 15) diff --git a/worker/README.md b/worker/README.md index 45135fc..5aa2718 100644 --- a/worker/README.md +++ b/worker/README.md @@ -61,8 +61,11 @@ Purpose: Current behavior: - extraction is page-based - temporary PDF file is written under `/tmp` +- PDFs larger than `MAX_FILE_SIZE_BYTES` or over `MAX_PDF_PAGE_COUNT` are marked `failed` +- scanned/image-only PDFs with no extractable text are marked `failed` - failures mark the document as `failed` - if the schema does not include `extracting`, the code falls back to `indexing` +- batch ingestion jobs carry optional `ingestion_run_id` metadata for grouped tracing ### `jobs/ingest_index.py` @@ -70,9 +73,9 @@ Purpose: - load extracted page text - split each page into chunks - insert chunk rows -- reserve token budget per chunk embedding -- call OpenAI embeddings -- insert vectors into `chunk_embeddings` +- reserve token budget per embedding batch +- call OpenAI embeddings in batches +- insert vectors into `chunk_embeddings` in batches - commit token usage - mark the document as `ready` or `indexed` @@ -80,6 +83,7 @@ Current chunking behavior: - page-bounded chunks - target chunk size: `500` tokens - overlap: `100` tokens +- embedding batch size is controlled by `EMBEDDING_BATCH_SIZE` - `tiktoken` when available, character approximation fallback otherwise Failure behavior: @@ -130,6 +134,10 @@ OPENAI_API_KEY=sk-... QUEUE_NAME=ingest_extract RESERVATION_TTL_SECONDS=600 EMBEDDING_MODEL=text-embedding-3-small +MAX_FILE_SIZE_BYTES=10485760 +MAX_PDF_PAGE_COUNT=10 +MIN_EXTRACTED_TEXT_CHARS=1 +EMBEDDING_BATCH_SIZE=32 ``` ## Run diff --git a/worker/jobs/ingest_callbacks.py b/worker/jobs/ingest_callbacks.py new file mode 100644 index 0000000..a9a62b2 --- /dev/null +++ b/worker/jobs/ingest_callbacks.py @@ -0,0 +1,106 @@ +from __future__ import annotations + +import logging +import uuid + +from sqlalchemy import text + +from app.core.ingestion_policy import ( + IngestionFailureCategory, + ingestion_failure_message, +) +from app.db.session import SessionLocal + + +logger = logging.getLogger(__name__) + + +def _failure_message(detail: str) -> str: + return ingestion_failure_message( + IngestionFailureCategory.TRANSIENT_INFRASTRUCTURE, detail + )[:2000] + + +def mark_ingestion_job_failed(job, connection, exc_type, exc_value, traceback) -> None: + meta = job.meta or {} + workspace_id = meta.get("workspace_id") + document_id = meta.get("document_id") + ingestion_run_id = meta.get("ingestion_run_id") + if not workspace_id or not document_id: + logger.warning( + "Ingestion failure callback skipped job without document metadata", + extra={"job_id": job.id, "queue": job.origin}, + ) + return + + try: + workspace_uuid = uuid.UUID(str(workspace_id)) + document_uuid = uuid.UUID(str(document_id)) + except ValueError: + logger.warning( + "Ingestion failure callback skipped invalid document metadata", + extra={ + "job_id": job.id, + "workspace_id": workspace_id, + "document_id": document_id, + }, + ) + return + + detail = f"{job.func_name} failed in queue {job.origin}: {exc_value}" + error_message = _failure_message(detail) + + with SessionLocal() as db: + row = ( + db.execute( + text( + """ + SELECT status, error_message + FROM documents + WHERE id = :document_id + AND workspace_id = :workspace_id + LIMIT 1 + """ + ), + {"workspace_id": workspace_uuid, "document_id": document_uuid}, + ) + .mappings() + .first() + ) + if row is None: + return + + # Normal job exceptions already set a structured failure before RQ runs + # callbacks. Only repair rows left in an active state by timeout/kill. + if row["status"] == "failed" and row["error_message"]: + return + + db.execute( + text( + """ + UPDATE documents + SET status = 'failed', + error_message = :error_message, + updated_at = NOW() + WHERE id = :document_id + AND workspace_id = :workspace_id + """ + ), + { + "workspace_id": workspace_uuid, + "document_id": document_uuid, + "error_message": error_message, + }, + ) + db.commit() + + logger.warning( + "Ingestion failure callback marked document failed", + extra={ + "job_id": job.id, + "queue": job.origin, + "workspace_id": workspace_id, + "document_id": document_id, + "ingestion_run_id": ingestion_run_id, + }, + ) diff --git a/worker/jobs/ingest_extract.py b/worker/jobs/ingest_extract.py index 7025e6d..772e0c6 100644 --- a/worker/jobs/ingest_extract.py +++ b/worker/jobs/ingest_extract.py @@ -6,10 +6,17 @@ from redis import Redis from rq import Queue +from rq.job import Callback from sqlalchemy import text from sqlalchemy.exc import IntegrityError from app.config import settings +from app.core.ingestion_policy import ( + IngestionFailure, + IngestionFailureCategory, + format_bytes, + ingestion_failure_message, +) from app.core.storage import download_object_bytes from app.db.session import SessionLocal @@ -20,6 +27,11 @@ logger = logging.getLogger(__name__) +INGESTION_FAILURE_CALLBACK = "jobs.ingest_callbacks.mark_ingestion_job_failed" + + +def _failure_message(category: IngestionFailureCategory, detail: str) -> str: + return ingestion_failure_message(category, detail)[:2000] def _set_document_status( @@ -51,6 +63,44 @@ def _set_document_status( db.commit() +def _cleanup_ingestion_artifacts( + *, workspace_id: uuid.UUID, document_id: uuid.UUID, include_pages: bool +) -> None: + with SessionLocal() as db: + db.execute( + text( + """ + DELETE FROM chunk_embeddings + WHERE workspace_id = :workspace_id + AND document_id = :document_id + """ + ), + {"workspace_id": workspace_id, "document_id": document_id}, + ) + db.execute( + text( + """ + DELETE FROM chunks + WHERE workspace_id = :workspace_id + AND document_id = :document_id + """ + ), + {"workspace_id": workspace_id, "document_id": document_id}, + ) + if include_pages: + db.execute( + text( + """ + DELETE FROM document_pages + WHERE workspace_id = :workspace_id + AND document_id = :document_id + """ + ), + {"workspace_id": workspace_id, "document_id": document_id}, + ) + db.commit() + + def _document_has_column(db, column_name: str) -> bool: result = db.execute( text( @@ -89,7 +139,13 @@ def _allowed_document_statuses(db) -> set[str]: return {part.strip().strip("'\"") for part in inside.split(",")} -def ingest_extract(workspace_id: str, document_id: str, bucket: str, storage_path: str) -> dict: +def ingest_extract( + workspace_id: str, + document_id: str, + bucket: str, + storage_path: str, + ingestion_run_id: str | None = None, +) -> dict: workspace_uuid = uuid.UUID(workspace_id) document_uuid = uuid.UUID(document_id) temp_path = Path(f"/tmp/{document_uuid}.pdf") @@ -97,14 +153,68 @@ def ingest_extract(workspace_id: str, document_id: str, bucket: str, storage_pat with SessionLocal() as db: allowed = _allowed_document_statuses(db) extracting_status = "extracting" if "extracting" in allowed else "indexing" - _set_document_status(workspace_id=workspace_uuid, document_id=document_uuid, status=extracting_status) + _set_document_status( + workspace_id=workspace_uuid, document_id=document_uuid, status=extracting_status + ) try: - file_bytes = download_object_bytes(bucket=bucket, path=storage_path) + try: + file_bytes = download_object_bytes(bucket=bucket, path=storage_path) + except Exception as exc: # noqa: BLE001 + raise IngestionFailure( + IngestionFailureCategory.UPLOAD_STORAGE, + "Unable to download the uploaded PDF from storage. Upload the file again and retry.", + ) from exc + + if len(file_bytes) > settings.MAX_FILE_SIZE_BYTES: + raise IngestionFailure( + IngestionFailureCategory.VALIDATION, + ( + f"PDF file size must be {format_bytes(settings.MAX_FILE_SIZE_BYTES)} or smaller. " + f"Uploaded object size is {format_bytes(len(file_bytes))}." + ), + ) + temp_path.write_bytes(file_bytes) - reader = PdfReader(str(temp_path)) + try: + reader = PdfReader(str(temp_path)) + except Exception as exc: # noqa: BLE001 + raise IngestionFailure( + IngestionFailureCategory.EXTRACTION, + "PDF could not be read. Upload a valid, unencrypted text-based PDF.", + ) from exc + + if getattr(reader, "is_encrypted", False): + raise IngestionFailure( + IngestionFailureCategory.UNSUPPORTED_CONTENT, + "Encrypted PDFs are not supported. Upload an unencrypted text-based PDF.", + ) + pages_total = len(reader.pages) + if pages_total < 1: + raise IngestionFailure( + IngestionFailureCategory.UNSUPPORTED_CONTENT, + "PDF does not contain any pages.", + ) + if pages_total > settings.MAX_PDF_PAGE_COUNT: + raise IngestionFailure( + IngestionFailureCategory.PAGE_LIMIT, + f"PDF has {pages_total} pages. Maximum supported page count is {settings.MAX_PDF_PAGE_COUNT}.", + ) + + extracted_pages: list[tuple[int, str]] = [] + total_text_chars = 0 + for page_index, page in enumerate(reader.pages): + page_text = page.extract_text() or "" + extracted_pages.append((page_index + 1, page_text)) + total_text_chars += len(page_text.strip()) + + if total_text_chars < settings.MIN_EXTRACTED_TEXT_CHARS: + raise IngestionFailure( + IngestionFailureCategory.UNSUPPORTED_CONTENT, + "No extractable text was found. Scanned/image-only PDFs are not supported.", + ) with SessionLocal() as db: # Idempotency: reruns overwrite prior extraction for this document. @@ -119,8 +229,16 @@ def ingest_extract(workspace_id: str, document_id: str, bucket: str, storage_pat {"workspace_id": workspace_uuid, "document_id": document_uuid}, ) - for page_index, page in enumerate(reader.pages): - page_text = page.extract_text() or "" + page_rows = [ + { + "workspace_id": workspace_uuid, + "document_id": document_uuid, + "page_number": page_number, + "content": page_text, + } + for page_number, page_text in extracted_pages + ] + if page_rows: db.execute( text( """ @@ -128,15 +246,15 @@ def ingest_extract(workspace_id: str, document_id: str, bucket: str, storage_pat VALUES (:workspace_id, :document_id, :page_number, :content) """ ), - { - "workspace_id": workspace_uuid, - "document_id": document_uuid, - "page_number": page_index + 1, - "content": page_text, - }, + page_rows, ) - update_fields = ["page_count = :page_count", "status = 'indexing'", "error_message = NULL", "updated_at = NOW()"] + update_fields = [ + "page_count = :page_count", + "status = 'indexing'", + "error_message = NULL", + "updated_at = NOW()", + ] params: dict[str, object] = { "workspace_id": workspace_uuid, "document_id": document_uuid, @@ -165,35 +283,81 @@ def ingest_extract(workspace_id: str, document_id: str, bucket: str, storage_pat "jobs.ingest_index.ingest_index", workspace_id=workspace_id, document_id=document_id, + ingestion_run_id=ingestion_run_id, + job_timeout=settings.INGEST_INDEX_JOB_TIMEOUT_SECONDS, + on_failure=Callback(INGESTION_FAILURE_CALLBACK, timeout=60), + meta={ + "workspace_id": workspace_id, + "document_id": document_id, + "ingestion_run_id": ingestion_run_id, + }, ) return { "document_id": document_id, + "ingestion_run_id": ingestion_run_id, "pages_total": pages_total, "status": "indexing", } + except IngestionFailure as exc: + logger.warning( + "ingest_extract rejected document", + extra={ + "workspace_id": workspace_id, + "document_id": document_id, + "ingestion_run_id": ingestion_run_id, + "failure_category": exc.category.value, + }, + ) + _cleanup_ingestion_artifacts( + workspace_id=workspace_uuid, document_id=document_uuid, include_pages=True + ) + _set_document_status( + workspace_id=workspace_uuid, + document_id=document_uuid, + status="failed", + error_message=str(exc)[:2000], + ) + raise except IntegrityError as exc: logger.exception( "ingest_extract status transition failed due to DB constraint", - extra={"workspace_id": workspace_id, "document_id": document_id}, + extra={ + "workspace_id": workspace_id, + "document_id": document_id, + "ingestion_run_id": ingestion_run_id, + }, ) _set_document_status( workspace_id=workspace_uuid, document_id=document_uuid, status="failed", - error_message=f"Status constraint mismatch: {exc.orig}", + error_message=_failure_message( + IngestionFailureCategory.TRANSIENT_INFRASTRUCTURE, + f"Database status constraint mismatch during extraction: {exc.orig}", + ), ) raise except Exception as exc: # noqa: BLE001 logger.exception( "ingest_extract failed", - extra={"workspace_id": workspace_id, "document_id": document_id}, + extra={ + "workspace_id": workspace_id, + "document_id": document_id, + "ingestion_run_id": ingestion_run_id, + }, + ) + _cleanup_ingestion_artifacts( + workspace_id=workspace_uuid, document_id=document_uuid, include_pages=True ) _set_document_status( workspace_id=workspace_uuid, document_id=document_uuid, status="failed", - error_message=str(exc)[:2000], + error_message=_failure_message( + IngestionFailureCategory.EXTRACTION, + f"Text extraction failed. Upload a valid text-based PDF and retry. Details: {exc}", + ), ) raise finally: diff --git a/worker/jobs/ingest_index.py b/worker/jobs/ingest_index.py index c118aea..2db359d 100644 --- a/worker/jobs/ingest_index.py +++ b/worker/jobs/ingest_index.py @@ -10,6 +10,10 @@ from app.config import settings, utc_today from app.core.errors import BudgetExceededError +from app.core.ingestion_policy import ( + IngestionFailureCategory, + ingestion_failure_message, +) from app.core.token_budget import commit_usage, release_tokens, reserve_tokens from app.db.session import SessionLocal @@ -24,7 +28,13 @@ OVERLAP_TOKENS = 100 -def _set_document_failed(workspace_id: uuid.UUID, document_id: uuid.UUID, error_message: str) -> None: +def _failure_message(category: IngestionFailureCategory, detail: str) -> str: + return ingestion_failure_message(category, detail)[:2000] + + +def _set_document_failed( + workspace_id: uuid.UUID, document_id: uuid.UUID, error_message: str +) -> None: with SessionLocal() as db: db.execute( text( @@ -46,6 +56,31 @@ def _set_document_failed(workspace_id: uuid.UUID, document_id: uuid.UUID, error_ db.commit() +def _cleanup_index_artifacts(workspace_id: uuid.UUID, document_id: uuid.UUID) -> None: + with SessionLocal() as db: + db.execute( + text( + """ + DELETE FROM chunk_embeddings + WHERE workspace_id = :workspace_id + AND document_id = :document_id + """ + ), + {"workspace_id": workspace_id, "document_id": document_id}, + ) + db.execute( + text( + """ + DELETE FROM chunks + WHERE workspace_id = :workspace_id + AND document_id = :document_id + """ + ), + {"workspace_id": workspace_id, "document_id": document_id}, + ) + db.commit() + + def _estimate_embedding_tokens(text_value: str) -> int: return max(1, int(math.ceil((len(text_value) / 4.0) * 1.1))) @@ -81,7 +116,11 @@ def _get_encoding(): return None -def chunk_text(text_value: str, chunk_size_tokens: int = CHUNK_SIZE_TOKENS, overlap_tokens: int = OVERLAP_TOKENS) -> list[str]: +def chunk_text( + text_value: str, + chunk_size_tokens: int = CHUNK_SIZE_TOKENS, + overlap_tokens: int = OVERLAP_TOKENS, +) -> list[str]: normalized = (text_value or "").strip() if not normalized: return [] @@ -126,7 +165,15 @@ def _embedding_to_vector_literal(embedding: list[float]) -> str: return "[" + ",".join(f"{value:.10f}" for value in embedding) + "]" -def ingest_index(workspace_id: str, document_id: str) -> dict: +def _batched(rows: list[dict[str, object]], batch_size: int): + size = max(1, batch_size) + for start in range(0, len(rows), size): + yield rows[start : start + size] + + +def ingest_index( + workspace_id: str, document_id: str, ingestion_run_id: str | None = None +) -> dict: workspace_uuid = uuid.UUID(workspace_id) document_uuid = uuid.UUID(document_id) usage_date = utc_today() @@ -135,18 +182,22 @@ def ingest_index(workspace_id: str, document_id: str) -> dict: try: with SessionLocal() as db: allowed_statuses = _allowed_document_statuses(db) - document_row = db.execute( - text( - """ + document_row = ( + db.execute( + text( + """ SELECT id, status FROM documents WHERE id = :document_id AND workspace_id = :workspace_id LIMIT 1 """ - ), - {"workspace_id": workspace_uuid, "document_id": document_uuid}, - ).mappings().first() + ), + {"workspace_id": workspace_uuid, "document_id": document_uuid}, + ) + .mappings() + .first() + ) if document_row is None: raise ValueError("Document not found for workspace") @@ -154,7 +205,9 @@ def ingest_index(workspace_id: str, document_id: str) -> dict: if "extracting" in allowed_statuses: accepted_current_statuses.add("extracting") if document_row["status"] not in accepted_current_statuses: - raise ValueError(f"Document status must be indexing or uploaded (got: {document_row['status']})") + raise ValueError( + f"Document status must be indexing or uploaded (got: {document_row['status']})" + ) db.execute( text( @@ -170,18 +223,22 @@ def ingest_index(workspace_id: str, document_id: str) -> dict: {"workspace_id": workspace_uuid, "document_id": document_uuid}, ) - pages = db.execute( - text( - """ + pages = ( + db.execute( + text( + """ SELECT page_number, content FROM document_pages WHERE workspace_id = :workspace_id AND document_id = :document_id ORDER BY page_number ASC """ - ), - {"workspace_id": workspace_uuid, "document_id": document_uuid}, - ).mappings().all() + ), + {"workspace_id": workspace_uuid, "document_id": document_uuid}, + ) + .mappings() + .all() + ) db.execute( text( @@ -220,40 +277,61 @@ def ingest_index(workspace_id: str, document_id: str) -> dict: "page_end": page_number, "chunk_index": chunk_index, "content": piece, - "content_hash": hashlib.sha256(piece.encode("utf-8")).hexdigest(), + "content_hash": hashlib.sha256( + piece.encode("utf-8") + ).hexdigest(), "token_count": _estimate_embedding_tokens(piece), } ) chunk_index += 1 + if not chunk_rows: + _set_document_failed( + workspace_uuid, + document_uuid, + _failure_message( + IngestionFailureCategory.UNSUPPORTED_CONTENT, + "No usable text chunks were created from the extracted PDF text.", + ), + ) + return { + "document_id": document_id, + "ingestion_run_id": ingestion_run_id, + "status": "failed", + "chunks_total": 0, + "embeddings_total": 0, + "embedding_tokens_used": 0, + } + with SessionLocal() as db: - for row in chunk_rows: - db.execute( - text( - """ - INSERT INTO chunks ( - id, workspace_id, document_id, page_start, page_end, - chunk_index, content, content_hash, token_count - ) - VALUES ( - :id, :workspace_id, :document_id, :page_start, :page_end, - :chunk_index, :content, :content_hash, :token_count - ) - """ - ), - row, - ) + db.execute( + text( + """ + INSERT INTO chunks ( + id, workspace_id, document_id, page_start, page_end, + chunk_index, content, content_hash, token_count + ) + VALUES ( + :id, :workspace_id, :document_id, :page_start, :page_end, + :chunk_index, :content, :content_hash, :token_count + ) + """ + ), + chunk_rows, + ) db.commit() - client = OpenAI(api_key=settings.OPENAI_API_KEY) + client = OpenAI( + api_key=settings.OPENAI_API_KEY, + timeout=settings.OPENAI_EMBEDDING_TIMEOUT_SECONDS, + ) model = settings.EMBEDDING_MODEL + embedding_batch_size = max(1, int(settings.EMBEDDING_BATCH_SIZE)) embedding_count = 0 total_embedding_tokens = 0 - for row in chunk_rows: - chunk_text_value = str(row["content"]) - estimated_tokens = _estimate_embedding_tokens(chunk_text_value) - + for batch_rows in _batched(chunk_rows, embedding_batch_size): + estimated_tokens = sum(int(row["token_count"]) for row in batch_rows) with SessionLocal() as db: reserve_tokens( db=db, @@ -265,9 +343,31 @@ def ingest_index(workspace_id: str, document_id: str) -> dict: db.commit() outstanding_reservations.append(estimated_tokens) - response = client.embeddings.create(model=model, input=chunk_text_value) - embedding = list(response.data[0].embedding) - vector_literal = _embedding_to_vector_literal(embedding) + response = client.embeddings.create( + model=model, + input=[str(row["content"]) for row in batch_rows], + ) + response_data = list(response.data) + if len(response_data) != len(batch_rows): + raise ValueError( + f"Embedding batch returned {len(response_data)} vectors for {len(batch_rows)} chunks" + ) + if all(getattr(item, "index", None) is not None for item in response_data): + response_data.sort(key=lambda item: int(item.index)) + + embedding_rows: list[dict[str, object]] = [] + for row, embedding_data in zip(batch_rows, response_data, strict=True): + embedding_rows.append( + { + "chunk_id": row["id"], + "workspace_id": workspace_uuid, + "document_id": document_uuid, + "embedding": _embedding_to_vector_literal( + list(embedding_data.embedding) + ), + "embedding_model": model, + } + ) with SessionLocal() as db: db.execute( @@ -277,13 +377,7 @@ def ingest_index(workspace_id: str, document_id: str) -> dict: VALUES (:chunk_id, :workspace_id, :document_id, CAST(:embedding AS vector), :embedding_model) """ ), - { - "chunk_id": row["id"], - "workspace_id": workspace_uuid, - "document_id": document_uuid, - "embedding": vector_literal, - "embedding_model": model, - }, + embedding_rows, ) commit_usage( db=db, @@ -294,11 +388,13 @@ def ingest_index(workspace_id: str, document_id: str) -> dict: db.commit() outstanding_reservations.pop() - embedding_count += 1 + embedding_count += len(batch_rows) total_embedding_tokens += estimated_tokens with SessionLocal() as db: - final_status = "indexed" if "indexed" in _allowed_document_statuses(db) else "ready" + final_status = ( + "indexed" if "indexed" in _allowed_document_statuses(db) else "ready" + ) db.execute( text( """ @@ -310,12 +406,17 @@ def ingest_index(workspace_id: str, document_id: str) -> dict: AND workspace_id = :workspace_id """ ), - {"workspace_id": workspace_uuid, "document_id": document_uuid, "final_status": final_status}, + { + "workspace_id": workspace_uuid, + "document_id": document_uuid, + "final_status": final_status, + }, ) db.commit() return { "document_id": document_id, + "ingestion_run_id": ingestion_run_id, "status": final_status, "chunks_total": len(chunk_rows), "embeddings_total": embedding_count, @@ -333,8 +434,22 @@ def ingest_index(workspace_id: str, document_id: str) -> dict: ) db.commit() except Exception: # noqa: BLE001 - logger.exception("Failed to release reserved tokens", extra={"document_id": document_id}) - _set_document_failed(workspace_uuid, document_uuid, "Insufficient token budget for embeddings") + logger.exception( + "Failed to release reserved tokens", + extra={ + "document_id": document_id, + "ingestion_run_id": ingestion_run_id, + }, + ) + _cleanup_index_artifacts(workspace_uuid, document_uuid) + _set_document_failed( + workspace_uuid, + document_uuid, + _failure_message( + IngestionFailureCategory.BUDGET, + "Insufficient token budget for embeddings. Retry after the workspace budget resets.", + ), + ) raise except Exception as exc: # noqa: BLE001 for reserved in reversed(outstanding_reservations): @@ -348,11 +463,29 @@ def ingest_index(workspace_id: str, document_id: str) -> dict: ) db.commit() except Exception: # noqa: BLE001 - logger.exception("Failed to release reserved tokens", extra={"document_id": document_id}) + logger.exception( + "Failed to release reserved tokens", + extra={ + "document_id": document_id, + "ingestion_run_id": ingestion_run_id, + }, + ) logger.exception( "ingest_index failed", - extra={"workspace_id": workspace_id, "document_id": document_id}, + extra={ + "workspace_id": workspace_id, + "document_id": document_id, + "ingestion_run_id": ingestion_run_id, + }, + ) + _cleanup_index_artifacts(workspace_uuid, document_uuid) + _set_document_failed( + workspace_uuid, + document_uuid, + _failure_message( + IngestionFailureCategory.INDEXING, + f"Indexing failed while creating chunks or embeddings. Details: {exc}", + ), ) - _set_document_failed(workspace_uuid, document_uuid, str(exc)) raise