Ctrl K
ComunidadGuías…

tests/test_followup.py

Rama main · réplica de lectura del repositorio original.

Archivo original ↗Preguntar a KiBot
"""
Tests for follow-up rounds: the coverage assessment of a round becomes the
search queries of the next round.
"""

from src.pipeline.followup import MAX_ASPECTS_PER_QUESTION, build_followup_queries
from src.pipeline.models import ResearchPlan, ResearchQuestion


def _plan():
    return ResearchPlan(questions=[
        ResearchQuestion(id="F1", question="How much power does the DGX B300 draw?",
                         search_terms=["DGX B300 power"], search_langs=["en", "de"],
                         search_terms_by_lang={"en": ["DGX B300 power"], "de": ["DGX B300 Stromaufnahme"]}),
        ResearchQuestion(id="F2", question="What does it cost?", search_terms=["DGX B300 price"],
                         search_langs=["en"], search_terms_by_lang={"en": ["DGX B300 price"]}),
        ResearchQuestion(id="F3", question="Who sells it in Germany?", search_terms=["DGX B300 reseller"],
                         search_langs=["en"], search_terms_by_lang={"en": ["DGX B300 reseller"]}),
        ResearchQuestion(id="F4", question="Is it rack-mountable?", search_terms=["DGX B300 rack"],
                         search_langs=["en"], search_terms_by_lang={"en": ["DGX B300 rack"]}),
    ])


def test_partial_question_gets_one_query_per_missing_aspect_and_language():
    fq = build_followup_queries(_plan(), {
        "F1": {"coverage": "partial", "missing_aspects": ["idle power", "- peak power."]},
    })
    assert len(fq) == 1 and fq[0]["question_id"] == "F1"
    assert fq[0]["search_terms"] == {
        "en": ["DGX B300 power idle power", "DGX B300 power peak power"],
        "de": ["DGX B300 Stromaufnahme idle power", "DGX B300 Stromaufnahme peak power"],
    }
    assert fq[0]["search_langs"] == ["en", "de"]


def test_answered_and_filter_blocked_questions_get_nothing():
    fq = build_followup_queries(_plan(), {
        "F1": {"coverage": "answered", "missing_aspects": ["x"]},
        "F2": {"coverage": "filter_blocked", "missing_aspects": ["y"]},
    })
    assert fq == []


def test_unanswered_without_aspects_searches_the_question_itself():
    fq = build_followup_queries(_plan(), {"F3": {"coverage": "unanswered", "missing_aspects": []}})
    assert fq[0]["search_terms"] == {"en": ["Who sells it in Germany?"]}


def test_partial_without_aspects_gets_nothing():
    """Nothing concrete to search for — better no query than a repeated one."""
    assert build_followup_queries(_plan(), {"F4": {"coverage": "partial", "missing_aspects": []}}) == []


def test_aspects_are_capped_and_queries_shortened():
    long_aspect = "one two three four five six seven eight nine ten eleven twelve"
    fq = build_followup_queries(_plan(), {
        "F2": {"coverage": "partial", "missing_aspects": [long_aspect, "b", "c", "d", "e"]},
    })
    terms = fq[0]["search_terms"]["en"]
    assert len(terms) == MAX_ASPECTS_PER_QUESTION
    assert all(len(t.split()) <= 10 for t in terms)


def test_questions_without_assessment_are_skipped():
    assert build_followup_queries(_plan(), {}) == []