corp-extractor 0.4.0__py3-none-any.whl → 0.9.0__py3-none-any.whl

statement_extractor/models/canonical.py

@@ -0,0 +1,182 @@
+"""
+Canonical models for the extraction pipeline.
+
+CanonicalMatch: Result of matching to a canonical form
+CanonicalEntity: Entity with canonical form from Stage 4
+"""
+
+from typing import Optional
+
+from pydantic import BaseModel, Field
+
+from .qualifiers import QualifiedEntity
+
+
+class CanonicalMatch(BaseModel):
+    """
+    Result of matching an entity to its canonical form in Stage 4.
+
+    Contains information about how the match was made and confidence level.
+    """
+    canonical_id: Optional[str] = Field(
+        None,
+        description="ID in canonical database (e.g., LEI, Wikidata QID)"
+    )
+    canonical_name: Optional[str] = Field(
+        None,
+        description="Canonical name/label"
+    )
+    match_method: str = Field(
+        ...,
+        description="How the match was made: 'identifier', 'name_exact', 'name_fuzzy', 'llm_verified'"
+    )
+    match_confidence: float = Field(
+        default=1.0,
+        ge=0.0,
+        le=1.0,
+        description="Confidence in the canonical match"
+    )
+    match_details: Optional[dict] = Field(
+        None,
+        description="Additional details about the match (e.g., fuzzy score, LLM reasoning)"
+    )
+
+    def is_high_confidence(self, threshold: float = 0.85) -> bool:
+        """Check if this is a high-confidence match."""
+        return self.match_confidence >= threshold
+
+
+class CanonicalEntity(BaseModel):
+    """
+    An entity with canonical form from Stage 4 (Canonicalization).
+
+    Contains the qualified entity plus its canonical match (if found)
+    and a fully qualified name (FQN) for display.
+    """
+    entity_ref: str = Field(..., description="Reference to the original ExtractedEntity")
+    qualified_entity: QualifiedEntity = Field(
+        ...,
+        description="The qualified entity from Stage 3"
+    )
+    canonical_match: Optional[CanonicalMatch] = Field(
+        None,
+        description="Canonical match if found"
+    )
+    fqn: str = Field(
+        ...,
+        description="Fully qualified name, e.g., 'AMAZON CORP INC (SEC-CIK,USA)'"
+    )
+
+    @property
+    def name(self) -> Optional[str]:
+        """Get the canonical/legal name if available."""
+        # Prefer legal_name from qualifiers (set by embedding qualifier)
+        if self.qualified_entity.qualifiers.legal_name:
+            return self.qualified_entity.qualifiers.legal_name
+        # Fall back to canonical match name
+        if self.canonical_match and self.canonical_match.canonical_name:
+            return self.canonical_match.canonical_name
+        return None
+
+    @property
+    def qualifiers_dict(self) -> Optional[dict[str, str]]:
+        """
+        Get qualifiers as a dict for serialization.
+
+        Returns a dict with keys like: legal_name, region, source, source_id
+        Only returns non-None values. Returns None if no qualifiers are set.
+        """
+        qualifiers = self.qualified_entity.qualifiers
+        identifiers = qualifiers.identifiers
+        result = {}
+
+        # Add legal name
+        if qualifiers.legal_name:
+            result["legal_name"] = qualifiers.legal_name
+
+        # Add region (prefer region, fall back to jurisdiction/country)
+        if qualifiers.region:
+            result["region"] = qualifiers.region
+        elif qualifiers.jurisdiction:
+            result["region"] = qualifiers.jurisdiction
+        elif qualifiers.country:
+            result["region"] = qualifiers.country
+
+        # Add source and source_id from identifiers
+        if "source" in identifiers:
+            result["source"] = identifiers["source"]
+        if "source_id" in identifiers:
+            result["source_id"] = identifiers["source_id"]
+
+        return result if result else None
+
+    @classmethod
+    def from_qualified(
+        cls,
+        qualified: QualifiedEntity,
+        canonical_match: Optional[CanonicalMatch] = None,
+        fqn: Optional[str] = None,
+    ) -> "CanonicalEntity":
+        """Create a CanonicalEntity from a QualifiedEntity."""
+        if fqn is None:
+            # Generate default FQN from qualifiers
+            fqn = cls._generate_fqn(qualified, canonical_match)
+
+        return cls(
+            entity_ref=qualified.entity_ref,
+            qualified_entity=qualified,
+            canonical_match=canonical_match,
+            fqn=fqn,
+        )
+
+    @staticmethod
+    def _generate_fqn(
+        qualified: QualifiedEntity,
+        canonical_match: Optional[CanonicalMatch] = None
+    ) -> str:
+        """
+        Generate a fully qualified name from qualifiers.
+
+        Examples:
+        - PERSON with role+org: "Tim Cook (CEO, Apple Inc)"
+        - ORG with canonical: "Apple Inc (AAPL)"
+        - PERSON with no qualifiers: "Tim Cook"
+        """
+        # Use canonical name if available, otherwise fall back to original text
+        if canonical_match and canonical_match.canonical_name:
+            base_name = canonical_match.canonical_name
+        else:
+            base_name = qualified.original_text
+
+        qualifiers = qualified.qualifiers
+        parts = []
+        seen = set()  # Track seen values to avoid duplicates
+
+        def add_part(value: str) -> None:
+            """Add a part if not already seen (case-insensitive)."""
+            if value and value.lower() not in seen:
+                parts.append(value)
+                seen.add(value.lower())
+
+        # Add role for PERSON entities
+        if qualifiers.role:
+            add_part(qualifiers.role)
+
+        # Add organization for PERSON entities
+        if qualifiers.org:
+            add_part(qualifiers.org)
+
+        # Add ticker for ORG entities
+        if "ticker" in qualifiers.identifiers:
+            add_part(qualifiers.identifiers["ticker"])
+
+        # Add jurisdiction if relevant
+        if qualifiers.jurisdiction and not qualifiers.org:
+            add_part(qualifiers.jurisdiction)
+
+        if parts:
+            return f"{base_name} ({', '.join(parts)})"
+        return base_name
+
+    class Config:
+        frozen = False  # Allow modification during pipeline stages

statement_extractor/models/document.py

@@ -0,0 +1,308 @@
+"""
+Document models for document-level processing.
+
+Document: A document with metadata, pages, and optional summary
+DocumentMetadata: Metadata about the document source
+DocumentPage: A single page within a document
+TextChunk: A chunk of text for processing with page tracking
+ChunkingConfig: Configuration for text chunking
+"""
+
+import uuid
+from typing import Any, Optional
+
+from pydantic import BaseModel, Field
+
+
+class DocumentMetadata(BaseModel):
+    """
+    Metadata about a document source.
+
+    Contains information about where the document came from and
+    who authored it, useful for generating citations.
+    """
+    url: Optional[str] = Field(None, description="URL source of the document")
+    title: Optional[str] = Field(None, description="Document title")
+    year: Optional[int] = Field(None, description="Publication year")
+    authors: list[str] = Field(default_factory=list, description="List of authors")
+    source_type: Optional[str] = Field(
+        None,
+        description="Type of source: 'pdf', 'webpage', 'text', etc."
+    )
+    custom: dict[str, Any] = Field(
+        default_factory=dict,
+        description="Custom metadata fields"
+    )
+
+    def format_citation(self, page_number: Optional[int] = None) -> str:
+        """
+        Format a citation string for this document.
+
+        Args:
+            page_number: Optional page number to include
+
+        Returns:
+            Citation string like "Title - Author, 2024, p. 5"
+        """
+        parts = []
+
+        if self.title:
+            parts.append(self.title)
+
+        if self.authors:
+            if len(self.authors) == 1:
+                parts.append(self.authors[0])
+            elif len(self.authors) == 2:
+                parts.append(f"{self.authors[0]} & {self.authors[1]}")
+            else:
+                parts.append(f"{self.authors[0]} et al.")
+
+        if self.year:
+            parts.append(str(self.year))
+
+        if page_number is not None:
+            parts.append(f"p. {page_number}")
+
+        return " - ".join(parts) if parts else ""
+
+
+class DocumentPage(BaseModel):
+    """
+    A single page within a document.
+
+    Tracks the page number and character offset for citation purposes.
+    """
+    page_number: int = Field(..., description="1-indexed page number")
+    text: str = Field(..., description="Text content of the page")
+    char_offset: int = Field(
+        ...,
+        description="Character offset of this page in the full document text"
+    )
+
+    @property
+    def char_end(self) -> int:
+        """Get the ending character offset of this page."""
+        return self.char_offset + len(self.text)
+
+
+class TextChunk(BaseModel):
+    """
+    A chunk of text for processing.
+
+    Contains the text along with position tracking for mapping
+    extracted statements back to their source pages.
+    """
+    chunk_index: int = Field(..., description="0-indexed chunk number")
+    text: str = Field(..., description="Chunk text content")
+    start_char: int = Field(..., description="Starting character offset in full document")
+    end_char: int = Field(..., description="Ending character offset in full document")
+    page_numbers: list[int] = Field(
+        default_factory=list,
+        description="Page numbers this chunk spans (1-indexed)"
+    )
+    token_count: int = Field(..., description="Number of tokens in this chunk")
+    overlap_chars: int = Field(
+        default=0,
+        description="Number of characters of overlap from previous chunk"
+    )
+    document_id: str = Field(..., description="ID of the source document")
+
+    @property
+    def primary_page(self) -> Optional[int]:
+        """Get the primary page number for this chunk (first page)."""
+        return self.page_numbers[0] if self.page_numbers else None
+
+
+class ChunkingConfig(BaseModel):
+    """
+    Configuration for document chunking.
+
+    Controls how documents are split into chunks for processing.
+    """
+    max_tokens: int = Field(
+        default=2000,
+        ge=100,
+        description="Maximum tokens per chunk (hard limit)"
+    )
+    target_tokens: int = Field(
+        default=1000,
+        ge=50,
+        description="Target tokens per chunk (soft limit, prefers to split here)"
+    )
+    overlap_tokens: int = Field(
+        default=100,
+        ge=0,
+        description="Tokens of overlap between consecutive chunks"
+    )
+    respect_page_boundaries: bool = Field(
+        default=True,
+        description="Try to split at page boundaries when possible"
+    )
+    respect_sentence_boundaries: bool = Field(
+        default=True,
+        description="Try to split at sentence boundaries when possible"
+    )
+
+
+class Document(BaseModel):
+    """
+    A document for processing through the extraction pipeline.
+
+    Contains the full text, optional page structure, metadata for citations,
+    and an optional summary for context.
+    """
+    document_id: str = Field(
+        default_factory=lambda: str(uuid.uuid4()),
+        description="Unique identifier for this document"
+    )
+    metadata: DocumentMetadata = Field(
+        default_factory=DocumentMetadata,
+        description="Document metadata for citations"
+    )
+    pages: list[DocumentPage] = Field(
+        default_factory=list,
+        description="List of pages (optional, for PDFs)"
+    )
+    full_text: str = Field(
+        default="",
+        description="Full text content of the document"
+    )
+    summary: Optional[str] = Field(
+        None,
+        description="Generated summary of the document"
+    )
+
+    @classmethod
+    def from_text(
+        cls,
+        text: str,
+        title: Optional[str] = None,
+        url: Optional[str] = None,
+        **metadata_kwargs,
+    ) -> "Document":
+        """
+        Create a document from plain text.
+
+        Args:
+            text: The document text
+            title: Optional document title
+            url: Optional source URL
+            **metadata_kwargs: Additional metadata fields
+
+        Returns:
+            Document instance
+        """
+        metadata = DocumentMetadata(
+            title=title,
+            url=url,
+            source_type="text",
+            **metadata_kwargs,
+        )
+        return cls(
+            metadata=metadata,
+            full_text=text,
+        )
+
+    @classmethod
+    def from_pages(
+        cls,
+        pages: list[str],
+        title: Optional[str] = None,
+        source_type: str = "pdf",
+        **metadata_kwargs,
+    ) -> "Document":
+        """
+        Create a document from a list of page texts.
+
+        Args:
+            pages: List of page text strings (0-indexed input, stored as 1-indexed)
+            title: Optional document title
+            source_type: Source type (default: "pdf")
+            **metadata_kwargs: Additional metadata fields
+
+        Returns:
+            Document instance
+        """
+        metadata = DocumentMetadata(
+            title=title,
+            source_type=source_type,
+            **metadata_kwargs,
+        )
+
+        # Build pages with character offsets
+        doc_pages = []
+        char_offset = 0
+
+        for i, page_text in enumerate(pages):
+            doc_pages.append(DocumentPage(
+                page_number=i + 1,  # 1-indexed
+                text=page_text,
+                char_offset=char_offset,
+            ))
+            char_offset += len(page_text)
+            if i < len(pages) - 1:
+                char_offset += 1  # Account for newline between pages
+
+        # Join pages with newlines for full text
+        full_text = "\n".join(pages)
+
+        return cls(
+            metadata=metadata,
+            pages=doc_pages,
+            full_text=full_text,
+        )
+
+    def get_page_at_char(self, char_offset: int) -> Optional[int]:
+        """
+        Get the page number containing a character offset.
+
+        Args:
+            char_offset: Character offset in full_text
+
+        Returns:
+            1-indexed page number, or None if no pages defined
+        """
+        if not self.pages:
+            return None
+
+        for page in self.pages:
+            if page.char_offset <= char_offset < page.char_end:
+                return page.page_number
+
+        # If past the last page, return last page
+        if char_offset >= self.pages[-1].char_end:
+            return self.pages[-1].page_number
+
+        return None
+
+    def get_pages_in_range(self, start_char: int, end_char: int) -> list[int]:
+        """
+        Get all page numbers that overlap with a character range.
+
+        Args:
+            start_char: Start character offset
+            end_char: End character offset
+
+        Returns:
+            List of 1-indexed page numbers
+        """
+        if not self.pages:
+            return []
+
+        page_numbers = []
+        for page in self.pages:
+            # Check if page overlaps with range
+            if page.char_offset < end_char and page.char_end > start_char:
+                page_numbers.append(page.page_number)
+
+        return page_numbers
+
+    @property
+    def page_count(self) -> int:
+        """Get the number of pages in the document."""
+        return len(self.pages)
+
+    @property
+    def char_count(self) -> int:
+        """Get the total character count."""
+        return len(self.full_text)

statement_extractor/models/entity.py

@@ -0,0 +1,102 @@
+"""
+Entity models for the extraction pipeline.
+
+ExtractedEntity represents entities identified during extraction with
+confidence scores and span information.
+
+Note: EntityType is imported from the original models.py for consistency.
+"""
+
+from typing import Optional, TYPE_CHECKING
+import uuid
+
+from pydantic import BaseModel, Field
+
+# Import EntityType from parent module to avoid duplication
+# This will be populated by __init__.py which loads from old models.py
+if TYPE_CHECKING:
+    from enum import Enum
+
+    class EntityType(str, Enum):
+        """Supported entity types for subjects and objects."""
+        ORG = "ORG"
+        PERSON = "PERSON"
+        GPE = "GPE"
+        LOC = "LOC"
+        PRODUCT = "PRODUCT"
+        EVENT = "EVENT"
+        WORK_OF_ART = "WORK_OF_ART"
+        LAW = "LAW"
+        DATE = "DATE"
+        MONEY = "MONEY"
+        PERCENT = "PERCENT"
+        QUANTITY = "QUANTITY"
+        UNKNOWN = "UNKNOWN"
+else:
+    # At runtime, we need to import it from somewhere
+    # Try the old models.py location first
+    try:
+        import importlib.util
+        from pathlib import Path
+        _models_py_path = Path(__file__).parent.parent / "models.py"
+        _spec = importlib.util.spec_from_file_location("_old_models", _models_py_path)
+        _old_models = importlib.util.module_from_spec(_spec)
+        _spec.loader.exec_module(_old_models)
+        EntityType = _old_models.EntityType
+    except Exception:
+        # Fallback to defining it here
+        from enum import Enum
+
+        class EntityType(str, Enum):
+            """Supported entity types for subjects and objects."""
+            ORG = "ORG"
+            PERSON = "PERSON"
+            GPE = "GPE"
+            LOC = "LOC"
+            PRODUCT = "PRODUCT"
+            EVENT = "EVENT"
+            WORK_OF_ART = "WORK_OF_ART"
+            LAW = "LAW"
+            DATE = "DATE"
+            MONEY = "MONEY"
+            PERCENT = "PERCENT"
+            QUANTITY = "QUANTITY"
+            UNKNOWN = "UNKNOWN"
+
+
+class ExtractedEntity(BaseModel):
+    """
+    An entity extracted from text with type and confidence information.
+
+    Used in Stage 2 (Extraction) and flows through subsequent stages.
+    """
+    text: str = Field(..., description="The entity text as extracted")
+    type: EntityType = Field(default=EntityType.UNKNOWN, description="The entity type")
+    span: Optional[tuple[int, int]] = Field(
+        None,
+        description="Character offsets (start, end) in source text"
+    )
+    confidence: float = Field(
+        default=1.0,
+        ge=0.0,
+        le=1.0,
+        description="Confidence score for this entity extraction"
+    )
+    entity_ref: str = Field(
+        default_factory=lambda: str(uuid.uuid4()),
+        description="Unique reference ID for tracking this entity through the pipeline"
+    )
+
+    def __str__(self) -> str:
+        return f"{self.text} ({self.type.value})"
+
+    def __hash__(self) -> int:
+        return hash(self.entity_ref)
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, ExtractedEntity):
+            return False
+        return self.entity_ref == other.entity_ref
+
+    class Config:
+        frozen = False  # Allow modification during pipeline stages