resolvekit 0.0.1__py3-none-any.whl

resolvekit/constraints/README.md

@@ -0,0 +1,102 @@
+# Constraints Module
+
+## Purpose
+
+The constraints module enforces Knowledge Graph (KG) and temporal constraints on candidate entities, filtering out invalid matches based on type, hierarchy, temporal validity, and group memberships.
+
+## Components
+
+### Constraint Validators
+
+1. **Type Validator** (`type_validator.py`)
+   - Enforces entity type constraints when provided
+   - Example: If `entity_type="country"`, filter out states/cities
+
+2. **Hierarchy Validator** (`hierarchy_validator.py`)
+   - Checks parent-child containment relationships
+   - Example: If `parent="country/FRA"`, only keep French subdivisions
+   - Validates full hierarchy paths
+
+3. **Temporal Validator** (`temporal_validator.py`)
+   - Checks validity ranges (valid_from, valid_until)
+   - Handles historical entity names
+   - Default: current date (pack build date)
+   - Supports `as_of` parameter for historical queries
+
+4. **Membership Validator** (`membership_validator.py`)
+   - Validates group memberships at specific dates
+   - Example: Check if entity was EU member in 2005
+   - Handles joining/leaving events (Brexit, EU expansions)
+
+### Helper Components
+
+- `constraint_engine.py`: Orchestrates constraint validation
+- `validity.py`: Temporal validity utilities
+- `hierarchy.py`: Hierarchy traversal utilities
+
+## Constraint Application Flow
+
+```python
+def apply_constraints(
+    candidates: list[Candidate],
+    entity_type: str | None,
+    parent: str | None,
+    at: date | None,
+    group: str | None
+) -> list[Candidate]:
+    """Apply all constraints to filter candidates."""
+
+    # Type constraint
+    if entity_type:
+        candidates = [c for c in candidates if c.type == entity_type]
+
+    # Hierarchy constraint
+    if parent:
+        candidates = [c for c in candidates if is_child_of(c.dcid, parent)]
+
+    # Temporal constraint
+    if at:
+        candidates = [c for c in candidates if is_valid_at(c, at)]
+
+    # Membership constraint
+    if group and at:
+        candidates = [c for c in candidates
+                     if is_member_at(c.dcid, group, at)]
+
+    # Mark validity in features
+    for c in candidates:
+        c.features.type_valid = (entity_type is None or c.type == entity_type)
+        c.features.parent_valid = (parent is None or is_child_of(c.dcid, parent))
+        c.features.date_valid = (at is None or is_valid_at(c, at))
+
+    return candidates
+```
+
+## Design Principles
+
+1. **Additive filtering**: Each constraint narrows candidate set
+2. **Feature tracking**: Record which constraints passed/failed for calibration
+3. **Graceful degradation**: Missing temporal data doesn't crash queries
+4. **Efficient lookups**: Pre-computed indexes for hierarchy/membership
+
+## Temporal Validity Model
+
+All temporal-aware tables include:
+- `valid_from`: Start date (ISO format) or NULL (always valid)
+- `valid_until`: End date (ISO format, exclusive) or NULL (still valid)
+
+Validity check:
+```python
+def is_valid_at(entity, as_of: date) -> bool:
+    """Check if entity is valid at given date."""
+    if entity.valid_from and as_of < entity.valid_from:
+        return False
+    if entity.valid_until and as_of >= entity.valid_until:
+        return False
+    return True
+```
+
+## Implementation Priority
+
+**Phase A** - Core resolver (type, hierarchy, basic temporal)
+**Phase B** - Full membership queries with temporal support

resolvekit/constraints/__init__.py

@@ -0,0 +1,17 @@
+"""Constraint validation module."""
+
+from resolvekit.constraints.constraint_engine import ConstraintEngine
+from resolvekit.constraints.hierarchy_validator import HierarchyValidator
+from resolvekit.constraints.membership_validator import MembershipValidator
+from resolvekit.constraints.protocols import Validator
+from resolvekit.constraints.temporal_validator import TemporalValidator
+from resolvekit.constraints.type_validator import TypeValidator
+
+__all__ = [
+    "ConstraintEngine",
+    "HierarchyValidator",
+    "MembershipValidator",
+    "TemporalValidator",
+    "TypeValidator",
+    "Validator",
+]

resolvekit/constraints/constraint_engine.py

@@ -0,0 +1,111 @@
+"""Constraint engine orchestrator."""
+
+from resolvekit.constraints.hierarchy_validator import HierarchyValidator
+from resolvekit.constraints.membership_validator import MembershipValidator
+from resolvekit.constraints.temporal_validator import TemporalValidator
+from resolvekit.constraints.type_validator import TypeValidator
+from resolvekit.data.entity_repository import EntityRepository
+from resolvekit.data.membership_repository import MembershipRepository
+from resolvekit.types import Candidate, MatchContext
+
+
+class ConstraintEngine:
+    """
+    Orchestrates constraint validation with hybrid filtering.
+
+    Hard filtering (removes candidates):
+    - Type constraint (when context.entity_type is set)
+    - Parent constraint (when context.parent_dcid is set)
+
+    Soft features (marks but doesn't filter):
+    - Temporal validity (when context.as_of is set)
+    - Membership validity (when context.group_dcid is set)
+    """
+
+    def __init__(
+        self,
+        entity_repo: EntityRepository,
+        membership_repo: MembershipRepository,
+        hierarchy_depth: int = 3,
+    ):
+        """
+        Initialize constraint engine.
+
+        Args:
+            entity_repo: Repository for entity lookups
+            membership_repo: Repository for membership queries
+            hierarchy_depth: Maximum depth for hierarchy traversal (default 3)
+        """
+        self.type_validator = TypeValidator()
+        self.hierarchy_validator = HierarchyValidator(entity_repo, hierarchy_depth)
+        self.temporal_validator = TemporalValidator()
+        self.membership_validator = MembershipValidator(membership_repo)
+        self.hierarchy_depth = hierarchy_depth
+
+    def apply_constraints(
+        self, candidates: list[Candidate], context: MatchContext | None
+    ) -> list[Candidate]:
+        """
+        Apply all constraints with hybrid filtering.
+
+        Args:
+            candidates: List of candidates to validate
+            context: Match context with optional constraints
+
+        Returns:
+            Validated/enriched candidates
+        """
+        # Fast path: empty candidate list
+        if not candidates:
+            return candidates
+
+        # Fast path: no context
+        if context is None:
+            return candidates
+
+        # Fast path: empty context (all fields None)
+        if self._is_empty_context(context):
+            return candidates
+
+        # HARD FILTERS (reduce candidate set)
+        # Order: cheapest first to minimize downstream work
+        candidates = self.type_validator.validate(candidates, context)
+        if not candidates:  # Early exit if all filtered
+            return candidates
+
+        # Early exit: Skip hierarchy validation if no parent constraint
+        if context.parent_dcid is not None:
+            candidates = self.hierarchy_validator.validate(candidates, context)
+            if not candidates:  # Early exit if all filtered
+                return candidates
+        else:
+            # Add default parent features even when constraint is absent.
+            # This ensures calibration model receives consistent feature vectors -
+            # all candidates have the same features regardless of which constraints
+            # are active. Default values indicate "no constraint" rather than "constraint failed".
+            for candidate in candidates:
+                candidate.features["parent_valid"] = True
+                candidate.features["parent_depth"] = 0
+
+        # SOFT FEATURES (enrich remaining candidates)
+        candidates = self.temporal_validator.validate(candidates, context)
+        candidates = self.membership_validator.validate(candidates, context)
+
+        return candidates
+
+    def _is_empty_context(self, context: MatchContext) -> bool:
+        """
+        Check if context has any constraints.
+
+        Args:
+            context: Match context to check
+
+        Returns:
+            True if all constraint fields are None
+        """
+        return (
+            context.entity_type is None
+            and context.parent_dcid is None
+            and context.as_of is None
+            and context.group_dcid is None
+        )

resolvekit/constraints/hierarchy_validator.py

@@ -0,0 +1,148 @@
+"""Hierarchy constraint validator."""
+
+from sqlalchemy import text
+from sqlmodel import Session
+
+from resolvekit.data.entity_repository import EntityRepository
+from resolvekit.types import Candidate, MatchContext
+from resolvekit.utils.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+class HierarchyValidator:
+    """
+    Validates parent-child hierarchy constraints.
+
+    Hard filtering: Removes candidates that are not descendants of context.parent_dcid
+    Features: Adds f_parent_valid and f_parent_depth
+    """
+
+    def __init__(self, entity_repo: EntityRepository, depth_limit: int = 3):
+        """
+        Initialize hierarchy validator.
+
+        Args:
+            entity_repo: Repository for entity lookups
+            depth_limit: Maximum depth to traverse up the hierarchy
+        """
+        self.entity_repo = entity_repo
+        self.depth_limit = depth_limit
+
+    def validate(
+        self, candidates: list[Candidate], context: MatchContext
+    ) -> list[Candidate]:
+        """
+        Validate hierarchy constraints.
+
+        Optimized to check each unique entity_dcid only once, even if multiple
+        candidates share the same entity.
+
+        Args:
+            candidates: List of candidates to validate
+            context: Match context with optional parent_dcid constraint
+
+        Returns:
+            Filtered candidates with f_parent_valid and f_parent_depth features
+        """
+        # Fast path: no parent constraint
+        if context.parent_dcid is None:
+            for candidate in candidates:
+                candidate.features["parent_valid"] = True
+                candidate.features["parent_depth"] = 0
+            return candidates
+
+        # Batch query: check all unique entities in single DB query
+        entity_dcids = list({c.entity.dcid for c in candidates})
+        entity_results = self._batch_check_descendants(
+            entity_dcids, context.parent_dcid
+        )
+
+        # Apply results to all candidates
+        for candidate in candidates:
+            is_valid, depth = entity_results[candidate.entity.dcid]
+            candidate.features["parent_valid"] = is_valid
+            candidate.features["parent_depth"] = depth
+
+        # Hard filter: remove invalid candidates
+        return [c for c in candidates if c.features["parent_valid"]]
+
+    def _batch_check_descendants(
+        self, entity_dcids: list[str], parent_dcid: str
+    ) -> dict[str, tuple[bool, int]]:
+        """
+        Check which entities are descendants of parent using recursive CTE.
+
+        Single database query to check all entities at once, dramatically faster
+        than sequential lookups for large candidate sets. Respects overlay
+        precedence - if an overlay changes an entity's parent_dcid, the overlay
+        version is used for hierarchy traversal.
+
+        Args:
+            entity_dcids: List of entity DCIDs to check
+            parent_dcid: Parent to find in ancestry
+
+        Returns:
+            Dict mapping entity_dcid -> (is_descendant, depth)
+        """
+        if not self.entity_repo.db.engine:
+            raise RuntimeError("Database not connected. Call connect() first.")
+
+        if not entity_dcids:
+            return {}
+
+        # Build placeholders for IN clause
+        placeholders = ", ".join(f":dcid_{i}" for i in range(len(entity_dcids)))
+        params = {f"dcid_{i}": dcid for i, dcid in enumerate(entity_dcids)}
+        params["parent"] = parent_dcid
+        params["depth_limit"] = self.depth_limit
+
+        # Build best_entities CTE - dedupe by precedence
+        # This ensures overlay parent_dcid overrides base parent_dcid
+        best_selects = ["SELECT dcid, parent_dcid, 0 AS precedence FROM main.entities"]
+        for schema_name, precedence in self.entity_repo.db.overlays:
+            best_selects.append(
+                f"SELECT dcid, parent_dcid, {precedence} AS precedence FROM {schema_name}.entities"
+            )
+
+        best_entities_query = "\nUNION ALL\n".join(best_selects)
+
+        # Build recursive CTE - track original candidate through traversal
+        # Uses best_entities to ensure we traverse using highest-precedence parent links
+        # ROW_NUMBER() deduplicates by dcid, keeping highest precedence row
+        sql = f"""
+        WITH RECURSIVE
+        best_entities AS (
+            SELECT dcid, parent_dcid
+            FROM (
+                SELECT dcid, parent_dcid,
+                       ROW_NUMBER() OVER (PARTITION BY dcid ORDER BY precedence DESC) as rn
+                FROM ({best_entities_query}) all_entities
+            ) ranked
+            WHERE rn = 1
+        ),
+        ancestors AS (
+            SELECT dcid AS original_dcid, dcid AS current_dcid, parent_dcid, 0 AS depth
+            FROM best_entities
+            WHERE dcid IN ({placeholders})
+
+            UNION ALL
+
+            SELECT a.original_dcid, e.dcid AS current_dcid, e.parent_dcid, a.depth + 1
+            FROM best_entities e
+            INNER JOIN ancestors a ON e.dcid = a.parent_dcid
+            WHERE a.depth < :depth_limit AND a.parent_dcid IS NOT NULL
+        )
+        SELECT original_dcid AS dcid, MIN(depth) + 1 AS depth
+        FROM ancestors
+        WHERE parent_dcid = :parent AND depth + 1 <= :depth_limit
+        GROUP BY original_dcid
+        """
+
+        # Execute query
+        with Session(self.entity_repo.db.engine) as session:
+            result = session.execute(text(sql), params)
+            valid_entities = {row.dcid: (True, row.depth) for row in result}
+
+        # Build full results map (invalid = False for entities not in results)
+        return {dcid: valid_entities.get(dcid, (False, 0)) for dcid in entity_dcids}

resolvekit/constraints/membership_validator.py

@@ -0,0 +1,60 @@
+"""Membership constraint validator."""
+
+from datetime import date
+
+from resolvekit.data.membership_repository import MembershipRepository
+from resolvekit.types import Candidate, MatchContext
+
+
+class MembershipValidator:
+    """
+    Validates group membership constraints.
+
+    Soft validation only: Adds f_membership_valid feature, does not filter.
+    """
+
+    def __init__(self, membership_repo: MembershipRepository):
+        """
+        Initialize membership validator.
+
+        Args:
+            membership_repo: Repository for membership queries
+        """
+        self.membership_repo = membership_repo
+
+    def validate(
+        self, candidates: list[Candidate], context: MatchContext
+    ) -> list[Candidate]:
+        """
+        Validate membership constraints.
+
+        Args:
+            candidates: List of candidates to validate
+            context: Match context with optional group_dcid constraint
+
+        Returns:
+            Candidates with f_membership_valid feature (no filtering)
+        """
+        # If no group constraint, mark all as valid (no constraint to check)
+        if context.group_dcid is None:
+            for candidate in candidates:
+                candidate.features["membership_valid"] = True
+            return candidates
+
+        # Use as_of date or default to today
+        as_of = context.as_of if context.as_of is not None else date.today()
+
+        # Batch query all candidates
+        entity_dcids = [c.entity.dcid for c in candidates]
+        members = self.membership_repo.check_memberships_batch(
+            entity_dcids=entity_dcids,
+            group_dcid=context.group_dcid,
+            as_of=as_of,
+        )
+
+        # Mark membership validity
+        for candidate in candidates:
+            candidate.features["membership_valid"] = candidate.entity.dcid in members
+
+        # Soft validation: return all candidates (don't filter)
+        return candidates

resolvekit/constraints/protocols.py

@@ -0,0 +1,33 @@
+"""Protocols for constraint validators."""
+
+from typing import Protocol
+
+from resolvekit.types import Candidate, MatchContext
+
+
+class Validator(Protocol):
+    """
+    Protocol for constraint validators.
+
+    Validators filter candidates or enrich them with validation features.
+    They may perform hard filtering (removing invalid candidates) or soft
+    validation (adding features for calibration to consider).
+    """
+
+    def validate(
+        self, candidates: list[Candidate], context: MatchContext
+    ) -> list[Candidate]:
+        """
+        Validate candidates against constraints.
+
+        May filter candidates (hard filtering) or enrich them with
+        validation features (soft validation).
+
+        Args:
+            candidates: List of candidates to validate
+            context: Match context with constraint parameters
+
+        Returns:
+            Validated/enriched candidates
+        """
+        ...

resolvekit/constraints/temporal_validator.py

@@ -0,0 +1,43 @@
+"""Temporal constraint validator."""
+
+from resolvekit.types import Candidate, MatchContext
+from resolvekit.utils.dates import is_valid_at
+
+
+class TemporalValidator:
+    """
+    Validates temporal validity of entities.
+
+    Soft validation only: Adds f_date_valid feature, does not filter.
+    Uses existing is_valid_at() utility from utils/dates.py.
+    """
+
+    def validate(
+        self, candidates: list[Candidate], context: MatchContext
+    ) -> list[Candidate]:
+        """
+        Validate temporal constraints.
+
+        Args:
+            candidates: List of candidates to validate
+            context: Match context with optional as_of date
+
+        Returns:
+            Candidates with f_date_valid feature (no filtering)
+        """
+        # If no as_of constraint, mark all as having no temporal constraint
+        if context.as_of is None:
+            for candidate in candidates:
+                candidate.features["date_valid"] = True
+            return candidates
+
+        # Check validity for each candidate
+        for candidate in candidates:
+            candidate.features["date_valid"] = is_valid_at(
+                as_of=context.as_of,
+                valid_from=candidate.entity.valid_from,
+                valid_until=candidate.entity.valid_until,
+            )
+
+        # Soft validation: return all candidates (don't filter)
+        return candidates

resolvekit/constraints/type_validator.py

@@ -0,0 +1,42 @@
+"""Type constraint validator."""
+
+from resolvekit.types import Candidate, MatchContext
+
+
+class TypeValidator:
+    """
+    Validates entity type constraints.
+
+    Hard filtering: Removes candidates that don't match context.entity_type
+    Feature: Adds f_type_valid boolean to all candidates
+    """
+
+    def validate(
+        self, candidates: list[Candidate], context: MatchContext
+    ) -> list[Candidate]:
+        """
+        Validate entity type constraints.
+
+        Args:
+            candidates: List of candidates to validate
+            context: Match context with optional entity_type constraint
+
+        Returns:
+            Filtered candidates with f_type_valid feature
+        """
+        # Add feature to all candidates
+        for candidate in candidates:
+            if context.entity_type is None:
+                candidate.features["type_valid"] = True
+            else:
+                candidate.features["type_valid"] = (
+                    candidate.entity.entity_type == context.entity_type
+                )
+
+        # Hard filter if type constraint is present
+        if context.entity_type is not None:
+            candidates = [
+                c for c in candidates if c.entity.entity_type == context.entity_type
+            ]
+
+        return candidates