"""Status polling endpoint for documents and reports.""" from collections import Counter from datetime import UTC, datetime from typing import Any from fastapi import APIRouter, Depends, HTTPException, Query, Request from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from app.config import settings from app.db.database import get_db from app.db.models import Document, IngestStatus, Report from app.models.schemas import ( DocumentBatchStatusRequest, DocumentBatchStatusResponse, DocumentStatusItem, ReportStatusResponse, ) router = APIRouter() def _utc_age_seconds(now_utc: datetime, then: datetime | None) -> float: """Compute age in seconds, tolerant of naive DB datetimes. SQLite commonly returns naive datetimes. We treat naive values as UTC. """ if then is None: return 0.0 if then.tzinfo is None: then = then.replace(tzinfo=UTC) return (now_utc - then).total_seconds() @router.get("/documents/{document_id}/info") async def document_info( document_id: str, request: Request, db: AsyncSession = Depends(get_db), ) -> dict[str, Any]: """Return document metadata including ingested chunk count. Args: document_id: UUID of the document. request: Provides ``state.tenant_id``. db: Injected database session. Returns: Dict with document metadata and chunk count. """ tenant_id: str = request.state.tenant_id doc = await db.get(Document, document_id) if doc is None or doc.tenant_id != tenant_id: raise HTTPException(status_code=404, detail="Document not found") from app.vectorstore.factory import get_vectorstore try: vs = get_vectorstore() chunk_count = vs.count_for_doc(document_id) except Exception: chunk_count = 0 return { "document_id": document_id, "filename": doc.filename, "status": doc.status.value, "chunk_count": chunk_count, "created_at": doc.created_at.isoformat(), "survey_level": doc.survey_level, } @router.get("/documents/{document_id}/status") async def document_status( document_id: str, request: Request, db: AsyncSession = Depends(get_db), ) -> dict[str, str]: """Return ingestion status of an uploaded document. Args: document_id: UUID of the document. request: Provides ``state.tenant_id``. db: Injected database session. Returns: Dict with ``document_id``, ``status``, and optional ``error``. Raises: HTTPException: 404 if not found or wrong tenant. """ tenant_id: str = request.state.tenant_id doc = await db.get(Document, document_id) if doc is None or doc.tenant_id != tenant_id: raise HTTPException(status_code=404, detail="Document not found") payload: dict[str, str] = { "document_id": document_id, "status": doc.status.value, } if doc.error_message: payload["error"] = doc.error_message return payload @router.get("/documents") async def list_documents( request: Request, db: AsyncSession = Depends(get_db), limit: int = Query(default=100, ge=1), offset: int = Query(default=0, ge=0), ) -> dict[str, Any]: """List uploaded documents for the tenant (newest first). Use pagination for large libraries; increase ``max_upload_batch_files`` / ``documents_list_max_limit`` via settings when you need higher throughput. """ tenant_id: str = request.state.tenant_id cap = settings.documents_list_max_limit limit = min(limit, cap) result = await db.execute( select(Document) .where(Document.tenant_id == tenant_id) .order_by(Document.created_at.desc()) .limit(limit) .offset(offset) ) rows = result.scalars().all() return { "tenant_id": tenant_id, "limit": limit, "offset": offset, "documents": [ { "document_id": d.id, "filename": d.filename, "status": d.status.value, "created_at": d.created_at.isoformat(), "survey_level": d.survey_level, "error": d.error_message, } for d in rows ], } @router.post("/documents/batch-status", response_model=DocumentBatchStatusResponse) async def documents_batch_status( body: DocumentBatchStatusRequest, request: Request, db: AsyncSession = Depends(get_db), ) -> DocumentBatchStatusResponse: """Return ingestion status for many document IDs in one request.""" tenant_id: str = request.state.tenant_id if not body.document_ids: return DocumentBatchStatusResponse( items=[], pending=0, processing=0, complete=0, failed=0, ) result = await db.execute( select(Document).where( Document.tenant_id == tenant_id, Document.id.in_(body.document_ids), ) ) found = {d.id: d for d in result.scalars().all()} items: list[DocumentStatusItem] = [] counts: Counter[str] = Counter() now = datetime.now(UTC) stale_cutoff_s = int(settings.ingest_timeout_seconds) mutated = False for did in body.document_ids: doc = found.get(did) if doc is None: items.append( DocumentStatusItem( document_id=did, status="not_found", filename="", error="Document not found or not owned by this tenant", ) ) continue st = doc.status.value if st == "processing": # If a worker gets stuck (network hang, loader deadlock), do not block # the UI forever: mark as failed after a safe timeout. age_s = _utc_age_seconds(now, doc.updated_at) if age_s > stale_cutoff_s: doc.status = IngestStatus.failed doc.error_message = ( f"Ingestion timed out after {stale_cutoff_s}s. " "The file may be too large/corrupted, or embedding/indexing may be unavailable." ) mutated = True st = "failed" # Treat unknown / not_found as failed for aggregates so the UI never # shows "0 failed" while every row is a terminal error (e.g. tenant mismatch). if st == "not_found": counts["failed"] += 1 elif st in ("pending", "processing", "complete", "failed"): counts[st] += 1 items.append( DocumentStatusItem( document_id=did, status=st, filename=doc.filename, error=doc.error_message, ) ) if mutated: await db.commit() return DocumentBatchStatusResponse( items=items, pending=counts["pending"], processing=counts["processing"], complete=counts["complete"], failed=counts["failed"], ) @router.get("/documents/tenant-chunk-summary") async def tenant_chunk_summary(request: Request) -> dict[str, str | int]: """Return total vector-store chunks for this tenant (all reference documents).""" tenant_id: str = request.state.tenant_id from app.vectorstore.factory import get_vectorstore try: vs = get_vectorstore() n = vs.count(tenant_id) except Exception: n = 0 return {"tenant_id": tenant_id, "indexed_chunk_count": n} @router.get("/reports/{report_id}/status", response_model=ReportStatusResponse) async def report_status( report_id: str, request: Request, db: AsyncSession = Depends(get_db), ) -> ReportStatusResponse: """Return generation status of a report. Args: report_id: UUID of the report. request: Provides ``state.tenant_id``. db: Injected database session. Returns: :class:`~app.models.schemas.ReportStatusResponse`. Raises: HTTPException: 404 if not found or wrong tenant. """ tenant_id: str = request.state.tenant_id report = await db.get(Report, report_id) if report is None or report.tenant_id != tenant_id: raise HTTPException(status_code=404, detail="Report not found") return ReportStatusResponse( report_id=report_id, status=report.status.value, created_at=report.created_at, updated_at=report.updated_at, error_message=report.error_message, survey_level=report.survey_level, )