[solarsys]

What this adds

pyhealth.autoimprove improves a PyHealth pipeline by guided search. You give it training and validation splits and your current configuration. A few agents take turns proposing configurations: the model, its constructor arguments and its training settings. Proposals come from an LLM (LLMProposer with ClaudeLLM, or any llm(system, prompt) -> str callable) or from random search (RandomProposer, no API key). The agents see each other's results and rotate through four strategies: explore, refine, crossover and contrast.

improver = AutoImprover(
    train, val, LLMProposer(ClaudeLLM()),
    baseline=PipelineConfig("RNN", {"hidden_dim": 128}, {"epochs": 10}),
    cv_folds=5, run_dir="runs/mortality",
)
result = improver.run()                  # resumes from run_dir
print(result.summary())
print(improver.evaluate_on_test(test))   # selected vs baseline, once, on test

Design

• Proposals are data, never code. Every proposal is a JSON config checked against a SearchSpace before anything runs. Unknown models are rejected, unknown arguments are dropped, out-of-range numbers are clipped, and the proposer is told what changed.
• Scoring: on the validation split, or by patient-grouped cross-validation over the training split (cv_folds). With CV, the validation split is used only for early stopping.
• Selection: with cv_folds, every trial is compared with the baseline fold by fold. The selection is the trial whose gain has the highest lower bound (mean minus promote_se standard errors). The baseline stays selected unless some bound is above 0. This guards against picking a config that won by luck on a small validation split.
• Model check: before searching, each model in the space takes one training step on a few training samples. Models that can't read the task's inputs are removed, for example XGBoost on code sequences. They are logged and listed in ImproveResult.dropped_models, so no trial is spent on them.
• Lessons: after each cycle the LLM proposer rewrites a short list of lessons from the results. Later prompts use them, and so do later runs via lessons_path. Lessons are stored per search, keyed by the inputs, the label and the models.
• Prompts hold aggregates only: the search space, the schema, split sizes, positive counts and scores. They never contain patient records.
• Test split: never used during the search. evaluate_on_test retrains the selected and baseline configs and compares them with patient-level bootstrap intervals and a paired difference.

The package is new and touches nothing outside pyhealth/autoimprove, apart from one line in docs/index.rst. anthropic is imported only when ClaudeLLM is used.

Tested on a second task: the benchmark's 48 h mortality task on the eICU demo

examples/autoimprove_eicu_ihm48.py runs the clinical prediction benchmark's T1 task, in-hospital mortality predicted 48 h after admission (ihm48_eicu@v1), on the open eICU-CRD demo from PhysioNet; no credentials are needed. The cohort is 1,608 hospital stays with 124 deaths, split 70/10/20 by patient. Each trial is scored by 5-fold CV, and APACHE IVa's predicted mortality on the same test stays is the clinical reference. Two input views were searched, each with random search and with Claude:

View, proposer	Trials	Selected	CV PR-AUC, baseline → selected	Test PR-AUC, selected vs baseline
Tabular (labs and vitals summaries), random	19	XGBoost, depth 3	0.216 → 0.250 (5/5 folds better)	0.303 vs 0.262
Tabular, Claude	19	XGBoost, depth 2, slow learning rate	0.216 → 0.282 (5/5)	0.356 vs 0.262
Hourly series plus codes, random	13	RETAIN	0.240 → 0.270 (4/5)	0.198 vs 0.143
Hourly series plus codes, Claude	10	GRU baseline kept	0.240	0.143

With 23 test deaths, none of the test differences is statistically significant. The demo is a check that the loop works on a new dataset, task and input types, not a benchmark result. APACHE IVa scored 0.422 test PR-AUC.

The eICU runs found four problems, all fixed in this PR:

• Transformer embedding size. The default space could draw an embedding_dim that is not divisible by heads, which fails an assert in attention. It is now a choice of powers of two.
• Bare assert messages. Failed trials recorded a bare AssertionError:. Errors now include the file, line and source line that raised (describe_error), so the proposer can act on them.
• Weight decay range. It was capped at 0.01, and Claude hit the cap on this small dataset. It is now log scale from 1e-6 to 0.1.
• Stopping rule. patience counted only cycles without a new selection. Because the selection rule is conservative, one search stopped after 10 trials while it was still finding better means. A cycle now also counts as progress when some trial beats the best score so far.

In the eICU example, LogisticRegression and MLP on the time-series view need #1315. Until it merges, the model check removes them from that view.

Tests

• tests/core/test_autoimprove.py: 22 tests. They cover the space checks, JSON parsing, the LLM loop with retries, lessons and resume, repeated proposals, failed trials as feedback, patient folds, CV scoring, the paired selection rule, scoped lessons, the model check, the stopping rule and the test comparison. All pass, as do the doctests.
• tools/check_pr_rules.py passes and ruff is clean on pyhealth/autoimprove.

🤖 Generated with Claude Code

https://claude.ai/code/session_015ntMLEmFsQEQJKTmPGGw4T

[github-actions]

Test Coverage

Total: 68.78% (17490 / 25428 lines)

Modified file	Coverage	Lines
pyhealth/autoimprove/__init__.py	100.00%	6 / 6
pyhealth/autoimprove/evaluate.py	92.07%	151 / 164
pyhealth/autoimprove/improver.py	95.63%	372 / 389
pyhealth/autoimprove/llm.py	90.62%	29 / 32
pyhealth/autoimprove/proposers.py	93.91%	108 / 115
pyhealth/autoimprove/space.py	95.31%	122 / 128

Workflow run · 79a36a7

[Logiquo]

/review

[github-actions]

Previous review

AI review: findings

Adds pyhealth.autoimprove, a guided search over models and training settings. An LLM (ClaudeLLM) or random proposer suggests configs, each is checked against a SearchSpace, scored on validation or by patient-grouped CV, and the selection uses paired fold gains against the baseline. evaluate_on_test retrains best and baseline once on the test split. It reuses the existing Trainer, get_dataloader, SampleDataset.subset, PatientSplit, bootstrap_ci and paired_bootstrap_diff and changes no base class or public API. I found no existing search, k-fold or LLM component that duplicates it. Blocking (medium): (1) pyhealth/autoimprove/improver.py:652, select_trial with cv_folds and no baseline uses the first finished trial as the reference, so the selection depends on trial order and the "baseline stays selected" rule has no meaning; (2) improver.py:531, reusing a run_dir with a different metric, cv_folds, baseline or space silently mixes trials and can compare against the old baseline. Non-blocking (low): exception text up to 300 chars goes verbatim into LLM prompts (evaluate.py:302), which the docs say never happens; patient_folds is not stratified (evaluate.py:160); one LLM refusal or API error aborts the run (improver.py:388); the Claude beta header, fallbacks and output_config are tested only against a fake client (llm.py:75); and the eICU example's docstring does not say LogisticRegression and MLP are dropped until #1315 merges. Tests: 22 unit tests on tiny synthetic data cover the loop, selection, CV folds, model check and the test comparison, but none cover the no-baseline CV path or resuming with changed settings, and the eICU results need PhysioNet data. Overall risk is moderate-low: the change is isolated, but the two selection and resume problems can silently return a wrong best config and should be fixed before merge.

_{Reviewed commit 1e5ed53}

[github-actions]

Previous review

AI Review

suggestion: need change
decision: need change

Thank you for this contribution. pyhealth.autoimprove adds a guarded, resumable search over models and training settings. It has a random-search option that needs no API key, and a test-split comparison that uses bootstrap intervals. The proposals-are-data design, the model check and the paired selection rule are all well thought out. Test coverage meets the bar: every added file is above 80% (80.8% to 100%), and the whole repository is at 68.6% against the 60% bar. You are close, but two problems can silently return a wrong "best" config.

To get this merged:

1. Selection with cv_folds and no baseline (select_trial in improver.py). With no baseline trial, the reference becomes done[0], the first finished trial. The selected config then depends on which trial happened to finish first, and "the baseline stays selected" has no meaning. Please do one of the following:

   • Require a baseline when cv_folds/promote_se is used, and raise a clear error if there isn't one.
   • Fall back to best-mean selection when there is no baseline.

   Add a test for this path.

2. Resuming with changed settings (_load in improver.py). Reusing a run_dir with a different metric, cv_folds, baseline or search space silently mixes old and new trials, and can compare new trials against the old baseline. Please store the run settings in run_dir, for example in a small meta.json. On resume, check them against the current settings. Raise an error or warn and start fresh if they differ. Add a test that resumes with a changed setting.

Optional suggestions:

• Error text of up to 300 characters goes verbatim into LLM prompts (describe_error), but the docs say prompts hold aggregates only. Either soften the docs or sanitise the error text.
• patient_folds is not stratified. With few positives, a fold could end up with none.
• A single LLM refusal or API error aborts the whole run. Consider treating it as a failed turn.
• The eICU example's docstring should mention that LogisticRegression and MLP are dropped on the sequence view until #1315 merges.

_{Reviewed commit 1e5ed53}

[solarsys]

Thanks for the review. All items are addressed in 011d115, each with a test (30 tests in test_autoimprove.py, all passing; the PR rules check passes).

Required

1. Selection without a baseline. select_trial no longer uses the first finished trial as the reference. Without a finished baseline trial with fold scores, there is nothing to compare with, so the best mean is selected. Ties go to the lower trial_id, so the result never depends on trial order. The prompt says "There is no baseline to compare with, so the selection is the config with the best mean score". AutoImprover also warns when promote_se is set without a baseline. Tests: test_without_a_baseline_the_best_mean_is_selected (both orders, plus a failed baseline) and test_cross_validation_without_a_baseline (an end-to-end run with cv_folds and no baseline).
2. Resuming with changed settings. run_dir/meta.json records the metric, cv_folds, promote_se, seed, baseline, search space (with ranges) and data (schema and sample counts).
   • Changed settings: resuming with any of them changed raises ValueError, naming the settings that differ.
   • Budgets: n_cycles, max_trials and the like may still change.
   • Old runs: a run saved before meta.json existed resumes with a warning, and its settings are recorded from then on.
   • Test: test_resuming_with_changed_settings_is_refused covers a changed metric, baseline and cv_folds, a longer budget, and a run without meta.json.

Optional

• Error text in prompts. Trial.error keeps the full text in trials.jsonl. Prompts get only Trial.error_brief: the exception type and the code line that raised it, without the message (describe_error(e, message=False)). The docs now say so. Test: test_error_messages_stay_out_of_prompts.
• Stratified folds. patient_folds deals patients with a positive label first, each to the fold with the fewest, so every fold gets positives when there are at least as many positive patients as folds. It warns when there are fewer. A NaN CV score now reads "undefined on a cross-validation fold". Test: test_patient_folds_are_stratified.
• LLM errors. A proposer error (an API error or a refusal) skips the turn with a warning. Three in a row stop the run with a RuntimeError, so a broken backend doesn't use up the budget. A failed reflection keeps the current lessons. Test: test_proposer_errors_skip_the_turn.
• Claude SDK. ClaudeLLM requires anthropic>=1.15 (MIN_ANTHROPIC) and says so in the docs and the ImportError. test_live_request sends a real request when ANTHROPIC_API_KEY is set and is skipped otherwise. It passed locally against the API with anthropic 1.15.0.
• eICU example.
  • The docstring says LogisticRegression and MLP are dropped on the sequence view until MLP and LogisticRegression: embed time series per step, not flattened #1315 merges.
  • It says InHospitalMortality48hEICU lives only in this example. It is not the same task as MortalityPredictionEICU.
  • The references to "spec section 8" and configs/models/gru.yaml now point to the benchmark plan, which is outside this repository.

The stacked drafts (#1317–#1321) are rebased onto this and apply the same no-baseline rule:

• autoimprove: explained runs (decision log, one-SE selection, attribution, report) #1317: one-SE selection and the report.
• autoimprove: feature search (alternating cycles, confirmation, pruning, review gate) #1320: confirmation, plus skipping a turn when feature proposals hit an LLM error.

[Logiquo]

/review

[github-actions]

Previous review

AI review: findings

Adds a new pyhealth/autoimprove package: a validated search space, random and LLM proposers (ClaudeLLM or any callable), a search loop with optional patient-grouped CV and paired fold-by-fold selection, a model check that drops models that cannot read the inputs, and evaluate_on_test. It also adds two example scripts and an API doc page, and changes only one line in existing files (docs/index.rst). It reuses existing components (Trainer, get_metrics_fn, bootstrap_ci, paired_bootstrap_diff, XGBoostModel, PatientSplit, BaseModel.set_pos_weight), and I found no existing tuning, patient-grouped k-fold or Anthropic client that it duplicates. Blocking (medium): the only test of the model-check removal path (tests/core/test_autoimprove.py:434) is skipped in CI, because the pixi test environment does not install xgboost, so the drop logic in improver.py:350 never runs there. Non-blocking (low): a reply whose model_params is not a dict raises TypeError, which proposers.py:182 does not catch, so it counts as an API failure; _positive_patients (evaluate.py:207) scans the whole training set on every construction; and the eICU example retrains the best config a second time (line 491). Tests cover the selection rule, resume, lessons scoping, retries and the test comparison, but not the non-binary test path or the balanced pos_weight path; I did not run the suite. Prompts were checked and carry only aggregates, schema and error source lines. Overall risk is low: the change is self-contained and the one blocking issue is test coverage.

_{Reviewed commit 011d115}

[github-actions]

Previous review

AI Review

suggestion: need change
decision: need change

Thank you for this contribution. pyhealth.autoimprove adds a guided search over models and training settings. Proposals are validated as data, trials can be scored by patient-grouped CV, and selection compares each trial with the baseline fold by fold. A model check drops models that can't read the task's inputs, and evaluate_on_test compares the selected and baseline configs once on test. The design is careful, and the eICU example is a good check on a second task.

Coverage meets the bar. The whole repository is at 68.7%, above the 60% minimum. Every added file is above 80%: __init__.py 100%, evaluate.py 90.6%, improver.py 92.4%, llm.py 90.6%, proposers.py 93.8% and space.py 95.1%. I didn't run the suite myself.

It is close, but one gap needs closing before merge.

To get this merged:

• Test the model-removal path in a way that runs in CI. test_models_that_cannot_read_the_inputs_are_removed is the only test of the dropping logic in AutoImprover._check_models. It carries skipUnless(find_spec("xgboost")), and the CI test environment doesn't install xgboost, so the test is skipped there. The behaviour this PR advertises (dropped models, the dropped_models field, the baseline-fails error, the "every model fails" error) is therefore not checked in CI. Please add a test that doesn't need xgboost. For example, add a ModelSpec whose target is a small class that raises in __init__ or forward. The test can then assert the following:
  • The model lands in dropped_models and is removed from improver.space.models.
  • summary() contains the "Not searched" line.
  • A failing baseline model raises ValueError.
  • A space where every model fails raises ValueError.

Optional suggestions:

• LLMProposer.propose only catches ValueError. A reply where model_params is not a dict (for example a list) makes PipelineConfig.from_dict raise TypeError, and the improver counts that as a backend failure. Catching TypeError too, so it triggers the "reply again" path, would be more robust.
• _positive_patients scans the whole training set each time patient_folds is called. Caching it would save time on large datasets.
• The eICU example retrains the selected config a second time just to get probabilities for the APACHE comparison. You could reuse the model that evaluate_on_test already trained.
• Add tests for the non-binary evaluate_on_test path and for pos_weight="balanced".

_{Reviewed commit 011d115}

[solarsys]

Thanks for the second review. Everything is addressed in 79a36a7. test_autoimprove.py now has 36 tests, and all pass when run with the CI command (unittest discover -t tests -s tests/core). The PR rules check passes too.

Required

• The model check runs in CI. The new test_models_that_fail_the_check_are_removed needs no xgboost. It adds a ModelSpec for a test class that raises in __init__ (UnreadableInputs) and asserts each of the following:

  • the model lands in dropped_models, with its error;
  • it is removed from improver.space.models, and every trial runs;
  • summary() has the "Not searched" line;
  • a failing baseline model raises ValueError;
  • a space where every model fails raises ValueError.

  The xgboost test is kept as an extra check.

Optional

• Malformed replies. PipelineConfig.from_dict raises TypeError ("model_params must be an object, not int") when model_params or train_params is not an object. LLMProposer.propose catches (TypeError, ValueError), so the LLM is asked again with that message instead of the turn counting as a backend failure. In reflect, a lessons value that is not a list keeps the current lessons. Tests: TestLLMProposer.
• Label scan. patient_folds reads labels for stratification only up to MAX_LABEL_SCAN (50,000) samples, the same cap as the positive count, which now shares the constant. Larger datasets get unstratified patient folds and an info log. Test: added to test_patient_folds_are_stratified.
• _label_mode. The label type is now resolved from PROCESSOR_REGISTRY, so the unbound call to BaseModel._resolve_mode is gone. test_label_type_sets_the_metric_and_the_space checks binary (as a string and as a processor class), multiclass, multilabel and regression against the expected default metric.
• Retraining in the eICU example. evaluate_on_test(..., return_predictions=True) also returns y_true, patient_ids and each config's y_prob. The example uses these for the ROC-AUC and APACHE comparison instead of training the best config again.
• Tests for the non-binary path and pos_weight="balanced". Writing the multiclass test found a bug. pos_weight stayed in the search space for multiclass and regression labels, but set_pos_weight raises for them, so every "balanced" proposal failed. AutoImprover now drops pos_weight from the space unless the label is binary or multilabel. A baseline that sets it gets a validation note instead of an error.
  • test_multiclass_search_and_test_scores runs a multiclass search with such a baseline and checks that no trial carries pos_weight. It then checks the evaluate_on_test output: estimates for best and baseline, no paired difference, and predictions of shape (n, 3).
  • test_balanced_pos_weight_uses_the_training_labels checks the weight against negatives / positives in the training split.

[Logiquo]

/review

[github-actions]

AI review: findings

The PR adds pyhealth.autoimprove, a guided search over PyHealth models and training settings. Proposals come from an LLM (ClaudeLLM or any callable) or from random search, each proposal is validated against a SearchSpace, and trials are scored on validation or by patient-grouped CV. Selection compares every trial to the baseline fold by fold, models that cannot read the inputs are dropped by a one-step check, and evaluate_on_test gives bootstrap intervals. It adds new files plus one docs index line and changes no base class, registry or existing API; it reuses Trainer, SampleDataset.subset, get_dataloader, set_pos_weight, bootstrap_ci and paired_bootstrap_diff. Nothing in the repo duplicates it: the splitter has no k-fold, so patient_folds is new, and bootstrap is reused rather than copied. I checked the default search-space argument names against the model constructors, the XGBoost fit path, the selection and stopping logic, and the prompt contents; I found no blocking (high or medium) problems. Non-blocking: the docs promise a warning below 30 positives that is skipped under cv_folds; lessons are keyed only by schema key names and models, so different tasks with the same key names share lessons; the reflect prompt grows with every trial; and the live ClaudeLLM request shape (beta header, fallbacks, output_config) is covered only by mocks since the live test is skipped by default. The description says the tests number 22, but the file has 36 test functions. The eICU example depends on the unmerged #1315, which the code comments state. Overall risk is low, and the tests cover the new behavior well for a package of this size.

_{Reviewed commit 79a36a7}

[github-actions]

[low] Docs promise a small-validation warning that is skipped under cv_folds

The text says the improver warns below 30 positives. In code, _warn_if_small_validation returns early when cv_folds is set (pyhealth/autoimprove/improver.py:395), so with CV no warning is logged. Either qualify the sentence ("without cv_folds") or drop the parenthetical.

[github-actions]

[low] Lessons are shared across different tasks that use the same schema key names

describe_search keys lessons by input and label key names and types plus the model list, not by task identity or label prevalence. Two different tasks that both use conditions and a label binary output will reuse each other's lessons, while the reflect prompt says they apply to 'THIS task'. This matches the documented design, but the risk is real. Consider adding a user-supplied task name to the fingerprint, or stating the limitation in the docs.

[github-actions]

[low] Reflection prompt includes every trial, so it grows without bound

LLMProposer.reflect formats all trials into the prompt (format_trials(trials)). With max_trials large, the prompt and its cost grow with each cycle. Consider sending only the top and worst trials plus recent ones, as build_prompt already does for proposals.

[github-actions]

[low] ClaudeLLM request shape is untested live, and fallback is on by default

The beta header server-side-fallback-2026-07-01, fallbacks, and output_config are checked only with mocked clients in tests/core/test_autoimprove.py. The one live test is skipped unless ANTHROPIC_API_KEY is set. With fallback=True (the default), a declined request may be answered by a different model without the caller knowing. Document the substitution where the data-use note appears, and record which model answered if the SDK exposes it.

[github-actions]

[low] Example depends on unmerged #1315 and cites an external plan in code

The docstring says LogisticRegression and MLP cannot read the time series until #1315 merges, and elsewhere the code cites 'the benchmark plan', which is not in this repository. The model check removes those models, so the example still runs, but the PR number and external references will go stale. Describe the limitation without the PR number, and drop the references to a document a reader cannot see.

[github-actions]

AI Review

suggestion: approve
decision: approve

Thank you for this contribution. pyhealth.autoimprove adds guided search over models and training settings. It includes proposals checked against a SearchSpace, patient-grouped CV with a paired selection rule, a model check before searching, and a final test comparison with bootstrap intervals. The eICU example and the docs page are useful additions.

It looks ready to merge.

Coverage: the whole repository is at 68.78%, above the 60% bar. Every added file is above the 80% per-file bar:

• __init__.py: 100%
• evaluate.py: 92.1%
• improver.py: 95.6%
• llm.py: 90.6%
• proposers.py: 93.9%
• space.py: 95.3%

The package is new and its only edit outside pyhealth/autoimprove is one line in docs/index.rst. It reuses the existing Trainer and bootstrap utilities rather than copying them.

Optional suggestions:

• The docs promise a warning when the validation split has fewer than 30 positives. That warning is skipped when cv_folds is set, so a short note in the docs would avoid confusion.
• Lessons are keyed by input and label key names plus the models. Two different tasks with the same key names would share lessons, so you could add something that identifies the task to the key.
• The reflect prompt lists every trial, so it will grow over long runs. Capping it to the top and bottom trials would keep it bounded.
• The live ClaudeLLM request shape (beta header, fallbacks, output_config) is only tested with mocks unless an API key is set. A periodic live run would catch SDK drift.
• The description says 22 tests, but the test file has more than that. It's worth updating the number.
• The eICU example depends on #1315, which you note clearly. Please keep that comment accurate once #1315 merges.

_{Reviewed commit 79a36a7}