test: cover the interwoven story, guidance popup, and view toggle

Rework the story tests for the block model and add coverage for the guidance dialog presets, the Report/Story toggle, and story reset.

Author: ghostiee-11

lumen/tests/ai/test_story.py

@@ -29,154 +29,249 @@ async def _execute(self, context, **kwargs):
         return [Markdown("B done")], {'text': 'B'}
 
 
-async def test_story_agent_write_story(llm):
-    from lumen.ai.agents.story import Story, StoryAgent
+class ChartAction(Action):
+    """Emits a pipeline-backed LumenEditor so the story catalog has real views."""
 
-    llm.set_responses([Story(chain_of_thought="arc", title="Olympics", markdown="# Olympics\nBody")])
-    agent = StoryAgent(llm=llm)
+    source = param.Parameter()
 
-    result = await agent.write_story("Section: Medals\n<data summary>", title="Olympic Report")
+    label = param.String(default="Chart")
 
-    assert isinstance(result, Story)
-    assert result.title == "Olympics"
-    assert result.markdown == "# Olympics\nBody"
+    async def _execute(self, context, **kwargs):
+        from lumen.ai.editors import LumenEditor
+        from lumen.pipeline import Pipeline
 
+        pipeline = Pipeline(source=self.source, table='tiny')
+        return [LumenEditor(component=pipeline, title=self.label)], {'text': self.label}
 
-async def test_story_agent_write_note(llm):
-    from lumen.ai.agents.story import SectionNote, StoryAgent
 
-    llm.set_responses([SectionNote(chain_of_thought="key", markdown="Men-only sports dominate.")])
-    agent = StoryAgent(llm=llm)
+def _nb_text(report):
+    return "".join(
+        "".join(cell["source"]) for cell in json.loads(report.to_notebook())["cells"]
+    )
 
-    result = await agent.write_note("Section: Gender\n<data summary>", title="Gender Exclusive")
 
-    assert isinstance(result, SectionNote)
-    assert "Men-only" in result.markdown
+def _kinds(report):
+    return [kind for kind, _ in report._story_blocks]
 
 
-async def test_build_report_context_includes_titles_and_text():
-    from lumen.ai.agents.story import (
-        build_report_context, build_section_context,
-    )
+async def test_story_agent_write_story(llm):
+    from lumen.ai.agents.story import Story, StoryAgent, StoryBlock
 
-    report = Report(
-        Section(A(order=[]), title='Section A'),
-        Section(B(order=[]), title='Section B'),
-        title='My Report',
-    )
-    await report.execute()
+    llm.set_responses([Story(
+        chain_of_thought="arc", title="Olympics",
+        blocks=[StoryBlock(prose="Intro."), StoryBlock(view=1), StoryBlock(prose="Wrap up.")],
+    )])
+    agent = StoryAgent(llm=llm)
 
-    ctx = await build_report_context([s for s in report if s.include_in_export])
-    assert "Section A" in ctx
-    assert "A done" in ctx
-    assert "Section B" in ctx
-    assert "B done" in ctx
+    result = await agent.write_story(
+        "### View 1: Medals\n<summary>", guidance="focus on trends", title="Olympic Report"
+    )
 
-    section_ctx = await build_section_context(report[0])
-    assert "A done" in section_ctx
-    assert "B done" not in section_ctx
+    assert isinstance(result, Story)
+    assert result.title == "Olympics"
+    assert [block.view for block in result.blocks] == [None, 1, None]
 
 
-async def test_build_section_context_summarizes_table_pipeline(tiny_source):
-    from lumen.ai.agents.story import build_view_context
+async def test_build_catalog_numbers_and_summarizes(tiny_source):
+    from lumen.ai.agents.story import build_catalog
     from lumen.ai.editors import LumenEditor
     from lumen.pipeline import Pipeline
 
-    pipeline = Pipeline(source=tiny_source, table='tiny')
-    editor = LumenEditor(component=pipeline, title='Tiny table')
+    editor = LumenEditor(component=Pipeline(source=tiny_source, table='tiny'), title='Tiny table')
 
-    ctx = await build_view_context(editor)
+    ctx = await build_catalog([editor])
 
-    # The table's columns should surface in the summary sent to the LLM.
+    # Views are numbered so the story can reference them, and the data surfaces.
+    assert "View 1" in ctx
     assert "category" in ctx
     assert "value" in ctx
 
 
-def _nb_text(report):
-    return "".join(
-        "".join(cell["source"]) for cell in json.loads(report.to_notebook())["cells"]
-    )
-
-
-async def test_report_annotate_inserts_story_and_notes(llm):
-    from lumen.ai.agents.story import SectionNote, Story
+async def test_report_annotate_interleaves_prose_and_view(llm, tiny_source):
+    from lumen.ai.agents.story import Story, StoryBlock
 
     report = Report(
-        Section(A(order=[]), title='Section A'),
-        Section(B(order=[]), title='Section B'),
+        Section(ChartAction(source=tiny_source, label='Chart A'), title='Section A'),
         title='My Report',
         llm=llm,
     )
     await report.execute()
-    llm.set_responses([
-        Story(chain_of_thought="c", title="The Big Picture", markdown="Overall narrative."),
-        SectionNote(chain_of_thought="c", markdown="Note for A."),
-        SectionNote(chain_of_thought="c", markdown="Note for B."),
-    ])
+    llm.set_responses([Story(chain_of_thought="c", title="The Big Picture", blocks=[
+        StoryBlock(prose="Lead-in prose."),
+        StoryBlock(view=1),
+        StoryBlock(prose="Closing prose."),
+    ])])
 
     await report._annotate_report()
 
     nb = _nb_text(report)
     assert "The Big Picture" in nb
-    assert "Overall narrative." in nb
-    assert "Note for A." in nb
-    assert "Note for B." in nb
+    # Prose is interleaved around the chart's reconstructed spec, blog-post style.
+    assert nb.index("Lead-in prose.") < nb.index("lm.Pipeline.from_spec") < nb.index("Closing prose.")
 
-    # The overall story and the notes must both survive HTML export too.
+    # The flow must survive HTML export too.
     html = report.to_html()
     assert "The Big Picture" in html
-    assert "Overall narrative." in html
-    assert "Note for A." in html
-    assert "Note for B." in html
+    assert "Lead-in prose." in html
+    assert "Closing prose." in html
 
 
-async def test_report_annotate_only_selected_sections(llm):
-    from lumen.ai.agents.story import SectionNote, Story
+async def test_report_annotate_appends_unreferenced_views(llm, tiny_source):
+    from lumen.ai.agents.story import Story, StoryBlock
 
     report = Report(
-        Section(A(order=[]), title='Section A'),
-        Section(B(order=[]), title='Section B'),
-        title='My Report',
+        Section(ChartAction(source=tiny_source, label='Chart A'), title='Section A'),
+        Section(ChartAction(source=tiny_source, label='Chart B'), title='Section B'),
+        title='R',
+        llm=llm,
+    )
+    await report.execute()
+    # The story only references the second chart; the first must not be dropped.
+    llm.set_responses([Story(chain_of_thought="c", title="S", blocks=[
+        StoryBlock(prose="Only about the second."),
+        StoryBlock(view=2),
+    ])])
+
+    await report._annotate_report()
+
+    assert _kinds(report).count("view") == 2
+
+
+async def test_report_annotate_only_selected_sections(llm, tiny_source):
+    from lumen.ai.agents.story import Story, StoryBlock
+
+    report = Report(
+        Section(ChartAction(source=tiny_source, label='Chart A'), title='Section A'),
+        Section(ChartAction(source=tiny_source, label='Chart B'), title='Section B'),
+        title='R',
         llm=llm,
     )
     await report.execute()
     report[1].include_in_export = False
-    llm.set_responses([
-        Story(chain_of_thought="c", title="S", markdown="overall"),
-        SectionNote(chain_of_thought="c", markdown="Note for A."),
-    ])
+    llm.set_responses([Story(chain_of_thought="c", title="S", blocks=[
+        StoryBlock(prose="p"), StoryBlock(view=1),
+    ])])
 
     await report._annotate_report()
 
-    nb = _nb_text(report)
-    assert "Note for A." in nb
-    assert "Note for B." not in nb
+    # Only the selected section's chart is catalogued.
+    assert _kinds(report).count("view") == 1
 
 
-async def test_report_annotate_is_idempotent(llm):
-    from lumen.ai.agents.story import SectionNote, Story
+async def test_report_annotate_is_idempotent(llm, tiny_source):
+    from lumen.ai.agents.story import Story, StoryBlock
 
-    report = Report(Section(A(order=[]), title='Section A'), title='My Report', llm=llm)
+    report = Report(
+        Section(ChartAction(source=tiny_source, label='Chart A'), title='Section A'),
+        title='R', llm=llm,
+    )
     await report.execute()
 
-    llm.set_responses([Story(chain_of_thought="c", title="S1", markdown="v1"), SectionNote(chain_of_thought="c", markdown="n1")])
+    llm.set_responses([Story(chain_of_thought="c", title="S1", blocks=[
+        StoryBlock(prose="v1"), StoryBlock(view=1),
+    ])])
     await report._annotate_report()
-    llm.set_responses([Story(chain_of_thought="c", title="S2", markdown="v2"), SectionNote(chain_of_thought="c", markdown="n2")])
+    llm.set_responses([Story(chain_of_thought="c", title="S2", blocks=[
+        StoryBlock(prose="v2"), StoryBlock(view=1),
+    ])])
     await report._annotate_report()
 
     nb = _nb_text(report)
-    assert "v2" in nb and "n2" in nb
-    assert "v1" not in nb and "n1" not in nb
+    assert "v2" in nb and "v1" not in nb
+    assert "S2" in nb and "S1" not in nb
 
 
-async def test_report_annotate_no_llm_is_noop():
-    report = Report(Section(A(order=[]), title='Section A'), title='My Report')
+async def test_report_annotate_no_llm_is_noop(tiny_source):
+    report = Report(
+        Section(ChartAction(source=tiny_source, label='Chart A'), title='Section A'),
+        title='R',
+    )
     await report.execute()
 
     await report._annotate_report()  # must not raise
 
-    nb = _nb_text(report)
-    assert "A done" in nb
+    # Without a story the chart still exports on its own.
+    assert "lm.Pipeline.from_spec" in _nb_text(report)
+
+
+async def test_report_story_dialog_presets_and_generate(llm, tiny_source):
+    from lumen.ai.agents.story import Story, StoryBlock
+
+    report = Report(
+        Section(ChartAction(source=tiny_source, label='Chart A'), title='Section A'),
+        title='R', llm=llm,
+    )
+    await report.execute()
+
+    # Clicking Annotate opens the guidance dialog instead of generating immediately.
+    report._open_story_dialog()
+    assert report._story_dialog.open is True
+
+    # A tone preset fills the guidance box.
+    report._apply_story_preset(text="Executive summary tone.")
+    assert report._story_guidance.value == "Executive summary tone."
+
+    # Generate runs the annotation with that guidance and closes the dialog.
+    llm.set_responses([Story(chain_of_thought="c", title="S", blocks=[
+        StoryBlock(prose="p"), StoryBlock(view=1),
+    ])])
+    await report._generate_story()
+
+    assert report._story_dialog.open is False
+    assert "prose" in _kinds(report)
+
+
+async def test_report_view_toggle_switches_story_and_sections(llm, tiny_source):
+    from lumen.ai.agents.story import Story, StoryBlock
+
+    report = Report(
+        Section(ChartAction(source=tiny_source, label='Chart A'), title='Section A'),
+        title='R', llm=llm,
+    )
+    await report.execute()
+    llm.set_responses([Story(chain_of_thought="c", title="Headline", blocks=[
+        StoryBlock(prose="Story prose."), StoryBlock(view=1),
+    ])])
+    await report._annotate_report()
+
+    # After annotate the story is shown and exported, and the toggle is visible.
+    assert report._show_story is True
+    assert report._view_toggle.visible is True
+    assert "Story prose." in _nb_text(report)
+
+    # Toggle to the section view: the story is kept, export follows the sections.
+    report._toggle_view()
+    assert report._show_story is False
+    assert report._story_blocks  # still remembered
+    nb_sections = _nb_text(report)
+    assert "Story prose." not in nb_sections
+    assert "lm.Pipeline.from_spec" in nb_sections
+
+    # Toggle again: the story returns without regenerating.
+    report._toggle_view()
+    assert report._show_story is True
+    assert "Story prose." in _nb_text(report)
+
+
+async def test_report_reset_discards_story(llm, tiny_source):
+    from lumen.ai.agents.story import Story, StoryBlock
+
+    report = Report(
+        Section(ChartAction(source=tiny_source, label='Chart A'), title='Section A'),
+        title='R', llm=llm,
+    )
+    await report.execute()
+    llm.set_responses([Story(chain_of_thought="c", title="H", blocks=[
+        StoryBlock(prose="p"), StoryBlock(view=1),
+    ])])
+    await report._annotate_report()
+    assert report._story_blocks
+
+    report.reset()
+
+    assert report._story_blocks == []
+    assert report._show_story is False
+    assert report._view_toggle.visible is False
 
 
 async def test_report_outline_defaults_empty():