ideas on breaking up functionality

merino/curated_recommendations/interest_picker.py

@@ -0,0 +1,121 @@
+"""Module for the New Tab 'interest picker' component that allows users to follow sections."""
+
+import random
+
+from merino.curated_recommendations.protocol import (
+    CuratedRecommendationsFeed,
+    InterestPickerSection,
+    InterestPicker,
+    Section,
+    SectionWithID
+)
+
+
+MIN_INITIALLY_VISIBLE_SECTION_COUNT = 3  # Minimum number of sections that are shown by default.
+MIN_INTEREST_PICKER_COUNT = 8  # Minimum number of items in the interest picker.
+
+
+def main(
+    feed: CuratedRecommendationsFeed,
+) -> tuple[list[SectionWithID], InterestPicker | None]:
+    # Get all available sections sorted by the order in which they should be shown.
+    sections = sort_sections(feed.get_sections())
+
+    # set initial visibility for sections
+    sections = set_section_initial_visibility(sections, MIN_INITIALLY_VISIBLE_SECTION_COUNT, MIN_INTEREST_PICKER_COUNT)
+
+    # see if we can create an interest picker based on section visibility
+    interest_picker = create_interest_picker(sections, MIN_INTEREST_PICKER_COUNT)
+
+    # if we were able to create an interest picker, re-order the sections
+    if interest_picker:
+        sections = _renumber_sections(sections, interest_picker.receivedFeedRank)
+
+    return sections, interest_picker
+
+
+def sort_sections(sections: list[SectionWithID]) -> list[SectionWithID]:
+    return sorted(sections, key=lambda tup: tup.section.receivedFeedRank)
+
+def set_section_initial_visibility(
+    sections: list[SectionWithID],
+    min_initially_visible_count: int,
+    min_interest_picker_count: int
+) -> list[tuple[Section, str]]:
+    visible_count = 0
+
+    # make sure all followed sections are initially visible
+    # and make sure we have at least min_initially_visible_count sections initially visible
+    for section, _ in sections:
+        if section.isFollowed or visible_count < min_initially_visible_count:
+            section.isInitiallyVisible = True
+            visible_count += 1
+        else:
+            section.isInitiallyVisible = False
+
+    # now see if we have enough non-initially visible sections to even display the interest picker
+    invisible_count = 0
+
+    for section, _ in sections:
+        if not section.isInitiallyVisible:
+            invisible_count += 1
+
+    # if we don't have enough invisible sections to satisfy making an interest picker, then
+    # set all sections to initially visible
+    if invisible_count < min_interest_picker_count:
+        for section, _ in sections:
+            section.isInitiallyVisible = True
+
+    return sections
+
+def create_interest_picker(
+    sections: list[SectionWithID],
+    min_interest_picker_count: int       
+) -> InterestPicker | None:
+    picker_sections = [
+        (section, section_id)
+        for (section, section_id) in sections
+        if not section.isInitiallyVisible
+    ]
+
+    if len(picker_sections) >= min_interest_picker_count:
+        interest_picker_rank = _get_interest_picker_rank(sections)
+
+        return InterestPicker(
+            receivedFeedRank=interest_picker_rank,
+            title="Follow topics to personalize your feed",
+            subtitle=(
+                "We will bring you personalized content, all while respecting your privacy. "
+                "You'll have powerful control over what content you see and what you don't."
+            ),
+            sections=[
+                InterestPickerSection(sectionId=section_id) for _, section_id in picker_sections
+            ],
+        )
+    else:
+        return None
+
+def _renumber_sections(sections: list[SectionWithID], picker_rank: int) -> list[SectionWithID]:
+    """Renumber section ranks, leaving a gap for the interest picker.
+
+    Increments ranks by 1 so that the section after the picker has rank picker_rank+1.
+    """
+    new_rank = 0
+    for section, _ in sections:
+        if new_rank == picker_rank:
+            # Skip the rank for the interest picker.
+            new_rank += 1
+        section.receivedFeedRank = new_rank
+        new_rank += 1
+
+    return sections
+
+
+def _get_interest_picker_rank(sections: list[SectionWithID]) -> int:
+    """Return a randomized rank for the interest picker.
+
+    If any section is followed, choose a random int in [2, 4]; otherwise, in [1, 3].
+    """
+    if any(section.isFollowed for section, _ in sections):
+        return random.randint(2, 4)
+    return random.randint(1, 3)

merino/curated_recommendations/protocol.py

@@ -7,6 +7,8 @@
 
 from pydantic import Field, field_validator, model_validator, BaseModel, ValidationInfo
 
+from collections import namedtuple
+
 from merino.curated_recommendations.corpus_backends.protocol import CorpusItem, Topic
 from merino.curated_recommendations.fakespot_backend.protocol import FakespotFeed
 
@@ -122,6 +124,7 @@ class CuratedRecommendationsRequest(BaseModel):
     # Allow any string value or null, because ExperimentName is not meant to be an exhaustive list.
     experimentName: ExperimentName | str | None = None
     experimentBranch: str | None = None
+    enableInterestPicker: bool = False
 
     @field_validator("topics", mode="before")
     def validate_topics(cls, values):
@@ -226,7 +229,10 @@ class Section(BaseModel):
     layout: Layout
     isFollowed: bool = False
     isBlocked: bool = False
+    isInitiallyVisible: bool = True
+
 
+SectionWithID = namedtuple('SectionWithID', ['section', 'ID'])
 
 class CuratedRecommendationsFeed(BaseModel):
     """Multiple lists of curated recommendations, that are currently in an experimental phase."""
@@ -268,14 +274,38 @@ def get_section_by_topic_id(self, serp_topic_id: str) -> Section | None:
                 return cast(Section, getattr(self, field_name, None))
         return None
 
+    def get_sections(self) -> list[SectionWithID]:
+        """Get a list of all sections as tuples, where each tuple is a Section and its ID."""
+        return [
+            SectionWithID(feed, str(model_field.alias))  # alias defines the section id
+            for field_name, model_field in self.model_fields.items()
+            if (feed := getattr(self, field_name)) is not None and type(feed) is Section
+        ]
+
     def set_topic_section(self, topic: Topic, section: Section):
         """Set a section for the given topic."""
         setattr(self, topic.name.lower(), section)
 
 
+class InterestPickerSection(BaseModel):
+    """Model representing a single section entry in the interest picker."""
+
+    sectionId: str
+
+
+class InterestPicker(BaseModel):
+    """Model representing the interest picker component for following sections."""
+
+    receivedFeedRank: int
+    title: str
+    subtitle: str
+    sections: list[InterestPickerSection]
+
+
 class CuratedRecommendationsResponse(BaseModel):
     """Response schema for a list of curated recommendations"""
 
     recommendedAt: int
     data: list[CuratedRecommendation]
     feeds: CuratedRecommendationsFeed | None = None
+    interestPicker: InterestPicker | None = None

merino/curated_recommendations/provider.py

@@ -15,6 +15,7 @@
 from merino.curated_recommendations.fakespot_backend.protocol import (
     FakespotBackend,
 )
+from merino.curated_recommendations.interest_picker import create_interest_picker
 from merino.curated_recommendations.layouts import (
     layout_4_medium,
     layout_4_large,
@@ -468,6 +469,11 @@ async def fetch(
                 curated_recommendations_request.sections, response.feeds
             )
 
+        if curated_recommendations_request.enableInterestPicker and response.feeds:
+            feeds, interest_picker = create_interest_picker(response.feeds)
+            response.feeds = feeds
+            response.interestPicker = interest_picker
+
         return response
 
     async def fetch_backup_recommendations(

tests/integration/api/v1/curated_recommendations/test_curated_recommendations.py

@@ -21,6 +21,7 @@
     get_provider,
     ConstantPrior,
     ExtendedExpirationCorpusBackend,
+    interest_picker,
 )
 from merino.curated_recommendations.corpus_backends.protocol import Topic, ScheduledSurfaceId
 from merino.curated_recommendations.engagement_backends.protocol import (
@@ -1463,6 +1464,31 @@ async def test_sections_feed_titles(self, locale, expected_titles):
             assert top_stories_section is not None
             assert top_stories_section["subtitle"] is None
 
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("enable_interest_picker", [True, False])
+    async def test_sections_interest_picker(self, enable_interest_picker, monkeypatch):
+        """Test the curated recommendations endpoint returns an interest picker when enabled"""
+        # The fixture data doesn't have enough sections for the interest picker to show up, so lower
+        # the minimum number of sections that it needs to have to 1.
+        monkeypatch.setattr(interest_picker, "MIN_INTEREST_PICKER_COUNT", 1)
+
+        async with AsyncClient(app=app, base_url="http://test") as ac:
+            response = await ac.post(
+                "/api/v1/curated-recommendations",
+                json={
+                    "locale": "en-US",
+                    "feeds": ["sections"],
+                    "enableInterestPicker": enable_interest_picker,
+                },
+            )
+            data = response.json()
+
+            interest_picker_response = data["interestPicker"]
+            if enable_interest_picker:
+                assert interest_picker_response is not None
+            else:
+                assert interest_picker_response is None
+
 
 class TestExtendedExpiration:
     """Test the behavior of the ExtendedExpiration experiment functionality."""

tests/unit/curated_recommendations/test_interest_picker.py

@@ -0,0 +1,105 @@
+"""Tests covering merino/curated_recommendations/interest_picker.py"""
+
+import pytest
+
+from merino.curated_recommendations.corpus_backends.protocol import Topic
+from merino.curated_recommendations.interest_picker import (
+    create_interest_picker,
+    MIN_INITIALLY_VISIBLE_SECTION_COUNT,
+    MIN_INTEREST_PICKER_COUNT,
+)
+from merino.curated_recommendations.layouts import layout_4_medium
+from merino.curated_recommendations.protocol import Section, CuratedRecommendationsFeed
+
+
+def generate_feed(section_count: int, followed_count: int = 0) -> CuratedRecommendationsFeed:
+    """Create a CuratedRecommendationsFeed populated with sections.
+
+    Args:
+        section_count (int): Number of sections to create.
+        followed_count (int, optional): Number of sections to follow. Defaults to 0.
+
+    Returns:
+        CuratedRecommendationsFeed: A feed instance containing the generated sections.
+    """
+    feed = CuratedRecommendationsFeed()
+
+    # Set top_stories_section first.
+    feed.top_stories_section = Section(
+        receivedFeedRank=0,
+        recommendations=[],  # No recommendations are added for dummy purposes.
+        title="Top Stories",
+        layout=layout_4_medium,
+    )
+
+    # Iterate over the Topic enum and add topic sections.
+    topics = list(Topic)[: section_count - 1]
+    for i, topic in enumerate(topics):
+        section = Section(
+            receivedFeedRank=i + 1,  # Ranks start after top_stories_section.
+            recommendations=[],
+            title=f"{topic.value.title()} Section",
+            layout=layout_4_medium,
+            isFollowed=i < followed_count,
+        )
+        feed.set_topic_section(topic, section)
+
+    return feed
+
+
+def test_no_sections():
+    """Test that if the feed has no sections, no interest picker is created."""
+    feed, interest_picker = create_interest_picker(CuratedRecommendationsFeed())
+    assert interest_picker is None
+
+
+def test_not_enough_sections():
+    """Test that the interest picker is not shown if insufficient sections are available."""
+    # Create 10 sections. With MIN_INITIALLY_VISIBLE_SECTION_COUNT = 3, there will be 10 - 3 = 7 sections
+    # eligible for the interest picker, which is less than the min of 8.
+    section_count = MIN_INTEREST_PICKER_COUNT + MIN_INITIALLY_VISIBLE_SECTION_COUNT - 1
+    feed, interest_picker = create_interest_picker(generate_feed(section_count))
+
+    # Interest picker should not be created.
+    assert interest_picker is None
+
+    # All sections must be visible.
+    for section, _ in feed.get_sections():
+        assert section.isInitiallyVisible is True
+
+
+@pytest.mark.parametrize("followed_count", list(range(7)))
+def test_interest_picker_is_created(followed_count: int):
+    """Test that the interest picker is created as expected, if enough sections are available."""
+    section_count = 15
+    feed = generate_feed(section_count, followed_count=followed_count)
+    feed, interest_picker = create_interest_picker(feed)
+
+    assert interest_picker is not None
+
+    # The picker is ranked randomly, with different ranges depending on if any sections are followed
+    min_picker_rank = 1 if followed_count == 0 else 2
+    max_picker_rank = min_picker_rank + 2
+    assert min_picker_rank <= interest_picker.receivedFeedRank <= max_picker_rank
+
+    # Verify that the first MIN_INITIALLY_VISIBLE_SECTION_COUNT sections are visible.
+    sections = feed.get_sections()
+    visible_sections = [section for section, _ in sections if section.isInitiallyVisible]
+    assert len(visible_sections) == max(MIN_INITIALLY_VISIBLE_SECTION_COUNT, 1 + followed_count)
+
+    # Verify that all followed sections are visible.
+    for section, _ in sections:
+        if section.isFollowed:
+            assert section.isInitiallyVisible is True
+
+    # Verify that receivedFeedRank (including on interestPicker) is numbered 0, 1, 2, etc.
+    ranks = [section.receivedFeedRank for section, _ in sections]
+    expected_ranks = [i for i in range(section_count + 1) if i != interest_picker.receivedFeedRank]
+    assert sorted(ranks) == expected_ranks
+
+    # Verify that the interest picker's sections include all sections not visible by default.
+    hidden_section_ids: list[str | None] = [
+        section_id for section, section_id in sections if not section.isInitiallyVisible
+    ]
+    picker_ids = [s.sectionId for s in interest_picker.sections]
+    assert set(picker_ids) == set(hidden_section_ids)