| """ |
| Data models for documents and document chunks. |
| """ |
| from typing import List, Dict, Any, Optional |
| import uuid |
| from PIL import Image |
|
|
|
|
| class RawDocument: |
| """Represents a raw document with extracted content.""" |
|
|
| def __init__( |
| self, |
| filename: str, |
| file_type: str, |
| pages: List[Dict[str, Any]], |
| raw_text: str, |
| raw_tables: List[Dict[str, Any]], |
| total_pages: int, |
| metadata: Optional[Dict[str, Any]] = None, |
| ): |
| """ |
| Initialize a RawDocument. |
| |
| Args: |
| filename: Name of the document file |
| file_type: Type of file (e.g., 'pdf', 'docx') |
| pages: List of page dictionaries with 'page_num' and 'text' keys |
| raw_text: Full extracted text from the document |
| raw_tables: List of tables extracted from the document |
| total_pages: Total number of pages in the document |
| metadata: Additional metadata (file_path, author, etc.) |
| """ |
| self.filename = filename |
| self.file_type = file_type |
| self.pages = pages |
| self.raw_text = raw_text |
| self.raw_tables = raw_tables |
| self.total_pages = total_pages |
| self.metadata = metadata or {} |
|
|
| def __repr__(self) -> str: |
| return f"RawDocument(filename={self.filename}, pages={self.total_pages})" |
|
|
|
|
| class DocumentChunk: |
| """Represents a chunk of document content with metadata.""" |
|
|
| def __init__( |
| self, |
| content: str, |
| chunk_type: str, |
| page_number: int, |
| metadata: Dict[str, Any] = None, |
| chunk_id: str = None, |
| ): |
| """ |
| Initialize a DocumentChunk. |
| |
| Args: |
| content: The text content of the chunk |
| chunk_type: Type of chunk (e.g., 'text', 'table') |
| page_number: Page number where this chunk appears |
| metadata: Additional metadata about the chunk |
| chunk_id: Unique identifier for the chunk (auto-generated if not provided) |
| """ |
| self.content = content |
| self.chunk_type = chunk_type |
| self.page_number = page_number |
| self.metadata = metadata or {} |
| self.chunk_id = chunk_id or str(uuid.uuid4()) |
|
|
| def __repr__(self) -> str: |
| return ( |
| f"DocumentChunk(type={self.chunk_type}, page={self.page_number}, " |
| f"length={len(self.content)})" |
| ) |
|
|
|
|
| class TableExtraction: |
| """Represents a table extracted from a document.""" |
|
|
| def __init__( |
| self, |
| headers: List[str], |
| rows: List[List[str]], |
| page_number: int, |
| schema_summary: str, |
| table_id: str = None, |
| ): |
| """ |
| Initialize a TableExtraction. |
| |
| Args: |
| headers: List of column headers |
| rows: List of rows, each containing cell values |
| page_number: Page number where this table appears |
| schema_summary: Summary description of the table schema |
| table_id: Unique identifier for the table (auto-generated if not provided) |
| """ |
| self.headers = headers |
| self.rows = rows |
| self.page_number = page_number |
| self.schema_summary = schema_summary |
| self.table_id = table_id or str(uuid.uuid4()) |
|
|
| def __repr__(self) -> str: |
| return ( |
| f"TableExtraction(columns={len(self.headers)}, " |
| f"rows={len(self.rows)}, page={self.page_number})" |
| ) |
|
|
|
|
| class ProcessedDocument: |
| """Represents a fully processed document with text chunks and tables.""" |
|
|
| def __init__( |
| self, |
| filename: str, |
| text_chunks: List[DocumentChunk], |
| tables: List["TableExtraction"], |
| total_pages: int, |
| file_type: str, |
| images: Optional[List["ImageExtraction"]] = None, |
| layout: Optional["LayoutExtraction"] = None, |
| metadata: Optional["MetadataExtraction"] = None, |
| ): |
| """ |
| Initialize a ProcessedDocument. |
| |
| Args: |
| filename: Name of the document file |
| text_chunks: List of text chunks extracted from the document |
| tables: List of tables extracted from the document |
| total_pages: Total number of pages in the document |
| file_type: Type of file (e.g., 'pdf', 'docx') |
| images: List of images extracted from the document (Phase 2) |
| layout: Layout information (Phase 2) |
| metadata: Document metadata (Phase 2) |
| """ |
| self.filename = filename |
| self.text_chunks = text_chunks |
| self.tables = tables |
| self.total_pages = total_pages |
| self.file_type = file_type |
| self.images = images or [] |
| self.layout = layout |
| self.metadata = metadata |
|
|
| def __repr__(self) -> str: |
| return ( |
| f"ProcessedDocument(filename={self.filename}, " |
| f"text_chunks={len(self.text_chunks)}, " |
| f"tables={len(self.tables)}, " |
| f"images={len(self.images)})" |
| ) |
|
|
|
|
| class ImageExtraction: |
| """Represents an image extracted from a document.""" |
|
|
| def __init__( |
| self, |
| image: Image.Image, |
| page_number: int, |
| image_index: int, |
| width: int, |
| height: int, |
| format: str, |
| image_id: str = None, |
| ): |
| """ |
| Initialize an ImageExtraction. |
| |
| Args: |
| image: PIL Image object |
| page_number: Page number where this image appears |
| image_index: Index of image on the page |
| width: Image width in pixels |
| height: Image height in pixels |
| format: Image format (png, jpg, etc.) |
| image_id: Unique identifier for the image (auto-generated if not provided) |
| """ |
| self.image = image |
| self.page_number = page_number |
| self.image_index = image_index |
| self.width = width |
| self.height = height |
| self.format = format |
| self.image_id = image_id or str(uuid.uuid4()) |
|
|
| def __repr__(self) -> str: |
| return ( |
| f"ImageExtraction(page={self.page_number}, " |
| f"size={self.width}x{self.height}, format={self.format})" |
| ) |
|
|
|
|
| class LayoutExtraction: |
| """Represents document layout and structure information.""" |
|
|
| def __init__( |
| self, |
| sections: List[Dict[str, Any]], |
| hierarchy: Dict[str, Any], |
| page_layouts: List[Dict[str, Any]], |
| total_pages: int, |
| ): |
| """ |
| Initialize a LayoutExtraction. |
| |
| Args: |
| sections: List of document sections with hierarchy info |
| hierarchy: Document hierarchy tree |
| page_layouts: Layout information per page |
| total_pages: Total number of pages |
| """ |
| self.sections = sections |
| self.hierarchy = hierarchy |
| self.page_layouts = page_layouts |
| self.total_pages = total_pages |
|
|
| def __repr__(self) -> str: |
| return ( |
| f"LayoutExtraction(sections={len(self.sections)}, " |
| f"pages={self.total_pages})" |
| ) |
|
|
|
|
| class MetadataExtraction: |
| """Represents document metadata.""" |
|
|
| def __init__( |
| self, |
| title: Optional[str] = None, |
| author: Optional[str] = None, |
| subject: Optional[str] = None, |
| keywords: Optional[List[str]] = None, |
| creator: Optional[str] = None, |
| producer: Optional[str] = None, |
| creation_date: Optional[str] = None, |
| modification_date: Optional[str] = None, |
| page_count: Optional[int] = None, |
| custom_properties: Optional[Dict[str, Any]] = None, |
| ): |
| """ |
| Initialize a MetadataExtraction. |
| |
| Args: |
| title: Document title |
| author: Document author |
| subject: Document subject |
| keywords: List of keywords |
| creator: Creator application |
| producer: Producer application |
| creation_date: Creation date |
| modification_date: Last modification date |
| page_count: Number of pages |
| custom_properties: Additional custom properties |
| """ |
| self.title = title |
| self.author = author |
| self.subject = subject |
| self.keywords = keywords or [] |
| self.creator = creator |
| self.producer = producer |
| self.creation_date = creation_date |
| self.modification_date = modification_date |
| self.page_count = page_count |
| self.custom_properties = custom_properties or {} |
|
|
| def __repr__(self) -> str: |
| return ( |
| f"MetadataExtraction(title={self.title}, " |
| f"author={self.author}, pages={self.page_count})" |
| ) |
|
|
