Spaces:
Paused
Paused
Mirrowel
commited on
Commit
·
2ccd2a1
1
Parent(s):
1980665
feat(client): ✨ add header fallback for custom_reasoning_budget
Browse filesAllow RotatingClient to source the custom reasoning budget flag from an HTTP header when the caller does not provide the flag in the request body.
- Prefer the explicit kwarg if present; fall back to request.headers only when missing
- Safely inspect request.headers and interpret a case-insensitive `true` as enabled
- Apply the header-based fallback in both the provider delegation path and the main execution path
- Preserve existing behavior when neither the kwarg nor the header is supplied
src/rotator_library/client.py
CHANGED
|
@@ -448,6 +448,18 @@ class RotatingClient:
|
|
| 448 |
if provider_plugin and provider_plugin.has_custom_logic():
|
| 449 |
lib_logger.debug(f"Provider '{provider}' has custom logic. Delegating call.")
|
| 450 |
litellm_kwargs["credential_identifier"] = current_cred
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 451 |
|
| 452 |
# The plugin handles the entire call, including retries on 401, etc.
|
| 453 |
# The main retry loop here is for key rotation on other errors.
|
|
@@ -631,8 +643,17 @@ class RotatingClient:
|
|
| 631 |
litellm_kwargs = self.all_providers.get_provider_kwargs(**kwargs.copy())
|
| 632 |
if "reasoning_effort" in kwargs:
|
| 633 |
litellm_kwargs["reasoning_effort"] = kwargs["reasoning_effort"]
|
|
|
|
| 634 |
if "custom_reasoning_budget" in kwargs:
|
| 635 |
litellm_kwargs["custom_reasoning_budget"] = kwargs["custom_reasoning_budget"]
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 636 |
|
| 637 |
# [NEW] Merge provider-specific params
|
| 638 |
if provider in self.litellm_provider_params:
|
|
|
|
| 448 |
if provider_plugin and provider_plugin.has_custom_logic():
|
| 449 |
lib_logger.debug(f"Provider '{provider}' has custom logic. Delegating call.")
|
| 450 |
litellm_kwargs["credential_identifier"] = current_cred
|
| 451 |
+
|
| 452 |
+
# Check body first for custom_reasoning_budget
|
| 453 |
+
if "custom_reasoning_budget" in kwargs:
|
| 454 |
+
litellm_kwargs["custom_reasoning_budget"] = kwargs["custom_reasoning_budget"]
|
| 455 |
+
else:
|
| 456 |
+
custom_budget_header = None
|
| 457 |
+
if request and hasattr(request, 'headers'):
|
| 458 |
+
custom_budget_header = request.headers.get("custom_reasoning_budget")
|
| 459 |
+
|
| 460 |
+
if custom_budget_header is not None:
|
| 461 |
+
is_budget_enabled = custom_budget_header.lower() == 'true'
|
| 462 |
+
litellm_kwargs["custom_reasoning_budget"] = is_budget_enabled
|
| 463 |
|
| 464 |
# The plugin handles the entire call, including retries on 401, etc.
|
| 465 |
# The main retry loop here is for key rotation on other errors.
|
|
|
|
| 643 |
litellm_kwargs = self.all_providers.get_provider_kwargs(**kwargs.copy())
|
| 644 |
if "reasoning_effort" in kwargs:
|
| 645 |
litellm_kwargs["reasoning_effort"] = kwargs["reasoning_effort"]
|
| 646 |
+
# Check body first for custom_reasoning_budget
|
| 647 |
if "custom_reasoning_budget" in kwargs:
|
| 648 |
litellm_kwargs["custom_reasoning_budget"] = kwargs["custom_reasoning_budget"]
|
| 649 |
+
else:
|
| 650 |
+
custom_budget_header = None
|
| 651 |
+
if request and hasattr(request, 'headers'):
|
| 652 |
+
custom_budget_header = request.headers.get("custom_reasoning_budget")
|
| 653 |
+
|
| 654 |
+
if custom_budget_header is not None:
|
| 655 |
+
is_budget_enabled = custom_budget_header.lower() == 'true'
|
| 656 |
+
litellm_kwargs["custom_reasoning_budget"] = is_budget_enabled
|
| 657 |
|
| 658 |
# [NEW] Merge provider-specific params
|
| 659 |
if provider in self.litellm_provider_params:
|