Mirrowel commited on
Commit
2ccd2a1
·
1 Parent(s): 1980665

feat(client): ✨ add header fallback for custom_reasoning_budget

Browse files

Allow 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

Files changed (1) hide show
  1. src/rotator_library/client.py +21 -0
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: