| from __future__ import annotations |
|
|
| import asyncio |
| import json |
| import random |
|
|
| import pytest |
| from fake_teacher import FakeTeacher |
|
|
| from deg.datagen.hard import ( |
| SIDECAR, |
| check_hard, |
| leaked, |
| make_nullable, |
| pick_targets, |
| split_planted, |
| ) |
| from deg.datagen.run import Runner, build_parser |
| from deg.datagen.validate import Reject, check_fill |
| from deg.schema import compile_schema |
|
|
| ORDER = {"type": "object", "additionalProperties": False, "properties": { |
| "order_id": {"type": "string", "description": "The order number"}, |
| "ship_date": {"type": "string", "format": "date", "description": "Date it shipped"}, |
| "total": {"type": "number", "description": "Order total"}, |
| "items": {"type": "array", "items": {"type": "string"}}, |
| "rush": {"type": "boolean"}}, "required": ["order_id", "rush"]} |
|
|
|
|
| def _checked(ctx, value): |
| plan = compile_schema(ORDER) |
| return plan, check_fill(plan, ctx, value) |
|
|
|
|
| def test_pick_targets_and_nullable(): |
| plan = compile_schema(ORDER) |
| for kind in ("absent", "distractor", "correction"): |
| got = pick_targets(plan, kind, random.Random(0)) |
| assert got and all("*" not in f.pointer and f.mode.startswith("extract.") for f in got) |
| out = make_nullable(ORDER, ["/order_id", "/total"]) |
| assert out["properties"]["order_id"]["type"] == ["string", "null"] |
| assert out["required"] == ["rush"] and ORDER["required"] == ["order_id", "rush"] |
| assert ORDER["properties"]["order_id"]["type"] == "string" |
| assert compile_schema(out).field("/order_id").nullable |
|
|
|
|
| def test_split_planted(): |
| doc, planted = split_planted(f"Hello there.\n{SIDECAR}\n{{\"/total\": \"$40.00\"}}") |
| assert doc == "Hello there." and planted == {"/total": "$40.00"} |
| assert split_planted("just a doc ") == ("just a doc", None) |
|
|
|
|
| def test_distractor_and_correction_checks(): |
| ctx = ("Order A-1001 shipped 2026-03-05. Total due $120.50. An earlier quote said $99.00. " |
| "Correction: the order number is actually A-1002.") |
| value = {"order_id": "A-1002", "ship_date": "2026-03-05", "total": 120.5, "rush": False} |
| plan, checked = _checked(ctx, value) |
| total, oid = plan.field("/total"), plan.field("/order_id") |
| got = check_hard("distractor", [total], ctx, checked.value, {"/total": "$99.00"}) |
| assert got["planted"] == {"/total": "$99.00"} |
| check_hard("correction", [oid], ctx, checked.value, {"/order_id": "A-1001"}) |
| with pytest.raises(Reject, match="distractor_chosen"): |
| check_hard("distractor", [total], ctx, checked.value, {"/total": 120.5}) |
| with pytest.raises(Reject, match="planted_not_in_context"): |
| check_hard("distractor", [total], ctx, checked.value, {"/total": "$77"}) |
| with pytest.raises(Reject, match="planted_missing"): |
| check_hard("correction", [oid], ctx, checked.value, None) |
| |
| ctx2 = "Order A-1002 placed. Earlier note: order A-1001. Shipped 2026-03-05, rush." |
| _, checked2 = _checked(ctx2, {"order_id": "A-1002", "rush": True}) |
| with pytest.raises(Reject, match="correction_order"): |
| check_hard("correction", [oid], ctx2, checked2.value, {"/order_id": "A-1001"}) |
|
|
|
|
| def test_absent_checks(): |
| schema = make_nullable(ORDER, ["/total"]) |
| plan = compile_schema(schema) |
| ctx = "Order A-1001 shipped 2026-03-05; no rush." |
| checked = check_fill(plan, ctx, {"order_id": "A-1001", "total": None, "rush": False}) |
| check_hard("absent", [plan.field("/total")], ctx, checked.value, None) |
| ctx3 = "Order A-1001 total 15. shipped 2026-03-05; no rush." |
| filled = check_fill(plan, ctx3, {"order_id": "A-1001", "total": 15, "rush": False}) |
| with pytest.raises(Reject, match="absent_filled"): |
| check_hard("absent", [plan.field("/total")], ctx3, filled.value, None) |
| f = [plan.field("/total")] |
| assert leaked(f, ctx, '{"/total": null}') is None |
| assert leaked(f, ctx, '{"/total": "not stated"}') is None |
| assert leaked(f, ctx3, '{"/total": "total 15"}') == "/total" |
| with pytest.raises(Reject, match="absent_check_json"): |
| leaked(f, ctx, "nope") |
|
|
|
|
| def _runner(tmp_path): |
| args = build_parser().parse_args(["--out", str(tmp_path / "o"), "--no-clef", |
| "--no-flash-labels", "--hard-frac", "1.0"]) |
| return Runner(args, [], None) |
|
|
|
|
| def _index_for(kind): |
| from deg.datagen.domains import sample_spec |
| return next(i for i in range(500) if sample_spec(0, i, 0.0, 1.0).hard == kind) |
|
|
|
|
| def test_runner_distractor_end_to_end(tmp_path): |
| runner = _runner(tmp_path) |
| state = {} |
| doc = ("Hi team, order A-1002 is confirmed and shipped 2026-03-05. Total charged: $120.50 " |
| "(the earlier quote of $99.00 expired). Not a rush order. Please file the " |
| "paperwork with the usual vendor and let me know if anything looks off. Thanks!") |
|
|
| def respond(prompt): |
| if prompt.startswith("Domain:"): |
| return json.dumps(ORDER) |
| if prompt.startswith("Write one realistic document"): |
| assert "HARD CASE (distractors)" in prompt |
| targets = [p for p in ("/order_id", "/ship_date", "/total") if f" * {p}:" in prompt] |
| planted = {"/order_id": "A-1001", "/ship_date": "2026-02-01", "/total": "$99.00"} |
| state["targets"] = {p: planted[p] for p in targets} |
| return doc + " Old ref A-1001 from 2026-02-01." |
| if prompt.startswith("Now output ONLY a JSON object"): |
| return json.dumps(state["targets"]) |
| if prompt.startswith("Fill the JSON Schema"): |
| return json.dumps({"order_id": "A-1002", "ship_date": "2026-03-05", |
| "total": 120.5, "items": [], "rush": False}) |
| raise AssertionError(prompt[:60]) |
|
|
| row, kept = asyncio.run(runner.example(_index_for("distractor"), FakeTeacher(respond))) |
| assert kept, row |
| assert row["hard"] == "distractor" and row["hard_case"]["planted"] |
| assert SIDECAR not in row["context"] |
|
|
|
|
| def test_runner_absent_rejects_leak(tmp_path): |
| runner = _runner(tmp_path) |
|
|
| def respond(prompt): |
| if prompt.startswith("Domain:"): |
| return json.dumps(ORDER) |
| if prompt.startswith("Write one realistic document"): |
| assert "HARD CASE (absent)" in prompt |
| return ("Order A-1002 shipped 2026-03-05, total 120.50, not urgent at all. " |
| "Please confirm receipt with the warehouse and update the ledger today.") |
| if prompt.startswith("Fill the JSON Schema"): |
| return json.dumps({"order_id": None, "ship_date": None, "total": None, |
| "rush": False}) |
| if prompt.startswith("For each field below"): |
| return json.dumps({"/order_id": "A-1002", "/ship_date": "2026-03-05", |
| "/total": "120.50"}) |
| raise AssertionError(prompt[:60]) |
|
|
| row, kept = asyncio.run(runner.example(_index_for("absent"), FakeTeacher(respond))) |
| assert not kept and row["reason"] in ("absent_leaked", "required_missing") |
| assert row["hard"] == "absent" |
|
|
|
|
| def test_planted_key_variants(): |
| from deg.datagen.hard import planted_for |
|
|
| assert planted_for({"dimensions.length": 4}, "/dimensions/length", 2) == 4 |
| assert planted_for({"dimensions/length": 4}, "/dimensions/length", 2) == 4 |
| assert planted_for({"Length": 4}, "/dimensions/length", 2) == 4 |
| assert planted_for({"whatever": 4}, "/dimensions/length", 1) == 4 |
| assert planted_for({"whatever": 4}, "/dimensions/length", 2) is None |
|
|
