DGXAI/driftcall-env
0
1"""Cell 07 — Procedural task-brief generator.2 3Implements docs/modules/task_generator.md. Pure, seeded, deterministic4expansion of a YAML template library into concrete ``GoalSpec`` briefs5for ``DriftCallEnv.reset()`` (DESIGN.md §4.2, §8.3, §8.4).6 7Contract: identical ``(seed, stage, language_weights)`` triples always8produce byte-identical ``GoalSpec.seed_utterance`` after NFC9normalization. No global mutable state; no ``random.random()``; no10``time.time()``; no ``hash()``. All stochastic choices thread through11``random.Random(stable_sub_seed(seed, tag))`` where ``stable_sub_seed``12uses ``hashlib.blake2b(digest_size=8)``.13"""14 15from __future__ import annotations16 17import hashlib18import random19import re20import string21import unicodedata22from collections.abc import Iterator, Mapping23from dataclasses import dataclass24from datetime import date, timedelta25from pathlib import Path26from typing import Any, Literal, cast27 28import yaml29 30from cells.step_04_models import GoalSpec31 32# ---------------------------------------------------------------------------33# Public literal types34# ---------------------------------------------------------------------------35 36LanguageCode = Literal["hi", "ta", "kn", "en", "hinglish"]37Domain = Literal["airline", "cab", "restaurant", "hotel"]38 39_LANGUAGE_CODES: frozenset[str] = frozenset({"hi", "ta", "kn", "en", "hinglish"})40_DOMAINS: frozenset[str] = frozenset({"airline", "cab", "restaurant", "hotel"})41_VALID_STAGES: frozenset[int] = frozenset({1, 2, 3})42 43# Fixed reference date for deterministic date sampling (task_generator.md §3.3).44_REFERENCE_DATE: date = date(2026, 4, 25)45_DATE_WINDOW_DAYS: int = 6046 47# SMS-length bound for ASR input (§3.6 invariant 7).48_MAX_UTTERANCE_LEN: int = 28049 50# Built-in slot conventions — §3.3 of task_generator.md. Templates may51# override by declaring slot_distributions explicitly; otherwise these52# name-based defaults apply.53_DATE_SLOT_NAMES: frozenset[str] = frozenset(54 {55 "when",56 "checkin",57 "checkout",58 "date",59 "departure",60 "arrival",61 "return_when",62 "new_when",63 }64)65_INTER_CITY_SLOT_NAMES: frozenset[str] = frozenset(66 {"from", "to", "city", "origin", "destination"}67)68_INTRA_CITY_SLOT_NAMES: frozenset[str] = frozenset({"pickup", "drop"})69 70# Default domain → city-code tuples (IATA-style). Authored here so the71# generator is self-contained without requiring the YAML library to72# declare a cities_by_domain block.73_DEFAULT_INTER_CITIES: tuple[str, ...] = (74 "HYD",75 "BLR",76 "DEL",77 "BOM",78 "MAA",79 "CCU",80 "PNQ",81 "AMD",82 "JAI",83 "GOI",84)85_DEFAULT_INTRA_CITIES: tuple[str, ...] = (86 "Koramangala",87 "Indiranagar",88 "Whitefield",89 "Andheri",90 "Bandra",91 "Powai",92 "Gurgaon",93 "Saket",94 "Banjara Hills",95 "Salt Lake",96)97_DEFAULT_CITIES_BY_DOMAIN: Mapping[Domain, tuple[str, ...]] = {98 "airline": _DEFAULT_INTER_CITIES,99 "hotel": _DEFAULT_INTER_CITIES,100 "restaurant": _DEFAULT_INTER_CITIES,101 "cab": _DEFAULT_INTRA_CITIES,102}103 104 105# ---------------------------------------------------------------------------106# Exception hierarchy (task_generator.md §5)107# ---------------------------------------------------------------------------108 109 110class TaskGeneratorError(Exception):111 """Base class for every failure raised by :mod:`step_07_task_generator`."""112 113 114class MissingSlotError(TaskGeneratorError):115 """Template variant references a ``{slot}`` placeholder not present in the filled SlotGrid."""116 117 118class InvalidLanguageError(TaskGeneratorError):119 """``language_weights`` contains a key outside :data:`LanguageCode`."""120 121 122class InvalidLanguageWeightError(TaskGeneratorError):123 """``language_weights`` is empty, has a negative value, sums off 1.0, or is all zero."""124 125 126class InvalidStageError(TaskGeneratorError):127 """``stage`` is not one of ``{1, 2, 3}``."""128 129 130class InvalidBudgetError(TaskGeneratorError):131 """Sampled numeric constraint falls outside the template's declared ``[low, high]`` range."""132 133 134class TemplateFileMissingError(TaskGeneratorError):135 """Template YAML file not found or unreadable."""136 137 138class TemplateSchemaError(TaskGeneratorError):139 """Template YAML present but fails schema validation."""140 141 142class UnicodeNormalizationError(TaskGeneratorError):143 """Rendered utterance fails NFC round-trip check (defensive)."""144 145 146class NoVariantForLanguageError(TaskGeneratorError):147 """Chosen template has no ``language_variants`` entry for the chosen language."""148 149 150# ---------------------------------------------------------------------------151# In-memory types (task_generator.md §4.2)152# ---------------------------------------------------------------------------153 154 155@dataclass(frozen=True)156class SlotDistribution:157 """Either an enum (``choices``) or a uniform numeric grid (``low``, ``high``, ``step``)."""158 159 kind: Literal["choices", "uniform", "date", "bool"]160 choices: tuple[str, ...] | None = None161 low: float | None = None162 high: float | None = None163 step: float | None = None164 165 166@dataclass(frozen=True)167class Template:168 template_id: str169 domain: Domain170 intent: str171 min_stage: Literal[1, 2, 3]172 required_slots: tuple[str, ...]173 optional_slots: tuple[str, ...]174 slot_distributions: Mapping[str, SlotDistribution]175 constraints_template: Mapping[str, SlotDistribution]176 drift_slot_tags: tuple[str, ...]177 language_variants: Mapping[LanguageCode, tuple[str, ...]]178 179 180@dataclass(frozen=True)181class TemplateLibrary:182 templates: tuple[Template, ...]183 cities_by_domain: Mapping[Domain, tuple[str, ...]]184 i18n: Mapping[LanguageCode, Mapping[str, str]]185 186 187@dataclass(frozen=True)188class SlotGrid:189 """Concrete slot values after expansion."""190 191 values: Mapping[str, object]192 193 194@dataclass(frozen=True)195class RawBrief:196 template_id: str197 domain: Domain198 intent: str199 slots: SlotGrid200 constraints: Mapping[str, object]201 language: LanguageCode202 203 204# ---------------------------------------------------------------------------205# Sub-seed helper (task_generator.md §3.1)206# ---------------------------------------------------------------------------207 208 209def stable_sub_seed(seed: int, tag: str) -> int:210 """Return a stable 64-bit integer derived from ``(seed, tag)``.211 212 Uses blake2b with ``digest_size=8`` so the formula is pinned and213 domain-separated across decision tags.214 """215 digest = hashlib.blake2b(f"{seed}:{tag}".encode(), digest_size=8).digest()216 return int.from_bytes(digest, "big")217 218 219# ---------------------------------------------------------------------------220# NFC helpers221# ---------------------------------------------------------------------------222 223 224def _nfc(text: str) -> str:225 return unicodedata.normalize("NFC", text)226 227 228def _assert_nfc(text: str, *, where: str) -> None:229 if not unicodedata.is_normalized("NFC", text):230 raise UnicodeNormalizationError(231 f"string at {where} failed NFC round-trip: {text!r}"232 )233 234 235# ---------------------------------------------------------------------------236# Template loader (task_generator.md §2.2, §3.4, §7 edge cases 1 & 8)237# ---------------------------------------------------------------------------238 239 240def _parse_distribution(raw: Mapping[str, Any], *, where: str) -> SlotDistribution:241 """Parse a single slot/constraint distribution block."""242 if "choices" in raw:243 choices = raw["choices"]244 if not isinstance(choices, list) or not choices:245 raise TemplateSchemaError(f"{where}: 'choices' must be non-empty list")246 norm_choices = tuple(_nfc(str(c)) for c in choices)247 return SlotDistribution(kind="choices", choices=norm_choices)248 if raw.get("distribution") == "uniform":249 for key in ("low", "high", "step"):250 if key not in raw:251 raise TemplateSchemaError(f"{where}: uniform missing '{key}'")252 low = float(raw["low"])253 high = float(raw["high"])254 step = float(raw["step"])255 if step <= 0:256 raise TemplateSchemaError(f"{where}: step must be > 0 (got {step})")257 if low > high:258 raise TemplateSchemaError(f"{where}: low > high ({low} > {high})")259 span = high - low260 # Grid must terminate cleanly at ``high`` (§7 edge case 8).261 # Use integer step check avoiding floating-point drift.262 ratio = span / step263 if abs(ratio - round(ratio)) > 1e-9:264 raise TemplateSchemaError(265 f"{where}: step grid misaligned "266 f"(low={low}, high={high}, step={step}) — (high-low) not divisible by step"267 )268 return SlotDistribution(kind="uniform", low=low, high=high, step=step)269 if raw.get("distribution") == "date":270 return SlotDistribution(kind="date")271 if raw.get("distribution") == "bool":272 return SlotDistribution(kind="bool")273 raise TemplateSchemaError(274 f"{where}: unrecognized distribution descriptor {dict(raw)!r}"275 )276 277 278def _parse_template(raw: Mapping[str, Any], *, where: str) -> Template:279 required_keys = (280 "template_id",281 "domain",282 "intent",283 "min_stage",284 "required_slots",285 "optional_slots",286 "constraints_template",287 "drift_slot_tags",288 "language_variants",289 )290 for key in required_keys:291 if key not in raw:292 raise TemplateSchemaError(f"{where}: missing required key {key!r}")293 294 template_id = _nfc(str(raw["template_id"]))295 domain_raw = str(raw["domain"])296 if domain_raw not in _DOMAINS:297 raise TemplateSchemaError(298 f"{where}: domain {domain_raw!r} not in {sorted(_DOMAINS)}"299 )300 min_stage = int(raw["min_stage"])301 if min_stage not in _VALID_STAGES:302 raise TemplateSchemaError(303 f"{where}: min_stage {min_stage} not in {sorted(_VALID_STAGES)}"304 )305 306 required_slots = tuple(_nfc(str(s)) for s in raw["required_slots"])307 optional_slots = tuple(_nfc(str(s)) for s in raw["optional_slots"])308 drift_slot_tags = tuple(_nfc(str(s)) for s in raw["drift_slot_tags"])309 310 slot_distributions_raw = raw.get("slot_distributions", {}) or {}311 slot_distributions: dict[str, SlotDistribution] = {}312 for name, block in slot_distributions_raw.items():313 slot_distributions[_nfc(str(name))] = _parse_distribution(314 block, where=f"{where}.slot_distributions.{name}"315 )316 317 constraints_template: dict[str, SlotDistribution] = {}318 for name, block in raw["constraints_template"].items():319 constraints_template[_nfc(str(name))] = _parse_distribution(320 block, where=f"{where}.constraints_template.{name}"321 )322 323 language_variants_raw = raw["language_variants"]324 if not isinstance(language_variants_raw, dict):325 raise TemplateSchemaError(f"{where}: language_variants must be a mapping")326 language_variants: dict[LanguageCode, tuple[str, ...]] = {}327 for lang, variants in language_variants_raw.items():328 if lang not in _LANGUAGE_CODES:329 raise TemplateSchemaError(330 f"{where}: language key {lang!r} not in {sorted(_LANGUAGE_CODES)}"331 )332 if not isinstance(variants, list) or not variants:333 raise TemplateSchemaError(334 f"{where}.language_variants.{lang}: must be non-empty list"335 )336 language_variants[cast("LanguageCode", lang)] = tuple(337 _nfc(str(v)) for v in variants338 )339 340 # Every template must have ≥ 1 variant per LanguageCode (§7 edge case 7).341 for code in _LANGUAGE_CODES:342 if code not in language_variants:343 raise TemplateSchemaError(344 f"{where}: language_variants missing required code {code!r}"345 )346 347 # Static placeholder scan (§7 edge case 1).348 declared_placeholders = (349 set(required_slots)350 | set(optional_slots)351 | set(constraints_template.keys())352 )353 for lang, variants in language_variants.items():354 for variant in variants:355 for placeholder in _iter_placeholders(variant):356 if placeholder not in declared_placeholders:357 raise TemplateSchemaError(358 f"{where}.language_variants.{lang}: variant references "359 f"undeclared placeholder {placeholder!r} in {variant!r}"360 )361 362 return Template(363 template_id=template_id,364 domain=cast("Domain", domain_raw),365 intent=_nfc(str(raw["intent"])),366 min_stage=cast("Literal[1, 2, 3]", min_stage),367 required_slots=required_slots,368 optional_slots=optional_slots,369 slot_distributions=slot_distributions,370 constraints_template=constraints_template,371 drift_slot_tags=drift_slot_tags,372 language_variants=language_variants,373 )374 375 376def _iter_placeholders(fmt: str) -> Iterator[str]:377 """Yield placeholder names in a format string (ignores literals)."""378 for _literal, field_name, _spec, _conv in string.Formatter().parse(fmt):379 if field_name is not None and field_name != "":380 yield field_name381 382 383def load_templates(384 path: str | Path = "data/task_briefs/templates.yaml",385 i18n_path: str | Path | None = None,386) -> TemplateLibrary:387 """Parse the template YAML file and return an in-memory :class:`TemplateLibrary`.388 389 ``i18n_path`` defaults to ``data/task_briefs/i18n.yaml`` alongside390 ``path``. All strings are NFC-normalized on read (§3.4).391 """392 templates_path = Path(path)393 if not templates_path.exists():394 raise TemplateFileMissingError(f"templates YAML not found: {templates_path}")395 396 if i18n_path is None:397 i18n_path = templates_path.parent / "i18n.yaml"398 i18n_path = Path(i18n_path)399 400 try:401 with templates_path.open("r", encoding="utf-8") as fh:402 raw_templates = yaml.safe_load(fh)403 except yaml.YAMLError as exc:404 raise TemplateSchemaError(f"templates YAML malformed: {exc}") from exc405 406 if raw_templates is None:407 raise TemplateSchemaError("templates YAML is empty")408 409 parsed_templates: list[Template] = []410 cities_by_domain: dict[Domain, tuple[str, ...]] = {}411 412 if isinstance(raw_templates, dict):413 tmpl_list = raw_templates.get("templates", [])414 raw_cities = raw_templates.get("cities_by_domain", {}) or {}415 for dom, lst in raw_cities.items():416 if dom not in _DOMAINS:417 raise TemplateSchemaError(f"cities_by_domain: bad domain {dom!r}")418 cities_by_domain[cast("Domain", dom)] = tuple(_nfc(str(c)) for c in lst)419 elif isinstance(raw_templates, list):420 tmpl_list = raw_templates421 else:422 raise TemplateSchemaError(423 f"templates YAML root must be list or mapping, got {type(raw_templates).__name__}"424 )425 426 if not isinstance(tmpl_list, list) or not tmpl_list:427 raise TemplateSchemaError("templates YAML must contain a non-empty list")428 429 for idx, raw in enumerate(tmpl_list):430 if not isinstance(raw, dict):431 raise TemplateSchemaError(432 f"templates[{idx}]: entry must be a mapping, got {type(raw).__name__}"433 )434 parsed_templates.append(_parse_template(raw, where=f"templates[{idx}]"))435 436 # i18n file is optional; if absent we use an empty mapping.437 _LANG_CODES: tuple[LanguageCode, ...] = ("hi", "ta", "kn", "en", "hinglish")438 i18n_data: dict[LanguageCode, dict[str, str]] = {code: {} for code in _LANG_CODES}439 if i18n_path.exists():440 try:441 with i18n_path.open("r", encoding="utf-8") as fh:442 raw_i18n = yaml.safe_load(fh) or {}443 except yaml.YAMLError as exc:444 raise TemplateSchemaError(f"i18n YAML malformed: {exc}") from exc445 if not isinstance(raw_i18n, dict):446 raise TemplateSchemaError("i18n YAML root must be a mapping")447 for lang, block in raw_i18n.items():448 if lang not in _LANGUAGE_CODES:449 raise TemplateSchemaError(450 f"i18n: language key {lang!r} not in {sorted(_LANGUAGE_CODES)}"451 )452 if not isinstance(block, dict):453 raise TemplateSchemaError(f"i18n.{lang}: must be a mapping")454 flat: dict[str, str] = {}455 _flatten_i18n(block, prefix="", out=flat)456 i18n_data[cast("LanguageCode", lang)] = {457 _nfc(str(k)): _nfc(str(v)) for k, v in flat.items()458 }459 460 return TemplateLibrary(461 templates=tuple(parsed_templates),462 cities_by_domain=cities_by_domain,463 i18n=i18n_data,464 )465 466 467def _flatten_i18n(block: Mapping[str, Any], *, prefix: str, out: dict[str, str]) -> None:468 """Flatten nested i18n dicts into dotted keys, NFC everything."""469 for k, v in block.items():470 key = f"{prefix}.{k}" if prefix else str(k)471 if isinstance(v, dict):472 _flatten_i18n(v, prefix=key, out=out)473 else:474 out[key] = str(v)475 476 477# ---------------------------------------------------------------------------478# Lazy singleton479# ---------------------------------------------------------------------------480 481_library_cache: TemplateLibrary | None = None482_library_override: TemplateLibrary | None = None483 484 485def _get_library() -> TemplateLibrary:486 """Return the process-wide TemplateLibrary, loading lazily."""487 if _library_override is not None:488 return _library_override489 global _library_cache490 if _library_cache is None:491 _library_cache = _load_default_library()492 return _library_cache493 494 495def _load_default_library() -> TemplateLibrary:496 """Try the production path, then fall back to the packaged inline library."""497 default_path = Path("data/task_briefs/templates.yaml")498 if default_path.exists():499 return load_templates(default_path)500 return _builtin_library()501 502 503def set_library_override(library: TemplateLibrary | None) -> None:504 """Test hook: pin :func:`_get_library` to a specific library (or clear)."""505 global _library_override506 _library_override = library507 508 509def reset_library_cache() -> None:510 """Test hook: clear the lazy cache so the next call reloads."""511 global _library_cache512 _library_cache = None513 514 515# ---------------------------------------------------------------------------516# Built-in library (fallback when data/ isn't authored yet)517# ---------------------------------------------------------------------------518 519 520def _builtin_library() -> TemplateLibrary:521 """Minimal 5-template library so the generator is self-contained during dev."""522 # Shared numeric grids.523 budget_flight = SlotDistribution(kind="uniform", low=3000.0, high=15000.0, step=500.0)524 budget_hotel = SlotDistribution(kind="uniform", low=2000.0, high=10000.0, step=500.0)525 budget_cab = SlotDistribution(kind="uniform", low=200.0, high=2000.0, step=50.0)526 budget_food = SlotDistribution(kind="uniform", low=200.0, high=1000.0, step=50.0)527 time_window = SlotDistribution(528 kind="choices", choices=("morning", "afternoon", "evening", "late_night")529 )530 date_dist = SlotDistribution(kind="date")531 veg_only = SlotDistribution(kind="bool")532 pax = SlotDistribution(kind="uniform", low=1.0, high=4.0, step=1.0)533 534 cities_inter = (535 "HYD",536 "BLR",537 "DEL",538 "BOM",539 "MAA",540 "CCU",541 "PNQ",542 "AMD",543 "JAI",544 "GOI",545 )546 cities_intra = (547 "Koramangala",548 "Indiranagar",549 "Whitefield",550 "Andheri",551 "Bandra",552 "Powai",553 "Gurgaon",554 "Saket",555 "Banjara Hills",556 "Salt Lake",557 )558 559 airline = Template(560 template_id="airline.book.fixture_v1",561 domain="airline",562 intent="book_flight",563 min_stage=1,564 required_slots=("from", "to", "when"),565 optional_slots=(),566 slot_distributions={567 "from": SlotDistribution(kind="choices", choices=cities_inter),568 "to": SlotDistribution(kind="choices", choices=cities_inter),569 "when": date_dist,570 },571 constraints_template={572 "budget_inr": budget_flight,573 "time_window": time_window,574 },575 drift_slot_tags=("price", "total_fare_inr"),576 language_variants={577 "hinglish": (578 "Bhai {when} ko {from} se {to} jaana hai, {budget_inr} rupees max, {time_window}",579 ),580 "hi": (581 "{when} को {from} से {to} जाना है, {budget_inr} रुपये से कम, {time_window}",582 ),583 "ta": (584 "{when} அன்று {from} லிருந்து {to} டிக்கெட் வேண்டும், {budget_inr} ரூபாய் கீழ், {time_window}",585 ),586 "kn": (587 "{when} ರಂದು {from} ಇಂದ {to} ಗೆ ಟಿಕೆಟ್ ಬೇಕು, {budget_inr} ರೂಪಾಯಿ ಒಳಗೆ, {time_window}",588 ),589 "en": (590 "Flight from {from} to {to} on {when}, under ₹{budget_inr}, {time_window}",591 ),592 },593 )594 595 cab = Template(596 template_id="cab.book.fixture_v1",597 domain="cab",598 intent="book_cab",599 min_stage=1,600 required_slots=("pickup", "drop", "when"),601 optional_slots=(),602 slot_distributions={603 "pickup": SlotDistribution(kind="choices", choices=cities_intra),604 "drop": SlotDistribution(kind="choices", choices=cities_intra),605 "when": date_dist,606 },607 constraints_template={608 "budget_inr": budget_cab,609 "vehicle_class": SlotDistribution(610 kind="choices", choices=("mini", "sedan", "suv")611 ),612 },613 drift_slot_tags=("fare_inr", "fare_breakdown"),614 language_variants={615 "hinglish": (616 "{when} ko {pickup} se {drop} cab chahiye, {budget_inr} ke andar, {vehicle_class}",617 ),618 "hi": (619 "{when} को {pickup} से {drop} कैब चाहिए, {budget_inr} के अंदर, {vehicle_class}",620 ),621 "ta": (622 "{when} அன்று {pickup} லிருந்து {drop} கேப், {budget_inr} கீழ், {vehicle_class}",623 ),624 "kn": (625 "{when} ರಂದು {pickup} ಇಂದ {drop} ಟ್ಯಾಕ್ಸಿ, {budget_inr} ಒಳಗೆ, {vehicle_class}",626 ),627 "en": (628 "Cab from {pickup} to {drop} on {when}, under ₹{budget_inr}, {vehicle_class}",629 ),630 },631 )632 633 restaurant = Template(634 template_id="restaurant.order.fixture_v1",635 domain="restaurant",636 intent="order_food",637 min_stage=2,638 required_slots=("city", "cuisine", "when"),639 optional_slots=(),640 slot_distributions={641 "city": SlotDistribution(kind="choices", choices=cities_inter),642 "cuisine": SlotDistribution(643 kind="choices", choices=("Biryani", "Dosa", "Pizza", "Thali", "Noodles")644 ),645 "when": date_dist,646 },647 constraints_template={648 "budget_inr": budget_food,649 "veg_only": veg_only,650 },651 drift_slot_tags=("min_order", "veg_filter"),652 language_variants={653 "hinglish": (654 "Bhai {when} ko {city} mein {cuisine} order karna hai, {budget_inr} ke andar, veg_only={veg_only}",655 ),656 "hi": (657 "{when} को {city} में {cuisine} ऑर्डर करना है, {budget_inr} के अंदर, veg_only={veg_only}",658 ),659 "ta": (660 "{when} அன்று {city} இல் {cuisine} ஆர்டர், {budget_inr} கீழ், veg_only={veg_only}",661 ),662 "kn": (663 "{when} ರಂದು {city} ನಲ್ಲಿ {cuisine} ಆರ್ಡರ್, {budget_inr} ಒಳಗೆ, veg_only={veg_only}",664 ),665 "en": (666 "Order {cuisine} in {city} on {when}, under ₹{budget_inr}, veg_only={veg_only}",667 ),668 },669 )670 671 hotel = Template(672 template_id="hotel.book.fixture_v1",673 domain="hotel",674 intent="book_hotel",675 min_stage=2,676 required_slots=("city", "checkin", "checkout"),677 optional_slots=(),678 slot_distributions={679 "city": SlotDistribution(kind="choices", choices=cities_inter),680 "checkin": date_dist,681 "checkout": date_dist,682 },683 constraints_template={684 "budget_inr": budget_hotel,685 "room_type": SlotDistribution(686 kind="choices", choices=("single", "double", "suite")687 ),688 },689 drift_slot_tags=("cancel_window", "gst_number"),690 language_variants={691 "hinglish": (692 "{city} mein {checkin} se {checkout} tak hotel chahiye, {budget_inr} per night, {room_type}",693 ),694 "hi": (695 "{city} में {checkin} से {checkout} तक होटल चाहिए, {budget_inr} प्रति रात, {room_type}",696 ),697 "ta": (698 "{city} இல் {checkin} முதல் {checkout} வரை ஹோட்டல், {budget_inr} ஒரு இரவு, {room_type}",699 ),700 "kn": (701 "{city} ನಲ್ಲಿ {checkin} ಇಂದ {checkout} ವರೆಗೆ ಹೋಟೆಲ್, {budget_inr} ಒಂದು ರಾತ್ರಿ, {room_type}",702 ),703 "en": (704 "Hotel in {city} from {checkin} to {checkout}, ₹{budget_inr} per night, {room_type}",705 ),706 },707 )708 709 # Stage-3 compound-constraint airline template — adds a third constraint.710 airline_compound = Template(711 template_id="airline.book.compound_v1",712 domain="airline",713 intent="book_flight",714 min_stage=3,715 required_slots=("from", "to", "when"),716 optional_slots=(),717 slot_distributions={718 "from": SlotDistribution(kind="choices", choices=cities_inter),719 "to": SlotDistribution(kind="choices", choices=cities_inter),720 "when": date_dist,721 },722 constraints_template={723 "budget_inr": budget_flight,724 "time_window": time_window,725 "passenger_count": pax,726 },727 drift_slot_tags=("price", "total_fare_inr", "passenger_count"),728 language_variants={729 "hinglish": (730 "{when} ko {from} se {to}, {passenger_count} log, {budget_inr} max, {time_window}",731 ),732 "hi": (733 "{when} को {from} से {to}, {passenger_count} लोग, {budget_inr} रुपये, {time_window}",734 ),735 "ta": (736 "{when} அன்று {from} லிருந்து {to}, {passenger_count} பேர், {budget_inr} ரூபாய், {time_window}",737 ),738 "kn": (739 "{when} ರಂದು {from} ಇಂದ {to}, {passenger_count} ಜನ, {budget_inr} ರೂಪಾಯಿ, {time_window}",740 ),741 "en": (742 "Flight {from} to {to} on {when} for {passenger_count} pax, ₹{budget_inr}, {time_window}",743 ),744 },745 )746 747 return TemplateLibrary(748 templates=(airline, cab, restaurant, hotel, airline_compound),749 cities_by_domain={750 "airline": cities_inter,751 "hotel": cities_inter,752 "cab": cities_intra,753 "restaurant": cities_inter,754 },755 i18n={756 "hi": {"cities.BLR": "बेंगलुरु", "cities.MAA": "चेन्नई"},757 "ta": {"cities.BLR": "பெங்களூரு", "cities.MAA": "சென்னை"},758 "kn": {"cities.BLR": "ಬೆಂಗಳೂರು", "cities.MAA": "ಚೆನ್ನೈ"},759 "en": {"cities.BLR": "Bengaluru"},760 "hinglish": {"cities.BLR": "Bengaluru"},761 },762 )763 764 765# ---------------------------------------------------------------------------766# Picker + expander (task_generator.md §2.2, §3.2, §3.3)767# ---------------------------------------------------------------------------768 769 770def _pick_domain(seed: int, library: TemplateLibrary, stage: int) -> Domain:771 """Pick uniformly from domains that have ≥ 1 eligible template at ``stage``."""772 available = sorted({t.domain for t in library.templates if t.min_stage <= stage})773 if not available:774 raise TemplateSchemaError(775 f"library has no templates eligible at stage={stage}"776 )777 rng = random.Random(stable_sub_seed(seed, "domain"))778 return rng.choice(available)779 780 781def _eligible_templates(782 library: TemplateLibrary,783 stage: int,784 domain: Domain,785) -> tuple[Template, ...]:786 return tuple(787 t for t in library.templates if t.domain == domain and t.min_stage <= stage788 )789 790 791def _pick_template(792 seed: int,793 stage: int,794 domain: Domain,795 library: TemplateLibrary,796) -> Template:797 eligible = _eligible_templates(library, stage, domain)798 if not eligible:799 raise TemplateSchemaError(800 f"no eligible templates for domain={domain!r} stage={stage}"801 )802 rng = random.Random(stable_sub_seed(seed, "template"))803 # Use sorted template_ids for deterministic ordering.804 ordered = tuple(sorted(eligible, key=lambda t: t.template_id))805 return rng.choice(ordered)806 807 808def _sample_slot_value(809 rng: random.Random,810 name: str,811 dist: SlotDistribution,812 *,813 template_id: str,814) -> object:815 if dist.kind == "choices":816 if not dist.choices:817 raise TemplateSchemaError(818 f"{template_id}.{name}: empty choices list"819 )820 return rng.choice(dist.choices)821 if dist.kind == "uniform":822 assert dist.low is not None and dist.high is not None and dist.step is not None823 steps = int(round((dist.high - dist.low) / dist.step))824 pick = rng.randint(0, steps)825 value = dist.low + pick * dist.step826 # Integer-ify when step + bounds are integral.827 if float(int(dist.step)) == dist.step and float(int(dist.low)) == dist.low:828 value = int(round(value))829 # Post-check (§7 edge case 3).830 lo = int(dist.low) if isinstance(value, int) else dist.low831 hi = int(dist.high) if isinstance(value, int) else dist.high832 if not (lo <= value <= hi):833 raise InvalidBudgetError(834 f"{template_id}.{name}: sampled {value} outside [{dist.low}, {dist.high}]"835 )836 return value837 if dist.kind == "date":838 offset = rng.randint(0, _DATE_WINDOW_DAYS - 1)839 return (_REFERENCE_DATE + timedelta(days=offset)).isoformat()840 if dist.kind == "bool":841 return bool(rng.getrandbits(1))842 raise TemplateSchemaError(843 f"{template_id}.{name}: unknown distribution kind {dist.kind!r}"844 )845 846 847def _resolve_slot_distribution(848 template: Template,849 name: str,850 library: TemplateLibrary,851) -> SlotDistribution | None:852 """Resolve a slot's distribution, preferring explicit declaration then conventions."""853 explicit = template.slot_distributions.get(name)854 if explicit is not None:855 return explicit856 # Constraints block can also declare slot distributions that double as fills.857 constraint = template.constraints_template.get(name)858 if constraint is not None:859 return constraint860 # Conventional fills by slot name.861 if name in _DATE_SLOT_NAMES:862 return SlotDistribution(kind="date")863 if name in _INTER_CITY_SLOT_NAMES:864 pool = library.cities_by_domain.get(template.domain) or _DEFAULT_CITIES_BY_DOMAIN.get(865 template.domain, _DEFAULT_INTER_CITIES866 )867 return SlotDistribution(kind="choices", choices=pool)868 if name in _INTRA_CITY_SLOT_NAMES:869 pool = library.cities_by_domain.get(template.domain) or _DEFAULT_INTRA_CITIES870 return SlotDistribution(kind="choices", choices=pool)871 return None872 873 874def _expand_slots(875 seed: int,876 template: Template,877 *,878 stage: int,879 library: TemplateLibrary,880) -> tuple[SlotGrid, dict[str, object]]:881 """Sample one concrete value per required slot; stage-aware constraint pick.882 883 Returns ``(SlotGrid, constraints_dict)``.884 """885 values: dict[str, object] = {}886 887 # Required slots — always sampled.888 for name in template.required_slots:889 dist = _resolve_slot_distribution(template, name, library)890 if dist is None:891 raise TemplateSchemaError(892 f"{template.template_id}: required slot {name!r} has no distribution "893 f"(declare in slot_distributions or use a conventional name)"894 )895 rng = random.Random(stable_sub_seed(seed, f"slot:{name}"))896 values[name] = _sample_slot_value(rng, name, dist, template_id=template.template_id)897 898 # Optional slots — included with probability 0.5 (seeded). Silently899 # skipped if no distribution resolves (template declares the slot as900 # available but does not wire a fill source).901 for name in template.optional_slots:902 dist = _resolve_slot_distribution(template, name, library)903 if dist is None:904 continue905 rng = random.Random(stable_sub_seed(seed, f"opt:{name}"))906 if rng.random() < 0.5:907 sub_rng = random.Random(stable_sub_seed(seed, f"slot:{name}"))908 values[name] = _sample_slot_value(909 sub_rng, name, dist, template_id=template.template_id910 )911 912 # Constraints — stage-aware sub-selection (§3.5).913 max_constraints = {1: 2, 2: 3, 3: 4}[stage]914 constraint_names = list(template.constraints_template.keys())915 # Stage 1: keep only the first max_constraints deterministically.916 # Stage 2/3: include all declared constraints up to max.917 kept = constraint_names[:max_constraints]918 constraints: dict[str, object] = {}919 for name in kept:920 dist = template.constraints_template[name]921 rng = random.Random(stable_sub_seed(seed, f"constraint:{name}"))922 value = _sample_slot_value(923 rng, name, dist, template_id=template.template_id924 )925 constraints[name] = value926 # Also mirror into slots so variant-format can reference {budget_inr}.927 values[name] = value928 929 # NFC-normalize any string leaves.930 for k, v in list(values.items()):931 if isinstance(v, str):932 values[k] = _nfc(v)933 for k, v in list(constraints.items()):934 if isinstance(v, str):935 constraints[k] = _nfc(v)936 937 return SlotGrid(values=values), constraints938 939 940# ---------------------------------------------------------------------------941# Language picker942# ---------------------------------------------------------------------------943 944 945def _validate_language_weights(language_weights: Mapping[str, float]) -> None:946 """Raise on any malformed input per §3.2."""947 if not isinstance(language_weights, Mapping) or len(language_weights) == 0:948 raise InvalidLanguageWeightError("language_weights is empty")949 950 bad_keys = [k for k in language_weights if k not in _LANGUAGE_CODES]951 if bad_keys:952 raise InvalidLanguageError(953 f"unsupported language key(s): {bad_keys} "954 f"(allowed: {sorted(_LANGUAGE_CODES)})"955 )956 957 for k, v in language_weights.items():958 if not isinstance(v, (int, float)) or isinstance(v, bool):959 raise InvalidLanguageWeightError(960 f"language_weights[{k!r}] must be numeric, got {type(v).__name__}"961 )962 if v < 0:963 raise InvalidLanguageWeightError(964 f"language_weights[{k!r}]={v} is negative"965 )966 967 total = sum(float(v) for v in language_weights.values())968 if abs(total - 1.0) > 1e-6:969 raise InvalidLanguageWeightError(970 f"language_weights sum {total!r} outside [1-1e-6, 1+1e-6]"971 )972 973 # Defensive all-zero check (§3.2 last bullet).974 if all(float(v) == 0.0 for v in language_weights.values()):975 raise InvalidLanguageWeightError(976 "language_weights are all zero (would have no population to sample)"977 )978 979 980def _pick_language(981 seed: int,982 language_weights: Mapping[LanguageCode, float],983) -> LanguageCode:984 rng = random.Random(stable_sub_seed(seed, "language"))985 # Deterministic ordering of keys for reproducibility across dict insertion orders.986 codes = sorted(language_weights.keys())987 weights = [float(language_weights[c]) for c in codes]988 chosen = rng.choices(codes, weights=weights, k=1)[0]989 return chosen990 991 992# ---------------------------------------------------------------------------993# Utterance formatter994# ---------------------------------------------------------------------------995 996 997_PLACEHOLDER_RE = re.compile(r"\{([a-zA-Z_][a-zA-Z0-9_]*)\}")998 999 1000def _format_utterance(1001 seed: int,1002 template: Template,1003 slots: SlotGrid,1004 language: LanguageCode,1005) -> str:1006 variants = template.language_variants.get(language)1007 if not variants:1008 raise NoVariantForLanguageError(1009 f"template {template.template_id!r} has no variants for language {language!r}"1010 )1011 rng = random.Random(stable_sub_seed(seed, "variant"))1012 chosen = rng.choice(tuple(variants))1013 1014 # Render by placeholder-by-placeholder substitution so a missing slot1015 # raises MissingSlotError with the exact field name rather than whatever1016 # ``str.format`` would surface.1017 def _repl(match: re.Match[str]) -> str:1018 name = match.group(1)1019 if name not in slots.values:1020 raise MissingSlotError(1021 f"template {template.template_id!r} variant references {{{name}}} "1022 f"but slot is unbound (slots={sorted(slots.values)})"1023 )1024 value = slots.values[name]1025 if isinstance(value, bool):1026 return "true" if value else "false"1027 if isinstance(value, float):1028 # Trim trailing zeros for cleanness, but keep determinism.1029 if value.is_integer():1030 return str(int(value))1031 return str(value)1032 return str(value)1033 1034 rendered = _PLACEHOLDER_RE.sub(_repl, chosen)1035 normalized = _nfc(rendered)1036 _assert_nfc(normalized, where=f"utterance({template.template_id}, {language})")1037 return normalized1038 1039 1040# ---------------------------------------------------------------------------1041# Primary entry point1042# ---------------------------------------------------------------------------1043 1044 1045def generate(1046 seed: int,1047 stage: Literal[1, 2, 3],1048 language_weights: Mapping[LanguageCode, float],1049) -> GoalSpec:1050 """Produce one :class:`GoalSpec` for episode ``seed`` at curriculum ``stage``.1051 1052 Determinism: identical ``(seed, stage, language_weights)`` ⇒ identical1053 ``GoalSpec`` after NFC normalization of ``seed_utterance``.1054 """1055 # Stage validation (cheapest first).1056 if stage not in _VALID_STAGES:1057 raise InvalidStageError(1058 f"stage must be in {sorted(_VALID_STAGES)}, got {stage!r}"1059 )1060 1061 _validate_language_weights(cast("Mapping[str, float]", language_weights))1062 1063 library = _get_library()1064 1065 domain = _pick_domain(seed, library, int(stage))1066 template = _pick_template(seed, int(stage), domain, library)1067 slot_grid, constraints = _expand_slots(1068 seed, template, stage=int(stage), library=library1069 )1070 language = _pick_language(seed, language_weights)1071 utterance = _format_utterance(seed, template, slot_grid, language)1072 1073 if len(utterance) > _MAX_UTTERANCE_LEN:1074 # Truncate is incorrect (breaks determinism/meaning). Raise so the1075 # template author shortens the variant.1076 raise TemplateSchemaError(1077 f"rendered utterance exceeds {_MAX_UTTERANCE_LEN} chars "1078 f"({len(utterance)}): {utterance!r}"1079 )1080 1081 # Slot dict exposed on GoalSpec should exclude constraint-named entries —1082 # those live in ``constraints``. ``required_slots`` + included optionals only.1083 slot_keys = set(template.required_slots) | set(template.optional_slots)1084 slots_out = {k: v for k, v in slot_grid.values.items() if k in slot_keys}1085 1086 return GoalSpec(1087 domain=template.domain,1088 intent=template.intent,1089 slots=slots_out,1090 constraints=constraints,1091 language=language,1092 seed_utterance=utterance,1093 )1094 1095 1096# ---------------------------------------------------------------------------1097# Variant enumerator (task_generator.md §2.2)1098# ---------------------------------------------------------------------------1099 1100 1101def enumerate_variants(1102 limit: int | None = None,1103 stage: int = 3,1104 language_weights: Mapping[LanguageCode, float] | None = None,1105) -> Iterator[GoalSpec]:1106 """Deterministic walk over the procedural grid."""1107 if stage not in _VALID_STAGES:1108 raise InvalidStageError(f"stage must be in {sorted(_VALID_STAGES)}, got {stage!r}")1109 if language_weights is None:1110 language_weights = {1111 "en": 0.2,1112 "hi": 0.2,1113 "ta": 0.2,1114 "kn": 0.2,1115 "hinglish": 0.2,1116 }1117 count = 01118 seed = 01119 while limit is None or count < limit:1120 yield generate(seed, cast("Literal[1, 2, 3]", stage), language_weights)1121 count += 11122 seed += 11123 1124 1125# ---------------------------------------------------------------------------1126# Test helpers (public so test modules can look up templates)1127# ---------------------------------------------------------------------------1128 1129 1130def _lookup_template_for_test(template_id: str) -> Template:1131 """Public-for-tests helper to resolve a template by ID."""1132 lib = _get_library()1133 for t in lib.templates:1134 if t.template_id == template_id:1135 return t1136 raise KeyError(template_id)1137 1138 1139__all__ = [1140 "Domain",1141 "InvalidBudgetError",1142 "InvalidLanguageError",1143 "InvalidLanguageWeightError",1144 "InvalidStageError",1145 "LanguageCode",1146 "MissingSlotError",1147 "NoVariantForLanguageError",1148 "RawBrief",1149 "SlotDistribution",1150 "SlotGrid",1151 "TaskGeneratorError",1152 "Template",1153 "TemplateFileMissingError",1154 "TemplateLibrary",1155 "TemplateSchemaError",1156 "UnicodeNormalizationError",1157 "_lookup_template_for_test",1158 "enumerate_variants",1159 "generate",1160 "load_templates",1161 "reset_library_cache",1162 "set_library_override",1163 "stable_sub_seed",1164]1165