tests/test_context_pack.py

import pytest

from app.services.context_pack import ContextPackService


class FakeMemory:
    async def query_prior_knowledge(self, query: str, project_id: str = "", mode: str = "concept") -> str:
        if mode == "profile":
            return (
                "Preferred name: Anshuman.\n"
                "Student prefers mechanism-first explanations.\n"
                "Irrelevant old profile note " + ("x" * 3000)
            )
        if mode == "project":
            return (
                "Project memory: this project studies Adam and Transformer optimization.\n"
                "The student previously confused bias correction with weight decay.\n"
                + ("project filler " * 400)
            )
        return ""


@pytest.mark.asyncio
async def test_identity_question_prioritizes_profile_over_paper_chunks():
    pack = await ContextPackService(memory=FakeMemory()).build(
        agent_id="chat",
        query="what is my name?",
        project_id="p1",
        paper_chunks=[
            {"source": "paper.pdf", "text": "Attention is all you need. " * 200},
        ],
        max_chars=1800,
    )

    rendered = pack.render()

    assert "Preferred name: Anshuman" in rendered
    assert "Attention is all you need" not in rendered
    assert pack.sections[0].kind == "student_profile"
    assert len(rendered) <= 1800


@pytest.mark.asyncio
async def test_paper_question_prioritizes_source_chunks_and_keeps_profile_tiny():
    pack = await ContextPackService(memory=FakeMemory()).build(
        agent_id="chat",
        query="explain Adam bias correction from the paper",
        project_id="p1",
        paper_chunks=[
            {"source": "adam.pdf", "text": "Adam bias correction divides moment estimates by one minus beta powers. " * 40},
            {"source": "other.pdf", "text": "A distracting unrelated appendix. " * 80},
        ],
        max_chars=2200,
    )

    rendered = pack.render()

    assert "Adam bias correction divides moment estimates" in rendered
    assert "Preferred name: Anshuman" in rendered
    assert rendered.index("SOURCE MATERIAL") < rendered.index("STUDENT PROFILE")
    assert len(rendered) <= 2200


@pytest.mark.asyncio
async def test_context_pack_compresses_before_final_budget_cutoff():
    pack = await ContextPackService(memory=FakeMemory()).build(
        agent_id="pair_buddy",
        query="bias correction",
        project_id="p1",
        paper_chunks=[
            {"source": "adam.pdf", "text": "Bias correction matters for early optimization steps. " + ("filler " * 2000)},
        ],
        recent_summaries=["Recent: student asked why bias correction matters."],
        max_chars=900,
    )

    rendered = pack.render()

    assert "Bias correction matters" in rendered
    assert "student asked why bias correction matters" in rendered
    assert "filler filler filler filler filler filler filler filler filler filler" not in rendered
    assert len(rendered) <= 900