structx 0.4.8__py3-none-any.whl

structx/extraction/core/model_utils.py

@@ -0,0 +1,143 @@
+"""
+Utility functions for model operations.
+"""
+
+from typing import Dict, List, Literal, Type
+
+from pydantic import BaseModel
+
+
+class ModelUtils:
+    """
+    Utility functions for working with Pydantic models.
+    """
+
+    @staticmethod
+    def extract_model_schema_info(model: Type[BaseModel]) -> str:
+        """
+        Extract model schema information for prompts.
+
+        Args:
+            model: Pydantic model to extract info from
+
+        Returns:
+            Formatted string with model field descriptions
+        """
+        model_schema = model.model_json_schema()
+        model_schema_info = ""
+
+        # Include field descriptions to help with extraction
+        for field, details in model_schema.get("properties", {}).items():
+            field_type = details.get("type", "unknown")
+            field_desc = details.get("description", "")
+            if "enum" in details:
+                field_desc += (
+                    f" Possible values: {', '.join(map(str, details['enum']))}"
+                )
+            model_schema_info += f"- {field} ({field_type}): {field_desc}\n"
+
+        return model_schema_info
+
+    @staticmethod
+    def create_model_context(
+        model: Type[BaseModel], instruction_type: Literal["text", "document"] = "text"
+    ) -> str:
+        """
+        Create contextual information for custom model extraction.
+
+        Args:
+            model: Pydantic model to create context for
+            instruction_type: Type of instruction ("text" or "document")
+
+        Returns:
+            Formatted context string for prompts
+        """
+        model_schema_info = ModelUtils.extract_model_schema_info(model)
+
+        if not model_schema_info.strip():
+            return ""
+
+        if instruction_type == "document":
+            return (
+                f"\nModel fields and descriptions:\n{model_schema_info}\n"
+                f"Ensure all applicable fields are populated with relevant information from the document."
+            )
+        else:
+            return (
+                f"\nModel fields and descriptions:\n{model_schema_info}\n"
+                f"Ensure all applicable fields are populated with relevant information from the text."
+            )
+
+    @staticmethod
+    def extract_field_characteristics(model: Type[BaseModel]) -> List[str]:
+        """
+        Extract data characteristics from model properties.
+
+        Args:
+            model: Pydantic model to analyze
+
+        Returns:
+            List of field characteristic descriptions
+        """
+        model_schema = model.model_json_schema()
+        model_properties = model_schema.get("properties", {})
+        data_characteristics = []
+
+        for prop_name, prop_info in model_properties.items():
+            prop_description = prop_info.get("description", "")
+            prop_type = prop_info.get("type", "")
+            enum_values = prop_info.get("enum", [])
+
+            if prop_description:
+                if enum_values:
+                    data_characteristics.append(
+                        f"{prop_name} ({prop_type}): {prop_description}. "
+                        f"Possible values: {', '.join(map(str, enum_values))}"
+                    )
+                else:
+                    data_characteristics.append(
+                        f"{prop_name} ({prop_type}): {prop_description}"
+                    )
+            else:
+                data_characteristics.append(f"{prop_name} ({prop_type})")
+
+        return data_characteristics
+
+    @staticmethod
+    def extract_structural_requirements(model: Type[BaseModel]) -> Dict[str, str]:
+        """
+        Extract structural requirements from model.
+
+        Args:
+            model: Pydantic model to analyze
+
+        Returns:
+            Dictionary mapping field names to types
+        """
+        model_schema = model.model_json_schema()
+        model_properties = model_schema.get("properties", {})
+        structural_requirements = {}
+
+        for prop_name, prop_info in model_properties.items():
+            if "type" in prop_info:
+                structural_requirements[prop_name] = prop_info["type"]
+
+        return structural_requirements
+
+    @staticmethod
+    def get_model_description(model: Type[BaseModel]) -> str:
+        """
+        Get a descriptive name for the model.
+
+        Args:
+            model: Pydantic model to describe
+
+        Returns:
+            Model description string
+        """
+        model_schema = model.model_json_schema()
+        return (
+            model_schema.get("description", "")
+            or model_schema.get("title", "")
+            or model.__name__
+        )

structx/extraction/engines/__init__.py

@@ -0,0 +1,5 @@
+"""Extraction engines for different processing strategies."""
+
+from .extraction_engine import ExtractionEngine
+
+__all__ = ["ExtractionEngine"]

structx/extraction/engines/extraction_engine.py

@@ -0,0 +1,221 @@
+"""
+Core extraction engine with different processing strategies.
+"""
+
+from typing import Any, Dict, List, Type, Union
+
+from loguru import logger
+from pydantic import BaseModel, Field, create_model
+
+from structx.core.exceptions import ExtractionError
+from structx.core.models import ExtractionGuide, QueryRefinement
+from structx.extraction.core.llm_core import LLMCore
+from structx.extraction.core.model_utils import ModelUtils
+from structx.utils.helpers import handle_errors
+from structx.utils.prompts import extraction_system_prompt, extraction_template
+from structx.utils.usage import ExtractionStep
+
+
+class ExtractionEngine:
+    """
+    Core extraction engine with different processing strategies.
+    """
+
+    def __init__(self, llm_core: LLMCore):
+        """
+        Initialize extraction engine.
+
+        Args:
+            llm_core: LLM core for completions
+        """
+        self.llm_core = llm_core
+
+    @handle_errors(error_message="Text extraction failed", error_type=ExtractionError)
+    def extract_with_model(
+        self,
+        text: str,
+        extraction_model: Type[BaseModel],
+        refined_query: QueryRefinement,
+        guide: ExtractionGuide,
+        is_custom_model: bool = False,
+    ) -> List[BaseModel]:
+        """
+        Extract data with enforced structure with retries and usage tracking.
+
+        Args:
+            text: Text to extract from
+            extraction_model: Pydantic model for extraction
+            refined_query: Refined query with details
+            guide: Extraction guide with patterns
+            is_custom_model: Whether this is a user-provided model
+
+        Returns:
+            List of extracted model instances
+        """
+        # Create a container model to wrap the list items
+        # this is necessary to be able to track token usage, when passing an iterable data model
+        # result._raw_response does not exist making usage calculations not possible
+        container_name = f"{extraction_model.__name__}Container"
+        container_model = create_model(
+            container_name,
+            __base__=BaseModel,
+            items=(
+                List[extraction_model],
+                Field(description=f"List of {extraction_model.__name__} items"),
+            ),
+        )
+
+        # Get model context for custom models
+        extra_context = ""
+        if is_custom_model:
+            extra_context = ModelUtils.create_model_context(extraction_model, "text")
+
+        # Use _perform_llm_completion with the container model
+        container = self.llm_core.complete_with_retry(
+            messages=[
+                {"role": "system", "content": extraction_system_prompt},
+                {
+                    "role": "user",
+                    "content": extraction_template.substitute(
+                        query=refined_query.refined_query,
+                        patterns=guide.structural_patterns,
+                        rules=(
+                            guide.relationship_rules + [extra_context]
+                            if extra_context
+                            else guide.relationship_rules
+                        ),
+                        text=text,
+                    ),
+                },
+            ],
+            response_model=container_model,
+            config=self.llm_core.config.extraction,
+            step=ExtractionStep.EXTRACTION,
+        )
+
+        # Return just the items
+        return container.items
+
+    @handle_errors(error_message="PDF extraction failed", error_type=ExtractionError)
+    def extract_with_multimodal_pdf(
+        self,
+        pdf_path: str,
+        extraction_model: Type[BaseModel],
+        refined_query: QueryRefinement,
+        guide: ExtractionGuide,
+        is_custom_model: bool = False,
+    ) -> List[BaseModel]:
+        """
+        Extract data from PDF using instructor's multimodal support.
+
+        Args:
+            pdf_path: Path to PDF file
+            extraction_model: Pydantic model for extraction
+            refined_query: Refined query with details
+            guide: Extraction guide with patterns
+            is_custom_model: Whether this is a user-provided model
+
+        Returns:
+            List of extracted model instances
+        """
+        try:
+            from instructor.multimodal import PDF
+        except ImportError:
+            raise ImportError(
+                "instructor multimodal support is required for PDF processing. "
+                "Install it with: pip install instructor[multimodal]"
+            )
+
+        # For multimodal PDF, we need a single wrapper model, not a container with multiple items
+        # This is because instructor's multimodal support expects a single response, not multiple tool calls
+        wrapper_name = f"{extraction_model.__name__}List"
+        wrapper_model = create_model(
+            wrapper_name,
+            __base__=BaseModel,
+            items=(
+                List[extraction_model],
+                Field(
+                    description=f"List of extracted {extraction_model.__name__} items from the document"
+                ),
+            ),
+        )
+
+        # Get model context for custom models
+        extra_context = ""
+        if is_custom_model:
+            extra_context = ModelUtils.create_model_context(
+                extraction_model, "document"
+            )
+
+        content = [
+            f"Extract structured information from this PDF document following these guidelines:\n\n"
+            f"Query: {refined_query.refined_query}\n"
+            f"Patterns: {guide.structural_patterns}\n"
+            f"Rules: {guide.relationship_rules + [extra_context] if extra_context else guide.relationship_rules}\n\n",
+            PDF.from_path(pdf_path),
+        ]
+
+        logger.info(
+            f"Extracting from PDF: {pdf_path} with query: {refined_query.refined_query}"
+        )
+
+        result = self.llm_core.complete_with_retry(
+            response_model=wrapper_model,
+            messages=[
+                {"role": "system", "content": extraction_system_prompt},
+                {
+                    "role": "user",
+                    "content": content,
+                },
+            ],
+            step=ExtractionStep.EXTRACTION,
+            config=self.llm_core.config.extraction,
+        )
+        logger.info(f"Completed extraction for PDF: {pdf_path}")
+        return result.items
+
+    def extract_from_row_data(
+        self,
+        row_data: Union[str, Dict[str, Any]],
+        extraction_model: Type[BaseModel],
+        refined_query: QueryRefinement,
+        guide: ExtractionGuide,
+        is_custom_model: bool = False,
+    ) -> List[BaseModel]:
+        """
+        Extract data from row data (either text or multimodal).
+
+        Args:
+            row_data: Row data to extract from
+            extraction_model: Pydantic model for extraction
+            refined_query: Refined query with details
+            guide: Extraction guide with patterns
+            is_custom_model: Whether this is a user-provided model
+
+        Returns:
+            List of extracted model instances
+        """
+        # Check if this is a multimodal PDF row
+        if (
+            isinstance(row_data, dict)
+            and row_data.get("multimodal")
+            and row_data.get("file_type") == "pdf"
+        ):
+            pdf_path = row_data.get("pdf_path")
+            return self.extract_with_multimodal_pdf(
+                pdf_path=pdf_path,
+                extraction_model=extraction_model,
+                refined_query=refined_query,
+                guide=guide,
+                is_custom_model=is_custom_model,
+            )
+        else:
+            # Handle regular text extraction
+            row_text = row_data if isinstance(row_data, str) else str(row_data)
+            return self.extract_with_model(
+                text=row_text,
+                extraction_model=extraction_model,
+                refined_query=refined_query,
+                guide=guide,
+                is_custom_model=is_custom_model,
+            )