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_pending_memory_promotion_discards_one_off_entries(monkeypatch, tmp_path):
    import app.services.student_memory as student_memory

    monkeypatch.setattr(student_memory, "MEMORY_ROOT", tmp_path)
    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"], text))

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

    service = StudentMemoryService()
    service.stage_pending_memory("p1", dataset="project_p1", session_id="p1", text="Student briefly mentioned DeepSeek once.", kind="project")
    service.stage_pending_memory("p1", dataset="research_profile", session_id="p1", text="Durable learner signal: Student prefers mechanism-first explanations.", kind="profile")
    result = await service.promote_pending_memory("p1")

    assert result["promoted"] == 1
    assert result["discarded"] == 1
    assert calls == [("research_profile", "Durable learner signal: Student prefers mechanism-first explanations.")]
    assert service.list_pending_memory("p1") == []


@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")

    assert calls == ["research_profile"]


@pytest.mark.asyncio
async def test_generic_session_summary_does_not_pollute_profile_memory(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"])

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

    await StudentMemoryService().push_session("p1", "Topic", [], [], "This session covered Adam and bias correction.")

    assert calls == []


@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_research_memory_reviewer_skill_bootstrap_uses_cognee_skill_content(monkeypatch):
    calls = []

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

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

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

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

    assert result["ok"] is True
    assert calls[0][1]["dataset_name"] == "project_p1"
    assert calls[0][1]["content_type"] == "skills"
    assert calls[0][1]["skill_name"] == "research-memory-reviewer"
    assert "research-memory-reviewer" in calls[0][0]


@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