mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +00:00
Fix runtime policy attachment initialization
Mark runtime-created policies and attachments initialized so global policy attachments created from the policy builder apply immediately without requiring a restart. Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
parent
0b87bdff3d
commit
536a24c5ca
3 changed files with 55 additions and 0 deletions
|
|
@ -220,6 +220,7 @@ class AttachmentRegistry:
|
|||
attachment: PolicyAttachment object to add
|
||||
"""
|
||||
self._attachments.append(attachment)
|
||||
self._initialized = True
|
||||
verbose_proxy_logger.debug(f"Added attachment for policy: {attachment.policy}")
|
||||
|
||||
def remove_attachments_for_policy(self, policy_name: str) -> int:
|
||||
|
|
|
|||
|
|
@ -226,6 +226,7 @@ class PolicyRegistry:
|
|||
policy: Policy object to add
|
||||
"""
|
||||
self._policies[policy_name] = policy
|
||||
self._initialized = True
|
||||
verbose_proxy_logger.debug(f"Added/updated policy: {policy_name}")
|
||||
|
||||
def remove_policy(self, policy_name: str) -> bool:
|
||||
|
|
|
|||
|
|
@ -2846,6 +2846,59 @@ async def test_add_guardrails_from_policy_engine_accepts_dynamic_policies_and_po
|
|||
assert "metadata" in data
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_api_created_global_policy_applies_to_new_key_without_restart():
|
||||
"""
|
||||
Regression: policies created at runtime via policy builder must apply
|
||||
immediately when attached globally, even if the server started with no
|
||||
initialized policy config.
|
||||
"""
|
||||
from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry
|
||||
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
|
||||
from litellm.types.proxy.policy_engine import (
|
||||
Policy,
|
||||
PolicyAttachment,
|
||||
PolicyGuardrails,
|
||||
)
|
||||
|
||||
data = {
|
||||
"model": "gpt-4",
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"metadata": {},
|
||||
}
|
||||
user_api_key_dict = UserAPIKeyAuth(api_key="test-key")
|
||||
|
||||
policy_registry = get_policy_registry()
|
||||
attachment_registry = get_attachment_registry()
|
||||
policy_registry._policies = {}
|
||||
policy_registry._initialized = False
|
||||
attachment_registry._attachments = []
|
||||
attachment_registry._initialized = False
|
||||
|
||||
try:
|
||||
policy_registry.add_policy(
|
||||
"runtime-global-policy",
|
||||
Policy(guardrails=PolicyGuardrails(add=["runtime-guardrail"])),
|
||||
)
|
||||
attachment_registry.add_attachment(
|
||||
PolicyAttachment(policy="runtime-global-policy", scope="*")
|
||||
)
|
||||
|
||||
await add_guardrails_from_policy_engine(
|
||||
data=data,
|
||||
metadata_variable_name="metadata",
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
assert "runtime-guardrail" in data["metadata"]["guardrails"]
|
||||
assert "runtime-global-policy" in data["metadata"]["applied_policies"]
|
||||
finally:
|
||||
policy_registry._policies = {}
|
||||
policy_registry._initialized = False
|
||||
attachment_registry._attachments = []
|
||||
attachment_registry._initialized = False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_guardrails_from_policy_engine_policy_version_by_id():
|
||||
"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue