diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 66db9b84dc8..11470ca7c29 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -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( diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 4d59c47aa3d..a75ad7dbdc0 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -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() )