57f154454d
* fix(recall): cap entity fanout in graph expansion to prevent slow queries On large banks, the entity co-occurrence self-join in _expand_combined() produces massive intermediate row counts when seeds reference high-fanout entities (e.g. an entity with 25K+ mentions). This causes recall latency to degrade significantly. Changes: - Replace unbounded entity self-join with LATERAL per-entity cap (graph_per_entity_limit, default 200), reducing intermediate rows from potentially millions to at most num_entities * 200 - Add ORDER BY unit_id DESC in LATERAL subquery for deterministic recency-biased sampling (rides the PK index, no extra sort) - Add timeout fallback (graph_expansion_timeout, default 10s) that drops entity expansion and falls back to semantic+causal only - Add composite index (entity_id, unit_id) on unit_entities for index-only scans in the LATERAL subquery - Merge 3 unmerged migration heads into one - Fix recall_perf.py dotenv override issue Unlike the approach in #895, this does NOT filter out hub entities entirely — all entities are kept but capped equally, preserving retrieval quality for queries about frequently-mentioned entities. Benchmarked on a 67K-unit bank (top entity = 25K mentions): - retrieval_graph: 0.337s → 0.055s (84% faster) - end-to-end recall: 0.912s → 0.519s (43% faster) * fix(tests): fix broken test_combined_scoring and test_reranking_proof_count - test_combined_scoring: replace MagicMock(spec=RetrievalResult) with real dataclass instances — MagicMock attributes returned nested mocks that failed on >= comparisons with int - test_reranking_proof_count: remove deleted `embedding` param from RetrievalResult constructor, use None for occurred_start/end to get neutral recency (datetime.now gave recency=1.0 which boosted scores) * refactor: rename config to link_expansion_ prefix, fix observation fanout - Rename GRAPH_PER_ENTITY_LIMIT → LINK_EXPANSION_PER_ENTITY_LIMIT and GRAPH_EXPANSION_TIMEOUT → LINK_EXPANSION_TIMEOUT to follow the convention that these are specific to the link_expansion graph retriever - Apply the same LATERAL per-entity cap to _expand_observations(), which had the same unbounded self-join through unit_entities * style: fix formatting in config.py
91 lines
3.5 KiB
Python
91 lines
3.5 KiB
Python
"""
|
|
Unit tests for proof_count boost in reranking.
|
|
"""
|
|
|
|
from datetime import datetime, timezone
|
|
import pytest
|
|
from uuid import uuid4
|
|
|
|
from hindsight_api.engine.search.types import RetrievalResult, MergedCandidate, ScoredResult
|
|
from hindsight_api.engine.search.reranking import apply_combined_scoring
|
|
|
|
UTC = timezone.utc
|
|
|
|
def create_mock_scored_result(proof_count: int | None = None, ce_score: float = 0.8) -> ScoredResult:
|
|
"""Helper to create a minimal ScoredResult suitable for scoring tests."""
|
|
retrieval = RetrievalResult(
|
|
id=str(uuid4()),
|
|
text="Test mock fact",
|
|
fact_type="observation" if proof_count is not None else "world",
|
|
document_id=str(uuid4()),
|
|
chunk_id=str(uuid4()),
|
|
proof_count=proof_count,
|
|
# Use None for neutral recency so only proof_count changes score
|
|
occurred_start=None,
|
|
occurred_end=None
|
|
)
|
|
candidate = MergedCandidate(
|
|
retrieval=retrieval,
|
|
rrf_score=0.1,
|
|
)
|
|
return ScoredResult(
|
|
candidate=candidate,
|
|
cross_encoder_score=ce_score,
|
|
cross_encoder_score_normalized=ce_score,
|
|
weight=ce_score,
|
|
)
|
|
|
|
def test_proof_count_neutral_when_none():
|
|
"""Test that when proof_count is None (e.g. non-observation), it gets neutral 0.5 norm."""
|
|
sr = create_mock_scored_result(proof_count=None, ce_score=0.8)
|
|
now = datetime.now(UTC)
|
|
|
|
apply_combined_scoring([sr], now, proof_count_alpha=0.1)
|
|
|
|
# Neutral multiplier means score shouldn't be boosted by proof_count
|
|
# Since recency is neutral (just created) and temporal is neutral, score should remain unchanged
|
|
assert sr.combined_score == pytest.approx(0.8, rel=1e-3)
|
|
|
|
def test_proof_count_neutral_at_one():
|
|
"""Test that proof_count=1 gives neutral multiplier."""
|
|
sr = create_mock_scored_result(proof_count=1, ce_score=0.8)
|
|
now = datetime.now(UTC)
|
|
|
|
apply_combined_scoring([sr], now, proof_count_alpha=0.1)
|
|
|
|
# proof_count=1 -> math.log(1) = 0 -> 0.5 + 0/10 = 0.5 (neutral) -> multiplier 1.0
|
|
assert sr.combined_score == pytest.approx(0.8, rel=1e-3)
|
|
|
|
def test_proof_count_increases_with_higher_counts():
|
|
"""Test that higher proof counts yield strictly higher scores."""
|
|
now = datetime.now(UTC)
|
|
|
|
# Create results with increasing proof counts
|
|
sr_5 = create_mock_scored_result(proof_count=5, ce_score=0.8)
|
|
sr_50 = create_mock_scored_result(proof_count=50, ce_score=0.8)
|
|
sr_100 = create_mock_scored_result(proof_count=100, ce_score=0.8)
|
|
|
|
# Process them
|
|
apply_combined_scoring([sr_5, sr_50, sr_100], now, proof_count_alpha=0.1)
|
|
|
|
# Assure scores strictly increase
|
|
assert sr_5.combined_score > 0.8
|
|
assert sr_50.combined_score > sr_5.combined_score
|
|
assert sr_100.combined_score > sr_50.combined_score
|
|
|
|
def test_proof_count_no_hardcoded_cap_at_100():
|
|
"""Test that proof_count continues to scale within the clamped [0, 1] range."""
|
|
now = datetime.now(UTC)
|
|
|
|
# Use values that stay below the clamp ceiling (proof_norm < 1.0)
|
|
# log(5)/10=0.16, log(20)/10=0.30, log(100)/10=0.46 → all below 0.5 headroom
|
|
sr_5 = create_mock_scored_result(proof_count=5, ce_score=0.8)
|
|
sr_20 = create_mock_scored_result(proof_count=20, ce_score=0.8)
|
|
sr_100 = create_mock_scored_result(proof_count=100, ce_score=0.8)
|
|
|
|
apply_combined_scoring([sr_5, sr_20, sr_100], now, proof_count_alpha=0.1)
|
|
|
|
# Must strictly increase within the valid range
|
|
assert sr_20.combined_score > sr_5.combined_score
|
|
assert sr_100.combined_score > sr_20.combined_score
|
|
|