File size: 9,243 Bytes
3a7eb07
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
"""
Document API Endpoints
Document upload, management, and analysis
"""

import logging
import uuid
from datetime import datetime
from typing import List, Optional

from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException, Query, status
from sqlalchemy import func
from sqlalchemy.orm import Session

from app.api.deps import audit_middleware, get_current_user, get_current_user_id
from app.core.exceptions import NotFoundError
from app.db.session import get_db
from app.models.audit import AuditAction, create_audit_log
from app.models.document import (
    ComplianceStatus,
    Document,
    DocumentCategory,
    DocumentStatus,
)
from app.models.user import User
from app.schemas.document import (
    DocumentAnalysisResponse,
    DocumentCreate,
    DocumentListResponse,
    DocumentResponse,
    DocumentUploadResponse,
)

logger = logging.getLogger(__name__)

router = APIRouter()


@router.get("", response_model=DocumentListResponse)
async def list_documents(
    page: int = Query(1, ge=1),
    page_size: int = Query(20, ge=1, le=100),
    category: Optional[DocumentCategory] = None,
    status: Optional[DocumentStatus] = None,
    search: Optional[str] = None,
    user_id: str = Depends(get_current_user_id),
    db: Session = Depends(get_db),
):
    """
    List user's documents with filtering and pagination.
    """
    query = db.query(Document).filter(Document.user_id == user_id)

    # Apply filters
    if category:
        query = query.filter(Document.category == category)
    if status:
        query = query.filter(Document.status == status)
    if search:
        # Escape SQL wildcard characters to prevent injection
        escaped_search = (
            search.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_")
        )
        safe_pattern = f"%{escaped_search}%"
        query = query.filter(
            Document.title.ilike(safe_pattern, escape="\\")
            | Document.file_name.ilike(safe_pattern, escape="\\")
        )

    # Get total count
    total = query.count()

    # Apply pagination
    offset = (page - 1) * page_size
    documents = (
        query.order_by(Document.created_at.desc()).offset(offset).limit(page_size).all()
    )

    # Calculate stats
    stats = _calculate_document_stats(db, user_id)

    return DocumentListResponse(
        documents=[DocumentResponse(**doc.to_dict()) for doc in documents],
        total=total,
        page=page,
        page_size=page_size,
        stats=stats,
    )


@router.post(
    "", response_model=DocumentUploadResponse, status_code=status.HTTP_202_ACCEPTED
)
async def create_document(
    request: DocumentCreate,
    user_id: str = Depends(get_current_user_id),
    db: Session = Depends(get_db),
):
    """
    Create a new document from UploadThing URL.
    Triggers background processing with AI analysis via task queue.
    """
    # Create document record
    document = Document(
        user_id=user_id,
        title=request.title or request.file_name.rsplit(".", 1)[0],
        file_name=request.file_name,
        file_size=request.file_size,
        file_type=request.file_type,
        file_url=request.file_url,
        status=DocumentStatus.PENDING,
        tags=request.tags or [],
    )

    db.add(document)
    db.flush()  # Flush to get the document.id without committing

    # Create audit log
    audit = create_audit_log(
        action=AuditAction.DOCUMENT_UPLOAD.value,
        user_id=user_id,
        resource_type="document",
        resource_id=str(document.id),
        details={
            "file_name": document.file_name,
            "file_size": document.file_size,
            "file_type": document.file_type,
        },
    )
    db.add(audit)

    # Commit both document and audit atomically
    db.commit()
    db.refresh(document)

    # Trigger background processing
    from app.services.queue import enqueue_document_task

    enqueue_document_task(str(document.id))

    logger.info(f"Document created and enqueued: {document.id} - {document.title}")

    return DocumentUploadResponse(
        id=str(document.id),
        title=document.title,
        file_name=document.file_name,
        status=DocumentStatus.PENDING,
        job_id=str(document.id),  # Using doc ID as job ID for now
        message="Document uploaded successfully. AI analysis queued.",
    )


@router.get("/{document_id}", response_model=DocumentResponse)
async def get_document(
    document_id: uuid.UUID,
    user_id: str = Depends(get_current_user_id),
    db: Session = Depends(get_db),
):
    """
    Get document details by ID.
    """
    document = (
        db.query(Document)
        .filter(Document.id == document_id, Document.user_id == user_id)
        .first()
    )

    if not document:
        raise NotFoundError("Document", document_id)

    return DocumentResponse(**document.to_dict())


@router.delete("/{document_id}")
async def delete_document(
    document_id: uuid.UUID,
    user_id: str = Depends(get_current_user_id),
    db: Session = Depends(get_db),
):
    """
    Delete a document and all associated data.
    """
    document = (
        db.query(Document)
        .filter(Document.id == document_id, Document.user_id == user_id)
        .first()
    )

    if not document:
        raise NotFoundError("Document", document_id)

    # Create audit log before deletion
    audit = create_audit_log(
        action=AuditAction.DOCUMENT_DELETE.value,
        user_id=user_id,
        resource_type="document",
        resource_id=str(document_id),
        details={"file_name": document.file_name},
    )
    db.add(audit)

    # Delete document (cascade will handle embeddings)
    db.delete(document)
    db.commit()

    logger.info(f"Document deleted: {document_id}")

    return {"success": True, "message": "Document deleted successfully"}


@router.get("/{document_id}/analysis", response_model=DocumentAnalysisResponse)
async def get_document_analysis(
    document_id: uuid.UUID,
    user_id: str = Depends(get_current_user_id),
    db: Session = Depends(get_db),
):
    """
    Get AI analysis results for a document.
    """
    document = (
        db.query(Document)
        .filter(Document.id == document_id, Document.user_id == user_id)
        .first()
    )

    if not document:
        raise NotFoundError("Document", document_id)

    if document.status != DocumentStatus.COMPLETED:
        return DocumentAnalysisResponse(
            document_id=str(document.id), status=document.status.value, analysis=None
        )

    analysis = {
        "category": document.category,
        "subcategory": document.subcategory,
        "classification_confidence": document.classification_confidence,
        "summary": document.summary,
        "key_points": document.key_points or [],
        "safety_score": document.safety_score,
        "compliance_status": document.compliance_status,
        "hazards_detected": document.hazards_detected or [],
        "safety_recommendations": document.safety_recommendations or [],
        "entities": {
            k: v if isinstance(v, list) else []
            for k, v in (document.entities or {}).items()
        },
    }

    return DocumentAnalysisResponse(
        document_id=str(document.id), status="completed", analysis=analysis
    )


@router.post("/{document_id}/reanalyze", status_code=status.HTTP_202_ACCEPTED)
async def reanalyze_document(
    document_id: uuid.UUID,
    user_id: str = Depends(get_current_user_id),
    db: Session = Depends(get_db),
):
    """
    Trigger re-analysis of a document.
    """
    document = (
        db.query(Document)
        .filter(Document.id == document_id, Document.user_id == user_id)
        .first()
    )

    if not document:
        raise NotFoundError("Document", document_id)

    # Reset status
    document.status = DocumentStatus.PENDING
    db.commit()

    # Trigger reprocessing
    from app.services.queue import enqueue_document_task

    enqueue_document_task(str(document.id))

    return {"success": True, "message": "Document reanalysis queued"}


def _calculate_document_stats(db: Session, user_id: str) -> dict:
    """Calculate aggregated statistics for user's documents"""
    # Category distribution
    category_counts = (
        db.query(Document.category, func.count(Document.id))
        .filter(Document.user_id == user_id, Document.category.isnot(None))
        .group_by(Document.category)
        .all()
    )

    by_category = {
        cat.value if cat else "other": count for cat, count in category_counts
    }

    # Status distribution
    status_counts = (
        db.query(Document.status, func.count(Document.id))
        .filter(Document.user_id == user_id)
        .group_by(Document.status)
        .all()
    )

    by_status = {
        status.value if status else "unknown": count for status, count in status_counts
    }

    # Average safety score
    avg_score = (
        db.query(func.avg(Document.safety_score))
        .filter(Document.user_id == user_id, Document.safety_score.isnot(None))
        .scalar()
    )

    return {
        "by_category": by_category,
        "by_status": by_status,
        "avg_safety_score": round(avg_score, 2) if avg_score else None,
    }