`expand_observations`' scoring join is O(C + U) as a hash join and O(U x C) as a nested loop, where C is the connected-source set and U the unnested candidate source ids. PostgreSQL picks between them from its row estimate for the `connected_sources` CTE, and that estimate was 1 against an actual ~3,700: the capped column came out of a LATERAL + LIMIT subquery, which carries no n_distinct statistic, so DISTINCT over it was estimated at 2 and the NOT EXISTS anti-join took that to 1. A 1-row inner side makes the nested loop look free, so it won on cost and lost by four orders of magnitude at runtime — 15s and ~15M rejected join rows on a realistically-shaped bank, matching the plans reported in the issue. Rank with row_number() instead. Identical output — same cap, same ordering, unit_id is unique — but the capped column now traces to unit_entities.unit_id, so the estimate comes from real statistics (207-3,449 against 2,242-4,193 actual) and the nested loop is priced honestly. Measured over 12 seed sets on the fixture below: p50 15,013ms -> 217ms, with the full scored set identical. The set-difference rewrite proposed in #3512 also clears the reported bank, but it leaves the estimate at 2 and survives only because a set-op prices the nested loop just above the hash join: 1.0-1.2x headroom against 1.6-1.8x here. The trade is that ranking reads every unit_entities row of a matched entity where the LATERAL stopped at per_entity_limit off the index: O(sum of degree) rather than O(entities x per_entity_limit). At parity up to ~12k-degree hubs, +50% traversal cost at 38k. Why the perf suite never caught it ---------------------------------- `recall-with-observations` measured 0.45s on the same query a realistically shaped bank runs in 15s. Two fixture properties were wrong, and neither alone reproduces the bug — measured on the suite's own bank: sources=113 sources=mean 2 old vocabulary 450ms 270ms new vocabulary 951ms 15,013ms - The entity vocabulary was a fixed 145 names at every scale, so degree grew with bank size instead of the entity count growing: 142 entities at median degree 40 with not one entity mentioned once, and every seed reaching 142 of 142 entities. It now grows with the corpus (1,354 entities, median degree 3, 449 mentioned once, seeds reaching 54). - Sources per observation was a constant. Real counts are long-tailed — the reported bank ran mean 1.7 / p95 4 — so it is now the mean of a Pareto draw. At mean 2 the fixture still emits observations carrying several hundred sources, keeping the array-length path from #3085 exercised. `recall-with-observations` at scale=large will step up when this lands: the suite can finally see this query. The SQL fix is in the same change so the dashboard moves once, not twice. Oracle's expand_observations has the same DISTINCT-over-LATERAL shape and is deliberately left alone — no Oracle instance was available to measure it, and its cardinality estimation differs. Documented in ops_oracle.py. Tests ----- - test_per_entity_cap_bounds_hub_traversal pins that the window ranks the same rows the LATERAL selected; it fails if the cap is widened or dropped. - test_perf_fixture_shape asserts the post-resolution entity graph keeps a long tail. It simulates the entity resolver's intra-batch fuzzy merge, because the tail names have to stay under the 0.5 pg_trgm threshold: a tail generated as "<stem> <counter>" scores 0.73 and the resolver collapsed 2,814 names to 159, silently restoring the flat graph the vocabulary exists to avoid.
This commit is contained in:
Nicolò Boschi
@@ -696,6 +696,21 @@ class OracleOps(DataAccessOps):
|
||||
# Entity expansion via observation_sources junction table.
|
||||
# Previously used JSON_TABLE to explode source_memory_ids CLOB. The junction
|
||||
# table approach uses standard SQL joins, identical to the PG backend.
|
||||
#
|
||||
# Two PostgreSQL fixes are deliberately NOT mirrored here, because neither
|
||||
# was measured against Oracle and both are tuned to PostgreSQL's planner:
|
||||
# - #3085 made PG score set-wise; the scoring below is still the
|
||||
# correlated per-observation COUNT(*). On Oracle that counts rows of
|
||||
# the indexed observation_sources junction table rather than scanning
|
||||
# an unpruned array, so it is a much weaker version of that problem.
|
||||
# - #3510 replaced PG's `DISTINCT` over a `LATERAL ... LIMIT` with a
|
||||
# row_number() window, because PostgreSQL cannot estimate the row count
|
||||
# of that shape and mis-planned the scoring join into a nested loop.
|
||||
# `connected_sources` below has the same shape, so the same collapse is
|
||||
# structurally possible, but Oracle's cardinality estimation differs and
|
||||
# no Oracle instance was available to measure it.
|
||||
# If observation recall is reported slow on Oracle, start by capturing the
|
||||
# plan for connected_sources and checking its estimated vs actual rows.
|
||||
from ..schema import fq_table
|
||||
|
||||
obs_sources_table = fq_table("observation_sources")
|
||||
|
||||
@@ -914,6 +914,27 @@ class PostgreSQLOps(DataAccessOps):
|
||||
# ~1.7B element comparisons, 2.6s of one saturated backend (issue #3085).
|
||||
# Unnesting once and hash-joining connected_sources makes the work linear in
|
||||
# the number of source ids instead.
|
||||
#
|
||||
# `connected_sources` caps each entity with row_number() rather than the
|
||||
# LATERAL + LIMIT that reads more naturally. Do not "simplify" it back
|
||||
# (issue #3510). The scoring join above is O(C + U) when planned as a hash
|
||||
# join and O(U x C) when planned as a nested loop — 15s and ~15M rejected
|
||||
# rows on a realistically-shaped bank — and PostgreSQL picks between them
|
||||
# from its row estimate for this CTE. Out of a LATERAL + LIMIT subquery the
|
||||
# capped column carries no n_distinct statistic, so DISTINCT over it was
|
||||
# estimated at 2 and the NOT EXISTS took that to 1 against an actual ~3,700;
|
||||
# a 1-row inner side makes the nested loop look free, so it won on cost and
|
||||
# lost by four orders of magnitude at runtime. Ranking with a window keeps
|
||||
# the column traceable to unit_entities.unit_id, so the estimate comes from
|
||||
# real statistics (207-3,449 against 2,242-4,193 actual) and the nested loop
|
||||
# is priced honestly.
|
||||
#
|
||||
# The trade is that this reads every unit_entities row of a matched entity
|
||||
# to rank it, where the LATERAL stopped at per_entity_limit off the index:
|
||||
# O(sum of degree) rather than O(entities x per_entity_limit). Measured at
|
||||
# parity up to ~12k-degree hubs and +50% traversal cost at 38k. If banks
|
||||
# grow hubs far past that, re-measure before assuming this is still the
|
||||
# right shape.
|
||||
|
||||
entity_rows = await conn.fetch(
|
||||
f"""
|
||||
@@ -930,17 +951,20 @@ class PostgreSQLOps(DataAccessOps):
|
||||
),
|
||||
connected_sources AS (
|
||||
SELECT DISTINCT t.unit_id AS source_id
|
||||
FROM source_entities se
|
||||
CROSS JOIN LATERAL (
|
||||
SELECT ue_target.unit_id
|
||||
FROM (
|
||||
SELECT
|
||||
ue_target.unit_id,
|
||||
row_number() OVER (
|
||||
PARTITION BY ue_target.entity_id
|
||||
ORDER BY ue_target.unit_id DESC
|
||||
) AS rn
|
||||
FROM {ue_table} ue_target
|
||||
WHERE ue_target.entity_id = se.entity_id
|
||||
ORDER BY ue_target.unit_id DESC
|
||||
LIMIT {per_entity_limit}
|
||||
JOIN source_entities se ON se.entity_id = ue_target.entity_id
|
||||
) t
|
||||
WHERE NOT EXISTS (
|
||||
SELECT 1 FROM seed_sources ss WHERE ss.source_id = t.unit_id
|
||||
)
|
||||
WHERE t.rn <= {per_entity_limit}
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM seed_sources ss WHERE ss.source_id = t.unit_id
|
||||
)
|
||||
),
|
||||
connected_array AS (
|
||||
SELECT array_agg(source_id) AS source_ids FROM connected_sources
|
||||
|
||||
@@ -163,3 +163,73 @@ async def test_wide_source_arrays_do_not_change_results(memory, request_context)
|
||||
assert scores[narrow] == scores[wide] == 1.0
|
||||
finally:
|
||||
await memory.delete_bank(bank_id=bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_per_entity_cap_bounds_hub_traversal(memory, request_context):
|
||||
"""A hub entity contributes at most ``per_entity_limit`` source facts (#3510).
|
||||
|
||||
The cap used to be a LATERAL + LIMIT and is now a row_number() window, because
|
||||
the planner cannot estimate DISTINCT over a LIMIT subquery and mis-planned the
|
||||
scoring join into a nested loop. The two forms have to select the *same* rows:
|
||||
the highest ``per_entity_limit`` unit_ids of each entity. Ranking is by unit_id
|
||||
descending, which is what the LATERAL ordered by, so a candidate built from the
|
||||
lowest ids of an over-cap entity must fall outside the cap and score nothing.
|
||||
"""
|
||||
from hindsight_api.engine.db.ops import UpdatedWindow
|
||||
from hindsight_api.engine.task_backend import fq_table
|
||||
|
||||
bank_id = f"test_obs_cap_{uuid.uuid4().hex[:8]}"
|
||||
per_entity_limit = 3
|
||||
try:
|
||||
pool = await memory._get_pool()
|
||||
backend = await memory._get_backend()
|
||||
mu, ue, ml = fq_table("memory_units"), fq_table("unit_entities"), fq_table("memory_links")
|
||||
|
||||
async with pool.acquire() as conn:
|
||||
await _ensure_bank(conn, bank_id)
|
||||
entity_id = uuid.uuid4()
|
||||
await conn.execute(
|
||||
f"INSERT INTO {fq_table('entities')} (id, bank_id, canonical_name) VALUES ($1, $2, $3)",
|
||||
entity_id,
|
||||
bank_id,
|
||||
"Hub",
|
||||
)
|
||||
# Six facts on one entity with ids we control, so "top 3 by unit_id
|
||||
# descending" is a known set rather than an accident of uuid4().
|
||||
facts = sorted(uuid.UUID(int=i) for i in range(1, 7))
|
||||
for fid in facts:
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {mu} (id, bank_id, text, fact_type, source_memory_ids, event_date)
|
||||
VALUES ($1, $2, $3, 'world', NULL, $4)
|
||||
""",
|
||||
fid,
|
||||
bank_id,
|
||||
f"hub fact {fid}",
|
||||
datetime.now(timezone.utc),
|
||||
)
|
||||
await conn.execute(f"INSERT INTO {ue} (unit_id, entity_id) VALUES ($1, $2)", fid, entity_id)
|
||||
|
||||
# The seed reaches the hub through its own source fact.
|
||||
seed = await _insert_unit(conn, mu, bank_id, "seed obs", "observation", [facts[0]])
|
||||
# Inside the cap: built from the highest ids. Outside: the lowest.
|
||||
inside = await _insert_unit(conn, mu, bank_id, "inside cap", "observation", facts[-3:])
|
||||
outside = await _insert_unit(conn, mu, bank_id, "outside cap", "observation", [facts[1]])
|
||||
|
||||
rows = await backend.ops.expand_observations(
|
||||
conn,
|
||||
mu,
|
||||
ue,
|
||||
ml,
|
||||
[seed],
|
||||
100,
|
||||
per_entity_limit,
|
||||
UpdatedWindow(after=None, before=None, first_param_index=3),
|
||||
)
|
||||
|
||||
scores = {r["id"] for r in rows.entity}
|
||||
assert inside in scores, "observations built from the top-ranked ids must be reachable"
|
||||
assert outside not in scores, "the per-entity cap must exclude ids ranked below it"
|
||||
finally:
|
||||
await memory.delete_bank(bank_id=bank_id, request_context=request_context)
|
||||
|
||||
@@ -42,6 +42,7 @@ import asyncio
|
||||
import os
|
||||
import statistics
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
from rich.console import Console
|
||||
@@ -54,7 +55,9 @@ console = Console()
|
||||
# Fact corpus
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# ~200 entity names spanning people, technologies, and places.
|
||||
# The head of the entity vocabulary: ~145 names spanning people, technologies,
|
||||
# and places. These are the recurring hubs; _build_vocabulary appends a long
|
||||
# tail sized to the corpus.
|
||||
ENTITIES = [
|
||||
# People
|
||||
"Alice Chen",
|
||||
@@ -206,9 +209,9 @@ ENTITIES = [
|
||||
"SSO",
|
||||
]
|
||||
|
||||
# ~300 fact templates. Placeholders {E0}..{E4} are filled with entity names
|
||||
# drawn from ENTITIES using a Zipf-like distribution so that ~20 entities
|
||||
# recur frequently across templates.
|
||||
# ~300 fact templates. Placeholders {E0}..{E4} are filled from the entity
|
||||
# vocabulary using a Zipf-like distribution: a few hubs recur across most
|
||||
# templates, and the tail grows with the corpus (see _build_vocabulary).
|
||||
FACT_TEMPLATES = [
|
||||
"{E0} deployed a new version of {E1} to {E2} on {E3}.",
|
||||
"{E0} reported a performance regression in {E1} affecting {E2}.",
|
||||
@@ -466,20 +469,207 @@ SCALES = {
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_entity_selector(seed: int = 42) -> "Callable[[int], list[str]]":
|
||||
"""Return a function that draws N entity names with Zipf-like distribution."""
|
||||
# The long tail of the entity vocabulary is spelled "<qualifier> <thing>" from two
|
||||
# disjoint word lists. Two such names share at most one word, which keeps their
|
||||
# pg_trgm similarity near 0.33 — under the 0.5 at which the entity resolver
|
||||
# fuzzy-merges names within a retain batch. That threshold is the constraint: a
|
||||
# tail built by appending a counter to a stem ("Ingest Worker 0014" /
|
||||
# "Ingest Worker 0024") scores 0.73, and the resolver collapses the whole tail
|
||||
# back into a handful of entities, which is the flat hub-only graph this
|
||||
# vocabulary exists to avoid. test_perf_fixture_shape.py guards that outcome.
|
||||
_TAIL_QUALIFIERS = [
|
||||
"Copper",
|
||||
"Basalt",
|
||||
"Amber",
|
||||
"Cobalt",
|
||||
"Quartz",
|
||||
"Saffron",
|
||||
"Jade",
|
||||
"Onyx",
|
||||
"Ivory",
|
||||
"Cedar",
|
||||
"Flint",
|
||||
"Garnet",
|
||||
"Hazel",
|
||||
"Indigo",
|
||||
"Juniper",
|
||||
"Krypton",
|
||||
"Larch",
|
||||
"Marble",
|
||||
"Nickel",
|
||||
"Opal",
|
||||
"Pewter",
|
||||
"Quill",
|
||||
"Rowan",
|
||||
"Sable",
|
||||
"Topaz",
|
||||
"Ultra",
|
||||
"Verdant",
|
||||
"Willow",
|
||||
"Xenon",
|
||||
"Yarrow",
|
||||
"Zephyr",
|
||||
"Alloy",
|
||||
"Bronze",
|
||||
"Cinder",
|
||||
"Dune",
|
||||
"Cinnabar",
|
||||
"Frost",
|
||||
"Granite",
|
||||
"Harbor",
|
||||
"Iron",
|
||||
"Kelp",
|
||||
"Lumen",
|
||||
"Mica",
|
||||
"Nimbus",
|
||||
"Ochre",
|
||||
"Pumice",
|
||||
"Meadow",
|
||||
"Ridge",
|
||||
"Slate",
|
||||
"Thistle",
|
||||
"Vellum",
|
||||
"Walnut",
|
||||
"Cactus",
|
||||
"Birch",
|
||||
"Coral",
|
||||
"Drift",
|
||||
"Elm",
|
||||
"Fjord",
|
||||
"Gypsum",
|
||||
"Heather",
|
||||
"Isle",
|
||||
"Jetty",
|
||||
"Knoll",
|
||||
"Loam",
|
||||
"Moss",
|
||||
"Nettle",
|
||||
"Orchid",
|
||||
"Prairie",
|
||||
"Reef",
|
||||
"Sorrel",
|
||||
"Tundra",
|
||||
"Vine",
|
||||
]
|
||||
_TAIL_THINGS = [
|
||||
"Falcon",
|
||||
"Kettle",
|
||||
"Lantern",
|
||||
"Anvil",
|
||||
"Beacon",
|
||||
"Compass",
|
||||
"Dagger",
|
||||
"Ferry",
|
||||
"Gauntlet",
|
||||
"Harrow",
|
||||
"Igloo",
|
||||
"Jigsaw",
|
||||
"Kayak",
|
||||
"Ledger",
|
||||
"Mallet",
|
||||
"Nomad",
|
||||
"Obelisk",
|
||||
"Pylon",
|
||||
"Quiver",
|
||||
"Rudder",
|
||||
"Satchel",
|
||||
"Trellis",
|
||||
"Urchin",
|
||||
"Vault",
|
||||
"Wagon",
|
||||
"Yoke",
|
||||
"Zither",
|
||||
"Abacus",
|
||||
"Bellows",
|
||||
"Chisel",
|
||||
"Domino",
|
||||
"Easel",
|
||||
"Fulcrum",
|
||||
"Gasket",
|
||||
"Hammock",
|
||||
"Inkwell",
|
||||
"Jockey",
|
||||
"Kiln",
|
||||
"Lattice",
|
||||
"Mosaic",
|
||||
"Nozzle",
|
||||
"Octave",
|
||||
"Paddle",
|
||||
"Quarto",
|
||||
"Ratchet",
|
||||
"Sundial",
|
||||
"Turret",
|
||||
"Ukulele",
|
||||
"Vise",
|
||||
"Whistle",
|
||||
"Yurt",
|
||||
"Bobbin",
|
||||
"Armoire",
|
||||
"Buckle",
|
||||
"Cistern",
|
||||
"Drawbridge",
|
||||
"Effigy",
|
||||
"Cauldron",
|
||||
"Grommet",
|
||||
"Hearth",
|
||||
"Ingot",
|
||||
"Javelin",
|
||||
"Kestrel",
|
||||
"Lintel",
|
||||
"Mandrel",
|
||||
"Nautilus",
|
||||
"Ottoman",
|
||||
"Placard",
|
||||
"Quorum",
|
||||
"Rivet",
|
||||
"Spindle",
|
||||
"Skillet",
|
||||
]
|
||||
|
||||
|
||||
def _tail_entity(index: int) -> str:
|
||||
"""Deterministic name for the *index*-th long-tail entity."""
|
||||
qualifier = _TAIL_QUALIFIERS[index % len(_TAIL_QUALIFIERS)]
|
||||
thing = _TAIL_THINGS[(index // len(_TAIL_QUALIFIERS)) % len(_TAIL_THINGS)]
|
||||
return f"{qualifier} {thing}"
|
||||
|
||||
|
||||
def _build_vocabulary(corpus_size: int) -> list[str]:
|
||||
"""Entity vocabulary for a corpus of *corpus_size* facts.
|
||||
|
||||
The vocabulary has to grow with the corpus. Real banks are long-tailed: the
|
||||
bank reported in #3510 held 8,872 entities over 11.5k facts — median degree 2,
|
||||
a third of them mentioned exactly once, top entity on 8.7% of entity links.
|
||||
|
||||
This fixture drew from a fixed 145-name list at every scale, which inverts
|
||||
that shape: degree grows with bank size instead of the entity count growing.
|
||||
At 5k facts it produced 142 entities at a *median* degree of 40 with not one
|
||||
entity mentioned once, so every observation seed reached the entire entity
|
||||
graph. Together with a constant sources-per-observation (see
|
||||
_insert_synthetic_observations) that kept the observation graph arm in a
|
||||
regime real banks are never in, and #3510 went unseen here.
|
||||
"""
|
||||
capacity = len(_TAIL_QUALIFIERS) * len(_TAIL_THINGS)
|
||||
tail = min(max(0, corpus_size - len(ENTITIES)), capacity)
|
||||
return [*ENTITIES, *(_tail_entity(i) for i in range(tail))]
|
||||
|
||||
|
||||
def _make_entity_selector(corpus_size: int, seed: int = 42) -> "Callable[[int], list[str]]":
|
||||
"""Return a function that draws N entity names with a Zipf-like distribution."""
|
||||
import random
|
||||
|
||||
rng = random.Random(seed)
|
||||
n = len(ENTITIES)
|
||||
# Weights: entity i gets weight 1/(i+1)
|
||||
weights = [1.0 / (i + 1) for i in range(n)]
|
||||
vocabulary = _build_vocabulary(corpus_size)
|
||||
# Weights: entity i gets weight 1/(i+1). Over a vocabulary that scales with
|
||||
# the corpus this gives the shape above — a few hubs carrying ~10% of links
|
||||
# each, and a majority of entities appearing once or twice.
|
||||
weights = [1.0 / (i + 1) for i in range(len(vocabulary))]
|
||||
|
||||
def pick(count: int) -> list[str]:
|
||||
seen = set()
|
||||
result = []
|
||||
while len(result) < count:
|
||||
choice = rng.choices(ENTITIES, weights=weights, k=1)[0]
|
||||
choice = rng.choices(vocabulary, weights=weights, k=1)[0]
|
||||
if choice not in seen:
|
||||
seen.add(choice)
|
||||
result.append(choice)
|
||||
@@ -488,20 +678,42 @@ def _make_entity_selector(seed: int = 42) -> "Callable[[int], list[str]]":
|
||||
return pick
|
||||
|
||||
|
||||
_pick_entities = _make_entity_selector()
|
||||
# Sized to the hub head until a suite declares its corpus via
|
||||
# configure_entity_vocabulary(), which every bank-populating path does.
|
||||
_pick_entities = _make_entity_selector(len(ENTITIES))
|
||||
|
||||
|
||||
def _fill_template(template: str) -> str:
|
||||
def configure_entity_vocabulary(corpus_size: int, seed: int = 42) -> None:
|
||||
"""Size the entity vocabulary to the corpus a suite is about to generate.
|
||||
|
||||
Must be called before building fact contents; a suite that forgets silently
|
||||
gets the 145-name head and the flat entity graph described in
|
||||
_build_vocabulary. Re-seeding here also makes each suite's corpus
|
||||
reproducible regardless of how many facts were generated before it.
|
||||
"""
|
||||
global _pick_entities
|
||||
_pick_entities = _make_entity_selector(corpus_size, seed)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class FilledFact:
|
||||
"""A fact rendered from a template, with the entities that went into it."""
|
||||
|
||||
text: str
|
||||
entities: list[str]
|
||||
|
||||
|
||||
def _fill_template(template: str) -> FilledFact:
|
||||
"""Replace {E0}..{E4} placeholders in a template with entity names."""
|
||||
placeholders = [f"{{E{i}}}" for i in range(5)]
|
||||
needed = sum(1 for p in placeholders if p in template)
|
||||
if needed == 0:
|
||||
return template
|
||||
return FilledFact(text=template, entities=[])
|
||||
entities = _pick_entities(needed)
|
||||
result = template
|
||||
for i, entity in enumerate(entities):
|
||||
result = result.replace(f"{{E{i}}}", entity)
|
||||
return result
|
||||
return FilledFact(text=result, entities=entities)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -530,11 +742,13 @@ def _make_fact_callback() -> tuple[Callable[[list[dict], str], Any], list[int]]:
|
||||
fact_type = _FACT_TYPE_CYCLE[call_counter[0] % len(_FACT_TYPE_CYCLE)]
|
||||
call_counter[0] += 1
|
||||
template = FACT_TEMPLATES[idx]
|
||||
fact_text = _fill_template(template)
|
||||
# Extract a few entity names from the filled text to populate entities
|
||||
# field (very rough — enough to drive entity link creation)
|
||||
entity_names = [e for e in ENTITIES if e in fact_text][:3]
|
||||
entities = [{"text": e} for e in entity_names]
|
||||
filled = _fill_template(template)
|
||||
fact_text = filled.text
|
||||
# Take the entities the template was actually filled with. Scanning the
|
||||
# vocabulary for substrings (what this did) is O(vocabulary) per fact and
|
||||
# mis-attributes once the vocabulary is large enough to hold names that
|
||||
# are substrings of one another.
|
||||
entities = [{"text": e} for e in filled.entities[:3]]
|
||||
return {
|
||||
"facts": [
|
||||
{
|
||||
@@ -590,15 +804,22 @@ async def _insert_synthetic_observations(pool: Any, bank_id: str, sources_per_ob
|
||||
observation with the same text, embedding, and tags, pointing back to that
|
||||
unit as a source fact.
|
||||
|
||||
``sources_per_observation`` controls how many source facts each observation
|
||||
carries. This is the dimension the observation graph arm actually scales on:
|
||||
``expand_observations`` unnests ``source_memory_ids`` for every seed and again
|
||||
for every candidate, so its cost grows with the length of those arrays — not
|
||||
with the observation count. Real banks grow them without bound, because
|
||||
consolidation appends a source id on every merge and never prunes (#1725); a
|
||||
reported bank ran at ~113 sources per observation (#3085). Leaving this at 1
|
||||
(as this fixture did originally) keeps the arm permanently in its cheapest
|
||||
regime and hides that whole class of regression.
|
||||
``sources_per_observation`` is the *mean* number of source facts an
|
||||
observation carries; the counts themselves are drawn long-tailed around it
|
||||
(see ``_draw_source_count``).
|
||||
|
||||
Two dimensions of the observation graph arm hang off this. Array *length*
|
||||
drives the scoring cost — ``expand_observations`` unnests
|
||||
``source_memory_ids`` for every seed and again for every candidate — and
|
||||
consolidation grows those arrays without pruning (#1725), so an aged bank can
|
||||
average ~113 (#3085). Array length also decides how much of the entity graph
|
||||
a handful of seeds reaches, which is what the graph arm's plan turns on
|
||||
(#3510).
|
||||
|
||||
This was a constant, which got both dimensions wrong at once: it erased the
|
||||
tail that #3085 cares about and the low floor that #3510 cares about. A mean
|
||||
of 2 with a long tail reproduces both — most observations carry a single
|
||||
source, while the widest carry several hundred.
|
||||
|
||||
Sources are drawn from a window of neighbouring facts so that source sets
|
||||
*overlap* between observations, which is what makes the `&&` candidate scan
|
||||
@@ -628,17 +849,39 @@ async def _insert_synthetic_observations(pool: Any, bank_id: str, sources_per_ob
|
||||
return 0
|
||||
|
||||
fact_ids = [row["id"] for row in rows]
|
||||
n_sources = max(1, min(sources_per_observation, len(fact_ids)))
|
||||
mean_sources = max(1, min(sources_per_observation, len(fact_ids)))
|
||||
rng = random.Random(42) # deterministic fixture
|
||||
|
||||
def _draw_source_count() -> int:
|
||||
"""How many source facts this observation carries.
|
||||
|
||||
Pinning every observation to one value was the wrong model. Consolidation
|
||||
grows source_memory_ids one merge at a time, so real counts are
|
||||
long-tailed: the bank in #3510 ran mean 1.7 / p95 4 / max 61, while the
|
||||
aged bank in #3085 averaged 113. A constant erases both ends — with every
|
||||
observation carrying the mean, a handful of seeds reaches almost the whole
|
||||
entity graph, which is a regime real banks are never in and is why this
|
||||
fixture measured 0.45s where a realistically-shaped bank measures 15s.
|
||||
|
||||
Pareto gives the right shape: most observations near 1, a thin tail of
|
||||
wide ones, scaled so the sample mean lands on ``mean_sources``.
|
||||
"""
|
||||
if mean_sources == 1:
|
||||
return 1
|
||||
# Pareto(a) has mean a/(a-1); rescale so the draw averages mean_sources.
|
||||
shape = 1.6
|
||||
draw = rng.paretovariate(shape) * mean_sources * (shape - 1) / shape
|
||||
return max(1, min(int(draw), len(fact_ids)))
|
||||
|
||||
def _sources_for(index: int) -> list:
|
||||
"""Own fact plus neighbours, so adjacent observations share source ids."""
|
||||
n_sources = _draw_source_count()
|
||||
if n_sources == 1:
|
||||
return [fact_ids[index]]
|
||||
window_size = min(len(fact_ids), n_sources * 3)
|
||||
start = min(max(0, index - window_size // 2), len(fact_ids) - window_size)
|
||||
window = [fid for fid in fact_ids[start : start + window_size] if fid != fact_ids[index]]
|
||||
return [fact_ids[index], *rng.sample(window, n_sources - 1)]
|
||||
return [fact_ids[index], *rng.sample(window, min(n_sources - 1, len(window)))]
|
||||
|
||||
inserted = 0
|
||||
for offset in range(0, len(rows), _BATCH_SIZE):
|
||||
@@ -767,8 +1010,9 @@ async def cmd_generate(
|
||||
engine._llm_config.set_response_callback(callback)
|
||||
|
||||
# Build all content items upfront
|
||||
configure_entity_vocabulary(total_items)
|
||||
all_contents: list[dict[str, Any]] = [
|
||||
{"content": _fill_template(FACT_TEMPLATES[i % len(FACT_TEMPLATES)])} for i in range(total_items)
|
||||
{"content": _fill_template(FACT_TEMPLATES[i % len(FACT_TEMPLATES)]).text} for i in range(total_items)
|
||||
]
|
||||
if event_date:
|
||||
# Stamp every item with the same event_date → mentioned_at clusters at one
|
||||
|
||||
@@ -55,6 +55,7 @@ from benchmarks.perf.recall_perf import (
|
||||
_make_fact_callback,
|
||||
_RRFReranker,
|
||||
_wait_for_operation,
|
||||
configure_entity_vocabulary,
|
||||
)
|
||||
|
||||
console = Console()
|
||||
@@ -68,7 +69,7 @@ SCALES: dict[str, dict[str, int]] = {
|
||||
"retain_items": 20,
|
||||
"recall_bank_size": 20,
|
||||
"recall_iterations": 5,
|
||||
"recall_obs_sources_per_observation": 4,
|
||||
"recall_obs_sources_per_observation": 2,
|
||||
"recall_concurrency": 1,
|
||||
"consolidation_items": 20,
|
||||
"graph_maintenance_bank_size": 20,
|
||||
@@ -83,7 +84,7 @@ SCALES: dict[str, dict[str, int]] = {
|
||||
"retain_items": 200,
|
||||
"recall_bank_size": 200,
|
||||
"recall_iterations": 20,
|
||||
"recall_obs_sources_per_observation": 8,
|
||||
"recall_obs_sources_per_observation": 2,
|
||||
"recall_concurrency": 4,
|
||||
"consolidation_items": 200,
|
||||
"graph_maintenance_bank_size": 200,
|
||||
@@ -98,7 +99,7 @@ SCALES: dict[str, dict[str, int]] = {
|
||||
"retain_items": 1_000,
|
||||
"recall_bank_size": 1_000,
|
||||
"recall_iterations": 50,
|
||||
"recall_obs_sources_per_observation": 32,
|
||||
"recall_obs_sources_per_observation": 3,
|
||||
"recall_concurrency": 8,
|
||||
"consolidation_items": 1_000,
|
||||
"graph_maintenance_bank_size": 1_000,
|
||||
@@ -113,12 +114,19 @@ SCALES: dict[str, dict[str, int]] = {
|
||||
"retain_items": 5_000,
|
||||
"recall_bank_size": 5_000,
|
||||
"recall_iterations": 100,
|
||||
# Source facts per synthetic observation. The observation graph arm's cost
|
||||
# scales with the length of source_memory_ids, which consolidation grows
|
||||
# without bound (#1725) — 113 is the average measured on the bank reported
|
||||
# in #3085. At the fixture's original value of 1 the arm never leaves its
|
||||
# cheapest regime, which is why a 39x blowup went unseen here.
|
||||
"recall_obs_sources_per_observation": 113,
|
||||
# *Mean* source facts per synthetic observation; the counts are drawn
|
||||
# long-tailed around it (see _insert_synthetic_observations). Real banks
|
||||
# sit low with a thin wide tail — the bank in #3510 ran mean 1.7, p95 4 —
|
||||
# and at this mean the fixture still produces observations carrying
|
||||
# several hundred sources, so the array-length cost path from #3085 stays
|
||||
# exercised without pinning every observation to it.
|
||||
#
|
||||
# This was a constant 113 (the mean of the aged bank in #3085). Pinning
|
||||
# every observation to that mean made a handful of seeds reach almost the
|
||||
# whole entity graph, so the observation graph arm sat permanently in a
|
||||
# regime real banks are never in: the suite measured 0.45s here where a
|
||||
# realistically-shaped bank measures 15s, and #3510 went unseen.
|
||||
"recall_obs_sources_per_observation": 2,
|
||||
"recall_concurrency": 16,
|
||||
"consolidation_items": 5_000,
|
||||
# Past the seqscan→HNSW crossover (~10k units) so this suite exercises
|
||||
@@ -150,7 +158,7 @@ SCALES: dict[str, dict[str, int]] = {
|
||||
"retain_items": 5_000,
|
||||
"recall_bank_size": 5_000,
|
||||
"recall_iterations": 10,
|
||||
"recall_obs_sources_per_observation": 113,
|
||||
"recall_obs_sources_per_observation": 2,
|
||||
"recall_concurrency": 4,
|
||||
"consolidation_items": 5_000,
|
||||
"graph_maintenance_bank_size": 15_000,
|
||||
@@ -402,7 +410,10 @@ async def _populate_bank(engine: Any, bank_id: str, size: int, event_date: str |
|
||||
|
||||
_attach_mock_callback(engine)
|
||||
|
||||
contents = [{"content": _fill_template(FACT_TEMPLATES[i % len(FACT_TEMPLATES)])} for i in range(size)]
|
||||
# Size the entity vocabulary to this corpus so the entity graph keeps a real
|
||||
# bank's long tail instead of turning every entity into a hub (#3510).
|
||||
configure_entity_vocabulary(size)
|
||||
contents = [{"content": _fill_template(FACT_TEMPLATES[i % len(FACT_TEMPLATES)]).text} for i in range(size)]
|
||||
if event_date:
|
||||
for item in contents:
|
||||
item["event_date"] = event_date
|
||||
@@ -693,7 +704,8 @@ async def run_retain_suite(scale_cfg: dict[str, int]) -> SuiteResult:
|
||||
await engine.initialize()
|
||||
_attach_mock_callback(engine)
|
||||
|
||||
contents = [{"content": _fill_template(FACT_TEMPLATES[i % len(FACT_TEMPLATES)])} for i in range(total_items)]
|
||||
configure_entity_vocabulary(total_items)
|
||||
contents = [{"content": _fill_template(FACT_TEMPLATES[i % len(FACT_TEMPLATES)]).text} for i in range(total_items)]
|
||||
request_context = RequestContext()
|
||||
|
||||
t0 = time.perf_counter()
|
||||
@@ -876,7 +888,8 @@ async def run_recall_with_observations_suite(scale_cfg: dict[str, int]) -> Suite
|
||||
engine._cross_encoder_reranker = _RRFReranker()
|
||||
|
||||
# Populate bank with facts then insert synthetic observations (1 per fact,
|
||||
# each carrying `sources_per_obs` source facts — see the scale config)
|
||||
# each carrying a long-tailed draw averaging `sources_per_obs` source
|
||||
# facts — see the scale config)
|
||||
await _populate_bank(engine, bank_id, bank_size)
|
||||
|
||||
pool = await engine._get_pool()
|
||||
@@ -888,7 +901,7 @@ async def run_recall_with_observations_suite(scale_cfg: dict[str, int]) -> Suite
|
||||
) as progress:
|
||||
progress.add_task("Inserting synthetic observations…")
|
||||
n_obs = await _insert_synthetic_observations(pool, bank_id, sources_per_obs)
|
||||
console.print(f" Inserted {n_obs:,} observations ({sources_per_obs} source facts each)")
|
||||
console.print(f" Inserted {n_obs:,} observations (mean {sources_per_obs} source facts each)")
|
||||
|
||||
request_context = RequestContext()
|
||||
durations: list[float] = []
|
||||
@@ -1202,7 +1215,8 @@ async def run_consolidation_suite(scale_cfg: dict[str, int]) -> SuiteResult:
|
||||
# This queues consolidation tasks (since observations are enabled) but
|
||||
# no worker is running so they just sit in the queue — we run consolidation
|
||||
# explicitly below.
|
||||
contents = [{"content": _fill_template(FACT_TEMPLATES[i % len(FACT_TEMPLATES)])} for i in range(total_items)]
|
||||
configure_entity_vocabulary(total_items)
|
||||
contents = [{"content": _fill_template(FACT_TEMPLATES[i % len(FACT_TEMPLATES)]).text} for i in range(total_items)]
|
||||
request_context = RequestContext()
|
||||
|
||||
with Progress(
|
||||
|
||||
@@ -0,0 +1,86 @@
|
||||
"""The perf fixture's entity-graph shape (issue #3510).
|
||||
|
||||
``recall-with-observations`` only exercises the observation graph arm's real
|
||||
failure mode if its synthetic bank is shaped like a real one. The property that
|
||||
does the work is the *post-resolution* entity graph: a few hubs over a long tail
|
||||
of rarely-mentioned entities, so a handful of observation seeds reaches a small
|
||||
fraction of the graph rather than all of it.
|
||||
|
||||
These tests assert that outcome rather than the mechanism. The mechanism has
|
||||
already failed in a non-obvious way once — tail names generated as
|
||||
"<stem> <counter>" scored 0.73 pg_trgm against their neighbours, so the entity
|
||||
resolver merged 2,814 generated names back down to 159 and the bank came out with
|
||||
the flat hub-only graph the vocabulary exists to avoid. Checking name distances
|
||||
directly would flag harmless pairs and still miss that outcome.
|
||||
"""
|
||||
|
||||
import statistics
|
||||
from collections import Counter
|
||||
|
||||
from benchmarks.perf.recall_perf import (
|
||||
ENTITIES,
|
||||
FACT_TEMPLATES,
|
||||
_build_vocabulary,
|
||||
_fill_template,
|
||||
configure_entity_vocabulary,
|
||||
)
|
||||
|
||||
# The entity resolver fuzzy-merges two names within a retain batch at or above
|
||||
# this pg_trgm similarity (EntityResolver.intrabatch_merge_similarity).
|
||||
_MERGE_SIMILARITY = 0.5
|
||||
_CORPUS = 5_000
|
||||
|
||||
|
||||
def _mentions_after_resolution(corpus_size: int) -> Counter:
|
||||
"""Entity mention counts as the bank would hold them, post fuzzy-merge.
|
||||
|
||||
Simulates the resolver's intra-batch merge so the shape assertions below
|
||||
describe the graph that actually lands in ``unit_entities``, not the names
|
||||
the generator emitted.
|
||||
"""
|
||||
from hindsight_api.engine.entity_resolver import _trigram_similarity
|
||||
|
||||
configure_entity_vocabulary(corpus_size)
|
||||
raw: Counter = Counter()
|
||||
for i in range(corpus_size):
|
||||
for entity in _fill_template(FACT_TEMPLATES[i % len(FACT_TEMPLATES)]).entities[:3]:
|
||||
raw[entity] += 1
|
||||
|
||||
canonical: dict[str, str] = {}
|
||||
for name in raw:
|
||||
for seen in canonical:
|
||||
if _trigram_similarity(name, seen) >= _MERGE_SIMILARITY:
|
||||
canonical[name] = canonical[seen]
|
||||
break
|
||||
else:
|
||||
canonical[name] = name
|
||||
|
||||
merged: Counter = Counter()
|
||||
for name, count in raw.items():
|
||||
merged[canonical[name]] += count
|
||||
return merged
|
||||
|
||||
|
||||
def test_vocabulary_grows_with_the_corpus():
|
||||
"""A fixed-size vocabulary makes degree scale with bank size, not entity count."""
|
||||
assert len(_build_vocabulary(len(ENTITIES))) == len(ENTITIES), "the hub head is the floor"
|
||||
assert len(_build_vocabulary(500)) < len(_build_vocabulary(5_000)), "vocabulary must grow with the corpus"
|
||||
|
||||
|
||||
def test_entity_graph_keeps_a_long_tail():
|
||||
"""Hubs over a tail of rare entities — the shape the #3510 bank had.
|
||||
|
||||
That bank ran median degree 2 with a third of its entities mentioned once and
|
||||
its top entity on 8.7% of links. The old fixed list produced median degree 40
|
||||
and *zero* entities mentioned once, so every seed reached the whole graph.
|
||||
"""
|
||||
merged = _mentions_after_resolution(_CORPUS)
|
||||
total = sum(merged.values())
|
||||
|
||||
assert len(merged) > 800, (
|
||||
f"only {len(merged)} entities survived resolution for {_CORPUS} facts — the tail is "
|
||||
"collapsing, most likely because generated names fuzzy-merge into each other"
|
||||
)
|
||||
assert statistics.median(merged.values()) <= 8, "the median entity must stay rare"
|
||||
assert sum(1 for c in merged.values() if c == 1) > 50, "a real bank has entities seen once"
|
||||
assert max(merged.values()) / total > 0.02, "and it still has hubs"
|
||||
Reference in New Issue
Block a user