Skip to content

Separate LLM-related utilies - #951

Open
AVHopp wants to merge 15 commits into
mainfrom
refactor/llm-utilities
Open

AVHopp wants to merge 15 commits into
mainfrom
refactor/llm-utilities

Conversation

@AVHopp

@AVHopp AVHopp commented Nov 3, 2026 •

Copy link
Copy Markdown
Collaborator

The goal of this PR is to separate LLM-related functionality that was introduced in #855 in such a way that it is usable outside of the context of the LLMRecommender.

For this, LLM-related functionalities which might be useful for other LLM-related parts or workflows (like the one described in #931) have been moved into a dedicated _llm subpackage. The original plan for this was the following structure:
1. query.py: sends one prompt to a LiteLLM model and returns the response text; also holds the LiteLLM argument keys that BayBE sets itself (model, messages, stream), so users cannot override them; expected functions are query_model and validate_litellm_args
2. recovery.py: parses the response. If the response is unusable, it retries exactly once with a corrective prompt; failed calls are never retried; expected functionality is query_with_recovery
3. extraction.py: finds the intended JSON array or object in free text (Markdown fences, prose, several blocks) and turns it into typed data via cattrs; expected functions are extract_json and structure_json
4. rendering.py: writes parameters, constraints and objectives into prompts (names, bounds, units, target direction) and renders Jinja templates with shared settings; expected functions are to_prompt_dict, to_json_lines, to_table_text and render_template

What stays with each feature:
1. Prompt and schema: Feature-specific prompt and corresponding schema
2. Response checks: Validation if the response is correct for the context of the feature
3. Model settings: Any feature-specific model settings
- baybe/_optional/llm.py, the llm extra and LLM_INSTALLED are already shared and stay unchanged.

TRANSPARENCY: The plan above was created with the help of AI and has not 100% been checked against the current implementation and might deviate in details. However, as the implementation basically followed this plan, I would not expect this. A clean PR description will be written at a later time.

Stack created with GitHub Stacks CLI • Give Feedback 💬

@AVHopp
AVHopp added this pull request to stack #952 November 3, 2026 16:03
@AVHopp AVHopp self-assigned this Nov 3, 2026
@AVHopp AVHopp added the repo Requires changes to the project configuration label Nov 3, 2026
@AVHopp AVHopp added this to the 0.16.0 milestone Nov 11, 2026
@AVHopp
AVHopp force-pushed the refactor/llm-utilities branch from 1b45007 to 2098af9 Compare November 11, 2026 16:48
@AVHopp
AVHopp requested a balanced review from Copilot November 12, 2026 08:02

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

🟡 Changes recommended

Reserved query parameters are not validated for standalone callers, causing unsupported streaming requests or misleading provider errors, and JSON extraction can silently return data from string literals instead of the actual nested payload.

2 open issues
What changed in this PR

Extracts shared LLM functionality introduced in #855 into baybe._llm, making it reusable beyond LLMRecommender.

Changes:

  • Separates querying, recovery, rendering, JSON extraction, and exceptions into shared modules.
  • Updates the recommender to delegate to those utilities and generalizes JSON extraction to support nested payloads.
  • Reorganizes tests around the shared helpers and retains only LLMError as a public exception re-export.
File Description
tests/​llm/​test_rendering.py Tests shared serialization and template rendering.
tests/​llm/​test_recovery.py Tests retry behavior and error propagation.
tests/​llm/​test_recommender.py Updates imports and moves shared-helper tests out.
tests/​llm/​test_query.py Tests provider-error wrapping independently.
tests/​llm/​test_prompts.py Moves rendering-helper tests to their own module.
tests/​llm/​test_parsing.py Updates imports and separates extraction tests.
tests/​llm/​test_extraction.py Tests generalized JSON extraction and structuring.
tests/​llm/​test_exceptions.py Verifies internal exceptions and the public re-export.
tests/​hypothesis_strategies/​recommenders.py Uses the shared reserved-key constant.
baybe/​recommenders/​pure/​llm/​llm.py Delegates querying and recovery to shared utilities.
baybe/​recommenders/​pure/​llm/​_prompts.py Uses shared rendering helpers.
baybe/​recommenders/​pure/​llm/​_parsing.py Uses shared JSON extraction and structuring.
baybe/​exceptions.py Re-exports only the LLM catch-all exception.
baybe/​_llm/​rendering.py Houses prompt serialization and template rendering.
baybe/​_llm/​recovery.py Provides a reusable one-retry recovery loop.
baybe/​_llm/​query.py Houses LiteLLM querying and argument validation.
baybe/​_llm/​extraction.py Extracts accepted JSON containers and structures results.
baybe/​_llm/​exceptions.py Houses LLM exceptions with generalized recovery guidance.
baybe/​_llm/​__init__.py Introduces the shared internal LLM package.

Review effort: Balanced

Comment thread baybe/_llm/extraction.py Resolved
Comment thread baybe/_llm/query.py Resolved
`_extract_json_array` already decoded the payload with `raw_decode` but returned
the matched substring, so `parse_llm_response` decoded it again. Move the function
to `baybe._llm.extraction` as `extract_json_array`, return the decoded value, and
remove the second `json.loads` and its unreachable `JSONDecodeError` branch.
Move the cattrs structuring and its `LLMMalformedResponseError` mapping out of
`parse_llm_response` into `baybe._llm.extraction.structure_json`. The target type is
an argument, so other consumers can convert extracted JSON with the same recovery
semantics.
The table renderer is inlined at its call sites instead of being moved, since it
only wraps `df.to_string(index=False)`
`LLMError` and `LLMResponseWarning` are needed outside the LLM
layer. Hence, we define them in `baybe/exceptions.py` instead of
re-exporting them from `baybe._llm.exceptions`, leaving the
remaining LLM-specific subclasses internal.
Previously, rejected values were rescanned as text,
so brackets inside string literals could be mistaken for the payload.
@AVHopp
AVHopp force-pushed the refactor/llm-utilities branch from 2098af9 to 070dd00 Compare November 12, 2026 15:50
@AVHopp
AVHopp marked this pull request as ready for review November 12, 2026 15:51

This branch has not been deployed

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

Labels

repo Requires changes to the project configuration

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants