scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-07 02:59:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

test(proxy): configure agent registry test doubles

Browse source
This commit is contained in:
AlisinaDevelo 2026-09-27 06:24:53 +02:00
parent 57a9211b1b
commit 0d0df83bfd
3 changed files with 7 additions and 8 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/proxy/hooks/max_budget_per_session_limiter.py
View file

@ -210,9 +210,7 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger):
def _make_cache_key(self, session_id: str, agent_id: str) -> str:
from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
stable_agent_id: Final = global_agent_registry.stable_agent_id(agent_id)
canonical_agent_id: Final = stable_agent_id if isinstance(stable_agent_id, str) else agent_id
scope: Final = json.dumps((canonical_agent_id, session_id), separators=(",", ":"))
scope: Final = json.dumps((global_agent_registry.stable_agent_id(agent_id), session_id), separators=(",", ":"))
return f"{{agent_session_budget:{scope}}}:spend"
async def _get_current_spend(self, cache_key: str) -> float:

4
litellm/proxy/hooks/max_iterations_limiter.py
View file

@ -184,9 +184,7 @@ class _PROXY_MaxIterationsHandler(CustomLogger):
return f"{{session_iterations:{session_id}}}:count"
from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
stable_agent_id: Final = global_agent_registry.stable_agent_id(agent_id)
canonical_agent_id: Final = stable_agent_id if isinstance(stable_agent_id, str) else agent_id
scope: Final = json.dumps((canonical_agent_id, session_id), separators=(",", ":"))
scope: Final = json.dumps((global_agent_registry.stable_agent_id(agent_id), session_id), separators=(",", ":"))
return f"{{agent_session_iterations:{scope}}}:count"
async def _increment_and_get(self, cache_key: str) -> int:

7
tests/test_litellm/proxy/hooks/test_proxy_rate_limit_provider_field.py
View file

@ -42,6 +42,7 @@ import litellm
from litellm.caching.caching import DualCache
from litellm.exceptions import RateLimitError
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
from litellm.proxy.hooks.batch_rate_limiter import (
BatchFileUsage,
_PROXY_BatchRateLimiter,
@ -60,7 +61,6 @@ from litellm.proxy.hooks.parallel_request_limiter import (
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
_PROXY_MaxParallelRequestsHandler_v3,
)
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
from litellm.proxy.hooks.rate_limiter_utils import (
PROXY_LLM_PROVIDER_FALLBACK,
resolve_llm_provider_for_rate_limit,
@ -68,7 +68,6 @@ from litellm.proxy.hooks.rate_limiter_utils import (
from litellm.proxy.utils import InternalUsageCache
from litellm.types.agents import AgentResponse
# ---------------------------------------------------------------------------
# Helper class itself
# ---------------------------------------------------------------------------
@ -855,6 +854,7 @@ async def test_max_iterations_limiter_populates_provider():
"litellm.proxy.agent_endpoints.agent_registry.global_agent_registry"
) as mock_registry:
mock_registry.get_agent_by_id.return_value = _make_iter_agent(max_iterations=1)
mock_registry.stable_agent_id.return_value = "agent-iter"
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
@ -896,6 +896,7 @@ async def test_max_iterations_limiter_unknown_model_falls_back():
"litellm.proxy.agent_endpoints.agent_registry.global_agent_registry"
) as mock_registry:
mock_registry.get_agent_by_id.return_value = _make_iter_agent(max_iterations=1)
mock_registry.stable_agent_id.return_value = "agent-iter"
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
@ -950,6 +951,7 @@ async def test_max_budget_per_session_limiter_populates_provider():
mock_registry.get_agent_by_id.return_value = _make_session_budget_agent(
max_budget=1.0
)
mock_registry.stable_agent_id.return_value = "agent-session-budget"
with patch.object(
handler, "_get_current_spend", new=AsyncMock(return_value=5.0)
):
@ -985,6 +987,7 @@ async def test_max_budget_per_session_limiter_unknown_model_falls_back():
mock_registry.get_agent_by_id.return_value = _make_session_budget_agent(
max_budget=1.0
)
mock_registry.stable_agent_id.return_value = "agent-session-budget"
with patch.object(
handler, "_get_current_spend", new=AsyncMock(return_value=5.0)
):