structx 0.4.8__py3-none-any.whl

structx/extraction/processors/model_operations.py

@@ -0,0 +1,233 @@
+"""
+Model operations including schema generation and custom model processing.
+"""
+
+import json
+from typing import List, Tuple, Type
+
+from loguru import logger
+from pydantic import BaseModel
+
+from structx.core.exceptions import ExtractionError
+from structx.core.models import ExtractionGuide, ExtractionRequest, QueryRefinement
+from structx.extraction.core.llm_core import LLMCore
+from structx.extraction.core.model_utils import ModelUtils
+from structx.extraction.generator import ModelGenerator
+from structx.extraction.processors.data_content import ContentAnalyzer
+from structx.utils.helpers import (
+    convert_pydantic_v1_to_v2,
+    handle_errors,
+    sanitize_regex_patterns,
+)
+from structx.utils.prompts import (
+    custom_model_guide_template,
+    refinement_system_prompt,
+    refinement_template,
+    schema_system_prompt,
+    schema_template,
+)
+from structx.utils.usage import ExtractionStep
+
+
+class ModelOperations:
+    """
+    Handles all model-related operations including schema generation and custom model processing.
+    """
+
+    def __init__(self, llm_core: LLMCore):
+        """
+        Initialize model operations.
+
+        Args:
+            llm_core: LLM core for completions
+        """
+        self.llm_core = llm_core
+
+    @handle_errors(error_message="Schema generation failed", error_type=ExtractionError)
+    def generate_extraction_schema(
+        self, sample_text: str, refined_query: QueryRefinement, guide: ExtractionGuide
+    ) -> ExtractionRequest:
+        """
+        Generate schema with enforced structure.
+
+        Args:
+            sample_text: Sample content for context
+            refined_query: Refined query with details
+            guide: Extraction guide with patterns
+
+        Returns:
+            Extraction request with model schema
+        """
+        return self.llm_core.complete(
+            messages=[
+                {"role": "system", "content": schema_system_prompt},
+                {
+                    "role": "user",
+                    "content": schema_template.substitute(
+                        refined_query=refined_query.refined_query,
+                        data_characteristics=refined_query.data_characteristics,
+                        structural_requirements=refined_query.structural_requirements,
+                        organization_principles=guide.organization_principles,
+                        sample_text=sample_text,
+                    ),
+                },
+            ],
+            response_model=ExtractionRequest,
+            config=self.llm_core.config.refinement,
+            step=ExtractionStep.SCHEMA_GENERATION,
+        )
+
+    def create_model_from_schema(
+        self, schema_request: ExtractionRequest
+    ) -> Type[BaseModel]:
+        """
+        Create Pydantic model from extraction request.
+
+        Args:
+            schema_request: Request containing model schema
+
+        Returns:
+            Generated Pydantic model class
+        """
+        extraction_model = ModelGenerator.from_extraction_request(schema_request)
+        logger.info("Generated Model Schema:")
+        logger.info(json.dumps(extraction_model.model_json_schema(), indent=2))
+        return extraction_model
+
+    def refine_existing_model(
+        self,
+        model: Type[BaseModel],
+        instructions: str,
+        model_name: str = None,
+    ) -> Type[BaseModel]:
+        """
+        Refine an existing data model based on natural language instructions.
+
+        Args:
+            model: Existing Pydantic model to refine
+            instructions: Natural language instructions for refinement
+            model_name: Optional name for the refined model
+
+        Returns:
+            A new refined Pydantic model
+        """
+        # Default model name if not provided
+        if model_name is None:
+            model_name = f"Refined{model.__name__}"
+
+        # Get the schema of the existing model
+        model_schema = model.model_json_schema()
+        model_schema_str = json.dumps(model_schema, indent=2)
+
+        # Generate schema for the refined model directly
+        extraction_request = self.llm_core.complete(
+            response_model=ExtractionRequest,
+            messages=[
+                {
+                    "role": "system",
+                    "content": refinement_system_prompt,
+                },
+                {
+                    "role": "user",
+                    "content": refinement_template.substitute(
+                        model_schema=model_schema_str,
+                        instructions=instructions,
+                    ),
+                },
+            ],
+            config=self.llm_core.config.refinement,
+            step=ExtractionStep.SCHEMA_GENERATION,
+        )
+
+        # Set the model name if specified
+        if model_name:
+            extraction_request.model_name = model_name
+
+        # Sanitize regex patterns to prevent validation errors
+        sanitized_request = sanitize_regex_patterns(extraction_request)
+
+        # Convert from v1 to v2 if needed and generate model
+        converted_request = convert_pydantic_v1_to_v2(sanitized_request)
+        refined_model = ModelGenerator.from_extraction_request(converted_request)
+
+        return refined_model
+
+    def generate_from_custom_model(
+        self,
+        model: Type[BaseModel],
+        query: str,
+        data_columns: List[str],
+    ) -> Tuple[QueryRefinement, ExtractionGuide]:
+        """
+        Generate refinement and guide from a provided custom model.
+
+        When a custom model is provided, we reverse engineer the refinement and guide
+        to match the model structure, rather than generating them from the query.
+
+        Args:
+            model: The provided custom model
+            query: The original query (used as context)
+            data_columns: Available columns in the dataset
+
+        Returns:
+            Tuple of refined_query and extraction_guide
+        """
+        # Get model schema and description
+        model_description = ModelUtils.get_model_description(model)
+        model_schema = model.model_json_schema()
+        model_properties = model_schema.get("properties", {})
+
+        # Extract data characteristics from the model properties
+        data_characteristics = ModelUtils.extract_field_characteristics(model)
+
+        # Extract structure requirements
+        structural_requirements = ModelUtils.extract_structural_requirements(model)
+
+        # Create a simplified query refinement with explicit field mapping instructions
+        model_fields = list(model_properties.keys())
+        refined_query = QueryRefinement(
+            refined_query=f"Extract {model_description} as specified in the provided model, filling all fields with relevant data from the appropriate columns. Original query: {query}",
+            data_characteristics=data_characteristics,
+            structural_requirements=structural_requirements,
+        )
+
+        # Create custom field-to-column mapping suggestions based on field names and data columns
+        field_descriptions = {}
+        for prop_name, prop_info in model_properties.items():
+            field_descriptions[prop_name] = {
+                "description": prop_info.get("description", ""),
+                "type": prop_info.get("type", ""),
+                "enum": prop_info.get("enum", []),
+            }
+
+        column_suggestions = ContentAnalyzer.suggest_column_mappings(
+            model_properties, data_columns, field_descriptions
+        )
+
+        # Generate guide with enhanced column mapping
+        guide_messages = [
+            {"role": "system", "content": "You are an extraction guide generator."},
+            {
+                "role": "user",
+                "content": custom_model_guide_template.substitute(
+                    data_characteristics=data_characteristics,
+                    available_columns=data_columns,
+                    model_fields=model_fields,
+                    column_suggestions=json.dumps(column_suggestions, indent=2),
+                ),
+            },
+        ]
+
+        guide = self.llm_core.complete(
+            messages=guide_messages,
+            response_model=ExtractionGuide,
+            config=self.llm_core.config.refinement,
+            step=ExtractionStep.GUIDE,
+        )
+
+        logger.info(f"Extraction Columns: {guide.target_columns}")
+        logger.info(
+            f"Generated refinement and guide from custom model: {model.__name__}"
+        )
+
+        return refined_query, guide

structx/extraction/result_manager.py

@@ -0,0 +1,125 @@
+"""
+Result management for extraction operations.
+"""
+
+from datetime import datetime
+from typing import Any, Dict, List
+
+import pandas as pd
+from loguru import logger
+from pydantic import BaseModel
+
+from structx.utils.helpers import flatten_extracted_data
+
+
+class ResultManager:
+    """
+    Manages extraction results, errors, and statistics.
+    """
+
+    @staticmethod
+    def initialize_results(
+        df: pd.DataFrame, extraction_model: type[BaseModel]
+    ) -> tuple[pd.DataFrame, List[Any], List[Dict]]:
+        """
+        Initialize result containers.
+
+        Args:
+            df: Original DataFrame
+            extraction_model: Pydantic model for extraction
+
+        Returns:
+            Tuple of (result_df, result_list, failed_rows)
+        """
+        result_df = df.copy()
+        result_list = []
+        failed_rows = []
+
+        # Initialize extraction columns
+        for field_name in extraction_model.model_fields:
+            result_df[field_name] = None
+        result_df["extraction_status"] = None
+
+        return result_df, result_list, failed_rows
+
+    @staticmethod
+    def update_dataframe(
+        result_df: pd.DataFrame,
+        items: List[BaseModel],
+        row_idx: int,
+        expand_nested: bool,
+    ) -> None:
+        """
+        Update DataFrame with extracted items.
+
+        Args:
+            result_df: DataFrame to update
+            items: Extracted model instances
+            row_idx: Row index to update
+            expand_nested: Whether to flatten nested structures
+        """
+        for i, item in enumerate(items):
+            # Flatten if needed
+            item_data = (
+                flatten_extracted_data(item.model_dump())
+                if expand_nested
+                else item.model_dump()
+            )
+
+            # For multiple items, append index to field names
+            if i > 0:
+                item_data = {f"{k}_{i}": v for k, v in item_data.items()}
+
+            # Update result dataframe
+            for field_name, value in item_data.items():
+                result_df.at[row_idx, field_name] = value
+
+        result_df.at[row_idx, "extraction_status"] = "Success"
+
+    @staticmethod
+    def handle_extraction_error(
+        result_df: pd.DataFrame,
+        failed_rows: List[Dict],
+        row_idx: int,
+        row_text: str,
+        error: Exception,
+    ) -> None:
+        """
+        Handle and log extraction errors.
+
+        Args:
+            result_df: DataFrame to update with error status
+            failed_rows: List to append failed row information
+            row_idx: Row index that failed
+            row_text: Text that failed to extract
+            error: Exception that occurred
+        """
+        failed_rows.append(
+            {
+                "index": row_idx,
+                "text": row_text,
+                "error": str(error),
+                "timestamp": datetime.now().isoformat(),
+            }
+        )
+        result_df.at[row_idx, "extraction_status"] = f"Failed: {str(error)}"
+
+    @staticmethod
+    def log_extraction_stats(total_rows: int, failed_rows: List[Dict]) -> None:
+        """
+        Log extraction statistics.
+
+        Args:
+            total_rows: Total number of rows processed
+            failed_rows: List of failed rows
+        """
+        success_count = total_rows - len(failed_rows)
+        logger.info("\nExtraction Statistics:")
+        logger.info(f"Total rows: {total_rows}")
+        logger.info(
+            f"Successfully processed: {success_count} "
+            f"({success_count/total_rows*100:.2f}%)"
+        )
+        logger.info(
+            f"Failed: {len(failed_rows)} " f"({len(failed_rows)/total_rows*100:.2f}%)"
+        )

structx/utils/constants.py

@@ -0,0 +1,36 @@
+from datetime import date, datetime, time
+from typing import Any, Dict, List, Optional
+
+TYPE_ALIASES = {
+    "string": "str",
+    "integer": "int",
+    "numeric": "float",
+    "number": "float",
+    "boolean": "bool",
+    "array": "List",
+    "object": "dict",
+    "date-time": "datetime",  # Add this mapping
+    "datetime": "datetime",  # Add this for consistency
+    "date": "date",  # Add this for completeness
+    "time": "time",  # Add this for completeness
+    "decimal": "float",  # Assuming decimal is represented as float
+}
+
+SAFE_TYPES = {
+    # Python built-ins
+    "str": str,
+    "int": int,
+    "float": float,
+    "bool": bool,
+    "list": list,
+    "dict": dict,
+    # Common types
+    "datetime": datetime,
+    "date": date,
+    "time": time,
+    # Generic type constructors
+    "List": List,
+    "Dict": Dict,
+    "Optional": Optional,
+    "Any": Any,
+}

structx/utils/file_reader.py

@@ -0,0 +1,298 @@
+import tempfile
+from pathlib import Path
+from typing import Any, Callable, Dict, List, Union
+
+import pandas as pd
+
+from structx.core.exceptions import FileError
+
+
+class FileReader:
+    """
+    Handles reading different file formats with a unified approach for unstructured documents.
+
+    For unstructured documents (TXT, DOCX, PDF), the default strategy is to convert
+    everything to PDF and use instructor's multimodal PDF support. This eliminates
+    the need for manual chunking and provides the best context preservation.
+    """
+
+    STRUCTURED_EXTENSIONS: Dict[
+        str, Callable[[Union[str, Path], Dict], pd.DataFrame]
+    ] = {
+        ".csv": pd.read_csv,
+        ".xlsx": pd.read_excel,
+        ".xls": pd.read_excel,
+        ".json": pd.read_json,
+        ".parquet": pd.read_parquet,
+        ".feather": pd.read_feather,
+    }
+
+    TEXT_EXTENSIONS: List[str] = [".txt", ".md", ".py", ".html", ".xml", ".log", ".rst"]
+    DOCUMENT_EXTENSIONS: List[str] = [".pdf", ".docx", ".doc"]
+
+    @staticmethod
+    def read_file(file_path: Union[str, Path], **kwargs: Any) -> pd.DataFrame:
+        """
+        Read a file and return its content based on the specified mode.
+
+        For unstructured documents (TXT, DOCX, PDF), the default approach is to
+        convert everything to PDF and use instructor's multimodal PDF support.
+        This eliminates the need for manual chunking and provides the best
+        context preservation.
+
+        Args:
+            file_path: Path to the file to read
+            **kwargs: Additional options for file reading including:
+                - mode: Reading mode - 'multimodal_pdf' (default), 'simple_text', or 'simple_pdf'
+                - use_multimodal: Use instructor's multimodal support (default: True)
+                - file_options: Additional options for reading the file
+
+        Returns:
+            pandas DataFrame with the appropriate structure for the specified mode
+
+        Raises:
+            FileError: If file cannot be read or processed
+        """
+        # Extract parameters from kwargs
+        mode = kwargs.get("mode", "multimodal_pdf")
+        use_multimodal = kwargs.get("use_multimodal", True)
+        file_options = kwargs.get("file_options", {})
+
+        # Handle legacy parameter structure
+        if use_multimodal and mode == "multimodal_pdf":
+            mode = "multimodal_pdf"
+        elif not use_multimodal:
+            mode = "simple_text"
+        try:
+            file_path = Path(file_path)
+            if not file_path.exists():
+                raise FileError(f"File not found: {file_path}")
+
+            file_extension = file_path.suffix.lower()
+
+            # Handle structured files (return DataFrame)
+            if file_extension in FileReader.STRUCTURED_EXTENSIONS:
+                read_func = FileReader.STRUCTURED_EXTENSIONS[file_extension]
+                return read_func(file_path, **file_options)
+
+            # Handle unstructured files
+            if (
+                file_extension in FileReader.TEXT_EXTENSIONS
+                or file_extension in FileReader.DOCUMENT_EXTENSIONS
+            ):
+                if mode == "multimodal_pdf":
+                    # Convert all unstructured documents to PDF for instructor's multimodal support
+                    pdf_path = FileReader._convert_to_pdf(file_path)
+                    # Return DataFrame with required structure for multimodal processing
+                    return pd.DataFrame(
+                        {
+                            "pdf_path": [str(pdf_path)],
+                            "source": [str(file_path)],
+                            "multimodal": [True],
+                            "file_type": ["pdf"],
+                        }
+                    )
+                elif mode == "simple_text":
+                    # Fallback: simple text reading with chunking
+                    return FileReader._read_as_text_chunks(file_path, kwargs)
+                elif mode == "simple_pdf":
+                    # Fallback: simple PDF reading (if it's already a PDF)
+                    if file_extension == ".pdf":
+                        return FileReader._read_pdf_chunks(file_path, kwargs)
+                    else:
+                        # Convert to PDF first, then read simply
+                        pdf_path = FileReader._convert_to_pdf(file_path)
+                        return FileReader._read_pdf_chunks(Path(pdf_path), kwargs)
+
+            raise FileError(f"Unsupported file type: {file_extension}")
+
+        except Exception as e:
+            raise FileError(f"Error reading file {file_path}: {str(e)}")
+
+    @staticmethod
+    def _convert_to_pdf(file_path: Path) -> str:
+        """
+        Convert any supported document to PDF using docling -> markdown -> PDF pipeline.
+
+        Returns the path to the generated PDF file for use with instructor's multimodal support.
+        """
+        try:
+            file_extension = file_path.suffix.lower()
+
+            # If it's already a PDF, return as-is
+            if file_extension == ".pdf":
+                return str(file_path)
+
+            # For simple text files, read directly
+            if file_extension in FileReader.TEXT_EXTENSIONS:
+                with open(file_path, "r", encoding="utf-8") as f:
+                    content = f.read()
+                return FileReader._markdown_to_pdf(content, file_path.stem)
+
+            # For document files, use docling to convert to markdown first
+            elif (
+                file_extension in FileReader.DOCUMENT_EXTENSIONS
+                and file_extension != ".pdf"
+            ):
+                try:
+                    from docling.document_converter import DocumentConverter
+
+                    converter = DocumentConverter()
+                    result = converter.convert(str(file_path))
+                    markdown_content = result.document.export_to_markdown()
+
+                    # Convert markdown to PDF
+                    return FileReader._markdown_to_pdf(markdown_content, file_path.stem)
+
+                except ImportError:
+                    raise FileError(
+                        f"docling not available for {file_extension} conversion"
+                    )
+            else:
+                raise FileError(
+                    f"Unsupported file type for conversion: {file_extension}"
+                )
+
+        except Exception as e:
+            raise FileError(f"Error converting {file_path} to PDF: {str(e)}")
+
+    @staticmethod
+    def _markdown_to_pdf(markdown_content: str, filename: str) -> str:
+        """Convert markdown content to PDF and return the path."""
+
+        import markdown
+        import weasyprint
+
+        # Convert markdown to HTML
+        md = markdown.Markdown(extensions=["extra", "codehilite"])
+        html_content = md.convert(markdown_content)
+
+        # Add basic CSS styling
+        html_with_css = f"""
+            <!DOCTYPE html>
+            <html>
+            <head>
+                <meta charset="utf-8">
+                <style>
+                    body {{ font-family: Arial, sans-serif; margin: 40px; line-height: 1.6; }}
+                    h1, h2, h3 {{ color: #333; }}
+                    pre {{ background-color: #f4f4f4; padding: 10px; border-radius: 5px; }}
+                    code {{ background-color: #f4f4f4; padding: 2px 4px; border-radius: 3px; }}
+                </style>
+            </head>
+            <body>
+                {html_content}
+            </body>
+            </html>
+            """
+
+        # Create temporary PDF file
+        with tempfile.NamedTemporaryFile(
+            delete=False, suffix=".pdf", prefix=f"{filename}_"
+        ) as tmp_file:
+            pdf_path = tmp_file.name
+
+        # Generate PDF with weasyprint
+        weasyprint.HTML(string=html_with_css).write_pdf(pdf_path)
+        return pdf_path
+
+    @staticmethod
+    def _read_as_text_chunks(file_path: Path, kwargs: Dict[str, Any]) -> pd.DataFrame:
+        """Simple text reading fallback with chunking."""
+        try:
+            file_extension = file_path.suffix.lower()
+            chunk_size = kwargs.get("chunk_size", 1000)
+            chunk_overlap = kwargs.get("chunk_overlap", 200)
+
+            if file_extension in FileReader.TEXT_EXTENSIONS:
+                with open(file_path, "r", encoding="utf-8") as f:
+                    content = f.read()
+            elif file_extension == ".docx":
+                try:
+                    from docx import Document
+
+                    doc = Document(file_path)
+                    content = "\n".join(
+                        [paragraph.text for paragraph in doc.paragraphs]
+                    )
+                except ImportError:
+                    raise FileError("python-docx not available for DOCX reading")
+            elif file_extension == ".pdf":
+                content = FileReader._extract_pdf_text(file_path)
+            else:
+                raise FileError(f"Cannot read {file_extension} as simple text")
+
+            # Simple chunking
+            chunks = []
+            for i in range(0, len(content), chunk_size - chunk_overlap):
+                chunks.append(content[i : i + chunk_size])
+
+            return pd.DataFrame(
+                {
+                    "text": chunks,
+                    "chunk_id": range(len(chunks)),
+                    "source": str(file_path),
+                    "processing_method": ["simple_text"] * len(chunks),
+                }
+            )
+
+        except Exception as e:
+            raise FileError(f"Error reading {file_path} as text: {str(e)}")
+
+    @staticmethod
+    def _read_pdf_chunks(file_path: Path, kwargs: Dict[str, Any]) -> pd.DataFrame:
+        """Simple PDF text extraction fallback with chunking."""
+        try:
+            chunk_size = kwargs.get("chunk_size", 1000)
+            chunk_overlap = kwargs.get("chunk_overlap", 200)
+
+            content = FileReader._extract_pdf_text(file_path)
+
+            # Simple chunking
+            chunks = []
+            for i in range(0, len(content), chunk_size - chunk_overlap):
+                chunks.append(content[i : i + chunk_size])
+
+            return pd.DataFrame(
+                {
+                    "text": chunks,
+                    "chunk_id": range(len(chunks)),
+                    "source": str(file_path),
+                    "processing_method": ["simple_pdf"] * len(chunks),
+                }
+            )
+
+        except Exception as e:
+            raise FileError(f"Error reading PDF {file_path}: {str(e)}")
+
+    @staticmethod
+    def _extract_pdf_text(file_path: Path) -> str:
+        """Extract text from PDF file using PyPDF2."""
+        try:
+            import PyPDF2
+
+            with open(file_path, "rb") as file:
+                pdf_reader = PyPDF2.PdfReader(file)
+                text = ""
+                for page in pdf_reader.pages:
+                    text += page.extract_text() + "\n"
+                return text
+
+        except ImportError:
+            raise FileError("PyPDF2 not available for simple PDF reading")
+        except Exception as e:
+            raise FileError(f"Error reading PDF {file_path}: {str(e)}")
+
+    @staticmethod
+    def get_file_type(file_path: Union[str, Path]) -> str:
+        """Get the type of file based on its extension"""
+        file_extension = Path(file_path).suffix.lower()
+
+        if file_extension in FileReader.STRUCTURED_EXTENSIONS:
+            return "structured"
+        elif file_extension in FileReader.TEXT_EXTENSIONS:
+            return "text"
+        elif file_extension in FileReader.DOCUMENT_EXTENSIONS:
+            return "document"
+        else:
+            return "unknown"