Coverage for scripts / capture_scaffold_fixtures.py: 100.00%
261 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-09-14 22:07 +0000
1#!/usr/bin/env python3
2"""Capture raw model output for the Mission scaffolder prompt.
4The scaffolder's sampling path is sensitive to the shapes a model
5emits — different families default to different Pythonic idioms
6(``r.get(...)``, comprehension dict-access, attribute walks). The
7fixture-replay test (`tests/test_scaffold_fixture_replay.py`)
8asserts that every captured response round-trips cleanly through
9``_parse_response`` -> ``_normalize_sampled_criteria`` ->
10``validate_criteria``. This script populates the fixture directory by calling
11the fixture directory by calling each Bedrock model on a fixed set of
12canonical directives.
14Usage:
16 # Capture every default model against every canonical directive
17 # (writes one JSON file per model under
18 # tests/fixtures/scaffold_responses/).
19 python3 scripts/capture_scaffold_fixtures.py
21 # Capture a single model.
22 python3 scripts/capture_scaffold_fixtures.py --model MODEL_ID
24 # Capture every uncaptured text-generation model line visible in the live
25 # Bedrock catalog. The command prints the candidate list and paid-call
26 # budget before making requests; one preferred inference profile represents
27 # each underlying model line. Discovery captures four models concurrently
28 # by default; use --workers to tune account throughput.
29 python3 scripts/capture_scaffold_fixtures.py --discover-all-models
31 # Review the same live-catalog selection without invoking any model.
32 python3 scripts/capture_scaffold_fixtures.py \
33 --discover-all-models --list-candidates
35 # Use a different region.
36 python3 scripts/capture_scaffold_fixtures.py --region us-west-2
38The script needs AWS credentials with ``bedrock:InvokeModel`` access
39to the listed models. Anthropic models — including the stock default —
40additionally require the one-time Anthropic first-time-use case form on
41the account; without it Bedrock answers ``FTUFormNotFilled`` and the
42capture for that model fails with that code (see
43``docs/CUSTOMIZATION.md``, Bedrock Model Selection).
45The configured default also consumes ``cdk.json``
46``context.bedrock.generation_reasoning``. At the stock ``high`` effort each capture
47can use substantially more billed output tokens and take longer;
48Claude models from Opus 4.7 onward additionally reject ``temperature``,
49``topP``, and ``topK``, which GCO omits for the canonical default.
50Failures (missing model access, transient ClientError) are reported
51per-model and never abort the run — every model that does succeed lands
52in the fixture directory and protects the validator path on every CI run
53thereafter.
54"""
56from __future__ import annotations
58import argparse
59import asyncio
60import datetime
61import hashlib
62import json
63import os
64import re
65import sys
66import tempfile
67from dataclasses import dataclass
68from pathlib import Path
69from typing import Any, cast
71# Mirror the path-injection pattern used throughout the Mission tree
72# so ``mission.*`` resolves regardless of how the script is launched.
73_REPO_ROOT = Path(__file__).resolve().parent.parent
74for _path in (str(_REPO_ROOT), str(_REPO_ROOT / "gco_mcp")):
75 if _path not in sys.path:
76 sys.path.insert(0, _path)
78import mission.criteria_scaffold as criteria_scaffold # noqa: E402
79from mission.sampling import ( # noqa: E402
80 BedrockSamplingBackend,
81 SamplingPrompt,
82 SamplingTransportError,
83)
84from mission.validation import MissionValidationError, validate_criteria # noqa: E402
86from gco.bedrock import ( # noqa: E402
87 BedrockFTUFormNotAcceptedError,
88 get_default_mission_model_id,
89)
91# ---------------------------------------------------------------------------
92# Defaults
93# ---------------------------------------------------------------------------
96@dataclass(frozen=True)
97class _Directive:
98 """A canonical directive paired with the allowlist used at scaffolding time.
100 The triplet covers the three template branches in
101 ``criteria_scaffold._classify_directive``: search-flavoured
102 directives (preferred shape: ``tool_call_succeeded``),
103 metric-flavoured directives (preferred shape:
104 ``metric_threshold``), and event-flavoured directives (preferred
105 shape: ``event``). Any model that handles all three shapes is
106 likely fine on the long tail.
107 """
109 slug: str
110 text: str
111 allowlist: tuple[str, ...]
114_DIRECTIVES: tuple[_Directive, ...] = (
115 _Directive(
116 slug="search_inference_docs",
117 text="Find documentation about inference endpoints.",
118 allowlist=("find_examples", "find_docs"),
119 ),
120 _Directive(
121 slug="metric_drive_loss",
122 text="Drive validation loss below 0.1 on the demo training tool.",
123 allowlist=("find_examples",),
124 ),
125 _Directive(
126 slug="event_goal_reached",
127 text="Wait for the training job to emit a goal_reached event.",
128 allowlist=("find_examples",),
129 ),
130)
133# Default models to capture against. Every entry is a Bedrock
134# inference-profile id that the calling principal must have invoke
135# access to. Add a model here and the next ``capture`` run picks it
136# up; failures (denied access, transient errors) are reported per-
137# model and never abort the run.
138#
139# The list intentionally spans families (Anthropic, Amazon Nova,
140# Meta Llama, Mistral, DeepSeek) and sizes (small / mid / large)
141# so the replay test stays representative of the long tail of
142# Pythonic emission shapes. When a new family or size lands in
143# Bedrock, add it here and re-run the capture script.
144_CURATED_MODELS: tuple[str, ...] = (
145 # Anthropic family — also the family the configured default belongs to.
146 # NOTE: Anthropic models require a one-time First-Time-Use (FTU) form per
147 # account/org before the first invoke; capture fails with
148 # ``FTUFormNotFilled`` until it is submitted.
149 "us.anthropic.claude-sonnet-4-5-20250929-v1:0",
150 "us.anthropic.claude-haiku-4-5-20251001-v1:0",
151 "us.anthropic.claude-opus-4-5-20251101-v1:0",
152 "us.anthropic.claude-3-haiku-20240307-v1:0",
153 # Amazon Nova family — first-party, no FTU form. The configured GCO
154 # default is prepended lazily by ``_default_models`` and deduplicated from
155 # this curated set before any paid calls are made, whichever family it
156 # belongs to.
157 "us.amazon.nova-pro-v1:0",
158 "us.amazon.nova-lite-v1:0",
159 "us.amazon.nova-micro-v1:0",
160 "us.amazon.nova-2-lite-v1:0",
161 # Meta Llama family — Llama 4 + recent Llama 3.
162 "us.meta.llama4-maverick-17b-instruct-v1:0",
163 "us.meta.llama4-scout-17b-instruct-v1:0",
164 "us.meta.llama3-3-70b-instruct-v1:0",
165 "us.meta.llama3-1-70b-instruct-v1:0",
166 # Mistral family — the visible text-instruction profile.
167 "us.mistral.pixtral-large-2502-v1:0",
168 # DeepSeek family.
169 "us.deepseek.r1-v1:0",
170)
173def _default_models() -> tuple[str, ...]:
174 """Return the configured Mission default plus the curated set, in stable order.
176 Scaffold fixtures replay Mission sampling responses, so the canonical
177 member of the set follows ``context.bedrock.mission_default_model_id``.
178 """
179 return tuple(dict.fromkeys((get_default_mission_model_id(), *_CURATED_MODELS)))
182_FIXTURE_DIR = _REPO_ROOT / "tests" / "fixtures" / "scaffold_responses"
184# Discovery is deliberately narrower than ``list-foundation-models``. Scaffold
185# captures require a normal text message and a generated text answer through
186# Converse; embeddings, rerankers, media transformers, speech-only models, and
187# safety classifiers do not satisfy that contract even when their catalog entry
188# mentions TEXT somewhere.
189_GEOGRAPHY_PROFILE_PREFIX_RE = re.compile(r"^(?:global|us|us-gov|eu|apac|jp|au|ca|sa|il|mx)\.")
190_NON_GENERATION_MODEL_FRAGMENTS = (
191 "embed",
192 "rerank",
193 "stable-",
194 "stability.",
195 "nova-2-sonic",
196 "twelvelabs.",
197 "safeguard",
198)
201def _base_model_id(model_id: str) -> str:
202 """Collapse a geography-scoped profile id to its foundation-model line."""
203 return _GEOGRAPHY_PROFILE_PREFIX_RE.sub("", model_id)
206def _is_text_generation_candidate(summary: dict[str, Any]) -> bool:
207 """Return whether a live catalog entry can plausibly serve this fixture."""
208 model_id = str(summary.get("modelId", ""))
209 return bool(
210 summary.get("modelLifecycle", {}).get("status") == "ACTIVE"
211 and "TEXT" in summary.get("inputModalities", [])
212 and "TEXT" in summary.get("outputModalities", [])
213 and not any(fragment in model_id for fragment in _NON_GENERATION_MODEL_FRAGMENTS)
214 )
217def _fixture_model_ids(output_dir: Path) -> frozenset[str]:
218 """Read exact model ids already represented under *output_dir*."""
219 model_ids: set[str] = set()
220 for path in sorted(output_dir.glob("*.json")):
221 payload = json.loads(path.read_text(encoding="utf-8"))
222 model_id = payload.get("model_id")
223 if not isinstance(model_id, str) or not model_id.strip():
224 raise ValueError(f"{path}: model_id must be a non-empty string")
225 model_ids.add(model_id.strip())
226 return frozenset(model_ids)
229def _profile_preference(profile_id: str) -> tuple[int, str]:
230 """Prefer a global profile, then US, then another geography."""
231 if profile_id.startswith("global."):
232 rank = 0
233 elif profile_id.startswith("us."):
234 rank = 1
235 else:
236 rank = 2
237 return rank, profile_id
240def _discover_all_models(
241 region: str,
242 output_dir: Path,
243 *,
244 bedrock_client: Any | None = None,
245) -> tuple[str, ...]:
246 """Discover every uncaptured text-generation candidate visible in Bedrock.
248 One inference profile represents each underlying model line (global is
249 preferred where available). Direct on-demand model ids are added only when
250 no profile represents that line. The configured/curated registry is folded
251 in first so legacy-but-intentional captures remain required candidates.
252 Existing fixtures suppress the entire underlying line, avoiding paid
253 recapture under a different geography prefix.
254 """
255 if bedrock_client is None:
256 import boto3
258 bedrock_client = boto3.client("bedrock", region_name=region)
260 model_page = bedrock_client.list_foundation_models(byOutputModality="TEXT")
261 active_models = {
262 summary["modelId"]: summary
263 for summary in model_page.get("modelSummaries", [])
264 if _is_text_generation_candidate(summary)
265 }
267 profiles: list[dict[str, Any]] = []
268 next_token: str | None = None
269 while True:
270 request: dict[str, Any] = {
271 "typeEquals": "SYSTEM_DEFINED",
272 "maxResults": 1000,
273 }
274 if next_token:
275 request["nextToken"] = next_token
276 page = bedrock_client.list_inference_profiles(**request)
277 profiles.extend(page.get("inferenceProfileSummaries", []))
278 next_token = page.get("nextToken")
279 if not next_token:
280 break
282 existing_bases = {_base_model_id(model_id) for model_id in _fixture_model_ids(output_dir)}
283 selected_by_base: dict[str, str] = {}
285 # Curated entries win exact-id selection, including intentional older
286 # profiles that no longer appear as ACTIVE foundation models.
287 for model_id in _default_models():
288 base = _base_model_id(model_id)
289 if base not in existing_bases:
290 selected_by_base.setdefault(base, model_id)
291 curated_bases = set(selected_by_base)
293 for profile in profiles:
294 if profile.get("status") != "ACTIVE":
295 continue
296 profile_id = profile.get("inferenceProfileId")
297 if not isinstance(profile_id, str) or any(
298 fragment in profile_id for fragment in _NON_GENERATION_MODEL_FRAGMENTS
299 ):
300 continue
301 base_ids = {
302 str(model.get("modelArn", "")).rsplit("/", 1)[-1] for model in profile.get("models", [])
303 }
304 for base in base_ids:
305 if base not in active_models or base in existing_bases or base in curated_bases:
306 continue
307 current = selected_by_base.get(base)
308 if current is None or _profile_preference(profile_id) < _profile_preference(current):
309 selected_by_base[base] = profile_id
311 represented_bases = existing_bases | set(selected_by_base)
312 for model_id, summary in active_models.items():
313 if model_id in represented_bases:
314 continue
315 if "ON_DEMAND" not in summary.get("inferenceTypesSupported", []):
316 continue
317 selected_by_base[model_id] = model_id
319 # Put Anthropic last: an account-wide FTU failure should not prevent other
320 # providers from being captured in the same broad run.
321 return tuple(
322 sorted(
323 selected_by_base.values(),
324 key=lambda model_id: ("anthropic." in model_id, model_id),
325 )
326 )
329# ---------------------------------------------------------------------------
330# Capture helpers
331# ---------------------------------------------------------------------------
334def _slug_for_model(model_id: str) -> str:
335 """Turn a Bedrock model id into a filesystem-safe slug.
337 ``us.anthropic.claude-haiku-4-5-20251001-v1:0`` ->
338 ``us_anthropic_claude_haiku_4_5_20251001_v1_0``. The replacement
339 is intentionally minimal — every non-alphanumeric becomes an
340 underscore — so two model ids with different metadata produce
341 different slugs.
342 """
343 out = []
344 prev_underscore = False
345 for ch in model_id:
346 if ch.isalnum():
347 out.append(ch)
348 prev_underscore = False
349 elif not prev_underscore:
350 out.append("_")
351 prev_underscore = True
352 return "".join(out).strip("_")
355class _PromptAdapter:
356 """Tiny stand-in for ``SamplingPrompt`` used by the scaffolder.
358 The Bedrock backend calls ``prompt.assemble()`` to render the
359 string it sends to Converse. We bypass the full ``SamplingPrompt``
360 constructor (which requires session-shaped data we don't have at
361 capture time) by giving the backend an object whose ``assemble``
362 returns the rendered scaffold prompt directly.
363 """
365 def __init__(self, text: str) -> None:
366 self._text = text
368 def assemble(self) -> str:
369 return self._text
372def _backend_for_capture(
373 model_id: str,
374 region: str,
375 *,
376 read_timeout_seconds: int | None = None,
377) -> BedrockSamplingBackend:
378 """Preserve canonical provenance and optionally bound one capture request."""
379 if model_id == get_default_mission_model_id():
380 backend = BedrockSamplingBackend.from_canonical_default(region=region)
381 else:
382 backend = BedrockSamplingBackend(model_id=model_id, region=region)
384 if read_timeout_seconds is not None:
385 import boto3
386 from botocore.config import Config
388 backend._client = boto3.Session().client(
389 "bedrock-runtime",
390 region_name=region,
391 config=Config(
392 connect_timeout=min(10, read_timeout_seconds),
393 read_timeout=read_timeout_seconds,
394 ),
395 )
396 return backend
399async def _capture_one(
400 backend: BedrockSamplingBackend,
401 directive: _Directive,
402) -> dict[str, Any]:
403 """Render the scaffold prompt and capture the raw model response.
405 Returns a dict carrying the prompt inputs, the rendered prompt digest, and
406 the untouched response so fixture provenance is auditable without storing
407 the full repeated prompt in every capture.
408 """
409 prompt_str = criteria_scaffold.build_scaffold_prompt(
410 directive.text,
411 allowlist=list(directive.allowlist),
412 )
413 # ``BedrockSamplingBackend.sample`` is typed against
414 # :class:`SamplingPrompt`, which requires session-shaped data we
415 # don't have at capture time. The backend only ever calls
416 # ``prompt.assemble()`` on its argument, so a duck-typed
417 # ``_PromptAdapter`` is sufficient at runtime; cast to satisfy
418 # mypy without weakening the production signature.
419 raw = await backend.sample(cast(SamplingPrompt, _PromptAdapter(prompt_str)))
421 # A successful Converse call is not enough for a positive playback
422 # fixture: the untouched raw text must also survive the exact production
423 # parse/normalize/validate path. Fail closed before the per-model file is
424 # published; a later live run can retry a model that emitted malformed or
425 # semantically ambiguous criteria.
426 parsed = criteria_scaffold._parse_response(raw)
427 if len(parsed) > criteria_scaffold.DEFAULT_MAX_CRITERIA:
428 parsed = parsed[: criteria_scaffold.DEFAULT_MAX_CRITERIA]
429 parsed = criteria_scaffold._normalize_sampled_criteria(parsed)
430 validated = validate_criteria(parsed)
431 criteria_scaffold._validate_sampled_criteria_context(
432 validated,
433 directive=directive.text,
434 allowlist=list(directive.allowlist),
435 )
437 return {
438 "prompt_directive": directive.text,
439 "prompt_allowlist": list(directive.allowlist),
440 "prompt_sha256": hashlib.sha256(prompt_str.encode("utf-8")).hexdigest(),
441 "raw_response": raw,
442 }
445async def _capture_model(
446 model_id: str,
447 region: str,
448 output_dir: Path,
449 *,
450 read_timeout_seconds: int | None = None,
451) -> bool:
452 """Capture all canonical directives for one model. Returns False on failure.
454 Every directive is written into the same per-model JSON file under
455 its ``slug`` key. A failure on one directive aborts the whole
456 model's capture so the fixture file is either written wholesale
457 or not at all — an incomplete fixture would silently weaken the
458 replay test.
459 """
460 backend = _backend_for_capture(
461 model_id,
462 region,
463 read_timeout_seconds=read_timeout_seconds,
464 )
465 captures: dict[str, dict[str, Any]] = {}
466 for directive in _DIRECTIVES:
467 try:
468 captures[directive.slug] = await _capture_one(backend, directive)
469 except BedrockFTUFormNotAcceptedError:
470 # Account-wide gate, not a per-model failure: every Anthropic model
471 # in the run would fail identically, so abort with the remediation
472 # instead of repeating it once per model.
473 raise
474 except SamplingTransportError as exc:
475 cause = f"; {exc.__cause__}" if exc.__cause__ is not None else ""
476 print(
477 f"[{model_id}] capture failed for {directive.slug!r}: {exc.code}: {exc}{cause}",
478 file=sys.stderr,
479 )
480 return False
481 except (MissionValidationError, ValueError) as exc:
482 details = (
483 f"; details={exc.details!r}" if isinstance(exc, MissionValidationError) else ""
484 )
485 print(
486 f"[{model_id}] capture rejected for {directive.slug!r}: "
487 f"{type(exc).__name__}: {exc}{details}",
488 file=sys.stderr,
489 )
490 return False
491 except Exception as exc: # noqa: BLE001 - surface and keep going
492 print(
493 f"[{model_id}] unexpected error for {directive.slug!r}: "
494 f"{type(exc).__name__}: {exc}",
495 file=sys.stderr,
496 )
497 return False
499 output_path = output_dir / f"{_slug_for_model(model_id)}.json"
500 payload = {
501 "model_id": model_id,
502 "region": region,
503 "captured_at": datetime.datetime.now(datetime.UTC).isoformat(),
504 "captures": captures,
505 }
506 serialized = json.dumps(payload, indent=2) + "\n"
507 temporary_path: Path | None = None
508 try:
509 with tempfile.NamedTemporaryFile(
510 mode="w",
511 encoding="utf-8",
512 dir=output_path.parent,
513 prefix=f".{output_path.name}.",
514 suffix=".tmp",
515 delete=False,
516 ) as handle:
517 temporary_path = Path(handle.name)
518 handle.write(serialized)
519 handle.flush()
520 os.fsync(handle.fileno())
521 assert temporary_path is not None
522 os.chmod(temporary_path, 0o600)
523 os.replace(temporary_path, output_path)
524 temporary_path = None
525 finally:
526 if temporary_path is not None:
527 temporary_path.unlink(missing_ok=True)
528 print(f"[{model_id}] wrote {output_path.relative_to(_REPO_ROOT)}")
529 return True
532# ---------------------------------------------------------------------------
533# CLI
534# ---------------------------------------------------------------------------
537def _positive_seconds(raw: str) -> int:
538 """Argparse type for a strictly positive request timeout."""
539 value = int(raw)
540 if value <= 0:
541 raise argparse.ArgumentTypeError("timeout must be greater than zero")
542 return value
545def _positive_workers(raw: str) -> int:
546 """Argparse type for a strictly positive model-worker count."""
547 value = int(raw)
548 if value <= 0:
549 raise argparse.ArgumentTypeError("workers must be greater than zero")
550 return value
553def _build_parser() -> argparse.ArgumentParser:
554 parser = argparse.ArgumentParser(
555 description=__doc__.split("\n", 1)[0] if __doc__ else "",
556 )
557 selection = parser.add_mutually_exclusive_group()
558 selection.add_argument(
559 "--model",
560 action="append",
561 dest="models",
562 default=None,
563 help=(
564 "Bedrock model id to capture against; repeatable. "
565 "Defaults to a curated cross-family set."
566 ),
567 )
568 selection.add_argument(
569 "--discover-all-models",
570 action="store_true",
571 help=(
572 "Query the live Bedrock catalog and capture every uncaptured "
573 "text-generation model line (one preferred profile per line)."
574 ),
575 )
576 parser.add_argument(
577 "--list-candidates",
578 action="store_true",
579 help="Print the selected models and paid-call budget without invoking Bedrock.",
580 )
581 parser.add_argument(
582 "--region",
583 default=os.environ.get("GCO_MISSION_BEDROCK_REGION", "us-east-1"),
584 help="Bedrock region (default: us-east-1).",
585 )
586 parser.add_argument(
587 "--read-timeout-seconds",
588 type=_positive_seconds,
589 default=None,
590 help=(
591 "Per-Converse read timeout. Discovery defaults to 300 seconds; "
592 "curated/explicit capture retains the normal backend timeout."
593 ),
594 )
595 parser.add_argument(
596 "--workers",
597 type=_positive_workers,
598 default=None,
599 help=(
600 "Models to capture concurrently. Discovery defaults to 4; "
601 "curated/explicit capture defaults to 1."
602 ),
603 )
604 parser.add_argument(
605 "--output-dir",
606 type=Path,
607 default=_FIXTURE_DIR,
608 help=("Directory to write fixtures into. Defaults to tests/fixtures/scaffold_responses/."),
609 )
610 return parser
613def _selected_models(args: argparse.Namespace) -> tuple[str, ...]:
614 """Resolve explicit, discovered, or curated model selection."""
615 if args.discover_all_models:
616 return _discover_all_models(args.region, args.output_dir)
617 if args.models:
618 return tuple(dict.fromkeys(args.models))
619 return _default_models()
622async def _main_async(args: argparse.Namespace) -> int:
623 models = _selected_models(args)
624 if args.discover_all_models or args.list_candidates:
625 print(
626 f"Selected {len(models)} model candidate(s); a capture run makes "
627 f"up to {len(models) * len(_DIRECTIVES)} paid Converse calls."
628 )
629 for model_id in models:
630 print(model_id)
631 if args.list_candidates:
632 return 0
633 if not models:
634 print("No uncaptured model candidates found.")
635 return 0
637 read_timeout_seconds = getattr(args, "read_timeout_seconds", None)
638 if read_timeout_seconds is None and args.discover_all_models:
639 read_timeout_seconds = 300
640 if args.discover_all_models:
641 print(f"Per-request read timeout: {read_timeout_seconds}s")
643 args.output_dir.mkdir(parents=True, exist_ok=True)
644 workers = getattr(args, "workers", None)
645 if workers is None:
646 workers = 4 if args.discover_all_models else 1
647 if args.discover_all_models or workers > 1:
648 print(f"Concurrent model workers: {workers}")
650 if workers == 1:
651 successes = 0
652 failures = 0
653 for model_id in models:
654 try:
655 ok = await _capture_model(
656 model_id,
657 args.region,
658 args.output_dir,
659 read_timeout_seconds=read_timeout_seconds,
660 )
661 except BedrockFTUFormNotAcceptedError as exc:
662 # Preserve the historical curated-run behavior: a serial run
663 # stops at the account-wide Anthropic prerequisite.
664 print(f"\n{exc}", file=sys.stderr)
665 print(f"Captured {successes} model(s) before aborting.", file=sys.stderr)
666 return 1
667 if ok:
668 successes += 1
669 else:
670 failures += 1
671 skipped = 0
672 else:
673 semaphore = asyncio.Semaphore(workers)
674 anthropic_ftu = asyncio.Event()
676 async def _capture_concurrently(
677 model_id: str,
678 ) -> bool | None | BedrockFTUFormNotAcceptedError:
679 async with semaphore:
680 if anthropic_ftu.is_set() and "anthropic." in model_id:
681 print(f"[{model_id}] skipped after Anthropic FTU failure", file=sys.stderr)
682 return None
683 try:
684 return await _capture_model(
685 model_id,
686 args.region,
687 args.output_dir,
688 read_timeout_seconds=read_timeout_seconds,
689 )
690 except BedrockFTUFormNotAcceptedError as exc:
691 anthropic_ftu.set()
692 return exc
693 except Exception as exc: # noqa: BLE001 - isolate broad model failures
694 print(
695 f"[{model_id}] unexpected capture failure: {type(exc).__name__}: {exc}",
696 file=sys.stderr,
697 )
698 return False
700 results = await asyncio.gather(*(_capture_concurrently(model_id) for model_id in models))
701 successes = sum(result is True for result in results)
702 failures = sum(
703 result is False or isinstance(result, BedrockFTUFormNotAcceptedError)
704 for result in results
705 )
706 skipped = sum(result is None for result in results)
707 ftu_error = next(
708 (result for result in results if isinstance(result, BedrockFTUFormNotAcceptedError)),
709 None,
710 )
711 if ftu_error is not None:
712 print(f"\n{ftu_error}", file=sys.stderr)
714 print(f"\nCaptured {successes} model(s); {failures} failed; {skipped} skipped.")
715 # A non-zero exit when *every* model failed is useful for cron
716 # wrappers; a partial-failure run still exits 0 so a denied or unsupported
717 # model doesn't stop successful fixtures from being committed.
718 return 0 if successes > 0 else 1
721def main() -> int:
722 args = _build_parser().parse_args()
723 return asyncio.run(_main_async(args))
726if __name__ == "__main__":
727 raise SystemExit(main())