Add LLMRecommender for LLM-based experimental design#855
Conversation
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
left a comment
There was a problem hiding this comment.
Converting discussion comments into review threads for easier inline discussion.
|
|
||
| prompt = self._construct_prompt(searchspace, batch_size, measurements) | ||
| response = completion( | ||
| model=self.model, |
There was a problem hiding this comment.
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:
-
Wrap with
LLMResponseError— catch broadly aroundcompletion()and re-raise asLLMResponseErrorwith an actionable message (e.g., "check your API credentials, network connection, and model identifier"). -
Let litellm errors propagate — since
litellmalready raises well-typed exceptions (AuthenticationError,RateLimitError,Timeout, etc.), BayBE could document that these are expected and let users handle them directly. -
Partial wrapping — only guard against structural response issues (empty
choices,Nonecontent) 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] |
There was a problem hiding this comment.
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] |
There was a problem hiding this comment.
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:
-
Wrap with try/except — catch exceptions from the callback and re-raise as
LLMResponseErrorwith a message pointing the user to their callback function. -
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.
-
Validate callback signature at construction time — use
inspect.signatureto verify the callback accepts apd.Seriesand returnsbool, 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 " |
There was a problem hiding this comment.
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:
-
Use existing
UnusedObjectWarning— semantically not a great fit, since the issue isn't an unused object but an insufficient result count. -
Define a new warning class — e.g.,
InsufficientSuggestionsWarning(UserWarning)inbaybe/exceptions.py. This gives users fine-grainedwarnings.filterwarnings()control over LLM-specific warnings. -
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:
InsufficientSuggestionsWarningLLMResponseWarning- Something more general that could cover future non-error LLM issues?
| ], | ||
| ) | ||
| def test_parse_llm_response_errors( | ||
| response_content, error_match, recommender, searchspace |
There was a problem hiding this comment.
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.""" | |||
There was a problem hiding this comment.
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
-
Hypothesis strategy in
tests/hypothesis_strategies/— a@st.compositefunction generating random validLLMRecommenderinstances with varyingmodel,experiment_description,objective_description,overflow_experiments, andlitellm_args. -
Serialization roundtrip test in
tests/serialization/usingassert_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:
-
Omit on serialization — register a cattrs hook that drops
is_feasible_experimentduring unstructuring and restores it toNoneon structuring. Document that this field does not survive serialization. This is pragmatic and lets the rest of the recommender serialize cleanly. -
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. -
Warn and omit — serialize everything except the callable, emit a warning that
is_feasible_experimentwas 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:
LLMRecommenderwithoutis_feasible_experimentround-trips correctlyLLMRecommenderwithrelated_dataround-trips correctly- The chosen behavior for
is_feasible_experimentis tested
Should this be addressed in this PR or as a follow-up?
There was a problem hiding this comment.
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): |
There was a problem hiding this comment.
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
-
Explicit exclusion list — maintain a set of excluded classes. More explicit but requires maintenance.
-
Provide a test fixture that constructs
LLMRecommenderwith 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. -
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?
|
@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
left a comment
There was a problem hiding this comment.
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 :)
|
|
||
|
|
||
| @define(slots=False) | ||
| class LLMRecommender(NonPredictiveRecommender): |
There was a problem hiding this comment.
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) |
There was a problem hiding this comment.
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))) |
There was a problem hiding this comment.
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: |
There was a problem hiding this comment.
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): |
There was a problem hiding this comment.
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.
|
|
||
| llm = [ | ||
| "jinja2>=3.1.0", | ||
| "litellm>=1.70.0,!=1.82.7,!=1.82.8", |
There was a problem hiding this comment.
Why not directly increase this to >=1.83? Any specific reason?
| @@ -0,0 +1,449 @@ | |||
| """Tests for the LLM-based recommender.""" | |||
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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))) |
There was a problem hiding this comment.
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
Summary
LLMRecommenderthat uses LiteLLM to query language models for experimental design suggestionsParameterDescriptionclassNonPredictiveRecommenderPorted from #561, adapted to the current codebase.
Key design decisions
LLMRecommenderinherits fromNonPredictiveRecommender→PureRecommender→RecommenderProtocol, consistent with the rest of the hierarchylitellmandjinja2are lazy-imported inside method bodies (not at module scope)SearchSpace.parametersand theirMeasurableMetadatamodel,experiment_description,objective_descriptionIncompatibilityErrorimmediately (no silent fallback)New files
baybe/_optional/llm.pylitellm/jinja2baybe/_optional/info.pyLLM_INSTALLEDflagbaybe/recommenders/pure/llm/tests/test_llm_recommender.pyChanges in latest review round (
cae5ee0c2)recommend()now rejectspending_experimentswithIncompatibleArgumentErrorand warns whenobjectiveis passed (matchingNonPredictiveRecommenderbehavior)overflow_experimentsvalidated>= 0;litellm_argsgets defensive copy (converter=dict) and reserved-key validation (model,messagesrejected)LLMResponseError; continuous parameter values checked withmath.isfinite(); discrete parameter values cast to canonical types (handlesintvsnp.float64mismatch from JSON)jinja2andlitellmhandled in separate try-blocks so the error message names the actually missing packagetest_naive_hybrid_serialization.pyuses_try_default_constructto skip recommenders that require constructor argumentsOpen discussion items (see comments)
completion()API calls (wrap vs. propagate litellm exceptions)is_feasible_experimentcallbackUserWarningvs. project-specific warning classis_feasible_experimentCallable serialization)test_naive_hybrid_serialization.py(try/except vs. explicit exclusion)Test plan
pytest tests/test_llm_recommender.py— 18 tests passpytest tests/serialization/test_naive_hybrid_serialization.py— 6 tests passruff check/ruff format— cleanmypy— no new errors in added filesimport baybedoes not triggerlitellm/jinja2import (lazy loading verified)