feat(llm): enhance llm_processor to support fallback models and update LitellmLLMProvider to utilize fallback_model_list #1739
Conversation
…e LitellmLLMProvider to utilize fallback_model_list
WalkthroughThe changes introduce a structured fallback mechanism for LLM model selection. The Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant llm_processor
participant LitellmLLMProvider
participant Router
User->>llm_processor: Provide model configuration(s)
llm_processor->>llm_processor: Build fallback_model_list for each model
llm_processor->>LitellmLLMProvider: Pass model config with fallback_model_list
User->>LitellmLLMProvider: Request completion
LitellmLLMProvider->>Router: Call acompletion with fallback_model_list
Router->>Router: Attempt completion with primary and fallback models
Router-->>LitellmLLMProvider: Return completion result
LitellmLLMProvider-->>User: Return completion or stream
Possibly related PRs
Suggested labels
Poem
✨ Finishing Touches
🧪 Generate Unit Tests
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (3)
wren-ai-service/src/providers/llm/litellm.py (1)
86-95: Remove commented code.The old implementation is no longer needed.
Apply this diff to remove the commented code:
- # completion: Union[ModelResponse] = await acompletion( - # model=self._model, - # api_key=self._api_key, - # api_base=self._api_base, - # api_version=self._api_version, - # timeout=self._timeout, - # messages=openai_formatted_messages, - # stream=streaming_callback is not None, - # **generation_kwargs, - # )wren-ai-service/src/providers/__init__.py (2)
82-96: Consider handling api_key configuration for fallback models.The fallback model parameters don't include
api_key_namewhich might be needed for authentication. Each fallback model might require its own API key configuration.Would you like me to help implement proper api_key handling for fallback models?
103-103: Remove commented code.Apply this diff to remove the commented line:
- # "fallbacks": model.get("fallbacks", []),
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
wren-ai-service/src/providers/__init__.py(1 hunks)wren-ai-service/src/providers/llm/litellm.py(4 hunks)
🧰 Additional context used
🪛 Ruff (0.11.9)
wren-ai-service/src/providers/llm/litellm.py
8-8: litellm.acompletion imported but unused
Remove unused import: litellm.acompletion
(F401)
🔇 Additional comments (1)
wren-ai-service/src/providers/llm/litellm.py (1)
96-101:❓ Verification inconclusive
Verify Router handles all required parameters.
The Router call omits
api_key,api_base,api_version, andtimeoutparameters that were present in the original implementation. Please verify that the Router properly handles these through thefallback_model_listconfiguration.Run the following script to check how Router handles these parameters:
🏁 Script executed:
#!/bin/bash # Description: Check Router class implementation for parameter handling # Search for Router class definition and its methods ast-grep --pattern 'class Router { $$$ }' # Search for Router's acompletion method to understand parameter handling rg -A 10 "def acompletion.*Router" # Search for how Router uses model_list configuration rg -A 5 "model_list.*Router|Router.*model_list"Length of output: 132
I wasn’t able to locate the
Routerimplementation or see how it handlesapi_key,api_base,api_version, andtimeout. Let’s deepen the search:
🏁 Script executed:
#!/bin/bash # 1. Find any Router class definitions rg -n "class .*Router" . # 2. Find where Router is imported rg -n "import .*Router" . # 3. Locate where self._router is initialized in litellm.py rg -n "self\. _router" -n wren-ai-service/src/providers/llm/litellm.py # 4. Search for fallback_model_list usage across the repo rg -n "fallback_model_list" . # 5. Search for any acompletion definitions rg -n "def acompletion" -n .Length of output: 4729
We still can’t find
Routerin our codebase—it’s coming from the externallitellmpackage. Let’s inspect its implementation to confirm how it handlesapi_key,api_base,api_version, andtimeout.
🏁 Script executed:
#!/bin/bash # Inspect Router in the installed litellm package python3 - << 'EOF' import inspect, litellm print("Router source:") print(inspect.getsource(litellm.Router)) print("\nacompletion signature:") print(inspect.signature(litellm.Router.acompletion)) EOFLength of output: 131
Verify external Router parameter support
We weren’t able to locate or inspect the
litellm.Routerimplementation in our sandbox, so it’s unclear whether it automatically picks up or falls back on the following parameters:
api_keyapi_baseapi_versiontimeoutPlease confirm that the external
Routerfrom thelitellmpackage handles these internally (e.g., via its default config or the providedfallback_model_list) to avoid unintentional omissions.
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (2)
wren-ai-service/src/providers/llm/litellm.py (2)
8-10: Remove leftover unusedModelResponseimport
ModelResponseis no longer referenced after switching toself._router.acompletion, so the import is dead code.-from litellm.types.utils import ModelResponse🧰 Tools
🪛 Ruff (0.11.9)
8-8: SyntaxError: Expected an identifier, but found a keyword 'from' that cannot be used here
8-8: SyntaxError: Simple statements must be separated by newlines or semicolons
8-8: SyntaxError: Simple statements must be separated by newlines or semicolons
9-9: SyntaxError: Expected an identifier, but found a keyword 'from' that cannot be used here
9-9: SyntaxError: Simple statements must be separated by newlines or semicolons
9-9: SyntaxError: Simple statements must be separated by newlines or semicolons
🪛 Pylint (3.3.7)
[error] 8-8: Parsing failed: 'invalid syntax (src.providers.llm.litellm, line 8)'
(E0001)
86-96: Delete the commented-outacompletionblockKeeping large commented sections clutters the file and can mislead readers. The git history already preserves this code.
- # completion: Union[ModelResponse] = await acompletion( - # model=self._model, - # api_key=self._api_key, - # api_base=self._api_base, - # api_version=self._api_version, - # timeout=self._timeout, - # messages=openai_formatted_messages, - # stream=streaming_callback is not None, - # **generation_kwargs, - # )
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
wren-ai-service/src/providers/llm/litellm.py(4 hunks)
🧰 Additional context used
🪛 Ruff (0.11.9)
wren-ai-service/src/providers/llm/litellm.py
8-8: SyntaxError: Expected an identifier, but found a keyword 'from' that cannot be used here
8-8: SyntaxError: Simple statements must be separated by newlines or semicolons
8-8: SyntaxError: Simple statements must be separated by newlines or semicolons
9-9: SyntaxError: Expected an identifier, but found a keyword 'from' that cannot be used here
9-9: SyntaxError: Simple statements must be separated by newlines or semicolons
9-9: SyntaxError: Simple statements must be separated by newlines or semicolons
🪛 Pylint (3.3.7)
wren-ai-service/src/providers/llm/litellm.py
[error] 8-8: Parsing failed: 'invalid syntax (src.providers.llm.litellm, line 8)'
(E0001)
…imports in LitellmLLMProvider
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
…itialized in LitellmLLMProvider
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
wren-ai-service/src/utils.py (2)
166-168: Unused variablegenerator_name– rename to_or remove.
generator_nameis captured but never used after the metadata switch tometa[0]["model"], so linters will flag this as an unused-variable warning.- result, generator_name = await func(*args, **kwargs) + # `generator_name` is intentionally ignored – rename to `_` to silence linters + result, _ = await func(*args, **kwargs)
169-174: Guard against missing or malformedmetaitems.
meta[0]is assumed to be a dict with"model"/"usage"keys.
A defensive approach avoidsAttributeErrorifmeta[0]isNoneor not a dict:- langfuse_context.update_current_observation( - model=meta[0].get("model"), - usage_details=meta[0].get("usage", {}), - ) + first = meta[0] or {} + langfuse_context.update_current_observation( + model=first.get("model"), + usage_details=first.get("usage", {}), + )
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
wren-ai-service/src/providers/llm/litellm.py(3 hunks)wren-ai-service/src/utils.py(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- wren-ai-service/src/providers/llm/litellm.py
🔇 Additional comments (1)
wren-ai-service/src/utils.py (1)
167-178: Potential API mismatch – wrapper now returns a single value.The decorated function still returns
(result, generator_name), but the wrapper only relaysresult.
If any call-sites relied on the two-tuple signature outside the decorator, this is a silent breaking change.Consider either:
- Forwarding the original tuple:
- return result + return result, _
- Or updating all call-sites and the original function to return just
result.Please verify usage across the codebase before merging.
…dynamically refactor trace_cost to ignore generator name in result unpacking
There was a problem hiding this comment.
Actionable comments posted: 1
🔭 Outside diff range comments (3)
wren-ai-service/src/pipelines/generation/sql_generation.py (1)
108-115:⚠️ Potential issueGuard against missing
"prompt"key
PromptBuilder.run()usually returns a dict with a"prompt"key, but if someone mutates the template or introduces a custom builder the key could be absent and we would passNoneto the LLM – Litellm raises an opaque validation error in that case.Add an early assertion for clarity:
@@ - return await generator(prompt=prompt.get("prompt")) + prompt_str = prompt.get("prompt") + if prompt_str is None: + raise ValueError("PromptBuilder did not return a 'prompt' key") + return await generator(prompt=prompt_str)wren-ai-service/src/pipelines/generation/sql_tables_extraction.py (1)
74-76:⚠️ Potential issue
post_processwill crash on malformed LLM output
orjson.loads(...)[ "tables" ]assumes:
extract_sql_tables["replies"][0]exists- The string is valid JSON
- A
"tables"key is presentAny deviation raises, returning a 500 to callers. Wrap in a try/except and fall back to
[], mirroring the pattern used inintent_classification.py.- return orjson.loads(extract_sql_tables.get("replies")[0])["tables"] + try: + return orjson.loads(extract_sql_tables.get("replies")[0])["tables"] + except Exception: + logger.warning("sql_tables_extraction – unable to parse tables from LLM reply") + return []wren-ai-service/src/pipelines/generation/sql_question.py (1)
70-72: 🛠️ Refactor suggestionGuard against malformed LLM JSON
orjson.loads(...)[ "question" ]will KeyError on bad output. Consider a safe fallback (e.g., original SQL) to avoid hard crashes.
🧹 Nitpick comments (11)
wren-ai-service/src/pipelines/generation/question_recommendation.py (1)
46-47: Consider tightening the type hints forgenerator
Anytells the type checker nothing. All thegeneratehelpers in this refactor expect an async callable that acceptsprompt=<str>and returns adict. Encoding that contract makes IDE auto-completion and static analysis happier:-from typing import Any +from typing import Awaitable, Callable, Dict -async def generate(prompt: dict, generator: Any) -> dict: +GeneratorT = Callable[..., Awaitable[Dict]] + +async def generate(prompt: dict, generator: GeneratorT) -> dict:wren-ai-service/src/pipelines/generation/sql_answer.py (1)
68-74:query_idshould be optional in both signature and typing
run()forwardsquery_idwhich may beNone, yetgenerate_answerdeclares the parameter as non-nullablestr. Static checkers will complain and users of the provider might not expectNone.-from typing import Any, Optional +from typing import Any, Optional -async def generate_answer( - prompt: dict, generator: Any, query_id: str -) -> dict: +async def generate_answer( + prompt: dict, + generator: Any, + query_id: Optional[str] = None, +) -> dict:No runtime behaviour changes, but the annotation now matches reality.
wren-ai-service/src/pipelines/generation/data_assistance.py (1)
74-79: Alignquery_idtyping with value passed in
run()usesquery_id or "", sodata_assistance()never receivesNone.
If an empty string really is the sentinel value, make that explicit with a default to avoid accidentalNonepropagation and clarify intent:-async def data_assistance( - prompt: dict, generator: Any, query_id: str +async def data_assistance( + prompt: dict, + generator: Any, + query_id: str = "", ) -> dict:wren-ai-service/src/pipelines/generation/relationship_recommendation.py (1)
59-60: Consistency: strengthen the generator’s type and fail fast on missing promptSame remark as for other pipelines – specify an async callable type and guard against a missing
"prompt"key to surface configuration errors early:-if prompt.get("prompt") is None: - raise KeyError('"prompt" key missing from prompt dict') -return await generator(prompt=prompt.get("prompt")) +prompt_text = prompt["prompt"] # KeyError if absent – explicit is better +return await generator(prompt=prompt_text)wren-ai-service/src/pipelines/generation/sql_regeneration.py (1)
127-128: Minor: add early validation of prompt dictionaryAll pipelines assume the prompt builder always returns a dict with
"prompt"present. A defensive check prevents an opaque downstream failure:- return await generator(prompt=prompt.get("prompt")) + try: + prompt_text = prompt["prompt"] + except KeyError as exc: + raise ValueError("Prompt builder did not return the expected key 'prompt'") from exc + return await generator(prompt=prompt_text)wren-ai-service/src/pipelines/generation/semantics_description.py (1)
69-72: Confirm that downstream nodes expect only the result dictThe return type changed from a tuple to a plain dict.
All downstream nodes (normalize, etc.) now access onlygenerate["replies"], so they are safe.
However, any external consumer that imported and invokedgeneratedirectly will silently break.Consider adding a short doc-string or a
TypedDict/pydanticmodel here to make the new contract explicit and let mypy/pyright catch misuse earlier.wren-ai-service/src/pipelines/generation/chart_adjustment.py (1)
111-118: Add minimal error handling for generator failuresAll new
generate_*helpers bubble any LLM exception straight up to Hamilton, which logs a cryptic stack-trace.
Consider catchingExceptionhere, logging the prompt & adjustment options, then re-raising, so ops can triage bad prompts faster.wren-ai-service/src/pipelines/generation/sql_tables_extraction.py (1)
68-70: Add defensive error handling around generator invocation
await generator(...)can raise provider/network exceptions that will propagate up the pipeline and bubble out of theAsyncDriver, causing the whole request to fail without context. Consider catching and re-raising with additional context or converting to a domain-specific error to improve observability.- return await generator(prompt=prompt.get("prompt")) + try: + return await generator(prompt=prompt.get("prompt")) + except Exception as exc: + logger.exception("extract_sql_tables – generator failed") + raisewren-ai-service/src/pipelines/generation/sql_question.py (1)
62-66: Mirror error-handling approach used elsewhereSame observation as in
sql_tables_extraction: wrap the generator call so failures surface with context.- return await generator(prompt=prompt.get("prompt")) + try: + return await generator(prompt=prompt.get("prompt")) + except Exception as exc: + logger.exception("generate_sql_question – generator failed") + raisewren-ai-service/src/pipelines/generation/chart_generation.py (1)
84-86: Same resiliency/observability remarks as other generatorsWrap the call in try/except for clearer logs and to keep pipeline failures consistent.
wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py (1)
321-329: Drop the unnecessaryelsebranch after an earlyreturnThe
elseis redundant and slightly indents the happy-path unnecessarily.async def filter_columns_in_tables( prompt: dict, table_columns_selection_generator: Any ) -> dict: if prompt: return await table_columns_selection_generator( prompt=prompt.get("prompt") ) - else: - return {} + return {}This keeps the function flatter and satisfies pylint R1705.
🧰 Tools
🪛 Pylint (3.3.7)
[refactor] 323-328: Unnecessary "else" after "return", remove the "else" and de-indent the code inside it
(R1705)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (21)
wren-ai-service/src/pipelines/generation/chart_adjustment.py(1 hunks)wren-ai-service/src/pipelines/generation/chart_generation.py(1 hunks)wren-ai-service/src/pipelines/generation/data_assistance.py(1 hunks)wren-ai-service/src/pipelines/generation/followup_sql_generation.py(1 hunks)wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py(1 hunks)wren-ai-service/src/pipelines/generation/intent_classification.py(1 hunks)wren-ai-service/src/pipelines/generation/misleading_assistance.py(1 hunks)wren-ai-service/src/pipelines/generation/question_recommendation.py(1 hunks)wren-ai-service/src/pipelines/generation/relationship_recommendation.py(1 hunks)wren-ai-service/src/pipelines/generation/semantics_description.py(1 hunks)wren-ai-service/src/pipelines/generation/sql_answer.py(1 hunks)wren-ai-service/src/pipelines/generation/sql_correction.py(1 hunks)wren-ai-service/src/pipelines/generation/sql_generation.py(1 hunks)wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py(1 hunks)wren-ai-service/src/pipelines/generation/sql_question.py(1 hunks)wren-ai-service/src/pipelines/generation/sql_regeneration.py(1 hunks)wren-ai-service/src/pipelines/generation/sql_tables_extraction.py(1 hunks)wren-ai-service/src/pipelines/generation/user_guide_assistance.py(1 hunks)wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py(1 hunks)wren-ai-service/src/providers/llm/litellm.py(3 hunks)wren-ai-service/src/utils.py(1 hunks)
✅ Files skipped from review due to trivial changes (1)
- wren-ai-service/src/pipelines/generation/followup_sql_generation.py
🚧 Files skipped from review as they are similar to previous changes (2)
- wren-ai-service/src/utils.py
- wren-ai-service/src/providers/llm/litellm.py
🧰 Additional context used
🧬 Code Graph Analysis (3)
wren-ai-service/src/pipelines/generation/sql_question.py (1)
wren-ai-service/src/pipelines/generation/sql_tables_extraction.py (1)
prompt(59-63)
wren-ai-service/src/pipelines/generation/question_recommendation.py (3)
wren-ai-service/src/pipelines/generation/relationship_recommendation.py (2)
generate(59-60)prompt(49-54)wren-ai-service/src/pipelines/generation/semantics_description.py (2)
generate(70-71)prompt(55-65)wren-ai-service/src/web/v1/routers/question_recommendation.py (1)
get(125-147)
wren-ai-service/src/pipelines/generation/relationship_recommendation.py (2)
wren-ai-service/src/pipelines/generation/question_recommendation.py (2)
generate(46-47)prompt(21-41)wren-ai-service/src/pipelines/generation/semantics_description.py (2)
generate(70-71)prompt(55-65)
🪛 Pylint (3.3.7)
wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py
[refactor] 323-328: Unnecessary "else" after "return", remove the "else" and de-indent the code inside it
(R1705)
🔇 Additional comments (5)
wren-ai-service/src/pipelines/generation/user_guide_assistance.py (1)
61-68: Double-check provider acceptsquery_id
llm_provider.get_generator()previously returned callables that accepted onlypromptand**generation_kwargs.
You now pass an extraquery_id. If the provider’s callable signature is stilldef __call__(prompt: str, **_), Python will swallow the kwarg but litellm’s router may not.Please verify the provider’s
__call__signature or wrap the call with:return await generator(prompt=prompt["prompt"], **({"query_id": query_id} if query_id else {}))wren-ai-service/src/pipelines/generation/intent_classification.py (1)
284-286:✅ Verification successful
Signature change LGTM – verify no stray “generator_name” usages remain
The function now matches the updated
trace_costcontract. Run a quick grep to ensure no remnants ofgenerator_nameare still passed toclassify_intent.
🏁 Script executed:
#!/bin/bash # Search for still-existing calls with a generator_name argument rg --line-number 'classify_intent\([^)]*generator_name' -g '*.py'Length of output: 67
🏁 Script executed:
#!/bin/bash # Search for any remaining occurrences of "generator_name" across Python files rg --line-number 'generator_name' -g '*.py'Length of output: 43
No stray “generator_name” references remain
A full search across all Python files returned zero occurrences ofgenerator_name. The updatedclassify_intentsignature aligns with thetrace_costcontract.wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py (1)
78-82:✅ Verification successful
Parameter order deviates from other pipelines – harmless but double-check decorator assumptions
Most generation functions take
(prompt, generator)only. Here we have(prompt, generator, query_id).
Iftrace_costintrospects positional arguments to locate the generator, confirm it still picks the second param and ignores extras; otherwise cost tracing could break.
🏁 Script executed:
#!/bin/bash # Inspect trace_cost implementation for positional assumptions rg --line-number 'def trace_cost' -A3 -B0Length of output: 309
🏁 Script executed:
#!/bin/bash rg --heading -n "def trace_cost" -A30 wren-ai-service/src/utils.pyLength of output: 1228
No action needed:
trace_costsafely handles extra positional arguments
Thetrace_costdecorator uses(*args, **kwargs)and merely forwards all parameters to the wrapped function without inspecting their positions. Addingquery_idaftergeneratorwon’t affect cost tracing.wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py (1)
88-96: Verify the updated generator signature at call-site
generate_sql_reasoning()now forwards onlypromptandquery_id.
Confirm that every generator produced byllm_provider.get_generator()still accepts those exact keyword arguments—nothing more, nothing less.
A silent mismatch will surface only at runtime as aTypeError.If any existing custom generator still expects
generator_name, wrap it or update its signature accordingly.wren-ai-service/src/pipelines/generation/misleading_assistance.py (1)
74-79: Double-check generator kw-arg compatibility aftergenerator_nameremovalAs in the other pipelines, the call now passes just
promptandquery_id.
Make sure the underlying generator factory was updated everywhere; otherwise this coroutine will raise at runtime.No functional issue spotted here, only a consistency check.
| @trace_cost | ||
| async def generate_sql_correction( | ||
| prompt: dict, generator: Any, generator_name: str | ||
| prompt: dict, generator: Any | ||
| ) -> dict: | ||
| return await generator(prompt=prompt.get("prompt")), generator_name | ||
| return await generator(prompt=prompt.get("prompt")) | ||
|
|
There was a problem hiding this comment.
🛠️ Refactor suggestion
Same nil-prompt edge case applies here
Replicate the defensive check suggested for sql_generation.py to avoid runtime errors if "prompt" is missing.
🤖 Prompt for AI Agents
In wren-ai-service/src/pipelines/generation/sql_correction.py around lines 70 to
75, add a defensive check to ensure the "prompt" key exists in the prompt
dictionary before calling generator. If "prompt" is missing or None, handle it
gracefully by returning an appropriate default or error response to prevent
runtime errors, similar to the fix applied in sql_generation.py.
This PR enhances
llm_processorto support a fallback mechanism, allowing it to automatically switch to alternative models in case the primary model fails.Summary by CodeRabbit
New Features
Improvements