structx 0.4.8__py3-none-any.whl

structx/__init__.py

@@ -0,0 +1,24 @@
+"""
+structx: Structured data extraction using LLMs
+"""
+
+from structx.core.config import ExtractionConfig, StepConfig
+from structx.core.models import (
+    ExtractionGuide,
+    ExtractionRequest,
+    ModelField,
+    QueryRefinement,
+)
+from structx.extraction.extractor import Extractor
+
+__version__ = "0.4.8"
+__all__ = [
+    "Extractor",
+    "ExtractionConfig",
+    "StepConfig",
+    "ModelField",
+    "QueryAnalysis",
+    "QueryRefinement",
+    "ExtractionGuide",
+    "ExtractionRequest",
+]

structx/core/config.py

@@ -0,0 +1,127 @@
+from pathlib import Path
+from typing import Annotated, Any, Dict, Optional, Union
+
+from omegaconf import OmegaConf
+from pydantic import BaseModel, Field
+
+from structx.utils.types import DictStrAny
+
+
+class StepConfig(BaseModel):
+    """Configuration for an individual extraction step"""
+
+    temperature: Optional[
+        Annotated[
+            float, Field(ge=0.0, le=1.0, description="Sampling temperature for LLM")
+        ]
+    ] = None
+    top_p: Optional[
+        Annotated[
+            float, Field(ge=0.0, le=1.0, description="Nucleus sampling parameter")
+        ]
+    ] = None
+    max_tokens: Optional[
+        Annotated[int, Field(gt=0, description="Maximum tokens in completion")]
+    ] = None
+
+    def model_dump(self, *args, **kwargs) -> Dict[str, Any]:
+        """
+        Return dictionary of non-None values merged with defaults
+        """
+        # Default values for each step
+        defaults = {
+            "temperature": 0.1,
+            "top_p": 0.1,
+            "max_tokens": 2000,
+        }
+
+        # Get current values, excluding None
+        current = {
+            k: v
+            for k, v in super().model_dump(*args, **kwargs).items()
+            if v is not None
+        }
+
+        # Merge with defaults, preferring current values
+        return {**defaults, **current}
+
+    class Config:
+        validate_assignment = True
+
+
+class ExtractionConfig:
+    """Configuration management for structx using OmegaConf and Pydantic"""
+
+    DEFAULT_CONFIG = {
+        "analysis": {"temperature": 0.2, "top_p": 0.1, "max_tokens": 2000},
+        "refinement": {"temperature": 0.1, "top_p": 0.05, "max_tokens": 2000},
+        "extraction": {"temperature": 0.0, "top_p": 0.1, "max_tokens": 8192},
+    }
+
+    def __init__(
+        self,
+        config: Optional[Union[Dict[str, Any], str]] = None,
+        config_path: Optional[Union[str, Path]] = None,
+    ):
+        """
+        Initialize configuration
+
+        Args:
+            config: Optional configuration dictionary or YAML string
+            config_path: Optional path to YAML configuration file
+        """
+        # Create base config from defaults
+        self.conf = OmegaConf.create(self.DEFAULT_CONFIG)
+
+        # Load from file if provided
+        if config_path:
+            file_conf = OmegaConf.load(config_path)
+            self.conf = OmegaConf.merge(self.conf, file_conf)
+
+        # Merge with provided config if any
+        if config:
+            if isinstance(config, str):
+                conf_to_merge = OmegaConf.create(config)
+            else:
+                conf_to_merge = OmegaConf.create(config)
+            self.conf = OmegaConf.merge(self.conf, conf_to_merge)
+
+        # Validate using Pydantic models
+        self._validate_config()
+
+    def _validate_config(self) -> None:
+        """Validate configuration using Pydantic models"""
+        for step in ["analysis", "refinement", "extraction"]:
+            step_config: DictStrAny = OmegaConf.to_container(self.conf.get(step, {}))  # type: ignore
+            # Validate using StepConfig model
+            StepConfig(**step_config)
+
+    @property
+    def analysis(self) -> DictStrAny:
+        """Get validated analysis step configuration"""
+        config: DictStrAny = OmegaConf.to_container(self.conf.analysis)  # type: ignore
+        return StepConfig(**config).model_dump(exclude_none=True)
+
+    @property
+    def refinement(self) -> DictStrAny:
+        """Get validated refinement step configuration"""
+        config: DictStrAny = OmegaConf.to_container(self.conf.refinement)  # type: ignore
+        return StepConfig(**config).model_dump(exclude_none=True)
+
+    @property
+    def extraction(self) -> DictStrAny:
+        """Get validated extraction step configuration"""
+        config: DictStrAny = OmegaConf.to_container(self.conf.extraction)  # type: ignore
+        return StepConfig(**config).model_dump(exclude_none=True)
+
+    def save(self, path: str) -> None:
+        """Save configuration to YAML file"""
+        OmegaConf.save(self.conf, path)
+
+    def __str__(self) -> str:
+        """String representation of configuration"""
+        return OmegaConf.to_yaml(self.conf)
+
+    def __repr__(self) -> str:
+        """Representation of configuration"""
+        return self.__str__()

structx/core/exceptions.py

@@ -0,0 +1,34 @@
+class StructXError(Exception):
+    """Base exception for all structx errors"""
+
+    pass
+
+
+class ConfigurationError(StructXError):
+    """Error in configuration"""
+
+    pass
+
+
+class ExtractionError(StructXError):
+    """Error during extraction process"""
+
+    pass
+
+
+class ValidationError(StructXError):
+    """Error in data validation"""
+
+    pass
+
+
+class ModelGenerationError(StructXError):
+    """Error in dynamic model generation"""
+
+    pass
+
+
+class FileError(StructXError):
+    """Error in file operations"""
+
+    pass

structx/core/models.py

@@ -0,0 +1,147 @@
+from dataclasses import dataclass
+from typing import Any, Dict, Generic, List, Optional, Type, Union
+
+import pandas as pd
+from pydantic import BaseModel, Field
+
+from structx.utils.types import T
+from structx.utils.usage import ExtractorUsage, UsageSummary
+
+
+class ModelField(BaseModel):
+    """Definition of a field in the extraction model"""
+
+    name: str = Field(description="Name of the field")
+    type: str = Field(description="Type of the field")
+    description: str = Field(description="Description of what this field represents")
+    validation: Optional[Dict[str, Any]] = Field(
+        default_factory=dict, description="Additional validation rules"
+    )
+    nested_fields: Optional[List["ModelField"]] = Field(
+        default=None, description="Fields for nested models"
+    )
+
+    class Config:
+        validate_assignment = True
+
+
+class QueryRefinement(BaseModel):
+    """Refined query with structural information"""
+
+    refined_query: str = Field(description="Expanded query with structure requirements")
+    data_characteristics: Optional[List[str]] = Field(
+        description="Characteristics of data to extract"
+    )
+    structural_requirements: Optional[Dict[str, Any]] = Field(
+        description="Requirements for data structure"
+    )
+
+
+class ExtractionGuide(BaseModel):
+    """Guide for structured extraction"""
+
+    target_columns: List[str] = Field(description="Columns to analyze")
+
+    structural_patterns: Optional[Dict[str, str]] = Field(
+        description="Patterns for structuring data"
+    )
+    relationship_rules: Optional[List[str]] = Field(
+        description="Rules for data relationships"
+    )
+    organization_principles: Optional[List[str]] = Field(
+        description="Principles for data organization"
+    )
+
+    class Config:
+        extra = "allow"  # Allow extra fields to be flexible with LLM responses
+
+
+class ExtractionRequest(BaseModel):
+    """Request for model generation"""
+
+    model_name: str = Field(description="Name for generated model")
+    model_description: str = Field(description="Description of model purpose")
+    fields: List[ModelField] = Field(description="Fields to extract")
+
+
+@dataclass
+class ExtractionResult(Generic[T]):
+    """
+    Container for extraction results.
+
+    Attributes:
+        data: Extracted data (DataFrame or list of model instances)
+        failed: DataFrame with failed extractions
+        model: Generated or provided model class
+        usage: Token usage information across all extraction steps
+    """
+
+    data: Union[pd.DataFrame, List[T]]
+    failed: pd.DataFrame
+    model: Type[T]
+    usage: Optional[ExtractorUsage] = None
+
+    @property
+    def success_count(self) -> int:
+        """Number of successful extractions"""
+        if isinstance(self.data, pd.DataFrame):
+            return len(self.data)
+        return len(self.data)
+
+    @property
+    def failure_count(self) -> int:
+        """Number of failed extractions"""
+        return len(self.failed)
+
+    @property
+    def success_rate(self) -> float:
+        """Success rate as a percentage"""
+        total = self.success_count + self.failure_count
+        return (self.success_count / total * 100) if total > 0 else 0
+
+    def get_token_usage(self, detailed: bool = False) -> Optional[UsageSummary]:
+        """
+        Get structured token usage information.
+
+        Provides a detailed breakdown of token usage across all steps of the
+        extraction process.
+
+        Args:
+            detailed: If True, includes detailed breakdown of each extraction step
+                    (useful for multi-document extraction)
+
+        Returns:
+            UsageSummary object with token usage information, or None if usage tracking
+            isn't available
+
+        Example:
+            ```python
+            result = extractor.extract(data, query)
+            usage = result.get_token_usage()
+            print(f"Total tokens: {usage.total_tokens}")
+
+            # Access step-specific usage
+            for step in usage.steps:
+                print(f"{step.name}: {step.tokens} tokens")
+            ```
+        """
+
+        if not self.usage:
+            return None
+
+        return self.usage.get_usage_summary(detailed=detailed)
+
+    def __repr__(self) -> str:
+        """String representation"""
+        return (
+            f"ExtractionResult(success={self.success_count}, "
+            f"failed={self.failure_count}, "
+            f"model={self.model.__name__})"
+        )
+
+    def __str__(self):
+        return self.__repr__()
+
+
+# Rebuild the model to ensure nested fields are properly defined
+ModelField.model_rebuild()

structx/extraction/__init__.py

@@ -0,0 +1,38 @@
+"""
+Extraction module for structured data extraction.
+
+This module provides a well-organized architecture for extracting structured data
+from various sources using LLMs with dynamic model generation.
+
+Structure:
+- core/: Core LLM operations and utilities
+- processors/: Data processing and model operations
+- engines/: Extraction engines for different strategies
+- extractor.py: Main orchestrator class
+- generator.py: Dynamic model generation
+- result_manager.py: Result management utilities
+"""
+
+# Core components
+from .core import LLMCore, ModelUtils
+
+# Engines
+from .engines import ExtractionEngine
+from .extractor import Extractor
+from .generator import ModelGenerator
+
+# Processors
+from .processors import ContentAnalyzer, DataProcessor, ModelOperations
+from .result_manager import ResultManager
+
+__all__ = [
+    "Extractor",
+    "ModelGenerator",
+    "ResultManager",
+    "LLMCore",
+    "ModelUtils",
+    "ContentAnalyzer",
+    "DataProcessor",
+    "ModelOperations",
+    "ExtractionEngine",
+]

structx/extraction/core/__init__.py

@@ -0,0 +1,6 @@
+"""Core components for extraction operations."""
+
+from .llm_core import LLMCore
+from .model_utils import ModelUtils
+
+__all__ = ["LLMCore", "ModelUtils"]

structx/extraction/core/llm_core.py

@@ -0,0 +1,206 @@
+"""
+Core LLM operations and query processing.
+"""
+
+import logging
+import threading
+from typing import Dict, List, Type
+
+from instructor import Instructor
+from loguru import logger
+from tenacity import (
+    after_log,
+    before_sleep_log,
+    retry,
+    retry_if_exception_type,
+    stop_after_attempt,
+    wait_exponential,
+)
+
+from structx.core.config import DictStrAny, ExtractionConfig
+from structx.core.exceptions import ExtractionError
+from structx.core.models import ExtractionGuide, QueryRefinement
+from structx.utils.helpers import handle_errors
+from structx.utils.prompts import (
+    guide_system_prompt,
+    guide_template,
+    query_refinement_system_prompt,
+    query_refinement_template,
+)
+from structx.utils.types import ResponseType
+from structx.utils.usage import ExtractionStep, ExtractorUsage, StepUsage
+
+
+class LLMCore:
+    """
+    Core LLM operations with retry logic, usage tracking, and query processing.
+    """
+
+    def __init__(
+        self,
+        client: Instructor,
+        model_name: str,
+        config: ExtractionConfig,
+        max_retries: int = 3,
+        min_wait: int = 1,
+        max_wait: int = 10,
+    ):
+        """
+        Initialize LLM core.
+
+        Args:
+            client: Instructor-patched client
+            model_name: Name of the model to use
+            config: Extraction configuration
+            max_retries: Maximum number of retries for extraction
+            min_wait: Minimum seconds to wait between retries
+            max_wait: Maximum seconds to wait between retries
+        """
+        self.client = client
+        self.model_name = model_name
+        self.config = config
+        self.max_retries = max_retries
+        self.min_wait = min_wait
+        self.max_wait = max_wait
+        self.usage_lock = threading.Lock()
+        self.usage = ExtractorUsage()
+
+    def reset_usage(self) -> None:
+        """Reset usage tracking."""
+        with self.usage_lock:
+            self.usage = ExtractorUsage()
+
+    def get_usage(self) -> ExtractorUsage:
+        """Get current usage statistics."""
+        with self.usage_lock:
+            return self.usage
+
+    def _create_retry_decorator(self):
+        """Create retry decorator with instance parameters."""
+        return retry(
+            stop=stop_after_attempt(self.max_retries),
+            wait=wait_exponential(
+                multiplier=self.min_wait, min=self.min_wait, max=self.max_wait
+            ),
+            retry=retry_if_exception_type(ExtractionError),
+            before_sleep=before_sleep_log(logger, logging.DEBUG),
+            after=after_log(logger, logging.DEBUG),
+        )
+
+    @handle_errors(error_message="LLM completion failed", error_type=ExtractionError)
+    def complete(
+        self,
+        messages: List[Dict[str, str]],
+        response_model: Type[ResponseType],
+        config: DictStrAny,
+        step: ExtractionStep,
+    ) -> ResponseType:
+        """
+        Perform LLM completion and track token usage.
+
+        Args:
+            messages: Messages for the completion
+            response_model: Pydantic model for response
+            config: Configuration for the completion
+            step: Step being performed for usage tracking
+
+        Returns:
+            Completion result
+        """
+        result, completion = self.client.chat.completions.create_with_completion(
+            model=self.model_name,
+            response_model=response_model,
+            messages=messages,
+            **config,
+        )
+
+        usage = StepUsage.from_completion(completion, step)
+
+        # Add to usage tracking if available (thread-safe)
+        if usage:
+            with self.usage_lock:
+                self.usage.add_step_usage(usage)
+            logger.debug(f"Step {step.value}: {usage.total_tokens} tokens used")
+
+        return result
+
+    def complete_with_retry(
+        self,
+        messages: List[Dict[str, str]],
+        response_model: Type[ResponseType],
+        config: DictStrAny,
+        step: ExtractionStep,
+    ) -> ResponseType:
+        """
+        Perform LLM completion with retry logic.
+
+        Args:
+            messages: Messages for the completion
+            response_model: Pydantic model for response
+            config: Configuration for the completion
+            step: Step being performed for usage tracking
+
+        Returns:
+            Completion result
+        """
+        retry_decorator = self._create_retry_decorator()
+
+        @retry_decorator
+        def _complete():
+            return self.complete(messages, response_model, config, step)
+
+        return _complete()
+
+    @handle_errors(error_message="Query refinement failed", error_type=ExtractionError)
+    def refine_query(self, query: str) -> QueryRefinement:
+        """
+        Refine and expand query with structural requirements.
+
+        Args:
+            query: Original query string
+
+        Returns:
+            Refined query with additional details
+        """
+        return self.complete(
+            messages=[
+                {"role": "system", "content": query_refinement_system_prompt},
+                {
+                    "role": "user",
+                    "content": query_refinement_template.substitute(query=query),
+                },
+            ],
+            response_model=QueryRefinement,
+            config=self.config.refinement,
+            step=ExtractionStep.REFINEMENT,
+        )
+
+    @handle_errors(error_message="Guide generation failed", error_type=ExtractionError)
+    def generate_extraction_guide(
+        self, refined_query: QueryRefinement, data_columns: List[str]
+    ) -> ExtractionGuide:
+        """
+        Generate extraction guide based on refined query.
+
+        Args:
+            refined_query: Refined query with details
+            data_columns: Available data columns
+
+        Returns:
+            Extraction guide with patterns and rules
+        """
+        return self.complete(
+            messages=[
+                {"role": "system", "content": guide_system_prompt},
+                {
+                    "role": "user",
+                    "content": guide_template.substitute(
+                        data_characteristics=refined_query.data_characteristics,
+                        available_columns=data_columns,
+                    ),
+                },
+            ],
+            response_model=ExtractionGuide,
+            config=self.config.refinement,
+            step=ExtractionStep.GUIDE,
+        )