resolvekit 0.0.1__py3-none-any.whl

resolvekit/types.py

@@ -0,0 +1,534 @@
+"""Type definitions for resolvekit."""
+
+from datetime import date
+from enum import StrEnum
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict, Field, field_serializer, field_validator
+from sqlalchemy import Column
+from sqlalchemy import String as SQLAString
+from sqlmodel import Field as SQLModelField
+from sqlmodel import SQLModel
+
+# ==============================================================================
+# Enums (for type safety and SQLModel compatibility)
+# ==============================================================================
+
+
+class EntityType(StrEnum):
+    """Entity type classifications."""
+
+    COUNTRY = "country"
+    ADMIN1 = "admin1"
+    ADMIN2 = "admin2"
+    ADMIN3 = "admin3"
+    ADMIN4 = "admin4"
+    CITY = "city"
+    ORGANIZATION = "organization"
+    GROUP = "group"
+    REGION = "region"
+    OTHER = "other"
+
+
+class AliasType(StrEnum):
+    """Alias type classifications."""
+
+    CANONICAL = "canonical"  # Official canonical name
+    ENDONYM = "endonym"  # Native name in local language
+    EXONYM = "exonym"  # Foreign name
+    ABBR = "abbr"  # Abbreviation
+    CODE = "code"  # Code identifier
+    HISTORICAL = "historical"  # Historical name (no longer current)
+    COLLOQUIAL = "colloquial"  # Informal/colloquial name
+
+
+class CodeSystem(StrEnum):
+    """Code system identifiers."""
+
+    DCID = "dcid"  # Data Commons ID (canonical)
+    ISO2 = "iso2"  # ISO 3166-1 alpha-2
+    ISO3 = "iso3"  # ISO 3166-1 alpha-3
+    ISO_NUMERIC = "iso_numeric"  # ISO 3166-1 numeric
+    ISO3166_2 = "iso3166_2"  # ISO 3166-2 subdivision codes
+    M49 = "m49"  # UN M49 codes
+    NUTS = "nuts"  # EU NUTS codes
+    LAU = "lau"  # EU LAU codes
+    PCODE = "pcode"  # OCHA P-codes
+    WIKIDATA = "wikidata"  # Wikidata QID
+    GEONAMES = "geonames"  # GeoNames ID
+    WB = "wb"  # World Bank codes
+    DAC = "dac"  # OECD DAC codes
+    FIPS = "fips"  # FIPS codes (deprecated but still used)
+
+
+class MatcherType(StrEnum):
+    """Matcher type classifications."""
+
+    EXACT_CODE = "exact_code"
+    CANONICAL_NAME = "canonical_name"
+    ALIAS_EXACT = "alias_exact"
+    FTS = "fts"
+    FUZZY = "fuzzy"
+    SEMANTIC = "semantic"
+
+
+class OutputFormat(StrEnum):
+    """Output format options."""
+
+    HUMAN = "human"
+    JSON = "json"
+    CSV = "csv"
+    TABLE = "table"
+
+
+class ExplanationMode(StrEnum):
+    """Explanation detail level."""
+
+    MINIMAL = "minimal"  # Just confidence scores, no explanation object
+    STANDARD = "standard"  # Matched features + top alternatives
+    FULL = "full"  # Complete debug trace with all stages
+
+
+class EntityRow(SQLModel, table=True):
+    """Database row representation of an entity (no computed fields)."""
+
+    __tablename__ = "entities"
+
+    # Configure model to use Enum values (not names) in database
+    model_config = {"use_enum_values": True}
+
+    dcid: str = SQLModelField(
+        primary_key=True, min_length=1, description="Data Commons ID"
+    )
+    canonical_name: str = SQLModelField(
+        min_length=1, description="Canonical entity name"
+    )
+    normalized_canonical: str = SQLModelField(
+        min_length=1,
+        index=True,
+        description="Normalized canonical name for efficient lookups",
+    )
+    entity_type: EntityType = SQLModelField(
+        sa_column=Column(SQLAString)
+    )  # Store as string
+    parent_dcid: str | None = SQLModelField(
+        default=None, foreign_key="entities.dcid", description="Parent entity DCID"
+    )
+    centroid_lat: float | None = SQLModelField(
+        default=None, ge=-90, le=90, description="Latitude"
+    )
+    centroid_lon: float | None = SQLModelField(
+        default=None, ge=-180, le=180, description="Longitude"
+    )
+    valid_from: date | None = None
+    valid_until: date | None = None
+
+    @field_validator("valid_from", "valid_until", mode="before")
+    @classmethod
+    def parse_date(cls, v):
+        """Parse date from various formats."""
+        if v is None or isinstance(v, date):
+            return v
+        if isinstance(v, str):
+            from datetime import datetime
+
+            return datetime.fromisoformat(v).date()
+        return v
+
+    @field_serializer("entity_type", when_used="always")
+    def serialize_entity_type(self, value: EntityType) -> str:
+        """Serialize EntityType enum to its string value for database storage."""
+        return value.value if isinstance(value, EntityType) else value
+
+
+class Entity(BaseModel):
+    """Represents a resolved entity."""
+
+    model_config = ConfigDict(frozen=False)
+
+    dcid: str = Field(..., min_length=1, description="Data Commons ID")
+    canonical_name: str = Field(..., min_length=1, description="Canonical entity name")
+    entity_type: EntityType
+    codes: dict[str, str] = Field(default_factory=dict, description="Code mappings")
+    parent_dcid: str | None = Field(None, description="Parent entity DCID")
+    centroid_lat: float | None = Field(None, ge=-90, le=90, description="Latitude")
+    centroid_lon: float | None = Field(None, ge=-180, le=180, description="Longitude")
+    valid_from: date | None = Field(None, description="Validity start date")
+    valid_until: date | None = Field(None, description="Validity end date")
+    provenance: dict[str, str] = Field(
+        default_factory=dict, description="Data source provenance"
+    )
+
+    @field_validator("valid_from", "valid_until", mode="before")
+    @classmethod
+    def parse_date(cls, v):
+        """Parse date from various formats."""
+        if v is None or isinstance(v, date):
+            return v
+        if isinstance(v, str):
+            from datetime import datetime
+
+            return datetime.fromisoformat(v).date()
+        return v
+
+    def __repr__(self) -> str:
+        """Technical representation."""
+        return f"Entity(dcid='{self.dcid}', name='{self.canonical_name}')"
+
+    def __str__(self) -> str:
+        """Human-friendly representation."""
+        lines = [
+            f"{self.canonical_name}",
+            f"  DCID: {self.dcid}",
+            f"  Type: {self.entity_type.value}",
+        ]
+
+        if self.codes:
+            lines.append("  Codes:")
+            for system, code in sorted(self.codes.items())[:5]:
+                lines.append(f"    {system}: {code}")
+
+        if self.parent_dcid:
+            lines.append(f"  Parent: {self.parent_dcid}")
+
+        return "\n".join(lines)
+
+
+class AliasRow(SQLModel, table=True):
+    """Database row representation of an alias."""
+
+    __tablename__ = "aliases"
+
+    # Configure model to use Enum values (not names) in database
+    model_config = {"use_enum_values": True}
+
+    alias_id: int = SQLModelField(
+        default=None, primary_key=True, ge=0, description="Unique alias ID"
+    )
+    entity_dcid: str = SQLModelField(
+        foreign_key="entities.dcid", min_length=1, description="Associated entity DCID"
+    )
+    alias_text: str = SQLModelField(min_length=1, description="Original alias text")
+    alias_norm: str = SQLModelField(
+        min_length=1, index=True, description="Normalized alias text"
+    )
+    language: str | None = SQLModelField(
+        default=None, max_length=2, description="ISO 639-1 language code"
+    )
+    alias_type: AliasType = SQLModelField(
+        default=AliasType.EXONYM, sa_column=Column(SQLAString)
+    )
+    valid_from: date | None = None
+    valid_until: date | None = None
+    source: str | None = None
+    alias_uid: str = SQLModelField(
+        min_length=1, unique=True, description="Unique alias identifier"
+    )
+
+    @field_validator("valid_from", "valid_until", mode="before")
+    @classmethod
+    def parse_date(cls, v):
+        """Parse date from various formats."""
+        if v is None or isinstance(v, date):
+            return v
+        if isinstance(v, str):
+            from datetime import datetime
+
+            return datetime.fromisoformat(v).date()
+        return v
+
+    @field_serializer("alias_type", when_used="always")
+    def serialize_alias_type(self, value: AliasType) -> str:
+        """Serialize AliasType enum to its string value for database storage."""
+        return value.value if isinstance(value, AliasType) else value
+
+
+class Alias(BaseModel):
+    """Represents an entity alias (business logic model)."""
+
+    model_config = ConfigDict(frozen=False)
+
+    alias_id: int = Field(..., ge=0)
+    entity_dcid: str = Field(..., min_length=1)
+    alias_text: str = Field(..., min_length=1)
+    alias_norm: str = Field(..., min_length=1)
+    language: str | None = Field(None, pattern=r"^[a-z]{2}$")
+    alias_type: AliasType = "exonym"
+    valid_from: date | None = None
+    valid_until: date | None = None
+    source: str | None = None
+
+
+class MatchContext(BaseModel):
+    """Optional filtering context for matching."""
+
+    entity_type: EntityType | None = None
+    parent_dcid: str | None = None
+    as_of: date | None = None
+    group_dcid: str | None = None  # Filter by group membership
+
+
+class Candidate(BaseModel):
+    """
+    A resolution candidate with features for calibration.
+
+    Attributes:
+        entity: Matched entity
+        score: Raw matcher score (0.0-1.0)
+        matcher_type: Which matcher produced this candidate
+        features: Feature dictionary for calibration
+        matched_alias: The specific alias text that matched (if applicable)
+
+    Common features:
+        - exact_code: bool - Matched via code lookup
+        - canonical_exact: bool - Exact canonical name match
+        - alias_exact: bool - Exact alias match
+        - fts_rank: int | None - FTS result rank
+        - fts_score: float | None - FTS BM25 score
+        - edit_similarity: float | None - Edit distance (0-1)
+        - trigram_jaccard: float | None - Trigram Jaccard similarity
+        - fuzzy_score: float | None - Combined fuzzy score
+        - code_system: str | None - Which code system matched
+        - matched_alias: str | None - Which alias matched
+    """
+
+    model_config = ConfigDict(frozen=False, arbitrary_types_allowed=True)
+
+    entity: Entity
+    score: float = Field(..., ge=0.0, le=1.0, description="Raw matcher score (0.0-1.0)")
+    matcher_type: MatcherType
+    features: dict[str, Any] = Field(
+        default_factory=dict, description="Feature dictionary for calibration"
+    )
+    matched_alias: str | None = Field(
+        None, description="The specific alias text that matched"
+    )
+
+    @property
+    def dcid(self) -> str:
+        """Get entity DCID."""
+        return self.entity.dcid
+
+
+class Explanation(BaseModel):
+    """Explains how a resolution was made."""
+
+    model_config = ConfigDict(frozen=False)
+
+    # STANDARD mode fields
+    matcher_used: MatcherType | None = None
+    features: dict[str, Any] = Field(default_factory=dict)
+
+    # Already exists
+    stages: list[str] = Field(default_factory=list)
+    candidates: list[Candidate] = Field(default_factory=list)
+    rules_applied: list[str] = Field(default_factory=list)
+    calibration: dict[str, float] = Field(default_factory=dict)
+    trace: dict = Field(default_factory=dict)
+
+    def __repr__(self) -> str:
+        """Technical representation."""
+        return (
+            f"Explanation(matcher={self.matcher_used}, features={len(self.features)})"
+        )
+
+    def __str__(self) -> str:
+        """Human-friendly scorecard."""
+        lines = ["Resolution Scorecard:", "=" * 50]
+
+        if self.matcher_used:
+            lines.append(f"Matcher: {self.matcher_used.value}")
+
+        # Show key features
+        if self.features:
+            lines.append("\nKey Features:")
+            for key, value in list(self.features.items())[:5]:
+                if isinstance(value, bool):
+                    lines.append(f"  {key}: {'✓' if value else '✗'}")
+                elif isinstance(value, float):
+                    lines.append(f"  {key}: {value:.3f}")
+                else:
+                    lines.append(f"  {key}: {value}")
+
+        # Show alternatives (candidates)
+        if self.candidates:
+            lines.append(f"\nAlternatives ({len(self.candidates)}):")
+            for i, cand in enumerate(self.candidates[:5], 1):
+                lines.append(
+                    f"  {i}. {cand.entity.canonical_name} (score: {cand.score:.3f})"
+                )
+
+        # FULL mode: show trace summary
+        if self.trace:
+            lines.append(f"\nTrace: {len(self.trace)} stages recorded")
+
+        return "\n".join(lines)
+
+
+class Resolution(BaseModel):
+    """Result of entity resolution."""
+
+    model_config = ConfigDict(frozen=False)
+
+    entity: Entity | None = None
+    confidence: float = Field(..., ge=0.0, le=1.0, description="Calibrated confidence")
+    alternatives: list[Entity] = Field(default_factory=list)
+    explanation: Explanation | None = None
+    conflicts: list[dict] | None = None
+
+    @property
+    def success(self) -> bool:
+        """Check if resolution was successful."""
+        return self.entity is not None
+
+    def __repr__(self) -> str:
+        """Technical representation."""
+        return (
+            f"Resolution(entity={self.entity.dcid if self.entity else None}, "
+            f"confidence={self.confidence:.3f}, "
+            f"alternatives={len(self.alternatives)})"
+        )
+
+    def __str__(self) -> str:
+        """Human-friendly output."""
+        if not self.success:
+            return f"No match found (confidence: {self.confidence:.3f})"
+
+        lines = [
+            f"Resolved to: {self.entity.canonical_name}",
+            f"  DCID: {self.entity.dcid}",
+            f"  Confidence: {self.confidence:.3f}",
+        ]
+
+        if self.alternatives:
+            lines.append(f"  Alternatives: {len(self.alternatives)}")
+            for i, alt in enumerate(self.alternatives[:3], 1):
+                lines.append(f"    {i}. {alt.canonical_name} ({alt.dcid})")
+
+        if self.explanation and self.explanation.matcher_used:
+            lines.append(f"  Matched via: {self.explanation.matcher_used.value}")
+
+        return "\n".join(lines)
+
+
+class ExtractedEntity(BaseModel):
+    """Entity extracted from text."""
+
+    model_config = ConfigDict(frozen=False)
+
+    text: str = Field(..., min_length=1)
+    span: tuple[int, int]
+    dcid: str = Field(..., min_length=1)
+    canonical_name: str = Field(..., min_length=1)
+    entity_type: EntityType
+    confidence: float = Field(..., ge=0.0, le=1.0)
+    context: str | None = None
+    method: str | None = None
+
+    @field_validator("span")
+    @classmethod
+    def validate_span(cls, v):
+        """Validate span is valid (start < end)."""
+        if len(v) != 2:
+            raise ValueError("span must be a tuple of (start, end)")
+        if v[0] > v[1]:
+            raise ValueError("span start must be <= end")
+        return v
+
+
+class CodeRow(SQLModel, table=True):
+    """Database row representation of a code mapping."""
+
+    __tablename__ = "codes"
+
+    # Configure model to use Enum values (not names) in database
+    model_config = {"use_enum_values": True}
+
+    entity_dcid: str = SQLModelField(
+        foreign_key="entities.dcid",
+        primary_key=True,
+        min_length=1,
+        description="Associated entity DCID",
+    )
+    code_system: CodeSystem = SQLModelField(
+        sa_column=Column(SQLAString, primary_key=True),
+        description="Code system identifier",
+    )
+    code_value: str = SQLModelField(min_length=1, description="Code value")
+    valid_from: date | None = None
+    valid_until: date | None = None
+    source: str | None = None
+
+    @field_validator("valid_from", "valid_until", mode="before")
+    @classmethod
+    def parse_date(cls, v):
+        """Parse date from various formats."""
+        if v is None or isinstance(v, date):
+            return v
+        if isinstance(v, str):
+            from datetime import datetime
+
+            return datetime.fromisoformat(v).date()
+        return v
+
+    @field_serializer("code_system", when_used="always")
+    def serialize_code_system(self, value: CodeSystem) -> str:
+        """Serialize CodeSystem enum to its string value for database storage."""
+        return value.value if isinstance(value, CodeSystem) else value
+
+
+class MembershipRow(SQLModel, table=True):
+    """Database row representation of a group membership."""
+
+    __tablename__ = "memberships"
+
+    id: int = SQLModelField(
+        default=None, primary_key=True, ge=0, description="Unique membership ID"
+    )
+    entity_dcid: str = SQLModelField(
+        foreign_key="entities.dcid", min_length=1, description="Member entity DCID"
+    )
+    group_dcid: str = SQLModelField(
+        foreign_key="entities.dcid", min_length=1, description="Group entity DCID"
+    )
+    valid_from: date
+    valid_until: date | None = None
+    source: str | None = None
+
+    @field_validator("valid_from", "valid_until", mode="before")
+    @classmethod
+    def parse_date(cls, v):
+        """Parse date from various formats."""
+        if v is None or isinstance(v, date):
+            return v
+        if isinstance(v, str):
+            from datetime import datetime
+
+            return datetime.fromisoformat(v).date()
+        return v
+
+
+class Membership(BaseModel):
+    """Group membership record (business logic model)."""
+
+    model_config = ConfigDict(frozen=False)
+
+    entity_dcid: str = Field(..., min_length=1)
+    group_dcid: str = Field(..., min_length=1)
+    valid_from: date
+    valid_until: date | None = None
+    source: str | None = None
+
+    @field_validator("valid_until")
+    @classmethod
+    def validate_dates(cls, v, info):
+        """Validate that valid_until > valid_from."""
+        if v is not None and "valid_from" in info.data and v <= info.data["valid_from"]:
+            raise ValueError("valid_until must be after valid_from")
+        return v
+
+    def is_valid_at(self, at: date) -> bool:
+        """Check if membership is valid at given date."""
+        if at < self.valid_from:
+            return False
+        return not (self.valid_until and at >= self.valid_until)