tests/test_student_memory.py

import inspect
import sys
import types
from types import SimpleNamespace

import pytest

from app.services.student_memory import StudentMemoryService


def install_fake_cognee(monkeypatch, **attrs):
    module = types.ModuleType("cognee")
    for key, value in attrs.items():
        setattr(module, key, value)
    monkeypatch.setitem(sys.modules, "cognee", module)
    return module


def fake_search_type():
    return SimpleNamespace(GRAPH_COMPLETION="graph", TEMPORAL="temporal", AGENTIC_COMPLETION="agentic")


@pytest.mark.asyncio
async def test_project_observation_is_quarantined_before_cognee_write(monkeypatch, tmp_path):
    import app.services.student_memory as student_memory

    monkeypatch.setattr(student_memory, "MEMORY_ROOT", tmp_path)
    calls = []

    async def remember(text, **kwargs):
        calls.append(("remember", kwargs["dataset_name"]))

    install_fake_cognee(monkeypatch, remember=remember)

    ok = await StudentMemoryService().stage_project_observation("p1", "Attention", ["student connected QK lookup"])

    assert ok is True
    assert calls == []
    pending = StudentMemoryService().list_pending_memory("p1")
    assert len(pending) == 1
    assert pending[0]["dataset"] == "project_p1"
    assert "student connected QK lookup" in pending[0]["text"]


@pytest.mark.asyncio
async def test_profile_write_failure_does_not_block_project_write(monkeypatch):
    calls = []

    class Datasets:
        async def list_datasets(self):
            return [SimpleNamespace(name="project_p1"), SimpleNamespace(name="research_profile")]

    async def remember(text, **kwargs):
        calls.append(kwargs["dataset_name"])
        if kwargs["dataset_name"] == "research_profile":
            raise RuntimeError("profile write failed")

    install_fake_cognee(monkeypatch, datasets=Datasets(), remember=remember)

    await StudentMemoryService().stage_profile_observation(
        "p1", "Student prefers concise explanations", attribution="explicit_student", confidence=1.0,
    )

    assert calls == ["research_profile"]


@pytest.mark.asyncio
async def test_recall_uses_only_context_and_falls_back_on_typeerror(monkeypatch):
    calls = []

    async def recall(**kwargs):
        calls.append(kwargs)
        if "only_context" in kwargs:
            raise TypeError("unexpected keyword")
        return ["memory context"]

    install_fake_cognee(monkeypatch, SearchType=fake_search_type(), recall=recall)

    result = await StudentMemoryService().query_prior_knowledge("attention", project_id="p1")

    assert "memory context" in result
    assert calls[0]["only_context"] is True
    assert calls[0]["feedback_influence"] == 0.35
    assert "only_context" not in calls[1]


@pytest.mark.asyncio
async def test_temporal_recall_uses_temporal_search_type(monkeypatch):
    calls = []

    async def recall(**kwargs):
        calls.append(kwargs)
        return ["changed over time"]

    install_fake_cognee(monkeypatch, SearchType=fake_search_type(), recall=recall)

    result = await StudentMemoryService().query_prior_knowledge("attention", project_id="p1", mode="temporal")

    assert "changed over time" in result
    assert calls[0]["query_type"] == "temporal"


@pytest.mark.asyncio
async def test_profile_recall_query_is_name_aware(monkeypatch):
    calls = []

    async def recall(**kwargs):
        calls.append(kwargs)
        return ["Preferred name: Anshuman"]

    install_fake_cognee(monkeypatch, SearchType=fake_search_type(), recall=recall)

    result = await StudentMemoryService().query_prior_knowledge("attention", project_id="p1", mode="profile")

    assert "Anshuman" in result
    assert "preferred name" in calls[0]["query_text"]
    assert "call me" in calls[0]["query_text"]


@pytest.mark.asyncio
async def test_style_feedback_is_profile_memory_not_native_weighting(monkeypatch):
    calls = []

    class Session:
        async def add_feedback(self, **kwargs):
            calls.append(("feedback", kwargs))
            return True

        async def add_frequency_weights(self, **kwargs):
            calls.append(("weights", kwargs))
            return True

    class Datasets:
        async def list_datasets(self):
            return [SimpleNamespace(name="research_profile")]

    async def remember(text, **kwargs):
        calls.append(("remember", kwargs))

    install_fake_cognee(monkeypatch, session=Session(), datasets=Datasets(), remember=remember)

    result = await StudentMemoryService().record_style_feedback("p1", "more concise")

    assert result == {"profile_memory": True}
    assert [call[0] for call in calls] == ["remember"]


@pytest.mark.asyncio
async def test_native_feedback_requires_cognee_recall_ids(monkeypatch):
    calls = []

    class Session:
        async def add_feedback(self, **kwargs):
            calls.append(("feedback", kwargs))
            return True

        async def add_frequency_weights(self, **kwargs):
            calls.append(("weights", kwargs))
            return True

    install_fake_cognee(monkeypatch, session=Session())

    result = await StudentMemoryService().record_feedback("p1", "style_feedback", 1, "more concise", ["n1"], ["e1"])

    assert result == {"feedback": False, "frequency_weights": False, "skipped": True}
    assert calls == []


@pytest.mark.asyncio
async def test_native_feedback_uses_cognee_recall_metadata(monkeypatch):
    calls = []

    class Session:
        async def add_feedback(self, **kwargs):
            calls.append(("feedback", kwargs))
            return True

        async def add_frequency_weights(self, **kwargs):
            calls.append(("weights", kwargs))
            return True

    install_fake_cognee(monkeypatch, session=Session())

    result = await StudentMemoryService().record_feedback(
        "p1",
        "qa1",
        1,
        "more concise",
        ["cg-node-1"],
        ["cg-edge-1"],
        cognee_native=True,
    )

    assert result == {"feedback": True, "frequency_weights": True}
    assert calls[0][1]["feedback_text"] == "more concise"
    assert calls[1][1]["node_ids"] == ["cg-node-1"]


@pytest.mark.asyncio
async def test_flush_project_can_distill_then_improve(monkeypatch):
    calls = []

    class Datasets:
        async def list_datasets(self):
            return [SimpleNamespace(name="project_p1"), SimpleNamespace(name="research_profile")]

    class Session:
        async def distill_session(self, **kwargs):
            calls.append(("distill", kwargs["dataset"]))

    async def improve(**kwargs):
        calls.append(("improve", kwargs["dataset"]))

    install_fake_cognee(monkeypatch, datasets=Datasets(), session=Session(), improve=improve)

    result = await StudentMemoryService().flush_project("p1", strategy="distill_then_improve")

    assert result == {"project_p1": True, "research_profile": True}
    assert calls == [
        ("distill", "project_p1"),
        ("improve", "project_p1"),
        ("distill", "research_profile"),
        ("improve", "research_profile"),
    ]


@pytest.mark.asyncio
async def test_flush_profile_only_improves_research_profile(monkeypatch):
    calls = []

    class Datasets:
        async def list_datasets(self):
            return [SimpleNamespace(name="project_p1"), SimpleNamespace(name="research_profile")]

    async def improve(**kwargs):
        calls.append(("improve", kwargs["dataset"]))

    install_fake_cognee(monkeypatch, datasets=Datasets(), improve=improve)

    result = await StudentMemoryService().flush_profile("p1")

    assert result is True
    assert calls == [("improve", "research_profile")]


@pytest.mark.asyncio
async def test_native_wrappers_tolerate_missing_cognee_apis(monkeypatch):
    install_fake_cognee(monkeypatch)
    service = StudentMemoryService()

    assert (await service.run_project_memify("p1"))["ok"] is False
    assert (await service.get_schema_inventory("p1"))["ok"] is False
    assert (await service.get_provenance("p1"))["ok"] is False
    assert (await service.export_memory("p1"))["ok"] is True


@pytest.mark.asyncio
async def test_forget_project_document_resets_project_memory_without_document_id(monkeypatch):
    calls = []

    async def forget(**kwargs):
        calls.append(("forget", kwargs))
        return {"ok": True}

    class Datasets:
        async def list_datasets(self):
            return [SimpleNamespace(name="project_p1")]

    async def add(text, dataset_name):
        calls.append(("add", dataset_name, text))

    install_fake_cognee(monkeypatch, datasets=Datasets(), forget=forget, add=add)

    result = await StudentMemoryService().forget_project_document("p1", "a" * 64)

    assert result["ok"] is True
    assert calls[0] == ("forget", {"dataset": "project_p1", "memory_only": True})
    assert all("document_id" not in call[1] for call in calls if call[0] == "forget")


@pytest.mark.asyncio
async def test_memory_liveness_reports_degraded_when_recall_fails(monkeypatch):
    class Datasets:
        async def list_datasets(self):
            return [SimpleNamespace(name="project_p1"), SimpleNamespace(name="research_profile")]

    async def add(text, dataset_name):
        return None

    async def remember(text, **kwargs):
        return None

    async def improve(**kwargs):
        return None

    async def recall(**kwargs):
        raise RuntimeError("recall broken")

    install_fake_cognee(
        monkeypatch,
        datasets=Datasets(),
        add=add,
        remember=remember,
        improve=improve,
        recall=recall,
        SearchType=fake_search_type(),
    )

    status = await StudentMemoryService().memory_liveness("p1", force=True)

    assert status["state"] == "degraded"
    assert status["checks"]["recall"] is False
    assert "recall broken" in status["last_error"]


@pytest.mark.asyncio
async def test_memory_status_skips_liveness_probe_by_default(monkeypatch):
    calls = []

    class Datasets:
        async def list_datasets(self):
            calls.append("list_datasets")
            return [SimpleNamespace(name="project_p1"), SimpleNamespace(name="research_profile")]

    async def get_schema_inventory(**kwargs):
        calls.append(("inventory", kwargs["dataset"]))
        return [{"name": "Claim"}]

    async def get_memory_provenance_graph(**kwargs):
        calls.append("provenance")
        return [], []

    async def export(**kwargs):
        calls.append(("export", kwargs["dataset"]))
        return []

    async def remember(*args, **kwargs):
        raise AssertionError("memory_status should not run liveness writes by default")

    async def improve(**kwargs):
        raise AssertionError("memory_status should not flush Cognee by default")

    async def recall(**kwargs):
        raise AssertionError("memory_status should not recall by default")

    install_fake_cognee(
        monkeypatch,
        datasets=Datasets(),
        get_schema_inventory=get_schema_inventory,
        get_memory_provenance_graph=get_memory_provenance_graph,
        export=export,
        remember=remember,
        improve=improve,
        recall=recall,
        SearchType=fake_search_type(),
    )

    status = await StudentMemoryService().memory_status("p1")

    assert status["state"] == "ready"
    assert status["liveness"] == {}
    assert "list_datasets" in calls
    assert ("inventory", "project_p1") in calls


@pytest.mark.asyncio
async def test_temporal_recall_falls_back_to_local_ledger(monkeypatch, tmp_path):
    monkeypatch.setattr("app.services.student_memory.MEMORY_ROOT", tmp_path)
    service = StudentMemoryService()
    service.record_temporal_event("p1", "commit", "Student connected Adam to sparse gradients")

    class Datasets:
        async def list_datasets(self):
            return [SimpleNamespace(name="project_p1")]

    async def recall(**kwargs):
        raise RuntimeError("No temporal graph")

    install_fake_cognee(monkeypatch, datasets=Datasets(), recall=recall, SearchType=fake_search_type())

    result = await service.query_prior_knowledge("Adam", project_id="p1", mode="temporal")

    assert "Temporal project memory" in result
    assert "sparse gradients" in result


def test_study_buddy_agent_no_longer_calls_missing_memory_remember():
    from app.agents.study_buddy_agent import StudyBuddyAgent

    source = inspect.getsource(StudyBuddyAgent.evaluate_and_ask_next)

    assert ".remember(" not in source
    assert "stage_project_observation" in source


@pytest.mark.asyncio
async def test_cross_project_recurrence_candidate_is_actually_staged(monkeypatch, tmp_path):
    from app.services.memory_promotion import MemoryPromotionGate
    from app.services.memory_candidates import MemoryCandidate, make_candidate_id

    calls = []

    async def remember(text, **kwargs):
        calls.append((text, kwargs))
        return {"remembered": True}

    install_fake_cognee(monkeypatch, remember=remember)
    gate = MemoryPromotionGate(ledger_path=tmp_path / "promotion_decisions.jsonl")
    service = StudentMemoryService()

    async def ensure_profile_dataset(observer=None):
        return True

    monkeypatch.setattr(service, "ensure_profile_dataset", ensure_profile_dataset)

    statement = "Student repeatedly struggles to interpret objective functions."

    first = MemoryCandidate(
        candidate_id=make_candidate_id("student", "project-a", "recurring_confusion", statement),
        destination="student", project_id="project-a", kind="recurring_confusion",
        statement=statement, attribution="idea_observer_profile_proposal", confidence=0.90,
        interaction_ids=["interaction-a"], evidence_ids=[],
    )
    first_decision = gate.evaluate_student(first)
    assert first_decision.promote is False

    second = MemoryCandidate(
        candidate_id=make_candidate_id("student", "project-b", "recurring_confusion", statement),
        destination="student", project_id="project-b", kind="recurring_confusion",
        statement=statement, attribution="idea_observer_profile_proposal", confidence=0.90,
        interaction_ids=["interaction-b"], evidence_ids=[],
    )
    second_decision = gate.evaluate_student(second)

    assert second_decision.promote is True
    assert second_decision.reason == "cross_project_recurrence"
    assert set(second_decision.supporting_projects) == {"project-a", "project-b"}

    staged = await service.stage_promoted_candidate(second)

    assert staged is True
    assert calls, "cognee.remember() was never called -- the inline gate rejected an already-approved candidate"


@pytest.mark.asyncio
async def test_stage_promoted_candidate_does_not_recheck_recurrence(monkeypatch):
    from app.services.memory_candidates import MemoryCandidate, make_candidate_id

    calls = []

    async def remember(text, **kwargs):
        calls.append((text, kwargs))
        return {"remembered": True}

    install_fake_cognee(monkeypatch, remember=remember)
    service = StudentMemoryService()

    async def ensure_profile_dataset(observer=None):
        return True

    monkeypatch.setattr(service, "ensure_profile_dataset", ensure_profile_dataset)

    candidate = MemoryCandidate(
        candidate_id=make_candidate_id("student", "project-b", "recurring_confusion", "Some inferred trait."),
        destination="student", project_id="project-b", kind="recurring_confusion",
        statement="Some inferred trait.", attribution="idea_observer_profile_proposal", confidence=0.6,
        interaction_ids=["interaction-b"], evidence_ids=[],
    )

    staged = await service.stage_promoted_candidate(candidate)

    assert staged is True
    assert len(calls) == 1


@pytest.mark.asyncio
async def test_stage_promoted_candidate_rejects_project_destination():
    from app.services.memory_candidates import MemoryCandidate, make_candidate_id

    service = StudentMemoryService()
    candidate = MemoryCandidate(
        candidate_id=make_candidate_id("project", "project-a", "project_observation", "Uses PyTorch."),
        destination="project", project_id="project-a", kind="project_observation",
        statement="Uses PyTorch.", attribution="idea_observer_interaction", confidence=0.8,
        interaction_ids=[], evidence_ids=[],
    )

    with pytest.raises(ValueError, match="destination='student'"):
        await service.stage_promoted_candidate(candidate)


@pytest.mark.asyncio
async def test_record_style_feedback_routes_through_promotion_gate(monkeypatch):
    from app.services.memory_promotion import MemoryPromotionGate

    calls = []
    gate_calls = []

    async def remember(text, **kwargs):
        calls.append((text, kwargs))
        return {"remembered": True}

    install_fake_cognee(monkeypatch, remember=remember)
    service = StudentMemoryService()

    async def ensure_profile_dataset(observer=None):
        return True

    monkeypatch.setattr(service, "ensure_profile_dataset", ensure_profile_dataset)

    original_evaluate_student = MemoryPromotionGate.evaluate_student

    def spying_evaluate_student(self, candidate):
        gate_calls.append(candidate)
        return original_evaluate_student(self, candidate)

    monkeypatch.setattr(MemoryPromotionGate, "evaluate_student", spying_evaluate_student)

    result = await service.record_style_feedback("project-a", "Give me less code and more diagrams.")

    assert result["profile_memory"] is True
    assert calls, "record_style_feedback did not reach cognee.remember()"
    assert len(gate_calls) == 1, "record_style_feedback must route through MemoryPromotionGate.evaluate_student()"
    assert gate_calls[0].destination == "student"


@pytest.mark.asyncio
async def test_record_style_feedback_promotes_immediately_as_explicit(monkeypatch):
    """A single-project explicit style-feedback call must promote immediately
    via the explicit_student fast path -- it should not require a second
    project's worth of recurrence."""
    calls = []

    async def remember(text, **kwargs):
        calls.append((text, kwargs))
        return {"remembered": True}

    install_fake_cognee(monkeypatch, remember=remember)
    service = StudentMemoryService()

    async def ensure_profile_dataset(observer=None):
        return True

    monkeypatch.setattr(service, "ensure_profile_dataset", ensure_profile_dataset)

    result = await service.record_style_feedback("only-one-project", "Be more direct with me.")

    assert result["profile_memory"] is True
    assert len(calls) == 1