127 lines
3.9 KiB
Python
127 lines
3.9 KiB
Python
|
# Data models
|
||
|
|
from pydantic import BaseModel, Field
|
||
|
|
from typing import List, Dict, Optional, Any, Union
|
||
|
|
from datetime import datetime
|
||
|
|
from enum import Enum
|
||
|
|
import uuid
|
||
|
|
|
||
|
|
class DocumentStatus(str, Enum):
|
||
|
|
"""Enum for document processing status."""
|
||
|
|
PENDING = "pending"
|
||
|
|
PROCESSING = "processing"
|
||
|
|
COMPLETED = "completed"
|
||
|
|
FAILED = "failed"
|
||
|
|
|
||
|
|
class ComplianceLevel(str, Enum):
|
||
|
|
"""Enum for compliance severity levels."""
|
||
|
|
CRITICAL = "critical"
|
||
|
|
MAJOR = "major"
|
||
|
|
MINOR = "minor"
|
||
|
|
INFO = "info"
|
||
|
|
|
||
|
|
class ComplianceIssue(BaseModel):
|
||
|
|
"""Model for compliance issues found in the document."""
|
||
|
|
id: str = Field(default_factory=lambda: str(uuid.uuid4()))
|
||
|
|
section: str
|
||
|
|
description: str
|
||
|
|
level: ComplianceLevel
|
||
|
|
line_number: Optional[int] = None
|
||
|
|
reasoning: str = "" # Detailed explanation of why this is an issue
|
||
|
|
standard_references: List[str] = [] # References to specific standards or requirements
|
||
|
|
recommendation: str
|
||
|
|
|
||
|
|
class DocumentMetadata(BaseModel):
|
||
|
|
"""Model for document metadata."""
|
||
|
|
filename: str
|
||
|
|
file_type: str
|
||
|
|
file_size: int # In bytes
|
||
|
|
upload_timestamp: datetime = Field(default_factory=datetime.now)
|
||
|
|
last_modified: Optional[datetime] = None
|
||
|
|
|
||
|
|
class DocumentEmbedding(BaseModel):
|
||
|
|
"""Model for document embeddings."""
|
||
|
|
embedding_id: str
|
||
|
|
embedding_model: str
|
||
|
|
vector_db: str
|
||
|
|
sections: Dict[str, str] # Section name to section ID in vector DB
|
||
|
|
|
||
|
|
class ComplianceReport(BaseModel):
|
||
|
|
"""Model for the generated compliance report."""
|
||
|
|
report_id: str = Field(default_factory=lambda: str(uuid.uuid4()))
|
||
|
|
document_id: str
|
||
|
|
timestamp: datetime = Field(default_factory=datetime.now)
|
||
|
|
compliance_score: float # 0.0 to 1.0
|
||
|
|
summary: str
|
||
|
|
issues: List[ComplianceIssue] = []
|
||
|
|
applied_standards: List[str] = [] # Standards used for analysis
|
||
|
|
|
||
|
|
@property
|
||
|
|
def critical_issues_count(self) -> int:
|
||
|
|
return sum(1 for issue in self.issues if issue.level == ComplianceLevel.CRITICAL)
|
||
|
|
|
||
|
|
@property
|
||
|
|
def major_issues_count(self) -> int:
|
||
|
|
return sum(1 for issue in self.issues if issue.level == ComplianceLevel.MAJOR)
|
||
|
|
|
||
|
|
@property
|
||
|
|
def minor_issues_count(self) -> int:
|
||
|
|
return sum(1 for issue in self.issues if issue.level == ComplianceLevel.MINOR)
|
||
|
|
|
||
|
|
@property
|
||
|
|
def info_issues_count(self) -> int:
|
||
|
|
return sum(1 for issue in self.issues if issue.level == ComplianceLevel.INFO)
|
||
|
|
|
||
|
|
class Document(BaseModel):
|
||
|
|
"""Model for document tracking."""
|
||
|
|
id: str = Field(default_factory=lambda: str(uuid.uuid4()))
|
||
|
|
metadata: DocumentMetadata
|
||
|
|
embedding: Optional[DocumentEmbedding] = None
|
||
|
|
status: DocumentStatus = DocumentStatus.PENDING
|
||
|
|
version: int = 1 # Incremented on resubmissions
|
||
|
|
reports: List[str] = [] # List of report IDs
|
||
|
|
|
||
|
|
class DocumentUploadResponse(BaseModel):
|
||
|
|
"""Response model for document uploads."""
|
||
|
|
document_id: str
|
||
|
|
filename: str
|
||
|
|
status: DocumentStatus
|
||
|
|
message: str
|
||
|
|
|
||
|
|
class DocumentAnalysisResponse(BaseModel):
|
||
|
|
"""Response model for document analysis retrieval."""
|
||
|
|
document_id: str
|
||
|
|
status: DocumentStatus
|
||
|
|
report: Optional[ComplianceReport] = None
|
||
|
|
message: str
|
||
|
|
|
||
|
|
|
||
|
|
class RequirementSeverity(str, Enum):
|
||
|
|
"""Enum for requirement severity levels."""
|
||
|
|
CRITICAL = "critical"
|
||
|
|
MAJOR = "major"
|
||
|
|
MINOR = "minor"
|
||
|
|
INFO = "info"
|
||
|
|
|
||
|
|
|
||
|
|
class Requirement(BaseModel):
|
||
|
|
"""Model for a compliance requirement."""
|
||
|
|
id: str = Field(default_factory=lambda: str(uuid.uuid4()))
|
||
|
|
description: str
|
||
|
|
severity: RequirementSeverity
|
||
|
|
details: Optional[str] = None
|
||
|
|
|
||
|
|
|
||
|
|
class Standard(BaseModel):
|
||
|
|
"""Model for a compliance standard."""
|
||
|
|
id: str = Field(default_factory=lambda: str(uuid.uuid4()))
|
||
|
|
name: str
|
||
|
|
description: str
|
||
|
|
requirements: List[Requirement] = []
|
||
|
|
|
||
|
|
|
||
|
|
class StandardUploadResponse(BaseModel):
|
||
|
|
"""Response model for standard uploads."""
|
||
|
|
standard_id: str
|
||
|
|
name: str
|
||
|
|
requirement_count: int
|
||
|
|
message: str
|