feat(core): add config-gated entity-aware ranking boost for hybrid search

Proper nouns in a query carry no extra weight against generic semantic
similarity, so documents about a different entity on the same topic can
outrank the document that actually names the queried entity (#951
cross-conversation confusion in the LoCoMo benchmark).

Add an optional, lexical-only re-scoring pass to hybrid fusion:

- Extract candidate entity terms from the query (capitalized / proper-noun
  tokens that are not common stopwords; trailing possessives stripped).
- Count how many distinct query entity terms appear in each fused
  candidate's entity name (title) or a relation row's linked entity names.
- Multiply matching candidates' fused scores by
  1 + weight * min(matches, max_terms), promoting entity-matching docs.

The boost runs over the full fused candidate set before the limit/offset
cut, so a matching doc below the cutoff can be promoted into the returned
window. It adds no model inference (index/lexical lookups only), so
per-query latency overhead is trivial, and only affects hybrid retrieval.

Behind three config flags, DEFAULT OFF pending LoCoMo benchmark validation:
search_entity_boost_enabled, search_entity_boost_weight,
search_entity_boost_max_terms. Documented in docs/semantic-search.md.

Tests: unit coverage for entity-term extraction and the boost math; a
hybrid-pipeline test showing reordering when enabled and unchanged ordering
when disabled; and a service-level integration test over a real DB with a
deterministic stub embedding provider proving an entity-matching doc
outranks a higher-similarity non-matching doc only when enabled.

Co-Authored-By: Claude <noreply@anthropic.com>
Signed-off-by: Drew Cain <groksrc@gmail.com>
This commit is contained in:
Drew Cain
2026-06-11 01:22:54 -05:00
parent d46c68806e
commit 9fac9a9416
8 changed files with 764 additions and 0 deletions
@@ -0,0 +1,227 @@
"""Service-level integration test for the entity-aware ranking boost (#951).
Drives a fully wired SearchService over a real database with a deterministic stub
embedding provider so vector similarity is controlled. Verifies that when the boost
is enabled, an entity-matching document outranks a higher-similarity non-matching
document, and that ordering is unchanged when the boost is disabled.
No model inference is involved: the stub provider returns fixed unit vectors, so the
test is fast and deterministic on both SQLite and Postgres.
"""
from __future__ import annotations
import math
from typing import Any
import pytest
import pytest_asyncio
from basic_memory.config import BasicMemoryConfig
from basic_memory.repository.entity_repository import EntityRepository
from basic_memory.repository.search_repository import SearchRepository
from basic_memory.schemas.base import Entity as EntitySchema
from basic_memory.schemas.search import SearchQuery, SearchRetrievalMode
from basic_memory.services.entity_service import EntityService
from basic_memory.services.file_service import FileService
from basic_memory.services.search_service import SearchService
# --- Deterministic stub embedding provider ---
_STUB_DIMENSIONS = 4
def _unit(vector: list[float]) -> list[float]:
norm = math.sqrt(sum(component * component for component in vector)) or 1.0
return [component / norm for component in vector]
class _StubEmbeddingProvider:
"""Maps known text fragments to fixed unit vectors for controlled similarity.
The query is engineered to sit closer (cosine) to the non-matching "hobbies"
document than to the gold "Joanna" document, reproducing the #951 failure where
generic semantic similarity outranks the entity-matching gold doc.
"""
model_name = "stub-entity-boost"
dimensions = _STUB_DIMENSIONS
def _vector_for(self, text: str) -> list[float]:
lowered = text.lower()
if "joanna" in lowered:
# Gold doc: shares some direction with the query but less than the decoy.
return _unit([0.6, 0.8, 0.0, 0.0])
if "hobbies" in lowered or "pastime" in lowered:
# Decoy doc: closest to the query direction.
return _unit([0.95, 0.31, 0.0, 0.0])
return _unit([0.0, 0.0, 1.0, 0.0])
async def embed_query(self, text: str) -> list[float]:
# Query direction is closest to the decoy vector above.
return _unit([0.97, 0.24, 0.0, 0.0])
async def embed_documents(self, texts: list[str]) -> list[list[float]]:
return [self._vector_for(text) for text in texts]
def runtime_log_attrs(self) -> dict[str, Any]:
return {}
# --- Fixtures ---
async def _build_search_service(
*,
session_maker,
test_project,
base_app_config: BasicMemoryConfig,
file_service: FileService,
entity_repository: EntityRepository,
boost_enabled: bool,
) -> SearchService:
"""Build a SearchService with semantic search + a deterministic stub provider."""
from basic_memory.repository.sqlite_search_repository import SQLiteSearchRepository
from basic_memory.repository.postgres_search_repository import PostgresSearchRepository
from basic_memory.config import DatabaseBackend
app_config = base_app_config.model_copy(
update={
"semantic_search_enabled": True,
"semantic_min_similarity": 0.0,
"search_entity_boost_enabled": boost_enabled,
"search_entity_boost_weight": 0.3,
"search_entity_boost_max_terms": 3,
}
)
provider = _StubEmbeddingProvider()
if app_config.database_backend == DatabaseBackend.POSTGRES: # pragma: no cover
search_repo: SearchRepository = PostgresSearchRepository(
session_maker,
project_id=test_project.id,
app_config=app_config,
embedding_provider=provider,
)
else:
repo = SQLiteSearchRepository(
session_maker,
project_id=test_project.id,
app_config=app_config,
)
repo._semantic_enabled = True
repo._embedding_provider = provider
repo._vector_dimensions = provider.dimensions
repo._vector_tables_initialized = False
search_repo = repo
service = SearchService(search_repo, entity_repository, file_service)
await service.init_search_index()
return service
@pytest_asyncio.fixture
async def boost_entities(
entity_service: EntityService,
):
"""Index two entities: a decoy 'hobbies' doc and the gold 'Joanna' doc."""
decoy, _ = await entity_service.create_or_update_entity(
EntitySchema(
title="Common Hobbies and Pastimes",
note_type="note",
directory="people",
content="A general overview of hobbies and pastimes people enjoy.",
)
)
gold, _ = await entity_service.create_or_update_entity(
EntitySchema(
title="Joanna",
note_type="note",
directory="people",
content="Notes about Joanna and what she likes to do.",
)
)
return decoy, gold
# --- Tests ---
async def _sync_vectors(service: SearchService, entity_ids: list[int]) -> None:
"""Embed the indexed entities via the stub provider."""
await service.sync_entity_vectors_batch(entity_ids)
@pytest.mark.asyncio
async def test_entity_boost_enabled_promotes_gold_doc(
session_maker,
test_project,
app_config,
file_service,
entity_repository,
boost_entities,
):
decoy, gold = boost_entities
service = await _build_search_service(
session_maker=session_maker,
test_project=test_project,
base_app_config=app_config,
file_service=file_service,
entity_repository=entity_repository,
boost_enabled=True,
)
# Re-index the entities through this service so vector tables exist for it.
for entity in (decoy, gold):
await service.index_entity(entity)
await _sync_vectors(service, [decoy.id, gold.id])
results = await service.search(
SearchQuery(
text="What are Joanna's hobbies?",
retrieval_mode=SearchRetrievalMode.HYBRID,
),
limit=10,
)
entity_ids = [r.entity_id for r in results]
assert gold.id in entity_ids and decoy.id in entity_ids
# With the boost on, the entity-matching gold doc ranks ahead of the
# higher-similarity decoy.
assert entity_ids.index(gold.id) < entity_ids.index(decoy.id)
@pytest.mark.asyncio
async def test_entity_boost_disabled_keeps_similarity_order(
session_maker,
test_project,
app_config,
file_service,
entity_repository,
boost_entities,
):
decoy, gold = boost_entities
service = await _build_search_service(
session_maker=session_maker,
test_project=test_project,
base_app_config=app_config,
file_service=file_service,
entity_repository=entity_repository,
boost_enabled=False,
)
for entity in (decoy, gold):
await service.index_entity(entity)
await _sync_vectors(service, [decoy.id, gold.id])
results = await service.search(
SearchQuery(
text="What are Joanna's hobbies?",
retrieval_mode=SearchRetrievalMode.HYBRID,
),
limit=10,
)
entity_ids = [r.entity_id for r in results]
assert gold.id in entity_ids and decoy.id in entity_ids
# With the boost off, the higher-similarity decoy ranks ahead of the gold doc.
assert entity_ids.index(decoy.id) < entity_ids.index(gold.id)
+109
View File
@@ -134,6 +134,115 @@ HYBRID_KWARGS: dict[str, Any] = dict(
)
def _hybrid_kwargs(**overrides: Any) -> dict[str, Any]:
"""Return HYBRID_KWARGS with overrides applied, typed as dict[str, Any].
Keeps the splat into the keyword-only _search_hybrid signature type-clean.
"""
merged: dict[str, Any] = {**HYBRID_KWARGS, **overrides}
return merged
@pytest.mark.asyncio
async def test_entity_boost_promotes_matching_doc_when_enabled():
"""With entity boost enabled, an entity-matching doc outranks a higher-similarity
non-matching doc.
Reproduces the #951 cross-conversation confusion: a generic same-topic document
(higher raw similarity) initially outranks the gold doc whose title names the
queried entity. Enabling the boost flips the order.
"""
repo = ConcreteSearchRepo()
repo._entity_boost_enabled = True
repo._entity_boost_weight = 0.15
repo._entity_boost_max_terms = 3
# Row 1: generic hobbies doc from the wrong conversation, higher vector similarity.
# Row 2: the gold doc whose title names the queried entity "Joanna".
fts_results = []
vector_results = [
FakeRow(id=1, score=0.80, title="Hobbies and pastimes"),
FakeRow(id=2, score=0.72, title="Joanna profile"),
]
with (
patch.object(repo, "search", new_callable=AsyncMock, return_value=fts_results),
patch.object(
repo, "_search_vector_only", new_callable=AsyncMock, return_value=vector_results
),
):
results = await repo._search_hybrid(
**_hybrid_kwargs(search_text="What are Joanna's hobbies?")
)
# Boost: row 2 -> 0.72 * 1.15 = 0.828 > row 1's 0.80
assert [r.id for r in results] == [2, 1]
assert results[0].score == pytest.approx(0.72 * 1.15, rel=1e-6)
assert results[1].score == pytest.approx(0.80, rel=1e-6)
@pytest.mark.asyncio
async def test_entity_boost_disabled_preserves_ordering():
"""With entity boost disabled (default), ordering matches pure similarity."""
repo = ConcreteSearchRepo()
# Defaults from the base class keep boosting off; assert explicitly.
assert repo._entity_boost_enabled is False
fts_results = []
vector_results = [
FakeRow(id=1, score=0.80, title="Hobbies and pastimes"),
FakeRow(id=2, score=0.72, title="Joanna profile"),
]
with (
patch.object(repo, "search", new_callable=AsyncMock, return_value=fts_results),
patch.object(
repo, "_search_vector_only", new_callable=AsyncMock, return_value=vector_results
),
):
results = await repo._search_hybrid(
**_hybrid_kwargs(search_text="What are Joanna's hobbies?")
)
# No boost: original similarity order is preserved, scores unchanged.
assert [r.id for r in results] == [1, 2]
assert results[0].score == pytest.approx(0.80, rel=1e-6)
assert results[1].score == pytest.approx(0.72, rel=1e-6)
@pytest.mark.asyncio
async def test_entity_boost_promotes_doc_into_limited_window():
"""Boosting runs before the limit cut, so a matching doc ranked below the cutoff
can be promoted into the returned window."""
repo = ConcreteSearchRepo()
repo._entity_boost_enabled = True
repo._entity_boost_weight = 0.6
repo._entity_boost_max_terms = 3
fts_results = []
# Three non-matching docs above the gold doc, which matches "Anthony".
vector_results = [
FakeRow(id=1, score=0.90, title="conversation six"),
FakeRow(id=2, score=0.85, title="conversation one"),
FakeRow(id=3, score=0.60, title="Anthony introduces himself"),
]
with (
patch.object(repo, "search", new_callable=AsyncMock, return_value=fts_results),
patch.object(
repo, "_search_vector_only", new_callable=AsyncMock, return_value=vector_results
),
):
results = await repo._search_hybrid(
**_hybrid_kwargs(search_text="Who is Anthony?", limit=1)
)
# Gold doc boost: 0.60 * 1.6 = 0.96 > row 1's 0.90, so it is promoted into the
# top-1 window even though it was ranked third before boosting.
assert len(results) == 1
assert results[0].id == 3
@pytest.mark.asyncio
async def test_high_fts_score_boosts_ranking():
"""FTS-only: a high normalized score should outscore a low normalized score."""
@@ -809,3 +809,182 @@ async def test_sync_entity_vectors_batch_logs_resolved_fastembed_runtime_setting
assert runtime_logs[0]["threads"] == 4
assert runtime_logs[0]["configured_parallel"] == 2
assert runtime_logs[0]["effective_parallel"] == 2
# --- Entity-aware ranking boost (#951) ---
def _make_index_row(
*,
row_id: int,
title: str,
row_type: str = SearchItemType.ENTITY.value,
) -> SearchIndexRow:
"""Build a real SearchIndexRow for entity-boost matching tests."""
now = datetime(2026, 1, 1)
return SearchIndexRow(
project_id=1,
id=row_id,
type=row_type,
file_path=f"notes/{row_id}.md",
created_at=now,
updated_at=now,
permalink=f"notes/{row_id}",
title=title,
)
class TestExtractQueryEntityTerms:
"""Verify proper-noun extraction from query strings."""
def test_extracts_single_proper_noun(self):
terms = SearchRepositoryBase._extract_query_entity_terms("What are Joanna's hobbies?")
assert terms == {"joanna"}
def test_extracts_multiple_proper_nouns(self):
terms = SearchRepositoryBase._extract_query_entity_terms(
"What symbolic gifts do Deborah and Jolene have from their mothers?"
)
assert terms == {"deborah", "jolene"}
def test_who_is_anthony(self):
terms = SearchRepositoryBase._extract_query_entity_terms("Who is Anthony?")
assert terms == {"anthony"}
def test_all_lowercase_query_has_no_entity_terms(self):
terms = SearchRepositoryBase._extract_query_entity_terms("what is the weather today")
assert terms == set()
def test_capitalized_stopword_at_sentence_start_is_ignored(self):
# "What" and "Who" are capitalized interrogatives, not entity names.
terms = SearchRepositoryBase._extract_query_entity_terms("What does Sarah like?")
assert terms == {"sarah"}
def test_all_caps_token_is_extracted(self):
terms = SearchRepositoryBase._extract_query_entity_terms("who works at NASA")
assert terms == {"nasa"}
def test_possessive_is_stripped(self):
terms = SearchRepositoryBase._extract_query_entity_terms("Joanna's mother")
assert terms == {"joanna"}
def test_single_capital_letter_is_ignored(self):
# "A" lowercases to a stopword; "X" is a non-stopword single letter that
# still carries no entity signal and must be dropped by the length guard.
terms = SearchRepositoryBase._extract_query_entity_terms("A is for X and Apple")
assert terms == {"apple"}
def test_empty_and_none_inputs(self):
assert SearchRepositoryBase._extract_query_entity_terms("") == set()
assert SearchRepositoryBase._extract_query_entity_terms(None) == set()
class TestRowEntityMatchCount:
"""Verify lexical match counting between query terms and row entity names."""
def test_title_match_counts_one(self):
row = _make_index_row(row_id=1, title="Joanna's Hobbies")
assert SearchRepositoryBase._row_entity_match_count(row, {"joanna"}) == 1
def test_relation_title_matches_both_endpoints(self):
row = _make_index_row(
row_id=2,
title="Deborah -> Jolene",
row_type=SearchItemType.RELATION.value,
)
count = SearchRepositoryBase._row_entity_match_count(row, {"deborah", "jolene"})
assert count == 2
def test_no_match_returns_zero(self):
row = _make_index_row(row_id=3, title="Anthony Profile")
assert SearchRepositoryBase._row_entity_match_count(row, {"joanna"}) == 0
def test_match_is_case_insensitive(self):
row = _make_index_row(row_id=4, title="JOANNA notes")
assert SearchRepositoryBase._row_entity_match_count(row, {"joanna"}) == 1
def test_empty_terms_returns_zero(self):
row = _make_index_row(row_id=5, title="Joanna")
assert SearchRepositoryBase._row_entity_match_count(row, set()) == 0
def test_missing_title_returns_zero(self):
row = _make_index_row(row_id=6, title="")
row.title = None
assert SearchRepositoryBase._row_entity_match_count(row, {"joanna"}) == 0
def test_distinct_terms_not_double_counted(self):
# A title containing the same term twice still counts as one distinct match.
row = _make_index_row(row_id=7, title="Joanna and Joanna")
assert SearchRepositoryBase._row_entity_match_count(row, {"joanna"}) == 1
class TestApplyEntityBoost:
"""Verify the entity-boost score math and gating."""
def _repo(self, *, enabled: bool, weight: float = 0.15, max_terms: int = 3) -> _ConcreteRepo:
repo = _ConcreteRepo()
repo._entity_boost_enabled = enabled
repo._entity_boost_weight = weight
repo._entity_boost_max_terms = max_terms
return repo
def test_disabled_returns_scores_unchanged(self):
repo = self._repo(enabled=False)
key = ("entity", 1)
rows = {key: _make_index_row(row_id=1, title="Joanna")}
scores = {key: 0.5}
assert repo._apply_entity_boost(scores, rows, {"joanna"}) == {key: 0.5}
def test_matching_row_is_boosted(self):
repo = self._repo(enabled=True, weight=0.2)
key = ("entity", 1)
rows = {key: _make_index_row(row_id=1, title="Joanna")}
boosted = repo._apply_entity_boost({key: 0.5}, rows, {"joanna"})
assert boosted[key] == pytest.approx(0.5 * 1.2)
def test_non_matching_row_unchanged(self):
repo = self._repo(enabled=True, weight=0.2)
key = ("entity", 1)
rows = {key: _make_index_row(row_id=1, title="Anthony")}
boosted = repo._apply_entity_boost({key: 0.5}, rows, {"joanna"})
assert boosted[key] == pytest.approx(0.5)
def test_boost_can_reorder_lower_scored_match_above_higher_non_match(self):
repo = self._repo(enabled=True, weight=0.5)
generic_key = ("entity", 1)
joanna_key = ("entity", 2)
rows = {
generic_key: _make_index_row(row_id=1, title="Generic topic about hobbies"),
joanna_key: _make_index_row(row_id=2, title="Joanna"),
}
# The generic row starts higher (0.6) but does not match; "Joanna" (0.5) matches.
boosted = repo._apply_entity_boost({generic_key: 0.6, joanna_key: 0.5}, rows, {"joanna"})
assert boosted[joanna_key] > boosted[generic_key]
def test_multiple_matches_scale_boost(self):
repo = self._repo(enabled=True, weight=0.1, max_terms=3)
key = ("relation", 1)
rows = {key: _make_index_row(row_id=1, title="Deborah -> Jolene")}
boosted = repo._apply_entity_boost({key: 1.0}, rows, {"deborah", "jolene"})
# Two matched terms -> 1 + 0.1 * 2 = 1.2
assert boosted[key] == pytest.approx(1.2)
def test_max_terms_caps_the_boost(self):
repo = self._repo(enabled=True, weight=0.1, max_terms=1)
key = ("relation", 1)
rows = {key: _make_index_row(row_id=1, title="Deborah -> Jolene")}
boosted = repo._apply_entity_boost({key: 1.0}, rows, {"deborah", "jolene"})
# Capped at 1 term -> 1 + 0.1 * 1 = 1.1
assert boosted[key] == pytest.approx(1.1)
def test_zero_weight_is_noop(self):
repo = self._repo(enabled=True, weight=0.0)
key = ("entity", 1)
rows = {key: _make_index_row(row_id=1, title="Joanna")}
assert repo._apply_entity_boost({key: 0.5}, rows, {"joanna"}) == {key: 0.5}
def test_empty_entity_terms_is_noop(self):
repo = self._repo(enabled=True, weight=0.2)
key = ("entity", 1)
rows = {key: _make_index_row(row_id=1, title="Joanna")}
assert repo._apply_entity_boost({key: 0.5}, rows, set()) == {key: 0.5}