713 lines
29 KiB
Python
713 lines
29 KiB
Python
"""DuplicateService — explainable clusters and reversible decisions.
|
|
|
|
Detection runs in two categories:
|
|
|
|
- **Content-identical** clusters join assets that share exact bytes (``sha256``) or
|
|
normalized pixels (``pixel_sha256``). These are high confidence and are decided
|
|
automatically (a recommended canonical, the rest variants) — unless a member pair
|
|
carries a ``not_duplicate`` negative link, which is a contradiction that reopens
|
|
the cluster for review instead of silently overriding the user.
|
|
- **Perceptual** clusters join assets whose perceptual hashes are within a distance
|
|
band (NEAR/SIMILAR). These are review candidates: never decided automatically, and
|
|
negative-linked pairs are suppressed so a rejected pair is not re-suggested.
|
|
|
|
Archived assets stay in both indexes (US06-03): a new active copy of an archived
|
|
original is recognised through its hashes even while the medium is offline, and
|
|
cluster review falls back to the retained protected preview plus hash evidence.
|
|
An exact/pixel match links straight to the archived canonical; a perceptual match
|
|
is a review candidate that names the medium to mount for a pixel-level decision.
|
|
|
|
Decisions (``canonical`` / ``not_duplicate`` / ``deferred``) persist with evidence,
|
|
use optimistic version checks, are reversible, and can never form a canonical cycle.
|
|
A new content-identical member of an already-decided cluster inherits the established
|
|
canonical relationship rather than reopening it.
|
|
|
|
Extracts photo_analyzer.cluster_duplicates/mark_duplicates (union-find over the
|
|
Hamming graph); path-keyed identity and automatic largest-file selection are
|
|
replaced by stable asset IDs and reviewable decisions (donor_ledger.yaml: pa-dupes).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import uuid
|
|
from collections import Counter
|
|
from dataclasses import dataclass, field
|
|
from datetime import datetime, timezone
|
|
from enum import Enum
|
|
|
|
from sqlalchemy import func, select
|
|
from sqlalchemy.orm import sessionmaker
|
|
|
|
from photo_pipeline.models import (
|
|
ArchiveLocation,
|
|
Asset,
|
|
DuplicateCluster,
|
|
DuplicateMember,
|
|
DuplicateNegativeLink,
|
|
Thumbnail,
|
|
)
|
|
from photo_pipeline.services import availability, hashing
|
|
|
|
NEAR_MAX = 5
|
|
SIMILAR_MAX = 10
|
|
|
|
|
|
class Method(str, Enum):
|
|
EXACT = "exact"
|
|
PIXEL = "pixel"
|
|
PERCEPTUAL = "perceptual"
|
|
|
|
|
|
class Confidence(str, Enum):
|
|
EXACT = "exact"
|
|
PIXEL = "pixel"
|
|
NEAR = "near"
|
|
SIMILAR = "similar"
|
|
|
|
|
|
class ClusterState(str, Enum):
|
|
OPEN = "open"
|
|
DECIDED = "decided"
|
|
DISMISSED = "dismissed"
|
|
REOPENED = "reopened"
|
|
DEFERRED = "deferred"
|
|
|
|
|
|
class Decision(str, Enum):
|
|
CANONICAL = "canonical"
|
|
NOT_DUPLICATE = "not_duplicate"
|
|
DEFERRED = "deferred"
|
|
|
|
|
|
class Role(str, Enum):
|
|
MEMBER = "member"
|
|
CANONICAL = "canonical"
|
|
VARIANT = "variant"
|
|
|
|
|
|
class DuplicateError(RuntimeError):
|
|
"""Invalid decision (unknown cluster, non-member canonical, would cycle)."""
|
|
|
|
|
|
class ConflictError(RuntimeError):
|
|
"""Optimistic version check failed; the cluster changed since it was read."""
|
|
|
|
|
|
@dataclass
|
|
class DetectionReport:
|
|
clusters: list[dict] = field(default_factory=list)
|
|
|
|
@property
|
|
def counts(self) -> dict[str, int]:
|
|
return dict(Counter(c["state"] for c in self.clusters))
|
|
|
|
|
|
def _pair(a: str, b: str) -> tuple[str, str]:
|
|
return (a, b) if a <= b else (b, a)
|
|
|
|
|
|
class _UnionFind:
|
|
def __init__(self, items):
|
|
self.parent = {x: x for x in items}
|
|
|
|
def find(self, x):
|
|
while self.parent[x] != x:
|
|
self.parent[x] = self.parent[self.parent[x]]
|
|
x = self.parent[x]
|
|
return x
|
|
|
|
def union(self, a, b):
|
|
ra, rb = self.find(a), self.find(b)
|
|
if ra != rb:
|
|
self.parent[rb] = ra
|
|
|
|
def groups(self) -> list[list[str]]:
|
|
out: dict[str, list[str]] = {}
|
|
for x in self.parent:
|
|
out.setdefault(self.find(x), []).append(x)
|
|
return [sorted(g) for g in out.values()]
|
|
|
|
|
|
class DuplicateService:
|
|
def __init__(self, session_factory: sessionmaker) -> None:
|
|
self._session_factory = session_factory
|
|
|
|
# ── perceptual hash backfill ───────────────────────────────────────────
|
|
def ensure_phashes(self) -> int:
|
|
"""Hash whatever is readable now — an archived asset keeps the hash it
|
|
already has, and gains one whenever its medium happens to be mounted."""
|
|
updated = 0
|
|
with self._session_factory() as session:
|
|
assets = session.execute(select(Asset)).scalars()
|
|
for asset in assets:
|
|
if asset.phash is not None and asset.phash_version == hashing.PHASH_VERSION:
|
|
continue
|
|
source = availability.readable_path(session, asset)
|
|
if source is None:
|
|
continue
|
|
value = hashing.safe_phash(str(source))
|
|
if value is not None:
|
|
asset.phash = value
|
|
asset.phash_version = hashing.PHASH_VERSION
|
|
updated += 1
|
|
session.commit()
|
|
return updated
|
|
|
|
def ensure_phash(self, asset_id: str, *, source=None) -> str | None:
|
|
"""Backfill one asset's perceptual hash while its bytes are still readable.
|
|
|
|
Archiving calls this before the original leaves — passing the archive copy
|
|
as ``source``, since the database does not point at it yet — because an
|
|
asset without a pHash silently drops out of the fuzzy index the moment its
|
|
medium is away.
|
|
"""
|
|
with self._session_factory() as session:
|
|
asset = session.get(Asset, asset_id)
|
|
if asset is None:
|
|
return None
|
|
if asset.phash is not None and asset.phash_version == hashing.PHASH_VERSION:
|
|
return asset.phash
|
|
source = source or availability.readable_path(session, asset)
|
|
if source is None:
|
|
return None
|
|
value = hashing.safe_phash(str(source))
|
|
if value is not None:
|
|
asset.phash = value
|
|
asset.phash_version = hashing.PHASH_VERSION
|
|
session.commit()
|
|
return value
|
|
|
|
# ── detection ──────────────────────────────────────────────────────────
|
|
def detect(self) -> DetectionReport:
|
|
self.ensure_phashes()
|
|
now = datetime.now(timezone.utc)
|
|
report = DetectionReport()
|
|
with self._session_factory() as session:
|
|
# Every known asset stays in the indexes, archived or not: a copy of an
|
|
# archived original must be recognised as a duplicate rather than
|
|
# treated as a new photo (concept §9, invariant 12).
|
|
assets = list(session.execute(select(Asset)).scalars())
|
|
by_id = {a.id: a for a in assets}
|
|
negatives = {
|
|
_pair(link.asset_a, link.asset_b)
|
|
for link in session.execute(select(DuplicateNegativeLink)).scalars()
|
|
}
|
|
clusters = list(session.execute(select(DuplicateCluster)).scalars())
|
|
members = {
|
|
c.id: list(
|
|
session.execute(
|
|
select(DuplicateMember).where(DuplicateMember.cluster_id == c.id)
|
|
).scalars()
|
|
)
|
|
for c in clusters
|
|
}
|
|
|
|
content_groups, content_pairs = self._content_groups(assets)
|
|
perceptual_groups = self._perceptual_groups(assets, content_pairs, negatives)
|
|
|
|
touched: list[DuplicateCluster] = []
|
|
for method, ids, distances in content_groups:
|
|
touched.append(
|
|
self._sync_content_cluster(
|
|
session, method, ids, by_id, negatives, clusters, members, now
|
|
)
|
|
)
|
|
for ids, distances in perceptual_groups:
|
|
touched.append(
|
|
self._sync_perceptual_cluster(
|
|
session, ids, distances, by_id, clusters, members, now
|
|
)
|
|
)
|
|
session.commit()
|
|
report.clusters = [self._snapshot(session, c.id) for c in touched]
|
|
return report
|
|
|
|
def _content_groups(self, assets):
|
|
"""Union-find over exact-byte and same-pixel edges. Returns groups of >=2
|
|
with method EXACT (all same bytes) or PIXEL, plus the set of joined pairs."""
|
|
uf = _UnionFind([a.id for a in assets])
|
|
pairs: set[tuple[str, str]] = set()
|
|
for key in ("current_sha256", "pixel_sha256"):
|
|
buckets: dict[str, list[str]] = {}
|
|
for a in assets:
|
|
value = getattr(a, key)
|
|
if value:
|
|
buckets.setdefault(value, []).append(a.id)
|
|
for group in buckets.values():
|
|
for other in group[1:]:
|
|
uf.union(group[0], other)
|
|
for i in range(len(group)):
|
|
for j in range(i + 1, len(group)):
|
|
pairs.add(_pair(group[i], group[j]))
|
|
by_id = {a.id: a for a in assets}
|
|
result = []
|
|
for group in uf.groups():
|
|
if len(group) < 2:
|
|
continue
|
|
shas = {by_id[i].current_sha256 for i in group}
|
|
method = Method.EXACT if len(shas) == 1 and None not in shas else Method.PIXEL
|
|
result.append((method, sorted(group), {}))
|
|
return result, pairs
|
|
|
|
def _perceptual_groups(self, assets, content_pairs, negatives):
|
|
"""Union-find over phash distance edges (<= SIMILAR_MAX), excluding pairs
|
|
that are already content-identical or negatively linked."""
|
|
hashed = [a for a in assets if a.phash]
|
|
uf = _UnionFind([a.id for a in hashed])
|
|
distances: dict[tuple[str, str], int] = {}
|
|
for i in range(len(hashed)):
|
|
for j in range(i + 1, len(hashed)):
|
|
a, b = hashed[i], hashed[j]
|
|
pair = _pair(a.id, b.id)
|
|
if pair in content_pairs or pair in negatives:
|
|
continue
|
|
dist = hashing.phash_distance(a.phash, b.phash)
|
|
if dist <= SIMILAR_MAX:
|
|
uf.union(a.id, b.id)
|
|
distances[pair] = dist
|
|
result = []
|
|
for group in uf.groups():
|
|
if len(group) < 2:
|
|
continue
|
|
group_dists = {
|
|
_pair(x, y): distances[_pair(x, y)]
|
|
for x in group
|
|
for y in group
|
|
if x < y and _pair(x, y) in distances
|
|
}
|
|
result.append((sorted(group), group_dists))
|
|
return result
|
|
|
|
# ── cluster upserts ──────────────────────────────────────────────────────
|
|
def _match_existing(self, ids, clusters, members, categories):
|
|
best = None
|
|
best_overlap = 0
|
|
for cluster in clusters:
|
|
if cluster.method not in categories:
|
|
continue
|
|
overlap = len(set(ids) & {m.asset_id for m in members[cluster.id]})
|
|
if overlap > best_overlap:
|
|
best, best_overlap = cluster, overlap
|
|
return best
|
|
|
|
def _sync_content_cluster(
|
|
self, session, method, ids, by_id, negatives, clusters, members, now
|
|
):
|
|
contradiction = any(
|
|
_pair(ids[i], ids[j]) in negatives
|
|
for i in range(len(ids))
|
|
for j in range(i + 1, len(ids))
|
|
)
|
|
confidence = Confidence.EXACT if method == Method.EXACT else Confidence.PIXEL
|
|
# Also consider a prior perceptual cluster over the same assets: if a member
|
|
# pair became byte-identical after a not_duplicate decision, reopen that same
|
|
# reviewed cluster rather than spawning a parallel one.
|
|
existing = self._match_existing(
|
|
ids,
|
|
clusters,
|
|
members,
|
|
{Method.EXACT.value, Method.PIXEL.value, Method.PERCEPTUAL.value},
|
|
)
|
|
|
|
if existing is None:
|
|
cluster = DuplicateCluster(
|
|
id=str(uuid.uuid4()),
|
|
method=method.value,
|
|
confidence=confidence.value,
|
|
state=ClusterState.OPEN.value,
|
|
version=1,
|
|
)
|
|
session.add(cluster)
|
|
session.flush()
|
|
self._replace_members(session, cluster, ids, by_id, distances={})
|
|
if contradiction:
|
|
cluster.state = ClusterState.REOPENED.value
|
|
else:
|
|
self._auto_decide(session, cluster, ids, by_id, now)
|
|
clusters.append(cluster)
|
|
members[cluster.id] = list(
|
|
session.execute(
|
|
select(DuplicateMember).where(DuplicateMember.cluster_id == cluster.id)
|
|
).scalars()
|
|
)
|
|
return cluster
|
|
|
|
cluster = existing
|
|
cluster.method = method.value
|
|
cluster.confidence = confidence.value
|
|
prior_members = {m.asset_id for m in members[cluster.id]}
|
|
added = set(ids) - prior_members
|
|
self._replace_members(session, cluster, ids, by_id, distances={})
|
|
|
|
if contradiction:
|
|
cluster.state = ClusterState.REOPENED.value
|
|
cluster.decision = None
|
|
cluster.canonical_asset_id = None
|
|
self._clear_asset_links(session, ids)
|
|
elif (
|
|
cluster.decision == Decision.CANONICAL.value
|
|
and cluster.canonical_asset_id in ids
|
|
):
|
|
# Consistent growth: new members inherit the established canonical.
|
|
self._apply_canonical(session, cluster, ids, cluster.canonical_asset_id)
|
|
cluster.state = ClusterState.DECIDED.value
|
|
elif cluster.state in (ClusterState.DISMISSED.value,) and added:
|
|
cluster.state = ClusterState.REOPENED.value
|
|
elif cluster.state in (ClusterState.OPEN.value, ClusterState.REOPENED.value):
|
|
self._auto_decide(session, cluster, ids, by_id, now)
|
|
cluster.version += 1
|
|
cluster.updated_at = now
|
|
members[cluster.id] = list(
|
|
session.execute(
|
|
select(DuplicateMember).where(DuplicateMember.cluster_id == cluster.id)
|
|
).scalars()
|
|
)
|
|
return cluster
|
|
|
|
def _sync_perceptual_cluster(self, session, ids, distances, by_id, clusters, members, now):
|
|
confidence = (
|
|
Confidence.NEAR
|
|
if distances and max(distances.values()) <= NEAR_MAX
|
|
else Confidence.SIMILAR
|
|
)
|
|
existing = self._match_existing(ids, clusters, members, {Method.PERCEPTUAL.value})
|
|
if existing is None:
|
|
cluster = DuplicateCluster(
|
|
id=str(uuid.uuid4()),
|
|
method=Method.PERCEPTUAL.value,
|
|
confidence=confidence.value,
|
|
state=ClusterState.OPEN.value,
|
|
version=1,
|
|
)
|
|
session.add(cluster)
|
|
session.flush()
|
|
self._replace_members(session, cluster, ids, by_id, distances)
|
|
clusters.append(cluster)
|
|
members[cluster.id] = list(
|
|
session.execute(
|
|
select(DuplicateMember).where(DuplicateMember.cluster_id == cluster.id)
|
|
).scalars()
|
|
)
|
|
return cluster
|
|
|
|
cluster = existing
|
|
prior = {m.asset_id for m in members[cluster.id]}
|
|
added = set(ids) - prior
|
|
cluster.confidence = confidence.value
|
|
self._replace_members(session, cluster, ids, by_id, distances)
|
|
if added and cluster.state in (
|
|
ClusterState.DECIDED.value,
|
|
ClusterState.DISMISSED.value,
|
|
ClusterState.DEFERRED.value,
|
|
):
|
|
cluster.state = ClusterState.REOPENED.value
|
|
cluster.version += 1
|
|
cluster.updated_at = now
|
|
members[cluster.id] = list(
|
|
session.execute(
|
|
select(DuplicateMember).where(DuplicateMember.cluster_id == cluster.id)
|
|
).scalars()
|
|
)
|
|
return cluster
|
|
|
|
def _replace_members(self, session, cluster, ids, by_id, distances):
|
|
current = {
|
|
m.asset_id: m
|
|
for m in session.execute(
|
|
select(DuplicateMember).where(DuplicateMember.cluster_id == cluster.id)
|
|
).scalars()
|
|
}
|
|
wanted = set(ids)
|
|
for asset_id, member in current.items():
|
|
if asset_id not in wanted:
|
|
session.delete(member)
|
|
for asset_id in ids:
|
|
asset = by_id[asset_id]
|
|
evidence = json.dumps(
|
|
{
|
|
"sha256": asset.current_sha256,
|
|
"pixel_sha256": asset.pixel_sha256,
|
|
"phash": asset.phash,
|
|
}
|
|
)
|
|
distance = next(
|
|
(d for p, d in distances.items() if asset_id in p), None
|
|
)
|
|
if asset_id in current:
|
|
current[asset_id].evidence = evidence
|
|
current[asset_id].distance = distance
|
|
else:
|
|
session.add(
|
|
DuplicateMember(
|
|
cluster_id=cluster.id,
|
|
asset_id=asset_id,
|
|
role=Role.MEMBER.value,
|
|
distance=distance,
|
|
evidence=evidence,
|
|
)
|
|
)
|
|
session.flush()
|
|
|
|
def _auto_decide(self, session, cluster, ids, by_id, now):
|
|
canonical = self._recommend_canonical(ids, by_id)
|
|
self._apply_canonical(session, cluster, ids, canonical)
|
|
cluster.decision = Decision.CANONICAL.value
|
|
cluster.canonical_asset_id = canonical
|
|
cluster.state = ClusterState.DECIDED.value
|
|
cluster.decided_at = now
|
|
|
|
@staticmethod
|
|
def _recommend_canonical(ids, by_id) -> str:
|
|
# ponytail: largest file, then the archived copy, then path as a
|
|
# deterministic tie-break. Archived wins ties because it is the reviewed,
|
|
# uploaded original — a fresh active copy must not demote it to a variant.
|
|
# The concept's richer policy (resolution, least recompression, metadata
|
|
# richness) lands with the review UI story.
|
|
return max(
|
|
ids,
|
|
key=lambda i: (
|
|
by_id[i].byte_size or 0,
|
|
by_id[i].availability_state in availability.ARCHIVED,
|
|
by_id[i].current_path or by_id[i].archive_path or "",
|
|
),
|
|
)
|
|
|
|
def _apply_canonical(self, session, cluster, ids, canonical_id):
|
|
for member in session.execute(
|
|
select(DuplicateMember).where(DuplicateMember.cluster_id == cluster.id)
|
|
).scalars():
|
|
asset = session.get(Asset, member.asset_id)
|
|
if member.asset_id == canonical_id:
|
|
member.role = Role.CANONICAL.value
|
|
asset.canonical_asset_id = None
|
|
else:
|
|
member.role = Role.VARIANT.value
|
|
asset.canonical_asset_id = canonical_id
|
|
self._assert_acyclic(session, ids)
|
|
|
|
@staticmethod
|
|
def _clear_asset_links(session, ids):
|
|
for asset_id in ids:
|
|
session.get(Asset, asset_id).canonical_asset_id = None
|
|
|
|
@staticmethod
|
|
def _assert_acyclic(session, ids):
|
|
for start in ids:
|
|
seen = set()
|
|
cur = start
|
|
while cur is not None:
|
|
if cur in seen:
|
|
raise DuplicateError(f"canonical link cycle at {cur}")
|
|
seen.add(cur)
|
|
asset = session.get(Asset, cur)
|
|
cur = asset.canonical_asset_id if asset else None
|
|
|
|
# ── reads for the review UI ────────────────────────────────────────────────
|
|
def list_clusters(
|
|
self, *, state: str | None = None, limit: int = 50, offset: int = 0
|
|
) -> dict:
|
|
limit = max(1, min(limit, 200))
|
|
offset = max(0, offset)
|
|
with self._session_factory() as session:
|
|
stmt = select(DuplicateCluster)
|
|
if state:
|
|
stmt = stmt.where(DuplicateCluster.state == state)
|
|
total = session.scalar(select(func.count()).select_from(stmt.subquery()))
|
|
rows = session.execute(
|
|
stmt.order_by(DuplicateCluster.created_at).limit(limit).offset(offset)
|
|
).scalars()
|
|
items = [self._snapshot(session, c.id) for c in rows]
|
|
return {"items": items, "total": int(total or 0), "limit": limit, "offset": offset}
|
|
|
|
def get_cluster(self, cluster_id: str) -> dict | None:
|
|
"""Cluster detail enriched with per-member asset evidence for comparison."""
|
|
with self._session_factory() as session:
|
|
cluster = session.get(DuplicateCluster, cluster_id)
|
|
if cluster is None:
|
|
return None
|
|
members = []
|
|
for member in session.execute(
|
|
select(DuplicateMember).where(DuplicateMember.cluster_id == cluster_id)
|
|
).scalars():
|
|
asset = session.get(Asset, member.asset_id)
|
|
try:
|
|
evidence = json.loads(member.evidence) if member.evidence else {}
|
|
except json.JSONDecodeError:
|
|
evidence = {}
|
|
members.append(
|
|
{
|
|
"asset_id": member.asset_id,
|
|
"role": member.role,
|
|
"distance": member.distance,
|
|
"evidence": evidence,
|
|
"current_path": asset.current_path if asset else None,
|
|
"byte_size": asset.byte_size if asset else None,
|
|
"phash": asset.phash if asset else None,
|
|
**self._offline_evidence(session, asset),
|
|
}
|
|
)
|
|
members.sort(key=lambda m: m["asset_id"])
|
|
# A full-resolution comparison of an offline original is impossible; the
|
|
# UI asks for that named medium instead of guessing (concept §9).
|
|
mount_required = sorted(
|
|
{m["archive_location"] for m in members if m["requires_mount"]}
|
|
)
|
|
return {
|
|
"id": cluster.id,
|
|
"method": cluster.method,
|
|
"confidence": cluster.confidence,
|
|
"state": cluster.state,
|
|
"decision": cluster.decision,
|
|
"canonical_asset_id": cluster.canonical_asset_id,
|
|
"version": cluster.version,
|
|
"requires_confirmation": cluster.method == Method.PERCEPTUAL.value,
|
|
"mount_required": mount_required,
|
|
"members": members,
|
|
}
|
|
|
|
def _offline_evidence(self, session, asset: Asset | None) -> dict:
|
|
"""What review can still rely on when a member's original is not readable."""
|
|
if asset is None:
|
|
return {
|
|
"availability_state": None,
|
|
"archive_location": None,
|
|
"archive_location_id": None,
|
|
"archive_path": None,
|
|
"preview": {"state": "missing", "protected": False},
|
|
"requires_mount": False,
|
|
}
|
|
location = (
|
|
session.get(ArchiveLocation, asset.archive_location_id)
|
|
if asset.archive_location_id
|
|
else None
|
|
)
|
|
preview = self._preview_evidence(session, asset.id)
|
|
archived = asset.availability_state in availability.ARCHIVED
|
|
return {
|
|
"availability_state": asset.availability_state,
|
|
"archive_location": location.name if location else None,
|
|
"archive_location_id": asset.archive_location_id,
|
|
"archive_path": asset.archive_path,
|
|
"preview": preview,
|
|
# Offline archived members can still be compared through their retained
|
|
# preview and hash evidence; only pixel-level review needs the medium.
|
|
"requires_mount": archived
|
|
and asset.availability_state == availability.ARCHIVED_OFFLINE
|
|
and bool(location),
|
|
}
|
|
|
|
@staticmethod
|
|
def _preview_evidence(session, asset_id: str) -> dict:
|
|
rows = list(
|
|
session.execute(select(Thumbnail).where(Thumbnail.asset_id == asset_id)).scalars()
|
|
)
|
|
ready = [r for r in rows if r.state == "ready" and r.path]
|
|
if ready:
|
|
best = max(ready, key=lambda r: (bool(r.protected), r.size or 0))
|
|
return {"state": "ready", "protected": bool(best.protected), "size": best.size}
|
|
if rows:
|
|
return {"state": "unsupported", "protected": False, "size": rows[0].size}
|
|
return {"state": "missing", "protected": False, "size": None}
|
|
|
|
# ── decisions ────────────────────────────────────────────────────────────
|
|
def decide(
|
|
self,
|
|
cluster_id: str,
|
|
decision: Decision | str,
|
|
*,
|
|
canonical_asset_id: str | None = None,
|
|
expected_version: int | None = None,
|
|
reason: str | None = None,
|
|
) -> dict:
|
|
decision = Decision(decision)
|
|
now = datetime.now(timezone.utc)
|
|
with self._session_factory() as session:
|
|
cluster = session.get(DuplicateCluster, cluster_id)
|
|
if cluster is None:
|
|
raise DuplicateError(f"unknown cluster {cluster_id}")
|
|
if expected_version is not None and cluster.version != expected_version:
|
|
raise ConflictError(
|
|
f"cluster {cluster_id} is at version {cluster.version}, "
|
|
f"expected {expected_version}"
|
|
)
|
|
member_ids = [
|
|
m.asset_id
|
|
for m in session.execute(
|
|
select(DuplicateMember).where(DuplicateMember.cluster_id == cluster_id)
|
|
).scalars()
|
|
]
|
|
self._reset_members(session, cluster, member_ids)
|
|
self._remove_negative_links(session, member_ids)
|
|
|
|
if decision == Decision.CANONICAL:
|
|
if canonical_asset_id not in member_ids:
|
|
raise DuplicateError("canonical_asset_id must be a cluster member")
|
|
self._apply_canonical(session, cluster, member_ids, canonical_asset_id)
|
|
cluster.decision = Decision.CANONICAL.value
|
|
cluster.canonical_asset_id = canonical_asset_id
|
|
cluster.state = ClusterState.DECIDED.value
|
|
elif decision == Decision.NOT_DUPLICATE:
|
|
self._add_negative_links(session, member_ids, reason)
|
|
cluster.decision = Decision.NOT_DUPLICATE.value
|
|
cluster.canonical_asset_id = None
|
|
cluster.state = ClusterState.DISMISSED.value
|
|
else: # DEFERRED
|
|
cluster.decision = Decision.DEFERRED.value
|
|
cluster.canonical_asset_id = None
|
|
cluster.state = ClusterState.DEFERRED.value
|
|
|
|
cluster.decided_at = now
|
|
cluster.updated_at = now
|
|
cluster.version += 1
|
|
session.commit()
|
|
return self._snapshot(session, cluster_id)
|
|
|
|
def _reset_members(self, session, cluster, member_ids):
|
|
for member in session.execute(
|
|
select(DuplicateMember).where(DuplicateMember.cluster_id == cluster.id)
|
|
).scalars():
|
|
member.role = Role.MEMBER.value
|
|
for asset_id in member_ids:
|
|
session.get(Asset, asset_id).canonical_asset_id = None
|
|
|
|
@staticmethod
|
|
def _add_negative_links(session, member_ids, reason):
|
|
existing = {
|
|
_pair(link.asset_a, link.asset_b)
|
|
for link in session.execute(select(DuplicateNegativeLink)).scalars()
|
|
}
|
|
for i in range(len(member_ids)):
|
|
for j in range(i + 1, len(member_ids)):
|
|
a, b = _pair(member_ids[i], member_ids[j])
|
|
if (a, b) not in existing:
|
|
session.add(
|
|
DuplicateNegativeLink(asset_a=a, asset_b=b, reason=reason)
|
|
)
|
|
|
|
@staticmethod
|
|
def _remove_negative_links(session, member_ids):
|
|
member_set = set(member_ids)
|
|
for link in session.execute(select(DuplicateNegativeLink)).scalars():
|
|
if link.asset_a in member_set and link.asset_b in member_set:
|
|
session.delete(link)
|
|
|
|
def _snapshot(self, session, cluster_id) -> dict:
|
|
cluster = session.get(DuplicateCluster, cluster_id)
|
|
members = session.execute(
|
|
select(DuplicateMember).where(DuplicateMember.cluster_id == cluster_id)
|
|
).scalars()
|
|
return {
|
|
"id": cluster.id,
|
|
"method": cluster.method,
|
|
"confidence": cluster.confidence,
|
|
"state": cluster.state,
|
|
"decision": cluster.decision,
|
|
"canonical_asset_id": cluster.canonical_asset_id,
|
|
"version": cluster.version,
|
|
"members": sorted(
|
|
({"asset_id": m.asset_id, "role": m.role, "distance": m.distance} for m in members),
|
|
key=lambda m: m["asset_id"],
|
|
),
|
|
}
|