basicmachines-co-basic-memory/tests/schema/test_inference.py

"""Tests for basic_memory.schema.inference -- schema inference from usage patterns."""

from basic_memory.schema.inference import (
    NoteData,
    ObservationData,
    RelationData,
    InferenceResult,
    infer_schema,
)

# Short aliases for test readability
Obs = ObservationData
Rel = RelationData


# --- Helpers ---


def _note(
    identifier: str,
    observations: list[ObservationData] | None = None,
    relations: list[RelationData] | None = None,
) -> NoteData:
    return NoteData(
        identifier=identifier,
        observations=observations or [],
        relations=relations or [],
    )


# --- Empty notes ---


class TestInferEmpty:
    def test_empty_notes_list(self):
        result = infer_schema("Person", [])
        assert isinstance(result, InferenceResult)
        assert result.notes_analyzed == 0
        assert result.field_frequencies == []
        assert result.suggested_schema == {}
        assert result.suggested_required == []
        assert result.suggested_optional == []
        assert result.excluded == []


# --- Single note ---


class TestInferSingleNote:
    def test_single_note_all_fields_at_100_percent(self):
        note = _note("n1", observations=[Obs("name", "Alice"), Obs("role", "Engineer")])
        result = infer_schema("Person", [note])

        assert result.notes_analyzed == 1
        assert len(result.suggested_required) == 2
        assert "name" in result.suggested_required
        assert "role" in result.suggested_required

    def test_single_note_with_relations(self):
        note = _note(
            "n1",
            observations=[Obs("name", "Alice")],
            relations=[Rel("works_at", "Acme")],
        )
        result = infer_schema("Person", [note])
        assert len(result.suggested_required) == 2
        assert "works_at" in result.suggested_required


# --- Frequency thresholds ---


class TestInferThresholds:
    def test_95_percent_required(self):
        """Fields at 95%+ become required."""
        notes = []
        for i in range(20):
            obs = [Obs("name", f"Person {i}")]
            if i < 19:
                obs.append(Obs("bio", f"Bio {i}"))
            notes.append(_note(f"n{i}", observations=obs))

        result = infer_schema("Person", notes)
        assert "name" in result.suggested_required
        assert "bio" in result.suggested_required  # 19/20 = 95%

    def test_below_95_percent_optional(self):
        """Fields between 25% and 95% become optional."""
        notes = []
        for i in range(10):
            obs = [Obs("name", f"Person {i}")]
            if i < 5:
                obs.append(Obs("role", f"Role {i}"))
            notes.append(_note(f"n{i}", observations=obs))

        result = infer_schema("Person", notes)
        assert "name" in result.suggested_required
        assert "role" in result.suggested_optional  # 50%

    def test_below_25_percent_excluded(self):
        """Fields below 25% are excluded."""
        notes = []
        for i in range(10):
            obs = [Obs("name", f"Person {i}")]
            if i < 2:
                obs.append(Obs("rare", f"Rare {i}"))
            notes.append(_note(f"n{i}", observations=obs))

        result = infer_schema("Person", notes)
        assert "rare" in result.excluded  # 20%
        assert "rare" not in result.suggested_required
        assert "rare" not in result.suggested_optional

    def test_custom_thresholds(self):
        """Custom thresholds override defaults."""
        notes = [_note(f"n{i}", observations=[Obs("field", f"val{i}")]) for i in range(3)]
        notes.append(_note("n3"))

        result = infer_schema(
            "Test",
            notes,
            required_threshold=0.80,
            optional_threshold=0.50,
        )
        assert "field" in result.suggested_optional  # 75% < 80%
        assert "field" not in result.suggested_required


# --- Array detection ---


class TestInferArrayDetection:
    def test_array_detected_when_multiple_per_note(self):
        """Category appearing multiple times in >50% of containing notes -> array."""
        notes = [
            _note("n0", observations=[Obs("tag", "python"), Obs("tag", "mcp")]),
            _note("n1", observations=[Obs("tag", "schema"), Obs("tag", "validation")]),
            _note("n2", observations=[Obs("tag", "ai"), Obs("tag", "llm")]),
            _note("n3", observations=[Obs("tag", "single")]),
        ]
        result = infer_schema("Project", notes)

        tag_freq = next(f for f in result.field_frequencies if f.name == "tag")
        assert tag_freq.is_array is True

    def test_single_value_not_array(self):
        notes = [_note(f"n{i}", observations=[Obs("name", f"Person {i}")]) for i in range(5)]
        result = infer_schema("Person", notes)
        name_freq = next(f for f in result.field_frequencies if f.name == "name")
        assert name_freq.is_array is False


# --- Relation frequency ---


class TestInferRelations:
    def test_relation_frequency(self):
        notes = [_note(f"n{i}", relations=[Rel("works_at", f"Org{i}")]) for i in range(3)]
        notes.append(_note("n3"))
        result = infer_schema("Person", notes)

        works_at = next(f for f in result.field_frequencies if f.name == "works_at")
        assert works_at.source == "relation"
        assert works_at.percentage == 0.75
        assert "works_at" in result.suggested_optional

    def test_relation_array_detection(self):
        notes = [
            _note("n0", relations=[Rel("knows", "Alice"), Rel("knows", "Bob")]),
            _note("n1", relations=[Rel("knows", "Charlie"), Rel("knows", "Dave")]),
            _note("n2", relations=[Rel("knows", "Eve")]),
        ]
        result = infer_schema("Person", notes)
        knows_freq = next(f for f in result.field_frequencies if f.name == "knows")
        assert knows_freq.is_array is True


# --- Sample values ---


class TestInferSampleValues:
    def test_sample_values_collected(self):
        notes = [_note(f"n{i}", observations=[Obs("name", f"Person {i}")]) for i in range(3)]
        result = infer_schema("Person", notes)
        name_freq = next(f for f in result.field_frequencies if f.name == "name")
        assert len(name_freq.sample_values) == 3

    def test_sample_values_capped_at_max(self):
        notes = [_note(f"n{i}", observations=[Obs("name", f"Person {i}")]) for i in range(10)]
        result = infer_schema("Person", notes, max_sample_values=5)
        name_freq = next(f for f in result.field_frequencies if f.name == "name")
        assert len(name_freq.sample_values) == 5


# --- Suggested schema dict ---


class TestInferSuggestedSchema:
    def test_required_field_key_format(self):
        """Required fields have bare name (no '?')."""
        notes = [_note("n1", observations=[Obs("name", "Alice")])]
        result = infer_schema("Person", notes)
        assert "name" in result.suggested_schema

    def test_optional_field_key_format(self):
        """Optional fields have '?' suffix."""
        notes = [
            _note("n0", observations=[Obs("name", "A"), Obs("role", "Eng")]),
            _note("n1", observations=[Obs("name", "B"), Obs("role", "PM")]),
            _note("n2", observations=[Obs("name", "C")]),
            _note("n3", observations=[Obs("name", "D")]),
        ]
        result = infer_schema("Person", notes)
        assert "role?" in result.suggested_schema

    def test_array_field_key_format(self):
        """Array fields have '(array)' suffix."""
        notes = [
            _note("n0", observations=[Obs("tag", "a"), Obs("tag", "b")]),
            _note("n1", observations=[Obs("tag", "c"), Obs("tag", "d")]),
        ]
        result = infer_schema("Project", notes)
        assert "tag(array)" in result.suggested_schema

    def test_excluded_fields_not_in_schema(self):
        """Fields below threshold not in suggested schema."""
        notes = [_note(f"n{i}", observations=[Obs("name", f"P{i}")]) for i in range(10)]
        notes[0] = _note("n0", observations=[Obs("name", "P0"), Obs("rare", "x")])
        result = infer_schema("Person", notes)
        for key in result.suggested_schema:
            assert "rare" not in key


# --- Target entity type on relations (bug fix) ---


class TestInferRelationTargetType:
    """Verify that relation target type comes from the individual relation,
    not from the source note's entity_type."""

    def test_target_type_from_relation_data(self):
        """Relations with target_entity_type produce correct schema suggestions."""
        notes = [
            _note(
                f"n{i}",
                relations=[Rel("works_at", f"Org{i}", target_entity_type="organization")],
            )
            for i in range(3)
        ]
        result = infer_schema("Person", notes)
        works_at = next(f for f in result.field_frequencies if f.name == "works_at")
        assert works_at.target_type == "organization"

    def test_target_type_none_when_unresolved(self):
        """Relations without target_entity_type produce None target_type."""
        notes = [_note(f"n{i}", relations=[Rel("knows", f"P{i}")]) for i in range(3)]
        result = infer_schema("Person", notes)
        knows = next(f for f in result.field_frequencies if f.name == "knows")
        assert knows.target_type is None

    def test_mixed_target_types_uses_most_common(self):
        """When relations point to different types, the most common wins."""
        notes = [
            _note("n0", relations=[Rel("works_at", "OrgA", target_entity_type="organization")]),
            _note("n1", relations=[Rel("works_at", "OrgB", target_entity_type="organization")]),
            _note("n2", relations=[Rel("works_at", "SchoolC", target_entity_type="school")]),
        ]
        result = infer_schema("Person", notes)
        works_at = next(f for f in result.field_frequencies if f.name == "works_at")
        assert works_at.target_type == "organization"