Harden document ingestion lifecycle

.env.example

@@ -8,6 +8,7 @@ INGESTION_SERVICE_URL="http://ingestion:8000"
 RETRIEVAL_SERVICE_URL="http://retrieval:8000"
 QUEUE_URL="amqp://guest:guest@rabbitmq:5672"
 INGESTION_QUEUE="document_ingestion"
+MAX_DOCUMENT_BYTES="1000000"
 MINIO_ENDPOINT="http://minio:9000"
 MINIO_ROOT_USER="grounded"
 MINIO_ROOT_PASSWORD="groundedsecret"

apps/gateway/src/main.rs

@@ -76,8 +76,12 @@ async fn main() {
         .route("/api/tenants", post(create_tenant))
         .route("/api/tenants/current", get(current_tenant))
         .route("/api/documents", get(list_documents).post(create_document))
-        .route("/api/documents/:document_id", get(get_document))
+        .route("/api/documents/:document_id", get(get_document).delete(delete_document))
         .route("/api/documents/:document_id/jobs", get(list_document_jobs))
+        .route(
+            "/api/documents/:document_id/jobs/:job_id/retry",
+            post(retry_document_job),
+        )
         .route("/api/questions", post(answer_question))
         .layer(CorsLayer::permissive())
         .layer(TraceLayer::new_for_http())
@@ -151,6 +155,20 @@ async fn get_document(
     .await
 }
 
+async fn delete_document(
+    State(state): State<AppState>,
+    headers: HeaderMap,
+    Path(document_id): Path<String>,
+) -> Result<impl IntoResponse, StatusCode> {
+    proxy_delete(
+        &state.client,
+        &state.services.ingestion,
+        &format!("/documents/{document_id}"),
+        &headers,
+    )
+    .await
+}
+
 async fn list_document_jobs(
     State(state): State<AppState>,
     headers: HeaderMap,
@@ -165,6 +183,21 @@ async fn list_document_jobs(
     .await
 }
 
+async fn retry_document_job(
+    State(state): State<AppState>,
+    headers: HeaderMap,
+    Path((document_id, job_id)): Path<(String, String)>,
+) -> Result<impl IntoResponse, StatusCode> {
+    proxy_json_with_headers(
+        &state.client,
+        &state.services.ingestion,
+        &format!("/documents/{document_id}/jobs/{job_id}/retry"),
+        Value::Null,
+        &headers,
+    )
+    .await
+}
+
 async fn register(
     State(state): State<AppState>,
     headers: HeaderMap,

apps/web/app/api/web/documents/[documentId]/jobs/[jobId]/retry/route.ts

@@ -0,0 +1,19 @@
+import { NextRequest } from "next/server";
+import { proxyPrivate } from "@/lib/server-gateway";
+
+type RouteContext = {
+  params: Promise<{
+    documentId: string;
+    jobId: string;
+  }>;
+};
+
+export async function POST(request: NextRequest, context: RouteContext) {
+  const { documentId, jobId } = await context.params;
+
+  return proxyPrivate({
+    method: "POST",
+    path: `/api/documents/${documentId}/jobs/${jobId}/retry`,
+    request
+  });
+}

apps/web/app/api/web/documents/[documentId]/route.ts

@@ -0,0 +1,18 @@
+import { NextRequest } from "next/server";
+import { proxyPrivate } from "@/lib/server-gateway";
+
+type RouteContext = {
+  params: Promise<{
+    documentId: string;
+  }>;
+};
+
+export async function DELETE(request: NextRequest, context: RouteContext) {
+  const { documentId } = await context.params;
+
+  return proxyPrivate({
+    method: "DELETE",
+    path: `/api/documents/${documentId}`,
+    request
+  });
+}

apps/web/lib/api.ts

@@ -144,6 +144,18 @@ export async function queueDocument(payload: DocumentRequest) {
   });
 }
 
+export async function deleteDocument(documentId: string) {
+  return request<DocumentSummary>(`/api/web/documents/${documentId}`, {
+    method: "DELETE"
+  });
+}
+
+export async function retryIngestionJob(documentId: string, jobId: string) {
+  return request(`/api/web/documents/${documentId}/jobs/${jobId}/retry`, {
+    method: "POST"
+  });
+}
+
 async function request<T>(path: string, options: RequestOptions): Promise<T> {
   const response = await fetch(path, {
     method: options.method,

apps/web/lib/http.ts

@@ -27,6 +27,10 @@ export function resolveErrorMessage(payload: unknown, status: number) {
       return payload.detail;
     }
 
+    if (isMessagePayload(payload.detail)) {
+      return payload.detail.message;
+    }
+
     if (Array.isArray(payload.detail)) {
       const message = payload.detail
         .map((item) => resolveValidationMessage(item))
@@ -50,6 +54,13 @@ function isErrorPayload(payload: unknown): payload is { detail?: unknown; messag
   return typeof payload === "object" && payload !== null;
 }
 
+function isMessagePayload(payload: unknown): payload is { message: string } {
+  return typeof payload === "object"
+    && payload !== null
+    && "message" in payload
+    && typeof (payload as { message?: unknown }).message === "string";
+}
+
 function resolveValidationMessage(item: unknown) {
   if (typeof item !== "object" || item === null || !("msg" in item)) {
     return "";

services/ingestion/app/config.py

@@ -14,6 +14,7 @@ class Settings(BaseSettings):
     minio_bucket: str = Field(default="grounded-documents")
     queue_url: str = Field(default="amqp://guest:guest@localhost:5672")
     ingestion_queue: str = Field(default="document_ingestion")
+    max_document_bytes: int = Field(default=1_000_000)
 
 
 @lru_cache

services/ingestion/app/errors.py

@@ -2,18 +2,38 @@
 
 
 class IngestionError:
+    @staticmethod
+    def response(status_code: int, message: str, code: str) -> HTTPException:
+        return HTTPException(status_code=status_code, detail={"message": message, "code": code})
+
     @staticmethod
     def invalid_token() -> HTTPException:
-        return HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid token")
+        return IngestionError.response(status.HTTP_401_UNAUTHORIZED, "Invalid token", "invalid_token")
 
     @staticmethod
     def tenant_required() -> HTTPException:
-        return HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="Tenant context is required")
+        return IngestionError.response(status.HTTP_403_FORBIDDEN, "Tenant context is required", "tenant_required")
 
     @staticmethod
     def document_not_found() -> HTTPException:
-        return HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Document was not found")
+        return IngestionError.response(status.HTTP_404_NOT_FOUND, "Document was not found", "document_not_found")
+
+    @staticmethod
+    def job_not_found() -> HTTPException:
+        return IngestionError.response(status.HTTP_404_NOT_FOUND, "Ingestion job was not found", "job_not_found")
 
     @staticmethod
     def unsupported_content_type() -> HTTPException:
-        return HTTPException(status_code=status.HTTP_415_UNSUPPORTED_MEDIA_TYPE, detail="Content type is not supported")
+        return IngestionError.response(status.HTTP_415_UNSUPPORTED_MEDIA_TYPE, "Content type is not supported", "unsupported_content_type")
+
+    @staticmethod
+    def document_too_large() -> HTTPException:
+        return IngestionError.response(status.HTTP_413_REQUEST_ENTITY_TOO_LARGE, "Document content is too large", "document_too_large")
+
+    @staticmethod
+    def duplicate_document() -> HTTPException:
+        return IngestionError.response(status.HTTP_409_CONFLICT, "Document content already exists in this tenant", "duplicate_document")
+
+    @staticmethod
+    def retry_not_allowed() -> HTTPException:
+        return IngestionError.response(status.HTTP_409_CONFLICT, "Only failed ingestion jobs can be retried", "retry_not_allowed")

services/ingestion/app/main.py

@@ -46,3 +46,22 @@ async def list_jobs(
     service: Annotated[DocumentService, Depends(document_service)],
 ) -> list[IngestionJobSummary]:
     return await service.list_jobs(claims["tid"], document_id)
+
+
+@app.delete("/documents/{document_id}")
+async def delete_document(
+    document_id: str,
+    claims: Annotated[dict, Depends(current_claims)],
+    service: Annotated[DocumentService, Depends(document_service)],
+) -> DocumentSummary:
+    return await service.delete_document(claims["tid"], document_id)
+
+
+@app.post("/documents/{document_id}/jobs/{job_id}/retry")
+async def retry_job(
+    document_id: str,
+    job_id: str,
+    claims: Annotated[dict, Depends(current_claims)],
+    service: Annotated[DocumentService, Depends(document_service)],
+) -> IngestionJobSummary:
+    return await service.retry_job(claims["tid"], document_id, job_id)

services/ingestion/app/repositories.py

@@ -29,6 +29,21 @@ async def create_document(self, tenant_id: str, user_id: str, title: str, filena
         )
         return dict(result.mappings().one())
 
+    async def find_document_by_checksum(self, tenant_id: str, checksum: str) -> dict | None:
+        result = await self.session.execute(
+            text(
+                'SELECT d."id", d."tenantId", d."title", d."filename", d."contentType", d."status", d."currentVersion", d."createdAt", d."updatedAt" '
+                'FROM "DocumentVersion" v '
+                'JOIN "Document" d ON d."id" = v."documentId" '
+                'WHERE v."tenantId" = :tenant_id AND v."checksum" = :checksum AND d."deletedAt" IS NULL '
+                'ORDER BY v."createdAt" DESC '
+                'LIMIT 1'
+            ),
+            {"tenant_id": tenant_id, "checksum": checksum},
+        )
+        row = result.mappings().first()
+        return dict(row) if row else None
+
     async def create_version(self, tenant_id: str, document_id: str, object_key: str, checksum: str, byte_size: int, extracted_text: str) -> dict:
         result = await self.session.execute(
             text(
@@ -101,6 +116,42 @@ async def list_jobs_for_document(self, tenant_id: str, document_id: str) -> list
         )
         return [dict(row) for row in result.mappings().all()]
 
+    async def find_job_for_document(self, tenant_id: str, document_id: str, job_id: str) -> dict | None:
+        result = await self.session.execute(
+            text(
+                'SELECT "id", "tenantId", "documentId", "documentVersionId", "status", "attempts", "errorCode", "errorMessage", "queuedAt", "startedAt", "finishedAt" '
+                'FROM "IngestionJob" '
+                'WHERE "tenantId" = :tenant_id AND "documentId" = :document_id AND "id" = :job_id'
+            ),
+            {"tenant_id": tenant_id, "document_id": document_id, "job_id": job_id},
+        )
+        row = result.mappings().first()
+        return dict(row) if row else None
+
+    async def soft_delete_document(self, tenant_id: str, document_id: str) -> dict | None:
+        result = await self.session.execute(
+            text(
+                'UPDATE "Document" '
+                'SET "deletedAt" = now(), "status" = CAST(:status AS "DocumentStatus"), "updatedAt" = now() '
+                'WHERE "tenantId" = :tenant_id AND "id" = :document_id AND "deletedAt" IS NULL '
+                'RETURNING "id", "tenantId", "title", "filename", "contentType", "status", "currentVersion", "createdAt", "updatedAt"'
+            ),
+            {"tenant_id": tenant_id, "document_id": document_id, "status": "deleted"},
+        )
+        row = result.mappings().first()
+        return dict(row) if row else None
+
+    async def create_retry_job(self, tenant_id: str, document_id: str, document_version_id: str) -> dict:
+        await self.session.execute(
+            text(
+                'UPDATE "Document" '
+                'SET "status" = CAST(:document_status AS "DocumentStatus"), "updatedAt" = now() '
+                'WHERE "tenantId" = :tenant_id AND "id" = :document_id AND "deletedAt" IS NULL'
+            ),
+            {"tenant_id": tenant_id, "document_id": document_id, "document_status": "queued"},
+        )
+        return await self.create_job(tenant_id, document_id, document_version_id)
+
     async def mark_publish_failed(self, tenant_id: str, document_id: str, job_id: str, error_message: str) -> None:
         await self.session.execute(
             text(

services/ingestion/app/services.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from hashlib import sha256
+from app.config import get_settings
 from app.errors import IngestionError
 from app.queue import IngestionQueue
 from app.repositories import DocumentRepository
@@ -17,16 +18,26 @@ def __init__(self, repository: DocumentRepository, storage: ObjectStorage, queue
         self.queue = queue
 
     async def create_document(self, tenant_id: str, user_id: str, request: DocumentCreateRequest) -> DocumentCreateResponse:
-        if request.content_type not in self.supported_content_types:
+        content_type = request.content_type.strip().lower()
+        if content_type not in self.supported_content_types:
             raise IngestionError.unsupported_content_type()
 
+        byte_size = len(request.content.encode("utf-8"))
+        if byte_size > get_settings().max_document_bytes:
+            raise IngestionError.document_too_large()
+
         title = request.title or request.filename
         checksum = sha256(request.content.encode("utf-8")).hexdigest()
+        duplicate = await self.repository.find_document_by_checksum(tenant_id, checksum)
+        if duplicate:
+            raise IngestionError.duplicate_document()
 
         try:
-            document = await self.repository.create_document(tenant_id, user_id, title, request.filename, request.content_type)
+            document = await self.repository.create_document(tenant_id, user_id, title, request.filename, content_type)
             object_key = f"{tenant_id}/{document['id']}/1/{request.filename}"
-            byte_size = self.storage.put_text(object_key, request.content, request.content_type)
+            stored_byte_size = self.storage.put_text(object_key, request.content, content_type)
+            if stored_byte_size != byte_size:
+                byte_size = stored_byte_size
             version = await self.repository.create_version(tenant_id, document["id"], object_key, checksum, byte_size, request.content)
             job = await self.repository.create_job(tenant_id, document["id"], version["id"])
             await self.repository.commit()
@@ -62,6 +73,36 @@ async def list_jobs(self, tenant_id: str, document_id: str) -> list[IngestionJob
             raise IngestionError.document_not_found()
         return [self.job_summary(job) for job in await self.repository.list_jobs_for_document(tenant_id, document_id)]
 
+    async def delete_document(self, tenant_id: str, document_id: str) -> DocumentSummary:
+        document = await self.repository.soft_delete_document(tenant_id, document_id)
+        if not document:
+            raise IngestionError.document_not_found()
+        await self.repository.commit()
+        return self.document_summary(document)
+
+    async def retry_job(self, tenant_id: str, document_id: str, job_id: str) -> IngestionJobSummary:
+        document = await self.repository.find_document(tenant_id, document_id)
+        if not document:
+            raise IngestionError.document_not_found()
+
+        job = await self.repository.find_job_for_document(tenant_id, document_id, job_id)
+        if not job:
+            raise IngestionError.job_not_found()
+        if job["status"] != "failed":
+            raise IngestionError.retry_not_allowed()
+
+        retry = await self.repository.create_retry_job(tenant_id, document_id, job["documentVersionId"])
+        await self.repository.commit()
+
+        try:
+            await self.queue.publish_job(retry["id"])
+        except Exception as exc:
+            await self.repository.mark_publish_failed(tenant_id, document_id, retry["id"], str(exc))
+            await self.repository.commit()
+            raise
+
+        return self.job_summary(retry)
+
     def document_summary(self, document: dict) -> DocumentSummary:
         return DocumentSummary(
             id=document["id"],

services/retrieval/app/repositories.py

@@ -18,7 +18,7 @@ async def load_chunks(self, tenant_id: str, chunk_ids: list[str]) -> list[dict]:
                 'SELECT c."id", c."tenantId", c."documentId", c."content", c."sourceStart", c."sourceEnd", d."title" AS "documentTitle" '
                 'FROM "DocumentChunk" c '
                 'JOIN "Document" d ON d."id" = c."documentId" '
-                'WHERE c."tenantId" = :tenant_id AND c."id" = ANY(:chunk_ids)'
+                'WHERE c."tenantId" = :tenant_id AND c."id" = ANY(:chunk_ids) AND d."deletedAt" IS NULL'
             ),
             {"tenant_id": tenant_id, "chunk_ids": chunk_ids},
         )