fleet-memory/hindsight-api/tests/test_mental_models.py

"""Tests for mental model functionality (v4 system)."""

import uuid

import pytest

from hindsight_api.engine.memory_engine import MemoryEngine


@pytest.fixture
async def memory_with_mission(memory: MemoryEngine, request_context):
    """Memory engine with a bank that has a mission set.

    Uses a unique bank_id to avoid conflicts between parallel tests.
    """
    # Use unique bank_id to avoid conflicts between parallel tests
    bank_id = f"test-mental-models-{uuid.uuid4().hex[:8]}"

    # Set up the bank with a mission
    await memory.set_bank_mission(
        bank_id=bank_id,
        mission="Be a PM for the engineering team",
        request_context=request_context,
    )

    # Add some test data
    await memory.retain_batch_async(
        bank_id=bank_id,
        contents=[
            {"content": "The team has daily standups at 9am where everyone shares their progress."},
            {"content": "Alice is the frontend engineer and specializes in React."},
            {"content": "Bob is the backend engineer and owns the API services."},
            {"content": "Sprint retrospectives happen every two weeks to discuss improvements."},
            {"content": "John is the tech lead and makes final decisions on architecture."},
        ],
        request_context=request_context,
    )

    # Wait for any background tasks from retain to complete
    await memory.wait_for_background_tasks()

    yield memory, bank_id

    # Cleanup
    await memory.delete_bank(bank_id, request_context=request_context)


class TestBankMission:
    """Test bank mission operations."""

    async def test_set_and_get_mission(self, memory: MemoryEngine, request_context):
        """Test setting and getting a bank's mission."""
        bank_id = f"test-mission-{uuid.uuid4().hex[:8]}"

        # Set mission
        result = await memory.set_bank_mission(
            bank_id=bank_id,
            mission="Track customer feedback",
            request_context=request_context,
        )

        assert result["bank_id"] == bank_id
        assert result["mission"] == "Track customer feedback"

        # Get mission via profile
        profile = await memory.get_bank_profile(bank_id=bank_id, request_context=request_context)
        assert profile["mission"] == "Track customer feedback"

        # Cleanup
        await memory.delete_bank(bank_id, request_context=request_context)


class TestRefreshMentalModels:
    """Test the main refresh_mental_models flow."""

    async def test_refresh_creates_structural_models(self, memory_with_mission, request_context):
        """Test that refresh creates structural models from the mission."""
        memory, bank_id = memory_with_mission

        # Refresh mental models (async - returns operation_id)
        result = await memory.refresh_mental_models(
            bank_id=bank_id,
            request_context=request_context,
        )

        # Check that we got an operation ID back
        assert "operation_id" in result
        assert result["status"] == "queued"

        # Wait for background task to complete
        await memory.wait_for_background_tasks()

        # Get the created models
        models = await memory.list_mental_models(
            bank_id=bank_id,
            request_context=request_context,
        )

        assert len(models) > 0

        # Check that structural models were created
        structural_models = [m for m in models if m["subtype"] == "structural"]
        assert len(structural_models) > 0

        # Check that models have the expected structure
        for model in models:
            assert "id" in model
            assert "name" in model
            assert "description" in model
            assert model["subtype"] in ["structural", "emergent"]

    async def test_refresh_without_mission_fails(self, memory: MemoryEngine, request_context):
        """Test that refresh fails when no mission is set."""
        bank_id = f"test-no-mission-refresh-{uuid.uuid4().hex[:8]}"

        # Add some data but don't set a mission
        await memory.retain_batch_async(
            bank_id=bank_id,
            contents=[
                {"content": "Alice is the frontend engineer."},
                {"content": "Bob is the backend engineer."},
            ],
            request_context=request_context,
        )

        # Wait for any background tasks from retain to complete
        await memory.wait_for_background_tasks()

        # Refresh mental models should fail without a mission
        with pytest.raises(ValueError) as exc_info:
            await memory.refresh_mental_models(
                bank_id=bank_id,
                request_context=request_context,
            )

        assert "no mission is set" in str(exc_info.value).lower()

        # Cleanup
        await memory.delete_bank(bank_id, request_context=request_context)


class TestMentalModelCRUD:
    """Test basic CRUD operations for mental models."""

    async def test_list_mental_models(self, memory_with_mission, request_context):
        """Test listing mental models."""
        memory, bank_id = memory_with_mission

        # Refresh to create models (async)
        await memory.refresh_mental_models(
            bank_id=bank_id,
            request_context=request_context,
        )
        await memory.wait_for_background_tasks()

        # List all models
        models = await memory.list_mental_models(
            bank_id=bank_id,
            request_context=request_context,
        )

        assert len(models) > 0

        # Test filtering by subtype
        structural_models = await memory.list_mental_models(
            bank_id=bank_id,
            subtype="structural",
            request_context=request_context,
        )

        assert all(m["subtype"] == "structural" for m in structural_models)

    async def test_get_mental_model(self, memory_with_mission, request_context):
        """Test getting a mental model by ID."""
        memory, bank_id = memory_with_mission

        # Refresh to create models (async)
        await memory.refresh_mental_models(
            bank_id=bank_id,
            request_context=request_context,
        )
        await memory.wait_for_background_tasks()

        # Get the created models
        models = await memory.list_mental_models(
            bank_id=bank_id,
            request_context=request_context,
        )

        # Get one by ID
        model_id = models[0]["id"]
        model = await memory.get_mental_model(
            bank_id=bank_id,
            model_id=model_id,
            request_context=request_context,
        )

        assert model is not None
        assert model["id"] == model_id

        # Test non-existent
        not_found = await memory.get_mental_model(
            bank_id=bank_id,
            model_id="non-existent",
            request_context=request_context,
        )
        assert not_found is None

    async def test_delete_mental_model(self, memory_with_mission, request_context):
        """Test deleting a mental model."""
        memory, bank_id = memory_with_mission

        # Refresh to create models (async)
        await memory.refresh_mental_models(
            bank_id=bank_id,
            request_context=request_context,
        )
        await memory.wait_for_background_tasks()

        # Get the created models
        models = await memory.list_mental_models(
            bank_id=bank_id,
            request_context=request_context,
        )

        # Delete one
        model_id = models[0]["id"]
        deleted = await memory.delete_mental_model(
            bank_id=bank_id,
            model_id=model_id,
            request_context=request_context,
        )
        assert deleted is True

        # Verify it's gone
        model = await memory.get_mental_model(
            bank_id=bank_id,
            model_id=model_id,
            request_context=request_context,
        )
        assert model is None

        # Delete non-existent returns False
        deleted_again = await memory.delete_mental_model(
            bank_id=bank_id,
            model_id=model_id,
            request_context=request_context,
        )
        assert deleted_again is False

    async def test_create_pinned_mental_model(self, memory: MemoryEngine, request_context):
        """Test creating a pinned mental model."""
        bank_id = f"test-pinned-{uuid.uuid4().hex[:8]}"

        # Ensure bank exists by getting its profile (auto-creates if needed)
        await memory.get_bank_profile(bank_id, request_context=request_context)

        # Create a pinned mental model
        model = await memory.create_mental_model(
            bank_id=bank_id,
            name="Product Roadmap",
            description="Key product priorities and upcoming features",
            tags=["project-x"],
            request_context=request_context,
        )

        assert model["name"] == "Product Roadmap"
        assert model["description"] == "Key product priorities and upcoming features"
        assert model["subtype"] == "pinned"
        assert model["tags"] == ["project-x"]
        assert model["id"] == "pinned-product-roadmap"

        # Verify it can be retrieved
        retrieved = await memory.get_mental_model(
            bank_id=bank_id,
            model_id=model["id"],
            request_context=request_context,
        )
        assert retrieved is not None
        assert retrieved["subtype"] == "pinned"

        # Cleanup
        await memory.delete_bank(bank_id, request_context=request_context)

    async def test_create_pinned_model_duplicate_fails(self, memory: MemoryEngine, request_context):
        """Test that creating a duplicate pinned model fails."""
        bank_id = f"test-pinned-dup-{uuid.uuid4().hex[:8]}"

        # Ensure bank exists
        await memory.get_bank_profile(bank_id, request_context=request_context)

        # Create first model
        await memory.create_mental_model(
            bank_id=bank_id,
            name="Test Model",
            description="First model",
            request_context=request_context,
        )

        # Try to create duplicate
        with pytest.raises(ValueError) as exc_info:
            await memory.create_mental_model(
                bank_id=bank_id,
                name="Test Model",
                description="Second model",
                request_context=request_context,
            )

        assert "already exists" in str(exc_info.value).lower()

        # Cleanup
        await memory.delete_bank(bank_id, request_context=request_context)

    async def test_pinned_models_survive_refresh(self, memory: MemoryEngine, request_context):
        """Test that pinned models are not deleted during refresh."""
        bank_id = f"test-pinned-refresh-{uuid.uuid4().hex[:8]}"

        # Set a mission
        await memory.set_bank_mission(
            bank_id=bank_id,
            mission="Track customer feedback",
            request_context=request_context,
        )

        # Create a pinned model
        pinned_model = await memory.create_mental_model(
            bank_id=bank_id,
            name="Key Customers",
            description="Important customers to track",
            request_context=request_context,
        )

        # Refresh mental models
        await memory.refresh_mental_models(
            bank_id=bank_id,
            request_context=request_context,
        )
        await memory.wait_for_background_tasks()

        # Verify pinned model still exists
        retrieved = await memory.get_mental_model(
            bank_id=bank_id,
            model_id=pinned_model["id"],
            request_context=request_context,
        )
        assert retrieved is not None
        assert retrieved["subtype"] == "pinned"
        assert retrieved["name"] == "Key Customers"

        # Cleanup
        await memory.delete_bank(bank_id, request_context=request_context)


class TestMentalModelRefresh:
    """Test mental model summary refresh functionality."""

    async def test_refresh_creates_models_with_summaries(self, memory_with_mission, request_context):
        """Test that refresh_mental_models creates models and generates summaries."""
        memory, bank_id = memory_with_mission

        # Refresh mental models (async - creates models and generates summaries)
        result = await memory.refresh_mental_models(
            bank_id=bank_id,
            request_context=request_context,
        )

        assert "operation_id" in result
        assert result["status"] == "queued"

        # Wait for background task to complete (includes summary generation)
        await memory.wait_for_background_tasks()

        # Get the created models
        models = await memory.list_mental_models(
            bank_id=bank_id,
            request_context=request_context,
        )

        assert len(models) > 0

        # After async refresh completes, models should have summaries generated
        for model in models:
            assert "id" in model
            assert "name" in model
            # Summaries should be generated now (unless no relevant facts found)
            # We don't strictly assert on summary presence since it depends on data

    async def test_refresh_nonexistent_mental_model(self, memory: MemoryEngine, request_context):
        """Test refreshing a non-existent mental model returns None."""
        bank_id = f"test-refresh-noexist-{uuid.uuid4().hex[:8]}"

        result = await memory.refresh_mental_model(
            bank_id=bank_id,
            model_id="does-not-exist",
            request_context=request_context,
        )

        assert result is None


class TestReflect:
    """Test reflect endpoint with mental models."""

    async def test_reflect_basic(self, memory_with_mission, request_context):
        """Test basic reflect query - reflect works even without mental models."""
        memory, bank_id = memory_with_mission

        # Run a reflect query
        result = await memory.reflect_async(
            bank_id=bank_id,
            query="Who are the team members?",
            request_context=request_context,
        )

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


class TestMentalModelLearnTool:
    """Test mental model learn tool - creates placeholders with background generation."""

    async def test_learn_creates_placeholder(self, memory: MemoryEngine, request_context):
        """Test that learn tool creates a placeholder mental model without observations."""
        bank_id = f"test-source-facts-{uuid.uuid4().hex[:8]}"

        # Add some test data
        await memory.retain_batch_async(
            bank_id=bank_id,
            contents=[
                {"content": "Alice is the team lead."},
                {"content": "Bob is the engineer."},
            ],
            request_context=request_context,
        )
        await memory.wait_for_background_tasks()

        # Directly use the learn tool to create a mental model placeholder
        from hindsight_api.engine.reflect.models import MentalModelInput
        from hindsight_api.engine.reflect.tools import tool_learn

        input_model = MentalModelInput(
            name="Team Members",
            description="Key team members and their roles",
        )

        pool = await memory._get_pool()
        async with pool.acquire() as conn:
            result = await tool_learn(conn, bank_id, input_model)

        assert result["status"] == "created"
        assert result["model_id"] == "team-members"
        assert result["name"] == "Team Members"
        assert result["pending_generation"] is True

        # Verify placeholder was stored in database with empty observations
        pool = await memory._get_pool()
        async with pool.acquire() as conn:
            row = await conn.fetchrow(
                "SELECT subtype, name, description, observations FROM mental_models WHERE id = $1 AND bank_id = $2",
                result["model_id"],
                bank_id,
            )

        assert row is not None
        assert row["subtype"] == "learned"
        assert row["name"] == "Team Members"
        assert row["description"] == "Key team members and their roles"
        # Observations should be empty - will be generated in background
        observations_data = row["observations"]
        # Handle both string and dict representations
        if isinstance(observations_data, str):
            import json
            observations_data = json.loads(observations_data) if observations_data else {}
        assert observations_data == {} or observations_data is None

        # Cleanup
        await memory.delete_bank(bank_id, request_context=request_context)

    async def test_learn_update_description(self, memory: MemoryEngine, request_context):
        """Test that updating a mental model updates the description."""
        bank_id = f"test-merge-facts-{uuid.uuid4().hex[:8]}"

        # Create bank by retaining some data
        await memory.retain_batch_async(
            bank_id=bank_id,
            contents=[{"content": "Test data"}],
            request_context=request_context,
        )
        await memory.wait_for_background_tasks()

        from hindsight_api.engine.reflect.models import MentalModelInput
        from hindsight_api.engine.reflect.tools import tool_learn

        # First create a placeholder
        input_model = MentalModelInput(
            name="Team Members",
            description="Initial description",
        )

        pool = await memory._get_pool()
        async with pool.acquire() as conn:
            result1 = await tool_learn(conn, bank_id, input_model)

        assert result1["status"] == "created"
        assert result1["pending_generation"] is True

        # Now update with new description
        input_model2 = MentalModelInput(
            name="Team Members",  # Same name = same ID
            description="Updated description with more context",
        )

        pool = await memory._get_pool()
        async with pool.acquire() as conn:
            result2 = await tool_learn(conn, bank_id, input_model2)

        assert result2["status"] == "updated"
        assert result2["model_id"] == "team-members"

        # Verify description was updated in database
        pool = await memory._get_pool()
        async with pool.acquire() as conn:
            row = await conn.fetchrow(
                "SELECT description FROM mental_models WHERE id = $1 AND bank_id = $2",
                result1["model_id"],
                bank_id,
            )

        assert row is not None
        assert row["description"] == "Updated description with more context"

        # Cleanup
        await memory.delete_bank(bank_id, request_context=request_context)


class TestMentalModelTags:
    """Test mental model tags functionality."""

    @pytest.fixture
    async def memory_with_mission_and_tags(self, memory: MemoryEngine, request_context):
        """Memory engine with a bank that has a mission set and tagged content."""
        bank_id = f"test-mm-tags-{uuid.uuid4().hex[:8]}"

        # Set up the bank with a mission
        await memory.set_bank_mission(
            bank_id=bank_id,
            mission="Be a PM for the engineering team",
            request_context=request_context,
        )

        # Add some test data
        await memory.retain_batch_async(
            bank_id=bank_id,
            contents=[
                {"content": "Alice is the frontend engineer."},
                {"content": "Bob is the backend engineer."},
            ],
            request_context=request_context,
        )
        await memory.wait_for_background_tasks()

        yield memory, bank_id

        # Cleanup
        await memory.delete_bank(bank_id, request_context=request_context)

    async def test_refresh_creates_models_with_tags(self, memory_with_mission_and_tags, request_context):
        """Test that refresh_mental_models creates models with specified tags."""
        memory, bank_id = memory_with_mission_and_tags

        # Refresh mental models with tags
        result = await memory.refresh_mental_models(
            bank_id=bank_id,
            tags=["project-alpha", "sprint-1"],
            request_context=request_context,
        )

        assert "operation_id" in result
        await memory.wait_for_background_tasks()

        # Get the created models
        models = await memory.list_mental_models(
            bank_id=bank_id,
            request_context=request_context,
        )

        assert len(models) > 0

        # All models should have the tags we specified
        for model in models:
            assert "tags" in model
            assert "project-alpha" in model["tags"]
            assert "sprint-1" in model["tags"]

    async def test_list_mental_models_filters_by_tags(self, memory_with_mission_and_tags, request_context):
        """Test that list_mental_models correctly filters by tags."""
        memory, bank_id = memory_with_mission_and_tags

        # Create models with different tags
        await memory.refresh_mental_models(
            bank_id=bank_id,
            tags=["project-alpha"],
            request_context=request_context,
        )
        await memory.wait_for_background_tasks()

        # Get all models
        all_models = await memory.list_mental_models(
            bank_id=bank_id,
            request_context=request_context,
        )
        assert len(all_models) > 0

        # Filter by tags - should return models with matching tags
        filtered_models = await memory.list_mental_models(
            bank_id=bank_id,
            tags=["project-alpha"],
            request_context=request_context,
        )
        assert len(filtered_models) == len(all_models)  # All models have this tag

        # Filter by non-existent tag - should only return untagged models (none here)
        # But since all models have tags, and the filter includes untagged,
        # we need to test with a mix
        empty_filtered = await memory.list_mental_models(
            bank_id=bank_id,
            tags=["non-existent-tag"],
            request_context=request_context,
        )
        # Should return empty since no models are untagged and none match
        # Actually, the logic includes untagged models, so let's verify the behavior
        # All our models have tags, so only checking for non-existent tag
        # should return nothing (since none match and none are untagged)

    async def test_untagged_models_included_in_filter(self, memory: MemoryEngine, request_context):
        """Test that untagged mental models are always included when filtering."""
        bank_id = f"test-untagged-{uuid.uuid4().hex[:8]}"

        # Set up bank with mission
        await memory.set_bank_mission(
            bank_id=bank_id,
            mission="Track projects",
            request_context=request_context,
        )

        # Add some data
        await memory.retain_batch_async(
            bank_id=bank_id,
            contents=[{"content": "Project Alpha is important."}],
            request_context=request_context,
        )
        await memory.wait_for_background_tasks()

        # First refresh without tags (creates untagged models)
        await memory.refresh_mental_models(
            bank_id=bank_id,
            request_context=request_context,  # No tags
        )
        await memory.wait_for_background_tasks()

        # Get all models (should be untagged)
        all_models = await memory.list_mental_models(
            bank_id=bank_id,
            request_context=request_context,
        )

        if len(all_models) > 0:
            # Verify models are untagged
            for model in all_models:
                assert model.get("tags", []) == []

            # Filter by any tag - untagged models should still be included
            filtered_models = await memory.list_mental_models(
                bank_id=bank_id,
                tags=["some-tag"],
                request_context=request_context,
            )
            # Untagged models should be included in the results
            assert len(filtered_models) == len(all_models)

        # Cleanup
        await memory.delete_bank(bank_id, request_context=request_context)

    async def test_tags_match_any(self, memory: MemoryEngine, request_context):
        """Test tags_match='any' returns models with at least one matching tag."""
        bank_id = f"test-tags-any-{uuid.uuid4().hex[:8]}"

        # Set up bank with mission
        await memory.set_bank_mission(
            bank_id=bank_id,
            mission="Track projects",
            request_context=request_context,
        )

        # Add data and create models with tags
        await memory.retain_batch_async(
            bank_id=bank_id,
            contents=[{"content": "Alice works on frontend."}],
            request_context=request_context,
        )
        await memory.wait_for_background_tasks()

        await memory.refresh_mental_models(
            bank_id=bank_id,
            tags=["tag-a", "tag-b"],
            request_context=request_context,
        )
        await memory.wait_for_background_tasks()

        # Filter with tags_match='any' - should match if any tag matches
        models = await memory.list_mental_models(
            bank_id=bank_id,
            tags=["tag-a", "tag-c"],  # tag-a matches, tag-c doesn't
            tags_match="any",
            request_context=request_context,
        )

        # Models with tag-a should be included
        for model in models:
            if model.get("tags"):
                # At least one of the filter tags should be in the model tags
                # OR model is untagged
                assert (
                    any(t in model["tags"] for t in ["tag-a", "tag-c"])
                    or model["tags"] == []
                )

        # Cleanup
        await memory.delete_bank(bank_id, request_context=request_context)

    async def test_reflect_with_tags_filter(self, memory_with_mission_and_tags, request_context):
        """Test that reflect filters memories by tags."""
        memory, bank_id = memory_with_mission_and_tags

        # Create mental models with tags
        await memory.refresh_mental_models(
            bank_id=bank_id,
            tags=["project-x"],
            request_context=request_context,
        )
        await memory.wait_for_background_tasks()

        # Reflect with matching tags
        result = await memory.reflect_async(
            bank_id=bank_id,
            query="Who are the engineers?",
            tags=["project-x"],
            request_context=request_context,
        )

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

        # Reflect with non-matching tags - should still work
        result2 = await memory.reflect_async(
            bank_id=bank_id,
            query="Who are the engineers?",
            tags=["different-project"],
            request_context=request_context,
        )

        assert result2.text is not None

    async def test_mental_model_response_includes_tags(self, memory_with_mission_and_tags, request_context):
        """Test that mental model responses include the tags field."""
        memory, bank_id = memory_with_mission_and_tags

        # Create models with tags
        await memory.refresh_mental_models(
            bank_id=bank_id,
            tags=["test-tag"],
            request_context=request_context,
        )
        await memory.wait_for_background_tasks()

        # Get models
        models = await memory.list_mental_models(
            bank_id=bank_id,
            request_context=request_context,
        )

        # Verify tags field is present in response
        for model in models:
            assert "tags" in model
            assert isinstance(model["tags"], list)

        # Get single model
        if models:
            model = await memory.get_mental_model(
                bank_id=bank_id,
                model_id=models[0]["id"],
                request_context=request_context,
            )
            assert "tags" in model
            assert isinstance(model["tags"], list)