"""Correlation graph builder — deterministic relationship graph. Builds a CorrelationGraph (nodes + edges) from evidence collected across multiple providers. No AI — only deterministic matching. """ from __future__ import annotations from dataclasses import dataclass, field from enum import Enum from typing import List from models.reports import ( CorrelationGraph, CorrelationNode, CorrelationEdge, ) class EdgeType(str, Enum): SAME_FACE = "same_face" SAME_OBJECT = "same_object" SAME_LOCATION = "same_location" SAME_CAMERA = "same_camera" SAME_HASH = "same_hash" SAME_EMBEDDING = "same_embedding" SAME_METADATA = "same_metadata" SAME_TIMESTAMP = "same_timestamp" @dataclass class Node: id: str node_type: str # "face" | "object" | "location" | "metadata" | "image" | "embedding" label: str = "" properties: dict = field(default_factory=dict) @dataclass class Edge: source: str target: str edge_type: EdgeType confidence: float = 1.0 evidence: str = "" class CorrelationGraphBuilder: """Builds a CorrelationGraph from nodes + edges.""" def __init__(self) -> None: self._nodes: list[Node] = [] self._edges: list[Edge] = [] def add_node(self, node: Node) -> None: self._nodes.append(node) def add_nodes(self, nodes: List[Node]) -> None: self._nodes.extend(nodes) def add_edge(self, edge: Edge) -> None: self._edges.append(edge) def add_edges(self, edges: List[Edge]) -> None: self._edges.extend(edges) def add_matches( self, edge_type: EdgeType, matches: List[tuple], ) -> None: """Add matched pairs as edges. Args: edge_type: the type of relationship. matches: list of (source_id, target_id, confidence, evidence_str) tuples. """ for source, target, confidence, evidence in matches: self._edges.append(Edge( source=source, target=target, edge_type=edge_type, confidence=confidence, evidence=evidence, )) def build(self, elapsed_ms: float = 0.0) -> CorrelationGraph: """Build the final CorrelationGraph.""" # Deduplicate nodes by ID seen_ids: set[str] = set() unique_nodes: list[Node] = [] for n in self._nodes: if n.id not in seen_ids: seen_ids.add(n.id) unique_nodes.append(n) # Deduplicate edges (same source+target+type) seen_edges: set[tuple[str, str, str]] = set() unique_edges: list[Edge] = [] for e in self._edges: key = (e.source, e.target, e.edge_type.value) reverse_key = (e.target, e.source, e.edge_type.value) if key not in seen_edges and reverse_key not in seen_edges: seen_edges.add(key) unique_edges.append(e) return CorrelationGraph( nodes=[ CorrelationNode( id=n.id, node_type=n.node_type, label=n.label, properties=n.properties, ) for n in unique_nodes ], edges=[ CorrelationEdge( source=e.source, target=e.target, edge_type=e.edge_type.value, confidence=e.confidence, evidence=e.evidence, ) for e in unique_edges ], num_nodes=len(unique_nodes), num_edges=len(unique_edges), elapsed_ms=round(elapsed_ms, 3), )