from app.pipeline.nodes.generate import _select_chunks_for_prompt def _chunk(title: str, section: str = "Overview") -> dict: return { "text": f"Info about {title}", "metadata": { "source_title": title, "section": section, "source_url": "", "source_type": "project", }, } def test_select_chunks_prefers_explicitly_mentioned_source_title() -> None: chunks = [ _chunk("Sorting Demo"), _chunk("EchoEcho"), _chunk("Donut.asm"), ] selected = _select_chunks_for_prompt( "Is the source code on GitHub up-to-date with the Sorting Demo live demo?", chunks, ) assert selected assert all(c["metadata"]["source_title"] == "Sorting Demo" for c in selected) def test_select_chunks_falls_back_to_top_ranked_when_no_title_match() -> None: chunks = [_chunk("A"), _chunk("B"), _chunk("C")] selected = _select_chunks_for_prompt("What technologies are used?", chunks) assert selected == chunks