Add LLMRecommender for LLM-based experimental design by tobiasploetz · Pull Request #855 · emdgroup/baybe

tobiasploetz · 2026-07-07T14:10:00Z

Add LLMRecommender for LLM-based experimental design

Summary

Adds an LLMRecommender that queries a large language model (via LiteLLM) to propose experiments. It is primarily intended for warm-starting a campaign: before enough data exists to train a surrogate, the model can use parameter semantics and domain/literature priors to suggest more sensible points than random sampling. It slots into the recommender hierarchy like any other pure recommender and is typically combined with a Bayesian recommender via a TwoPhaseMetaRecommender.

Ported from #561 and substantially reworked for the current codebase.

What it does

Builds a natural-language prompt from the search space (parameter names, types, ranges/allowed values, and any metadata descriptions/units), the textual experiment_description / objective_description, the structured objective, and any collected measurements / pending_experiments.
Sends the prompt to the configured model through LiteLLM, parses the returned JSON into candidate experiments, and validates them against the search space.
Recovers from malformed responses via a one-shot correction prompt, and tolerates models that wrap JSON in Markdown fences or surrounding prose.

Key design decisions

Inherits directly from PureRecommender. It legitimately consumes measurements/pending_experiments for the prompt but is neither non-predictive nor Bayesian. recommend() is fully overridden and builds the result directly from the model response, so the recommend* hooks are intentionally not implemented.
Credentials are not a field. LiteLLM reads provider keys from environment variables (e.g. OPENAI_API_KEY, ANTHROPIC_API_KEY) selected by the model prefix — so there is deliberately no API-key attribute.
Scope kept tight. There were some fields that were not really necessary for the first version like related_data or a feasibility check. These have been removed.
Module layout. Prompt construction (_prompts.py) and response parsing/validation (_parsing.py) live in dedicated submodules; llm.py holds the recommender class. litellm/jinja2 are lazy-imported so import baybe stays light.

Testing

Unit tests are fully mocked — no real LLM calls in CI (@pytest.mark.skipif(not LLM_INSTALLED)); they skip cleanly under coretest.
Covers success paths, prompt inclusion of measurements/pending/objective, the full set of parse-error branches, code-fence/prose stripping, error wrapping, and a serialization roundtrip.
LLMRecommender is excluded from the automatic get_subclasses test enumerations (it has required constructor args) via an explicit, documented helper in conftest.py.

Docs & example

New "LLM Recommenders" section in docs/components/recommenders.md.
Runnable example (examples/LLM/llm_recommender.py) comparing an LLM warm-start (then Bayesian) against BayBE's default recommender and a random baseline on the direct-arylation dataset; it mocks the LLM when no credentials are present, and the committed comparison images are pre-generated (the docs build is mocked).

tobiasploetz

Converting discussion comments into review threads for easier inline discussion.

AVHopp · 2026-07-08T09:29:15Z

@Scienfitz @AdrianSosic @kalama-ai I will do the first round of reviews/feedback, feel free to ignore until we move it out of draft :)

AVHopp

First round of comments. I hope that impementing those will solve the seralization issues as well as the most important design choice such that we can then see the CI in action. Note that I did not yet investigate the tests or the inner workings of the class in detail, but I think that doing small iterations like this is better than me providing too many comments :)

Copilot

Pull request overview

This PR introduces an LLMRecommender to BayBE’s pure recommender stack, enabling LLM-driven experiment suggestions (via LiteLLM) for warm-starting campaigns before surrogate-based Bayesian optimization becomes effective.

Changes:

Add LLMRecommender implementation with prompt construction and robust JSON response parsing (including one-shot recovery).
Add optional dependency group llm (Jinja2 + LiteLLM) plus optional-import plumbing and mypy config.
Add tests, documentation, and an end-to-end example (with pre-generated comparison plots) and adjust iteration tests to avoid auto-instantiating non-default-constructible recommenders.

Reviewed changes

Copilot reviewed 19 out of 22 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
`baybe/recommenders/pure/llm/llm.py`	Implements `LLMRecommender` with LiteLLM call, recovery flow, and prompt/parse integration.
`baybe/recommenders/pure/llm/_prompts.py`	Builds Jinja2 prompts from search space + descriptions + optional measurements/objective.
`baybe/recommenders/pure/llm/_parsing.py`	Extracts/validates JSON suggestions into a DataFrame.
`baybe/recommenders/pure/llm/__init__.py`	Exposes `LLMRecommender` from the `llm` subpackage.
`baybe/recommenders/pure/__init__.py`	Re-exports `LLMRecommender` in pure recommenders.
`baybe/recommenders/__init__.py`	Re-exports `LLMRecommender` at the recommender package level.
`baybe/_optional/llm.py`	Adds optional imports for Jinja2 Template + LiteLLM completion.
`baybe/_optional/info.py`	Adds `LLM_INSTALLED` check based on Jinja2 + LiteLLM presence.
`baybe/exceptions.py`	Adds `LLMResponseError` and `LLMResponseWarning`.
`pyproject.toml`	Adds `baybe[llm]` extra with `jinja2` + `litellm`.
`mypy.ini`	Adds mypy ignore section for `litellm.*`.
`tests/test_llm_recommender.py`	Adds mocked unit tests for prompting, parsing branches, recovery, and warnings.
`tests/serialization/test_llm_recommender_serialization.py`	Adds serialization roundtrip test via meta recommender wrapper.
`tests/conftest.py`	Adds helper to exclude non-default-constructible recommenders from subclass enumeration tests.
`tests/test_iterations.py`	Switches subclass enumeration to the new default-constructible helper.
`docs/components/recommenders.md`	Documents the new LLM recommender and the `baybe[llm]` extra.
`examples/LLM/llm_recommender.py`	Adds runnable warm-start example with mock fallback for CI/offline runs.
`examples/LLM/LLM_Header.md`	Adds header page for the LLM examples section.
`examples/LLM/llm_recommender_light.svg`	Adds pre-generated plot asset (light theme).
`examples/LLM/llm_recommender_dark.svg`	Adds pre-generated plot asset (dark theme).
`CHANGELOG.md`	Adds changelog entry for `LLMRecommender`.

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

Scienfitz

currently does not use all available metadata

AdrianSosic

First round of comments, more coming

AdrianSosic · 2026-08-27T12:41:30Z

+OPTIMIZATION TARGETS:
+{% for target in objective.targets %}
+Target: {{ target.name }}
+{% if target.metadata.description is not none %}


I'm not familiar with jinja logic, but are the conditions supposed to be python code? I.e. why none and not None? Perhaps you can enlighten me how to read this stuff 😅

Let me summon @tobiasploetz here in case he knows an easy and quick win, I'll investigate.

none seems to be the way that this is written in Jinja, even though it appears that None would also be fine. So I'd say we keep this.

@AdrianSosic please resolve if you agree with this if not tell me what to do differently/verify.

AdrianSosic · 2026-08-27T12:45:22Z

+{% if pending_experiments is not none and not pending_experiments.empty %}
+PENDING EXPERIMENTS:
+The following experiments have already been proposed and are awaiting results.
+Do not recommend these again.


I know, perhaps it doesn't even make a difference for the LLM, but perhaps we should at least state it correctly in our prompt. Instead of saying do not recommend again, how about something like or similar? Feel free to adjust or let me know if this is bullshit

Suggested change

Do not recommend these again.

Take these into account when generating your recommendations. In particular, only consider re-recommending parts of them if you have a strong reason to do so, for instance, when you are convinced that a second measurement is needed for noise reduction.

I think the prompt-tuning can only be done after some experimentation, and I like your block here. Makes more sense imo so I'll probably just add that

In general, prompt tuning should be done with a proper evaluation set of tasks. The proposed modification looks very task-dependent.

In this case, I'd say any sort of prompt-changing that actually changes content is deferred from now. Is that fine for you @AdrianSosic ?

Let's use this thread to talk about the prompt and prompt tuning.

I realized that some prompt-tuning is actually necessary, as some of the models are really having issues with e.g. recommending the same point over and over, even if they are told not to do to. Once I have finished testing, we can start discussing the actual prompt here.

The handling of the recovery as well as the "forbidden configurations" has changed quite a bit. Most of this has been done in 3b35a86, but it might make more sense/be easier to take a holistic look at how this is handled.

AdrianSosic

I'm sorry to be the showstopper here, but I think the PR requires a severe structural overhaul 😬 Much of the functionality is scattered / repeated (comments are pointing to several such places) and there are many smaller issues to fix as well. Haven't looked at the examples and tests yet, but we've already several dozen comments. Could it be that the PR was largely generated using LLMs – at least I can see the typical issues here...

If you want, I'm happy to have a general brainstorming session first, but let me know how you wish to proceed and how I can help best 👌🏼

AdrianSosic · 2026-08-27T15:16:44Z

+
+
+def build_prompt(
+    searchspace: SearchSpace,


RE signature: any particular reason why you chose this layout (i.e. ordering and positional vs kwonly)? I'm asking because the signature has a large overlap with other signatures we use in a recommendation context (e.g. RecommenderProtocol) and there we have established a sort of canonical ordering of the arguments. Would be nice if we could replicate that here

Nope, and has been changed in 9d9adcf. Please let me know if the corresponding comment in the docstring is helping or should be removed.

AdrianSosic · 2026-08-28T11:06:09Z

I'm sorry to be the showstopper here, but I think the PR requires a severe structural overhaul 😬 Much of the functionality is scattered / repeated (comments are pointing to several such places) and there are many smaller issues to fix as well. Haven't looked at the examples and tests yet, but we've already several dozen comments. Could it be that the PR was largely generated using LLMs – at least I can see the typical issues here...

If you want, I'm happy to have a general brainstorming session first, but let me know how you wish to proceed and how I can help best 👌🏼

@tobiasploetz: I'm talking to @AVHopp, in case this wasn't clear 🙃

AVHopp · 2026-08-29T08:59:08Z

@AdrianSosic thank you for the detailed review :) This is exactly what I expected. As far as I know, this PR was created using a lot of agentic assistance, both on my and @tobiasploetz end. So a lot of the things that you uncovered were to be expected, and I will take care of those.

I think that there are some major points that we just need to align on to see how to shape this PR. The obvious one is everything around type safety, the other one is validation, which also touches upon the question of what exactly the PR should ship (e.g. do we want to restrict to discrete spaces for now). But I think it makes most sense to just discuss this in the corresponding threads.

While these discussions are happening, I will start working on the more "mechanical" parts of the feedback - nomenclature, references and so on.

Scienfitz · 2026-09-28T13:37:14Z

+    )
+
+
+def make_llm_alternating_recommender(


please lets introduce the two convenience recommenders in a format thats consistent with the other - a class

imo no need to make a factory here, did you consider jut deriving a special subclass from sequential and 2phase recommender?

Implemented in eb4768a, please resolve if this structure is now as you imagined it.

Scienfitz · 2026-09-28T13:50:50Z

+{{ objective.metadata.description }}
+
+{% endif %}
+OPTIMIZATION TARGETS:


please correct me if im wrong, but doesnt this objective section severly lack information still?

Example: The basic target informations like whether to maximize or minimize or transformations are not included. So the only info the llm can go off is objective.metadata. But the structure of single targets and how they are combined (which is objective and targets non-metadata info) should always be included too.

This leads for instance to the identical prompt for these two situations: NumericalTarget("y") and NumericalTarget("y", minimize=True). We should definitely not rely on the objective metadata correting this. For all we know obejctive metadat is allowed to be empty (just warned).

Handling of obejctives has been improved in 0d90d46 and they are now handled in a similar way that we handle parameters.

Question maybe worth discussing: Do we also want to share all of the information? For example, we now also share the transformation, not sure if this is actually needed.

Scienfitz · 2026-09-28T13:53:47Z

+                error=initial_error,
+                original_response=content,
+            )
+            recovery_content = self._query_model(recovery_prompt)


this sends a fresh query and dropped all the context of the original call (all important info like obejctive and contetx etc).

Is this intended? Is this not a problem because it is guaranteed that the LLM endpoint holds that context? How does sucha. query here work? OD we have to take care of carrying the context or not? Depending on that this might have a very bad influence here for the recovery call @tobiasploetz

Recovery mechanism has been changed in general in 60afb19. It seems like every call is stateless (or at least one should assume so), hence this recovery mechanism now works by fully sending the original response together with what went wrong. This has the nice effect of us only needing a single template and function, so in my opinion this is quite an elegant solution if I may say so :) Resolve/comment if you agree/disagree.

Scienfitz · 2026-09-28T15:06:47Z

+                **self.litellm_args,
+            )
+        except Exception as e:
+            raise LLMResponseError(


strictly speaking there was no response, wouldnt be calling it an auth error or similar be more precise and more useful for filtering?

Has been changed in 77b4b2e where we improve general error handling, so please resolve here if this is to your liking.

Without this, `mypyc` would not correctly recognize the type of the recommenders list. Instead of an `ignore` comment, I added this typed converter instead.

…Constraint in tests

Testing revealed that models often repeat recommendations, even if they should not do so. Since we cannot directly hand the corresponding flags easily down to the recommender, we attempt to include forbidden configurations directly in the prompt.

… order Reorder `make_prompt` to (batch_size, searchspace, objective, measurements, pending_experiments), matching `RecommenderProtocol.recommend`, and make the LLM-specific experiment_description keyword-only.

This is more in-line with our usage patterns. We use our `to_tuple` converter as we cannot use a plain `tuple` due to python/mypy#8417.

Fixes two small issues: First, the prompt now only propagates `active_values` instead of all values. Second, it fixes a leak in `np` representation for numerical values.

There was a request to fully rewrite those, hence the current version was deleted.

Since `litellm` seems to be typed since this version this bump allows us to remove the package from previous exceptions regarding typing.

All LLM-related exceptions now carry an `LLM` prefix in their class names. The hierarchy was extended by differentiating between general call errors (with authentication errors as one special subcase), and response errors.

The `recommend` call now stashes necessary context for creating the prompts and delegates to `_recommend_hybrid`. As a consequence, validation and preprocessing of the `super` call is used and has been removed.

It has been decided that those should be added in a follow-up PR in order to decrease the scope of this one

Previously, some of the context and information was not used in the recovery prompt. To fix this, there is now a single general prompt that includes the full context, with an optional block containing recovery-related information if necessary.

We did not carry all of the relevant information for targets and objectives in the prompt. This has now been fixed by including views of both targets and objectives in the prompt, similar to the way we do it for parameters.

It seems like newer versions of `mypy` require more detailed type information that the simple `type` annotation we had in `find_subclass`

Scienfitz mentioned this pull request Jul 8, 2026

Feature: Add LLMRecommender #561

Closed

tobiasploetz commented Jul 8, 2026

View reviewed changes

AVHopp assigned tobiasploetz Jul 8, 2026

AVHopp added the enhancement Expand / change existing functionality label Jul 8, 2026

AVHopp reviewed Jul 8, 2026

View reviewed changes

AVHopp self-assigned this Jul 29, 2026

AVHopp force-pushed the feature/llm-recommender branch from e419a9d to f62c14b Compare July 29, 2026 10:52

AVHopp force-pushed the feature/llm-recommender branch from 8184fa5 to 84ed98c Compare August 11, 2026 14:49

AVHopp marked this pull request as ready for review August 11, 2026 15:08

AVHopp requested a review from AdrianSosic as a code owner August 11, 2026 15:08

Copilot AI lite review requested due to automatic review settings August 11, 2026 15:08

Copilot started reviewing on behalf of AVHopp August 11, 2026 15:09 View session

Copilot AI reviewed Aug 11, 2026

View reviewed changes

Comment thread baybe/recommenders/pure/llm/llm.py Outdated

Comment thread baybe/recommenders/pure/llm/_parsing.py Outdated

Scienfitz requested changes Aug 18, 2026

View reviewed changes

AVHopp force-pushed the feature/llm-recommender branch 4 times, most recently from b1c8cb4 to 5630f51 Compare August 23, 2026 17:38

AdrianSosic requested changes Aug 27, 2026

View reviewed changes

AdrianSosic requested changes Aug 28, 2026

View reviewed changes

Scienfitz added this to the 0.16.0 milestone Sep 1, 2026

AVHopp marked this pull request as draft September 3, 2026 05:45

AVHopp force-pushed the feature/llm-recommender branch from 1aa662f to 641592c Compare September 7, 2026 06:41

AVHopp marked this pull request as ready for review September 10, 2026 07:09

AVHopp reviewed Sep 10, 2026

View reviewed changes

Comment thread examples/LLM/llm_recommender.py Outdated

AVHopp force-pushed the feature/llm-recommender branch from e36fce4 to 9a17935 Compare September 11, 2026 12:52

Scienfitz requested changes Sep 28, 2026

View reviewed changes

AVHopp added 21 commits October 8, 2026 16:22

Add type guard

031f307

Add typed converter

2c27aa6

Without this, `mypyc` would not correctly recognize the type of the recommenders list. Instead of an `ignore` comment, I added this typed converter instead.

Fix typing

578afa4

Use DiscreteSelectionConstraint instead of deprecated DiscreteExclude…

0a66350

…Constraint in tests

Align make_prompt signature with the canonical recommender argument…

9d9adcf

… order Reorder `make_prompt` to (batch_size, searchspace, objective, measurements, pending_experiments), matching `RecommenderProtocol.recommend`, and make the LLM-specific experiment_description keyword-only.

Validate type of litellm_args keys

b1647ca

Enable handling of batch sizes >1

7b254f2

Move CHANGELOG entry to Unreleased section

ab29162

Improve wording and handling of forbidden configurations

3b35a86

Improve docstrings

41c7145

Change type of recommenders to tuple in SequentialMetaRecommender

4a455d9

This is more in-line with our usage patterns. We use our `to_tuple` converter as we cannot use a plain `tuple` due to python/mypy#8417.

Improve and fix parameter representation in prompts

6af4ea0

Fixes two small issues: First, the prompt now only propagates `active_values` instead of all values. Second, it fixes a leak in `np` representation for numerical values.

Delete LLMRecommender tests

089c36b

There was a request to fully rewrite those, hence the current version was deleted.

Bump litellm version to 1.97.0

bf3f3b8

Since `litellm` seems to be typed since this version this bump allows us to remove the package from previous exceptions regarding typing.

Rework and improve LLM-related exceptions and error handling

77b4b2e

All LLM-related exceptions now carry an `LLM` prefix in their class names. The hierarchy was extended by differentiating between general call errors (with authentication errors as one special subcase), and response errors.

Implement _recommend_hybrid instead of overriding recommend

1a3ec1a

The `recommend` call now stashes necessary context for creating the prompts and delegates to `_recommend_hybrid`. As a consequence, validation and preprocessing of the `super` call is used and has been removed.

Remove examples and docs on LLMs

81484b2

It has been decided that those should be added in a follow-up PR in order to decrease the scope of this one

Create convenience LLM-based meta recommenders

eb4768a

Provide full context in recovery prompt

60afb19

Previously, some of the context and information was not used in the recovery prompt. To fix this, there is now a single general prompt that includes the full context, with an optional block containing recovery-related information if necessary.

Improve handling of targets and objectives in prompt

0d90d46

We did not carry all of the relevant information for targets and objectives in the prompt. This has now been fixed by including views of both targets and objectives in the prompt, similar to the way we do it for parameters.

AVHopp force-pushed the feature/llm-recommender branch from f06c49a to 0d90d46 Compare October 8, 2026 14:22

AVHopp added 8 commits October 9, 2026 09:56

Simplify the JSON array extraction

233f862

Remove dead checks

a0e0ece

Move model query inside of try block

a2b78e6

Improve handling of point matching

88eaf2d

Force mode to be cyclic in LLMAlternatingRecommender

75ca275

Improve spacing in prompt

762fd06

Add tests for LLM-related code parts

a3980df

Fix typing issues in get_- and find-subclass methods

c60cbb8

It seems like newer versions of `mypy` require more detailed type information that the simple `type` annotation we had in `find_subclass`

	Do not recommend these again.
	Take these into account when generating your recommendations. In particular, only consider re-recommending parts of them if you have a strong reason to do so, for instance, when you are convinced that a second measurement is needed for noise reduction.

		)


		def make_llm_alternating_recommender(

Conversation

tobiasploetz commented Jul 7, 2026 • edited by AVHopp Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

tobiasploetz left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

AVHopp commented Jul 8, 2026

Uh oh!

AVHopp left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Scienfitz left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

AdrianSosic left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

AVHopp Oct 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

tobiasploetz commented Jul 7, 2026 •

edited by AVHopp

Loading

AVHopp Oct 8, 2026 •

edited

Loading

AVHopp Sep 7, 2026 •

edited

Loading