0pt/services/api/workers/entity_resolution.py

from __future__ import annotations

from dataclasses import dataclass, field
from enum import Enum
from typing import Literal
from urllib.parse import urlparse
import re
import unicodedata

from rapidfuzz.fuzz import ratio

EntityType = Literal["organization", "person"]


class ResolutionDecision(str, Enum):
    EXACT_LINK = "exact_link"
    REVIEW = "review"
    DISTINCT = "distinct"


@dataclass(frozen=True)
class EntityCandidate:
    entity_id: str
    entity_type: EntityType
    name: str
    source_ref: str
    registration_country: str | None = None
    registration_number: str | None = None
    email: str | None = None
    website: str | None = None
    organization_id: str | None = None
    city: str | None = None
    metadata: dict[str, object] = field(default_factory=dict)


@dataclass(frozen=True)
class ResolutionResult:
    left_id: str
    right_id: str
    decision: ResolutionDecision
    score: float
    reasons: tuple[str, ...]
    requires_human_review: bool
    algorithm_version: str = "entity-resolution-v1"


def _norm(value: str | None) -> str:
    if not value:
        return ""
    normalized = unicodedata.normalize("NFKD", value)
    ascii_text = "".join(char for char in normalized if not unicodedata.combining(char))
    return re.sub(r"[^a-z0-9]+", " ", ascii_text.lower()).strip()


def _norm_registration_number(value: str | None) -> str:
    return re.sub(r"[^0-9a-z]+", "", (value or "").lower())


def _norm_email(value: str | None) -> str:
    return (value or "").strip().lower()


def _domain(value: str | None) -> str:
    if not value:
        return ""
    raw = value.strip().lower()
    if "@" in raw and "://" not in raw:
        raw = raw.rsplit("@", 1)[-1]
    if "://" not in raw:
        raw = f"https://{raw}"
    host = (urlparse(raw).hostname or "").lower()
    return host.removeprefix("www.")


def _exact_organization(left: EntityCandidate, right: EntityCandidate) -> ResolutionResult | None:
    left_number = _norm_registration_number(left.registration_number)
    right_number = _norm_registration_number(right.registration_number)
    left_country = _norm(left.registration_country)
    right_country = _norm(right.registration_country)
    if left_number and right_number and left_country and right_country:
        if left_country == right_country and left_number == right_number:
            return ResolutionResult(
                left.entity_id,
                right.entity_id,
                ResolutionDecision.EXACT_LINK,
                1.0,
                ("same registration country and registration number",),
                False,
            )
        if left_country == right_country and left_number != right_number:
            return ResolutionResult(
                left.entity_id,
                right.entity_id,
                ResolutionDecision.DISTINCT,
                0.0,
                ("different registration numbers in the same country",),
                False,
            )
    return None


def _exact_person(left: EntityCandidate, right: EntityCandidate) -> ResolutionResult | None:
    left_email = _norm_email(left.email)
    right_email = _norm_email(right.email)
    if left_email and right_email:
        if left_email == right_email:
            return ResolutionResult(
                left.entity_id,
                right.entity_id,
                ResolutionDecision.EXACT_LINK,
                1.0,
                ("same normalized email address",),
                False,
            )
    return None


def resolve_entity_pair(left: EntityCandidate, right: EntityCandidate) -> ResolutionResult:
    if left.entity_type != right.entity_type:
        return ResolutionResult(
            left.entity_id,
            right.entity_id,
            ResolutionDecision.DISTINCT,
            0.0,
            ("different entity types",),
            False,
        )

    exact = _exact_organization(left, right) if left.entity_type == "organization" else _exact_person(left, right)
    if exact is not None:
        return exact

    left_name = _norm(left.name)
    right_name = _norm(right.name)
    if not left_name or not right_name:
        return ResolutionResult(
            left.entity_id,
            right.entity_id,
            ResolutionDecision.REVIEW,
            0.0,
            ("missing normalized name prevents deterministic comparison",),
            True,
        )

    name_similarity = ratio(left_name, right_name) / 100.0
    reasons = [f"name similarity {name_similarity:.3f}"]
    score = name_similarity * 0.65

    if left.entity_type == "organization":
        left_domain = _domain(left.website or left.email)
        right_domain = _domain(right.website or right.email)
        if left_domain and right_domain:
            if left_domain == right_domain:
                score += 0.25
                reasons.append("same normalized domain")
            else:
                score -= 0.1
                reasons.append("different normalized domains")
        if _norm(left.registration_country) and _norm(left.registration_country) == _norm(right.registration_country):
            score += 0.05
            reasons.append("same registration country")
        if _norm(left.city) and _norm(left.city) == _norm(right.city):
            score += 0.05
            reasons.append("same city")
    else:
        if left.organization_id and right.organization_id:
            if left.organization_id == right.organization_id:
                score += 0.25
                reasons.append("same known organization")
            else:
                score -= 0.1
                reasons.append("different known organizations")
        left_domain = _domain(left.email)
        right_domain = _domain(right.email)
        if left_domain and right_domain and left_domain == right_domain:
            score += 0.1
            reasons.append("same email domain")

    score = max(0.0, min(1.0, score))

    # Only strong identifiers auto-link. Similarity-based matches remain operator-review candidates.
    if score >= 0.72:
        decision = ResolutionDecision.REVIEW
        review = True
    elif score <= 0.3:
        decision = ResolutionDecision.DISTINCT
        review = False
    else:
        decision = ResolutionDecision.REVIEW
        review = True

    return ResolutionResult(
        left.entity_id,
        right.entity_id,
        decision,
        round(score, 6),
        tuple(reasons),
        review,
    )


def rank_resolution_candidates(
    query: EntityCandidate,
    candidates: list[EntityCandidate],
    *,
    limit: int = 20,
) -> list[ResolutionResult]:
    results = [resolve_entity_pair(query, candidate) for candidate in candidates if candidate.entity_id != query.entity_id]
    decision_order = {
        ResolutionDecision.EXACT_LINK: 0,
        ResolutionDecision.REVIEW: 1,
        ResolutionDecision.DISTINCT: 2,
    }
    return sorted(
        results,
        key=lambda result: (decision_order[result.decision], -result.score, result.right_id),
    )[: max(0, limit)]