Skip to content

Commit bd49da8

Browse files
committed
feat(api): generate Experiments REST bindings
Implements step 3 of #683 and fixes #639. This adds the second public generated REST resource: - `BraintrustOpenApiClient.experiments`: all 10 operations selected by the Experiments OpenAPI tag, including `get_experiment_id_summarize` - `braintrust.api.types`: public experiment request and response types Codegen now partitions models across multiple resources. Models reached by one resource remain in its resource-specific module, while shared definitions are emitted once in `models/common.py` and imported explicitly. Against the pinned specification, Projects plus Experiments generates 15 operations and 45 reachable component schemas. Logical POST reads use a reviewed `safe_reads` allowlist, while generated GETs retain mechanical `SAFE_READ` classification and writes remain non-retrying unless explicitly classified. `Experiment.summarize()` now uses the generated summarize binding. Successful and intentionally skipped summaries are represented by `SummarySuccess` and `SummarySkipped`, with `comparison` as the primary result. The deprecated read-only `scores` and `metrics` bridges remain serialized for compatibility. Summary retrieval errors are no longer swallowed: transient failures retry through the policy-aware transport and final failures raise typed API errors. Structured summaries support tagged deep deserialization and legacy payloads containing only top-level score and metric maps. Coverage includes deterministic multi-resource codegen, retry-policy validation, exact wire behavior for all generated Experiments methods, additive responses, retry exhaustion, framework propagation, static and runtime typing, structured-summary round trips, and real-backend VCR flows for implicit and explicit comparison selection.
1 parent fcb5dcc commit bd49da8

24 files changed

Lines changed: 3131 additions & 340 deletions

openapi/README.md

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -15,12 +15,13 @@ make check-api-client-codegen
1515

1616
The check regenerates into a temporary directory and does not modify the worktree. Endpoint bindings
1717
are rolled out explicitly through `endpoint_generator.generated_tags`. The current rollout supports
18-
exactly one selected OpenAPI tag and emits its operation registry and resource class together in
19-
`projects.py`, with reachable types in `models/projects.py`; unreachable models are omitted. Add
20-
explicit cross-resource model partitioning before selecting a second tag. Public resource method and
21-
inline response type names are derived mechanically from each `operationId`, and generated methods
22-
forward request fields and parameters without implicit defaults. Writes that are safe to retry are
23-
listed declaratively in `endpoint_generator.idempotent_writes`; reads and all other writes use
18+
the Projects and Experiments tags and emits one operation registry/resource class per tag. Reachable
19+
models used by one resource live in that resource's model module; models shared by multiple resources
20+
live once in `models/common.py` and are imported explicitly. Unreachable models are omitted. Public
21+
resource method and inline response type names are derived mechanically from each `operationId`, and
22+
generated methods forward request fields and parameters without implicit defaults. Logical POST reads
23+
that are safe to retry are listed in `endpoint_generator.safe_reads`, while verified idempotent writes
24+
are listed in `endpoint_generator.idempotent_writes`; GET/HEAD reads and all other writes use
2425
mechanical retry defaults.
2526

2627
To fetch the configured upstream commit explicitly:

openapi/config.json

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,11 @@
3030
"endpoint_generator": {
3131
"schema_version": 1,
3232
"generated_tags": [
33-
"Projects"
33+
"Projects",
34+
"Experiments"
35+
],
36+
"safe_reads": [
37+
"postExperimentIdFetch"
3438
],
3539
"idempotent_writes": [
3640
"postProject"

py/scripts/openapi_codegen.py

Lines changed: 172 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
"""Shared validation and generation helpers for the pinned Braintrust OpenAPI spec."""
22

3+
import ast
34
import copy
45
import difflib
56
import hashlib
@@ -289,19 +290,119 @@ def _with_inline_models(
289290
return model_spec
290291

291292

292-
def _single_model_module(operations: Sequence[GeneratedOperation]) -> str:
293-
tags = {operation.tag for operation in operations}
294-
if len(tags) != 1:
295-
raise CodegenError(
296-
"Model generation currently requires exactly one generated OpenAPI tag; "
297-
"add explicit cross-resource model partitioning before enabling another tag"
293+
_NON_MODEL_ANNOTATION_NAMES = {"Any", "Literal", "Mapping", "None", "Sequence"}
294+
295+
296+
def _operation_annotation_names(operation: GeneratedOperation) -> Set[str]:
297+
annotation_names: Set[str] = set()
298+
for type_name in [
299+
operation.request_body_type,
300+
operation.response_type,
301+
*(parameter.type_name for parameter in operation.parameters),
302+
]:
303+
if type_name:
304+
annotation_names.update(re.findall(r"\b[A-Z][A-Za-z0-9_]*\b", type_name))
305+
return annotation_names
306+
307+
308+
def _operation_model_roots(operations: Sequence[GeneratedOperation]) -> Dict[str, Set[str]]:
309+
roots: Dict[str, Set[str]] = {}
310+
for operation in operations:
311+
roots.setdefault(operation.tag, set()).update(
312+
_operation_annotation_names(operation) - _NON_MODEL_ANNOTATION_NAMES
298313
)
299-
return _snake_case(next(iter(tags)))
314+
return roots
315+
316+
317+
def _partition_model_source(
318+
source: str, operations: Sequence[GeneratedOperation]
319+
) -> Tuple[Dict[str, str], Dict[str, str]]:
320+
"""Partition one deterministic model-generator output by resource dependency closure.
321+
322+
Definitions reached by more than one generated tag live in ``common.py``. Resource-specific
323+
modules import those shared definitions explicitly, avoiding duplicate runtime type identities.
324+
"""
325+
tree = ast.parse(source)
326+
imports: List[ast.stmt] = []
327+
definitions: List[Tuple[str, List[ast.stmt]]] = []
328+
for node in tree.body:
329+
if isinstance(node, (ast.Import, ast.ImportFrom)):
330+
imports.append(node)
331+
continue
332+
if isinstance(node, ast.ClassDef):
333+
names = [node.name]
334+
elif isinstance(node, (ast.Assign, ast.AnnAssign)):
335+
targets = node.targets if isinstance(node, ast.Assign) else [node.target]
336+
names = [target.id for target in targets if isinstance(target, ast.Name)]
337+
elif isinstance(node, ast.Expr) and isinstance(node.value, ast.Constant) and isinstance(node.value.value, str):
338+
if definitions:
339+
definitions[-1][1].append(node)
340+
continue
341+
else:
342+
raise CodegenError(f"Unsupported generated model statement: {type(node).__name__}")
343+
if len(names) != 1:
344+
raise CodegenError("Generated model definitions must bind exactly one public name")
345+
definitions.append((names[0], [node]))
346+
347+
definition_names = {name for name, _ in definitions}
348+
dependencies: Dict[str, Set[str]] = {}
349+
for name, nodes in definitions:
350+
dependencies[name] = {
351+
child.id
352+
for node in nodes
353+
for child in ast.walk(node)
354+
if isinstance(child, ast.Name) and child.id in definition_names and child.id != name
355+
}
300356

357+
owners: Dict[str, Set[str]] = {name: set() for name in definition_names}
358+
for tag, roots in _operation_model_roots(operations).items():
359+
pending = list(roots)
360+
seen: Set[str] = set()
361+
while pending:
362+
name = pending.pop()
363+
if name in seen:
364+
continue
365+
if name not in definition_names:
366+
raise CodegenError(f"Generated resource {tag!r} references unknown model {name!r}")
367+
seen.add(name)
368+
owners[name].add(tag)
369+
pending.extend(dependencies[name])
370+
371+
unreachable = sorted(name for name, tags in owners.items() if not tags)
372+
if unreachable:
373+
raise CodegenError(f"Generated models are unreachable from resource methods: {unreachable}")
374+
375+
common_names = {name for name, tags in owners.items() if len(tags) > 1}
376+
module_for_name = {
377+
name: "common" if name in common_names else _snake_case(next(iter(tags))) for name, tags in owners.items()
378+
}
301379

302-
def _model_modules(spec: Mapping[str, Any], module: str) -> Dict[str, str]:
303-
schemas = spec.get("components", {}).get("schemas", {})
304-
return {_python_type_name(name): module for name in schemas}
380+
def source_for(node: ast.stmt) -> str:
381+
segment = ast.get_source_segment(source, node)
382+
if segment is None:
383+
raise CodegenError(f"Could not recover generated model source for {type(node).__name__}")
384+
return segment
385+
386+
import_source = "\n".join(source_for(node) for node in imports)
387+
bodies: Dict[str, List[str]] = {}
388+
for name, nodes in definitions:
389+
bodies.setdefault(module_for_name[name], []).append("\n".join(source_for(node) for node in nodes))
390+
391+
modules: Dict[str, str] = {}
392+
for module, blocks in sorted(bodies.items()):
393+
common_imports = sorted(
394+
dependency
395+
for name, _ in definitions
396+
if module_for_name[name] == module
397+
for dependency in dependencies[name]
398+
if dependency in common_names
399+
)
400+
sections = [import_source]
401+
if common_imports and module != "common":
402+
sections.append(f"from .common import {', '.join(dict.fromkeys(common_imports))}")
403+
sections.append("\n\n".join(blocks))
404+
modules[module] = "\n\n".join(section for section in sections if section) + "\n"
405+
return modules, module_for_name
305406

306407

307408
def generate_tree(output_root: Path, config: Mapping[str, Any], spec: Mapping[str, Any]) -> ValidationReport:
@@ -310,21 +411,29 @@ def generate_tree(output_root: Path, config: Mapping[str, Any], spec: Mapping[st
310411
operations, inline_models = _collect_generated_operations(spec, config)
311412
selected_spec = _slice_model_spec(spec, {operation.operation_id for operation in operations})
312413
model_spec = _with_inline_models(selected_spec, inline_models)
313-
model_module = _single_model_module(operations)
314-
model_modules = _model_modules(model_spec, model_module)
315414
output_root.mkdir(parents=True, exist_ok=True)
316415
selected_spec_path = output_root.parent / "selected-spec.json"
416+
monolithic_models_path = output_root.parent / "models.py"
317417
selected_spec_path.write_text(
318418
json.dumps(model_spec, sort_keys=True, separators=(",", ":"), ensure_ascii=False) + "\n", encoding="utf-8"
319419
)
320420
try:
321-
_generate_models(selected_spec_path, output_root / "models" / f"{model_module}.py", config)
421+
_generate_models(selected_spec_path, monolithic_models_path, config)
422+
model_sources, model_modules = _partition_model_source(monolithic_models_path.read_text(), operations)
322423
finally:
323424
selected_spec_path.unlink(missing_ok=True)
425+
monolithic_models_path.unlink(missing_ok=True)
426+
model_paths = []
427+
for module, body in model_sources.items():
428+
model_path = output_root / "models" / f"{module}.py"
429+
model_path.parent.mkdir(parents=True, exist_ok=True)
430+
_write_generated_file(model_path, body, config)
431+
model_paths.append(model_path)
324432
_write_generated_file(output_root / "__init__.py", _GENERATED_INIT_BODY, config)
325-
_write_generated_file(output_root / "models" / "__init__.py", '"""Generated private model types."""\n', config)
433+
model_init_path = output_root / "models" / "__init__.py"
434+
_write_generated_file(model_init_path, _model_package_source(model_modules), config)
326435
resource_files = _generate_resources(output_root, operations, model_modules, config)
327-
_format_generated_files(resource_files)
436+
_format_generated_files([*model_paths, model_init_path, *resource_files])
328437
return report
329438

330439

@@ -459,6 +568,20 @@ def _generated_header(config: Mapping[str, Any], content_hash: str) -> str:
459568
'''
460569

461570

571+
def _model_package_source(model_modules: Mapping[str, str]) -> str:
572+
by_module: Dict[str, List[str]] = {}
573+
for name, module in model_modules.items():
574+
by_module.setdefault(module, []).append(name)
575+
576+
lines = ['"""Generated private model types with stable package-level imports."""', ""]
577+
for module, names in sorted(by_module.items()):
578+
lines.append(f"from .{module} import {', '.join(sorted(names))}")
579+
lines.extend(["", "", "__all__ = ["])
580+
lines.extend(f" {name!r}," for name in sorted(model_modules))
581+
lines.extend(["]", ""])
582+
return "\n".join(lines)
583+
584+
462585
def _validate_selected_operations(
463586
operations: Sequence[Tuple[str, str, Any, Mapping[str, Any], Mapping[str, Any]]],
464587
endpoint: Mapping[str, Any],
@@ -474,11 +597,31 @@ def _validate_selected_operations(
474597
if missing_tags:
475598
raise CodegenError(f"endpoint_generator.generated_tags contains unknown tags: {sorted(missing_tags)}")
476599

600+
safe_reads = set(endpoint["safe_reads"])
601+
stale_safe_reads = safe_reads - set(supported)
602+
if stale_safe_reads:
603+
operation_id = sorted(stale_safe_reads)[0]
604+
raise CodegenError(f"endpoint_generator.safe_reads references non-generated operation {operation_id!r}")
605+
non_post_safe_reads = sorted(
606+
operation_id for method, _, operation_id, _, _ in operations if operation_id in safe_reads and method != "post"
607+
)
608+
if non_post_safe_reads:
609+
raise CodegenError(
610+
f"endpoint_generator.safe_reads must reference POST operations; got {non_post_safe_reads[0]!r}"
611+
)
612+
477613
idempotent_writes = set(endpoint["idempotent_writes"])
478614
stale_idempotent_writes = idempotent_writes - set(supported)
479615
if stale_idempotent_writes:
480616
operation_id = sorted(stale_idempotent_writes)[0]
481617
raise CodegenError(f"endpoint_generator.idempotent_writes references non-generated operation {operation_id!r}")
618+
overlapping_retry_modes = safe_reads & idempotent_writes
619+
if overlapping_retry_modes:
620+
operation_id = sorted(overlapping_retry_modes)[0]
621+
raise CodegenError(
622+
f"Operation {operation_id!r} cannot appear in both endpoint_generator.safe_reads "
623+
"and endpoint_generator.idempotent_writes"
624+
)
482625
non_writes = sorted(
483626
operation_id
484627
for method, _, operation_id, _, _ in operations
@@ -488,8 +631,8 @@ def _validate_selected_operations(
488631
raise CodegenError(f"endpoint_generator.idempotent_writes references read operation {non_writes[0]!r}")
489632

490633

491-
def _operation_retry_mode(method: str, operation_id: str, idempotent_writes: Set[str]) -> str:
492-
if method in {"get", "head"}:
634+
def _operation_retry_mode(method: str, operation_id: str, safe_reads: Set[str], idempotent_writes: Set[str]) -> str:
635+
if method in {"get", "head"} or operation_id in safe_reads:
493636
return "SAFE_READ"
494637
if operation_id in idempotent_writes:
495638
return "IDEMPOTENT_WRITE"
@@ -500,6 +643,7 @@ def _collect_generated_operations(
500643
spec: Mapping[str, Any], config: Mapping[str, Any]
501644
) -> Tuple[List[GeneratedOperation], List[Tuple[str, Mapping[str, Any]]]]:
502645
endpoint = _endpoint_config(config)
646+
safe_reads = set(endpoint["safe_reads"])
503647
idempotent_writes = set(endpoint["idempotent_writes"])
504648
operations: List[GeneratedOperation] = []
505649
inline_models: Dict[str, Mapping[str, Any]] = {}
@@ -528,7 +672,7 @@ def _collect_generated_operations(
528672
response_type=response_type,
529673
success_statuses=statuses,
530674
json_success_statuses=json_statuses,
531-
retry_mode=_operation_retry_mode(method, operation_id, idempotent_writes),
675+
retry_mode=_operation_retry_mode(method, operation_id, safe_reads, idempotent_writes),
532676
)
533677
)
534678
return operations, list(inline_models.items())
@@ -663,18 +807,10 @@ def _generate_resources(
663807
def _resource_module_source(
664808
tag: str, operations: Sequence[GeneratedOperation], model_modules: Mapping[str, str]
665809
) -> str:
666-
annotation_names: Set[str] = set()
667-
for operation in operations:
668-
for type_name in [
669-
operation.request_body_type,
670-
operation.response_type,
671-
*(parameter.type_name for parameter in operation.parameters),
672-
]:
673-
if type_name:
674-
annotation_names.update(re.findall(r"\b[A-Z][A-Za-z0-9_]*\b", type_name))
810+
annotation_names = set().union(*(_operation_annotation_names(operation) for operation in operations))
675811
collections_imports = sorted(annotation_names & {"Mapping", "Sequence"})
676812
typing_imports = sorted(annotation_names & {"Any", "Literal"})
677-
model_type_names = annotation_names - {"Any", "Literal", "Mapping", "None", "Sequence"}
813+
model_type_names = annotation_names - _NON_MODEL_ANNOTATION_NAMES
678814
model_imports: Dict[str, Set[str]] = {}
679815
for type_name in model_type_names:
680816
module = model_modules.get(type_name)
@@ -834,13 +970,14 @@ def _endpoint_config(config: Mapping[str, Any]) -> Mapping[str, Any]:
834970
or len(generated_tags) != len(set(generated_tags))
835971
):
836972
raise CodegenError("endpoint_generator.generated_tags must be a unique list of non-empty strings")
837-
idempotent_writes = endpoint.get("idempotent_writes")
838-
if (
839-
not isinstance(idempotent_writes, list)
840-
or not all(isinstance(value, str) and value for value in idempotent_writes)
841-
or len(idempotent_writes) != len(set(idempotent_writes))
842-
):
843-
raise CodegenError("endpoint_generator.idempotent_writes must be a unique list of non-empty strings")
973+
for key in ("safe_reads", "idempotent_writes"):
974+
values = endpoint.get(key)
975+
if (
976+
not isinstance(values, list)
977+
or not all(isinstance(value, str) and value for value in values)
978+
or len(values) != len(set(values))
979+
):
980+
raise CodegenError(f"endpoint_generator.{key} must be a unique list of non-empty strings")
844981
for key in ("supported_request_media_types", "supported_response_media_types", "supported_success_statuses"):
845982
values = endpoint.get(key)
846983
if not isinstance(values, list) or not values or not all(isinstance(value, str) for value in values):

0 commit comments

Comments
 (0)