From ea1fc1986be6cea80a144a0fc9478fa006611f0b Mon Sep 17 00:00:00 2001 From: yucheng Date: Sun, 27 Sep 2026 01:04:44 +0000 Subject: [PATCH] refactor(guardrails): inject the http handler into the agent 365 initializer instead of assigning it after construction Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../guardrails/guardrail_hooks/agent_365/__init__.py | 9 ++++++++- .../proxy/guardrails/guardrail_hooks/test_agent_365.py | 7 +++---- 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/agent_365/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/agent_365/__init__.py index 12ecb7f619f..2a8c6479ae6 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/agent_365/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/agent_365/__init__.py @@ -6,10 +6,16 @@ from litellm.types.guardrails import SupportedGuardrailIntegrations from .agent_365 import Agent365Guardrail if TYPE_CHECKING: + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.types.guardrails import Guardrail, LitellmParams -def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail") -> Agent365Guardrail: +def initialize_guardrail( + litellm_params: "LitellmParams", + guardrail: "Guardrail", + *, + async_handler: "AsyncHTTPHandler | None" = None, +) -> Agent365Guardrail: import litellm from litellm.secret_managers.main import get_secret_str @@ -57,6 +63,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" client_secret=client_secret, request_timeout=litellm_params.timeout if litellm_params.timeout is not None else 10.0, unreachable_fallback=litellm_params.unreachable_fallback, + async_handler=async_handler, event_hook=litellm_params.mode, default_on=litellm_params.default_on, ) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_agent_365.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_agent_365.py index 99b45272754..23131321938 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_agent_365.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_agent_365.py @@ -257,11 +257,10 @@ class TestInitializeGuardrail: resource_app_id="00000000-0000-0000-0000-000000000000", agent_id="yaml-agent", ) - with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): - guardrail: Final = initialize_guardrail(params, {"guardrail_name": "a365-stale"}) - assert "ignoring api_base, resource_app_id, agent_id" in caplog.text handler: Final = FakeHandler([_token_response(), _allow_response()]) - guardrail.async_handler = handler + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + guardrail: Final = initialize_guardrail(params, {"guardrail_name": "a365-stale"}, async_handler=handler) + assert "ignoring api_base, resource_app_id, agent_id" in caplog.text await _run(guardrail, _mcp_data()) token_call, evaluate_call = handler.calls assert token_call.url == TOKEN_URL