Skip to content

Add LLMRecommender for LLM-based experimental design#855

Draft
tobiasploetz wants to merge 9 commits into
mainfrom
feature/llm-recommender
Draft

Add LLMRecommender for LLM-based experimental design#855
tobiasploetz wants to merge 9 commits into
mainfrom
feature/llm-recommender

Conversation

@tobiasploetz

@tobiasploetz tobiasploetz commented Jul 7, 2026

Copy link
Copy Markdown
Contributor

Summary

  • Add LLMRecommender that uses LiteLLM to query language models for experimental design suggestions
  • Auto-extracts parameter descriptions from the search space using the metadata system (PR Add flexible metadata system #580), eliminating the need for a separate ParameterDescription class
  • Includes recovery mechanism for malformed LLM responses and optional feasibility filtering
  • Integrates into the recommender hierarchy via NonPredictiveRecommender

Ported from #561, adapted to the current codebase.

Key design decisions

  • LLMRecommender inherits from NonPredictiveRecommenderPureRecommenderRecommenderProtocol, consistent with the rest of the hierarchy
  • litellm and jinja2 are lazy-imported inside method bodies (not at module scope)
  • Parameter info (name, type, bounds/values, description, unit) is extracted directly from SearchSpace.parameters and their MeasurableMetadata
  • Recovery mechanism: if the initial LLM response is malformed, a correction prompt is sent; separate try/except blocks for API errors vs. parsing errors
  • Feasibility filtering: warns when fewer experiments pass than requested, returns available feasible experiments
  • Field validators enforce non-empty strings for model, experiment_description, objective_description
  • Unsupported parameter types raise IncompatibilityError immediately (no silent fallback)

New files

File Purpose
baybe/_optional/llm.py Guarded imports for litellm/jinja2
baybe/_optional/info.py LLM_INSTALLED flag
baybe/recommenders/pure/llm/ LLM recommender package
tests/test_llm_recommender.py 18 unit tests (all mocked, no real LLM calls)

Changes in latest review round (cae5ee0c2)

  • Base class contract: recommend() now rejects pending_experiments with IncompatibleArgumentError and warns when objective is passed (matching NonPredictiveRecommender behavior)
  • Input validation: overflow_experiments validated >= 0; litellm_args gets defensive copy (converter=dict) and reserved-key validation (model, messages rejected)
  • Parse robustness: Empty suggestion arrays raise LLMResponseError; continuous parameter values checked with math.isfinite(); discrete parameter values cast to canonical types (handles int vs np.float64 mismatch from JSON)
  • Optional imports: jinja2 and litellm handled in separate try-blocks so the error message names the actually missing package
  • Test guard: test_naive_hybrid_serialization.py uses _try_default_construct to skip recommenders that require constructor arguments

Open discussion items (see comments)

  1. Error handling strategy for completion() API calls (wrap vs. propagate litellm exceptions)
  2. Feasibility filtering behavior when zero experiments pass (warn vs. raise)
  3. Error handling for user-provided is_feasible_experiment callback
  4. Warning classes: bare UserWarning vs. project-specific warning class
  5. Missing parametrized test cases for parse error branches
  6. Missing hypothesis strategy and serialization roundtrip test (incl. is_feasible_experiment Callable serialization)
  7. Guard approach in test_naive_hybrid_serialization.py (try/except vs. explicit exclusion)

Test plan

  • pytest tests/test_llm_recommender.py — 18 tests pass
  • pytest tests/serialization/test_naive_hybrid_serialization.py — 6 tests pass
  • ruff check / ruff format — clean
  • mypy — no new errors in added files
  • import baybe does not trigger litellm/jinja2 import (lazy loading verified)
  • Pre-commit hooks pass
  • CI pipeline

tobiasploetz and others added 8 commits July 7, 2026 14:04
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Implements a recommender that uses LiteLLM to query language models for
experimental design suggestions. Leverages the metadata system to
auto-extract parameter descriptions from the search space.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
- Inherit from NonPredictiveRecommender instead of RecommenderProtocol
- Add field validators for required string fields
- Improve error handling in _attempt_recovery (split broad except)
- Raise IncompatibilityError for unsupported parameter types
- Add batch_size count validation with warning
- Improve error messages with invalid/allowed values detail
- Add Raises section to recommend() docstring
- Fix dict type annotations to dict[str, Any]
- Add LLMRecommender to pure/__init__.py for consistent hierarchy
- Refactor tests: parametrize error cases, use batch_size=3,
  add NumericalDiscreteParameter, test related_data/recovery_model/
  format_instructions, test feasibility edge cases

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
- Add pending_experiments rejection and objective warning in recommend()
  to match NonPredictiveRecommender contract
- Guard test_naive_hybrid_serialization against classes with required args
- Split optional imports to give accurate error for jinja2 vs litellm
- Add overflow_experiments ge(0) validator
- Add litellm_args defensive copy and reserved-key validation
- Add empty suggestions guard in _parse_llm_response
- Add math.isfinite check for continuous parameter values
- Add discrete parameter canonical type coercion

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

@tobiasploetz tobiasploetz left a comment

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Converting discussion comments into review threads for easier inline discussion.


prompt = self._construct_prompt(searchspace, batch_size, measurements)
response = completion(
model=self.model,

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Discussion: Error handling for completion() API calls

recommend() and _attempt_recovery()

The primary completion() call in recommend() has no error handling. Network failures, authentication errors, rate limits, timeouts, and API errors all propagate as raw litellm exceptions with potentially confusing stack traces.

Example: An expired API key produces a traceback through litellm internals instead of a clear BayBE error.

Options to consider:

  1. Wrap with LLMResponseError — catch broadly around completion() and re-raise as LLMResponseError with an actionable message (e.g., "check your API credentials, network connection, and model identifier").

  2. Let litellm errors propagate — since litellm already raises well-typed exceptions (AuthenticationError, RateLimitError, Timeout, etc.), BayBE could document that these are expected and let users handle them directly.

  3. Partial wrapping — only guard against structural response issues (empty choices, None content) but let API-level errors from litellm propagate as-is.

Option 2 has the advantage of not hiding the original error type, which is useful for retry logic. Option 1 follows BayBE's "explicit errors" principle more closely. What's the preferred approach?


if self.is_feasible_experiment is not None:
feasible_mask = output.apply(self.is_feasible_experiment, axis=1)
feasible_experiments = output[feasible_mask]

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Discussion: Feasibility filtering behavior when too few experiments pass

When is_feasible_experiment is provided and fewer than batch_size experiments pass the filter, the current behavior (after the fix in this round) is:

  • Emit a warning when len(feasible) < batch_size
  • Return whatever feasible experiments exist (possibly empty)

The alternative would be to raise (e.g., InfeasibilityError) when zero feasible experiments are found, following the "no partial results" principle strictly.

Trade-off:

  • Raising is consistent with BayBE's philosophy and prevents downstream code from silently operating on fewer experiments than expected.
  • Warning + partial return is more practical for LLM-based recommendations where the user may want to use whatever good suggestions were generated, and can always request more.

The current implementation warns. Should it instead raise when zero feasible experiments are returned? Or always raise when len(feasible) < batch_size?


if self.is_feasible_experiment is not None:
feasible_mask = output.apply(self.is_feasible_experiment, axis=1)
feasible_experiments = output[feasible_mask]

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Discussion: Error handling for user-provided is_feasible_experiment callback

The user-provided is_feasible_experiment callable is invoked via DataFrame.apply() with no error handling. If this callback raises (e.g., KeyError for a missing column, TypeError for unexpected types), the error surfaces through pandas apply() internals with a confusing stack trace.

Options:

  1. Wrap with try/except — catch exceptions from the callback and re-raise as LLMResponseError with a message pointing the user to their callback function.

  2. Leave unwrapped — the traceback still shows the user's function in the call stack, and wrapping could hide useful debugging information. This is the standard Python approach for user-provided callables.

  3. Validate callback signature at construction time — use inspect.signature to verify the callback accepts a pd.Series and returns bool, catching obvious mistakes early.

Option 2 is simpler and follows the principle of least surprise for Python developers. Option 1 gives better error messages for less experienced users. Thoughts?

)
if objective is not None:
warnings.warn(
f"'{self.recommend.__name__}' was called with an explicit objective "

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Discussion: Warning classes for LLM-specific warnings

The two warnings.warn() calls (LLM returned fewer suggestions, feasibility check shortfall) currently use the implicit UserWarning default. Per project conventions: "Use specific warning classes, not bare Warning."

Options:

  1. Use existing UnusedObjectWarning — semantically not a great fit, since the issue isn't an unused object but an insufficient result count.

  2. Define a new warning class — e.g., InsufficientSuggestionsWarning(UserWarning) in baybe/exceptions.py. This gives users fine-grained warnings.filterwarnings() control over LLM-specific warnings.

  3. Reuse a generic existing class — The closest existing option might be SearchSpaceMatchWarning, but it's semantically wrong for LLM output issues.

Option 2 seems cleanest. What should the warning class be named? Some options:

  • InsufficientSuggestionsWarning
  • LLMResponseWarning
  • Something more general that could cover future non-error LLM issues?

],
)
def test_parse_llm_response_errors(
response_content, error_match, recommender, searchspace

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Discussion: Missing parametrized test cases in test_parse_llm_response_errors

The current parametrization covers 5 error cases but misses several error branches in _parse_llm_response. These guard against real LLM failure modes and should be covered:

Missing case Error message Example input
Response is a dict, not a list "Response must be a JSON array" json.dumps({"key": "value"})
Suggestion is not a dict "Each suggestion must be a JSON object" json.dumps(["a string"])
parameters field is not a dict "Parameters must be a JSON object" json.dumps([{"explanation": "x", "parameters": [1, 2]}])
Non-numeric continuous value "Non-finite or non-numeric values" json.dumps([{"explanation": "x", "parameters": {"temperature": "hot", ...}}])
Missing explanation field "must contain an 'explanation' field" json.dumps([{"parameters": {...}}])
Empty suggestions array "empty array with no suggestions" json.dumps([])

These are straightforward pytest.param additions to the existing parametrized test. Should they be added in this PR or as a follow-up?

@@ -0,0 +1,449 @@
"""Tests for the LLM-based recommender."""

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Discussion: Missing hypothesis strategy and serialization roundtrip test

Convention: tests/AGENTS.md requires a hypothesis strategy and a @given roundtrip serialization test for every new user-facing class. LLMRecommender is exported in __all__ and is missing both.

What's needed

  1. Hypothesis strategy in tests/hypothesis_strategies/ — a @st.composite function generating random valid LLMRecommender instances with varying model, experiment_description, objective_description, overflow_experiments, and litellm_args.

  2. Serialization roundtrip test in tests/serialization/ using assert_roundtrip_consistency.

The hard part: is_feasible_experiment is not serializable

The is_feasible_experiment: Callable[[pd.Series], bool] | None field cannot be serialized to JSON with cattrs. No custom hook is registered. If a user puts an LLMRecommender with this field set into a Campaign and calls to_json(), it will fail with an unhelpful error.

Options:

  1. Omit on serialization — register a cattrs hook that drops is_feasible_experiment during unstructuring and restores it to None on structuring. Document that this field does not survive serialization. This is pragmatic and lets the rest of the recommender serialize cleanly.

  2. Raise explicitly — if is_feasible_experiment is not None, raise a clear error on serialization attempt (e.g., NotImplementedError("LLMRecommender with is_feasible_experiment cannot be serialized")). More strict but avoids silent data loss.

  3. Warn and omit — serialize everything except the callable, emit a warning that is_feasible_experiment was dropped. Middle ground.

The related_data: pd.DataFrame | None field likely works via the existing pickle+base64 hook, but this also needs verification in a roundtrip test.

Recommendation

Option 1 (omit on serialization) seems most consistent with how BayBE handles other non-serializable state. The roundtrip test should then verify:

  • LLMRecommender without is_feasible_experiment round-trips correctly
  • LLMRecommender with related_data round-trips correctly
  • The chosen behavior for is_feasible_experiment is tested

Should this be addressed in this PR or as a follow-up?

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This will hopefully resolve itself once the is_feasible_experiment field is gone. If there is still an issue here, we should re-evaluate then.

from tests.serialization.utils import assert_roundtrip_consistency


def _try_default_construct(cls):

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Discussion: Guard for recommenders with required args

The existing code at module level assumes all NonPredictiveRecommender subclasses can be default-constructed with cls(). LLMRecommender has three required fields (model, experiment_description, objective_description), so cls() raises TypeError at test collection time, breaking the entire serialization test module.

Current fix

Added a _try_default_construct(cls) helper that wraps cls() in a try/except TypeError. This silently excludes LLMRecommender from the naive hybrid serialization parametrization.

Alternatives

  1. Explicit exclusion list — maintain a set of excluded classes. More explicit but requires maintenance.

  2. Provide a test fixture that constructs LLMRecommender with required args — add it to the parametrized list manually. Would actually test the naive hybrid roundtrip with an LLM recommender, but may not make semantic sense.

  3. Keep the current approach — generic, forward-compatible for any future recommender with required args, and doesn't need maintenance.

Option 3 (current) seems most pragmatic. The dedicated LLMRecommender serialization test (from the separate discussion) will cover its roundtrip independently. Thoughts?

@AVHopp

AVHopp commented Jul 8, 2026

Copy link
Copy Markdown
Collaborator

@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 AVHopp added the enhancement Expand / change existing functionality label Jul 8, 2026

@AVHopp AVHopp left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 :)

Comment thread baybe/_optional/info.py


@define(slots=False)
class LLMRecommender(NonPredictiveRecommender):

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

One core aspect we need to discuss: Is NonPredictiveRecommender the correct base class, or should this rather inherit from the PureRecommender?

In our perspective, a NonPredictiveRecommender does not use any data that is provided to it for training - see the following check that is done within the recommend of that class:

if (measurements is not None) and not measurements.empty:
            warnings.warn(
                f"'{self.recommend.__name__}' was called with a non-empty "
                f"set of measurements but '{self.__class__.__name__}' does not "
                f"utilize any training data, meaning that the argument is ignored.",
                UnusedObjectWarning,
            )

However, the recommend that you re-implement here via the override actually makes use of the measurement, so you technically violate this. This is also the reason why you have to fully override the recommend function instead of e.g. only writing a new _recommend_hybrid which is the function for "this is how to do stuff in a hybrid space". So I think this recommender should simply inherit from PureRecommender and then override recommend

similar experiments.
"""

is_feasible_experiment: Callable[[pd.Series], bool] | None = field(default=None)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As discussed via chat, this should be removed by now. This should hopefully also resolve the issue with the serializability.

filter for feasibility. Only feasible experiments are returned.
"""

overflow_experiments: int = field(default=0, validator=(instance_of(int), ge(0)))

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Usage of this will change and it will be checked if these experiments will be part of the search space. We could maybe also investigate if we simply always want to have a certain multiple of batch_size be requested. Also, this should have a reasonable default, like maybe min(2, 0.1*batch_size) or similar?

"""
from baybe._optional.llm import completion

if pending_experiments is not None:

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If you inherit from PureRecommender, then you can actually use pending_experiments and have those also being acknowledged in the prompt.

from tests.serialization.utils import assert_roundtrip_consistency


def _try_default_construct(cls):

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a way to complex solution. In my opinion, it is simply not possible to have reasonable defaults for the LLMRecommender. I would actually prefer to just have a recommender here that assign some manual values to those fields and just add that manually to the list of all recommenders, given that this is only for testing serialization anyway.

Comment thread pyproject.toml

llm = [
"jinja2>=3.1.0",
"litellm>=1.70.0,!=1.82.7,!=1.82.8",

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why not directly increase this to >=1.83? Any specific reason?

@@ -0,0 +1,449 @@
"""Tests for the LLM-based recommender."""

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This will hopefully resolve itself once the is_feasible_experiment field is gone. If there is still an issue here, we should re-evaluate then.

Comment thread baybe/_optional/llm.py

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Different general comment: Before we un-draft this at some point, make sure that the PR descirption will have a human-readable and understandable summary of what this PR contains and provides. It seems like it is currently used as an implementation guideline, which is fine for now but needs to be changed later.

compatibility: ClassVar[SearchSpaceType] = SearchSpaceType.HYBRID
# See base class.

model: str = field(validator=(instance_of(str), min_len(1)))

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not an expert in how to use litellm, but I was surprised that users do not seem to provide an API Key or something. Can you explain to me here what users will need to provide and how this works "in the back"? We should probably explain this in a user guide or so later, but for now it is sufficient if I understand it :D

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement Expand / change existing functionality

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants