97 lines
2.9 KiB
Python
97 lines
2.9 KiB
Python
"""
|
|
Deduplication logic for retain pipeline.
|
|
|
|
Checks for duplicate facts using semantic similarity and temporal proximity.
|
|
"""
|
|
import logging
|
|
from datetime import datetime
|
|
from typing import List
|
|
from collections import defaultdict
|
|
|
|
from .types import ProcessedFact
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
async def check_duplicates_batch(
|
|
conn,
|
|
bank_id: str,
|
|
facts: List[ProcessedFact],
|
|
duplicate_checker_fn
|
|
) -> List[bool]:
|
|
"""
|
|
Check which facts are duplicates using batched time-window queries.
|
|
|
|
Groups facts by 12-hour time buckets to efficiently check for duplicates
|
|
within a 24-hour window.
|
|
|
|
Args:
|
|
conn: Database connection
|
|
bank_id: Bank identifier
|
|
facts: List of ProcessedFact objects to check
|
|
duplicate_checker_fn: Async function(conn, bank_id, texts, embeddings, date, time_window_hours)
|
|
that returns List[bool] indicating duplicates
|
|
|
|
Returns:
|
|
List of boolean flags (same length as facts) indicating if each fact is a duplicate
|
|
"""
|
|
if not facts:
|
|
return []
|
|
|
|
# Group facts by event_date (rounded to 12-hour buckets) for efficient batching
|
|
time_buckets = defaultdict(list)
|
|
for idx, fact in enumerate(facts):
|
|
# Use occurred_start as the representative date
|
|
fact_date = fact.occurred_start
|
|
# Round to 12-hour bucket to group similar times
|
|
bucket_key = fact_date.replace(
|
|
hour=(fact_date.hour // 12) * 12,
|
|
minute=0,
|
|
second=0,
|
|
microsecond=0
|
|
)
|
|
time_buckets[bucket_key].append((idx, fact))
|
|
|
|
# Process each bucket in batch
|
|
all_is_duplicate = [False] * len(facts)
|
|
|
|
for bucket_date, bucket_items in time_buckets.items():
|
|
indices = [item[0] for item in bucket_items]
|
|
texts = [item[1].fact_text for item in bucket_items]
|
|
embeddings = [item[1].embedding for item in bucket_items]
|
|
|
|
# Check duplicates for this time bucket
|
|
dup_flags = await duplicate_checker_fn(
|
|
conn,
|
|
bank_id,
|
|
texts,
|
|
embeddings,
|
|
bucket_date,
|
|
time_window_hours=24
|
|
)
|
|
|
|
# Map results back to original indices
|
|
for idx, is_dup in zip(indices, dup_flags):
|
|
all_is_duplicate[idx] = is_dup
|
|
|
|
return all_is_duplicate
|
|
|
|
|
|
def filter_duplicates(
|
|
facts: List[ProcessedFact],
|
|
is_duplicate_flags: List[bool]
|
|
) -> List[ProcessedFact]:
|
|
"""
|
|
Filter out duplicate facts based on duplicate flags.
|
|
|
|
Args:
|
|
facts: List of ProcessedFact objects
|
|
is_duplicate_flags: Boolean flags indicating which facts are duplicates
|
|
|
|
Returns:
|
|
List of non-duplicate facts
|
|
"""
|
|
if len(facts) != len(is_duplicate_flags):
|
|
raise ValueError(f"Mismatch between facts ({len(facts)}) and flags ({len(is_duplicate_flags)})")
|
|
|
|
return [fact for fact, is_dup in zip(facts, is_duplicate_flags) if not is_dup]
|