fix: skip fuzzy entity resolution for user-defined label entities (#1558)

Entity resolution was merging distinct multivalue label entities (e.g.,
"use:use-001" and "use:use-002") because their high string similarity
(~0.91) combined with temporal proximity exceeded the 0.6 merge threshold.

Tags were stored correctly (direct string storage on memory_units) but
entity links in unit_entities only contained a subset because both values
resolved to the same entity ID.

Fix: when entity_labels are configured, label entities use exact
case-insensitive matching only — no fuzzy scoring. Their canonical names
are user-defined and must not be normalized.

hindsight-api-slim/hindsight_api/engine/entity_resolver.py

@@ -16,7 +16,15 @@ from typing import Any, Final
 
 from .db_utils import acquire_with_retry
 from .memory_engine import fq_table
-from .retain.entity_labels import build_labels_lookup as _build_labels_lookup_from_config
+from .retain.entity_labels import (
+    build_labels_lookup as _build_labels_lookup_from_config,
+)
+from .retain.entity_labels import (
+    is_label_entity as _is_label_entity,
+)
+from .retain.entity_labels import (
+    parse_entity_labels as _parse_entity_labels,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -228,14 +236,15 @@ class EntityResolver:
             return []
 
         taxonomy_lookup = self._build_labels_lookup(entity_labels)
+        labels_cfg = _parse_entity_labels(entity_labels)
         if conn is None:
             async with acquire_with_retry(self.pool) as conn:
                 return await self._resolve_entities_batch_impl(
-                    conn, bank_id, entities_data, context, unit_event_date, taxonomy_lookup
+                    conn, bank_id, entities_data, context, unit_event_date, taxonomy_lookup, labels_cfg
                 )
         else:
             return await self._resolve_entities_batch_impl(
-                conn, bank_id, entities_data, context, unit_event_date, taxonomy_lookup
+                conn, bank_id, entities_data, context, unit_event_date, taxonomy_lookup, labels_cfg
             )
 
     async def _resolve_entities_batch_impl(
@@ -246,13 +255,16 @@ class EntityResolver:
         context: str,
         unit_event_date,
         taxonomy_lookup: set[str] | None = None,
+        labels_cfg=None,
     ) -> list[str]:
         if self.entity_lookup == "trigram":
             # Route to backend-specific fuzzy strategy.
             # Non-PG backends (Oracle) use UTL_MATCH instead of pg_trgm.
             backend_strategy = self._ops.get_entity_resolution_strategy()
             if backend_strategy == "oracle_fuzzy":
-                return await self._resolve_entities_batch_oracle_fuzzy(conn, bank_id, entities_data, unit_event_date)
+                return await self._resolve_entities_batch_oracle_fuzzy(
+                    conn, bank_id, entities_data, unit_event_date, taxonomy_lookup, labels_cfg
+                )
             # Auto-detect pg_trgm availability on first call and fall back to
             # "full" strategy if the extension is not installed.  See #626.
             if not self._pg_trgm_checked:
@@ -266,12 +278,24 @@ class EntityResolver:
                         "https://github.com/vectorize-io/hindsight/issues/626"
                     )
                     self.entity_lookup = "full"
-                    return await self._resolve_entities_batch_full(conn, bank_id, entities_data, unit_event_date)
-            return await self._resolve_entities_batch_trigram(conn, bank_id, entities_data, unit_event_date)
-        return await self._resolve_entities_batch_full(conn, bank_id, entities_data, unit_event_date)
+                    return await self._resolve_entities_batch_full(
+                        conn, bank_id, entities_data, unit_event_date, taxonomy_lookup, labels_cfg
+                    )
+            return await self._resolve_entities_batch_trigram(
+                conn, bank_id, entities_data, unit_event_date, taxonomy_lookup, labels_cfg
+            )
+        return await self._resolve_entities_batch_full(
+            conn, bank_id, entities_data, unit_event_date, taxonomy_lookup, labels_cfg
+        )
 
     async def _resolve_entities_batch_full(
-        self, conn, bank_id: str, entities_data: list[dict], unit_event_date
+        self,
+        conn,
+        bank_id: str,
+        entities_data: list[dict],
+        unit_event_date,
+        taxonomy_lookup: set[str] | None = None,
+        labels_cfg=None,
     ) -> list[str]:
         """Original strategy: load all bank entities then match in Python."""
         # Query ALL candidates for this bank
@@ -338,11 +362,24 @@ class EntityResolver:
             all_candidates[entity_text] = matching
 
         return await self._resolve_from_candidates(
-            conn, bank_id, entities_data, unit_event_date, all_candidates, cooccurrence_map
+            conn,
+            bank_id,
+            entities_data,
+            unit_event_date,
+            all_candidates,
+            cooccurrence_map,
+            taxonomy_lookup,
+            labels_cfg,
         )
 
     async def _resolve_entities_batch_trigram(
-        self, conn, bank_id: str, entities_data: list[dict], unit_event_date
+        self,
+        conn,
+        bank_id: str,
+        entities_data: list[dict],
+        unit_event_date,
+        taxonomy_lookup: set[str] | None = None,
+        labels_cfg=None,
     ) -> list[str]:
         """
         Trigram strategy: fetch only similar candidates per entity name using pg_trgm.
@@ -418,11 +455,24 @@ class EntityResolver:
                     cooccurrence_map[eid2].add(id_to_name[eid1])
 
         return await self._resolve_from_candidates(
-            conn, bank_id, entities_data, unit_event_date, all_candidates, cooccurrence_map
+            conn,
+            bank_id,
+            entities_data,
+            unit_event_date,
+            all_candidates,
+            cooccurrence_map,
+            taxonomy_lookup,
+            labels_cfg,
         )
 
     async def _resolve_entities_batch_oracle_fuzzy(
-        self, conn: Any, bank_id: str, entities_data: list[dict], unit_event_date: datetime | None
+        self,
+        conn: Any,
+        bank_id: str,
+        entities_data: list[dict],
+        unit_event_date: datetime | None,
+        taxonomy_lookup: set[str] | None = None,
+        labels_cfg=None,
     ) -> list[str]:
         """
         Oracle strategy: fetch similar candidates using UTL_MATCH.JARO_WINKLER_SIMILARITY.
@@ -506,7 +556,14 @@ class EntityResolver:
                     cooccurrence_map[eid2].add(id_to_name[eid1])
 
         return await self._resolve_from_candidates(
-            conn, bank_id, entities_data, unit_event_date, all_candidates, cooccurrence_map
+            conn,
+            bank_id,
+            entities_data,
+            unit_event_date,
+            all_candidates,
+            cooccurrence_map,
+            taxonomy_lookup,
+            labels_cfg,
         )
 
     async def _resolve_from_candidates(
@@ -517,6 +574,8 @@ class EntityResolver:
         unit_event_date,
         all_candidates: dict[str, list],
         cooccurrence_map: dict[str, set[str]],
+        taxonomy_lookup: set[str] | None = None,
+        labels_cfg=None,
     ) -> list[str]:
         """Shared scoring + upsert logic used by both lookup strategies."""
 
@@ -533,11 +592,34 @@ class EntityResolver:
 
             candidates = all_candidates.get(entity_text, [])
 
+            # Label entities (from entity_labels config) use exact matching only.
+            # Their canonical names are user-defined (e.g., "use:use-001"),
+            # so fuzzy resolution must NOT merge distinct label values that
+            # happen to be textually similar (GH-1558).
+            is_label = bool(
+                labels_cfg and taxonomy_lookup and _is_label_entity(entity_text, labels_cfg, taxonomy_lookup)
+            )
+
             if not candidates:
                 # Will create new entity
                 entities_to_create.append(_EntityToCreate(idx=idx, name=entity_text, event_date=entity_event_date))
                 continue
 
+            if is_label:
+                # Exact case-insensitive match only for label entities
+                exact_match = None
+                entity_text_lower = entity_text.lower()
+                for candidate_id, canonical_name, metadata, last_seen, mention_count in candidates:
+                    if canonical_name.lower() == entity_text_lower:
+                        exact_match = candidate_id
+                        break
+                if exact_match:
+                    entity_ids[idx] = exact_match
+                    entities_to_update.append(_EntityStat(entity_id=exact_match, event_date=entity_event_date))
+                else:
+                    entities_to_create.append(_EntityToCreate(idx=idx, name=entity_text, event_date=entity_event_date))
+                continue
+
             # Score candidates
             best_candidate = None
             best_score = 0.0

hindsight-api-slim/tests/test_entity_labels.py

@@ -1894,3 +1894,354 @@ def test_duplicate_entity_strings_deduplicated():
 
     texts = [e.text for e in validated]
     assert texts == ["person:name:Alice"]  # only once
+
+
+# ─── GH-1558: multivalue tag entities missing from unit_entities ────────────
+
+
+def test_inject_label_tags_multivalue_all_tags_added():
+    """GH-1558 reproducer (unit-level): all multivalue entities with tag=True end up in tags."""
+    from unittest.mock import MagicMock
+
+    from hindsight_api.engine.retain.fact_extraction import _inject_label_tags
+    from hindsight_api.engine.retain.types import ExtractedFact
+
+    config = MagicMock()
+    config.entity_labels = [
+        {
+            "key": "use",
+            "type": "multi-values",
+            "tag": True,
+            "values": [
+                {"value": "use-001"},
+                {"value": "use-002"},
+                {"value": "use-003"},
+            ],
+        },
+    ]
+
+    fact = ExtractedFact(
+        fact_text="System references use-001 and use-002",
+        fact_type="world",
+        entities=["use:use-001", "use:use-002"],
+        tags=[],
+    )
+    _inject_label_tags([fact], config)
+
+    # Both label entities should be present in tags
+    assert "use:use-001" in fact.tags
+    assert "use:use-002" in fact.tags
+    assert len(fact.tags) == 2
+
+
+@pytest.mark.asyncio
+async def test_retain_multivalue_tag_entities_all_stored(memory, request_context):
+    """
+    GH-1558 reproducer (integration): retain content referencing multiple values
+    of a multi-values entity label with tag=True.
+
+    Verify that ALL multivalue entities appear in BOTH:
+    - memory_units.tags (the tags column)
+    - unit_entities table (the entity links)
+
+    The original bug: tags are added correctly, but unit_entities only stores
+    a subset (typically the first entity).
+    """
+    from hindsight_api.engine.memory_engine import fq_table
+
+    bank_id = f"test-1558-multivalue-tag-{uuid.uuid4().hex[:8]}"
+    try:
+        await memory.get_bank_profile(bank_id=bank_id, request_context=request_context)
+
+        # Configure entity labels matching the bug report scenario:
+        # - multi-values type
+        # - tag=True
+        # - entities_allow_free_form=False
+        await memory._config_resolver.update_bank_config(
+            bank_id=bank_id,
+            updates={
+                "entity_labels": [
+                    {
+                        "key": "use",
+                        "description": "Use case identifier for this section",
+                        "type": "multi-values",
+                        "tag": True,
+                        "values": [
+                            {"value": "use-001", "description": "First use case"},
+                            {"value": "use-002", "description": "Second use case"},
+                            {"value": "use-003", "description": "Third use case"},
+                        ],
+                    }
+                ],
+                "entities_allow_free_form": False,
+                "retain_extraction_mode": "verbose",
+            },
+            context=request_context,
+        )
+
+        # Content that explicitly references multiple use case identifiers
+        # in a way that a single fact should capture both
+        unit_ids = await memory.retain_async(
+            bank_id=bank_id,
+            content=(
+                "## System Integration Notes (use-001, use-002)\n\n"
+                "This section covers both use-001 and use-002 use cases. "
+                "The integration between use-001 (authentication flow) and "
+                "use-002 (authorization flow) requires careful coordination. "
+                "Both use-001 and use-002 must be tested together."
+            ),
+            request_context=request_context,
+        )
+
+        assert len(unit_ids) > 0, "Should have extracted at least one fact"
+
+        async with memory._pool.acquire() as conn:
+            # Check entities in unit_entities table
+            entity_rows = await conn.fetch(
+                f"""
+                SELECT e.canonical_name
+                FROM {fq_table("unit_entities")} ue
+                JOIN {fq_table("entities")} e ON e.id = ue.entity_id
+                WHERE ue.unit_id = ANY($1::uuid[])
+                """,
+                [u for u in unit_ids],
+            )
+            entity_names = {r["canonical_name"].lower() for r in entity_rows}
+
+            # Check tags on memory_units
+            tag_rows = await conn.fetch(
+                f"""
+                SELECT id, tags
+                FROM {fq_table("memory_units")}
+                WHERE id = ANY($1::uuid[])
+                """,
+                [u for u in unit_ids],
+            )
+            all_tags = set()
+            for row in tag_rows:
+                if row["tags"]:
+                    all_tags.update(t.lower() for t in row["tags"])
+
+        # Filter to use:* entities/tags
+        use_entities = {n for n in entity_names if n.startswith("use:")}
+        use_tags = {t for t in all_tags if t.startswith("use:")}
+
+        # The core assertion from GH-1558: tags and entities should match
+        # Tags show both but entities only show a subset → BUG
+        assert len(use_tags) >= 2, (
+            f"Expected at least 2 use:* tags. Got: {use_tags}"
+        )
+        assert len(use_entities) >= 2, (
+            f"GH-1558 BUG: Expected at least 2 use:* entities in unit_entities, "
+            f"but only got {len(use_entities)}: {use_entities}. "
+            f"Tags correctly show: {use_tags}"
+        )
+        # Every tag should also be an entity
+        missing_entities = use_tags - use_entities
+        assert len(missing_entities) == 0, (
+            f"GH-1558 BUG: Tags {use_tags} were added but entities are missing: {missing_entities}. "
+            f"Entities found: {use_entities}"
+        )
+    finally:
+        await memory.delete_bank(bank_id, request_context=request_context)
+
+
+@pytest.mark.asyncio
+async def test_retain_multivalue_tag_entities_second_retain(memory, request_context):
+    """
+    GH-1558 reproducer (second retain): entity resolution with existing entities.
+
+    On a second retain, entity resolution tries to match new entity names against
+    existing entities in the bank. With very similar names like "use:use-001" and
+    "use:use-002", the SequenceMatcher similarity is ~0.91 which combined with
+    temporal proximity could exceed the 0.6 merge threshold, causing both to
+    resolve to the same entity ID.
+    """
+    from hindsight_api.engine.memory_engine import fq_table
+
+    bank_id = f"test-1558-second-{uuid.uuid4().hex[:8]}"
+    try:
+        await memory.get_bank_profile(bank_id=bank_id, request_context=request_context)
+
+        await memory._config_resolver.update_bank_config(
+            bank_id=bank_id,
+            updates={
+                "entity_labels": [
+                    {
+                        "key": "use",
+                        "description": "Use case identifier",
+                        "type": "multi-values",
+                        "tag": True,
+                        "values": [
+                            {"value": "use-001", "description": "First use case"},
+                            {"value": "use-002", "description": "Second use case"},
+                        ],
+                    }
+                ],
+                "entities_allow_free_form": False,
+                "retain_extraction_mode": "verbose",
+            },
+            context=request_context,
+        )
+
+        # First retain: creates entities in the bank
+        await memory.retain_async(
+            bank_id=bank_id,
+            content=(
+                "## Authentication Flow (use-001)\n\n"
+                "The authentication flow use-001 handles user login via OAuth2."
+            ),
+            request_context=request_context,
+        )
+
+        # Second retain: references BOTH use-001 and use-002
+        # Entity resolution now has existing entities to match against
+        unit_ids_2 = await memory.retain_async(
+            bank_id=bank_id,
+            content=(
+                "## Integration Notes (use-001, use-002)\n\n"
+                "This section covers the integration between use-001 (authentication) "
+                "and use-002 (authorization). Both use-001 and use-002 are required."
+            ),
+            request_context=request_context,
+        )
+
+        assert len(unit_ids_2) > 0
+
+        async with memory._pool.acquire() as conn:
+            entity_rows = await conn.fetch(
+                f"""
+                SELECT e.canonical_name
+                FROM {fq_table("unit_entities")} ue
+                JOIN {fq_table("entities")} e ON e.id = ue.entity_id
+                WHERE ue.unit_id = ANY($1::uuid[])
+                """,
+                [u for u in unit_ids_2],
+            )
+            entity_names = {r["canonical_name"].lower() for r in entity_rows}
+
+            tag_rows = await conn.fetch(
+                f"""
+                SELECT id, tags
+                FROM {fq_table("memory_units")}
+                WHERE id = ANY($1::uuid[])
+                """,
+                [u for u in unit_ids_2],
+            )
+            all_tags = set()
+            for row in tag_rows:
+                if row["tags"]:
+                    all_tags.update(t.lower() for t in row["tags"])
+
+        use_entities = {n for n in entity_names if n.startswith("use:")}
+        use_tags = {t for t in all_tags if t.startswith("use:")}
+
+        assert len(use_tags) >= 2, (
+            f"Expected at least 2 use:* tags on second retain. Got: {use_tags}"
+        )
+        assert len(use_entities) >= 2, (
+            f"GH-1558 BUG: On second retain, expected at least 2 use:* entities "
+            f"but only got {len(use_entities)}: {use_entities}. "
+            f"Tags correctly show: {use_tags}. "
+            f"Entity resolution may be merging similar names."
+        )
+        missing = use_tags - use_entities
+        assert len(missing) == 0, (
+            f"GH-1558 BUG: Tags present but entities missing after second retain: {missing}"
+        )
+    finally:
+        await memory.delete_bank(bank_id, request_context=request_context)
+
+
+@pytest.mark.asyncio
+async def test_entity_resolution_does_not_merge_distinct_label_values(memory, request_context):
+    """
+    GH-1558 reproducer (deterministic): directly test that entity resolution
+    keeps distinct label values separate even when their names are very similar.
+
+    "use:use-001" and "use:use-002" have SequenceMatcher similarity of ~0.91.
+    With the 0.6 merge threshold and temporal/co-occurrence boosts, the resolver
+    might incorrectly merge them into a single entity.
+    """
+    from hindsight_api.engine.memory_engine import fq_table
+    from hindsight_api.engine.retain.entity_processing import resolve_entities
+    from hindsight_api.engine.retain.types import EntityRef, ProcessedFact
+
+    bank_id = f"test-1558-resolve-{uuid.uuid4().hex[:8]}"
+    try:
+        await memory.get_bank_profile(bank_id=bank_id, request_context=request_context)
+
+        # First, insert a "use:use-001" entity into the bank so that
+        # entity resolution has an existing entity to match against
+        async with memory._pool.acquire() as conn:
+            await conn.execute(
+                f"""
+                INSERT INTO {fq_table("entities")} (bank_id, canonical_name, first_seen, last_seen, mention_count)
+                VALUES ($1, $2, now(), now(), 1)
+                ON CONFLICT DO NOTHING
+                """,
+                bank_id,
+                "use:use-001",
+            )
+
+        # Now resolve entities for a fact that has BOTH use:use-001 and use:use-002
+        from datetime import datetime, timezone
+
+        now = datetime.now(timezone.utc)
+        facts = [
+            ProcessedFact(
+                fact_text="Integration between use-001 and use-002",
+                fact_type="world",
+                embedding=[0.0] * 384,
+                occurred_start=now,
+                occurred_end=None,
+                mentioned_at=now,
+                context="",
+                metadata={},
+                entities=[
+                    EntityRef(name="use:use-001"),
+                    EntityRef(name="use:use-002"),
+                ],
+                content_index=0,
+                tags=["use:use-001", "use:use-002"],
+            )
+        ]
+
+        # Use placeholder unit IDs
+        placeholder_unit_ids = [str(uuid.uuid4())]
+
+        entity_labels = [
+            {
+                "key": "use",
+                "description": "Use case identifier",
+                "type": "multi-values",
+                "tag": True,
+                "values": [
+                    {"value": "use-001"},
+                    {"value": "use-002"},
+                ],
+            }
+        ]
+
+        async with memory._pool.acquire() as conn:
+            resolved_entity_ids, entity_to_unit, unit_to_entity_ids = await resolve_entities(
+                entity_resolver=memory.entity_resolver,
+                conn=conn,
+                bank_id=bank_id,
+                unit_ids=placeholder_unit_ids,
+                facts=facts,
+                entity_labels=entity_labels,
+            )
+
+        # We should get 2 DISTINCT entity IDs, not the same ID twice
+        assert len(resolved_entity_ids) == 2, (
+            f"Expected 2 resolved entity IDs, got {len(resolved_entity_ids)}"
+        )
+        unique_ids = set(resolved_entity_ids)
+        assert len(unique_ids) == 2, (
+            f"GH-1558 BUG: Entity resolution merged 'use:use-001' and 'use:use-002' "
+            f"into the same entity ID. Got IDs: {resolved_entity_ids}. "
+            f"These are distinct label values and must NOT be merged."
+        )
+    finally:
+        await memory.delete_bank(bank_id, request_context=request_context)