vectorize-io__hindsight/hindsight-api-slim/tests/test_graph_entity_fanout_cap.py

"""
Tests for LATERAL entity fanout cap in graph expansion.

Verifies that the per-entity LIMIT in _expand_combined prevents high-fanout
entities from exploding the self-join, while still returning entity-based
graph results.
"""

import asyncio
from datetime import datetime, timezone

import pytest

# These read `result.trace["retrieval_results"]` and assert an entry per ARM (method_name ==
# "graph"). That structure is produced by the engine running the arms itself; a store that answers
# a whole recall in one hop runs them internally and reports phases ("arms", "fuse", "trim"), not
# one entry per arm -- so there is nothing here to assert against for such a store, and the arm
# behaviour is covered by the extension's own recall suite instead.
pytestmark = pytest.mark.memory_backend_incompatible


@pytest.mark.asyncio
@pytest.mark.timeout(1200)
async def test_high_fanout_entity_returns_results(memory, request_context):
    """
    A high-fanout entity (appearing in many facts) should still produce
    graph retrieval results — the LATERAL cap limits rows per entity but
    does not drop the entity entirely.
    """
    bank_id = f"test_fanout_cap_{datetime.now(timezone.utc).timestamp()}"

    try:
        # Create many facts sharing one common entity ("Acme Corp") plus
        # a few with a unique entity so we can query for the unique one
        # and verify graph expansion finds siblings via "Acme Corp".
        contents = [
            # Target: unique entity "Zara" shares "Acme Corp" with the rest
            {
                "content": "Zara joined Acme Corp as a senior engineer last month",
                "context": "hr update",
                "entities": [{"text": "Zara"}, {"text": "Acme Corp"}],
            },
        ]
        # Add many facts that all share "Acme Corp" — creates a high-fanout entity
        for i in range(60):
            contents.append(
                {
                    "content": f"Employee {i} completed onboarding at Acme Corp in department {i % 5}",
                    "context": "hr update",
                    "entities": [{"text": f"Employee {i}"}, {"text": "Acme Corp"}],
                }
            )

        await memory.retain_batch_async(
            bank_id=bank_id,
            contents=contents,
            request_context=request_context,
        )

        from hindsight_api.engine.memory_engine import Budget

        # Query for "Zara" — semantic search finds Zara's fact as a seed,
        # then graph expansion should find other Acme Corp facts via the
        # shared entity, even though "Acme Corp" has 60+ mentions.
        result = await memory.recall_async(
            bank_id=bank_id,
            query="Zara",
            budget=Budget.HIGH,
            max_tokens=4096,
            enable_trace=True,
            request_context=request_context,
            _quiet=True,
        )

        assert result.results is not None
        assert len(result.results) > 0

        # Verify graph retrieval ran and found results
        retrieval_results = result.trace.get("retrieval_results", [])
        graph_results = [r for r in retrieval_results if r.get("method_name") == "graph"]
        assert len(graph_results) > 0, "Graph retrieval should have run"

        # At least one graph result should contain Acme Corp content
        # (found via shared entity, not just semantic similarity)
        all_texts = [r.text for r in result.results]
        acme_found = any("Acme Corp" in t for t in all_texts)
        assert acme_found, "Should find Acme Corp facts via entity graph expansion"

    finally:
        await memory.delete_bank(bank_id, request_context=request_context)


@pytest.mark.asyncio
async def test_entity_expansion_timeout_fallback(memory, request_context):
    """
    When graph_expansion_timeout is set very low, entity expansion should
    time out gracefully and fall back to semantic+causal links only,
    rather than failing the entire recall.
    """
    bank_id = f"test_timeout_fallback_{datetime.now(timezone.utc).timestamp()}"

    try:
        await memory.retain_batch_async(
            bank_id=bank_id,
            contents=[
                {
                    "content": "Alice works on the backend API at TechCorp",
                    "context": "team info",
                    "entities": [{"text": "Alice"}, {"text": "TechCorp"}],
                },
                {
                    "content": "Bob maintains the frontend at TechCorp",
                    "context": "team info",
                    "entities": [{"text": "Bob"}, {"text": "TechCorp"}],
                },
            ],
            request_context=request_context,
        )

        from hindsight_api.config import _get_raw_config
        from hindsight_api.engine.memory_engine import Budget

        config = _get_raw_config()
        original_timeout = config.link_expansion_timeout

        try:
            # Set an impossibly low timeout to force the fallback path
            config.link_expansion_timeout = 0.0001

            result = await memory.recall_async(
                bank_id=bank_id,
                query="Alice",
                budget=Budget.MID,
                max_tokens=2048,
                enable_trace=True,
                request_context=request_context,
                _quiet=True,
            )

            # Recall should succeed even when entity expansion times out
            assert result.results is not None
            assert len(result.results) > 0

            # Alice should still be found via semantic search
            result_texts = [r.text for r in result.results]
            alice_found = any("Alice" in t for t in result_texts)
            assert alice_found, "Should find Alice via semantic search despite graph timeout"
        finally:
            config.link_expansion_timeout = original_timeout

    finally:
        await memory.delete_bank(bank_id, request_context=request_context)


@pytest.mark.asyncio
@pytest.mark.timeout(1200)
async def test_per_entity_limit_caps_expansion(memory, request_context):
    """
    With graph_per_entity_limit set to a small value, entity expansion should
    still work but return fewer results from high-fanout entities.
    """
    bank_id = f"test_per_entity_limit_{datetime.now(timezone.utc).timestamp()}"

    try:
        # Create facts with a shared entity
        contents = [
            {
                "content": "Lead engineer Dana oversees the Widgets project at MegaCorp",
                "context": "project info",
                "entities": [{"text": "Dana"}, {"text": "MegaCorp"}],
            },
        ]
        for i in range(30):
            contents.append(
                {
                    "content": f"MegaCorp hired contractor {i} for the Q4 push",
                    "context": "hiring info",
                    "entities": [{"text": f"Contractor {i}"}, {"text": "MegaCorp"}],
                }
            )

        await memory.retain_batch_async(
            bank_id=bank_id,
            contents=contents,
            request_context=request_context,
        )

        from hindsight_api.config import _get_raw_config
        from hindsight_api.engine.memory_engine import Budget

        config = _get_raw_config()
        original_limit = config.link_expansion_per_entity_limit

        try:
            # Set a very small per-entity limit
            config.link_expansion_per_entity_limit = 5

            result = await memory.recall_async(
                bank_id=bank_id,
                query="Dana",
                budget=Budget.HIGH,
                max_tokens=4096,
                enable_trace=True,
                request_context=request_context,
                _quiet=True,
            )

            # Recall should succeed with the cap
            assert result.results is not None
            assert len(result.results) > 0

            # Graph retrieval should have run
            retrieval_results = result.trace.get("retrieval_results", [])
            graph_results = [r for r in retrieval_results if r.get("method_name") == "graph"]
            assert len(graph_results) > 0, "Graph retrieval should have run"
        finally:
            config.link_expansion_per_entity_limit = original_limit

    finally:
        await memory.delete_bank(bank_id, request_context=request_context)