scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 207

fix(proxy): socket guardrail picks read the key's team metadata when the proxy has no database

Browse source
This commit is contained in:
mateo-berri 2026-10-29 23:45:41 -07:00
parent 18d96358d6
commit 3878bda759
2 changed files with 21 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
litellm/proxy/response_api_endpoints/endpoints.py
View file

@ -2203,6 +2203,9 @@ async def _key_team_object(user_api_key_dict: UserAPIKeyAuth) -> LiteLLM_TeamTab
team_id: Final = user_api_key_dict.team_id
if team_id is None or team_id == UI_TEAM_ID:
return None
team_from_key: Final = LiteLLM_TeamTable(team_id=team_id, metadata=user_api_key_dict.team_metadata)
if prisma_client is None:
return team_from_key
try:
return await get_team_object(
team_id=team_id,
@ -2212,7 +2215,7 @@ async def _key_team_object(user_api_key_dict: UserAPIKeyAuth) -> LiteLLM_TeamTab
proxy_logging_obj=proxy_logging_obj,
)
except HTTPException:
return LiteLLM_TeamTable(team_id=team_id, metadata=user_api_key_dict.team_metadata)
return team_from_key
def _responses_ws_frame_guard(

26
tests/unit/proxy/response_api_endpoints/test_endpoints.py
View file

@ -3067,8 +3067,22 @@ class TestResponsesWSFirstFrameModelAuth:
[(True, ["hi"]), (False, [])],
ids=["open", "locked"],
)
async def test_frame_guard_reads_the_key_s_team_metadata_when_the_team_row_is_gone(
self, modify_guardrails: bool, picked_inputs: list[str]
@pytest.mark.parametrize(
"team_lookup",
[
lambda: patch(
"litellm.proxy.response_api_endpoints.endpoints.get_team_object",
AsyncMock(side_effect=HTTPException(status_code=404, detail={"error": "Team doesn't exist"})),
),
lambda: patch("litellm.proxy.proxy_server.prisma_client", None),
],
ids=["team-row-gone", "no-database"],
)
async def test_frame_guard_reads_the_key_s_team_metadata_when_the_team_cannot_be_loaded(
self,
modify_guardrails: bool,
picked_inputs: list[str],
team_lookup: Callable[[], contextlib.AbstractContextManager[object]],
):
from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard
@ -3080,13 +3094,7 @@ class TestResponsesWSFirstFrameModelAuth:
)
frame = json.dumps({"type": "response.create", "input": "hi", "guardrails": ["frame-pick"]})
with (
_guardrails_on(frame_pick) as proxy_logging_obj,
patch(
"litellm.proxy.response_api_endpoints.endpoints.get_team_object",
AsyncMock(side_effect=HTTPException(status_code=404, detail={"error": "Team doesn't exist"})),
),
):
with _guardrails_on(frame_pick) as proxy_logging_obj, team_lookup():
guard = _responses_ws_frame_guard(
{"model": "gpt-4o-mini", "litellm_metadata": {}}, key, proxy_logging_obj, MagicMock()
)