llm-ingestion-okf/tests/test_subquestions.py

"""Several sub-questions in ONE call (v1.1 order C, C2).

A broad question is asked best as two to four narrow ones in the collection's
own words. Until C2 that meant one call per sub-question and a merge done by
whoever was asking -- the search gate did it itself, in its own code. Now
`consume.build_multi_payload` (and `okf consume --question A --question B`,
and `okf_ask` with `questions`) ranks each sub-question on ONE load of the
bundle and interleaves the deliveries: first excerpt of each sub-question in
turn, then the second, duplicates dropped, cut at `k`. Every excerpt says
which sub-questions it answered; the first index is the one whose text it
carries, since a large concept is delivered as the passage its OWN
sub-question found.

The expected merge below is written from the definition, over the SINGLE
payloads, so the product's merge is held against an independent reading of
the same rule and not against itself.
"""

from __future__ import annotations

import json
import sys
from collections.abc import Sequence
from pathlib import Path

import pytest

from llm_ingestion_okf import consume, contract_check, mcp_server
from llm_ingestion_okf import skill as okf_skill

TOOLS = Path(__file__).resolve().parent.parent / "tools"
if str(TOOLS) not in sys.path:
    sys.path.insert(0, str(TOOLS))

import okf_retrieval_gate as retrieval  # noqa: E402

HEATING = "How is the cabin heated in winter?"
ENGINE = "How is the boat engine serviced?"
APPLES = "When are the garden apples picked?"
UNCOVERED = "zzqx vvkw"


def _concept(slug: str, title: str, body: str) -> retrieval.ConceptSpec:
    return retrieval.ConceptSpec(slug=slug, title=title, body=body)


@pytest.fixture(scope="module")
def bundle(tmp_path_factory: pytest.TempPathFactory) -> Path:
    spec = retrieval.BundleSpec(
        "subquestions-synthetic",
        (
            retrieval.DocumentSpec(
                "cabin",
                "cabin.md",
                (
                    _concept(
                        "heating", "Heating", "The cabin is heated by a wood stove in winter."
                    ),
                    _concept("water", "Water", "Water comes from the well; the cabin pipes drain."),
                    _concept("roof", "Roof", "The cabin roof is cleared of snow in winter."),
                ),
            ),
            retrieval.DocumentSpec(
                "boat",
                "boat.md",
                (
                    _concept("engine", "Engine", "The boat engine is serviced every spring."),
                    _concept("sails", "Sails", "The boat sails are dried before storage."),
                    _concept("anchor", "Anchor", "The anchor chain of the boat is checked."),
                ),
            ),
            retrieval.DocumentSpec(
                "garden",
                "garden.md",
                (
                    _concept("apples", "Apples", "The garden apples are picked in September."),
                    _concept("roses", "Roses", "The roses in the garden are pruned in March."),
                ),
            ),
        ),
    )
    return retrieval.build_bundle(tmp_path_factory.mktemp("subquestions") / "bundle", spec)


def _ids(payload: dict[str, object]) -> list[str]:
    excerpts = payload["excerpts"]
    assert isinstance(excerpts, list)
    return [str(excerpt["concept_id"]) for excerpt in excerpts]


def _interleaved(lists: Sequence[Sequence[str]], cap: int) -> list[str]:
    """The rule, written from its definition: position by position, each
    sub-question in turn, a concept already taken skipped, stop at `cap`."""
    out: list[str] = []
    for position in range(max(len(ids) for ids in lists)):
        for ids in lists:
            if position < len(ids) and ids[position] not in out and len(out) < cap:
                out.append(ids[position])
    return out


def test_the_interleave_takes_turns_skips_what_is_taken_and_cuts_at_k() -> None:
    first = [{"concept_id": "a"}, {"concept_id": "b"}, {"concept_id": "c"}]
    second = [{"concept_id": "b"}, {"concept_id": "d"}]
    merged = consume.interleave([first, second], k=3, limit=consume.DEFAULT_LIMIT)
    assert [(excerpt["concept_id"], named) for excerpt, named in merged] == [
        ("a", [0]),
        ("b", [1, 0]),
        ("d", [1]),
    ]
    assert consume.interleave([], k=5, limit=consume.DEFAULT_LIMIT) == []


def test_the_interleave_never_spends_more_than_the_limit() -> None:
    large = {"concept_id": "a", "text": "x" * 2_000}
    small = {"concept_id": "b", "text": "y"}
    limit = consume.excerpt_weight(small) + 10
    merged = consume.interleave([[large], [small]], k=8, limit=limit)
    assert [excerpt["concept_id"] for excerpt, _ in merged] == ["b"]


def test_deliveries_are_interleaved_deduplicated_and_cut_at_k(bundle: Path) -> None:
    questions = [HEATING, ENGINE, APPLES]
    singles = [_ids(consume.build_payload(bundle, question=q, k=3)) for q in questions]
    # The premise: each sub-question reaches something the others do not, or
    # an interleave and a concatenation could not be told apart.
    assert len({ids[0] for ids in singles}) == 3
    payload = consume.build_multi_payload(bundle, questions=questions, k=4)
    assert _ids(payload) == _interleaved(singles, 4)


def test_every_excerpt_names_the_subquestions_it_answered(bundle: Path) -> None:
    questions = [HEATING, ENGINE, "What happens to the cabin in winter?"]
    singles = [consume.build_payload(bundle, question=q, k=3) for q in questions]
    payload = consume.build_multi_payload(bundle, questions=questions, k=3)
    excerpts = payload["excerpts"]
    assert isinstance(excerpts, list) and excerpts
    shared = 0
    for excerpt in excerpts:
        named = excerpt["subquestions"]
        answered = [i for i, single in enumerate(singles) if excerpt["concept_id"] in _ids(single)]
        assert sorted(named) == answered
        shared += len(named) > 1
        # The text is the one its FIRST sub-question delivered.
        placing = singles[named[0]]["excerpts"]
        assert isinstance(placing, list)
        original = next(e for e in placing if e["concept_id"] == excerpt["concept_id"])
        assert {key: value for key, value in excerpt.items() if key != "subquestions"} == original
    assert shared >= 1, "the premise: two sub-questions reach one concept"


def test_one_question_is_the_single_payload_byte_for_byte(bundle: Path) -> None:
    assert consume.serialise(
        consume.build_multi_payload(bundle, questions=[HEATING])
    ) == consume.serialise(consume.build_payload(bundle, question=HEATING))


def test_the_same_subquestions_give_the_same_bytes(bundle: Path) -> None:
    questions = [HEATING, ENGINE, APPLES, UNCOVERED]
    first = consume.serialise(consume.build_multi_payload(bundle, questions=questions))
    second = consume.serialise(consume.build_multi_payload(bundle, questions=questions))
    assert first == second


def test_the_denominators_close_and_the_contract_checker_accepts_it(bundle: Path) -> None:
    payload = consume.build_multi_payload(bundle, questions=[HEATING, ENGINE], k=3)
    counts = payload["denominators"]
    assert isinstance(counts, dict)
    assert counts["considered"] == counts["withheld"] + counts["delivered"] == 8
    withheld = payload["withheld"]
    assert isinstance(withheld, dict)
    assert withheld["total"] == counts["withheld"]
    assert sum(withheld["by_rule"].values()) == withheld["total"]
    report = contract_check.check(okf_skill.render_generic(), payload)
    assert report.findings == ()


def test_a_concept_another_subquestion_delivered_and_the_cut_dropped_is_below_k(
    bundle: Path,
) -> None:
    questions = [HEATING, ENGINE]
    singles = [_ids(consume.build_payload(bundle, question=q, k=3)) for q in questions]
    payload = consume.build_multi_payload(bundle, questions=questions, k=3)
    dropped = {cid for ids in singles for cid in ids} - set(_ids(payload))
    assert dropped, "the premise: the merge's cut drops something a sub-question delivered"
    withheld = payload["withheld"]
    assert isinstance(withheld, dict)
    rules = {entry["concept_id"]: entry["rule"] for entry in withheld["nearest"]}
    assert all(rules[cid] == "below_k" for cid in dropped)


def test_the_payload_states_each_subquestion_and_its_coverage(bundle: Path) -> None:
    payload = consume.build_multi_payload(bundle, questions=[HEATING, UNCOVERED])
    assert payload["questions"] == [HEATING, UNCOVERED]
    assert "question" not in payload
    coverage = payload["coverage"]
    assert isinstance(coverage, dict)
    per = coverage["subquestions"]
    assert [entry["question"] for entry in per] == [HEATING, UNCOVERED]
    assert [entry["weak"] for entry in per] == [False, True]
    assert per[1]["absent_terms"] == ["zzqx", "vvkw"]
    assert coverage["weak_subquestions"] == [1]
    # One sub-question the collection covers: the whole is not read as uncovered.
    assert coverage["weak"] is False


def test_every_subquestion_weak_makes_the_whole_weak(bundle: Path) -> None:
    coverage = consume.build_multi_payload(bundle, questions=[UNCOVERED, "qqzv wwkx"])["coverage"]
    assert isinstance(coverage, dict)
    assert coverage["weak"] is True
    assert coverage["weak_subquestions"] == [0, 1]


def test_no_question_is_refused(bundle: Path) -> None:
    with pytest.raises(consume.ConsumeError) as raised:
        consume.build_multi_payload(bundle, questions=[])
    assert raised.value.code == "question_missing"
    with pytest.raises(consume.ConsumeError) as raised:
        consume.build_multi_payload(bundle, questions=[HEATING, "  "])
    assert raised.value.code == "question_missing"


def test_the_command_line_takes_the_question_more_than_once(bundle: Path, tmp_path: Path) -> None:
    out = tmp_path / "payload.json"
    code = consume.main(
        [str(bundle), "--question", HEATING, "--question", ENGINE, "--out", str(out)]
    )
    assert code == 0
    written = out.read_text(encoding="utf-8")
    assert written == consume.serialise(
        consume.build_multi_payload(bundle, questions=[HEATING, ENGINE])
    )


def test_okf_ask_takes_several_questions_in_one_call(bundle: Path) -> None:
    surface = mcp_server.build_surface(bundle=bundle, roots=())
    result = mcp_server.call_ask(surface, {"questions": [HEATING, ENGINE]})
    assert result["questions"] == [HEATING, ENGINE]
    payload = result["answers"][0]["payload"]
    assert json.dumps(payload, sort_keys=True) == json.dumps(
        consume.build_multi_payload(bundle, questions=[HEATING, ENGINE]), sort_keys=True
    )


def test_okf_ask_refuses_both_forms_at_once(bundle: Path) -> None:
    surface = mcp_server.build_surface(bundle=bundle, roots=())
    with pytest.raises(mcp_server.ToolError) as raised:
        mcp_server.call_ask(surface, {"question": HEATING, "questions": [ENGINE]})
    assert raised.value.code == "question_ambiguous"


def test_four_subquestions_over_large_concepts_stay_under_a_tool_reply(tmp_path: Path) -> None:
    """The worst case the passage cut exists for: every delivered concept is far
    over `PASSAGE_CHARS`. A client keeps a tool reply of 25 000 tokens; at a
    pessimistic two bytes a token that is 50 000 bytes of text."""
    words = ("stove", "engine", "apples", "roof")
    spec = retrieval.BundleSpec(
        "large-concepts",
        tuple(
            retrieval.DocumentSpec(
                f"doc-{word}",
                f"doc-{word}.md",
                tuple(
                    retrieval.ConceptSpec(
                        slug=f"part-{n}",
                        title=f"{word.title()} part {n}",
                        body=f"The {word} is described here in detail. ",
                        repeat=400,
                    )
                    for n in range(3)
                ),
            )
            for word in words
        ),
    )
    bundle = retrieval.build_bundle(tmp_path / "bundle", spec)
    surface = mcp_server.build_surface(bundle=bundle, roots=())
    result = mcp_server.call_ask(
        surface, {"questions": [f"How is the {w} described?" for w in words]}
    )
    payload = result["answers"][0]["payload"]
    assert len(payload["excerpts"]) == consume.DEFAULT_K
    assert all(len(excerpt["text"]) > consume.PASSAGE_CHARS // 2 for excerpt in payload["excerpts"])
    text = mcp_server._tool_result(result)["content"][0]["text"]
    assert len(text.encode("utf-8")) < 50_000