fix(agents): wire FallbackModel so a primary 503 retries the fallback (#183)#184
Conversation
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
✨ Finishing Touches🧪 Generate unit tests (beta)
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. Comment |
Reviewer's GuideIntroduces a shared helper to construct the agent model as a PydanticAI FallbackModel (primary + fallback) and wires it into both agents, with API-key validation, graceful degradation when the fallback is unusable, and tests to cover the main behaviors. Sequence diagram for primary model failure falling back to agent_fallback_modelsequenceDiagram
actor User
participant ExperimentAgent
participant FallbackModel
participant PrimaryProvider
participant FallbackProvider
User->>ExperimentAgent: create_experiment_agent()
ExperimentAgent->>ExperimentAgent: build_agent_model_with_fallback()
ExperimentAgent->>FallbackModel: Agent(model=FallbackModel)
User->>ExperimentAgent: agent.run()
ExperimentAgent->>FallbackModel: call
FallbackModel->>PrimaryProvider: request
alt primary succeeds
PrimaryProvider-->>FallbackModel: response
FallbackModel-->>ExperimentAgent: result
else ModelAPIError (HTTP 5xx, rate limit)
PrimaryProvider-->>FallbackModel: ModelAPIError
FallbackModel->>FallbackProvider: retry_request
FallbackProvider-->>FallbackModel: response
FallbackModel-->>ExperimentAgent: result
end
Flow diagram for build_agent_model_with_fallback construction logicflowchart TD
A[build_agent_model_with_fallback] --> B[get_model_identifier]
B --> C[validate_api_key_for_model primary_id]
C --> D[build_agent_model primary_id]
D --> E[get_fallback_model]
E --> F{fallback_id is empty<br/>or fallback_id == primary_id}
F -->|yes| G[return primary]
F -->|no| H[validate_api_key_for_model fallback_id]
H --> I[build_agent_model fallback_id]
I --> J[logger.info agents.fallback_enabled]
J --> K[return FallbackModel primary fallback]
H -->|ValueError| L[logger.warning agents.fallback_disabled]
L --> G
File-Level Changes
Possibly linked issues
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
There was a problem hiding this comment.
Hey - I've found 2 issues
Prompt for AI Agents
Please address the comments from this code review:
## Individual Comments
### Comment 1
<location path="app/features/agents/tests/test_base.py" line_range="67-76" />
<code_context>
validate_api_key_for_model("ollama:llama3.1")
+def test_build_agent_model_with_fallback_wraps_primary_and_fallback():
+ """A distinct, key-backed fallback yields a FallbackModel(primary, fallback)."""
+ settings = get_settings()
+ settings.agent_default_model = "anthropic:claude-sonnet-4-5"
+ settings.agent_fallback_model = "openai:gpt-4o"
+ settings.anthropic_api_key = "test-anthropic-key"
+ settings.openai_api_key = "test-openai-key"
+
+ model = build_agent_model_with_fallback()
+
+ assert isinstance(model, FallbackModel)
+
+
</code_context>
<issue_to_address>
**suggestion (testing):** Strengthen the FallbackModel test by asserting that the primary and fallback are wired as expected, not just that a FallbackModel instance is returned.
This test could still pass if the primary and fallback models were swapped or misconfigured. To make it more robust, assert the internal configuration of the `FallbackModel` via its public API so that the primary maps to `settings.agent_default_model` and the fallback to `settings.agent_fallback_model`. If that isn’t possible, add a behavior-focused test using stub models that confirms the primary is tried first and the fallback is used when the primary raises `ModelAPIError`.
Suggested implementation:
```python
def test_build_agent_model_with_fallback_wraps_primary_and_fallback():
"""A distinct, key-backed fallback yields a FallbackModel with correctly wired primary and fallback."""
settings = get_settings()
settings.agent_default_model = "anthropic:claude-sonnet-4-5"
settings.agent_fallback_model = "openai:gpt-4o"
settings.anthropic_api_key = "test-anthropic-key"
settings.openai_api_key = "test-openai-key"
model = build_agent_model_with_fallback()
# Still ensure the wrapper type is correct
assert isinstance(model, FallbackModel)
# Verify that the primary model corresponds to the agent_default_model
primary = model.primary
# Different model classes may expose their identifier differently; prefer a public attribute if available.
primary_id = getattr(primary, "model", None) or getattr(primary, "model_name", None) or getattr(primary, "model_id", None)
assert primary_id == settings.agent_default_model
# Verify that the fallback model corresponds to the agent_fallback_model
fallback = model.fallback
fallback_id = getattr(fallback, "model", None) or getattr(fallback, "model_name", None) or getattr(fallback, "model_id", None)
assert fallback_id == settings.agent_fallback_model
```
To make this test pass robustly, ensure that:
1. `FallbackModel` exposes the wrapped models via public attributes or properties named `primary` and `fallback`. If different names are used, update the test accordingly.
2. The underlying model instances on `primary` and `fallback` expose a public identifier attribute such as `model`, `model_name`, or `model_id` that matches the configured `"provider:model"` string (e.g., `"anthropic:claude-sonnet-4-5"` and `"openai:gpt-4o"`). If the public API uses a different attribute or a method (e.g., `primary.name` or `primary.info.model`), adjust the `primary_id` / `fallback_id` extraction to match that API.
3. If no such public attributes exist, consider adding a small, public, read-only property on the model classes (or on `FallbackModel` itself) that exposes the configured model string so tests can assert wiring without relying on internals.
</issue_to_address>
### Comment 2
<location path="app/features/agents/tests/test_base.py" line_range="75-77" />
<code_context>
+ assert isinstance(model, FallbackModel)
+
+
+def test_build_agent_model_with_fallback_primary_only_when_fallback_key_missing():
+ """With no API key for the fallback provider, the primary is returned alone."""
+ settings = get_settings()
+ settings.agent_default_model = "anthropic:claude-sonnet-4-5"
+ settings.agent_fallback_model = "openai:gpt-4o"
+ settings.anthropic_api_key = "test-anthropic-key"
+ settings.openai_api_key = ""
+
+ model = build_agent_model_with_fallback()
+
+ assert model == "anthropic:claude-sonnet-4-5"
+
+
</code_context>
<issue_to_address>
**suggestion (testing):** Add a test for the fail-fast behavior when the primary model API key is missing in build_agent_model_with_fallback.
The current tests cover a missing fallback key and the primary==fallback case. Since `build_agent_model_with_fallback()` also calls `validate_api_key_for_model(primary_id)`, please add a test that sets `agent_default_model` to a non-Ollama provider, clears its API key, and asserts that `build_agent_model_with_fallback()` raises `ValueError`. This will lock in the fail-fast behavior and protect against future removal of the primary key validation.
```suggestion
model = build_agent_model_with_fallback()
assert isinstance(model, FallbackModel)
def test_build_agent_model_with_fallback_raises_when_primary_api_key_missing():
"""Fail fast when primary model uses non-Ollama provider and its API key is missing."""
settings = get_settings()
settings.agent_default_model = "anthropic:claude-sonnet-4-5"
settings.agent_fallback_model = "openai:gpt-4o"
settings.anthropic_api_key = ""
settings.openai_api_key = "test-openai-key"
with pytest.raises(ValueError):
build_agent_model_with_fallback()
```
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
| def test_build_agent_model_with_fallback_wraps_primary_and_fallback(): | ||
| """A distinct, key-backed fallback yields a FallbackModel(primary, fallback).""" | ||
| settings = get_settings() | ||
| settings.agent_default_model = "anthropic:claude-sonnet-4-5" | ||
| settings.agent_fallback_model = "openai:gpt-4o" | ||
| settings.anthropic_api_key = "test-anthropic-key" | ||
| settings.openai_api_key = "test-openai-key" | ||
|
|
||
| model = build_agent_model_with_fallback() | ||
|
|
There was a problem hiding this comment.
suggestion (testing): Strengthen the FallbackModel test by asserting that the primary and fallback are wired as expected, not just that a FallbackModel instance is returned.
This test could still pass if the primary and fallback models were swapped or misconfigured. To make it more robust, assert the internal configuration of the FallbackModel via its public API so that the primary maps to settings.agent_default_model and the fallback to settings.agent_fallback_model. If that isn’t possible, add a behavior-focused test using stub models that confirms the primary is tried first and the fallback is used when the primary raises ModelAPIError.
Suggested implementation:
def test_build_agent_model_with_fallback_wraps_primary_and_fallback():
"""A distinct, key-backed fallback yields a FallbackModel with correctly wired primary and fallback."""
settings = get_settings()
settings.agent_default_model = "anthropic:claude-sonnet-4-5"
settings.agent_fallback_model = "openai:gpt-4o"
settings.anthropic_api_key = "test-anthropic-key"
settings.openai_api_key = "test-openai-key"
model = build_agent_model_with_fallback()
# Still ensure the wrapper type is correct
assert isinstance(model, FallbackModel)
# Verify that the primary model corresponds to the agent_default_model
primary = model.primary
# Different model classes may expose their identifier differently; prefer a public attribute if available.
primary_id = getattr(primary, "model", None) or getattr(primary, "model_name", None) or getattr(primary, "model_id", None)
assert primary_id == settings.agent_default_model
# Verify that the fallback model corresponds to the agent_fallback_model
fallback = model.fallback
fallback_id = getattr(fallback, "model", None) or getattr(fallback, "model_name", None) or getattr(fallback, "model_id", None)
assert fallback_id == settings.agent_fallback_modelTo make this test pass robustly, ensure that:
FallbackModelexposes the wrapped models via public attributes or properties namedprimaryandfallback. If different names are used, update the test accordingly.- The underlying model instances on
primaryandfallbackexpose a public identifier attribute such asmodel,model_name, ormodel_idthat matches the configured"provider:model"string (e.g.,"anthropic:claude-sonnet-4-5"and"openai:gpt-4o"). If the public API uses a different attribute or a method (e.g.,primary.nameorprimary.info.model), adjust theprimary_id/fallback_idextraction to match that API. - If no such public attributes exist, consider adding a small, public, read-only property on the model classes (or on
FallbackModelitself) that exposes the configured model string so tests can assert wiring without relying on internals.
| model = build_agent_model_with_fallback() | ||
|
|
||
| assert isinstance(model, FallbackModel) |
There was a problem hiding this comment.
suggestion (testing): Add a test for the fail-fast behavior when the primary model API key is missing in build_agent_model_with_fallback.
The current tests cover a missing fallback key and the primary==fallback case. Since build_agent_model_with_fallback() also calls validate_api_key_for_model(primary_id), please add a test that sets agent_default_model to a non-Ollama provider, clears its API key, and asserts that build_agent_model_with_fallback() raises ValueError. This will lock in the fail-fast behavior and protect against future removal of the primary key validation.
| model = build_agent_model_with_fallback() | |
| assert isinstance(model, FallbackModel) | |
| model = build_agent_model_with_fallback() | |
| assert isinstance(model, FallbackModel) | |
| def test_build_agent_model_with_fallback_raises_when_primary_api_key_missing(): | |
| """Fail fast when primary model uses non-Ollama provider and its API key is missing.""" | |
| settings = get_settings() | |
| settings.agent_default_model = "anthropic:claude-sonnet-4-5" | |
| settings.agent_fallback_model = "openai:gpt-4o" | |
| settings.anthropic_api_key = "" | |
| settings.openai_api_key = "test-openai-key" | |
| with pytest.raises(ValueError): | |
| build_agent_model_with_fallback() |
Summary
An agent chat against
gemini-3-flash-previewdied on a transient503 UNAVAILABLEfrom Google. The configuredagent_fallback_modelnever engaged —get_fallback_model()was dead code, and both agents builtAgent(model=<primary only>).This wires PydanticAI's
FallbackModelinto both agents, so aModelAPIError(HTTP 5xx, rate limit) on the primary transparently retriesagent_fallback_model.Closes #183.
Changes
app/features/agents/agents/base.pybuild_agent_model_with_fallback()— returnsFallbackModel(primary, fallback), or the primary alone when no fallback is configured / it equals the primary / its provider has no API key (logged)app/features/agents/agents/experiment.pyapp/features/agents/agents/rag_assistant.pyapp/features/agents/tests/test_base.pyBehaviour
FallbackModel's defaultfallback_on=(ModelAPIError,)already covers the 503 — no custom predicate needed.agents.fallback_disabledand runs primary-only rather than crashing at construction.Modelobjects) and cloud identifiers (strings) both pass through.openai:gpt-4o.Verification
ruff check+ruff format --check— cleanmypy+pyright(changed modules) — 0 errorspytest -m "not integration"— 1093 passed (incl. 3 new)Summary by Sourcery
Wire agents to use a primary model wrapped with an optional fallback model to improve resilience to transient provider failures.
Bug Fixes:
Enhancements:
Tests: