tests/test_recommendation_service.py

import pytest

from app.services.citation_graph import CitationGraph
from app.services.citation_topology import ResolvedSeed
from app.services.scholar_service import AuthorCandidate, AuthorResolution
from app.services.recommendation_service import RecommendationService
from app.services.verified_paper_discovery import VerifiedDiscoveryResult

QUERY = (
    "i want the paper Revisiting reinforce style optimization for learning from human "
    "feedback in llms. also who is its authors?"
)


class FakeClient:
    """Stands in for CerebrasClient -> returns a bare, unopinionated plan.

    The deterministic contract (_apply_deterministic_contract) is what's under test, and it
    runs *after* this LLM-shaped plan, so a minimal stub is enough here.
    """

    def structured_complete(self, messages, output_model, **kwargs):
        return output_model()


def _empty_graph() -> CitationGraph:
    return CitationGraph(project_id="project-1")


def test_single_paper_seed_extracted_from_natural_phrasing():
    seed = RecommendationService._single_paper_seed(QUERY)
    assert seed == "Revisiting reinforce style optimization for learning from human feedback in llms"


def test_single_paper_seed_ignores_plural_paper_requests():
    assert RecommendationService._single_paper_seed("recommend me some papers about attention") is None
    assert RecommendationService._single_paper_seed("find papers on contrastive learning") is None


def test_context_paper_seed_resolves_this_paper_to_the_selection():
    selection = "Naturalreasoning: Reasoning in the wild with 2.8m challenging questions"
    assert RecommendationService._context_paper_seed("/paper need this paper", selection) == selection
    assert RecommendationService._context_paper_seed("need this paper", selection) == selection
    assert RecommendationService._context_paper_seed("i want that paper", selection) == selection


def test_context_paper_seed_requires_both_reference_and_selection():
    assert RecommendationService._context_paper_seed("need this paper", None) is None
    assert RecommendationService._context_paper_seed("recommend some papers on attention", "some selection") is None


@pytest.mark.asyncio
async def test_recommend_resolves_exactly_one_named_paper(tmp_path):
    resolved = {
        "title": "Revisiting REINFORCE-Style Optimization for Learning from Human Feedback in LLMs",
        "authors": "Jane Doe, John Smith",
        "year": 2025,
        "cited_by": 3,
        "url": "https://example.com/reinforce-revisited",
        "abstract": "We revisit REINFORCE-style optimization...",
        "doi": "",
        "arxiv_id": "2501.00000",
        "openalex_id": "W123",
        "provider": "openalex",
        "can_acquire": True,
    }

    async def fake_title_resolver(title: str):
        assert "reinforce" in title.casefold()
        return dict(resolved)

    async def fake_fetcher(*args, **kwargs):
        raise AssertionError("generic discovery search should not run for a resolved single-paper request")

    service = RecommendationService(
        fetcher=fake_fetcher,
        title_resolver=fake_title_resolver,
        client=FakeClient(),
        root=str(tmp_path),
    )

    papers, plan = await service.recommend(
        "project-1",
        QUERY,
        _empty_graph(),
        [],
        use_project_context=False,
    )

    assert plan.result_limit == 1
    assert len(papers) == 1
    assert papers[0]["title"] == resolved["title"]
    assert papers[0]["authors"] == resolved["authors"]
    assert papers[0]["why_recommended"].startswith("The exact paper you asked for")


@pytest.mark.asyncio
async def test_recommend_falls_back_to_search_when_paper_unresolved(tmp_path):
    async def fake_title_resolver(title: str):
        return None

    async def fake_fetcher(query: str, n: int = 5, year_start=None, year_end=None):
        return [{
            "title": "Some Closest Match",
            "authors": "Someone",
            "year": 2024,
            "cited_by": 1,
            "url": "https://example.com/closest",
            "abstract": "",
        }]

    service = RecommendationService(
        fetcher=fake_fetcher,
        title_resolver=fake_title_resolver,
        client=FakeClient(),
        root=str(tmp_path),
    )

    papers, plan = await service.recommend(
        "project-1",
        QUERY,
        _empty_graph(),
        [],
        use_project_context=False,
    )

    assert plan.single_paper_lookup is False
    assert any("Could not confidently match" in assumption for assumption in plan.assumptions)
    assert len(papers) >= 1


@pytest.mark.asyncio
async def test_recommend_uses_selection_context_for_this_paper_requests(tmp_path):
    """Reproduces the reported bug: highlighting a paper title then typing "need this paper"
    must resolve to that highlighted paper, not generic project-graph topics."""
    selection = "Naturalreasoning: Reasoning in the wild with 2.8m challenging questions"
    resolved = {
        "title": "NaturalReasoning: Reasoning in the Wild with 2.8 Million Challenging Questions",
        "authors": "Weizhe Yuan, Jane Doe",
        "year": 2025,
        "cited_by": 7,
        "url": "https://example.com/naturalreasoning",
        "abstract": "",
    }

    async def fake_title_resolver(title: str):
        assert title == selection
        return dict(resolved)

    async def fake_fetcher(*args, **kwargs):
        raise AssertionError("generic discovery search should not run once selection context resolves the paper")

    service = RecommendationService(
        fetcher=fake_fetcher,
        title_resolver=fake_title_resolver,
        client=FakeClient(),
        root=str(tmp_path),
    )

    papers, plan = await service.recommend(
        "project-1",
        "need this paper",
        _empty_graph(),
        ["Reinforcement Learning with Verifiable Rewards and Retentive Networks"],
        use_project_context=False,
        selection_text=selection,
    )

    assert plan.result_limit == 1
    assert len(papers) == 1
    assert papers[0]["title"] == resolved["title"]


@pytest.mark.asyncio
async def test_latest_author_query_returns_only_latest_verified_work(tmp_path):
    calls = []

    async def fake_author_resolver(name: str):
        assert name == "Jonas Gehring"
        return AuthorResolution(
            status="resolved",
            candidates=[AuthorCandidate(openalex_id="A1", name=name)],
        )

    async def fake_author_fetcher(author_id: str, **kwargs):
        calls.append((author_id, kwargs))
        return [{
            "title": "Newest verified work",
            "authors": "Jonas Gehring",
            "year": 2026,
            "cited_by": 4,
            "url": "https://example.com/latest",
            "abstract": "",
            "openalex_id": "W1",
        }]

    service = RecommendationService(
        author_resolver=fake_author_resolver,
        author_fetcher=fake_author_fetcher,
        client=FakeClient(),
        root=str(tmp_path),
    )
    outcome = await service.recommend_or_clarify(
        project_id="project-1",
        query="I need the latest paper of Jonas Gehring",
        citation_graph=_empty_graph(),
        graph_labels=["Large Language Models"],
        use_project_context=True,
    )

    assert outcome.kind == "recommendations"
    assert len(outcome.papers) == 1
    assert outcome.papers[0]["title"] == "Newest verified work"
    assert calls == [("A1", {"n": 1, "year_start": None, "year_end": None, "ranking": "publication_date"})]


@pytest.mark.asyncio
async def test_ambiguous_author_returns_provider_backed_clarification(tmp_path):
    async def fake_author_resolver(name: str):
        return AuthorResolution(
            status="ambiguous",
            candidates=[
                AuthorCandidate(openalex_id="A1", name=name, works_count=10),
                AuthorCandidate(openalex_id="A2", name=name, works_count=4),
            ],
        )

    service = RecommendationService(
        author_resolver=fake_author_resolver,
        client=FakeClient(),
        root=str(tmp_path / "recommendations"),
    )
    service.clarification_store.root = str(tmp_path / "clarifications")
    service.clarification_store.__init__(root=service.clarification_store.root)

    outcome = await service.recommend_or_clarify(
        project_id="project-1",
        query="latest paper by Alex Kim",
        citation_graph=_empty_graph(),
        graph_labels=[],
        use_project_context=False,
    )

    assert outcome.kind == "clarification_required"
    assert [choice.resolved_values["author_id"] for choice in outcome.choices] == ["A1", "A2"]


@pytest.mark.asyncio
async def test_top_papers_of_year_dispatches_provider_citation_ranking(tmp_path):
    calls = []

    async def fake_ranked_fetcher(**kwargs):
        calls.append(kwargs)
        return [{
            "title": f"Ranked paper {index}", "authors": "Author", "year": 2026,
            "cited_by": 100 - index, "url": f"https://example.com/{index}", "abstract": "",
        } for index in range(5)]

    service = RecommendationService(
        ranked_fetcher=fake_ranked_fetcher,
        client=FakeClient(),
        root=str(tmp_path),
    )
    outcome = await service.recommend_or_clarify(
        project_id="project-1",
        query="What are the top 5 papers of 2026 by citation?",
        citation_graph=_empty_graph(),
        graph_labels=[],
        use_project_context=False,
    )

    assert outcome.kind == "recommendations"
    assert len(outcome.papers) == 5
    assert calls == [{"n": 5, "year_start": 2026, "year_end": 2026, "ranking": "citation_count"}]

    second = await service.recommend_or_clarify(
        project_id="project-1",
        query="What are the top 5 papers of 2026 by citation?",
        citation_graph=_empty_graph(),
        graph_labels=[],
        use_project_context=False,
    )
    assert len(second.papers) == 5
    assert len(calls) == 1


@pytest.mark.asyncio
@pytest.mark.parametrize(
    ("query", "expected_direction"),
    [
        ("top 5 papers that cite Attention Is All You Need", "cites_seed"),
        ("top 5 papers Attention Is All You Need cites", "cited_by_seed"),
    ],
)
async def test_citation_queries_dispatch_verified_direction(tmp_path, query, expected_direction):
    class FakeVerifiedDiscovery:
        def __init__(self):
            self.calls = []

        async def discover(self, **kwargs):
            self.calls.append(kwargs)
            return VerifiedDiscoveryResult(
                candidates=[{
                    "title": "Verified related paper", "authors": "Author", "year": 2025,
                    "cited_by": 10, "url": "https://example.com/verified", "abstract": "",
                    "openalex_id": "W2", "citation_verified": True,
                }],
                seed=ResolvedSeed(openalex_id="W1", title="Attention Is All You Need"),
                exhausted=True,
                provider_requests=2,
            )

    discovery = FakeVerifiedDiscovery()
    service = RecommendationService(
        verified_discovery=discovery,
        client=FakeClient(),
        root=str(tmp_path),
    )
    outcome = await service.recommend_or_clarify(
        project_id="project-1",
        query=query,
        citation_graph=_empty_graph(),
        graph_labels=["unrelated project topic"],
        use_project_context=True,
    )

    assert outcome.kind == "recommendations"
    assert discovery.calls[0]["seed_title"] == "Attention Is All You Need"
    assert discovery.calls[0]["citation_constraint"] == expected_direction