From 6d65ac3e75a10fcdc08e7fb5ab11cf4b18247fe4 Mon Sep 17 00:00:00 2001 From: Oliver Fei Date: Tue, 27 Oct 2026 17:05:43 -0400 Subject: [PATCH] fix(guardrails): resolve TrendAI API key through secret manager --- .../guardrail_hooks/trendai/trendai.py | 3 +- .../guardrail_hooks/trendai/test_trendai.py | 43 +++++++++++++++++++ 2 files changed, 45 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/trendai/trendai.py b/litellm/proxy/guardrails/guardrail_hooks/trendai/trendai.py index f21f18f2711..5c1c276e89a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/trendai/trendai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/trendai/trendai.py @@ -26,6 +26,7 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # legacy client factory has an untyped params map httpxSpecialProvider, ) +from litellm.secret_managers.main import get_secret_str from litellm.types.guardrails import GuardrailEventHooks, Mode from litellm.types.utils import GenericGuardrailAPIInputs, GuardrailStatus @@ -89,7 +90,7 @@ class TrendAIGuardrail(CustomGuardrail): default_on: bool = False, streaming_buffer_until_moderated: bool = True, ) -> None: - resolved_api_key: Final = api_key or os.environ.get("TMV1_API_KEY") + resolved_api_key: Final = api_key or get_secret_str("TMV1_API_KEY") if not resolved_api_key: raise ValueError( "Trend AI Guard requires an API key. Pass api_key or set the TMV1_API_KEY environment variable." diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/trendai/test_trendai.py b/tests/unit/proxy/guardrails/guardrail_hooks/trendai/test_trendai.py index 44fb52ee546..2e16a77eb55 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/trendai/test_trendai.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/trendai/test_trendai.py @@ -3,7 +3,9 @@ import json from collections.abc import AsyncIterator, Callable, Mapping, Sequence from contextlib import aclosing from copy import deepcopy +from dataclasses import dataclass from functools import reduce +from types import MappingProxyType from typing import Final, Literal import httpx @@ -32,6 +34,7 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import UnifiedLLMGuardrails from litellm.proxy.guardrails.guardrail_registry import InMemoryGuardrailHandler from litellm.types.guardrails import GuardrailEventHooks, LitellmParams +from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem from litellm.types.utils import ( CallTypes, Choices, @@ -132,6 +135,46 @@ def test_explicit_configuration_takes_precedence(monkeypatch: pytest.MonkeyPatch assert guardrail.app_name == "configured-app" +@dataclass(frozen=True, slots=True) +class _SecretStore: + secrets: Mapping[str, str] + + def sync_read_secret(self, secret_name: str) -> str | None: + return self.secrets.get(secret_name) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("api_key", [None, "explicit-key"]) +async def test_outbound_scan_uses_secret_manager_fallback_unless_key_is_explicit( + api_key: str | None, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("TMV1_API_KEY", "environment-key") + monkeypatch.setattr( + litellm, "secret_manager_client", _SecretStore(MappingProxyType({"TMV1_API_KEY": "managed-key"})) + ) + monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.HASHICORP_VAULT) + monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(access_mode="read_only")) + expected_key: Final = api_key or "managed-key" + + def respond(request: httpx.Request) -> httpx.Response: + if request.headers.get("Authorization") != f"Bearer {expected_key}": + return httpx.Response(401) + return httpx.Response(200, json={"action": "allow"}) + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client: + guardrail: Final = TrendAIGuardrail( + api_key=api_key, + api_base="https://guard.example.com", + guardrail_name="trendai", + event_hook=GuardrailEventHooks.pre_call, + async_handler=client, + ) + inputs: Final[GenericGuardrailAPIInputs] = {"texts": ["hello"]} + result: Final = await guardrail.apply_guardrail(inputs=inputs, request_data={}, input_type="request") + + assert result == inputs + + def test_negative_stream_overlap_is_rejected() -> None: with pytest.raises(ValueError, match="stream_overlap_size"): _guardrail(stream_overlap_size=-1)