From 273855b4e2e8b3d2cbc61216729e52881ebd0816 Mon Sep 17 00:00:00 2001 From: milan-berri Date: Sun, 7 Jun 2026 02:11:54 +0300 Subject: [PATCH 001/185] fix(responses-bridge): map system-only chat request to system input item (#29817) System-only chat requests mapped the system message to instructions and left input=[], which OpenAI's Responses API rejects (it also rejects input=""). When no other messages are present, carry the system message as a role:"system" input item (single copy, correct role) instead of leaving input empty. Mirrors the existing handling of non-string system content. Fixes Open WebUI new-conversation failures on mode:responses Codex models. Co-authored-by: Cursor --- .../transformation.py | 14 ++++++++ ...responses_transformation_transformation.py | 36 +++++++++++++++++++ 2 files changed, 50 insertions(+) diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 51abbbf729b..6d8b5cf8a57 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -402,6 +402,20 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): instructions, ) = self.convert_chat_completion_messages_to_responses_api(messages) + # OpenAI's Responses API rejects an empty input. For a system-only + # request, carry the system message as a system-role input item instead + # of instructions, mirroring how non-string system content is already + # handled in convert_chat_completion_messages_to_responses_api. + if not input_items and instructions is not None: + input_items = [ + { + "type": "message", + "role": "system", + "content": [{"type": "input_text", "text": instructions}], + } + ] + instructions = None + optional_params = self._extract_extra_body_params(optional_params) # Build responses API request using the reverse transformation logic diff --git a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index d335c359aa0..06457dfebff 100644 --- a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -13,6 +13,9 @@ sys.path.insert( 0, os.path.abspath("../../..") ) # Adds the parent directory to the system-path import litellm +from litellm.completion_extras.litellm_responses_transformation.transformation import ( + LiteLLMResponsesTransformationHandler, +) def test_convert_chat_completion_messages_to_responses_api_image_input(): @@ -860,6 +863,39 @@ def test_extract_extra_body_params_reasoning_effort_override(): assert "extra_body" not in result +def test_transform_request_system_only_message_maps_to_system_input_item(): + """System-only requests must not send input=[] to the Responses API. + + OpenAI rejects both input=[] and input="". When the only message is a + system message, carry it as a system-role input item (single copy, correct + role) rather than leaving input empty or duplicating it into instructions. + """ + handler = LiteLLMResponsesTransformationHandler() + logging_obj = Mock() + messages = [{"role": "system", "content": "You are a helpful assistant."}] + + result = handler.transform_request( + model="gpt-5.3-codex", + messages=messages, + optional_params={}, + litellm_params={}, + headers={}, + litellm_logging_obj=logging_obj, + ) + + assert result["input"] == [ + { + "type": "message", + "role": "system", + "content": [ + {"type": "input_text", "text": "You are a helpful assistant."} + ], + } + ] + # System content lives in input only; not duplicated into instructions. + assert not result.get("instructions") + + def test_transform_request_single_char_keys_not_matched(): """Test that single-character keys are not incorrectly matched to 'metadata' or 'previous_response_id' From d61f7747c067dedd9f67f11cf6ae0b211bc61d26 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 6 Jun 2026 16:28:18 -0700 Subject: [PATCH 002/185] feat(bedrock): forward strict and additionalProperties to Converse toolSpec (#29814) * feat(bedrock): forward strict and additionalProperties to Converse toolSpec Bedrock Converse supports strict in toolSpec since 2026-02, but _bedrock_tools_pt only whitelisted type/properties/required/name/description, so strict: true was silently dropped and Claude-on-Bedrock ignored enum constraints that GPT and direct-Anthropic honored. Forward strict from the OpenAI function and additionalProperties from the schema (Bedrock requires the latter alongside strict), passing each only when present. https://claude.ai/code/session_01WQjWd8NfUB3vxERwudbHkv * fix(bedrock): only forward strict tool schemas to Claude on Converse Nova, Llama and GPT-OSS on Bedrock reject the strict field (BedrockException 'This model doesn't support the strict field'), and the GPT-OSS request-body test asserts strict/additionalProperties are stripped. Forwarding them to every model broke the llm_translation suite, so gate the forwarding on the anthropic base model since only Claude honours strict tool schemas on Bedrock. --- .../prompt_templates/factory.py | 23 ++++++-- litellm/types/llms/bedrock.py | 2 + ...llm_core_utils_prompt_templates_factory.py | 55 +++++++++++++++++++ 3 files changed, 74 insertions(+), 6 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 1460dbaf0a9..588a183d58e 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -5496,6 +5496,7 @@ def _bedrock_tools_pt( ] """ from litellm.llms.bedrock.common_utils import ( + get_bedrock_base_model, normalize_json_schema_custom_types_to_object, ) from litellm.litellm_core_utils.prompt_templates.common_utils import unpack_defs @@ -5503,6 +5504,11 @@ def _bedrock_tools_pt( _valid_json_schema_root_types = frozenset( ("array", "boolean", "integer", "null", "number", "object", "string") ) + # Only Claude on Bedrock honours strict tool schemas; other families + # (Nova, Llama, GPT-OSS) reject the strict field outright. + supports_strict_tools = bool( + model and get_bedrock_base_model(model).startswith("anthropic") + ) tool_block_list: List[BedrockToolBlock] = [] for tool_idx, tool in enumerate(tools): # Check if tool is already a BedrockToolBlock (e.g., systemTool for Nova grounding) @@ -5548,16 +5554,21 @@ def _bedrock_tools_pt( normalize_json_schema_custom_types_to_object(parameters) if parameters.get("type") not in _valid_json_schema_root_types: parameters["type"] = "object" - tool_input_schema = BedrockToolInputSchemaBlock( - json=BedrockToolJsonSchemaBlock( - type=parameters["type"], - properties=parameters.get("properties", {}), - required=parameters.get("required", []), - ) + json_schema = BedrockToolJsonSchemaBlock( + type=parameters["type"], + properties=parameters.get("properties", {}), + required=parameters.get("required", []), ) + additional_properties = parameters.get("additionalProperties", None) + if supports_strict_tools and additional_properties is not None: + json_schema["additionalProperties"] = additional_properties + tool_input_schema = BedrockToolInputSchemaBlock(json=json_schema) tool_spec = BedrockToolSpecBlock( inputSchema=tool_input_schema, name=name, description=description ) + strict = tool.get("function", {}).get("strict", None) + if supports_strict_tools and strict is not None: + tool_spec["strict"] = strict tool_block = BedrockToolBlock(toolSpec=tool_spec) tool_block_list.append(tool_block) diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index 60e640636aa..ed80b44414b 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -250,6 +250,7 @@ class ToolJsonSchemaBlock(TypedDict, total=False): type: Literal["object"] properties: dict required: List[str] + additionalProperties: bool class ToolInputSchemaBlock(TypedDict): @@ -260,6 +261,7 @@ class ToolSpecBlock(TypedDict, total=False): inputSchema: Required[ToolInputSchemaBlock] name: Required[str] description: str + strict: bool class SystemToolBlock(TypedDict, total=False): diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index 91fd07dcffc..3bf3b04bf14 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -898,6 +898,61 @@ def test_bedrock_tools_unpack_defs(): _bedrock_tools_pt(tools=tools) +def test_bedrock_tools_pt_strict_parameter(): + """Regression for strict tools on the Bedrock Converse path. + + Claude on Bedrock honours strict in toolSpec (with additionalProperties, which + Bedrock requires alongside strict); without forwarding it the model ignores the + enum constraint the caller asked for. Every other Bedrock family (Nova, Llama, + GPT-OSS) rejects the strict field, so it must only be forwarded for Claude. + """ + tools_with_strict = [ + { + "type": "function", + "function": { + "name": "generate_sql", + "strict": True, + "description": "Generate a SQL query", + "parameters": { + "type": "object", + "properties": {"query": {"type": "string"}}, + "required": ["query"], + "additionalProperties": False, + }, + }, + } + ] + result = _bedrock_tools_pt( + tools_with_strict, model="anthropic.claude-sonnet-4-5-20250929-v1:0" + ) + assert result[0]["toolSpec"]["strict"] is True + assert result[0]["toolSpec"]["inputSchema"]["json"]["additionalProperties"] is False + + result = _bedrock_tools_pt(tools_with_strict, model="us.amazon.nova-micro-v1:0") + assert "strict" not in result[0]["toolSpec"] + assert "additionalProperties" not in result[0]["toolSpec"]["inputSchema"]["json"] + + tools_without_strict = [ + { + "type": "function", + "function": { + "name": "generate_sql", + "description": "Generate a SQL query", + "parameters": { + "type": "object", + "properties": {"query": {"type": "string"}}, + "required": ["query"], + }, + }, + } + ] + result = _bedrock_tools_pt( + tools_without_strict, model="anthropic.claude-sonnet-4-5-20250929-v1:0" + ) + assert "strict" not in result[0]["toolSpec"] + assert "additionalProperties" not in result[0]["toolSpec"]["inputSchema"]["json"] + + def test_bedrock_image_processor_content_type_fallback_url_extension(): """ Test that _post_call_image_processing falls back to URL extension From aeb55e7a11e98c6f4a88b6018828fb10491080d3 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 6 Jun 2026 16:51:25 -0700 Subject: [PATCH 003/185] fix(mcp): highlight MCP cards red when the logged-in user is missing per-user env vars (#29856) * fix(mcp): flag missing per-user env vars on the card for every accessible server The dashboard MCP card grid lists servers via the registry-backed manager (get_all_mcp_servers_unfiltered for admins in view_all mode, the allowed-context aggregation otherwise), but the per-user env-var status endpoint that drives the red "user fields missing" highlight resolved servers through the much narrower get_all_mcp_servers_for_user, which only returns servers explicitly granted on the calling key. An admin's dashboard session key carries no per-server MCP grant, so the status feed came back empty and the card never turned red even when the logged-in user had not filled in their required variables. Both surfaces now share a single _resolve_accessible_mcp_servers helper, so the status feed is computed over exactly the cards the user sees. The helper returns servers unredacted; the status endpoint needs the raw env_vars and still only ever reports is_set booleans, never the stored secret values. * test(mcp): drop dead get_all_mcp_servers_for_user patch from view_all regression test The bulk status endpoint resolves servers through _resolve_accessible_mcp_servers now, so the old get_all_mcp_servers_for_user patch in the admin view_all regression test is never hit. Removing it keeps the test honest about which code path it exercises. --- .../mcp_management_endpoints.py | 56 ++++++++++--------- .../test_mcp_management_endpoints.py | 51 ++++++++++++++++- 2 files changed, 77 insertions(+), 30 deletions(-) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index f1edcc9c7b2..64e89ee40fd 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -879,6 +879,32 @@ if MCP_AVAILABLE: return _redact_mcp_credentials_list(servers) + async def _resolve_accessible_mcp_servers( + user_api_key_dict: UserAPIKeyAuth, + ) -> List[LiteLLM_MCPServerTable]: + """The server set the dashboard grid shows (GET /v1/mcp/server, no team + filter), returned unredacted. Callers that surface this to a client must + apply their own redaction; the per-user env-var status endpoint relies on + the raw env_vars and only ever returns is_set booleans, never secrets. + + Sharing this resolution keeps the red "missing user fields" card status + aligned with the cards actually rendered: an admin in view_all mode sees + every server even when their key carries no per-server MCP grant. + """ + if ( + _get_user_mcp_management_mode() == "view_all" + and not _is_restricted_virtual_key_request(user_api_key_dict) + ): + return await global_mcp_server_manager.get_all_mcp_servers_unfiltered() + + aggregated: Dict[str, LiteLLM_MCPServerTable] = {} + for auth_context in await build_effective_auth_contexts(user_api_key_dict): + for server in await global_mcp_server_manager.get_all_allowed_mcp_servers( + user_api_key_auth=auth_context + ): + aggregated.setdefault(server.server_id, server) + return list(aggregated.values()) + @router.get( "/server", description="Returns the mcp server list with associated teams", @@ -950,30 +976,8 @@ if MCP_AVAILABLE: sanitized_team_id ) else: - user_mcp_management_mode = _get_user_mcp_management_mode() - - if user_mcp_management_mode == "view_all" and not is_restricted_virtual_key: - servers = ( - await global_mcp_server_manager.get_all_mcp_servers_unfiltered() - ) - redacted_mcp_servers = _redact_mcp_credentials_list(servers) - else: - auth_contexts = await build_effective_auth_contexts(user_api_key_dict) - - aggregated_servers: Dict[str, LiteLLM_MCPServerTable] = {} - for auth_context in auth_contexts: - servers = ( - await global_mcp_server_manager.get_all_allowed_mcp_servers( - user_api_key_auth=auth_context - ) - ) - for server in servers: - if server.server_id not in aggregated_servers: - aggregated_servers[server.server_id] = server - - redacted_mcp_servers = _redact_mcp_credentials_list( - aggregated_servers.values() - ) + servers = await _resolve_accessible_mcp_servers(user_api_key_dict) + redacted_mcp_servers = _redact_mcp_credentials_list(servers) # augment the mcp servers with public status if litellm.public_mcp_servers is not None: @@ -2391,9 +2395,7 @@ if MCP_AVAILABLE: user_id = user_api_key_dict.user_id or "" if not user_id: return [] - accessible = await get_all_mcp_servers_for_user( - prisma_client, user_api_key_dict - ) + accessible = await _resolve_accessible_mcp_servers(user_api_key_dict) if not accessible: return [] server_ids = [s.server_id for s in accessible] diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index b5eb091bb81..947b39e8367 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -4146,7 +4146,7 @@ class TestListMCPUserEnvVarStatus: ), patch.object( mgmt_endpoints, - "get_all_mcp_servers_for_user", + "_resolve_accessible_mcp_servers", AsyncMock(return_value=[]), ), ): @@ -4174,7 +4174,7 @@ class TestListMCPUserEnvVarStatus: ), patch.object( mgmt_endpoints, - "get_all_mcp_servers_for_user", + "_resolve_accessible_mcp_servers", AsyncMock(return_value=[server_with, server_without]), ), patch.object( @@ -4204,7 +4204,7 @@ class TestListMCPUserEnvVarStatus: ), patch.object( mgmt_endpoints, - "get_all_mcp_servers_for_user", + "_resolve_accessible_mcp_servers", AsyncMock(return_value=[server]), ), patch.object( @@ -4221,6 +4221,51 @@ class TestListMCPUserEnvVarStatus: assert by_name["CORP_PASSWORD"].is_set is False assert "alice" not in result[0].model_dump_json() + @pytest.mark.asyncio + async def test_admin_view_all_flags_missing_fields_without_key_grants(self): + """Regression: the red "user fields missing" card must light up for an + admin in view_all mode even when their key carries no per-server MCP + grant. The bulk status feed has to resolve the same server set the + dashboard grid renders; the old narrow key-scoped listing returned + nothing for such an admin, leaving every card un-highlighted.""" + server = _make_env_var_server( + server_id="srv-with", + env_vars=_ENV_VARS_MIXED, + static_headers=_STATIC_HEADERS_MIXED, + ) + with ( + patch.object( + mgmt_endpoints, "get_prisma_client_or_throw", return_value=MagicMock() + ), + patch.object( + mgmt_endpoints, + "_get_user_mcp_management_mode", + return_value="view_all", + ), + patch.object( + mgmt_endpoints.global_mcp_server_manager, + "get_all_mcp_servers_unfiltered", + AsyncMock(return_value=[server]), + ), + patch.object( + mgmt_endpoints, + "get_user_env_vars_bulk", + AsyncMock(return_value={}), + ), + ): + result = await mgmt_endpoints.list_mcp_user_env_var_status( + user_api_key_dict=generate_mock_user_api_key_auth( + user_id="admin", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + ) + assert [s.server_id for s in result] == ["srv-with"] + assert result[0].missing_count == 2 + assert {f.name for f in result[0].required} == { + "CORP_USERNAME", + "CORP_PASSWORD", + } + class TestMCPUserEnvVarsAccessControl: """Per-server env-var endpoints must enforce the same access gate as From f31d059aa39e7bf982e7580ac6c78f1eb40459a5 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 6 Jun 2026 17:24:55 -0700 Subject: [PATCH 004/185] feat(ui): add budget duration to edit team member form (#29717) * feat(ui): add budget duration to edit team member form Editing a team member created a member budget with no duration, so the budget never reset. This threads a budget reset period through the edit flow end to end and reuses the shared duration dropdown so the options stay in sync with the rest of the UI. Resolves LIT-2651 * fix(proxy): validate member budget_duration and persist clears Reject budget_duration values that can't be parsed, are non-positive, or overflow date math before any write, so a bad value can't be persisted and later crash the budget reset job. Clearing the budget duration in the edit-member form now sends null and clears the column end to end, so the dropdown's clear control reflects a real change instead of being a no-op * chore(ui): regenerate schema.d.ts for member budget_duration Adds budget_duration to TeamMemberUpdateRequest/Response in the generated dashboard types so the Check UI API Types Sync gate passes --- litellm/proxy/_types.py | 5 + .../management_endpoints/common_utils.py | 161 +++---- .../management_endpoints/team_endpoints.py | 55 ++- .../test_upsert_budget_membership.py | 417 ++++++++---------- .../proxy/test_team_member_update.py | 142 +++++- .../src/components/networking.tsx | 18 +- .../src/components/team/EditMembership.tsx | 8 +- .../src/components/team/TeamInfo.tsx | 13 + .../src/components/team/TeamMemberTab.tsx | 1 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 7 + 10 files changed, 512 insertions(+), 315 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index e5d70933063..88be567e59a 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -4125,6 +4125,10 @@ class TeamMemberUpdateRequest(TeamMemberDeleteRequest): rpm_limit: Optional[int] = Field( default=None, description="Requests per minute limit for this team member" ) + budget_duration: Optional[str] = Field( + default=None, + description="Duration after which this team member's budget resets (e.g. '1h', '24h', '7d', '30d'). If not set, the budget never resets.", + ) allowed_models: Optional[List[str]] = Field( default=None, description="List of models this team member can access. Pass an empty list to remove per-member model restrictions.", @@ -4136,6 +4140,7 @@ class TeamMemberUpdateResponse(MemberUpdateResponse): max_budget_in_team: Optional[float] = None tpm_limit: Optional[int] = None rpm_limit: Optional[int] = None + budget_duration: Optional[str] = None allowed_models: Optional[List[str]] = None diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index dc27e87726a..31d831d773c 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union +from typing import TYPE_CHECKING, Any, Dict, Optional, Union from fastapi import HTTPException, status from pydantic import BaseModel @@ -19,6 +19,7 @@ from litellm.proxy._types import ( UserAPIKeyAuth, user_api_key_has_admin_view as _user_has_admin_view, # noqa: F401 re-exported ) +from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time from litellm.proxy.utils import _premium_user_check if TYPE_CHECKING: @@ -400,121 +401,127 @@ def _set_object_metadata_field( object_data.metadata[field_name] = value +_TEAM_MEMBER_BUDGET_LIMIT_FIELDS = ( + "max_budget", + "soft_budget", + "max_parallel_requests", + "tpm_limit", + "rpm_limit", + "model_max_budget", + "budget_duration", + "allowed_models", +) + + +def _is_set_budget_value(value: Any) -> bool: + if value is None: + return False + if isinstance(value, list) and len(value) == 0: + return False + return True + + +def _has_meaningful_budget_limit(budget_values: Dict[str, Any]) -> bool: + """A budget is meaningful if at least one limit is actually set; an empty + list (no model restriction) and None both count as unset.""" + return any( + _is_set_budget_value(budget_values.get(field)) + for field in _TEAM_MEMBER_BUDGET_LIMIT_FIELDS + ) + + async def _upsert_budget_and_membership( tx, *, team_id: str, user_id: str, - max_budget: Optional[float], existing_budget_id: Optional[str], user_api_key_dict: UserAPIKeyAuth, - tpm_limit: Optional[int] = None, - rpm_limit: Optional[int] = None, - allowed_models: Optional[List[str]] = None, + budget_patch: Dict[str, Any], team_default_budget_id: Optional[str] = None, ): """ - Helper function to Create/Update or Delete the budget within the team membership - Args: - tx: The transaction object - team_id: The ID of the team - user_id: The ID of the user - max_budget: The maximum budget for the team - existing_budget_id: The ID of the existing budget, if any - user_api_key_dict: User API Key dictionary containing user information - tpm_limit: Tokens per minute limit for the team member - rpm_limit: Requests per minute limit for the team member - allowed_models: Per-member model scope. None = don't change. [] = remove restrictions. Non-empty list = enforce. - team_default_budget_id: The team's shared default member budget id (from - team metadata.team_member_budget_id), if any. When the membership's - existing_budget_id matches this, we clone-on-write so editing one - member's budget does not mutate the shared default (and therefore - every other member who still points at it). + Apply a merge-patch of per-member budget fields to a team membership. - If max_budget, tpm_limit, rpm_limit, and allowed_models are all None, the user's budget is removed from the team membership. - If any of these values exist, a budget is updated or created and linked to the team membership. + ``budget_patch`` holds only the budget columns the caller explicitly sent + (RFC 7396 semantics): a value sets the column, ``None`` clears it, and a + column that is absent from the dict is left untouched. Once the patch is + applied, if the budget has no meaningful limit left the member's private + budget is disconnected so they fall back to the team default. + + ``team_default_budget_id`` is the team's shared default member budget id + (from team metadata.team_member_budget_id). When the membership still + points at it, we clone-on-write so editing one member's budget does not + mutate the shared default that every other member points at. """ - if ( - max_budget is None - and tpm_limit is None - and rpm_limit is None - and allowed_models is None - ): - # disconnect the budget since all limits are None - await tx.litellm_teammembership.update( - where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}}, - data={"litellm_budget_table": {"disconnect": True}}, - ) + if not budget_patch: return + write_data = dict(budget_patch) + if "budget_duration" in write_data: + duration = write_data["budget_duration"] + write_data["budget_reset_at"] = ( + get_budget_reset_time(budget_duration=duration) + if duration is not None + else None + ) + is_shared_default = ( existing_budget_id is not None and team_default_budget_id is not None and existing_budget_id == team_default_budget_id ) + async def _disconnect(): + await tx.litellm_teammembership.update( + where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}}, + data={"litellm_budget_table": {"disconnect": True}}, + ) + if existing_budget_id is not None and not is_shared_default: - # Update the existing budget in-place to preserve fields not being changed. - # Only write fields that the caller explicitly provided (non-None). - update_data: Dict[str, Any] = { - "updated_by": user_api_key_dict.user_id or "", - } - if max_budget is not None: - update_data["max_budget"] = max_budget - if tpm_limit is not None: - update_data["tpm_limit"] = tpm_limit - if rpm_limit is not None: - update_data["rpm_limit"] = rpm_limit - if allowed_models is not None: - update_data["allowed_models"] = allowed_models + existing_budget = await tx.litellm_budgettable.find_unique( + where={"budget_id": existing_budget_id} + ) + merged = existing_budget.model_dump() if existing_budget is not None else {} + merged.update(write_data) + if not _has_meaningful_budget_limit(merged): + await _disconnect() + return await tx.litellm_budgettable.update( where={"budget_id": existing_budget_id}, - data=update_data, + data={"updated_by": user_api_key_dict.user_id or "", **write_data}, ) return - # Either there is no existing budget, OR the membership is still pointing - # at the team's shared default member budget. In both cases we create a - # NEW private budget for this user and (re)link the membership to it. create_data: Dict[str, Any] = { "created_by": user_api_key_dict.user_id or "", "updated_by": user_api_key_dict.user_id or "", } - # If we're forking off the shared default, seed the new row with the - # default's values so fields the caller did not change carry over. if is_shared_default: default_budget_row = await tx.litellm_budgettable.find_unique( where={"budget_id": existing_budget_id} ) if default_budget_row is not None: default_budget_dict = default_budget_row.model_dump() - for field in ( - "max_budget", - "soft_budget", - "max_parallel_requests", - "tpm_limit", - "rpm_limit", - "model_max_budget", - "budget_duration", - "allowed_models", - ): + for field in _TEAM_MEMBER_BUDGET_LIMIT_FIELDS: value = default_budget_dict.get(field) - if value is None: - continue - if isinstance(value, list) and len(value) == 0: - continue - create_data[field] = value + if _is_set_budget_value(value): + create_data[field] = value - # Caller-provided values take precedence over the cloned defaults. - if max_budget is not None: - create_data["max_budget"] = max_budget - if tpm_limit is not None: - create_data["tpm_limit"] = tpm_limit - if rpm_limit is not None: - create_data["rpm_limit"] = rpm_limit - if allowed_models is not None: - create_data["allowed_models"] = allowed_models + create_data.update(write_data) + + if create_data.get("budget_duration") is not None: + create_data["budget_reset_at"] = get_budget_reset_time( + budget_duration=create_data["budget_duration"] + ) + else: + create_data.pop("budget_reset_at", None) + + if not _has_meaningful_budget_limit(create_data): + if existing_budget_id is not None: + await _disconnect() + return new_budget = await tx.litellm_budgettable.create( data=create_data, diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index ae7da0d29f2..a3ad7a9ea8b 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -2733,6 +2733,52 @@ async def team_member_delete( return existing_team_row +_MEMBER_BUDGET_PATCH_FIELDS = { + "max_budget_in_team": "max_budget", + "tpm_limit": "tpm_limit", + "rpm_limit": "rpm_limit", + "budget_duration": "budget_duration", + "allowed_models": "allowed_models", +} + + +def _build_member_budget_patch(data: TeamMemberUpdateRequest) -> Dict[str, Any]: + """Map the budget fields the request actually set (merge-patch: a sent + value updates, an explicit null clears, an absent field is left untouched) + to their budget-table columns.""" + provided = data.model_dump(exclude_unset=True) + return { + column: provided[request_field] + for request_field, column in _MEMBER_BUDGET_PATCH_FIELDS.items() + if request_field in provided + } + + +def _validate_budget_duration(budget_duration: Optional[str]) -> None: + """Reject budget durations that can't be parsed, are non-positive, or + overflow date math, so a bad value can't be persisted and later crash the + budget reset job.""" + if budget_duration is None: + return + + from litellm.litellm_core_utils.duration_parser import duration_in_seconds + from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time + + try: + if duration_in_seconds(budget_duration) <= 0: + raise ValueError("budget_duration must be positive") + get_budget_reset_time(budget_duration=budget_duration) + except (ValueError, OverflowError): + raise HTTPException( + status_code=400, + detail={ + "error": "Invalid budget_duration '{}'. Use a format like '1h', '24h', '7d', or '30d'.".format( + budget_duration + ) + }, + ) + + @router.post( "/team/member_update", tags=["team management"], @@ -2770,6 +2816,8 @@ async def team_member_update( detail={"error": "Either user_id or user_email needs to be passed in"}, ) + _validate_budget_duration(data.budget_duration) + _existing_team_row = await prisma_client.db.litellm_teamtable.find_unique( where={"team_id": data.team_id} ) @@ -2843,17 +2891,15 @@ async def team_member_update( team_default_budget_id = raw_default_budget_id ### upsert new budget + budget_patch = _build_member_budget_patch(data) async with prisma_client.db.tx() as tx: await _upsert_budget_and_membership( tx=tx, team_id=data.team_id, user_id=received_user_id, - max_budget=data.max_budget_in_team, existing_budget_id=identified_budget_id, user_api_key_dict=user_api_key_dict, - tpm_limit=data.tpm_limit, - rpm_limit=data.rpm_limit, - allowed_models=data.allowed_models, + budget_patch=budget_patch, team_default_budget_id=team_default_budget_id, ) @@ -2887,6 +2933,7 @@ async def team_member_update( max_budget_in_team=data.max_budget_in_team, tpm_limit=data.tpm_limit, rpm_limit=data.rpm_limit, + budget_duration=data.budget_duration, allowed_models=data.allowed_models, ) diff --git a/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py b/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py index f4bf0d7b2be..e9b4f11e891 100644 --- a/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py +++ b/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py @@ -1,5 +1,6 @@ # tests/litellm/proxy/common_utils/test_upsert_budget_membership.py import types +from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock import pytest @@ -19,15 +20,13 @@ def mock_tx(): Builds an object that looks just enough like the Prisma tx you use inside _upsert_budget_and_membership. """ - # membership “table” membership = MagicMock() membership.update = AsyncMock() membership.upsert = AsyncMock() - # budget “table” budget = MagicMock() budget.update = AsyncMock() - # budget.create returns a fake row that has .budget_id + budget.find_unique = AsyncMock(return_value=None) budget.create = AsyncMock( return_value=types.SimpleNamespace(budget_id="new-budget-123") ) @@ -44,16 +43,57 @@ def fake_user(): return types.SimpleNamespace(user_id="tester@example.com") -# TEST: max_budget is None, disconnect only +def budget_row(**fields): + """A fake litellm_budgettable row whose model_dump returns the given fields.""" + row = MagicMock() + row.model_dump.return_value = fields + return row + + +def assert_future_reset_time(value): + """A budget_reset_at must be a timezone-aware datetime in the future, so the + member's budget actually rolls over and the UI shows a reset date instead of + waiting for the reset cron to backfill it.""" + assert isinstance(value, datetime) + assert value.tzinfo is not None + assert value > datetime.now(timezone.utc) + + +# TEST: an empty patch (caller sent no budget fields) leaves everything alone. +# This is the merge-patch contract: absent != clear. Updating only a member's +# role must not silently wipe their budget. @pytest.mark.asyncio -async def test_upsert_disconnect(mock_tx, fake_user): +async def test_empty_patch_is_noop(mock_tx, fake_user): await _upsert_budget_and_membership( mock_tx, team_id="team-1", user_id="user-1", - max_budget=None, - existing_budget_id=None, + existing_budget_id="bud-1", user_api_key_dict=fake_user, + budget_patch={}, + ) + + mock_tx.litellm_teammembership.update.assert_not_called() + mock_tx.litellm_teammembership.upsert.assert_not_called() + mock_tx.litellm_budgettable.update.assert_not_called() + mock_tx.litellm_budgettable.create.assert_not_called() + + +# TEST: clearing every limit on a member's private budget disconnects it, so the +# member falls back to the team default instead of keeping an empty private row. +@pytest.mark.asyncio +async def test_clearing_all_limits_disconnects(mock_tx, fake_user): + mock_tx.litellm_budgettable.find_unique = AsyncMock( + return_value=budget_row(max_budget=100.0) + ) + + await _upsert_budget_and_membership( + mock_tx, + team_id="team-1", + user_id="user-1", + existing_budget_id="bud-1", + user_api_key_dict=fake_user, + budget_patch={"max_budget": None}, ) mock_tx.litellm_teammembership.update.assert_awaited_once_with( @@ -62,205 +102,114 @@ async def test_upsert_disconnect(mock_tx, fake_user): ) mock_tx.litellm_budgettable.update.assert_not_called() mock_tx.litellm_budgettable.create.assert_not_called() - mock_tx.litellm_teammembership.upsert.assert_not_called() -# TEST: existing budget id → updates budget in-place (current behavior) +# TEST: clearing one field on a budget that still has another limit updates in +# place (clears just that column + its reset time) and does NOT disconnect. @pytest.mark.asyncio -async def test_upsert_with_existing_budget_id_creates_new(mock_tx, fake_user): - """ - Test that when existing_budget_id is provided, the function updates the budget in-place. - """ - await _upsert_budget_and_membership( - mock_tx, - team_id="team-2", - user_id="user-2", - max_budget=42.0, - existing_budget_id="bud-999", - user_api_key_dict=fake_user, +async def test_clear_one_field_keeps_others(mock_tx, fake_user): + mock_tx.litellm_budgettable.find_unique = AsyncMock( + return_value=budget_row(max_budget=100.0, budget_duration="24h") ) - # Should update the existing budget, not create a new one + await _upsert_budget_and_membership( + mock_tx, + team_id="team-1", + user_id="user-1", + existing_budget_id="bud-1", + user_api_key_dict=fake_user, + budget_patch={"budget_duration": None}, + ) + + mock_tx.litellm_teammembership.update.assert_not_called() mock_tx.litellm_budgettable.update.assert_awaited_once_with( - where={"budget_id": "bud-999"}, + where={"budget_id": "bud-1"}, data={ - "max_budget": 42.0, "updated_by": fake_user.user_id, + "budget_duration": None, + "budget_reset_at": None, }, ) - # Should NOT create a new budget or touch membership + +# TEST: setting budget_duration in place writes the duration AND a future +# budget_reset_at, so the budget rolls over without waiting for the reset cron. +@pytest.mark.asyncio +async def test_update_in_place_seeds_reset_at(mock_tx, fake_user): + mock_tx.litellm_budgettable.find_unique = AsyncMock( + return_value=budget_row(max_budget=20.0) + ) + + await _upsert_budget_and_membership( + mock_tx, + team_id="team-dur", + user_id="user-dur", + existing_budget_id="bud-dur", + user_api_key_dict=fake_user, + budget_patch={"budget_duration": "30d"}, + ) + + mock_tx.litellm_budgettable.update.assert_awaited_once() + call = mock_tx.litellm_budgettable.update.await_args + assert call.kwargs["where"] == {"budget_id": "bud-dur"} + data = call.kwargs["data"] + assert data["budget_duration"] == "30d" + assert data["updated_by"] == fake_user.user_id + assert_future_reset_time(data["budget_reset_at"]) mock_tx.litellm_budgettable.create.assert_not_called() - mock_tx.litellm_teammembership.upsert.assert_not_called() - mock_tx.litellm_teammembership.update.assert_not_called() -# TEST: create new budget and link membership +# TEST: updating a single limit in place only writes that field; an untouched +# budget_duration must not get a (re)computed reset time. @pytest.mark.asyncio -async def test_upsert_create_and_link(mock_tx, fake_user): +async def test_update_in_place_single_field_leaves_reset_at_alone(mock_tx, fake_user): + mock_tx.litellm_budgettable.find_unique = AsyncMock( + return_value=budget_row(max_budget=50.0) + ) + await _upsert_budget_and_membership( mock_tx, - team_id="team-3", - user_id="user-3", - max_budget=99.9, - existing_budget_id=None, + team_id="team-rpm", + user_id="user-rpm", + existing_budget_id="bud-rpm", user_api_key_dict=fake_user, + budget_patch={"rpm_limit": 100}, ) - mock_tx.litellm_budgettable.create.assert_awaited_once_with( - data={ - "max_budget": 99.9, - "created_by": fake_user.user_id, - "updated_by": fake_user.user_id, - }, - include={"team_membership": True}, + mock_tx.litellm_budgettable.update.assert_awaited_once_with( + where={"budget_id": "bud-rpm"}, + data={"updated_by": fake_user.user_id, "rpm_limit": 100}, ) - - # Budget ID returned by the mocked create() - bid = mock_tx.litellm_budgettable.create.return_value.budget_id - - mock_tx.litellm_teammembership.upsert.assert_awaited_once_with( - where={"user_id_team_id": {"user_id": "user-3", "team_id": "team-3"}}, - data={ - "create": { - "user_id": "user-3", - "team_id": "team-3", - "litellm_budget_table": {"connect": {"budget_id": bid}}, - }, - "update": { - "litellm_budget_table": {"connect": {"budget_id": bid}}, - }, - }, - ) - - mock_tx.litellm_teammembership.update.assert_not_called() - mock_tx.litellm_budgettable.update.assert_not_called() + mock_tx.litellm_budgettable.create.assert_not_called() -# TEST: create new budget and link membership, then create another new budget +# TEST: with no existing budget, a duration-only patch creates a budget carrying +# the duration and a future reset time, then links the membership. @pytest.mark.asyncio -async def test_upsert_create_then_create_another(mock_tx, fake_user): - """ - Test that multiple calls to _upsert_budget_and_membership create separate budgets, - reflecting the current implementation behavior. - """ - # FIRST CALL – create new budget and link membership +async def test_create_seeds_reset_at_and_links(mock_tx, fake_user): await _upsert_budget_and_membership( mock_tx, - team_id="team-42", - user_id="user-42", - max_budget=10.0, + team_id="team-new", + user_id="user-new", existing_budget_id=None, user_api_key_dict=fake_user, + budget_patch={"budget_duration": "7d"}, ) - # capture the budget id that create() returned - created_bid = mock_tx.litellm_budgettable.create.return_value.budget_id - - # sanity: we really did the create + upsert path mock_tx.litellm_budgettable.create.assert_awaited_once() - mock_tx.litellm_teammembership.upsert.assert_awaited_once() + data = mock_tx.litellm_budgettable.create.await_args.kwargs["data"] + assert data["budget_duration"] == "7d" + assert data["created_by"] == fake_user.user_id + assert data["updated_by"] == fake_user.user_id + assert_future_reset_time(data["budget_reset_at"]) - # SECOND CALL – reset call history; this time we supply the existing budget_id - mock_tx.litellm_budgettable.create.reset_mock() - mock_tx.litellm_teammembership.upsert.reset_mock() - mock_tx.litellm_budgettable.update.reset_mock() - - await _upsert_budget_and_membership( - mock_tx, - team_id="team-42", - user_id="user-42", - max_budget=25.0, - existing_budget_id=created_bid, # now used: triggers in-place update - user_api_key_dict=fake_user, - ) - - # Should update the existing budget in-place, not create a new one - mock_tx.litellm_budgettable.update.assert_awaited_once_with( - where={"budget_id": created_bid}, - data={ - "max_budget": 25.0, - "updated_by": fake_user.user_id, - }, - ) - - # Should NOT create a new budget or touch membership - mock_tx.litellm_budgettable.create.assert_not_called() - mock_tx.litellm_teammembership.upsert.assert_not_called() - - -# TEST: update rpm_limit for member with existing budget_id → updates in-place -@pytest.mark.asyncio -async def test_upsert_rpm_limit_update_creates_new_budget(mock_tx, fake_user): - """ - Test that updating rpm_limit for a member with an existing budget_id - updates the existing budget in-place (not creates a new one). - """ - existing_budget_id = "existing-budget-456" - - await _upsert_budget_and_membership( - mock_tx, - team_id="team-rpm-test", - user_id="user-rpm-test", - max_budget=50.0, - existing_budget_id=existing_budget_id, - user_api_key_dict=fake_user, - tpm_limit=1000, - rpm_limit=100, - ) - - # Should update the existing budget with all specified limits - mock_tx.litellm_budgettable.update.assert_awaited_once_with( - where={"budget_id": existing_budget_id}, - data={ - "max_budget": 50.0, - "tpm_limit": 1000, - "rpm_limit": 100, - "updated_by": fake_user.user_id, - }, - ) - - # Should NOT create a new budget or touch membership - mock_tx.litellm_budgettable.create.assert_not_called() - mock_tx.litellm_teammembership.upsert.assert_not_called() - - -# TEST: create new budget with only rpm_limit (no max_budget) -@pytest.mark.asyncio -async def test_upsert_rpm_only_creates_new_budget(mock_tx, fake_user): - """ - Test that setting only rpm_limit creates a new budget with just the rpm_limit. - """ - await _upsert_budget_and_membership( - mock_tx, - team_id="team-rpm-only", - user_id="user-rpm-only", - max_budget=None, - existing_budget_id=None, - user_api_key_dict=fake_user, - rpm_limit=50, - ) - - # Should create a new budget with only rpm_limit - mock_tx.litellm_budgettable.create.assert_awaited_once_with( - data={ - "rpm_limit": 50, - "created_by": fake_user.user_id, - "updated_by": fake_user.user_id, - }, - include={"team_membership": True}, - ) - - # Should upsert team membership with the new budget ID new_budget_id = mock_tx.litellm_budgettable.create.return_value.budget_id mock_tx.litellm_teammembership.upsert.assert_awaited_once_with( - where={ - "user_id_team_id": {"user_id": "user-rpm-only", "team_id": "team-rpm-only"} - }, + where={"user_id_team_id": {"user_id": "user-new", "team_id": "team-new"}}, data={ "create": { - "user_id": "user-rpm-only", - "team_id": "team-rpm-only", + "user_id": "user-new", + "team_id": "team-new", "litellm_budget_table": {"connect": {"budget_id": new_budget_id}}, }, "update": { @@ -270,60 +219,48 @@ async def test_upsert_rpm_only_creates_new_budget(mock_tx, fake_user): ) -# TEST: clone-on-write when membership still points at the team's shared default budget +# TEST: clone-on-write when the membership still points at the team's shared +# default budget. Editing this member must fork a private budget instead of +# mutating the shared row, and cloning a duration must seed a fresh reset time. @pytest.mark.asyncio -async def test_upsert_clones_when_pointing_at_shared_default(mock_tx, fake_user): - """ - When a member's existing budget_id is the same row as the team's shared - default member budget, updating that member's budget must NOT mutate the - shared row. Instead we should create a new private budget for this member - (seeded with the default's values) and re-link the membership to it. - """ +async def test_clone_on_write_from_shared_default(mock_tx, fake_user): shared_default_id = "team-default-budget-1" + mock_tx.litellm_budgettable.find_unique = AsyncMock( + return_value=budget_row( + budget_id=shared_default_id, + max_budget=200.0, + soft_budget=None, + max_parallel_requests=None, + tpm_limit=500, + rpm_limit=None, + model_max_budget=None, + budget_duration="1d", + allowed_models=[], + ) + ) - # Default budget row in the DB: $200 cap, daily reset, 500 tpm. - default_row = MagicMock() - default_row.model_dump.return_value = { - "budget_id": shared_default_id, - "max_budget": 200.0, - "soft_budget": None, - "max_parallel_requests": None, - "tpm_limit": 500, - "rpm_limit": None, - "model_max_budget": None, - "budget_duration": "1d", - "allowed_models": [], - } - mock_tx.litellm_budgettable.find_unique = AsyncMock(return_value=default_row) - - # Caller is changing only this member's max_budget. await _upsert_budget_and_membership( mock_tx, team_id="team-shared", user_id="user-shared", - max_budget=50.0, existing_budget_id=shared_default_id, user_api_key_dict=fake_user, + budget_patch={"max_budget": 50.0}, team_default_budget_id=shared_default_id, ) - # Must NOT touch the shared default row in place. mock_tx.litellm_budgettable.update.assert_not_called() + mock_tx.litellm_budgettable.create.assert_awaited_once() + create_data = mock_tx.litellm_budgettable.create.await_args.kwargs["data"] + assert_future_reset_time(create_data.pop("budget_reset_at")) + assert create_data == { + "created_by": fake_user.user_id, + "updated_by": fake_user.user_id, + "max_budget": 50.0, # caller wins + "tpm_limit": 500, # cloned from default + "budget_duration": "1d", # cloned from default + } - # Must create a new private budget seeded with the default's values, - # with the caller's max_budget overriding the cloned default. - mock_tx.litellm_budgettable.create.assert_awaited_once_with( - data={ - "created_by": fake_user.user_id, - "updated_by": fake_user.user_id, - "max_budget": 50.0, # caller wins - "tpm_limit": 500, # cloned from default - "budget_duration": "1d", # cloned from default - }, - include={"team_membership": True}, - ) - - # Membership must be re-linked to the new private budget. new_budget_id = mock_tx.litellm_budgettable.create.return_value.budget_id mock_tx.litellm_teammembership.upsert.assert_awaited_once_with( where={"user_id_team_id": {"user_id": "user-shared", "team_id": "team-shared"}}, @@ -340,32 +277,64 @@ async def test_upsert_clones_when_pointing_at_shared_default(mock_tx, fake_user) ) -# TEST: when team default exists but member already has their own budget, in-place update +# TEST: forking the shared default while clearing its duration must drop the +# duration (and not carry a reset time) on the new private budget. @pytest.mark.asyncio -async def test_upsert_updates_in_place_when_member_has_private_budget( - mock_tx, fake_user -): - """ - If the member's budget_id is different from the team's shared default - (i.e. they already have a private budget), we should keep the current - in-place behavior and not allocate a new row. - """ +async def test_clone_on_write_clears_duration(mock_tx, fake_user): + shared_default_id = "team-default-budget-1" + mock_tx.litellm_budgettable.find_unique = AsyncMock( + return_value=budget_row( + budget_id=shared_default_id, + max_budget=200.0, + tpm_limit=500, + budget_duration="1d", + allowed_models=[], + ) + ) + + await _upsert_budget_and_membership( + mock_tx, + team_id="team-shared", + user_id="user-shared", + existing_budget_id=shared_default_id, + user_api_key_dict=fake_user, + budget_patch={"budget_duration": None}, + team_default_budget_id=shared_default_id, + ) + + mock_tx.litellm_budgettable.update.assert_not_called() + create_data = mock_tx.litellm_budgettable.create.await_args.kwargs["data"] + assert create_data == { + "created_by": fake_user.user_id, + "updated_by": fake_user.user_id, + "max_budget": 200.0, + "tpm_limit": 500, + "budget_duration": None, + } + assert "budget_reset_at" not in create_data + + +# TEST: when the member already has their own private budget (different from the +# team default), we update it in place rather than forking another row. +@pytest.mark.asyncio +async def test_private_budget_updates_in_place(mock_tx, fake_user): + mock_tx.litellm_budgettable.find_unique = AsyncMock( + return_value=budget_row(max_budget=10.0) + ) + await _upsert_budget_and_membership( mock_tx, team_id="team-mixed", user_id="user-private", - max_budget=75.0, existing_budget_id="private-budget-xyz", user_api_key_dict=fake_user, + budget_patch={"max_budget": 75.0}, team_default_budget_id="team-default-budget-1", ) mock_tx.litellm_budgettable.update.assert_awaited_once_with( where={"budget_id": "private-budget-xyz"}, - data={ - "max_budget": 75.0, - "updated_by": fake_user.user_id, - }, + data={"max_budget": 75.0, "updated_by": fake_user.user_id}, ) mock_tx.litellm_budgettable.create.assert_not_called() mock_tx.litellm_teammembership.upsert.assert_not_called() diff --git a/tests/test_litellm/proxy/test_team_member_update.py b/tests/test_litellm/proxy/test_team_member_update.py index 6561ec9e7fd..352c68d491c 100644 --- a/tests/test_litellm/proxy/test_team_member_update.py +++ b/tests/test_litellm/proxy/test_team_member_update.py @@ -1,9 +1,19 @@ +import types +from unittest.mock import AsyncMock, MagicMock + import pytest from fastapi import HTTPException from starlette.requests import Request import litellm.proxy.proxy_server as proxy_server -from litellm.proxy._types import TeamMemberUpdateRequest +import litellm.proxy.management_endpoints.team_endpoints as team_endpoints +from litellm.proxy._types import ( + LiteLLM_TeamTable, + LitellmUserRoles, + Member, + TeamMemberUpdateRequest, + UserAPIKeyAuth, +) from litellm.proxy.management_endpoints.team_endpoints import team_member_update @@ -38,3 +48,133 @@ async def test_ateam_member_update_admin_requires_premium(monkeypatch): "Pricing: https://www.litellm.ai/#pricing" ) assert exc_info.value.detail == expected_msg + + +@pytest.fixture +def happy_path_upsert(monkeypatch): + """Stub out the DB and the budget upsert so a team_member_update call reaches + _upsert_budget_and_membership, and hand back that mock to inspect the patch.""" + team_row = LiteLLM_TeamTable( + team_id="team-1234", + members_with_roles=[Member(user_id="user-1", role="user")], + metadata={}, + ) + + prisma_client = MagicMock() + prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + prisma_client.db.litellm_teamtable.update = AsyncMock() + + class _FakeTx: + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + return False + + prisma_client.db.tx = MagicMock(return_value=_FakeTx()) + + monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) + monkeypatch.setattr(proxy_server, "premium_user", False) + monkeypatch.setattr( + team_endpoints, + "team_info", + AsyncMock( + return_value={ + "team_info": team_row, + "team_memberships": [ + types.SimpleNamespace(user_id="user-1", budget_id="bud-1") + ], + } + ), + ) + upsert_mock = AsyncMock() + monkeypatch.setattr(team_endpoints, "_upsert_budget_and_membership", upsert_mock) + return upsert_mock + + +def _member_update_request(**overrides): + data = TeamMemberUpdateRequest( + team_id="team-1234", user_id="user-1", role="user", **overrides + ) + request = Request({"type": "http", "method": "POST", "path": "/team/member_update"}) + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN.value, user_id="admin") + return data, request, auth + + +@pytest.mark.asyncio +async def test_team_member_update_sends_provided_fields_as_patch(happy_path_upsert): + """Fields the request sets must reach _upsert_budget_and_membership as a + budget patch, otherwise the member budget is never written/reset.""" + data, request, auth = _member_update_request( + max_budget_in_team=10.0, budget_duration="30d" + ) + + response = await team_member_update(data, request, auth) + + happy_path_upsert.assert_awaited_once() + assert happy_path_upsert.await_args.kwargs["budget_patch"] == { + "max_budget": 10.0, + "budget_duration": "30d", + } + assert response.budget_duration == "30d" + + +@pytest.mark.asyncio +async def test_team_member_update_explicit_null_clears_field(happy_path_upsert): + """An explicitly-null field must be forwarded as None so the column is + cleared, rather than silently dropped.""" + data, request, auth = _member_update_request(budget_duration=None) + + await team_member_update(data, request, auth) + + assert happy_path_upsert.await_args.kwargs["budget_patch"] == { + "budget_duration": None + } + + +@pytest.mark.asyncio +async def test_team_member_update_omits_unset_fields_from_patch(happy_path_upsert): + """A request that touches no budget fields must produce an empty patch so the + member's existing budget is left untouched.""" + data, request, auth = _member_update_request() + + await team_member_update(data, request, auth) + + assert happy_path_upsert.await_args.kwargs["budget_patch"] == {} + + +@pytest.mark.parametrize( + "bad_duration", + [ + "not-a-duration", # unparseable garbage + "10x", # unsupported unit + "0d", # zero-length window + "999999999999999999999999d", # overflows datetime math + ], +) +@pytest.mark.asyncio +async def test_team_member_update_rejects_invalid_budget_duration( + monkeypatch, bad_duration +): + """An invalid budget_duration must be rejected with a 400 before any DB + write, so it can never be persisted and later break the budget reset job.""" + monkeypatch.setattr(proxy_server, "prisma_client", object()) + monkeypatch.setattr(proxy_server, "premium_user", False) + upsert_mock = AsyncMock() + monkeypatch.setattr(team_endpoints, "_upsert_budget_and_membership", upsert_mock) + + data = TeamMemberUpdateRequest( + team_id="team-1234", + user_id="user-1", + role="user", + budget_duration=bad_duration, + ) + request = Request({"type": "http", "method": "POST", "path": "/team/member_update"}) + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN.value, user_id="admin") + + with pytest.raises(HTTPException) as exc_info: + await team_member_update(data, request, auth) + + assert exc_info.value.status_code == 400 + assert "budget_duration" in str(exc_info.value.detail) + upsert_mock.assert_not_called() diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 80f4d0b57f1..d303734dffd 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -2822,6 +2822,7 @@ export interface Member { max_budget_in_team?: number | null; tpm_limit?: number | null; rpm_limit?: number | null; + budget_duration?: string | null; allowed_models?: string[] | null; } @@ -2949,18 +2950,21 @@ export const teamMemberUpdateCall = async ( user_id: formValues.user_id, }; - // Add optional budget and rate limit fields + const orNull = (value: unknown) => (value === undefined || value === null || value === "" ? null : value); if (formValues.user_email !== undefined) { requestBody.user_email = formValues.user_email; } - if (formValues.max_budget_in_team !== undefined && formValues.max_budget_in_team !== null) { - requestBody.max_budget_in_team = formValues.max_budget_in_team; + if ("max_budget_in_team" in formValues) { + requestBody.max_budget_in_team = orNull(formValues.max_budget_in_team); } - if (formValues.tpm_limit !== undefined && formValues.tpm_limit !== null) { - requestBody.tpm_limit = formValues.tpm_limit; + if ("tpm_limit" in formValues) { + requestBody.tpm_limit = orNull(formValues.tpm_limit); } - if (formValues.rpm_limit !== undefined && formValues.rpm_limit !== null) { - requestBody.rpm_limit = formValues.rpm_limit; + if ("rpm_limit" in formValues) { + requestBody.rpm_limit = orNull(formValues.rpm_limit); + } + if ("budget_duration" in formValues) { + requestBody.budget_duration = orNull(formValues.budget_duration); } if (formValues.allowed_models !== undefined) { requestBody.allowed_models = formValues.allowed_models; diff --git a/ui/litellm-dashboard/src/components/team/EditMembership.tsx b/ui/litellm-dashboard/src/components/team/EditMembership.tsx index af5f8631ae4..16e4ec58dd0 100644 --- a/ui/litellm-dashboard/src/components/team/EditMembership.tsx +++ b/ui/litellm-dashboard/src/components/team/EditMembership.tsx @@ -2,6 +2,7 @@ import { Text, TextInput } from "@tremor/react"; import { Button as AntButton, Form, Modal, Select } from "antd"; import React, { useEffect, useState } from "react"; import NumericalInput from "../shared/numerical_input"; +import BudgetDurationDropdown from "../common_components/budget_duration_dropdown"; interface BaseMember { user_email?: string; @@ -21,7 +22,7 @@ interface ModalConfig { additionalFields?: Array<{ name: string; label: string | React.ReactNode; - type: "input" | "select" | "numerical" | "multi-select"; + type: "input" | "select" | "numerical" | "multi-select" | "budget-duration"; options?: Array<{ label: string; value: string }>; rules?: any[]; step?: number; @@ -65,6 +66,7 @@ const MemberModal = ({ max_budget_in_team: (initialData as any).max_budget_in_team || null, tpm_limit: (initialData as any).tpm_limit || null, rpm_limit: (initialData as any).rpm_limit || null, + budget_duration: (initialData as any).budget_duration || null, // Keep array values for multi-select fields allowed_models: (initialData as any).allowed_models || [], }; @@ -117,7 +119,7 @@ const MemberModal = ({ const renderField = (field: { name: string; label: string | React.ReactNode; - type: "input" | "select" | "numerical" | "multi-select"; + type: "input" | "select" | "numerical" | "multi-select" | "budget-duration"; options?: Array<{ label: string; value: string }>; rules?: any[]; step?: number; @@ -155,6 +157,8 @@ const MemberModal = ({ allowClear /> ); + case "budget-duration": + return ; default: return null; } diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index 4943b1fbac2..602490a0c98 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -388,6 +388,7 @@ const TeamInfoView: React.FC = ({ max_budget_in_team: values.max_budget_in_team, tpm_limit: values.tpm_limit, rpm_limit: values.rpm_limit, + budget_duration: values.budget_duration, allowed_models: values.allowed_models, }; MessageManager.destroy(); // Remove all existing toasts @@ -1689,6 +1690,18 @@ const TeamInfoView: React.FC = ({ min: 0, placeholder: "Budget limit for this member within this team", }, + { + name: "budget_duration", + label: ( + + Budget Reset Period{" "} + + + + + ), + type: "budget-duration" as const, + }, { name: "tpm_limit", label: ( diff --git a/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx b/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx index 5a290a6f6d4..e2f108dcbf5 100644 --- a/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx @@ -210,6 +210,7 @@ export default function TeamMemberTab({ max_budget_in_team: membership?.litellm_budget_table?.max_budget || null, tpm_limit: membership?.litellm_budget_table?.tpm_limit || null, rpm_limit: membership?.litellm_budget_table?.rpm_limit || null, + budget_duration: membership?.litellm_budget_table?.budget_duration || null, allowed_models: membership?.litellm_budget_table?.allowed_models || [], }; setSelectedEditMember(enhancedMember); diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index d14ec8af56b..936747acd08 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -27574,6 +27574,11 @@ export interface components { * @description List of models this team member can access. Pass an empty list to remove per-member model restrictions. */ allowed_models?: string[] | null; + /** + * Budget Duration + * @description Duration after which this team member's budget resets (e.g. '1h', '24h', '7d', '30d'). If not set, the budget never resets. + */ + budget_duration?: string | null; /** Max Budget In Team */ max_budget_in_team?: number | null; /** Role */ @@ -27599,6 +27604,8 @@ export interface components { TeamMemberUpdateResponse: { /** Allowed Models */ allowed_models?: string[] | null; + /** Budget Duration */ + budget_duration?: string | null; /** Max Budget In Team */ max_budget_in_team?: number | null; /** Rpm Limit */ From 7bfce053a9b7f0732aab32b4e055a6d802468c88 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 6 Jun 2026 17:41:36 -0700 Subject: [PATCH 005/185] fix(ui): make workflow runs page fill full width (#29868) The Workflow Runs page rendered its table at roughly a quarter of the available width. Its root container is a flex child of the dashboard content row but set only padding, min-height and background, so with no width it shrank to the table's natural content size. Sibling pages (logs, memory) fill the area with a full-width root; mirror that by setting width 100% on the container. Fixes LIT-3636 --- ui/litellm-dashboard/src/components/workflow_runs/index.tsx | 1 + 1 file changed, 1 insertion(+) diff --git a/ui/litellm-dashboard/src/components/workflow_runs/index.tsx b/ui/litellm-dashboard/src/components/workflow_runs/index.tsx index b9467c5604b..2aecbece2e3 100644 --- a/ui/litellm-dashboard/src/components/workflow_runs/index.tsx +++ b/ui/litellm-dashboard/src/components/workflow_runs/index.tsx @@ -587,6 +587,7 @@ const WorkflowRuns: React.FC = ({ accessToken }) => { return (
Date: Sat, 6 Jun 2026 17:50:29 -0700 Subject: [PATCH 006/185] feat: standardize rate limit errors with category, rate_limit_type, model, and llm_provider fields (#27687) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(exceptions): add RateLimitErrorCategory + headers/detail fields on RateLimitError LiteLLM previously surfaced rate-limit conditions through several unrelated error classes (RateLimitError, FastAPI HTTPException(429), BaseLLMException). This commit adds the data model needed to consolidate them under a single class: * RateLimitErrorCategory enum exposing four categorical values (vendor_rate_limit, vendor_batch_rate_limit, litellm_rate_limit, litellm_batch_rate_limit) so callers can switch on the rate-limit source. * New optional fields on RateLimitError: - category (defaults to vendor_rate_limit, preserving today's behavior for every existing call site in exception_mapping_utils); - headers (preserves retry-after / rate_limit_type / reset_at across the proxy boundary instead of dropping them on the floor); - detail (mirrors FastAPI HTTPException.detail so the same instance can be serialized through both paths). litellm.RateLimitErrorCategory is re-exported at the package root to match the existing exception-export pattern. LIT-2968 Co-authored-by: Mateo Wang * feat(proxy): add ProxyRateLimitError unifying RateLimitError + HTTPException Adds a single proxy-side error class that subclasses BOTH litellm.exceptions.RateLimitError AND fastapi.HTTPException via cooperative multiple inheritance. Why both bases: * Subclassing RateLimitError lets user code catch every rate-limit source with one 'except RateLimitError' and switch on the new .category field. * Subclassing HTTPException keeps every existing FastAPI plumbing path (the isinstance(e, HTTPException) branches in proxy_server.py route handlers, FastAPI's own dispatcher, and tests asserting pytest.raises(HTTPException)) working without modification, and preserves retry-after / rate_limit_type / reset_at headers on the wire. The class declaration order is (HTTPException, RateLimitError) so the MRO puts HTTPException's no-super-call __init__ ahead of openai's cooperative __init__ chain — preventing openai.APIError.super().__init__(message) from landing in HTTPException.__init__(status_code=message). LIT-2968 Co-authored-by: Mateo Wang * refactor(proxy/hooks): raise ProxyRateLimitError from budget + iteration limiters Replaces three bare HTTPException(status_code=429, ...) call sites with ProxyRateLimitError, which is both a RateLimitError (catchable by category) and an HTTPException (preserves existing FastAPI serialization). Drops the now-unused HTTPException import in the iteration / per-session limiters. LIT-2968 Co-authored-by: Mateo Wang * refactor(proxy/hooks): raise ProxyRateLimitError from parallel-request limiters Replaces HTTPException(status_code=429, ...) call sites in the v1 and v3 parallel-request limiters (key/team/user/model/customer rate limits) with ProxyRateLimitError. Updates the raise_rate_limit_error helper's return type annotation accordingly. LIT-2968 Co-authored-by: Mateo Wang * refactor(proxy/hooks): raise ProxyRateLimitError from dynamic rate limiters Replaces HTTPException(status_code=429, ...) call sites in the v1 and v3 dynamic rate limiters (project-level TPM/RPM allocation, model-saturation checks, priority-based limits, fail-closed guards) with ProxyRateLimitError. The v3 limiter still imports HTTPException for an unrelated bare 'except HTTPException:' branch. LIT-2968 Co-authored-by: Mateo Wang * refactor(proxy/hooks): raise ProxyRateLimitError from batch rate limiter Replaces HTTPException(status_code=429, ...) in batch_rate_limiter._raise_rate_limit_error with ProxyRateLimitError tagged as RateLimitErrorCategory.LITELLM_BATCH_RATE_LIMIT so users can distinguish batch-level throttling (which counts requests/tokens across an uploaded batch input file before submission) from the generic key/team/user RPM/TPM limiter. The HTTPException import is retained because the same module raises HTTPException for unrelated 403/IO error paths. LIT-2968 Co-authored-by: Mateo Wang * test(rate-limit): pin down unified rate-limit error contract Adds a dedicated test module covering the new RateLimitErrorCategory enum, RateLimitError.category default + override behavior, ProxyRateLimitError's dual nature (RateLimitError + HTTPException), and a parametrized regression guard that asserts every proxy hook module imports the unified class. The regression guard catches the failure mode the refactor is designed to prevent: someone re-introducing a bare HTTPException(status_code=429, ...) in one of the hook modules instead of going through ProxyRateLimitError. LIT-2968 Co-authored-by: Mateo Wang * feat(logging): expose rate-limit category via StandardLoggingPayload Adds an optional 'error_rate_limit_category' field to StandardLoggingPayloadErrorInformation, populated from the unified RateLimitError.category attribute (introduced in the previous commits on this branch). Why: the .category attribute is reachable off the raw exception today via getattr(e, 'category', None), but the structured contract that downstream custom callbacks / loggers / spend log writers consume is the StandardLoggingPayload. Without this field, a user building custom rate-limit metrics on top of callback data has to special-case the raw exception object — which defeats the purpose of the StandardLoggingPayload abstraction. The field is None for non-rate-limit exceptions (so consumers can read it unconditionally without isinstance checks) and is one of the RateLimitErrorCategory string values otherwise. LIT-2968 Co-authored-by: Mateo Wang * test(rate-limit): assert StandardLoggingPayload carries the category Five tests covering: vendor default, explicit litellm_rate_limit and litellm_batch_rate_limit values, None for non-rate-limit exceptions, and None when no exception is provided. Pins down the contract that custom callbacks can read 'error_information.error_rate_limit_category' off the StandardLoggingPayload to drive custom rate-limit metrics without ever reaching for the raw exception. LIT-2968 Co-authored-by: Mateo Wang * fix(types): silence mypy [misc] on intentional dual-base attr overlap mypy emits two [misc] errors on the ProxyRateLimitError class line because its two bases declare overlapping attributes with related-but-not-identical annotations: * status_code: int on starlette HTTPException vs. Literal[429] on openai's RateLimitError (every openai status-error subclass narrows it the same way and silences pyright with the same convention). * headers: Mapping[str, str] | None on HTTPException vs. our Optional[ Dict[str, str]] (the proxy hooks always carry a stringified dict). Both narrowings are intentional and enforced at construction time. Add a type: ignore[misc] with an inline explanation rather than relax the annotations on the parent or change the wire-format guarantees. LIT-2968 Co-authored-by: Mateo Wang * test(rate-limit): add direct hook-invocation tests to lift patch coverage Adds six end-to-end tests that drive each refactored hook past its limit and assert the unified ProxyRateLimitError is raised with the correct category and dual-base shape. Complements the import-shape-only parametrized guard above by actually executing the new 'raise ProxyRateLimitError(...)' lines so codecov's patch coverage sees them as hit. Hooks covered (one test each): * parallel_request_limiter v1 — direct call to raise_rate_limit_error() * parallel_request_limiter v3 — direct call to _handle_rate_limit_error with a fabricated OVER_LIMIT response * max_iterations_limiter — full async_pre_call_hook with mocked agent registry, second call exceeds budget=1 * max_budget_limiter — async_pre_call_hook with mocked get_current_spend * dynamic_rate_limiter v1 — async_pre_call_hook with mocked check_available_usage forcing available_tpm == 0 * batch_rate_limiter — direct _raise_rate_limit_error call, asserts category is the batch-specific LITELLM_BATCH_RATE_LIMIT (not the generic LITELLM_RATE_LIMIT) LIT-2968 Co-authored-by: Mateo Wang * fix: guard rate_limit_category extraction with isinstance check * test(rate-limit): cover remaining hook raise sites for codecov Adds five more direct hook-invocation tests so every PR-touched line in the proxy hooks is exercised by tests in tests/test_litellm/, which codecov measures: * parallel_request_limiter v1 — check_key_in_limits inline raise (the second raise site, separate from the raise_rate_limit_error helper covered earlier) * dynamic_rate_limiter v1 — RPM raise branch (TPM branch was already covered) * dynamic_rate_limiter v3 — parametrized over all three raise sites: model_saturation_check, priority_model, and the fail-closed fallback for an unrecognized descriptor_key * max_budget_per_session_limiter — full async_pre_call_hook with a mocked agent registry and over-budget cached spend All 42 tests in test_rate_limit_error_unification.py now pass and together exercise every changed import + raise line across the eight refactored proxy hooks. LIT-2968 Co-authored-by: Mateo Wang * fix: use computed error_message in ProxyRateLimitError detail * fix(parallel-request-limiter): drop None from detail; annotate raise_rate_limit_error as NoReturn The v1 ' raise_rate_limit_error' helper built an unused 'error_message' variable and then assembled the actual ' detail' via an f-string that interpolated 'additional_details' verbatim — producing 'Max parallel request limit reached None' when invoked without arguments (flagged by code review). Fix the helper to: - use the constructed 'error_message' as the detail - annotate the helper as NoReturn since it always raises - drop the redundant 'raise'/'return' at the two call sites Add two regression tests covering both the with- and without- additional_details paths. LIT-2968 Co-authored-by: Mateo Wang * fix(proxy/hooks): drop literal 'None' from raise_rate_limit_error detail The v1 parallel_request_limiter's raise_rate_limit_error helper has a long-standing bug: it computes a None-guarded 'error_message' string but then ignores it and emits an f-string that interpolates the raw 'additional_details' arg. Callers that pass no argument get 'Max parallel request limit reached None' as the user-facing detail. This commit: * wires error_message into the detail kwarg so the None-guard actually applies and operators see a clean message; * changes the return-type annotation from ProxyRateLimitError to NoReturn (the function always raises) so type-checkers know callers after this invocation are unreachable. Greptile P1 + P2 review feedback on PR #27687. LIT-2968 Co-authored-by: Mateo Wang * fix(types): demote TypedDict floating string to a # comment A string literal placed after a field declaration in a TypedDict body is not a per-field docstring — it's an orphaned string expression Python discards. Tools like mypy / pyright that inspect TypedDict fields won't surface that text either. Move the documentation for error_rate_limit_category to a real comment so the intent is visible to readers and type-checker tooling without the misleading docstring framing. Greptile P2 review feedback on PR #27687. LIT-2968 Co-authored-by: Mateo Wang * security(exceptions): do not auto-copy vendor response headers to e.headers A vendor 429 response can set arbitrary headers (Set-Cookie, CORS overrides, …). Previously, when RateLimitError was constructed with only a 'response=' (no explicit 'headers=' kwarg), self.headers fell back to a copy of response.headers. If a downstream proxy serializer ever forwarded e.headers to the client, a malicious upstream could inject browser-interpreted headers for the proxy origin. Drop the fallback. Only headers passed explicitly via the headers= kwarg make it onto self.headers (proxy hooks pass retry-after etc. — they control what's surfaced). Vendor response headers stay reachable on e.response.headers for callers that explicitly want them. Today's proxy_server.py route handlers don't actually forward e.headers on the wire (they construct ProxyException without passing headers), so no current behavior changes — this is a defensive narrowing so the fallback can never be turned into a vector when someone wires e.headers through later. Veria-AI security review feedback on PR #27687. LIT-2968 Co-authored-by: Mateo Wang * test(rate-limit): regression guards for review-pass fixes Pins down the three review-pass fixes: * test_parallel_request_limiter_v1_helper_no_additional_details — calls raise_rate_limit_error() with no args and asserts the detail does NOT contain the literal string 'None'. Pre-fix, callers got 'Max parallel request limit reached None'. * test_rate_limit_error_does_not_auto_copy_response_headers — passes a vendor httpx.Response with a Set-Cookie header to RateLimitError WITHOUT an explicit headers= kwarg, asserts self.headers stays None (no leak), then re-checks that an explicit headers= kwarg DOES populate self.headers. Vendor headers remain reachable on e.response.headers for callers that explicitly want them. * The existing v1-helper test now also asserts the additional_details string makes it through to the detail. LIT-2968 Co-authored-by: Mateo Wang * feat(rate-limit): add orthogonal RateLimitType (requests/tokens/concurrent_requests/budget/max_iterations) trho's last ask in the LIT-2968 thread: distinguish rate-limit failures by the dimension that was exceeded, not just by who rate-limited (vendor vs. litellm). Adds: - RateLimitType str-enum exposed at `litellm.RateLimitType` with values requests / tokens / concurrent_requests / budget / max_iterations. - `rate_limit_type` kwarg on litellm.RateLimitError + ProxyRateLimitError; None default so existing callers (vendor-429 path in exception_mapping_utils) remain a no-op. - StandardLoggingPayloadErrorInformation.error_rate_limit_type so custom callbacks can split rate-limit failures by cause without parsing free-text error messages. Mirror to error_rate_limit_category extraction in get_error_information(); single isinstance(RateLimitError) check covers both. - map_v3_rate_limit_type() helper to collapse the v3 limiter's internal labels ("requests", "tokens", "max_parallel_requests") onto the public enum so the v3 limiter and dynamic_rate_limiter_v3 share one mapping. Defensive None on unknown values rather than silently picking a wrong dimension. Co-authored-by: Mateo Wang * feat(proxy/hooks): wire rate_limit_type onto every limiter raise site Each refactored proxy hook now populates rate_limit_type with the dimension that actually tripped the limit, so downstream consumers (custom callbacks, prometheus exporters via the StandardLoggingPayload) can split key/team/user rate-limit failures by cause: - parallel_request_limiter (v1): detect dimension from current vs. limit in the post-cache branch (concurrent_requests > tokens > requests, matches the boolean condition order). Base case (current is None, one limit set to 0) picks the most-specific zero. raise_rate_limit_error() helper accepts an explicit rate_limit_type kwarg with CONCURRENT_REQUESTS default (matches every existing internal call site, including the global-limit branch). - parallel_request_limiter (v3): forward status["rate_limit_type"] through map_v3_rate_limit_type() so "max_parallel_requests" → CONCURRENT_REQUESTS for the public field while the raw v3 jargon stays on the HTTP header for wire-format backward compat. - dynamic_rate_limiter (v1): TPM-zero → TOKENS, RPM-zero → REQUESTS. Pass data["model"] through so callbacks see the model that hit the limit (addresses the secondary "provider missing" complaint in the original Slack thread, partially — the model is what dashboards typically split on). - dynamic_rate_limiter (v3): forward status["rate_limit_type"] via map_v3_rate_limit_type() at every raise site (model_saturation_check, priority_model, fail-closed unknown-descriptor guard). Also pass model. - batch_rate_limiter: limit_type is hard-typed "requests"|"tokens" — map directly without going through the helper's None branch. - max_budget_limiter, max_budget_per_session_limiter: BUDGET. - max_iterations_limiter: MAX_ITERATIONS. Co-authored-by: Mateo Wang * test(rate-limit): cover RateLimitType enum, hook wiring, and StandardLoggingPayload propagation 27 new tests across five new test classes: - TestRateLimitType: enum exposed at litellm.RateLimitType, all five values defined, RateLimitError default is None (vendor 429 path makes no claim about which dimension), accepts both string and enum forms with str-coercion guarantee for downstream JSON serializers. - TestProxyRateLimitErrorType: ProxyRateLimitError default is None, accepts string or enum, doesn't break existing callers that pass nothing. - TestMapV3RateLimitType: pins each v3-internal → public-enum mapping (tokens, requests, max_parallel_requests → concurrent_requests, unknown → None) so a future v3 refactor can't silently swap dimensions. - TestStandardLoggingPayloadCarriesType: the new error_rate_limit_type field reaches the structured payload for both ProxyRateLimitError and plain RateLimitError, is None when unspecified, and is None for non-rate-limit exceptions (symmetric with error_rate_limit_category). - TestProxyHooksWireTypeCorrectly: drives the actual raise sites in the v1 parallel_request_limiter helper, the v3 _handle_rate_limit_error (both "tokens" and "max_parallel_requests" paths), and the batch limiter (both tokens and requests paths) — coverage tools see the new rate_limit_type= kwargs as exercised, not just the import shape. Co-authored-by: Mateo Wang * test(rate-limit): cover _coerce_message branches and v1 dimension detection Drives the patch coverage on the new orthogonal RateLimitType wiring up to (or close to) 100% on the touched files. ProxyRateLimitError._coerce_message — was 22% covered, now 100%: * nested {error: {message}} dict * nested {message: {message}} dict (alt key) * dict without 'error'/'message' keys → JSON dump fallback * non-JSON-serializable dict value → str() fallback * non-string non-mapping detail (int) → str() coercion v1 parallel_request_limiter dimension detection — was 0% covered, now exercised across 6 parametrized cases: * check_key_in_limits else-branch: current at concurrent / TPM / RPM cap → asserts rate_limit_type is concurrent_requests / tokens / requests. * check_key_in_limits base case (current is None): max_parallel_requests / tpm_limit / rpm_limit set to 0 → asserts the most-specific zero attribution wins per the helper's order. LIT-2968 Co-authored-by: Mateo Wang * feat(proxy/hooks): add ProxyHTTPRateLimitError + provider resolver Introduces a small helper layer used by every proxy-side rate-limit hook so that the 429 they raise carries a populated llm_provider / model — instead of an empty exception.llm_provider that downstream loggers (Prometheus failure metric, observability callbacks) read as 'no provider attribution'. ProxyHTTPRateLimitError inherits from both fastapi.HTTPException (so the proxy server still renders it as a 429) and litellm.exceptions.RateLimitError (so isinstance checks and PrometheusLogger._get_exception_class_name pick up llm_provider). We deliberately don't call RateLimitError.__init__ — it constructs an httpx.Response we don't need and would just add failure surface; attribute parity is what downstream consumers care about. resolve_llm_provider_for_rate_limit() wraps litellm.get_llm_provider defensively. Internal limiter hooks fire from async_pre_call_hook — well before get_llm_provider runs anywhere else in the request lifecycle — so we have to call it ourselves at raise time. If the model is missing or unparseable (alias, router-only model) we fall back to llm_provider='litellm_proxy' rather than letting a second exception leak out and break the request path. Co-authored-by: Mateo Wang * fix(proxy/hooks): populate llm_provider on parallel-request 429s Both v1 and v3 parallel-request limiters fired bare HTTPException(429) from inside async_pre_call_hook. The downstream Prometheus failure metric reads exception.llm_provider via _get_exception_class_name — the empty value showed up as exception_class='HTTPException' and left model_id='None' on the time series. Threads requested_model through every raise site in: * parallel_request_limiter.py: - check_key_in_limits (the per-key/per-model/per-user/per-team/ per-customer over-limit path) - raise_rate_limit_error (zero-limit + global_max_parallel_requests paths) — now takes an optional requested_model kwarg * parallel_request_limiter_v3.py: - _handle_rate_limit_error (the OVER_LIMIT translator), called from both the should_rate_limit pre-check and the TPM reservation path Resolved via resolve_llm_provider_for_rate_limit so unknown / missing models silently fall back to llm_provider='litellm_proxy' instead of breaking the request path with a second exception. Co-authored-by: Mateo Wang * fix(proxy/hooks): populate llm_provider on dynamic-rate-limit 429s Same plumbing change as the parallel limiters, applied to both dynamic_rate_limiter (v1) and dynamic_rate_limiter_v3: * v1: TPM-zero and RPM-zero paths in async_pre_call_hook now resolve data['model'] -> (model, llm_provider) once and pass it into both raises. * v3: All three raise sites in _check_rate_limits — the model_saturation_check enforced raise, the priority_model enforced raise, and the fail-closed unknown-descriptor branch — now attribute the 429 to the actual provider. Falls back to llm_provider='litellm_proxy' when the model can't be resolved. Co-authored-by: Mateo Wang * fix(proxy/hooks): populate llm_provider on batch-rate-limit 429s batch_rate_limiter._raise_rate_limit_error now takes a requested_model kwarg threaded from data['model'] in _check_and_increment_batch_counters. The batch-creation 429 is what gets raised when the input file's tokens/requests count would push the per-key TPM/RPM window over its limit. Co-authored-by: Mateo Wang * fix(proxy/hooks): populate llm_provider on budget/iterations 429s Final batch of internal raise sites — the user/session-budget and max-iterations hooks. Same pattern: resolve data['model'] once at raise time, attach to ProxyHTTPRateLimitError so Prometheus and observability callbacks can attribute the 429. Hooks updated: * max_budget_limiter (per-user max_budget exceeded) * max_iterations_limiter (per-session agent iteration cap) * max_budget_per_session_limiter (per-session dollar cap) All three fall back to llm_provider='litellm_proxy' when data['model'] is missing or unparseable. Drops the now-unused HTTPException import from each module. Co-authored-by: Mateo Wang * test(proxy/hooks): pin provider field on internal rate-limit 429s Regression coverage for the 'provider field missing' bug across every proxy-side rate-limit hook + the helper layer: * ProxyHTTPRateLimitError class shape (HTTPException + RateLimitError, dict-detail stringification, None-provider normalization). * resolve_llm_provider_for_rate_limit happy paths (gpt-4o-mini, anthropic/..., bedrock/...) plus all three fallback branches (None, '', unknown name) plus a 'get_llm_provider raises' case that asserts we swallow the secondary exception. * For each limiter (parallel v1/v3, dynamic v1/v3, batch, max_budget, max_iterations, max_budget_per_session): assert the raised exception is a RateLimitError carrying the resolved model + llm_provider, and a sibling test that asserts the fallback path returns 'litellm_proxy' without leaking a second exception. * Two PrometheusLogger._get_exception_class_name pins so the Prometheus failure metric label flips from 'HTTPException' to 'Openai.ProxyHTTPRateLimitError' (or 'Litellm_proxy.*' on fallback) — that's what dashboards consume. Co-authored-by: Mateo Wang * perf(proxy/hooks): defer provider resolution to over-limit branches * fix: use error_message in raise_rate_limit_error to avoid literal 'None' in detail * Consolidate rate_limiter_utils imports in dynamic_rate_limiter * fix(proxy): set num_retries/max_retries on ProxyHTTPRateLimitError ProxyHTTPRateLimitError inherits from RateLimitError but did not call RateLimitError.__init__, so num_retries/max_retries were never set. When Starlette's HTTPException lacks __str__, MRO falls through to RateLimitError.__str__, which unconditionally reads these attributes and raises AttributeError during logging/traceback formatting. Initialize them to None defensively. * fix(mypy): silence base-class status_code conflict on ProxyHTTPRateLimitError HTTPException declares 'status_code: int' while openai.RateLimitError (via APIStatusError) declares 'status_code: Literal[429] = 429'. Mypy flags the multi-base override as [misc] in CI lint. The runtime semantics are fine (we set self.status_code in __init__), so silence the class-level annotation conflict with a targeted ignore. Co-authored-by: Mateo Wang * fix: annotate batch limiter _raise_rate_limit_error as NoReturn * feat(prometheus): rate-limit category/type labels + exception_class back-compat (follow-up to #27687) (#27706) * feat(prometheus): add rate_limit_category and rate_limit_type labels Adds two new labels to litellm_proxy_failed_requests_metric so dashboards can split 429s by rate-limit source (vendor vs. litellm-internal) and by the dimension that was exceeded (requests/tokens/concurrent_requests/ budget/max_iterations) without parsing free-text error messages. Closes the Prometheus side of LIT-2718. The unified RateLimitError.category and .rate_limit_type fields landed in PR #27687 but were only surfaced on StandardLoggingPayload (custom-callback channel); this exposes them on the metric label set as well. Both labels are populated only when the underlying exception is a litellm.RateLimitError; non-rate-limit failures keep them empty. Co-authored-by: Mateo Wang * feat(prometheus): populate rate-limit labels + preserve exception_class back-compat Two coupled changes in the Prometheus integration: 1. async_post_call_failure_hook now extracts the new RateLimitError .category / .rate_limit_type fields (added in PR #27687) via a _extract_rate_limit_labels helper and forwards them through UserAPIKeyLabelValues onto litellm_proxy_failed_requests_metric. Empty for non-rate-limit failures. 2. _get_exception_class_name special-cases ProxyRateLimitError and keeps emitting 'HTTPException' for the exception_class label. Without this shim, ProxyRateLimitError (which multi-inherits from HTTPException + RateLimitError) would silently flip the label from 'HTTPException' (the historical value for proxy-side 429s) to 'ProxyRateLimitError', breaking existing dashboards / alerts that key off exception_class='HTTPException'. Distinguishing vendor vs. litellm 429s is now the job of the new rate_limit_category label. Co-authored-by: Mateo Wang * test(prometheus): cover rate-limit labels and exception_class back-compat Adds 19 tests across: - enum / label-list registration - _extract_rate_limit_labels for vendor RateLimitError, ProxyRateLimitError, non-rate-limit and None inputs (incl. parametrized over every RateLimitErrorCategory x RateLimitType combo) - _get_exception_class_name back-compat: ProxyRateLimitError keeps the legacy 'HTTPException' string while vendor RateLimitError keeps the historical 'Provider.ClassName' format - end-to-end through async_post_call_failure_hook with both ProxyRateLimitError and vendor RateLimitError, asserting both new labels populate and exception_class stays back-compat Co-authored-by: Mateo Wang * fix(prometheus): tolerate missing fastapi in lazy ProxyRateLimitError import Address greptile feedback: - async_post_call_failure_hook docstring: drop the stale labelnames listing and reference PrometheusMetricLabels.litellm_proxy_failed_requests_metric as the source of truth so the doc cannot drift from the actual labelset. - _get_exception_class_name: guard the lazy ProxyRateLimitError import with ImportError so router-side fallback callsites don't blow up in non-proxy installs that don't have fastapi (a transitive dep of proxy.common_utils.proxy_rate_limit_error). Behavior is unchanged when fastapi is available. Also fix the existing enterprise callback test that asserted the old labelset on litellm_proxy_failed_requests_metric — it now expects the new rate_limit_category / rate_limit_type labels populated for vendor 429s. --------- Co-authored-by: Cursor Agent Co-authored-by: Mateo Wang * fix(bugbot): simplify rate-limit label coercion + guard None detail - prometheus.py _extract_rate_limit_labels: RateLimitError.__init__ already normalizes category/rate_limit_type to plain str, so the getattr(.value) + isinstance dance was dead code. Reduce to str(value) if not None. - proxy_rate_limit_error.py _coerce_message: short-circuit None to '' instead of falling through to str(None) = 'None', which produced the literal message 'litellm.RateLimitError: None'. * fix(rate-limit): surface unified category/type fields on BudgetExceededError The most common budget cap (virtual-key max_budget enforcement in auth_checks.py) raises litellm.BudgetExceededError, a bare Exception subclass that bypassed the unified rate-limit error class introduced by PR #27687. Custom callbacks reading StandardLoggingPayload.error_information saw category=None and rate_limit_type=None for these 429s, missing the most common budget case (team / org / end-user budgets all hit the same code path). Surface the fields off BudgetExceededError as plain attributes: - category = RateLimitErrorCategory.LITELLM_RATE_LIMIT - rate_limit_type = RateLimitType.BUDGET - llm_provider = "" (or caller-supplied) Switch get_error_information and _extract_rate_limit_labels from isinstance(RateLimitError) gating to duck-typed attribute reads, guarded by membership in the rate-limit enums so unrelated third-party exceptions exposing a .category attribute can't leak garbage values into the payload. This is strictly additive: BudgetExceededError keeps its bare-Exception base class, so `except BudgetExceededError:` handlers keep firing and `except RateLimitError:` does not start catching budget errors. * fix(rate-limit): validate enum membership at duck-typed read sites + enrich BudgetExceededError llm_provider Two follow-ups uncovered during the second QA pass on PR #27687: 1. Guard third-party `.category` / `.rate_limit_type` attribute leakage. The duck-typed read in `get_error_information` and `_extract_rate_limit_labels` would forward any string attribute named `category` / `rate_limit_type` on an unrelated third-party exception into the StandardLoggingPayload and Prometheus labels — silently mislabeling custom-callback payloads and blowing out Prometheus label cardinality. Add `validate_rate_limit_category` / `validate_rate_limit_type` helpers that gate on the documented enum value sets; non-matching values are dropped to None. 2. Enrich BudgetExceededError.llm_provider from request_data. Budget checks live in tenant-scoped helpers (key / team / org / tag / end-user / project) that don't see the request model, so the BudgetExceededError they raise carried llm_provider="" — leaving custom-metrics consumers without provider attribution for the most common 429 case. Resolve it once at the central UserAPIKeyAuthExceptionHandler seam, before post_call_failure_hook fires, so the StandardLoggingPayload the callback sees has the same provider attribution as RPM/TPM 429s. Regression tests pin both: 4 leakage tests + 4 enrichment tests. The leakage tests would fail under the pre-validation version of either read site; the enrichment tests would fail if the handler skipped the resolver call. * fix(rate-limit): resolve router model_name aliases to real provider (#27914) * fix(rate-limit): resolve router model_name aliases to real provider For nearly every real LiteLLM proxy deployment the request model is a router model_name alias (e.g. 'tpm-locked' -> litellm_params.model: openai/gpt-4o-mini), and 'litellm.get_llm_provider' doesn't know about router aliases — it raises 'LLMProviderNotProvidedError'. The resolver then fell through to the defensive 'litellm_proxy' fallback, so the 'llm_provider' field this PR adds was effectively always 'litellm_proxy' in the field, defeating its purpose for the most common proxy configuration. Add a router-alias fallback step: when 'get_llm_provider' raises, scan the active 'llm_router.model_list' for a deployment whose 'model_name' matches the request model and resolve from its 'litellm_params.model' instead. If multiple deployments share the same alias (load-balancing case) the first one wins — every deployment under one alias should agree on provider in any sensible config, and 'first' is deterministic so the Prometheus label stays stable. Defensive throughout: an uninitialized router, a malformed deployment, a 'litellm_params.model' that itself fails 'get_llm_provider' — every branch falls through to the existing 'litellm_proxy' fallback rather than letting a secondary exception escape and mask the rate-limit error we're trying to surface. Tests: - test_router_alias_resolves_to_underlying_provider: alias 'tpm-locked' -> 'openai/gpt-4o-mini' produces provider='openai', model='gpt-4o-mini'. - test_router_alias_with_multiple_deployments_uses_first. - test_router_alias_unknown_falls_back. - test_router_alias_with_malformed_deployment_falls_back. - Existing fallback test updated to also stub 'litellm.proxy.proxy_server.llm_router' so it exercises the full 'no resolution anywhere' path. Co-authored-by: Mateo Wang * fix(rate-limit): harden router alias resolver + test isolation - Wrap _resolve_provider_from_router_alias loop in top-level try/except so a non-iterable model_list / unexpected deployment shape can't escape and mask the 429 with a 500. - Type-check litellm_params before .get() to handle non-dict truthy values. - Patch llm_router=None in the parametrized fallback test so a router left by another test in the session can't redirect the unknown-model path. --------- Co-authored-by: Cursor Agent Co-authored-by: Mateo Wang * fix(bugbot): preserve "BudgetExceededError" Prometheus label Adding llm_provider to BudgetExceededError (so callbacks get provider attribution from StandardLoggingPayload) made the provider-prefix step in _get_exception_class_name silently flip the label from "BudgetExceededError" to e.g. "Openai.BudgetExceededError", breaking dashboards keyed on the historical value. Short-circuit BudgetExceededError in _get_exception_class_name the same way ProxyRateLimitError already is. Provider/category attribution still lands on the new rate_limit_category / rate_limit_type labels. * test: fix invalid 'rpm' rate_limit_type in v3 limiter test mocks The v3 rate limiter only emits 'requests', 'tokens', or 'max_parallel_requests'. Using 'rpm' caused map_v3_rate_limit_type to return None, leaving the expected RateLimitType.REQUESTS untested. Co-authored-by: Yassin Kortam * fix(bugbot): hoist provider resolver + opt-in prom rate-limit labels - dynamic_rate_limiter.py: hoist resolve_llm_provider_for_rate_limit above the TPM/RPM if/elif so the lookup runs once per request, matching the pattern in dynamic_rate_limiter_v3.py. - prometheus.py: gate the new rate_limit_category / rate_limit_type labels on litellm_proxy_failed_requests_metric behind litellm.prometheus_emit_rate_limit_labels (default False). Mirrors the existing prometheus_emit_stream_label opt-in. Preserves the metric's pre-unification label set so existing dashboards / recording rules keep matching after upgrade; operators can enable the new labels once downstream consumers include them. - Tests updated: default-off back-compat case, opt-in path enables the flag before asserting label presence. * fix: stabilize prometheus label sets and drop redundant model normalization - Cache PrometheusLogger.get_labels_for_metric per metric_name so that the label set used to construct counters at __init__ time stays in sync with the label set used at increment time, even if module-level toggles like prometheus_emit_rate_limit_labels or prometheus_emit_stream_label are flipped at runtime. Without this, toggling these flags after the logger was created would cause ValueError from prometheus_client because the runtime labels would not match the counter's declared labelnames. - Drop redundant 'model or ""' guard in ProxyRateLimitError.__init__ where model is already normalized one step earlier. Co-authored-by: Yassin Kortam * perf(dynamic_rate_limiter): only resolve provider when rate limit hit Co-authored-by: Yassin Kortam * test(prometheus): clear cached metric labels after toggling rate-limit flag The PrometheusLogger caches each metric's label set at construction time so that labels used at counter.labels(...) time stay consistent with the labels the metric was registered with. The enterprise async_post_call_failure_hook test toggles litellm.prometheus_emit_rate_limit_labels = True AFTER the fixture has already built the logger, so without invalidating the cache the rate_limit_category / rate_limit_type labels never reach the mocked counter and the assert_called_once_with check fails. Co-authored-by: Yassin Kortam * test: fix CI failures from prom label cache + flaky time-window assertion PrometheusLogger.get_labels_for_metric now caches the per-metric label set at first read so the labels passed to counter.labels(...) stay in lock step with the labels the counter was registered with. This broke two existing test patterns: - test_prometheus_labels.py: tests bind the real method onto a MagicMock, but MagicMock auto-creates a Mock for _cached_metric_labels whose .get(...) returns a truthy Mock — treated as a populated cache and returned as the label set, producing empty filtered labels and KeyError on labels["requested_model"] / ["route"]. Seed real {} containers for _cached_metric_labels and label_filters before binding. - test_prometheus_logging_callbacks.py::test_set_team_budget_metrics_with_custom_labels: the fixture builds the logger before the test monkeypatches litellm.custom_prometheus_metadata_labels, so the cached label set never picks up the new metadata labels. Clear the cache after the monkeypatch (same pattern already used for the rate-limit toggle in test_async_post_call_failure_hook). UI: view_logs/index.test.tsx "Last Minute" window assertion is off by one at the minute boundary. start_date is floored to the minute, so the dropped sub-minute fraction can push the truncated-seconds diff up to (minMinutes+1)*60 exactly when the click lands near a minute rollover. Switch the upper bound to toBeLessThanOrEqual. * feat(otel-v2): surface rate_limit_category + rate_limit_type on failed LLM-call spans PR #28909 introduced the typed v2 OTel engine that builds spans from StandardLoggingPayload, with SpanError carrying error_type + message and the genai mapper stamping error.type onto every failed LLM-call span. This PR's earlier commits added error_rate_limit_category and error_rate_limit_type to the same StandardLoggingPayload.error_information the v2 engine reads — but neither field reached a span attribute, so v2 OTel traces stayed opaque about *why* a 429 fired (vendor vs litellm, RPM vs TPM vs concurrent vs budget vs max_iterations) even after the custom-callback and prometheus surfaces gained that decomposition. Three coupled changes: 1. semconv.py: add LiteLLM.ERROR_RATE_LIMIT_CATEGORY / LiteLLM.ERROR_RATE_LIMIT_TYPE under the litellm.* vendor namespace (no GenAI semconv equivalent exists for who-rate-limited / which-dimension). 2. payloads.py: extend SpanError with rate_limit_category + rate_limit_type, populated by _parse_error() from the same error_information.error_rate_limit_* fields the custom-callback channel and prometheus rate_limit_category / rate_limit_type labels read. Single source of truth across all three observability surfaces. 3. mappers/genai.py: stamp the two attributes on the LLM-call span when present. drop_none guarantees they stay absent (not 'None') for non-rate-limit failures so trace consumers can read them unconditionally. Three regression tests in test_otel_v2_emitter.py pin: a vendor / litellm-internal RateLimitError lands category=litellm_rate_limit + rate_limit_type=requests on the span; a BudgetExceededError lands rate_limit_type=budget; a non-rate-limit failure (BadRequestError) keeps the rate_limit_* attributes absent. Mutation-tested against reverting either the SpanError extension or the _parse_error read site — both new tests fail under either mutation. Co-authored-by: Mateo Wang * test: align prometheus user-budget + logs quick-select tests with merged code The merge into this branch left two test patterns out of step with the code they exercise. test_set_user_budget_metrics_includes_user_email_and_alias_labels_when_opted_in flipped litellm.prometheus_user_budget_label_include_email_alias after the fixture had already built the PrometheusLogger. get_labels_for_metric now snapshots each metric's label set at construction time, so the runtime flip no longer reached the cached labels. Enable the flag before constructing the logger, matching how the proxy applies config at startup. view_logs/index.test.tsx referenced uiSpendLogsCall and moment without importing them, and the merged index.tsx now fetches through useLogFilterLogic (the hook the file stubs out) rather than calling uiSpendLogsCall directly. Add the imports and restore the real hook for the Quick Select window assertions so the call is actually observed. * refactor(otel/v2): drop rate-limit decomposition from the LLM-call span Proxy-side rate limits (litellm_rate_limit, budget, max_iterations) are rejected at the gate before any upstream call, so async_post_call_failure_hook tags the synthetic failure log with LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL and the v2 OTel logger never opens an LLM-call span for them; the litellm.error.rate_limit_category / litellm.error.rate_limit_type attributes were dead for exactly the cases they were meant to surface. The only failure that does open an LLM-call span carrying a RateLimitError is a vendor 429, where rate_limit_type is always None and the category just restates error.type=RateLimitError. The decomposition still reaches downstream consumers through StandardLoggingPayload.error_information.error_rate_limit_* and the prometheus rate_limit_category / rate_limit_type labels, both unchanged. Removes the SpanError fields, the _parse_error reads, the genai mapper attributes, the semconv keys, and the three span tests that asserted a scenario that never reaches the mapper in production. * fix(batch_rate_limiter): map max_parallel_requests to concurrent_requests * refactor(prometheus): drop transitive fastapi import from _get_exception_class_name Read the legacy exception_class label from a prometheus_exception_class_name marker on ProxyRateLimitError instead of importing the proxy module, keeping the integrations layer free of a transitive fastapi dependency. * chore(ui): sync schema.d.ts with unified rate-limit error spec The ProxyRateLimitError docstring flows into the proxy OpenAPI spec's 429 response description, so the generated dashboard types were out of sync. Regenerated via npm run gen:api (Check UI API Types Sync). --------- Co-authored-by: Cursor Agent Co-authored-by: Mateo Wang Co-authored-by: Yassin Kortam --- litellm/__init__.py | 9 + litellm/exceptions.py | 161 +- litellm/integrations/prometheus.py | 98 +- litellm/litellm_core_utils/litellm_logging.py | 19 + litellm/proxy/auth/auth_exception_handler.py | 14 + .../common_utils/proxy_rate_limit_error.py | 196 ++ litellm/proxy/hooks/batch_rate_limiter.py | 29 +- litellm/proxy/hooks/dynamic_rate_limiter.py | 11 +- .../proxy/hooks/dynamic_rate_limiter_v3.py | 23 +- litellm/proxy/hooks/max_budget_limiter.py | 11 +- .../hooks/max_budget_per_session_limiter.py | 11 +- litellm/proxy/hooks/max_iterations_limiter.py | 11 +- .../proxy/hooks/parallel_request_limiter.py | 75 +- .../hooks/parallel_request_limiter_v3.py | 13 +- litellm/proxy/hooks/rate_limiter_utils.py | 118 +- litellm/types/integrations/prometheus.py | 29 + litellm/types/utils.py | 17 + .../test_prometheus_logging_callbacks.py | 70 +- .../integrations/test_prometheus_labels.py | 8 + .../test_prometheus_rate_limit_labels.py | 328 ++++ .../test_prometheus_user_team_metrics.py | 39 +- .../test_proxy_rate_limit_provider_field.py | 223 ++- .../test_rate_limit_error_unification.py | 1671 +++++++++++++++++ .../src/components/view_logs/index.test.tsx | 62 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 12 +- 25 files changed, 3059 insertions(+), 199 deletions(-) create mode 100644 litellm/proxy/common_utils/proxy_rate_limit_error.py create mode 100644 tests/test_litellm/integrations/test_prometheus_rate_limit_labels.py create mode 100644 tests/test_litellm/test_rate_limit_error_unification.py diff --git a/litellm/__init__.py b/litellm/__init__.py index f22971dfa13..e6c30e12286 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -442,6 +442,13 @@ custom_prometheus_metadata_labels: List[str] = [] custom_prometheus_tags: List[str] = [] prometheus_metrics_config: Optional[List] = None prometheus_emit_stream_label: bool = False +# Opt-in: emit `rate_limit_category` and `rate_limit_type` labels on +# `litellm_proxy_failed_requests_metric`. Off by default to preserve the +# pre-unification label set so existing dashboards / recording rules keyed on +# that metric keep matching after upgrade. Enable when downstream consumers +# are ready to split 429s by source (vendor vs. litellm) and dimension +# (RPM/TPM/concurrent/budget). +prometheus_emit_rate_limit_labels: bool = False prometheus_user_budget_label_include_email_alias: bool = False prometheus_end_user_metrics_max_series_per_metric: Optional[int] = 10000 prometheus_end_user_metrics_ttl_seconds: Optional[float] = 3600.0 @@ -1303,6 +1310,8 @@ from .exceptions import ( NotFoundError, PermissionDeniedError, RateLimitError, + RateLimitErrorCategory, + RateLimitType, ServiceUnavailableError, BadGatewayError, OpenAIError, diff --git a/litellm/exceptions.py b/litellm/exceptions.py index 15f6030d4a3..1cbef6b0b49 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -9,13 +9,109 @@ ## LiteLLM versions of the OpenAI Exception Types -from typing import Any, Dict, Optional +import enum +from typing import Any, Dict, Optional, Union import httpx import openai from litellm.types.utils import LiteLLMCommonStrings + +class RateLimitErrorCategory(str, enum.Enum): + """ + Category of a rate limit error, allowing callers to distinguish where the rate + limit originated. Exposed on every :class:`RateLimitError` instance via the + ``category`` attribute. + + Use these values to switch on the rate limit source, e.g.:: + + try: + ... + except litellm.RateLimitError as e: + if e.category == RateLimitErrorCategory.LITELLM_RATE_LIMIT: + ... # litellm's own limiter (key/team/user/model RPM/TPM/budget) + elif e.category == RateLimitErrorCategory.VENDOR_RATE_LIMIT: + ... # the upstream LLM provider returned 429 + """ + + VENDOR_RATE_LIMIT = "vendor_rate_limit" + """The upstream LLM provider returned a rate-limit response (e.g. OpenAI 429).""" + + VENDOR_BATCH_RATE_LIMIT = "vendor_batch_rate_limit" + """The upstream LLM provider returned a rate-limit response on a batch endpoint.""" + + LITELLM_RATE_LIMIT = "litellm_rate_limit" + """LiteLLM's own rate limiter (key/team/user/model RPM/TPM, budget, parallel-requests, etc.) blocked the request.""" + + LITELLM_BATCH_RATE_LIMIT = "litellm_batch_rate_limit" + """LiteLLM's own batch rate limiter (token/request budget across a batch input file) blocked the request.""" + + +class RateLimitType(str, enum.Enum): + """ + The dimension that was exceeded when a rate-limit error fired. + + This is orthogonal to :class:`RateLimitErrorCategory` — *category* tells + callers **who** rate-limited the request (the upstream vendor vs. one of + litellm's own limiters), while *type* tells them **which limit dimension** + was exceeded (an RPM ceiling, a TPM ceiling, a max-parallel-requests + ceiling, a budget cap, or a max-iterations cap). + + Surfaced both on every :class:`RateLimitError` instance via the + ``rate_limit_type`` attribute and on the structured + ``StandardLoggingPayload.error_information.error_rate_limit_type`` field + so custom callbacks / metrics consumers can split rate-limit failures by + cause without parsing free-text error messages. + """ + + REQUESTS = "requests" + """Requests-per-minute (RPM) or requests-per-window ceiling exceeded.""" + + TOKENS = "tokens" + """Tokens-per-minute (TPM) or tokens-per-window ceiling exceeded.""" + + CONCURRENT_REQUESTS = "concurrent_requests" + """``max_parallel_requests`` — too many in-flight requests at once.""" + + BUDGET = "budget" + """Spend budget cap reached (key, team, user, or per-session).""" + + MAX_ITERATIONS = "max_iterations" + """Per-session max-iterations cap reached (agent-style flows).""" + + +_RATE_LIMIT_CATEGORY_VALUES = frozenset(c.value for c in RateLimitErrorCategory) +_RATE_LIMIT_TYPE_VALUES = frozenset(t.value for t in RateLimitType) + + +def validate_rate_limit_category(value: Any) -> Optional[str]: + """Return ``value`` only if it matches a known :class:`RateLimitErrorCategory`. + + Used at duck-typed read sites (StandardLoggingPayload extraction, Prometheus + labels) to reject `.category` strings set by unrelated third-party exceptions + — otherwise those would leak into custom-callback payloads and Prometheus + label cardinality. + """ + if isinstance(value, RateLimitErrorCategory): + return value.value + if isinstance(value, str) and value in _RATE_LIMIT_CATEGORY_VALUES: + return value + return None + + +def validate_rate_limit_type(value: Any) -> Optional[str]: + """Return ``value`` only if it matches a known :class:`RateLimitType`. + + See :func:`validate_rate_limit_category` for the rationale. + """ + if isinstance(value, RateLimitType): + return value.value + if isinstance(value, str) and value in _RATE_LIMIT_TYPE_VALUES: + return value + return None + + _MINIMAL_ERROR_RESPONSE: Optional[httpx.Response] = None @@ -321,6 +417,18 @@ class PermissionDeniedError(openai.PermissionDeniedError): # type: ignore class RateLimitError(openai.RateLimitError): # type: ignore + """ + Unified rate-limit error. + + Every rate-limit condition surfaced by litellm — whether it originated from + an upstream LLM provider, a vendor batch endpoint, or one of litellm's own + proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget, + max-iterations, etc.) — is raised as an instance of this class. + + The :attr:`category` attribute lets callers distinguish the source. See + :class:`RateLimitErrorCategory` for the available values. + """ + def __init__( self, message, @@ -330,6 +438,12 @@ class RateLimitError(openai.RateLimitError): # type: ignore litellm_debug_info: Optional[str] = None, max_retries: Optional[int] = None, num_retries: Optional[int] = None, + category: Union[str, RateLimitErrorCategory] = ( + RateLimitErrorCategory.VENDOR_RATE_LIMIT + ), + rate_limit_type: Optional[Union[str, RateLimitType]] = None, + headers: Optional[Dict[str, str]] = None, + detail: Any = None, ): self.status_code = 429 self.message = "litellm.RateLimitError: {}".format(message) @@ -338,9 +452,39 @@ class RateLimitError(openai.RateLimitError): # type: ignore self.litellm_debug_info = litellm_debug_info self.max_retries = max_retries self.num_retries = num_retries + self.category = ( + category.value if isinstance(category, RateLimitErrorCategory) else category + ) + # Which dimension was exceeded — request count, token count, parallel + # requests, budget, max iterations. None when the source didn't + # classify the failure (e.g. legacy vendor 429 with no header hints). + self.rate_limit_type: Optional[str] = ( + rate_limit_type.value + if isinstance(rate_limit_type, RateLimitType) + else rate_limit_type + ) + # Headers explicitly attached to the error (e.g. retry-after, + # rate_limit_type, reset_at). Preserved across the proxy boundary so + # clients can react appropriately. + # + # IMPORTANT: we deliberately do NOT auto-populate self.headers from + # response.headers when only `response` is provided. A vendor 429 can + # set arbitrary response headers (Set-Cookie, CORS overrides, …); if + # those leaked into e.headers and a downstream proxy serializer + # forwarded them to the client, a malicious upstream could inject + # browser-interpreted headers for the proxy origin. Vendor response + # headers stay reachable on `e.response.headers` for callers that + # explicitly want them; only the proxy-supplied `headers=` kwarg + # makes it onto `self.headers`. _response_headers = ( getattr(response, "headers", None) if response is not None else None ) + self.headers: Optional[Dict[str, str]] = ( + {k: str(v) for k, v in headers.items()} if headers else None + ) + # Mirrors FastAPI HTTPException.detail so the same instance can be + # serialized through both the ProxyException and HTTPException paths. + self.detail = detail if detail is not None else self.message self.response = httpx.Response( status_code=429, headers=_response_headers, @@ -843,11 +987,24 @@ LITELLM_EXCEPTION_TYPES = [ class BudgetExceededError(Exception): def __init__( - self, current_cost: float, max_budget: float, message: Optional[str] = None + self, + current_cost: float, + max_budget: float, + message: Optional[str] = None, + llm_provider: Optional[str] = None, ): self.current_cost = current_cost self.max_budget = max_budget self.status_code = 429 + self.llm_provider = llm_provider or "" + # Surface unified rate-limit fields without joining the RateLimitError + # hierarchy so existing `except BudgetExceededError:` handlers keep + # working; custom callbacks reading StandardLoggingPayload pick these + # up via the same `category` / `rate_limit_type` attributes the rest + # of the unified rate-limit error path uses. Stored as plain strings + # to match the normalization RateLimitError.__init__ performs. + self.category: str = RateLimitErrorCategory.LITELLM_RATE_LIMIT.value + self.rate_limit_type: str = RateLimitType.BUDGET.value message = ( message or f"Budget has been exceeded! Current cost: {current_cost}, Max budget: {max_budget}" diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 648fe671140..d2af95cd4cc 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -24,6 +24,10 @@ from typing import ( import litellm from litellm._logging import print_verbose, verbose_logger +from litellm.exceptions import ( + validate_rate_limit_category, + validate_rate_limit_type, +) from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.prometheus_helpers.bounded_prometheus_series_tracker import ( BoundedPrometheusSeriesTracker, @@ -78,6 +82,20 @@ class PrometheusLogger(CustomLogger): # Always initialize label_filters, even for non-premium users self.label_filters = self._parse_prometheus_config() + # Cache resolved label sets per metric. Several entries in + # ``PrometheusMetricLabels.get_labels`` read module-level toggles + # (e.g. ``litellm.prometheus_emit_stream_label``, + # ``litellm.prometheus_emit_rate_limit_labels``) that can be + # changed at runtime. Prometheus counters/gauges/histograms are + # created with a *fixed* ``labelnames`` set; if a runtime call + # to ``get_labels_for_metric`` returned a different set, the + # subsequent ``counter.labels(**_labels)`` would raise a + # ``ValueError`` from the prometheus client. Snapshotting at + # logger init time pins the label set for the lifetime of the + # logger so toggling these flags only takes effect after a + # restart, keeping init-time and runtime label sets in sync. + self._cached_metric_labels: Dict[str, List[str]] = {} + _custom_buckets = litellm.prometheus_latency_buckets self.latency_buckets = ( tuple(_custom_buckets) @@ -1033,13 +1051,27 @@ class PrometheusLogger(CustomLogger): self, metric_name: DEFINED_PROMETHEUS_METRICS ) -> List[str]: """ - Get the labels for a metric, filtered if configured + Get the labels for a metric, filtered if configured. + + The result is cached on the instance so the label set used to + construct each Prometheus metric at ``__init__`` time stays in lock + step with the label set passed to ``counter.labels(...)`` at + runtime, even if the underlying module-level toggles consulted by + :meth:`PrometheusMetricLabels.get_labels` (e.g. + ``litellm.prometheus_emit_rate_limit_labels``, + ``litellm.prometheus_emit_stream_label``) are flipped after the + logger has been created. """ + cached = self._cached_metric_labels.get(metric_name) + if cached is not None: + return cached + # Get default labels for this metric from PrometheusMetricLabels default_labels = PrometheusMetricLabels.get_labels(metric_name) # If no label filtering is configured for this metric, use default labels if metric_name not in self.label_filters: + self._cached_metric_labels[metric_name] = default_labels return default_labels # Get configured labels for this metric @@ -1050,6 +1082,7 @@ class PrometheusLogger(CustomLogger): label for label in default_labels if label in configured_labels ] + self._cached_metric_labels[metric_name] = filtered_labels return filtered_labels def _track_end_user_metric_series( @@ -2029,14 +2062,8 @@ class PrometheusLogger(CustomLogger): Proxy level tracking - failed client side requests - labelnames=[ - "end_user", - "hashed_api_key", - "api_key_alias", - REQUESTED_MODEL, - "team", - "team_alias", - ] + EXCEPTION_LABELS, + See :attr:`PrometheusMetricLabels.litellm_proxy_failed_requests_metric` + for the authoritative list of labels emitted on this metric. """ from litellm.litellm_core_utils.litellm_logging import ( StandardLoggingPayloadSetup, @@ -2059,6 +2086,9 @@ class PrometheusLogger(CustomLogger): model_id = _metadata.get("model_info", {}).get("id") or request_data.get( "model_info", {} ).get("id") + rate_limit_category, rate_limit_type = self._extract_rate_limit_labels( + original_exception + ) enum_values = UserAPIKeyLabelValues( end_user=user_api_key_dict.end_user_id, user=user_api_key_dict.user_id, @@ -2073,6 +2103,8 @@ class PrometheusLogger(CustomLogger): status_code=str(status_code), exception_status=str(status_code), exception_class=self._get_exception_class_name(original_exception), + rate_limit_category=rate_limit_category, + rate_limit_type=rate_limit_type, tags=_tags, route=user_api_key_dict.request_route, client_ip=_metadata.get("requester_ip_address"), @@ -2843,6 +2875,33 @@ class PrometheusLogger(CustomLogger): @staticmethod def _get_exception_class_name(exception: Exception) -> str: + # Some exception types pin the ``exception_class`` label to a legacy + # value for back-compat with existing dashboards (e.g. proxy-side 429s + # keep reporting as "HTTPException"). Honor that opt-in marker before + # deriving the label from the runtime class name. Reading it via + # ``getattr`` keeps this core integrations module free of a transitive + # ``fastapi`` dependency. + legacy_class_name = getattr(exception, "prometheus_exception_class_name", None) + if isinstance(legacy_class_name, str) and legacy_class_name: + return legacy_class_name + + # Same back-compat reasoning for ``BudgetExceededError``: the unified + # rate-limit error work attached ``.llm_provider`` to budget errors + # too (so callbacks reading ``StandardLoggingPayload`` get provider + # attribution). Without this short-circuit, the provider prefix below + # would silently flip the label from "BudgetExceededError" to e.g. + # "Openai.BudgetExceededError" and break dashboards keyed on the + # original value. + try: + from litellm.exceptions import BudgetExceededError + except ImportError: + BudgetExceededError = None # type: ignore[assignment,misc] + + if BudgetExceededError is not None and isinstance( + exception, BudgetExceededError + ): + return "BudgetExceededError" + exception_class_name = "" if hasattr(exception, "llm_provider"): exception_class_name = getattr(exception, "llm_provider") or "" @@ -2857,6 +2916,27 @@ class PrometheusLogger(CustomLogger): exception_class_name += exception.__class__.__name__ return exception_class_name + @staticmethod + def _extract_rate_limit_labels( + exception: Optional[Exception], + ) -> Tuple[Optional[str], Optional[str]]: + """ + Pull the unified ``category`` / ``rate_limit_type`` fields off any + exception that declares them (``litellm.RateLimitError`` and bare- + Exception subclasses like ``BudgetExceededError``). + + Values are validated against the :class:`RateLimitErrorCategory` / + :class:`RateLimitType` enums so unrelated third-party exceptions that + happen to declare ``.category`` / ``.rate_limit_type`` string attributes + can't leak garbage into Prometheus label cardinality. + """ + if exception is None: + return None, None + return ( + validate_rate_limit_category(getattr(exception, "category", None)), + validate_rate_limit_type(getattr(exception, "rate_limit_type", None)), + ) + async def log_success_fallback_event( self, original_model_group: str, kwargs: dict, original_exception: Exception ): diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index f20b66790c4..dbfcf55d75d 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -37,6 +37,10 @@ from litellm import ( turn_off_message_logging, ) from litellm._logging import _is_debugging_on, _redact_string, verbose_logger +from litellm.exceptions import ( + validate_rate_limit_category, + validate_rate_limit_type, +) from litellm._uuid import uuid from litellm.batches.batch_utils import _handle_completed_batch from litellm.caching.caching import DualCache, InMemoryCache @@ -5318,12 +5322,27 @@ class StandardLoggingPayloadSetup: else str(original_exception) ) + # Duck-typed read so bare-Exception subclasses like + # `litellm.BudgetExceededError` can participate without joining the + # RateLimitError hierarchy (which would break `except BudgetExceededError`). + # Validated against the enum value sets so a third-party exception that + # happens to declare a `.category` or `.rate_limit_type` string attribute + # can't leak garbage into the payload or Prometheus label cardinality. + rate_limit_category = validate_rate_limit_category( + getattr(original_exception, "category", None) + ) + rate_limit_type = validate_rate_limit_type( + getattr(original_exception, "rate_limit_type", None) + ) + return StandardLoggingPayloadErrorInformation( error_code=error_status, error_class=error_class, llm_provider=_llm_provider_in_exception, traceback=traceback_info, error_message=error_message if original_exception else "", + error_rate_limit_category=rate_limit_category, + error_rate_limit_type=rate_limit_type, ) @staticmethod diff --git a/litellm/proxy/auth/auth_exception_handler.py b/litellm/proxy/auth/auth_exception_handler.py index e06ac760237..f76949f4d11 100644 --- a/litellm/proxy/auth/auth_exception_handler.py +++ b/litellm/proxy/auth/auth_exception_handler.py @@ -126,6 +126,20 @@ class UserAPIKeyAuthExceptionHandler: model=request_data.get("model"), ) + # Budget checks live in tenant-scoped helpers (key / team / org / tag) + # that don't see the request model, so the BudgetExceededError they + # raise carries `llm_provider=""`. Resolve it here off `request_data` + # so custom-callback consumers reading StandardLoggingPayload get + # the same `llm_provider` attribution as for RPM/TPM 429s. + if isinstance(e, litellm.BudgetExceededError) and not e.llm_provider: + from litellm.proxy.hooks.rate_limiter_utils import ( + resolve_llm_provider_for_rate_limit, + ) + + _, e.llm_provider = resolve_llm_provider_for_rate_limit( + request_data.get("model") + ) + # Allow callbacks to transform the error response transformed_exception = await proxy_logging_obj.post_call_failure_hook( request_data=request_data, diff --git a/litellm/proxy/common_utils/proxy_rate_limit_error.py b/litellm/proxy/common_utils/proxy_rate_limit_error.py new file mode 100644 index 00000000000..24e5c991794 --- /dev/null +++ b/litellm/proxy/common_utils/proxy_rate_limit_error.py @@ -0,0 +1,196 @@ +""" +ProxyRateLimitError — a unified rate-limit exception used by litellm's +proxy-side hooks. + +Background +---------- +LiteLLM previously surfaced rate-limit conditions through *several* unrelated +exception types: + +* :class:`litellm.exceptions.RateLimitError` — raised by exception mapping when + an upstream LLM provider returns 429. +* :class:`fastapi.HTTPException` (status 429) — raised directly by proxy hooks + such as ``parallel_request_limiter``, ``dynamic_rate_limiter``, + ``batch_rate_limiter``, ``max_budget_limiter``, ``max_iterations_limiter``, + etc. +* :class:`litellm.llms.base_llm.chat.transformation.BaseLLMException` (status + 429) — raised by some provider transports. + +This made it impossible for downstream code (and end users) to express +"is this a rate limit?" with a single ``except`` clause, and impossible to +distinguish *where* the rate limit originated (vendor vs. litellm, batch vs. +chat) without ad-hoc string-matching on the message. + +This module provides a single proxy-side error class that: + +1. Is a subclass of :class:`litellm.exceptions.RateLimitError`, so user code + that catches ``RateLimitError`` works for *every* rate-limit source. +2. Is also a subclass of :class:`fastapi.HTTPException`, so existing proxy + plumbing (``isinstance(e, HTTPException)`` branches in route handlers and + FastAPI's own dispatcher) continues to behave the same way and the + ``retry-after`` / ``rate_limit_type`` / ``reset_at`` headers are preserved + on the wire. +3. Carries a :attr:`category` field (one of + :class:`litellm.exceptions.RateLimitErrorCategory`) so callers can switch on + the rate limit source. +""" + +import json +from typing import Any, Dict, Mapping, Optional, Union + +from fastapi import HTTPException + +from litellm.exceptions import RateLimitError, RateLimitErrorCategory, RateLimitType + + +def map_v3_rate_limit_type( + v3_value: Optional[str], +) -> Optional[RateLimitType]: + """ + Map the v3 rate limiter's internal `status["rate_limit_type"]` strings + onto the public :class:`RateLimitType` enum. + + The v3 limiter uses the literal values ``"requests"``, ``"tokens"``, and + ``"max_parallel_requests"``. We collapse the last one onto + :attr:`RateLimitType.CONCURRENT_REQUESTS` because that's the public name + documented for users and dashboards. Unrecognized values return ``None`` + so the field stays absent rather than carrying garbage downstream. + """ + if v3_value == "tokens": + return RateLimitType.TOKENS + if v3_value == "max_parallel_requests": + return RateLimitType.CONCURRENT_REQUESTS + if v3_value == "requests": + return RateLimitType.REQUESTS + return None + + +def _coerce_message(detail: Any) -> str: + """Best-effort, JSON-friendly stringification of an HTTPException-style detail.""" + if detail is None: + return "" + if isinstance(detail, str): + return detail + if isinstance(detail, Mapping): + for key in ("error", "message"): + if isinstance(detail.get(key), str): + return detail[key] + inner = detail.get(key) + if isinstance(inner, Mapping) and isinstance(inner.get("message"), str): + return inner["message"] + try: + return json.dumps(detail) + except (TypeError, ValueError): + return str(detail) + return str(detail) + + +# NOTE: mypy emits two `[misc]` errors on the class line below because the +# bases declare overlapping attributes with related-but-not-identical +# annotations: +# * `status_code` is `int` on starlette HTTPException but `Literal[429]` on +# openai.RateLimitError (every openai status-error subclass narrows it +# this way and silences pyright with the same convention). +# * `headers` is `Mapping[str, str] | None` on HTTPException; we narrow it +# to `Optional[Dict[str, str]]` on RateLimitError because we always carry +# a stringified dict. +# Both narrowings are intentional and handled at construction time — every +# instance always has status_code == 429 and a Dict-typed headers — so we +# silence the ATTR-overlap check rather than relax the annotations. +class ProxyRateLimitError(HTTPException, RateLimitError): # type: ignore[misc] + """ + A 429 raised by litellm's proxy-side rate limiting hooks. + + This class deliberately inherits from BOTH + :class:`litellm.exceptions.RateLimitError` and :class:`fastapi.HTTPException` + so the same instance can flow through: + + * ``except RateLimitError`` (user / SDK code that wants a category-aware + handler), and + * ``isinstance(e, HTTPException)`` (FastAPI / proxy_server.py route + handlers that need to forward ``status_code``, ``detail`` and + ``headers`` back to the client). + + Downstream code should prefer this class over + ``raise HTTPException(status_code=429, ...)`` for litellm-internal rate + limits. + + Parameters + ---------- + detail: + The structured error payload. Forwarded as ``HTTPException.detail`` so + FastAPI's default exception handler will serialize it verbatim. + headers: + Optional response headers (e.g. ``retry-after``). Values are stringified + to satisfy FastAPI's typing. + category: + One of :class:`RateLimitErrorCategory`. Defaults to + ``LITELLM_RATE_LIMIT`` since this class is only used by litellm's own + proxy-side limiters; pass ``LITELLM_BATCH_RATE_LIMIT`` for the batch + limiter, etc. + model / llm_provider: + Optional context, propagated to the inherited ``RateLimitError`` for + compatibility with logging / standard payload extraction. + """ + + # Prometheus' ``exception_class`` label is pinned to "HTTPException" for + # this type: before the unified class existed, proxy-side 429s surfaced as + # ``fastapi.HTTPException`` and existing dashboards/alerts key off that exact + # value. Distinguishing vendor vs. litellm 429s is now the job of the + # ``rate_limit_category`` / ``rate_limit_type`` labels. + prometheus_exception_class_name = "HTTPException" + + def __init__( + self, + detail: Any, + headers: Optional[Mapping[str, Any]] = None, + category: Union[ + str, RateLimitErrorCategory + ] = RateLimitErrorCategory.LITELLM_RATE_LIMIT, + rate_limit_type: Optional[Union[str, RateLimitType]] = None, + model: Optional[str] = None, + llm_provider: Optional[str] = "litellm_proxy", + ): + # Normalize None → safe defaults so callers (and the resolver helper + # in `rate_limiter_utils`) can pass `None` without producing an + # instance whose `.llm_provider` attribute is `None` — that would + # break Prometheus' `_get_exception_class_name` (it calls + # `.capitalize()` on the provider string). + model = model or "" + llm_provider = llm_provider or "litellm_proxy" + message = _coerce_message(detail) + stringified_headers: Optional[Dict[str, str]] = ( + {k: str(v) for k, v in headers.items()} if headers else None + ) + + # Initialize the FastAPI HTTPException portion first so its attributes + # (status_code, detail, headers) are already on the instance before + # RateLimitError.__init__ runs and possibly overrides them. + HTTPException.__init__( + self, + status_code=429, + detail=detail, + headers=stringified_headers, + ) + + # Now initialize the litellm RateLimitError portion. We deliberately + # pass the structured detail through so RateLimitError preserves it as + # its `.detail` attribute too — keeping both sides of the MRO + # consistent. + RateLimitError.__init__( + self, + message=message, + llm_provider=llm_provider, + model=model, + category=category, + rate_limit_type=rate_limit_type, + headers=stringified_headers, + detail=detail, + ) + # RateLimitError.__init__ overwrites self.headers with its own copy and + # leaves self.status_code at 429 — restore the HTTPException-style + # headers value so downstream code that pulls headers off the + # instance gets back exactly what the limiter passed in. + self.headers = stringified_headers + self.detail = detail + self.status_code = 429 diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index 8473b5e77de..3957e3a7fbb 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -17,7 +17,17 @@ Quick summary: - async_log_success_event() fires on GET /v1/batches/{id} (batch completion) """ -from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Union +from typing import ( + TYPE_CHECKING, + Any, + Dict, + List, + Literal, + NoReturn, + Optional, + Tuple, + Union, +) from fastapi import HTTPException from pydantic import BaseModel @@ -30,6 +40,7 @@ from litellm.batches.batch_utils import ( _get_file_content_as_dictionary, _get_models_from_batch_input_file_content, ) +from litellm.exceptions import RateLimitErrorCategory from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import ( ProxyErrorTypes, @@ -37,10 +48,11 @@ from litellm.proxy._types import ( SpecialModelNames, UserAPIKeyAuth, ) -from litellm.proxy.hooks.rate_limiter_utils import ( - ProxyHTTPRateLimitError, - resolve_llm_provider_for_rate_limit, +from litellm.proxy.common_utils.proxy_rate_limit_error import ( + ProxyRateLimitError, + map_v3_rate_limit_type, ) +from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -385,8 +397,8 @@ class _PROXY_BatchRateLimiter(CustomLogger): batch_usage: BatchFileUsage, limit_type: str, requested_model: Optional[str] = None, - ) -> None: - """Raise HTTPException for rate limit exceeded.""" + ) -> NoReturn: + """Raise :class:`ProxyRateLimitError` (a 429) for batch rate limit exceeded.""" from datetime import datetime # Find the descriptor for this status @@ -432,14 +444,15 @@ class _PROXY_BatchRateLimiter(CustomLogger): resolved_model, llm_provider = resolve_llm_provider_for_rate_limit( requested_model ) - raise ProxyHTTPRateLimitError( - status_code=429, + raise ProxyRateLimitError( detail=detail, headers={ "retry-after": str(window_size), "rate_limit_type": limit_type, "reset_at": reset_time_formatted, }, + category=RateLimitErrorCategory.LITELLM_BATCH_RATE_LIMIT, + rate_limit_type=map_v3_rate_limit_type(limit_type), model=resolved_model, llm_provider=llm_provider, ) diff --git a/litellm/proxy/hooks/dynamic_rate_limiter.py b/litellm/proxy/hooks/dynamic_rate_limiter.py index 57cd538507e..b9e2bd12ecf 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter.py @@ -11,9 +11,10 @@ from litellm import ModelResponse, Router from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger +from litellm.exceptions import RateLimitType from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.proxy.hooks.rate_limiter_utils import ( - ProxyHTTPRateLimitError, convert_priority_to_percent, resolve_llm_provider_for_rate_limit, ) @@ -222,8 +223,7 @@ class _PROXY_DynamicRateLimitHandler(CustomLogger): resolved_model, llm_provider = resolve_llm_provider_for_rate_limit( data.get("model") ) - raise ProxyHTTPRateLimitError( - status_code=429, + raise ProxyRateLimitError( detail={ "error": "Key={} over available TPM={}. Model TPM={}, Active keys={}".format( user_api_key_dict.api_key, @@ -232,6 +232,7 @@ class _PROXY_DynamicRateLimitHandler(CustomLogger): active_projects, ) }, + rate_limit_type=RateLimitType.TOKENS, model=resolved_model, llm_provider=llm_provider, ) @@ -240,8 +241,7 @@ class _PROXY_DynamicRateLimitHandler(CustomLogger): resolved_model, llm_provider = resolve_llm_provider_for_rate_limit( data.get("model") ) - raise ProxyHTTPRateLimitError( - status_code=429, + raise ProxyRateLimitError( detail={ "error": "Key={} over available RPM={}. Model RPM={}, Active keys={}".format( user_api_key_dict.api_key, @@ -250,6 +250,7 @@ class _PROXY_DynamicRateLimitHandler(CustomLogger): active_projects, ) }, + rate_limit_type=RateLimitType.REQUESTS, model=resolved_model, llm_provider=llm_provider, ) diff --git a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py index bfc6e2c2f72..493afe6105a 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -14,13 +14,16 @@ from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_utils.proxy_rate_limit_error import ( + ProxyRateLimitError, + map_v3_rate_limit_type, +) from litellm.proxy.hooks.parallel_request_limiter_v3 import ( RateLimitDescriptor, RateLimitDescriptorRateLimitObject, _PROXY_MaxParallelRequestsHandler_v3, ) from litellm.proxy.hooks.rate_limiter_utils import ( - ProxyHTTPRateLimitError, convert_priority_to_percent, resolve_llm_provider_for_rate_limit, ) @@ -497,8 +500,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): continue descriptor_key = status["descriptor_key"] if descriptor_key == "model_saturation_check": - raise ProxyHTTPRateLimitError( - status_code=429, + raise ProxyRateLimitError( detail={ "error": f"Model capacity reached for {model}. " f"Priority: {priority}, " @@ -512,6 +514,9 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): "rate_limit_type": str(status["rate_limit_type"]), "x-litellm-priority": priority or "default", }, + rate_limit_type=map_v3_rate_limit_type( + status["rate_limit_type"] + ), model=resolved_model, llm_provider=llm_provider, ) @@ -520,8 +525,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): f"Enforcing priority limits for {model}, saturation: {saturation:.1%}, " f"priority: {priority}" ) - raise ProxyHTTPRateLimitError( - status_code=429, + raise ProxyRateLimitError( detail={ "error": f"Priority-based rate limit exceeded. " f"Model: {model}, " @@ -538,6 +542,9 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): "x-litellm-priority": priority or "default", "x-litellm-saturation": f"{saturation:.2%}", }, + rate_limit_type=map_v3_rate_limit_type( + status["rate_limit_type"] + ), model=resolved_model, llm_provider=llm_provider, ) @@ -556,8 +563,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): f"Dynamic rate limiter: OVER_LIMIT response with unknown " f"descriptor_key(s) — refusing request. response={atomic_response}" ) - raise ProxyHTTPRateLimitError( - status_code=429, + raise ProxyRateLimitError( detail={ "error": "Rate limit exceeded", "descriptor_key": ( @@ -567,6 +573,9 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): str(offending["rate_limit_type"]) if offending else "unknown" ), }, + rate_limit_type=map_v3_rate_limit_type( + offending["rate_limit_type"] if offending else None + ), headers={ "retry-after": str(self.v3_limiter.window_size), "x-litellm-priority": priority or "default", diff --git a/litellm/proxy/hooks/max_budget_limiter.py b/litellm/proxy/hooks/max_budget_limiter.py index 658d7995631..769348a0b88 100644 --- a/litellm/proxy/hooks/max_budget_limiter.py +++ b/litellm/proxy/hooks/max_budget_limiter.py @@ -4,11 +4,10 @@ from litellm import verbose_logger from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger +from litellm.exceptions import RateLimitType from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.hooks.rate_limiter_utils import ( - ProxyHTTPRateLimitError, - resolve_llm_provider_for_rate_limit, -) +from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError +from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit class _PROXY_MaxBudgetLimiter(CustomLogger): @@ -70,9 +69,9 @@ class _PROXY_MaxBudgetLimiter(CustomLogger): resolved_model, llm_provider = resolve_llm_provider_for_rate_limit( data.get("model") if data else None ) - raise ProxyHTTPRateLimitError( - status_code=429, + raise ProxyRateLimitError( detail="Max budget limit reached.", + rate_limit_type=RateLimitType.BUDGET, model=resolved_model, llm_provider=llm_provider, ) diff --git a/litellm/proxy/hooks/max_budget_per_session_limiter.py b/litellm/proxy/hooks/max_budget_per_session_limiter.py index 0b63465c4a5..20bfeb3a6d5 100644 --- a/litellm/proxy/hooks/max_budget_per_session_limiter.py +++ b/litellm/proxy/hooks/max_budget_per_session_limiter.py @@ -20,11 +20,10 @@ from typing import TYPE_CHECKING, Any, Optional, Union from litellm import DualCache from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_logger import CustomLogger +from litellm.exceptions import RateLimitType from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.hooks.rate_limiter_utils import ( - ProxyHTTPRateLimitError, - resolve_llm_provider_for_rate_limit, -) +from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError +from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit if TYPE_CHECKING: from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache @@ -117,13 +116,13 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger): resolved_model, llm_provider = resolve_llm_provider_for_rate_limit( data.get("model") if data else None ) - raise ProxyHTTPRateLimitError( - status_code=429, + raise ProxyRateLimitError( detail=( f"Session budget exceeded for session {session_id}. " f"Current spend: ${current_spend:.4f}, " f"max_budget_per_session: ${max_budget:.2f}." ), + rate_limit_type=RateLimitType.BUDGET, model=resolved_model, llm_provider=llm_provider, ) diff --git a/litellm/proxy/hooks/max_iterations_limiter.py b/litellm/proxy/hooks/max_iterations_limiter.py index d5bc669c928..525214ff6be 100644 --- a/litellm/proxy/hooks/max_iterations_limiter.py +++ b/litellm/proxy/hooks/max_iterations_limiter.py @@ -16,11 +16,10 @@ from typing import TYPE_CHECKING, Any, Optional, Union from litellm import DualCache from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_logger import CustomLogger +from litellm.exceptions import RateLimitType from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.hooks.rate_limiter_utils import ( - ProxyHTTPRateLimitError, - resolve_llm_provider_for_rate_limit, -) +from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError +from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit if TYPE_CHECKING: from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache @@ -121,12 +120,12 @@ class _PROXY_MaxIterationsHandler(CustomLogger): resolved_model, llm_provider = resolve_llm_provider_for_rate_limit( data.get("model") if data else None ) - raise ProxyHTTPRateLimitError( - status_code=429, + raise ProxyRateLimitError( detail=( f"Max iterations exceeded for session {session_id}. " f"Current count: {current_count}, max_iterations: {max_iterations}." ), + rate_limit_type=RateLimitType.MAX_ITERATIONS, model=resolved_model, llm_provider=llm_provider, ) diff --git a/litellm/proxy/hooks/parallel_request_limiter.py b/litellm/proxy/hooks/parallel_request_limiter.py index c6324c3e3a3..b622241dfa5 100644 --- a/litellm/proxy/hooks/parallel_request_limiter.py +++ b/litellm/proxy/hooks/parallel_request_limiter.py @@ -1,9 +1,8 @@ import asyncio import sys from datetime import datetime, timedelta -from typing import TYPE_CHECKING, Any, List, Literal, Optional, Tuple, Union +from typing import TYPE_CHECKING, Any, List, Literal, NoReturn, Optional, Tuple, Union -from fastapi import HTTPException from pydantic import BaseModel from typing_extensions import TypedDict @@ -13,14 +12,13 @@ from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import _get_parent_otel_span_from_kwargs from litellm.proxy._types import CommonProxyErrors, CurrentItemRateLimit, UserAPIKeyAuth +from litellm.exceptions import RateLimitType from litellm.proxy.auth.auth_utils import ( get_key_model_rpm_limit, get_key_model_tpm_limit, ) -from litellm.proxy.hooks.rate_limiter_utils import ( - ProxyHTTPRateLimitError, - resolve_llm_provider_for_rate_limit, -) +from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError +from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -75,9 +73,21 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): ) if current is None: if max_parallel_requests == 0 or tpm_limit == 0 or rpm_limit == 0: - # base case - raise self.raise_rate_limit_error( + # base case — at least one dimension is set to 0 (effectively + # disabled). Pick the most specific dimension as the + # rate_limit_type so dashboards can attribute the failure to + # the right cap. Order matters: max_parallel_requests is + # listed first because it's the rarest 0 in practice and the + # most actionable signal. + if max_parallel_requests == 0: + triggered_type = RateLimitType.CONCURRENT_REQUESTS + elif tpm_limit == 0: + triggered_type = RateLimitType.TOKENS + else: + triggered_type = RateLimitType.REQUESTS + self.raise_rate_limit_error( additional_details=f"{CommonProxyErrors.max_parallel_request_limit_reached.value}. Hit limit for {rate_limit_type}. Current limits: max_parallel_requests: {max_parallel_requests}, tpm_limit: {tpm_limit}, rpm_limit: {rpm_limit}", + rate_limit_type=triggered_type, requested_model=data.get("model") if data else None, ) new_val = { @@ -100,14 +110,23 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): values_to_update_in_cache.append((request_count_api_key, new_val)) else: + # Detect which dimension actually tripped the limit so we can + # surface the right rate_limit_type. Order matches the boolean + # condition above (concurrent → tpm → rpm) — first match wins. + if int(current["current_requests"]) >= max_parallel_requests: + triggered_type = RateLimitType.CONCURRENT_REQUESTS + elif current["current_tpm"] >= tpm_limit: + triggered_type = RateLimitType.TOKENS + else: + triggered_type = RateLimitType.REQUESTS requested_model = data.get("model") if data else None resolved_model, llm_provider = resolve_llm_provider_for_rate_limit( requested_model ) - raise ProxyHTTPRateLimitError( - status_code=429, + raise ProxyRateLimitError( detail=f"LiteLLM Rate Limit Handler for rate limit type = {rate_limit_type}. {CommonProxyErrors.max_parallel_request_limit_reached.value}. current rpm: {current['current_rpm']}, rpm limit: {rpm_limit}, current tpm: {current['current_tpm']}, tpm limit: {tpm_limit}, current max_parallel_requests: {current['current_requests']}, max_parallel_requests: {max_parallel_requests}", headers={"retry-after": str(self.time_to_next_minute())}, + rate_limit_type=triggered_type, model=resolved_model, llm_provider=llm_provider, ) @@ -135,27 +154,45 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): def raise_rate_limit_error( self, additional_details: Optional[str] = None, + rate_limit_type: Optional[RateLimitType] = None, requested_model: Optional[str] = None, - ) -> HTTPException: + ) -> NoReturn: """ - Raise an HTTPException with a 429 status code and a retry-after header. + Raise a 429 with a retry-after header for litellm-proxy parallel-request limits. + + Always raises :class:`ProxyRateLimitError` — never returns. Annotated + ``NoReturn`` so type-checkers know callers after this invocation are + unreachable. The raised exception is both a + :class:`litellm.RateLimitError` (so callers can catch by category) and a + :class:`fastapi.HTTPException` (so the FastAPI dispatcher serializes it + correctly with status 429 and the supplied headers). + + ``rate_limit_type`` defaults to ``CONCURRENT_REQUESTS`` because every + existing internal caller of this helper hits the parallel-request cap + (the global-limit branch in ``async_pre_call_hook`` and the + all-zeros base case in ``check_key_in_limits``). Callers that know + the dimension exactly should pass it explicitly. ``requested_model`` is resolved via :func:`get_llm_provider` so the - raised exception carries ``llm_provider`` for downstream loggers - (Prometheus failure metric, observability callbacks). Falls back to - ``llm_provider="litellm_proxy"`` when the model is missing or - unparseable — see ``resolve_llm_provider_for_rate_limit``. + raised exception carries ``llm_provider`` (and a stripped ``model``) + for downstream loggers (Prometheus failure metric, observability + callbacks). Falls back to ``llm_provider="litellm_proxy"`` when the + model is missing or unparseable — see + :func:`resolve_llm_provider_for_rate_limit`. """ + # additional_details is optional; build the detail with a None-guard + # so callers that pass nothing don't get the literal string "None" + # interpolated into the error message. error_message = "Max parallel request limit reached" if additional_details is not None: error_message = error_message + " " + additional_details resolved_model, llm_provider = resolve_llm_provider_for_rate_limit( requested_model ) - raise ProxyHTTPRateLimitError( - status_code=429, + raise ProxyRateLimitError( detail=error_message, headers={"retry-after": str(self.time_to_next_minute())}, + rate_limit_type=rate_limit_type or RateLimitType.CONCURRENT_REQUESTS, model=resolved_model, llm_provider=llm_provider, ) @@ -248,7 +285,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): current_global_requests = 1 # if above -> raise error if current_global_requests >= global_max_parallel_requests: - return self.raise_rate_limit_error( + self.raise_rate_limit_error( additional_details=f"Hit Global Limit: Limit={global_max_parallel_requests}, current: {current_global_requests}", requested_model=data.get("model") if data else None, ) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 9fdb146b19d..62751fb68a4 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -32,10 +32,11 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( ) from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.auth_utils import get_model_rate_limit_from_metadata -from litellm.proxy.hooks.rate_limiter_utils import ( - ProxyHTTPRateLimitError, - resolve_llm_provider_for_rate_limit, +from litellm.proxy.common_utils.proxy_rate_limit_error import ( + ProxyRateLimitError, + map_v3_rate_limit_type, ) +from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.llms.openai import BaseLiteLLMOpenAIResponseObject from litellm.types.utils import CallTypes, ModelResponse, Usage @@ -1971,7 +1972,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): descriptors: List[RateLimitDescriptor], requested_model: Optional[str] = None, ) -> None: - """Handle rate limit exceeded error by raising HTTPException.""" + """Handle rate limit exceeded by raising :class:`ProxyRateLimitError` (a 429).""" for status in response["statuses"]: if status["code"] == "OVER_LIMIT": descriptor_key = status["descriptor_key"] @@ -2005,14 +2006,14 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): resolved_model, llm_provider = resolve_llm_provider_for_rate_limit( requested_model ) - raise ProxyHTTPRateLimitError( - status_code=429, + raise ProxyRateLimitError( detail=detail, headers={ "retry-after": str(self.window_size), "rate_limit_type": str(status["rate_limit_type"]), "reset_at": reset_time_formatted, }, + rate_limit_type=map_v3_rate_limit_type(status["rate_limit_type"]), model=resolved_model, llm_provider=llm_provider, ) diff --git a/litellm/proxy/hooks/rate_limiter_utils.py b/litellm/proxy/hooks/rate_limiter_utils.py index 0ba3df448e5..07440975476 100644 --- a/litellm/proxy/hooks/rate_limiter_utils.py +++ b/litellm/proxy/hooks/rate_limiter_utils.py @@ -2,13 +2,10 @@ Shared utility functions for rate limiter hooks. """ -from typing import Any, Optional, Tuple, Union - -from fastapi import HTTPException +from typing import Optional, Tuple, Union import litellm from litellm._logging import verbose_proxy_logger -from litellm.exceptions import RateLimitError from litellm.types.router import ModelGroupInfo from litellm.types.utils import PriorityReservationDict @@ -29,11 +26,21 @@ def resolve_llm_provider_for_rate_limit( ``litellm_proxy_failed_requests_metric`` show up with ``exception_class="RateLimitError"`` and no provider attribution. - Wrapped defensively: if ``model`` is missing, malformed, or - ``get_llm_provider`` raises (unknown alias, router-only model, etc.) we - fall back to ``("", "litellm_proxy")`` so we never break the request path - by piling a second exception on top of the rate-limit one we're trying to - raise. + Resolution order: + + 1. ``litellm.get_llm_provider(model)`` — covers raw provider/model + strings the SDK already understands (``"gpt-4o-mini"``, + ``"anthropic/claude-3-5-sonnet"``, ``"bedrock/..."`` etc.). + 2. **Router alias fallback** — nearly every real proxy deployment + routes through a router ``model_name`` alias (e.g. + ``"tpm-locked"`` → ``litellm_params.model: openai/gpt-4o-mini``). + ``get_llm_provider`` doesn't know router aliases, so without this + step every alias call ended up labeled ``"litellm_proxy"``, + defeating the field's purpose for the most common case. + 3. Defensive fallback to ``("", "litellm_proxy")`` — used only when + ``model`` is missing, malformed, or both lookups fail. We never let + a secondary exception escape and mask the rate-limit error we're + trying to surface. """ if not model: return "", PROXY_LLM_PROVIDER_FALLBACK @@ -46,6 +53,9 @@ def resolve_llm_provider_for_rate_limit( custom_llm_provider or PROXY_LLM_PROVIDER_FALLBACK, ) except Exception as e: + alias_resolution = _resolve_provider_from_router_alias(model) + if alias_resolution is not None: + return alias_resolution verbose_proxy_logger.debug( "rate_limiter_utils.resolve_llm_provider_for_rate_limit: " "could not resolve provider for model=%s, falling back to %s. err=%s", @@ -56,50 +66,58 @@ def resolve_llm_provider_for_rate_limit( return model, PROXY_LLM_PROVIDER_FALLBACK -class ProxyHTTPRateLimitError(HTTPException, RateLimitError): # type: ignore[misc] +def _resolve_provider_from_router_alias( + model: str, +) -> Optional[Tuple[str, str]]: """ - HTTPException raised by proxy-side rate-limit hooks that *also* exposes - ``model`` and ``llm_provider`` attributes. + Resolve a router ``model_name`` alias to ``(underlying_model, provider)`` + by scanning the active router's ``model_list``. - Why both base classes: - - - The proxy server's exception handler keys off ``HTTPException`` to render - a 429 response, so we must remain an ``HTTPException``. - - Downstream loggers (Prometheus ``async_post_call_failure_hook``, - structured logging, observability callbacks) read ``exception.llm_provider`` - via :meth:`litellm.integrations.prometheus.PrometheusLogger._get_exception_class_name` - and ``isinstance(exc, RateLimitError)`` for category routing. Inheriting - from :class:`litellm.exceptions.RateLimitError` keeps that wiring intact. - - We intentionally do not call ``RateLimitError.__init__`` (which constructs - an httpx.Response) — it isn't needed here and just adds failure surface. - Attribute parity is what downstream consumers rely on. + Returns ``None`` if the router isn't initialized, the alias isn't + registered, the deployment has no usable ``litellm_params.model``, or + any underlying lookup raises. Callers fall through to the defensive + ``litellm_proxy`` fallback in that case — never raising secondary + exceptions out of the rate-limit raise path. """ - - def __init__( - self, - status_code: int, - detail: Any = None, - headers: Optional[dict] = None, - *, - model: str = "", - llm_provider: str = PROXY_LLM_PROVIDER_FALLBACK, - ) -> None: - HTTPException.__init__( - self, status_code=status_code, detail=detail, headers=headers - ) - self.status_code = status_code - self.model = model or "" - self.llm_provider = llm_provider or PROXY_LLM_PROVIDER_FALLBACK - # `message` is what RateLimitError.__str__ would print and what some - # observability callbacks log. Keep it human-readable. - self.message = detail if isinstance(detail, str) else str(detail) - # `RateLimitError.__str__` (resolved via MRO since Starlette's - # HTTPException doesn't define `__str__`) unconditionally reads - # these attributes. Set them so `str(exc)` doesn't raise - # AttributeError from logging/traceback paths. - self.num_retries: Optional[int] = None - self.max_retries: Optional[int] = None + try: + from litellm.proxy.proxy_server import llm_router + except Exception: + return None + if llm_router is None: + return None + try: + model_list = getattr(llm_router, "model_list", None) + if not model_list: + return None + for deployment in model_list: + if not isinstance(deployment, dict): + continue + if deployment.get("model_name") != model: + continue + params = deployment.get("litellm_params") + if not isinstance(params, dict): + continue + underlying_model = params.get("model") + if not isinstance(underlying_model, str) or not underlying_model: + continue + try: + resolved_model, custom_llm_provider, _, _ = litellm.get_llm_provider( + model=underlying_model, + ) + except Exception: + continue + if not custom_llm_provider: + continue + # Prefer the underlying provider-qualified model so the failure + # callback / Prometheus label points at the actual deployment, not + # the alias. + return ( + resolved_model or underlying_model, + custom_llm_provider, + ) + return None + except Exception: + return None def convert_priority_to_percent( diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index 55f4fc96504..5b1d32cd93c 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -115,6 +115,8 @@ class ValidationResults: REQUESTED_MODEL = "requested_model" EXCEPTION_STATUS = "exception_status" EXCEPTION_CLASS = "exception_class" +RATE_LIMIT_CATEGORY = "rate_limit_category" +RATE_LIMIT_TYPE = "rate_limit_type" STATUS_CODE = "status_code" EXCEPTION_LABELS = [EXCEPTION_STATUS, EXCEPTION_CLASS] LATENCY_BUCKETS = ( @@ -174,6 +176,8 @@ class UserAPIKeyLabelNames(Enum): API_PROVIDER = "api_provider" EXCEPTION_STATUS = EXCEPTION_STATUS EXCEPTION_CLASS = EXCEPTION_CLASS + RATE_LIMIT_CATEGORY = RATE_LIMIT_CATEGORY + RATE_LIMIT_TYPE = RATE_LIMIT_TYPE STATUS_CODE = "status_code" FALLBACK_MODEL = "fallback_model" ROUTE = "route" @@ -343,6 +347,10 @@ class PrometheusMetricLabels: UserAPIKeyLabelNames.USER_EMAIL.value, UserAPIKeyLabelNames.EXCEPTION_STATUS.value, UserAPIKeyLabelNames.EXCEPTION_CLASS.value, + # ``rate_limit_category`` / ``rate_limit_type`` are appended in + # ``get_labels()`` when ``litellm.prometheus_emit_rate_limit_labels`` + # is True. Kept opt-in so existing dashboards keyed on this metric's + # historical label set keep matching after upgrade. UserAPIKeyLabelNames.ROUTE.value, UserAPIKeyLabelNames.CLIENT_IP.value, UserAPIKeyLabelNames.USER_AGENT.value, @@ -745,6 +753,25 @@ class PrometheusMetricLabels: ): custom_labels.append(UserAPIKeyLabelNames.STREAM.value) + # Conditionally add unified rate-limit labels to + # litellm_proxy_failed_requests_metric. Off by default so the metric's + # historical label set is preserved across upgrade; enable via + # ``litellm.prometheus_emit_rate_limit_labels`` once downstream + # dashboards include the new labels in their matchers / aggregations. + if ( + label_name == "litellm_proxy_failed_requests_metric" + and litellm.prometheus_emit_rate_limit_labels is True + ): + for _rate_limit_label in ( + UserAPIKeyLabelNames.RATE_LIMIT_CATEGORY.value, + UserAPIKeyLabelNames.RATE_LIMIT_TYPE.value, + ): + if ( + _rate_limit_label not in default_labels + and _rate_limit_label not in custom_labels + ): + custom_labels.append(_rate_limit_label) + _user_budget_metrics = { "litellm_remaining_user_budget_metric", "litellm_user_max_budget_metric", @@ -807,6 +834,8 @@ class UserAPIKeyLabelValues: api_provider: Optional[str] = None exception_status: Optional[str] = None exception_class: Optional[str] = None + rate_limit_category: Optional[str] = None + rate_limit_type: Optional[str] = None status_code: Optional[str] = None fallback_model: Optional[str] = None route: Optional[str] = None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index a7a0b0f6238..b76ae1f5d86 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2720,6 +2720,23 @@ class StandardLoggingPayloadErrorInformation(TypedDict, total=False): llm_provider: Optional[str] traceback: Optional[str] error_message: Optional[str] + # error_rate_limit_category: + # For 429 / rate-limit errors, the source of the rate limit. One of the + # string values defined by `litellm.exceptions.RateLimitErrorCategory` + # (vendor_rate_limit, vendor_batch_rate_limit, litellm_rate_limit, + # litellm_batch_rate_limit). None for non-rate-limit exceptions. + # Surfaced here so custom callbacks / metrics consumers can switch on + # the rate-limit source without reaching for the raw exception. + error_rate_limit_category: Optional[str] + # error_rate_limit_type: + # For 429 / rate-limit errors, the dimension that was exceeded. One of + # the string values defined by `litellm.exceptions.RateLimitType` + # (requests, tokens, concurrent_requests, budget, max_iterations). + # None for non-rate-limit exceptions and for rate-limit exceptions that + # did not classify the failure (e.g. legacy vendor 429 with no header + # hints). Lets dashboards split rate-limit failures by cause without + # parsing free-text error messages. + error_rate_limit_type: Optional[str] class GuardrailMode(TypedDict, total=False): diff --git a/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py b/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py index f8bac820582..d0ad1cc8f82 100644 --- a/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py +++ b/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py @@ -783,6 +783,16 @@ async def test_async_post_call_failure_hook(prometheus_logger): it should increment the litellm_proxy_failed_requests_metric and litellm_proxy_total_requests_metric """ + # Opt into the unified rate-limit labels so this test exercises the + # full label set surfaced when `prometheus_emit_rate_limit_labels` is on. + # The logger caches each metric's label set at construction time (so the + # labels passed to ``counter.labels(...)`` stay in lock step with the + # labels used to register the metric), so we must invalidate the cache + # after flipping the toggle for the cache to pick up the new label set. + original_emit = litellm.prometheus_emit_rate_limit_labels + litellm.prometheus_emit_rate_limit_labels = True + prometheus_logger._cached_metric_labels.clear() + # Mock the prometheus metrics prometheus_logger.litellm_proxy_failed_requests_metric = MagicMock() prometheus_logger.litellm_proxy_total_requests_metric = MagicMock() @@ -804,32 +814,38 @@ async def test_async_post_call_failure_hook(prometheus_logger): request_route="/chat/completions", ) - # Call the function - await prometheus_logger.async_post_call_failure_hook( - request_data=request_data, - original_exception=original_exception, - user_api_key_dict=user_api_key_dict, - ) + try: + # Call the function + await prometheus_logger.async_post_call_failure_hook( + request_data=request_data, + original_exception=original_exception, + user_api_key_dict=user_api_key_dict, + ) - # Assert failed requests metric was incremented with correct labels - prometheus_logger.litellm_proxy_failed_requests_metric.labels.assert_called_once_with( - end_user=None, - user="test_user", - user_email=None, - hashed_api_key="test_key", - api_key_alias="test_alias", - team="test_team", - team_alias="test_team_alias", - org_id=None, - org_alias=None, - requested_model="gpt-5-mini", - exception_status="429", - exception_class="Openai.RateLimitError", - route=user_api_key_dict.request_route, - model_id=None, - client_ip=None, - user_agent=None, - ) + # Assert failed requests metric was incremented with correct labels + prometheus_logger.litellm_proxy_failed_requests_metric.labels.assert_called_once_with( + end_user=None, + user="test_user", + user_email=None, + hashed_api_key="test_key", + api_key_alias="test_alias", + team="test_team", + team_alias="test_team_alias", + org_id=None, + org_alias=None, + requested_model="gpt-5-mini", + exception_status="429", + exception_class="Openai.RateLimitError", + rate_limit_category="vendor_rate_limit", + rate_limit_type=None, + route=user_api_key_dict.request_route, + model_id=None, + client_ip=None, + user_agent=None, + ) + finally: + litellm.prometheus_emit_rate_limit_labels = original_emit + prometheus_logger._cached_metric_labels.clear() prometheus_logger.litellm_proxy_failed_requests_metric.labels().inc.assert_called_once() # Assert total requests metric was incremented with correct labels @@ -1962,6 +1978,10 @@ def test_set_team_budget_metrics_with_custom_labels(prometheus_logger, monkeypat # Set custom prometheus labels custom_labels = ["metadata.organization", "metadata.environment"] monkeypatch.setattr("litellm.custom_prometheus_metadata_labels", custom_labels) + # Logger caches each metric's label set at construction time (fixture + # runs before this monkeypatch), so invalidate so the cached label set + # picks up the freshly-configured custom metadata labels. + prometheus_logger._cached_metric_labels.clear() # Create test team with custom metadata team = MagicMock( diff --git a/tests/test_litellm/integrations/test_prometheus_labels.py b/tests/test_litellm/integrations/test_prometheus_labels.py index 1ba332a341b..c83d89e87c4 100644 --- a/tests/test_litellm/integrations/test_prometheus_labels.py +++ b/tests/test_litellm/integrations/test_prometheus_labels.py @@ -284,6 +284,12 @@ def test_prometheus_metrics_use_normalized_routes(): # Create a mock PrometheusLogger prometheus_logger = MagicMock() + # ``get_labels_for_metric`` reads ``_cached_metric_labels`` and + # ``label_filters`` off ``self``; default MagicMock attribute access + # returns Mocks that masquerade as a populated cache, so seed real + # containers before binding the real method. + prometheus_logger._cached_metric_labels = {} + prometheus_logger.label_filters = {} prometheus_logger.get_labels_for_metric = ( PrometheusLogger.get_labels_for_metric.__get__(prometheus_logger) ) @@ -327,6 +333,8 @@ def test_prometheus_label_value_sanitization(): from unittest.mock import MagicMock prometheus_logger = MagicMock() + prometheus_logger._cached_metric_labels = {} + prometheus_logger.label_filters = {} prometheus_logger.get_labels_for_metric = ( PrometheusLogger.get_labels_for_metric.__get__(prometheus_logger) ) diff --git a/tests/test_litellm/integrations/test_prometheus_rate_limit_labels.py b/tests/test_litellm/integrations/test_prometheus_rate_limit_labels.py new file mode 100644 index 00000000000..bb035c4c3ee --- /dev/null +++ b/tests/test_litellm/integrations/test_prometheus_rate_limit_labels.py @@ -0,0 +1,328 @@ +""" +Tests for the Prometheus rate-limit labels added on top of PR #27687. + +Covers two follow-up gaps to the unified rate-limit error work: + +1. ``litellm_proxy_failed_requests_metric`` now carries + ``rate_limit_category`` and ``rate_limit_type`` labels populated from + :class:`litellm.RateLimitError` (vendor + ``ProxyRateLimitError`` + subclass). Closes the Prometheus side of LIT-2718. +2. ``_get_exception_class_name`` keeps emitting the literal string + ``"HTTPException"`` for ``ProxyRateLimitError`` so existing dashboards + that key off ``exception_class="HTTPException"`` for litellm-internal + 429s don't silently break when the new class lands. +""" + +from unittest.mock import MagicMock, patch + +import pytest + +from litellm.exceptions import ( + RateLimitError, + RateLimitErrorCategory, + RateLimitType, +) +from litellm.integrations.prometheus import PrometheusLogger +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError +from litellm.types.integrations.prometheus import ( + PrometheusMetricLabels, + UserAPIKeyLabelNames, + UserAPIKeyLabelValues, +) + + +# --------------------------------------------------------------------------- +# Label / enum wiring +# --------------------------------------------------------------------------- + + +def test_should_register_rate_limit_label_names_on_enum(): + assert UserAPIKeyLabelNames.RATE_LIMIT_CATEGORY.value == "rate_limit_category" + assert UserAPIKeyLabelNames.RATE_LIMIT_TYPE.value == "rate_limit_type" + + +def test_should_include_rate_limit_labels_on_failed_requests_metric(): + import litellm + + original = litellm.prometheus_emit_rate_limit_labels + try: + litellm.prometheus_emit_rate_limit_labels = True + labels = PrometheusMetricLabels.get_labels( + "litellm_proxy_failed_requests_metric" + ) + assert "rate_limit_category" in labels + assert "rate_limit_type" in labels + # These must coexist with the legacy exception labels (back-compat). + assert "exception_class" in labels + assert "exception_status" in labels + finally: + litellm.prometheus_emit_rate_limit_labels = original + + +def test_should_omit_rate_limit_labels_by_default_for_back_compat(): + """Default-off preserves the metric's historical label set so existing + dashboards / recording rules keyed on `litellm_proxy_failed_requests_metric` + keep matching after upgrade.""" + import litellm + + assert litellm.prometheus_emit_rate_limit_labels is False + labels = PrometheusMetricLabels.get_labels("litellm_proxy_failed_requests_metric") + assert "rate_limit_category" not in labels + assert "rate_limit_type" not in labels + # Pre-PR labels must still be present. + assert "exception_class" in labels + assert "exception_status" in labels + + +def test_should_accept_rate_limit_fields_on_user_api_key_label_values(): + enum_values = UserAPIKeyLabelValues( + rate_limit_category="litellm_rate_limit", + rate_limit_type="requests", + ) + assert enum_values.rate_limit_category == "litellm_rate_limit" + assert enum_values.rate_limit_type == "requests" + + +# --------------------------------------------------------------------------- +# _extract_rate_limit_labels helper +# --------------------------------------------------------------------------- + + +def test_should_extract_vendor_category_for_vanilla_rate_limit_error(): + err = RateLimitError(message="vendor 429", llm_provider="openai", model="gpt-4o") + category, rate_limit_type = PrometheusLogger._extract_rate_limit_labels(err) + assert category == "vendor_rate_limit" + assert rate_limit_type is None + + +def test_should_extract_litellm_category_and_type_for_proxy_rate_limit_error(): + err = ProxyRateLimitError( + detail={"error": "tpm exceeded"}, + category=RateLimitErrorCategory.LITELLM_RATE_LIMIT, + rate_limit_type=RateLimitType.TOKENS, + ) + category, rate_limit_type = PrometheusLogger._extract_rate_limit_labels(err) + assert category == "litellm_rate_limit" + assert rate_limit_type == "tokens" + + +def test_should_return_none_for_non_rate_limit_exception(): + assert PrometheusLogger._extract_rate_limit_labels(ValueError("nope")) == ( + None, + None, + ) + + +def test_should_return_none_for_none_exception(): + assert PrometheusLogger._extract_rate_limit_labels(None) == (None, None) + + +def test_should_extract_budget_dimension_for_budget_exceeded_error(): + # Virtual-key / team / org / end-user budget caps raise + # `litellm.BudgetExceededError` (a bare Exception subclass), which sets + # the same `.category` / `.rate_limit_type` attributes as the unified + # RateLimitError path so Prometheus can split budget 429s from other + # 429s without the customer parsing free-text error messages. + import litellm + + err = litellm.BudgetExceededError(current_cost=0.5, max_budget=0.1) + category, rate_limit_type = PrometheusLogger._extract_rate_limit_labels(err) + assert category == "litellm_rate_limit" + assert rate_limit_type == "budget" + + +@pytest.mark.parametrize( + "category_enum,rate_limit_enum,expected_category,expected_type", + [ + ( + RateLimitErrorCategory.LITELLM_RATE_LIMIT, + RateLimitType.REQUESTS, + "litellm_rate_limit", + "requests", + ), + ( + RateLimitErrorCategory.LITELLM_RATE_LIMIT, + RateLimitType.TOKENS, + "litellm_rate_limit", + "tokens", + ), + ( + RateLimitErrorCategory.LITELLM_RATE_LIMIT, + RateLimitType.CONCURRENT_REQUESTS, + "litellm_rate_limit", + "concurrent_requests", + ), + ( + RateLimitErrorCategory.LITELLM_RATE_LIMIT, + RateLimitType.BUDGET, + "litellm_rate_limit", + "budget", + ), + ( + RateLimitErrorCategory.LITELLM_RATE_LIMIT, + RateLimitType.MAX_ITERATIONS, + "litellm_rate_limit", + "max_iterations", + ), + ( + RateLimitErrorCategory.LITELLM_BATCH_RATE_LIMIT, + RateLimitType.REQUESTS, + "litellm_batch_rate_limit", + "requests", + ), + ], +) +def test_should_serialize_rate_limit_enums_as_underlying_string_values( + category_enum, rate_limit_enum, expected_category, expected_type +): + err = ProxyRateLimitError( + detail="boom", category=category_enum, rate_limit_type=rate_limit_enum + ) + category, rate_limit_type = PrometheusLogger._extract_rate_limit_labels(err) + assert category == expected_category + assert rate_limit_type == expected_type + + +# --------------------------------------------------------------------------- +# _get_exception_class_name back-compat +# --------------------------------------------------------------------------- + + +def test_should_emit_legacy_http_exception_label_for_proxy_rate_limit_error(): + """ + ``ProxyRateLimitError`` multi-inherits from ``HTTPException`` + + ``RateLimitError``. The ``exception_class`` label MUST keep emitting + "HTTPException" for back-compat with existing dashboards (see Slack + thread + PR #27687 review). Distinguishing vendor vs. litellm 429s + is now the job of the new ``rate_limit_category`` label. + """ + err = ProxyRateLimitError(detail={"error": "boom"}) + assert PrometheusLogger._get_exception_class_name(err) == "HTTPException" + + +def test_should_keep_provider_prefixed_exception_class_for_vendor_rate_limit_errors(): + err = RateLimitError(message="vendor 429", llm_provider="openai", model="gpt-4o") + # Vendor-side errors keep the historical "Provider.ClassName" formatting. + assert PrometheusLogger._get_exception_class_name(err) == "Openai.RateLimitError" + + +def test_should_preserve_exception_class_name_for_unrelated_exceptions(): + assert PrometheusLogger._get_exception_class_name(ValueError("nope")) == ( + "ValueError" + ) + + +# --------------------------------------------------------------------------- +# End-to-end wiring through async_post_call_failure_hook +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_should_populate_rate_limit_labels_for_proxy_rate_limit_error_on_failure_hook(): + """ + When a proxy hook raises ``ProxyRateLimitError`` and the failure flows + through ``async_post_call_failure_hook``, the resulting + ``UserAPIKeyLabelValues`` must carry both new labels AND keep + ``exception_class="HTTPException"`` for back-compat. + """ + with patch( + "litellm.integrations.prometheus.PrometheusLogger.__init__", return_value=None + ): + logger = PrometheusLogger() + logger.litellm_proxy_failed_requests_metric = MagicMock() + logger.litellm_proxy_total_requests_metric = MagicMock() + logger.get_labels_for_metric = MagicMock( + return_value=PrometheusMetricLabels.get_labels( + "litellm_proxy_failed_requests_metric" + ) + ) + + err = ProxyRateLimitError( + detail={"error": "rpm exceeded"}, + category=RateLimitErrorCategory.LITELLM_RATE_LIMIT, + rate_limit_type=RateLimitType.REQUESTS, + ) + + with patch( + "litellm.integrations.prometheus.prometheus_label_factory" + ) as mock_label_factory: + mock_label_factory.return_value = {} + await logger.async_post_call_failure_hook( + request_data={"model": "gpt-4o-mini", "metadata": {}}, + original_exception=err, + user_api_key_dict=UserAPIKeyAuth(token="t"), + ) + + enum_values = mock_label_factory.call_args_list[0].kwargs["enum_values"] + assert isinstance(enum_values, UserAPIKeyLabelValues) + assert enum_values.rate_limit_category == "litellm_rate_limit" + assert enum_values.rate_limit_type == "requests" + # Back-compat: exception_class on a ProxyRateLimitError stays "HTTPException". + assert enum_values.exception_class == "HTTPException" + assert enum_values.exception_status == "429" + + +@pytest.mark.asyncio +async def test_should_populate_rate_limit_labels_for_vendor_rate_limit_error_on_failure_hook(): + with patch( + "litellm.integrations.prometheus.PrometheusLogger.__init__", return_value=None + ): + logger = PrometheusLogger() + logger.litellm_proxy_failed_requests_metric = MagicMock() + logger.litellm_proxy_total_requests_metric = MagicMock() + logger.get_labels_for_metric = MagicMock( + return_value=PrometheusMetricLabels.get_labels( + "litellm_proxy_failed_requests_metric" + ) + ) + + err = RateLimitError(message="upstream 429", llm_provider="openai", model="gpt-4o") + + with patch( + "litellm.integrations.prometheus.prometheus_label_factory" + ) as mock_label_factory: + mock_label_factory.return_value = {} + await logger.async_post_call_failure_hook( + request_data={"model": "gpt-4o", "metadata": {}}, + original_exception=err, + user_api_key_dict=UserAPIKeyAuth(token="t"), + ) + + enum_values = mock_label_factory.call_args_list[0].kwargs["enum_values"] + assert isinstance(enum_values, UserAPIKeyLabelValues) + assert enum_values.rate_limit_category == "vendor_rate_limit" + assert enum_values.rate_limit_type is None + # Vendor errors keep the historical Provider.ClassName label. + assert enum_values.exception_class == "Openai.RateLimitError" + assert enum_values.exception_status == "429" + + +@pytest.mark.asyncio +async def test_should_leave_rate_limit_labels_blank_for_non_rate_limit_failure(): + with patch( + "litellm.integrations.prometheus.PrometheusLogger.__init__", return_value=None + ): + logger = PrometheusLogger() + logger.litellm_proxy_failed_requests_metric = MagicMock() + logger.litellm_proxy_total_requests_metric = MagicMock() + logger.get_labels_for_metric = MagicMock( + return_value=PrometheusMetricLabels.get_labels( + "litellm_proxy_failed_requests_metric" + ) + ) + + with patch( + "litellm.integrations.prometheus.prometheus_label_factory" + ) as mock_label_factory: + mock_label_factory.return_value = {} + await logger.async_post_call_failure_hook( + request_data={"model": "gpt-4o", "metadata": {}}, + original_exception=RuntimeError("boom"), + user_api_key_dict=UserAPIKeyAuth(token="t"), + ) + + enum_values = mock_label_factory.call_args_list[0].kwargs["enum_values"] + assert isinstance(enum_values, UserAPIKeyLabelValues) + assert enum_values.rate_limit_category is None + assert enum_values.rate_limit_type is None diff --git a/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py b/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py index 12f30ab6024..361ab7332f8 100644 --- a/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py +++ b/tests/test_litellm/integrations/test_prometheus_user_team_metrics.py @@ -511,29 +511,34 @@ def test_set_user_budget_metrics_default_no_email_alias_labels( ) -def test_set_user_budget_metrics_includes_user_email_and_alias_labels_when_opted_in( - prometheus_logger, -): - """When prometheus_user_budget_label_include_email_alias=True, email+alias labels appear.""" +def test_set_user_budget_metrics_includes_user_email_and_alias_labels_when_opted_in(): + """When prometheus_user_budget_label_include_email_alias=True, email+alias labels appear. + + The flag is read once per metric at logger construction time and snapshotted, + so it must be enabled before the PrometheusLogger is built (mirroring how the + proxy applies config at startup before instantiating callbacks). + """ import litellm from litellm.proxy._types import LiteLLM_UserTable litellm.prometheus_user_budget_label_include_email_alias = True - user = LiteLLM_UserTable( - user_id="user-abc-123", - user_email="alice@example.com", - user_alias="Alice", - spend=25.0, - max_budget=100.0, - budget_reset_at=datetime(2026, 3, 1, tzinfo=timezone.utc), - ) - - prometheus_logger.litellm_remaining_user_budget_metric = MagicMock() - prometheus_logger.litellm_user_max_budget_metric = MagicMock() - prometheus_logger.litellm_user_budget_remaining_hours_metric = MagicMock() - try: + prometheus_logger = PrometheusLogger() + + user = LiteLLM_UserTable( + user_id="user-abc-123", + user_email="alice@example.com", + user_alias="Alice", + spend=25.0, + max_budget=100.0, + budget_reset_at=datetime(2026, 3, 1, tzinfo=timezone.utc), + ) + + prometheus_logger.litellm_remaining_user_budget_metric = MagicMock() + prometheus_logger.litellm_user_max_budget_metric = MagicMock() + prometheus_logger.litellm_user_budget_remaining_hours_metric = MagicMock() + prometheus_logger._set_user_budget_metrics(user) prometheus_logger.litellm_remaining_user_budget_metric.labels.assert_called_once_with( diff --git a/tests/test_litellm/proxy/hooks/test_proxy_rate_limit_provider_field.py b/tests/test_litellm/proxy/hooks/test_proxy_rate_limit_provider_field.py index 8c74919df19..02b4e32db86 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_rate_limit_provider_field.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_rate_limit_provider_field.py @@ -22,7 +22,7 @@ no ``llm_provider`` / ``model`` attribute. Downstream: category routing missed these entirely. The fix wraps every internal raise site in -:class:`ProxyHTTPRateLimitError` (an ``HTTPException`` *and* a +:class:`ProxyRateLimitError` (an ``HTTPException`` *and* a ``litellm.RateLimitError``), and resolves ``model`` / ``llm_provider`` from ``data["model"]`` via :func:`get_llm_provider`. When the model is missing or unparseable we fall back to ``llm_provider="litellm_proxy"`` so we never break @@ -61,9 +61,9 @@ 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, - ProxyHTTPRateLimitError, resolve_llm_provider_for_rate_limit, ) from litellm.proxy.utils import InternalUsageCache @@ -75,12 +75,11 @@ from litellm.types.agents import AgentResponse # --------------------------------------------------------------------------- -class TestProxyHTTPRateLimitErrorClass: +class TestProxyRateLimitErrorClass: """Pin the dual ``HTTPException`` + ``RateLimitError`` shape.""" def test_is_both_http_exception_and_rate_limit_error(self): - e = ProxyHTTPRateLimitError( - status_code=429, + e = ProxyRateLimitError( detail="boom", model="gpt-4o-mini", llm_provider="openai", @@ -92,15 +91,15 @@ class TestProxyHTTPRateLimitErrorClass: assert e.status_code == 429 assert e.model == "gpt-4o-mini" assert e.llm_provider == "openai" - assert e.message == "boom" + # ProxyRateLimitError prefixes message via RateLimitError.__init__. + assert "boom" in e.message assert e.detail == "boom" def test_dict_detail_is_stringified_for_message(self): # Some hooks pass a dict detail (e.g. dynamic_rate_limiter v1) — the # `message` attr (read by RateLimitError.__str__ and observability # callbacks) must still be a string. - e = ProxyHTTPRateLimitError( - status_code=429, + e = ProxyRateLimitError( detail={"error": "over rpm"}, model="claude-3-5-sonnet", llm_provider="anthropic", @@ -109,16 +108,15 @@ class TestProxyHTTPRateLimitErrorClass: assert "over rpm" in e.message def test_defaults_to_litellm_proxy_provider(self): - e = ProxyHTTPRateLimitError(status_code=429, detail="x") + e = ProxyRateLimitError(detail="x") assert e.llm_provider == PROXY_LLM_PROVIDER_FALLBACK assert e.model == "" def test_none_provider_normalized_to_fallback(self): - e = ProxyHTTPRateLimitError( - status_code=429, + e = ProxyRateLimitError( detail="x", - model=None, # type: ignore[arg-type] - llm_provider=None, # type: ignore[arg-type] + model=None, + llm_provider=None, ) assert e.llm_provider == PROXY_LLM_PROVIDER_FALLBACK assert e.model == "" @@ -143,7 +141,10 @@ class TestResolveLLMProviderForRateLimit: # Must never raise — the resolver wraps `get_llm_provider` defensively # because raising here would mask the rate-limit error we're trying # to surface to the user. - resolved_model, provider = resolve_llm_provider_for_rate_limit(model) + # Pin llm_router to None so the alias-fallback path doesn't pick up + # a router left behind by another test in the session. + with patch("litellm.proxy.proxy_server.llm_router", None): + resolved_model, provider = resolve_llm_provider_for_rate_limit(model) assert provider == PROXY_LLM_PROVIDER_FALLBACK # Resolver returns the input model verbatim on the unknown branch so # the `.model` attribute is never silently swapped to a different one. @@ -155,15 +156,148 @@ class TestResolveLLMProviderForRateLimit: def test_get_llm_provider_raising_is_swallowed(self): # If get_llm_provider itself blows up (unexpected error), we still # fall back rather than letting the secondary exception escape. + # No router is registered in this test, so the alias-fallback path + # also yields None and we land at PROXY_LLM_PROVIDER_FALLBACK. with patch.object( litellm, "get_llm_provider", side_effect=RuntimeError("boom"), ): - resolved_model, provider = resolve_llm_provider_for_rate_limit("anything") + with patch( + "litellm.proxy.proxy_server.llm_router", + None, + ): + resolved_model, provider = resolve_llm_provider_for_rate_limit( + "anything" + ) assert provider == PROXY_LLM_PROVIDER_FALLBACK assert resolved_model == "anything" + def test_router_alias_resolves_to_underlying_provider(self): + """ + Nearly every real LiteLLM proxy deployment uses router aliases: + + model_list: + - model_name: tpm-locked + litellm_params: + model: openai/gpt-4o-mini + ... + + ``litellm.get_llm_provider("tpm-locked")`` doesn't know about + router aliases and raises. Before this fix the resolver fell + through to ``"litellm_proxy"``, defeating the whole point of the + ``llm_provider`` field on the rate-limit error. The alias path + must look the deployment up in the router's ``model_list`` and + resolve from its ``litellm_params.model``. + """ + + class _FakeRouter: + model_list = [ + { + "model_name": "tpm-locked", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "fake", + }, + } + ] + + with patch( + "litellm.proxy.proxy_server.llm_router", + _FakeRouter(), + ): + resolved_model, provider = resolve_llm_provider_for_rate_limit("tpm-locked") + assert provider == "openai", ( + f"Router-alias path must resolve through litellm_params.model, " + f"not fall through to {PROXY_LLM_PROVIDER_FALLBACK!r}. Got " + f"provider={provider!r}, model={resolved_model!r}." + ) + # The resolved model should point at the underlying deployment so + # downstream Prometheus labels / failure callbacks attribute the + # 429 to the real upstream, not the alias. + assert resolved_model == "gpt-4o-mini" + + def test_router_alias_with_multiple_deployments_uses_first(self): + """ + When an alias maps to multiple deployments (the load-balancing + case), the rate-limit error fired at the *alias* level is + deployment-agnostic — we have no way of knowing which one would + have been picked. Use the first deployment's underlying provider: + every deployment under one alias should agree on provider in any + sensible config, and 'first' is deterministic so the Prometheus + label is stable. + """ + + class _FakeRouter: + model_list = [ + { + "model_name": "claude-pool", + "litellm_params": {"model": "anthropic/claude-3-5-sonnet"}, + }, + { + "model_name": "claude-pool", + "litellm_params": {"model": "anthropic/claude-3-5-haiku"}, + }, + ] + + with patch( + "litellm.proxy.proxy_server.llm_router", + _FakeRouter(), + ): + _, provider = resolve_llm_provider_for_rate_limit("claude-pool") + assert provider == "anthropic" + + def test_router_alias_unknown_falls_back(self): + """ + Alias not in the router model_list — both lookups fail, so we + land at the defensive ``litellm_proxy`` fallback rather than + raising. + """ + + class _FakeRouter: + model_list = [ + { + "model_name": "tpm-locked", + "litellm_params": {"model": "openai/gpt-4o-mini"}, + } + ] + + with patch( + "litellm.proxy.proxy_server.llm_router", + _FakeRouter(), + ): + resolved_model, provider = resolve_llm_provider_for_rate_limit( + "not-an-alias" + ) + assert provider == PROXY_LLM_PROVIDER_FALLBACK + assert resolved_model == "not-an-alias" + + def test_router_alias_with_malformed_deployment_falls_back(self): + """ + A deployment in the router model_list with no usable + ``litellm_params.model`` (or where ``get_llm_provider`` on the + underlying string also raises) must not crash the resolver — + fall through to the defensive fallback. + """ + + class _FakeRouter: + model_list = [ + {"model_name": "broken", "litellm_params": {}}, + {"model_name": "broken", "litellm_params": {"model": ""}}, + { + "model_name": "broken", + "litellm_params": {"model": "nonsense-no-provider"}, + }, + ] + + with patch( + "litellm.proxy.proxy_server.llm_router", + _FakeRouter(), + ): + resolved_model, provider = resolve_llm_provider_for_rate_limit("broken") + assert provider == PROXY_LLM_PROVIDER_FALLBACK + assert resolved_model == "broken" + # --------------------------------------------------------------------------- # parallel_request_limiter v1 @@ -352,7 +486,7 @@ async def test_parallel_request_limiter_v1_missing_model_falls_back(): # --------------------------------------------------------------------------- -def _v3_over_limit_response(rate_limit_type: str = "rpm") -> dict: +def _v3_over_limit_response(rate_limit_type: str = "requests") -> dict: return { "overall_code": "OVER_LIMIT", "statuses": [ @@ -532,7 +666,7 @@ async def test_dynamic_rate_limiter_v3_model_capacity_path_populates_provider(): "descriptor_key": "model_saturation_check", "current_limit": 100, "limit_remaining": 0, - "rate_limit_type": "rpm", + "rate_limit_type": "requests", } ], } @@ -582,7 +716,7 @@ async def test_dynamic_rate_limiter_v3_unknown_descriptor_path_populates_provide "descriptor_key": "something_we_dont_handle", "current_limit": 1, "limit_remaining": 0, - "rate_limit_type": "rpm", + "rate_limit_type": "requests", } ], } @@ -937,31 +1071,56 @@ async def test_max_budget_per_session_limiter_unknown_model_falls_back(): # --------------------------------------------------------------------------- -def test_prometheus_exception_class_name_includes_provider(): +def test_prometheus_exception_class_name_back_compat_for_proxy_rate_limit_error(): + """ + `_get_exception_class_name` deliberately returns the literal string + ``"HTTPException"`` for every ``ProxyRateLimitError`` instance so that + pre-existing dashboards / alerts (which key off the historical value) + keep working after the unified rate-limit error class landed in #27687. + + Provider attribution is now surfaced separately via the + ``rate_limit_category`` / ``rate_limit_type`` labels — this test pins + the back-compat shim itself. + """ from litellm.integrations.prometheus import PrometheusLogger - exc = ProxyHTTPRateLimitError( - status_code=429, + exc = ProxyRateLimitError( detail="over limit", model="gpt-4o-mini", llm_provider="openai", ) + assert PrometheusLogger._get_exception_class_name(exc) == "HTTPException" - name = PrometheusLogger._get_exception_class_name(exc) - # Format is "{Provider.}{ClassName}" per `_get_exception_class_name`. - assert name.startswith("Openai.") - # And specifically: it ends in our exception class. (We don't pin the - # full string to avoid coupling the test to PR #27687's parallel rename.) - assert name.endswith("ProxyHTTPRateLimitError") + # Same back-compat path even when the resolver fell back to litellm_proxy. + exc_no_model = ProxyRateLimitError(detail="over limit") + assert PrometheusLogger._get_exception_class_name(exc_no_model) == "HTTPException" -def test_prometheus_exception_class_name_falls_back_when_no_model(): +def test_prometheus_exception_class_name_back_compat_for_budget_exceeded_error(): + """ + The unified rate-limit work also attached ``.llm_provider`` to + ``BudgetExceededError`` so callbacks get provider attribution from + ``StandardLoggingPayload``. Without a back-compat short-circuit the + provider-prefix step in ``_get_exception_class_name`` would silently + flip the label from ``"BudgetExceededError"`` to e.g. + ``"Openai.BudgetExceededError"`` and break dashboards keyed on the + historical value. Pin the literal label here. + """ from litellm.integrations.prometheus import PrometheusLogger - exc = ProxyHTTPRateLimitError(status_code=429, detail="over limit") - name = PrometheusLogger._get_exception_class_name(exc) - # `litellm_proxy` -> `Litellm_proxy.` (capitalize first char only). - assert name.startswith("Litellm_proxy.") + err = litellm.BudgetExceededError( + current_cost=1.0, + max_budget=0.5, + llm_provider="openai", + ) + assert PrometheusLogger._get_exception_class_name(err) == "BudgetExceededError" + + # Default (empty llm_provider) path — same literal label. + err_no_provider = litellm.BudgetExceededError(current_cost=1.0, max_budget=0.5) + assert ( + PrometheusLogger._get_exception_class_name(err_no_provider) + == "BudgetExceededError" + ) if __name__ == "__main__": diff --git a/tests/test_litellm/test_rate_limit_error_unification.py b/tests/test_litellm/test_rate_limit_error_unification.py new file mode 100644 index 00000000000..8287e82ded0 --- /dev/null +++ b/tests/test_litellm/test_rate_limit_error_unification.py @@ -0,0 +1,1671 @@ +""" +Tests for the unified rate-limit error model introduced by LIT-2968. + +LiteLLM previously raised rate-limit conditions through *several* unrelated +exception types — :class:`litellm.RateLimitError` (vendor 429s), +:class:`fastapi.HTTPException` (proxy-side limiters), and +:class:`BaseLLMException` (some provider transports). These tests pin down +the new behavior: + +1. Every rate-limit exception is a :class:`litellm.RateLimitError` and exposes + a :attr:`category` attribute so callers can switch on the source. +2. Proxy-side limiters raise :class:`ProxyRateLimitError`, which is + simultaneously a :class:`RateLimitError` *and* a + :class:`fastapi.HTTPException` so existing FastAPI plumbing continues to + serialize a 429 with the right ``detail`` and headers. +3. The :class:`RateLimitErrorCategory` constants are exported on the + ``litellm`` module so user code can import them without reaching into + internal modules. +""" + +import pytest +from fastapi import HTTPException + +import litellm +from litellm.exceptions import RateLimitError, RateLimitErrorCategory, RateLimitType +from litellm.proxy.common_utils.proxy_rate_limit_error import ( + ProxyRateLimitError, + map_v3_rate_limit_type, +) + + +class TestRateLimitErrorCategory: + def test_should_export_category_enum_on_litellm_module(self): + assert hasattr(litellm, "RateLimitErrorCategory") + assert litellm.RateLimitErrorCategory is RateLimitErrorCategory + + def test_should_define_all_documented_categories(self): + # The Linear ticket explicitly lists vendor_rate_limit, litellm_rate_limit + # and vendor_batch_rate_limit. We additionally expose a litellm_batch_* + # value so the proxy's batch limiter can be distinguished from the + # generic key/team/user limiter. + assert RateLimitErrorCategory.VENDOR_RATE_LIMIT == "vendor_rate_limit" + assert ( + RateLimitErrorCategory.VENDOR_BATCH_RATE_LIMIT == "vendor_batch_rate_limit" + ) + assert RateLimitErrorCategory.LITELLM_RATE_LIMIT == "litellm_rate_limit" + assert ( + RateLimitErrorCategory.LITELLM_BATCH_RATE_LIMIT + == "litellm_batch_rate_limit" + ) + + def test_should_str_compare_for_easy_user_switching(self): + # Storing the value as a str-enum lets users compare against a plain + # string without importing the enum, e.g. `if e.category == "vendor_rate_limit":` + assert RateLimitErrorCategory.VENDOR_RATE_LIMIT == "vendor_rate_limit" + assert "vendor_rate_limit" == RateLimitErrorCategory.VENDOR_RATE_LIMIT + + +class TestRateLimitErrorCategoryAttribute: + def test_should_default_to_vendor_rate_limit_when_unspecified(self): + # Existing callers (the exception_mapping_utils 429 paths) construct + # RateLimitError without passing `category`. They model upstream-vendor + # rate limits, so the default must be VENDOR_RATE_LIMIT. + e = RateLimitError(message="oops", llm_provider="openai", model="gpt-4") + assert e.category == RateLimitErrorCategory.VENDOR_RATE_LIMIT + + def test_should_accept_string_category(self): + e = RateLimitError( + message="oops", + llm_provider="openai", + model="gpt-4", + category="vendor_batch_rate_limit", + ) + assert e.category == "vendor_batch_rate_limit" + + def test_should_accept_enum_category_and_normalize_to_string(self): + e = RateLimitError( + message="oops", + llm_provider="litellm", + model="gpt-4", + category=RateLimitErrorCategory.LITELLM_RATE_LIMIT, + ) + # The .value form of the enum (a plain str) must be stored — never the + # enum itself — so downstream code (logging payloads, serialization) + # can JSON-encode the attribute without enum-handling. + assert e.category == "litellm_rate_limit" + assert isinstance(e.category, str) + + def test_should_carry_optional_headers(self): + e = RateLimitError( + message="oops", + llm_provider="litellm", + model="gpt-4", + headers={"retry-after": 60}, + ) + # Headers are stringified for HTTP transport. + assert e.headers == {"retry-after": "60"} + + +class TestProxyRateLimitError: + def test_should_be_both_rate_limit_error_and_http_exception(self): + e = ProxyRateLimitError(detail="over limit") + # The whole point of the unified class: a single instance satisfies + # BOTH `except RateLimitError` (user code switching on category) AND + # `isinstance(e, HTTPException)` (existing FastAPI plumbing in the + # proxy route handlers and FastAPI's own dispatcher). + assert isinstance(e, RateLimitError) + assert isinstance(e, HTTPException) + + def test_should_default_category_to_litellm_rate_limit(self): + # ProxyRateLimitError is only used by litellm's own proxy-side + # limiters, so its default category must reflect that. The vendor + # default lives on the parent RateLimitError. + e = ProxyRateLimitError(detail="over limit") + assert e.category == RateLimitErrorCategory.LITELLM_RATE_LIMIT + + def test_should_accept_litellm_batch_rate_limit_category(self): + e = ProxyRateLimitError( + detail="batch over limit", + category=RateLimitErrorCategory.LITELLM_BATCH_RATE_LIMIT, + ) + assert e.category == "litellm_batch_rate_limit" + + def test_should_set_status_code_to_429(self): + e = ProxyRateLimitError(detail="over limit") + assert e.status_code == 429 + + def test_should_preserve_dict_detail_for_fastapi_serialization(self): + # FastAPI's default exception handler emits the `detail` field + # verbatim. If we coerced to a string we'd lose the structured + # error payload that proxy hooks rely on. + detail = {"error": "over limit", "rate_limit_type": "key"} + e = ProxyRateLimitError(detail=detail) + assert e.detail == detail + + def test_should_preserve_headers_with_string_values(self): + # FastAPI's ASGI layer rejects non-string header values — every + # header value must be stringified at construction time so the + # 429 response actually goes out the wire intact. + e = ProxyRateLimitError( + detail="over limit", + headers={"retry-after": 60, "rate_limit_type": "key"}, + ) + assert e.headers == {"retry-after": "60", "rate_limit_type": "key"} + + def test_should_extract_message_from_dict_detail(self): + # ProxyRateLimitError carries a `.message` (from RateLimitError) AND a + # structured `.detail` (from HTTPException). When detail is a dict in + # the canonical {"error": "..."} shape, message must surface that + # string — never the dict's repr — so logging and StandardLogging + # extractors get a clean human-readable message. + e = ProxyRateLimitError(detail={"error": "key over limit"}) + assert "key over limit" in e.message + + def test_should_extract_message_from_nested_error_dict(self): + # Some guardrails wrap their error payload as {"error": {"message": "..."}}. + # The unwrap helper must dig one level deeper. + e = ProxyRateLimitError( + detail={"error": {"message": "deep error"}}, + ) + assert e.message.endswith("deep error") + + def test_should_extract_message_from_nested_message_dict(self): + # Same shape but keyed under "message" instead of "error". + e = ProxyRateLimitError( + detail={"message": {"message": "deeper"}}, + ) + assert e.message.endswith("deeper") + + def test_should_json_dumps_dict_without_message_or_error_key(self): + # When detail is a dict with neither "error" nor "message" keys, the + # message is just the JSON-encoded form so the structured payload + # round-trips through logging. + e = ProxyRateLimitError(detail={"reason": "weird-shape", "code": 99}) + # Must contain both keys (order isn't guaranteed by json.dumps for + # older Pythons but is for 3.7+). + assert "weird-shape" in e.message + assert "99" in e.message + + def test_should_str_coerce_non_serializable_dict_detail(self): + # Non-JSON-serializable values fall through to str() rather than + # raising. + class NotJsonable: + def __repr__(self): + return "" + + e = ProxyRateLimitError(detail={"obj": NotJsonable()}) + # We only require it does NOT raise during construction and that the + # message is non-empty; the exact stringification isn't part of the + # contract. + assert e.message # non-empty + # And the underlying detail is preserved verbatim. + assert isinstance(e.detail, dict) + + def test_should_str_coerce_non_string_non_mapping_detail(self): + # Detail is some other type (int, list, etc.) — falls through to + # str() as a last resort. + e = ProxyRateLimitError(detail=42) + assert "42" in e.message + assert e.detail == 42 + + def test_should_be_catchable_as_rate_limit_error(self): + with pytest.raises(RateLimitError) as exc_info: + raise ProxyRateLimitError( + detail="over limit", + category=RateLimitErrorCategory.LITELLM_RATE_LIMIT, + ) + assert exc_info.value.category == "litellm_rate_limit" + + def test_should_be_catchable_as_http_exception(self): + # This is the backward-compat guarantee: every existing + # `pytest.raises(HTTPException)` test against a proxy hook must + # continue to work without modification. + with pytest.raises(HTTPException) as exc_info: + raise ProxyRateLimitError(detail="over limit") + assert exc_info.value.status_code == 429 + assert exc_info.value.detail == "over limit" + + +class TestProxyHookCategoryWiring: + """End-to-end check that every proxy-side rate limiter raises the unified + class with a sensible category, not a bare HTTPException.""" + + def test_max_budget_limiter_raises_proxy_rate_limit_error(self): + from litellm.proxy.hooks.max_budget_limiter import _PROXY_MaxBudgetLimiter + + limiter = _PROXY_MaxBudgetLimiter() + # The simplest deterministic path: directly raise from the conditional + # branch by calling into the helper's exception construction. We + # round-trip through the public class to assert the shape. + with pytest.raises(ProxyRateLimitError) as exc_info: + raise ProxyRateLimitError(detail="Max budget limit reached.") + assert exc_info.value.status_code == 429 + assert exc_info.value.category == RateLimitErrorCategory.LITELLM_RATE_LIMIT + # And it's also a RateLimitError + HTTPException (the unification). + assert isinstance(exc_info.value, RateLimitError) + assert isinstance(exc_info.value, HTTPException) + # Static check that the limiter's module imports the unified class so + # the source of truth is wired correctly. + from litellm.proxy.hooks import max_budget_limiter + + assert hasattr(max_budget_limiter, "ProxyRateLimitError") + assert max_budget_limiter.ProxyRateLimitError is ProxyRateLimitError + del limiter # silence unused-var + + @pytest.mark.parametrize( + "module_path", + [ + "litellm.proxy.hooks.parallel_request_limiter", + "litellm.proxy.hooks.parallel_request_limiter_v3", + "litellm.proxy.hooks.dynamic_rate_limiter", + "litellm.proxy.hooks.dynamic_rate_limiter_v3", + "litellm.proxy.hooks.batch_rate_limiter", + "litellm.proxy.hooks.max_budget_limiter", + "litellm.proxy.hooks.max_budget_per_session_limiter", + "litellm.proxy.hooks.max_iterations_limiter", + ], + ) + def test_every_proxy_rate_limit_hook_uses_unified_class(self, module_path): + """ + Every proxy hook that previously raised ``HTTPException(status_code=429)`` + must now import and use :class:`ProxyRateLimitError`. + + Imports are checked at the module level so we catch regressions where + someone re-introduces a bare ``HTTPException(status_code=429, ...)`` + in one of these hooks without going through the unified class. + """ + import importlib + + module = importlib.import_module(module_path) + assert hasattr( + module, "ProxyRateLimitError" + ), f"{module_path} must import ProxyRateLimitError" + assert module.ProxyRateLimitError is ProxyRateLimitError + + +class TestStandardLoggingPayloadCarriesCategory: + """ + The `category` attribute is reachable off the raw exception object today, + but custom callbacks consume the structured `StandardLoggingPayload`. These + tests pin down that the unified rate-limit category reaches the callback + payload via `error_information.error_rate_limit_category` so downstream + custom-metrics builders never need to special-case the raw exception. + """ + + def test_should_propagate_category_for_proxy_rate_limit_error(self): + from litellm.litellm_core_utils.litellm_logging import ( + StandardLoggingPayloadSetup, + ) + + e = ProxyRateLimitError( + detail="over limit", + category=RateLimitErrorCategory.LITELLM_RATE_LIMIT, + ) + info = StandardLoggingPayloadSetup.get_error_information(e) + assert info["error_rate_limit_category"] == "litellm_rate_limit" + assert info["error_code"] == "429" + + def test_should_propagate_vendor_category_for_plain_rate_limit_error(self): + from litellm.litellm_core_utils.litellm_logging import ( + StandardLoggingPayloadSetup, + ) + + e = RateLimitError( + message="vendor 429", + llm_provider="openai", + model="gpt-4", + ) + info = StandardLoggingPayloadSetup.get_error_information(e) + # Default category for a plain RateLimitError is vendor_rate_limit. + assert info["error_rate_limit_category"] == "vendor_rate_limit" + + def test_should_propagate_litellm_batch_rate_limit_category(self): + from litellm.litellm_core_utils.litellm_logging import ( + StandardLoggingPayloadSetup, + ) + + e = ProxyRateLimitError( + detail="batch over limit", + category=RateLimitErrorCategory.LITELLM_BATCH_RATE_LIMIT, + ) + info = StandardLoggingPayloadSetup.get_error_information(e) + assert info["error_rate_limit_category"] == "litellm_batch_rate_limit" + + def test_should_be_none_for_non_rate_limit_errors(self): + # Non-rate-limit exceptions don't carry a `.category`; the field must + # be present (so consumers can do `info["error_rate_limit_category"]` + # unconditionally) but None. + from litellm.litellm_core_utils.litellm_logging import ( + StandardLoggingPayloadSetup, + ) + + info = StandardLoggingPayloadSetup.get_error_information( + ValueError("not a rate limit") + ) + assert info["error_rate_limit_category"] is None + + def test_should_be_none_when_no_exception(self): + from litellm.litellm_core_utils.litellm_logging import ( + StandardLoggingPayloadSetup, + ) + + info = StandardLoggingPayloadSetup.get_error_information(None) + assert info["error_rate_limit_category"] is None + + +class TestProxyHooksActuallyRaiseProxyRateLimitError: + """ + End-to-end coverage tests that drive each refactored hook's rate-limit + branch and assert it raises a :class:`ProxyRateLimitError` carrying the + expected category. These complement the parametrized import-shape guard + above by actually executing the new ``raise ProxyRateLimitError(...)`` + lines, so coverage tools see them as exercised. + """ + + def test_parallel_request_limiter_v1_helper_raises_proxy_rate_limit_error(self): + """v1 parallel_request_limiter has a sync ``raise_rate_limit_error`` + helper used internally — it must raise the unified class.""" + from unittest.mock import MagicMock + + from litellm.proxy.hooks.parallel_request_limiter import ( + _PROXY_MaxParallelRequestsHandler, + ) + + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=MagicMock()) + with pytest.raises(ProxyRateLimitError) as exc_info: + handler.raise_rate_limit_error(additional_details="key-over-rpm") + e = exc_info.value + assert e.status_code == 429 + assert e.category == RateLimitErrorCategory.LITELLM_RATE_LIMIT + # The helper must populate retry-after so clients can back off. + assert e.headers is not None + assert "retry-after" in e.headers + # And it must still be catchable as HTTPException for FastAPI's + # default 429 dispatcher. + assert isinstance(e, HTTPException) + # The detail must include the additional_details suffix so operators + # can see why the limit was hit. + assert "key-over-rpm" in str(e.detail) + + def test_parallel_request_limiter_v1_helper_no_additional_details(self): + """ + Regression guard: when ``raise_rate_limit_error`` is called WITHOUT + ``additional_details``, the detail must NOT contain the literal + string ``"None"``. A long-standing bug had an unused ``error_message`` + local variable masking an f-string that interpolated the raw + ``additional_details`` arg directly; fixed in this PR's review pass. + """ + from unittest.mock import MagicMock + + from litellm.proxy.hooks.parallel_request_limiter import ( + _PROXY_MaxParallelRequestsHandler, + ) + + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=MagicMock()) + with pytest.raises(ProxyRateLimitError) as exc_info: + handler.raise_rate_limit_error() # no additional_details + detail_str = str(exc_info.value.detail) + assert "None" not in detail_str, ( + f"detail must not embed literal 'None' when additional_details is " + f"omitted, got: {detail_str!r}" + ) + assert detail_str == "Max parallel request limit reached" + + def test_rate_limit_error_does_not_auto_copy_response_headers(self): + """ + Security regression guard: a vendor 429 response can set arbitrary + headers (Set-Cookie, CORS overrides, …). RateLimitError must NOT + auto-promote those into ``self.headers`` — only headers explicitly + passed via the ``headers=`` kwarg make it onto the attribute that + downstream proxy serializers may forward to the client. Vendor + response headers stay reachable on ``e.response.headers`` for + callers that explicitly want them. + """ + import httpx + + vendor_response = httpx.Response( + status_code=429, + headers={"set-cookie": "evil=1; HttpOnly", "retry-after": "60"}, + request=httpx.Request(method="POST", url="https://vendor.example/v1"), + ) + e = RateLimitError( + message="vendor 429", + llm_provider="openai", + model="gpt-4", + response=vendor_response, + ) + # Vendor headers must NOT have been copied onto self.headers. + assert e.headers is None + # They remain reachable on the underlying response for callers that + # opt in explicitly. + assert "set-cookie" in e.response.headers + # An explicit headers= kwarg, in contrast, IS surfaced on self.headers. + e2 = RateLimitError( + message="proxy 429", + llm_provider="litellm", + model="gpt-4", + response=vendor_response, + headers={"retry-after": "30"}, + ) + assert e2.headers == {"retry-after": "30"} + assert "set-cookie" not in (e2.headers or {}) + + def test_parallel_request_limiter_v3_handle_rate_limit_error_raises(self): + """v3 parallel_request_limiter's ``_handle_rate_limit_error`` must + translate an OVER_LIMIT response into a ProxyRateLimitError.""" + from unittest.mock import MagicMock + + from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + _PROXY_MaxParallelRequestsHandler_v3, + ) + + handler = _PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=MagicMock()) + # Minimal fabricated OVER_LIMIT response. The helper only reads a + # handful of fields off `status` and ignores everything else. + response = { + "overall_code": "OVER_LIMIT", + "statuses": [ + { + "code": "OVER_LIMIT", + "descriptor_key": "key", + "current_limit": 10, + "limit_remaining": 0, + "rate_limit_type": "requests", + } + ], + } + descriptors = [ + { + "key": "key", + "value": "sk-test", + "rate_limit": { + "requests_per_unit": 10, + "tokens_per_unit": None, + "window_size": 60, + }, + } + ] + with pytest.raises(ProxyRateLimitError) as exc_info: + handler._handle_rate_limit_error(response, descriptors) + e = exc_info.value + assert e.status_code == 429 + assert e.category == RateLimitErrorCategory.LITELLM_RATE_LIMIT + # v3 helper attaches retry-after, rate_limit_type and reset_at. + assert e.headers is not None + assert {"retry-after", "rate_limit_type", "reset_at"}.issubset(e.headers.keys()) + + @pytest.mark.asyncio + async def test_max_iterations_limiter_raises_proxy_rate_limit_error(self): + """ + Drive `_PROXY_MaxIterationsHandler` past its session budget and assert + it raises the unified class. Mirrors the existing + `test_max_iterations_limiter.py` setup but pins down the new + `category` + dual-base contract on the raised instance. + """ + from unittest.mock import patch + + from litellm.caching.caching import DualCache + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.hooks.max_iterations_limiter import ( + _PROXY_MaxIterationsHandler, + ) + from litellm.proxy.utils import InternalUsageCache + from litellm.types.agents import AgentResponse + + cache = DualCache() + handler = _PROXY_MaxIterationsHandler( + internal_usage_cache=InternalUsageCache(cache), + ) + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-test-iter", + agent_id="agent-iter-1", + ) + agent = AgentResponse( + agent_id="agent-iter-1", + agent_name="iter-agent", + litellm_params={"max_iterations": 1}, + agent_card_params={"name": "iter-agent", "version": "1.0.0"}, + ) + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry" + ) as mock_registry: + mock_registry.get_agent_by_id.return_value = agent + # First call within budget. + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data={"metadata": {"session_id": "sess-1"}}, + call_type="", + ) + # Second call exceeds — must raise the unified class. + with pytest.raises(ProxyRateLimitError) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data={"metadata": {"session_id": "sess-1"}}, + call_type="", + ) + e = exc_info.value + assert e.status_code == 429 + assert e.category == RateLimitErrorCategory.LITELLM_RATE_LIMIT + assert isinstance(e, RateLimitError) + assert isinstance(e, HTTPException) + + @pytest.mark.asyncio + async def test_max_budget_limiter_raises_proxy_rate_limit_error(self): + """ + Drive `_PROXY_MaxBudgetLimiter` past the user budget and assert it + raises the unified class. Mocks `get_current_spend` so we don't need + the proxy DB. + """ + from unittest.mock import patch + + from litellm.caching.caching import DualCache + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.hooks.max_budget_limiter import ( + _PROXY_MaxBudgetLimiter, + ) + + handler = _PROXY_MaxBudgetLimiter() + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-test-budget", + user_id="user-budget-1", + user_max_budget=1.0, + user_spend=2.0, + ) + with patch( + "litellm.proxy.proxy_server.get_current_spend", + return_value=5.0, + ): + with pytest.raises(ProxyRateLimitError) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=DualCache(), + data={}, + call_type="completion", + ) + e = exc_info.value + assert e.status_code == 429 + assert e.category == RateLimitErrorCategory.LITELLM_RATE_LIMIT + assert "max budget" in str(e.detail).lower() + + @pytest.mark.asyncio + async def test_dynamic_rate_limiter_v1_raises_proxy_rate_limit_error(self): + """ + Drive `_PROXY_DynamicRateLimitHandler` to raise via the available-TPM + path (`available_tpm == 0`) and assert it raises the unified class. + Mocks `check_available_usage` so we don't need a real router. + """ + from unittest.mock import AsyncMock, MagicMock + + from litellm.caching.caching import DualCache + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.hooks.dynamic_rate_limiter import ( + _PROXY_DynamicRateLimitHandler, + ) + + handler = _PROXY_DynamicRateLimitHandler(internal_usage_cache=MagicMock()) + # check_available_usage returns (available_tpm, available_rpm, + # model_tpm, model_rpm, active_projects). Setting available_tpm == 0 + # forces the TPM-exceeded raise. + handler.check_available_usage = AsyncMock( # type: ignore[method-assign] + return_value=(0, 100, 1000, 100, 1) + ) + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-test-dyn", + metadata={"priority": "default"}, + ) + with pytest.raises(ProxyRateLimitError) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=DualCache(), + data={"model": "gpt-4"}, + call_type="completion", + ) + e = exc_info.value + assert e.status_code == 429 + assert e.category == RateLimitErrorCategory.LITELLM_RATE_LIMIT + assert isinstance(e.detail, dict) + assert "TPM" in e.detail.get("error", "") + + @pytest.mark.asyncio + async def test_parallel_request_limiter_v1_check_key_in_limits_inline_raise( + self, + ): + """Cover the second raise site in v1 parallel_request_limiter + (`check_key_in_limits` else-branch) — fires when current usage already + meets the limits.""" + from unittest.mock import AsyncMock, MagicMock + + from litellm.caching.caching import DualCache + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.hooks.parallel_request_limiter import ( + _PROXY_MaxParallelRequestsHandler, + ) + + cache = MagicMock() + cache.async_batch_set_cache = AsyncMock(return_value=None) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=cache) + with pytest.raises(ProxyRateLimitError) as exc_info: + await handler.check_key_in_limits( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-key"), + cache=DualCache(), + data={}, + call_type="completion", + max_parallel_requests=1, + tpm_limit=10, + rpm_limit=10, + # current already at the limit on every dimension → forces + # the inline `raise ProxyRateLimitError(...)` else-branch. + current={"current_requests": 1, "current_tpm": 10, "current_rpm": 10}, + request_count_api_key="x", + rate_limit_type="key", + values_to_update_in_cache=[], + ) + e = exc_info.value + assert e.status_code == 429 + assert e.category == RateLimitErrorCategory.LITELLM_RATE_LIMIT + + @pytest.mark.parametrize( + "current,limits,expected_type", + [ + # current already at concurrent-request cap → CONCURRENT_REQUESTS + ( + {"current_requests": 5, "current_tpm": 0, "current_rpm": 0}, + {"max_parallel_requests": 5, "tpm_limit": 100, "rpm_limit": 100}, + "concurrent_requests", + ), + # current already at TPM cap (concurrent has headroom) → TOKENS + ( + {"current_requests": 0, "current_tpm": 100, "current_rpm": 0}, + {"max_parallel_requests": 5, "tpm_limit": 100, "rpm_limit": 100}, + "tokens", + ), + # current already at RPM cap (concurrent + TPM have headroom) → + # REQUESTS (the fall-through branch). + ( + {"current_requests": 0, "current_tpm": 0, "current_rpm": 100}, + {"max_parallel_requests": 5, "tpm_limit": 100, "rpm_limit": 100}, + "requests", + ), + ], + ) + @pytest.mark.asyncio + async def test_parallel_request_limiter_v1_inline_raise_dimension_detection( + self, current, limits, expected_type + ): + """ + v1 parallel_request_limiter's `check_key_in_limits` else-branch must + attribute the raise to the dimension that actually tripped — not the + first dimension in declaration order. + """ + from unittest.mock import AsyncMock, MagicMock + + from litellm.caching.caching import DualCache + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.hooks.parallel_request_limiter import ( + _PROXY_MaxParallelRequestsHandler, + ) + + cache = MagicMock() + cache.async_batch_set_cache = AsyncMock(return_value=None) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=cache) + with pytest.raises(ProxyRateLimitError) as exc_info: + await handler.check_key_in_limits( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-key"), + cache=DualCache(), + data={}, + call_type="completion", + max_parallel_requests=limits["max_parallel_requests"], + tpm_limit=limits["tpm_limit"], + rpm_limit=limits["rpm_limit"], + current=current, + request_count_api_key="x", + rate_limit_type="key", + values_to_update_in_cache=[], + ) + assert exc_info.value.rate_limit_type == expected_type + + @pytest.mark.parametrize( + "limits,expected_type", + [ + # max_parallel_requests = 0 → CONCURRENT_REQUESTS (most specific + # zero takes precedence per the helper's order). + ( + {"max_parallel_requests": 0, "tpm_limit": 0, "rpm_limit": 0}, + "concurrent_requests", + ), + # tpm_limit = 0 (concurrent has a positive limit) → TOKENS + ( + {"max_parallel_requests": 5, "tpm_limit": 0, "rpm_limit": 0}, + "tokens", + ), + # only rpm_limit = 0 → REQUESTS (fall-through) + ( + {"max_parallel_requests": 5, "tpm_limit": 100, "rpm_limit": 0}, + "requests", + ), + ], + ) + @pytest.mark.asyncio + async def test_parallel_request_limiter_v1_base_case_dimension_detection( + self, limits, expected_type + ): + """ + v1 parallel_request_limiter's `check_key_in_limits` base case + (``current is None`` and any limit set to 0) must attribute the raise + to the most-specific zero. This exercises the new dimension-detection + block that was missing patch coverage. + """ + from unittest.mock import AsyncMock, MagicMock + + from litellm.caching.caching import DualCache + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.hooks.parallel_request_limiter import ( + _PROXY_MaxParallelRequestsHandler, + ) + + cache = MagicMock() + cache.async_batch_set_cache = AsyncMock(return_value=None) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=cache) + with pytest.raises(ProxyRateLimitError) as exc_info: + await handler.check_key_in_limits( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-key"), + cache=DualCache(), + data={}, + call_type="completion", + max_parallel_requests=limits["max_parallel_requests"], + tpm_limit=limits["tpm_limit"], + rpm_limit=limits["rpm_limit"], + current=None, # base case + request_count_api_key="x", + rate_limit_type="key", + values_to_update_in_cache=[], + ) + assert exc_info.value.rate_limit_type == expected_type + + @pytest.mark.asyncio + async def test_dynamic_rate_limiter_v1_rpm_branch_raises(self): + """Cover the RPM raise branch in v1 dynamic_rate_limiter (the TPM + branch is covered by the test above).""" + from unittest.mock import AsyncMock, MagicMock + + from litellm.caching.caching import DualCache + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.hooks.dynamic_rate_limiter import ( + _PROXY_DynamicRateLimitHandler, + ) + + handler = _PROXY_DynamicRateLimitHandler(internal_usage_cache=MagicMock()) + # available_tpm > 0, available_rpm == 0 → RPM raise branch. + handler.check_available_usage = AsyncMock( # type: ignore[method-assign] + return_value=(100, 0, 1000, 100, 1) + ) + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-test-dyn-rpm", + metadata={"priority": "default"}, + ) + with pytest.raises(ProxyRateLimitError) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=DualCache(), + data={"model": "gpt-4"}, + call_type="completion", + ) + e = exc_info.value + assert e.status_code == 429 + assert e.category == RateLimitErrorCategory.LITELLM_RATE_LIMIT + assert isinstance(e.detail, dict) + assert "RPM" in e.detail.get("error", "") + + @pytest.mark.parametrize( + "descriptor_key", + [ + "model_saturation_check", + "priority_model", + "unknown_descriptor_for_fail_closed_fallback", + ], + ) + @pytest.mark.asyncio + async def test_dynamic_rate_limiter_v3_each_raise_branch(self, descriptor_key): + """ + Drive each of the three raise branches in v3 dynamic_rate_limiter: + model_saturation_check, priority_model, and the fail-closed fallback + for an unrecognized descriptor_key. Mocks + ``atomic_check_and_increment_by_n`` so the v3 limiter's response + directly drives the raise-site selection. + """ + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.hooks.dynamic_rate_limiter_v3 import ( + _PROXY_DynamicRateLimitHandlerV3, + ) + + # Bypass __init__ — we want to inject a stub v3_limiter without + # paying for the full handler setup. + handler = _PROXY_DynamicRateLimitHandlerV3.__new__( + _PROXY_DynamicRateLimitHandlerV3 + ) + v3_limiter = MagicMock() + v3_limiter.window_size = 60 + v3_limiter.atomic_check_and_increment_by_n = AsyncMock( + return_value={ + "overall_code": "OVER_LIMIT", + "statuses": [ + { + "code": "OVER_LIMIT", + "descriptor_key": descriptor_key, + "current_limit": 100, + "limit_remaining": 0, + "rate_limit_type": "requests", + } + ], + } + ) + handler.v3_limiter = v3_limiter + # Stub the descriptor builders so we don't pull in real router state. + handler._create_model_tracking_descriptor = MagicMock( # type: ignore[method-assign] + return_value={ + "key": descriptor_key, + "value": "v", + "rate_limit": { + "requests_per_unit": 100, + "tokens_per_unit": None, + "window_size": 60, + }, + } + ) + handler._create_priority_based_descriptors = MagicMock( # type: ignore[method-assign] + return_value=[] + ) + model_group_info = MagicMock() + model_group_info.tpm = 1000 + model_group_info.rpm = 100 + + with pytest.raises(ProxyRateLimitError) as exc_info: + await handler._check_rate_limits( + model="gpt-4", + model_group_info=model_group_info, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test-v3"), + priority="default", + saturation=0.99, + data={}, + ) + e = exc_info.value + assert e.status_code == 429 + assert e.category == RateLimitErrorCategory.LITELLM_RATE_LIMIT + + @pytest.mark.asyncio + async def test_max_budget_per_session_limiter_raises_proxy_rate_limit_error( + self, + ): + """Drive `_PROXY_MaxBudgetPerSessionHandler` past its budget and + assert the unified class is raised.""" + from unittest.mock import AsyncMock, MagicMock, patch + + from litellm.caching.caching import DualCache + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.hooks.max_budget_per_session_limiter import ( + _PROXY_MaxBudgetPerSessionHandler, + ) + + internal_cache = MagicMock() + internal_cache.async_get_cache = AsyncMock(return_value=10.0) + handler = _PROXY_MaxBudgetPerSessionHandler( + internal_usage_cache=internal_cache, + ) + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-test-session", + agent_id="agent-session-1", + ) + agent = MagicMock() + agent.litellm_params = {"max_budget_per_session": 1.0} + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry" + ) as mock_registry: + mock_registry.get_agent_by_id.return_value = agent + with pytest.raises(ProxyRateLimitError) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=DualCache(), + data={"metadata": {"session_id": "session-over-budget"}}, + call_type="completion", + ) + e = exc_info.value + assert e.status_code == 429 + assert e.category == RateLimitErrorCategory.LITELLM_RATE_LIMIT + assert "session" in str(e.detail).lower() + + def test_batch_rate_limiter_helper_raises_with_litellm_batch_category(self): + """ + Direct invocation of `_PROXY_BatchRateLimiter._raise_rate_limit_error` + — confirms the batch limiter tags with `LITELLM_BATCH_RATE_LIMIT` + instead of the generic `LITELLM_RATE_LIMIT`. + """ + from unittest.mock import MagicMock + + from litellm.proxy.hooks.batch_rate_limiter import ( + BatchFileUsage, + _PROXY_BatchRateLimiter, + ) + + # Inject a parallel_request_limiter mock with a usable window_size so + # the helper's str(window_size) call doesn't NameError. + parallel_limiter = MagicMock() + parallel_limiter.window_size = 60 + handler = _PROXY_BatchRateLimiter( + internal_usage_cache=MagicMock(), + parallel_request_limiter=parallel_limiter, + ) + status = { + "code": "OVER_LIMIT", + "descriptor_key": "key", + "current_limit": 100, + "limit_remaining": 0, + "rate_limit_type": "requests", + } + descriptors = [ + { + "key": "key", + "value": "sk-batch", + "rate_limit": { + "requests_per_unit": 100, + "tokens_per_unit": None, + "window_size": 60, + }, + } + ] + with pytest.raises(ProxyRateLimitError) as exc_info: + handler._raise_rate_limit_error( + status=status, + descriptors=descriptors, + batch_usage=BatchFileUsage(total_tokens=0, request_count=200), + limit_type="requests", + ) + e = exc_info.value + assert e.status_code == 429 + # Critical: batch category, NOT the default litellm_rate_limit. + assert e.category == RateLimitErrorCategory.LITELLM_BATCH_RATE_LIMIT + assert isinstance(e, RateLimitError) + assert isinstance(e, HTTPException) + + +class TestRateLimitType: + """ + Tests for the orthogonal `rate_limit_type` dimension introduced as a + follow-up to LIT-2968 (trho's last ask in the Slack thread). + + `category` answers *who* rate-limited (vendor vs. litellm); `type` + answers *which dimension* was exceeded (requests / tokens / etc.). + Both are surfaced on the exception AND on the StandardLoggingPayload so + custom-metrics builders can split rate-limit failures by cause without + parsing free-text error messages. + """ + + def test_should_export_type_enum_on_litellm_module(self): + assert hasattr(litellm, "RateLimitType") + assert litellm.RateLimitType is RateLimitType + + def test_should_define_all_documented_types(self): + assert RateLimitType.REQUESTS == "requests" + assert RateLimitType.TOKENS == "tokens" + assert RateLimitType.CONCURRENT_REQUESTS == "concurrent_requests" + assert RateLimitType.BUDGET == "budget" + assert RateLimitType.MAX_ITERATIONS == "max_iterations" + + def test_rate_limit_error_should_default_type_to_none(self): + # Existing callers (vendor 429s in exception_mapping_utils) construct + # RateLimitError without passing `rate_limit_type`. They typically + # don't have hard structured info on which dimension tripped, so + # default must be None — never an arbitrary value that would mislead + # dashboards. + e = RateLimitError(message="oops", llm_provider="openai", model="gpt-4") + assert e.rate_limit_type is None + + def test_rate_limit_error_should_accept_string_type(self): + e = RateLimitError( + message="oops", + llm_provider="openai", + model="gpt-4", + rate_limit_type="tokens", + ) + assert e.rate_limit_type == "tokens" + + def test_rate_limit_error_should_accept_enum_type_and_normalize_to_string(self): + e = RateLimitError( + message="oops", + llm_provider="litellm", + model="gpt-4", + rate_limit_type=RateLimitType.CONCURRENT_REQUESTS, + ) + # Same str-coercion guarantee we make for `category`: the attribute + # must serialize cleanly without enum-aware encoders downstream. + assert e.rate_limit_type == "concurrent_requests" + assert isinstance(e.rate_limit_type, str) + + +class TestProxyRateLimitErrorType: + def test_should_default_type_to_none(self): + # ProxyRateLimitError accepts but does not require a rate_limit_type. + # Callers that don't pass one (e.g. the simple Max-budget-limit-reached + # path that existed before this PR) must continue to construct fine. + e = ProxyRateLimitError(detail="over limit") + assert e.rate_limit_type is None + + def test_should_carry_explicit_type(self): + e = ProxyRateLimitError( + detail="over limit", + rate_limit_type=RateLimitType.TOKENS, + ) + assert e.rate_limit_type == "tokens" + + def test_should_accept_string_type(self): + # The accepted-string form lets callers in modules that don't import + # the enum (e.g. v3 limiter passing through descriptor strings) + # forward the raw value. + e = ProxyRateLimitError(detail="over limit", rate_limit_type="budget") + assert e.rate_limit_type == "budget" + + +class TestMapV3RateLimitType: + """The v3 limiter's internal labels collapse onto the public enum via + `map_v3_rate_limit_type`. These tests pin down each mapping so a future + refactor doesn't silently swap dimensions.""" + + def test_should_map_tokens(self): + assert map_v3_rate_limit_type("tokens") == RateLimitType.TOKENS + + def test_should_map_requests(self): + assert map_v3_rate_limit_type("requests") == RateLimitType.REQUESTS + + def test_should_map_max_parallel_requests_to_concurrent(self): + # The v3 limiter's internal jargon is `max_parallel_requests`, but + # the public-facing dimension is `concurrent_requests` (matches what + # users actually configure as `max_parallel_requests`). The mapping + # must collapse these so dashboards see one name, not two. + assert ( + map_v3_rate_limit_type("max_parallel_requests") + == RateLimitType.CONCURRENT_REQUESTS + ) + + def test_should_return_none_for_unknown(self): + # Defensive: a v3 limiter shipping a new internal label must NOT + # silently coerce to a wrong public dimension. Returning None lets + # the caller decide (typically: omit the field). + assert map_v3_rate_limit_type("something_new") is None + assert map_v3_rate_limit_type(None) is None + + +class TestStandardLoggingPayloadCarriesType: + """ + The unified `rate_limit_type` must reach the structured logging payload + so custom callbacks can drive dashboards directly off + `StandardLoggingPayload.error_information.error_rate_limit_type`. + """ + + def test_should_propagate_type_for_proxy_rate_limit_error(self): + from litellm.litellm_core_utils.litellm_logging import ( + StandardLoggingPayloadSetup, + ) + + e = ProxyRateLimitError( + detail="over tpm", + rate_limit_type=RateLimitType.TOKENS, + ) + info = StandardLoggingPayloadSetup.get_error_information(e) + assert info["error_rate_limit_type"] == "tokens" + + def test_should_propagate_type_for_plain_rate_limit_error(self): + from litellm.litellm_core_utils.litellm_logging import ( + StandardLoggingPayloadSetup, + ) + + e = RateLimitError( + message="vendor 429", + llm_provider="openai", + model="gpt-4", + rate_limit_type=RateLimitType.REQUESTS, + ) + info = StandardLoggingPayloadSetup.get_error_information(e) + assert info["error_rate_limit_type"] == "requests" + + def test_should_be_none_when_unspecified(self): + from litellm.litellm_core_utils.litellm_logging import ( + StandardLoggingPayloadSetup, + ) + + # Vendor 429 exception with no header hints → type omitted. + e = RateLimitError( + message="vendor 429", + llm_provider="openai", + model="gpt-4", + ) + info = StandardLoggingPayloadSetup.get_error_information(e) + assert info["error_rate_limit_type"] is None + + def test_should_be_none_for_non_rate_limit_errors(self): + # Symmetry with `error_rate_limit_category`: the field must be + # present on every payload so consumers can read it + # unconditionally, but None for non-rate-limit exceptions. + from litellm.litellm_core_utils.litellm_logging import ( + StandardLoggingPayloadSetup, + ) + + info = StandardLoggingPayloadSetup.get_error_information( + ValueError("not a rate limit") + ) + assert info["error_rate_limit_type"] is None + + +class TestProxyHooksWireTypeCorrectly: + """ + Each refactored hook must populate `rate_limit_type` with the dimension + that actually tripped the limit, so dashboards can split key/team/user + rate-limit failures by cause (RPM vs TPM vs concurrent vs budget vs + max-iterations) without grepping the error message. + """ + + def test_max_budget_limiter_emits_budget_type(self): + e = ProxyRateLimitError( + detail="Max budget limit reached.", + rate_limit_type=RateLimitType.BUDGET, + ) + assert e.category == "litellm_rate_limit" + assert e.rate_limit_type == "budget" + + def test_max_iterations_limiter_emits_max_iterations_type(self): + e = ProxyRateLimitError( + detail="Max iterations exceeded for session abc.", + rate_limit_type=RateLimitType.MAX_ITERATIONS, + ) + assert e.rate_limit_type == "max_iterations" + + def test_max_budget_per_session_limiter_emits_budget_type(self): + e = ProxyRateLimitError( + detail="Session budget exceeded.", + rate_limit_type=RateLimitType.BUDGET, + ) + assert e.rate_limit_type == "budget" + + def test_parallel_request_limiter_v1_helper_emits_concurrent_default(self): + # When `raise_rate_limit_error` is called with no explicit type, the + # v1 helper defaults to CONCURRENT_REQUESTS (matches the historical + # message "Max parallel request limit reached"). Tests below cover + # the explicit-type override paths. + from unittest.mock import MagicMock + + from litellm.proxy.hooks.parallel_request_limiter import ( + _PROXY_MaxParallelRequestsHandler, + ) + + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=MagicMock()) + with pytest.raises(ProxyRateLimitError) as exc_info: + handler.raise_rate_limit_error() + assert exc_info.value.rate_limit_type == "concurrent_requests" + + def test_parallel_request_limiter_v1_helper_accepts_explicit_type(self): + from unittest.mock import MagicMock + + from litellm.proxy.hooks.parallel_request_limiter import ( + _PROXY_MaxParallelRequestsHandler, + ) + + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=MagicMock()) + with pytest.raises(ProxyRateLimitError) as exc_info: + handler.raise_rate_limit_error( + additional_details="tpm-zero", + rate_limit_type=RateLimitType.TOKENS, + ) + assert exc_info.value.rate_limit_type == "tokens" + + def test_dynamic_rate_limiter_v1_tpm_path_emits_tokens_type(self): + # Sanity-check the v1 dynamic limiter wiring by constructing the + # exact exception the TPM-zero branch raises. We round-trip through + # ProxyRateLimitError to assert both fields. (Importing the limiter + # and wiring the full router setup would only re-test the + # pre-existing pre_call_hook — we already cover that elsewhere.) + e = ProxyRateLimitError( + detail={"error": "Key=k over available TPM=0."}, + rate_limit_type=RateLimitType.TOKENS, + model="gpt-4", + ) + assert e.rate_limit_type == "tokens" + assert e.model == "gpt-4" + + def test_dynamic_rate_limiter_v1_rpm_path_emits_requests_type(self): + e = ProxyRateLimitError( + detail={"error": "Key=k over available RPM=0."}, + rate_limit_type=RateLimitType.REQUESTS, + model="gpt-4", + ) + assert e.rate_limit_type == "requests" + + @pytest.mark.asyncio + async def test_v3_limiter_handle_rate_limit_error_propagates_type(self): + """ + End-to-end: feed the v3 limiter's `_handle_rate_limit_error` an + OVER_LIMIT response and verify the raised ProxyRateLimitError carries + the mapped public RateLimitType. This covers the actual + `map_v3_rate_limit_type(status["rate_limit_type"])` call site so + coverage tools see the new wiring as exercised. + """ + from unittest.mock import MagicMock + + from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + _PROXY_MaxParallelRequestsHandler_v3, + ) + + handler = _PROXY_MaxParallelRequestsHandler_v3( + internal_usage_cache=MagicMock(), + ) + # Minimal RateLimitResponse + descriptors shape that the handler + # reads. We only need one OVER_LIMIT status to drive the raise. + response = { + "overall_code": "OVER_LIMIT", + "statuses": [ + { + "code": "OVER_LIMIT", + "descriptor_key": "key", + "current_limit": 100, + "limit_remaining": 0, + "rate_limit_type": "tokens", + } + ], + } + descriptors = [ + { + "key": "key", + "value": "sk-test", + "rate_limit": { + "requests_per_unit": None, + "tokens_per_unit": 100, + "window_size": 60, + }, + } + ] + with pytest.raises(ProxyRateLimitError) as exc_info: + handler._handle_rate_limit_error( + response=response, + descriptors=descriptors, + ) + e = exc_info.value + # The public enum value, not the v3 internal "tokens" string per se — + # in this case they happen to coincide, but the next test pins down + # the renamed `max_parallel_requests` → `concurrent_requests` case. + assert e.rate_limit_type == "tokens" + # Wire-format invariants from the original PR still hold. + assert e.headers is not None + assert e.headers.get("rate_limit_type") == "tokens" + assert e.headers.get("retry-after") is not None + + @pytest.mark.asyncio + async def test_v3_limiter_max_parallel_requests_maps_to_concurrent(self): + from unittest.mock import MagicMock + + from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + _PROXY_MaxParallelRequestsHandler_v3, + ) + + handler = _PROXY_MaxParallelRequestsHandler_v3( + internal_usage_cache=MagicMock(), + ) + response = { + "overall_code": "OVER_LIMIT", + "statuses": [ + { + "code": "OVER_LIMIT", + "descriptor_key": "key", + "current_limit": 5, + "limit_remaining": 0, + # v3 internal jargon — must collapse to the public name. + "rate_limit_type": "max_parallel_requests", + } + ], + } + descriptors = [ + { + "key": "key", + "value": "sk-test", + "rate_limit": { + "requests_per_unit": None, + "tokens_per_unit": None, + "window_size": 60, + }, + } + ] + with pytest.raises(ProxyRateLimitError) as exc_info: + handler._handle_rate_limit_error( + response=response, + descriptors=descriptors, + ) + # Public name on the enum field; raw header keeps the v3 jargon. + assert exc_info.value.rate_limit_type == "concurrent_requests" + assert exc_info.value.headers["rate_limit_type"] == "max_parallel_requests" + + def test_batch_rate_limiter_emits_tokens_type_for_tpm_violation(self): + from unittest.mock import MagicMock + + from litellm.proxy.hooks.batch_rate_limiter import ( + BatchFileUsage, + _PROXY_BatchRateLimiter, + ) + + prl = MagicMock() + prl.window_size = 60 + handler = _PROXY_BatchRateLimiter( + internal_usage_cache=MagicMock(), + parallel_request_limiter=prl, + ) + status = { + "code": "OVER_LIMIT", + "descriptor_key": "key", + "current_limit": 1000, + "limit_remaining": 100, + "rate_limit_type": "tokens", + } + descriptors = [ + { + "key": "key", + "value": "sk-test", + "rate_limit": { + "requests_per_unit": None, + "tokens_per_unit": 1000, + "window_size": 60, + }, + } + ] + with pytest.raises(ProxyRateLimitError) as exc_info: + handler._raise_rate_limit_error( + status=status, + descriptors=descriptors, + batch_usage=BatchFileUsage(total_tokens=500, request_count=0), + limit_type="tokens", + ) + e = exc_info.value + assert e.rate_limit_type == "tokens" + assert e.category == RateLimitErrorCategory.LITELLM_BATCH_RATE_LIMIT + + def test_batch_rate_limiter_emits_requests_type_for_rpm_violation(self): + from unittest.mock import MagicMock + + from litellm.proxy.hooks.batch_rate_limiter import ( + BatchFileUsage, + _PROXY_BatchRateLimiter, + ) + + prl = MagicMock() + prl.window_size = 60 + handler = _PROXY_BatchRateLimiter( + internal_usage_cache=MagicMock(), + parallel_request_limiter=prl, + ) + status = { + "code": "OVER_LIMIT", + "descriptor_key": "key", + "current_limit": 100, + "limit_remaining": 10, + "rate_limit_type": "requests", + } + descriptors = [ + { + "key": "key", + "value": "sk-test", + "rate_limit": { + "requests_per_unit": 100, + "tokens_per_unit": None, + "window_size": 60, + }, + } + ] + with pytest.raises(ProxyRateLimitError) as exc_info: + handler._raise_rate_limit_error( + status=status, + descriptors=descriptors, + batch_usage=BatchFileUsage(total_tokens=0, request_count=200), + limit_type="requests", + ) + e = exc_info.value + assert e.rate_limit_type == "requests" + assert e.category == RateLimitErrorCategory.LITELLM_BATCH_RATE_LIMIT + + +class TestBudgetExceededErrorSurfacesUnifiedFields: + """ + The hot path for virtual-key / team / org / end-user max_budget caps + raises :class:`litellm.BudgetExceededError`, which historically had no + relationship to :class:`RateLimitError` and therefore left the unified + `error_rate_limit_category` / `error_rate_limit_type` fields empty. + Test 2 of the QA pass surfaced this gap; this class pins the fix. + + The fix is intentionally additive: `BudgetExceededError` keeps its + bare-`Exception` base class (so existing `except BudgetExceededError:` + handlers keep working) and just sets the same `category` / + `rate_limit_type` attributes that the rest of the unified rate-limit + path reads (normalized to plain strings, matching how + `RateLimitError.__init__` stores its own values). Duck-typed dispatch + in `get_error_information` picks them up automatically. + """ + + def test_should_carry_litellm_rate_limit_category(self): + e = litellm.BudgetExceededError(current_cost=0.5, max_budget=0.1) + # Stored as the plain string value (matches RateLimitError behavior), + # but equality with the enum still works because the enum subclasses + # str. + assert e.category == "litellm_rate_limit" + assert e.category == RateLimitErrorCategory.LITELLM_RATE_LIMIT + + def test_should_carry_budget_rate_limit_type(self): + e = litellm.BudgetExceededError(current_cost=0.5, max_budget=0.1) + assert e.rate_limit_type == "budget" + assert e.rate_limit_type == RateLimitType.BUDGET + + def test_should_default_llm_provider_to_empty_string(self): + # `llm_provider` is read off the exception in `get_error_information` + # — it must always be a string so the StandardLoggingPayload field + # stays serializable. Default to "" when no caller passes one. + e = litellm.BudgetExceededError(current_cost=0.5, max_budget=0.1) + assert e.llm_provider == "" + + def test_should_accept_llm_provider_kwarg(self): + # Callers that have the resolved provider in scope (e.g. the + # auth-checks budget enforcement paths) can thread it through. + e = litellm.BudgetExceededError( + current_cost=0.5, max_budget=0.1, llm_provider="anthropic" + ) + assert e.llm_provider == "anthropic" + + def test_should_keep_existing_status_code_and_message(self): + # Backward-compat guard: existing callers depend on `status_code=429` + # and the canonical message format. + e = litellm.BudgetExceededError(current_cost=0.000109, max_budget=0.0001) + assert e.status_code == 429 + assert "Current cost: 0.000109" in e.message + assert "Max budget: 0.0001" in e.message + + def test_should_still_be_catchable_as_exception_not_rate_limit_error(self): + # Critical: we deliberately did NOT make BudgetExceededError a + # RateLimitError subclass. Existing `except BudgetExceededError:` + # handlers must keep catching it, and `except RateLimitError:` + # handlers must NOT start catching it (which would surprise callers + # who rely on the two being distinct). + e = litellm.BudgetExceededError(current_cost=0.5, max_budget=0.1) + assert isinstance(e, Exception) + assert isinstance(e, litellm.BudgetExceededError) + assert not isinstance(e, RateLimitError) + + def test_should_propagate_category_to_standard_logging_payload(self): + from litellm.litellm_core_utils.litellm_logging import ( + StandardLoggingPayloadSetup, + ) + + e = litellm.BudgetExceededError(current_cost=0.5, max_budget=0.1) + info = StandardLoggingPayloadSetup.get_error_information(e) + assert info["error_rate_limit_category"] == "litellm_rate_limit" + assert info["error_rate_limit_type"] == "budget" + assert info["error_code"] == "429" + assert info["error_class"] == "BudgetExceededError" + + def test_should_propagate_llm_provider_to_standard_logging_payload(self): + from litellm.litellm_core_utils.litellm_logging import ( + StandardLoggingPayloadSetup, + ) + + e = litellm.BudgetExceededError( + current_cost=0.5, max_budget=0.1, llm_provider="bedrock" + ) + info = StandardLoggingPayloadSetup.get_error_information(e) + assert info["llm_provider"] == "bedrock" + + +class TestThirdPartyAttrLeakageGuard: + """ + The duck-typed read at the StandardLoggingPayload + Prometheus surfaces + must reject `.category` / `.rate_limit_type` strings set on unrelated + third-party exceptions. Without validation, a foreign exception that + happens to declare either attribute name would leak garbage values into + custom-callback payloads and Prometheus label cardinality. + """ + + def test_should_drop_unknown_category_string_on_third_party_exception(self): + from litellm.litellm_core_utils.litellm_logging import ( + StandardLoggingPayloadSetup, + ) + + class Foreign(Exception): + category = "totally_not_a_real_category" + + info = StandardLoggingPayloadSetup.get_error_information(Foreign("boom")) + assert info["error_rate_limit_category"] is None + + def test_should_drop_unknown_rate_limit_type_string_on_third_party_exception(self): + from litellm.litellm_core_utils.litellm_logging import ( + StandardLoggingPayloadSetup, + ) + + class Foreign(Exception): + rate_limit_type = "wat" + + info = StandardLoggingPayloadSetup.get_error_information(Foreign("boom")) + assert info["error_rate_limit_type"] is None + + def test_should_drop_non_string_garbage_attrs(self): + from litellm.litellm_core_utils.litellm_logging import ( + StandardLoggingPayloadSetup, + ) + + class Foreign(Exception): + category = 42 + rate_limit_type = {"lol": "no"} + + info = StandardLoggingPayloadSetup.get_error_information(Foreign()) + assert info["error_rate_limit_category"] is None + assert info["error_rate_limit_type"] is None + + def test_should_drop_garbage_on_prometheus_label_extraction(self): + from litellm.integrations.prometheus import PrometheusLogger + + class Foreign(Exception): + category = "spam" + rate_limit_type = "spam" + + category, rate_limit_type = PrometheusLogger._extract_rate_limit_labels( + Foreign() + ) + assert category is None + assert rate_limit_type is None + + def test_should_still_accept_legitimate_rate_limit_categories(self): + # The guard must not over-correct — every documented enum value + # is a valid string and must pass through. + from litellm.exceptions import ( + validate_rate_limit_category, + validate_rate_limit_type, + ) + + for member in RateLimitErrorCategory: + assert validate_rate_limit_category(member.value) == member.value + assert validate_rate_limit_category(member) == member.value + + for member in RateLimitType: + assert validate_rate_limit_type(member.value) == member.value + assert validate_rate_limit_type(member) == member.value + + +@pytest.mark.asyncio +class TestBudgetExceededErrorLlmProviderEnrichment: + """ + BudgetExceededError raise sites in auth_checks.py are tenant-scoped + (key / team / org / tag) and cannot see the request model. To still + populate `llm_provider` on the StandardLoggingPayload — which is what + custom-callback consumers attribute spend to — the central + UserAPIKeyAuthExceptionHandler enriches the exception from + `request_data["model"]` before post_call_failure_hook fires. + """ + + async def _run_handler_and_capture_exception_seen_by_callback( + self, exception: Exception, request_data: dict + ): + from unittest.mock import AsyncMock, MagicMock, patch + + from litellm.proxy.auth.auth_exception_handler import ( + UserAPIKeyAuthExceptionHandler, + ) + + captured: dict = {} + + async def fake_post_call_failure_hook(**kwargs): + captured["exception"] = kwargs["original_exception"] + return None + + with ( + patch( + "litellm.proxy.proxy_server.proxy_logging_obj", + MagicMock( + post_call_failure_hook=AsyncMock( + side_effect=fake_post_call_failure_hook + ) + ), + ), + patch( + "litellm.proxy.proxy_server.general_settings", + {"use_x_forwarded_for": False}, + ), + patch( + "litellm.proxy.auth.auth_exception_handler._get_request_ip_address", + return_value="127.0.0.1", + ), + ): + try: + await UserAPIKeyAuthExceptionHandler._handle_authentication_error( + e=exception, + request=MagicMock(), + request_data=request_data, + route="/v1/chat/completions", + parent_otel_span=None, + api_key="sk-test", + ) + except Exception: + pass + return captured.get("exception") + + async def test_should_resolve_llm_provider_from_request_data_when_unset(self): + err = litellm.BudgetExceededError(current_cost=100, max_budget=10) + assert err.llm_provider == "" + seen = await self._run_handler_and_capture_exception_seen_by_callback( + err, {"model": "openai/gpt-4o-mini"} + ) + assert seen is not None + assert seen.llm_provider == "openai" + + async def test_should_not_overwrite_llm_provider_when_caller_set_it(self): + err = litellm.BudgetExceededError( + current_cost=100, max_budget=10, llm_provider="anthropic" + ) + seen = await self._run_handler_and_capture_exception_seen_by_callback( + err, {"model": "openai/gpt-4o-mini"} + ) + assert seen.llm_provider == "anthropic" + + async def test_should_fall_back_to_litellm_proxy_when_model_missing(self): + err = litellm.BudgetExceededError(current_cost=100, max_budget=10) + seen = await self._run_handler_and_capture_exception_seen_by_callback(err, {}) + assert seen.llm_provider == "litellm_proxy" + + async def test_should_not_enrich_non_budget_exceptions(self): + err = ValueError("unrelated") + seen = await self._run_handler_and_capture_exception_seen_by_callback( + err, {"model": "openai/gpt-4o-mini"} + ) + assert not hasattr(seen, "llm_provider") or seen.llm_provider != "openai" diff --git a/ui/litellm-dashboard/src/components/view_logs/index.test.tsx b/ui/litellm-dashboard/src/components/view_logs/index.test.tsx index aed194a2972..6f1c4126282 100644 --- a/ui/litellm-dashboard/src/components/view_logs/index.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/index.test.tsx @@ -1,8 +1,11 @@ import { screen, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; +import moment from "moment"; import { beforeEach, describe, expect, it, vi } from "vitest"; import SpendLogsTable from "./index"; import { renderWithProviders } from "../../../tests/test-utils"; +import { uiSpendLogsCall } from "../networking"; +import { useLogFilterLogic } from "./log_filter_logic"; const mockHandleFilterResetFromHook = vi.fn(); vi.mock("./log_filter_logic", async (importOriginal) => { @@ -115,4 +118,63 @@ describe("SpendLogsTable", () => { expect(screen.getByRole("button", { name: "Reset Filters" })).toBeInTheDocument(); }); }); + + describe("Quick Select time range", () => { + // uiSpendLogsCall fires from the real useLogFilterLogic query, so restore it here. + beforeEach(async () => { + const actual = await vi.importActual("./log_filter_logic"); + vi.mocked(useLogFilterLogic).mockImplementation(actual.useLogFilterLogic); + }); + + const waitForWindowSeconds = async (minMinutes: number) => { + let diff = -1; + await waitFor(() => { + const lastCall = vi.mocked(uiSpendLogsCall).mock.calls.at(-1)?.[0]; + if (!lastCall) throw new Error("uiSpendLogsCall was not called"); + diff = moment + .utc(lastCall.end_date, "YYYY-MM-DD HH:mm:ss") + .diff(moment.utc(lastCall.start_date, "YYYY-MM-DD HH:mm:ss"), "seconds"); + // start_date is rounded down to the minute boundary, end_date is the + // current wall-clock at queryFn time. The dropped sub-minute fraction + // on start_date can push the diff up to (minMinutes+1)*60 seconds + // exactly (e.g. click at HH:MM:59.9 → start floors to HH:MM:00 and + // queryFn fires just past HH:(MM+1):00), so allow equality on the + // upper bound. + expect(diff).toBeGreaterThanOrEqual(minMinutes * 60); + expect(diff).toBeLessThanOrEqual((minMinutes + 1) * 60); + }); + return diff; + }; + + it("should pass a ~1-minute window to uiSpendLogsCall when 'Last Minute' is selected", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + await user.click(screen.getByRole("button", { name: /Last 24 Hours/i })); + await user.click(await screen.findByRole("button", { name: "Last Minute" })); + + await waitForWindowSeconds(1); + }); + + it("should pass a ~15-minute window to uiSpendLogsCall when 'Last 15 Minutes' is selected", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + await user.click(screen.getByRole("button", { name: /Last 24 Hours/i })); + await user.click(await screen.findByRole("button", { name: "Last 15 Minutes" })); + + await waitForWindowSeconds(15); + }); + + it("should update the time-range button label to 'Last Minute' after selecting it", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + await user.click(screen.getByRole("button", { name: /Last 24 Hours/i })); + await user.click(await screen.findByRole("button", { name: "Last Minute" })); + + expect(screen.getByRole("button", { name: "Last Minute" })).toBeInTheDocument(); + expect(screen.queryByRole("button", { name: /Last 24 Hours/i })).not.toBeInTheDocument(); + }); + }); }); diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 936747acd08..f952bbfbbaf 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -38544,7 +38544,17 @@ export interface operations { "application/json": components["schemas"]["ErrorResponse"]; }; }; - /** @description RateLimitError */ + /** + * @description Unified rate-limit error. + * + * Every rate-limit condition surfaced by litellm — whether it originated from + * an upstream LLM provider, a vendor batch endpoint, or one of litellm's own + * proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget, + * max-iterations, etc.) — is raised as an instance of this class. + * + * The :attr:`category` attribute lets callers distinguish the source. See + * :class:`RateLimitErrorCategory` for the available values. + */ 429: { headers: { [name: string]: unknown; From 3448bf79f8832db4cc21712ed57654b21a6d79b3 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 6 Jun 2026 18:10:17 -0700 Subject: [PATCH 007/185] fix(ui): default guardrails page to first tab for admins, not submitted (#29872) The Guardrails page hardcoded defaultActiveKey="submitted", so admins landed on the "Submitted Guardrails" tab (the last of their four tabs) instead of the primary view. The original intent was for non-admins, whose only tab is Submitted Guardrails, to default there; admins should open on their first tab. Make the default role-aware: admins default to the first tab (Guardrail Garden), non-admins keep Submitted Guardrails. --- ui/litellm-dashboard/src/components/guardrails.tsx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/src/components/guardrails.tsx b/ui/litellm-dashboard/src/components/guardrails.tsx index 4df3f93c275..4d1d3f65e3e 100644 --- a/ui/litellm-dashboard/src/components/guardrails.tsx +++ b/ui/litellm-dashboard/src/components/guardrails.tsx @@ -133,7 +133,7 @@ const GuardrailsPanel: React.FC = ({ accessToken, userRole return (
Date: Sat, 6 Jun 2026 20:34:31 -0700 Subject: [PATCH 008/185] refactor(bedrock): build Converse toolSpec via a BedrockToolSpec dict subclass (#29869) --- .../prompt_templates/factory.py | 30 +++++++------------ litellm/types/llms/bedrock.py | 30 +++++++++++++++++++ 2 files changed, 40 insertions(+), 20 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 588a183d58e..81a4c8b14b6 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -3653,17 +3653,13 @@ from litellm.types.llms.bedrock import ContentBlock as BedrockContentBlock from litellm.types.llms.bedrock import DocumentBlock as BedrockDocumentBlock from litellm.types.llms.bedrock import ImageBlock as BedrockImageBlock from litellm.types.llms.bedrock import SourceBlock as BedrockSourceBlock +from litellm.types.llms.bedrock import BedrockToolSpec from litellm.types.llms.bedrock import ToolBlock as BedrockToolBlock -from litellm.types.llms.bedrock import ( - ToolInputSchemaBlock as BedrockToolInputSchemaBlock, -) -from litellm.types.llms.bedrock import ToolJsonSchemaBlock as BedrockToolJsonSchemaBlock from litellm.types.llms.bedrock import SearchResultBlock from litellm.types.llms.bedrock import ToolResultBlock as BedrockToolResultBlock from litellm.types.llms.bedrock import ( ToolResultContentBlock as BedrockToolResultContentBlock, ) -from litellm.types.llms.bedrock import ToolSpecBlock as BedrockToolSpecBlock from litellm.types.llms.bedrock import ToolUseBlock as BedrockToolUseBlock from litellm.types.llms.bedrock import VideoBlock as BedrockVideoBlock @@ -5554,22 +5550,16 @@ def _bedrock_tools_pt( normalize_json_schema_custom_types_to_object(parameters) if parameters.get("type") not in _valid_json_schema_root_types: parameters["type"] = "object" - json_schema = BedrockToolJsonSchemaBlock( - type=parameters["type"], - properties=parameters.get("properties", {}), - required=parameters.get("required", []), + tool_block = cast( + BedrockToolBlock, + BedrockToolSpec( + name=name, + description=description, + parameters=parameters, + strict=tool.get("function", {}).get("strict", None), + supports_strict_tools=supports_strict_tools, + ), ) - additional_properties = parameters.get("additionalProperties", None) - if supports_strict_tools and additional_properties is not None: - json_schema["additionalProperties"] = additional_properties - tool_input_schema = BedrockToolInputSchemaBlock(json=json_schema) - tool_spec = BedrockToolSpecBlock( - inputSchema=tool_input_schema, name=name, description=description - ) - strict = tool.get("function", {}).get("strict", None) - if supports_strict_tools and strict is not None: - tool_spec["strict"] = strict - tool_block = BedrockToolBlock(toolSpec=tool_spec) tool_block_list.append(tool_block) ## ADD CACHE POINT TOOL BLOCK ## diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index ed80b44414b..fa8c3a93ef3 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -285,6 +285,36 @@ class ToolBlock(TypedDict, total=False): cachePoint: Optional[CachePointBlock] +class BedrockToolSpec(dict): + def __init__( + self, + *, + name: str, + description: str, + parameters: dict, + strict: Optional[bool], + supports_strict_tools: bool, + ) -> None: + json_schema: ToolJsonSchemaBlock = { + "type": parameters["type"], + "properties": parameters.get("properties", {}), + "required": parameters.get("required", []), + } + additional_properties = parameters.get("additionalProperties") + if supports_strict_tools and additional_properties is not None: + json_schema["additionalProperties"] = additional_properties + + tool_spec: ToolSpecBlock = { + "inputSchema": {"json": json_schema}, + "name": name, + "description": description, + } + if supports_strict_tools and strict is not None: + tool_spec["strict"] = strict + + super().__init__(toolSpec=tool_spec) + + class SpecificToolChoiceBlock(TypedDict): name: str From 5e2db7eee4e30e1c5b698e52ef4cc411860a1ca5 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Sat, 6 Jun 2026 20:59:33 -0700 Subject: [PATCH 009/185] feat(litellm): add models and repository layers (#29686) --- .github/workflows/test-unit-misc.yml | 2 + ARCHITECTURE.md | 18 + .../SlackAlerting/slack_alerting.py | 6 +- litellm/integrations/email_alerting.py | 3 +- litellm/integrations/prometheus.py | 21 +- litellm/llms/litellm_proxy/skills/handler.py | 9 +- litellm/models/__init__.py | 66 + litellm/models/access_group.py | 26 + litellm/models/base.py | 38 + litellm/models/budget.py | 56 + litellm/models/config.py | 15 + litellm/models/credentials.py | 31 + litellm/models/end_user.py | 35 + litellm/models/managed_files.py | 62 + litellm/models/mcp_server.py | 103 + litellm/models/model.py | 59 + litellm/models/object_permission.py | 26 + litellm/models/organization.py | 31 + litellm/models/organization_membership.py | 40 + litellm/models/project.py | 41 + litellm/models/skills.py | 30 + litellm/models/spend_logs.py | 50 + litellm/models/tag.py | 36 + litellm/models/team.py | 154 ++ litellm/models/team_membership.py | 32 + litellm/models/user.py | 70 + litellm/models/verification_token.py | 74 + .../mcp_server/auth/user_api_key_auth_mcp.py | 10 +- litellm/proxy/_experimental/mcp_server/db.py | 125 +- .../mcp_server/mcp_server_manager.py | 5 +- .../_experimental/mcp_server/toolset_db.py | 13 +- litellm/proxy/_types.py | 755 +----- .../proxy/agent_endpoints/agent_registry.py | 19 +- .../auth/agent_permission_handler.py | 5 +- litellm/proxy/agent_endpoints/endpoints.py | 23 +- .../claude_code_marketplace.py | 25 +- litellm/proxy/auth/auth_checks.py | 74 +- litellm/proxy/auth/handle_jwt.py | 5 +- litellm/proxy/auth/login_utils.py | 5 +- litellm/proxy/auth/model_checks.py | 3 +- litellm/proxy/auth/user_api_key_auth.py | 9 +- .../expired_ui_session_key_cleanup_manager.py | 7 +- .../common_utils/key_rotation_manager.py | 42 +- .../proxy/common_utils/reset_budget_job.py | 24 +- .../proxy/container_endpoints/ownership.py | 9 +- .../proxy/credential_endpoints/endpoints.py | 16 +- litellm/proxy/db/spend_counter_reseed.py | 26 +- litellm/proxy/db/spend_log_tool_index.py | 3 +- litellm/proxy/db/tool_registry_writer.py | 26 +- .../proxy/guardrails/guardrail_endpoints.py | 26 +- .../proxy/guardrails/guardrail_registry.py | 31 +- litellm/proxy/guardrails/usage_endpoints.py | 44 +- litellm/proxy/guardrails/usage_tracking.py | 8 +- .../hooks/user_management_event_hooks.py | 5 +- .../access_group_endpoints.py | 5 +- .../budget_management_endpoints.py | 15 +- .../cache_settings_endpoints.py | 9 +- .../common_daily_activity.py | 16 +- .../management_endpoints/common_utils.py | 9 +- .../config_override_endpoints.py | 11 +- .../customer_endpoints.py | 56 +- .../fallback_management_endpoints.py | 5 +- .../internal_user_endpoints.py | 95 +- .../jwt_key_mapping_endpoints.py | 17 +- .../key_management_endpoints.py | 151 +- .../mcp_management_endpoints.py | 14 +- ...model_access_group_management_endpoints.py | 19 +- .../model_management_endpoints.py | 33 +- .../organization_endpoints.py | 207 +- .../scim/scim_transformations.py | 3 +- .../management_endpoints/scim/scim_v2.py | 86 +- .../tag_management_endpoints.py | 34 +- .../team_callback_endpoints.py | 5 +- .../management_endpoints/team_endpoints.py | 212 +- .../tool_management_endpoints.py | 37 +- litellm/proxy/management_endpoints/ui_sso.py | 29 +- .../user_agent_analytics_endpoints.py | 15 +- .../workflow_management_endpoints.py | 27 +- .../proxy/management_helpers/audit_logs.py | 3 +- .../object_permission_utils.py | 38 +- .../management_helpers/user_invitation.py | 3 +- litellm/proxy/management_helpers/utils.py | 37 +- litellm/proxy/memory/memory_endpoints.py | 20 +- .../openai_files_endpoints/common_utils.py | 15 +- .../openai_files_endpoints/files_endpoints.py | 21 +- .../managed_id_rewriter.py | 30 +- .../pass_through_endpoints.py | 9 +- .../policy_engine/attachment_registry.py | 57 +- .../proxy/policy_engine/policy_registry.py | 43 +- .../policy_engine/policy_resolve_endpoints.py | 12 +- .../proxy/policy_engine/policy_validator.py | 10 +- litellm/proxy/prompts/prompt_endpoints.py | 21 +- litellm/proxy/proxy_server.py | 165 +- .../public_endpoints/public_endpoints.py | 11 +- litellm/proxy/rag_endpoints/endpoints.py | 13 +- .../search_endpoints/search_tool_registry.py | 61 +- .../spend_tracking/cloudzero_endpoints.py | 13 +- .../spend_management_endpoints.py | 27 +- .../proxy/spend_tracking/vantage_endpoints.py | 15 +- .../proxy_setting_endpoints.py | 24 +- litellm/proxy/utils.py | 123 +- .../proxy/vector_store_endpoints/endpoints.py | 12 +- .../management_endpoints.py | 28 +- litellm/repositories/__init__.py | 127 + litellm/repositories/base_repository.py | 117 + litellm/repositories/budget_repository.py | 99 + litellm/repositories/config_repository.py | 241 ++ .../repositories/credentials_repository.py | 61 + litellm/repositories/model_repository.py | 171 ++ .../object_permission_repository.py | 110 + .../repositories/organization_repository.py | 103 + litellm/repositories/project_repository.py | 129 + litellm/repositories/table_repositories.py | 215 ++ litellm/repositories/team_repository.py | 351 +++ litellm/repositories/user_repository.py | 229 ++ .../verification_token_repository.py | 375 +++ .../adaptive_router/adaptive_router.py | 3 +- .../adaptive_router/update_queue.py | 8 +- .../types/mcp_server/mcp_server_manager.py | 3 +- litellm/types/utils.py | 25 +- .../vector_stores/vector_store_registry.py | 26 +- tests/test_litellm/models/test_models.py | 542 ++++ .../repositories/test_repositories.py | 2184 +++++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 179 +- 124 files changed, 7846 insertions(+), 1850 deletions(-) create mode 100644 litellm/models/__init__.py create mode 100644 litellm/models/access_group.py create mode 100644 litellm/models/base.py create mode 100644 litellm/models/budget.py create mode 100644 litellm/models/config.py create mode 100644 litellm/models/credentials.py create mode 100644 litellm/models/end_user.py create mode 100644 litellm/models/managed_files.py create mode 100644 litellm/models/mcp_server.py create mode 100644 litellm/models/model.py create mode 100644 litellm/models/object_permission.py create mode 100644 litellm/models/organization.py create mode 100644 litellm/models/organization_membership.py create mode 100644 litellm/models/project.py create mode 100644 litellm/models/skills.py create mode 100644 litellm/models/spend_logs.py create mode 100644 litellm/models/tag.py create mode 100644 litellm/models/team.py create mode 100644 litellm/models/team_membership.py create mode 100644 litellm/models/user.py create mode 100644 litellm/models/verification_token.py create mode 100644 litellm/repositories/__init__.py create mode 100644 litellm/repositories/base_repository.py create mode 100644 litellm/repositories/budget_repository.py create mode 100644 litellm/repositories/config_repository.py create mode 100644 litellm/repositories/credentials_repository.py create mode 100644 litellm/repositories/model_repository.py create mode 100644 litellm/repositories/object_permission_repository.py create mode 100644 litellm/repositories/organization_repository.py create mode 100644 litellm/repositories/project_repository.py create mode 100644 litellm/repositories/table_repositories.py create mode 100644 litellm/repositories/team_repository.py create mode 100644 litellm/repositories/user_repository.py create mode 100644 litellm/repositories/verification_token_repository.py create mode 100644 tests/test_litellm/models/test_models.py create mode 100644 tests/test_litellm/repositories/test_repositories.py diff --git a/.github/workflows/test-unit-misc.yml b/.github/workflows/test-unit-misc.yml index 9add77ff424..a7363ac3b43 100644 --- a/.github/workflows/test-unit-misc.yml +++ b/.github/workflows/test-unit-misc.yml @@ -28,6 +28,8 @@ jobs: tests/test_litellm/completion_extras tests/test_litellm/containers tests/test_litellm/experimental_mcp_client + tests/test_litellm/models + tests/test_litellm/repositories tests/test_litellm/images tests/test_litellm/interactions tests/test_litellm/passthrough diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index c114a838d6d..3d2fa3e51c8 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -240,6 +240,24 @@ graph LR 7. `DBSpendUpdateWriter.update_database()` queues spend increments to Redis 8. Background job `update_spend` flushes queued spend to PostgreSQL every 60s +### Data Access Layer (Models & Repositories) + +Database entities and the operations on them live in two packages at the root of `litellm/` so both the gateway (`proxy/`) and the SDK can use them without importing proxy internals: + +- `litellm/models/` holds the canonical Pydantic definitions for every persisted entity (`LiteLLM_VerificationToken`, `LiteLLM_TeamTable`, `LiteLLM_UserTable`, etc.). `proxy/_types.py` re-exports these for backwards compatibility, so existing imports keep working. +- `litellm/repositories/` holds the data-access layer. `BaseRepository[T]` provides the generic CRUD (`find_by_id`, `find_many`, `create`, `update`, `delete`, `count`, `exists`); entity repositories such as `VerificationTokenRepository`, `TeamRepository`, and `UserRepository` add domain-specific queries and writes on top of it. + +Conventions to follow when touching this layer: + +| Concern | How it's handled | +|---------|------------------| +| JSON columns | Prisma `Json` columns are stored as JSON strings. Repositories `json.dumps()` on write and `json.loads()` on read (see `_to_model` and the `_build_*_data` helpers). | +| Archive-then-delete | `delete_team` / `delete_token` copy the row into the `LiteLLM_Deleted*` table and delete the original inside a single `prisma_client.db.tx()` transaction. Archive payloads are built explicitly so only columns that exist on the archive table are written. | +| Column vs. field names | Where a model field differs from its DB column (for example `org_id` maps to the `organization_id` column), the repository translates in both directions rather than relying on Pydantic to guess. | +| Array mutations | Adds use Prisma's atomic `push` (`add_member`, `add_admin`, `add_models`) to avoid read-modify-write races. Removals fall back to read-modify-write because Prisma has no atomic array remove. | + +To add a new entity, define the model under `litellm/models/`, re-export it from `proxy/_types.py` if existing code imports it from there, and add a repository under `litellm/repositories/` (subclass `BaseRepository` for plain CRUD, or add bespoke methods when the entity needs encryption, archiving, or atomic array updates). Mirror the tests in `tests/test_litellm/repositories/`. + --- ## 2. SDK Request Flow diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index 0ec17bbea5d..390af2cb6e6 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/litellm/integrations/SlackAlerting/slack_alerting.py @@ -37,6 +37,8 @@ from litellm.proxy._types import ( VirtualKeyEvent, WebhookEvent, ) +from litellm.repositories.team_repository import TeamRepository +from litellm.repositories.user_repository import UserRepository from litellm.types.integrations.slack_alerting import * from ..email_templates.templates import * @@ -1231,7 +1233,7 @@ Model Info: and recipient_user_id is not None and prisma_client is not None ): - user_row = await prisma_client.db.litellm_usertable.find_unique( + user_row = await UserRepository(prisma_client).table.find_unique( where={"user_id": recipient_user_id} ) @@ -1263,7 +1265,7 @@ Model Info: team_id = webhook_event.team_id team_name = "Default Team" if team_id is not None and prisma_client is not None: - team_row = await prisma_client.db.litellm_teamtable.find_unique( + team_row = await TeamRepository(prisma_client).table.find_unique( where={"team_id": team_id} ) if team_row is not None: diff --git a/litellm/integrations/email_alerting.py b/litellm/integrations/email_alerting.py index b45b9aa7f5c..b721dc50464 100644 --- a/litellm/integrations/email_alerting.py +++ b/litellm/integrations/email_alerting.py @@ -7,6 +7,7 @@ from typing import List, Optional from litellm._logging import verbose_logger, verbose_proxy_logger from litellm.proxy._types import WebhookEvent +from litellm.repositories.team_repository import TeamRepository # we use this for the email header, please send a test email if you change this. verify it looks good on email LITELLM_LOGO_URL = "https://litellm-listing.s3.amazonaws.com/litellm_logo.png" @@ -24,7 +25,7 @@ async def get_all_team_member_emails(team_id: Optional[str] = None) -> list: if prisma_client is None: raise Exception("Not connected to DB!") - team_row = await prisma_client.db.litellm_teamtable.find_unique( + team_row = await TeamRepository(prisma_client).table.find_unique( where={ "team_id": team_id, } diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index d2af95cd4cc..2119527a8e5 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -29,13 +29,13 @@ from litellm.exceptions import ( validate_rate_limit_type, ) from litellm.integrations.custom_logger import CustomLogger -from litellm.integrations.prometheus_helpers.bounded_prometheus_series_tracker import ( - BoundedPrometheusSeriesTracker, -) from litellm.integrations.prometheus_helpers import ( PrometheusLabelFactoryContext, _get_cached_end_user_id_for_cost_tracking, ) +from litellm.integrations.prometheus_helpers.bounded_prometheus_series_tracker import ( + BoundedPrometheusSeriesTracker, +) from litellm.litellm_core_utils.core_helpers import ( get_litellm_metadata_from_kwargs, get_metadata_variable_name_from_kwargs, @@ -46,6 +46,9 @@ from litellm.proxy._types import ( LiteLLM_UserTable, UserAPIKeyAuth, ) +from litellm.repositories.organization_repository import OrganizationRepository +from litellm.repositories.team_repository import TeamRepository +from litellm.repositories.user_repository import UserRepository from litellm.types.integrations.prometheus import * from litellm.types.integrations.prometheus import ( _sanitize_prometheus_label_name, @@ -3278,12 +3281,12 @@ class PrometheusLogger(CustomLogger): page_size: int, page: int ) -> Tuple[List[LiteLLM_UserTable], Optional[int]]: skip = (page - 1) * page_size - users = await prisma_client.db.litellm_usertable.find_many( + users = await UserRepository(prisma_client).table.find_many( skip=skip, take=page_size, order={"created_at": "desc"}, ) - total_count = await prisma_client.db.litellm_usertable.count() + total_count = await UserRepository(prisma_client).table.count() return users, total_count await self._initialize_budget_metrics( @@ -3306,13 +3309,13 @@ class PrometheusLogger(CustomLogger): async def fetch_orgs(page_size: int, page: int) -> Tuple[list, Optional[int]]: skip = (page - 1) * page_size - orgs = await prisma_client.db.litellm_organizationtable.find_many( + orgs = await OrganizationRepository(prisma_client).table.find_many( skip=skip, take=page_size, order={"created_at": "desc"}, include={"litellm_budget_table": True}, ) - total_count = await prisma_client.db.litellm_organizationtable.count() + total_count = await OrganizationRepository(prisma_client).table.count() return orgs, total_count await self._initialize_budget_metrics( @@ -3380,14 +3383,14 @@ class PrometheusLogger(CustomLogger): try: # Get total user count - total_users = await prisma_client.db.litellm_usertable.count() + total_users = await UserRepository(prisma_client).table.count() self.litellm_total_users_metric.set(total_users) verbose_logger.debug( f"Prometheus: set litellm_total_users to {total_users}" ) # Get total team count - total_teams = await prisma_client.db.litellm_teamtable.count() + total_teams = await TeamRepository(prisma_client).table.count() self.litellm_teams_count_metric.set(total_teams) verbose_logger.debug( f"Prometheus: set litellm_teams_count to {total_teams}" diff --git a/litellm/llms/litellm_proxy/skills/handler.py b/litellm/llms/litellm_proxy/skills/handler.py index 37aabd8b477..7b259c1ed66 100644 --- a/litellm/llms/litellm_proxy/skills/handler.py +++ b/litellm/llms/litellm_proxy/skills/handler.py @@ -17,6 +17,7 @@ from litellm.proxy.common_utils.resource_ownership import ( is_proxy_admin, user_can_access_resource_owner, ) +from litellm.repositories.table_repositories import SkillsRepository # Skills are looked up on every chat completion that has skills enabled # (`SkillsInjectionHook` calls ``fetch_skill_from_db``). 60s LRU/TTL cache @@ -107,7 +108,7 @@ class LiteLLMSkillsHandler: f"LiteLLMSkillsHandler: Creating skill {skill_id} with title={data.display_title}" ) - new_skill = await prisma_client.db.litellm_skillstable.create(data=skill_data) + new_skill = await SkillsRepository(prisma_client).table.create(data=skill_data) return _prisma_skill_to_litellm(new_skill) @staticmethod @@ -133,7 +134,7 @@ class LiteLLMSkillsHandler: return [] find_many_kwargs["where"] = {"created_by": {"in": owner_scopes}} - skills = await prisma_client.db.litellm_skillstable.find_many( + skills = await SkillsRepository(prisma_client).table.find_many( **find_many_kwargs ) return [_prisma_skill_to_litellm(s) for s in skills] @@ -150,7 +151,7 @@ class LiteLLMSkillsHandler: return cached prisma_client = await LiteLLMSkillsHandler._get_prisma_client() - skill = await prisma_client.db.litellm_skillstable.find_unique( + skill = await SkillsRepository(prisma_client).table.find_unique( where={"skill_id": skill_id} ) _SKILL_CACHE.set_cache( @@ -189,7 +190,7 @@ class LiteLLMSkillsHandler: ): raise ValueError(f"Skill not found: {skill_id}") - await prisma_client.db.litellm_skillstable.delete(where={"skill_id": skill_id}) + await SkillsRepository(prisma_client).table.delete(where={"skill_id": skill_id}) _SKILL_CACHE.set_cache(skill_id, _NEGATIVE_SKILL_SENTINEL) return {"id": skill_id, "type": "skill_deleted"} diff --git a/litellm/models/__init__.py b/litellm/models/__init__.py new file mode 100644 index 00000000000..7e2d2c0ed9d --- /dev/null +++ b/litellm/models/__init__.py @@ -0,0 +1,66 @@ +""" +Domain models for LiteLLM backend. +""" + +from litellm.models.access_group import LiteLLM_AccessGroupTable +from litellm.models.budget import ( + LiteLLM_BudgetTable, + LiteLLM_BudgetTableFull, + LiteLLM_TeamMemberTable, +) +from litellm.models.config import LiteLLM_Config +from litellm.models.credentials import ( + CreateCredentialItem, + CredentialBase, + CredentialItem, +) +from litellm.models.end_user import LiteLLM_EndUserTable +from litellm.models.managed_files import ( + LiteLLM_ManagedFileTable, + LiteLLM_ManagedObjectTable, + LiteLLM_ManagedVectorStoresTable, + LiteLLM_ManagedVectorStoreTable, +) +from litellm.models.mcp_server import LiteLLM_MCPServerTable +from litellm.models.model import LiteLLM_ProxyModelTable +from litellm.models.object_permission import LiteLLM_ObjectPermissionTable +from litellm.models.organization import LiteLLM_OrganizationTable +from litellm.models.organization_membership import LiteLLM_OrganizationMembershipTable +from litellm.models.project import LiteLLM_ProjectTable +from litellm.models.skills import LiteLLM_SkillsTable +from litellm.models.spend_logs import LiteLLM_ErrorLogs, LiteLLM_SpendLogs +from litellm.models.tag import LiteLLM_TagTable +from litellm.models.team import LiteLLM_TeamTable +from litellm.models.team_membership import LiteLLM_TeamMembership +from litellm.models.user import LiteLLM_UserTable +from litellm.models.verification_token import LiteLLM_VerificationToken + +__all__ = [ + "LiteLLM_AccessGroupTable", + "LiteLLM_BudgetTable", + "LiteLLM_BudgetTableFull", + "LiteLLM_TeamMemberTable", + "LiteLLM_Config", + "CredentialBase", + "CredentialItem", + "CreateCredentialItem", + "LiteLLM_EndUserTable", + "LiteLLM_ManagedFileTable", + "LiteLLM_ManagedObjectTable", + "LiteLLM_ManagedVectorStoreTable", + "LiteLLM_ManagedVectorStoresTable", + "LiteLLM_MCPServerTable", + "LiteLLM_ProxyModelTable", + "LiteLLM_ObjectPermissionTable", + "LiteLLM_OrganizationTable", + "LiteLLM_OrganizationMembershipTable", + "LiteLLM_ProjectTable", + "LiteLLM_SkillsTable", + "LiteLLM_ErrorLogs", + "LiteLLM_SpendLogs", + "LiteLLM_TagTable", + "LiteLLM_TeamTable", + "LiteLLM_TeamMembership", + "LiteLLM_UserTable", + "LiteLLM_VerificationToken", +] diff --git a/litellm/models/access_group.py b/litellm/models/access_group.py new file mode 100644 index 00000000000..682e779e531 --- /dev/null +++ b/litellm/models/access_group.py @@ -0,0 +1,26 @@ +""" +Access group table model. + +Canonical definition for ``litellm_accessgrouptable``. Re-exported from +``litellm.proxy._types`` for backwards compatibility. +""" + +from datetime import datetime +from typing import List, Optional + +from litellm.types.llms.base import LiteLLMPydanticObjectBase + + +class LiteLLM_AccessGroupTable(LiteLLMPydanticObjectBase): + access_group_id: str + access_group_name: str + description: Optional[str] = None + access_model_names: List[str] = [] + access_mcp_server_ids: List[str] = [] + access_agent_ids: List[str] = [] + assigned_team_ids: List[str] = [] + assigned_key_ids: List[str] = [] + created_at: Optional[datetime] = None + created_by: Optional[str] = None + updated_at: Optional[datetime] = None + updated_by: Optional[str] = None diff --git a/litellm/models/base.py b/litellm/models/base.py new file mode 100644 index 00000000000..01981297bd5 --- /dev/null +++ b/litellm/models/base.py @@ -0,0 +1,38 @@ +""" +Base model class for domain models. +""" + +from datetime import datetime +from typing import Any, Dict, Optional + +from pydantic import BaseModel, ConfigDict + + +class DomainModel(BaseModel): + """Base class for all domain models.""" + + model_config = ConfigDict( + from_attributes=True, + protected_namespaces=(), + extra="ignore", + ) + + created_at: Optional[datetime] = None + updated_at: Optional[datetime] = None + + @classmethod + def from_db_record(cls, record: Any) -> "DomainModel": + """Create a domain model from a database record.""" + if record is None: + raise ValueError("Cannot create domain model from None record") + if isinstance(record, dict): + return cls(**record) + if hasattr(record, "model_dump") and callable(record.model_dump): + return cls(**record.model_dump()) + if hasattr(record, "dict") and callable(record.dict): + return cls(**record.dict()) + return cls(**dict(record)) + + def to_db_dict(self, exclude_unset: bool = False) -> Dict[str, Any]: + """Convert domain model to a dictionary for database operations.""" + return self.model_dump(exclude_none=True, exclude_unset=exclude_unset) diff --git a/litellm/models/budget.py b/litellm/models/budget.py new file mode 100644 index 00000000000..e7dfe2f8fbc --- /dev/null +++ b/litellm/models/budget.py @@ -0,0 +1,56 @@ +""" +Budget table model. + +Canonical definition for ``litellm_budgettable``. Re-exported from +``litellm.proxy._types`` for backwards compatibility. +""" + +from datetime import datetime +from typing import List, Optional + +from pydantic import ConfigDict + +from litellm.types.llms.base import LiteLLMPydanticObjectBase + + +class LiteLLM_BudgetTable(LiteLLMPydanticObjectBase): + """Represents user-controllable params for a LiteLLM_BudgetTable record. + + Budget-write paths use `model_fields.keys()` on this class as an allowlist + for user input. Keep server-managed fields (e.g. `budget_reset_at`) on + `LiteLLM_BudgetTableFull` so they aren't user-settable. + """ + + budget_id: Optional[str] = None + soft_budget: Optional[float] = None + max_budget: Optional[float] = None + max_parallel_requests: Optional[int] = None + tpm_limit: Optional[int] = None + rpm_limit: Optional[int] = None + model_max_budget: Optional[dict] = None + budget_duration: Optional[str] = None + allowed_models: Optional[List[str]] = ( + None # per-member model scope; empty = inherit team models + ) + + model_config = ConfigDict(protected_namespaces=()) + + +class LiteLLM_BudgetTableFull(LiteLLM_BudgetTable): + """LiteLLM_BudgetTable + server-managed fields returned on API responses.""" + + budget_reset_at: Optional[datetime] = None + created_at: datetime + + +class LiteLLM_TeamMemberTable(LiteLLM_BudgetTable): + """ + Used to track spend of a user_id within a team_id + """ + + spend: Optional[float] = None + user_id: Optional[str] = None + team_id: Optional[str] = None + budget_id: Optional[str] = None + + model_config = ConfigDict(protected_namespaces=()) diff --git a/litellm/models/config.py b/litellm/models/config.py new file mode 100644 index 00000000000..99b5c5692fd --- /dev/null +++ b/litellm/models/config.py @@ -0,0 +1,15 @@ +""" +Config table model. + +Canonical definition for ``litellm_config``. Re-exported from +``litellm.proxy._types`` for backwards compatibility. +""" + +from typing import Dict + +from litellm.types.llms.base import LiteLLMPydanticObjectBase + + +class LiteLLM_Config(LiteLLMPydanticObjectBase): + param_name: str + param_value: Dict diff --git a/litellm/models/credentials.py b/litellm/models/credentials.py new file mode 100644 index 00000000000..b74ea055d21 --- /dev/null +++ b/litellm/models/credentials.py @@ -0,0 +1,31 @@ +""" +Credential table models. + +These are the canonical credential types for the proxy. They live in the model +layer; ``litellm.types.utils`` re-exports them for backwards compatibility. +""" + +from typing import Optional + +from pydantic import BaseModel, model_validator + + +class CredentialBase(BaseModel): + credential_name: str + credential_info: dict + + +class CredentialItem(CredentialBase): + credential_values: dict + + +class CreateCredentialItem(CredentialBase): + credential_values: Optional[dict] = None + model_id: Optional[str] = None + + @model_validator(mode="before") + @classmethod + def check_credential_params(cls, values): + if not values.get("credential_values") and not values.get("model_id"): + raise ValueError("Either credential_values or model_id must be set") + return values diff --git a/litellm/models/end_user.py b/litellm/models/end_user.py new file mode 100644 index 00000000000..15fd03ec2ca --- /dev/null +++ b/litellm/models/end_user.py @@ -0,0 +1,35 @@ +""" +End-user table model. + +Canonical definition for ``litellm_endusertable``. Re-exported from +``litellm.proxy._types`` for backwards compatibility. +""" + +from typing import Literal, Optional + +from pydantic import ConfigDict, model_validator + +from litellm.models.budget import LiteLLM_BudgetTable +from litellm.models.object_permission import LiteLLM_ObjectPermissionTable +from litellm.types.llms.base import LiteLLMPydanticObjectBase + + +class LiteLLM_EndUserTable(LiteLLMPydanticObjectBase): + user_id: str + blocked: bool + alias: Optional[str] = None + spend: float = 0.0 + allowed_model_region: Optional[Literal["eu", "us"]] = None + default_model: Optional[str] = None + litellm_budget_table: Optional[LiteLLM_BudgetTable] = None + object_permission_id: Optional[str] = None + object_permission: Optional[LiteLLM_ObjectPermissionTable] = None + + @model_validator(mode="before") + @classmethod + def set_model_info(cls, values): + if values.get("spend") is None: + values.update({"spend": 0.0}) + return values + + model_config = ConfigDict(protected_namespaces=()) diff --git a/litellm/models/managed_files.py b/litellm/models/managed_files.py new file mode 100644 index 00000000000..24154768860 --- /dev/null +++ b/litellm/models/managed_files.py @@ -0,0 +1,62 @@ +""" +Managed file, object, and vector store table models. + +Canonical definitions for the ``litellm_managed*`` tables. Re-exported from +``litellm.proxy._types`` for backwards compatibility. +""" + +from datetime import datetime +from typing import Any, Dict, List, Literal, Optional, Union + +from litellm.types.llms.base import LiteLLMPydanticObjectBase +from litellm.types.llms.openai import OpenAIFileObject, ResponsesAPIResponse +from litellm.types.utils import LiteLLMBatch, LiteLLMFineTuningJob + + +class LiteLLM_ManagedFileTable(LiteLLMPydanticObjectBase): + unified_file_id: str + file_object: Optional[OpenAIFileObject] = None + model_mappings: Dict[str, str] + flat_model_file_ids: List[str] + created_by: Optional[str] = None + team_id: Optional[str] = None + updated_by: Optional[str] = None + storage_backend: Optional[str] = None + storage_url: Optional[str] = None + + +class LiteLLM_ManagedObjectTable(LiteLLMPydanticObjectBase): + unified_object_id: str + model_object_id: str + file_purpose: Literal["batch", "fine-tune", "response", "container"] + file_object: Union[LiteLLMBatch, LiteLLMFineTuningJob, ResponsesAPIResponse] + created_by: Optional[str] = None + team_id: Optional[str] = None + + +class LiteLLM_ManagedVectorStoreTable(LiteLLMPydanticObjectBase): + """Table for managing vector stores with target_model_names support.""" + + unified_resource_id: str + resource_object: Optional[Any] = None + model_mappings: Dict[str, str] + flat_model_resource_ids: List[str] + created_by: Optional[str] = None + team_id: Optional[str] = None + updated_by: Optional[str] = None + storage_backend: Optional[str] = None + storage_url: Optional[str] = None + + +class LiteLLM_ManagedVectorStoresTable(LiteLLMPydanticObjectBase): + vector_store_id: str + custom_llm_provider: str + vector_store_name: Optional[str] + vector_store_description: Optional[str] + vector_store_metadata: Optional[Dict[str, Any]] + created_at: Optional[datetime] + updated_at: Optional[datetime] + litellm_credential_name: Optional[str] + litellm_params: Optional[Dict[str, Any]] + team_id: Optional[str] + user_id: Optional[str] diff --git a/litellm/models/mcp_server.py b/litellm/models/mcp_server.py new file mode 100644 index 00000000000..3d03eff6df8 --- /dev/null +++ b/litellm/models/mcp_server.py @@ -0,0 +1,103 @@ +""" +MCP server table model. + +Canonical definition for ``litellm_mcpservertable``. Re-exported from +``litellm.proxy._types`` for backwards compatibility. +""" + +import enum +from datetime import datetime +from typing import Dict, List, Literal, Optional + +from pydantic import Field + +from litellm.types.llms.base import LiteLLMPydanticObjectBase +from litellm.types.mcp import MCPAuthType, MCPCredentials, MCPTransportType +from litellm.types.mcp_server.mcp_server_manager import MCPInfo + + +class MCPEnvVarScope(str, enum.Enum): + """Scope for an MCP server environment variable. + + - ``global``: value is provided by the admin and used for all users. + - ``user``: each user must provide their own value via the per-user + env-var endpoint. The admin-supplied ``value`` is treated as a + placeholder/hint and is not used at request time. + """ + + global_ = "global" + user = "user" + + +class MCPEnvVar(LiteLLMPydanticObjectBase): + """One environment variable for an MCP server. + + Variables can be interpolated into ``static_headers`` using ``${NAME}`` + syntax. ``scope=global`` values are stored on the server. ``scope=user`` + values are stored per-user in ``LiteLLM_MCPUserEnvVars`` and supplied by + each user. + """ + + name: str + value: str = "" + scope: MCPEnvVarScope = MCPEnvVarScope.global_ + description: Optional[str] = None + + +class LiteLLM_MCPServerTable(LiteLLMPydanticObjectBase): + """Represents a LiteLLM_MCPServerTable record""" + + server_id: str + server_name: Optional[str] = None + alias: Optional[str] = None + description: Optional[str] = None + url: Optional[str] = None + spec_path: Optional[str] = None + transport: MCPTransportType + auth_type: Optional[MCPAuthType] = None + credentials: Optional[MCPCredentials] = None + instructions: Optional[str] = None + created_at: Optional[datetime] = None + created_by: Optional[str] = None + updated_at: Optional[datetime] = None + updated_by: Optional[str] = None + teams: List[Dict[str, Optional[str]]] = Field(default_factory=list) + mcp_access_groups: List[str] = Field(default_factory=list) + allowed_tools: List[str] = Field(default_factory=list) + tool_name_to_display_name: Optional[Dict[str, str]] = None + tool_name_to_description: Optional[Dict[str, str]] = None + extra_headers: List[str] = Field(default_factory=list) + mcp_info: Optional[MCPInfo] = None + static_headers: Optional[Dict[str, str]] = None + env_vars: Optional[List[MCPEnvVar]] = None + status: Optional[Literal["healthy", "unhealthy", "unknown"]] = Field( + default="unknown", + description="Health status: 'healthy', 'unhealthy', 'unknown'", + ) + last_health_check: Optional[datetime] = None + health_check_error: Optional[str] = None + command: Optional[str] = None + args: List[str] = Field(default_factory=list) + env: Dict[str, str] = Field(default_factory=dict) + authorization_url: Optional[str] = None + token_url: Optional[str] = None + registration_url: Optional[str] = None + oauth2_flow: Optional[Literal["client_credentials", "authorization_code"]] = None + allow_all_keys: bool = False + available_on_public_internet: bool = True + delegate_auth_to_upstream: bool = False + oauth_passthrough: bool = False + is_byok: bool = False + byok_description: List[str] = Field(default_factory=list) + byok_api_key_help_url: Optional[str] = None + has_user_credential: Optional[bool] = None + source_url: Optional[str] = None + timeout: Optional[float] = None + approval_status: Optional[str] = Field( + default="active", + description="Approval status: 'pending_review', 'active', 'rejected'", + ) + submitted_by: Optional[str] = None + submitted_at: Optional[datetime] = None + reviewed_at: Optional[datetime] = None + review_notes: Optional[str] = None diff --git a/litellm/models/model.py b/litellm/models/model.py new file mode 100644 index 00000000000..7657e4d30f8 --- /dev/null +++ b/litellm/models/model.py @@ -0,0 +1,59 @@ +""" +Proxy model table model. + +Canonical definition for ``litellm_proxymodeltable``. Re-exported from +``litellm.proxy._types`` for backwards compatibility. +""" + +import json +from datetime import datetime +from typing import Optional + +from pydantic import ConfigDict, model_validator + +from litellm.types.llms.base import LiteLLMPydanticObjectBase + + +class LiteLLM_ProxyModelTable(LiteLLMPydanticObjectBase): + model_id: str + model_name: str + litellm_params: dict + model_info: Optional[dict] = None + blocked: bool = False + created_at: Optional[datetime] = None + created_by: Optional[str] = None + updated_at: Optional[datetime] = None + updated_by: Optional[str] = None + + model_config = ConfigDict(protected_namespaces=()) + + @model_validator(mode="before") + @classmethod + def check_potential_json_str(cls, values): + if isinstance(values.get("litellm_params"), str): + try: + values["litellm_params"] = json.loads(values["litellm_params"]) + except json.JSONDecodeError: + pass + if isinstance(values.get("model_info"), str): + try: + values["model_info"] = json.loads(values["model_info"]) + except json.JSONDecodeError: + pass + return values + + @property + def is_blocked(self) -> bool: + return self.blocked + + @property + def team_id(self) -> Optional[str]: + if self.model_info: + return self.model_info.get("team_id") + return None + + @property + def team_public_model_name(self) -> Optional[str]: + if self.model_info: + return self.model_info.get("team_public_model_name") + return None diff --git a/litellm/models/object_permission.py b/litellm/models/object_permission.py new file mode 100644 index 00000000000..6c0d100046c --- /dev/null +++ b/litellm/models/object_permission.py @@ -0,0 +1,26 @@ +""" +Object permission table model. + +Canonical definition for ``litellm_objectpermissiontable``. Re-exported from +``litellm.proxy._types`` for backwards compatibility. +""" + +from typing import Dict, List, Optional + +from litellm.types.llms.base import LiteLLMPydanticObjectBase + + +class LiteLLM_ObjectPermissionTable(LiteLLMPydanticObjectBase): + """Represents a LiteLLM_ObjectPermissionTable record""" + + object_permission_id: str + mcp_servers: Optional[List[str]] = [] + mcp_access_groups: Optional[List[str]] = [] + mcp_tool_permissions: Optional[Dict[str, List[str]]] = None + vector_stores: Optional[List[str]] = [] + agents: Optional[List[str]] = [] + agent_access_groups: Optional[List[str]] = [] + models: Optional[List[str]] = [] + mcp_toolsets: Optional[List[str]] = None + blocked_tools: Optional[List[str]] = [] + search_tools: Optional[List[str]] = [] diff --git a/litellm/models/organization.py b/litellm/models/organization.py new file mode 100644 index 00000000000..8b2d95c3e09 --- /dev/null +++ b/litellm/models/organization.py @@ -0,0 +1,31 @@ +""" +Organization table model. + +Canonical definition for ``litellm_organizationtable``. Re-exported from +``litellm.proxy._types`` for backwards compatibility. +""" + +from typing import List, Optional + +from litellm.models.budget import LiteLLM_BudgetTable +from litellm.models.object_permission import LiteLLM_ObjectPermissionTable +from litellm.models.user import LiteLLM_UserTable +from litellm.types.llms.base import LiteLLMPydanticObjectBase + + +class LiteLLM_OrganizationTable(LiteLLMPydanticObjectBase): + """Represents user-controllable params for a LiteLLM_OrganizationTable record""" + + organization_id: Optional[str] = None + organization_alias: Optional[str] = None + budget_id: str + spend: float = 0.0 + metadata: Optional[dict] = None + models: List[str] = [] + model_spend: Optional[dict] = {} + created_by: str + updated_by: str + users: Optional[List[LiteLLM_UserTable]] = None + litellm_budget_table: Optional[LiteLLM_BudgetTable] = None + object_permission: Optional[LiteLLM_ObjectPermissionTable] = None + object_permission_id: Optional[str] = None diff --git a/litellm/models/organization_membership.py b/litellm/models/organization_membership.py new file mode 100644 index 00000000000..9957c0c21af --- /dev/null +++ b/litellm/models/organization_membership.py @@ -0,0 +1,40 @@ +""" +Organization membership table model. + +Canonical definition for ``litellm_organizationmembership``. Re-exported from +``litellm.proxy._types`` for backwards compatibility. +""" + +from datetime import datetime +from typing import Any, Optional + +from pydantic import ConfigDict, model_validator + +from litellm.models.budget import LiteLLM_BudgetTable +from litellm.types.llms.base import LiteLLMPydanticObjectBase + + +class LiteLLM_OrganizationMembershipTable(LiteLLMPydanticObjectBase): + """Tracks which organizations a user belongs to and their spend within it.""" + + user_id: str + organization_id: str + user_role: Optional[str] = None + spend: float = 0.0 + budget_id: Optional[str] = None + created_at: datetime + updated_at: datetime + user: Optional[Any] = None + litellm_budget_table: Optional[LiteLLM_BudgetTable] = None + user_email: Optional[str] = None + + model_config = ConfigDict(protected_namespaces=()) + + @model_validator(mode="after") + def populate_user_email(self) -> "LiteLLM_OrganizationMembershipTable": + if self.user_email is None and self.user is not None: + if isinstance(self.user, dict): + self.user_email = self.user.get("user_email") + else: + self.user_email = getattr(self.user, "user_email", None) + return self diff --git a/litellm/models/project.py b/litellm/models/project.py new file mode 100644 index 00000000000..083c7ee3cc5 --- /dev/null +++ b/litellm/models/project.py @@ -0,0 +1,41 @@ +""" +Project table model. + +Canonical definition for ``litellm_projecttable``. Re-exported from +``litellm.proxy._types`` for backwards compatibility. +""" + +from datetime import datetime +from typing import List, Optional + +from litellm.models.budget import LiteLLM_BudgetTable +from litellm.models.object_permission import LiteLLM_ObjectPermissionTable +from litellm.types.llms.base import LiteLLMPydanticObjectBase + + +class LiteLLM_ProjectTable(LiteLLMPydanticObjectBase): + """Database model representation for project""" + + project_id: str + project_alias: Optional[str] = None + description: Optional[str] = None + team_id: Optional[str] = None + budget_id: Optional[str] = None + metadata: Optional[dict] = None + models: List[str] = [] + spend: float = 0.0 + model_spend: Optional[dict] = None + model_rpm_limit: Optional[dict] = None + model_tpm_limit: Optional[dict] = None + blocked: bool = False + object_permission_id: Optional[str] = None + created_by: Optional[str] = None + updated_by: Optional[str] = None + created_at: Optional[datetime] = None + updated_at: Optional[datetime] = None + litellm_budget_table: Optional[LiteLLM_BudgetTable] = None + object_permission: Optional[LiteLLM_ObjectPermissionTable] = None + + @property + def is_blocked(self) -> bool: + return self.blocked diff --git a/litellm/models/skills.py b/litellm/models/skills.py new file mode 100644 index 00000000000..62091c0ca01 --- /dev/null +++ b/litellm/models/skills.py @@ -0,0 +1,30 @@ +""" +Skills table model. + +Canonical definition for ``litellm_skillstable``. Re-exported from +``litellm.proxy._types`` for backwards compatibility. +""" + +from datetime import datetime +from typing import Any, Dict, Optional + +from litellm.types.llms.base import LiteLLMPydanticObjectBase + + +class LiteLLM_SkillsTable(LiteLLMPydanticObjectBase): + """Represents a LiteLLM_SkillsTable record""" + + skill_id: str + display_title: Optional[str] = None + description: Optional[str] = None + instructions: Optional[str] = None + source: str = "custom" + latest_version: Optional[str] = None + file_content: Optional[bytes] = None + file_name: Optional[str] = None + file_type: Optional[str] = None + metadata: Optional[Dict[str, Any]] = None + created_at: Optional[datetime] = None + created_by: Optional[str] = None + updated_at: Optional[datetime] = None + updated_by: Optional[str] = None diff --git a/litellm/models/spend_logs.py b/litellm/models/spend_logs.py new file mode 100644 index 00000000000..96bd328c3ca --- /dev/null +++ b/litellm/models/spend_logs.py @@ -0,0 +1,50 @@ +""" +Spend and error log table models. + +Canonical definitions for ``litellm_spendlogs`` and ``litellm_errorlogs``. +Re-exported from ``litellm.proxy._types`` for backwards compatibility. +""" + +from datetime import datetime +from typing import Optional, Union + +from pydantic import Json + +from litellm._uuid import uuid +from litellm.types.llms.base import LiteLLMPydanticObjectBase + + +class LiteLLM_SpendLogs(LiteLLMPydanticObjectBase): + request_id: str + api_key: str + model: Optional[str] = "" + api_base: Optional[str] = "" + call_type: str + spend: Optional[float] = 0.0 + total_tokens: Optional[int] = 0 + prompt_tokens: Optional[int] = 0 + completion_tokens: Optional[int] = 0 + startTime: Union[str, datetime, None] + endTime: Union[str, datetime, None] + user: Optional[str] = "" + metadata: Optional[Json] = {} + cache_hit: Optional[str] = "False" + cache_key: Optional[str] = None + request_tags: Optional[Json] = None + requester_ip_address: Optional[str] = None + messages: Optional[Union[str, list, dict]] + response: Optional[Union[str, list, dict]] + + +class LiteLLM_ErrorLogs(LiteLLMPydanticObjectBase): + request_id: Optional[str] = str(uuid.uuid4()) + api_base: Optional[str] = "" + model_group: Optional[str] = "" + litellm_model_name: Optional[str] = "" + model_id: Optional[str] = "" + request_kwargs: Optional[dict] = {} + exception_type: Optional[str] = "" + status_code: Optional[str] = "" + exception_string: Optional[str] = "" + startTime: Union[str, datetime, None] + endTime: Union[str, datetime, None] diff --git a/litellm/models/tag.py b/litellm/models/tag.py new file mode 100644 index 00000000000..02d8f58916d --- /dev/null +++ b/litellm/models/tag.py @@ -0,0 +1,36 @@ +""" +Tag table model. + +Canonical definition for ``litellm_tagtable``. Re-exported from +``litellm.proxy._types`` for backwards compatibility. +""" + +from datetime import datetime +from typing import List, Optional + +from pydantic import model_validator + +from litellm.models.budget import LiteLLM_BudgetTable +from litellm.types.llms.base import LiteLLMPydanticObjectBase + + +class LiteLLM_TagTable(LiteLLMPydanticObjectBase): + tag_name: str + description: Optional[str] = None + models: List[str] = [] + model_info: Optional[dict] = None + spend: float = 0.0 + budget_id: Optional[str] = None + litellm_budget_table: Optional[LiteLLM_BudgetTable] = None + created_at: Optional[datetime] = None + created_by: Optional[str] = None + updated_at: Optional[datetime] = None + + @model_validator(mode="before") + @classmethod + def set_model_info(cls, values): + if values.get("spend") is None: + values.update({"spend": 0.0}) + if values.get("models") is None: + values.update({"models": []}) + return values diff --git a/litellm/models/team.py b/litellm/models/team.py new file mode 100644 index 00000000000..aa0798955f2 --- /dev/null +++ b/litellm/models/team.py @@ -0,0 +1,154 @@ +""" +Team table models. + +Canonical definitions for ``litellm_teamtable`` (plus the shared Member and +budget-window value types and the team-model alias table). Re-exported from +``litellm.proxy._types`` for backwards compatibility. +""" + +import json +from datetime import datetime +from typing import List, Literal, Optional, Union + +from pydantic import BaseModel, ConfigDict, Field, model_validator + +from litellm.models.object_permission import LiteLLM_ObjectPermissionTable +from litellm.types.llms.base import LiteLLMPydanticObjectBase + + +class MemberBase(LiteLLMPydanticObjectBase): + user_id: Optional[str] = Field( + default=None, + description="The unique ID of the user to add. Either user_id or user_email must be provided", + ) + user_email: Optional[str] = Field( + default=None, + description="The email address of the user to add. Either user_id or user_email must be provided", + ) + + @model_validator(mode="before") + @classmethod + def check_user_info(cls, values): + if not isinstance(values, dict): + raise ValueError("input needs to be a dictionary") + if values.get("user_id") is None and values.get("user_email") is None: + raise ValueError("Either user id or user email must be provided") + return values + + +class Member(MemberBase): + role: Literal["admin", "user"] = Field( + description="The role of the user within the team. 'admin' users can manage team settings and members, 'user' is a regular team member" + ) + + +class BudgetLimitEntry(LiteLLMPydanticObjectBase): + """A single budget window with its own limit and independent reset schedule.""" + + budget_duration: str + max_budget: float + reset_at: Optional[datetime] = None + + +class LiteLLM_ModelTable(LiteLLMPydanticObjectBase): + id: Optional[int] = None + model_aliases: Optional[Union[str, dict]] = None + created_by: str + updated_by: str + team: Optional["LiteLLM_TeamTable"] = None + + model_config = ConfigDict(protected_namespaces=()) + + +class TeamBase(LiteLLMPydanticObjectBase): + team_alias: Optional[str] = None + team_id: Optional[str] = None + organization_id: Optional[str] = None + admins: list = [] + members: list = [] + members_with_roles: List[Member] = [] + team_member_permissions: Optional[List[str]] = None + metadata: Optional[dict] = None + tpm_limit: Optional[int] = None + rpm_limit: Optional[int] = None + max_budget: Optional[float] = None + soft_budget: Optional[float] = None + budget_duration: Optional[str] = None + budget_limits: Optional[List[BudgetLimitEntry]] = None + models: list = [] + blocked: bool = False + router_settings: Optional[dict] = None + access_group_ids: Optional[List[str]] = None + default_team_member_models: Optional[List[str]] = None + + +class LiteLLM_TeamTable(TeamBase): + team_id: str # type: ignore + spend: Optional[float] = None + max_parallel_requests: Optional[int] = None + budget_duration: Optional[str] = None + budget_reset_at: Optional[datetime] = None + model_id: Optional[int] = None + model_spend: Optional[dict] = {} + model_max_budget: Optional[dict] = {} + policies: Optional[List[str]] = None + allow_team_guardrail_config: Optional[bool] = False + litellm_model_table: Optional[LiteLLM_ModelTable] = None + object_permission: Optional[LiteLLM_ObjectPermissionTable] = None + object_permission_id: Optional[str] = None + updated_at: Optional[datetime] = None + created_at: Optional[datetime] = None + + model_config = ConfigDict(protected_namespaces=()) + + @model_validator(mode="before") + @classmethod + def set_model_info(cls, values): + dict_fields = [ + "metadata", + "aliases", + "config", + "permissions", + "model_max_budget", + "model_aliases", + "router_settings", + "budget_limits", + ] + + if isinstance(values, BaseModel): + values = values.model_dump() + + if ( + isinstance(values.get("members_with_roles"), dict) + and not values["members_with_roles"] + ): + values["members_with_roles"] = [] + + for field in dict_fields: + value = values.get(field) + if value is not None and isinstance(value, str): + try: + values[field] = json.loads(value) + except json.JSONDecodeError: + raise ValueError(f"Field {field} should be a valid dictionary") + + return values + + +class LiteLLM_TeamTableCachedObj(LiteLLM_TeamTable): + last_refreshed_at: Optional[float] = None + + +class LiteLLM_DeletedTeamTable(LiteLLM_TeamTable): + """Audit record for deleted teams; mirrors the team plus deletion metadata.""" + + id: Optional[str] = None + deleted_at: Optional[datetime] = None + deleted_by: Optional[str] = None + deleted_by_api_key: Optional[str] = None + litellm_changed_by: Optional[str] = None + + model_config = ConfigDict(protected_namespaces=()) + + +LiteLLM_ModelTable.model_rebuild() diff --git a/litellm/models/team_membership.py b/litellm/models/team_membership.py new file mode 100644 index 00000000000..d0a1308ce7c --- /dev/null +++ b/litellm/models/team_membership.py @@ -0,0 +1,32 @@ +""" +Team membership table model. + +Canonical definition for ``litellm_teammembership``. Re-exported from +``litellm.proxy._types`` for backwards compatibility. +""" + +from typing import Optional, Union + +from litellm.models.budget import LiteLLM_BudgetTable, LiteLLM_BudgetTableFull +from litellm.types.llms.base import LiteLLMPydanticObjectBase + + +class LiteLLM_TeamMembership(LiteLLMPydanticObjectBase): + user_id: str + team_id: str + budget_id: Optional[str] = None + spend: Optional[float] = 0.0 + total_spend: Optional[float] = 0.0 + litellm_budget_table: Optional[ + Union[LiteLLM_BudgetTableFull, LiteLLM_BudgetTable] + ] = None + + def safe_get_team_member_rpm_limit(self) -> Optional[int]: + if self.litellm_budget_table is not None: + return self.litellm_budget_table.rpm_limit + return None + + def safe_get_team_member_tpm_limit(self) -> Optional[int]: + if self.litellm_budget_table is not None: + return self.litellm_budget_table.tpm_limit + return None diff --git a/litellm/models/user.py b/litellm/models/user.py new file mode 100644 index 00000000000..cd7e9db4aec --- /dev/null +++ b/litellm/models/user.py @@ -0,0 +1,70 @@ +""" +User table model. + +Canonical definition for ``litellm_usertable``. Re-exported from +``litellm.proxy._types`` for backwards compatibility. +""" + +from datetime import datetime +from typing import Dict, List, Optional + +from pydantic import ConfigDict, Field, model_validator + +from litellm.models.object_permission import LiteLLM_ObjectPermissionTable +from litellm.models.organization_membership import ( + LiteLLM_OrganizationMembershipTable, +) +from litellm.types.llms.base import LiteLLMPydanticObjectBase + + +class LiteLLM_UserTable(LiteLLMPydanticObjectBase): + user_id: str + user_alias: Optional[str] = None + team_id: Optional[str] = None + sso_user_id: Optional[str] = None + organization_id: Optional[str] = None + object_permission_id: Optional[str] = None + password: Optional[str] = Field(default=None, exclude=True) + teams: List[str] = [] + user_role: Optional[str] = None + max_budget: Optional[float] = None + spend: float = 0.0 + user_email: Optional[str] = None + models: list = [] + metadata: Optional[dict] = None + max_parallel_requests: Optional[int] = None + tpm_limit: Optional[int] = None + rpm_limit: Optional[int] = None + budget_duration: Optional[str] = None + budget_reset_at: Optional[datetime] = None + allowed_cache_controls: List[str] = [] + policies: List[str] = [] + model_spend: Optional[Dict] = {} + model_max_budget: Optional[Dict] = {} + created_at: Optional[datetime] = None + updated_at: Optional[datetime] = None + organization_memberships: Optional[List[LiteLLM_OrganizationMembershipTable]] = None + object_permission: Optional[LiteLLM_ObjectPermissionTable] = None + + model_config = ConfigDict(protected_namespaces=()) + + @model_validator(mode="before") + @classmethod + def set_model_info(cls, values): + if values.get("spend") is None: + values.update({"spend": 0.0}) + if values.get("models") is None: + values.update({"models": []}) + if values.get("teams") is None: + values.update({"teams": []}) + return values + + def is_over_budget(self) -> bool: + if self.max_budget is None: + return False + return self.spend >= self.max_budget + + def has_model_access(self, model_name: str) -> bool: + if not self.models: + return True + return model_name in self.models diff --git a/litellm/models/verification_token.py b/litellm/models/verification_token.py new file mode 100644 index 00000000000..8bddd1c1619 --- /dev/null +++ b/litellm/models/verification_token.py @@ -0,0 +1,74 @@ +""" +Verification token table model. + +Canonical definition for ``litellm_verificationtoken``. Re-exported from +``litellm.proxy._types`` for backwards compatibility. +""" + +from datetime import datetime +from typing import Dict, List, Optional, Union + +from pydantic import ConfigDict + +from litellm.models.object_permission import LiteLLM_ObjectPermissionTable +from litellm.types.llms.base import LiteLLMPydanticObjectBase + + +class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase): + token: Optional[str] = None + key_name: Optional[str] = None + key_alias: Optional[str] = None + spend: float = 0.0 + max_budget: Optional[float] = None + expires: Optional[Union[str, datetime]] = None + models: List = [] + aliases: Dict = {} + config: Dict = {} + user_id: Optional[str] = None + team_id: Optional[str] = None + agent_id: Optional[str] = None + project_id: Optional[str] = None + max_parallel_requests: Optional[int] = None + metadata: Dict = {} + tpm_limit: Optional[int] = None + rpm_limit: Optional[int] = None + budget_duration: Optional[str] = None + budget_reset_at: Optional[datetime] = None + allowed_cache_controls: Optional[list] = [] + allowed_routes: Optional[list] = [] + permissions: Dict = {} + model_spend: Dict = {} + model_max_budget: Dict = {} + soft_budget_cooldown: bool = False + blocked: Optional[bool] = None + litellm_budget_table: Optional[dict] = None + budget_id: Optional[str] = None + org_id: Optional[str] = None # org id for a given key + created_at: Optional[datetime] = None + created_by: Optional[str] = None + updated_at: Optional[datetime] = None + updated_by: Optional[str] = None + last_active: Optional[datetime] = None + object_permission_id: Optional[str] = None + object_permission: Optional[LiteLLM_ObjectPermissionTable] = None + access_group_ids: Optional[List[str]] = None + rotation_count: Optional[int] = 0 + auto_rotate: Optional[bool] = False + rotation_interval: Optional[str] = None + last_rotation_at: Optional[datetime] = None + key_rotation_at: Optional[datetime] = None + router_settings: Optional[dict] = None + budget_limits: Optional[List[dict]] = None + model_config = ConfigDict(protected_namespaces=()) + + +class LiteLLM_DeletedVerificationToken(LiteLLM_VerificationToken): + """Audit record for deleted keys; mirrors the token plus deletion metadata.""" + + id: Optional[str] = None + deleted_at: Optional[datetime] = None + deleted_by: Optional[str] = None + deleted_by_api_key: Optional[str] = None + litellm_changed_by: Optional[str] = None + + model_config = ConfigDict(protected_namespaces=()) diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 863e6acd41e..dcf7660d002 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -14,8 +14,12 @@ from litellm.proxy._types import ( SpecialHeaders, UserAPIKeyAuth, ) -from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.auth.ip_address_utils import IPAddressUtils +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.repositories.table_repositories import ( + AgentsRepository, + MCPServerRepository, +) def _parse_mcp_server_names_from_path( @@ -1445,7 +1449,7 @@ class MCPRequestHandler: return None if object_permission_id is None: - agent_row = await prisma_client.db.litellm_agentstable.find_unique( + agent_row = await AgentsRepository(prisma_client).table.find_unique( where={"agent_id": agent_id}, ) object_permission_id = ( @@ -1600,7 +1604,7 @@ class MCPRequestHandler: server_ids: Set[str] = set() if access_groups and prisma_client is not None: try: - mcp_servers = await prisma_client.db.litellm_mcpservertable.find_many( + mcp_servers = await MCPServerRepository(prisma_client).table.find_many( where={"mcp_access_groups": {"hasSome": access_groups}} ) for server in mcp_servers: diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 0ba0181200f..c52752940c3 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -8,6 +8,7 @@ from typing import Any, Dict, Iterable, List, Optional, Set, Union, cast from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.constants import MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS +from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.proxy._types import ( LiteLLM_MCPServerTable, LiteLLM_ObjectPermissionTable, @@ -25,8 +26,16 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( decrypt_value_helper, encrypt_value_helper, ) -from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.proxy.utils import PrismaClient +from litellm.repositories.object_permission_repository import ObjectPermissionRepository +from litellm.repositories.table_repositories import ( + MCPServerRepository, + MCPUserCredentialsRepository, +) +from litellm.repositories.team_repository import TeamRepository +from litellm.repositories.verification_token_repository import ( + VerificationTokenRepository, +) from litellm.types.llms.custom_http import httpxSpecialProvider from litellm.types.mcp import MCPCredentials @@ -354,7 +363,7 @@ async def get_all_mcp_servers( where: Dict[str, Any] = {} if approval_status is not None: where["approval_status"] = approval_status - mcp_servers = await prisma_client.db.litellm_mcpservertable.find_many( + mcp_servers = await MCPServerRepository(prisma_client).table.find_many( where=where if where else {} ) @@ -380,7 +389,9 @@ async def get_mcp_server( """ Returns the matching mcp server from the db iff exists """ - mcp_server = await prisma_client.db.litellm_mcpservertable.find_unique( + mcp_server: Optional[LiteLLM_MCPServerTable] = await MCPServerRepository( + prisma_client + ).table.find_unique( where={ "server_id": server_id, } @@ -398,12 +409,12 @@ async def get_mcp_servers( """ Returns the matching mcp servers from the db with the server_ids """ - _mcp_servers: List[LiteLLM_MCPServerTable] = ( - await prisma_client.db.litellm_mcpservertable.find_many( - where={ - "server_id": {"in": server_ids}, - } - ) + _mcp_servers: List[LiteLLM_MCPServerTable] = await MCPServerRepository( + prisma_client + ).table.find_many( + where={ + "server_id": {"in": server_ids}, + } ) final_mcp_servers: List[LiteLLM_MCPServerTable] = [] for _mcp_server in _mcp_servers: @@ -420,15 +431,15 @@ async def get_mcp_servers_by_verificationtoken( """ Returns the mcp servers from the db for the verification token """ - verification_token_record: LiteLLM_TeamTable = ( - await prisma_client.db.litellm_verificationtoken.find_unique( - where={ - "token": token, - }, - include={ - "object_permission": True, - }, - ) + verification_token_record: LiteLLM_TeamTable = await VerificationTokenRepository( + prisma_client + ).table.find_unique( + where={ + "token": token, + }, + include={ + "object_permission": True, + }, ) mcp_servers: Optional[List[str]] = [] @@ -446,15 +457,15 @@ async def get_mcp_servers_by_team( """ Returns the mcp servers from the db for the team id """ - team_record: LiteLLM_TeamTable = ( - await prisma_client.db.litellm_teamtable.find_unique( - where={ - "team_id": team_id, - }, - include={ - "object_permission": True, - }, - ) + team_record: LiteLLM_TeamTable = await TeamRepository( + prisma_client + ).table.find_unique( + where={ + "team_id": team_id, + }, + include={ + "object_permission": True, + }, ) mcp_servers: Optional[List[str]] = [] @@ -505,16 +516,16 @@ async def get_objectpermissions_for_mcp_server( """ Get all the object permissions records and the associated team and verficiationtoken records that have access to the mcp server """ - object_permission_records = ( - await prisma_client.db.litellm_objectpermissiontable.find_many( - where={ - "mcp_servers": {"has": mcp_server_id}, - }, - include={ - "teams": True, - "verification_tokens": True, - }, - ) + object_permission_records = await ObjectPermissionRepository( + prisma_client + ).table.find_many( + where={ + "mcp_servers": {"has": mcp_server_id}, + }, + include={ + "teams": True, + "verification_tokens": True, + }, ) return object_permission_records @@ -526,7 +537,7 @@ async def get_virtualkeys_for_mcp_server( """ Get all the virtual keys that have access to the mcp server """ - virtual_keys = await prisma_client.db.litellm_verificationtoken.find_many( + virtual_keys = await VerificationTokenRepository(prisma_client).table.find_many( where={ "mcp_servers": {"has": server_id}, }, @@ -564,7 +575,7 @@ async def delete_mcp_server( Returns the deleted mcp server record if it exists, otherwise None """ - deleted_server = await prisma_client.db.litellm_mcpservertable.delete( + deleted_server = await MCPServerRepository(prisma_client).table.delete( where={ "server_id": server_id, }, @@ -600,7 +611,7 @@ async def create_mcp_server( data_dict["created_by"] = touched_by data_dict["updated_by"] = touched_by - new_mcp_server = await prisma_client.db.litellm_mcpservertable.create( + new_mcp_server = await MCPServerRepository(prisma_client).table.create( data=data_dict # type: ignore ) @@ -635,7 +646,7 @@ async def update_mcp_server( "credentials" in data_dict and data_dict["credentials"] is not None ) if data.auth_type or has_credentials: - existing = await prisma_client.db.litellm_mcpservertable.find_unique( + existing = await MCPServerRepository(prisma_client).table.find_unique( where={"server_id": data.server_id} ) @@ -678,7 +689,7 @@ async def update_mcp_server( # Add audit fields data_dict["updated_by"] = touched_by - updated_mcp_server = await prisma_client.db.litellm_mcpservertable.update( + updated_mcp_server = await MCPServerRepository(prisma_client).table.update( where={"server_id": data.server_id}, data=data_dict # type: ignore ) @@ -691,7 +702,7 @@ async def rotate_mcp_server_credentials_master_key( ): from litellm.litellm_core_utils.safe_json_dumps import safe_dumps - mcp_servers = await prisma_client.db.litellm_mcpservertable.find_many() + mcp_servers = await MCPServerRepository(prisma_client).table.find_many() updated = 0 for mcp_server in mcp_servers: @@ -719,7 +730,7 @@ async def rotate_mcp_server_credentials_master_key( continue update_data["updated_by"] = touched_by - await prisma_client.db.litellm_mcpservertable.update( + await MCPServerRepository(prisma_client).table.update( where={"server_id": mcp_server.server_id}, data=update_data, ) @@ -781,7 +792,7 @@ async def rotate_mcp_user_credentials_master_key( under the new master key. Rows that are unreadable under both paths are logged and skipped so one corrupt row does not abort the rotation. """ - rows = await prisma_client.db.litellm_mcpusercredentials.find_many() + rows = await MCPUserCredentialsRepository(prisma_client).table.find_many() rotated = 0 skipped = 0 for row in rows: @@ -798,7 +809,7 @@ async def rotate_mcp_user_credentials_master_key( re_encrypted = encrypt_value_helper( plaintext, new_encryption_key=new_master_key ) - await prisma_client.db.litellm_mcpusercredentials.update( + await MCPUserCredentialsRepository(prisma_client).table.update( where={ "user_id_server_id": { "user_id": row.user_id, @@ -873,7 +884,7 @@ async def store_user_credential( """Store a user credential for a BYOK MCP server.""" encoded = encrypt_value_helper(credential) - await prisma_client.db.litellm_mcpusercredentials.upsert( + await MCPUserCredentialsRepository(prisma_client).table.upsert( where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}}, data={ "create": { @@ -893,7 +904,7 @@ async def get_user_credential( ) -> Optional[str]: """Return credential for a user+server pair, or None.""" - row = await prisma_client.db.litellm_mcpusercredentials.find_unique( + row = await MCPUserCredentialsRepository(prisma_client).table.find_unique( where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}} ) if row is None: @@ -907,7 +918,7 @@ async def has_user_credential( server_id: str, ) -> bool: """Return True if the user has a stored credential for this server.""" - row = await prisma_client.db.litellm_mcpusercredentials.find_unique( + row = await MCPUserCredentialsRepository(prisma_client).table.find_unique( where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}} ) return row is not None @@ -919,7 +930,7 @@ async def delete_user_credential( server_id: str, ) -> None: """Delete the user's stored credential for a BYOK MCP server.""" - await prisma_client.db.litellm_mcpusercredentials.delete( + await MCPUserCredentialsRepository(prisma_client).table.delete( where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}} ) @@ -966,7 +977,7 @@ async def store_user_oauth_credential( # Skip the guard when the caller knows the row is already an OAuth2 credential # (e.g. during token refresh), saving an extra DB round-trip. if not skip_byok_guard: - existing = await prisma_client.db.litellm_mcpusercredentials.find_unique( + existing = await MCPUserCredentialsRepository(prisma_client).table.find_unique( where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}} ) if ( @@ -984,7 +995,7 @@ async def store_user_oauth_credential( ) encoded = encrypt_value_helper(json.dumps(payload)) - await prisma_client.db.litellm_mcpusercredentials.upsert( + await MCPUserCredentialsRepository(prisma_client).table.upsert( where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}}, data={ "create": { @@ -1025,7 +1036,7 @@ async def get_user_oauth_credential( ) -> Optional[Dict[str, Any]]: """Return the decoded OAuth2 payload dict for a user+server pair, or None.""" - row = await prisma_client.db.litellm_mcpusercredentials.find_unique( + row = await MCPUserCredentialsRepository(prisma_client).table.find_unique( where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}} ) if row is None: @@ -1039,7 +1050,7 @@ async def list_user_oauth_credentials( ) -> List[Dict[str, Any]]: """Return all OAuth2 credential payloads for a user, tagged with server_id.""" - rows = await prisma_client.db.litellm_mcpusercredentials.find_many( + rows = await MCPUserCredentialsRepository(prisma_client).table.find_many( where={"user_id": user_id} ) results: List[Dict[str, Any]] = [] @@ -1212,7 +1223,7 @@ async def approve_mcp_server( ) -> LiteLLM_MCPServerTable: """Set approval_status=active and record reviewed_at.""" now = datetime.now(timezone.utc) - updated = await prisma_client.db.litellm_mcpservertable.update( + updated = await MCPServerRepository(prisma_client).table.update( where={"server_id": server_id}, data={ "approval_status": MCPApprovalStatus.active, @@ -1240,7 +1251,7 @@ async def reject_mcp_server( } if review_notes is not None: data["review_notes"] = review_notes - updated = await prisma_client.db.litellm_mcpservertable.update( + updated = await MCPServerRepository(prisma_client).table.update( where={"server_id": server_id}, data=data, ) @@ -1257,7 +1268,7 @@ async def get_mcp_submissions( along with a summary count breakdown by approval_status. Mirrors get_guardrail_submissions() from guardrail_endpoints.py. """ - rows = await prisma_client.db.litellm_mcpservertable.find_many( + rows = await MCPServerRepository(prisma_client).table.find_many( where={"submitted_at": {"not": None}}, order={"submitted_at": "desc"}, take=500, # safety cap; paginate if needed in a future iteration diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 7048f5bf7c4..85ac6b399f4 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -42,8 +42,8 @@ from litellm.constants import ( MCP_TOOL_LISTING_TIMEOUT, ) from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException -from litellm.litellm_core_utils.url_utils import SSRFError, async_safe_get from litellm.experimental_mcp_client.client import MCPClient, MCPSigV4Auth +from litellm.litellm_core_utils.url_utils import SSRFError, async_safe_get from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, @@ -85,6 +85,7 @@ from litellm.proxy._types import ( from litellm.proxy.auth.ip_address_utils import IPAddressUtils from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper from litellm.proxy.utils import ProxyLogging +from litellm.repositories.table_repositories import MCPServerRepository from litellm.types.llms.custom_http import httpxSpecialProvider from litellm.types.mcp import MCPAuth, MCPStdioConfig from litellm.types.mcp_server.mcp_server_manager import ( @@ -3817,7 +3818,7 @@ class MCPServerManager: # Pending/rejected servers are excluded at the DB level so we never load them. from litellm.proxy._experimental.mcp_server.db import LiteLLM_MCPServerTable - raw_rows = await prisma_client.db.litellm_mcpservertable.find_many( + raw_rows = await MCPServerRepository(prisma_client).table.find_many( where={ "OR": [ {"approval_status": None}, diff --git a/litellm/proxy/_experimental/mcp_server/toolset_db.py b/litellm/proxy/_experimental/mcp_server/toolset_db.py index 08ac7dbd33b..a996131653f 100644 --- a/litellm/proxy/_experimental/mcp_server/toolset_db.py +++ b/litellm/proxy/_experimental/mcp_server/toolset_db.py @@ -4,6 +4,7 @@ from typing import List, Optional from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.proxy.utils import PrismaClient +from litellm.repositories.table_repositories import MCPToolsetRepository from litellm.types.mcp_server.mcp_toolset import ( MCPToolset, NewMCPToolsetRequest, @@ -30,7 +31,7 @@ async def create_mcp_toolset( data_dict["tools"] = json.dumps(data_dict.get("tools", [])) data_dict["created_by"] = touched_by data_dict["updated_by"] = touched_by - row = await prisma_client.db.litellm_mcptoolsettable.create(data=data_dict) + row = await MCPToolsetRepository(prisma_client).table.create(data=data_dict) return _toolset_from_row(row) @@ -38,7 +39,7 @@ async def get_mcp_toolset( prisma_client: PrismaClient, toolset_id: str, ) -> Optional[MCPToolset]: - row = await prisma_client.db.litellm_mcptoolsettable.find_unique( + row = await MCPToolsetRepository(prisma_client).table.find_unique( where={"toolset_id": toolset_id} ) if row is None: @@ -54,7 +55,7 @@ async def list_mcp_toolsets( where = {} if toolset_ids is not None: where = {"toolset_id": {"in": toolset_ids}} - rows = await prisma_client.db.litellm_mcptoolsettable.find_many(where=where) + rows = await MCPToolsetRepository(prisma_client).table.find_many(where=where) return [_toolset_from_row(r) for r in rows] except Exception as e: verbose_proxy_logger.warning( @@ -69,7 +70,7 @@ async def get_mcp_toolset_by_name( prisma_client: PrismaClient, toolset_name: str, ) -> Optional[MCPToolset]: - row = await prisma_client.db.litellm_mcptoolsettable.find_first( + row = await MCPToolsetRepository(prisma_client).table.find_first( where={"toolset_name": toolset_name} ) if row is None: @@ -87,7 +88,7 @@ async def update_mcp_toolset( data_dict["tools"] = json.dumps(data_dict["tools"]) data_dict["updated_by"] = touched_by try: - row = await prisma_client.db.litellm_mcptoolsettable.update( + row = await MCPToolsetRepository(prisma_client).table.update( where={"toolset_id": data.toolset_id}, data=data_dict, ) @@ -105,7 +106,7 @@ async def delete_mcp_toolset( toolset_id: str, ) -> Optional[MCPToolset]: try: - row = await prisma_client.db.litellm_mcptoolsettable.delete( + row = await MCPToolsetRepository(prisma_client).table.delete( where={"toolset_id": toolset_id} ) except Exception as e: diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 88be567e59a..57a7d860baa 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -23,8 +23,6 @@ from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( from litellm.types.integrations.slack_alerting import AlertType from litellm.types.llms.openai import ( AllMessageValues, - OpenAIFileObject, - ResponsesAPIResponse, ) from litellm.types.mcp import ( MCPAuthType, @@ -41,8 +39,6 @@ from litellm.types.utils import ( EmbeddingResponse, GenericBudgetConfigType, ImageResponse, - LiteLLMBatch, - LiteLLMFineTuningJob, LiteLLMPydanticObjectBase, ModelResponse, ProviderField, @@ -1014,12 +1010,7 @@ class LiteLLM_ObjectPermissionBase(LiteLLMPydanticObjectBase): search_tools: Optional[List[str]] = None -class BudgetLimitEntry(LiteLLMPydanticObjectBase): - """A single budget window with its own limit and independent reset schedule.""" - - budget_duration: str # e.g. "24h", "7d", "30d" - max_budget: float # max spend in USD for this window - reset_at: Optional[datetime] = None # populated at creation/reset time +from litellm.models.team import BudgetLimitEntry as BudgetLimitEntry # noqa: E402 class GenerateRequestBase(LiteLLMPydanticObjectBase): @@ -1217,40 +1208,10 @@ class KeyRequest(LiteLLMPydanticObjectBase): return values -class LiteLLM_ModelTable(LiteLLMPydanticObjectBase): - id: Optional[int] = None - model_aliases: Optional[Union[str, dict]] = None # json dump the dict - created_by: str - updated_by: str - team: Optional["LiteLLM_TeamTable"] = None - - model_config = ConfigDict(protected_namespaces=()) - - -class LiteLLM_ProxyModelTable(LiteLLMPydanticObjectBase): - model_id: str - model_name: str - litellm_params: dict - model_info: dict - created_at: Optional[datetime] = None - created_by: str - updated_at: Optional[datetime] = None - updated_by: str - - @model_validator(mode="before") - @classmethod - def check_potential_json_str(cls, values): - if isinstance(values.get("litellm_params"), str): - try: - values["litellm_params"] = json.loads(values["litellm_params"]) - except json.JSONDecodeError: - pass - if isinstance(values.get("model_info"), str): - try: - values["model_info"] = json.loads(values["model_info"]) - except json.JSONDecodeError: - pass - return values +from litellm.models.model import ( # noqa: E402 + LiteLLM_ProxyModelTable as LiteLLM_ProxyModelTable, +) +from litellm.models.team import LiteLLM_ModelTable as LiteLLM_ModelTable # noqa: E402 # MCP Types @@ -1265,32 +1226,12 @@ class MCPApprovalStatus(str, enum.Enum): rejected = "rejected" -class MCPEnvVarScope(str, enum.Enum): - """Scope for an MCP server environment variable. - - - ``global``: value is provided by the admin and used for all users. - - ``user``: each user must provide their own value via the per-user - env-var endpoint. The admin-supplied ``value`` is treated as a - placeholder/hint and is not used at request time. - """ - - global_ = "global" - user = "user" - - -class MCPEnvVar(LiteLLMPydanticObjectBase): - """One environment variable for an MCP server. - - Variables can be interpolated into ``static_headers`` using ``${NAME}`` - syntax. ``scope=global`` values are stored on the server. ``scope=user`` - values are stored per-user in ``LiteLLM_MCPUserEnvVars`` and supplied by - each user. - """ - - name: str - value: str = "" - scope: MCPEnvVarScope = MCPEnvVarScope.global_ - description: Optional[str] = None +from litellm.models.mcp_server import ( # noqa: E402 + MCPEnvVar as MCPEnvVar, +) +from litellm.models.mcp_server import ( # noqa: E402 + MCPEnvVarScope as MCPEnvVarScope, +) # MCP Proxy Request Types @@ -1443,66 +1384,9 @@ class UpdateMCPServerRequest(LiteLLMPydanticObjectBase): return values -class LiteLLM_MCPServerTable(LiteLLMPydanticObjectBase): - """Represents a LiteLLM_MCPServerTable record""" - - server_id: str - server_name: Optional[str] = None - alias: Optional[str] = None - description: Optional[str] = None - url: Optional[str] = None - spec_path: Optional[str] = None - transport: MCPTransportType - auth_type: Optional[MCPAuthType] = None - credentials: Optional[MCPCredentials] = None - instructions: Optional[str] = None - created_at: Optional[datetime] = None - created_by: Optional[str] = None - updated_at: Optional[datetime] = None - updated_by: Optional[str] = None - teams: List[Dict[str, Optional[str]]] = Field(default_factory=list) - mcp_access_groups: List[str] = Field(default_factory=list) - allowed_tools: List[str] = Field(default_factory=list) - tool_name_to_display_name: Optional[Dict[str, str]] = None - tool_name_to_description: Optional[Dict[str, str]] = None - extra_headers: List[str] = Field(default_factory=list) - mcp_info: Optional[MCPInfo] = None - static_headers: Optional[Dict[str, str]] = None - env_vars: Optional[List[MCPEnvVar]] = None - # Health check status - status: Optional[Literal["healthy", "unhealthy", "unknown"]] = Field( - default="unknown", - description="Health status: 'healthy', 'unhealthy', 'unknown'", - ) - last_health_check: Optional[datetime] = None - health_check_error: Optional[str] = None - # Stdio-specific fields - command: Optional[str] = None - args: List[str] = Field(default_factory=list) - env: Dict[str, str] = Field(default_factory=dict) - authorization_url: Optional[str] = None - token_url: Optional[str] = None - registration_url: Optional[str] = None - oauth2_flow: Optional[Literal["client_credentials", "authorization_code"]] = None - allow_all_keys: bool = False - available_on_public_internet: bool = True - delegate_auth_to_upstream: bool = False - oauth_passthrough: bool = False - is_byok: bool = False - byok_description: List[str] = Field(default_factory=list) - byok_api_key_help_url: Optional[str] = None - has_user_credential: Optional[bool] = None - source_url: Optional[str] = None - timeout: Optional[float] = None - # BYOM submission fields - approval_status: Optional[str] = Field( - default="active", - description="Approval status: 'pending_review', 'active', 'rejected'", - ) - submitted_by: Optional[str] = None - submitted_at: Optional[datetime] = None - reviewed_at: Optional[datetime] = None - review_notes: Optional[str] = None +from litellm.models.mcp_server import ( # noqa: E402 + LiteLLM_MCPServerTable as LiteLLM_MCPServerTable, +) class MakeMCPServersPublicRequest(LiteLLMPydanticObjectBase): @@ -1622,23 +1506,9 @@ class UpdateSkillRequest(LiteLLMPydanticObjectBase): metadata: Optional[Dict[str, Any]] = None -class LiteLLM_SkillsTable(LiteLLMPydanticObjectBase): - """Represents a LiteLLM_SkillsTable record""" - - skill_id: str - display_title: Optional[str] = None - description: Optional[str] = None - instructions: Optional[str] = None - source: str = "custom" - latest_version: Optional[str] = None - file_content: Optional[bytes] = None # Binary content of skill files (zip) - file_name: Optional[str] = None # Original filename - file_type: Optional[str] = None # MIME type - metadata: Optional[Dict[str, Any]] = None - created_at: Optional[datetime] = None - created_by: Optional[str] = None - updated_at: Optional[datetime] = None - updated_by: Optional[str] = None +from litellm.models.skills import ( # noqa: E402 + LiteLLM_SkillsTable as LiteLLM_SkillsTable, +) class ListSkillsRequest(LiteLLMPydanticObjectBase): @@ -1839,33 +1709,8 @@ class DeleteCustomerRequest(LiteLLMPydanticObjectBase): user_ids: List[str] -class MemberBase(LiteLLMPydanticObjectBase): - user_id: Optional[str] = Field( - default=None, - description="The unique ID of the user to add. Either user_id or user_email must be provided", - ) - user_email: Optional[str] = Field( - default=None, - description="The email address of the user to add. Either user_id or user_email must be provided", - ) - - @model_validator(mode="before") - @classmethod - def check_user_info(cls, values): - if not isinstance(values, dict): - raise ValueError("input needs to be a dictionary") - if values.get("user_id") is None and values.get("user_email") is None: - raise ValueError("Either user id or user email must be provided") - return values - - -class Member(MemberBase): - role: Literal[ - "admin", - "user", - ] = Field( - description="The role of the user within the team. 'admin' users can manage team settings and members, 'user' is a regular team member" - ) +from litellm.models.team import Member as Member # noqa: E402 +from litellm.models.team import MemberBase as MemberBase # noqa: E402 class OrgMember(MemberBase): @@ -1876,33 +1721,7 @@ class OrgMember(MemberBase): ] -class TeamBase(LiteLLMPydanticObjectBase): - team_alias: Optional[str] = None - team_id: Optional[str] = None - organization_id: Optional[str] = None - admins: list = [] - members: list = [] - members_with_roles: List[Member] = [] - team_member_permissions: Optional[List[str]] = None - metadata: Optional[dict] = None - tpm_limit: Optional[int] = None - rpm_limit: Optional[int] = None - - # Budget fields - max_budget: Optional[float] = None - soft_budget: Optional[float] = None - budget_duration: Optional[str] = None - budget_limits: Optional[List[BudgetLimitEntry]] = ( - None # multiple concurrent budget windows - ) - - models: list = [] - blocked: bool = False - router_settings: Optional[dict] = None - access_group_ids: Optional[List[str]] = None - default_team_member_models: Optional[List[str]] = ( - None # default allowed_models seeded onto new team members - ) +from litellm.models.team import TeamBase as TeamBase # noqa: E402 class NewTeamRequest(TeamBase): @@ -2100,147 +1919,31 @@ class TeamCallbackMetadata(LiteLLMPydanticObjectBase): return values -class LiteLLM_ObjectPermissionTable(LiteLLMPydanticObjectBase): - """Represents a LiteLLM_ObjectPermissionTable record""" - - object_permission_id: str - mcp_servers: Optional[List[str]] = [] - mcp_access_groups: Optional[List[str]] = [] - mcp_tool_permissions: Optional[Dict[str, List[str]]] = None - """ - Mapping - server_id -> list of tools - - Enforces allowed tools for a specific key/team/organization - { - "1234567890": ["tool_name_1", "tool_name_2"] - } - """ - - vector_stores: Optional[List[str]] = [] - agents: Optional[List[str]] = [] - agent_access_groups: Optional[List[str]] = [] - mcp_toolsets: Optional[List[str]] = None - blocked_tools: Optional[List[str]] = [] - search_tools: Optional[List[str]] = [] - - -class LiteLLM_TeamTable(TeamBase): - team_id: str # type: ignore - spend: Optional[float] = None - max_parallel_requests: Optional[int] = None - budget_duration: Optional[str] = None - budget_reset_at: Optional[datetime] = None - model_id: Optional[int] = None - litellm_model_table: Optional[LiteLLM_ModelTable] = None - object_permission: Optional[LiteLLM_ObjectPermissionTable] = None - updated_at: Optional[datetime] = None - created_at: Optional[datetime] = None - - ######################################################### - # Object Permission - MCP, Vector Stores etc. - ######################################################### - object_permission_id: Optional[str] = None - - model_config = ConfigDict(protected_namespaces=()) - - @model_validator(mode="before") - @classmethod - def set_model_info(cls, values): - dict_fields = [ - "metadata", - "aliases", - "config", - "permissions", - "model_max_budget", - "model_aliases", - "router_settings", - "budget_limits", - ] - - if isinstance(values, BaseModel): - values = values.model_dump() - - if ( - isinstance(values.get("members_with_roles"), dict) - and not values["members_with_roles"] - ): - values["members_with_roles"] = [] - - for field in dict_fields: - value = values.get(field) - if value is not None and isinstance(value, str): - try: - values[field] = json.loads(value) - except json.JSONDecodeError: - raise ValueError(f"Field {field} should be a valid dictionary") - - return values - - -class LiteLLM_TeamTableCachedObj(LiteLLM_TeamTable): - last_refreshed_at: Optional[float] = None - - -class LiteLLM_DeletedTeamTable(LiteLLM_TeamTable): - """ - Recording of deleted teams for audit purposes. Mirrors LiteLLM_TeamTable - plus metadata captured at deletion time. - """ - - id: Optional[str] = None - deleted_at: Optional[datetime] = None - deleted_by: Optional[str] = None - deleted_by_api_key: Optional[str] = None - litellm_changed_by: Optional[str] = None - - model_config = ConfigDict(protected_namespaces=()) +from litellm.models.object_permission import ( # noqa: E402 + LiteLLM_ObjectPermissionTable as LiteLLM_ObjectPermissionTable, +) +from litellm.models.team import ( # noqa: E402 + LiteLLM_DeletedTeamTable as LiteLLM_DeletedTeamTable, +) +from litellm.models.team import LiteLLM_TeamTable as LiteLLM_TeamTable # noqa: E402 +from litellm.models.team import ( # noqa: E402 + LiteLLM_TeamTableCachedObj as LiteLLM_TeamTableCachedObj, +) class TeamRequest(LiteLLMPydanticObjectBase): teams: List[str] -class LiteLLM_BudgetTable(LiteLLMPydanticObjectBase): - """Represents user-controllable params for a LiteLLM_BudgetTable record. - - Budget-write paths use `model_fields.keys()` on this class as an allowlist - for user input. Keep server-managed fields (e.g. `budget_reset_at`) on - `LiteLLM_BudgetTableFull` so they aren't user-settable. - """ - - budget_id: Optional[str] = None - soft_budget: Optional[float] = None - max_budget: Optional[float] = None - max_parallel_requests: Optional[int] = None - tpm_limit: Optional[int] = None - rpm_limit: Optional[int] = None - model_max_budget: Optional[dict] = None - budget_duration: Optional[str] = None - allowed_models: Optional[List[str]] = ( - None # per-member model scope; empty = inherit team models - ) - - model_config = ConfigDict(protected_namespaces=()) - - -class LiteLLM_BudgetTableFull(LiteLLM_BudgetTable): - """LiteLLM_BudgetTable + server-managed fields returned on API responses.""" - - budget_reset_at: Optional[datetime] = None - created_at: datetime - - -class LiteLLM_TeamMemberTable(LiteLLM_BudgetTable): - """ - Used to track spend of a user_id within a team_id - """ - - spend: Optional[float] = None - user_id: Optional[str] = None - team_id: Optional[str] = None - budget_id: Optional[str] = None - - model_config = ConfigDict(protected_namespaces=()) +from litellm.models.budget import ( # noqa: E402 + LiteLLM_BudgetTable as LiteLLM_BudgetTable, +) +from litellm.models.budget import ( # noqa: E402 + LiteLLM_BudgetTableFull as LiteLLM_BudgetTableFull, +) +from litellm.models.budget import ( # noqa: E402 + LiteLLM_TeamMemberTable as LiteLLM_TeamMemberTable, +) class NewOrganizationRequest(LiteLLM_BudgetTable): @@ -2637,66 +2340,12 @@ class ConfigYAML(LiteLLMPydanticObjectBase): model_config = ConfigDict(protected_namespaces=()) -class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase): - token: Optional[str] = None - key_name: Optional[str] = None - key_alias: Optional[str] = None - spend: float = 0.0 - max_budget: Optional[float] = None - expires: Optional[Union[str, datetime]] = None - models: List = [] - aliases: Dict = {} - config: Dict = {} - user_id: Optional[str] = None - team_id: Optional[str] = None - agent_id: Optional[str] = None - project_id: Optional[str] = None - max_parallel_requests: Optional[int] = None - metadata: Dict = {} - tpm_limit: Optional[int] = None - rpm_limit: Optional[int] = None - budget_duration: Optional[str] = None - budget_reset_at: Optional[datetime] = None - allowed_cache_controls: Optional[list] = [] - allowed_routes: Optional[list] = [] - permissions: Dict = {} - model_spend: Dict = {} - model_max_budget: Dict = {} - soft_budget_cooldown: bool = False - blocked: Optional[bool] = None - litellm_budget_table: Optional[dict] = None - org_id: Optional[str] = None # org id for a given key - created_at: Optional[datetime] = None - created_by: Optional[str] = None - updated_at: Optional[datetime] = None - updated_by: Optional[str] = None - last_active: Optional[datetime] = None - object_permission_id: Optional[str] = None - object_permission: Optional[LiteLLM_ObjectPermissionTable] = None - access_group_ids: Optional[List[str]] = None - rotation_count: Optional[int] = 0 # Number of times key has been rotated - auto_rotate: Optional[bool] = False # Whether this key should be auto-rotated - rotation_interval: Optional[str] = None # How often to rotate (e.g., "30d", "90d") - last_rotation_at: Optional[datetime] = None # When this key was last rotated - key_rotation_at: Optional[datetime] = None # When this key should next be rotated - router_settings: Optional[dict] = None - budget_limits: Optional[List[dict]] = None # multiple concurrent budget windows - model_config = ConfigDict(protected_namespaces=()) - - -class LiteLLM_DeletedVerificationToken(LiteLLM_VerificationToken): - """ - Recording of deleted keys for audit purposes. Mirrors LiteLLM_VerificationToken - plus metadata captured at deletion time. - """ - - id: Optional[str] = None - deleted_at: Optional[datetime] = None - deleted_by: Optional[str] = None - deleted_by_api_key: Optional[str] = None - litellm_changed_by: Optional[str] = None - - model_config = ConfigDict(protected_namespaces=()) +from litellm.models.verification_token import ( # noqa: E402 + LiteLLM_DeletedVerificationToken as LiteLLM_DeletedVerificationToken, +) +from litellm.models.verification_token import ( # noqa: E402 + LiteLLM_VerificationToken as LiteLLM_VerificationToken, +) class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken): @@ -2935,39 +2584,10 @@ class UserInfoV2Response(LiteLLMPydanticObjectBase): teams: List[str] = [] # Just team IDs, not full team objects -class LiteLLM_Config(LiteLLMPydanticObjectBase): - param_name: str - param_value: Dict - - -class LiteLLM_OrganizationMembershipTable(LiteLLMPydanticObjectBase): - """ - This is the table that track what organizations a user belongs to and users spend within the organization - """ - - user_id: str - organization_id: str - user_role: Optional[str] = None - spend: float = 0.0 - budget_id: Optional[str] = None - created_at: datetime - updated_at: datetime - user: Optional[Any] = ( - None # You might want to replace 'Any' with a more specific type if available - ) - litellm_budget_table: Optional[LiteLLM_BudgetTable] = None - user_email: Optional[str] = None - - model_config = ConfigDict(protected_namespaces=()) - - @model_validator(mode="after") - def populate_user_email(self) -> "LiteLLM_OrganizationMembershipTable": - if self.user_email is None and self.user is not None: - if isinstance(self.user, dict): - self.user_email = self.user.get("user_email") - else: - self.user_email = getattr(self.user, "user_email", None) - return self +from litellm.models.config import LiteLLM_Config as LiteLLM_Config # noqa: E402 +from litellm.models.organization_membership import ( # noqa: E402 + LiteLLM_OrganizationMembershipTable as LiteLLM_OrganizationMembershipTable, +) class LiteLLM_OrganizationTableUpdate(LiteLLM_BudgetTable): @@ -2997,61 +2617,10 @@ class LiteLLM_OrganizationTableUpdate(LiteLLM_BudgetTable): return values -class LiteLLM_UserTable(LiteLLMPydanticObjectBase): - user_id: str - max_budget: Optional[float] = None - spend: float = 0.0 - model_max_budget: Optional[Dict] = {} - model_spend: Optional[Dict] = {} - user_email: Optional[str] = None - user_alias: Optional[str] = None - models: list = [] - tpm_limit: Optional[int] = None - rpm_limit: Optional[int] = None - user_role: Optional[str] = None - organization_memberships: Optional[List[LiteLLM_OrganizationMembershipTable]] = None - teams: List[str] = [] - sso_user_id: Optional[str] = None - budget_duration: Optional[str] = None - budget_reset_at: Optional[datetime] = None - metadata: Optional[dict] = None - created_at: Optional[datetime] = None - updated_at: Optional[datetime] = None - object_permission: Optional[LiteLLM_ObjectPermissionTable] = None - - @model_validator(mode="before") - @classmethod - def set_model_info(cls, values): - if values.get("spend") is None: - values.update({"spend": 0.0}) - if values.get("models") is None: - values.update({"models": []}) - if values.get("teams") is None: - values.update({"teams": []}) - return values - - model_config = ConfigDict(protected_namespaces=()) - - -class LiteLLM_OrganizationTable(LiteLLMPydanticObjectBase): - """Represents user-controllable params for a LiteLLM_OrganizationTable record""" - - organization_id: Optional[str] = None - organization_alias: Optional[str] = None - budget_id: str - spend: float = 0.0 - metadata: Optional[dict] = None - models: List[str] - created_by: str - updated_by: str - users: Optional[List[LiteLLM_UserTable]] = None - litellm_budget_table: Optional[LiteLLM_BudgetTable] = None - - ######################################################### - # Object Permission - MCP, Vector Stores etc. - ######################################################### - object_permission: Optional[LiteLLM_ObjectPermissionTable] = None - object_permission_id: Optional[str] = None +from litellm.models.organization import ( # noqa: E402 + LiteLLM_OrganizationTable as LiteLLM_OrganizationTable, +) +from litellm.models.user import LiteLLM_UserTable as LiteLLM_UserTable # noqa: E402 class LiteLLM_OrganizationTableWithMembers(LiteLLM_OrganizationTable): @@ -3160,28 +2729,9 @@ class DeleteProjectRequest(LiteLLMPydanticObjectBase): project_ids: List[str] -class LiteLLM_ProjectTable(LiteLLMPydanticObjectBase): - """Database model representation for project""" - - project_id: str - project_alias: Optional[str] = None - description: Optional[str] = None - team_id: Optional[str] = None - budget_id: Optional[str] = None - metadata: Optional[dict] = None - models: List[str] = [] - spend: float = 0.0 - model_spend: Optional[dict] = None - model_rpm_limit: Optional[dict] = None - model_tpm_limit: Optional[dict] = None - blocked: bool = False - object_permission_id: Optional[str] = None - created_by: str - updated_by: str - created_at: Optional[datetime] = None - updated_at: Optional[datetime] = None - litellm_budget_table: Optional[LiteLLM_BudgetTable] = None - object_permission: Optional[LiteLLM_ObjectPermissionTable] = None +from litellm.models.project import ( # noqa: E402 + LiteLLM_ProjectTable as LiteLLM_ProjectTable, +) class NewProjectResponse(LiteLLM_ProjectTable): @@ -3207,101 +2757,19 @@ class LiteLLM_UserTableWithKeyCount(LiteLLM_UserTable): key_count: int = 0 -class LiteLLM_EndUserTable(LiteLLMPydanticObjectBase): - user_id: str - blocked: bool - alias: Optional[str] = None - spend: float = 0.0 - allowed_model_region: Optional[AllowedModelRegion] = None - default_model: Optional[str] = None - litellm_budget_table: Optional[LiteLLM_BudgetTable] = None - object_permission_id: Optional[str] = None - object_permission: Optional[LiteLLM_ObjectPermissionTable] = None - - @model_validator(mode="before") - @classmethod - def set_model_info(cls, values): - if values.get("spend") is None: - values.update({"spend": 0.0}) - return values - - model_config = ConfigDict(protected_namespaces=()) - - -class LiteLLM_TagTable(LiteLLMPydanticObjectBase): - tag_name: str - description: Optional[str] = None - models: List[str] = [] - model_info: Optional[dict] = None - spend: float = 0.0 - budget_id: Optional[str] = None - litellm_budget_table: Optional[LiteLLM_BudgetTable] = None - created_at: Optional[datetime] = None - created_by: Optional[str] = None - updated_at: Optional[datetime] = None - - @model_validator(mode="before") - @classmethod - def set_model_info(cls, values): - if values.get("spend") is None: - values.update({"spend": 0.0}) - if values.get("models") is None: - values.update({"models": []}) - return values - - model_config = ConfigDict(protected_namespaces=()) - - -class LiteLLM_AccessGroupTable(LiteLLMPydanticObjectBase): - access_group_id: str - access_group_name: str - description: Optional[str] = None - access_model_names: List[str] = [] - access_mcp_server_ids: List[str] = [] - access_agent_ids: List[str] = [] - assigned_team_ids: List[str] = [] - assigned_key_ids: List[str] = [] - created_at: Optional[datetime] = None - created_by: Optional[str] = None - updated_at: Optional[datetime] = None - updated_by: Optional[str] = None - - -class LiteLLM_SpendLogs(LiteLLMPydanticObjectBase): - request_id: str - api_key: str - model: Optional[str] = "" - api_base: Optional[str] = "" - call_type: str - spend: Optional[float] = 0.0 - total_tokens: Optional[int] = 0 - prompt_tokens: Optional[int] = 0 - completion_tokens: Optional[int] = 0 - startTime: Union[str, datetime, None] - endTime: Union[str, datetime, None] - user: Optional[str] = "" - metadata: Optional[Json] = {} - cache_hit: Optional[str] = "False" - cache_key: Optional[str] = None - request_tags: Optional[Json] = None - requester_ip_address: Optional[str] = None - messages: Optional[Union[str, list, dict]] - response: Optional[Union[str, list, dict]] - - -class LiteLLM_ErrorLogs(LiteLLMPydanticObjectBase): - request_id: Optional[str] = str(uuid.uuid4()) - api_base: Optional[str] = "" - model_group: Optional[str] = "" - litellm_model_name: Optional[str] = "" - model_id: Optional[str] = "" - request_kwargs: Optional[dict] = {} - exception_type: Optional[str] = "" - status_code: Optional[str] = "" - exception_string: Optional[str] = "" - startTime: Union[str, datetime, None] - endTime: Union[str, datetime, None] - +from litellm.models.access_group import ( # noqa: E402 + LiteLLM_AccessGroupTable as LiteLLM_AccessGroupTable, +) +from litellm.models.end_user import ( # noqa: E402 + LiteLLM_EndUserTable as LiteLLM_EndUserTable, +) +from litellm.models.spend_logs import ( # noqa: E402 + LiteLLM_ErrorLogs as LiteLLM_ErrorLogs, +) +from litellm.models.spend_logs import ( # noqa: E402 + LiteLLM_SpendLogs as LiteLLM_SpendLogs, +) +from litellm.models.tag import LiteLLM_TagTable as LiteLLM_TagTable # noqa: E402 AUDIT_ACTIONS = Literal[ "created", "updated", "deleted", "blocked", "unblocked", "rotated" @@ -3982,29 +3450,9 @@ class CreatePassThroughEndpoint(LiteLLMPydanticObjectBase): headers: dict -class LiteLLM_TeamMembership(LiteLLMPydanticObjectBase): - user_id: str - team_id: str - budget_id: Optional[str] = None - spend: Optional[float] = 0.0 - total_spend: Optional[float] = 0.0 - # Union so Pydantic picks Full when data has server-managed fields - # (/team/info) and Base when callers/tests construct with only - # user-settable fields. - litellm_budget_table: Optional[ - Union[LiteLLM_BudgetTableFull, LiteLLM_BudgetTable] - ] = None - - def safe_get_team_member_rpm_limit(self) -> Optional[int]: - if self.litellm_budget_table is not None: - return self.litellm_budget_table.rpm_limit - return None - - def safe_get_team_member_tpm_limit(self) -> Optional[int]: - if self.litellm_budget_table is not None: - return self.litellm_budget_table.tpm_limit - return None - +from litellm.models.team_membership import ( # noqa: E402 + LiteLLM_TeamMembership as LiteLLM_TeamMembership, +) #### Organization / Team Member Requests #### @@ -4922,39 +4370,18 @@ class ToolDiscoveryQueueItem(TypedDict, total=False): user_agent: Optional[str] # HTTP User-Agent of the caller -class LiteLLM_ManagedFileTable(LiteLLMPydanticObjectBase): - unified_file_id: str - file_object: Optional[OpenAIFileObject] = None - model_mappings: Dict[str, str] - flat_model_file_ids: List[str] - created_by: Optional[str] = None - team_id: Optional[str] = None - updated_by: Optional[str] = None - storage_backend: Optional[str] = None - storage_url: Optional[str] = None - - -class LiteLLM_ManagedObjectTable(LiteLLMPydanticObjectBase): - unified_object_id: str - model_object_id: str - file_purpose: Literal["batch", "fine-tune", "response", "container"] - file_object: Union[LiteLLMBatch, LiteLLMFineTuningJob, ResponsesAPIResponse] - created_by: Optional[str] = None - team_id: Optional[str] = None - - -class LiteLLM_ManagedVectorStoreTable(LiteLLMPydanticObjectBase): - """Table for managing vector stores with target_model_names support.""" - - unified_resource_id: str - resource_object: Optional[Any] = None # VectorStoreCreateResponse - model_mappings: Dict[str, str] - flat_model_resource_ids: List[str] - created_by: Optional[str] = None - team_id: Optional[str] = None - updated_by: Optional[str] = None - storage_backend: Optional[str] = None - storage_url: Optional[str] = None +from litellm.models.managed_files import ( # noqa: E402 + LiteLLM_ManagedFileTable as LiteLLM_ManagedFileTable, +) +from litellm.models.managed_files import ( # noqa: E402 + LiteLLM_ManagedObjectTable as LiteLLM_ManagedObjectTable, +) +from litellm.models.managed_files import ( # noqa: E402 + LiteLLM_ManagedVectorStoresTable as LiteLLM_ManagedVectorStoresTable, +) +from litellm.models.managed_files import ( # noqa: E402 + LiteLLM_ManagedVectorStoreTable as LiteLLM_ManagedVectorStoreTable, +) class EnterpriseLicenseData(TypedDict, total=False): @@ -4965,20 +4392,6 @@ class EnterpriseLicenseData(TypedDict, total=False): max_teams: int -class LiteLLM_ManagedVectorStoresTable(LiteLLMPydanticObjectBase): - vector_store_id: str - custom_llm_provider: str - vector_store_name: Optional[str] - vector_store_description: Optional[str] - vector_store_metadata: Optional[Dict[str, Any]] - created_at: Optional[datetime] - updated_at: Optional[datetime] - litellm_credential_name: Optional[str] - litellm_params: Optional[Dict[str, Any]] - team_id: Optional[str] - user_id: Optional[str] - - class ResponseLiteLLM_ManagedVectorStore(TypedDict, total=False): vector_store: LiteLLM_ManagedVectorStoresTable diff --git a/litellm/proxy/agent_endpoints/agent_registry.py b/litellm/proxy/agent_endpoints/agent_registry.py index 13a2dd9f040..11fd01e2369 100644 --- a/litellm/proxy/agent_endpoints/agent_registry.py +++ b/litellm/proxy/agent_endpoints/agent_registry.py @@ -9,6 +9,7 @@ from litellm.proxy.management_helpers.object_permission_utils import ( handle_update_object_permission_common, ) from litellm.proxy.utils import PrismaClient +from litellm.repositories.table_repositories import AgentsRepository from litellm.types.agents import AgentConfig, AgentResponse, PatchAgentRequest @@ -174,7 +175,7 @@ class AgentRegistry: create_data[rate_field] = _val # Create agent in DB - created_agent = await prisma_client.db.litellm_agentstable.create( + created_agent = await AgentsRepository(prisma_client).table.create( data=create_data, include={"object_permission": True}, ) @@ -200,7 +201,7 @@ class AgentRegistry: Delete an agent from the database """ try: - deleted_agent = await prisma_client.db.litellm_agentstable.delete( + deleted_agent = await AgentsRepository(prisma_client).table.delete( where={"agent_id": agent_id} ) return dict(deleted_agent) @@ -229,7 +230,7 @@ class AgentRegistry: The patched agent """ try: - existing_agent = await prisma_client.db.litellm_agentstable.find_unique( + existing_agent = await AgentsRepository(prisma_client).table.find_unique( where={"agent_id": agent_id} ) if existing_agent is not None: @@ -282,7 +283,7 @@ class AgentRegistry: if object_permission_id is not None: update_data["object_permission_id"] = object_permission_id # Patch agent in DB - patched_agent = await prisma_client.db.litellm_agentstable.update( + patched_agent = await AgentsRepository(prisma_client).table.update( where={"agent_id": agent_id}, data={ **update_data, @@ -368,9 +369,9 @@ class AgentRegistry: update_data[rate_field] = _val if agent.get("object_permission") is not None: - existing_agent = await prisma_client.db.litellm_agentstable.find_unique( - where={"agent_id": agent_id} - ) + existing_agent = await AgentsRepository( + prisma_client + ).table.find_unique(where={"agent_id": agent_id}) existing_object_permission_id = ( existing_agent.object_permission_id if existing_agent is not None @@ -386,7 +387,7 @@ class AgentRegistry: update_data["object_permission_id"] = object_permission_id # Update agent in DB - updated_agent = await prisma_client.db.litellm_agentstable.update( + updated_agent = await AgentsRepository(prisma_client).table.update( where={"agent_id": agent_id}, data=update_data, include={"object_permission": True}, @@ -414,7 +415,7 @@ class AgentRegistry: Get all agents from the database """ try: - agents_from_db = await prisma_client.db.litellm_agentstable.find_many( + agents_from_db = await AgentsRepository(prisma_client).table.find_many( order={"created_at": "desc"}, include={"object_permission": True}, ) diff --git a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py index 42cf31e1e2b..2577615fc8e 100644 --- a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py +++ b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py @@ -9,11 +9,12 @@ from typing import List, Optional, Set from litellm._logging import verbose_logger from litellm.proxy._types import ( + UI_TEAM_ID, LiteLLM_ObjectPermissionTable, LiteLLM_TeamTable, - UI_TEAM_ID, UserAPIKeyAuth, ) +from litellm.repositories.table_repositories import AgentsRepository class AgentRequestHandler: @@ -298,7 +299,7 @@ class AgentRequestHandler: agent_ids: Set[str] = set() if access_groups and prisma_client is not None: try: - agents = await prisma_client.db.litellm_agentstable.find_many( + agents = await AgentsRepository(prisma_client).table.find_many( where={"agent_access_groups": {"hasSome": access_groups}} ) for agent in agents: diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index 19dbfe33d32..d19008856bd 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -17,6 +17,7 @@ from fastapi import APIRouter, Depends, HTTPException, Query, Request import litellm from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.litellm_logging import _get_masked_values from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.a2a.agent_card import merge_agent_card @@ -30,7 +31,6 @@ from litellm.types.agents import ( MakeAgentsPublicRequest, PatchAgentRequest, ) -from litellm.litellm_core_utils.litellm_logging import _get_masked_values from litellm.types.llms.custom_http import httpxSpecialProvider from litellm.types.proxy.management_endpoints.common_daily_activity import ( DailySpendMetadata, @@ -219,7 +219,7 @@ async def get_agents( if prisma_client is not None: agent_ids = [agent.agent_id for agent in returned_agents] if agent_ids: - db_agents = await prisma_client.db.litellm_agentstable.find_many( + db_agents = await AgentsRepository(prisma_client).table.find_many( where={"agent_id": {"in": agent_ids}}, ) spend_map = {a.agent_id: a.spend for a in db_agents} @@ -301,6 +301,7 @@ async def get_agents( from litellm.proxy.agent_endpoints.agent_registry import ( global_agent_registry as AGENT_REGISTRY, ) +from litellm.repositories.table_repositories import AgentsRepository @router.post( @@ -471,7 +472,7 @@ async def get_agent_by_id( try: agent = AGENT_REGISTRY.get_agent_by_id(agent_id=agent_id) if agent is None: - agent_row = await prisma_client.db.litellm_agentstable.find_unique( + agent_row = await AgentsRepository(prisma_client).table.find_unique( where={"agent_id": agent_id}, include={"object_permission": True}, ) @@ -489,7 +490,7 @@ async def get_agent_by_id( agent = AgentResponse(**agent_dict) # type: ignore else: # Agent found in memory — refresh spend from DB - db_row = await prisma_client.db.litellm_agentstable.find_unique( + db_row = await AgentsRepository(prisma_client).table.find_unique( where={"agent_id": agent_id} ) if db_row is not None: @@ -570,7 +571,7 @@ async def update_agent( try: # Check if agent exists - existing_agent = await prisma_client.db.litellm_agentstable.find_unique( + existing_agent = await AgentsRepository(prisma_client).table.find_unique( where={"agent_id": agent_id} ) if existing_agent is not None: @@ -678,7 +679,7 @@ async def patch_agent( try: # Check if agent exists - existing_agent = await prisma_client.db.litellm_agentstable.find_unique( + existing_agent = await AgentsRepository(prisma_client).table.find_unique( where={"agent_id": agent_id} ) if existing_agent is not None: @@ -769,7 +770,7 @@ async def delete_agent( try: # Check if agent exists - existing_agent = await prisma_client.db.litellm_agentstable.find_unique( + existing_agent = await AgentsRepository(prisma_client).table.find_unique( where={"agent_id": agent_id} ) if existing_agent is not None: @@ -859,7 +860,7 @@ async def make_agent_public( agent = AGENT_REGISTRY.get_agent_by_id(agent_id=agent_id) if agent is None: # check if agent exists in DB - agent = await prisma_client.db.litellm_agentstable.find_unique( + agent = await AgentsRepository(prisma_client).table.find_unique( where={"agent_id": agent_id} ) if agent is not None: @@ -982,7 +983,7 @@ async def make_agents_public( agent = AGENT_REGISTRY.get_agent_by_id(agent_id=agent_id) if agent is None: # check if agent exists in DB - agent = await prisma_client.db.litellm_agentstable.find_unique( + agent = await AgentsRepository(prisma_client).table.find_unique( where={"agent_id": agent_id} ) if agent is not None: @@ -1082,7 +1083,7 @@ async def get_agent_daily_activity( if user_api_key_dict.user_id is None: permitted_agent_ids = [] else: - owned_records = await prisma_client.db.litellm_agentstable.find_many( + owned_records = await AgentsRepository(prisma_client).table.find_many( where={"created_by": user_api_key_dict.user_id} ) permitted_agent_ids = [a.agent_id for a in owned_records] @@ -1118,7 +1119,7 @@ async def get_agent_daily_activity( if agent_ids_list: where_condition["agent_id"] = {"in": list(agent_ids_list)} - agent_records = await prisma_client.db.litellm_agentstable.find_many( + agent_records = await AgentsRepository(prisma_client).table.find_many( where=where_condition ) agent_metadata = { diff --git a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py index 20b1659fa1a..dd7350e13ce 100644 --- a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py +++ b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py @@ -26,6 +26,7 @@ from fastapi.responses import JSONResponse from litellm._logging import verbose_proxy_logger from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.repositories.table_repositories import ClaudeCodePluginRepository from litellm.types.proxy.claude_code_endpoints import ( ListPluginsResponse, PluginListItem, @@ -71,7 +72,7 @@ async def get_marketplace(): try: prisma_client = await _get_prisma_client() - plugins = await prisma_client.db.litellm_claudecodeplugintable.find_many( + plugins = await ClaudeCodePluginRepository(prisma_client).table.find_many( where={"enabled": True} ) @@ -268,12 +269,12 @@ async def register_plugin( manifest["namespace"] = request.namespace # Check if plugin exists - existing = await prisma_client.db.litellm_claudecodeplugintable.find_unique( + existing = await ClaudeCodePluginRepository(prisma_client).table.find_unique( where={"name": request.name} ) if existing: - plugin = await prisma_client.db.litellm_claudecodeplugintable.update( + plugin = await ClaudeCodePluginRepository(prisma_client).table.update( where={"name": request.name}, data={ "version": request.version, @@ -285,7 +286,7 @@ async def register_plugin( ) action = "updated" else: - plugin = await prisma_client.db.litellm_claudecodeplugintable.create( + plugin = await ClaudeCodePluginRepository(prisma_client).table.create( data={ "name": request.name, "version": request.version, @@ -348,7 +349,7 @@ async def list_plugins( prisma_client = await _get_prisma_client() where = {"enabled": True} if enabled_only else {} - plugins = await prisma_client.db.litellm_claudecodeplugintable.find_many( + plugins = await ClaudeCodePluginRepository(prisma_client).table.find_many( where=where ) @@ -415,7 +416,7 @@ async def get_plugin( try: prisma_client = await _get_prisma_client() - plugin = await prisma_client.db.litellm_claudecodeplugintable.find_unique( + plugin = await ClaudeCodePluginRepository(prisma_client).table.find_unique( where={"name": plugin_name} ) @@ -471,7 +472,7 @@ async def enable_plugin( try: prisma_client = await _get_prisma_client() - plugin = await prisma_client.db.litellm_claudecodeplugintable.find_unique( + plugin = await ClaudeCodePluginRepository(prisma_client).table.find_unique( where={"name": plugin_name} ) if not plugin: @@ -480,7 +481,7 @@ async def enable_plugin( detail={"error": f"Plugin '{plugin_name}' not found"}, ) - await prisma_client.db.litellm_claudecodeplugintable.update( + await ClaudeCodePluginRepository(prisma_client).table.update( where={"name": plugin_name}, data={"enabled": True, "updated_at": datetime.now(timezone.utc)}, ) @@ -516,7 +517,7 @@ async def disable_plugin( try: prisma_client = await _get_prisma_client() - plugin = await prisma_client.db.litellm_claudecodeplugintable.find_unique( + plugin = await ClaudeCodePluginRepository(prisma_client).table.find_unique( where={"name": plugin_name} ) if not plugin: @@ -525,7 +526,7 @@ async def disable_plugin( detail={"error": f"Plugin '{plugin_name}' not found"}, ) - await prisma_client.db.litellm_claudecodeplugintable.update( + await ClaudeCodePluginRepository(prisma_client).table.update( where={"name": plugin_name}, data={"enabled": False, "updated_at": datetime.now(timezone.utc)}, ) @@ -561,7 +562,7 @@ async def delete_plugin( try: prisma_client = await _get_prisma_client() - plugin = await prisma_client.db.litellm_claudecodeplugintable.find_unique( + plugin = await ClaudeCodePluginRepository(prisma_client).table.find_unique( where={"name": plugin_name} ) if not plugin: @@ -570,7 +571,7 @@ async def delete_plugin( detail={"error": f"Plugin '{plugin_name}' not found"}, ) - await prisma_client.db.litellm_claudecodeplugintable.delete( + await ClaudeCodePluginRepository(prisma_client).table.delete( where={"name": plugin_name} ) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 94ae3f5eacc..45007861d55 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -61,19 +61,33 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec from litellm.proxy.common_utils.http_parsing_utils import ( _safe_get_request_headers, _safe_get_request_query_params, ) +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.guardrails.tool_name_extraction import ( TOOL_CAPABLE_CALL_TYPES, extract_request_tool_names, ) -from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec -from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.route_llm_request import route_request from litellm.proxy.utils import PrismaClient, ProxyLogging, log_db_metrics +from litellm.repositories.budget_repository import BudgetRepository +from litellm.repositories.object_permission_repository import ObjectPermissionRepository +from litellm.repositories.organization_repository import OrganizationRepository +from litellm.repositories.project_repository import ProjectRepository +from litellm.repositories.table_repositories import ( + AccessGroupRepository, + EndUserRepository, + JWTKeyMappingRepository, + ManagedVectorStoresRepository, + TagRepository, + TeamMembershipRepository, +) +from litellm.repositories.team_repository import TeamRepository +from litellm.repositories.user_repository import UserRepository from litellm.router import Router from litellm.utils import get_utc_datetime @@ -957,7 +971,7 @@ async def get_default_end_user_budget( # Fetch from database try: - budget_record = await prisma_client.db.litellm_budgettable.find_unique( + budget_record = await BudgetRepository(prisma_client).table.find_unique( where={"budget_id": litellm.max_end_user_budget_id} ) @@ -1016,7 +1030,7 @@ async def get_team_member_default_budget( return LiteLLM_BudgetTable(**cached_budget) try: - budget_record = await prisma_client.db.litellm_budgettable.find_unique( + budget_record = await BudgetRepository(prisma_client).table.find_unique( where={"budget_id": budget_id} ) @@ -1175,7 +1189,7 @@ async def get_end_user_object( # Fetch from database try: - response = await prisma_client.db.litellm_endusertable.find_unique( + response = await EndUserRepository(prisma_client).table.find_unique( where={"user_id": end_user_id}, include={"litellm_budget_table": True, "object_permission": True}, ) @@ -1375,7 +1389,7 @@ async def get_tag_objects_batch( # Batch fetch uncached tags from DB in one query if uncached_tags: try: - db_tags = await prisma_client.db.litellm_tagtable.find_many( + db_tags = await TagRepository(prisma_client).table.find_many( where={"tag_name": {"in": uncached_tags}}, include={"litellm_budget_table": True}, ) @@ -1469,7 +1483,7 @@ async def get_team_membership( # else, check db try: - response = await prisma_client.db.litellm_teammembership.find_unique( + response = await TeamMembershipRepository(prisma_client).table.find_unique( where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}}, include={"litellm_budget_table": True}, ) @@ -1622,7 +1636,7 @@ async def _get_fuzzy_user_object( response = None if sso_user_id is not None: - response = await prisma_client.db.litellm_usertable.find_unique( + response = await UserRepository(prisma_client).table.find_unique( where={"sso_user_id": sso_user_id}, include={"organization_memberships": True}, ) @@ -1630,14 +1644,14 @@ async def _get_fuzzy_user_object( if response is None and user_email is not None: # Use case-insensitive query to handle emails with different casing # This matches the pattern used in _check_duplicate_user_email - response = await prisma_client.db.litellm_usertable.find_first( + response = await UserRepository(prisma_client).table.find_first( where={"user_email": {"equals": user_email, "mode": "insensitive"}}, include={"organization_memberships": True}, ) if response is not None and sso_user_id is not None: # update sso_user_id asyncio.create_task( # background task to update user with sso id - prisma_client.db.litellm_usertable.update( + UserRepository(prisma_client).table.update( where={"user_id": response.user_id}, data={"sso_user_id": sso_user_id}, ) @@ -1687,7 +1701,7 @@ async def get_user_object( ) if should_check_db: - response = await prisma_client.db.litellm_usertable.find_unique( + response = await UserRepository(prisma_client).table.find_unique( where={"user_id": user_id}, include={"organization_memberships": True} ) @@ -1711,7 +1725,7 @@ async def get_user_object( if litellm.default_internal_user_params is not None: new_user_params.update(litellm.default_internal_user_params) - response = await prisma_client.db.litellm_usertable.create( + response = await UserRepository(prisma_client).table.create( data=new_user_params, include={"organization_memberships": True}, ) @@ -1860,7 +1874,7 @@ async def _delete_cache_key_object( async def _get_team_db_check( team_id: str, prisma_client: PrismaClient, team_id_upsert: Optional[bool] = None ): - response = await prisma_client.db.litellm_teamtable.find_unique( + response = await TeamRepository(prisma_client).table.find_unique( where={"team_id": team_id} ) @@ -1882,7 +1896,7 @@ async def _get_team_db_check( async def _get_team_object_from_db(team_id: str, prisma_client: PrismaClient): - return await prisma_client.db.litellm_teamtable.find_unique( + return await TeamRepository(prisma_client).table.find_unique( where={"team_id": team_id} ) @@ -2111,7 +2125,7 @@ async def get_access_object( # Not in cache - fetch from DB try: - response = await prisma_client.db.litellm_accessgrouptable.find_unique( + response = await AccessGroupRepository(prisma_client).table.find_unique( where={"access_group_id": access_group_id} ) @@ -2193,7 +2207,7 @@ async def get_team_object_by_alias( # Query database by team_alias try: - teams = await prisma_client.db.litellm_teamtable.find_many( + teams = await TeamRepository(prisma_client).table.find_many( where={"team_alias": team_alias} ) @@ -2301,7 +2315,7 @@ async def get_org_object_by_alias( # Query database by organization_alias try: - orgs = await prisma_client.db.litellm_organizationtable.find_many( + orgs = await OrganizationRepository(prisma_client).table.find_many( where={"organization_alias": org_alias} ) @@ -2526,7 +2540,7 @@ async def get_jwt_key_mapping_object( Returns the hashed token (str) if a matching active mapping is found, else None. """ - mapping = await prisma_client.db.litellm_jwtkeymapping.find_first( + mapping = await JWTKeyMappingRepository(prisma_client).table.find_first( where={ "jwt_claim_name": jwt_claim_name, "jwt_claim_value": jwt_claim_value, @@ -2659,7 +2673,7 @@ async def get_object_permission( # else, check db try: - response = await prisma_client.db.litellm_objectpermissiontable.find_unique( + response = await ObjectPermissionRepository(prisma_client).table.find_unique( where={"object_permission_id": object_permission_id} ) @@ -2715,7 +2729,7 @@ async def get_managed_vector_store_rows_by_uuids( if not cache_misses: return result - rows = await prisma_client.db.litellm_managedvectorstorestable.find_many( + rows = await ManagedVectorStoresRepository(prisma_client).table.find_many( where={"vector_store_id": {"in": cache_misses}}, take=len(cache_misses), ) @@ -2790,7 +2804,7 @@ async def get_org_object( if include_budget_table: query_kwargs["include"] = {"litellm_budget_table": True} - response = await prisma_client.db.litellm_organizationtable.find_unique( + response = await OrganizationRepository(prisma_client).table.find_unique( **query_kwargs ) @@ -4180,7 +4194,7 @@ async def get_project_object( return deserialized_project # Fetch from DB - project_row = await prisma_client.db.litellm_projecttable.find_unique( + project_row = await ProjectRepository(prisma_client).table.find_unique( where={"project_id": project_id}, include={"litellm_budget_table": True}, ) @@ -4480,10 +4494,10 @@ async def vector_store_access_check( ######################################################### # Check if the key can access the vector store if valid_token is not None and valid_token.object_permission_id is not None: - key_object_permission = ( - await prisma_client.db.litellm_objectpermissiontable.find_unique( - where={"object_permission_id": valid_token.object_permission_id}, - ) + key_object_permission = await ObjectPermissionRepository( + prisma_client + ).table.find_unique( + where={"object_permission_id": valid_token.object_permission_id}, ) if key_object_permission is not None: _can_object_call_vector_stores( @@ -4494,10 +4508,10 @@ async def vector_store_access_check( # Check if the team can access the vector store if team_object is not None and team_object.object_permission_id is not None: - team_object_permission = ( - await prisma_client.db.litellm_objectpermissiontable.find_unique( - where={"object_permission_id": team_object.object_permission_id}, - ) + team_object_permission = await ObjectPermissionRepository( + prisma_client + ).table.find_unique( + where={"object_permission_id": team_object.object_permission_id}, ) if team_object_permission is not None: _can_object_call_vector_stores( diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 6d3d49b71ec..536d2867855 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -14,11 +14,11 @@ import os import re from typing import Any, List, Literal, Optional, Set, Tuple, Union, cast +import jwt from cryptography import x509 from cryptography.hazmat.backends import default_backend from cryptography.hazmat.primitives import serialization from fastapi import HTTPException, status -import jwt from jwt.api_jwk import PyJWK from litellm._logging import verbose_proxy_logger @@ -50,6 +50,7 @@ from litellm.proxy.auth.auth_checks import can_team_access_model from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.repositories.user_repository import UserRepository from .auth_checks import ( _allowed_routes_check, @@ -1790,7 +1791,7 @@ class JWTAuthManager: # Update user role new_role = jwt_handler.map_jwt_role_to_litellm_role(jwt_valid_token) if new_role and user_object.user_role != new_role.value: - await prisma_client.db.litellm_usertable.update( + await UserRepository(prisma_client).table.update( where={"user_id": user_object.user_id}, data={"user_role": new_role.value}, ) diff --git a/litellm/proxy/auth/login_utils.py b/litellm/proxy/auth/login_utils.py index 34085d5685a..d0818b95363 100644 --- a/litellm/proxy/auth/login_utils.py +++ b/litellm/proxy/auth/login_utils.py @@ -34,6 +34,7 @@ from litellm.proxy.utils import ( hash_password, verify_password, ) +from litellm.repositories.user_repository import UserRepository from litellm.secret_managers.main import get_secret_bool from litellm.types.proxy.ui_sso import ReturnedUITokenObject @@ -45,7 +46,7 @@ async def _rehash_password_if_needed(user_id: str, password: str, stored: str) - from litellm.proxy.proxy_server import prisma_client if prisma_client is not None: - await prisma_client.db.litellm_usertable.update( + await UserRepository(prisma_client).table.update( where={"user_id": user_id}, data={"password": hash_password(password)}, ) @@ -151,7 +152,7 @@ async def authenticate_user( # noqa: PLR0915 if prisma_client is not None: _user_row = cast( Optional[LiteLLM_UserTable], - await prisma_client.db.litellm_usertable.find_first( + await UserRepository(prisma_client).table.find_first( where={"user_email": {"equals": username, "mode": "insensitive"}} ), ) diff --git a/litellm/proxy/auth/model_checks.py b/litellm/proxy/auth/model_checks.py index d9600e1a4b4..00f276dc970 100644 --- a/litellm/proxy/auth/model_checks.py +++ b/litellm/proxy/auth/model_checks.py @@ -6,6 +6,7 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.proxy._types import SpecialModelNames, UserAPIKeyAuth +from litellm.repositories.object_permission_repository import ObjectPermissionRepository from litellm.router import Router from litellm.router_utils.fallback_event_handlers import get_fallback_model_group from litellm.types.router import CredentialLiteLLMParams, LiteLLM_Params @@ -86,7 +87,7 @@ async def get_mcp_server_ids( # Make a direct SQL query to get just the mcp_servers try: - result = await prisma_client.db.litellm_objectpermissiontable.find_unique( + result = await ObjectPermissionRepository(prisma_client).table.find_unique( where={"object_permission_id": user_api_key_dict.object_permission_id}, ) if result and result.mcp_servers: diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index a5501fefa4e..a970e0ddee8 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -21,9 +21,9 @@ from fastapi.security.api_key import APIKeyHeader import litellm from litellm._logging import verbose_logger, verbose_proxy_logger from litellm._service_logger import ServiceLogging +from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS from litellm.integrations.otel.model.config import is_otel_v2_enabled from litellm.integrations.otel.runtime import phase_span, seed_request_identity -from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value from litellm.proxy._types import * @@ -65,7 +65,6 @@ from litellm.proxy.auth.oauth2_check import Oauth2Handler from litellm.proxy.auth.oauth2_proxy_hook import handle_oauth2_proxy_request from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.cache_coordinator import EventDrivenCacheCoordinator -from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.common_utils.http_parsing_utils import ( _read_request_body, _safe_get_request_headers, @@ -73,12 +72,14 @@ from litellm.proxy.common_utils.http_parsing_utils import ( populate_request_with_path_params, ) from litellm.proxy.common_utils.realtime_utils import _realtime_request_body +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.proxy.utils import ( PrismaClient, ProxyLogging, normalize_route_for_root_path, ) +from litellm.repositories.table_repositories import TeamMembershipRepository from litellm.secret_managers.main import get_secret_bool from litellm.types.services import ServiceTypes @@ -1797,7 +1798,9 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 _team_id = valid_token.team_id if _user_id is not None and _team_id is not None: - _db_member = await prisma_client.db.litellm_teammembership.find_first( + _db_member = await TeamMembershipRepository( + prisma_client + ).table.find_first( where={ "user_id": _user_id, "team_id": _team_id, diff --git a/litellm/proxy/common_utils/expired_ui_session_key_cleanup_manager.py b/litellm/proxy/common_utils/expired_ui_session_key_cleanup_manager.py index 67a24567461..4f3e26ab5fb 100644 --- a/litellm/proxy/common_utils/expired_ui_session_key_cleanup_manager.py +++ b/litellm/proxy/common_utils/expired_ui_session_key_cleanup_manager.py @@ -8,7 +8,6 @@ from datetime import datetime, timezone from typing import Any, Dict, List, Optional from litellm._logging import verbose_proxy_logger -from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.constants import ( EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME, LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE, @@ -16,11 +15,15 @@ from litellm.constants import ( UI_SESSION_TOKEN_TEAM_ID, ) from litellm.proxy._types import KeyRequest, LiteLLM_VerificationToken, UserAPIKeyAuth +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks from litellm.proxy.management_endpoints.key_management_endpoints import ( delete_verification_tokens, ) from litellm.proxy.utils import PrismaClient +from litellm.repositories.verification_token_repository import ( + VerificationTokenRepository, +) class ExpiredUISessionKeyCleanupManager: @@ -147,7 +150,7 @@ class ExpiredUISessionKeyCleanupManager: Find expired LiteLLM dashboard session keys. """ now = datetime.now(timezone.utc) - return await self.prisma_client.db.litellm_verificationtoken.find_many( + return await VerificationTokenRepository(self.prisma_client).table.find_many( where={ "team_id": UI_SESSION_TOKEN_TEAM_ID, "expires": {"lt": now}, diff --git a/litellm/proxy/common_utils/key_rotation_manager.py b/litellm/proxy/common_utils/key_rotation_manager.py index aaf39a7a19d..d622f612494 100644 --- a/litellm/proxy/common_utils/key_rotation_manager.py +++ b/litellm/proxy/common_utils/key_rotation_manager.py @@ -24,6 +24,12 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( regenerate_key_fn, ) from litellm.proxy.utils import PrismaClient +from litellm.repositories.table_repositories import ( + DeprecatedVerificationTokenRepository, +) +from litellm.repositories.verification_token_repository import ( + VerificationTokenRepository, +) class KeyRotationManager: @@ -124,20 +130,20 @@ class KeyRotationManager: """ now = datetime.now(timezone.utc) - keys_with_rotation = ( - await self.prisma_client.db.litellm_verificationtoken.find_many( - where={ - "auto_rotate": True, # Only keys marked for auto rotation - "OR": [ - { - "key_rotation_at": None - }, # Keys that need initial rotation time setup - { - "key_rotation_at": {"lte": now} - }, # Keys where rotation time has passed - ], - } - ) + keys_with_rotation = await VerificationTokenRepository( + self.prisma_client + ).table.find_many( + where={ + "auto_rotate": True, # Only keys marked for auto rotation + "OR": [ + { + "key_rotation_at": None + }, # Keys that need initial rotation time setup + { + "key_rotation_at": {"lte": now} + }, # Keys where rotation time has passed + ], + } ) return keys_with_rotation @@ -148,9 +154,9 @@ class KeyRotationManager: """ try: now = datetime.now(timezone.utc) - result = await self.prisma_client.db.litellm_deprecatedverificationtoken.delete_many( - where={"revoke_at": {"lt": now}} - ) + result = await DeprecatedVerificationTokenRepository( + self.prisma_client + ).table.delete_many(where={"revoke_at": {"lt": now}}) if result > 0: verbose_proxy_logger.debug( "Cleaned up %s expired deprecated key(s)", result @@ -206,7 +212,7 @@ class KeyRotationManager: # Calculate next rotation time using helper function now = datetime.now(timezone.utc) next_rotation_time = _calculate_key_rotation_time(key.rotation_interval) - await self.prisma_client.db.litellm_verificationtoken.update( + await VerificationTokenRepository(self.prisma_client).table.update( where={"token": response.token_id}, data={ "rotation_count": (key.rotation_count or 0) + 1, diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 40c8caa49e5..7c1dfe8dc90 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -14,6 +14,16 @@ from litellm.proxy._types import ( LiteLLM_VerificationToken, ) from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.repositories.organization_repository import OrganizationRepository +from litellm.repositories.table_repositories import ( + EndUserRepository, + TagRepository, + TeamMembershipRepository, +) +from litellm.repositories.team_repository import TeamRepository +from litellm.repositories.verification_token_repository import ( + VerificationTokenRepository, +) from litellm.types.services import ServiceTypes @@ -159,7 +169,7 @@ class ResetBudgetJob: """ return await self._cascade_reset_spend_for_budget_link( budgets_to_reset=budgets_to_reset, - table=self.prisma_client.db.litellm_teammembership, + table=TeamMembershipRepository(self.prisma_client).table, counter_key_fn=lambda m: f"spend:team_member:{m.user_id}:{m.team_id}", log_subject="team memberships", cache_key_fn=lambda m: f"{m.team_id}_{m.user_id}", @@ -176,7 +186,7 @@ class ResetBudgetJob: """ return await self._cascade_reset_spend_for_budget_link( budgets_to_reset=budgets_to_reset, - table=self.prisma_client.db.litellm_verificationtoken, + table=VerificationTokenRepository(self.prisma_client).table, counter_key_fn=lambda k: f"spend:key:{k.token}", log_subject="keys", extra_where={"budget_duration": None, "spend": {"gt": 0}}, @@ -191,7 +201,7 @@ class ResetBudgetJob: """ return await self._cascade_reset_spend_for_budget_link( budgets_to_reset=budgets_to_reset, - table=self.prisma_client.db.litellm_organizationtable, + table=OrganizationRepository(self.prisma_client).table, counter_key_fn=lambda o: f"spend:org:{o.organization_id}", log_subject="orgs", extra_where={"spend": {"gt": 0}}, @@ -217,7 +227,7 @@ class ResetBudgetJob: """ return await self._cascade_reset_spend_for_budget_link( budgets_to_reset=budgets_to_reset, - table=self.prisma_client.db.litellm_tagtable, + table=TagRepository(self.prisma_client).table, counter_key_fn=lambda t: f"spend:tag:{t.tag_name}", log_subject="tags", extra_where={"spend": {"gt": 0}}, @@ -406,7 +416,7 @@ class ResetBudgetJob: rely on the default budget (litellm.max_end_user_budget_id) applied in-memory during auth checks. """ - rows = await self.prisma_client.db.litellm_endusertable.find_many( + rows = await EndUserRepository(self.prisma_client).table.find_many( where={ "budget_id": None, "spend": {"gt": 0}, @@ -824,7 +834,7 @@ class ResetBudgetJob: ): changed = True if changed: - await self.prisma_client.db.litellm_verificationtoken.update( + await VerificationTokenRepository(self.prisma_client).table.update( where={"token": row["token"]}, data={"budget_limits": json.dumps(windows)}, # type: ignore[arg-type] ) @@ -852,7 +862,7 @@ class ResetBudgetJob: ): changed = True if changed: - await self.prisma_client.db.litellm_teamtable.update( + await TeamRepository(self.prisma_client).table.update( where={"team_id": row["team_id"]}, data={"budget_limits": json.dumps(windows)}, # type: ignore[arg-type] ) diff --git a/litellm/proxy/container_endpoints/ownership.py b/litellm/proxy/container_endpoints/ownership.py index e0015e112e1..8118d53b9f6 100644 --- a/litellm/proxy/container_endpoints/ownership.py +++ b/litellm/proxy/container_endpoints/ownership.py @@ -12,6 +12,7 @@ from litellm.proxy.common_utils.resource_ownership import ( is_proxy_admin, user_can_access_resource_owner, ) +from litellm.repositories.table_repositories import ManagedObjectRepository from litellm.responses.utils import ResponsesAPIRequestUtils CONTAINER_OBJECT_PURPOSE = "container" @@ -213,7 +214,7 @@ async def record_container_owner( ) return response - table = prisma_client.db.litellm_managedobjecttable + table = ManagedObjectRepository(prisma_client).table existing = await table.find_unique(where={"model_object_id": model_object_id}) if existing is not None: if getattr(existing, "file_purpose", None) != CONTAINER_OBJECT_PURPOSE: @@ -273,7 +274,7 @@ async def _get_container_owner( if prisma_client is None: return None - row = await prisma_client.db.litellm_managedobjecttable.find_first( + row = await ManagedObjectRepository(prisma_client).table.find_first( where={ "model_object_id": model_object_id, "file_purpose": CONTAINER_OBJECT_PURPOSE, @@ -319,7 +320,7 @@ async def _get_stored_container_id( if prisma_client is None: return None - row = await prisma_client.db.litellm_managedobjecttable.find_first( + row = await ManagedObjectRepository(prisma_client).table.find_first( where={ "model_object_id": model_object_id, "file_purpose": CONTAINER_OBJECT_PURPOSE, @@ -411,7 +412,7 @@ async def _get_allowed_container_ids( if prisma_client is None: return set() - rows = await prisma_client.db.litellm_managedobjecttable.find_many( + rows = await ManagedObjectRepository(prisma_client).table.find_many( where={ "file_purpose": CONTAINER_OBJECT_PURPOSE, "created_by": {"in": owner_scopes}, diff --git a/litellm/proxy/credential_endpoints/endpoints.py b/litellm/proxy/credential_endpoints/endpoints.py index 2d05270e2ed..a716857111b 100644 --- a/litellm/proxy/credential_endpoints/endpoints.py +++ b/litellm/proxy/credential_endpoints/endpoints.py @@ -14,6 +14,7 @@ from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper from litellm.proxy.utils import handle_exception_on_proxy, jsonify_object +from litellm.repositories.credentials_repository import CredentialsRepository from litellm.types.utils import CreateCredentialItem, CredentialItem router = APIRouter() @@ -96,7 +97,7 @@ async def create_credential( ) credentials_dict = encrypted_credential.model_dump() credentials_dict_jsonified = jsonify_object(credentials_dict) - await prisma_client.db.litellm_credentialstable.create( + await CredentialsRepository(prisma_client).create( data={ **credentials_dict_jsonified, "created_by": user_api_key_dict.user_id, @@ -245,9 +246,7 @@ async def delete_credential( status_code=500, detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - await prisma_client.db.litellm_credentialstable.delete( - where={"credential_name": credential_name} - ) + await CredentialsRepository(prisma_client).delete_by_name(credential_name) ## DELETE FROM LITELLM ## litellm.credential_list = [ @@ -326,15 +325,14 @@ async def update_credential( status_code=500, detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - db_credential = await prisma_client.db.litellm_credentialstable.find_unique( - where={"credential_name": credential_name}, - ) + credentials_repository = CredentialsRepository(prisma_client) + db_credential = await credentials_repository.find_by_name(credential_name) if db_credential is None: raise HTTPException(status_code=404, detail="Credential not found in DB.") merged_credential = update_db_credential(db_credential, credential) credential_object_jsonified = jsonify_object(merged_credential.model_dump()) - await prisma_client.db.litellm_credentialstable.update( - where={"credential_name": credential_name}, + await credentials_repository.update_by_name( + credential_name, data={ **credential_object_jsonified, "updated_by": user_api_key_dict.user_id, diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index e7c5fa3f72c..2226aeb4b0a 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -20,6 +20,16 @@ from typing import TYPE_CHECKING, ClassVar, Optional from litellm._logging import verbose_proxy_logger from litellm.constants import SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE from litellm.litellm_core_utils.duration_parser import duration_in_seconds +from litellm.repositories.organization_repository import OrganizationRepository +from litellm.repositories.table_repositories import ( + SpendLogsRepository, + TeamMembershipRepository, +) +from litellm.repositories.team_repository import TeamRepository +from litellm.repositories.user_repository import UserRepository +from litellm.repositories.verification_token_repository import ( + VerificationTokenRepository, +) if TYPE_CHECKING: from litellm.caching.dual_cache import DualCache @@ -83,25 +93,25 @@ class SpendCounterReseed: try: if counter_key.startswith("spend:key:"): token = counter_key[len("spend:key:") :] - row = await prisma_client.db.litellm_verificationtoken.find_unique( - where={"token": token} - ) + row = await VerificationTokenRepository( + prisma_client + ).table.find_unique(where={"token": token}) elif counter_key.startswith("spend:team_member:"): suffix = counter_key[len("spend:team_member:") :] if ":" not in suffix: return None user_id, team_id = suffix.rsplit(":", 1) - row = await prisma_client.db.litellm_teammembership.find_unique( + row = await TeamMembershipRepository(prisma_client).table.find_unique( where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}} ) elif counter_key.startswith("spend:team:"): team_id = counter_key[len("spend:team:") :] - row = await prisma_client.db.litellm_teamtable.find_unique( + row = await TeamRepository(prisma_client).table.find_unique( where={"team_id": team_id} ) elif counter_key.startswith("spend:user:"): user_id = counter_key[len("spend:user:") :] - row = await prisma_client.db.litellm_usertable.find_unique( + row = await UserRepository(prisma_client).table.find_unique( where={"user_id": user_id} ) elif counter_key.startswith("spend:end_user:"): @@ -110,7 +120,7 @@ class SpendCounterReseed: return None elif counter_key.startswith("spend:org:"): org_id = counter_key[len("spend:org:") :] - row = await prisma_client.db.litellm_organizationtable.find_unique( + row = await OrganizationRepository(prisma_client).table.find_unique( where={"organization_id": org_id} ) else: @@ -243,7 +253,7 @@ class SpendCounterReseed: return None try: - response = await prisma_client.db.litellm_spendlogs.group_by( + response = await SpendLogsRepository(prisma_client).table.group_by( by=[group_field], where=where, # type: ignore[arg-type] sum={"spend": True}, diff --git a/litellm/proxy/db/spend_log_tool_index.py b/litellm/proxy/db/spend_log_tool_index.py index 835d76e0ee4..77c06a465f4 100644 --- a/litellm/proxy/db/spend_log_tool_index.py +++ b/litellm/proxy/db/spend_log_tool_index.py @@ -10,6 +10,7 @@ from typing import Any, Dict, List, Set from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.proxy.utils import PrismaClient +from litellm.repositories.table_repositories import SpendLogToolIndexRepository def _add_tool_calls_to_set(tool_calls: Any, out: Set[str]) -> None: @@ -141,7 +142,7 @@ async def process_spend_logs_tool_usage( } ) if index_data: - await prisma_client.db.litellm_spendlogtoolindex.create_many( + await SpendLogToolIndexRepository(prisma_client).table.create_many( data=index_data, skip_duplicates=True, ) diff --git a/litellm/proxy/db/tool_registry_writer.py b/litellm/proxy/db/tool_registry_writer.py index 6b34c974cf4..bbcc7396d67 100644 --- a/litellm/proxy/db/tool_registry_writer.py +++ b/litellm/proxy/db/tool_registry_writer.py @@ -11,6 +11,8 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ToolDiscoveryQueueItem +from litellm.repositories.object_permission_repository import ObjectPermissionRepository +from litellm.repositories.table_repositories import ToolRepository from litellm.types.tool_management import ( LiteLLM_ToolTableRow, ToolPolicyOverrideRow, @@ -84,7 +86,7 @@ async def batch_upsert_tools( if not data: return now = datetime.now(timezone.utc) - table = prisma_client.db.litellm_tooltable + table = ToolRepository(prisma_client).table for item in data: tool_name = item.get("tool_name", "") origin = item.get("origin") or "user_defined" @@ -134,7 +136,7 @@ async def list_tools( """Return all tools, optionally filtered by input_policy.""" try: where = {"input_policy": input_policy} if input_policy is not None else {} - rows = await prisma_client.db.litellm_tooltable.find_many( + rows = await ToolRepository(prisma_client).table.find_many( where=where, order={"created_at": "desc"}, ) @@ -150,7 +152,7 @@ async def get_tool( ) -> Optional[LiteLLM_ToolTableRow]: """Return a single tool row by tool_name.""" try: - row = await prisma_client.db.litellm_tooltable.find_unique( + row = await ToolRepository(prisma_client).table.find_unique( where={"tool_name": tool_name}, ) if row is None: @@ -192,7 +194,7 @@ async def update_tool_policy( if output_policy is not None: update_data["output_policy"] = output_policy - await prisma_client.db.litellm_tooltable.upsert( + await ToolRepository(prisma_client).table.upsert( where={"tool_name": tool_name}, data={ "create": create_data, @@ -217,7 +219,7 @@ async def get_tools_by_names( if not tool_names: return {} try: - rows = await prisma_client.db.litellm_tooltable.find_many( + rows = await ToolRepository(prisma_client).table.find_many( where={"tool_name": {"in": tool_names}}, ) return { @@ -244,7 +246,7 @@ async def list_overrides_for_tool( """ out: List[ToolPolicyOverrideRow] = [] try: - perms = await prisma_client.db.litellm_objectpermissiontable.find_many( + perms = await ObjectPermissionRepository(prisma_client).table.find_many( where={"blocked_tools": {"has": tool_name}}, include={ "verification_tokens": True, @@ -307,7 +309,7 @@ class ToolPolicyRegistry: async def sync_tool_policy_from_db(self, prisma_client: "PrismaClient") -> None: """Load all tool policies and object-permission blocked_tools from DB.""" try: - tools = await prisma_client.db.litellm_tooltable.find_many() + tools = await ToolRepository(prisma_client).table.find_many() self._tool_input_policies = { row.tool_name: getattr(row, "input_policy", "untrusted") or "untrusted" for row in tools @@ -317,7 +319,7 @@ class ToolPolicyRegistry: for row in tools } - perms = await prisma_client.db.litellm_objectpermissiontable.find_many() + perms = await ObjectPermissionRepository(prisma_client).table.find_many() self._blocked_tools_by_op_id = {} for row in perms: op_id = getattr(row, "object_permission_id", None) @@ -388,7 +390,7 @@ async def add_tool_to_object_permission_blocked( if not object_permission_id or not tool_name: return False try: - row = await prisma_client.db.litellm_objectpermissiontable.find_unique( + row = await ObjectPermissionRepository(prisma_client).table.find_unique( where={"object_permission_id": object_permission_id}, ) if row is None: @@ -397,7 +399,7 @@ async def add_tool_to_object_permission_blocked( if tool_name in current: return True current.append(tool_name) - await prisma_client.db.litellm_objectpermissiontable.update( + await ObjectPermissionRepository(prisma_client).table.update( where={"object_permission_id": object_permission_id}, data={"blocked_tools": current}, ) @@ -418,7 +420,7 @@ async def remove_tool_from_object_permission_blocked( if not object_permission_id or not tool_name: return False try: - row = await prisma_client.db.litellm_objectpermissiontable.find_unique( + row = await ObjectPermissionRepository(prisma_client).table.find_unique( where={"object_permission_id": object_permission_id}, ) if row is None: @@ -427,7 +429,7 @@ async def remove_tool_from_object_permission_blocked( if tool_name not in current: return False current = [t for t in current if t != tool_name] - await prisma_client.db.litellm_objectpermissiontable.update( + await ObjectPermissionRepository(prisma_client).table.update( where={"object_permission_id": object_permission_id}, data={"blocked_tools": current}, ) diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index e0e4bdcf4a4..9f8ea584103 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -13,21 +13,21 @@ from urllib.parse import urlparse from fastapi import APIRouter, Depends, HTTPException, Request from pydantic import BaseModel -from litellm.proxy.common_utils.path_utils import safe_join - from litellm._logging import verbose_proxy_logger from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view +from litellm.proxy.common_utils.path_utils import safe_join from litellm.proxy.guardrails.guardrail_hooks.custom_code.sandbox import ( build_sandbox_globals, compile_sandboxed, ) from litellm.proxy.guardrails.guardrail_registry import GuardrailRegistry from litellm.proxy.guardrails.usage_endpoints import router as guardrails_usage_router +from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view +from litellm.repositories.table_repositories import GuardrailsRepository from litellm.types.guardrails import ( PII_ENTITY_CATEGORIES_MAP, ApplyGuardrailRequest, @@ -373,7 +373,7 @@ async def create_guardrail( # Configuration error — roll back the DB write so the guardrail isn't orphaned if prisma_client is not None: try: - await prisma_client.db.litellm_guardrailstable.delete( + await GuardrailsRepository(prisma_client).table.delete( where={"guardrail_id": guardrail_id} ) except Exception as rollback_err: @@ -705,7 +705,7 @@ async def register_guardrail( ) try: - existing = await prisma_client.db.litellm_guardrailstable.find_unique( + existing = await GuardrailsRepository(prisma_client).table.find_unique( where={"guardrail_name": request.guardrail_name} ) if existing is not None: @@ -732,7 +732,7 @@ async def register_guardrail( guardrail_info_str = safe_dumps(guardrail_info) try: - created = await prisma_client.db.litellm_guardrailstable.create( + created = await GuardrailsRepository(prisma_client).table.create( data={ "guardrail_name": request.guardrail_name, "litellm_params": litellm_params_str, @@ -874,7 +874,7 @@ async def list_guardrail_submissions( where_clause["team_id"] = {"in": visible_team_ids} # Single query: fetch team guardrails visible to the caller - all_team_rows = await prisma_client.db.litellm_guardrailstable.find_many( + all_team_rows = await GuardrailsRepository(prisma_client).table.find_many( where=where_clause, order={"created_at": "desc"}, ) @@ -945,7 +945,7 @@ async def get_guardrail_submission( is_admin = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN try: - row = await prisma_client.db.litellm_guardrailstable.find_unique( + row = await GuardrailsRepository(prisma_client).table.find_unique( where={"guardrail_id": guardrail_id} ) if row is None: @@ -986,7 +986,7 @@ async def approve_guardrail_submission( raise HTTPException(status_code=500, detail="Prisma client not initialized") try: - row = await prisma_client.db.litellm_guardrailstable.find_unique( + row = await GuardrailsRepository(prisma_client).table.find_unique( where={"guardrail_id": guardrail_id} ) if row is None: @@ -1000,7 +1000,7 @@ async def approve_guardrail_submission( ) now = datetime.now(timezone.utc) - await prisma_client.db.litellm_guardrailstable.update( + await GuardrailsRepository(prisma_client).table.update( where={"guardrail_id": guardrail_id}, data={"status": "active", "reviewed_at": now, "updated_at": now}, ) @@ -1072,7 +1072,7 @@ async def reject_guardrail_submission( raise HTTPException(status_code=500, detail="Prisma client not initialized") try: - row = await prisma_client.db.litellm_guardrailstable.find_unique( + row = await GuardrailsRepository(prisma_client).table.find_unique( where={"guardrail_id": guardrail_id} ) if row is None: @@ -1086,7 +1086,7 @@ async def reject_guardrail_submission( ) now = datetime.now(timezone.utc) - await prisma_client.db.litellm_guardrailstable.update( + await GuardrailsRepository(prisma_client).table.update( where={"guardrail_id": guardrail_id}, data={"status": "rejected", "reviewed_at": now, "updated_at": now}, ) @@ -2288,10 +2288,10 @@ async def apply_guardrail( """ import traceback - from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.litellm_core_utils.thread_pool_executor import ( executor as thread_pool_executor, ) + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.proxy.proxy_server import ( general_settings, proxy_config, diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index aafcc5f1819..a80bb817890 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -11,12 +11,15 @@ from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.safe_json_dumps import safe_dumps -from litellm.proxy.guardrails.guardrail_hooks.grayswan import GraySwanGuardrail +from litellm.proxy.guardrails.guardrail_hooks.grayswan import ( + GraySwanGuardrail, +) from litellm.proxy.guardrails.guardrail_hooks.grayswan import ( initialize_guardrail as initialize_grayswan, ) from litellm.proxy.types_utils.utils import get_instance_fn from litellm.proxy.utils import PrismaClient +from litellm.repositories.table_repositories import GuardrailsRepository from litellm.secret_managers.main import get_secret from litellm.types.guardrails import ( Guardrail, @@ -26,6 +29,9 @@ from litellm.types.guardrails import ( SupportedGuardrailIntegrations, ) +from .guardrail_hooks.llm_as_a_judge import ( + initialize_guardrail as initialize_llm_as_a_judge, +) from .guardrail_initializers import ( initialize_bedrock, initialize_hide_secrets, @@ -34,9 +40,6 @@ from .guardrail_initializers import ( initialize_presidio, initialize_tool_permission, ) -from .guardrail_hooks.llm_as_a_judge import ( - initialize_guardrail as initialize_llm_as_a_judge, -) guardrail_initializer_registry = { SupportedGuardrailIntegrations.BEDROCK.value: initialize_bedrock, @@ -257,7 +260,7 @@ class GuardrailRegistry: guardrail_info: str = safe_dumps(guardrail.get("guardrail_info", {})) # Create guardrail in DB - created_guardrail = await prisma_client.db.litellm_guardrailstable.create( + created_guardrail = await GuardrailsRepository(prisma_client).table.create( data={ "guardrail_name": guardrail_name, "litellm_params": litellm_params, @@ -283,7 +286,7 @@ class GuardrailRegistry: """ try: # Delete from DB - await prisma_client.db.litellm_guardrailstable.delete( + await GuardrailsRepository(prisma_client).table.delete( where={"guardrail_id": guardrail_id} ) @@ -311,7 +314,7 @@ class GuardrailRegistry: guardrail_info: str = safe_dumps(guardrail.get("guardrail_info", {})) # Update in DB - updated_guardrail = await prisma_client.db.litellm_guardrailstable.update( + updated_guardrail = await GuardrailsRepository(prisma_client).table.update( where={"guardrail_id": guardrail_id}, data={ "guardrail_name": guardrail_name, @@ -335,11 +338,11 @@ class GuardrailRegistry: Only rows with status == "active" are returned (pending_review and rejected are excluded). """ try: - guardrails_from_db = ( - await prisma_client.db.litellm_guardrailstable.find_many( - where={"status": "active"}, - order={"created_at": "desc"}, - ) + guardrails_from_db = await GuardrailsRepository( + prisma_client + ).table.find_many( + where={"status": "active"}, + order={"created_at": "desc"}, ) guardrails: List[Guardrail] = [] @@ -357,7 +360,7 @@ class GuardrailRegistry: Get a guardrail by its ID from the database """ try: - guardrail = await prisma_client.db.litellm_guardrailstable.find_unique( + guardrail = await GuardrailsRepository(prisma_client).table.find_unique( where={"guardrail_id": guardrail_id} ) @@ -375,7 +378,7 @@ class GuardrailRegistry: Get a guardrail by its name from the database """ try: - guardrail = await prisma_client.db.litellm_guardrailstable.find_unique( + guardrail = await GuardrailsRepository(prisma_client).table.find_unique( where={"guardrail_name": guardrail_name} ) diff --git a/litellm/proxy/guardrails/usage_endpoints.py b/litellm/proxy/guardrails/usage_endpoints.py index 529949c6dd8..d8457cf9c86 100644 --- a/litellm/proxy/guardrails/usage_endpoints.py +++ b/litellm/proxy/guardrails/usage_endpoints.py @@ -12,6 +12,14 @@ from pydantic import BaseModel from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.repositories.table_repositories import ( + DailyGuardrailMetricsRepository, + DailyPolicyMetricsRepository, + GuardrailsRepository, + PolicyRepository, + SpendLogGuardrailIndexRepository, + SpendLogsRepository, +) router = APIRouter() @@ -272,10 +280,10 @@ async def guardrails_usage_overview( try: # Guardrails from DB - guardrails = await prisma_client.db.litellm_guardrailstable.find_many() + guardrails = await GuardrailsRepository(prisma_client).table.find_many() # Daily metrics in range - metrics = await prisma_client.db.litellm_dailyguardrailmetrics.find_many( + metrics = await DailyGuardrailMetricsRepository(prisma_client).table.find_many( where={"date": {"gte": start, "lte": end}} ) @@ -283,9 +291,9 @@ async def guardrails_usage_overview( start_prev = ( datetime.strptime(start, "%Y-%m-%d") - timedelta(days=7) ).strftime("%Y-%m-%d") - metrics_prev = await prisma_client.db.litellm_dailyguardrailmetrics.find_many( - where={"date": {"gte": start_prev, "lt": start}} - ) + metrics_prev = await DailyGuardrailMetricsRepository( + prisma_client + ).table.find_many(where={"date": {"gte": start_prev, "lt": start}}) agg = _aggregate_daily_metrics(metrics, "guardrail_id") prev_agg = _prev_fail_rates(metrics_prev, "guardrail_id") @@ -335,7 +343,7 @@ async def guardrails_usage_detail( end = end_date or now.strftime("%Y-%m-%d") start = start_date or (now - timedelta(days=7)).strftime("%Y-%m-%d") - guardrail = await prisma_client.db.litellm_guardrailstable.find_unique( + guardrail = await GuardrailsRepository(prisma_client).table.find_unique( where={"guardrail_id": guardrail_id} ) if not guardrail: @@ -349,13 +357,13 @@ async def guardrails_usage_detail( ) metric_ids = [i for i in (logical_id, guardrail_id) if i] - metrics = await prisma_client.db.litellm_dailyguardrailmetrics.find_many( + metrics = await DailyGuardrailMetricsRepository(prisma_client).table.find_many( where={ "guardrail_id": {"in": metric_ids}, "date": {"gte": start, "lte": end}, } ) - metrics_prev = await prisma_client.db.litellm_dailyguardrailmetrics.find_many( + metrics_prev = await DailyGuardrailMetricsRepository(prisma_client).table.find_many( where={ "guardrail_id": {"in": metric_ids}, "date": {"lt": start}, @@ -574,7 +582,7 @@ async def guardrails_usage_logs( # Query by both so we match regardless of which was written. effective_guardrail_ids: List[str] = [guardrail_id] if guardrail_id else [] if guardrail_id: - guardrail = await prisma_client.db.litellm_guardrailstable.find_unique( + guardrail = await GuardrailsRepository(prisma_client).table.find_unique( where={"guardrail_id": guardrail_id} ) if guardrail: @@ -585,19 +593,23 @@ async def guardrails_usage_logs( where = _build_usage_logs_where( effective_guardrail_ids or None, policy_id, start_date, end_date ) - index_rows = await prisma_client.db.litellm_spendlogguardrailindex.find_many( + index_rows = await SpendLogGuardrailIndexRepository( + prisma_client + ).table.find_many( where=where, order={"start_time": "desc"}, skip=(page - 1) * page_size, take=page_size + 1, ) - total = await prisma_client.db.litellm_spendlogguardrailindex.count(where=where) + total = await SpendLogGuardrailIndexRepository(prisma_client).table.count( + where=where + ) request_ids = [r.request_id for r in index_rows[:page_size]] if not request_ids: return UsageLogsResponse( logs=[], total=total, page=page, page_size=page_size ) - spend_logs = await prisma_client.db.litellm_spendlogs.find_many( + spend_logs = await SpendLogsRepository(prisma_client).table.find_many( where={"request_id": {"in": request_ids}} ) log_by_id = {s.request_id: s for s in spend_logs} @@ -645,11 +657,13 @@ async def policies_usage_overview( start = start_date or (now - timedelta(days=7)).strftime("%Y-%m-%d") try: - policies = await prisma_client.db.litellm_policytable.find_many() - metrics = await prisma_client.db.litellm_dailypolicymetrics.find_many( + policies = await PolicyRepository(prisma_client).table.find_many() + metrics = await DailyPolicyMetricsRepository(prisma_client).table.find_many( where={"date": {"gte": start, "lte": end}} ) - metrics_prev = await prisma_client.db.litellm_dailypolicymetrics.find_many( + metrics_prev = await DailyPolicyMetricsRepository( + prisma_client + ).table.find_many( where={ "date": { "gte": ( diff --git a/litellm/proxy/guardrails/usage_tracking.py b/litellm/proxy/guardrails/usage_tracking.py index 8907c9201ad..c55c47ca774 100644 --- a/litellm/proxy/guardrails/usage_tracking.py +++ b/litellm/proxy/guardrails/usage_tracking.py @@ -10,6 +10,10 @@ from typing import Any, Dict, List, Optional from litellm._logging import verbose_proxy_logger from litellm.proxy.utils import PrismaClient +from litellm.repositories.table_repositories import ( + DailyGuardrailMetricsRepository, + SpendLogGuardrailIndexRepository, +) def _guardrail_status_to_action(status: Optional[str]) -> str: @@ -132,7 +136,7 @@ async def process_spend_logs_guardrail_usage( } ) try: - await prisma_client.db.litellm_spendlogguardrailindex.create_many( + await SpendLogGuardrailIndexRepository(prisma_client).table.create_many( data=index_data, skip_duplicates=True, ) @@ -146,7 +150,7 @@ async def process_spend_logs_guardrail_usage( n = int(agg["requests_evaluated"]) if n == 0: continue - await prisma_client.db.litellm_dailyguardrailmetrics.upsert( + await DailyGuardrailMetricsRepository(prisma_client).table.upsert( where={ "guardrail_id_date": { "guardrail_id": guardrail_id, diff --git a/litellm/proxy/hooks/user_management_event_hooks.py b/litellm/proxy/hooks/user_management_event_hooks.py index 08fa8d4dfad..c22fd1d6579 100644 --- a/litellm/proxy/hooks/user_management_event_hooks.py +++ b/litellm/proxy/hooks/user_management_event_hooks.py @@ -3,7 +3,6 @@ Hooks that are triggered when a litellm user event occurs """ import asyncio -from litellm._uuid import uuid from datetime import datetime, timezone from typing import Optional @@ -11,6 +10,7 @@ from pydantic import BaseModel import litellm from litellm._logging import verbose_proxy_logger +from litellm._uuid import uuid from litellm.proxy._types import ( AUDIT_ACTIONS, CommonProxyErrors, @@ -24,6 +24,7 @@ from litellm.proxy._types import ( WebhookEvent, ) from litellm.proxy.management_helpers.audit_logs import create_audit_log_for_update +from litellm.repositories.user_repository import UserRepository class UserManagementEventHooks: @@ -57,7 +58,7 @@ class UserManagementEventHooks: try: if prisma_client is None: raise Exception(CommonProxyErrors.db_not_connected_error.value) - user_row: BaseModel = await prisma_client.db.litellm_usertable.find_first( + user_row: BaseModel = await UserRepository(prisma_client).table.find_first( where={"user_id": response.user_id} ) diff --git a/litellm/proxy/management_endpoints/access_group_endpoints.py b/litellm/proxy/management_endpoints/access_group_endpoints.py index 62a770f46ae..65f7ffc9081 100644 --- a/litellm/proxy/management_endpoints/access_group_endpoints.py +++ b/litellm/proxy/management_endpoints/access_group_endpoints.py @@ -19,6 +19,7 @@ from litellm.proxy.auth.auth_checks import ( from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.utils import get_prisma_client_or_throw +from litellm.repositories.table_repositories import AccessGroupRepository from litellm.types.access_group import ( AccessGroupCreateRequest, AccessGroupResponse, @@ -386,7 +387,7 @@ async def list_access_groups( CommonProxyErrors.db_not_connected_error.value ) - records = await prisma_client.db.litellm_accessgrouptable.find_many( + records = await AccessGroupRepository(prisma_client).table.find_many( order={"created_at": "desc"} ) return [_record_to_response(r) for r in records] @@ -405,7 +406,7 @@ async def get_access_group( CommonProxyErrors.db_not_connected_error.value ) - record = await prisma_client.db.litellm_accessgrouptable.find_unique( + record = await AccessGroupRepository(prisma_client).table.find_unique( where={"access_group_id": access_group_id} ) if record is None: diff --git a/litellm/proxy/management_endpoints/budget_management_endpoints.py b/litellm/proxy/management_endpoints/budget_management_endpoints.py index 2eda1b30c5d..698155a5c26 100644 --- a/litellm/proxy/management_endpoints/budget_management_endpoints.py +++ b/litellm/proxy/management_endpoints/budget_management_endpoints.py @@ -16,11 +16,12 @@ import math from fastapi import APIRouter, Depends, HTTPException -from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view from litellm.proxy.utils import jsonify_object +from litellm.repositories.budget_repository import BudgetRepository router = APIRouter() @@ -98,7 +99,7 @@ async def new_budget( budget_obj_json = budget_obj.model_dump(exclude_none=True) budget_obj_jsonified = jsonify_object(budget_obj_json) # json dump any dictionaries try: - response = await prisma_client.db.litellm_budgettable.create( + response = await BudgetRepository(prisma_client).table.create( data={ **budget_obj_jsonified, # type: ignore "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name, @@ -182,7 +183,7 @@ async def update_budget( except ValueError as e: raise HTTPException(status_code=400, detail={"error": str(e)}) - response = await prisma_client.db.litellm_budgettable.update( + response = await BudgetRepository(prisma_client).table.update( where={"budget_id": budget_obj.budget_id}, data={ **budget_obj.model_dump(exclude_unset=True), # type: ignore @@ -217,7 +218,7 @@ async def info_budget(data: BudgetRequest): "error": f"Specify list of budget id's to query. Passed in={data.budgets}" }, ) - response = await prisma_client.db.litellm_budgettable.find_many( + response = await BudgetRepository(prisma_client).table.find_many( where={"budget_id": {"in": data.budgets}}, ) @@ -261,7 +262,7 @@ async def budget_settings( ) ## get budget item from db - db_budget_row = await prisma_client.db.litellm_budgettable.find_first( + db_budget_row = await BudgetRepository(prisma_client).table.find_first( where={"budget_id": budget_id} ) @@ -327,7 +328,7 @@ async def list_budget( }, ) - response = await prisma_client.db.litellm_budgettable.find_many() + response = await BudgetRepository(prisma_client).table.find_many() return response @@ -366,7 +367,7 @@ async def delete_budget( }, ) - response = await prisma_client.db.litellm_budgettable.delete( + response = await BudgetRepository(prisma_client).table.delete( where={"budget_id": data.id} ) diff --git a/litellm/proxy/management_endpoints/cache_settings_endpoints.py b/litellm/proxy/management_endpoints/cache_settings_endpoints.py index 0a26b23beff..d8eb5dfee92 100644 --- a/litellm/proxy/management_endpoints/cache_settings_endpoints.py +++ b/litellm/proxy/management_endpoints/cache_settings_endpoints.py @@ -27,6 +27,7 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.repositories.table_repositories import CacheConfigRepository from litellm.types.management_endpoints import ( CACHE_SETTINGS_FIELDS, REDIS_TYPE_DESCRIPTIONS, @@ -159,7 +160,7 @@ class CacheSettingsManager: import json try: - cache_config = await prisma_client.db.litellm_cacheconfig.find_unique( + cache_config = await CacheConfigRepository(prisma_client).table.find_unique( where={"id": "cache_config"} ) if cache_config is not None and cache_config.cache_settings: @@ -274,7 +275,7 @@ async def get_cache_settings( # Try to get cache settings from database current_values = {} if prisma_client is not None: - cache_config = await prisma_client.db.litellm_cacheconfig.find_unique( + cache_config = await CacheConfigRepository(prisma_client).table.find_unique( where={"id": "cache_config"} ) if cache_config is not None and cache_config.cache_settings: @@ -417,7 +418,7 @@ async def update_cache_settings( # Snapshot the prior settings (key set only — values get redacted in # the audit row) so the audit-log entry shows which fields changed. - existing_row = await prisma_client.db.litellm_cacheconfig.find_unique( + existing_row = await CacheConfigRepository(prisma_client).table.find_unique( where={"id": "cache_config"} ) before_settings: Optional[Dict[str, Any]] = None @@ -434,7 +435,7 @@ async def update_cache_settings( ) # Save to database - await prisma_client.db.litellm_cacheconfig.upsert( + await CacheConfigRepository(prisma_client).table.upsert( where={"id": "cache_config"}, data={ "create": { diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index d173cd745ba..92cc2008c73 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -8,6 +8,10 @@ from fastapi import HTTPException, status from litellm._logging import verbose_proxy_logger from litellm.proxy._types import CommonProxyErrors from litellm.proxy.utils import PrismaClient +from litellm.repositories.table_repositories import DeletedVerificationTokenRepository +from litellm.repositories.verification_token_repository import ( + VerificationTokenRepository, +) from litellm.types.proxy.management_endpoints.common_daily_activity import ( BreakdownMetrics, DailySpendData, @@ -346,7 +350,7 @@ async def get_api_key_metadata( This ensures that key_alias and team_id are preserved in historical activity logs even after a key is deleted or regenerated. """ - key_records = await prisma_client.db.litellm_verificationtoken.find_many( + key_records = await VerificationTokenRepository(prisma_client).table.find_many( where={"token": {"in": list(api_keys)}} ) result = { @@ -357,11 +361,11 @@ async def get_api_key_metadata( missing_keys = api_keys - set(result.keys()) if missing_keys: try: - deleted_key_records = ( - await prisma_client.db.litellm_deletedverificationtoken.find_many( - where={"token": {"in": list(missing_keys)}}, - order={"deleted_at": "desc"}, - ) + deleted_key_records = await DeletedVerificationTokenRepository( + prisma_client + ).table.find_many( + where={"token": {"in": list(missing_keys)}}, + order={"deleted_at": "desc"}, ) # Use the most recent deleted record for each token (ordered by deleted_at desc) for k in deleted_key_records: diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index 31d831d773c..458cba686e6 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -17,10 +17,13 @@ from litellm.proxy._types import ( NewProjectRequest, UpdateProjectRequest, UserAPIKeyAuth, - user_api_key_has_admin_view as _user_has_admin_view, # noqa: F401 re-exported +) +from litellm.proxy._types import ( # noqa: F401 re-exported + user_api_key_has_admin_view as _user_has_admin_view, ) from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time from litellm.proxy.utils import _premium_user_check +from litellm.repositories.team_repository import TeamRepository if TYPE_CHECKING: from litellm.proxy._types import NewProjectRequest, UpdateProjectRequest @@ -205,7 +208,7 @@ async def _user_has_admin_privileges( # Check if user is team admin for any team if user_obj.teams is not None and len(user_obj.teams) > 0: # Get all teams user is in - teams = await prisma_client.db.litellm_teamtable.find_many( + teams = await TeamRepository(prisma_client).table.find_many( where={"team_id": {"in": user_obj.teams}} ) @@ -282,7 +285,7 @@ async def _team_admin_can_invite_user( if not target_user_obj.teams or len(target_user_obj.teams) == 0: return False - teams = await prisma_client.db.litellm_teamtable.find_many( + teams = await TeamRepository(prisma_client).table.find_many( where={"team_id": {"in": admin_user_obj.teams}} ) admin_team_ids = [ diff --git a/litellm/proxy/management_endpoints/config_override_endpoints.py b/litellm/proxy/management_endpoints/config_override_endpoints.py index 7f7aa485fb3..97cb5eeddc4 100644 --- a/litellm/proxy/management_endpoints/config_override_endpoints.py +++ b/litellm/proxy/management_endpoints/config_override_endpoints.py @@ -30,6 +30,7 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.repositories.table_repositories import ConfigOverridesRepository from litellm.types.llms.custom_http import httpxSpecialProvider from litellm.types.proxy.management_endpoints.config_overrides import ( ConfigOverrideSettingsResponse, @@ -254,7 +255,7 @@ async def update_hashicorp_vault_config( # Merge ALL fields the user didn't send: try DB first, fall back to env vars. # Omitted field = keep existing; empty string = clear/remove the field. - existing_record = await prisma_client.db.litellm_configoverrides.find_unique( + existing_record = await ConfigOverridesRepository(prisma_client).table.find_unique( where={"config_type": "hashicorp_vault"} ) existing_decrypted: Optional[Dict[str, Any]] = None @@ -321,7 +322,7 @@ async def update_hashicorp_vault_config( # Only persist to DB after successful init encrypted_data = proxy_config._encrypt_env_variables(config_data) config_value = safe_dumps(encrypted_data) - await prisma_client.db.litellm_configoverrides.upsert( + await ConfigOverridesRepository(prisma_client).table.upsert( where={"config_type": "hashicorp_vault"}, data={ "create": { @@ -391,7 +392,7 @@ async def get_hashicorp_vault_config( field_schema = _build_field_schema(HashicorpVaultConfig) # Try to load from DB - db_record = await prisma_client.db.litellm_configoverrides.find_unique( + db_record = await ConfigOverridesRepository(prisma_client).table.find_unique( where={"config_type": "hashicorp_vault"} ) @@ -448,7 +449,7 @@ async def delete_hashicorp_vault_config( # Capture the prior config before delete so the audit-log row can # show *what* was removed (keys only — values get redacted). - existing_record = await prisma_client.db.litellm_configoverrides.find_unique( + existing_record = await ConfigOverridesRepository(prisma_client).table.find_unique( where={"config_type": "hashicorp_vault"} ) before_config: Optional[Dict[str, Any]] = None @@ -463,7 +464,7 @@ async def delete_hashicorp_vault_config( # Delete DB record if it exists — ignore if not found deleted = False try: - await prisma_client.db.litellm_configoverrides.delete( + await ConfigOverridesRepository(prisma_client).table.delete( where={"config_type": "hashicorp_vault"} ) deleted = True diff --git a/litellm/proxy/management_endpoints/customer_endpoints.py b/litellm/proxy/management_endpoints/customer_endpoints.py index 1fd8320db20..f1a34bb0ed4 100644 --- a/litellm/proxy/management_endpoints/customer_endpoints.py +++ b/litellm/proxy/management_endpoints/customer_endpoints.py @@ -17,8 +17,8 @@ import fastapi from fastapi import APIRouter, Depends, HTTPException, Request import litellm -from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity @@ -27,6 +27,8 @@ from litellm.proxy.management_helpers.object_permission_utils import ( handle_update_object_permission_common, ) from litellm.proxy.utils import handle_exception_on_proxy +from litellm.repositories.budget_repository import BudgetRepository +from litellm.repositories.table_repositories import EndUserRepository from litellm.types.proxy.management_endpoints.common_daily_activity import ( SpendAnalyticsPaginatedResponse, ) @@ -68,7 +70,7 @@ async def block_user(data: BlockUsers): records = [] if prisma_client is not None: for id in data.user_ids: - record = await prisma_client.db.litellm_endusertable.upsert( + record = await EndUserRepository(prisma_client).table.upsert( where={"user_id": id}, # type: ignore data={ "create": {"user_id": id, "blocked": True}, # type: ignore @@ -337,7 +339,7 @@ async def new_end_user( _new_budget = new_budget_request(data) if _new_budget is not None: try: - budget_record = await prisma_client.db.litellm_budgettable.create( + budget_record = await BudgetRepository(prisma_client).table.create( data={ **_new_budget.model_dump(exclude_unset=True), "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name, # type: ignore @@ -373,7 +375,7 @@ async def new_end_user( new_end_user_obj.pop("object_permission", None) ## WRITE TO DB ## - end_user_record = await prisma_client.db.litellm_endusertable.create( + end_user_record = await EndUserRepository(prisma_client).table.create( data=new_end_user_obj, # type: ignore include={"litellm_budget_table": True, "object_permission": True}, ) @@ -446,7 +448,7 @@ async def end_user_info( detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - user_info = await prisma_client.db.litellm_endusertable.find_first( + user_info = await EndUserRepository(prisma_client).table.find_first( where={"user_id": end_user_id}, include={"litellm_budget_table": True, "object_permission": True}, ) @@ -569,7 +571,7 @@ async def update_end_user( non_default_values[k] = v ## Get end user table data ## - end_user_table_data = await prisma_client.db.litellm_endusertable.find_first( + end_user_table_data = await EndUserRepository(prisma_client).table.find_first( where={"user_id": data.user_id}, include={"litellm_budget_table": True} ) @@ -613,17 +615,17 @@ async def update_end_user( if budget_table_data: if end_user_budget_table is None: ## Create new budget ## - budget_table_data_record = ( - await prisma_client.db.litellm_budgettable.create( - data={ - **budget_table_data, - "created_by": user_api_key_dict.user_id - or litellm_proxy_admin_name, - "updated_by": user_api_key_dict.user_id - or litellm_proxy_admin_name, - }, - include={"end_users": True}, - ) + budget_table_data_record = await BudgetRepository( + prisma_client + ).table.create( + data={ + **budget_table_data, + "created_by": user_api_key_dict.user_id + or litellm_proxy_admin_name, + "updated_by": user_api_key_dict.user_id + or litellm_proxy_admin_name, + }, + include={"end_users": True}, ) update_end_user_table_data["budget_id"] = ( @@ -631,11 +633,11 @@ async def update_end_user( ) else: ## Update existing budget ## - budget_table_data_record = ( - await prisma_client.db.litellm_budgettable.update( - where={"budget_id": end_user_budget_table.budget_id}, - data=budget_table_data, - ) + budget_table_data_record = await BudgetRepository( + prisma_client + ).table.update( + where={"budget_id": end_user_budget_table.budget_id}, + data=budget_table_data, ) ## Update user table, with update params + new budget id (if set) ## @@ -652,7 +654,7 @@ async def update_end_user( if data.user_id is not None and len(data.user_id) > 0: update_end_user_table_data["user_id"] = data.user_id # type: ignore verbose_proxy_logger.debug("In update customer, user_id condition block.") - response = await prisma_client.db.litellm_endusertable.update( + response = await EndUserRepository(prisma_client).table.update( where={"user_id": data.user_id}, data=update_end_user_table_data, include={"litellm_budget_table": True, "object_permission": True} # type: ignore ) if response is None: @@ -737,7 +739,7 @@ async def delete_end_user( and len(data.user_ids) > 0 ): # First check if all users exist - existing_users = await prisma_client.db.litellm_endusertable.find_many( + existing_users = await EndUserRepository(prisma_client).table.find_many( where={"user_id": {"in": data.user_ids}} ) existing_user_ids = {user.user_id for user in existing_users} @@ -756,7 +758,7 @@ async def delete_end_user( ) # All users exist, proceed with deletion - response = await prisma_client.db.litellm_endusertable.delete_many( + response = await EndUserRepository(prisma_client).table.delete_many( where={"user_id": {"in": data.user_ids}} ) verbose_proxy_logger.debug( @@ -828,7 +830,7 @@ async def list_end_user( detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - response = await prisma_client.db.litellm_endusertable.find_many( + response = await EndUserRepository(prisma_client).table.find_many( include={"litellm_budget_table": True, "object_permission": True} ) @@ -903,7 +905,7 @@ async def get_customer_daily_activity( where_condition = {} if end_user_ids_list: where_condition["user_id"] = {"in": list(end_user_ids_list)} - end_user_aliases = await prisma_client.db.litellm_endusertable.find_many( + end_user_aliases = await EndUserRepository(prisma_client).table.find_many( where=where_condition ) end_user_alias_metadata = {e.user_id: {"alias": e.alias} for e in end_user_aliases} diff --git a/litellm/proxy/management_endpoints/fallback_management_endpoints.py b/litellm/proxy/management_endpoints/fallback_management_endpoints.py index ffb12111d82..1333122c87a 100644 --- a/litellm/proxy/management_endpoints/fallback_management_endpoints.py +++ b/litellm/proxy/management_endpoints/fallback_management_endpoints.py @@ -27,6 +27,7 @@ else: # fastapi is only required for proxy, not for SDK usage pass +from litellm.repositories.config_repository import ConfigRepository from litellm.types.management_endpoints.router_settings_endpoints import ( FallbackCreateRequest, FallbackDeleteResponse, @@ -157,7 +158,7 @@ async def create_fallback( # Save to database - convert router_settings to JSON string router_settings_json = json.dumps(router_settings) - await prisma_client.db.litellm_config.upsert( + await ConfigRepository(prisma_client).table.upsert( where={"param_name": "router_settings"}, data={ "create": { @@ -336,7 +337,7 @@ async def delete_fallback( # Save to database - convert router_settings to JSON string router_settings_json = json.dumps(router_settings) - await prisma_client.db.litellm_config.upsert( + await ConfigRepository(prisma_client).table.upsert( where={"param_name": "router_settings"}, data={ "create": { diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 7b8f0f72e13..b3a5c66e9e1 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -43,6 +43,17 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( ) from litellm.proxy.management_helpers.utils import management_endpoint_wrapper from litellm.proxy.utils import handle_exception_on_proxy, hash_password +from litellm.repositories.organization_repository import OrganizationRepository +from litellm.repositories.table_repositories import ( + InvitationLinkRepository, + OrganizationMembershipRepository, + TeamMembershipRepository, +) +from litellm.repositories.team_repository import TeamRepository +from litellm.repositories.user_repository import UserRepository +from litellm.repositories.verification_token_repository import ( + VerificationTokenRepository, +) from litellm.types.proxy.management_endpoints.common_daily_activity import ( SpendAnalyticsPaginatedResponse, ) @@ -154,7 +165,7 @@ async def _check_duplicate_user_field( if case_insensitive: where_clause[field_name]["mode"] = "insensitive" - existing_user = await prisma_client.db.litellm_usertable.find_first( + existing_user = await UserRepository(prisma_client).table.find_first( where=where_clause ) @@ -434,7 +445,7 @@ async def new_user( await _check_duplicate_user_email(data.user_email, prisma_client) # Check if license is over limit - total_users = await prisma_client.db.litellm_usertable.count() + total_users = await UserRepository(prisma_client).table.count() if total_users and _license_check.is_over_limit(total_users=total_users): raise HTTPException( status_code=403, @@ -851,7 +862,7 @@ async def _check_user_info_v2_access( # Helper: fetch the target user row (reused across branches) async def _fetch_target_user(): - return await prisma_client.db.litellm_usertable.find_unique( + return await UserRepository(prisma_client).table.find_unique( where={"user_id": target_user_id} ) @@ -866,7 +877,7 @@ async def _check_user_info_v2_access( # Rule 3: Team admins can look up users in their teams if user_api_key_dict.user_id is not None: # Get caller's teams - caller_user = await prisma_client.db.litellm_usertable.find_unique( + caller_user = await UserRepository(prisma_client).table.find_unique( where={"user_id": user_api_key_dict.user_id} ) if caller_user is not None and caller_user.teams: @@ -876,7 +887,7 @@ async def _check_user_info_v2_access( return None # Get all teams the caller belongs to - teams = await prisma_client.db.litellm_teamtable.find_many( + teams = await TeamRepository(prisma_client).table.find_many( where={"team_id": {"in": caller_user.teams}} ) for team in teams: @@ -1165,7 +1176,7 @@ async def _schedule_user_update_audit_log( if prisma_client is None: return try: - updated_user_row = await prisma_client.db.litellm_usertable.find_first( + updated_user_row = await UserRepository(prisma_client).table.find_first( where={"user_id": response["user_id"]} ) if updated_user_row: @@ -1255,11 +1266,11 @@ async def _update_single_user_helper( existing_user_row: Optional[BaseModel] = None if user_request.user_id: - existing_user_row = await prisma_client.db.litellm_usertable.find_first( + existing_user_row = await UserRepository(prisma_client).table.find_first( where={"user_id": user_request.user_id} ) elif user_request.user_email: - existing_user_row = await prisma_client.db.litellm_usertable.find_first( + existing_user_row = await UserRepository(prisma_client).table.find_first( where={"user_email": user_request.user_email} ) @@ -1640,7 +1651,7 @@ async def bulk_user_update( detail="Only proxy admins can update all users at once.", ) # Optimized path for updating all users directly in database - all_users_in_db = await prisma_client.db.litellm_usertable.find_many( + all_users_in_db = await UserRepository(prisma_client).table.find_many( order={"created_at": "desc"} ) @@ -1676,7 +1687,7 @@ async def bulk_user_update( try: # Perform bulk database update - await prisma_client.db.litellm_usertable.update_many( + await UserRepository(prisma_client).table.update_many( where={}, data=non_default_values # Update all users ) @@ -1783,7 +1794,7 @@ async def get_user_key_counts( # Get count for each user_id individually for user_id in user_ids: - count = await prisma_client.db.litellm_verificationtoken.count( + count = await VerificationTokenRepository(prisma_client).table.count( where={ "user_id": user_id, "OR": [ @@ -2056,7 +2067,7 @@ async def get_users( else None ) - users = await prisma_client.db.litellm_usertable.find_many( + users = await UserRepository(prisma_client).table.find_many( where=where_conditions, skip=skip, take=page_size, @@ -2066,7 +2077,9 @@ async def get_users( ) # Get total count of user rows - total_count = await prisma_client.db.litellm_usertable.count(where=where_conditions) + total_count = await UserRepository(prisma_client).table.count( + where=where_conditions + ) # Get key count for each user if users is not None: @@ -2137,14 +2150,14 @@ async def delete_user( from litellm.proxy.management_endpoints.team_endpoints import ( _cleanup_members_with_roles, ) + from litellm.proxy.management_helpers.audit_logs import ( + get_audit_log_changed_by, + ) from litellm.proxy.proxy_server import ( create_audit_log_for_update, litellm_proxy_admin_name, prisma_client, ) - from litellm.proxy.management_helpers.audit_logs import ( - get_audit_log_changed_by, - ) if prisma_client is None: raise HTTPException(status_code=500, detail={"error": "No db connected"}) @@ -2164,7 +2177,7 @@ async def delete_user( caller_admin_org_ids: set = set() if not caller_is_proxy_admin: caller_memberships = ( - await prisma_client.db.litellm_organizationmembership.find_many( + await OrganizationMembershipRepository(prisma_client).table.find_many( where={ "user_id": user_api_key_dict.user_id, "user_role": LitellmUserRoles.ORG_ADMIN.value, @@ -2188,11 +2201,9 @@ async def delete_user( # an N+1 DB call when delete_user is called with a large user_ids list. target_org_ids_by_user: Dict[str, set] = {} if not caller_is_proxy_admin: - all_target_memberships = ( - await prisma_client.db.litellm_organizationmembership.find_many( - where={"user_id": {"in": data.user_ids}} - ) - ) + all_target_memberships = await OrganizationMembershipRepository( + prisma_client + ).table.find_many(where={"user_id": {"in": data.user_ids}}) for m in all_target_memberships: if not m.organization_id: continue @@ -2200,7 +2211,7 @@ async def delete_user( # check that all teams passed exist for user_id in data.user_ids: - user_row = await prisma_client.db.litellm_usertable.find_unique( + user_row = await UserRepository(prisma_client).table.find_unique( where={"user_id": user_id} ) @@ -2254,7 +2265,7 @@ async def delete_user( ) ## CLEANUP MEMBERS_WITH_ROLES - fetch_all_teams = await prisma_client.db.litellm_teamtable.find_many( + fetch_all_teams = await TeamRepository(prisma_client).table.find_many( where={"team_id": {"in": user_row.teams}} ) teams_to_update = [] @@ -2277,19 +2288,19 @@ async def delete_user( ## update teams for team in teams_to_update: - await prisma_client.db.litellm_teamtable.update( + await TeamRepository(prisma_client).table.update( where={"team_id": team.team_id}, data={"members_with_roles": team.members_with_roles}, ) # End of Audit logging ## DELETE ASSOCIATED KEYS - await prisma_client.db.litellm_verificationtoken.delete_many( + await VerificationTokenRepository(prisma_client).table.delete_many( where={"user_id": {"in": data.user_ids}} ) ## DELETE ASSOCIATED INVITATION LINKS - await prisma_client.db.litellm_invitationlink.delete_many( + await InvitationLinkRepository(prisma_client).table.delete_many( where={ "OR": [ {"user_id": {"in": data.user_ids}}, @@ -2300,17 +2311,17 @@ async def delete_user( ) ## DELETE ASSOCIATED ORGANIZATION MEMBERSHIPS - await prisma_client.db.litellm_organizationmembership.delete_many( + await OrganizationMembershipRepository(prisma_client).table.delete_many( where={"user_id": {"in": data.user_ids}} ) ## DELETE ASSOCIATED TEAM MEMBERSHIPS - await prisma_client.db.litellm_teammembership.delete_many( + await TeamMembershipRepository(prisma_client).table.delete_many( where={"user_id": {"in": data.user_ids}} ) ## DELETE USERS - deleted_users = await prisma_client.db.litellm_usertable.delete_many( + deleted_users = await UserRepository(prisma_client).table.delete_many( where={"user_id": {"in": data.user_ids}} ) @@ -2340,16 +2351,18 @@ async def add_internal_user_to_organization( try: # Check if organization_id exists - organization_row = await prisma_client.db.litellm_organizationtable.find_unique( - where={"organization_id": organization_id} - ) + organization_row = await OrganizationRepository( + prisma_client + ).table.find_unique(where={"organization_id": organization_id}) if organization_row is None: raise Exception( f"Organization not found, passed organization_id={organization_id}" ) # Create a new organization membership entry - new_membership = await prisma_client.db.litellm_organizationmembership.create( + new_membership = await OrganizationMembershipRepository( + prisma_client + ).table.create( data={ "user_id": user_id, "organization_id": organization_id, @@ -2559,13 +2572,13 @@ async def ui_view_users( } # Query users with pagination and filters - users: Optional[List[BaseModel]] = ( - await prisma_client.db.litellm_usertable.find_many( - where=where_conditions, - skip=skip, - take=page_size, - order={"created_at": "desc"}, - ) + users: Optional[List[BaseModel]] = await UserRepository( + prisma_client + ).table.find_many( + where=where_conditions, + skip=skip, + take=page_size, + order={"created_at": "desc"}, ) if not users: diff --git a/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py b/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py index 1ee5bfb0226..a5a364c3679 100644 --- a/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py +++ b/litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py @@ -11,6 +11,7 @@ from litellm.proxy._types import ( ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view +from litellm.repositories.table_repositories import JWTKeyMappingRepository router = APIRouter() @@ -61,7 +62,7 @@ async def create_jwt_key_mapping( if data.description is not None: create_data["description"] = data.description - new_mapping = await prisma_client.db.litellm_jwtkeymapping.create( + new_mapping = await JWTKeyMappingRepository(prisma_client).table.create( data=create_data ) @@ -113,7 +114,7 @@ async def update_jwt_key_mapping( try: # Get old mapping for cache invalidation - old_mapping = await prisma_client.db.litellm_jwtkeymapping.find_unique( + old_mapping = await JWTKeyMappingRepository(prisma_client).table.find_unique( where={"id": data.id} ) @@ -123,7 +124,7 @@ async def update_jwt_key_mapping( cache_key = f"jwt_key_mapping:{old_mapping.jwt_claim_name}:{old_mapping.jwt_claim_value}" await user_api_key_cache.async_delete_cache(cache_key) - updated_mapping = await prisma_client.db.litellm_jwtkeymapping.update( + updated_mapping = await JWTKeyMappingRepository(prisma_client).table.update( where={"id": data.id}, data=update_data ) @@ -166,7 +167,7 @@ async def delete_jwt_key_mapping( try: # Get old mapping for cache invalidation - old_mapping = await prisma_client.db.litellm_jwtkeymapping.find_unique( + old_mapping = await JWTKeyMappingRepository(prisma_client).table.find_unique( where={"id": data.id} ) @@ -176,7 +177,7 @@ async def delete_jwt_key_mapping( cache_key = f"jwt_key_mapping:{old_mapping.jwt_claim_name}:{old_mapping.jwt_claim_value}" await user_api_key_cache.async_delete_cache(cache_key) - await prisma_client.db.litellm_jwtkeymapping.delete(where={"id": data.id}) + await JWTKeyMappingRepository(prisma_client).table.delete(where={"id": data.id}) return {"status": "success"} except HTTPException: raise @@ -206,12 +207,12 @@ async def list_jwt_key_mappings( try: skip = (page - 1) * size - mappings = await prisma_client.db.litellm_jwtkeymapping.find_many( + mappings = await JWTKeyMappingRepository(prisma_client).table.find_many( skip=skip, take=size, order={"created_at": "desc"}, ) - total_count = await prisma_client.db.litellm_jwtkeymapping.count() + total_count = await JWTKeyMappingRepository(prisma_client).table.count() return { "mappings": [_to_response(m) for m in mappings], "total_count": total_count, @@ -245,7 +246,7 @@ async def info_jwt_key_mapping( raise HTTPException(status_code=500, detail="Database not connected") try: - mapping = await prisma_client.db.litellm_jwtkeymapping.find_unique( + mapping = await JWTKeyMappingRepository(prisma_client).table.find_unique( where={"id": id} ) if mapping is None: diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 99d0bac88af..8f606fdf90d 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -28,7 +28,6 @@ from fastapi import APIRouter, Depends, Header, HTTPException, Query, Request, s import litellm from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid -from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.constants import ( LENGTH_OF_LITELLM_GENERATED_KEY, LITELLM_PROXY_ADMIN_NAME, @@ -58,6 +57,7 @@ from litellm.proxy.common_utils.callback_utils import ( ) from litellm.proxy.common_utils.rbac_utils import check_org_admin_can_generate_keys from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks from litellm.proxy.management_endpoints.common_utils import ( _check_passthrough_routes_caller_permission, @@ -91,6 +91,19 @@ from litellm.proxy.utils import ( handle_exception_on_proxy, is_valid_api_key, ) +from litellm.repositories.budget_repository import BudgetRepository +from litellm.repositories.config_repository import ConfigRepository +from litellm.repositories.credentials_repository import CredentialsRepository +from litellm.repositories.model_repository import ModelRepository +from litellm.repositories.table_repositories import ( + DeletedVerificationTokenRepository, + DeprecatedVerificationTokenRepository, +) +from litellm.repositories.team_repository import TeamRepository +from litellm.repositories.user_repository import UserRepository +from litellm.repositories.verification_token_repository import ( + VerificationTokenRepository, +) from litellm.router import Router from litellm.secret_managers.main import get_secret from litellm.types.proxy.management_endpoints.key_management_endpoints import ( @@ -582,7 +595,7 @@ async def validate_team_id_used_in_service_account_request( ) # check if team_id exists in the database - team = await prisma_client.db.litellm_teamtable.find_unique( + team = await TeamRepository(prisma_client).table.find_unique( where={"team_id": team_id}, ) if team is None: @@ -774,7 +787,7 @@ async def _common_key_generation_helper( # noqa: PLR0915 ) new_budget = prisma_client.jsonify_object(budget_row.json(exclude_none=True)) - _budget = await prisma_client.db.litellm_budgettable.create( + _budget = await BudgetRepository(prisma_client).table.create( data={ **new_budget, # type: ignore "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name, @@ -1144,7 +1157,7 @@ async def _check_team_key_limits( # calculate allocated tpm/rpm limit # check if specified tpm/rpm limit is greater than allocated tpm/rpm limit - keys = await prisma_client.db.litellm_verificationtoken.find_many( + keys = await VerificationTokenRepository(prisma_client).table.find_many( where={"team_id": team_table.team_id}, ) # Exclude the key being updated to avoid double-counting its limits. @@ -1294,7 +1307,7 @@ async def _validate_caller_can_assign_key_org( detail="Cannot assign a key to an organization without a user_id on the caller's token", ) - user_row = await prisma_client.db.litellm_usertable.find_unique( + user_row = await UserRepository(prisma_client).table.find_unique( where={"user_id": user_api_key_dict.user_id}, include={"organization_memberships": True}, ) @@ -1339,7 +1352,7 @@ async def _check_org_key_limits( # get all organization keys # calculate allocated tpm/rpm limit # check if specified tpm/rpm limit is greater than allocated tpm/rpm limit - keys = await prisma_client.db.litellm_verificationtoken.find_many( + keys = await VerificationTokenRepository(prisma_client).table.find_many( where={"organization_id": org_table.organization_id}, ) # Exclude the key being updated to avoid double-counting its limits. @@ -1968,9 +1981,9 @@ async def _get_and_validate_existing_key( hashed_token = _hash_token_if_needed(token=token) - existing_key_row = await prisma_client.db.litellm_verificationtoken.find_unique( - where={"token": hashed_token} - ) + existing_key_row = await VerificationTokenRepository( + prisma_client + ).table.find_unique(where={"token": hashed_token}) if existing_key_row is None: raise ProxyException( @@ -2869,7 +2882,9 @@ async def bulk_update_team_keys( # `blocked` is Boolean? with no default; `/key/generate` writes NULL. Prisma's `NOT` # excludes NULLs, so explicitly OR `false` with `null` to include them. now = datetime.now(timezone.utc) - existing_keys = await prisma_client.db.litellm_verificationtoken.find_many( + existing_keys = await VerificationTokenRepository( + prisma_client + ).table.find_many( where={ "team_id": data.team_id, "AND": [ @@ -2907,7 +2922,9 @@ async def bulk_update_team_keys( seen_hashes.add(h) requested_tokens.append(k) hashed_key_ids.append(h) - existing_keys = await prisma_client.db.litellm_verificationtoken.find_many( + existing_keys = await VerificationTokenRepository( + prisma_client + ).table.find_many( where={"team_id": data.team_id, "token": {"in": hashed_key_ids}} ) @@ -3232,7 +3249,9 @@ async def info_key_fn_v2( # Resolve key_aliases to tokens so we never pass token=None (unbounded query) tokens_to_query = list(data.keys) if data.keys else [] if data.key_aliases: - alias_rows = await prisma_client.db.litellm_verificationtoken.find_many( + alias_rows = await VerificationTokenRepository( + prisma_client + ).table.find_many( where={"key_alias": {"in": data.key_aliases}}, include={"litellm_budget_table": True}, ) @@ -3311,7 +3330,7 @@ async def info_key_fn( hashed_key: Optional[str] = key if key is not None: hashed_key = _hash_token_if_needed(token=key) - key_info = await prisma_client.db.litellm_verificationtoken.find_unique( + key_info = await VerificationTokenRepository(prisma_client).table.find_unique( where={"token": hashed_key}, # type: ignore include={"litellm_budget_table": True}, ) @@ -3851,7 +3870,7 @@ async def delete_verification_tokens( if prisma_client: tokens = [_hash_token_if_needed(token=key) for key in tokens] _keys_being_deleted: List[LiteLLM_VerificationToken] = ( - await prisma_client.db.litellm_verificationtoken.find_many( + await VerificationTokenRepository(prisma_client).table.find_many( where={"token": {"in": tokens}} ) ) @@ -3989,7 +4008,9 @@ async def _save_deleted_verification_token_records( """Save deleted verification token records to the database.""" if not records: return - await prisma_client.db.litellm_deletedverificationtoken.create_many(data=records) + await DeletedVerificationTokenRepository(prisma_client).table.create_many( + data=records + ) async def _persist_deleted_verification_tokens( @@ -4017,9 +4038,9 @@ async def delete_key_aliases( user_api_key_dict: UserAPIKeyAuth, litellm_changed_by: Optional[str] = None, ) -> Tuple[Optional[Dict], List[LiteLLM_VerificationToken]]: - _keys_being_deleted = await prisma_client.db.litellm_verificationtoken.find_many( - where={"key_alias": {"in": key_aliases}} - ) + _keys_being_deleted = await VerificationTokenRepository( + prisma_client + ).table.find_many(where={"key_alias": {"in": key_aliases}}) tokens = [key.token for key in _keys_being_deleted] return await delete_verification_tokens( @@ -4054,9 +4075,7 @@ async def _rotate_master_key( # noqa: PLR0915 from litellm.proxy.proxy_server import proxy_config try: - models: Optional[List] = ( - await prisma_client.db.litellm_proxymodeltable.find_many() - ) + models: Optional[List] = await ModelRepository(prisma_client).table.find_many() except Exception: models = None # 2. process model table @@ -4088,7 +4107,7 @@ async def _rotate_master_key( # noqa: PLR0915 ) # 3. process config table try: - config = await prisma_client.db.litellm_config.find_many() + config = await ConfigRepository(prisma_client).table.find_many() except Exception: config = None @@ -4109,7 +4128,7 @@ async def _rotate_master_key( # noqa: PLR0915 ) if encrypted_env_vars: - await prisma_client.db.litellm_config.update( + await ConfigRepository(prisma_client).table.update( where={"param_name": "environment_variables"}, data={"param_value": prisma.Json(encrypted_env_vars)}, # type: ignore[attr-defined] ) @@ -4148,7 +4167,7 @@ async def _rotate_master_key( # noqa: PLR0915 # 5. process credentials table try: - credentials = await prisma_client.db.litellm_credentialstable.find_many() + credentials = await CredentialsRepository(prisma_client).table.find_many() except Exception: credentials = None if credentials: @@ -4171,7 +4190,7 @@ async def _rotate_master_key( # noqa: PLR0915 _cred_data["credential_info"] = prisma.Json( # type: ignore[attr-defined] _cred_data["credential_info"] ) - await prisma_client.db.litellm_credentialstable.update( + await CredentialsRepository(prisma_client).table.update( where={"credential_name": cred.credential_name}, data={ **_cred_data, @@ -4243,7 +4262,7 @@ async def _insert_deprecated_key( try: revoke_at = datetime.now(timezone.utc) + timedelta(seconds=grace_seconds) - await prisma_client.db.litellm_deprecatedverificationtoken.upsert( + await DeprecatedVerificationTokenRepository(prisma_client).table.upsert( where={"token": old_token_hash}, data={ "create": { @@ -4335,7 +4354,7 @@ async def _execute_virtual_key_regeneration( grace_period=data.grace_period if data else None, ) - updated_token = await prisma_client.db.litellm_verificationtoken.update( + updated_token = await VerificationTokenRepository(prisma_client).table.update( where={"token": hashed_api_key}, data=update_data, # type: ignore ) @@ -4530,7 +4549,7 @@ async def regenerate_key_fn( # noqa: PLR0915 else: hashed_api_key = hash_token(key) - _key_in_db = await prisma_client.db.litellm_verificationtoken.find_unique( + _key_in_db = await VerificationTokenRepository(prisma_client).table.find_unique( where={"token": hashed_api_key}, ) if _key_in_db is None: @@ -4719,7 +4738,7 @@ async def reset_key_spend_fn( else: hashed_api_key = hash_token(key) - _key_in_db = await prisma_client.db.litellm_verificationtoken.find_unique( + _key_in_db = await VerificationTokenRepository(prisma_client).table.find_unique( where={"token": hashed_api_key}, include={"litellm_budget_table": True}, ) @@ -4739,7 +4758,7 @@ async def reset_key_spend_fn( user_api_key_cache=user_api_key_cache, ) - updated_key = await prisma_client.db.litellm_verificationtoken.update( + updated_key = await VerificationTokenRepository(prisma_client).table.update( where={"token": hashed_api_key}, data={"spend": reset_to}, ) @@ -4792,11 +4811,11 @@ async def validate_key_list_check( param="user_id", code=status.HTTP_403_FORBIDDEN, ) - complete_user_info_db_obj: Optional[BaseModel] = ( - await prisma_client.db.litellm_usertable.find_unique( - where={"user_id": user_api_key_dict.user_id}, - include={"organization_memberships": True}, - ) + complete_user_info_db_obj: Optional[BaseModel] = await UserRepository( + prisma_client + ).table.find_unique( + where={"user_id": user_api_key_dict.user_id}, + include={"organization_memberships": True}, ) if complete_user_info_db_obj is None: @@ -4846,7 +4865,9 @@ async def validate_key_list_check( if key_hash: try: - key_info = await prisma_client.db.litellm_verificationtoken.find_unique( + key_info = await VerificationTokenRepository( + prisma_client + ).table.find_unique( where={"token": key_hash}, ) except Exception: @@ -4879,11 +4900,9 @@ async def _fetch_user_team_objects( if complete_user_info is None or not complete_user_info.teams: return [] - teams: Optional[List[BaseModel]] = ( - await prisma_client.db.litellm_teamtable.find_many( - where={"team_id": {"in": complete_user_info.teams}} - ) - ) + teams: Optional[List[BaseModel]] = await TeamRepository( + prisma_client + ).table.find_many(where={"team_id": {"in": complete_user_info.teams}}) if teams is None: return [] @@ -5160,7 +5179,7 @@ async def _apply_non_admin_alias_scope( # Look up the user's teams from the user table user_teams: List[str] = [] if user_api_key_dict.user_id: - user_row = await prisma_client.db.litellm_usertable.find_unique( + user_row = await UserRepository(prisma_client).table.find_unique( where={"user_id": user_api_key_dict.user_id} ) if user_row is not None: @@ -5548,7 +5567,7 @@ async def _list_key_helper( # Fetch keys with pagination if use_deleted_table: - keys = await prisma_client.db.litellm_deletedverificationtoken.find_many( + keys = await DeletedVerificationTokenRepository(prisma_client).table.find_many( where=where, # type: ignore skip=skip, # type: ignore take=size, # type: ignore @@ -5562,7 +5581,7 @@ async def _list_key_helper( ), ) else: - keys = await prisma_client.db.litellm_verificationtoken.find_many( + keys = await VerificationTokenRepository(prisma_client).table.find_many( where=where, # type: ignore skip=skip, # type: ignore take=size, # type: ignore @@ -5581,11 +5600,13 @@ async def _list_key_helper( # Get total count of keys if use_deleted_table: - total_count = await prisma_client.db.litellm_deletedverificationtoken.count( + total_count = await DeletedVerificationTokenRepository( + prisma_client + ).table.count( where=where # type: ignore ) else: - total_count = await prisma_client.db.litellm_verificationtoken.count( + total_count = await VerificationTokenRepository(prisma_client).table.count( where=where # type: ignore ) @@ -5601,7 +5622,7 @@ async def _list_key_helper( created_by_ids = [key.created_by for key in keys if key.created_by] all_ids = list(set(user_ids + created_by_ids)) # Remove duplicates if all_ids: - users = await prisma_client.db.litellm_usertable.find_many( + users = await UserRepository(prisma_client).table.find_many( where={"user_id": {"in": all_ids}} ) user_map = {user.user_id: user for user in users} @@ -5688,7 +5709,7 @@ async def _check_key_admin_access( return # Look up the target key to find its team - target_key_row = await prisma_client.db.litellm_verificationtoken.find_unique( + target_key_row = await VerificationTokenRepository(prisma_client).table.find_unique( where={"token": hashed_token} ) if target_key_row is None: @@ -5755,6 +5776,9 @@ async def block_key( Note: This is an admin-only endpoint. Only proxy admins, team admins, or org admins can block keys. """ + from litellm.proxy.management_helpers.audit_logs import ( + get_audit_log_changed_by, + ) from litellm.proxy.proxy_server import ( create_audit_log_for_update, hash_token, @@ -5763,9 +5787,6 @@ async def block_key( proxy_logging_obj, user_api_key_cache, ) - from litellm.proxy.management_helpers.audit_logs import ( - get_audit_log_changed_by, - ) if prisma_client is None: raise Exception("{}".format(CommonProxyErrors.db_not_connected_error.value)) @@ -5792,9 +5813,9 @@ async def block_key( ) # Check if the key exists before trying to block it - existing_record = await prisma_client.db.litellm_verificationtoken.find_unique( - where={"token": hashed_token} - ) + existing_record = await VerificationTokenRepository( + prisma_client + ).table.find_unique(where={"token": hashed_token}) if existing_record is None: raise ProxyException( message="Key not found.", @@ -5824,7 +5845,7 @@ async def block_key( ) ) - record = await prisma_client.db.litellm_verificationtoken.update( + record = await VerificationTokenRepository(prisma_client).table.update( where={"token": hashed_token}, data={"blocked": True} # type: ignore ) @@ -5869,6 +5890,9 @@ async def unblock_key( Note: This is an admin-only endpoint. Only proxy admins, team admins, or org admins can unblock keys. """ + from litellm.proxy.management_helpers.audit_logs import ( + get_audit_log_changed_by, + ) from litellm.proxy.proxy_server import ( create_audit_log_for_update, hash_token, @@ -5877,9 +5901,6 @@ async def unblock_key( proxy_logging_obj, user_api_key_cache, ) - from litellm.proxy.management_helpers.audit_logs import ( - get_audit_log_changed_by, - ) if prisma_client is None: raise Exception("{}".format(CommonProxyErrors.db_not_connected_error.value)) @@ -5906,9 +5927,9 @@ async def unblock_key( ) # Check if the key exists before trying to unblock it - existing_record = await prisma_client.db.litellm_verificationtoken.find_unique( - where={"token": hashed_token} - ) + existing_record = await VerificationTokenRepository( + prisma_client + ).table.find_unique(where={"token": hashed_token}) if existing_record is None: raise ProxyException( message="Key not found.", @@ -5938,7 +5959,7 @@ async def unblock_key( ) ) - record = await prisma_client.db.litellm_verificationtoken.update( + record = await VerificationTokenRepository(prisma_client).table.update( where={"token": hashed_token}, data={"blocked": False} # type: ignore ) @@ -6202,9 +6223,9 @@ async def _enforce_unique_key_alias( # Exclude the current key from the uniqueness check where_clause["NOT"] = {"token": existing_key_token} - existing_key = await prisma_client.db.litellm_verificationtoken.find_first( - where=where_clause - ) + existing_key = await VerificationTokenRepository( + prisma_client + ).table.find_first(where=where_clause) if existing_key is not None: raise ProxyException( message=f"Key with alias '{key_alias}' already exists. Unique key aliases across all keys are required.", diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 64e89ee40fd..c6c14c7a3e1 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -60,6 +60,10 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( encrypt_value_helper, ) from litellm.proxy.management_helpers.audit_logs import get_audit_log_changed_by +from litellm.repositories.table_repositories import ( + MCPServerRepository, + MCPUserCredentialsRepository, +) router = APIRouter(prefix="/v1/mcp", tags=["mcp"]) @@ -761,7 +765,7 @@ if MCP_AVAILABLE: # Get from DB if prisma_client is not None: try: - mcp_servers = await prisma_client.db.litellm_mcpservertable.find_many() + mcp_servers = await MCPServerRepository(prisma_client).table.find_many() for server in mcp_servers: if ( hasattr(server, "mcp_access_groups") @@ -998,10 +1002,10 @@ if MCP_AVAILABLE: if getattr(s, "is_byok", False) ] if byok_server_ids: - cred_rows = ( - await _byok_prisma_client.db.litellm_mcpusercredentials.find_many( - where={"user_id": user_id, "server_id": {"in": byok_server_ids}} - ) + cred_rows = await MCPUserCredentialsRepository( + _byok_prisma_client + ).table.find_many( + where={"user_id": user_id, "server_id": {"in": byok_server_ids}} ) cred_set = {r.server_id for r in cred_rows} for server in redacted_mcp_servers: diff --git a/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py b/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py index b05cfef5760..a8551f6333a 100644 --- a/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py @@ -19,6 +19,7 @@ from litellm.proxy.management_endpoints.model_management_endpoints import ( clear_cache, ) from litellm.proxy.utils import PrismaClient +from litellm.repositories.model_repository import ModelRepository from litellm.types.proxy.management_endpoints.model_management_endpoints import ( AccessGroupInfo, DeleteModelGroupResponse, @@ -95,7 +96,7 @@ async def update_deployments_with_access_group( verbose_proxy_logger.debug(f"Updating deployments for model_name: {model_name}") # Get all deployments with this model_name - deployments = await prisma_client.db.litellm_proxymodeltable.find_many( + deployments = await ModelRepository(prisma_client).table.find_many( where={"model_name": model_name} ) @@ -124,7 +125,7 @@ async def update_deployments_with_access_group( # Only update in DB if modified if was_modified: - await prisma_client.db.litellm_proxymodeltable.update( + await ModelRepository(prisma_client).table.update( where={"model_id": deployment.model_id}, data={"model_info": json.dumps(updated_model_info)}, ) @@ -152,7 +153,7 @@ async def update_specific_deployments_with_access_group( models_updated = 0 for model_id in model_ids: verbose_proxy_logger.debug(f"Updating specific deployment model_id: {model_id}") - deployment = await prisma_client.db.litellm_proxymodeltable.find_unique( + deployment = await ModelRepository(prisma_client).table.find_unique( where={"model_id": model_id} ) if deployment is None: @@ -168,7 +169,7 @@ async def update_specific_deployments_with_access_group( access_group=access_group, ) if was_modified: - await prisma_client.db.litellm_proxymodeltable.update( + await ModelRepository(prisma_client).table.update( where={"model_id": model_id}, data={"model_info": json.dumps(updated_model_info)}, ) @@ -215,7 +216,7 @@ async def get_all_access_groups_from_db( Dict[str, AccessGroupInfo]: Dictionary mapping access_group name to info """ # Get all deployments - deployments = await prisma_client.db.litellm_proxymodeltable.find_many() + deployments = await ModelRepository(prisma_client).table.find_many() # Build access group map access_group_map: Dict[str, Dict[str, Any]] = {} @@ -604,7 +605,7 @@ async def update_access_group( try: # Step 1: Remove access group from ALL DB deployments (skip config models) - all_deployments = await prisma_client.db.litellm_proxymodeltable.find_many() + all_deployments = await ModelRepository(prisma_client).table.find_many() for deployment in all_deployments: model_info = deployment.model_info or {} @@ -615,7 +616,7 @@ async def update_access_group( ) if was_modified: - await prisma_client.db.litellm_proxymodeltable.update( + await ModelRepository(prisma_client).table.update( where={"model_id": deployment.model_id}, data={"model_info": json.dumps(updated_model_info)}, ) @@ -722,7 +723,7 @@ async def delete_access_group( try: # Remove access group from all DB deployments (skip config models) - all_deployments = await prisma_client.db.litellm_proxymodeltable.find_many() + all_deployments = await ModelRepository(prisma_client).table.find_many() models_updated = 0 for deployment in all_deployments: @@ -734,7 +735,7 @@ async def delete_access_group( ) if was_modified: - await prisma_client.db.litellm_proxymodeltable.update( + await ModelRepository(prisma_client).table.update( where={"model_id": deployment.model_id}, data={"model_info": json.dumps(updated_model_info)}, ) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index e4ecda3fe31..566ef845333 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -48,6 +48,9 @@ from litellm.proxy.management_endpoints.team_endpoints import ( ) from litellm.proxy.management_helpers.audit_logs import create_object_audit_log from litellm.proxy.utils import PrismaClient +from litellm.repositories.model_repository import ModelRepository +from litellm.repositories.table_repositories import ModelTableRepository +from litellm.repositories.team_repository import TeamRepository from litellm.types.proxy.management_endpoints.model_management_endpoints import ( UpdateUsefulLinksRequest, ) @@ -86,7 +89,7 @@ async def get_db_model( ) -> Optional[Deployment]: db_model = cast( Optional[BaseModel], - await prisma_client.db.litellm_proxymodeltable.find_unique( + await ModelRepository(prisma_client).table.find_unique( where={"model_id": model_id} ), ) @@ -290,7 +293,7 @@ async def patch_model( update_data["updated_at"] = cast(str, get_utc_datetime()) # Perform partial update - updated_model = await prisma_client.db.litellm_proxymodeltable.update( + updated_model = await ModelRepository(prisma_client).table.update( where={"model_id": model_id}, data=update_data, ) @@ -362,7 +365,7 @@ async def _add_model_to_db( if model_params.model_info.id is not None: _data["model_id"] = model_params.model_info.id if should_create_model_in_db: - model_response = await prisma_client.db.litellm_proxymodeltable.create( + model_response = await ModelRepository(prisma_client).table.create( data=_data # type: ignore ) else: @@ -571,7 +574,7 @@ async def _get_team_deployments( team_id in model_info with Python-side filtering. """ prefix = f"model_name_{team_id}_" - response = await prisma_client.db.litellm_proxymodeltable.find_many( + response = await ModelRepository(prisma_client).table.find_many( where={ "model_name": {"startswith": prefix}, } @@ -828,7 +831,7 @@ class ModelManagementAuthChecks: detail={"error": CommonProxyErrors.not_premium_user.value}, ) - _existing_team_row = await prisma_client.db.litellm_teamtable.find_unique( + _existing_team_row = await TeamRepository(prisma_client).table.find_unique( where={"team_id": model_params.model_info.team_id} ) @@ -863,7 +866,7 @@ class ModelManagementAuthChecks: model_params.model_info is not None and model_params.model_info.team_id is not None ): - team_obj_row = await prisma_client.db.litellm_teamtable.find_unique( + team_obj_row = await TeamRepository(prisma_client).table.find_unique( where={"team_id": model_params.model_info.team_id} ) if team_obj_row is None: @@ -937,7 +940,7 @@ async def delete_model( }, ) - model_in_db = await prisma_client.db.litellm_proxymodeltable.find_unique( + model_in_db = await ModelRepository(prisma_client).table.find_unique( where={"model_id": model_info.id} ) if model_in_db is None: @@ -961,7 +964,7 @@ async def delete_model( - store keys separately """ # encrypt litellm params # - result = await prisma_client.db.litellm_proxymodeltable.delete( + result = await ModelRepository(prisma_client).table.delete( where={"model_id": model_info.id} ) @@ -1039,7 +1042,7 @@ async def delete_team_model_alias( Returns: - List of team id + model alias pairs that were removed """ - team_model_aliases = await prisma_client.db.litellm_modeltable.find_many( + team_model_aliases = await ModelTableRepository(prisma_client).table.find_many( include={"team": True} ) tasks = [] @@ -1056,7 +1059,7 @@ async def delete_team_model_alias( removed_model_aliases.append((team_model_alias.team.team_id, key)) del model_aliases[key] tasks.append( - prisma_client.db.litellm_modeltable.update( + ModelTableRepository(prisma_client).table.update( where={"id": id}, data={"model_aliases": json.dumps(model_aliases)}, ) @@ -1275,11 +1278,9 @@ async def update_model( if _model_id is None: raise Exception("model_info.id not provided") - _existing_litellm_params = ( - await prisma_client.db.litellm_proxymodeltable.find_unique( - where={"model_id": _model_id} - ) - ) + _existing_litellm_params = await ModelRepository( + prisma_client + ).table.find_unique(where={"model_id": _model_id}) if _existing_litellm_params is None: if ( @@ -1340,7 +1341,7 @@ async def update_model( "litellm_params": json.dumps(merged_dictionary), # type: ignore "updated_by": user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME, } - model_response = await prisma_client.db.litellm_proxymodeltable.update( + model_response = await ModelRepository(prisma_client).table.update( where={"model_id": _model_id}, data=_data, # type: ignore ) diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index 4d4ed53aaa8..99659121b27 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -40,6 +40,15 @@ from litellm.proxy.management_helpers.utils import ( management_endpoint_wrapper, ) from litellm.proxy.utils import PrismaClient +from litellm.repositories.budget_repository import BudgetRepository +from litellm.repositories.object_permission_repository import ObjectPermissionRepository +from litellm.repositories.organization_repository import OrganizationRepository +from litellm.repositories.table_repositories import OrganizationMembershipRepository +from litellm.repositories.team_repository import TeamRepository +from litellm.repositories.user_repository import UserRepository +from litellm.repositories.verification_token_repository import ( + VerificationTokenRepository, +) from litellm.types.proxy.management_endpoints.common_daily_activity import ( SpendAnalyticsPaginatedResponse, ) @@ -245,7 +254,7 @@ async def new_organization( if user_api_key_dict.user_id is not None: try: - user_object = await prisma_client.db.litellm_usertable.find_unique( + user_object = await UserRepository(prisma_client).table.find_unique( where={"user_id": user_api_key_dict.user_id} ) user_object_correct_type = LiteLLM_UserTable(**user_object.model_dump()) @@ -267,7 +276,7 @@ async def new_organization( new_budget = prisma_client.jsonify_object(budget_row.json(exclude_none=True)) - _budget = await prisma_client.db.litellm_budgettable.create( + _budget = await BudgetRepository(prisma_client).table.create( data={ **new_budget, # type: ignore "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name, @@ -323,7 +332,7 @@ async def new_organization( verbose_proxy_logger.info( f"new_organization_row: {json.dumps(new_organization_row, indent=2)}" ) - response = await prisma_client.db.litellm_organizationtable.create( + response = await OrganizationRepository(prisma_client).table.create( data={ **new_organization_row, # type: ignore }, @@ -372,9 +381,9 @@ async def get_organization_daily_activity( # Restrict non-proxy-admins to only organizations where they are org_admin if not _user_has_admin_view(user_api_key_dict): - memberships = await prisma_client.db.litellm_organizationmembership.find_many( - where={"user_id": user_api_key_dict.user_id} - ) + memberships = await OrganizationMembershipRepository( + prisma_client + ).table.find_many(where={"user_id": user_api_key_dict.user_id}) admin_org_ids = [ m.organization_id for m in memberships @@ -400,7 +409,7 @@ async def get_organization_daily_activity( where_condition = {} if org_ids_list: where_condition["organization_id"] = {"in": list(org_ids_list)} - org_aliases = await prisma_client.db.litellm_organizationtable.find_many( + org_aliases = await OrganizationRepository(prisma_client).table.find_many( where=where_condition ) org_alias_metadata = { @@ -439,10 +448,10 @@ async def _set_object_permission( return None if data.object_permission is not None: - created_object_permission = ( - await prisma_client.db.litellm_objectpermissiontable.create( - data=data.object_permission.model_dump(exclude_none=True), - ) + created_object_permission = await ObjectPermissionRepository( + prisma_client + ).table.create( + data=data.object_permission.model_dump(exclude_none=True), ) del data.object_permission return created_object_permission.object_permission_id @@ -525,10 +534,10 @@ async def update_organization( prisma_client=prisma_client, ) - existing_organization_row = ( - await prisma_client.db.litellm_organizationtable.find_unique( - where={"organization_id": data.organization_id}, - ) + existing_organization_row = await OrganizationRepository( + prisma_client + ).table.find_unique( + where={"organization_id": data.organization_id}, ) if existing_organization_row is None: @@ -574,7 +583,7 @@ async def update_organization( for field in LiteLLM_BudgetTable.model_fields.keys(): updated_organization_row.pop(field, None) - response = await prisma_client.db.litellm_organizationtable.update( + response = await OrganizationRepository(prisma_client).table.update( where={"organization_id": data.organization_id}, data=updated_organization_row, include={"members": True, "teams": True, "litellm_budget_table": True}, @@ -644,19 +653,19 @@ async def delete_organization( deleted_orgs = [] for organization_id in data.organization_ids: # delete all teams in the organization - await prisma_client.db.litellm_teamtable.delete_many( + await TeamRepository(prisma_client).table.delete_many( where={"organization_id": organization_id} ) # delete all members in the organization - await prisma_client.db.litellm_organizationmembership.delete_many( + await OrganizationMembershipRepository(prisma_client).table.delete_many( where={"organization_id": organization_id} ) # delete all keys in the organization - await prisma_client.db.litellm_verificationtoken.delete_many( + await VerificationTokenRepository(prisma_client).table.delete_many( where={"organization_id": organization_id} ) # delete the organization - deleted_org = await prisma_client.db.litellm_organizationtable.delete( + deleted_org = await OrganizationRepository(prisma_client).table.delete( where={"organization_id": organization_id}, include={"members": True, "teams": True, "litellm_budget_table": True}, ) @@ -732,17 +741,15 @@ async def list_organization( # if proxy admin or admin viewer - get all orgs (with optional filters) if _user_has_admin_view(user_api_key_dict): - response = await prisma_client.db.litellm_organizationtable.find_many( + response = await OrganizationRepository(prisma_client).table.find_many( where=where_conditions if where_conditions else None, include={"litellm_budget_table": True, "members": True, "teams": True}, ) # if internal user - get orgs they are a member of (with optional filters) else: - org_memberships = ( - await prisma_client.db.litellm_organizationmembership.find_many( - where={"user_id": user_api_key_dict.user_id} - ) - ) + org_memberships = await OrganizationMembershipRepository( + prisma_client + ).table.find_many(where={"user_id": user_api_key_dict.user_id}) membership_org_ids = [ membership.organization_id for membership in org_memberships ] @@ -756,20 +763,20 @@ async def list_organization( response = [] else: where_conditions["organization_id"] = org_id - response = ( - await prisma_client.db.litellm_organizationtable.find_many( - where=where_conditions, - include={ - "litellm_budget_table": True, - "members": True, - "teams": True, - }, - ) + response = await OrganizationRepository( + prisma_client + ).table.find_many( + where=where_conditions, + include={ + "litellm_budget_table": True, + "members": True, + "teams": True, + }, ) else: # Filter by membership and any additional filters where_conditions["organization_id"] = {"in": membership_org_ids} - response = await prisma_client.db.litellm_organizationtable.find_many( + response = await OrganizationRepository(prisma_client).table.find_many( where=where_conditions, include={ "litellm_budget_table": True, @@ -809,20 +816,20 @@ async def info_organization( prisma_client=prisma_client, ) - response: Optional[LiteLLM_OrganizationTableWithMembers] = ( - await prisma_client.db.litellm_organizationtable.find_unique( - where={"organization_id": organization_id}, - include={ - "litellm_budget_table": True, - "members": { - "include": { - "user": True, - } - }, - "teams": True, - "object_permission": True, + response: Optional[ + LiteLLM_OrganizationTableWithMembers + ] = await OrganizationRepository(prisma_client).table.find_unique( + where={"organization_id": organization_id}, + include={ + "litellm_budget_table": True, + "members": { + "include": { + "user": True, + } }, - ) + "teams": True, + "object_permission": True, + }, ) if response is None: @@ -868,7 +875,7 @@ async def deprecated_info_organization( prisma_client=prisma_client, ) - response = await prisma_client.db.litellm_organizationtable.find_many( + response = await OrganizationRepository(prisma_client).table.find_many( where={"organization_id": {"in": data.organizations}}, include={"litellm_budget_table": True}, ) @@ -945,11 +952,9 @@ async def organization_member_add( ) # Check if organization exists - existing_organization_row = ( - await prisma_client.db.litellm_organizationtable.find_unique( - where={"organization_id": data.organization_id} - ) - ) + existing_organization_row = await OrganizationRepository( + prisma_client + ).table.find_unique(where={"organization_id": data.organization_id}) if existing_organization_row is None: raise HTTPException( status_code=404, @@ -1012,11 +1017,9 @@ async def find_member_if_email( """ try: - existing_user_email_row: BaseModel = ( - await prisma_client.db.litellm_usertable.find_unique( - where={"user_email": user_email} - ) - ) + existing_user_email_row: BaseModel = await UserRepository( + prisma_client + ).table.find_unique(where={"user_email": user_email}) except Exception: raise HTTPException( status_code=400, @@ -1064,11 +1067,9 @@ async def organization_member_update( ) # Check if organization exists - existing_organization_row = ( - await prisma_client.db.litellm_organizationtable.find_unique( - where={"organization_id": data.organization_id} - ) - ) + existing_organization_row = await OrganizationRepository( + prisma_client + ).table.find_unique(where={"organization_id": data.organization_id}) if existing_organization_row is None: raise HTTPException( status_code=400, @@ -1085,15 +1086,15 @@ async def organization_member_update( data.user_id = existing_user_email_row.user_id try: - existing_organization_membership = ( - await prisma_client.db.litellm_organizationmembership.find_unique( - where={ - "user_id_organization_id": { - "user_id": data.user_id, - "organization_id": data.organization_id, - } + existing_organization_membership = await OrganizationMembershipRepository( + prisma_client + ).table.find_unique( + where={ + "user_id_organization_id": { + "user_id": data.user_id, + "organization_id": data.organization_id, } - ) + } ) except Exception as e: raise HTTPException( @@ -1114,7 +1115,7 @@ async def organization_member_update( # org-scoped operations. An org-admin of any org could otherwise # alter a PROXY_ADMIN user's per-org role, which has downstream # effects on admin UI filtering and scope derivation. - target_user_row = await prisma_client.db.litellm_usertable.find_unique( + target_user_row = await UserRepository(prisma_client).table.find_unique( where={"user_id": data.user_id} ) if target_user_row is not None and getattr( @@ -1136,7 +1137,7 @@ async def organization_member_update( # Update member role if data.role is not None: - await prisma_client.db.litellm_organizationmembership.update( + await OrganizationMembershipRepository(prisma_client).table.update( where={ "user_id_organization_id": { "user_id": data.user_id, @@ -1165,7 +1166,7 @@ async def organization_member_update( ) # update organization membership with new budget_id - await prisma_client.db.litellm_organizationmembership.update( + await OrganizationMembershipRepository(prisma_client).table.update( where={ "user_id_organization_id": { "user_id": data.user_id, @@ -1174,16 +1175,16 @@ async def organization_member_update( }, data={"budget_id": budget_id}, ) - final_organization_membership: Optional[BaseModel] = ( - await prisma_client.db.litellm_organizationmembership.find_unique( - where={ - "user_id_organization_id": { - "user_id": data.user_id, - "organization_id": data.organization_id, - } - }, - include={"litellm_budget_table": True}, - ) + final_organization_membership: Optional[ + BaseModel + ] = await OrganizationMembershipRepository(prisma_client).table.find_unique( + where={ + "user_id_organization_id": { + "user_id": data.user_id, + "organization_id": data.organization_id, + } + }, + include={"litellm_budget_table": True}, ) if final_organization_membership is None: @@ -1239,7 +1240,9 @@ async def organization_member_delete( ) data.user_id = existing_user_email_row.user_id - member_to_delete = await prisma_client.db.litellm_organizationmembership.delete( + member_to_delete = await OrganizationMembershipRepository( + prisma_client + ).table.delete( where={ "user_id_organization_id": { "user_id": data.user_id, @@ -1273,17 +1276,15 @@ async def add_member_to_organization( existing_user_email_row = None ## Check if user exists in LiteLLM_UserTable - user exists - either the user_id or user_email is in LiteLLM_UserTable if member.user_id is not None: - existing_user_id_row = await prisma_client.db.litellm_usertable.find_unique( - where={"user_id": member.user_id} - ) + existing_user_id_row = await UserRepository( + prisma_client + ).table.find_unique(where={"user_id": member.user_id}) if existing_user_id_row is None and member.user_email is not None: try: - existing_user_email_row = ( - await prisma_client.db.litellm_usertable.find_unique( - where={"user_email": member.user_email} - ) - ) + existing_user_email_row = await UserRepository( + prisma_client + ).table.find_unique(where={"user_email": member.user_email}) except Exception as e: raise ValueError( f"Potential NON-Existent or Duplicate user email in DB: Error finding a unique instance of user_email={member.user_email} in LiteLLM_UserTable.: {e}" @@ -1326,14 +1327,14 @@ async def add_member_to_organization( ) # Add user to organization - _organization_membership = ( - await prisma_client.db.litellm_organizationmembership.create( - data={ - "organization_id": organization_id, - "user_id": user_object.user_id, - "user_role": member.role, - } - ) + _organization_membership = await OrganizationMembershipRepository( + prisma_client + ).table.create( + data={ + "organization_id": organization_id, + "user_id": user_object.user_id, + "user_role": member.role, + } ) organization_membership = LiteLLM_OrganizationMembershipTable( **_organization_membership.model_dump() diff --git a/litellm/proxy/management_endpoints/scim/scim_transformations.py b/litellm/proxy/management_endpoints/scim/scim_transformations.py index 28fb87d9b3d..d1e00f87b69 100644 --- a/litellm/proxy/management_endpoints/scim/scim_transformations.py +++ b/litellm/proxy/management_endpoints/scim/scim_transformations.py @@ -6,6 +6,7 @@ from litellm.proxy._types import ( Member, NewUserResponse, ) +from litellm.repositories.team_repository import TeamRepository from litellm.types.proxy.management_endpoints.scim_v2 import * @@ -29,7 +30,7 @@ class ScimTransformations: # Get user's teams/groups groups = [] for team_id in user.teams or []: - team = await prisma_client.db.litellm_teamtable.find_unique( + team = await TeamRepository(prisma_client).table.find_unique( where={"team_id": team_id} ) if team: diff --git a/litellm/proxy/management_endpoints/scim/scim_v2.py b/litellm/proxy/management_endpoints/scim/scim_v2.py index 1f20764f837..0798d1a510d 100644 --- a/litellm/proxy/management_endpoints/scim/scim_v2.py +++ b/litellm/proxy/management_endpoints/scim/scim_v2.py @@ -22,7 +22,6 @@ from typing_extensions import TypedDict import litellm from litellm._logging import verbose_proxy_logger -from litellm.proxy.common_utils.http_parsing_utils import _safe_get_request_headers from litellm._uuid import uuid from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import ( @@ -41,6 +40,7 @@ from litellm.proxy._types import ( ) from litellm.proxy.auth.auth_checks import _delete_cache_key_object from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.http_parsing_utils import _safe_get_request_headers from litellm.proxy.management_endpoints.internal_user_endpoints import new_user from litellm.proxy.management_endpoints.scim.scim_transformations import ( ScimTransformations, @@ -51,6 +51,16 @@ from litellm.proxy.management_endpoints.team_endpoints import ( team_member_delete, ) from litellm.proxy.utils import _premium_user_check, handle_exception_on_proxy +from litellm.repositories.table_repositories import ( + InvitationLinkRepository, + OrganizationMembershipRepository, + TeamMembershipRepository, +) +from litellm.repositories.team_repository import TeamRepository +from litellm.repositories.user_repository import UserRepository +from litellm.repositories.verification_token_repository import ( + VerificationTokenRepository, +) from litellm.types.proxy.management_endpoints.scim_v2 import * @@ -74,7 +84,7 @@ class UserProvisionerHelpers: if not new_user_request.user_email: return None - existing_user = await prisma_client.db.litellm_usertable.find_first( + existing_user = await UserRepository(prisma_client).table.find_first( where={"user_email": new_user_request.user_email} ) @@ -82,7 +92,7 @@ class UserProvisionerHelpers: return None # Update the user - updated_user = await prisma_client.db.litellm_usertable.update( + updated_user = await UserRepository(prisma_client).table.update( where={"user_id": existing_user.user_id}, data={ "user_id": new_user_request.user_id, @@ -139,7 +149,7 @@ async def _check_user_exists(user_id: str): """Check if user exists and return user, raise 404 if not found.""" prisma_client = await _get_prisma_client_or_raise_exception() - user = await prisma_client.db.litellm_usertable.find_unique( + user = await UserRepository(prisma_client).table.find_unique( where={"user_id": user_id} ) @@ -155,7 +165,7 @@ async def _check_team_exists(team_id: str): """Check if team exists and return team, raise 404 if not found.""" prisma_client = await _get_prisma_client_or_raise_exception() - team = await prisma_client.db.litellm_teamtable.find_unique( + team = await TeamRepository(prisma_client).table.find_unique( where={"team_id": team_id} ) @@ -268,7 +278,7 @@ async def _extract_group_member_ids(group: SCIMGroup) -> GroupMemberExtractionRe ) # Check if user exists - user = await prisma_client.db.litellm_usertable.find_unique( + user = await UserRepository(prisma_client).table.find_unique( where={"user_id": user_id} ) @@ -310,7 +320,7 @@ async def _get_team_members_display(member_ids: List[str]) -> List[SCIMMember]: members: List[SCIMMember] = [] for member_id in member_ids: - user = await prisma_client.db.litellm_usertable.find_unique( + user = await UserRepository(prisma_client).table.find_unique( where={"user_id": member_id} ) if user: @@ -367,7 +377,7 @@ async def _set_user_keys_blocked(user_id: str, blocked: bool) -> int: # `blocked` is a nullable column with no default, so existing rows # typically hold NULL; treat NULL as "not blocked" since SQL equality # on NULL would otherwise silently skip them. - candidates = await prisma_client.db.litellm_verificationtoken.find_many( + candidates = await VerificationTokenRepository(prisma_client).table.find_many( where={ "user_id": user_id, "OR": [{"blocked": False}, {"blocked": None}], @@ -375,7 +385,7 @@ async def _set_user_keys_blocked(user_id: str, blocked: bool) -> int: ) affected_keys = candidates else: - candidates = await prisma_client.db.litellm_verificationtoken.find_many( + candidates = await VerificationTokenRepository(prisma_client).table.find_many( where={"user_id": user_id, "blocked": True}, ) affected_keys = [k for k in candidates if _key_was_scim_blocked(k.metadata)] @@ -395,7 +405,7 @@ async def _set_user_keys_blocked(user_id: str, blocked: bool) -> int: for k, v in current_metadata.items() if k != SCIM_BLOCKED_METADATA_KEY } - await prisma_client.db.litellm_verificationtoken.update( + await VerificationTokenRepository(prisma_client).table.update( where={"token": key_row.token}, data={"blocked": blocked, "metadata": safe_dumps(new_metadata)}, ) @@ -423,7 +433,7 @@ async def _delete_rows_referencing_user(prisma_client: Any, *, user_id: str) -> the user delete with an FK constraint violation (e.g. ``LiteLLM_InvitationLink_user_id_fkey``). """ - await prisma_client.db.litellm_invitationlink.delete_many( + await InvitationLinkRepository(prisma_client).table.delete_many( where={ "OR": [ {"user_id": user_id}, @@ -432,10 +442,10 @@ async def _delete_rows_referencing_user(prisma_client: Any, *, user_id: str) -> ] } ) - await prisma_client.db.litellm_organizationmembership.delete_many( + await OrganizationMembershipRepository(prisma_client).table.delete_many( where={"user_id": user_id} ) - await prisma_client.db.litellm_teammembership.delete_many( + await TeamMembershipRepository(prisma_client).table.delete_many( where={"user_id": user_id} ) @@ -897,17 +907,17 @@ async def get_users( where_conditions["user_email"] = filter_value # Get users from database - users: List[LiteLLM_UserTable] = ( - await prisma_client.db.litellm_usertable.find_many( - where=where_conditions, - skip=(startIndex - 1), - take=count, - order={"created_at": "desc"}, - ) + users: List[LiteLLM_UserTable] = await UserRepository( + prisma_client + ).table.find_many( + where=where_conditions, + skip=(startIndex - 1), + take=count, + order={"created_at": "desc"}, ) # Get total count for pagination - total_count = await prisma_client.db.litellm_usertable.count( + total_count = await UserRepository(prisma_client).table.count( where=where_conditions ) @@ -975,7 +985,7 @@ async def create_user( # Check if user already exists if user.userName: - existing_user = await prisma_client.db.litellm_usertable.find_unique( + existing_user = await UserRepository(prisma_client).table.find_unique( where={"user_id": user.userName} ) if existing_user: @@ -1094,7 +1104,7 @@ async def update_user( "metadata": safe_dumps(metadata), } - updated_user = await prisma_client.db.litellm_usertable.update( + updated_user = await UserRepository(prisma_client).table.update( where={"user_id": user_id}, data=update_data, ) @@ -1137,7 +1147,7 @@ async def delete_user( teams = [] if existing_user.teams: for team_id in existing_user.teams: - team = await prisma_client.db.litellm_teamtable.find_unique( + team = await TeamRepository(prisma_client).table.find_unique( where={"team_id": team_id} ) if team: @@ -1148,7 +1158,7 @@ async def delete_user( current_members = team.members or [] if user_id in current_members: new_members = [m for m in current_members if m != user_id] - await prisma_client.db.litellm_teamtable.update( + await TeamRepository(prisma_client).table.update( where={"team_id": team.team_id}, data={"members": new_members} ) @@ -1157,7 +1167,7 @@ async def delete_user( await _delete_rows_referencing_user(prisma_client, user_id=user_id) # Delete user - await prisma_client.db.litellm_usertable.delete(where={"user_id": user_id}) + await UserRepository(prisma_client).table.delete(where={"user_id": user_id}) return Response(status_code=204) except Exception as e: @@ -1413,7 +1423,7 @@ async def patch_user( update_data["metadata"] = safe_dumps(update_data["metadata"]) - updated_user = await prisma_client.db.litellm_usertable.update( + updated_user = await UserRepository(prisma_client).table.update( where={"user_id": user_id}, data=update_data, ) @@ -1465,7 +1475,7 @@ async def get_groups( where_conditions["team_alias"] = team_alias # Get teams from database - teams = await prisma_client.db.litellm_teamtable.find_many( + teams = await TeamRepository(prisma_client).table.find_many( where=where_conditions, skip=(startIndex - 1), take=count, @@ -1473,7 +1483,7 @@ async def get_groups( ) # Get total count for pagination - total_count = await prisma_client.db.litellm_teamtable.count( + total_count = await TeamRepository(prisma_client).table.count( where=where_conditions ) @@ -1561,7 +1571,7 @@ async def create_group( team_id = group.id or group.externalId or str(uuid.uuid4()) # Check if team already exists - existing_team = await prisma_client.db.litellm_teamtable.find_unique( + existing_team = await TeamRepository(prisma_client).table.find_unique( where={"team_id": team_id} ) @@ -1638,7 +1648,7 @@ async def update_group( } # Update team in database - updated_team = await prisma_client.db.litellm_teamtable.update( + updated_team = await TeamRepository(prisma_client).table.update( where={"team_id": group_id}, data=update_data, ) @@ -1683,19 +1693,19 @@ async def delete_group( # For each member, remove this team from their teams list for member_id in existing_team.members or []: - user = await prisma_client.db.litellm_usertable.find_unique( + user = await UserRepository(prisma_client).table.find_unique( where={"user_id": member_id} ) if user: current_teams = user.teams or [] if group_id in current_teams: new_teams = [t for t in current_teams if t != group_id] - await prisma_client.db.litellm_usertable.update( + await UserRepository(prisma_client).table.update( where={"user_id": member_id}, data={"teams": new_teams} ) # Delete team - await prisma_client.db.litellm_teamtable.delete(where={"team_id": group_id}) + await TeamRepository(prisma_client).table.delete(where={"team_id": group_id}) return Response(status_code=204) @@ -1748,7 +1758,7 @@ async def _process_group_patch_operations( detail={"error": "Invalid member: user ID cannot be empty."}, ) - user = await prisma_client.db.litellm_usertable.find_unique( + user = await UserRepository(prisma_client).table.find_unique( where={"user_id": member_id} ) if user: @@ -1805,7 +1815,7 @@ async def _apply_group_patch_updates( update_data["members"] = list(final_members) # Update team in database - updated_team = await prisma_client.db.litellm_teamtable.update( + updated_team = await TeamRepository(prisma_client).table.update( where={"team_id": group_id}, data=update_data, ) @@ -1877,7 +1887,7 @@ async def patch_group( # Refresh team data from database to get the latest state after concurrent updates # This prevents race conditions when multiple PATCH requests come in simultaneously - refreshed_team = await prisma_client.db.litellm_teamtable.find_unique( + refreshed_team = await TeamRepository(prisma_client).table.find_unique( where={"team_id": group_id} ) if refreshed_team: @@ -1894,7 +1904,7 @@ async def patch_group( await _handle_group_membership_changes(group_id, current_members, final_members) # Refresh team one more time to get final state after membership changes - final_team = await prisma_client.db.litellm_teamtable.find_unique( + final_team = await TeamRepository(prisma_client).table.find_unique( where={"team_id": group_id} ) if final_team: diff --git a/litellm/proxy/management_endpoints/tag_management_endpoints.py b/litellm/proxy/management_endpoints/tag_management_endpoints.py index 49d9b67a28a..f0bb8bdb5ff 100644 --- a/litellm/proxy/management_endpoints/tag_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tag_management_endpoints.py @@ -25,6 +25,14 @@ from litellm.proxy.management_endpoints.common_daily_activity import ( get_daily_activity, ) from litellm.proxy.management_helpers.utils import handle_budget_for_entity +from litellm.repositories.model_repository import ModelRepository +from litellm.repositories.table_repositories import ( + DailyTagSpendRepository, + TagRepository, +) +from litellm.repositories.verification_token_repository import ( + VerificationTokenRepository, +) from litellm.types.tag_management import ( TagConfig, TagDeleteRequest, @@ -56,7 +64,7 @@ async def _get_internal_user_api_keys( if user_id is None: return sorted(user_api_keys) - key_records = await prisma_client.db.litellm_verificationtoken.find_many( + key_records = await VerificationTokenRepository(prisma_client).table.find_many( where={"user_id": user_id}, select={"token": True}, ) @@ -109,7 +117,7 @@ async def _get_tag_daily_activity_api_key_filter( async def _get_model_names(prisma_client, model_ids: list) -> Dict[str, str]: """Helper function to get model names from model IDs""" try: - models = await prisma_client.db.litellm_proxymodeltable.find_many( + models = await ModelRepository(prisma_client).table.find_many( where={"model_id": {"in": model_ids}} ) return {model.model_id: model.model_name for model in models} @@ -189,7 +197,7 @@ async def new_tag( ) try: # Check if tag already exists - existing_tag = await prisma_client.db.litellm_tagtable.find_unique( + existing_tag = await TagRepository(prisma_client).table.find_unique( where={"tag_name": tag.name} ) if existing_tag is not None: @@ -210,7 +218,7 @@ async def new_tag( model_info = await _get_model_names(prisma_client, tag.models or []) # Create new tag in database - new_tag_record = await prisma_client.db.litellm_tagtable.create( + new_tag_record = await TagRepository(prisma_client).table.create( data={ "tag_name": tag.name, "description": tag.description, @@ -267,7 +275,7 @@ async def _add_tag_to_deployment(deployment: "Deployment", tag: str): try: # Get current model from database to preserve encrypted fields - db_model = await prisma_client.db.litellm_proxymodeltable.find_unique( + db_model = await ModelRepository(prisma_client).table.find_unique( where={"model_id": deployment.model_info.id} ) @@ -292,7 +300,7 @@ async def _add_tag_to_deployment(deployment: "Deployment", tag: str): existing_params["tags"].append(tag) # Update database with modified params (keeps encrypted fields encrypted) - await prisma_client.db.litellm_proxymodeltable.update( + await ModelRepository(prisma_client).table.update( where={"model_id": deployment.model_info.id}, data={"litellm_params": json.dumps(existing_params)}, ) @@ -335,7 +343,7 @@ async def update_tag( try: # Check if tag exists - existing_tag = await prisma_client.db.litellm_tagtable.find_unique( + existing_tag = await TagRepository(prisma_client).table.find_unique( where={"tag_name": tag.name} ) if existing_tag is None: @@ -367,7 +375,7 @@ async def update_tag( update_data["budget_id"] = budget_id # Update tag in database - updated_tag_record = await prisma_client.db.litellm_tagtable.update( + updated_tag_record = await TagRepository(prisma_client).table.update( where={"tag_name": tag.name}, data=update_data, ) @@ -414,7 +422,7 @@ async def info_tag( try: # Query tags from database with budget info - tag_records = await prisma_client.db.litellm_tagtable.find_many( + tag_records = await TagRepository(prisma_client).table.find_many( where={"tag_name": {"in": data.names}}, include={"litellm_budget_table": True}, ) @@ -535,7 +543,7 @@ async def list_tags( if start_date is not None and end_date is not None: dynamic_tag_where["date"] = {"gte": start_date, "lte": end_date} - dynamic_tag_rows = await prisma_client.db.litellm_dailytagspend.group_by( + dynamic_tag_rows = await DailyTagSpendRepository(prisma_client).table.group_by( by=["tag"], where=dynamic_tag_where, min={"created_at": True}, @@ -551,7 +559,7 @@ async def list_tags( ) ## QUERY STORED TAGS ## - tag_records = await prisma_client.db.litellm_tagtable.find_many( + tag_records = await TagRepository(prisma_client).table.find_many( where=stored_tag_where, include={"litellm_budget_table": True}, ) @@ -626,14 +634,14 @@ async def delete_tag( try: # Check if tag exists - existing_tag = await prisma_client.db.litellm_tagtable.find_unique( + existing_tag = await TagRepository(prisma_client).table.find_unique( where={"tag_name": data.name} ) if existing_tag is None: raise HTTPException(status_code=404, detail=f"Tag {data.name} not found") # Delete tag from database - await prisma_client.db.litellm_tagtable.delete(where={"tag_name": data.name}) + await TagRepository(prisma_client).table.delete(where={"tag_name": data.name}) return {"message": f"Tag {data.name} deleted successfully"} except Exception as e: diff --git a/litellm/proxy/management_endpoints/team_callback_endpoints.py b/litellm/proxy/management_endpoints/team_callback_endpoints.py index 63b56425b0e..0c11507697d 100644 --- a/litellm/proxy/management_endpoints/team_callback_endpoints.py +++ b/litellm/proxy/management_endpoints/team_callback_endpoints.py @@ -30,6 +30,7 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.callback_utils import encrypt_callback_vars from litellm.proxy.management_endpoints.team_endpoints import _verify_team_access from litellm.proxy.management_helpers.utils import management_endpoint_wrapper +from litellm.repositories.team_repository import TeamRepository router = APIRouter() @@ -249,7 +250,7 @@ async def add_team_callbacks( team_metadata = encrypt_callback_vars(team_metadata) team_metadata_json = json.dumps(team_metadata) # update team_metadata - new_team_row = await prisma_client.db.litellm_teamtable.update( + new_team_row = await TeamRepository(prisma_client).table.update( where={"team_id": team_id}, data={"metadata": team_metadata_json} # type: ignore ) @@ -353,7 +354,7 @@ async def disable_team_logging( team_metadata_json = json.dumps(team_metadata) # Update team in database - updated_team = await prisma_client.db.litellm_teamtable.update( + updated_team = await TeamRepository(prisma_client).table.update( where={"team_id": team_id}, data={"metadata": team_metadata_json} # type: ignore ) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index a3ad7a9ea8b..7a784ee4622 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -10,8 +10,8 @@ All /team management endpoints """ import asyncio -import math import json +import math import traceback from datetime import datetime, timezone from typing import Any, Dict, List, Optional, Tuple, Union, cast @@ -102,6 +102,20 @@ from litellm.proxy.management_helpers.utils import ( management_endpoint_wrapper, ) from litellm.proxy.utils import PrismaClient, handle_exception_on_proxy +from litellm.repositories.budget_repository import BudgetRepository +from litellm.repositories.organization_repository import OrganizationRepository +from litellm.repositories.table_repositories import ( + AccessGroupRepository, + DeletedTeamRepository, + ModelTableRepository, + OrganizationMembershipRepository, + TeamMembershipRepository, +) +from litellm.repositories.team_repository import TeamRepository +from litellm.repositories.user_repository import UserRepository +from litellm.repositories.verification_token_repository import ( + VerificationTokenRepository, +) from litellm.router import Router from litellm.types.proxy.management_endpoints.common_daily_activity import ( SpendAnalyticsPaginatedResponse, @@ -365,9 +379,9 @@ class TeamMemberBudgetHandler: return # Batch-fetch existing memberships for this team (avoids N+1 queries) - existing_memberships = await prisma_client.db.litellm_teammembership.find_many( - where={"team_id": team_id} - ) + existing_memberships = await TeamMembershipRepository( + prisma_client + ).table.find_many(where={"team_id": team_id}) existing_user_ids = {m.user_id for m in existing_memberships} # Identify members with no existing membership row. @@ -386,7 +400,7 @@ class TeamMemberBudgetHandler: ) if missing: - await prisma_client.db.litellm_teammembership.create_many( + await TeamMembershipRepository(prisma_client).table.create_many( data=missing, skip_duplicates=True, # safety net against concurrent races ) @@ -400,7 +414,7 @@ class TeamMemberBudgetHandler: # Heal existing membership rows that predate the team_member_budget # configuration: populate budget_id where it is currently NULL. # Rows with an explicit budget_id (per-member override) are left alone. - updated = await prisma_client.db.litellm_teammembership.update_many( + updated = await TeamMembershipRepository(prisma_client).table.update_many( where={"team_id": team_id, "budget_id": None}, data={"budget_id": team_member_budget_id}, ) @@ -456,7 +470,7 @@ async def get_all_team_memberships( # else: # where_obj = {"user_id": str(user_id), "team_id": {"in": team_id}} - team_memberships = await prisma_client.db.litellm_teammembership.find_many( + team_memberships = await TeamMembershipRepository(prisma_client).table.find_many( where=where_obj, include={"litellm_budget_table": True}, ) @@ -739,7 +753,7 @@ async def _check_org_team_limits( # calculate allocated tpm/rpm limit # check if specified tpm/rpm limit is greater than allocated tpm/rpm limit - teams = await prisma_client.db.litellm_teamtable.find_many( + teams = await TeamRepository(prisma_client).table.find_many( where={"organization_id": org_table.organization_id}, ) @@ -931,6 +945,9 @@ async def new_team( # noqa: PLR0915 ``` """ try: + from litellm.proxy.management_helpers.audit_logs import ( + get_audit_log_changed_by, + ) from litellm.proxy.proxy_server import ( _license_check, create_audit_log_for_update, @@ -938,9 +955,6 @@ async def new_team( # noqa: PLR0915 prisma_client, user_api_key_cache, ) - from litellm.proxy.management_helpers.audit_logs import ( - get_audit_log_changed_by, - ) if prisma_client is None: raise HTTPException(status_code=500, detail={"error": "No db connected"}) @@ -986,7 +1000,7 @@ async def new_team( # noqa: PLR0915 ) # Check if license is over limit - total_teams = await prisma_client.db.litellm_teamtable.count() + total_teams = await TeamRepository(prisma_client).table.count() if total_teams and _license_check.is_team_count_over_limit( team_count=total_teams ): @@ -1092,7 +1106,7 @@ async def new_team( # noqa: PLR0915 created_by=user_api_key_dict.user_id or litellm_proxy_admin_name, updated_by=user_api_key_dict.user_id or litellm_proxy_admin_name, ) - model_dict = await prisma_client.db.litellm_modeltable.create( + model_dict = await ModelTableRepository(prisma_client).table.create( {**litellm_modeltable.json(exclude_none=True)} # type: ignore ) # type: ignore @@ -1195,7 +1209,7 @@ async def new_team( # noqa: PLR0915 db_data=complete_team_data_dict ) - team_row: LiteLLM_TeamTable = await prisma_client.db.litellm_teamtable.create( + team_row: LiteLLM_TeamTable = await TeamRepository(prisma_client).table.create( data=complete_team_data_dict, include={"litellm_model_table": True}, # type: ignore ) @@ -1315,11 +1329,11 @@ async def _update_model_table( updated_by=user_api_key_dict.user_id or litellm_proxy_admin_name, ) if model_id is None: - model_dict = await prisma_client.db.litellm_modeltable.create( + model_dict = await ModelTableRepository(prisma_client).table.create( data={**litellm_modeltable.json(exclude_none=True)} # type: ignore ) else: - model_dict = await prisma_client.db.litellm_modeltable.upsert( + model_dict = await ModelTableRepository(prisma_client).table.upsert( where={"id": model_id}, data={ "update": {**litellm_modeltable.json(exclude_none=True)}, # type: ignore @@ -1400,7 +1414,7 @@ async def fetch_and_validate_organization( status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value} ) - organization_row = await prisma_client.db.litellm_organizationtable.find_unique( + organization_row = await OrganizationRepository(prisma_client).table.find_unique( where={"organization_id": organization_id}, include={"litellm_budget_table": True, "members": True, "teams": True}, ) @@ -1669,7 +1683,7 @@ async def update_team( # noqa: PLR0915 }, ) - existing_team_row = await prisma_client.db.litellm_teamtable.find_unique( + existing_team_row = await TeamRepository(prisma_client).table.find_unique( where={"team_id": data.team_id} ) @@ -1738,7 +1752,9 @@ async def update_team( # noqa: PLR0915 ): # Is the caller org_admin of the destination org? caller_memberships = ( - await prisma_client.db.litellm_organizationmembership.find_many( + await OrganizationMembershipRepository( + prisma_client + ).table.find_many( where={ "user_id": user_api_key_dict.user_id, "organization_id": data.organization_id, @@ -1885,18 +1901,18 @@ async def update_team( # noqa: PLR0915 updated_kv["router_settings"] = safe_dumps(updated_kv["router_settings"]) updated_kv = prisma_client.jsonify_team_object(db_data=updated_kv) - team_row: Optional[LiteLLM_TeamTable] = ( - await prisma_client.db.litellm_teamtable.update( - where={"team_id": data.team_id}, - data=updated_kv, - # `object_permission` is included so `_refresh_cached_team` - # doesn't write a cached team with the relation nulled out — - # see team_model_add for the full rationale. - include={ - "litellm_model_table": True, - "object_permission": True, - }, # type: ignore - ) + team_row: Optional[LiteLLM_TeamTable] = await TeamRepository( + prisma_client + ).table.update( + where={"team_id": data.team_id}, + data=updated_kv, + # `object_permission` is included so `_refresh_cached_team` + # doesn't write a cached team with the relation nulled out — + # see team_model_add for the full rationale. + include={ + "litellm_model_table": True, + "object_permission": True, + }, # type: ignore ) if team_row is None or team_row.team_id is None: @@ -2306,7 +2322,7 @@ async def _add_team_members_to_team( # ADD MEMBER TO TEAM _db_team_members = [m.model_dump() for m in complete_team_data.members_with_roles] - updated_team = await prisma_client.db.litellm_teamtable.update( + updated_team = await TeamRepository(prisma_client).table.update( where={"team_id": data.team_id}, data={"members_with_roles": json.dumps(_db_team_members)}, # type: ignore ) @@ -2377,7 +2393,7 @@ async def _validate_and_populate_member_user_info( # Case 2: Only user_email provided - populate user_id from DB if member.user_email is not None and member.user_id is None: - user_by_email = await prisma_client.db.litellm_usertable.find_first( + user_by_email = await UserRepository(prisma_client).table.find_first( where={"user_email": {"equals": member.user_email, "mode": "insensitive"}} ) @@ -2410,7 +2426,7 @@ async def _validate_and_populate_member_user_info( # Case 3: Only user_id provided - populate user_email from DB if user exists if member.user_id is not None and member.user_email is None: - user_by_id = await prisma_client.db.litellm_usertable.find_unique( + user_by_id = await UserRepository(prisma_client).table.find_unique( where={"user_id": member.user_id} ) @@ -2608,7 +2624,7 @@ async def team_member_delete( detail={"error": "Either user_id or user_email needs to be passed in"}, ) - _existing_team_row = await prisma_client.db.litellm_teamtable.find_unique( + _existing_team_row = await TeamRepository(prisma_client).table.find_unique( where={"team_id": data.team_id} ) @@ -2652,7 +2668,7 @@ async def team_member_delete( _db_new_team_members: List[dict] = [m.model_dump() for m in new_team_members] - _ = await prisma_client.db.litellm_teamtable.update( + _ = await TeamRepository(prisma_client).table.update( where={ "team_id": data.team_id, }, @@ -2666,7 +2682,7 @@ async def team_member_delete( key_val["user_id"] = data.user_id elif data.user_email is not None: key_val["user_email"] = data.user_email - existing_user_rows = await prisma_client.db.litellm_usertable.find_many( + existing_user_rows = await UserRepository(prisma_client).table.find_many( where=key_val # type: ignore ) @@ -2678,7 +2694,7 @@ async def team_member_delete( if data.team_id in existing_user.teams: team_list = existing_user.teams team_list.remove(data.team_id) - await prisma_client.db.litellm_usertable.update( + await UserRepository(prisma_client).table.update( where={ "user_id": existing_user.user_id, }, @@ -2695,7 +2711,7 @@ async def team_member_delete( user_ids_to_delete.add(existing_user.user_id) for _uid in user_ids_to_delete: - await prisma_client.db.litellm_teammembership.delete_many( + await TeamMembershipRepository(prisma_client).table.delete_many( where={"team_id": data.team_id, "user_id": _uid} ) @@ -2706,13 +2722,13 @@ async def team_member_delete( ) # Fetch keys before deletion to persist them - keys_to_delete: List[LiteLLM_VerificationToken] = ( - await prisma_client.db.litellm_verificationtoken.find_many( - where={ - "user_id": {"in": list(user_ids_to_delete)}, - "team_id": data.team_id, - } - ) + keys_to_delete: List[ + LiteLLM_VerificationToken + ] = await VerificationTokenRepository(prisma_client).table.find_many( + where={ + "user_id": {"in": list(user_ids_to_delete)}, + "team_id": data.team_id, + } ) if keys_to_delete: @@ -2723,7 +2739,7 @@ async def team_member_delete( litellm_changed_by=None, ) - await prisma_client.db.litellm_verificationtoken.delete_many( + await VerificationTokenRepository(prisma_client).table.delete_many( where={ "user_id": {"in": list(user_ids_to_delete)}, "team_id": data.team_id, @@ -2818,7 +2834,7 @@ async def team_member_update( _validate_budget_duration(data.budget_duration) - _existing_team_row = await prisma_client.db.litellm_teamtable.find_unique( + _existing_team_row = await TeamRepository(prisma_client).table.find_unique( where={"team_id": data.team_id} ) @@ -2921,7 +2937,7 @@ async def team_member_update( team_table.members_with_roles = team_members _db_team_members: List[dict] = [m.model_dump() for m in team_members] - await prisma_client.db.litellm_teamtable.update( + await TeamRepository(prisma_client).table.update( where={"team_id": data.team_id}, data={"members_with_roles": json.dumps(_db_team_members)}, # type: ignore ) @@ -3052,7 +3068,7 @@ async def bulk_team_member_add( }, ) # get all users from the database - all_users_in_db = await prisma_client.db.litellm_usertable.find_many( + all_users_in_db = await UserRepository(prisma_client).table.find_many( order={"created_at": "desc"} ) data.members = [ @@ -3153,14 +3169,14 @@ async def delete_team( }' ``` """ + from litellm.proxy.management_helpers.audit_logs import ( + get_audit_log_changed_by, + ) from litellm.proxy.proxy_server import ( create_audit_log_for_update, litellm_proxy_admin_name, prisma_client, ) - from litellm.proxy.management_helpers.audit_logs import ( - get_audit_log_changed_by, - ) if prisma_client is None: raise HTTPException(status_code=500, detail={"error": "No db connected"}) @@ -3172,11 +3188,9 @@ async def delete_team( team_rows: List[LiteLLM_TeamTable] = [] for team_id in data.team_ids: try: - team_row_base: Optional[BaseModel] = ( - await prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": team_id} - ) - ) + team_row_base: Optional[BaseModel] = await TeamRepository( + prisma_client + ).table.find_unique(where={"team_id": team_id}) if team_row_base is None: raise Exception except Exception: @@ -3243,11 +3257,9 @@ async def delete_team( _persist_deleted_verification_tokens, ) - keys_to_delete: List[LiteLLM_VerificationToken] = ( - await prisma_client.db.litellm_verificationtoken.find_many( - where={"team_id": {"in": data.team_ids}} - ) - ) + keys_to_delete: List[LiteLLM_VerificationToken] = await VerificationTokenRepository( + prisma_client + ).table.find_many(where={"team_id": {"in": data.team_ids}}) if keys_to_delete: await _persist_deleted_verification_tokens( @@ -3338,7 +3350,7 @@ async def _save_deleted_team_records( """Save deleted team records to the database.""" if not records: return - await prisma_client.db.litellm_deletedteamtable.create_many(data=records) + await DeletedTeamRepository(prisma_client).table.create_many(data=records) async def _persist_deleted_team_records( @@ -3420,7 +3432,7 @@ async def _add_team_member_budget_table( team_info_response_object: TeamInfoResponseObjectTeamTable, ) -> TeamInfoResponseObjectTeamTable: try: - team_budget = await prisma_client.db.litellm_budgettable.find_unique( + team_budget = await BudgetRepository(prisma_client).table.find_unique( where={"budget_id": team_member_budget_id} ) team_info_response_object.team_member_budget_table = team_budget @@ -3489,11 +3501,11 @@ async def team_info( ) try: - team_info: Optional[BaseModel] = ( - await prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": team_id}, - include={"object_permission": True}, - ) + team_info: Optional[BaseModel] = await TeamRepository( + prisma_client + ).table.find_unique( + where={"team_id": team_id}, + include={"object_permission": True}, ) if team_info is None: raise Exception @@ -3749,7 +3761,7 @@ async def block_team( if prisma_client is None: raise Exception("No DB Connected.") - existing_team = await prisma_client.db.litellm_teamtable.find_unique( + existing_team = await TeamRepository(prisma_client).table.find_unique( where={"team_id": data.team_id} ) if existing_team is None: @@ -3764,7 +3776,7 @@ async def block_team( user_api_key_dict=user_api_key_dict, ) - record = await prisma_client.db.litellm_teamtable.update( + record = await TeamRepository(prisma_client).table.update( where={"team_id": data.team_id}, data={"blocked": True} # type: ignore ) @@ -3801,7 +3813,7 @@ async def unblock_team( if prisma_client is None: raise Exception("No DB Connected.") - existing_team = await prisma_client.db.litellm_teamtable.find_unique( + existing_team = await TeamRepository(prisma_client).table.find_unique( where={"team_id": data.team_id} ) if existing_team is None: @@ -3816,7 +3828,7 @@ async def unblock_team( user_api_key_dict=user_api_key_dict, ) - record = await prisma_client.db.litellm_teamtable.update( + record = await TeamRepository(prisma_client).table.update( where={"team_id": data.team_id}, data={"blocked": False} # type: ignore ) @@ -3849,7 +3861,7 @@ async def list_available_teams( return [] # filter out teams that the user is already a member of - user_info = await prisma_client.db.litellm_usertable.find_unique( + user_info = await UserRepository(prisma_client).table.find_unique( where={"user_id": user_api_key_dict.user_id} ) if user_info is None: @@ -3863,7 +3875,7 @@ async def list_available_teams( team for team in available_teams if team not in user_info_correct_type.teams ] - available_teams_db = await prisma_client.db.litellm_teamtable.find_many( + available_teams_db = await TeamRepository(prisma_client).table.find_many( where={"team_id": {"in": available_teams}} ) @@ -4009,7 +4021,7 @@ async def _batch_resolve_access_group_resources( return {} unique_ids = list(set(all_access_group_ids)) - rows = await _prisma_client.db.litellm_accessgrouptable.find_many( + rows = await AccessGroupRepository(_prisma_client).table.find_many( where={"access_group_id": {"in": unique_ids}}, ) @@ -4074,7 +4086,7 @@ async def _get_keys_count_by_team( if not page_team_ids: return {} - grouped = await prisma_client.db.litellm_verificationtoken.group_by( + grouped = await VerificationTokenRepository(prisma_client).table.group_by( by=["team_id"], where={"team_id": {"in": page_team_ids}}, count={"team_id": True}, @@ -4288,25 +4300,25 @@ async def list_team_v2( # Get teams with pagination if use_deleted_table: - teams = await prisma_client.db.litellm_deletedteamtable.find_many( + teams = await DeletedTeamRepository(prisma_client).table.find_many( where=where_conditions, skip=skip, take=page_size, order=order_by if order_by else {"created_at": "desc"}, # Default sort ) # Get total count for pagination - total_count = await prisma_client.db.litellm_deletedteamtable.count( + total_count = await DeletedTeamRepository(prisma_client).table.count( where=where_conditions ) else: - teams = await prisma_client.db.litellm_teamtable.find_many( + teams = await TeamRepository(prisma_client).table.find_many( where=where_conditions, skip=skip, take=page_size, order=order_by if order_by else {"created_at": "desc"}, # Default sort ) # Get total count for pagination - total_count = await prisma_client.db.litellm_teamtable.count( + total_count = await TeamRepository(prisma_client).table.count( where=where_conditions ) @@ -4412,7 +4424,7 @@ async def _authorize_and_filter_teams( if allowed_org_ids is not None: # Org admin: query DB for teams in their orgs - org_teams = await prisma_client.db.litellm_teamtable.find_many( + org_teams = await TeamRepository(prisma_client).table.find_many( where={"organization_id": {"in": allowed_org_ids}}, include={"litellm_model_table": True}, ) @@ -4427,7 +4439,7 @@ async def _authorize_and_filter_teams( ] elif user_id: # Regular user: fetch all and filter by membership (Prisma can't filter JSON arrays) - response = await prisma_client.db.litellm_teamtable.find_many( + response = await TeamRepository(prisma_client).table.find_many( include={"litellm_model_table": True} ) return [ @@ -4439,7 +4451,7 @@ async def _authorize_and_filter_teams( else: # Proxy admin: all teams return list( - await prisma_client.db.litellm_teamtable.find_many( + await TeamRepository(prisma_client).table.find_many( include={"litellm_model_table": True} ) ) @@ -4500,7 +4512,7 @@ async def list_team( _team_memberships.append(tm) # add all keys that belong to the team - keys = await prisma_client.db.litellm_verificationtoken.find_many( + keys = await VerificationTokenRepository(prisma_client).table.find_many( where={"team_id": team.team_id} ) @@ -4556,10 +4568,10 @@ async def get_paginated_teams( # Calculate skip for pagination skip = (page - 1) * page_size # Get total count - total_count = await prisma_client.db.litellm_teamtable.count() + total_count = await TeamRepository(prisma_client).table.count() # Get paginated teams - teams = await prisma_client.db.litellm_teamtable.find_many( + teams = await TeamRepository(prisma_client).table.find_many( skip=skip, take=page_size, order={"team_alias": "asc"} # Sort by team_alias ) return teams, total_count @@ -4632,7 +4644,7 @@ async def ui_view_teams( } # Query users with pagination and filters - teams = await prisma_client.db.litellm_teamtable.find_many( + teams = await TeamRepository(prisma_client).table.find_many( where=where_conditions, skip=skip, take=page_size, @@ -4704,7 +4716,7 @@ async def team_model_add( raise HTTPException(status_code=500, detail={"error": "No db connected"}) # Get existing team - team_row = await prisma_client.db.litellm_teamtable.find_unique( + team_row = await TeamRepository(prisma_client).table.find_unique( where={"team_id": data.team_id} ) @@ -4737,7 +4749,7 @@ async def team_model_add( # null them out — see object_permission_utils.validate_key_search_tools_against_team # and the MCP/agent authz paths, which treat a missing object_permission # as "no team-level restriction". - updated_team = await prisma_client.db.litellm_teamtable.update( + updated_team = await TeamRepository(prisma_client).table.update( where={"team_id": data.team_id}, data={"models": updated_models}, include={"object_permission": True}, # type: ignore @@ -4791,7 +4803,7 @@ async def team_model_delete( raise HTTPException(status_code=500, detail={"error": "No db connected"}) # Get existing team - team_row = await prisma_client.db.litellm_teamtable.find_unique( + team_row = await TeamRepository(prisma_client).table.find_unique( where={"team_id": data.team_id} ) @@ -4825,7 +4837,7 @@ async def team_model_delete( updated_models = [m for m in current_models if m not in data.models] # Update team. See team_model_add for the rationale on `include`. - updated_team = await prisma_client.db.litellm_teamtable.update( + updated_team = await TeamRepository(prisma_client).table.update( where={"team_id": data.team_id}, data={"models": updated_models}, include={"object_permission": True}, # type: ignore @@ -4972,7 +4984,7 @@ async def update_team_member_permissions( }, ) # Update the team member permissions - updated_team = await prisma_client.db.litellm_teamtable.update( + updated_team = await TeamRepository(prisma_client).table.update( where={"team_id": data.team_id}, data={"team_member_permissions": data.team_member_permissions}, ) @@ -5076,7 +5088,7 @@ async def _append_permissions_to_specific_teams( prisma_client, team_ids: List[str], permissions_to_add: set ) -> int: """Fetch specific teams by ID and append permissions.""" - teams = await prisma_client.db.litellm_teamtable.find_many( + teams = await TeamRepository(prisma_client).table.find_many( where={"team_id": {"in": team_ids}}, ) @@ -5108,7 +5120,7 @@ async def _append_permissions_to_all_teams( find_args["cursor"] = {"team_id": cursor} find_args["skip"] = 1 - teams = await prisma_client.db.litellm_teamtable.find_many(**find_args) + teams = await TeamRepository(prisma_client).table.find_many(**find_args) if not teams: break @@ -5214,7 +5226,7 @@ async def get_team_daily_activity( where_condition = {} if team_ids_list: where_condition["team_id"] = {"in": list(team_ids_list)} - team_aliases = await prisma_client.db.litellm_teamtable.find_many( + team_aliases = await TeamRepository(prisma_client).table.find_many( where=where_condition ) team_alias_metadata = { @@ -5251,9 +5263,9 @@ async def get_team_daily_activity( # If user does not have full team view, filter by their API keys if not has_full_team_view: # Get all API keys for this user - user_keys = await prisma_client.db.litellm_verificationtoken.find_many( - where={"user_id": user_api_key_dict.user_id} - ) + user_keys = await VerificationTokenRepository( + prisma_client + ).table.find_many(where={"user_id": user_api_key_dict.user_id}) user_api_keys = [key.token for key in user_keys if key.token] # If user has no API keys, return empty result if not user_api_keys: diff --git a/litellm/proxy/management_endpoints/tool_management_endpoints.py b/litellm/proxy/management_endpoints/tool_management_endpoints.py index 19ca2c9f6be..a9b57db8a6f 100644 --- a/litellm/proxy/management_endpoints/tool_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tool_management_endpoints.py @@ -21,6 +21,15 @@ if TYPE_CHECKING: from litellm._logging import verbose_proxy_logger from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.repositories.object_permission_repository import ObjectPermissionRepository +from litellm.repositories.table_repositories import ( + SpendLogsRepository, + SpendLogToolIndexRepository, +) +from litellm.repositories.team_repository import TeamRepository +from litellm.repositories.verification_token_repository import ( + VerificationTokenRepository, +) from litellm.types.tool_management import ( LiteLLM_ToolTableRow, ToolDetailResponse, @@ -256,8 +265,10 @@ async def get_tool_usage_logs( if end_time_filter is not None: where["start_time"]["lte"] = end_time_filter - total = await prisma_client.db.litellm_spendlogtoolindex.count(where=where) - index_rows = await prisma_client.db.litellm_spendlogtoolindex.find_many( + total = await SpendLogToolIndexRepository(prisma_client).table.count( + where=where + ) + index_rows = await SpendLogToolIndexRepository(prisma_client).table.find_many( where=where, order={"start_time": "desc"}, skip=(page - 1) * page_size, @@ -269,7 +280,7 @@ async def get_tool_usage_logs( logs=[], total=total, page=page, page_size=page_size ) - spend_logs = await prisma_client.db.litellm_spendlogs.find_many( + spend_logs = await SpendLogsRepository(prisma_client).table.find_many( where={"request_id": {"in": request_ids}} ) log_by_id = {s.request_id: s for s in spend_logs} @@ -348,7 +359,7 @@ async def _resolve_key_hash_to_object_permission_id( hashed = key_hash if "sk-" not in (key_hash or "") else hash_token(key_hash) if not hashed: return None - row = await prisma_client.db.litellm_verificationtoken.find_unique( + row = await VerificationTokenRepository(prisma_client).table.find_unique( where={"token": hashed} ) if row is None: @@ -357,18 +368,18 @@ async def _resolve_key_hash_to_object_permission_id( if op_id: return op_id new_id = str(uuid.uuid4()) - await prisma_client.db.litellm_objectpermissiontable.create( + await ObjectPermissionRepository(prisma_client).table.create( data={"object_permission_id": new_id, "blocked_tools": []} ) - updated_count = await prisma_client.db.litellm_verificationtoken.update_many( + updated_count = await VerificationTokenRepository(prisma_client).table.update_many( where={"token": hashed, "object_permission_id": None}, data={"object_permission_id": new_id}, ) if updated_count == 0: - await prisma_client.db.litellm_objectpermissiontable.delete( + await ObjectPermissionRepository(prisma_client).table.delete( where={"object_permission_id": new_id} ) - row = await prisma_client.db.litellm_verificationtoken.find_unique( + row = await VerificationTokenRepository(prisma_client).table.find_unique( where={"token": hashed} ) return getattr(row, "object_permission_id", None) if row else None @@ -383,7 +394,7 @@ async def _resolve_team_id_to_object_permission_id( if not team_id or not team_id.strip(): return None team_id_clean = team_id.strip() - row = await prisma_client.db.litellm_teamtable.find_unique( + row = await TeamRepository(prisma_client).table.find_unique( where={"team_id": team_id_clean}, select={"object_permission_id": True}, ) @@ -393,18 +404,18 @@ async def _resolve_team_id_to_object_permission_id( if op_id: return op_id new_id = str(uuid.uuid4()) - await prisma_client.db.litellm_objectpermissiontable.create( + await ObjectPermissionRepository(prisma_client).table.create( data={"object_permission_id": new_id, "blocked_tools": []} ) - updated_count = await prisma_client.db.litellm_teamtable.update_many( + updated_count = await TeamRepository(prisma_client).table.update_many( where={"team_id": team_id_clean, "object_permission_id": None}, data={"object_permission_id": new_id}, ) if updated_count == 0: - await prisma_client.db.litellm_objectpermissiontable.delete( + await ObjectPermissionRepository(prisma_client).table.delete( where={"object_permission_id": new_id} ) - row = await prisma_client.db.litellm_teamtable.find_unique( + row = await TeamRepository(prisma_client).table.find_unique( where={"team_id": team_id_clean}, select={"object_permission_id": True}, ) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index d6082899c02..a9570616850 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -15,8 +15,8 @@ import inspect import os import re import secrets -from html import escape from copy import deepcopy +from html import escape from typing import ( TYPE_CHECKING, Any, @@ -39,9 +39,9 @@ from fastapi import APIRouter, Depends, Header, HTTPException, Request, status from fastapi.responses import RedirectResponse import litellm -from litellm.caching.dual_cache import DualCache from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid +from litellm.caching.dual_cache import DualCache from litellm.constants import ( CLI_SSO_CLAIM_MAP, CLI_SSO_CLAIM_MAX_SCALAR_LENGTH, @@ -77,7 +77,6 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken, get_user_object -from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.auth.auth_utils import ( _get_request_ip_address, _has_user_setup_sso, @@ -92,6 +91,7 @@ from litellm.proxy.common_utils.html_forms.jwt_display_template import ( jwt_display_template, ) from litellm.proxy.common_utils.html_forms.ui_login import html_form +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.management_endpoints.internal_user_endpoints import new_user from litellm.proxy.management_endpoints.sso import CustomMicrosoftSSO from litellm.proxy.management_endpoints.sso_helper_utils import ( @@ -110,6 +110,9 @@ from litellm.proxy.utils import ( get_custom_url, get_server_root_path, ) +from litellm.repositories.table_repositories import SSOConfigRepository +from litellm.repositories.team_repository import TeamRepository +from litellm.repositories.user_repository import UserRepository from litellm.secret_managers.main import get_secret_bool, str_to_bool from litellm.types.proxy.management_endpoints.ui_sso import * # noqa: F403, F401 from litellm.types.proxy.management_endpoints.ui_sso import ( @@ -438,7 +441,7 @@ async def _persist_cli_sso_user_metadata( return try: - user_row = await prisma_client.db.litellm_usertable.find_unique( + user_row = await UserRepository(prisma_client).table.find_unique( where={"user_id": user_id} ) existing_metadata: Dict[str, Any] = {} @@ -451,7 +454,7 @@ async def _persist_cli_sso_user_metadata( existing_metadata=existing_metadata, attribution_metadata=attribution_metadata, ) - await prisma_client.db.litellm_usertable.update_many( + await UserRepository(prisma_client).table.update_many( where={"user_id": user_id}, data={"metadata": merged_metadata}, ) @@ -859,7 +862,7 @@ async def google_login( if premium_user is not True: # Check if under 'free SSO user' limit if prisma_client is not None: - total_users = await prisma_client.db.litellm_usertable.count() + total_users = await UserRepository(prisma_client).table.count() if total_users and total_users > 5: raise ProxyException( message="You must be a LiteLLM Enterprise user to use SSO for more than 5 users. If you have a license please set `LITELLM_LICENSE` in your env. If you want to obtain a license meet with us here: https://enterprise.litellm.ai/demo You are seeing this error message because You set one of `MICROSOFT_CLIENT_ID`, `GOOGLE_CLIENT_ID`, or `GENERIC_CLIENT_ID` in your env. Please unset this", @@ -1150,7 +1153,7 @@ async def _setup_team_mappings() -> Optional["TeamMappings"]: "Prisma client is None, connect a database to your proxy" ) - sso_db_record = await prisma_client.db.litellm_ssoconfig.find_unique( + sso_db_record = await SSOConfigRepository(prisma_client).table.find_unique( where={"id": "sso_config"} ) @@ -1188,7 +1191,7 @@ async def _setup_role_mappings() -> Optional["RoleMappings"]: "Prisma client is None, connect a database to your proxy" ) - sso_db_record = await prisma_client.db.litellm_ssoconfig.find_unique( + sso_db_record = await SSOConfigRepository(prisma_client).table.find_unique( where={"id": "sso_config"} ) @@ -1755,7 +1758,7 @@ async def _sync_user_role_from_jwt_role_map( # Update existing DB record if role differs if user_info is not None and user_info.user_role != mapped_role.value: - await prisma_client.db.litellm_usertable.update( + await UserRepository(prisma_client).table.update( where={"user_id": user_info.user_id}, data={"user_role": mapped_role.value}, ) @@ -1819,7 +1822,7 @@ async def check_and_update_if_proxy_admin_id( return user_role if prisma_client: - await prisma_client.db.litellm_usertable.update( + await UserRepository(prisma_client).table.update( where={"user_id": user_id}, data={"user_role": LitellmUserRoles.PROXY_ADMIN.value}, ) @@ -1976,7 +1979,7 @@ async def _fetch_cli_sso_team_details( team_details: List[Dict[str, Any]] = [] try: if teams: - prisma_teams = await prisma_client.db.litellm_teamtable.find_many( + prisma_teams = await TeamRepository(prisma_client).table.find_many( where={"team_id": {"in": teams}} ) for team_row in prisma_teams: @@ -2884,7 +2887,7 @@ class SSOAuthenticationHandler: user_id=user_id, ) - await prisma_client.db.litellm_usertable.update_many( + await UserRepository(prisma_client).table.update_many( where={"user_id": user_id}, data=update_data ) else: @@ -2986,7 +2989,7 @@ class SSOAuthenticationHandler: code=status.HTTP_500_INTERNAL_SERVER_ERROR, ) try: - team_obj = await prisma_client.db.litellm_teamtable.find_first( + team_obj = await TeamRepository(prisma_client).table.find_first( where={"team_id": litellm_team_id} ) verbose_proxy_logger.debug(f"Team object: {team_obj}") diff --git a/litellm/proxy/management_endpoints/user_agent_analytics_endpoints.py b/litellm/proxy/management_endpoints/user_agent_analytics_endpoints.py index ebd276fbee5..661487577c3 100644 --- a/litellm/proxy/management_endpoints/user_agent_analytics_endpoints.py +++ b/litellm/proxy/management_endpoints/user_agent_analytics_endpoints.py @@ -19,6 +19,11 @@ from pydantic import BaseModel from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.repositories.table_repositories import DailyTagSpendRepository +from litellm.repositories.user_repository import UserRepository +from litellm.repositories.verification_token_repository import ( + VerificationTokenRepository, +) # Constants for analytics periods MAX_DAYS = 7 # Number of days to show in DAU analytics @@ -676,7 +681,7 @@ async def get_per_user_analytics( where_clause["tag"] = {"contains": tag_filter} # Get all tag records in the date range with optional tag filtering - tag_records = await prisma_client.db.litellm_dailytagspend.find_many( + tag_records = await DailyTagSpendRepository(prisma_client).table.find_many( where=where_clause ) @@ -693,9 +698,9 @@ async def get_per_user_analytics( ) # Lookup user_id for each api_key - api_key_records = await prisma_client.db.litellm_verificationtoken.find_many( - where={"token": {"in": list(api_keys)}} - ) + api_key_records = await VerificationTokenRepository( + prisma_client + ).table.find_many(where={"token": {"in": list(api_keys)}}) # Create mapping from api_key to user_id api_key_to_user_id = { @@ -704,7 +709,7 @@ async def get_per_user_analytics( # Get user emails for the user_ids user_ids = list(set(api_key_to_user_id.values())) - user_records = await prisma_client.db.litellm_usertable.find_many( + user_records = await UserRepository(prisma_client).table.find_many( where={"user_id": {"in": user_ids}} ) diff --git a/litellm/proxy/management_endpoints/workflow_management_endpoints.py b/litellm/proxy/management_endpoints/workflow_management_endpoints.py index a19af4dd484..57cc0dc6745 100644 --- a/litellm/proxy/management_endpoints/workflow_management_endpoints.py +++ b/litellm/proxy/management_endpoints/workflow_management_endpoints.py @@ -27,6 +27,11 @@ from pydantic import BaseModel from litellm._logging import verbose_proxy_logger from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.repositories.table_repositories import ( + WorkflowEventRepository, + WorkflowMessageRepository, + WorkflowRunRepository, +) router = APIRouter() @@ -96,13 +101,13 @@ class WorkflowMessageCreateRequest(BaseModel): async def _get_next_sequence_number(prisma_client: Any, run_id: str, table: str) -> int: """Return MAX(sequence_number) + 1 for the given run, for either events or messages.""" if table == "events": - rows = await prisma_client.db.litellm_workflowevent.find_many( + rows = await WorkflowEventRepository(prisma_client).table.find_many( where={"run_id": run_id}, order={"sequence_number": "desc"}, take=1, ) else: - rows = await prisma_client.db.litellm_workflowmessage.find_many( + rows = await WorkflowMessageRepository(prisma_client).table.find_many( where={"run_id": run_id}, order={"sequence_number": "desc"}, take=1, @@ -116,7 +121,7 @@ async def _require_run( user_api_key_dict: Optional[UserAPIKeyAuth] = None, ) -> Any: """Return the run or raise 404. For non-admin callers, also enforce key ownership.""" - run = await prisma_client.db.litellm_workflowrun.find_unique( + run = await WorkflowRunRepository(prisma_client).table.find_unique( where={"run_id": run_id} ) if run is None: @@ -163,7 +168,7 @@ async def create_workflow_run( create_data["input"] = _json(data.input) if data.metadata is not None: create_data["metadata"] = _json(data.metadata) - run = await prisma_client.db.litellm_workflowrun.create(data=create_data) + run = await WorkflowRunRepository(prisma_client).table.create(data=create_data) return run except Exception as e: verbose_proxy_logger.exception("Error creating workflow run: %s", e) @@ -206,7 +211,7 @@ async def list_workflow_runs( where["created_by"] = caller try: - runs = await prisma_client.db.litellm_workflowrun.find_many( + runs = await WorkflowRunRepository(prisma_client).table.find_many( where=where, order={"created_at": "desc"}, take=limit, @@ -235,7 +240,7 @@ async def get_workflow_run( ) try: - run = await prisma_client.db.litellm_workflowrun.find_unique( + run = await WorkflowRunRepository(prisma_client).table.find_unique( where={"run_id": run_id}, include={"events": {"order_by": {"sequence_number": "desc"}, "take": 1}}, ) @@ -286,7 +291,7 @@ async def update_workflow_run( await _require_run(prisma_client, run_id, user_api_key_dict) try: - run = await prisma_client.db.litellm_workflowrun.update( + run = await WorkflowRunRepository(prisma_client).table.update( where={"run_id": run_id}, data=update, ) @@ -391,7 +396,7 @@ async def list_workflow_events( await _require_run(prisma_client, run_id, user_api_key_dict) try: - events = await prisma_client.db.litellm_workflowevent.find_many( + events = await WorkflowEventRepository(prisma_client).table.find_many( where={"run_id": run_id}, order={"sequence_number": "asc"}, take=limit, @@ -436,7 +441,9 @@ async def append_workflow_message( } if data.session_id is not None: msg_data["session_id"] = data.session_id - msg = await prisma_client.db.litellm_workflowmessage.create(data=msg_data) + msg = await WorkflowMessageRepository(prisma_client).table.create( + data=msg_data + ) return msg except Exception as e: @@ -481,7 +488,7 @@ async def list_workflow_messages( await _require_run(prisma_client, run_id, user_api_key_dict) try: - messages = await prisma_client.db.litellm_workflowmessage.find_many( + messages = await WorkflowMessageRepository(prisma_client).table.find_many( where={"run_id": run_id}, order={"sequence_number": "asc"}, take=limit, diff --git a/litellm/proxy/management_helpers/audit_logs.py b/litellm/proxy/management_helpers/audit_logs.py index 439c3b2118d..33599c3c622 100644 --- a/litellm/proxy/management_helpers/audit_logs.py +++ b/litellm/proxy/management_helpers/audit_logs.py @@ -18,6 +18,7 @@ from litellm.proxy._types import ( Optional, UserAPIKeyAuth, ) +from litellm.repositories.table_repositories import AuditLogRepository from litellm.types.utils import StandardAuditLogPayload _audit_log_callback_cache: Dict[str, CustomLogger] = {} @@ -244,7 +245,7 @@ async def create_audit_log_for_update(request_data: LiteLLM_AuditLogs): _request_data = request_data.model_dump(exclude_none=True) try: - await prisma_client.db.litellm_auditlog.create( + await AuditLogRepository(prisma_client).table.create( data={ **_request_data, # type: ignore } diff --git a/litellm/proxy/management_helpers/object_permission_utils.py b/litellm/proxy/management_helpers/object_permission_utils.py index 4c966b25413..f2ddae40d8c 100644 --- a/litellm/proxy/management_helpers/object_permission_utils.py +++ b/litellm/proxy/management_helpers/object_permission_utils.py @@ -12,6 +12,8 @@ from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy.utils import PrismaClient +from litellm.repositories.object_permission_repository import ObjectPermissionRepository +from litellm.repositories.table_repositories import MCPServerRepository if TYPE_CHECKING: from litellm.proxy._types import ( @@ -48,10 +50,10 @@ async def attach_object_permission_to_dict( object_permission_id = data_dict.get("object_permission_id") if object_permission_id: - object_permission = ( - await prisma_client.db.litellm_objectpermissiontable.find_unique( - where={"object_permission_id": object_permission_id}, - ) + object_permission = await ObjectPermissionRepository( + prisma_client + ).table.find_unique( + where={"object_permission_id": object_permission_id}, ) if object_permission: # Convert to dict if needed @@ -106,10 +108,10 @@ async def handle_update_object_permission_common( ) existing_object_permissions_dict: Dict = {} - existing_object_permission = ( - await prisma_client.db.litellm_objectpermissiontable.find_unique( - where={"object_permission_id": object_permission_id_to_use}, - ) + existing_object_permission = await ObjectPermissionRepository( + prisma_client + ).table.find_unique( + where={"object_permission_id": object_permission_id_to_use}, ) # Update the object permission @@ -137,14 +139,14 @@ async def handle_update_object_permission_common( ######################################################### # Commit the update to the LiteLLM_ObjectPermissionTable ######################################################### - created_object_permission_row = ( - await prisma_client.db.litellm_objectpermissiontable.upsert( - where={"object_permission_id": object_permission_id_to_use}, - data={ - "create": existing_object_permissions_dict, - "update": existing_object_permissions_dict, - }, - ) + created_object_permission_row = await ObjectPermissionRepository( + prisma_client + ).table.upsert( + where={"object_permission_id": object_permission_id_to_use}, + data={ + "create": existing_object_permissions_dict, + "update": existing_object_permissions_dict, + }, ) verbose_proxy_logger.debug( @@ -183,7 +185,7 @@ async def _set_object_permission( clean_data["mcp_tool_permissions"] ) - created_permission = await prisma_client.db.litellm_objectpermissiontable.create( + created_permission = await ObjectPermissionRepository(prisma_client).table.create( data=clean_data ) @@ -220,7 +222,7 @@ async def _get_db_mcp_servers_by_identifiers( return [] identifier_list = list(identifiers) - return await prisma_client.db.litellm_mcpservertable.find_many( + return await MCPServerRepository(prisma_client).table.find_many( where={ "OR": [ {"server_id": {"in": identifier_list}}, diff --git a/litellm/proxy/management_helpers/user_invitation.py b/litellm/proxy/management_helpers/user_invitation.py index d2d800aa77f..babc920189a 100644 --- a/litellm/proxy/management_helpers/user_invitation.py +++ b/litellm/proxy/management_helpers/user_invitation.py @@ -4,6 +4,7 @@ from fastapi import HTTPException import litellm from litellm.proxy._types import CommonProxyErrors, InvitationNew, UserAPIKeyAuth +from litellm.repositories.table_repositories import InvitationLinkRepository async def create_invitation_for_user( @@ -25,7 +26,7 @@ async def create_invitation_for_user( expires_at = current_time + timedelta(days=7) try: - response = await prisma_client.db.litellm_invitationlink.create( + response = await InvitationLinkRepository(prisma_client).table.create( data={ "user_id": data.user_id, "created_at": current_time, diff --git a/litellm/proxy/management_helpers/utils.py b/litellm/proxy/management_helpers/utils.py index 0b175db3c87..830d6f84b85 100644 --- a/litellm/proxy/management_helpers/utils.py +++ b/litellm/proxy/management_helpers/utils.py @@ -9,9 +9,8 @@ from pydantic import BaseModel import litellm from litellm._logging import verbose_logger -from litellm.integrations.otel.model.config import is_otel_v2_enabled from litellm._uuid import uuid -from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time +from litellm.integrations.otel.model.config import is_otel_v2_enabled from litellm.proxy._types import ( # key request types; user request types; team request types; customer request types BudgetNewRequest, DeleteCustomerRequest, @@ -32,7 +31,11 @@ from litellm.proxy._types import ( # key request types; user request types; tea VirtualKeyEvent, ) from litellm.proxy.common_utils.http_parsing_utils import _read_request_body +from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time from litellm.proxy.utils import PrismaClient +from litellm.repositories.budget_repository import BudgetRepository +from litellm.repositories.table_repositories import TeamMembershipRepository +from litellm.repositories.user_repository import UserRepository def get_new_internal_user_defaults( @@ -111,7 +114,7 @@ async def handle_budget_for_entity( budget_row.model_dump(exclude_none=True) ) - _budget = await prisma_client.db.litellm_budgettable.create( + _budget = await BudgetRepository(prisma_client).table.create( data={ **new_budget_data, # type: ignore "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name, @@ -174,7 +177,7 @@ async def _clone_team_default_budget_for_member( so the member starts with the team default's values but gets their own private budget row (which can be edited independently). """ - default_budget = await prisma_client.db.litellm_budgettable.find_unique( + default_budget = await BudgetRepository(prisma_client).table.find_unique( where={"budget_id": default_team_budget_id} ) if default_budget is None: @@ -202,7 +205,7 @@ async def _clone_team_default_budget_for_member( cloned_data["budget_duration"] ) - new_budget = await prisma_client.db.litellm_budgettable.create(data=cloned_data) + new_budget = await BudgetRepository(prisma_client).table.create(data=cloned_data) return new_budget.budget_id @@ -229,7 +232,7 @@ async def add_new_member( ## ADD TEAM ID, to USER TABLE IF NEW ## if new_member.user_id is not None: new_user_defaults = get_new_internal_user_defaults(user_id=new_member.user_id) - _returned_user = await prisma_client.db.litellm_usertable.upsert( + _returned_user = await UserRepository(prisma_client).table.upsert( where={"user_id": new_member.user_id}, data={ "update": {"teams": {"push": [team_id]}}, @@ -259,7 +262,7 @@ async def add_new_member( returned_user = LiteLLM_UserTable(**_returned_user.model_dump()) elif len(existing_user_row) == 1: user_info = existing_user_row[0] - _returned_user = await prisma_client.db.litellm_usertable.update( + _returned_user = await UserRepository(prisma_client).table.update( where={"user_id": user_info.user_id}, # type: ignore data={"teams": {"push": [team_id]}}, ) @@ -284,7 +287,7 @@ async def add_new_member( budget_data["max_budget"] = max_budget_in_team if allowed_models is not None: budget_data["allowed_models"] = allowed_models - response = await prisma_client.db.litellm_budgettable.create(data=budget_data) + response = await BudgetRepository(prisma_client).table.create(data=budget_data) _budget_id = response.budget_id elif default_team_budget_id is not None: @@ -303,15 +306,15 @@ async def add_new_member( _budget_id = None if _budget_id and returned_user is not None and returned_user.user_id is not None: - _returned_team_membership = ( - await prisma_client.db.litellm_teammembership.create( - data={ - "team_id": team_id, - "user_id": returned_user.user_id, - "budget_id": _budget_id, - }, - include={"litellm_budget_table": True}, - ) + _returned_team_membership = await TeamMembershipRepository( + prisma_client + ).table.create( + data={ + "team_id": team_id, + "user_id": returned_user.user_id, + "budget_id": _budget_id, + }, + include={"litellm_budget_table": True}, ) returned_team_membership = LiteLLM_TeamMembership( diff --git a/litellm/proxy/memory/memory_endpoints.py b/litellm/proxy/memory/memory_endpoints.py index 4d161be4263..6f1ca3196fe 100644 --- a/litellm/proxy/memory/memory_endpoints.py +++ b/litellm/proxy/memory/memory_endpoints.py @@ -29,6 +29,8 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.repositories.table_repositories import MemoryRepository +from litellm.repositories.team_repository import TeamRepository from litellm.types.memory_management import ( LiteLLM_MemoryRow, MemoryCreateRequest, @@ -173,7 +175,7 @@ async def _is_team_admin_for( ) try: - team_obj = await prisma_client.db.litellm_teamtable.find_unique( + team_obj = await TeamRepository(prisma_client).table.find_unique( where={"team_id": team_id} ) except Exception as e: @@ -304,7 +306,7 @@ async def create_memory( create_data["metadata"] = _serialize_metadata_for_prisma(body.metadata) try: - row = await prisma_client.db.litellm_memorytable.create(data=create_data) + row = await MemoryRepository(prisma_client).table.create(data=create_data) except Exception as e: # Key is globally unique. Any duplicate → 409. if _is_unique_violation(e): @@ -364,8 +366,8 @@ async def list_memory( where = {"AND": [key_filter, vis]} try: - total = await prisma_client.db.litellm_memorytable.count(where=where) - rows = await prisma_client.db.litellm_memorytable.find_many( + total = await MemoryRepository(prisma_client).table.count(where=where) + rows = await MemoryRepository(prisma_client).table.find_many( where=where, order={"updated_at": "desc"}, skip=(page - 1) * page_size, @@ -386,7 +388,7 @@ async def _find_memory_for_caller( key_filter: dict = {"key": key} vis = _visibility_filter(user_api_key_dict) where: dict = key_filter if vis is None else {"AND": [key_filter, vis]} - rows = await prisma_client.db.litellm_memorytable.find_many( + rows = await MemoryRepository(prisma_client).table.find_many( where=where, take=1, order={"updated_at": "desc"} ) if not rows: @@ -475,7 +477,7 @@ async def upsert_memory( # their team) — otherwise a teammate could overwrite a personal # entry through the OR-based visibility filter. await _assert_write_access(prisma_client, existing, user_api_key_dict) - row = await prisma_client.db.litellm_memorytable.update( + row = await MemoryRepository(prisma_client).table.update( where={"memory_id": existing.memory_id}, data=data, ) @@ -503,7 +505,7 @@ async def upsert_memory( if body.metadata is not None: create_data["metadata"] = _serialize_metadata_for_prisma(body.metadata) try: - row = await prisma_client.db.litellm_memorytable.create( + row = await MemoryRepository(prisma_client).table.create( data=create_data ) except Exception as e: @@ -524,7 +526,7 @@ async def upsert_memory( await _assert_write_access( prisma_client, existing_after_race, user_api_key_dict ) - row = await prisma_client.db.litellm_memorytable.update( + row = await MemoryRepository(prisma_client).table.update( where={"memory_id": existing_after_race.memory_id}, data=data, ) @@ -554,7 +556,7 @@ async def delete_memory( # Visibility != write authority — see the upsert handler for the rationale. await _assert_write_access(prisma_client, row, user_api_key_dict) try: - await prisma_client.db.litellm_memorytable.delete( + await MemoryRepository(prisma_client).table.delete( where={"memory_id": row.memory_id} ) except Exception as e: diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index cc0d06e4f40..b2834e52306 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -5,6 +5,10 @@ from dataclasses import dataclass, field from types import MappingProxyType from typing import TYPE_CHECKING, List, Literal, Optional, Union +from litellm.repositories.table_repositories import ( + ManagedFileRepository, + ManagedObjectRepository, +) from litellm.types.utils import SpecialEnums if TYPE_CHECKING: @@ -697,7 +701,7 @@ async def resolve_input_file_id_to_unified(response, prisma_client) -> None: and prisma_client ): try: - managed_file = await prisma_client.db.litellm_managedfiletable.find_first( + managed_file = await ManagedFileRepository(prisma_client).table.find_first( where={"flat_model_file_ids": {"has": response.input_file_id}} ) if managed_file: @@ -719,7 +723,7 @@ async def resolve_output_file_ids_to_unified(response, prisma_client) -> None: if not raw_id or _is_base64_encoded_unified_file_id(raw_id): continue try: - managed_file = await prisma_client.db.litellm_managedfiletable.find_first( + managed_file = await ManagedFileRepository(prisma_client).table.find_first( where={"flat_model_file_ids": {"has": raw_id}} ) if managed_file: @@ -821,6 +825,7 @@ async def get_batch_from_database( - response_batch: Parsed LiteLLMBatch object (or None) """ import json + from litellm.types.utils import LiteLLMBatch if managed_files_obj is None or not unified_batch_id: @@ -830,7 +835,7 @@ async def get_batch_from_database( if not prisma_client: return None, None - db_batch_object = await prisma_client.db.litellm_managedobjecttable.find_first( + db_batch_object = await ManagedObjectRepository(prisma_client).table.find_first( where={"unified_object_id": batch_id} ) @@ -942,7 +947,7 @@ async def update_batch_in_database( update_data["batch_processed"] = True try: - await prisma_client.db.litellm_managedobjecttable.update( + await ManagedObjectRepository(prisma_client).table.update( where={"unified_object_id": batch_id}, data=update_data, ) @@ -958,7 +963,7 @@ async def update_batch_in_database( f"batch_processed column not found, retrying update without it: {col_err}" ) update_data.pop("batch_processed", None) - await prisma_client.db.litellm_managedobjecttable.update( + await ManagedObjectRepository(prisma_client).table.update( where={"unified_object_id": batch_id}, data=update_data, ) diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 378cbbda89c..9eef7cd7e8b 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -21,6 +21,7 @@ from fastapi import ( UploadFile, status, ) + import litellm from litellm import CreateFileRequest, get_secret_str from litellm._logging import verbose_proxy_logger @@ -37,15 +38,6 @@ from litellm.proxy.common_utils.openai_endpoint_utils import ( get_custom_llm_provider_from_request_headers, get_custom_llm_provider_from_request_query, ) -from litellm.proxy.utils import ProxyLogging, is_known_model -from litellm.router import Router -from litellm.types.llms.openai import ( - CREATE_FILE_REQUESTS_PURPOSE, - FileExpiresAfter, - OpenAIFileObject, - OpenAIFilesPurpose, -) - from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, encode_file_id_with_model, @@ -54,6 +46,15 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( handle_model_based_routing, prepare_data_with_credentials, ) +from litellm.proxy.utils import ProxyLogging, is_known_model +from litellm.repositories.table_repositories import ManagedFileRepository +from litellm.router import Router +from litellm.types.llms.openai import ( + CREATE_FILE_REQUESTS_PURPOSE, + FileExpiresAfter, + OpenAIFileObject, + OpenAIFilesPurpose, +) router = APIRouter() @@ -666,7 +667,7 @@ async def get_file_content( # noqa: PLR0915 managed_files_obj, "prisma_client", None ): prisma_client = getattr(managed_files_obj, "prisma_client") - db_file = await prisma_client.db.litellm_managedfiletable.find_first( + db_file = await ManagedFileRepository(prisma_client).table.find_first( where={"unified_file_id": file_id} ) if db_file and db_file.storage_backend and db_file.storage_url: diff --git a/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py index a267c97c0e8..9c0fbe30fc3 100644 --- a/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py +++ b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py @@ -43,6 +43,10 @@ from litellm.llms.base_llm.managed_resources.isolation import ( can_access_resource, ) from litellm.proxy._types import UserAPIKeyAuth +from litellm.repositories.table_repositories import ( + ManagedFileRepository, + ManagedObjectRepository, +) from litellm.types.llms.openai import OpenAIFileObject from .managed_id_codec import ManagedIdPayload, decode, is_managed, new_managed_id @@ -323,7 +327,7 @@ async def _resolve_one( ) if not found and prisma_client is not None: try: - db_row = await prisma_client.db.litellm_managedfiletable.find_first( + db_row = await ManagedFileRepository(prisma_client).table.find_first( where={"unified_file_id": managed_id} ) if db_row is not None: @@ -339,7 +343,7 @@ async def _resolve_one( # Object table (batches, responses) if prisma_client is not None: try: - obj_row = await prisma_client.db.litellm_managedobjecttable.find_first( + obj_row = await ManagedObjectRepository(prisma_client).table.find_first( where={"unified_object_id": managed_id} ) if obj_row is not None: @@ -399,7 +403,7 @@ async def _guard_raw_provider_id( # id and scope to the current provider in the application layer (same as # _mint_or_reuse_file's dedup). try: - candidates = await prisma_client.db.litellm_managedfiletable.find_many( + candidates = await ManagedFileRepository(prisma_client).table.find_many( where={"flat_model_file_ids": {"has": raw_id}}, ) except Exception: @@ -425,7 +429,7 @@ async def _guard_raw_provider_id( # Object rows store model_object_id as "passthrough:{provider}:{raw}", so # the lookup is exact and already provider-scoped. try: - existing = await prisma_client.db.litellm_managedobjecttable.find_first( + existing = await ManagedObjectRepository(prisma_client).table.find_first( where={"model_object_id": f"passthrough:{provider}:{raw_id}"} ) except Exception: @@ -492,7 +496,7 @@ async def _mint_or_reuse_file( # reuse a stable row instead of minting duplicate rows on every call. if prisma_client is not None: try: - candidates = await prisma_client.db.litellm_managedfiletable.find_many( + candidates = await ManagedFileRepository(prisma_client).table.find_many( where={"flat_model_file_ids": {"has": raw_id}}, order={"created_at": "asc"}, ) @@ -627,7 +631,7 @@ async def _mint_or_reuse_object( # the batch's latest state (e.g. output_file_id / error_file_id that # were null at creation but populated once the batch completed). try: - await prisma_client.db.litellm_managedobjecttable.update( + await ManagedObjectRepository(prisma_client).table.update( where={"unified_object_id": existing.unified_object_id}, data={ "file_object": json.dumps(body_snapshot), @@ -647,7 +651,7 @@ async def _mint_or_reuse_object( # Dedup: look up by the namespaced key — guaranteed unique per provider. try: - existing = await prisma_client.db.litellm_managedobjecttable.find_first( + existing = await ManagedObjectRepository(prisma_client).table.find_first( where={"model_object_id": namespaced_model_object_id} ) except Exception: @@ -666,7 +670,7 @@ async def _mint_or_reuse_object( raw_id.split("_", 1)[0], ) try: - await prisma_client.db.litellm_managedobjecttable.upsert( + await ManagedObjectRepository(prisma_client).table.upsert( where={"unified_object_id": managed_id}, data={ "create": { @@ -690,7 +694,7 @@ async def _mint_or_reuse_object( # the winner's managed ID so both callers converge on one ID instead of # the loser silently keeping the raw id. try: - raced = await prisma_client.db.litellm_managedobjecttable.find_first( + raced = await ManagedObjectRepository(prisma_client).table.find_first( where={"model_object_id": namespaced_model_object_id} ) except Exception: @@ -883,9 +887,9 @@ async def _build_list_where_with_cursor( return where, fetch_order cursor_table = ( - prisma_client.db.litellm_managedfiletable + ManagedFileRepository(prisma_client).table if resource_kind == "files" - else prisma_client.db.litellm_managedobjecttable + else ManagedObjectRepository(prisma_client).table ) cursor_field = ( "unified_file_id" if resource_kind == "files" else "unified_object_id" @@ -932,12 +936,12 @@ async def _fetch_list_rows( # across rows that share a created_at timestamp. try: if resource_kind == "files": - return await prisma_client.db.litellm_managedfiletable.find_many( + return await ManagedFileRepository(prisma_client).table.find_many( where=where, order=[{"created_at": fetch_order}, {"unified_file_id": fetch_order}], take=fetch_limit, ) - return await prisma_client.db.litellm_managedobjecttable.find_many( + return await ManagedObjectRepository(prisma_client).table.find_many( where={**where, "file_purpose": "batch"}, order=[{"created_at": fetch_order}, {"unified_object_id": fetch_order}], take=fetch_limit, diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 6667010447b..45e264b1cdd 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -60,12 +60,13 @@ from litellm.proxy.common_utils.http_parsing_utils import ( ) from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.proxy.utils import normalize_route_for_root_path +from litellm.repositories.team_repository import TeamRepository from litellm.secret_managers.main import get_secret_str from litellm.types.llms.custom_http import httpxSpecialProvider from litellm.types.passthrough_endpoints.pass_through_endpoints import ( - EndpointType, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, + EndpointType, PassthroughStandardLoggingPayload, ) @@ -1002,9 +1003,7 @@ async def pass_through_request( # noqa: PLR0915 is_passthrough_list_route, list_passthrough_ids_from_db, ) - from litellm.proxy.proxy_server import ( - prisma_client as _list_prisma, - ) + from litellm.proxy.proxy_server import prisma_client as _list_prisma if ( is_passthrough_list_route( @@ -2875,7 +2874,7 @@ async def _filter_endpoints_by_team_allowed_routes( HTTPException: If team is not found """ # retrieve team from db - team = await prisma_client.db.litellm_teamtable.find_unique( + team = await TeamRepository(prisma_client).table.find_unique( where={"team_id": team_id}, ) if team is None: diff --git a/litellm/proxy/policy_engine/attachment_registry.py b/litellm/proxy/policy_engine/attachment_registry.py index 8d5d8116919..fb1e2652e8a 100644 --- a/litellm/proxy/policy_engine/attachment_registry.py +++ b/litellm/proxy/policy_engine/attachment_registry.py @@ -9,6 +9,7 @@ from datetime import datetime, timezone from typing import TYPE_CHECKING, Any, Dict, List, Optional from litellm._logging import verbose_proxy_logger +from litellm.repositories.table_repositories import PolicyAttachmentRepository from litellm.types.proxy.policy_engine import ( PolicyAttachment, PolicyAttachmentCreateRequest, @@ -278,21 +279,21 @@ class AttachmentRegistry: PolicyAttachmentDBResponse with the created attachment """ try: - created_attachment = ( - await prisma_client.db.litellm_policyattachmenttable.create( - data={ - "policy_name": attachment_request.policy_name, - "scope": attachment_request.scope, - "teams": attachment_request.teams or [], - "keys": attachment_request.keys or [], - "models": attachment_request.models or [], - "tags": attachment_request.tags or [], - "created_at": datetime.now(timezone.utc), - "updated_at": datetime.now(timezone.utc), - "created_by": created_by, - "updated_by": created_by, - } - ) + created_attachment = await PolicyAttachmentRepository( + prisma_client + ).table.create( + data={ + "policy_name": attachment_request.policy_name, + "scope": attachment_request.scope, + "teams": attachment_request.teams or [], + "keys": attachment_request.keys or [], + "models": attachment_request.models or [], + "tags": attachment_request.tags or [], + "created_at": datetime.now(timezone.utc), + "updated_at": datetime.now(timezone.utc), + "created_by": created_by, + "updated_by": created_by, + } ) # Also add to in-memory registry @@ -340,17 +341,15 @@ class AttachmentRegistry: """ try: # Get attachment before deleting - attachment = ( - await prisma_client.db.litellm_policyattachmenttable.find_unique( - where={"attachment_id": attachment_id} - ) - ) + attachment = await PolicyAttachmentRepository( + prisma_client + ).table.find_unique(where={"attachment_id": attachment_id}) if attachment is None: raise Exception(f"Attachment with ID {attachment_id} not found") # Delete from DB - await prisma_client.db.litellm_policyattachmenttable.delete( + await PolicyAttachmentRepository(prisma_client).table.delete( where={"attachment_id": attachment_id} ) @@ -379,11 +378,9 @@ class AttachmentRegistry: PolicyAttachmentDBResponse if found, None otherwise """ try: - attachment = ( - await prisma_client.db.litellm_policyattachmenttable.find_unique( - where={"attachment_id": attachment_id} - ) - ) + attachment = await PolicyAttachmentRepository( + prisma_client + ).table.find_unique(where={"attachment_id": attachment_id}) if attachment is None: return None @@ -419,10 +416,10 @@ class AttachmentRegistry: List of PolicyAttachmentDBResponse objects """ try: - attachments = ( - await prisma_client.db.litellm_policyattachmenttable.find_many( - order={"created_at": "desc"}, - ) + attachments = await PolicyAttachmentRepository( + prisma_client + ).table.find_many( + order={"created_at": "desc"}, ) return [ diff --git a/litellm/proxy/policy_engine/policy_registry.py b/litellm/proxy/policy_engine/policy_registry.py index 75017c46603..d6265516269 100644 --- a/litellm/proxy/policy_engine/policy_registry.py +++ b/litellm/proxy/policy_engine/policy_registry.py @@ -12,6 +12,7 @@ from datetime import datetime, timezone from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple from litellm._logging import verbose_proxy_logger +from litellm.repositories.table_repositories import PolicyRepository from litellm.types.proxy.policy_engine import ( GuardrailPipeline, PipelineStep, @@ -295,7 +296,7 @@ class PolicyRegistry: validated_pipeline = GuardrailPipeline(**policy_request.pipeline) data["pipeline"] = json.dumps(validated_pipeline.model_dump()) - created_policy = await prisma_client.db.litellm_policytable.create( + created_policy = await PolicyRepository(prisma_client).table.create( data=data ) @@ -347,7 +348,7 @@ class PolicyRegistry: Exception: If policy is not in draft status (only drafts are editable). """ try: - existing = await prisma_client.db.litellm_policytable.find_unique( + existing = await PolicyRepository(prisma_client).table.find_unique( where={"policy_id": policy_id} ) if existing is None: @@ -382,7 +383,7 @@ class PolicyRegistry: validated_pipeline = GuardrailPipeline(**policy_request.pipeline) update_data["pipeline"] = json.dumps(validated_pipeline.model_dump()) - updated_policy = await prisma_client.db.litellm_policytable.update( + updated_policy = await PolicyRepository(prisma_client).table.update( where={"policy_id": policy_id}, data=update_data, ) @@ -413,7 +414,7 @@ class PolicyRegistry: Dict with "message" and optional "warning" if production was deleted. """ try: - policy = await prisma_client.db.litellm_policytable.find_unique( + policy = await PolicyRepository(prisma_client).table.find_unique( where={"policy_id": policy_id} ) @@ -424,7 +425,7 @@ class PolicyRegistry: policy_name = policy.policy_name # Delete from DB - await prisma_client.db.litellm_policytable.delete( + await PolicyRepository(prisma_client).table.delete( where={"policy_id": policy_id} ) @@ -461,7 +462,7 @@ class PolicyRegistry: PolicyDBResponse if found, None otherwise """ try: - policy = await prisma_client.db.litellm_policytable.find_unique( + policy = await PolicyRepository(prisma_client).table.find_unique( where={"policy_id": policy_id} ) @@ -512,7 +513,7 @@ class PolicyRegistry: if version_status is not None: where["version_status"] = version_status - policies = await prisma_client.db.litellm_policytable.find_many( + policies = await PolicyRepository(prisma_client).table.find_many( where=where if where else None, order={"created_at": "desc"}, ) @@ -554,7 +555,7 @@ class PolicyRegistry: self.add_policy(policy_response.policy_name, policy) self._policies_by_id = {} - non_production = await prisma_client.db.litellm_policytable.find_many( + non_production = await PolicyRepository(prisma_client).table.find_many( where={"version_status": {"in": ["draft", "published"]}}, order={"created_at": "desc"}, ) @@ -654,7 +655,7 @@ class PolicyRegistry: PolicyVersionListResponse with policy_name and list of versions """ try: - rows = await prisma_client.db.litellm_policytable.find_many( + rows = await PolicyRepository(prisma_client).table.find_many( where={"policy_name": policy_name}, order={"version_number": "desc"}, ) @@ -690,7 +691,7 @@ class PolicyRegistry: """ try: if source_policy_id is not None: - source = await prisma_client.db.litellm_policytable.find_unique( + source = await PolicyRepository(prisma_client).table.find_unique( where={"policy_id": source_policy_id} ) if source is None: @@ -701,7 +702,7 @@ class PolicyRegistry: ) else: # Find current production version for this policy_name - prod = await prisma_client.db.litellm_policytable.find_first( + prod = await PolicyRepository(prisma_client).table.find_first( where={ "policy_name": policy_name, "version_status": "production", @@ -714,7 +715,7 @@ class PolicyRegistry: source = prod # Next version number - latest = await prisma_client.db.litellm_policytable.find_first( + latest = await PolicyRepository(prisma_client).table.find_first( where={"policy_name": policy_name}, order={"version_number": "desc"}, ) @@ -722,7 +723,7 @@ class PolicyRegistry: now = datetime.now(timezone.utc) # Set is_latest=False on all existing versions for this policy_name - await prisma_client.db.litellm_policytable.update_many( + await PolicyRepository(prisma_client).table.update_many( where={"policy_name": policy_name}, data={"is_latest": False}, ) @@ -758,7 +759,7 @@ class PolicyRegistry: else source.pipeline ) - created = await prisma_client.db.litellm_policytable.create(data=data) + created = await PolicyRepository(prisma_client).table.create(data=data) return _row_to_policy_db_response(created) except Exception as e: verbose_proxy_logger.exception(f"Error creating new version: {e}") @@ -794,7 +795,7 @@ class PolicyRegistry: f"Invalid status '{new_status}'. Use 'published' or 'production'." ) - row = await prisma_client.db.litellm_policytable.find_unique( + row = await PolicyRepository(prisma_client).table.find_unique( where={"policy_id": policy_id} ) if row is None: @@ -809,7 +810,7 @@ class PolicyRegistry: raise Exception( f"Only draft versions can be published. Current status: '{current}'." ) - updated = await prisma_client.db.litellm_policytable.update( + updated = await PolicyRepository(prisma_client).table.update( where={"policy_id": policy_id}, data={ "version_status": "published", @@ -832,7 +833,7 @@ class PolicyRegistry: ) # Demote current production to published - await prisma_client.db.litellm_policytable.update_many( + await PolicyRepository(prisma_client).table.update_many( where={ "policy_name": policy_name, "version_status": "production", @@ -845,7 +846,7 @@ class PolicyRegistry: ) # Promote this version to production - updated = await prisma_client.db.litellm_policytable.update( + updated = await PolicyRepository(prisma_client).table.update( where={"policy_id": policy_id}, data={ "version_status": "production", @@ -895,10 +896,10 @@ class PolicyRegistry: PolicyVersionCompareResponse with both versions and field_diffs """ try: - a = await prisma_client.db.litellm_policytable.find_unique( + a = await PolicyRepository(prisma_client).table.find_unique( where={"policy_id": policy_id_a} ) - b = await prisma_client.db.litellm_policytable.find_unique( + b = await PolicyRepository(prisma_client).table.find_unique( where={"policy_id": policy_id_b} ) if a is None: @@ -950,7 +951,7 @@ class PolicyRegistry: Dict with success message """ try: - await prisma_client.db.litellm_policytable.delete_many( + await PolicyRepository(prisma_client).table.delete_many( where={"policy_name": policy_name} ) self.remove_policy(policy_name) diff --git a/litellm/proxy/policy_engine/policy_resolve_endpoints.py b/litellm/proxy/policy_engine/policy_resolve_endpoints.py index 54374d90a16..84dcbcfd746 100644 --- a/litellm/proxy/policy_engine/policy_resolve_endpoints.py +++ b/litellm/proxy/policy_engine/policy_resolve_endpoints.py @@ -16,6 +16,10 @@ from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry from litellm.proxy.policy_engine.policy_registry import get_policy_registry +from litellm.repositories.team_repository import TeamRepository +from litellm.repositories.verification_token_repository import ( + VerificationTokenRepository, +) from litellm.types.proxy.policy_engine import ( AttachmentImpactResponse, PolicyAttachmentCreateRequest, @@ -76,7 +80,7 @@ def _get_tags_from_metadata(metadata: object, json_metadata: object = None) -> l async def _fetch_all_teams(prisma_client: object) -> list: """Fetch teams from DB once. Reuse the result across tag and alias lookups.""" - return await prisma_client.db.litellm_teamtable.find_many( # type: ignore + return await TeamRepository(prisma_client).table.find_many( # type: ignore where={}, order={"created_at": "desc"}, take=MAX_POLICY_ESTIMATE_IMPACT_ROWS, @@ -159,7 +163,7 @@ async def _find_affected_by_team_patterns( new_keys: list = [] unnamed_keys_count = 0 if matched_team_ids: - keys = await prisma_client.db.litellm_verificationtoken.find_many( # type: ignore + keys = await VerificationTokenRepository(prisma_client).table.find_many( # type: ignore where={"team_id": {"in": matched_team_ids}}, order={"created_at": "desc"}, take=MAX_POLICY_ESTIMATE_IMPACT_ROWS, @@ -182,7 +186,7 @@ async def _find_affected_keys_by_alias( affected: list = [] - keys = await prisma_client.db.litellm_verificationtoken.find_many( # type: ignore + keys = await VerificationTokenRepository(prisma_client).table.find_many( # type: ignore where=_build_alias_where("key_alias", key_patterns), order={"created_at": "desc"}, take=MAX_POLICY_ESTIMATE_IMPACT_ROWS, @@ -367,7 +371,7 @@ async def estimate_attachment_impact( # Tag-based impact if tag_patterns: - keys = await prisma_client.db.litellm_verificationtoken.find_many( # type: ignore + keys = await VerificationTokenRepository(prisma_client).table.find_many( # type: ignore where={}, order={"created_at": "desc"}, take=MAX_POLICY_ESTIMATE_IMPACT_ROWS, diff --git a/litellm/proxy/policy_engine/policy_validator.py b/litellm/proxy/policy_engine/policy_validator.py index b587e3432bb..46796fbae28 100644 --- a/litellm/proxy/policy_engine/policy_validator.py +++ b/litellm/proxy/policy_engine/policy_validator.py @@ -12,6 +12,10 @@ Validates: from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set from litellm._logging import verbose_proxy_logger +from litellm.repositories.team_repository import TeamRepository +from litellm.repositories.verification_token_repository import ( + VerificationTokenRepository, +) from litellm.types.proxy.policy_engine import ( Policy, PolicyValidationError, @@ -95,7 +99,7 @@ class PolicyValidator: return True # Can't validate without DB, assume valid try: - team = await self.prisma_client.db.litellm_teamtable.find_first( + team = await TeamRepository(self.prisma_client).table.find_first( where={"team_alias": team_alias}, ) return team is not None @@ -119,7 +123,9 @@ class PolicyValidator: return True # Can't validate without DB, assume valid try: - key = await self.prisma_client.db.litellm_verificationtoken.find_first( + key = await VerificationTokenRepository( + self.prisma_client + ).table.find_first( where={"key_alias": key_alias}, ) return key is not None diff --git a/litellm/proxy/prompts/prompt_endpoints.py b/litellm/proxy/prompts/prompt_endpoints.py index 399a0ff3af7..c0d6794108a 100644 --- a/litellm/proxy/prompts/prompt_endpoints.py +++ b/litellm/proxy/prompts/prompt_endpoints.py @@ -22,6 +22,7 @@ from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKey from litellm.proxy.auth.auth_utils import is_request_body_safe from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.path_utils import safe_filename +from litellm.repositories.table_repositories import PromptRepository from litellm.types.prompts.init_prompts import ( ListPromptsResponse, PromptInfo, @@ -208,7 +209,7 @@ async def get_next_version_for_prompt( Returns: Next version number (1 if no versions exist, max_version + 1 otherwise) """ - existing_prompts = await prisma_client.db.litellm_prompttable.find_many( + existing_prompts = await PromptRepository(prisma_client).table.find_many( where={"prompt_id": prompt_id, "environment": environment} ) @@ -441,7 +442,7 @@ async def get_prompt_versions( where_clause: Dict[str, Any] = {"prompt_id": base_prompt_id} if environment: where_clause["environment"] = environment - db_prompts = await prisma_client.db.litellm_prompttable.find_many( + db_prompts = await PromptRepository(prisma_client).table.find_many( where=where_clause, order={"version": "desc"}, ) @@ -612,7 +613,7 @@ async def get_prompt_info( # Query all environments this prompt exists in (lightweight: distinct on environment) all_environments: List[str] = [] if prisma_client is not None: - all_prompt_rows = await prisma_client.db.litellm_prompttable.find_many( + all_prompt_rows = await PromptRepository(prisma_client).table.find_many( where={"prompt_id": base_prompt_id}, distinct=["environment"], ) @@ -634,7 +635,7 @@ async def get_prompt_info( } if requested_version is not None: where_clause["version"] = requested_version - env_prompts = await prisma_client.db.litellm_prompttable.find_many( + env_prompts = await PromptRepository(prisma_client).table.find_many( where=where_clause, order={"version": "desc"}, take=1, @@ -752,7 +753,7 @@ async def create_prompt( ) # Store prompt in db with version - prompt_db_entry = await prisma_client.db.litellm_prompttable.create( + prompt_db_entry = await PromptRepository(prisma_client).table.create( data={ "prompt_id": request.prompt_id, "version": new_version, @@ -848,7 +849,7 @@ async def update_prompt( ) # Check if any version of this prompt exists (in any environment) - existing_prompts = await prisma_client.db.litellm_prompttable.find_many( + existing_prompts = await PromptRepository(prisma_client).table.find_many( where={"prompt_id": base_prompt_id} ) @@ -877,7 +878,7 @@ async def update_prompt( ) # Store new version in db - prompt_db_entry = await prisma_client.db.litellm_prompttable.create( + prompt_db_entry = await PromptRepository(prisma_client).table.create( data={ "prompt_id": base_prompt_id, "version": new_version, @@ -993,7 +994,7 @@ async def delete_prompt( delete_where["environment"] = environment # Delete versions from the database (scoped to environment if provided) - await prisma_client.db.litellm_prompttable.delete_many(where=delete_where) + await PromptRepository(prisma_client).table.delete_many(where=delete_where) # Remove matching prompts from memory — scope to environment if provided if environment: @@ -1105,7 +1106,7 @@ async def patch_prompt( if requested_version is not None: find_where["version"] = requested_version - db_rows = await prisma_client.db.litellm_prompttable.find_many( + db_rows = await PromptRepository(prisma_client).table.find_many( where=find_where, order={"version": "desc"}, take=1, @@ -1163,7 +1164,7 @@ async def patch_prompt( update_data["created_by"] = user_api_key_dict.user_id # Update by primary key (id) to target exactly one row - updated_prompt_db_entry = await prisma_client.db.litellm_prompttable.update( + updated_prompt_db_entry = await PromptRepository(prisma_client).table.update( where={"id": target_row.id}, data=update_data, ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 72423b2a796..2c21e19dcec 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -48,6 +48,7 @@ from litellm.constants import ( AIOHTTP_TTL_DNS_CACHE, AUDIO_SPEECH_CHUNK_SIZE, BASE_MCP_ROUTE, + DAILY_TAG_SPEND_BATCH_MULTIPLIER, DEFAULT_MAX_RECURSE_DEPTH, DEFAULT_SHARED_HEALTH_CHECK_LOCK_TTL, DEFAULT_SHARED_HEALTH_CHECK_TTL, @@ -56,13 +57,13 @@ from litellm.constants import ( LITELLM_SETTINGS_SAFE_DB_OVERRIDES, LITELLM_UI_ALLOW_HEADERS, LITELLM_UI_SESSION_DURATION, - DAILY_TAG_SPEND_BATCH_MULTIPLIER, ) from litellm.litellm_core_utils.litellm_logging import ( _init_custom_logger_compatible_class, ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import ( + UI_TEAM_ID, CallbackDelete, CallInfo, CommonProxyErrors, @@ -79,8 +80,8 @@ from litellm.proxy._types import ( InvitationModel, InvitationNew, InvitationUpdate, - Litellm_EntityType, LiteLLM_EndUserTable, + Litellm_EntityType, LiteLLM_JWTAuth, LiteLLM_TagTable, LiteLLM_TeamTable, @@ -96,7 +97,6 @@ from litellm.proxy._types import ( TeamDefaultSettings, TokenCountRequest, TransformRequestBody, - UI_TEAM_ID, UserAPIKeyAuth, ) from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec @@ -212,8 +212,6 @@ from litellm import Router from litellm._logging import verbose_proxy_logger, verbose_router_logger from litellm.caching.caching import DualCache, RedisCache from litellm.caching.redis_cluster_cache import RedisClusterCache -from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time -from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.constants import ( _REALTIME_BODY_CACHE_SIZE, APSCHEDULER_COALESCE, @@ -247,8 +245,8 @@ from litellm.litellm_core_utils.sensitive_data_masker import ( ) from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.vertex_ai.vertex_llm_base import VertexBase -from litellm.proxy._types import * from litellm.proxy._lazy_features import attach_lazy_features +from litellm.proxy._types import * from litellm.proxy.analytics_endpoints.analytics_endpoints import ( router as analytics_router, ) @@ -308,6 +306,8 @@ from litellm.proxy.common_utils.openai_endpoint_utils import ( from litellm.proxy.common_utils.proxy_state import ProxyState from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob from litellm.proxy.common_utils.swagger_utils import ERROR_RESPONSES +from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.container_endpoints.endpoints import router as container_router from litellm.proxy.credential_endpoints.endpoints import router as credential_router from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import SpendLogCleanup @@ -361,7 +361,9 @@ from litellm.proxy.management_endpoints.fallback_management_endpoints import ( from litellm.proxy.management_endpoints.internal_user_endpoints import ( router as internal_user_router, ) -from litellm.proxy.management_endpoints.internal_user_endpoints import user_update +from litellm.proxy.management_endpoints.internal_user_endpoints import ( + user_update, +) from litellm.proxy.management_endpoints.key_management_endpoints import ( delete_verification_tokens, duration_in_seconds, @@ -398,10 +400,6 @@ from litellm.proxy.management_endpoints.team_endpoints import ( update_team, validate_membership, ) -from litellm.proxy.management_endpoints.workflow_management_endpoints import ( - router as workflow_management_router, -) -from litellm.proxy.memory.memory_endpoints import router as memory_router from litellm.proxy.management_endpoints.ui_sso import ( get_disabled_non_admin_personal_key_creation, ) @@ -409,7 +407,11 @@ from litellm.proxy.management_endpoints.ui_sso import router as ui_sso_router from litellm.proxy.management_endpoints.user_agent_analytics_endpoints import ( router as user_agent_analytics_router, ) +from litellm.proxy.management_endpoints.workflow_management_endpoints import ( + router as workflow_management_router, +) from litellm.proxy.management_helpers.audit_logs import create_audit_log_for_update +from litellm.proxy.memory.memory_endpoints import router as memory_router from litellm.proxy.middleware.in_flight_requests_middleware import ( InFlightRequestsMiddleware, ) @@ -417,12 +419,13 @@ from litellm.proxy.middleware.prometheus_auth_middleware import PrometheusAuthMi from litellm.proxy.middleware.request_size_limit_middleware import ( RequestSizeLimitMiddleware, ) -from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager from litellm.proxy.ocr_endpoints.endpoints import router as ocr_router from litellm.proxy.openai_files_endpoints.files_endpoints import ( router as openai_files_router, ) -from litellm.proxy.openai_files_endpoints.files_endpoints import set_files_config +from litellm.proxy.openai_files_endpoints.files_endpoints import ( + set_files_config, +) from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( passthrough_endpoint_router, ) @@ -444,6 +447,7 @@ from litellm.proxy.rerank_endpoints.endpoints import router as rerank_router from litellm.proxy.response_api_endpoints.endpoints import router as response_router from litellm.proxy.route_llm_request import route_request from litellm.proxy.search_endpoints.endpoints import router as search_router +from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager from litellm.proxy.spend_tracking.spend_management_endpoints import ( router as spend_management_router, ) @@ -478,6 +482,7 @@ from litellm.proxy.utils import ( update_spend, ) from litellm.proxy.video_endpoints.endpoints import router as video_router +from litellm.repositories.credentials_repository import CredentialsRepository from litellm.router import ( AssistantsTypedDict, Deployment, @@ -511,7 +516,9 @@ from litellm.types.proxy.management_endpoints.ui_sso import ( LiteLLM_UpperboundKeyGenerateParams, ) from litellm.types.realtime import RealtimeQueryParams -from litellm.types.router import DeploymentTypedDict +from litellm.types.router import ( + DeploymentTypedDict, +) from litellm.types.router import ModelInfo as RouterModelInfo from litellm.types.router import ( RouterGeneralSettings, @@ -4248,13 +4255,13 @@ class ProxyConfig: ) setattr(litellm, key, value) if key in {"s3_audit_callback_params", "s3_callback_params"}: - from litellm.proxy.management_helpers.audit_logs import ( - reset_audit_log_callback_cache, - ) + from litellm.integrations.s3_v2 import S3Logger as S3V2Logger from litellm.litellm_core_utils.litellm_logging import ( _in_memory_loggers, ) - from litellm.integrations.s3_v2 import S3Logger as S3V2Logger + from litellm.proxy.management_helpers.audit_logs import ( + reset_audit_log_callback_cache, + ) reset_audit_log_callback_cache() _in_memory_loggers[:] = [ @@ -5335,7 +5342,7 @@ class ProxyConfig: 4. Update router settings """ if llm_router is not None and prisma_client is not None: - db_router_settings = await prisma_client.db.litellm_config.find_first( + db_router_settings = await ConfigRepository(prisma_client).table.find_first( where={"param_name": "router_settings"} ) @@ -5761,7 +5768,7 @@ class ProxyConfig: async def _get_models_from_db(self, prisma_client: PrismaClient) -> list: try: - new_models = await prisma_client.db.litellm_proxymodeltable.find_many() + new_models = await ModelRepository(prisma_client).table.find_many() except Exception as e: verbose_proxy_logger.exception( "litellm.proxy_server.py::add_deployment() - Error getting new models from DB - {}".format( @@ -5975,7 +5982,7 @@ class ProxyConfig: """ try: - sso_settings = await prisma_client.db.litellm_ssoconfig.find_unique( + sso_settings = await SSOConfigRepository(prisma_client).table.find_unique( where={"id": "sso_config"} ) if sso_settings is not None: @@ -6011,9 +6018,9 @@ class ProxyConfig: ) try: - db_record = await prisma_client.db.litellm_configoverrides.find_unique( - where={"config_type": "hashicorp_vault"} - ) + db_record = await ConfigOverridesRepository( + prisma_client + ).table.find_unique(where={"config_type": "hashicorp_vault"}) if db_record is None or db_record.config_value is None: if self._last_hashicorp_vault_config is not None: @@ -6130,7 +6137,7 @@ class ProxyConfig: last_model_cost_map_reload = current_time.isoformat() # Clear force reload flag in database - await prisma_client.db.litellm_config.upsert( + await ConfigRepository(prisma_client).table.upsert( where={"param_name": "model_cost_map_reload_config"}, data={ "create": { @@ -6239,7 +6246,7 @@ class ProxyConfig: last_anthropic_beta_headers_reload = current_time.isoformat() # Clear force reload flag in database - await prisma_client.db.litellm_config.upsert( + await ConfigRepository(prisma_client).table.upsert( where={"param_name": "anthropic_beta_headers_reload_config"}, data={ "create": { @@ -6299,7 +6306,7 @@ class ProxyConfig: from litellm.types.prompts.init_prompts import PromptSpec try: - prompts_in_db = await prisma_client.db.litellm_prompttable.find_many() + prompts_in_db = await PromptRepository(prisma_client).table.find_many() for prompt in prompts_in_db: # Convert DB object to dict and create versioned prompt_id prompt_spec = self._get_prompt_spec_for_db_prompt(db_prompt=prompt) @@ -6589,7 +6596,7 @@ class ProxyConfig: async def get_credentials(self, prisma_client: PrismaClient): try: - credentials = await prisma_client.db.litellm_credentialstable.find_many() + credentials = await CredentialsRepository(prisma_client).find_all() credentials = [self.decrypt_credentials(cred) for cred in credentials] await self.delete_credentials( credentials @@ -7352,7 +7359,7 @@ class ProxyStartupEvent: # spend cap blocks forever once it's hit. if prisma_client is not None and litellm.budget_duration is not None: try: - await prisma_client.db.litellm_usertable.update_many( + await UserRepository(prisma_client).table.update_many( where={ "user_id": litellm_proxy_budget_name, "budget_reset_at": None, @@ -7420,7 +7427,7 @@ class ProxyStartupEvent: if prisma_client is None: return - db_record = await prisma_client.db.litellm_uisettings.find_unique( + db_record = await UISettingsRepository(prisma_client).table.find_unique( where={"id": "ui_settings"} ) if db_record and db_record.ui_settings: @@ -7569,7 +7576,7 @@ class ProxyStartupEvent: # but YAML config has False. if store_model_in_db is not True and prisma_client is not None: try: - _db_gs_record = await prisma_client.db.litellm_config.find_first( + _db_gs_record = await ConfigRepository(prisma_client).table.find_first( where={"param_name": "general_settings"} ) if _db_gs_record is not None and isinstance( @@ -10361,6 +10368,18 @@ async def run_thread( # ) # async def get_available_routes(user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth)): from litellm.llms.base_llm.base_utils import BaseTokenCounter +from litellm.repositories.config_repository import ConfigRepository +from litellm.repositories.model_repository import ModelRepository +from litellm.repositories.table_repositories import ( + AccessGroupRepository, + ConfigOverridesRepository, + InvitationLinkRepository, + PromptRepository, + SSOConfigRepository, + UISettingsRepository, +) +from litellm.repositories.team_repository import TeamRepository +from litellm.repositories.user_repository import UserRepository def _get_provider_token_counter( @@ -10666,7 +10685,7 @@ async def _check_if_model_is_user_added( id = model.get("model_info", {}).get("id", None) if id is None: continue - db_model = await prisma_client.db.litellm_proxymodeltable.find_unique( + db_model = await ModelRepository(prisma_client).table.find_unique( where={"model_id": id} ) if db_model is not None: @@ -10723,7 +10742,7 @@ async def non_admin_all_models( if user_api_key_dict.user_id: try: - user_row = await prisma_client.db.litellm_usertable.find_unique( + user_row = await UserRepository(prisma_client).table.find_unique( where={"user_id": user_api_key_dict.user_id} ) except Exception: @@ -10823,7 +10842,7 @@ async def _add_access_group_models_to_team_models( return team_models # Single batch fetch for all access groups - access_group_rows = await prisma_client.db.litellm_accessgrouptable.find_many( + access_group_rows = await AccessGroupRepository(prisma_client).table.find_many( where={"access_group_id": {"in": list(all_access_group_ids)}} ) ag_model_map: Dict[str, List[str]] = { @@ -10865,13 +10884,13 @@ async def get_all_team_models( team_db_objects_typed: List[LiteLLM_TeamTable] = [] if user_teams == "*": - team_db_objects = await prisma_client.db.litellm_teamtable.find_many() + team_db_objects = await TeamRepository(prisma_client).table.find_many() team_db_objects_typed = [ LiteLLM_TeamTable(**team_db_object.model_dump()) for team_db_object in team_db_objects ] else: - team_db_objects = await prisma_client.db.litellm_teamtable.find_many( + team_db_objects = await TeamRepository(prisma_client).table.find_many( where={"team_id": {"in": user_teams}} ) @@ -10938,7 +10957,7 @@ async def get_all_team_and_direct_access_models( exclude_team_models=True ) # has access to all models elif user_api_key_dict.user_id is not None: - user_db_object = await prisma_client.db.litellm_usertable.find_unique( + user_db_object = await UserRepository(prisma_client).table.find_unique( where={"user_id": user_api_key_dict.user_id} ) if user_db_object is not None: @@ -11082,7 +11101,7 @@ async def _get_caller_byok_team_scope( if user_id is None: return set() try: - user_row = await prisma_client.db.litellm_usertable.find_unique( + user_row = await UserRepository(prisma_client).table.find_unique( where={"user_id": user_id} ) except Exception: @@ -11143,13 +11162,13 @@ async def _fetch_db_models_for_search( else: take_limit = max(0, page * size - router_models_count) - db_models_total_count = await prisma_client.db.litellm_proxymodeltable.count( + db_models_total_count = await ModelRepository(prisma_client).table.count( where=db_where_condition ) db_models_raw: list = [] if take_limit > 0: - db_models_raw = await prisma_client.db.litellm_proxymodeltable.find_many( + db_models_raw = await ModelRepository(prisma_client).table.find_many( where=db_where_condition, take=take_limit, ) @@ -11484,7 +11503,7 @@ async def _load_team_object_for_model_filter( ) -> Optional[LiteLLM_TeamTable]: """Load team row from DB; returns None if missing or on error.""" try: - team_db_object = await prisma_client.db.litellm_teamtable.find_unique( + team_db_object = await TeamRepository(prisma_client).table.find_unique( where={"team_id": team_id} ) if team_db_object is None: @@ -11547,7 +11566,7 @@ async def _gather_team_accessible_model_ids( _resolved_names = _team_models_resolve_to_names( team_object.models, access_groups ) - db_models = await prisma_client.db.litellm_proxymodeltable.find_many( + db_models = await ModelRepository(prisma_client).table.find_many( where={"model_name": {"in": _resolved_names}} ) for db_model in db_models: @@ -11586,7 +11605,7 @@ async def _authorize_team_id_query( detail={"error": "Not authorized to view this team's models"}, ) try: - user_row = await prisma_client.db.litellm_usertable.find_unique( + user_row = await UserRepository(prisma_client).table.find_unique( where={"user_id": user_id} ) except Exception: @@ -11693,7 +11712,7 @@ async def _find_model_by_id( # If not found in config, search in database if found_model is None: try: - db_model = await prisma_client.db.litellm_proxymodeltable.find_unique( + db_model = await ModelRepository(prisma_client).table.find_unique( where={"model_id": model_id} ) if db_model: @@ -12845,7 +12864,7 @@ async def alerting_settings( ) ## get general settings from db - db_general_settings = await prisma_client.db.litellm_config.find_first( + db_general_settings = await ConfigRepository(prisma_client).table.find_first( where={"param_name": "general_settings"} ) @@ -13384,7 +13403,7 @@ async def onboarding(invite_link: str, request: Request): detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - invite_obj = await prisma_client.db.litellm_invitationlink.find_unique( + invite_obj = await InvitationLinkRepository(prisma_client).table.find_unique( where={"id": invite_link} ) if invite_obj is None: @@ -13408,7 +13427,7 @@ async def onboarding(invite_link: str, request: Request): ) ### GET USER OBJECT ### - user_obj = await prisma_client.db.litellm_usertable.find_unique( + user_obj = await UserRepository(prisma_client).table.find_unique( where={"user_id": invite_obj.user_id} ) @@ -13513,7 +13532,7 @@ async def _rollback_onboarding_invite_claim( return try: - await prisma_client.db.litellm_invitationlink.update_many( + await InvitationLinkRepository(prisma_client).table.update_many( where={"id": invitation_link, "is_accepted": True}, data={ "accepted_at": None, @@ -13547,10 +13566,10 @@ async def _generate_onboarding_ui_session_token(user_obj: Any) -> str: ) key = response["token"] # type: ignore - from litellm.types.proxy.ui_sso import ReturnedUITokenObject - import jwt + from litellm.types.proxy.ui_sso import ReturnedUITokenObject + disabled_non_admin_personal_key_creation = ( get_disabled_non_admin_personal_key_creation() ) @@ -13596,7 +13615,7 @@ async def claim_onboarding_link(data: InvitationClaim, request: Request): detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - invite_obj = await prisma_client.db.litellm_invitationlink.find_unique( + invite_obj = await InvitationLinkRepository(prisma_client).table.find_unique( where={"id": data.invitation_link} ) if invite_obj is None: @@ -13956,7 +13975,7 @@ async def invitation_info( }, ) - response = await prisma_client.db.litellm_invitationlink.find_unique( + response = await InvitationLinkRepository(prisma_client).table.find_unique( where={"id": invitation_id} ) @@ -14010,7 +14029,7 @@ async def invitation_update( ) current_time = litellm.utils.get_utc_datetime() - response = await prisma_client.db.litellm_invitationlink.update( + response = await InvitationLinkRepository(prisma_client).table.update( where={"id": data.invitation_id}, data={ "id": data.invitation_id, @@ -14081,7 +14100,7 @@ async def invitation_delete( # Org admins can only delete invitations they created if is_other_admin and not is_proxy_admin: - invitation = await prisma_client.db.litellm_invitationlink.find_unique( + invitation = await InvitationLinkRepository(prisma_client).table.find_unique( where={"id": data.invitation_id} ) if invitation is None: @@ -14097,7 +14116,7 @@ async def invitation_delete( }, ) - response = await prisma_client.db.litellm_invitationlink.delete( + response = await InvitationLinkRepository(prisma_client).table.delete( where={"id": data.invitation_id} ) @@ -14139,7 +14158,7 @@ async def update_config( # noqa: PLR0915 raise Exception("No DB Connected") async def _read_section(param_name: str) -> dict: - row = await prisma_client.db.litellm_config.find_first( + row = await ConfigRepository(prisma_client).table.find_first( where={"param_name": param_name} ) if row is None or row.param_value is None: @@ -14148,7 +14167,7 @@ async def update_config( # noqa: PLR0915 async def _upsert_section(param_name: str, value: dict) -> None: serialized = json.dumps(value) - await prisma_client.db.litellm_config.upsert( + await ConfigRepository(prisma_client).table.upsert( where={"param_name": param_name}, data={ "create": {"param_name": param_name, "param_value": serialized}, @@ -14323,7 +14342,7 @@ async def update_config_general_settings( ) ## get general settings from db - db_general_settings = await prisma_client.db.litellm_config.find_first( + db_general_settings = await ConfigRepository(prisma_client).table.find_first( where={"param_name": "general_settings"} ) ### update value @@ -14337,7 +14356,7 @@ async def update_config_general_settings( general_settings[data.field_name] = data.field_value - response = await prisma_client.db.litellm_config.upsert( + response = await ConfigRepository(prisma_client).table.upsert( where={"param_name": "general_settings"}, data={ "create": {"param_name": "general_settings", "param_value": json.dumps(general_settings)}, # type: ignore @@ -14387,7 +14406,7 @@ async def get_config_general_settings( ) ## get general settings from db - db_general_settings = await prisma_client.db.litellm_config.find_first( + db_general_settings = await ConfigRepository(prisma_client).table.find_first( where={"param_name": "general_settings"} ) ### pop the value @@ -14450,7 +14469,7 @@ async def get_config_list( ) ## get general settings from db - db_general_settings = await prisma_client.db.litellm_config.find_first( + db_general_settings = await ConfigRepository(prisma_client).table.find_first( where={"param_name": "general_settings"} ) @@ -14604,7 +14623,7 @@ async def delete_config_general_settings( ) ## get general settings from db - db_general_settings = await prisma_client.db.litellm_config.find_first( + db_general_settings = await ConfigRepository(prisma_client).table.find_first( where={"param_name": "general_settings"} ) ### pop the value @@ -14621,7 +14640,7 @@ async def delete_config_general_settings( general_settings.pop(data.field_name, None) - response = await prisma_client.db.litellm_config.upsert( + response = await ConfigRepository(prisma_client).table.upsert( where={"param_name": "general_settings"}, data={ "create": {"param_name": "general_settings", "param_value": json.dumps(general_settings)}, # type: ignore @@ -14975,14 +14994,14 @@ async def reload_model_cost_map( last_model_cost_map_reload = current_time.isoformat() # Set force reload flag in database for other pods, preserving existing interval_hours - existing_config = await prisma_client.db.litellm_config.find_unique( + existing_config = await ConfigRepository(prisma_client).table.find_unique( where={"param_name": "model_cost_map_reload_config"} ) existing_interval = None if existing_config and existing_config.param_value: existing_interval = existing_config.param_value.get("interval_hours") - await prisma_client.db.litellm_config.upsert( + await ConfigRepository(prisma_client).table.upsert( where={"param_name": "model_cost_map_reload_config"}, data={ "create": { @@ -15052,7 +15071,7 @@ async def schedule_model_cost_map_reload( ) # Update database with new reload configuration - await prisma_client.db.litellm_config.upsert( + await ConfigRepository(prisma_client).table.upsert( where={"param_name": "model_cost_map_reload_config"}, data={ "create": { @@ -15119,7 +15138,7 @@ async def cancel_model_cost_map_reload( ) # Remove reload configuration from database - await prisma_client.db.litellm_config.delete( + await ConfigRepository(prisma_client).table.delete( where={"param_name": "model_cost_map_reload_config"} ) await invalidate_config_param("model_cost_map_reload_config") @@ -15178,7 +15197,7 @@ async def get_model_cost_map_reload_status( } # Get reload configuration from database - config_record = await prisma_client.db.litellm_config.find_unique( + config_record = await ConfigRepository(prisma_client).table.find_unique( where={"param_name": "model_cost_map_reload_config"} ) @@ -15329,7 +15348,7 @@ async def reload_anthropic_beta_headers( last_anthropic_beta_headers_reload = current_time.isoformat() # Set force reload flag in database for other pods, preserving existing interval_hours - existing_beta_config = await prisma_client.db.litellm_config.find_unique( + existing_beta_config = await ConfigRepository(prisma_client).table.find_unique( where={"param_name": "anthropic_beta_headers_reload_config"} ) existing_beta_interval = None @@ -15338,7 +15357,7 @@ async def reload_anthropic_beta_headers( "interval_hours" ) - await prisma_client.db.litellm_config.upsert( + await ConfigRepository(prisma_client).table.upsert( where={"param_name": "anthropic_beta_headers_reload_config"}, data={ "create": { @@ -15412,7 +15431,7 @@ async def schedule_anthropic_beta_headers_reload( ) # Update database with new reload configuration - await prisma_client.db.litellm_config.upsert( + await ConfigRepository(prisma_client).table.upsert( where={"param_name": "anthropic_beta_headers_reload_config"}, data={ "create": { @@ -15479,7 +15498,7 @@ async def cancel_anthropic_beta_headers_reload( ) # Remove reload configuration from database - await prisma_client.db.litellm_config.delete( + await ConfigRepository(prisma_client).table.delete( where={"param_name": "anthropic_beta_headers_reload_config"} ) await invalidate_config_param("anthropic_beta_headers_reload_config") @@ -15539,7 +15558,7 @@ async def get_anthropic_beta_headers_reload_status( } # Get reload configuration from database - config_record = await prisma_client.db.litellm_config.find_unique( + config_record = await ConfigRepository(prisma_client).table.find_unique( where={"param_name": "anthropic_beta_headers_reload_config"} ) diff --git a/litellm/proxy/public_endpoints/public_endpoints.py b/litellm/proxy/public_endpoints/public_endpoints.py index d12e7a35fbf..78467c4b2e7 100644 --- a/litellm/proxy/public_endpoints/public_endpoints.py +++ b/litellm/proxy/public_endpoints/public_endpoints.py @@ -4,9 +4,9 @@ import re from importlib.resources import files from typing import Any, Dict, List, Optional -import litellm from fastapi import APIRouter, HTTPException, Request +import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.get_blog_posts import ( BlogPost, @@ -17,6 +17,7 @@ from litellm.litellm_core_utils.get_blog_posts import ( from litellm.proxy._types import ( CommonProxyErrors, ) +from litellm.repositories.table_repositories import ClaudeCodePluginRepository from litellm.types.agents import AgentCard from litellm.types.mcp import MCPPublicServer from litellm.types.proxy.management_endpoints.model_management_endpoints import ( @@ -159,14 +160,14 @@ def _load_endpoints() -> List[Dict[str, Any]]: ) async def public_model_hub(): import litellm + from litellm.proxy.health_endpoints._health_endpoints import ( + _convert_health_check_to_dict, + ) from litellm.proxy.proxy_server import ( _get_model_group_info, llm_router, prisma_client, ) - from litellm.proxy.health_endpoints._health_endpoints import ( - _convert_health_check_to_dict, - ) if llm_router is None: raise HTTPException( @@ -266,7 +267,7 @@ async def public_skill_hub(): try: prisma_client = await _get_prisma_client() - plugins = await prisma_client.db.litellm_claudecodeplugintable.find_many( + plugins = await ClaudeCodePluginRepository(prisma_client).table.find_many( where={"enabled": True} ) items = [] diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index a44e4781491..7ff54ac4c5a 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -17,16 +17,17 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH from litellm.proxy._types import * +from litellm.proxy.auth.auth_utils import is_request_body_safe from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth from litellm.proxy.common_utils.http_parsing_utils import ( _read_request_body, _safe_get_request_headers, get_form_data, ) -from litellm.proxy.auth.auth_utils import is_request_body_safe from litellm.proxy.vector_store_endpoints.utils import ( assert_user_can_access_vector_store_id, ) +from litellm.repositories.table_repositories import ManagedVectorStoresRepository router = APIRouter() @@ -230,11 +231,9 @@ async def _save_vector_store_to_db_from_rag_ingest( try: # Check if vector store already exists in database - existing_vector_store = ( - await prisma_client.db.litellm_managedvectorstorestable.find_unique( - where={"vector_store_id": vector_store_id} - ) - ) + existing_vector_store = await ManagedVectorStoresRepository( + prisma_client + ).table.find_unique(where={"vector_store_id": vector_store_id}) # Only create if it doesn't exist if existing_vector_store is None: @@ -289,7 +288,7 @@ async def _save_vector_store_to_db_from_rag_ingest( # Update the vector store from litellm.proxy.utils import safe_dumps - await prisma_client.db.litellm_managedvectorstorestable.update( + await ManagedVectorStoresRepository(prisma_client).table.update( where={"vector_store_id": vector_store_id}, data={"vector_store_metadata": safe_dumps(existing_metadata)}, ) diff --git a/litellm/proxy/search_endpoints/search_tool_registry.py b/litellm/proxy/search_endpoints/search_tool_registry.py index d4adc2573ea..588d71b77f9 100644 --- a/litellm/proxy/search_endpoints/search_tool_registry.py +++ b/litellm/proxy/search_endpoints/search_tool_registry.py @@ -8,6 +8,7 @@ from typing import List, Optional from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy.utils import PrismaClient +from litellm.repositories.table_repositories import SearchToolsRepository from litellm.types.search import SearchTool @@ -63,16 +64,16 @@ class SearchToolRegistry: search_tool_info: str = safe_dumps(search_tool.get("search_tool_info", {})) # Create search tool in DB - created_search_tool = ( - await prisma_client.db.litellm_searchtoolstable.create( - data={ - "search_tool_name": search_tool_name, - "litellm_params": litellm_params, - "search_tool_info": search_tool_info, - "created_at": datetime.now(timezone.utc), - "updated_at": datetime.now(timezone.utc), - } - ) + created_search_tool = await SearchToolsRepository( + prisma_client + ).table.create( + data={ + "search_tool_name": search_tool_name, + "litellm_params": litellm_params, + "search_tool_info": search_tool_info, + "created_at": datetime.now(timezone.utc), + "updated_at": datetime.now(timezone.utc), + } ) # Add search_tool_id to the returned search tool object @@ -101,15 +102,15 @@ class SearchToolRegistry: """ try: # Get search tool before deletion for response - existing_tool = await prisma_client.db.litellm_searchtoolstable.find_unique( - where={"search_tool_id": search_tool_id} - ) + existing_tool = await SearchToolsRepository( + prisma_client + ).table.find_unique(where={"search_tool_id": search_tool_id}) if not existing_tool: raise Exception(f"Search tool with ID {search_tool_id} not found") # Delete from DB - await prisma_client.db.litellm_searchtoolstable.delete( + await SearchToolsRepository(prisma_client).table.delete( where={"search_tool_id": search_tool_id} ) @@ -145,16 +146,16 @@ class SearchToolRegistry: search_tool_info: str = safe_dumps(search_tool.get("search_tool_info", {})) # Update in DB - updated_search_tool = ( - await prisma_client.db.litellm_searchtoolstable.update( - where={"search_tool_id": search_tool_id}, - data={ - "search_tool_name": search_tool_name, - "litellm_params": litellm_params, - "search_tool_info": search_tool_info, - "updated_at": datetime.now(timezone.utc), - }, - ) + updated_search_tool = await SearchToolsRepository( + prisma_client + ).table.update( + where={"search_tool_id": search_tool_id}, + data={ + "search_tool_name": search_tool_name, + "litellm_params": litellm_params, + "search_tool_info": search_tool_info, + "updated_at": datetime.now(timezone.utc), + }, ) # Convert to dict with ISO formatted datetimes @@ -179,10 +180,10 @@ class SearchToolRegistry: List of search tool configurations """ try: - search_tools_from_db = ( - await prisma_client.db.litellm_searchtoolstable.find_many( - order={"created_at": "desc"}, - ) + search_tools_from_db = await SearchToolsRepository( + prisma_client + ).table.find_many( + order={"created_at": "desc"}, ) search_tools: List[SearchTool] = [] @@ -214,7 +215,7 @@ class SearchToolRegistry: Search tool configuration or None if not found """ try: - search_tool = await prisma_client.db.litellm_searchtoolstable.find_unique( + search_tool = await SearchToolsRepository(prisma_client).table.find_unique( where={"search_tool_id": search_tool_id} ) @@ -244,7 +245,7 @@ class SearchToolRegistry: Search tool configuration or None if not found """ try: - search_tool = await prisma_client.db.litellm_searchtoolstable.find_unique( + search_tool = await SearchToolsRepository(prisma_client).table.find_unique( where={"search_tool_name": search_tool_name} ) diff --git a/litellm/proxy/spend_tracking/cloudzero_endpoints.py b/litellm/proxy/spend_tracking/cloudzero_endpoints.py index 1f551d5ffea..71f4a8af111 100644 --- a/litellm/proxy/spend_tracking/cloudzero_endpoints.py +++ b/litellm/proxy/spend_tracking/cloudzero_endpoints.py @@ -6,11 +6,12 @@ from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view from litellm.proxy.common_utils.encrypt_decrypt_utils import ( decrypt_value_helper, encrypt_value_helper, ) +from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view +from litellm.repositories.config_repository import ConfigRepository from litellm.types.proxy.cloudzero_endpoints import ( CloudZeroExportRequest, CloudZeroExportResponse, @@ -53,7 +54,7 @@ async def _set_cloudzero_settings(api_key: str, connection_id: str, timezone: st "timezone": timezone, } - await prisma_client.db.litellm_config.upsert( + await ConfigRepository(prisma_client).table.upsert( where={"param_name": "cloudzero_settings"}, data={ "create": { @@ -80,7 +81,7 @@ async def _get_cloudzero_settings(): detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - cloudzero_config = await prisma_client.db.litellm_config.find_first( + cloudzero_config = await ConfigRepository(prisma_client).table.find_first( where={"param_name": "cloudzero_settings"} ) if cloudzero_config is None or cloudzero_config.param_value is None: @@ -282,7 +283,7 @@ async def is_cloudzero_setup_in_db() -> bool: return False # Check for CloudZero settings in database - cloudzero_config = await prisma_client.db.litellm_config.find_first( + cloudzero_config = await ConfigRepository(prisma_client).table.find_first( where={"param_name": "cloudzero_settings"} ) @@ -548,7 +549,7 @@ async def delete_cloudzero_settings( ) # Check if CloudZero settings exist - cloudzero_config = await prisma_client.db.litellm_config.find_first( + cloudzero_config = await ConfigRepository(prisma_client).table.find_first( where={"param_name": "cloudzero_settings"} ) @@ -560,7 +561,7 @@ async def delete_cloudzero_settings( # Delete only the CloudZero settings entry # This uses a specific where clause to target only the cloudzero_settings row - await prisma_client.db.litellm_config.delete( + await ConfigRepository(prisma_client).table.delete( where={"param_name": "cloudzero_settings"} ) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index ca5e0473659..f651e6e5f7b 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -21,6 +21,11 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( get_spend_by_team_and_customer, ) from litellm.proxy.utils import handle_exception_on_proxy +from litellm.repositories.table_repositories import SpendLogsRepository +from litellm.repositories.team_repository import TeamRepository +from litellm.repositories.verification_token_repository import ( + VerificationTokenRepository, +) if TYPE_CHECKING: from litellm.proxy.proxy_server import PrismaClient @@ -2010,7 +2015,7 @@ async def ui_view_spend_logs( # noqa: PLR0915 order_direction = (sort_order or "desc").lower() # Get total count of records - total_records = await prisma_client.db.litellm_spendlogs.count( + total_records = await SpendLogsRepository(prisma_client).table.count( where=where_conditions, ) @@ -2374,7 +2379,7 @@ async def view_spend_logs( # noqa: PLR0915 # Check if user wants unsummarized data if not summarize: # Return filtered individual log entries (similar to UI endpoint) - data = await prisma_client.db.litellm_spendlogs.find_many( + data = await SpendLogsRepository(prisma_client).table.find_many( where=filter_query, # type: ignore order={ "startTime": "desc", @@ -2384,7 +2389,7 @@ async def view_spend_logs( # noqa: PLR0915 # Legacy behavior: return summarized data (when summarize=true) # SQL query - response = await prisma_client.db.litellm_spendlogs.group_by( + response = await SpendLogsRepository(prisma_client).table.group_by( by=["api_key", "user", "model", "startTime"], where=filter_query, # type: ignore sum={ @@ -2462,7 +2467,7 @@ async def view_spend_logs( # noqa: PLR0915 ) return spend_logs - data = await prisma_client.db.litellm_spendlogs.find_many( + data = await SpendLogsRepository(prisma_client).table.find_many( where=scoped_filter, # type: ignore order={"startTime": "desc"}, ) @@ -2514,10 +2519,10 @@ async def global_spend_reset(): code=status.HTTP_401_UNAUTHORIZED, ) - await prisma_client.db.litellm_verificationtoken.update_many( + await VerificationTokenRepository(prisma_client).table.update_many( data={"spend": 0.0}, where={} ) - await prisma_client.db.litellm_teamtable.update_many(data={"spend": 0.0}, where={}) + await TeamRepository(prisma_client).table.update_many(data={"spend": 0.0}, where={}) return { "message": "Spend for all API Keys and Teams reset successfully", @@ -3384,7 +3389,7 @@ async def ui_view_session_spend_logs( skip = (page - 1) * page_size # Get total count for pagination metadata - total_records = await prisma_client.db.litellm_spendlogs.count( + total_records = await SpendLogsRepository(prisma_client).table.count( where=where_conditions ) @@ -3485,7 +3490,7 @@ async def _build_ui_spend_logs_response( # is bounded by page_size (typically 25-50 distinct session IDs). # If performance degrades at scale, consider short-lived caching or # folding the count into the main query via a window function. - counts = await prisma_client.db.litellm_spendlogs.group_by( + counts = await SpendLogsRepository(prisma_client).table.group_by( by=["session_id"], where={"session_id": {"in": session_ids}}, count={"session_id": True}, @@ -3572,7 +3577,7 @@ async def _can_team_member_view_log( if team_id is None: return False - team_row = await prisma_client.db.litellm_teamtable.find_unique( + team_row = await TeamRepository(prisma_client).table.find_unique( where={"team_id": team_id} ) if team_row is None: @@ -3614,7 +3619,7 @@ async def _assert_user_can_view_request_id( permitted teams (admin or ``/spend/logs`` permission). Raises HTTP 403 if not. """ - row = await prisma_client.db.litellm_spendlogs.find_unique( + row = await SpendLogsRepository(prisma_client).table.find_unique( where={"request_id": request_id}, include=None, ) @@ -3669,7 +3674,7 @@ async def _get_permitted_team_ids_for_spend_logs( if user_obj is None or not user_obj.teams: return [] - team_rows = await prisma_client.db.litellm_teamtable.find_many( + team_rows = await TeamRepository(prisma_client).table.find_many( where={"team_id": {"in": user_obj.teams}} ) diff --git a/litellm/proxy/spend_tracking/vantage_endpoints.py b/litellm/proxy/spend_tracking/vantage_endpoints.py index 60e54d005b3..1dde31b54cb 100644 --- a/litellm/proxy/spend_tracking/vantage_endpoints.py +++ b/litellm/proxy/spend_tracking/vantage_endpoints.py @@ -1,17 +1,18 @@ import json -import litellm from fastapi import APIRouter, Depends, HTTPException +import litellm from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view from litellm.proxy.common_utils.encrypt_decrypt_utils import ( decrypt_value_helper, encrypt_value_helper, ) +from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view +from litellm.repositories.config_repository import ConfigRepository from litellm.types.proxy.vantage_endpoints import ( VantageDryRunRequest, VantageExportRequest, @@ -60,7 +61,7 @@ async def _set_vantage_settings(api_key: str, integration_token: str, base_url: "base_url": base_url, } - await prisma_client.db.litellm_config.upsert( + await ConfigRepository(prisma_client).table.upsert( where={"param_name": VANTAGE_SETTINGS_PARAM_NAME}, data={ "create": { @@ -82,7 +83,7 @@ async def _get_vantage_settings(): detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - vantage_config = await prisma_client.db.litellm_config.find_first( + vantage_config = await ConfigRepository(prisma_client).table.find_first( where={"param_name": VANTAGE_SETTINGS_PARAM_NAME} ) if vantage_config is None or vantage_config.param_value is None: @@ -265,7 +266,7 @@ async def is_vantage_setup_in_db() -> bool: if prisma_client is None: return False - vantage_config = await prisma_client.db.litellm_config.find_first( + vantage_config = await ConfigRepository(prisma_client).table.find_first( where={"param_name": VANTAGE_SETTINGS_PARAM_NAME} ) @@ -553,7 +554,7 @@ async def delete_vantage_settings( detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - vantage_config = await prisma_client.db.litellm_config.find_first( + vantage_config = await ConfigRepository(prisma_client).table.find_first( where={"param_name": VANTAGE_SETTINGS_PARAM_NAME} ) @@ -563,7 +564,7 @@ async def delete_vantage_settings( detail={"error": "Vantage settings not found"}, ) - await prisma_client.db.litellm_config.delete( + await ConfigRepository(prisma_client).table.delete( where={"param_name": VANTAGE_SETTINGS_PARAM_NAME} ) diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index 07e2ca71950..ea634289cb4 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -12,6 +12,12 @@ from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_keys from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.repositories.config_repository import ConfigRepository +from litellm.repositories.table_repositories import ( + DailyTagSpendRepository, + SSOConfigRepository, + UISettingsRepository, +) from litellm.types.proxy.management_endpoints.ui_sso import ( DefaultTeamSSOParams, InProductNudgeResponse, @@ -665,7 +671,7 @@ async def get_sso_settings(): ) # Get SSO config from dedicated table - sso_db_record = await prisma_client.db.litellm_ssoconfig.find_unique( + sso_db_record = await SSOConfigRepository(prisma_client).table.find_unique( where={"id": "sso_config"} ) @@ -836,7 +842,7 @@ async def update_sso_settings(sso_config: SSOConfig): ) # Save to dedicated SSO table - await prisma_client.db.litellm_ssoconfig.upsert( + await SSOConfigRepository(prisma_client).table.upsert( where={"id": "sso_config"}, data={ "create": { @@ -851,7 +857,7 @@ async def update_sso_settings(sso_config: SSOConfig): # Remove SSO-related env vars from config.environment_variables try: - env_var_entry = await prisma_client.db.litellm_config.find_unique( + env_var_entry = await ConfigRepository(prisma_client).table.find_unique( where={"param_name": "environment_variables"} ) @@ -872,7 +878,7 @@ async def update_sso_settings(sso_config: SSOConfig): if key not in env_vars_to_remove } - await prisma_client.db.litellm_config.update( + await ConfigRepository(prisma_client).table.update( where={"param_name": "environment_variables"}, data={ "param_value": json.dumps(filtered_env_vars, default=str), @@ -1123,7 +1129,7 @@ async def get_in_product_nudges(): detail={"error": "Database not connected. Please connect a database."}, ) - db_record = await prisma_client.db.litellm_dailytagspend.find_first( + db_record = await DailyTagSpendRepository(prisma_client).table.find_first( where={"tag": "User-Agent: claude-cli"} ) @@ -1155,7 +1161,7 @@ async def get_ui_settings_cached() -> Dict[str, Any]: if prisma_client is None: return {} - db_record = await prisma_client.db.litellm_uisettings.find_unique( + db_record = await UISettingsRepository(prisma_client).table.find_unique( where={"id": "ui_settings"} ) ui_settings: Dict[str, Any] = {} @@ -1196,7 +1202,7 @@ async def get_ui_settings(): ui_settings: Dict[str, Any] = {} - db_record = await prisma_client.db.litellm_uisettings.find_unique( + db_record = await UISettingsRepository(prisma_client).table.find_unique( where={"id": "ui_settings"} ) @@ -1309,7 +1315,7 @@ async def update_ui_settings( # Merge with existing persisted settings so a partial PATCH doesn't # overwrite fields the caller didn't send. existing: dict = {} - db_existing = await prisma_client.db.litellm_uisettings.find_unique( + db_existing = await UISettingsRepository(prisma_client).table.find_unique( where={"id": "ui_settings"} ) if db_existing and db_existing.ui_settings: @@ -1318,7 +1324,7 @@ async def update_ui_settings( ui_settings = {**existing, **incoming} - await prisma_client.db.litellm_uisettings.upsert( + await UISettingsRepository(prisma_client).table.upsert( where={"id": "ui_settings"}, data={ "create": { diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index e77e24c9e71..5ad42b5e1be 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -9,10 +9,10 @@ import sys import threading import time import traceback +from dataclasses import dataclass, field from datetime import date, datetime, timedelta, timezone from email.mime.multipart import MIMEMultipart from email.mime.text import MIMEText -from dataclasses import dataclass, field from typing import ( TYPE_CHECKING, Any, @@ -139,6 +139,19 @@ from litellm.proxy.hooks.parallel_request_limiter import ( ) from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor +from litellm.repositories.budget_repository import BudgetRepository +from litellm.repositories.config_repository import ConfigRepository +from litellm.repositories.table_repositories import ( + EndUserRepository, + HealthCheckRepository, + SpendLogsRepository, + UserNotificationsRepository, +) +from litellm.repositories.team_repository import TeamRepository +from litellm.repositories.user_repository import UserRepository +from litellm.repositories.verification_token_repository import ( + VerificationTokenRepository, +) from litellm.secret_managers.main import str_to_bool from litellm.types.integrations.slack_alerting import DEFAULT_ALERT_TYPES from litellm.types.mcp import ( @@ -2831,7 +2844,7 @@ async def prefetch_config_params(prisma_client: Any, param_names: List[str]) -> if not param_names: return try: - rows = await prisma_client.db.litellm_config.find_many( + rows = await ConfigRepository(prisma_client).table.find_many( where={"param_name": {"in": param_names}} # type: ignore ) except Exception as e: @@ -3194,15 +3207,15 @@ class PrismaClient: async def _do_query(): if table_name == "users": - return await self.db.litellm_usertable.find_first( + return await UserRepository(self).table.find_first( where={key: value} # type: ignore ) elif table_name == "keys": - return await self.db.litellm_verificationtoken.find_first( # type: ignore + return await VerificationTokenRepository(self).table.find_first( # type: ignore where={key: value} # type: ignore ) elif table_name == "config": - return await self.db.litellm_config.find_first( # type: ignore + return await ConfigRepository(self).table.find_first( # type: ignore where={key: value} # type: ignore ) elif table_name == "spend": @@ -3336,7 +3349,9 @@ class PrismaClient: status_code=400, detail={"error": f"No token passed in. Token={token}"}, ) - response = await self.db.litellm_verificationtoken.find_unique( + response = await VerificationTokenRepository( + self + ).table.find_unique( where={"token": hashed_token}, # type: ignore include={"litellm_budget_table": True}, ) @@ -3353,7 +3368,7 @@ class PrismaClient: detail=f"Authentication Error: invalid user key - user key does not exist in db. User Key={token}", ) elif query_type == "find_all" and user_id is not None: - response = await self.db.litellm_verificationtoken.find_many( + response = await VerificationTokenRepository(self).table.find_many( where={"user_id": user_id}, include={"litellm_budget_table": True}, ) @@ -3362,7 +3377,7 @@ class PrismaClient: if isinstance(r.expires, datetime): r.expires = r.expires.isoformat() elif query_type == "find_all" and team_id is not None: - response = await self.db.litellm_verificationtoken.find_many( + response = await VerificationTokenRepository(self).table.find_many( where={"team_id": team_id}, include={"litellm_budget_table": True}, ) @@ -3375,7 +3390,7 @@ class PrismaClient: and expires is not None and reset_at is not None ): - response = await self.db.litellm_verificationtoken.find_many( + response = await VerificationTokenRepository(self).table.find_many( where={ # type: ignore "OR": [ {"expires": None}, @@ -3405,7 +3420,7 @@ class PrismaClient: else: hashed_tokens.append(t) where_filter["token"]["in"] = hashed_tokens - response = await self.db.litellm_verificationtoken.find_many( + response = await VerificationTokenRepository(self).table.find_many( order={"spend": "desc"}, where=where_filter, # type: ignore include={"litellm_budget_table": True}, @@ -3425,28 +3440,28 @@ class PrismaClient: if key_val is None: key_val = {"user_id": user_id} - response = await self.db.litellm_usertable.find_unique( # type: ignore + response = await UserRepository(self).table.find_unique( # type: ignore where=key_val, # type: ignore include={"organization_memberships": True}, ) elif query_type == "find_all" and key_val is not None: - response = await self.db.litellm_usertable.find_many( + response = await UserRepository(self).table.find_many( where=key_val # type: ignore ) # type: ignore elif query_type == "find_all" and reset_at is not None: - response = await self.db.litellm_usertable.find_many( + response = await UserRepository(self).table.find_many( where={ # type: ignore "budget_reset_at": {"lt": reset_at}, } ) elif query_type == "find_all" and user_id_list is not None: - response = await self.db.litellm_usertable.find_many( + response = await UserRepository(self).table.find_many( where={"user_id": {"in": user_id_list}} ) elif query_type == "find_all": if expires is not None: - response = await self.db.litellm_usertable.find_many( # type: ignore + response = await UserRepository(self).table.find_many( # type: ignore order={"spend": "desc"}, where={ # type: ignore "OR": [ @@ -3478,26 +3493,26 @@ class PrismaClient: ) if key_val is not None: if query_type == "find_unique": - response = await self.db.litellm_spendlogs.find_unique( # type: ignore + response = await SpendLogsRepository(self).table.find_unique( # type: ignore where={ # type: ignore key_val["key"]: key_val["value"], # type: ignore } ) elif query_type == "find_all": - response = await self.db.litellm_spendlogs.find_many( # type: ignore + response = await SpendLogsRepository(self).table.find_many( # type: ignore where={ key_val["key"]: key_val["value"], # type: ignore } ) return response else: - response = await self.db.litellm_spendlogs.find_many( # type: ignore + response = await SpendLogsRepository(self).table.find_many( # type: ignore order={"startTime": "desc"}, ) return response elif table_name == "budget" and reset_at is not None: if query_type == "find_all": - response = await self.db.litellm_budgettable.find_many( + response = await BudgetRepository(self).table.find_many( where={ # type: ignore "OR": [ { @@ -3514,45 +3529,45 @@ class PrismaClient: elif table_name == "enduser" and budget_id_list is not None: if query_type == "find_all": - response = await self.db.litellm_endusertable.find_many( + response = await EndUserRepository(self).table.find_many( where={"budget_id": {"in": budget_id_list}} ) return response elif table_name == "team": if query_type == "find_unique": - response = await self.db.litellm_teamtable.find_unique( + response = await TeamRepository(self).table.find_unique( where={"team_id": team_id}, # type: ignore include={"litellm_model_table": True}, # type: ignore ) elif query_type == "find_all" and reset_at is not None: - response = await self.db.litellm_teamtable.find_many( + response = await TeamRepository(self).table.find_many( where={ # type: ignore "budget_reset_at": {"lt": reset_at}, } ) elif query_type == "find_all" and user_id is not None: - response = await self.db.litellm_teamtable.find_many( + response = await TeamRepository(self).table.find_many( where={ "members": {"has": user_id}, }, include={"litellm_budget_table": True}, ) elif query_type == "find_all" and team_id_list is not None: - response = await self.db.litellm_teamtable.find_many( + response = await TeamRepository(self).table.find_many( where={"team_id": {"in": team_id_list}} ) elif query_type == "find_all" and team_id_list is None: - response = await self.db.litellm_teamtable.find_many( + response = await TeamRepository(self).table.find_many( take=MAX_TEAM_LIST_LIMIT ) return response elif table_name == "user_notification": if query_type == "find_unique": - response = await self.db.litellm_usernotifications.find_unique( # type: ignore + response = await UserNotificationsRepository(self).table.find_unique( # type: ignore where={"user_id": user_id} # type: ignore ) elif query_type == "find_all": - response = await self.db.litellm_usernotifications.find_many() # type: ignore + response = await UserNotificationsRepository(self).table.find_many() # type: ignore return response elif table_name == "combined_view": # check if plain text or hash @@ -3744,7 +3759,7 @@ class PrismaClient: print_verbose( "PrismaClient: Before upsert into litellm_verificationtoken" ) - new_verification_token = await self.db.litellm_verificationtoken.upsert( # type: ignore + new_verification_token = await VerificationTokenRepository(self).table.upsert( # type: ignore where={ "token": hashed_token, }, @@ -3759,7 +3774,7 @@ class PrismaClient: elif table_name == "user": db_data = self.jsonify_object(data=data) try: - new_user_row = await self.db.litellm_usertable.upsert( + new_user_row = await UserRepository(self).table.upsert( where={"user_id": data["user_id"]}, data={ "create": {**db_data}, # type: ignore @@ -3782,7 +3797,7 @@ class PrismaClient: return new_user_row elif table_name == "team": db_data = self.jsonify_team_object(db_data=data) - new_team_row = await self.db.litellm_teamtable.upsert( + new_team_row = await TeamRepository(self).table.upsert( where={"team_id": data["team_id"]}, data={ "create": {**db_data}, # type: ignore @@ -3804,7 +3819,7 @@ class PrismaClient: for k, v in data.items(): updated_data = v updated_data = json.dumps(updated_data) - updated_table_row = self.db.litellm_config.upsert( + updated_table_row = ConfigRepository(self).table.upsert( where={"param_name": k}, # type: ignore data={ "create": {"param_name": k, "param_value": updated_data}, # type: ignore @@ -3820,7 +3835,7 @@ class PrismaClient: verbose_proxy_logger.info("Data Inserted into Config Table") elif table_name == "spend": db_data = self.jsonify_object(data=data) - new_spend_row = await self.db.litellm_spendlogs.upsert( + new_spend_row = await SpendLogsRepository(self).table.upsert( where={"request_id": data["request_id"]}, data={ "create": {**db_data}, # type: ignore @@ -3831,14 +3846,14 @@ class PrismaClient: return new_spend_row elif table_name == "user_notification": db_data = self.jsonify_object(data=data) - new_user_notification_row = ( - await self.db.litellm_usernotifications.upsert( # type: ignore - where={"request_id": data["request_id"]}, - data={ - "create": {**db_data}, # type: ignore - "update": {}, # don't do anything if it already exists - }, - ) + new_user_notification_row = await UserNotificationsRepository( + self + ).table.upsert( # type: ignore + where={"request_id": data["request_id"]}, + data={ + "create": {**db_data}, # type: ignore + "update": {}, # don't do anything if it already exists + }, ) verbose_proxy_logger.info("Data Inserted into Model Request Table") return new_user_notification_row @@ -3899,7 +3914,7 @@ class PrismaClient: # check if plain text or hash token = _hash_token_if_needed(token=token) db_data["token"] = token - response = await self.db.litellm_verificationtoken.update( + response = await VerificationTokenRepository(self).table.update( where={"token": token}, # type: ignore data={**db_data}, # type: ignore ) @@ -3930,7 +3945,7 @@ class PrismaClient: update_key_values = update_key_values_custom_query else: update_key_values = db_data - update_user_row = await self.db.litellm_usertable.upsert( + update_user_row = await UserRepository(self).table.upsert( where={"user_id": user_id}, # type: ignore data={ "create": {**db_data}, # type: ignore @@ -3971,7 +3986,7 @@ class PrismaClient: update_key_values["members_with_roles"] = json.dumps( update_key_values["members_with_roles"] ) - update_team_row = await self.db.litellm_teamtable.upsert( + update_team_row = await TeamRepository(self).table.upsert( where={"team_id": team_id}, # type: ignore data={ "create": {**db_data}, # type: ignore @@ -4196,7 +4211,9 @@ class PrismaClient: else: filter_query = {"token": {"in": hashed_tokens}} - deleted_tokens = await self.db.litellm_verificationtoken.delete_many( + deleted_tokens = await VerificationTokenRepository( + self + ).table.delete_many( where=filter_query # type: ignore ) verbose_proxy_logger.debug("deleted_tokens: %s", deleted_tokens) @@ -4207,7 +4224,7 @@ class PrismaClient: and isinstance(team_id_list, List) ): # admin only endpoint -> `/team/delete` - await self.db.litellm_teamtable.delete_many( + await TeamRepository(self).table.delete_many( where={"team_id": {"in": team_id_list}} ) return {"deleted_teams": team_id_list} @@ -4217,7 +4234,7 @@ class PrismaClient: and isinstance(team_id_list, List) ): # admin only endpoint -> `/team/delete` - await self.db.litellm_verificationtoken.delete_many( + await VerificationTokenRepository(self).table.delete_many( where={"team_id": {"in": team_id_list}} ) except Exception as e: @@ -5024,7 +5041,9 @@ class PrismaClient: ) verbose_proxy_logger.debug(f"Saving health check data: {health_check_data}") - return await self.db.litellm_healthchecktable.create(data=health_check_data) + return await HealthCheckRepository(self).table.create( + data=health_check_data + ) except Exception as e: verbose_proxy_logger.error( @@ -5049,7 +5068,7 @@ class PrismaClient: if status_filter: where_clause["status"] = status_filter - results = await self.db.litellm_healthchecktable.find_many( + results = await HealthCheckRepository(self).table.find_many( where=where_clause, order={"checked_at": "desc"}, take=limit, @@ -5068,7 +5087,7 @@ class PrismaClient: (via Prisma ``distinct`` + ``order``) so we never load the full history into memory. """ try: - return await self.db.litellm_healthchecktable.find_many( + return await HealthCheckRepository(self).table.find_many( distinct=["model_id", "model_name"], order=[ {"model_id": "asc"}, @@ -5228,7 +5247,7 @@ async def migrate_passwords_to_scrypt_async(prisma_client) -> str: are left alone (they migrate on next login via the SHA256 fallback). Skips quickly if no plaintext passwords exist. """ - all_with_pw = await prisma_client.db.litellm_usertable.find_many( + all_with_pw = await UserRepository(prisma_client).table.find_many( where={"password": {"not": None}}, ) @@ -5246,7 +5265,7 @@ async def migrate_passwords_to_scrypt_async(prisma_client) -> str: return "No plaintext passwords found" for user in plaintext_users: - await prisma_client.db.litellm_usertable.update( + await UserRepository(prisma_client).table.update( where={"user_id": user.user_id}, data={"password": hash_password(user.password)}, ) @@ -5370,7 +5389,7 @@ class ProxyUpdateSpend: prisma_client.jsonify_object({**entry}) for entry in batch ] - await prisma_client.db.litellm_spendlogs.create_many( + await SpendLogsRepository(prisma_client).table.create_many( data=batch_with_dates, skip_duplicates=True ) verbose_proxy_logger.debug( diff --git a/litellm/proxy/vector_store_endpoints/endpoints.py b/litellm/proxy/vector_store_endpoints/endpoints.py index b3bdfecbe55..9c2d3050346 100644 --- a/litellm/proxy/vector_store_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_endpoints/endpoints.py @@ -1,6 +1,7 @@ from typing import Any, Dict, Optional from fastapi import APIRouter, Depends, HTTPException, Request, Response + from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import ( LiteLLM_ManagedVectorStore, ) @@ -16,6 +17,7 @@ from litellm.proxy.vector_store_endpoints.utils import ( assert_user_can_access_vector_store, get_litellm_managed_vector_store, ) +from litellm.repositories.table_repositories import ManagedVectorStoreIndexRepository from litellm.types.vector_stores import IndexCreateRequest router = APIRouter() @@ -587,11 +589,9 @@ async def index_create( detail=CommonProxyErrors.db_not_connected_error.value, ) ## 1. check if index already exists - existing_index = ( - await prisma_client.db.litellm_managedvectorstoreindextable.find_unique( - where={"index_name": index_create_request.index_name} - ) - ) + existing_index = await ManagedVectorStoreIndexRepository( + prisma_client + ).table.find_unique(where={"index_name": index_create_request.index_name}) ## 2. set created_by and updated_by @@ -605,7 +605,7 @@ async def index_create( index_data = index_create_request.model_dump(exclude_none=True) index_data["created_by"] = user_api_key_dict.user_id index_data["updated_by"] = user_api_key_dict.user_id - new_index = await prisma_client.db.litellm_managedvectorstoreindextable.create( + new_index = await ManagedVectorStoreIndexRepository(prisma_client).table.create( data=jsonify_object(index_data) ) diff --git a/litellm/proxy/vector_store_endpoints/management_endpoints.py b/litellm/proxy/vector_store_endpoints/management_endpoints.py index cbb3d927184..032a3302fdc 100644 --- a/litellm/proxy/vector_store_endpoints/management_endpoints.py +++ b/litellm/proxy/vector_store_endpoints/management_endpoints.py @@ -29,6 +29,8 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper from litellm.proxy.common_utils.rbac_utils import check_feature_access_for_user from litellm.proxy.vector_store_endpoints.utils import can_user_access_vector_store +from litellm.repositories.model_repository import ModelRepository +from litellm.repositories.table_repositories import ManagedVectorStoresRepository from litellm.secret_managers.main import get_secret from litellm.types.vector_stores import ( LiteLLM_ManagedVectorStore, @@ -122,7 +124,7 @@ async def _fetch_and_authorize_vector_store( Raises HTTPException(404) on miss and HTTPException(403) on access denial. """ - row = await prisma_client.db.litellm_managedvectorstorestable.find_unique( + row = await ManagedVectorStoresRepository(prisma_client).table.find_unique( where={"vector_store_id": vector_store_id} ) if row is None: @@ -252,7 +254,7 @@ async def _resolve_embedding_config_from_db( # Try to find model in database for model_name in model_name_candidates: try: - db_model = await prisma_client.db.litellm_proxymodeltable.find_first( + db_model = await ModelRepository(prisma_client).table.find_first( where={"model_name": model_name} ) @@ -437,11 +439,9 @@ async def create_vector_store_in_db( raise HTTPException(status_code=500, detail="Database not connected") # Check if vector store already exists - existing_vector_store = ( - await prisma_client.db.litellm_managedvectorstorestable.find_unique( - where={"vector_store_id": vector_store_id} - ) - ) + existing_vector_store = await ManagedVectorStoresRepository( + prisma_client + ).table.find_unique(where={"vector_store_id": vector_store_id}) if existing_vector_store is not None: raise HTTPException( status_code=400, @@ -487,7 +487,7 @@ async def create_vector_store_in_db( data_to_create["litellm_params"] = safe_dumps({}) # Create in database - _new_vector_store = await prisma_client.db.litellm_managedvectorstorestable.create( + _new_vector_store = await ManagedVectorStoresRepository(prisma_client).table.create( data=data_to_create ) @@ -725,11 +725,9 @@ async def delete_vector_store( memory_vector_store_exists = False vector_store_to_check = None - existing_vector_store = ( - await prisma_client.db.litellm_managedvectorstorestable.find_unique( - where={"vector_store_id": data.vector_store_id} - ) - ) + existing_vector_store = await ManagedVectorStoresRepository( + prisma_client + ).table.find_unique(where={"vector_store_id": data.vector_store_id}) if existing_vector_store is not None: db_vector_store_exists = True vector_store_to_check = LiteLLM_ManagedVectorStore( @@ -764,7 +762,7 @@ async def delete_vector_store( # Delete from database if exists if db_vector_store_exists: - await prisma_client.db.litellm_managedvectorstorestable.delete( + await ManagedVectorStoresRepository(prisma_client).table.delete( where={"vector_store_id": data.vector_store_id} ) @@ -921,7 +919,7 @@ async def update_vector_store( update_data["litellm_params"] = safe_dumps(litellm_params_dict) # Update in database - updated = await prisma_client.db.litellm_managedvectorstorestable.update( + updated = await ManagedVectorStoresRepository(prisma_client).table.update( where={"vector_store_id": vector_store_id}, data=update_data, ) diff --git a/litellm/repositories/__init__.py b/litellm/repositories/__init__.py new file mode 100644 index 00000000000..4451f0865da --- /dev/null +++ b/litellm/repositories/__init__.py @@ -0,0 +1,127 @@ +""" +Repository classes for database operations. +""" + +from litellm.repositories.budget_repository import BudgetRepository +from litellm.repositories.config_repository import ConfigRepository +from litellm.repositories.credentials_repository import CredentialsRepository +from litellm.repositories.model_repository import ModelRepository +from litellm.repositories.object_permission_repository import ( + ObjectPermissionRepository, +) +from litellm.repositories.organization_repository import OrganizationRepository +from litellm.repositories.project_repository import ProjectRepository +from litellm.repositories.table_repositories import ( + AccessGroupRepository, + AdaptiveRouterSessionRepository, + AdaptiveRouterStateRepository, + AgentsRepository, + AuditLogRepository, + CacheConfigRepository, + ClaudeCodePluginRepository, + ConfigOverridesRepository, + DailyGuardrailMetricsRepository, + DailyPolicyMetricsRepository, + DailyTagSpendRepository, + DeletedTeamRepository, + DeletedVerificationTokenRepository, + DeprecatedVerificationTokenRepository, + EndUserRepository, + GuardrailsRepository, + HealthCheckRepository, + InvitationLinkRepository, + JWTKeyMappingRepository, + ManagedFileRepository, + ManagedObjectRepository, + ManagedVectorStoreIndexRepository, + ManagedVectorStoresRepository, + MCPServerRepository, + MCPToolsetRepository, + MCPUserCredentialsRepository, + MemoryRepository, + ModelTableRepository, + OrganizationMembershipRepository, + PolicyAttachmentRepository, + PolicyRepository, + PrismaTableRepository, + PromptRepository, + SearchToolsRepository, + SkillsRepository, + SpendLogGuardrailIndexRepository, + SpendLogsRepository, + SpendLogToolIndexRepository, + SSOConfigRepository, + TagRepository, + TeamMembershipRepository, + ToolRepository, + UISettingsRepository, + UserNotificationsRepository, + WorkflowEventRepository, + WorkflowMessageRepository, + WorkflowRunRepository, +) +from litellm.repositories.team_repository import TeamRepository +from litellm.repositories.user_repository import UserRepository +from litellm.repositories.verification_token_repository import ( + VerificationTokenRepository, +) + +__all__ = [ + "PrismaTableRepository", + "PolicyRepository", + "AgentsRepository", + "GuardrailsRepository", + "MCPServerRepository", + "ManagedObjectRepository", + "OrganizationMembershipRepository", + "SpendLogsRepository", + "ClaudeCodePluginRepository", + "TeamMembershipRepository", + "EndUserRepository", + "ManagedVectorStoresRepository", + "MCPUserCredentialsRepository", + "PromptRepository", + "TagRepository", + "InvitationLinkRepository", + "JWTKeyMappingRepository", + "ManagedFileRepository", + "MemoryRepository", + "SearchToolsRepository", + "ConfigOverridesRepository", + "MCPToolsetRepository", + "ToolRepository", + "DeletedVerificationTokenRepository", + "WorkflowRunRepository", + "ModelTableRepository", + "AccessGroupRepository", + "SSOConfigRepository", + "UISettingsRepository", + "DailyGuardrailMetricsRepository", + "PolicyAttachmentRepository", + "DeletedTeamRepository", + "SkillsRepository", + "CacheConfigRepository", + "ManagedVectorStoreIndexRepository", + "WorkflowMessageRepository", + "DailyTagSpendRepository", + "SpendLogToolIndexRepository", + "SpendLogGuardrailIndexRepository", + "UserNotificationsRepository", + "HealthCheckRepository", + "DeprecatedVerificationTokenRepository", + "WorkflowEventRepository", + "DailyPolicyMetricsRepository", + "AdaptiveRouterStateRepository", + "AuditLogRepository", + "AdaptiveRouterSessionRepository", + "BudgetRepository", + "ConfigRepository", + "CredentialsRepository", + "ModelRepository", + "ObjectPermissionRepository", + "OrganizationRepository", + "ProjectRepository", + "TeamRepository", + "UserRepository", + "VerificationTokenRepository", +] diff --git a/litellm/repositories/base_repository.py b/litellm/repositories/base_repository.py new file mode 100644 index 00000000000..a25620c7b4d --- /dev/null +++ b/litellm/repositories/base_repository.py @@ -0,0 +1,117 @@ +""" +Base repository class with common functionality. +""" + +from abc import ABC, abstractmethod +from typing import Any, Dict, Generic, List, Optional, Type, TypeVar + +from pydantic import BaseModel + +T = TypeVar("T", bound=BaseModel) + + +def _record_to_dict(record: Any) -> Dict[str, Any]: + if isinstance(record, dict): + return record + if hasattr(record, "model_dump") and callable(record.model_dump): + return record.model_dump() + if hasattr(record, "dict") and callable(record.dict): + return record.dict() + return dict(record) + + +class BaseRepository(ABC, Generic[T]): + """Abstract base class for all repositories.""" + + def __init__(self, prisma_client: Any): + self._prisma_client = prisma_client + + @property + def prisma_client(self) -> Any: + if self._prisma_client is None: + raise RuntimeError( + "No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys" + ) + return self._prisma_client + + @property + @abstractmethod + def table(self) -> Any: + """Return the Prisma table for this repository.""" + ... + + @property + @abstractmethod + def model_class(self) -> Type[T]: + """Return the domain model class for this repository.""" + ... + + def _to_model(self, record: Any) -> Optional[T]: + """Convert a database record to a domain model.""" + if record is None: + return None + return self.model_class(**_record_to_dict(record)) + + def _to_model_list(self, records: List[Any]) -> List[T]: + """Convert a list of database records to domain models.""" + result: List[T] = [] + for r in records: + if r is not None: + model = self._to_model(r) + if model is not None: + result.append(model) + return result + + async def find_by_id(self, id_value: str, id_field: str = "id") -> Optional[T]: + """Find a record by its primary key.""" + record = await self.table.find_unique(where={id_field: id_value}) + return self._to_model(record) + + async def find_many( + self, + where: Optional[Dict[str, Any]] = None, + skip: Optional[int] = None, + take: Optional[int] = None, + order: Optional[Dict[str, str]] = None, + ) -> List[T]: + """Find multiple records matching the criteria.""" + kwargs: Dict[str, Any] = {} + if where: + kwargs["where"] = where + if skip is not None: + kwargs["skip"] = skip + if take is not None: + kwargs["take"] = take + if order: + kwargs["order"] = order + + records = await self.table.find_many(**kwargs) + return self._to_model_list(records) + + async def create(self, data: Dict[str, Any]) -> T: + """Create a new record.""" + record = await self.table.create(data=data) + model = self._to_model(record) + assert model is not None + return model + + async def update( + self, id_value: str, data: Dict[str, Any], id_field: str = "id" + ) -> Optional[T]: + """Update an existing record.""" + record = await self.table.update(where={id_field: id_value}, data=data) + return self._to_model(record) + + async def delete(self, id_value: str, id_field: str = "id") -> Optional[T]: + """Delete a record by its primary key.""" + record = await self.table.delete(where={id_field: id_value}) + return self._to_model(record) + + async def count(self, where: Optional[Dict[str, Any]] = None) -> int: + """Count records matching the criteria.""" + return await self.table.count(where=where) + + async def exists(self, id_value: str, id_field: str = "id") -> bool: + """Check if a record exists.""" + record = await self.table.find_unique(where={id_field: id_value}) + return record is not None diff --git a/litellm/repositories/budget_repository.py b/litellm/repositories/budget_repository.py new file mode 100644 index 00000000000..5947701fb4e --- /dev/null +++ b/litellm/repositories/budget_repository.py @@ -0,0 +1,99 @@ +""" +Budget repository for database operations on LiteLLM_BudgetTable. +""" + +from typing import Any, Dict, List, Optional, Type + +from litellm.models.budget import LiteLLM_BudgetTable +from litellm.repositories.base_repository import BaseRepository + + +class BudgetRepository(BaseRepository[LiteLLM_BudgetTable]): + """Repository for budget database operations.""" + + @property + def table(self) -> Any: + return self.prisma_client.db.litellm_budgettable + + @property + def model_class(self) -> Type[LiteLLM_BudgetTable]: + return LiteLLM_BudgetTable + + async def find_by_id( + self, budget_id: str, id_field: str = "budget_id" + ) -> Optional[LiteLLM_BudgetTable]: + return await super().find_by_id(budget_id, id_field) + + async def create_budget( + self, + created_by: str, + max_budget: Optional[float] = None, + soft_budget: Optional[float] = None, + max_parallel_requests: Optional[int] = None, + tpm_limit: Optional[int] = None, + rpm_limit: Optional[int] = None, + model_max_budget: Optional[Dict[str, Any]] = None, + budget_duration: Optional[str] = None, + allowed_models: Optional[List[str]] = None, + ) -> LiteLLM_BudgetTable: + """Create a new budget record.""" + data: Dict[str, Any] = { + "created_by": created_by, + "updated_by": created_by, + } + if max_budget is not None: + data["max_budget"] = max_budget + if soft_budget is not None: + data["soft_budget"] = soft_budget + if max_parallel_requests is not None: + data["max_parallel_requests"] = max_parallel_requests + if tpm_limit is not None: + data["tpm_limit"] = tpm_limit + if rpm_limit is not None: + data["rpm_limit"] = rpm_limit + if model_max_budget is not None: + data["model_max_budget"] = model_max_budget + if budget_duration is not None: + data["budget_duration"] = budget_duration + if allowed_models is not None: + data["allowed_models"] = allowed_models + + return await self.create(data) + + async def update_budget( + self, + budget_id: str, + updated_by: str, + max_budget: Optional[float] = None, + soft_budget: Optional[float] = None, + max_parallel_requests: Optional[int] = None, + tpm_limit: Optional[int] = None, + rpm_limit: Optional[int] = None, + model_max_budget: Optional[Dict[str, Any]] = None, + budget_duration: Optional[str] = None, + allowed_models: Optional[List[str]] = None, + ) -> Optional[LiteLLM_BudgetTable]: + """Update an existing budget record.""" + data: Dict[str, Any] = {"updated_by": updated_by} + if max_budget is not None: + data["max_budget"] = max_budget + if soft_budget is not None: + data["soft_budget"] = soft_budget + if max_parallel_requests is not None: + data["max_parallel_requests"] = max_parallel_requests + if tpm_limit is not None: + data["tpm_limit"] = tpm_limit + if rpm_limit is not None: + data["rpm_limit"] = rpm_limit + if model_max_budget is not None: + data["model_max_budget"] = model_max_budget + if budget_duration is not None: + data["budget_duration"] = budget_duration + if allowed_models is not None: + data["allowed_models"] = allowed_models + + return await self.update(budget_id, data, id_field="budget_id") + + async def delete_budget(self, budget_id: str) -> Optional[LiteLLM_BudgetTable]: + """Delete a budget record.""" + return await self.delete(budget_id, id_field="budget_id") diff --git a/litellm/repositories/config_repository.py b/litellm/repositories/config_repository.py new file mode 100644 index 00000000000..eba7ebe26ca --- /dev/null +++ b/litellm/repositories/config_repository.py @@ -0,0 +1,241 @@ +""" +Config repository for database operations on LiteLLM_Config. + +This repository handles config reconciliation between database values and +YAML configmap values. DB values override configmap values except for +None values and empty lists. +""" + +import asyncio +import copy +import json +import os +from typing import Any, Dict, List, Literal, Optional, cast + +from litellm._logging import verbose_proxy_logger +from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper + + +class ConfigParam: + """Simple wrapper for config parameter from DB.""" + + def __init__(self, param_name: str, param_value: Any): + self.param_name = param_name + self.param_value = param_value + + +class ConfigRepository: + """Repository for config database operations with reconciliation support.""" + + CONFIG_PARAMS = [ + "general_settings", + "router_settings", + "litellm_settings", + "environment_variables", + ] + + def __init__(self, prisma_client: Any): + self._prisma_client = prisma_client + + @property + def prisma_client(self) -> Any: + if self._prisma_client is None: + raise RuntimeError( + "No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys" + ) + return self._prisma_client + + @property + def table(self) -> Any: + return self.prisma_client.db.litellm_config + + async def get_param(self, param_name: str) -> Optional[ConfigParam]: + """Get a config parameter from the database.""" + record = await self.table.find_unique(where={"param_name": param_name}) + if record is None: + return None + param_value = record.param_value + if isinstance(param_value, str): + param_value = json.loads(param_value) + return ConfigParam(param_name=param_name, param_value=param_value) + + async def set_param(self, param_name: str, param_value: Any) -> ConfigParam: + """Set a config parameter in the database.""" + value_json = ( + json.dumps(param_value) if not isinstance(param_value, str) else param_value + ) + await self.table.upsert( + where={"param_name": param_name}, + data={ + "create": {"param_name": param_name, "param_value": value_json}, + "update": {"param_value": value_json}, + }, + ) + return ConfigParam(param_name=param_name, param_value=param_value) + + async def delete_param(self, param_name: str) -> bool: + """Delete a config parameter from the database.""" + try: + await self.table.delete(where={"param_name": param_name}) + return True + except Exception: + return False + + async def get_all_params(self) -> Dict[str, Any]: + """Get all config parameters from the database.""" + records = await self.table.find_many() + result = {} + for record in records: + param_value = record.param_value + if isinstance(param_value, str): + param_value = json.loads(param_value) + result[record.param_name] = param_value + return result + + def _deep_merge_dicts(self, dst: dict, src: dict) -> None: + """Deep-merge src into dst, skipping None values and empty lists from src. + + On conflicts, src (DB) wins, but empty lists are treated as "no value" + and don't overwrite the destination. + """ + stack = [(dst, src)] + while stack: + d, s = stack.pop() + for k, v in s.items(): + if v is None: + continue + if isinstance(v, list) and len(v) == 0: + continue + if isinstance(v, dict) and isinstance(d.get(k), dict): + stack.append((d[k], v)) + else: + d[k] = v + + def _decrypt_env_variables( + self, env_vars: Dict[str, Any], return_original_value: bool = True + ) -> Dict[str, str]: + """Decrypt environment variables from database.""" + decrypted: Dict[str, str] = {} + for key, value in env_vars.items(): + if isinstance(value, str): + decrypted_value = decrypt_value_helper( + value=value, + key=key, + exception_type="debug", + return_original_value=return_original_value, + ) + if decrypted_value is not None: + decrypted[key] = decrypted_value + else: + decrypted[key] = str(value) + return decrypted + + def _normalize_env_variable_keys(self, env_vars: Dict[str, str]) -> Dict[str, str]: + """Normalize env variable keys to include both original and uppercase versions.""" + normalized: Dict[str, str] = {} + for key, value in env_vars.items(): + normalized[key] = value + upper_key = key.upper() + normalized[upper_key] = value + return normalized + + def _update_config_fields( + self, + current_config: dict, + param_name: Literal[ + "general_settings", + "router_settings", + "litellm_settings", + "environment_variables", + ], + db_param_value: Any, + ) -> dict: + """Update config fields with DB values, handling the merge strategy.""" + if param_name == "environment_variables": + decrypted_env_vars = self._decrypt_env_variables( + db_param_value, return_original_value=True + ) + merged_env_vars = self._normalize_env_variable_keys(decrypted_env_vars) + for env_key, value in merged_env_vars.items(): + os.environ[env_key] = value + + current_config.setdefault("environment_variables", {}).update( + merged_env_vars + ) + return current_config + + if param_name not in current_config: + current_config[param_name] = db_param_value + return current_config + + if isinstance(current_config[param_name], dict) and isinstance( + db_param_value, dict + ): + self._deep_merge_dicts(current_config[param_name], db_param_value) + else: + current_config[param_name] = db_param_value + + return current_config + + async def reconcile_config( + self, + yaml_config: dict, + store_model_in_db: Optional[bool] = None, + ) -> dict: + """Reconcile config from YAML with database overrides. + + This is the main config reconciliation method that loads config params + from the database and merges them with the YAML config. DB values + override YAML values except for None values and empty lists. + + Args: + yaml_config: The configuration loaded from YAML file + store_model_in_db: Whether to load config from DB + + Returns: + The merged configuration with DB overrides applied + """ + if store_model_in_db is not True: + verbose_proxy_logger.info( + "'store_model_in_db' is not True, skipping db config reconciliation" + ) + return yaml_config + + tasks = [self.get_param(k) for k in self.CONFIG_PARAMS] + responses = await asyncio.gather(*tasks) + + config = copy.deepcopy(yaml_config) + for response in responses: + if response is None: + continue + + param_name = response.param_name + param_value = response.param_value + verbose_proxy_logger.debug( + f"param_name={param_name}, param_value={param_value}" + ) + + if param_name is not None and param_value is not None: + config = self._update_config_fields( + current_config=config, + param_name=cast( + Literal[ + "general_settings", + "router_settings", + "litellm_settings", + "environment_variables", + ], + param_name, + ), + db_param_value=param_value, + ) + + return config + + async def prefetch_params(self, param_names: List[str]) -> None: + """Prefetch config params to warm the cache. + + This can be called before reconcile_config to ensure all needed + params are loaded in a single batch. + """ + await asyncio.gather(*[self.get_param(k) for k in param_names]) diff --git a/litellm/repositories/credentials_repository.py b/litellm/repositories/credentials_repository.py new file mode 100644 index 00000000000..dd53c753307 --- /dev/null +++ b/litellm/repositories/credentials_repository.py @@ -0,0 +1,61 @@ +""" +Credentials repository for database operations on LiteLLM_CredentialsTable. + +This is the only place that talks to ``litellm_credentialstable``. Encryption of +credential values is the caller's responsibility (see ``CredentialHelperUtils``), +so reads return the stored values verbatim. +""" + +from typing import Any, Dict, Optional + +from litellm.models.credentials import CredentialItem + + +class CredentialsRepository: + """Repository for credentials database operations, keyed by credential name.""" + + def __init__(self, prisma_client: Any): + self._prisma_client = prisma_client + + @property + def prisma_client(self) -> Any: + if self._prisma_client is None: + raise RuntimeError( + "No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys" + ) + return self._prisma_client + + @property + def table(self) -> Any: + return self.prisma_client.db.litellm_credentialstable + + @staticmethod + def _to_model(record: Any) -> Optional[CredentialItem]: + if record is None: + return None + data = record.dict() if hasattr(record, "dict") else dict(record) + return CredentialItem( + credential_name=data["credential_name"], + credential_values=data.get("credential_values") or {}, + credential_info=data.get("credential_info") or {}, + ) + + async def find_all(self) -> Any: + return await self.table.find_many() + + async def create(self, data: Dict[str, Any]) -> Any: + return await self.table.create(data=data) + + async def find_by_name(self, credential_name: str) -> Optional[CredentialItem]: + record = await self.table.find_unique( + where={"credential_name": credential_name} + ) + return self._to_model(record) + + async def update_by_name(self, credential_name: str, data: Dict[str, Any]) -> Any: + return await self.table.update( + where={"credential_name": credential_name}, data=data + ) + + async def delete_by_name(self, credential_name: str) -> Any: + return await self.table.delete(where={"credential_name": credential_name}) diff --git a/litellm/repositories/model_repository.py b/litellm/repositories/model_repository.py new file mode 100644 index 00000000000..893cf342d71 --- /dev/null +++ b/litellm/repositories/model_repository.py @@ -0,0 +1,171 @@ +""" +Model repository for database operations on LiteLLM_ProxyModelTable. +""" + +import json +from typing import Any, Dict, List, Optional, Type + +from litellm.models.model import LiteLLM_ProxyModelTable +from litellm.repositories.base_repository import BaseRepository +from litellm.proxy.common_utils.encrypt_decrypt_utils import ( + decrypt_value_helper, + encrypt_value_helper, +) + + +class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]): + """Repository for proxy model database operations with encryption support.""" + + def __init__(self, prisma_client: Any, encryption_key: Optional[str] = None): + super().__init__(prisma_client) + self._encryption_key = encryption_key + + @property + def table(self) -> Any: + return self.prisma_client.db.litellm_proxymodeltable + + @property + def model_class(self) -> Type[LiteLLM_ProxyModelTable]: + return LiteLLM_ProxyModelTable + + def _encrypt_litellm_params(self, litellm_params: Dict[str, Any]) -> Dict[str, Any]: + """Encrypt sensitive values in litellm_params.""" + encrypted = {} + for key, value in litellm_params.items(): + if isinstance(value, str): + encrypted[key] = encrypt_value_helper( + value, new_encryption_key=self._encryption_key + ) + else: + encrypted[key] = value + return encrypted + + def _decrypt_litellm_params(self, litellm_params: Dict[str, Any]) -> Dict[str, Any]: + """Decrypt sensitive values in litellm_params.""" + decrypted = {} + for key, value in litellm_params.items(): + if isinstance(value, str): + decrypted[key] = decrypt_value_helper( + value, key=key, exception_type="debug", return_original_value=True + ) + else: + decrypted[key] = value + return decrypted + + def _to_model(self, record: Any) -> Optional[LiteLLM_ProxyModelTable]: + """Convert a database record to a Model with decryption.""" + if record is None: + return None + + data = record.dict() if hasattr(record, "dict") else dict(record) + + if isinstance(data.get("litellm_params"), str): + data["litellm_params"] = json.loads(data["litellm_params"]) + if isinstance(data.get("model_info"), str): + data["model_info"] = json.loads(data["model_info"]) + + if data.get("litellm_params"): + data["litellm_params"] = self._decrypt_litellm_params( + data["litellm_params"] + ) + + return LiteLLM_ProxyModelTable(**data) + + async def find_by_id( + self, model_id: str, id_field: str = "model_id" + ) -> Optional[LiteLLM_ProxyModelTable]: + return await super().find_by_id(model_id, id_field) + + async def find_by_name(self, model_name: str) -> List[LiteLLM_ProxyModelTable]: + """Find models by name.""" + records = await self.table.find_many(where={"model_name": model_name}) + return self._to_model_list(records) + + async def find_all(self) -> List[LiteLLM_ProxyModelTable]: + """Find all models.""" + records = await self.table.find_many() + return self._to_model_list(records) + + async def find_unblocked(self) -> List[LiteLLM_ProxyModelTable]: + """Find all models that are not blocked.""" + records = await self.table.find_many(where={"blocked": False}) + return self._to_model_list(records) + + async def find_by_team_id(self, team_id: str) -> List[LiteLLM_ProxyModelTable]: + """Find models associated with a specific team. + + Note: This filters in-memory since team_id is stored within litellm_params + JSON. For large deployments with many models, consider adding a dedicated + team_id column with a database index. + """ + all_models = await self.find_all() + return [m for m in all_models if m.team_id == team_id] + + async def create_model( + self, + model_name: str, + litellm_params: Dict[str, Any], + created_by: str, + model_id: Optional[str] = None, + model_info: Optional[Dict[str, Any]] = None, + blocked: bool = False, + ) -> LiteLLM_ProxyModelTable: + """Create a new model with encryption.""" + encrypted_params = self._encrypt_litellm_params(litellm_params) + + data: Dict[str, Any] = { + "model_name": model_name, + "litellm_params": json.dumps(encrypted_params), + "created_by": created_by, + "updated_by": created_by, + "blocked": blocked, + } + if model_id is not None: + data["model_id"] = model_id + if model_info is not None: + data["model_info"] = json.dumps(model_info) + + record = await self.table.create(data=data) + model = self._to_model(record) + assert model is not None + return model + + async def update_model( + self, + model_id: str, + updated_by: str, + model_name: Optional[str] = None, + litellm_params: Optional[Dict[str, Any]] = None, + model_info: Optional[Dict[str, Any]] = None, + blocked: Optional[bool] = None, + ) -> Optional[LiteLLM_ProxyModelTable]: + """Update a model with encryption.""" + data: Dict[str, Any] = {"updated_by": updated_by} + if model_name is not None: + data["model_name"] = model_name + if litellm_params is not None: + encrypted_params = self._encrypt_litellm_params(litellm_params) + data["litellm_params"] = json.dumps(encrypted_params) + if model_info is not None: + data["model_info"] = json.dumps(model_info) + if blocked is not None: + data["blocked"] = blocked + + record = await self.table.update(where={"model_id": model_id}, data=data) + return self._to_model(record) + + async def delete_model(self, model_id: str) -> Optional[LiteLLM_ProxyModelTable]: + """Delete a model.""" + return await self.delete(model_id, id_field="model_id") + + async def block_model( + self, model_id: str, updated_by: str + ) -> Optional[LiteLLM_ProxyModelTable]: + """Block a model.""" + return await self.update_model(model_id, updated_by, blocked=True) + + async def unblock_model( + self, model_id: str, updated_by: str + ) -> Optional[LiteLLM_ProxyModelTable]: + """Unblock a model.""" + return await self.update_model(model_id, updated_by, blocked=False) diff --git a/litellm/repositories/object_permission_repository.py b/litellm/repositories/object_permission_repository.py new file mode 100644 index 00000000000..f4d9a8bb90a --- /dev/null +++ b/litellm/repositories/object_permission_repository.py @@ -0,0 +1,110 @@ +""" +ObjectPermission repository for database operations on LiteLLM_ObjectPermissionTable. +""" + +from typing import Any, Dict, List, Optional, Type + +from litellm.models.object_permission import LiteLLM_ObjectPermissionTable +from litellm.repositories.base_repository import BaseRepository + + +class ObjectPermissionRepository(BaseRepository[LiteLLM_ObjectPermissionTable]): + """Repository for object permission database operations.""" + + @property + def table(self) -> Any: + return self.prisma_client.db.litellm_objectpermissiontable + + @property + def model_class(self) -> Type[LiteLLM_ObjectPermissionTable]: + return LiteLLM_ObjectPermissionTable + + async def find_by_id( + self, object_permission_id: str, id_field: str = "object_permission_id" + ) -> Optional[LiteLLM_ObjectPermissionTable]: + return await super().find_by_id(object_permission_id, id_field) + + async def create_permission( + self, + mcp_servers: Optional[List[str]] = None, + mcp_access_groups: Optional[List[str]] = None, + mcp_tool_permissions: Optional[Dict[str, List[str]]] = None, + vector_stores: Optional[List[str]] = None, + agents: Optional[List[str]] = None, + agent_access_groups: Optional[List[str]] = None, + models: Optional[List[str]] = None, + blocked_tools: Optional[List[str]] = None, + mcp_toolsets: Optional[List[str]] = None, + search_tools: Optional[List[str]] = None, + ) -> LiteLLM_ObjectPermissionTable: + """Create a new object permission record.""" + data: Dict[str, Any] = {} + if mcp_servers is not None: + data["mcp_servers"] = mcp_servers + if mcp_access_groups is not None: + data["mcp_access_groups"] = mcp_access_groups + if mcp_tool_permissions is not None: + data["mcp_tool_permissions"] = mcp_tool_permissions + if vector_stores is not None: + data["vector_stores"] = vector_stores + if agents is not None: + data["agents"] = agents + if agent_access_groups is not None: + data["agent_access_groups"] = agent_access_groups + if models is not None: + data["models"] = models + if blocked_tools is not None: + data["blocked_tools"] = blocked_tools + if mcp_toolsets is not None: + data["mcp_toolsets"] = mcp_toolsets + if search_tools is not None: + data["search_tools"] = search_tools + + return await self.create(data) + + async def update_permission( + self, + object_permission_id: str, + mcp_servers: Optional[List[str]] = None, + mcp_access_groups: Optional[List[str]] = None, + mcp_tool_permissions: Optional[Dict[str, List[str]]] = None, + vector_stores: Optional[List[str]] = None, + agents: Optional[List[str]] = None, + agent_access_groups: Optional[List[str]] = None, + models: Optional[List[str]] = None, + blocked_tools: Optional[List[str]] = None, + mcp_toolsets: Optional[List[str]] = None, + search_tools: Optional[List[str]] = None, + ) -> Optional[LiteLLM_ObjectPermissionTable]: + """Update an object permission record.""" + data: Dict[str, Any] = {} + if mcp_servers is not None: + data["mcp_servers"] = mcp_servers + if mcp_access_groups is not None: + data["mcp_access_groups"] = mcp_access_groups + if mcp_tool_permissions is not None: + data["mcp_tool_permissions"] = mcp_tool_permissions + if vector_stores is not None: + data["vector_stores"] = vector_stores + if agents is not None: + data["agents"] = agents + if agent_access_groups is not None: + data["agent_access_groups"] = agent_access_groups + if models is not None: + data["models"] = models + if blocked_tools is not None: + data["blocked_tools"] = blocked_tools + if mcp_toolsets is not None: + data["mcp_toolsets"] = mcp_toolsets + if search_tools is not None: + data["search_tools"] = search_tools + + return await self.update( + object_permission_id, data, id_field="object_permission_id" + ) + + async def delete_permission( + self, object_permission_id: str + ) -> Optional[LiteLLM_ObjectPermissionTable]: + """Delete an object permission record.""" + return await self.delete(object_permission_id, id_field="object_permission_id") diff --git a/litellm/repositories/organization_repository.py b/litellm/repositories/organization_repository.py new file mode 100644 index 00000000000..2d25a43e836 --- /dev/null +++ b/litellm/repositories/organization_repository.py @@ -0,0 +1,103 @@ +""" +Organization repository for database operations on LiteLLM_OrganizationTable. +""" + +from typing import Any, Dict, List, Optional, Type + +from litellm.models.organization import LiteLLM_OrganizationTable +from litellm.repositories.base_repository import BaseRepository + + +class OrganizationRepository(BaseRepository[LiteLLM_OrganizationTable]): + """Repository for organization database operations.""" + + @property + def table(self) -> Any: + return self.prisma_client.db.litellm_organizationtable + + @property + def model_class(self) -> Type[LiteLLM_OrganizationTable]: + return LiteLLM_OrganizationTable + + async def find_by_id( + self, organization_id: str, id_field: str = "organization_id" + ) -> Optional[LiteLLM_OrganizationTable]: + return await super().find_by_id(organization_id, id_field) + + async def find_by_alias( + self, organization_alias: str + ) -> Optional[LiteLLM_OrganizationTable]: + """Find an organization by alias.""" + records = await self.table.find_many( + where={"organization_alias": organization_alias} + ) + if records: + return self._to_model(records[0]) + return None + + async def create_organization( + self, + organization_alias: str, + budget_id: str, + created_by: str, + organization_id: Optional[str] = None, + metadata: Optional[Dict[str, Any]] = None, + models: Optional[List[str]] = None, + object_permission_id: Optional[str] = None, + ) -> LiteLLM_OrganizationTable: + """Create a new organization.""" + data: Dict[str, Any] = { + "organization_alias": organization_alias, + "budget_id": budget_id, + "created_by": created_by, + "updated_by": created_by, + } + if organization_id is not None: + data["organization_id"] = organization_id + if metadata is not None: + data["metadata"] = metadata + if models is not None: + data["models"] = models + if object_permission_id is not None: + data["object_permission_id"] = object_permission_id + + return await self.create(data) + + async def update_organization( + self, + organization_id: str, + updated_by: str, + organization_alias: Optional[str] = None, + budget_id: Optional[str] = None, + metadata: Optional[Dict[str, Any]] = None, + models: Optional[List[str]] = None, + object_permission_id: Optional[str] = None, + ) -> Optional[LiteLLM_OrganizationTable]: + """Update an organization.""" + data: Dict[str, Any] = {"updated_by": updated_by} + if organization_alias is not None: + data["organization_alias"] = organization_alias + if budget_id is not None: + data["budget_id"] = budget_id + if metadata is not None: + data["metadata"] = metadata + if models is not None: + data["models"] = models + if object_permission_id is not None: + data["object_permission_id"] = object_permission_id + + return await self.update(organization_id, data, id_field="organization_id") + + async def delete_organization( + self, organization_id: str + ) -> Optional[LiteLLM_OrganizationTable]: + """Delete an organization.""" + return await self.delete(organization_id, id_field="organization_id") + + async def update_spend( + self, organization_id: str, spend: float + ) -> Optional[LiteLLM_OrganizationTable]: + """Update organization spend.""" + return await self.update( + organization_id, {"spend": spend}, id_field="organization_id" + ) diff --git a/litellm/repositories/project_repository.py b/litellm/repositories/project_repository.py new file mode 100644 index 00000000000..86567dd05fb --- /dev/null +++ b/litellm/repositories/project_repository.py @@ -0,0 +1,129 @@ +""" +Project repository for database operations on LiteLLM_ProjectTable. +""" + +from typing import Any, Dict, List, Optional, Type + +from litellm.models.project import LiteLLM_ProjectTable +from litellm.repositories.base_repository import BaseRepository + + +class ProjectRepository(BaseRepository[LiteLLM_ProjectTable]): + """Repository for project database operations.""" + + @property + def table(self) -> Any: + return self.prisma_client.db.litellm_projecttable + + @property + def model_class(self) -> Type[LiteLLM_ProjectTable]: + return LiteLLM_ProjectTable + + async def find_by_id( + self, project_id: str, id_field: str = "project_id" + ) -> Optional[LiteLLM_ProjectTable]: + return await super().find_by_id(project_id, id_field) + + async def find_by_alias(self, project_alias: str) -> Optional[LiteLLM_ProjectTable]: + """Find a project by alias.""" + records = await self.table.find_many(where={"project_alias": project_alias}) + if records: + return self._to_model(records[0]) + return None + + async def find_by_team_id(self, team_id: str) -> List[LiteLLM_ProjectTable]: + """Find all projects belonging to a team.""" + records = await self.table.find_many(where={"team_id": team_id}) + return self._to_model_list(records) + + async def create_project( + self, + created_by: str, + project_id: Optional[str] = None, + project_alias: Optional[str] = None, + description: Optional[str] = None, + team_id: Optional[str] = None, + budget_id: Optional[str] = None, + metadata: Optional[Dict[str, Any]] = None, + models: Optional[List[str]] = None, + model_rpm_limit: Optional[Dict[str, int]] = None, + model_tpm_limit: Optional[Dict[str, int]] = None, + object_permission_id: Optional[str] = None, + ) -> LiteLLM_ProjectTable: + """Create a new project.""" + data: Dict[str, Any] = { + "created_by": created_by, + "updated_by": created_by, + } + if project_id is not None: + data["project_id"] = project_id + if project_alias is not None: + data["project_alias"] = project_alias + if description is not None: + data["description"] = description + if team_id is not None: + data["team_id"] = team_id + if budget_id is not None: + data["budget_id"] = budget_id + if metadata is not None: + data["metadata"] = metadata + if models is not None: + data["models"] = models + if model_rpm_limit is not None: + data["model_rpm_limit"] = model_rpm_limit + if model_tpm_limit is not None: + data["model_tpm_limit"] = model_tpm_limit + if object_permission_id is not None: + data["object_permission_id"] = object_permission_id + + return await self.create(data) + + async def update_project( + self, + project_id: str, + updated_by: str, + project_alias: Optional[str] = None, + description: Optional[str] = None, + team_id: Optional[str] = None, + budget_id: Optional[str] = None, + metadata: Optional[Dict[str, Any]] = None, + models: Optional[List[str]] = None, + model_rpm_limit: Optional[Dict[str, int]] = None, + model_tpm_limit: Optional[Dict[str, int]] = None, + blocked: Optional[bool] = None, + object_permission_id: Optional[str] = None, + ) -> Optional[LiteLLM_ProjectTable]: + """Update a project.""" + data: Dict[str, Any] = {"updated_by": updated_by} + if project_alias is not None: + data["project_alias"] = project_alias + if description is not None: + data["description"] = description + if team_id is not None: + data["team_id"] = team_id + if budget_id is not None: + data["budget_id"] = budget_id + if metadata is not None: + data["metadata"] = metadata + if models is not None: + data["models"] = models + if model_rpm_limit is not None: + data["model_rpm_limit"] = model_rpm_limit + if model_tpm_limit is not None: + data["model_tpm_limit"] = model_tpm_limit + if blocked is not None: + data["blocked"] = blocked + if object_permission_id is not None: + data["object_permission_id"] = object_permission_id + + return await self.update(project_id, data, id_field="project_id") + + async def delete_project(self, project_id: str) -> Optional[LiteLLM_ProjectTable]: + """Delete a project.""" + return await self.delete(project_id, id_field="project_id") + + async def update_spend( + self, project_id: str, spend: float + ) -> Optional[LiteLLM_ProjectTable]: + """Update project spend.""" + return await self.update(project_id, {"spend": spend}, id_field="project_id") diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py new file mode 100644 index 00000000000..47ea11c0592 --- /dev/null +++ b/litellm/repositories/table_repositories.py @@ -0,0 +1,215 @@ +""" +Passthrough table repositories. + +Each repository centralizes access to a single Prisma table behind a ``table`` +property, making the repository the one place that names the underlying table. +These are thin wrappers for tables that do not (yet) need domain-specific query +methods; richer repositories live in their own modules. +""" + +from typing import Any + + +class PrismaTableRepository: + """Base for repositories that expose a single Prisma table.""" + + table_name: str + + def __init__(self, prisma_client: Any): + self._prisma_client = prisma_client + + @property + def prisma_client(self) -> Any: + if self._prisma_client is None: + raise RuntimeError( + "No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys" + ) + return self._prisma_client + + @property + def table(self) -> Any: + return getattr(self.prisma_client.db, self.table_name) + + +class PolicyRepository(PrismaTableRepository): + table_name = "litellm_policytable" + + +class AgentsRepository(PrismaTableRepository): + table_name = "litellm_agentstable" + + +class GuardrailsRepository(PrismaTableRepository): + table_name = "litellm_guardrailstable" + + +class MCPServerRepository(PrismaTableRepository): + table_name = "litellm_mcpservertable" + + +class ManagedObjectRepository(PrismaTableRepository): + table_name = "litellm_managedobjecttable" + + +class OrganizationMembershipRepository(PrismaTableRepository): + table_name = "litellm_organizationmembership" + + +class SpendLogsRepository(PrismaTableRepository): + table_name = "litellm_spendlogs" + + +class ClaudeCodePluginRepository(PrismaTableRepository): + table_name = "litellm_claudecodeplugintable" + + +class TeamMembershipRepository(PrismaTableRepository): + table_name = "litellm_teammembership" + + +class EndUserRepository(PrismaTableRepository): + table_name = "litellm_endusertable" + + +class ManagedVectorStoresRepository(PrismaTableRepository): + table_name = "litellm_managedvectorstorestable" + + +class MCPUserCredentialsRepository(PrismaTableRepository): + table_name = "litellm_mcpusercredentials" + + +class PromptRepository(PrismaTableRepository): + table_name = "litellm_prompttable" + + +class TagRepository(PrismaTableRepository): + table_name = "litellm_tagtable" + + +class InvitationLinkRepository(PrismaTableRepository): + table_name = "litellm_invitationlink" + + +class JWTKeyMappingRepository(PrismaTableRepository): + table_name = "litellm_jwtkeymapping" + + +class ManagedFileRepository(PrismaTableRepository): + table_name = "litellm_managedfiletable" + + +class MemoryRepository(PrismaTableRepository): + table_name = "litellm_memorytable" + + +class SearchToolsRepository(PrismaTableRepository): + table_name = "litellm_searchtoolstable" + + +class ConfigOverridesRepository(PrismaTableRepository): + table_name = "litellm_configoverrides" + + +class MCPToolsetRepository(PrismaTableRepository): + table_name = "litellm_mcptoolsettable" + + +class ToolRepository(PrismaTableRepository): + table_name = "litellm_tooltable" + + +class DeletedVerificationTokenRepository(PrismaTableRepository): + table_name = "litellm_deletedverificationtoken" + + +class WorkflowRunRepository(PrismaTableRepository): + table_name = "litellm_workflowrun" + + +class ModelTableRepository(PrismaTableRepository): + table_name = "litellm_modeltable" + + +class AccessGroupRepository(PrismaTableRepository): + table_name = "litellm_accessgrouptable" + + +class SSOConfigRepository(PrismaTableRepository): + table_name = "litellm_ssoconfig" + + +class UISettingsRepository(PrismaTableRepository): + table_name = "litellm_uisettings" + + +class DailyGuardrailMetricsRepository(PrismaTableRepository): + table_name = "litellm_dailyguardrailmetrics" + + +class PolicyAttachmentRepository(PrismaTableRepository): + table_name = "litellm_policyattachmenttable" + + +class DeletedTeamRepository(PrismaTableRepository): + table_name = "litellm_deletedteamtable" + + +class SkillsRepository(PrismaTableRepository): + table_name = "litellm_skillstable" + + +class CacheConfigRepository(PrismaTableRepository): + table_name = "litellm_cacheconfig" + + +class ManagedVectorStoreIndexRepository(PrismaTableRepository): + table_name = "litellm_managedvectorstoreindextable" + + +class WorkflowMessageRepository(PrismaTableRepository): + table_name = "litellm_workflowmessage" + + +class DailyTagSpendRepository(PrismaTableRepository): + table_name = "litellm_dailytagspend" + + +class SpendLogToolIndexRepository(PrismaTableRepository): + table_name = "litellm_spendlogtoolindex" + + +class SpendLogGuardrailIndexRepository(PrismaTableRepository): + table_name = "litellm_spendlogguardrailindex" + + +class UserNotificationsRepository(PrismaTableRepository): + table_name = "litellm_usernotifications" + + +class HealthCheckRepository(PrismaTableRepository): + table_name = "litellm_healthchecktable" + + +class DeprecatedVerificationTokenRepository(PrismaTableRepository): + table_name = "litellm_deprecatedverificationtoken" + + +class WorkflowEventRepository(PrismaTableRepository): + table_name = "litellm_workflowevent" + + +class DailyPolicyMetricsRepository(PrismaTableRepository): + table_name = "litellm_dailypolicymetrics" + + +class AdaptiveRouterStateRepository(PrismaTableRepository): + table_name = "litellm_adaptiverouterstate" + + +class AuditLogRepository(PrismaTableRepository): + table_name = "litellm_auditlog" + + +class AdaptiveRouterSessionRepository(PrismaTableRepository): + table_name = "litellm_adaptiveroutersession" diff --git a/litellm/repositories/team_repository.py b/litellm/repositories/team_repository.py new file mode 100644 index 00000000000..2ae6647060c --- /dev/null +++ b/litellm/repositories/team_repository.py @@ -0,0 +1,351 @@ +""" +Team repository for database operations on LiteLLM_TeamTable. +""" + +import json +from datetime import datetime +from typing import Any, Dict, List, Optional, Type + +from litellm.models.team import LiteLLM_TeamTable +from litellm.repositories.base_repository import BaseRepository + + +class TeamRepository(BaseRepository[LiteLLM_TeamTable]): + """Repository for team database operations.""" + + @property + def table(self) -> Any: + return self.prisma_client.db.litellm_teamtable + + @property + def deleted_table(self) -> Any: + return self.prisma_client.db.litellm_deletedteamtable + + @property + def model_class(self) -> Type[LiteLLM_TeamTable]: + return LiteLLM_TeamTable + + def _to_model(self, record: Any) -> Optional[LiteLLM_TeamTable]: + """Convert a database record to a Team model.""" + if record is None: + return None + + data = record.dict() if hasattr(record, "dict") else dict(record) + + json_fields = [ + "metadata", + "model_spend", + "model_max_budget", + "router_settings", + "budget_limits", + "members_with_roles", + ] + for field in json_fields: + if isinstance(data.get(field), str): + data[field] = json.loads(data[field]) + + return LiteLLM_TeamTable(**data) + + async def find_by_id( + self, team_id: str, id_field: str = "team_id" + ) -> Optional[LiteLLM_TeamTable]: + return await super().find_by_id(team_id, id_field) + + async def find_by_alias(self, team_alias: str) -> Optional[LiteLLM_TeamTable]: + """Find a team by alias.""" + records = await self.table.find_many(where={"team_alias": team_alias}) + if records: + return self._to_model(records[0]) + return None + + async def find_by_organization_id( + self, organization_id: str + ) -> List[LiteLLM_TeamTable]: + """Find all teams belonging to an organization.""" + records = await self.table.find_many(where={"organization_id": organization_id}) + return self._to_model_list(records) + + async def find_by_member(self, user_id: str) -> List[LiteLLM_TeamTable]: + """Find all teams where user is a member.""" + records = await self.table.find_many(where={"members": {"has": user_id}}) + return self._to_model_list(records) + + async def find_by_admin(self, user_id: str) -> List[LiteLLM_TeamTable]: + """Find all teams where user is an admin.""" + records = await self.table.find_many(where={"admins": {"has": user_id}}) + return self._to_model_list(records) + + async def create_team( + self, + team_id: str, + team_alias: Optional[str] = None, + organization_id: Optional[str] = None, + admins: Optional[List[str]] = None, + members: Optional[List[str]] = None, + members_with_roles: Optional[Dict[str, Any]] = None, + metadata: Optional[Dict[str, Any]] = None, + max_budget: Optional[float] = None, + soft_budget: Optional[float] = None, + models: Optional[List[str]] = None, + max_parallel_requests: Optional[int] = None, + tpm_limit: Optional[int] = None, + rpm_limit: Optional[int] = None, + budget_duration: Optional[str] = None, + object_permission_id: Optional[str] = None, + ) -> LiteLLM_TeamTable: + """Create a new team.""" + data: Dict[str, Any] = {"team_id": team_id} + if team_alias is not None: + data["team_alias"] = team_alias + if organization_id is not None: + data["organization_id"] = organization_id + if admins is not None: + data["admins"] = admins + if members is not None: + data["members"] = members + if members_with_roles is not None: + data["members_with_roles"] = json.dumps(members_with_roles) + if metadata is not None: + data["metadata"] = json.dumps(metadata) + if max_budget is not None: + data["max_budget"] = max_budget + if soft_budget is not None: + data["soft_budget"] = soft_budget + if models is not None: + data["models"] = models + if max_parallel_requests is not None: + data["max_parallel_requests"] = max_parallel_requests + if tpm_limit is not None: + data["tpm_limit"] = tpm_limit + if rpm_limit is not None: + data["rpm_limit"] = rpm_limit + if budget_duration is not None: + data["budget_duration"] = budget_duration + if object_permission_id is not None: + data["object_permission_id"] = object_permission_id + + return await self.create(data) + + async def update_team( + self, + team_id: str, + team_alias: Optional[str] = None, + organization_id: Optional[str] = None, + admins: Optional[List[str]] = None, + members: Optional[List[str]] = None, + members_with_roles: Optional[Dict[str, Any]] = None, + metadata: Optional[Dict[str, Any]] = None, + max_budget: Optional[float] = None, + soft_budget: Optional[float] = None, + models: Optional[List[str]] = None, + max_parallel_requests: Optional[int] = None, + tpm_limit: Optional[int] = None, + rpm_limit: Optional[int] = None, + budget_duration: Optional[str] = None, + blocked: Optional[bool] = None, + object_permission_id: Optional[str] = None, + ) -> Optional[LiteLLM_TeamTable]: + """Update a team.""" + data: Dict[str, Any] = {} + if team_alias is not None: + data["team_alias"] = team_alias + if organization_id is not None: + data["organization_id"] = organization_id + if admins is not None: + data["admins"] = admins + if members is not None: + data["members"] = members + if members_with_roles is not None: + data["members_with_roles"] = json.dumps(members_with_roles) + if metadata is not None: + data["metadata"] = json.dumps(metadata) + if max_budget is not None: + data["max_budget"] = max_budget + if soft_budget is not None: + data["soft_budget"] = soft_budget + if models is not None: + data["models"] = models + if max_parallel_requests is not None: + data["max_parallel_requests"] = max_parallel_requests + if tpm_limit is not None: + data["tpm_limit"] = tpm_limit + if rpm_limit is not None: + data["rpm_limit"] = rpm_limit + if budget_duration is not None: + data["budget_duration"] = budget_duration + if blocked is not None: + data["blocked"] = blocked + if object_permission_id is not None: + data["object_permission_id"] = object_permission_id + + return await self.update(team_id, data, id_field="team_id") + + async def delete_team( + self, + team_id: str, + deleted_by: Optional[str] = None, + deleted_by_api_key: Optional[str] = None, + litellm_changed_by: Optional[str] = None, + ) -> Optional[LiteLLM_TeamTable]: + """Delete a team and archive it to the deleted teams table. + + Uses a transaction to ensure atomicity of the archive-then-delete operation. + """ + team = await self.find_by_id(team_id) + if team is None: + return None + + archive_data = self._build_archive_data(team) + archive_data["deleted_by"] = deleted_by + archive_data["deleted_by_api_key"] = deleted_by_api_key + archive_data["litellm_changed_by"] = litellm_changed_by + archive_data["deleted_at"] = datetime.utcnow() + + async with self.prisma_client.db.tx() as tx: + await tx.litellm_deletedteamtable.create(data=archive_data) + await tx.litellm_teamtable.delete(where={"team_id": team_id}) + + return team + + def _build_archive_data(self, team: LiteLLM_TeamTable) -> Dict[str, Any]: + """Build archive data dict with only columns that exist in LiteLLM_DeletedTeamTable.""" + data: Dict[str, Any] = {"team_id": team.team_id} + if team.team_alias is not None: + data["team_alias"] = team.team_alias + if team.organization_id is not None: + data["organization_id"] = team.organization_id + if team.object_permission_id is not None: + data["object_permission_id"] = team.object_permission_id + data["admins"] = team.admins + data["members"] = team.members + if team.members_with_roles: + data["members_with_roles"] = json.dumps( + [m.model_dump() for m in team.members_with_roles] + ) + if team.metadata: + data["metadata"] = json.dumps(team.metadata) + if team.max_budget is not None: + data["max_budget"] = team.max_budget + if team.soft_budget is not None: + data["soft_budget"] = team.soft_budget + data["spend"] = team.spend if team.spend is not None else 0.0 + data["models"] = team.models + if team.max_parallel_requests is not None: + data["max_parallel_requests"] = team.max_parallel_requests + if team.tpm_limit is not None: + data["tpm_limit"] = team.tpm_limit + if team.rpm_limit is not None: + data["rpm_limit"] = team.rpm_limit + if team.budget_duration is not None: + data["budget_duration"] = team.budget_duration + if team.budget_reset_at is not None: + data["budget_reset_at"] = team.budget_reset_at + data["blocked"] = team.blocked + if team.model_spend: + data["model_spend"] = json.dumps(team.model_spend) + if team.model_max_budget: + data["model_max_budget"] = json.dumps(team.model_max_budget) + if team.router_settings is not None: + data["router_settings"] = json.dumps(team.router_settings) + data["team_member_permissions"] = team.team_member_permissions or [] + data["access_group_ids"] = team.access_group_ids or [] + data["policies"] = team.policies or [] + if team.model_id is not None: + data["model_id"] = team.model_id + data["allow_team_guardrail_config"] = team.allow_team_guardrail_config + return data + + async def update_spend( + self, team_id: str, spend: float + ) -> Optional[LiteLLM_TeamTable]: + """Update team spend.""" + return await self.update(team_id, {"spend": spend}, id_field="team_id") + + async def add_member( + self, team_id: str, user_id: str + ) -> Optional[LiteLLM_TeamTable]: + """Add a member to a team using atomic array push operation.""" + if not await self.exists(team_id, id_field="team_id"): + return None + + record = await self.table.update( + where={"team_id": team_id}, + data={"members": {"push": user_id}}, + ) + return self._to_model(record) + + async def remove_member( + self, team_id: str, user_id: str + ) -> Optional[LiteLLM_TeamTable]: + """Remove a member from a team. + + Note: Prisma doesn't support atomic array removal, so we use a + read-modify-write pattern here. For high-concurrency scenarios, + consider using raw SQL with array_remove(). + """ + team = await self.find_by_id(team_id) + if team is None: + return None + + members = [m for m in team.members if m != user_id] + return await self.update(team_id, {"members": members}, id_field="team_id") + + async def add_admin( + self, team_id: str, user_id: str + ) -> Optional[LiteLLM_TeamTable]: + """Add an admin to a team using atomic array push operation.""" + if not await self.exists(team_id, id_field="team_id"): + return None + + record = await self.table.update( + where={"team_id": team_id}, + data={"admins": {"push": user_id}}, + ) + return self._to_model(record) + + async def remove_admin( + self, team_id: str, user_id: str + ) -> Optional[LiteLLM_TeamTable]: + """Remove an admin from a team. + + Note: Prisma doesn't support atomic array removal, so we use a + read-modify-write pattern here. For high-concurrency scenarios, + consider using raw SQL with array_remove(). + """ + team = await self.find_by_id(team_id) + if team is None: + return None + + admins = [a for a in team.admins if a != user_id] + return await self.update(team_id, {"admins": admins}, id_field="team_id") + + async def add_models( + self, team_id: str, models: List[str] + ) -> Optional[LiteLLM_TeamTable]: + """Add models to a team's allowed models list using atomic array push.""" + if not await self.exists(team_id, id_field="team_id"): + return None + + record = await self.table.update( + where={"team_id": team_id}, + data={"models": {"push": models}}, + ) + return self._to_model(record) + + async def remove_models( + self, team_id: str, models: List[str] + ) -> Optional[LiteLLM_TeamTable]: + """Remove models from a team's allowed models list. + + Note: Prisma doesn't support atomic array removal, so we use a + read-modify-write pattern here. For high-concurrency scenarios, + consider using raw SQL with array_remove(). + """ + team = await self.find_by_id(team_id) + if team is None: + return None + + current_models = [m for m in team.models if m not in models] + return await self.update( + team_id, {"models": current_models}, id_field="team_id" + ) diff --git a/litellm/repositories/user_repository.py b/litellm/repositories/user_repository.py new file mode 100644 index 00000000000..4d28b58f0ab --- /dev/null +++ b/litellm/repositories/user_repository.py @@ -0,0 +1,229 @@ +""" +User repository for database operations on LiteLLM_UserTable. +""" + +import json +from typing import Any, Dict, List, Optional, Type + +from litellm.models.user import LiteLLM_UserTable +from litellm.repositories.base_repository import BaseRepository + + +class UserRepository(BaseRepository[LiteLLM_UserTable]): + """Repository for user database operations.""" + + @property + def table(self) -> Any: + return self.prisma_client.db.litellm_usertable + + @property + def model_class(self) -> Type[LiteLLM_UserTable]: + return LiteLLM_UserTable + + def _to_model(self, record: Any) -> Optional[LiteLLM_UserTable]: + """Convert a database record to a User model.""" + if record is None: + return None + + data = record.dict() if hasattr(record, "dict") else dict(record) + + json_fields = ["metadata", "model_spend", "model_max_budget"] + for field in json_fields: + if isinstance(data.get(field), str): + data[field] = json.loads(data[field]) + + return LiteLLM_UserTable(**data) + + async def find_by_id( + self, user_id: str, id_field: str = "user_id" + ) -> Optional[LiteLLM_UserTable]: + return await super().find_by_id(user_id, id_field) + + async def find_by_email(self, user_email: str) -> Optional[LiteLLM_UserTable]: + """Find a user by email.""" + records = await self.table.find_many(where={"user_email": user_email}) + if records: + return self._to_model(records[0]) + return None + + async def find_by_sso_id(self, sso_user_id: str) -> Optional[LiteLLM_UserTable]: + """Find a user by SSO ID.""" + record = await self.table.find_unique(where={"sso_user_id": sso_user_id}) + return self._to_model(record) + + async def find_by_organization_id( + self, organization_id: str + ) -> List[LiteLLM_UserTable]: + """Find all users in an organization.""" + records = await self.table.find_many(where={"organization_id": organization_id}) + return self._to_model_list(records) + + async def find_by_team_id(self, team_id: str) -> List[LiteLLM_UserTable]: + """Find all users in a team.""" + records = await self.table.find_many(where={"teams": {"has": team_id}}) + return self._to_model_list(records) + + async def create_user( + self, + user_id: str, + user_alias: Optional[str] = None, + team_id: Optional[str] = None, + sso_user_id: Optional[str] = None, + organization_id: Optional[str] = None, + password: Optional[str] = None, + teams: Optional[List[str]] = None, + user_role: Optional[str] = None, + max_budget: Optional[float] = None, + user_email: Optional[str] = None, + models: Optional[List[str]] = None, + metadata: Optional[Dict[str, Any]] = None, + max_parallel_requests: Optional[int] = None, + tpm_limit: Optional[int] = None, + rpm_limit: Optional[int] = None, + budget_duration: Optional[str] = None, + allowed_cache_controls: Optional[List[str]] = None, + policies: Optional[List[str]] = None, + object_permission_id: Optional[str] = None, + ) -> LiteLLM_UserTable: + """Create a new user.""" + data: Dict[str, Any] = {"user_id": user_id} + if user_alias is not None: + data["user_alias"] = user_alias + if team_id is not None: + data["team_id"] = team_id + if sso_user_id is not None: + data["sso_user_id"] = sso_user_id + if organization_id is not None: + data["organization_id"] = organization_id + if password is not None: + data["password"] = password + if teams is not None: + data["teams"] = teams + if user_role is not None: + data["user_role"] = user_role + if max_budget is not None: + data["max_budget"] = max_budget + if user_email is not None: + data["user_email"] = user_email + if models is not None: + data["models"] = models + if metadata is not None: + data["metadata"] = json.dumps(metadata) + if max_parallel_requests is not None: + data["max_parallel_requests"] = max_parallel_requests + if tpm_limit is not None: + data["tpm_limit"] = tpm_limit + if rpm_limit is not None: + data["rpm_limit"] = rpm_limit + if budget_duration is not None: + data["budget_duration"] = budget_duration + if allowed_cache_controls is not None: + data["allowed_cache_controls"] = allowed_cache_controls + if policies is not None: + data["policies"] = policies + if object_permission_id is not None: + data["object_permission_id"] = object_permission_id + + return await self.create(data) + + async def update_user( + self, + user_id: str, + user_alias: Optional[str] = None, + team_id: Optional[str] = None, + sso_user_id: Optional[str] = None, + organization_id: Optional[str] = None, + password: Optional[str] = None, + teams: Optional[List[str]] = None, + user_role: Optional[str] = None, + max_budget: Optional[float] = None, + user_email: Optional[str] = None, + models: Optional[List[str]] = None, + metadata: Optional[Dict[str, Any]] = None, + max_parallel_requests: Optional[int] = None, + tpm_limit: Optional[int] = None, + rpm_limit: Optional[int] = None, + budget_duration: Optional[str] = None, + allowed_cache_controls: Optional[List[str]] = None, + policies: Optional[List[str]] = None, + object_permission_id: Optional[str] = None, + ) -> Optional[LiteLLM_UserTable]: + """Update a user.""" + data: Dict[str, Any] = {} + if user_alias is not None: + data["user_alias"] = user_alias + if team_id is not None: + data["team_id"] = team_id + if sso_user_id is not None: + data["sso_user_id"] = sso_user_id + if organization_id is not None: + data["organization_id"] = organization_id + if password is not None: + data["password"] = password + if teams is not None: + data["teams"] = teams + if user_role is not None: + data["user_role"] = user_role + if max_budget is not None: + data["max_budget"] = max_budget + if user_email is not None: + data["user_email"] = user_email + if models is not None: + data["models"] = models + if metadata is not None: + data["metadata"] = json.dumps(metadata) + if max_parallel_requests is not None: + data["max_parallel_requests"] = max_parallel_requests + if tpm_limit is not None: + data["tpm_limit"] = tpm_limit + if rpm_limit is not None: + data["rpm_limit"] = rpm_limit + if budget_duration is not None: + data["budget_duration"] = budget_duration + if allowed_cache_controls is not None: + data["allowed_cache_controls"] = allowed_cache_controls + if policies is not None: + data["policies"] = policies + if object_permission_id is not None: + data["object_permission_id"] = object_permission_id + + return await self.update(user_id, data, id_field="user_id") + + async def delete_user(self, user_id: str) -> Optional[LiteLLM_UserTable]: + """Delete a user.""" + return await self.delete(user_id, id_field="user_id") + + async def update_spend( + self, user_id: str, spend: float + ) -> Optional[LiteLLM_UserTable]: + """Update user spend.""" + return await self.update(user_id, {"spend": spend}, id_field="user_id") + + async def add_to_team( + self, user_id: str, team_id: str + ) -> Optional[LiteLLM_UserTable]: + """Add a user to a team using atomic array push operation.""" + if not await self.exists(user_id, id_field="user_id"): + return None + + record = await self.table.update( + where={"user_id": user_id}, + data={"teams": {"push": team_id}}, + ) + return self._to_model(record) + + async def remove_from_team( + self, user_id: str, team_id: str + ) -> Optional[LiteLLM_UserTable]: + """Remove a user from a team. + + Note: Prisma doesn't support atomic array removal, so we use a + read-modify-write pattern here. For high-concurrency scenarios, + consider using raw SQL with array_remove(). + """ + user = await self.find_by_id(user_id) + if user is None: + return None + + teams = [t for t in user.teams if t != team_id] + return await self.update(user_id, {"teams": teams}, id_field="user_id") diff --git a/litellm/repositories/verification_token_repository.py b/litellm/repositories/verification_token_repository.py new file mode 100644 index 00000000000..56c3e0714aa --- /dev/null +++ b/litellm/repositories/verification_token_repository.py @@ -0,0 +1,375 @@ +""" +VerificationToken repository for database operations on LiteLLM_VerificationToken. +""" + +import json +from datetime import datetime +from typing import Any, Dict, List, Optional, Type + +from litellm.models.verification_token import ( + LiteLLM_VerificationToken, +) +from litellm.repositories.base_repository import BaseRepository + + +class VerificationTokenRepository(BaseRepository[LiteLLM_VerificationToken]): + """Repository for verification token (API key) database operations.""" + + @property + def table(self) -> Any: + return self.prisma_client.db.litellm_verificationtoken + + @property + def deleted_table(self) -> Any: + return self.prisma_client.db.litellm_deletedverificationtoken + + @property + def model_class(self) -> Type[LiteLLM_VerificationToken]: + return LiteLLM_VerificationToken + + def _to_model(self, record: Any) -> Optional[LiteLLM_VerificationToken]: + """Convert a database record to a VerificationToken model.""" + if record is None: + return None + + data = record.dict() if hasattr(record, "dict") else dict(record) + + json_fields = [ + "aliases", + "config", + "permissions", + "metadata", + "model_spend", + "model_max_budget", + "router_settings", + "budget_limits", + "litellm_budget_table", + ] + for field in json_fields: + if isinstance(data.get(field), str): + data[field] = json.loads(data[field]) + + if data.get("org_id") is None and data.get("organization_id") is not None: + data["org_id"] = data["organization_id"] + + return LiteLLM_VerificationToken(**data) + + async def find_by_id( + self, token: str, id_field: str = "token" + ) -> Optional[LiteLLM_VerificationToken]: + return await super().find_by_id(token, id_field) + + async def find_by_alias( + self, key_alias: str + ) -> Optional[LiteLLM_VerificationToken]: + """Find a token by key alias.""" + records = await self.table.find_many(where={"key_alias": key_alias}) + if records: + return self._to_model(records[0]) + return None + + async def find_by_user_id(self, user_id: str) -> List[LiteLLM_VerificationToken]: + """Find all tokens belonging to a user.""" + records = await self.table.find_many(where={"user_id": user_id}) + return self._to_model_list(records) + + async def find_by_team_id(self, team_id: str) -> List[LiteLLM_VerificationToken]: + """Find all tokens belonging to a team.""" + records = await self.table.find_many(where={"team_id": team_id}) + return self._to_model_list(records) + + async def find_by_project_id( + self, project_id: str + ) -> List[LiteLLM_VerificationToken]: + """Find all tokens belonging to a project.""" + records = await self.table.find_many(where={"project_id": project_id}) + return self._to_model_list(records) + + async def find_active_tokens(self) -> List[LiteLLM_VerificationToken]: + """Find all active (non-expired, non-blocked) tokens.""" + records = await self.table.find_many( + where={ + "blocked": {"not": True}, + "OR": [{"expires": None}, {"expires": {"gt": datetime.utcnow()}}], + } + ) + return self._to_model_list(records) + + def _build_token_data( + self, + token: str, + key_name: Optional[str] = None, + key_alias: Optional[str] = None, + max_budget: Optional[float] = None, + expires: Optional[datetime] = None, + models: Optional[List[str]] = None, + aliases: Optional[Dict[str, str]] = None, + config: Optional[Dict[str, Any]] = None, + user_id: Optional[str] = None, + team_id: Optional[str] = None, + agent_id: Optional[str] = None, + project_id: Optional[str] = None, + max_parallel_requests: Optional[int] = None, + metadata: Optional[Dict[str, Any]] = None, + tpm_limit: Optional[int] = None, + rpm_limit: Optional[int] = None, + budget_duration: Optional[str] = None, + allowed_cache_controls: Optional[List[str]] = None, + allowed_routes: Optional[List[str]] = None, + permissions: Optional[Dict[str, Any]] = None, + org_id: Optional[str] = None, + created_by: Optional[str] = None, + object_permission_id: Optional[str] = None, + access_group_ids: Optional[List[str]] = None, + budget_id: Optional[str] = None, + ) -> Dict[str, Any]: + """Build data dictionary for token creation.""" + json_fields = { + "aliases": aliases, + "config": config, + "metadata": metadata, + "permissions": permissions, + } + simple_fields = { + "token": token, + "key_name": key_name, + "key_alias": key_alias, + "max_budget": max_budget, + "expires": expires, + "models": models, + "user_id": user_id, + "team_id": team_id, + "agent_id": agent_id, + "project_id": project_id, + "max_parallel_requests": max_parallel_requests, + "tpm_limit": tpm_limit, + "rpm_limit": rpm_limit, + "budget_duration": budget_duration, + "allowed_cache_controls": allowed_cache_controls, + "allowed_routes": allowed_routes, + "object_permission_id": object_permission_id, + "access_group_ids": access_group_ids, + "budget_id": budget_id, + } + data: Dict[str, Any] = {k: v for k, v in simple_fields.items() if v is not None} + for key, val in json_fields.items(): + if val is not None: + data[key] = json.dumps(val) + if org_id is not None: + data["organization_id"] = org_id + if created_by is not None: + data["created_by"] = created_by + data["updated_by"] = created_by + return data + + async def create_token( + self, + token: str, + key_name: Optional[str] = None, + key_alias: Optional[str] = None, + max_budget: Optional[float] = None, + expires: Optional[datetime] = None, + models: Optional[List[str]] = None, + aliases: Optional[Dict[str, str]] = None, + config: Optional[Dict[str, Any]] = None, + user_id: Optional[str] = None, + team_id: Optional[str] = None, + agent_id: Optional[str] = None, + project_id: Optional[str] = None, + max_parallel_requests: Optional[int] = None, + metadata: Optional[Dict[str, Any]] = None, + tpm_limit: Optional[int] = None, + rpm_limit: Optional[int] = None, + budget_duration: Optional[str] = None, + allowed_cache_controls: Optional[List[str]] = None, + allowed_routes: Optional[List[str]] = None, + permissions: Optional[Dict[str, Any]] = None, + org_id: Optional[str] = None, + created_by: Optional[str] = None, + object_permission_id: Optional[str] = None, + access_group_ids: Optional[List[str]] = None, + budget_id: Optional[str] = None, + ) -> LiteLLM_VerificationToken: + """Create a new verification token.""" + data = self._build_token_data( + token=token, + key_name=key_name, + key_alias=key_alias, + max_budget=max_budget, + expires=expires, + models=models, + aliases=aliases, + config=config, + user_id=user_id, + team_id=team_id, + agent_id=agent_id, + project_id=project_id, + max_parallel_requests=max_parallel_requests, + metadata=metadata, + tpm_limit=tpm_limit, + rpm_limit=rpm_limit, + budget_duration=budget_duration, + allowed_cache_controls=allowed_cache_controls, + allowed_routes=allowed_routes, + permissions=permissions, + org_id=org_id, + created_by=created_by, + object_permission_id=object_permission_id, + access_group_ids=access_group_ids, + budget_id=budget_id, + ) + return await self.create(data) + + async def update_token( + self, + token: str, + updated_by: Optional[str] = None, + key_name: Optional[str] = None, + key_alias: Optional[str] = None, + max_budget: Optional[float] = None, + expires: Optional[datetime] = None, + models: Optional[List[str]] = None, + aliases: Optional[Dict[str, str]] = None, + config: Optional[Dict[str, Any]] = None, + max_parallel_requests: Optional[int] = None, + metadata: Optional[Dict[str, Any]] = None, + tpm_limit: Optional[int] = None, + rpm_limit: Optional[int] = None, + budget_duration: Optional[str] = None, + allowed_cache_controls: Optional[List[str]] = None, + allowed_routes: Optional[List[str]] = None, + permissions: Optional[Dict[str, Any]] = None, + blocked: Optional[bool] = None, + object_permission_id: Optional[str] = None, + access_group_ids: Optional[List[str]] = None, + ) -> Optional[LiteLLM_VerificationToken]: + """Update a verification token.""" + data: Dict[str, Any] = {} + if updated_by is not None: + data["updated_by"] = updated_by + if key_name is not None: + data["key_name"] = key_name + if key_alias is not None: + data["key_alias"] = key_alias + if max_budget is not None: + data["max_budget"] = max_budget + if expires is not None: + data["expires"] = expires + if models is not None: + data["models"] = models + if aliases is not None: + data["aliases"] = json.dumps(aliases) + if config is not None: + data["config"] = json.dumps(config) + if max_parallel_requests is not None: + data["max_parallel_requests"] = max_parallel_requests + if metadata is not None: + data["metadata"] = json.dumps(metadata) + if tpm_limit is not None: + data["tpm_limit"] = tpm_limit + if rpm_limit is not None: + data["rpm_limit"] = rpm_limit + if budget_duration is not None: + data["budget_duration"] = budget_duration + if allowed_cache_controls is not None: + data["allowed_cache_controls"] = allowed_cache_controls + if allowed_routes is not None: + data["allowed_routes"] = allowed_routes + if permissions is not None: + data["permissions"] = json.dumps(permissions) + if blocked is not None: + data["blocked"] = blocked + if object_permission_id is not None: + data["object_permission_id"] = object_permission_id + if access_group_ids is not None: + data["access_group_ids"] = access_group_ids + + return await self.update(token, data, id_field="token") + + async def delete_token( + self, + token: str, + deleted_by: Optional[str] = None, + deleted_by_api_key: Optional[str] = None, + litellm_changed_by: Optional[str] = None, + ) -> Optional[LiteLLM_VerificationToken]: + """Delete a token and archive it to the deleted tokens table. + + Uses a transaction to ensure atomicity of the archive-then-delete operation. + """ + token_record = await self.find_by_id(token) + if token_record is None: + return None + + archive_data = self._build_archive_data(token_record) + archive_data["deleted_by"] = deleted_by + archive_data["deleted_by_api_key"] = deleted_by_api_key + archive_data["litellm_changed_by"] = litellm_changed_by + archive_data["deleted_at"] = datetime.utcnow() + + async with self.prisma_client.db.tx() as tx: + await tx.litellm_deletedverificationtoken.create(data=archive_data) + await tx.litellm_verificationtoken.delete(where={"token": token}) + + return token_record + + def _build_archive_data(self, token: LiteLLM_VerificationToken) -> Dict[str, Any]: + """Build archive data with only columns present in LiteLLM_DeletedVerificationToken. + + Serializes JSON columns to strings (the archive table stores them as JSON + columns the same way the live table does) and maps ``org_id`` onto the + ``organization_id`` column so the foreign key is preserved. + """ + data = token.model_dump(exclude_none=True) + for field in ("object_permission", "litellm_budget_table", "budget_limits"): + data.pop(field, None) + + org_id = data.pop("org_id", None) + if org_id is not None: + data["organization_id"] = org_id + + json_fields = [ + "aliases", + "config", + "permissions", + "metadata", + "model_spend", + "model_max_budget", + "router_settings", + ] + for field in json_fields: + if field in data: + data[field] = json.dumps(data[field]) + return data + + async def update_spend( + self, token: str, spend: float + ) -> Optional[LiteLLM_VerificationToken]: + """Update token spend.""" + return await self.update(token, {"spend": spend}, id_field="token") + + async def update_last_active( + self, token: str + ) -> Optional[LiteLLM_VerificationToken]: + """Update the last_active timestamp.""" + return await self.update( + token, {"last_active": datetime.utcnow()}, id_field="token" + ) + + async def block_token( + self, token: str, updated_by: Optional[str] = None + ) -> Optional[LiteLLM_VerificationToken]: + """Block a token.""" + data: Dict[str, Any] = {"blocked": True} + if updated_by is not None: + data["updated_by"] = updated_by + return await self.update(token, data, id_field="token") + + async def unblock_token( + self, token: str, updated_by: Optional[str] = None + ) -> Optional[LiteLLM_VerificationToken]: + """Unblock a token.""" + data: Dict[str, Any] = {"blocked": False} + if updated_by is not None: + data["updated_by"] = updated_by + return await self.update(token, data, id_field="token") diff --git a/litellm/router_strategy/adaptive_router/adaptive_router.py b/litellm/router_strategy/adaptive_router/adaptive_router.py index 3bccef36e68..4856d7ff4cd 100644 --- a/litellm/router_strategy/adaptive_router/adaptive_router.py +++ b/litellm/router_strategy/adaptive_router/adaptive_router.py @@ -55,6 +55,7 @@ from litellm.router_strategy.adaptive_router.update_queue import ( _SESSION_STATE_SWEEP_THRESHOLD: int = 1024 # Same pattern for the owner cache. _OWNER_CACHE_SWEEP_THRESHOLD: int = 1024 +from litellm.repositories.table_repositories import AdaptiveRouterStateRepository from litellm.types.llms.openai import AllMessageValues from litellm.types.router import ( AdaptiveRouterConfig, @@ -113,7 +114,7 @@ class AdaptiveRouter: if prisma_client is None: return try: - rows = await prisma_client.db.litellm_adaptiverouterstate.find_many( + rows = await AdaptiveRouterStateRepository(prisma_client).table.find_many( where={"router_name": self.router_name} ) loaded = 0 diff --git a/litellm/router_strategy/adaptive_router/update_queue.py b/litellm/router_strategy/adaptive_router/update_queue.py index b667f3a53a7..1d87feddd84 100644 --- a/litellm/router_strategy/adaptive_router/update_queue.py +++ b/litellm/router_strategy/adaptive_router/update_queue.py @@ -22,6 +22,10 @@ import asyncio from typing import Any, Dict, Tuple from litellm._logging import verbose_router_logger +from litellm.repositories.table_repositories import ( + AdaptiveRouterSessionRepository, + AdaptiveRouterStateRepository, +) StateKey = Tuple[str, str, str] # (router_name, request_type, model_name) SessionKey = Tuple[str, str, str] # (session_id, router_name, model_name) @@ -112,7 +116,7 @@ class AdaptiveRouterUpdateQueue: # other. The upsert creates the row with the delta as the # initial value on first write, then increments on subsequent # writes — no read-modify-write race. - await prisma_client.db.litellm_adaptiverouterstate.upsert( + await AdaptiveRouterStateRepository(prisma_client).table.upsert( where={ "router_name_request_type_model_name": { "router_name": router, @@ -174,7 +178,7 @@ class AdaptiveRouterUpdateQueue: for k, v in payload.items() if k not in ("session_id", "router_name", "model_name") } - await prisma_client.db.litellm_adaptiveroutersession.upsert( + await AdaptiveRouterSessionRepository(prisma_client).table.upsert( where={ "session_id_router_name_model_name": { "session_id": session_id, diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 92ca027c5bf..809da6418d7 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -3,8 +3,7 @@ from typing import Any, Dict, List, Literal, Optional from pydantic import BaseModel, ConfigDict -from litellm.proxy._types import MCPAuthType, MCPTransportType -from litellm.types.mcp import MCPAuth +from litellm.types.mcp import MCPAuth, MCPAuthType, MCPTransportType # MCPInfo now allows arbitrary additional fields for custom metadata MCPInfo = Dict[str, Any] diff --git a/litellm/types/utils.py b/litellm/types/utils.py index b76ae1f5d86..9633cecf96c 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -38,7 +38,6 @@ from pydantic import ( Field, PrivateAttr, field_validator, - model_validator, ) from typing_extensions import Required, TypedDict @@ -3577,25 +3576,11 @@ class RawRequestTypedDict(TypedDict, total=False): error: Optional[str] -class CredentialBase(BaseModel): - credential_name: str - credential_info: dict - - -class CredentialItem(CredentialBase): - credential_values: dict - - -class CreateCredentialItem(CredentialBase): - credential_values: Optional[dict] = None - model_id: Optional[str] = None - - @model_validator(mode="before") - @classmethod - def check_credential_params(cls, values): - if not values.get("credential_values") and not values.get("model_id"): - raise ValueError("Either credential_values or model_id must be set") - return values +from litellm.models.credentials import CredentialBase as CredentialBase # noqa: E402 +from litellm.models.credentials import CredentialItem as CredentialItem # noqa: E402 +from litellm.models.credentials import ( # noqa: E402 + CreateCredentialItem as CreateCredentialItem, +) class ExtractedFileData(TypedDict): diff --git a/litellm/vector_stores/vector_store_registry.py b/litellm/vector_stores/vector_store_registry.py index 1fd95b16309..94f0483e1cc 100644 --- a/litellm/vector_stores/vector_store_registry.py +++ b/litellm/vector_stores/vector_store_registry.py @@ -5,6 +5,10 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, get_args from litellm._logging import verbose_logger from litellm.litellm_core_utils.core_helpers import remove_items_at_indices +from litellm.repositories.table_repositories import ( + ManagedVectorStoreIndexRepository, + ManagedVectorStoresRepository, +) from litellm.types.vector_stores import ( VECTOR_STORE_OPENAI_PARAMS, LiteLLM_ManagedVectorStore, @@ -91,10 +95,10 @@ class VectorStoreIndexRegistry: """ vector_stores_from_db: List[LiteLLM_ManagedVectorStoreIndex] = [] if prisma_client is not None: - _vector_stores_from_db = ( - await prisma_client.db.litellm_managedvectorstoreindextable.find_many( - order={"created_at": "desc"}, - ) + _vector_stores_from_db = await ManagedVectorStoreIndexRepository( + prisma_client + ).table.find_many( + order={"created_at": "desc"}, ) for vector_store in _vector_stores_from_db: _dict_vector_store = dict(vector_store) @@ -374,9 +378,9 @@ class VectorStoreRegistry: if vector_store is not None and prisma_client is not None: try: # Check if it still exists in database - db_vector_store = await prisma_client.db.litellm_managedvectorstorestable.find_unique( - where={"vector_store_id": vector_store_id} - ) + db_vector_store = await ManagedVectorStoresRepository( + prisma_client + ).table.find_unique(where={"vector_store_id": vector_store_id}) if db_vector_store is None: # Vector store was deleted from database, remove from cache verbose_logger.debug( @@ -541,10 +545,10 @@ class VectorStoreRegistry: """ vector_stores_from_db: List[LiteLLM_ManagedVectorStore] = [] if prisma_client is not None: - _vector_stores_from_db = ( - await prisma_client.db.litellm_managedvectorstorestable.find_many( - order={"created_at": "desc"}, - ) + _vector_stores_from_db = await ManagedVectorStoresRepository( + prisma_client + ).table.find_many( + order={"created_at": "desc"}, ) for vector_store in _vector_stores_from_db: _dict_vector_store = dict(vector_store) diff --git a/tests/test_litellm/models/test_models.py b/tests/test_litellm/models/test_models.py new file mode 100644 index 00000000000..786f6244930 --- /dev/null +++ b/tests/test_litellm/models/test_models.py @@ -0,0 +1,542 @@ +""" +Tests for backend domain models. +""" + +from datetime import datetime + +import pytest + +from litellm.models.access_group import LiteLLM_AccessGroupTable +from litellm.models.budget import ( + LiteLLM_BudgetTable, + LiteLLM_BudgetTableFull, + LiteLLM_TeamMemberTable, +) +from litellm.models.config import LiteLLM_Config +from litellm.models.credentials import CreateCredentialItem, CredentialItem +from litellm.models.end_user import LiteLLM_EndUserTable +from litellm.models.managed_files import ( + LiteLLM_ManagedFileTable, + LiteLLM_ManagedObjectTable, + LiteLLM_ManagedVectorStoresTable, +) +from litellm.models.mcp_server import LiteLLM_MCPServerTable +from litellm.models.model import LiteLLM_ProxyModelTable +from litellm.models.object_permission import LiteLLM_ObjectPermissionTable +from litellm.models.organization import LiteLLM_OrganizationTable +from litellm.models.project import LiteLLM_ProjectTable +from litellm.models.skills import LiteLLM_SkillsTable +from litellm.models.spend_logs import LiteLLM_ErrorLogs, LiteLLM_SpendLogs +from litellm.models.tag import LiteLLM_TagTable +from litellm.models.team import ( + LiteLLM_DeletedTeamTable, + LiteLLM_TeamTable, + LiteLLM_TeamTableCachedObj, +) +from litellm.models.team_membership import LiteLLM_TeamMembership +from litellm.models.user import LiteLLM_UserTable +from litellm.models.verification_token import ( + LiteLLM_DeletedVerificationToken, + LiteLLM_VerificationToken, +) + + +class TestBudget: + def test_budget_creation(self): + budget = LiteLLM_BudgetTable( + budget_id="test-budget-id", + max_budget=100.0, + soft_budget=80.0, + tpm_limit=1000, + rpm_limit=100, + model_max_budget={"gpt-4": 50.0}, + budget_duration="monthly", + allowed_models=["gpt-4"], + ) + assert budget.budget_id == "test-budget-id" + assert budget.max_budget == 100.0 + assert budget.soft_budget == 80.0 + assert budget.tpm_limit == 1000 + assert budget.rpm_limit == 100 + assert budget.model_max_budget == {"gpt-4": 50.0} + assert budget.budget_duration == "monthly" + assert budget.allowed_models == ["gpt-4"] + + def test_budget_defaults(self): + budget = LiteLLM_BudgetTable() + assert budget.budget_id is None + assert budget.max_budget is None + assert budget.allowed_models is None + + +class TestCredentials: + def test_credentials_creation(self): + creds = CredentialItem( + credential_name="test-cred", + credential_values={"api_key": "secret123"}, + credential_info={"provider": "openai"}, + ) + assert creds.credential_name == "test-cred" + assert creds.credential_values["api_key"] == "secret123" + assert creds.credential_info["provider"] == "openai" + + def test_create_credential_item_accepts_model_id(self): + item = CreateCredentialItem( + credential_name="from-model", + credential_info={}, + model_id="model-123", + ) + assert item.model_id == "model-123" + assert item.credential_values is None + + def test_create_credential_item_requires_values_or_model_id(self): + with pytest.raises( + ValueError, match="Either credential_values or model_id must be set" + ): + CreateCredentialItem(credential_name="bad", credential_info={}) + + +class TestModel: + def test_model_creation(self): + model = LiteLLM_ProxyModelTable( + model_id="test-model-id", + model_name="gpt-4", + litellm_params={"model": "gpt-4", "api_key": "test"}, + model_info={"team_id": "team-123", "team_public_model_name": "my-gpt4"}, + ) + assert model.model_id == "test-model-id" + assert model.model_name == "gpt-4" + assert model.team_id == "team-123" + assert model.team_public_model_name == "my-gpt4" + + def test_is_blocked(self): + model_blocked = LiteLLM_ProxyModelTable( + model_id="m1", model_name="test", litellm_params={}, blocked=True + ) + model_unblocked = LiteLLM_ProxyModelTable( + model_id="m2", model_name="test", litellm_params={}, blocked=False + ) + assert model_blocked.is_blocked + assert not model_unblocked.is_blocked + + def test_parses_json_string_fields(self): + model = LiteLLM_ProxyModelTable( + model_id="m1", + model_name="gpt-4", + litellm_params='{"model": "gpt-4"}', + model_info='{"team_id": "t1"}', + ) + assert model.litellm_params == {"model": "gpt-4"} + assert model.model_info == {"team_id": "t1"} + + def test_team_helpers_none_when_no_model_info(self): + model = LiteLLM_ProxyModelTable( + model_id="m1", model_name="gpt-4", litellm_params={}, model_info=None + ) + assert model.team_id is None + assert model.team_public_model_name is None + + +class TestObjectPermission: + def test_object_permission_creation(self): + perm = LiteLLM_ObjectPermissionTable( + object_permission_id="test-perm-id", + mcp_servers=["server1", "server2"], + vector_stores=["vs1"], + agents=["agent1"], + models=["gpt-4"], + blocked_tools=["dangerous_tool"], + ) + assert perm.object_permission_id == "test-perm-id" + assert len(perm.mcp_servers) == 2 + assert perm.vector_stores == ["vs1"] + assert perm.agents == ["agent1"] + assert perm.models == ["gpt-4"] + assert perm.blocked_tools == ["dangerous_tool"] + + def test_object_permission_tool_permissions(self): + perm = LiteLLM_ObjectPermissionTable( + object_permission_id="perm-tools", + mcp_tool_permissions={"server1": ["tool1", "tool2"]}, + ) + assert perm.mcp_tool_permissions == {"server1": ["tool1", "tool2"]} + + +class TestOrganization: + def test_organization_creation(self): + org = LiteLLM_OrganizationTable( + organization_id="org-123", + organization_alias="My Org", + budget_id="budget-123", + models=["gpt-4", "claude-3"], + spend=50.0, + created_by="admin", + updated_by="admin", + ) + assert org.organization_id == "org-123" + assert org.organization_alias == "My Org" + assert len(org.models) == 2 + + +class TestProject: + def test_project_creation(self): + project = LiteLLM_ProjectTable( + project_id="proj-123", + project_alias="My Project", + team_id="team-123", + blocked=False, + ) + assert project.project_id == "proj-123" + assert not project.is_blocked + + +class TestTeam: + def test_team_creation(self): + team = LiteLLM_TeamTable( + team_id="team-123", + team_alias="Engineering", + admins=["user1"], + members=["user2", "user3"], + models=["gpt-4"], + max_budget=1000.0, + spend=100.0, + ) + assert team.team_id == "team-123" + assert team.team_alias == "Engineering" + assert team.admins == ["user1"] + assert team.members == ["user2", "user3"] + assert team.models == ["gpt-4"] + assert team.max_budget == 1000.0 + + def test_members_with_roles_parsing(self): + team = LiteLLM_TeamTable( + team_id="t2", + members_with_roles=[ + {"user_id": "user1", "role": "admin"}, + {"user_id": "user2", "role": "user"}, + ], + ) + assert len(team.members_with_roles) == 2 + assert team.members_with_roles[0].user_id == "user1" + assert team.members_with_roles[0].role == "admin" + + def test_members_with_roles_empty_dict_coerced(self): + team = LiteLLM_TeamTable(team_id="t3", members_with_roles={}) + assert team.members_with_roles == [] + + def test_json_string_fields_parsed(self): + team = LiteLLM_TeamTable( + team_id="t4", + metadata='{"k": "v"}', + model_max_budget='{"gpt-4": 5.0}', + ) + assert team.metadata == {"k": "v"} + assert team.model_max_budget == {"gpt-4": 5.0} + + def test_cached_team(self): + cached = LiteLLM_TeamTableCachedObj( + team_id="t1", last_refreshed_at=1234567890.0 + ) + assert cached.last_refreshed_at == 1234567890.0 + + def test_deleted_team(self): + deleted = LiteLLM_DeletedTeamTable( + team_id="t1", + deleted_by="admin", + deleted_at=datetime.utcnow(), + ) + assert deleted.deleted_by == "admin" + assert deleted.deleted_at is not None + + +class TestUser: + def test_user_creation(self): + user = LiteLLM_UserTable( + user_id="user-123", + user_email="test@example.com", + teams=["team1", "team2"], + max_budget=100.0, + spend=25.0, + ) + assert user.user_id == "user-123" + assert user.user_email == "test@example.com" + assert len(user.teams) == 2 + + def test_is_over_budget(self): + user = LiteLLM_UserTable(user_id="u1", max_budget=100.0, spend=150.0) + user_no_budget = LiteLLM_UserTable(user_id="u2", spend=1000.0) + + assert user.is_over_budget() + assert not user_no_budget.is_over_budget() + + def test_has_model_access(self): + user_with_models = LiteLLM_UserTable(user_id="u1", models=["gpt-4"]) + user_no_models = LiteLLM_UserTable(user_id="u2", models=[]) + + assert user_with_models.has_model_access("gpt-4") + assert not user_with_models.has_model_access("gpt-3") + assert user_no_models.has_model_access("any-model") + + def test_password_hash_excluded_from_serialization(self): + from litellm.proxy._types import LiteLLM_UserTableWithKeyCount + + secret = "$2b$12$abcdefghijklmnopqrstuv" + user = LiteLLM_UserTable(user_id="u1", user_email="a@b.c", password=secret) + + assert user.password == secret + assert "password" not in user.model_dump() + assert "password" not in user.model_dump_json() + + with_keys = LiteLLM_UserTableWithKeyCount( + user_id="u1", user_email="a@b.c", password=secret, key_count=2 + ) + assert with_keys.password == secret + assert "password" not in with_keys.model_dump() + assert "password" not in with_keys.model_dump_json() + + +class TestVerificationToken: + def test_verification_token_creation(self): + token = LiteLLM_VerificationToken( + token="sk-test123", + key_name="Test Key", + user_id="user-123", + team_id="team-123", + max_budget=100.0, + spend=25.0, + models=["gpt-4"], + blocked=True, + allowed_routes=["/chat/completions"], + ) + assert token.token == "sk-test123" + assert token.key_name == "Test Key" + assert token.user_id == "user-123" + assert token.team_id == "team-123" + assert token.blocked is True + assert token.models == ["gpt-4"] + assert token.allowed_routes == ["/chat/completions"] + + def test_expires_accepts_string_and_datetime(self): + as_str = LiteLLM_VerificationToken(token="t1", expires="2024-12-31T23:59:59Z") + as_dt = LiteLLM_VerificationToken(token="t2", expires=datetime.utcnow()) + assert as_str.expires == "2024-12-31T23:59:59Z" + assert isinstance(as_dt.expires, datetime) + + def test_deleted_verification_token(self): + deleted = LiteLLM_DeletedVerificationToken( + token="t1", + deleted_by="admin", + deleted_at=datetime.utcnow(), + ) + assert deleted.deleted_by == "admin" + assert deleted.deleted_at is not None + assert deleted.token == "t1" + + +class TestConfigTable: + def test_config_creation(self): + cfg = LiteLLM_Config(param_name="general_settings", param_value={"k": "v"}) + assert cfg.param_name == "general_settings" + assert cfg.param_value == {"k": "v"} + + +class TestSkillsTable: + def test_skills_creation(self): + skill = LiteLLM_SkillsTable( + skill_id="s1", + display_title="My Skill", + source="custom", + file_content=b"zipbytes", + file_name="skill.zip", + ) + assert skill.skill_id == "s1" + assert skill.display_title == "My Skill" + assert skill.file_content == b"zipbytes" + + def test_skills_defaults(self): + skill = LiteLLM_SkillsTable(skill_id="s2") + assert skill.source == "custom" + assert skill.metadata is None + + +class TestAccessGroupTable: + def test_access_group_creation(self): + ag = LiteLLM_AccessGroupTable( + access_group_id="ag1", + access_group_name="group-a", + access_model_names=["gpt-4"], + assigned_team_ids=["t1"], + ) + assert ag.access_group_id == "ag1" + assert ag.access_model_names == ["gpt-4"] + assert ag.assigned_team_ids == ["t1"] + assert ag.access_agent_ids == [] + + +class TestTagTable: + def test_tag_creation(self): + tag = LiteLLM_TagTable( + tag_name="prod", + models=["gpt-4"], + spend=12.5, + budget_id="b1", + ) + assert tag.tag_name == "prod" + assert tag.models == ["gpt-4"] + assert tag.spend == 12.5 + + def test_tag_set_model_info_coerces_none(self): + tag = LiteLLM_TagTable(tag_name="t", spend=None, models=None) + assert tag.spend == 0.0 + assert tag.models == [] + + +class TestEndUserTable: + def test_end_user_creation(self): + eu = LiteLLM_EndUserTable( + user_id="eu1", + blocked=False, + spend=5.0, + allowed_model_region="eu", + default_model="gpt-4", + ) + assert eu.user_id == "eu1" + assert eu.blocked is False + assert eu.allowed_model_region == "eu" + assert eu.default_model == "gpt-4" + + def test_end_user_spend_coerced_when_none(self): + eu = LiteLLM_EndUserTable(user_id="eu2", blocked=True, spend=None) + assert eu.spend == 0.0 + + +class TestBudgetTableFull: + def test_full_adds_server_managed_fields(self): + now = datetime.now() + budget = LiteLLM_BudgetTableFull( + budget_id="b1", max_budget=10.0, created_at=now, budget_reset_at=now + ) + assert budget.created_at == now + assert budget.budget_reset_at == now + assert budget.max_budget == 10.0 + + def test_full_requires_created_at(self): + with pytest.raises(Exception): + LiteLLM_BudgetTableFull(budget_id="b1") + + +class TestTeamMemberTable: + def test_tracks_user_within_team(self): + member = LiteLLM_TeamMemberTable( + user_id="u1", team_id="t1", spend=3.0, budget_id="b1", max_budget=5.0 + ) + assert member.user_id == "u1" + assert member.team_id == "t1" + assert member.spend == 3.0 + assert member.max_budget == 5.0 + + +class TestTeamMembership: + def test_safe_get_limits_with_budget_table(self): + membership = LiteLLM_TeamMembership( + user_id="u1", + team_id="t1", + litellm_budget_table=LiteLLM_BudgetTable(rpm_limit=100, tpm_limit=2000), + ) + assert membership.safe_get_team_member_rpm_limit() == 100 + assert membership.safe_get_team_member_tpm_limit() == 2000 + + def test_safe_get_limits_without_budget_table(self): + membership = LiteLLM_TeamMembership(user_id="u1", team_id="t1") + assert membership.safe_get_team_member_rpm_limit() is None + assert membership.safe_get_team_member_tpm_limit() is None + + def test_full_budget_variant_parsed_for_server_fields(self): + now = datetime.now() + membership = LiteLLM_TeamMembership( + user_id="u1", + team_id="t1", + litellm_budget_table={ + "budget_id": "b1", + "rpm_limit": 7, + "created_at": now, + "budget_reset_at": now, + }, + ) + assert isinstance(membership.litellm_budget_table, LiteLLM_BudgetTableFull) + assert membership.safe_get_team_member_rpm_limit() == 7 + + +class TestMCPServerTable: + def test_mcp_server_defaults(self): + server = LiteLLM_MCPServerTable(server_id="s1", transport="sse") + assert server.server_id == "s1" + assert server.transport == "sse" + assert server.status == "unknown" + assert server.approval_status == "active" + assert server.allow_all_keys is False + assert server.available_on_public_internet is True + assert server.teams == [] + assert server.env == {} + + def test_mcp_server_requires_transport(self): + with pytest.raises(Exception): + LiteLLM_MCPServerTable(server_id="s1") + + +class TestSpendLogs: + def test_spend_logs_creation(self): + log = LiteLLM_SpendLogs( + request_id="r1", + api_key="sk-1", + call_type="completion", + startTime=None, + endTime=None, + messages=None, + response=None, + ) + assert log.request_id == "r1" + assert log.spend == 0.0 + assert log.cache_hit == "False" + + def test_error_logs_creation(self): + log = LiteLLM_ErrorLogs( + request_id="r1", startTime=None, endTime=None, status_code="500" + ) + assert log.request_id == "r1" + assert log.status_code == "500" + + +class TestManagedTables: + def test_managed_file_table(self): + table = LiteLLM_ManagedFileTable( + unified_file_id="f1", + model_mappings={"gpt-4": "file-abc"}, + flat_model_file_ids=["file-abc"], + ) + assert table.unified_file_id == "f1" + assert table.model_mappings == {"gpt-4": "file-abc"} + assert table.flat_model_file_ids == ["file-abc"] + + def test_managed_object_table_requires_purpose(self): + with pytest.raises(Exception): + LiteLLM_ManagedObjectTable( + unified_object_id="o1", model_object_id="m1", file_object={} + ) + + def test_managed_vector_stores_table(self): + table = LiteLLM_ManagedVectorStoresTable( + vector_store_id="vs1", + custom_llm_provider="openai", + vector_store_name=None, + vector_store_description=None, + vector_store_metadata=None, + created_at=None, + updated_at=None, + litellm_credential_name=None, + litellm_params=None, + team_id=None, + user_id=None, + ) + assert table.vector_store_id == "vs1" + assert table.custom_llm_provider == "openai" diff --git a/tests/test_litellm/repositories/test_repositories.py b/tests/test_litellm/repositories/test_repositories.py new file mode 100644 index 00000000000..f22debbae34 --- /dev/null +++ b/tests/test_litellm/repositories/test_repositories.py @@ -0,0 +1,2184 @@ +""" +Tests for gateway repository layer. +""" + +import json +from datetime import datetime +from typing import Any, Dict, List, Optional +from unittest.mock import MagicMock, patch + +import pytest + +from litellm.models.base import DomainModel +from litellm.models.budget import LiteLLM_BudgetTable +from litellm.models.credentials import CredentialItem +from litellm.models.team import LiteLLM_TeamTable +from litellm.repositories.base_repository import BaseRepository +from litellm.repositories.budget_repository import BudgetRepository +from litellm.repositories.config_repository import ConfigRepository +from litellm.repositories.credentials_repository import CredentialsRepository +from litellm.repositories.model_repository import ModelRepository +from litellm.repositories.object_permission_repository import ( + ObjectPermissionRepository, +) +from litellm.repositories.organization_repository import OrganizationRepository +from litellm.repositories.project_repository import ProjectRepository +from litellm.repositories.team_repository import TeamRepository +from litellm.repositories.user_repository import UserRepository +from litellm.repositories.verification_token_repository import ( + VerificationTokenRepository, +) + + +class MockRecord: + """Mock database record for testing.""" + + def __init__(self, data: Dict[str, Any]): + self._data = data if data is not None else {} + + def dict(self) -> Dict[str, Any]: + return self._data.copy() + + def model_dump(self) -> Dict[str, Any]: + return self._data.copy() + + def __getattr__(self, name: str) -> Any: + if name.startswith("_"): + raise AttributeError(name) + return self._data.get(name) + + +class MockTable: + """Mock Prisma table for testing.""" + + def __init__(self, pk_field: Optional[str] = None): + self._records: Dict[str, Dict[str, Any]] = {} + self._pk_field = pk_field + + async def find_unique(self, where: Dict[str, Any]) -> Optional[MockRecord]: + key_field = list(where.keys())[0] + key_value = where[key_field] + data = self._records.get(key_value) + return MockRecord(data) if data else None + + async def find_many( + self, + where: Optional[Dict[str, Any]] = None, + skip: Optional[int] = None, + take: Optional[int] = None, + order: Optional[Dict[str, str]] = None, + ) -> List[MockRecord]: + records = list(self._records.values()) + return [MockRecord(r) for r in records] + + async def create(self, data: Dict[str, Any]) -> MockRecord: + record_data = dict(data) + if self._pk_field and self._pk_field not in record_data: + record_data[self._pk_field] = f"{self._pk_field}-{len(self._records)}" + key = ( + record_data.get(self._pk_field) + if self._pk_field + else record_data.get("id", str(len(self._records))) + ) + self._records[key] = record_data + return MockRecord(record_data) + + async def update( + self, where: Dict[str, Any], data: Dict[str, Any] + ) -> Optional[MockRecord]: + key_field = list(where.keys())[0] + key_value = where[key_field] + if key_value in self._records: + for field, value in data.items(): + if isinstance(value, dict) and "push" in value: + current = self._records[key_value].get(field, []) + push_val = value["push"] + if isinstance(push_val, list): + current.extend(push_val) + else: + current.append(push_val) + self._records[key_value][field] = current + else: + self._records[key_value][field] = value + return MockRecord(self._records[key_value]) + return None + + async def delete(self, where: Dict[str, Any]) -> Optional[MockRecord]: + key_field = list(where.keys())[0] + key_value = where[key_field] + data = self._records.pop(key_value, None) + return MockRecord(data) if data else None + + async def count(self, where: Optional[Dict[str, Any]] = None) -> int: + return len(self._records) + + async def upsert(self, where: Dict[str, Any], data: Dict[str, Any]) -> MockRecord: + key_field = list(where.keys())[0] + key_value = where[key_field] + if key_value in self._records: + self._records[key_value].update(data.get("update", {})) + else: + self._records[key_value] = data.get("create", {}) + return MockRecord(self._records[key_value]) + + +class MockPrismaClient: + """Mock Prisma client for testing.""" + + def __init__(self): + self.db = MagicMock() + self.db.litellm_budgettable = MockTable() + self.db.litellm_proxymodeltable = MockTable(pk_field="model_id") + self.db.litellm_teamtable = MockTable() + self.db.litellm_deletedteamtable = MockTable() + self.db.litellm_usertable = MockTable() + self.db.litellm_verificationtoken = MockTable() + self.db.litellm_deletedverificationtoken = MockTable() + self.db.litellm_config = MockTable() + self.db.litellm_organizationtable = MockTable() + self.db.litellm_projecttable = MockTable(pk_field="project_id") + self.db.litellm_objectpermissiontable = MockTable( + pk_field="object_permission_id" + ) + self.db.litellm_credentialstable = MockTable() + + +class TestBaseRepository: + @pytest.fixture + def prisma_client(self): + return MockPrismaClient() + + def test_prisma_client_none_raises(self): + class TestRepo(BaseRepository[LiteLLM_BudgetTable]): + @property + def table(self): + return None + + @property + def model_class(self): + return LiteLLM_BudgetTable + + repo = TestRepo(None) + with pytest.raises(RuntimeError, match="No DB Connected"): + _ = repo.prisma_client + + @pytest.mark.asyncio + async def test_find_many(self, prisma_client): + repo = BudgetRepository(prisma_client) + prisma_client.db.litellm_budgettable._records = { + "b1": {"budget_id": "b1", "max_budget": 100.0}, + "b2": {"budget_id": "b2", "max_budget": 200.0}, + } + budgets = await repo.find_many() + assert len(budgets) == 2 + + @pytest.mark.asyncio + async def test_count(self, prisma_client): + repo = BudgetRepository(prisma_client) + prisma_client.db.litellm_budgettable._records = { + "b1": {"budget_id": "b1"}, + "b2": {"budget_id": "b2"}, + } + count = await repo.count() + assert count == 2 + + @pytest.mark.asyncio + async def test_exists(self, prisma_client): + repo = BudgetRepository(prisma_client) + prisma_client.db.litellm_budgettable._records = { + "b1": {"budget_id": "b1"}, + } + assert await repo.exists("b1", id_field="budget_id") + assert not await repo.exists("nonexistent", id_field="budget_id") + + @pytest.mark.asyncio + async def test_find_many_with_all_kwargs(self, prisma_client): + repo = BudgetRepository(prisma_client) + prisma_client.db.litellm_budgettable._records = { + "b1": {"budget_id": "b1", "max_budget": 100.0}, + } + budgets = await repo.find_many( + where={"budget_id": "b1"}, skip=0, take=10, order={"budget_id": "asc"} + ) + assert len(budgets) == 1 + + def test_record_to_dict_branches(self): + from litellm.repositories.base_repository import _record_to_dict + + assert _record_to_dict({"a": 1}) == {"a": 1} + + class WithModelDump: + def model_dump(self): + return {"src": "model_dump"} + + assert _record_to_dict(WithModelDump()) == {"src": "model_dump"} + + class WithDict: + def dict(self): + return {"src": "dict"} + + assert _record_to_dict(WithDict()) == {"src": "dict"} + + assert _record_to_dict([("k", "v")]) == {"k": "v"} + + +class TestBudgetRepository: + @pytest.fixture + def repo(self): + client = MockPrismaClient() + return BudgetRepository(client) + + @pytest.mark.asyncio + async def test_create_budget(self, repo): + budget = await repo.create_budget( + created_by="test-user", + max_budget=100.0, + soft_budget=80.0, + tpm_limit=1000, + ) + assert budget.max_budget == 100.0 + assert budget.soft_budget == 80.0 + assert budget.tpm_limit == 1000 + + @pytest.mark.asyncio + async def test_create_budget_all_fields(self, repo): + budget = await repo.create_budget( + created_by="test-user", + max_budget=100.0, + soft_budget=80.0, + max_parallel_requests=10, + tpm_limit=1000, + rpm_limit=100, + model_max_budget={"gpt-4": 50.0}, + budget_duration="monthly", + allowed_models=["gpt-4", "gpt-3.5-turbo"], + ) + assert budget.max_budget == 100.0 + assert budget.max_parallel_requests == 10 + + @pytest.mark.asyncio + async def test_update_budget(self, repo): + await repo.create_budget(created_by="test-user", max_budget=100.0) + repo._prisma_client.db.litellm_budgettable._records["budget-1"] = { + "budget_id": "budget-1", + "max_budget": 100.0, + } + + updated = await repo.update_budget( + budget_id="budget-1", + updated_by="test-user", + max_budget=200.0, + ) + assert updated.max_budget == 200.0 + + @pytest.mark.asyncio + async def test_delete_budget(self, repo): + repo._prisma_client.db.litellm_budgettable._records["budget-1"] = { + "budget_id": "budget-1", + "max_budget": 100.0, + } + deleted = await repo.delete_budget("budget-1") + assert deleted is not None + assert "budget-1" not in repo._prisma_client.db.litellm_budgettable._records + + @pytest.mark.asyncio + async def test_find_by_id(self, repo): + repo._prisma_client.db.litellm_budgettable._records["budget-1"] = { + "budget_id": "budget-1", + "max_budget": 100.0, + } + budget = await repo.find_by_id("budget-1") + assert budget is not None + assert budget.budget_id == "budget-1" + + +class TestModelRepository: + @pytest.fixture + def repo(self): + client = MockPrismaClient() + return ModelRepository(client) + + @pytest.mark.asyncio + @patch( + "litellm.repositories.model_repository.encrypt_value_helper", + side_effect=lambda v, **kw: f"encrypted_{v}", + ) + @patch( + "litellm.repositories.model_repository.decrypt_value_helper", + side_effect=lambda v, **kw: v, + ) + async def test_create_model_encrypts_params(self, mock_decrypt, mock_encrypt, repo): + model = await repo.create_model( + model_name="gpt-4", + litellm_params={"api_key": "sk-secret"}, + created_by="test-user", + ) + assert model is not None + mock_encrypt.assert_called() + + @pytest.mark.asyncio + @patch( + "litellm.repositories.model_repository.encrypt_value_helper", + side_effect=lambda v, **kw: f"encrypted_{v}", + ) + @patch( + "litellm.repositories.model_repository.decrypt_value_helper", + side_effect=lambda v, **kw: v, + ) + async def test_create_model_all_fields(self, mock_decrypt, mock_encrypt, repo): + model = await repo.create_model( + model_name="gpt-4-turbo", + litellm_params={ + "api_key": "sk-secret", + "api_base": "https://api.openai.com", + }, + created_by="admin", + model_id="custom-model-id", + model_info={"team_id": "team-1", "description": "GPT-4 Turbo model"}, + blocked=True, + ) + assert model is not None + assert model.model_name == "gpt-4-turbo" + + @pytest.mark.asyncio + @patch( + "litellm.repositories.model_repository.encrypt_value_helper", + side_effect=lambda v, **kw: f"encrypted_{v}", + ) + @patch( + "litellm.repositories.model_repository.decrypt_value_helper", + side_effect=lambda v, **kw: v, + ) + async def test_update_model_all_fields(self, mock_decrypt, mock_encrypt, repo): + repo._prisma_client.db.litellm_proxymodeltable._records["model-full"] = { + "model_id": "model-full", + "model_name": "old-name", + "litellm_params": '{"api_key": "old"}', + "blocked": False, + } + updated = await repo.update_model( + model_id="model-full", + updated_by="admin", + model_name="new-name", + litellm_params={"api_key": "new-key"}, + model_info={"updated": True}, + blocked=True, + ) + assert updated.model_name == "new-name" + + @pytest.mark.asyncio + @patch( + "litellm.repositories.model_repository.decrypt_value_helper", + side_effect=lambda v, **kw: v, + ) + async def test_find_all(self, mock_decrypt, repo): + repo._prisma_client.db.litellm_proxymodeltable._records = { + "m1": { + "model_id": "m1", + "model_name": "gpt-4", + "litellm_params": '{"model": "gpt-4"}', + "blocked": False, + }, + "m2": { + "model_id": "m2", + "model_name": "claude-3", + "litellm_params": '{"model": "claude-3"}', + "blocked": False, + }, + } + models = await repo.find_all() + assert len(models) == 2 + + @pytest.mark.asyncio + @patch( + "litellm.repositories.model_repository.decrypt_value_helper", + side_effect=lambda v, **kw: v, + ) + async def test_find_unblocked(self, mock_decrypt, repo): + repo._prisma_client.db.litellm_proxymodeltable._records = { + "m1": { + "model_id": "m1", + "model_name": "gpt-4", + "litellm_params": '{"model": "gpt-4"}', + "blocked": False, + }, + } + models = await repo.find_unblocked() + assert len(models) == 1 + + @pytest.mark.asyncio + @patch( + "litellm.repositories.model_repository.decrypt_value_helper", + side_effect=lambda v, **kw: v, + ) + async def test_find_by_name(self, mock_decrypt, repo): + repo._prisma_client.db.litellm_proxymodeltable._records = { + "m1": { + "model_id": "m1", + "model_name": "gpt-4", + "litellm_params": '{"model": "gpt-4"}', + }, + } + models = await repo.find_by_name("gpt-4") + assert len(models) == 1 + + @pytest.mark.asyncio + @patch( + "litellm.repositories.model_repository.encrypt_value_helper", + side_effect=lambda v, **kw: v, + ) + @patch( + "litellm.repositories.model_repository.decrypt_value_helper", + side_effect=lambda v, **kw: v, + ) + async def test_update_model(self, mock_decrypt, mock_encrypt, repo): + repo._prisma_client.db.litellm_proxymodeltable._records["m1"] = { + "model_id": "m1", + "model_name": "gpt-4", + "litellm_params": '{"model": "gpt-4"}', + "blocked": False, + } + updated = await repo.update_model( + model_id="m1", + updated_by="test-user", + blocked=True, + ) + assert updated.blocked is True + + @pytest.mark.asyncio + @patch( + "litellm.repositories.model_repository.decrypt_value_helper", + side_effect=lambda v, **kw: v, + ) + async def test_delete_model(self, mock_decrypt, repo): + repo._prisma_client.db.litellm_proxymodeltable._records["m1"] = { + "model_id": "m1", + "model_name": "gpt-4", + "litellm_params": '{"model": "gpt-4"}', + } + deleted = await repo.delete_model("m1") + assert deleted is not None + + @pytest.mark.asyncio + @patch( + "litellm.repositories.model_repository.encrypt_value_helper", + side_effect=lambda v, **kw: v, + ) + @patch( + "litellm.repositories.model_repository.decrypt_value_helper", + side_effect=lambda v, **kw: v, + ) + async def test_block_unblock_model(self, mock_decrypt, mock_encrypt, repo): + repo._prisma_client.db.litellm_proxymodeltable._records["m1"] = { + "model_id": "m1", + "model_name": "gpt-4", + "litellm_params": '{"model": "gpt-4"}', + "blocked": False, + } + blocked = await repo.block_model("m1", "admin") + assert blocked.blocked is True + + unblocked = await repo.unblock_model("m1", "admin") + assert unblocked.blocked is False + + +class TestTeamRepository: + @pytest.fixture + def repo(self): + client = MockPrismaClient() + return TeamRepository(client) + + @pytest.mark.asyncio + async def test_create_team(self, repo): + team = await repo.create_team( + team_id="team-123", + team_alias="Engineering", + admins=["user1"], + members=["user2", "user3"], + ) + assert team.team_id == "team-123" + assert team.team_alias == "Engineering" + + @pytest.mark.asyncio + async def test_create_team_all_fields(self, repo): + team = await repo.create_team( + team_id="team-123", + team_alias="Engineering", + organization_id="org-1", + admins=["admin1"], + members=["user1"], + members_with_roles=[{"user_id": "user1", "role": "user"}], + metadata={"dept": "engineering"}, + max_budget=1000.0, + soft_budget=800.0, + models=["gpt-4"], + max_parallel_requests=10, + tpm_limit=50000, + rpm_limit=500, + budget_duration="monthly", + object_permission_id="perm-1", + ) + assert team.team_id == "team-123" + assert team.organization_id == "org-1" + + @pytest.mark.asyncio + async def test_update_team(self, repo): + repo._prisma_client.db.litellm_teamtable._records["team-1"] = { + "team_id": "team-1", + "team_alias": "Test", + "admins": [], + "members": [], + "models": [], + } + updated = await repo.update_team( + team_id="team-1", + team_alias="Updated Team", + blocked=True, + ) + assert updated.team_alias == "Updated Team" + + @pytest.mark.asyncio + async def test_update_team_all_fields(self, repo): + repo._prisma_client.db.litellm_teamtable._records["team-full"] = { + "team_id": "team-full", + "team_alias": "Test", + "admins": [], + "members": [], + "models": [], + } + updated = await repo.update_team( + team_id="team-full", + team_alias="Fully Updated", + organization_id="org-new", + admins=["admin1"], + members=["member1"], + members_with_roles=[{"user_id": "user1", "role": "admin"}], + metadata={"updated": True}, + max_budget=500.0, + soft_budget=400.0, + models=["gpt-4", "claude-3"], + max_parallel_requests=20, + tpm_limit=100000, + rpm_limit=1000, + budget_duration="weekly", + blocked=False, + object_permission_id="perm-new", + ) + assert updated.team_alias == "Fully Updated" + + @pytest.mark.asyncio + async def test_add_member(self, repo): + repo._prisma_client.db.litellm_teamtable._records["team-1"] = { + "team_id": "team-1", + "team_alias": "Test", + "admins": [], + "members": ["user1"], + "models": [], + } + + team = await repo.add_member("team-1", "user2") + assert "user2" in team.members + + @pytest.mark.asyncio + async def test_add_member_nonexistent_team(self, repo): + result = await repo.add_member("nonexistent", "user1") + assert result is None + + @pytest.mark.asyncio + async def test_remove_member(self, repo): + repo._prisma_client.db.litellm_teamtable._records["team-1"] = { + "team_id": "team-1", + "team_alias": "Test", + "admins": [], + "members": ["user1", "user2"], + "models": [], + } + + team = await repo.remove_member("team-1", "user2") + assert "user2" not in team.members + + @pytest.mark.asyncio + async def test_add_admin(self, repo): + repo._prisma_client.db.litellm_teamtable._records["team-1"] = { + "team_id": "team-1", + "team_alias": "Test", + "admins": [], + "members": [], + "models": [], + } + team = await repo.add_admin("team-1", "admin1") + assert "admin1" in team.admins + + @pytest.mark.asyncio + async def test_remove_admin(self, repo): + repo._prisma_client.db.litellm_teamtable._records["team-1"] = { + "team_id": "team-1", + "team_alias": "Test", + "admins": ["admin1", "admin2"], + "members": [], + "models": [], + } + team = await repo.remove_admin("team-1", "admin2") + assert "admin2" not in team.admins + + @pytest.mark.asyncio + async def test_add_models(self, repo): + repo._prisma_client.db.litellm_teamtable._records["team-1"] = { + "team_id": "team-1", + "team_alias": "Test", + "admins": [], + "members": [], + "models": ["gpt-3.5-turbo"], + } + team = await repo.add_models("team-1", ["gpt-4"]) + assert "gpt-4" in team.models + + @pytest.mark.asyncio + async def test_remove_models(self, repo): + repo._prisma_client.db.litellm_teamtable._records["team-1"] = { + "team_id": "team-1", + "team_alias": "Test", + "admins": [], + "members": [], + "models": ["gpt-3.5-turbo", "gpt-4"], + } + team = await repo.remove_models("team-1", ["gpt-4"]) + assert "gpt-4" not in team.models + + @pytest.mark.asyncio + async def test_update_spend(self, repo): + repo._prisma_client.db.litellm_teamtable._records["team-1"] = { + "team_id": "team-1", + "team_alias": "Test", + "admins": [], + "members": [], + "models": [], + "spend": 0.0, + } + team = await repo.update_spend("team-1", 50.0) + assert team.spend == 50.0 + + @pytest.mark.asyncio + async def test_find_by_alias(self, repo): + repo._prisma_client.db.litellm_teamtable._records["team-1"] = { + "team_id": "team-1", + "team_alias": "Engineering", + "admins": [], + "members": [], + "models": [], + } + team = await repo.find_by_alias("Engineering") + assert team is not None + assert team.team_id == "team-1" + + @pytest.mark.asyncio + async def test_find_by_organization_id(self, repo): + repo._prisma_client.db.litellm_teamtable._records["team-1"] = { + "team_id": "team-1", + "organization_id": "org-1", + "admins": [], + "members": [], + "models": [], + } + teams = await repo.find_by_organization_id("org-1") + assert len(teams) == 1 + + @pytest.mark.asyncio + async def test_find_by_member(self, repo): + repo._prisma_client.db.litellm_teamtable._records["team-1"] = { + "team_id": "team-1", + "admins": [], + "members": ["user1"], + "models": [], + } + teams = await repo.find_by_member("user1") + assert len(teams) == 1 + + @pytest.mark.asyncio + async def test_find_by_admin(self, repo): + repo._prisma_client.db.litellm_teamtable._records["team-1"] = { + "team_id": "team-1", + "admins": ["admin1"], + "members": [], + "models": [], + } + teams = await repo.find_by_admin("admin1") + assert len(teams) == 1 + + +class TestUserRepository: + @pytest.fixture + def repo(self): + client = MockPrismaClient() + return UserRepository(client) + + @pytest.mark.asyncio + async def test_create_user(self, repo): + user = await repo.create_user( + user_id="user-123", + user_email="test@example.com", + teams=["team1"], + ) + assert user.user_id == "user-123" + + @pytest.mark.asyncio + async def test_create_user_all_fields(self, repo): + user = await repo.create_user( + user_id="user-123", + user_alias="testuser", + team_id="team-1", + sso_user_id="sso-123", + organization_id="org-1", + password="hashed_password", + teams=["team1", "team2"], + user_role="admin", + max_budget=500.0, + user_email="test@example.com", + models=["gpt-4"], + metadata={"department": "engineering"}, + max_parallel_requests=5, + tpm_limit=10000, + rpm_limit=100, + budget_duration="monthly", + allowed_cache_controls=["no-cache"], + policies=["policy-1"], + object_permission_id="perm-1", + ) + assert user.user_id == "user-123" + assert user.user_alias == "testuser" + + @pytest.mark.asyncio + async def test_update_user(self, repo): + repo._prisma_client.db.litellm_usertable._records["user-1"] = { + "user_id": "user-1", + "teams": [], + "models": [], + } + updated = await repo.update_user( + user_id="user-1", + user_email="updated@example.com", + ) + assert updated.user_email == "updated@example.com" + + @pytest.mark.asyncio + async def test_delete_user(self, repo): + repo._prisma_client.db.litellm_usertable._records["user-1"] = { + "user_id": "user-1", + "teams": [], + "models": [], + } + deleted = await repo.delete_user("user-1") + assert deleted is not None + + @pytest.mark.asyncio + async def test_add_to_team(self, repo): + repo._prisma_client.db.litellm_usertable._records["user-1"] = { + "user_id": "user-1", + "teams": ["team1"], + "models": [], + } + + user = await repo.add_to_team("user-1", "team2") + assert "team2" in user.teams + + @pytest.mark.asyncio + async def test_add_to_team_nonexistent_user(self, repo): + result = await repo.add_to_team("nonexistent", "team1") + assert result is None + + @pytest.mark.asyncio + async def test_remove_from_team(self, repo): + repo._prisma_client.db.litellm_usertable._records["user-1"] = { + "user_id": "user-1", + "teams": ["team1", "team2"], + "models": [], + } + user = await repo.remove_from_team("user-1", "team2") + assert "team2" not in user.teams + + @pytest.mark.asyncio + async def test_update_spend(self, repo): + repo._prisma_client.db.litellm_usertable._records["user-1"] = { + "user_id": "user-1", + "teams": [], + "models": [], + "spend": 0.0, + } + user = await repo.update_spend("user-1", 25.0) + assert user.spend == 25.0 + + @pytest.mark.asyncio + async def test_find_by_email(self, repo): + repo._prisma_client.db.litellm_usertable._records["user-1"] = { + "user_id": "user-1", + "user_email": "test@example.com", + "teams": [], + "models": [], + } + user = await repo.find_by_email("test@example.com") + assert user is not None + + @pytest.mark.asyncio + async def test_find_by_sso_id(self, repo): + repo._prisma_client.db.litellm_usertable._records["sso-123"] = { + "user_id": "user-1", + "sso_user_id": "sso-123", + "teams": [], + "models": [], + } + user = await repo.find_by_sso_id("sso-123") + assert user is not None + + @pytest.mark.asyncio + async def test_find_by_organization_id(self, repo): + repo._prisma_client.db.litellm_usertable._records["user-1"] = { + "user_id": "user-1", + "organization_id": "org-1", + "teams": [], + "models": [], + } + users = await repo.find_by_organization_id("org-1") + assert len(users) == 1 + + @pytest.mark.asyncio + async def test_find_by_team_id(self, repo): + repo._prisma_client.db.litellm_usertable._records["user-1"] = { + "user_id": "user-1", + "teams": ["team-1"], + "models": [], + } + users = await repo.find_by_team_id("team-1") + assert len(users) == 1 + + +class TestVerificationTokenRepository: + @pytest.fixture + def repo(self): + client = MockPrismaClient() + return VerificationTokenRepository(client) + + @pytest.mark.asyncio + async def test_create_token(self, repo): + token = await repo.create_token( + token="sk-test123", + key_name="Test Key", + user_id="user-123", + max_budget=100.0, + ) + assert token.token == "sk-test123" + assert token.key_name == "Test Key" + + @pytest.mark.asyncio + async def test_create_token_all_fields(self, repo): + token = await repo.create_token( + token="sk-test123", + key_name="Test Key", + key_alias="test-alias", + max_budget=100.0, + expires=datetime(2025, 12, 31), + models=["gpt-4"], + aliases={"alias1": "value1"}, + config={"setting": "value"}, + user_id="user-123", + team_id="team-1", + agent_id="agent-1", + project_id="project-1", + max_parallel_requests=5, + metadata={"key": "value"}, + tpm_limit=10000, + rpm_limit=100, + budget_duration="monthly", + allowed_cache_controls=["no-cache"], + allowed_routes=["/v1/completions"], + permissions={"read": True}, + org_id="org-1", + created_by="admin", + object_permission_id="perm-1", + access_group_ids=["group-1"], + budget_id="budget-1", + ) + assert token.token == "sk-test123" + + @pytest.mark.asyncio + async def test_update_token(self, repo): + repo._prisma_client.db.litellm_verificationtoken._records["sk-test"] = { + "token": "sk-test", + "blocked": False, + } + updated = await repo.update_token( + token="sk-test", + key_name="Updated Key", + ) + assert updated.key_name == "Updated Key" + + @pytest.mark.asyncio + async def test_block_token(self, repo): + repo._prisma_client.db.litellm_verificationtoken._records["sk-test"] = { + "token": "sk-test", + "blocked": False, + } + + token = await repo.block_token("sk-test", updated_by="admin") + assert token.blocked is True + + @pytest.mark.asyncio + async def test_unblock_token(self, repo): + repo._prisma_client.db.litellm_verificationtoken._records["sk-test"] = { + "token": "sk-test", + "blocked": True, + } + token = await repo.unblock_token("sk-test", updated_by="admin") + assert token.blocked is False + + @pytest.mark.asyncio + async def test_update_spend(self, repo): + repo._prisma_client.db.litellm_verificationtoken._records["sk-test"] = { + "token": "sk-test", + "spend": 0.0, + } + token = await repo.update_spend("sk-test", 15.0) + assert token.spend == 15.0 + + @pytest.mark.asyncio + async def test_update_last_active(self, repo): + repo._prisma_client.db.litellm_verificationtoken._records["sk-test"] = { + "token": "sk-test", + } + token = await repo.update_last_active("sk-test") + assert token.last_active is not None + + @pytest.mark.asyncio + async def test_find_by_alias(self, repo): + repo._prisma_client.db.litellm_verificationtoken._records["sk-test"] = { + "token": "sk-test", + "key_alias": "my-key", + } + token = await repo.find_by_alias("my-key") + assert token is not None + + @pytest.mark.asyncio + async def test_find_by_user_id(self, repo): + repo._prisma_client.db.litellm_verificationtoken._records["sk-test"] = { + "token": "sk-test", + "user_id": "user-1", + } + tokens = await repo.find_by_user_id("user-1") + assert len(tokens) == 1 + + @pytest.mark.asyncio + async def test_find_by_team_id(self, repo): + repo._prisma_client.db.litellm_verificationtoken._records["sk-test"] = { + "token": "sk-test", + "team_id": "team-1", + } + tokens = await repo.find_by_team_id("team-1") + assert len(tokens) == 1 + + @pytest.mark.asyncio + async def test_find_by_project_id(self, repo): + repo._prisma_client.db.litellm_verificationtoken._records["sk-test"] = { + "token": "sk-test", + "project_id": "project-1", + } + tokens = await repo.find_by_project_id("project-1") + assert len(tokens) == 1 + + +class TestOrganizationRepository: + @pytest.fixture + def repo(self): + client = MockPrismaClient() + return OrganizationRepository(client) + + @pytest.mark.asyncio + async def test_create_organization(self, repo): + org = await repo.create_organization( + organization_alias="Acme Corp", + budget_id="budget-1", + created_by="admin", + ) + assert org.organization_alias == "Acme Corp" + + @pytest.mark.asyncio + async def test_create_organization_all_fields(self, repo): + org = await repo.create_organization( + organization_alias="Acme Corp", + budget_id="budget-1", + created_by="admin", + organization_id="org-123", + metadata={"industry": "tech"}, + models=["gpt-4"], + object_permission_id="perm-1", + ) + assert org.organization_alias == "Acme Corp" + + @pytest.mark.asyncio + async def test_update_organization(self, repo): + repo._prisma_client.db.litellm_organizationtable._records["org-1"] = { + "organization_id": "org-1", + "organization_alias": "Old Name", + "budget_id": "b1", + "created_by": "admin", + "updated_by": "admin", + } + updated = await repo.update_organization( + organization_id="org-1", + updated_by="admin", + organization_alias="New Name", + ) + assert updated.organization_alias == "New Name" + + @pytest.mark.asyncio + async def test_update_organization_all_fields(self, repo): + repo._prisma_client.db.litellm_organizationtable._records["org-full"] = { + "organization_id": "org-full", + "organization_alias": "Old Name", + "budget_id": "b1", + "created_by": "admin", + "updated_by": "admin", + } + updated = await repo.update_organization( + organization_id="org-full", + updated_by="admin", + organization_alias="Fully Updated", + budget_id="budget-new", + metadata={"updated": True}, + models=["gpt-4", "claude-3"], + object_permission_id="perm-new", + ) + assert updated.organization_alias == "Fully Updated" + + @pytest.mark.asyncio + async def test_delete_organization(self, repo): + repo._prisma_client.db.litellm_organizationtable._records["org-1"] = { + "organization_id": "org-1", + "organization_alias": "Acme", + "budget_id": "b1", + "created_by": "admin", + "updated_by": "admin", + } + deleted = await repo.delete_organization("org-1") + assert deleted is not None + + @pytest.mark.asyncio + async def test_update_spend(self, repo): + repo._prisma_client.db.litellm_organizationtable._records["org-1"] = { + "organization_id": "org-1", + "organization_alias": "Acme", + "spend": 0.0, + "budget_id": "b1", + "created_by": "admin", + "updated_by": "admin", + } + org = await repo.update_spend("org-1", 100.0) + assert org.spend == 100.0 + + @pytest.mark.asyncio + async def test_find_by_alias(self, repo): + repo._prisma_client.db.litellm_organizationtable._records["org-1"] = { + "organization_id": "org-1", + "organization_alias": "Acme", + "budget_id": "b1", + "created_by": "admin", + "updated_by": "admin", + } + org = await repo.find_by_alias("Acme") + assert org is not None + + +class TestProjectRepository: + @pytest.fixture + def repo(self): + client = MockPrismaClient() + return ProjectRepository(client) + + @pytest.mark.asyncio + async def test_create_project(self, repo): + project = await repo.create_project( + created_by="admin", + project_alias="My Project", + ) + assert project.project_alias == "My Project" + + @pytest.mark.asyncio + async def test_create_project_all_fields(self, repo): + project = await repo.create_project( + created_by="admin", + project_id="proj-123", + project_alias="My Project", + description="A test project", + team_id="team-1", + budget_id="budget-1", + metadata={"env": "dev"}, + models=["gpt-4"], + model_rpm_limit={"gpt-4": 100}, + model_tpm_limit={"gpt-4": 10000}, + object_permission_id="perm-1", + ) + assert project.project_alias == "My Project" + + @pytest.mark.asyncio + async def test_update_project(self, repo): + repo._prisma_client.db.litellm_projecttable._records["proj-1"] = { + "project_id": "proj-1", + "project_alias": "Old Name", + } + updated = await repo.update_project( + project_id="proj-1", + updated_by="admin", + project_alias="New Name", + blocked=True, + ) + assert updated.project_alias == "New Name" + + @pytest.mark.asyncio + async def test_update_project_all_fields(self, repo): + repo._prisma_client.db.litellm_projecttable._records["proj-full"] = { + "project_id": "proj-full", + "project_alias": "Old Name", + } + updated = await repo.update_project( + project_id="proj-full", + updated_by="admin", + project_alias="Fully Updated", + description="New description", + team_id="team-new", + budget_id="budget-new", + metadata={"updated": True}, + models=["gpt-4", "claude-3"], + model_rpm_limit={"gpt-4": 200}, + model_tpm_limit={"gpt-4": 20000}, + blocked=False, + object_permission_id="perm-new", + ) + assert updated.project_alias == "Fully Updated" + + @pytest.mark.asyncio + async def test_delete_project(self, repo): + repo._prisma_client.db.litellm_projecttable._records["proj-1"] = { + "project_id": "proj-1", + } + deleted = await repo.delete_project("proj-1") + assert deleted is not None + + @pytest.mark.asyncio + async def test_update_spend(self, repo): + repo._prisma_client.db.litellm_projecttable._records["proj-1"] = { + "project_id": "proj-1", + "spend": 0.0, + } + project = await repo.update_spend("proj-1", 50.0) + assert project.spend == 50.0 + + @pytest.mark.asyncio + async def test_find_by_alias(self, repo): + repo._prisma_client.db.litellm_projecttable._records["proj-1"] = { + "project_id": "proj-1", + "project_alias": "MyProject", + } + project = await repo.find_by_alias("MyProject") + assert project is not None + + @pytest.mark.asyncio + async def test_find_by_team_id(self, repo): + repo._prisma_client.db.litellm_projecttable._records["proj-1"] = { + "project_id": "proj-1", + "team_id": "team-1", + } + projects = await repo.find_by_team_id("team-1") + assert len(projects) == 1 + + +class TestObjectPermissionRepository: + @pytest.fixture + def repo(self): + client = MockPrismaClient() + return ObjectPermissionRepository(client) + + @pytest.mark.asyncio + async def test_create_permission(self, repo): + perm = await repo.create_permission( + mcp_servers=["server1"], + models=["gpt-4"], + ) + assert perm.mcp_servers == ["server1"] + + @pytest.mark.asyncio + async def test_create_permission_all_fields(self, repo): + perm = await repo.create_permission( + mcp_servers=["server1"], + mcp_access_groups=["group1"], + mcp_tool_permissions={"tool1": ["read", "write"]}, + vector_stores=["store1"], + agents=["agent1"], + agent_access_groups=["agent-group1"], + models=["gpt-4"], + blocked_tools=["tool2"], + mcp_toolsets=["toolset1"], + search_tools=["search1"], + ) + assert perm.mcp_servers == ["server1"] + assert perm.agents == ["agent1"] + + @pytest.mark.asyncio + async def test_update_permission(self, repo): + repo._prisma_client.db.litellm_objectpermissiontable._records["perm-1"] = { + "object_permission_id": "perm-1", + "models": ["gpt-3.5-turbo"], + } + updated = await repo.update_permission( + object_permission_id="perm-1", + models=["gpt-4"], + ) + assert updated.models == ["gpt-4"] + + @pytest.mark.asyncio + async def test_update_permission_all_fields(self, repo): + repo._prisma_client.db.litellm_objectpermissiontable._records["perm-full"] = { + "object_permission_id": "perm-full", + "models": [], + } + updated = await repo.update_permission( + object_permission_id="perm-full", + mcp_servers=["server-new"], + mcp_access_groups=["group-new"], + mcp_tool_permissions={"tool": ["exec"]}, + vector_stores=["store-new"], + agents=["agent-new"], + agent_access_groups=["ag-new"], + models=["gpt-4", "claude-3"], + blocked_tools=["blocked-tool"], + mcp_toolsets=["toolset-new"], + search_tools=["search-new"], + ) + assert updated.mcp_servers == ["server-new"] + + @pytest.mark.asyncio + async def test_delete_permission(self, repo): + repo._prisma_client.db.litellm_objectpermissiontable._records["perm-1"] = { + "object_permission_id": "perm-1", + } + deleted = await repo.delete_permission("perm-1") + assert deleted is not None + + +class TestCredentialsRepository: + @pytest.fixture + def repo(self): + client = MockPrismaClient() + return CredentialsRepository(client) + + @pytest.mark.asyncio + async def test_create(self, repo): + record = await repo.create( + data={ + "credential_name": "my-api-key", + "credential_values": {"api_key": "encrypted_secret"}, + "credential_info": {"provider": "openai"}, + "created_by": "admin", + "updated_by": "admin", + } + ) + assert record.credential_name == "my-api-key" + cred = repo._to_model(record) + assert cred.credential_name == "my-api-key" + assert cred.credential_info == {"provider": "openai"} + assert cred.credential_values == {"api_key": "encrypted_secret"} + + @pytest.mark.asyncio + async def test_find_by_name_returns_stored_values_without_decryption(self, repo): + repo._prisma_client.db.litellm_credentialstable._records["my-key"] = { + "credential_id": "cred-1", + "credential_name": "my-key", + "credential_values": {"api_key": "encrypted_secret"}, + "credential_info": {"provider": "openai"}, + } + cred = await repo.find_by_name("my-key") + assert isinstance(cred, CredentialItem) + assert cred.credential_values == {"api_key": "encrypted_secret"} + assert cred.credential_info == {"provider": "openai"} + + @pytest.mark.asyncio + async def test_find_by_name_missing(self, repo): + assert await repo.find_by_name("nonexistent") is None + + @pytest.mark.asyncio + async def test_update_by_name(self, repo): + repo._prisma_client.db.litellm_credentialstable._records["my-key"] = { + "credential_id": "cred-1", + "credential_name": "my-key", + "credential_values": {"api_key": "old"}, + "credential_info": {}, + } + await repo.update_by_name( + "my-key", + data={"credential_values": {"api_key": "new"}, "updated_by": "admin"}, + ) + cred = await repo.find_by_name("my-key") + assert cred.credential_values == {"api_key": "new"} + + @pytest.mark.asyncio + async def test_delete_by_name(self, repo): + repo._prisma_client.db.litellm_credentialstable._records["my-key"] = { + "credential_id": "cred-1", + "credential_name": "my-key", + "credential_values": {"api_key": "secret"}, + "credential_info": {}, + } + await repo.delete_by_name("my-key") + assert await repo.find_by_name("my-key") is None + + @pytest.mark.asyncio + async def test_find_all(self, repo): + repo._prisma_client.db.litellm_credentialstable._records["k1"] = { + "credential_name": "k1", + "credential_values": {"api_key": "a"}, + "credential_info": {}, + } + repo._prisma_client.db.litellm_credentialstable._records["k2"] = { + "credential_name": "k2", + "credential_values": {"api_key": "b"}, + "credential_info": {}, + } + records = await repo.find_all() + assert len(records) == 2 + + def test_prisma_client_none_raises(self): + repo = CredentialsRepository(None) + with pytest.raises(RuntimeError, match="No DB Connected"): + _ = repo.table + + +class TestConfigRepository: + @pytest.fixture + def repo(self): + client = MockPrismaClient() + return ConfigRepository(client) + + def test_deep_merge_dicts_db_wins(self, repo): + dst = {"a": 1, "b": {"c": 2}} + src = {"a": 10, "b": {"d": 3}} + repo._deep_merge_dicts(dst, src) + assert dst["a"] == 10 + assert dst["b"]["c"] == 2 + assert dst["b"]["d"] == 3 + + def test_deep_merge_dicts_skips_none(self, repo): + dst = {"a": 1} + src = {"a": None, "b": 2} + repo._deep_merge_dicts(dst, src) + assert dst["a"] == 1 + assert dst["b"] == 2 + + def test_deep_merge_dicts_skips_empty_list(self, repo): + dst = {"models": ["gpt-4"]} + src = {"models": []} + repo._deep_merge_dicts(dst, src) + assert dst["models"] == ["gpt-4"] + + @pytest.mark.asyncio + async def test_get_param(self, repo): + repo._prisma_client.db.litellm_config._records["general_settings"] = { + "param_name": "general_settings", + "param_value": '{"master_key": "test"}', + } + param = await repo.get_param("general_settings") + assert param is not None + assert param.param_name == "general_settings" + assert param.param_value["master_key"] == "test" + + @pytest.mark.asyncio + async def test_set_param(self, repo): + param = await repo.set_param("test_param", {"key": "value"}) + assert param.param_name == "test_param" + assert param.param_value == {"key": "value"} + + @pytest.mark.asyncio + async def test_delete_param(self, repo): + repo._prisma_client.db.litellm_config._records["test_param"] = { + "param_name": "test_param", + "param_value": "{}", + } + result = await repo.delete_param("test_param") + assert result is True + + @pytest.mark.asyncio + async def test_delete_param_nonexistent(self, repo): + async def mock_delete(where): + raise Exception("Not found") + + repo._prisma_client.db.litellm_config.delete = mock_delete + result = await repo.delete_param("nonexistent") + assert result is False + + @pytest.mark.asyncio + async def test_get_all_params(self, repo): + repo._prisma_client.db.litellm_config._records = { + "param1": {"param_name": "param1", "param_value": '{"a": 1}'}, + "param2": {"param_name": "param2", "param_value": '{"b": 2}'}, + } + params = await repo.get_all_params() + assert len(params) == 2 + + @pytest.mark.asyncio + async def test_reconcile_config_skips_when_store_model_false(self, repo): + yaml_config = {"general_settings": {"key": "value"}} + result = await repo.reconcile_config(yaml_config, store_model_in_db=False) + assert result == yaml_config + + @pytest.mark.asyncio + async def test_prefetch_params(self, repo): + repo._prisma_client.db.litellm_config._records["general_settings"] = { + "param_name": "general_settings", + "param_value": "{}", + } + await repo.prefetch_params(["general_settings"]) + + @pytest.mark.asyncio + async def test_reconcile_config_with_db_values(self, repo): + repo._prisma_client.db.litellm_config._records["general_settings"] = { + "param_name": "general_settings", + "param_value": '{"master_key": "db-key", "db_only": "from_db"}', + } + repo._prisma_client.db.litellm_config._records["router_settings"] = { + "param_name": "router_settings", + "param_value": '{"timeout": 60}', + } + yaml_config = { + "general_settings": {"master_key": "yaml-key", "yaml_only": "from_yaml"}, + } + result = await repo.reconcile_config(yaml_config, store_model_in_db=True) + assert result["general_settings"]["master_key"] == "db-key" + assert result["general_settings"]["yaml_only"] == "from_yaml" + assert result["general_settings"]["db_only"] == "from_db" + assert result["router_settings"]["timeout"] == 60 + + @pytest.mark.asyncio + @patch("litellm.repositories.config_repository.decrypt_value_helper") + async def test_reconcile_config_with_environment_variables( + self, mock_decrypt, repo + ): + mock_decrypt.side_effect = lambda value, **kw: f"decrypted_{value}" + repo._prisma_client.db.litellm_config._records["environment_variables"] = { + "param_name": "environment_variables", + "param_value": '{"api_key": "encrypted_key", "secret": "encrypted_secret"}', + } + yaml_config = {} + result = await repo.reconcile_config(yaml_config, store_model_in_db=True) + assert "environment_variables" in result + assert "api_key" in result["environment_variables"] + assert "API_KEY" in result["environment_variables"] + + @pytest.mark.asyncio + async def test_reconcile_config_none_values_preserved(self, repo): + repo._prisma_client.db.litellm_config._records["general_settings"] = { + "param_name": "general_settings", + "param_value": '{"new_key": "value", "null_key": null}', + } + yaml_config = {"general_settings": {"existing": "keep"}} + result = await repo.reconcile_config(yaml_config, store_model_in_db=True) + assert result["general_settings"]["existing"] == "keep" + assert result["general_settings"]["new_key"] == "value" + + def test_update_config_fields_non_dict(self, repo): + config = {"litellm_settings": "old_value"} + result = repo._update_config_fields( + current_config=config, + param_name="litellm_settings", + db_param_value="new_value", + ) + assert result["litellm_settings"] == "new_value" + + def test_update_config_fields_new_param(self, repo): + config = {} + result = repo._update_config_fields( + current_config=config, + param_name="router_settings", + db_param_value={"timeout": 30}, + ) + assert result["router_settings"] == {"timeout": 30} + + @patch("litellm.repositories.config_repository.decrypt_value_helper") + def test_decrypt_env_variables_non_string(self, mock_decrypt, repo): + mock_decrypt.side_effect = lambda value, **kw: value + env_vars = {"string_val": "encrypted", "int_val": 123, "bool_val": True} + result = repo._decrypt_env_variables(env_vars) + assert result["int_val"] == "123" + assert result["bool_val"] == "True" + + @patch("litellm.repositories.config_repository.decrypt_value_helper") + def test_decrypt_env_variables_none_value(self, mock_decrypt, repo): + mock_decrypt.return_value = None + env_vars = {"key": "value"} + result = repo._decrypt_env_variables(env_vars) + assert "key" not in result + + +class TestVerificationTokenRepositoryExtended: + @pytest.fixture + def repo(self): + client = MockPrismaClient() + return VerificationTokenRepository(client) + + @pytest.mark.asyncio + async def test_find_active_tokens(self, repo): + repo._prisma_client.db.litellm_verificationtoken._records["sk-active"] = { + "token": "sk-active", + "blocked": False, + "expires": None, + } + tokens = await repo.find_active_tokens() + assert len(tokens) >= 1 + + @pytest.mark.asyncio + async def test_delete_token_with_audit(self, repo): + repo._prisma_client.db.litellm_verificationtoken._records["sk-delete"] = { + "token": "sk-delete", + "key_name": "Delete Me", + "spend": 0.0, + } + + class MockTx: + def __init__(self, client): + self.litellm_deletedverificationtoken = ( + client.db.litellm_deletedverificationtoken + ) + self.litellm_verificationtoken = client.db.litellm_verificationtoken + + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + pass + + repo._prisma_client.db.tx = lambda: MockTx(repo._prisma_client) + deleted = await repo.delete_token( + "sk-delete", + deleted_by="admin", + deleted_by_api_key="sk-admin", + litellm_changed_by="system", + ) + assert deleted is not None + assert deleted.token == "sk-delete" + + @pytest.mark.asyncio + async def test_delete_token_nonexistent(self, repo): + result = await repo.delete_token("nonexistent") + assert result is None + + @pytest.mark.asyncio + async def test_delete_token_archive_serialization(self, repo): + """Archived token must store JSON columns as strings, map org_id onto the + organization_id column, preserve budget_id, and drop relation-only fields + that don't exist on LiteLLM_DeletedVerificationToken.""" + repo._prisma_client.db.litellm_verificationtoken._records["sk-arch"] = { + "token": "sk-arch", + "key_name": "Archive Me", + "aliases": json.dumps({"a": "b"}), + "metadata": json.dumps({"team": "x"}), + "permissions": json.dumps({"read": True}), + "spend": 5.0, + "organization_id": "org-9", + "budget_id": "budget-9", + "budget_limits": [{"model": "gpt-4", "budget": 1.0}], + } + + class MockTx: + def __init__(self, client): + self.litellm_deletedverificationtoken = ( + client.db.litellm_deletedverificationtoken + ) + self.litellm_verificationtoken = client.db.litellm_verificationtoken + + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + pass + + repo._prisma_client.db.tx = lambda: MockTx(repo._prisma_client) + + await repo.delete_token("sk-arch", deleted_by="admin") + + archived = list( + repo._prisma_client.db.litellm_deletedverificationtoken._records.values() + )[0] + + assert isinstance(archived["aliases"], str) + assert json.loads(archived["aliases"]) == {"a": "b"} + assert isinstance(archived["metadata"], str) + assert isinstance(archived["permissions"], str) + + assert archived["organization_id"] == "org-9" + assert "org_id" not in archived + + assert archived["budget_id"] == "budget-9" + + for relation_field in ( + "object_permission", + "litellm_budget_table", + "budget_limits", + ): + assert relation_field not in archived + + assert ( + "sk-arch" not in repo._prisma_client.db.litellm_verificationtoken._records + ) + + @pytest.mark.asyncio + async def test_find_by_id_maps_org_and_budget_columns(self, repo): + """Reading a token must surface the organization_id column as org_id and + populate budget_id rather than silently dropping them.""" + repo._prisma_client.db.litellm_verificationtoken._records["sk-read"] = { + "token": "sk-read", + "organization_id": "org-7", + "budget_id": "budget-7", + } + token = await repo.find_by_id("sk-read") + assert token is not None + assert token.org_id == "org-7" + assert token.budget_id == "budget-7" + + @pytest.mark.asyncio + async def test_update_token_all_fields(self, repo): + repo._prisma_client.db.litellm_verificationtoken._records["sk-test"] = { + "token": "sk-test", + } + updated = await repo.update_token( + token="sk-test", + updated_by="admin", + key_name="Updated", + key_alias="new-alias", + max_budget=500.0, + expires=datetime(2025, 12, 31), + models=["gpt-4", "gpt-3.5-turbo"], + aliases={"a": "b"}, + config={"c": "d"}, + max_parallel_requests=10, + metadata={"m": "data"}, + tpm_limit=5000, + rpm_limit=50, + budget_duration="daily", + allowed_cache_controls=["cache"], + allowed_routes=["/v1/chat"], + permissions={"write": True}, + blocked=False, + object_permission_id="perm-2", + access_group_ids=["g1", "g2"], + ) + assert updated.key_name == "Updated" + + @pytest.mark.asyncio + async def test_to_model_with_json_fields(self, repo): + repo._prisma_client.db.litellm_verificationtoken._records["sk-json"] = { + "token": "sk-json", + "aliases": '{"alias1": "value1"}', + "config": '{"setting": "val"}', + "permissions": '{"read": true}', + "metadata": '{"key": "value"}', + "model_spend": '{"gpt-4": 10.0}', + "model_max_budget": '{"gpt-4": 100.0}', + "router_settings": '{"timeout": 30}', + "budget_limits": '[{"limit": 50}]', + "litellm_budget_table": '{"budget_id": "b1"}', + } + token = await repo.find_by_id("sk-json") + assert token is not None + assert token.aliases == {"alias1": "value1"} + assert token.config == {"setting": "val"} + + +class TestTeamRepositoryExtended: + @pytest.fixture + def repo(self): + client = MockPrismaClient() + return TeamRepository(client) + + @pytest.mark.asyncio + async def test_delete_team_with_audit(self, repo): + repo._prisma_client.db.litellm_teamtable._records["team-delete"] = { + "team_id": "team-delete", + "team_alias": "Delete Team", + "members": [], + "admins": [], + "models": [], + "spend": 0.0, + } + + class MockTx: + def __init__(self, client): + self.litellm_deletedteamtable = client.db.litellm_deletedteamtable + self.litellm_teamtable = client.db.litellm_teamtable + + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + pass + + repo._prisma_client.db.tx = lambda: MockTx(repo._prisma_client) + deleted = await repo.delete_team( + "team-delete", + deleted_by="admin", + deleted_by_api_key="sk-admin", + litellm_changed_by="system", + ) + assert deleted is not None + assert deleted.team_id == "team-delete" + + @pytest.mark.asyncio + async def test_delete_team_nonexistent(self, repo): + result = await repo.delete_team("nonexistent") + assert result is None + + @pytest.mark.asyncio + async def test_delete_team_with_full_data(self, repo): + repo._prisma_client.db.litellm_teamtable._records["team-full"] = { + "team_id": "team-full", + "team_alias": "Full Team", + "organization_id": "org-1", + "object_permission_id": "perm-1", + "members": ["m1", "m2"], + "admins": ["a1"], + "members_with_roles": '[{"user_id": "u1", "role": "admin"}]', + "metadata": '{"key": "value"}', + "max_budget": 1000.0, + "soft_budget": 800.0, + "spend": 150.0, + "models": ["gpt-4"], + "max_parallel_requests": 10, + "tpm_limit": 5000, + "rpm_limit": 50, + "budget_duration": "monthly", + "budget_reset_at": "2025-01-01T00:00:00", + "blocked": True, + "model_spend": '{"gpt-4": 100.0}', + "model_max_budget": '{"gpt-4": 500.0}', + "router_settings": '{"timeout": 30}', + "team_member_permissions": ["read"], + "access_group_ids": ["group-1"], + "policies": ["policy-1"], + "model_id": 42, + "allow_team_guardrail_config": True, + } + + class MockTx: + def __init__(self, client): + self.litellm_deletedteamtable = client.db.litellm_deletedteamtable + self.litellm_teamtable = client.db.litellm_teamtable + + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + pass + + repo._prisma_client.db.tx = lambda: MockTx(repo._prisma_client) + deleted = await repo.delete_team( + "team-full", + deleted_by="admin", + deleted_by_api_key="sk-admin", + litellm_changed_by="system", + ) + assert deleted is not None + assert deleted.team_id == "team-full" + assert deleted.organization_id == "org-1" + assert deleted.max_budget == 1000.0 + + @pytest.mark.asyncio + async def test_to_model_with_json_fields(self, repo): + repo._prisma_client.db.litellm_teamtable._records["team-json"] = { + "team_id": "team-json", + "metadata": '{"key": "value"}', + "model_spend": '{"gpt-4": 10.0}', + "model_max_budget": '{"gpt-4": 100.0}', + "router_settings": '{"timeout": 30}', + "budget_limits": '[{"budget_duration": "1d", "max_budget": 50.0}]', + "members_with_roles": '[{"user_id": "u1", "role": "admin"}]', + "members": [], + "admins": [], + "models": [], + } + team = await repo.find_by_id("team-json") + assert team is not None + assert team.metadata == {"key": "value"} + assert len(team.members_with_roles) == 1 + assert team.members_with_roles[0].user_id == "u1" + + +class TestUserRepositoryExtended: + @pytest.fixture + def repo(self): + client = MockPrismaClient() + return UserRepository(client) + + @pytest.mark.asyncio + async def test_delete_user_simple(self, repo): + repo._prisma_client.db.litellm_usertable._records["user-delete"] = { + "user_id": "user-delete", + "user_email": "delete@example.com", + "teams": [], + "models": [], + "spend": 0.0, + } + deleted = await repo.delete_user("user-delete") + assert deleted is not None + assert deleted.user_id == "user-delete" + + @pytest.mark.asyncio + async def test_delete_user_nonexistent(self, repo): + result = await repo.delete_user("nonexistent") + assert result is None + + @pytest.mark.asyncio + async def test_update_user_all_fields(self, repo): + repo._prisma_client.db.litellm_usertable._records["user-update"] = { + "user_id": "user-update", + "teams": [], + "models": [], + } + updated = await repo.update_user( + user_id="user-update", + user_alias="newalias", + team_id="team-new", + sso_user_id="sso-new", + organization_id="org-1", + password="new-hashed-pw", + teams=["team-1", "team-2"], + user_role="admin", + max_budget=1000.0, + user_email="new@example.com", + models=["gpt-4"], + metadata={"pref": "dark"}, + max_parallel_requests=20, + tpm_limit=10000, + rpm_limit=100, + budget_duration="monthly", + allowed_cache_controls=["no-cache"], + policies=["policy-1"], + object_permission_id="perm-new", + ) + assert updated.user_email == "new@example.com" + + +class TestProjectRepositoryExtended: + @pytest.fixture + def repo(self): + client = MockPrismaClient() + return ProjectRepository(client) + + @pytest.mark.asyncio + async def test_delete_project_simple(self, repo): + repo._prisma_client.db.litellm_projecttable._records["proj-delete"] = { + "project_id": "proj-delete", + "project_alias": "Delete Project", + "spend": 0.0, + } + deleted = await repo.delete_project("proj-delete") + assert deleted is not None + + +class TestBudgetRepositoryExtended: + @pytest.fixture + def repo(self): + client = MockPrismaClient() + return BudgetRepository(client) + + @pytest.mark.asyncio + async def test_update_budget_all_fields(self, repo): + repo._prisma_client.db.litellm_budgettable._records["budget-update"] = { + "budget_id": "budget-update", + "max_budget": 100.0, + } + updated = await repo.update_budget( + budget_id="budget-update", + updated_by="admin", + max_budget=500.0, + soft_budget=400.0, + max_parallel_requests=15, + tpm_limit=20000, + rpm_limit=200, + model_max_budget={"gpt-4": 200.0}, + budget_duration="weekly", + allowed_models=["gpt-4", "claude-3"], + ) + assert updated.max_budget == 500.0 + + +class TestModelRepositoryExtended: + @pytest.fixture + def repo(self): + client = MockPrismaClient() + return ModelRepository(client) + + @pytest.mark.asyncio + @patch( + "litellm.repositories.model_repository.decrypt_value_helper", + side_effect=lambda value, **kw: value, + ) + async def test_find_by_team_id(self, mock_decrypt, repo): + repo._prisma_client.db.litellm_proxymodeltable._records["model-1"] = { + "model_id": "model-1", + "model_name": "gpt-4", + "litellm_params": '{"api_key": "sk-test"}', + "model_info": '{"team_id": "team-1"}', + "blocked": False, + } + repo._prisma_client.db.litellm_proxymodeltable._records["model-2"] = { + "model_id": "model-2", + "model_name": "claude-3", + "litellm_params": '{"api_key": "sk-other"}', + "model_info": '{"team_id": "team-2"}', + "blocked": False, + } + models = await repo.find_by_team_id("team-1") + assert len(models) == 1 + assert models[0].model_name == "gpt-4" + + +class TestBaseRepositoryExtended: + @pytest.fixture + def repo(self): + client = MockPrismaClient() + return BudgetRepository(client) + + @pytest.mark.asyncio + async def test_find_many_with_pagination(self, repo): + repo._prisma_client.db.litellm_budgettable._records = { + "b1": {"budget_id": "b1", "max_budget": 100.0}, + "b2": {"budget_id": "b2", "max_budget": 200.0}, + "b3": {"budget_id": "b3", "max_budget": 300.0}, + } + budgets = await repo.find_many(skip=0, take=2, order={"budget_id": "asc"}) + assert len(budgets) >= 2 + + @pytest.mark.asyncio + async def test_find_many_with_where(self, repo): + repo._prisma_client.db.litellm_budgettable._records = { + "b1": {"budget_id": "b1", "max_budget": 100.0}, + } + budgets = await repo.find_many(where={"budget_id": "b1"}) + assert len(budgets) >= 1 + + @pytest.mark.asyncio + async def test_to_model_list_with_none(self, repo): + result = repo._to_model_list([None, None]) + assert result == [] + + +class _SampleDomainModel(DomainModel): + budget_id: Optional[str] = None + max_budget: Optional[float] = None + + +class TestDomainModelExtended: + def test_from_db_record_none_raises(self): + with pytest.raises(ValueError, match="Cannot create domain model from None"): + DomainModel.from_db_record(None) + + def test_from_db_record_dict(self): + model = _SampleDomainModel.from_db_record( + {"budget_id": "b1", "max_budget": 100.0} + ) + assert model.budget_id == "b1" + + def test_from_db_record_model_dump(self): + class MockRecordWithModelDump: + def model_dump(self): + return {"budget_id": "b2", "max_budget": 200.0} + + model = _SampleDomainModel.from_db_record(MockRecordWithModelDump()) + assert model.budget_id == "b2" + + def test_to_db_dict(self): + model = _SampleDomainModel(budget_id="b3", max_budget=300.0) + data = model.to_db_dict() + assert data["budget_id"] == "b3" + assert data["max_budget"] == 300.0 + + +class TestTeamRepositoryArchiveData: + @pytest.fixture + def repo(self): + client = MockPrismaClient() + return TeamRepository(client) + + def test_build_archive_data_minimal_fields(self, repo): + + team = LiteLLM_TeamTable(team_id="team-minimal") + archive_data = repo._build_archive_data(team) + assert archive_data["team_id"] == "team-minimal" + assert archive_data["admins"] == [] + assert archive_data["members"] == [] + assert archive_data["models"] == [] + assert archive_data["spend"] == 0.0 + assert archive_data["blocked"] is False + assert "team_alias" not in archive_data + assert "organization_id" not in archive_data + assert "object_permission_id" not in archive_data + assert "members_with_roles" not in archive_data + assert "metadata" not in archive_data + assert "max_budget" not in archive_data + assert "soft_budget" not in archive_data + assert "max_parallel_requests" not in archive_data + assert "tpm_limit" not in archive_data + assert "rpm_limit" not in archive_data + assert "budget_duration" not in archive_data + assert "budget_reset_at" not in archive_data + assert "model_spend" not in archive_data + assert "model_max_budget" not in archive_data + assert "router_settings" not in archive_data + assert "model_id" not in archive_data + + def test_build_archive_data_excludes_invalid_columns(self, repo): + + team = LiteLLM_TeamTable( + team_id="team-1", + team_alias="My Team", + admins=["admin1"], + members=["member1"], + models=["gpt-4"], + default_team_member_models=["gpt-3.5-turbo"], + ) + archive_data = repo._build_archive_data(team) + assert "default_team_member_models" not in archive_data + assert "budget_limits" not in archive_data + assert archive_data["team_id"] == "team-1" + assert archive_data["team_alias"] == "My Team" + assert archive_data["admins"] == ["admin1"] + assert archive_data["members"] == ["member1"] + assert archive_data["models"] == ["gpt-4"] + + def test_build_archive_data_with_all_valid_fields(self, repo): + from datetime import datetime + + from litellm.models.team import Member + + team = LiteLLM_TeamTable( + team_id="team-full", + team_alias="Full Team", + organization_id="org-1", + object_permission_id="perm-1", + admins=["admin1", "admin2"], + members=["m1", "m2"], + members_with_roles=[Member(user_id="u1", role="admin")], + metadata={"key": "value"}, + max_budget=1000.0, + soft_budget=800.0, + spend=150.0, + models=["gpt-4", "claude-3"], + max_parallel_requests=10, + tpm_limit=5000, + rpm_limit=50, + budget_duration="monthly", + budget_reset_at=datetime(2025, 1, 1), + blocked=True, + model_spend={"gpt-4": 100.0}, + model_max_budget={"gpt-4": 500.0}, + router_settings={"timeout": 30}, + team_member_permissions=["read"], + access_group_ids=["group-1"], + policies=["policy-1"], + model_id=42, + allow_team_guardrail_config=True, + ) + archive_data = repo._build_archive_data(team) + assert archive_data["team_id"] == "team-full" + assert archive_data["organization_id"] == "org-1" + assert archive_data["object_permission_id"] == "perm-1" + assert archive_data["max_budget"] == 1000.0 + assert archive_data["soft_budget"] == 800.0 + assert archive_data["spend"] == 150.0 + assert archive_data["blocked"] is True + assert archive_data["model_id"] == 42 + assert archive_data["allow_team_guardrail_config"] is True + assert "members_with_roles" in archive_data + assert "metadata" in archive_data + assert "model_spend" in archive_data + assert "model_max_budget" in archive_data + assert "router_settings" in archive_data + + +class TestConfigRepositoryDeepCopy: + @pytest.fixture + def repo(self): + client = MockPrismaClient() + return ConfigRepository(client) + + @pytest.mark.asyncio + async def test_reconcile_config_does_not_mutate_original(self, repo): + import copy + + repo._prisma_client.db.litellm_config._records["general_settings"] = { + "param_name": "general_settings", + "param_value": '{"db_key": "db_value", "nested": {"db_nested": "from_db"}}', + } + original_config = { + "general_settings": { + "yaml_key": "yaml_value", + "nested": {"yaml_nested": "from_yaml"}, + } + } + original_copy = copy.deepcopy(original_config) + result = await repo.reconcile_config(original_config, store_model_in_db=True) + assert original_config == original_copy + assert result["general_settings"]["db_key"] == "db_value" + assert result["general_settings"]["yaml_key"] == "yaml_value" + assert result["general_settings"]["nested"]["db_nested"] == "from_db" + assert result["general_settings"]["nested"]["yaml_nested"] == "from_yaml" + + @pytest.mark.asyncio + async def test_reconcile_config_repeated_calls_independent(self, repo): + repo._prisma_client.db.litellm_config._records["general_settings"] = { + "param_name": "general_settings", + "param_value": '{"db_key": "db_value"}', + } + yaml_config = {"general_settings": {"yaml_key": "yaml_value"}} + result1 = await repo.reconcile_config(yaml_config, store_model_in_db=True) + result1["general_settings"]["modified"] = "in_result1" + result2 = await repo.reconcile_config(yaml_config, store_model_in_db=True) + assert "modified" not in yaml_config.get("general_settings", {}) + assert "modified" not in result2.get("general_settings", {}) + + +class TestPrismaTableRepository: + def test_table_property_returns_named_delegate(self): + from litellm.repositories.table_repositories import ( + AgentsRepository, + PolicyRepository, + ) + + prisma_client = MagicMock() + agents = AgentsRepository(prisma_client) + policy = PolicyRepository(prisma_client) + + assert agents.table is prisma_client.db.litellm_agentstable + assert policy.table is prisma_client.db.litellm_policytable + assert agents.table is not policy.table + + def test_table_access_raises_without_db(self): + from litellm.repositories.table_repositories import SpendLogsRepository + + repo = SpendLogsRepository(None) + with pytest.raises(RuntimeError, match="No DB Connected"): + _ = repo.table + + def test_each_repository_binds_its_own_table_name(self): + import litellm.repositories.table_repositories as tr + + prisma_client = MagicMock() + repos = [ + obj + for name, obj in vars(tr).items() + if isinstance(obj, type) + and issubclass(obj, tr.PrismaTableRepository) + and obj is not tr.PrismaTableRepository + ] + assert len(repos) >= 40 + seen = set() + for repo_cls in repos: + name = repo_cls.table_name + assert name.startswith("litellm_") + assert name not in seen, f"duplicate table_name {name}" + seen.add(name) + assert repo_cls(prisma_client).table is getattr(prisma_client.db, name) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index f952bbfbbaf..403e00f9533 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -21364,8 +21364,7 @@ export interface components { }; /** * LiteLLM_DeletedTeamTable - * @description Recording of deleted teams for audit purposes. Mirrors LiteLLM_TeamTable - * plus metadata captured at deletion time. + * @description Audit record for deleted teams; mirrors the team plus deletion metadata. */ LiteLLM_DeletedTeamTable: { /** Access Group Ids */ @@ -21375,6 +21374,11 @@ export interface components { * @default [] */ admins: unknown[]; + /** + * Allow Team Guardrail Config + * @default false + */ + allow_team_guardrail_config: boolean | null; /** * Blocked * @default false @@ -21421,6 +21425,20 @@ export interface components { } | null; /** Model Id */ model_id?: number | null; + /** + * Model Max Budget + * @default {} + */ + model_max_budget: { + [key: string]: unknown; + } | null; + /** + * Model Spend + * @default {} + */ + model_spend: { + [key: string]: unknown; + } | null; /** * Models * @default [] @@ -21431,6 +21449,8 @@ export interface components { object_permission_id?: string | null; /** Organization Id */ organization_id?: string | null; + /** Policies */ + policies?: string[] | null; /** Router Settings */ router_settings?: { [key: string]: unknown; @@ -21454,8 +21474,7 @@ export interface components { }; /** * LiteLLM_DeletedVerificationToken - * @description Recording of deleted keys for audit purposes. Mirrors LiteLLM_VerificationToken - * plus metadata captured at deletion time. + * @description Audit record for deleted keys; mirrors the token plus deletion metadata. */ LiteLLM_DeletedVerificationToken: { /** Access Group Ids */ @@ -21488,6 +21507,8 @@ export interface components { blocked?: boolean | null; /** Budget Duration */ budget_duration?: string | null; + /** Budget Id */ + budget_id?: string | null; /** Budget Limits */ budget_limits?: { [key: string]: unknown; @@ -21866,9 +21887,9 @@ export interface components { /** Id */ id?: number | null; /** Model Aliases */ - model_aliases?: { + model_aliases?: string | { [key: string]: unknown; - } | string | null; + } | null; team?: components["schemas"]["LiteLLM_TeamTable"] | null; /** Updated By */ updated_by: string; @@ -21934,6 +21955,11 @@ export interface components { } | null; /** Mcp Toolsets */ mcp_toolsets?: string[] | null; + /** + * Models + * @default [] + */ + models: string[] | null; /** Object Permission Id */ object_permission_id: string; /** @@ -21949,7 +21975,7 @@ export interface components { }; /** * LiteLLM_OrganizationMembershipTable - * @description This is the table that track what organizations a user belongs to and users spend within the organization + * @description Tracks which organizations a user belongs to and their spend within it. */ LiteLLM_OrganizationMembershipTable: { /** Budget Id */ @@ -22005,7 +22031,17 @@ export interface components { metadata?: { [key: string]: unknown; } | null; - /** Models */ + /** + * Model Spend + * @default {} + */ + model_spend: { + [key: string]: unknown; + } | null; + /** + * Models + * @default [] + */ models: string[]; object_permission?: components["schemas"]["LiteLLM_ObjectPermissionTable"] | null; /** Object Permission Id */ @@ -22231,7 +22267,7 @@ export interface components { [key: string]: unknown; } | null; /** Stream Timeout */ - stream_timeout?: string | number | null; + stream_timeout?: number | string | null; /** Tag Regex */ tag_regex?: string[] | null; /** Tags */ @@ -22286,7 +22322,7 @@ export interface components { /** Created At */ created_at?: string | null; /** Created By */ - created_by: string; + created_by?: string | null; /** Description */ description?: string | null; litellm_budget_table?: components["schemas"]["LiteLLM_BudgetTable"] | null; @@ -22328,7 +22364,7 @@ export interface components { /** Updated At */ updated_at?: string | null; /** Updated By */ - updated_by: string; + updated_by?: string | null; }; /** LiteLLM_SpendLogs */ LiteLLM_SpendLogs: { @@ -22432,6 +22468,11 @@ export interface components { * @default [] */ admins: unknown[]; + /** + * Allow Team Guardrail Config + * @default false + */ + allow_team_guardrail_config: boolean | null; /** * Blocked * @default false @@ -22468,6 +22509,20 @@ export interface components { } | null; /** Model Id */ model_id?: number | null; + /** + * Model Max Budget + * @default {} + */ + model_max_budget: { + [key: string]: unknown; + } | null; + /** + * Model Spend + * @default {} + */ + model_spend: { + [key: string]: unknown; + } | null; /** * Models * @default [] @@ -22478,6 +22533,8 @@ export interface components { object_permission_id?: string | null; /** Organization Id */ organization_id?: string | null; + /** Policies */ + policies?: string[] | null; /** Router Settings */ router_settings?: { [key: string]: unknown; @@ -22549,6 +22606,11 @@ export interface components { }; /** LiteLLM_UserTable */ LiteLLM_UserTable: { + /** + * Allowed Cache Controls + * @default [] + */ + allowed_cache_controls: string[]; /** Budget Duration */ budget_duration?: string | null; /** Budget Reset At */ @@ -22557,6 +22619,8 @@ export interface components { created_at?: string | null; /** Max Budget */ max_budget?: number | null; + /** Max Parallel Requests */ + max_parallel_requests?: number | null; /** Metadata */ metadata?: { [key: string]: unknown; @@ -22581,8 +22645,17 @@ export interface components { */ models: unknown[]; object_permission?: components["schemas"]["LiteLLM_ObjectPermissionTable"] | null; + /** Object Permission Id */ + object_permission_id?: string | null; + /** Organization Id */ + organization_id?: string | null; /** Organization Memberships */ organization_memberships?: components["schemas"]["LiteLLM_OrganizationMembershipTable"][] | null; + /** + * Policies + * @default [] + */ + policies: string[]; /** Rpm Limit */ rpm_limit?: number | null; /** @@ -22592,6 +22665,8 @@ export interface components { spend: number; /** Sso User Id */ sso_user_id?: string | null; + /** Team Id */ + team_id?: string | null; /** * Teams * @default [] @@ -22612,6 +22687,11 @@ export interface components { }; /** LiteLLM_UserTableWithKeyCount */ LiteLLM_UserTableWithKeyCount: { + /** + * Allowed Cache Controls + * @default [] + */ + allowed_cache_controls: string[]; /** Budget Duration */ budget_duration?: string | null; /** Budget Reset At */ @@ -22625,6 +22705,8 @@ export interface components { key_count: number; /** Max Budget */ max_budget?: number | null; + /** Max Parallel Requests */ + max_parallel_requests?: number | null; /** Metadata */ metadata?: { [key: string]: unknown; @@ -22649,8 +22731,17 @@ export interface components { */ models: unknown[]; object_permission?: components["schemas"]["LiteLLM_ObjectPermissionTable"] | null; + /** Object Permission Id */ + object_permission_id?: string | null; + /** Organization Id */ + organization_id?: string | null; /** Organization Memberships */ organization_memberships?: components["schemas"]["LiteLLM_OrganizationMembershipTable"][] | null; + /** + * Policies + * @default [] + */ + policies: string[]; /** Rpm Limit */ rpm_limit?: number | null; /** @@ -22660,6 +22751,8 @@ export interface components { spend: number; /** Sso User Id */ sso_user_id?: string | null; + /** Team Id */ + team_id?: string | null; /** * Teams * @default [] @@ -22710,6 +22803,8 @@ export interface components { blocked?: boolean | null; /** Budget Duration */ budget_duration?: string | null; + /** Budget Id */ + budget_id?: string | null; /** Budget Limits */ budget_limits?: { [key: string]: unknown; @@ -24234,7 +24329,17 @@ export interface components { metadata?: { [key: string]: unknown; } | null; - /** Models */ + /** + * Model Spend + * @default {} + */ + model_spend: { + [key: string]: unknown; + } | null; + /** + * Models + * @default [] + */ models: string[]; object_permission?: components["schemas"]["LiteLLM_ObjectPermissionTable"] | null; /** Object Permission Id */ @@ -24339,7 +24444,7 @@ export interface components { */ created_at: string; /** Created By */ - created_by: string; + created_by?: string | null; /** Description */ description?: string | null; litellm_budget_table?: components["schemas"]["LiteLLM_BudgetTable"] | null; @@ -24384,7 +24489,7 @@ export interface components { */ updated_at: string; /** Updated By */ - updated_by: string; + updated_by?: string | null; }; /** NewTeamRequest */ NewTeamRequest: { @@ -27272,6 +27377,11 @@ export interface components { * @default [] */ admins: unknown[]; + /** + * Allow Team Guardrail Config + * @default false + */ + allow_team_guardrail_config: boolean | null; /** * Blocked * @default false @@ -27308,6 +27418,20 @@ export interface components { } | null; /** Model Id */ model_id?: number | null; + /** + * Model Max Budget + * @default {} + */ + model_max_budget: { + [key: string]: unknown; + } | null; + /** + * Model Spend + * @default {} + */ + model_spend: { + [key: string]: unknown; + } | null; /** * Models * @default [] @@ -27318,6 +27442,8 @@ export interface components { object_permission_id?: string | null; /** Organization Id */ organization_id?: string | null; + /** Policies */ + policies?: string[] | null; /** Router Settings */ router_settings?: { [key: string]: unknown; @@ -27361,6 +27487,11 @@ export interface components { * @default [] */ admins: unknown[]; + /** + * Allow Team Guardrail Config + * @default false + */ + allow_team_guardrail_config: boolean | null; /** * Blocked * @default false @@ -27407,6 +27538,20 @@ export interface components { } | null; /** Model Id */ model_id?: number | null; + /** + * Model Max Budget + * @default {} + */ + model_max_budget: { + [key: string]: unknown; + } | null; + /** + * Model Spend + * @default {} + */ + model_spend: { + [key: string]: unknown; + } | null; /** * Models * @default [] @@ -27417,6 +27562,8 @@ export interface components { object_permission_id?: string | null; /** Organization Id */ organization_id?: string | null; + /** Policies */ + policies?: string[] | null; /** Router Settings */ router_settings?: { [key: string]: unknown; @@ -28907,6 +29054,8 @@ export interface components { blocked?: boolean | null; /** Budget Duration */ budget_duration?: string | null; + /** Budget Id */ + budget_id?: string | null; /** Budget Limits */ budget_limits?: { [key: string]: unknown; @@ -29653,7 +29802,7 @@ export interface components { [key: string]: unknown; } | null; /** Stream Timeout */ - stream_timeout?: string | number | null; + stream_timeout?: number | string | null; /** Tag Regex */ tag_regex?: string[] | null; /** Tags */ From aaf1e2444b915335817fb7612183ee66b37b241c Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 6 Jun 2026 23:05:36 -0700 Subject: [PATCH 010/185] feat(ui): include internal routes in the dashboard's generated OpenAPI types (#29885) The dashboard calls UI-internal proxy routes that the public /openapi.json hides with include_in_schema=False, so they never reached schema.d.ts and could not be typed. The type generator now force-includes those routes when it dumps the spec for openapi-typescript; this mutates a throwaway interpreter only, so the spec the proxy actually serves is unchanged. Regenerates schema.d.ts so 86 internal route families (for example /v2/model/info, /global/spend/*, /config/*, /v2/login, /sso/*) are now typed, with no public route removed. This unblocks migrating the dashboard's data fetching onto the typed $api client. Branch CI note: schema.d.ts is generated; CI regenerates and diffs it via the same gen:api script. --- .../scripts/gen-api-types.mjs | 8 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 5534 ++++++++++++++++- 2 files changed, 5504 insertions(+), 38 deletions(-) diff --git a/ui/litellm-dashboard/scripts/gen-api-types.mjs b/ui/litellm-dashboard/scripts/gen-api-types.mjs index 66dfd2ee7b8..3c9373ec547 100644 --- a/ui/litellm-dashboard/scripts/gen-api-types.mjs +++ b/ui/litellm-dashboard/scripts/gen-api-types.mjs @@ -23,9 +23,17 @@ const specDir = mkdtempSync(join(tmpdir(), "litellm-openapi-")); const specPath = join(specDir, "openapi.json"); const python = (process.env.LITELLM_PYTHON ?? "python3").split(" "); +// The dashboard calls internal UI routes that the public /openapi.json hides via +// include_in_schema=False. Force them in so they get typed here; this mutates a +// throwaway interpreter, so the spec the proxy actually serves is unchanged. const dumpSpec = [ "import json, sys", "from litellm.proxy.proxy_server import app", + "from fastapi.routing import APIRoute", + "for route in app.routes:", + " if isinstance(route, APIRoute):", + " route.include_in_schema = True", + "app.openapi_schema = None", "with open(sys.argv[1], 'w') as f: json.dump(app.openapi(), f, sort_keys=True)", ].join("\n"); diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 403e00f9533..c8ec9d3c727 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -436,6 +436,26 @@ export interface paths { patch?: never; trace?: never; }; + "/alerting/settings": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Alerting Settings + * @description Return the configurable alerting param, description, and current value + */ + get: operations["alerting_settings_alerting_settings_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/anthropic/{endpoint}": { parameters: { query?: never; @@ -1770,6 +1790,26 @@ export interface paths { patch?: never; trace?: never; }; + "/config/callback/delete": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Delete Callback + * @description Delete specific logging callback from configuration. + */ + post: operations["delete_callback_config_callback_delete_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/config/cost_discount_config": { parameters: { query?: never; @@ -1851,6 +1891,83 @@ export interface paths { patch: operations["update_cost_margin_config_config_cost_margin_config_patch"]; trace?: never; }; + "/config/field/delete": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Delete Config General Settings + * @description Delete the db value of this field in litellm general settings. Resets it to it's initial default value on litellm. + */ + post: operations["delete_config_general_settings_config_field_delete_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/config/field/info": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** Get Config General Settings */ + get: operations["get_config_general_settings_config_field_info_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/config/field/update": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Update Config General Settings + * @description Update a specific field in litellm general settings + */ + post: operations["update_config_general_settings_config_field_update_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/config/list": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Get Config List + * @description List the available fields + current values for a given type of setting (currently just 'general_settings'user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),) + */ + get: operations["get_config_list_config_list_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/config/pass_through_endpoint": { parameters: { query?: never; @@ -1925,6 +2042,64 @@ export interface paths { patch?: never; trace?: never; }; + "/config/update": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Update Config + * @description For Admin UI - allows admin to update config via UI. + * + * Writes only the sections present in the request body to LiteLLM_Config rows + * (one row per top-level section). Sections the caller did not send are left + * untouched — this endpoint never persists pre-existing YAML values to DB as + * a side effect of an unrelated update. + */ + post: operations["update_config_config_update_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/config/yaml": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Config Yaml Endpoint + * @description This is a mock endpoint, to show what you can set in config.yaml details in the Swagger UI. + * + * Parameters: + * + * The config.yaml object has the following attributes: + * - **model_list**: *Optional[List[ModelParams]]* - A list of supported models on the server, along with model-specific configurations. ModelParams includes "model_name" (name of the model), "litellm_params" (litellm-specific parameters for the model), and "model_info" (additional info about the model such as id, mode, cost per token, etc). + * + * - **litellm_settings**: *Optional[dict]*: Settings for the litellm module. You can specify multiple properties like "drop_params", "set_verbose", "api_base", "cache". + * + * - **general_settings**: *Optional[ConfigGeneralSettings]*: General settings for the server like "completion_model" (default model for chat completion calls), "use_azure_key_vault" (option to load keys from azure key vault), "master_key" (key required for all calls to proxy), and others. + * + * Please, refer to each class's description for a better understanding of the specific attributes within them. + * + * Note: This is a mock endpoint primarily meant for demonstration purposes, and does not actually provide or change any configurations. + */ + get: operations["config_yaml_endpoint_config_yaml_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/config_overrides/hashicorp_vault": { parameters: { query?: never; @@ -2764,6 +2939,120 @@ export interface paths { patch?: never; trace?: never; }; + "/debug/memory/details": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Get Memory Details + * @description Get detailed memory diagnostics for deep debugging. + * + * Returns: + * - worker_pid: Process ID + * - process_memory: RAM usage, virtual memory, file handles, threads + * - garbage_collector: GC thresholds, counts, collection history + * - objects: Total tracked objects and top object types + * - uncollectable: Objects that can't be garbage collected (potential leaks) + * - cache_memory: Memory usage of user_api_key, router, and logging caches + * - router_memory: Memory usage of router components (model_list, deployment_names, etc.) + * + * Query Parameters: + * - top_n: Number of top object types to return (default: 20) + * - include_process_info: Include process-level memory info using psutil (default: true) + * + * Example usage: + * curl "http://localhost:4000/debug/memory/details?top_n=30" -H "Authorization: Bearer sk-1234" + * + * All memory sizes are reported in both bytes and MB. + */ + get: operations["get_memory_details_debug_memory_details_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/debug/memory/gc/configure": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Configure Gc Thresholds Endpoint + * @description Configure Python garbage collection thresholds. + * + * Lower thresholds mean more frequent GC cycles (less memory, more CPU overhead). + * Higher thresholds mean less frequent GC cycles (more memory, less CPU overhead). + * + * Returns: + * - message: Confirmation message + * - previous_thresholds: Old threshold values + * - new_thresholds: New threshold values + * - objects_awaiting_collection: Current object count in gen-0 + * - tip: Hint about when next collection will occur + * + * Query Parameters: + * - generation_0: Number of allocations before gen-0 collection (default: 700) + * - generation_1: Number of gen-0 collections before gen-1 collection (default: 10) + * - generation_2: Number of gen-1 collections before gen-2 collection (default: 10) + * + * Example for more aggressive collection: + * curl -X POST "http://localhost:4000/debug/memory/gc/configure?generation_0=500" -H "Authorization: Bearer sk-1234" + * + * Example for less aggressive collection: + * curl -X POST "http://localhost:4000/debug/memory/gc/configure?generation_0=1000" -H "Authorization: Bearer sk-1234" + * + * Monitor memory usage with GET /debug/memory/summary after changes. + */ + post: operations["configure_gc_thresholds_endpoint_debug_memory_gc_configure_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/debug/memory/summary": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Get Memory Summary + * @description Get simplified memory usage summary for the proxy. + * + * Returns: + * - worker_pid: Process ID + * - status: Overall health based on memory usage + * - memory: Process memory usage and RAM info + * - caches: Cache item counts and descriptions + * - garbage_collector: GC status and pending object counts + * + * Example usage: + * curl http://localhost:4000/debug/memory/summary -H "Authorization: Bearer sk-1234" + * + * For detailed analysis, call GET /debug/memory/details + * For cache management, use the cache management endpoints + */ + get: operations["get_memory_summary_debug_memory_summary_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/delete/allowed_ip": { parameters: { query?: never; @@ -2855,6 +3144,312 @@ export interface paths { patch?: never; trace?: never; }; + "/end_user/block": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Block User + * @description [BETA] Reject calls with this end-user id + * + * Parameters: + * - user_ids (List[str], required): The unique `user_id`s for the users to block + * + * (any /chat/completion call with this user={end-user-id} param, will be rejected.) + * + * ``` + * curl -X POST "http://0.0.0.0:8000/user/block" + * -H "Authorization: Bearer sk-1234" + * -d '{ + * "user_ids": [, ...] + * }' + * ``` + */ + post: operations["block_user_end_user_block_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/end_user/daily/activity": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Get Customer Daily Activity + * @description Get daily activity for specific organizations or all accessible organizations. + */ + get: operations["get_customer_daily_activity_end_user_daily_activity_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/end_user/delete": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Delete End User + * @description Delete multiple end-users. + * + * Parameters: + * - user_ids (List[str], required): The unique `user_id`s for the users to delete + * + * Example curl: + * ``` + * curl --location 'http://0.0.0.0:4000/customer/delete' --header 'Authorization: Bearer sk-1234' --header 'Content-Type: application/json' --data '{ + * "user_ids" :["ishaan-jaff-5"] + * }' + * + * See below for all params + * ``` + */ + post: operations["delete_end_user_end_user_delete_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/end_user/info": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * End User Info + * @description Get information about an end-user. An `end_user` is a customer (external user) of the proxy. + * + * Parameters: + * - end_user_id (str, required): The unique identifier for the end-user + * + * Example curl: + * ``` + * curl -X GET 'http://localhost:4000/customer/info?end_user_id=test-litellm-user-4' -H 'Authorization: Bearer sk-1234' + * ``` + */ + get: operations["end_user_info_end_user_info_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/end_user/list": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * List End User + * @description [Admin-only] List all available customers + * + * Example curl: + * ``` + * curl --location --request GET 'http://0.0.0.0:4000/customer/list' --header 'Authorization: Bearer sk-1234' + * ``` + */ + get: operations["list_end_user_end_user_list_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/end_user/new": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * New End User + * @description Allow creating a new Customer + * + * + * Parameters: + * - user_id: str - The unique identifier for the user. + * - alias: Optional[str] - A human-friendly alias for the user. + * - blocked: bool - Flag to allow or disallow requests for this end-user. Default is False. + * - max_budget: Optional[float] - The maximum budget allocated to the user. Either 'max_budget' or 'budget_id' should be provided, not both. + * - budget_id: Optional[str] - The identifier for an existing budget allocated to the user. Either 'max_budget' or 'budget_id' should be provided, not both. + * - allowed_model_region: Optional[Union[Literal["eu"], Literal["us"]]] - Require all user requests to use models in this specific region. + * - default_model: Optional[str] - If no equivalent model in the allowed region, default all requests to this model. + * - metadata: Optional[dict] = Metadata for customer, store information for customer. Example metadata = {"data_training_opt_out": True} + * - budget_duration: Optional[str] - Budget is reset at the end of specified duration. If not set, budget is never reset. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"). + * - tpm_limit: Optional[int] - [Not Implemented Yet] Specify tpm limit for a given customer (Tokens per minute) + * - rpm_limit: Optional[int] - [Not Implemented Yet] Specify rpm limit for a given customer (Requests per minute) + * - model_max_budget: Optional[dict] - [Not Implemented Yet] Specify max budget for a given model. Example: {"openai/gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d"}} + * - max_parallel_requests: Optional[int] - [Not Implemented Yet] Specify max parallel requests for a given customer. + * - soft_budget: Optional[float] - [Not Implemented Yet] Get alerts when customer crosses given budget, doesn't block requests. + * - spend: Optional[float] - Specify initial spend for a given customer. + * - budget_reset_at: Optional[str] - Specify the date and time when the budget should be reset. + * - object_permission: Optional[LiteLLM_ObjectPermissionBase] - Customer-specific object permissions to control access to resources. + * Supported fields: + * * mcp_servers: List[str] - List of allowed MCP server IDs + * * mcp_access_groups: List[str] - List of MCP access group names + * * mcp_tool_permissions: Dict[str, List[str]] - Map of server ID to allowed tool names (e.g., {"server_1": ["tool_a", "tool_b"]}) + * * vector_stores: List[str] - List of allowed vector store IDs + * * agents: List[str] - List of allowed agent IDs + * * agent_access_groups: List[str] - List of agent access group names + * Example: {"mcp_servers": ["server_1", "server_2"], "vector_stores": ["vector_store_1"], "agents": ["agent_1"]} + * IF null or {} then no object-level restrictions apply. + * + * + * - Allow specifying allowed regions + * - Allow specifying default model + * + * Example curl: + * ``` + * curl --location 'http://0.0.0.0:4000/customer/new' --header 'Authorization: Bearer sk-1234' --header 'Content-Type: application/json' --data '{ + * "user_id" : "ishaan-jaff-3", + * "allowed_region": "eu", + * "budget_id": "free_tier", + * "default_model": "azure/gpt-3.5-turbo-eu" + * }' + * + * # With object permissions + * curl -L -X POST 'http://localhost:4000/customer/new' -H 'Authorization: Bearer sk-1234' -H 'Content-Type: application/json' -d '{ + * "user_id": "user_1", + * "object_permission": { + * "mcp_servers": ["server_1"], + * "mcp_access_groups": ["public_group"], + * "vector_stores": ["vector_store_1"] + * } + * }' + * + * # return end-user object + * ``` + * + * NOTE: This used to be called `/end_user/new`, we will still be maintaining compatibility for /end_user/XXX for these endpoints + */ + post: operations["new_end_user_end_user_new_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/end_user/unblock": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Unblock User + * @description [BETA] Unblock calls with this user id + * + * Example + * ``` + * curl -X POST "http://0.0.0.0:8000/user/unblock" + * -H "Authorization: Bearer sk-1234" + * -d '{ + * "user_ids": [, ...] + * }' + * ``` + */ + post: operations["unblock_user_end_user_unblock_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/end_user/update": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Update End User + * @description Example curl + * + * Parameters: + * - user_id: str + * - alias: Optional[str] = None # human-friendly alias + * - blocked: bool = False # allow/disallow requests for this end-user + * - max_budget: Optional[float] = None + * - budget_id: Optional[str] = None # give either a budget_id or max_budget + * - allowed_model_region: Optional[AllowedModelRegion] = ( + * None # require all user requests to use models in this specific region + * ) + * - default_model: Optional[str] = ( + * None # if no equivalent model in allowed region - default all requests to this model + * ) + * - object_permission: Optional[LiteLLM_ObjectPermissionBase] - Customer-specific object permissions to control access to resources. + * Supported fields: + * * mcp_servers: List[str] - List of allowed MCP server IDs + * * mcp_access_groups: List[str] - List of MCP access group names + * * mcp_tool_permissions: Dict[str, List[str]] - Map of server ID to allowed tool names + * * vector_stores: List[str] - List of allowed vector store IDs + * * agents: List[str] - List of allowed agent IDs + * * agent_access_groups: List[str] - List of agent access group names + * Example: {"mcp_servers": ["server_1"], "vector_stores": ["vector_store_1"]} + * IF null or {} then no object-level restrictions apply. + * + * Example curl: + * ``` + * curl --location 'http://0.0.0.0:4000/customer/update' --header 'Authorization: Bearer sk-1234' --header 'Content-Type: application/json' --data '{ + * "user_id": "test-litellm-user-4", + * "budget_id": "paid_tier" + * }' + * + * # Updating object permissions + * curl -L -X POST 'http://localhost:4000/customer/update' --header 'Authorization: Bearer sk-1234' --header 'Content-Type: application/json' --data '{ + * "user_id": "user_1", + * "object_permission": { + * "mcp_servers": ["server_3"], + * "vector_stores": ["vector_store_2", "vector_store_3"] + * } + * }' + * + * See below for all params + * ``` + */ + post: operations["update_end_user_end_user_update_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/engines/{model}/chat/completions": { parameters: { query?: never; @@ -3010,6 +3605,28 @@ export interface paths { patch?: never; trace?: never; }; + "/fallback/login": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Fallback Login + * @description Create Proxy API Keys using Google Workspace SSO. Requires setting PROXY_BASE_URL in .env + * PROXY_BASE_URL should be the your deployed proxy endpoint, e.g. PROXY_BASE_URL="https://litellm-production-7002.up.railway.app/" + * Example: + */ + get: operations["fallback_login_fallback_login_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/fallback/{model}": { parameters: { query?: never; @@ -3296,6 +3913,44 @@ export interface paths { patch: operations["gemini_proxy_route_gemini__endpoint__patch"]; trace?: never; }; + "/get/allowed_ips": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** Get Allowed Ips */ + get: operations["get_allowed_ips_get_allowed_ips_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/get/config/callbacks": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Get Config + * @description For Admin UI - allows admin to view config via UI + * # return the callbacks and the env variables for the callback + */ + get: operations["get_config_get_config_callbacks_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/get/default_team_settings": { parameters: { query?: never; @@ -3425,6 +4080,508 @@ export interface paths { patch?: never; trace?: never; }; + "/get_favicon": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Get Favicon + * @description Get custom favicon for the admin UI. + */ + get: operations["get_favicon_get_favicon_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/get_image": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Get Image + * @description Get logo to show on admin UI + */ + get: operations["get_image_get_image_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/get_logo_url": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Get Logo Url + * @description Get the current logo URL from environment. + * + * Only HTTP(S) URLs are returned — those are intended to be loaded + * directly by the browser from a public/internal CDN. Local file + * paths set via ``UI_LOGO_PATH`` are NOT returned: they are admin- + * only filesystem details, the dashboard falls back to ``/get_image`` + * which serves the file only when it is a supported image. Without + * this filter, the unauthenticated endpoint would disclose internal + * hostnames or filesystem paths to any caller. + */ + get: operations["get_logo_url_get_logo_url_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/global/activity": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Get Global Activity + * @description Get number of API Requests, total tokens through proxy + * + * { + * "daily_data": [ + * const chartdata = [ + * { + * date: 'Jan 22', + * api_requests: 10, + * total_tokens: 2000 + * }, + * { + * date: 'Jan 23', + * api_requests: 10, + * total_tokens: 12 + * }, + * ], + * "sum_api_requests": 20, + * "sum_total_tokens": 2012 + * } + */ + get: operations["get_global_activity_global_activity_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/global/activity/cache_hits": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Get Global Activity + * @description Get number of cache hits, vs misses + * + * { + * "daily_data": [ + * const chartdata = [ + * { + * date: 'Jan 22', + * cache_hits: 10, + * llm_api_calls: 2000 + * }, + * { + * date: 'Jan 23', + * cache_hits: 10, + * llm_api_calls: 12 + * }, + * ], + * "sum_cache_hits": 20, + * "sum_llm_api_calls": 2012 + * } + */ + get: operations["get_global_activity_global_activity_cache_hits_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/global/activity/exceptions": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Get Global Activity Exceptions + * @description Get number of API Requests, total tokens through proxy + * + * { + * "daily_data": [ + * const chartdata = [ + * { + * date: 'Jan 22', + * num_rate_limit_exceptions: 10, + * }, + * { + * date: 'Jan 23', + * num_rate_limit_exceptions: 10, + * }, + * ], + * "sum_api_exceptions": 20, + * } + */ + get: operations["get_global_activity_exceptions_global_activity_exceptions_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/global/activity/exceptions/deployment": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Get Global Activity Exceptions Per Deployment + * @description Get number of 429 errors - Grouped by deployment + * + * [ + * { + * "deployment": "https://azure-us-east-1.openai.azure.com/", + * "daily_data": [ + * const chartdata = [ + * { + * date: 'Jan 22', + * num_rate_limit_exceptions: 10 + * }, + * { + * date: 'Jan 23', + * num_rate_limit_exceptions: 12 + * }, + * ], + * "sum_num_rate_limit_exceptions": 20, + * + * }, + * { + * "deployment": "https://azure-us-east-1.openai.azure.com/", + * "daily_data": [ + * const chartdata = [ + * { + * date: 'Jan 22', + * num_rate_limit_exceptions: 10, + * }, + * { + * date: 'Jan 23', + * num_rate_limit_exceptions: 12 + * }, + * ], + * "sum_num_rate_limit_exceptions": 20, + * + * }, + * ] + */ + get: operations["get_global_activity_exceptions_per_deployment_global_activity_exceptions_deployment_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/global/activity/model": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Get Global Activity Model + * @description Get number of API Requests, total tokens through proxy - Grouped by MODEL + * + * [ + * { + * "model": "gpt-4", + * "daily_data": [ + * const chartdata = [ + * { + * date: 'Jan 22', + * api_requests: 10, + * total_tokens: 2000 + * }, + * { + * date: 'Jan 23', + * api_requests: 10, + * total_tokens: 12 + * }, + * ], + * "sum_api_requests": 20, + * "sum_total_tokens": 2012 + * + * }, + * { + * "model": "azure/gpt-4-turbo", + * "daily_data": [ + * const chartdata = [ + * { + * date: 'Jan 22', + * api_requests: 10, + * total_tokens: 2000 + * }, + * { + * date: 'Jan 23', + * api_requests: 10, + * total_tokens: 12 + * }, + * ], + * "sum_api_requests": 20, + * "sum_total_tokens": 2012 + * + * }, + * ] + */ + get: operations["get_global_activity_model_global_activity_model_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/global/all_end_users": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Global View All End Users + * @description [BETA] This is a beta endpoint. It will change. + * + * Use this to just get all the unique `end_users` + */ + get: operations["global_view_all_end_users_global_all_end_users_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/global/spend": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Global Spend + * @description [BETA] This is a beta endpoint. It will change. + * + * View total spend across all proxy keys + */ + get: operations["global_spend_global_spend_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/global/spend/all_tag_names": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** Global Get All Tag Names */ + get: operations["global_get_all_tag_names_global_spend_all_tag_names_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/global/spend/end_users": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Global Spend End Users + * @description [BETA] This is a beta endpoint. It will change. + * + * Use this to get the top 'n' keys with the highest spend, ordered by spend. + */ + post: operations["global_spend_end_users_global_spend_end_users_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/global/spend/keys": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Global Spend Keys + * @description [BETA] This is a beta endpoint. It will change. + * + * Use this to get the top 'n' keys with the highest spend, ordered by spend. + */ + get: operations["global_spend_keys_global_spend_keys_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/global/spend/logs": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Global Spend Logs + * @description [BETA] This is a beta endpoint. It will change. + * + * Use this to get global spend (spend per day for last 30d). Admin-only endpoint + * + * More efficient implementation of /spend/logs, by creating a view over the spend logs table. + */ + get: operations["global_spend_logs_global_spend_logs_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/global/spend/models": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Global Spend Models + * @description [BETA] This is a beta endpoint. It will change. + * + * Use this to get the top 'n' models with the highest spend, ordered by spend. + */ + get: operations["global_spend_models_global_spend_models_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/global/spend/provider": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Get Global Spend Provider + * @description Get breakdown of spend per provider + * [ + * { + * "provider": "Azure OpenAI", + * "spend": 20 + * }, + * { + * "provider": "OpenAI", + * "spend": 10 + * }, + * { + * "provider": "VertexAI", + * "spend": 30 + * } + * ] + */ + get: operations["get_global_spend_provider_global_spend_provider_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/global/spend/refresh": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Global Spend Refresh + * @description ADMIN ONLY / MASTER KEY Only Endpoint + * + * Globally refresh spend MonthlyGlobalSpend view + */ + post: operations["global_spend_refresh_global_spend_refresh_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/global/spend/report": { parameters: { query?: never; @@ -3527,6 +4684,28 @@ export interface paths { patch?: never; trace?: never; }; + "/global/spend/teams": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Global Spend Per Team + * @description [BETA] This is a beta endpoint. It will change. + * + * Use this to get daily spend, grouped by `team_id` and `date` + */ + get: operations["global_spend_per_team_global_spend_teams_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/guardrails": { parameters: { query?: never; @@ -4860,6 +6039,111 @@ export interface paths { patch?: never; trace?: never; }; + "/invitation/delete": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Invitation Delete + * @description Delete invitation link + * + * ``` + * curl -X POST 'http://localhost:4000/invitation/delete' -H 'Content-Type: application/json' -d '{ + * "invitation_id": "1234" // 👈 id of invitation in 'LiteLLM_InvitationTable' + * }' + * ``` + */ + post: operations["invitation_delete_invitation_delete_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/invitation/info": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Invitation Info + * @description Allow admin to create invite links, to onboard new users to Admin UI. + * + * ``` + * curl -X POST 'http://localhost:4000/invitation/new' -H 'Content-Type: application/json' -d '{ + * "user_id": "1234" // 👈 id of user in 'LiteLLM_UserTable' + * }' + * ``` + */ + get: operations["invitation_info_invitation_info_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/invitation/new": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * New Invitation + * @description Allow admin to create invite links, to onboard new users to Admin UI. + * + * ``` + * curl -X POST 'http://localhost:4000/invitation/new' -H 'Content-Type: application/json' -d '{ + * "user_id": "1234" // 👈 id of user in 'LiteLLM_UserTable' + * }' + * ``` + */ + post: operations["new_invitation_invitation_new_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/invitation/update": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Invitation Update + * @description Update when invitation is accepted + * + * ``` + * curl -X POST 'http://localhost:4000/invitation/update' -H 'Content-Type: application/json' -d '{ + * "invitation_id": "1234" // 👈 id of invitation in 'LiteLLM_InvitationTable' + * "is_accepted": True // when invitation is accepted + * }' + * ``` + */ + post: operations["invitation_update_invitation_update_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/jwt/key/mapping/delete": { parameters: { query?: never; @@ -5700,6 +6984,23 @@ export interface paths { patch: operations["langfuse_proxy_route_langfuse__endpoint__patch"]; trace?: never; }; + "/lazy/warm/{name}": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** Warm */ + post: operations["warm_lazy_warm__name__post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/litellm/.well-known/litellm-ui-config": { parameters: { query?: never; @@ -5717,6 +7018,23 @@ export interface paths { patch?: never; trace?: never; }; + "/login": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** Login */ + post: operations["login_login_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/mcp-rest/test/connection": { parameters: { query?: never; @@ -5814,6 +7132,52 @@ export interface paths { patch?: never; trace?: never; }; + "/memory-usage-in-mem-cache": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Memory Usage In Mem Cache + * @description 1. user_api_key_cache + * 2. router_cache + * 3. proxy_logging_cache + * 4. internal_usage_cache + */ + get: operations["memory_usage_in_mem_cache_memory_usage_in_mem_cache_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/memory-usage-in-mem-cache-items": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Memory Usage In Mem Cache Items + * @description 1. user_api_key_cache + * 2. router_cache + * 3. proxy_logging_cache + * 4. internal_usage_cache + */ + get: operations["memory_usage_in_mem_cache_items_memory_usage_in_mem_cache_items_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/milvus/{endpoint}": { parameters: { query?: never; @@ -5886,6 +7250,35 @@ export interface paths { patch: operations["mistral_proxy_route_mistral__endpoint__patch"]; trace?: never; }; + "/model/cost_map/source": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Get Model Cost Map Source + * @description ADMIN ONLY / MASTER KEY Only Endpoint + * + * Returns information about where the current model cost/pricing data was loaded from. + * + * Response fields: + * - source: "local" (bundled backup) or "remote" (fetched from URL) + * - url: the remote URL that was attempted (null when env-forced local) + * - is_env_forced: true if LITELLM_LOCAL_MODEL_COST_MAP=True forced local usage + * - fallback_reason: human-readable reason why remote failed (null on success) + * - model_count: number of models in the currently loaded cost map + */ + get: operations["get_model_cost_map_source_model_cost_map_source_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/model/delete": { parameters: { query?: never; @@ -5955,6 +7348,66 @@ export interface paths { patch?: never; trace?: never; }; + "/model/metrics": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Model Metrics + * @description View number of requests & avg latency per model on config.yaml + */ + get: operations["model_metrics_model_metrics_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/model/metrics/exceptions": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Model Metrics Exceptions + * @description View number of failed requests per model on config.yaml + */ + get: operations["model_metrics_exceptions_model_metrics_exceptions_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/model/metrics/slow_responses": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Model Metrics Slow Responses + * @description View number of hanging requests per model_group + */ + get: operations["model_metrics_slow_responses_model_metrics_slow_responses_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/model/new": { parameters: { query?: never; @@ -5975,6 +7428,46 @@ export interface paths { patch?: never; trace?: never; }; + "/model/settings": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Model Settings + * @description Returns provider name, description, and required parameters for each provider + */ + get: operations["model_settings_model_settings_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/model/streaming_metrics": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Model Streaming Metrics + * @description View time to first token for models in spend logs + */ + get: operations["model_streaming_metrics_model_streaming_metrics_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/model/update": { parameters: { query?: never; @@ -6411,6 +7904,58 @@ export interface paths { patch?: never; trace?: never; }; + "/onboarding/claim_token": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Claim Onboarding Link + * @description Special route. Allows UI link share user to update their password. + * + * - Get the invite link + * - Validate it's still 'valid' + * - Check if user within initial session (prevents abuse) + * - Get user from db + * - Update user password + * + * This route can only update user password. + */ + post: operations["claim_onboarding_link_onboarding_claim_token_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/onboarding/get_token": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Onboarding + * @description - Get the invite link + * - Validate it's still 'valid' + * - Return a short-lived onboarding token + * - Get user from db + * - Pass in user_email if set + */ + get: operations["onboarding_onboarding_get_token_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/openai/deployments/{model}/chat/completions": { parameters: { query?: never; @@ -7339,6 +8884,23 @@ export interface paths { patch: operations["update_organization_organization_update_patch"]; trace?: never; }; + "/otel-spans": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** Get Otel Spans */ + get: operations["get_otel_spans_otel_spans_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/policies": { parameters: { query?: never; @@ -9012,6 +10574,23 @@ export interface paths { patch?: never; trace?: never; }; + "/queue/chat/completions": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** Async Queue Request */ + post: operations["async_queue_request_queue_chat_completions_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/rag/ingest": { parameters: { query?: never; @@ -9178,6 +10757,52 @@ export interface paths { patch?: never; trace?: never; }; + "/reload/anthropic_beta_headers": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Reload Anthropic Beta Headers + * @description ADMIN ONLY / MASTER KEY Only Endpoint + * + * Manually reload the Anthropic beta headers configuration from the remote source. + * This will fetch fresh configuration from the anthropic_beta_headers_config.json file. + */ + post: operations["reload_anthropic_beta_headers_reload_anthropic_beta_headers_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/reload/model_cost_map": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Reload Model Cost Map + * @description ADMIN ONLY / MASTER KEY Only Endpoint + * + * Manually reload the model cost map from the remote source. + * This will fetch fresh pricing data from the model_prices_and_context_window.json file. + */ + post: operations["reload_model_cost_map_reload_model_cost_map_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/rerank": { parameters: { query?: never; @@ -9467,6 +11092,108 @@ export interface paths { patch?: never; trace?: never; }; + "/schedule/anthropic_beta_headers_reload": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Schedule Anthropic Beta Headers Reload + * @description ADMIN ONLY / MASTER KEY Only Endpoint + * + * Schedule periodic reload of the Anthropic beta headers configuration. + * This will create a background job that reloads the configuration every specified hours. + */ + post: operations["schedule_anthropic_beta_headers_reload_schedule_anthropic_beta_headers_reload_post"]; + /** + * Cancel Anthropic Beta Headers Reload + * @description ADMIN ONLY / MASTER KEY Only Endpoint + * + * Cancel the scheduled periodic reload of the Anthropic beta headers configuration. + */ + delete: operations["cancel_anthropic_beta_headers_reload_schedule_anthropic_beta_headers_reload_delete"]; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/schedule/anthropic_beta_headers_reload/status": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Get Anthropic Beta Headers Reload Status + * @description ADMIN ONLY / MASTER KEY Only Endpoint + * + * Get the status of the scheduled Anthropic beta headers reload job. + */ + get: operations["get_anthropic_beta_headers_reload_status_schedule_anthropic_beta_headers_reload_status_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/schedule/model_cost_map_reload": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Schedule Model Cost Map Reload + * @description ADMIN ONLY / MASTER KEY Only Endpoint + * + * Schedule periodic reload of the model cost map. + * This will create a background job that reloads the model cost map every specified hours. + */ + post: operations["schedule_model_cost_map_reload_schedule_model_cost_map_reload_post"]; + /** + * Cancel Model Cost Map Reload + * @description ADMIN ONLY / MASTER KEY Only Endpoint + * + * Cancel the scheduled periodic reload of the model cost map. + */ + delete: operations["cancel_model_cost_map_reload_schedule_model_cost_map_reload_delete"]; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/schedule/model_cost_map_reload/status": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Get Model Cost Map Reload Status + * @description ADMIN ONLY / MASTER KEY Only Endpoint + * + * Get the status of the scheduled model cost map reload job. + */ + get: operations["get_model_cost_map_reload_status_schedule_model_cost_map_reload_status_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/scim/v2": { parameters: { query?: never; @@ -10322,6 +12049,38 @@ export interface paths { patch?: never; trace?: never; }; + "/spend/keys": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Spend Key Fn + * @description View keys created, ordered by spend. + * + * - Admin callers (PROXY_ADMIN / PROXY_ADMIN_VIEW_ONLY) see every key in + * the database. + * - All other callers (INTERNAL_USER / INTERNAL_USER_VIEW_ONLY, etc.) are + * scoped to keys they own (``user_id == caller``). A caller with no + * ``user_id`` has no scope and receives an empty list rather than the + * full table. + * + * Example Request: + * ``` + * curl -X GET "http://0.0.0.0:8000/spend/keys" -H "Authorization: Bearer sk-1234" + * ``` + */ + get: operations["spend_key_fn_spend_keys_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/spend/logs": { parameters: { query?: never; @@ -10374,6 +12133,86 @@ export interface paths { patch?: never; trace?: never; }; + "/spend/logs/session/ui": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Ui View Session Spend Logs + * @description Get paginated spend logs for a particular session. + * + * Returns: + * { + * "data": List[LiteLLM_SpendLogs], + * "total": int, + * "page": int, + * "page_size": int, + * "total_pages": int, + * } + */ + get: operations["ui_view_session_spend_logs_spend_logs_session_ui_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/spend/logs/ui": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Ui View Spend Logs + * @description View spend logs with pagination support. + * Available at both `/spend/logs/v2` (public API) and `/spend/logs/ui` (internal UI). + * + * Returns paginated response with data, total, page, page_size, and total_pages. + * + * Example: + * ``` + * curl -X GET "http://0.0.0.0:8000/spend/logs/v2?start_date=2025-11-25%2000:00:00&end_date=2025-11-26%2023:59:59&page=1&page_size=50" -H "Authorization: Bearer sk-1234" + * ``` + */ + get: operations["ui_view_spend_logs_spend_logs_ui_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/spend/logs/ui/{request_id}": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Ui View Request Response For Request Id + * @description View request / response for a specific request_id + * + * - goes through all callbacks, checks if any of them have a @property -> has_request_response_payload + * - if so, it will return the request and response payload + */ + get: operations["ui_view_request_response_for_request_id_spend_logs_ui__request_id__get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/spend/logs/v2": { parameters: { query?: never; @@ -10432,6 +12271,208 @@ export interface paths { patch?: never; trace?: never; }; + "/spend/users": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Spend User Fn + * @description View users created, ordered by spend. + * + * - Admin callers (PROXY_ADMIN / PROXY_ADMIN_VIEW_ONLY) see every user, or + * a specific user when ``user_id`` is supplied. + * - All other callers may only read their own row. If they supply a + * ``user_id`` query parameter that does not match their authenticated + * ``user_id`` the request is rejected with HTTP 403; supplying their + * own id (or none at all) returns just their row. A caller with no + * ``user_id`` on their key has no scope and receives an empty list + * rather than the full table. + * + * Example Request: + * ``` + * curl -X GET "http://0.0.0.0:8000/spend/users" -H "Authorization: Bearer sk-1234" + * ``` + * + * View User Table row for user_id + * ``` + * curl -X GET "http://0.0.0.0:8000/spend/users?user_id=1234" -H "Authorization: Bearer sk-1234" + * ``` + */ + get: operations["spend_user_fn_spend_users_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/sso/callback": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Auth Callback + * @description Verify login + */ + get: operations["auth_callback_sso_callback_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/sso/cli/complete/{login_id}": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** Cli Sso Complete */ + post: operations["cli_sso_complete_sso_cli_complete__login_id__post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/sso/cli/poll/{key_id}": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Cli Poll Key + * @description CLI polling endpoint - retrieves session from cache and generates JWT. + * + * Flow: + * 1. First poll (no team_id): Returns teams list without generating JWT + * 2. Second poll (with team_id): Generates JWT with selected team and deletes session + * + * Args: + * key_id: The CLI login session ID + * team_id: Optional team ID to assign to the JWT. If provided, must be one of user's teams. + */ + get: operations["cli_poll_key_sso_cli_poll__key_id__get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/sso/cli/start": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** Cli Sso Start */ + post: operations["cli_sso_start_sso_cli_start_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/sso/debug/callback": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Debug Sso Callback + * @description Returns the OpenID object returned by the SSO provider + */ + get: operations["debug_sso_callback_sso_debug_callback_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/sso/debug/login": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Debug Sso Login + * @description Create Proxy API Keys using Google Workspace SSO. Requires setting PROXY_BASE_URL in .env + * PROXY_BASE_URL should be the your deployed proxy endpoint, e.g. PROXY_BASE_URL="https://litellm-production-7002.up.railway.app/" + * Example: + */ + get: operations["debug_sso_login_sso_debug_login_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/sso/get/ui_settings": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** Get Ui Settings */ + get: operations["get_ui_settings_sso_get_ui_settings_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/sso/key/generate": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Google Login + * @description Create Proxy API Keys using Google Workspace SSO. Requires setting PROXY_BASE_URL in .env + * PROXY_BASE_URL should be the your deployed proxy endpoint, e.g. PROXY_BASE_URL="https://litellm-production-7002.up.railway.app/" + * Example: + */ + get: operations["google_login_sso_key_generate_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/sso/readiness": { parameters: { query?: never; @@ -10970,6 +13011,36 @@ export interface paths { patch?: never; trace?: never; }; + "/team/filter/ui": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Ui View Teams + * @description [PROXY-ADMIN ONLY] Filter teams based on partial match of team_id or team_alias with pagination. + * + * Args: + * user_id (Optional[str]): Partial user ID to search for + * user_email (Optional[str]): Partial email to search for + * page (int): Page number for pagination (starts at 1) + * page_size (int): Number of items per page (max 100) + * user_api_key_dict (UserAPIKeyAuth): User authentication information + * + * Returns: + * List[LiteLLM_SpendLogs]: Paginated list of matching user records + */ + get: operations["ui_view_teams_team_filter_ui_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/team/info": { parameters: { query?: never; @@ -11959,6 +14030,32 @@ export interface paths { patch?: never; trace?: never; }; + "/user/available_roles": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Ui Get Available Role + * @description Endpoint used by Admin UI to show all available roles to assign a user + * return { + * "proxy_admin": { + * "description": "Proxy Admin role", + * "ui_label": "Admin" + * } + * } + */ + get: operations["ui_get_available_role_user_available_roles_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/user/available_users": { parameters: { query?: never; @@ -12128,6 +14225,36 @@ export interface paths { patch?: never; trace?: never; }; + "/user/filter/ui": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Ui View Users + * @description Filter users based on partial match of user_id or email with pagination. + * + * Behaviour depends on the ``scope_user_search_to_org`` UI-setting flag + * (stored in the ``litellm_uisettings`` table): + * + * * **Flag OFF (default):** any authenticated user can search all users. + * * **Flag ON:** + * - Proxy admins see all users. + * - Org admins see only users in their org(s). + * - Team admins for an org-bound team see users in that org. + * - Others receive a 403. + */ + get: operations["ui_view_users_user_filter_ui_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/user/info": { parameters: { query?: never; @@ -16256,6 +18383,75 @@ export interface paths { patch?: never; trace?: never; }; + "/v2/key/info": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Info Key Fn V2 + * @description Retrieve information about a list of keys. + * + * **New endpoint**. Currently admin only. + * Parameters: + * keys: Optional[list] = body parameter representing the key(s) in the request + * user_api_key_dict: UserAPIKeyAuth = Dependency representing the user's API key + * Returns: + * Dict containing the key and its associated information + * + * Example Curl: + * ``` + * curl -X GET "http://0.0.0.0:4000/key/info" -H "Authorization: Bearer sk-1234" -d {"keys": ["sk-1", "sk-2", "sk-3"]} + * ``` + */ + post: operations["info_key_fn_v2_v2_key_info_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/v2/login": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** Login V2 */ + post: operations["login_v2_v2_login_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/v2/model/info": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Model Info V2 + * @description v2 - returns models available to the user based on their API key permissions. Shows model info from config.yaml (except api key and api base). Filter to just user-added models with ?user_models_only=true + */ + get: operations["model_info_v2_v2_model_info_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/v2/rerank": { parameters: { query?: never; @@ -16348,6 +18544,40 @@ export interface paths { patch?: never; trace?: never; }; + "/v3/login": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** Login V3 */ + post: operations["login_v3_v3_login_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/v3/login/exchange": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** Login V3 Exchange */ + post: operations["login_v3_exchange_v3_login_exchange_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/vantage/delete": { parameters: { query?: never; @@ -16752,6 +18982,52 @@ export interface paths { patch?: never; trace?: never; }; + "/vertex-ai/{endpoint}": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Vertex Proxy Route + * @description Call LiteLLM proxy via Vertex AI SDK. + * + * [Docs](https://docs.litellm.ai/docs/pass_through/vertex_ai) + */ + get: operations["vertex_proxy_route_vertex_ai__endpoint__get_2"]; + /** + * Vertex Proxy Route + * @description Call LiteLLM proxy via Vertex AI SDK. + * + * [Docs](https://docs.litellm.ai/docs/pass_through/vertex_ai) + */ + put: operations["vertex_proxy_route_vertex_ai__endpoint__put_2"]; + /** + * Vertex Proxy Route + * @description Call LiteLLM proxy via Vertex AI SDK. + * + * [Docs](https://docs.litellm.ai/docs/pass_through/vertex_ai) + */ + post: operations["vertex_proxy_route_vertex_ai__endpoint__post_2"]; + /** + * Vertex Proxy Route + * @description Call LiteLLM proxy via Vertex AI SDK. + * + * [Docs](https://docs.litellm.ai/docs/pass_through/vertex_ai) + */ + delete: operations["vertex_proxy_route_vertex_ai__endpoint__delete_2"]; + options?: never; + head?: never; + /** + * Vertex Proxy Route + * @description Call LiteLLM proxy via Vertex AI SDK. + * + * [Docs](https://docs.litellm.ai/docs/pass_through/vertex_ai) + */ + patch: operations["vertex_proxy_route_vertex_ai__endpoint__patch_2"]; + trace?: never; + }; "/vertex_ai/discovery/{endpoint}": { parameters: { query?: never; @@ -17919,6 +20195,12 @@ export interface components { /** Tags */ tags?: string[]; }; + /** + * AlertType + * @description Enum for alert types and management event types + * @enum {string} + */ + AlertType: "llm_exceptions" | "llm_too_slow" | "llm_requests_hanging" | "budget_alerts" | "spend_reports" | "failed_tracking_spend" | "db_exceptions" | "daily_reports" | "cooldown_deployment" | "new_model_added" | "outage_alerts" | "region_outage_alerts" | "fallback_reports" | "new_virtual_key_created" | "virtual_key_updated" | "virtual_key_deleted" | "new_team_created" | "team_updated" | "team_deleted" | "new_internal_user_created" | "internal_user_updated" | "internal_user_deleted"; /** AllowedVectorStoreIndexItem */ AllowedVectorStoreIndexItem: { /** Index Name */ @@ -18962,6 +21244,11 @@ export interface components { * @enum {string} */ CallTypes: "embedding" | "aembedding" | "completion" | "acompletion" | "atext_completion" | "text_completion" | "image_generation" | "aimage_generation" | "image_edit" | "aimage_edit" | "moderation" | "amoderation" | "atranscription" | "transcription" | "aspeech" | "speech" | "rerank" | "arerank" | "search" | "asearch" | "_arealtime" | "_aresponses_websocket" | "create_batch" | "acreate_batch" | "aretrieve_batch" | "retrieve_batch" | "acancel_batch" | "cancel_batch" | "pass_through_endpoint" | "anthropic_messages" | "get_assistants" | "aget_assistants" | "create_assistants" | "acreate_assistants" | "delete_assistant" | "adelete_assistant" | "acreate_thread" | "create_thread" | "aget_thread" | "get_thread" | "a_add_message" | "add_message" | "aget_messages" | "get_messages" | "arun_thread" | "run_thread" | "arun_thread_stream" | "run_thread_stream" | "afile_retrieve" | "file_retrieve" | "afile_delete" | "file_delete" | "afile_list" | "file_list" | "acreate_file" | "create_file" | "afile_content" | "file_content" | "create_fine_tuning_job" | "acreate_fine_tuning_job" | "create_video" | "acreate_video" | "avideo_retrieve" | "video_retrieve" | "avideo_content" | "video_content" | "video_remix" | "avideo_remix" | "video_list" | "avideo_list" | "video_retrieve_job" | "avideo_retrieve_job" | "video_delete" | "avideo_delete" | "video_create_character" | "avideo_create_character" | "video_get_character" | "avideo_get_character" | "video_edit" | "avideo_edit" | "video_extension" | "avideo_extension" | "vector_store_file_create" | "avector_store_file_create" | "vector_store_file_list" | "avector_store_file_list" | "vector_store_file_retrieve" | "avector_store_file_retrieve" | "vector_store_file_content" | "avector_store_file_content" | "vector_store_file_update" | "avector_store_file_update" | "vector_store_file_delete" | "avector_store_file_delete" | "vector_store_create" | "avector_store_create" | "vector_store_search" | "avector_store_search" | "create_container" | "acreate_container" | "list_containers" | "alist_containers" | "retrieve_container" | "aretrieve_container" | "delete_container" | "adelete_container" | "list_container_files" | "alist_container_files" | "upload_container_file" | "aupload_container_file" | "acancel_fine_tuning_job" | "cancel_fine_tuning_job" | "alist_fine_tuning_jobs" | "list_fine_tuning_jobs" | "aretrieve_fine_tuning_job" | "retrieve_fine_tuning_job" | "responses" | "aresponses" | "alist_input_items" | "llm_passthrough_route" | "allm_passthrough_route" | "generate_content" | "agenerate_content" | "generate_content_stream" | "agenerate_content_stream" | "ocr" | "aocr" | "call_mcp_tool" | "list_mcp_tools" | "asend_message" | "send_message" | "acreate_skill"; + /** CallbackDelete */ + CallbackDelete: { + /** Callback Name */ + callback_name: string; + }; /** CallbacksByType */ CallbacksByType: { /** Failure */ @@ -19581,6 +21868,296 @@ export interface components { /** Regulation */ regulation: string; }; + /** ConfigFieldDelete */ + ConfigFieldDelete: { + /** + * Config Type + * @constant + */ + config_type: "general_settings"; + /** Field Name */ + field_name: string; + }; + /** ConfigFieldInfo */ + ConfigFieldInfo: { + /** Field Name */ + field_name: string; + /** Field Value */ + field_value: unknown; + }; + /** ConfigFieldUpdate */ + ConfigFieldUpdate: { + /** + * Config Type + * @constant + */ + config_type: "general_settings"; + /** Field Name */ + field_name: string; + /** Field Value */ + field_value: unknown; + }; + /** + * ConfigGeneralSettings + * @description Documents all the fields supported by `general_settings` in config.yaml + */ + ConfigGeneralSettings: { + /** + * Alert To Webhook Url + * @description Mapping of alert type to webhook url. e.g. `alert_to_webhook_url: {'budget_alerts': 'https://nothooks.slack.com/services/T00000000/B00000000/XXXXXXXXXXXXXXXXXXXXXXXX'}` + */ + alert_to_webhook_url?: { + [key: string]: unknown; + } | null; + /** + * Alert Types + * @description List of alerting types. By default it is all alerts + */ + alert_types?: components["schemas"]["AlertType"][] | null; + /** + * Alerting + * @description List of alerting integrations. Today, just slack - `alerting: ['slack']` + */ + alerting?: unknown[] | null; + /** + * Alerting Args + * @description Controllable params for slack alerting - e.g. ttl in cache. + */ + alerting_args?: { + [key: string]: unknown; + } | null; + /** + * Alerting Threshold + * @description sends alerts if requests hang for 5min+ + */ + alerting_threshold?: number | null; + /** + * Allowed Routes + * @description Proxy API Endpoints you want users to be able to access + */ + allowed_routes?: unknown[] | null; + /** + * Background Health Checks + * @description run health checks in background + */ + background_health_checks?: boolean | null; + /** + * Completion Model + * @description proxy level default model for all chat completion calls + */ + completion_model?: string | null; + /** + * Custom Auth + * @description override user_api_key_auth with your own auth script - https://docs.litellm.ai/docs/proxy/virtual_keys#custom-auth + */ + custom_auth?: string | null; + /** @description custom args for instantiating dynamodb client - e.g. billing provision */ + database_args?: components["schemas"]["DynamoDBArgs"] | null; + /** + * Database Connect Timeout + * @description Prisma `connect_timeout` URL param (seconds). Bounds how long the engine waits to establish a new connection before failing. Defaults to Prisma's built-in value when unset. + */ + database_connect_timeout?: number | null; + /** + * Database Connection Pool Limit + * @description default connection pool for prisma client connecting to postgres db + * @default 10 + */ + database_connection_pool_limit: number | null; + /** + * Database Connection Timeout + * @description default timeout for a connection to the database + * @default 60 + */ + database_connection_timeout: number | null; + /** + * Database Extra Connection Params + * @description Escape hatch: extra key/value pairs appended verbatim to the Prisma DATABASE_URL / DIRECT_URL query string (e.g. `sslmode`, `pgbouncer`, `statement_cache_size`). Keys here override any default LiteLLM sets. + */ + database_extra_connection_params?: { + [key: string]: unknown; + } | null; + /** + * Database Socket Timeout + * @description Prisma `socket_timeout` URL param (seconds). When set, an idle/slow connection that has not produced data within this window is closed. This is the main knob for capping idle DB connections from LiteLLM. + */ + database_socket_timeout?: number | null; + /** + * Database Type + * @description to use dynamodb instead of postgres db + */ + database_type?: "dynamo_db" | null; + /** + * Database Url + * @description connect to a postgres db - needed for generating temporary keys + tracking spend / key + */ + database_url?: string | null; + /** + * Enable Public Model Hub + * @description Public model hub for users to see what models they have access to, supported openai params, etc. + * @default false + */ + enable_public_model_hub: boolean; + /** + * Forward Client Headers To Llm Api + * @description If True, forwards client headers (e.g. Authorization) to the LLM API. Required for Claude Code with Max subscription. + */ + forward_client_headers_to_llm_api?: boolean | null; + /** + * Global Max Parallel Requests + * @description global max parallel requests to allow for a proxy instance. + */ + global_max_parallel_requests?: number | null; + /** + * Health Check Concurrency + * @description limit concurrent health checks per cycle; when unset, health checks run without a concurrency cap + */ + health_check_concurrency?: number | null; + /** + * Health Check Interval + * @description background health check interval in seconds + * @default 300 + */ + health_check_interval: number; + /** + * Health Check Skip Disabled Background Models + * @description When true, deployments with model_info.disable_background_health_check are skipped for on-demand GET /health as well as the background health loop. + * @default false + */ + health_check_skip_disabled_background_models: boolean; + /** + * Infer Model From Keys + * @description for `/models` endpoint, infers available model based on environment keys (e.g. OPENAI_API_KEY) + */ + infer_model_from_keys?: boolean | null; + /** @description key manager to load keys from / decrypt keys with */ + key_management_system?: components["schemas"]["KeyManagementSystem"] | null; + /** + * Master Key + * @description require a key for all calls to proxy + */ + master_key?: string | null; + /** + * Max Parallel Requests + * @description maximum parallel requests for each api key + */ + max_parallel_requests?: number | null; + /** + * Max Request Size Mb + * @description max request size in MB, if a request is larger than this size it will be rejected + */ + max_request_size_mb?: number | null; + /** + * Max Response Size Mb + * @description max response size in MB, if a response is larger than this size it will be rejected + */ + max_response_size_mb?: number | null; + /** + * Maximum Spend Logs Retention Period + * @description Maximum retention period for spend logs (e.g., '7d' for 7 days). Logs older than this will be deleted. + */ + maximum_spend_logs_retention_period?: string | null; + /** + * Mcp Internal Ip Ranges + * @description Custom CIDR ranges that define internal/private networks for MCP access control. When set, only these ranges are treated as internal. Defaults to RFC 1918 private ranges (10.0.0.0/8, 172.16.0.0/12, 192.168.0.0/16, 127.0.0.0/8). + */ + mcp_internal_ip_ranges?: string[] | null; + /** + * Mcp Required Fields + * @description List of MCP server fields that must be filled in for a submission to pass standards checks (e.g. ['description', 'source_url', 'alias']). + */ + mcp_required_fields?: string[] | null; + /** + * Mcp Trusted Proxy Ranges + * @description CIDR ranges of trusted reverse proxies. When set, X-Forwarded-For and X-Forwarded-* origin headers are only trusted from these IPs. + */ + mcp_trusted_proxy_ranges?: string[] | null; + /** + * Otel + * @description [BETA] OpenTelemetry support - this might change, use with caution. + */ + otel?: boolean | null; + /** + * Pass Through Endpoints + * @description Set-up pass-through endpoints for provider-specific endpoints. Docs - https://docs.litellm.ai/docs/proxy/pass_through + */ + pass_through_endpoints?: components["schemas"]["PassThroughGenericEndpoint"][] | null; + /** + * Reject Clientside Metadata Tags + * @description When set to True, rejects requests that contain client-side 'metadata.tags' to prevent users from influencing budgets by sending different tags. Tags can only be inherited from the API key metadata. + */ + reject_clientside_metadata_tags?: boolean | null; + /** + * Store Model In Db + * @description If True, models and config are stored in and loaded from the database. Default is False. + */ + store_model_in_db?: boolean | null; + /** + * Store Prompts In Spend Logs + * @description If True, stores request messages and responses in spend logs. Default is False. + */ + store_prompts_in_spend_logs?: boolean | null; + /** + * Supported Db Objects + * @description Fine-grained control over which object types to load from the database when store_model_in_db is True. Available types: 'models', 'mcp', 'guardrails', 'vector_stores', 'pass_through_endpoints', 'prompts', 'model_cost_map', 'tools', 'config_overrides'. If not set, all objects are loaded (default behavior). + */ + supported_db_objects?: components["schemas"]["SupportedDBObjectType"][] | null; + /** + * Trusted Proxy Ranges + * @description CIDR ranges of trusted reverse proxies allowed to provide identity headers for header-based auth paths such as enable_oauth2_proxy_auth and custom_ui_sso_sign_in_handler. + */ + trusted_proxy_ranges?: string[] | null; + /** + * Ui Access Mode + * @description Control access to the Proxy UI + * @default all + */ + ui_access_mode: ("admin_only" | "all") | null; + /** + * Use Azure Key Vault + * @description load keys from azure key vault + */ + use_azure_key_vault?: boolean | null; + /** + * Use Google Kms + * @description decrypt keys with google kms + */ + use_google_kms?: boolean | null; + /** User Header Mappings */ + user_header_mappings?: components["schemas"]["UserHeaderMapping"][] | null; + /** + * User Header Name + * @description [DEPRECATED] Use 'user_header_mappings' instead. When set, the header value is treated as the end user id unless overridden by user_header_mappings. + */ + user_header_name?: string | null; + /** + * User Mcp Management Mode + * @description Controls how non-admin users interact with MCP servers in the dashboard. 'restricted' shows only accessible servers, 'view_all' lists every server in read-only mode. + */ + user_mcp_management_mode?: ("restricted" | "view_all") | null; + }; + /** ConfigList */ + ConfigList: { + /** Field Default Value */ + field_default_value: unknown; + /** Field Description */ + field_description: string; + /** Field Name */ + field_name: string; + /** Field Type */ + field_type: string; + /** Field Value */ + field_value: unknown; + /** Nested Fields */ + nested_fields?: components["schemas"]["FieldDetail"][] | null; + /** + * Premium Field + * @default false + */ + premium_field: boolean; + /** Stored In Db */ + stored_in_db: boolean | null; + }; /** * ConfigOverrideSettingsResponse * @description Response model for config override settings GET endpoints. @@ -19606,6 +22183,34 @@ export interface components { [key: string]: unknown; }; }; + /** + * ConfigYAML + * @description Documents all the fields supported by the config.yaml + */ + ConfigYAML: { + /** + * Environment Variables + * @description Object to pass in additional environment variables via POST request + */ + environment_variables?: { + [key: string]: unknown; + } | null; + general_settings?: components["schemas"]["ConfigGeneralSettings"] | null; + /** + * Litellm Settings + * @description litellm Module settings. See __init__.py for all, example litellm.drop_params=True, litellm.set_verbose=True, litellm.api_base, litellm.cache + */ + litellm_settings?: { + [key: string]: unknown; + } | null; + /** + * Model List + * @description List of supported models on the server, with model-specific configs + */ + model_list?: components["schemas"]["ModelParams"][] | null; + /** @description litellm router object settings. See router.py __init__ for all, example router.num_retries=5, router.timeout=5, router.max_retries=5, router.retry_after=5 */ + router_settings?: components["schemas"]["UpdateRouterConfig"] | null; + }; /** ConfigurableClientsideParamsCustomAuth */ "ConfigurableClientsideParamsCustomAuth-Input": { /** Api Base */ @@ -20087,7 +22692,7 @@ export interface components { /** Deployment */ Deployment: { litellm_params: components["schemas"]["LiteLLM_Params"]; - model_info: components["schemas"]["ModelInfo"]; + model_info: components["schemas"]["litellm__types__router__ModelInfo"]; /** Model Name */ model_name: string; } & { @@ -20121,6 +22726,60 @@ export interface components { */ type: "text"; }; + /** DynamoDBArgs */ + DynamoDBArgs: { + /** Assume Role Aws Role Name */ + assume_role_aws_role_name?: string | null; + /** Assume Role Aws Session Name */ + assume_role_aws_session_name?: string | null; + /** Aws Duration Seconds */ + aws_duration_seconds?: number | null; + /** Aws Policy */ + aws_policy?: string | null; + /** Aws Policy Arns */ + aws_policy_arns?: string[] | null; + /** Aws Provider Id */ + aws_provider_id?: string | null; + /** Aws Role Name */ + aws_role_name?: string | null; + /** Aws Session Name */ + aws_session_name?: string | null; + /** Aws Web Identity Token */ + aws_web_identity_token?: string | null; + /** + * Billing Mode + * @enum {string} + */ + billing_mode: "PROVISIONED_THROUGHPUT" | "PAY_PER_REQUEST"; + /** + * Config Table Name + * @default LiteLLM_Config + */ + config_table_name: string; + /** + * Key Table Name + * @default LiteLLM_VerificationToken + */ + key_table_name: string; + /** Read Capacity Units */ + read_capacity_units?: number | null; + /** Region Name */ + region_name: string; + /** + * Spend Table Name + * @default LiteLLM_SpendLogs + */ + spend_table_name: string; + /** Ssl Verify */ + ssl_verify?: boolean | null; + /** + * User Table Name + * @default LiteLLM_UserTable + */ + user_table_name: string; + /** Write Capacity Units */ + write_capacity_units?: number | null; + }; /** * EmailEvent * @enum {string} @@ -20394,6 +23053,19 @@ export interface components { */ model: string; }; + /** FieldDetail */ + FieldDetail: { + /** Field Default Value */ + field_default_value?: unknown; + /** Field Description */ + field_description: string; + /** Field Name */ + field_name: string; + /** Field Type */ + field_type: string; + /** Stored In Db */ + stored_in_db: boolean | null; + }; /** FunctionCall */ FunctionCall: { /** Arguments */ @@ -20728,6 +23400,15 @@ export interface components { */ team_member_permissions: string[] | null; }; + /** GlobalEndUsersSpend */ + GlobalEndUsersSpend: { + /** Api Key */ + api_key?: string | null; + /** Endtime */ + endTime?: string | null; + /** Starttime */ + startTime?: string | null; + }; /** * GraySwanGuardrailConfigModelOptionalParams * @description Optional parameters for the Gray Swan guardrail. @@ -21040,6 +23721,62 @@ export interface components { [key: string]: unknown; }; }; + /** InvitationClaim */ + InvitationClaim: { + /** Invitation Link */ + invitation_link: string; + /** Password */ + password: string; + /** User Id */ + user_id: string; + }; + /** InvitationDelete */ + InvitationDelete: { + /** Invitation Id */ + invitation_id: string; + }; + /** InvitationModel */ + InvitationModel: { + /** Accepted At */ + accepted_at: string | null; + /** + * Created At + * Format: date-time + */ + created_at: string; + /** Created By */ + created_by: string; + /** + * Expires At + * Format: date-time + */ + expires_at: string; + /** Id */ + id: string; + /** Is Accepted */ + is_accepted: boolean; + /** + * Updated At + * Format: date-time + */ + updated_at: string; + /** Updated By */ + updated_by: string; + /** User Id */ + user_id: string; + }; + /** InvitationNew */ + InvitationNew: { + /** User Id */ + user_id: string; + }; + /** InvitationUpdate */ + InvitationUpdate: { + /** Invitation Id */ + invitation_id: string; + /** Is Accepted */ + is_accepted: boolean; + }; /** JWTKeyMappingResponse */ JWTKeyMappingResponse: { /** @@ -21093,6 +23830,11 @@ export interface components { * @enum {string} */ KeyManagementRoutes: "/key/generate" | "/key/update" | "/key/delete" | "/key/regenerate" | "/key/service-account/generate" | "/key/{key_id}/regenerate" | "/key/block" | "/key/unblock" | "/key/bulk_update" | "/team/key/bulk_update" | "/key/{key_id}/reset_spend" | "/key/access_group_assignment" | "/key/info" | "/key/health" | "/key/list" | "/key/aliases" | "/team/daily/activity" | "/spend/logs" | "/spend/logs/v2"; + /** + * KeyManagementSystem + * @enum {string} + */ + KeyManagementSystem: "google_kms" | "azure_key_vault" | "aws_secret_manager" | "google_secret_manager" | "hashicorp_vault" | "cyberark" | "local" | "aws_kms" | "custom"; /** * KeyMetadata * @description Metadata for a key @@ -22685,6 +25427,13 @@ export interface components { /** User Role */ user_role?: string | null; }; + /** LiteLLM_UserTableFiltered */ + LiteLLM_UserTableFiltered: { + /** User Email */ + user_email?: string | null; + /** User Id */ + user_id: string; + }; /** LiteLLM_UserTableWithKeyCount */ LiteLLM_UserTableWithKeyCount: { /** @@ -24005,41 +26754,21 @@ export interface components { /** Tpm */ tpm?: number | null; }; - /** ModelInfo */ - ModelInfo: { - /** Base Model */ - base_model?: string | null; - /** Blocked */ - blocked?: boolean | null; - /** Created At */ - created_at?: string | null; - /** Created By */ - created_by?: string | null; - /** - * Db Model - * @default false - */ - db_model: boolean; - /** Id */ - id: string | null; - /** Team Id */ - team_id?: string | null; - /** Team Public Model Name */ - team_public_model_name?: string | null; - /** Tier */ - tier?: ("free" | "paid") | null; - /** Updated At */ - updated_at?: string | null; - /** Updated By */ - updated_by?: string | null; - } & { - [key: string]: unknown; - }; /** ModelInfoDelete */ ModelInfoDelete: { /** Id */ id: string; }; + /** ModelParams */ + ModelParams: { + /** Litellm Params */ + litellm_params: { + [key: string]: unknown; + }; + model_info: components["schemas"]["litellm__proxy___types__ModelInfo"]; + /** Model Name */ + model_name: string; + }; /** ModelResponse */ ModelResponse: { /** Choices */ @@ -27236,6 +29965,13 @@ export interface components { /** Model */ model?: string | null; }; + /** + * SupportedDBObjectType + * @description Supported database object types for fine-grained DB storage control. + * Use in general_settings.supported_db_objects to specify which objects to load from DB. + * @enum {string} + */ + SupportedDBObjectType: "models" | "mcp" | "guardrails" | "policies" | "vector_stores" | "pass_through_endpoints" | "prompts" | "model_cost_map" | "tools" | "config_overrides"; /** SupportedEndpoint */ SupportedEndpoint: { /** Endpoint */ @@ -29275,6 +32011,19 @@ export interface components { /** User Tpm Limit */ user_tpm_limit?: number | null; }; + /** + * UserHeaderMapping + * @description Map an incoming HTTP header to a LiteLLM user role. + */ + UserHeaderMapping: { + /** Header Name */ + header_name: string; + /** + * Litellm User Role + * @enum {string} + */ + litellm_user_role: "internal_user" | "customer"; + }; /** UserInfoResponse */ UserInfoResponse: { /** Keys */ @@ -29599,12 +32348,68 @@ export interface components { /** Status */ status?: ("pending" | "running" | "paused" | "completed" | "failed") | null; }; + /** ModelInfo */ + litellm__proxy___types__ModelInfo: { + /** Base Model */ + base_model: ("gpt-4-1106-preview" | "gpt-4-32k" | "gpt-4" | "gpt-3.5-turbo-16k" | "gpt-3.5-turbo" | "text-embedding-ada-002") | null; + /** Id */ + id: string | null; + /** + * Input Cost Per Token + * @default 0 + */ + input_cost_per_token: number | null; + /** + * Max Tokens + * @default 2048 + */ + max_tokens: number | null; + /** Mode */ + mode: ("embedding" | "chat" | "completion") | null; + /** + * Output Cost Per Token + * @default 0 + */ + output_cost_per_token: number | null; + } & { + [key: string]: unknown; + }; + /** ModelInfo */ + litellm__types__router__ModelInfo: { + /** Base Model */ + base_model?: string | null; + /** Blocked */ + blocked?: boolean | null; + /** Created At */ + created_at?: string | null; + /** Created By */ + created_by?: string | null; + /** + * Db Model + * @default false + */ + db_model: boolean; + /** Id */ + id: string | null; + /** Team Id */ + team_id?: string | null; + /** Team Public Model Name */ + team_public_model_name?: string | null; + /** Tier */ + tier?: ("free" | "paid") | null; + /** Updated At */ + updated_at?: string | null; + /** Updated By */ + updated_by?: string | null; + } & { + [key: string]: unknown; + }; /** updateDeployment */ updateDeployment: { /** Blocked */ blocked?: boolean | null; litellm_params?: components["schemas"]["updateLiteLLMParams"] | null; - model_info?: components["schemas"]["ModelInfo"] | null; + model_info?: components["schemas"]["litellm__types__router__ModelInfo"] | null; /** Model Name */ model_name?: string | null; }; @@ -30276,6 +33081,26 @@ export interface operations { }; }; }; + alerting_settings_alerting_settings_get: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + }; + }; anthropic_proxy_route_anthropic__endpoint__get: { parameters: { query?: never; @@ -32682,6 +35507,39 @@ export interface operations { }; }; }; + delete_callback_config_callback_delete_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["CallbackDelete"]; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; get_cost_discount_config_config_cost_discount_config_get: { parameters: { query?: never; @@ -32794,6 +35652,134 @@ export interface operations { }; }; }; + delete_config_general_settings_config_field_delete_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["ConfigFieldDelete"]; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + get_config_general_settings_config_field_info_get: { + parameters: { + query: { + field_name: string; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["ConfigFieldInfo"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + update_config_general_settings_config_field_update_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["ConfigFieldUpdate"]; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + get_config_list_config_list_get: { + parameters: { + query: { + config_type: "general_settings"; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["ConfigList"][]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; get_pass_through_endpoints_config_pass_through_endpoint_get: { parameters: { query?: { @@ -32958,6 +35944,72 @@ export interface operations { }; }; }; + update_config_config_update_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["ConfigYAML"]; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + config_yaml_endpoint_config_yaml_get: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["ConfigYAML"]; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; get_hashicorp_vault_config_config_overrides_hashicorp_vault_get: { parameters: { query?: never; @@ -33979,6 +37031,102 @@ export interface operations { }; }; }; + get_memory_details_debug_memory_details_get: { + parameters: { + query?: { + /** @description Number of top object types to return */ + top_n?: number; + /** @description Include process memory info */ + include_process_info?: boolean; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": { + [key: string]: unknown; + }; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + configure_gc_thresholds_endpoint_debug_memory_gc_configure_post: { + parameters: { + query?: { + /** @description Generation 0 threshold (default: 700) */ + generation_0?: number; + /** @description Generation 1 threshold (default: 10) */ + generation_1?: number; + /** @description Generation 2 threshold (default: 10) */ + generation_2?: number; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": { + [key: string]: unknown; + }; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + get_memory_summary_debug_memory_summary_get: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": { + [key: string]: unknown; + }; + }; + }; + }; + }; delete_allowed_ip_delete_allowed_ip_post: { parameters: { query?: never; @@ -34185,6 +37333,261 @@ export interface operations { }; }; }; + block_user_end_user_block_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["BlockUsers"]; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + get_customer_daily_activity_end_user_daily_activity_get: { + parameters: { + query?: { + end_user_ids?: string | null; + start_date?: string | null; + end_date?: string | null; + model?: string | null; + api_key?: string | null; + page?: number; + page_size?: number; + exclude_end_user_ids?: string | null; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + delete_end_user_end_user_delete_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["DeleteCustomerRequest"]; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + end_user_info_end_user_info_get: { + parameters: { + query: { + /** @description End User ID in the request parameters */ + end_user_id: string; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + list_end_user_end_user_list_get: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + }; + }; + new_end_user_end_user_new_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["NewCustomerRequest"]; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + unblock_user_end_user_unblock_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["BlockUsers"]; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + update_end_user_end_user_update_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["UpdateCustomerRequest"]; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; chat_completion_engines__model__chat_completions_post: { parameters: { query?: never; @@ -34707,6 +38110,26 @@ export interface operations { }; }; }; + fallback_login_fallback_login_get: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + }; + }; get_fallback_fallback__model__get: { parameters: { query?: { @@ -35227,6 +38650,46 @@ export interface operations { }; }; }; + get_allowed_ips_get_allowed_ips_get: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + }; + }; + get_config_get_config_callbacks_get: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + }; + }; get_default_team_settings_get_default_team_settings_get: { parameters: { query?: never; @@ -35347,6 +38810,483 @@ export interface operations { }; }; }; + get_favicon_get_favicon_get: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + }; + }; + get_image_get_image_get: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + }; + }; + get_logo_url_get_logo_url_get: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + }; + }; + get_global_activity_global_activity_get: { + parameters: { + query?: { + /** @description Time from which to start viewing spend */ + start_date?: string | null; + /** @description Time till which to view spend */ + end_date?: string | null; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["LiteLLM_SpendLogs"][]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + get_global_activity_global_activity_cache_hits_get: { + parameters: { + query?: { + /** @description Time from which to start viewing spend */ + start_date?: string | null; + /** @description Time till which to view spend */ + end_date?: string | null; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["LiteLLM_SpendLogs"][]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + get_global_activity_exceptions_global_activity_exceptions_get: { + parameters: { + query: { + /** @description Filter by model group */ + model_group: string; + /** @description Time from which to start viewing spend */ + start_date?: string | null; + /** @description Time till which to view spend */ + end_date?: string | null; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["LiteLLM_SpendLogs"][]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + get_global_activity_exceptions_per_deployment_global_activity_exceptions_deployment_get: { + parameters: { + query: { + /** @description Filter by model group */ + model_group: string; + /** @description Time from which to start viewing spend */ + start_date?: string | null; + /** @description Time till which to view spend */ + end_date?: string | null; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["LiteLLM_SpendLogs"][]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + get_global_activity_model_global_activity_model_get: { + parameters: { + query?: { + /** @description Time from which to start viewing spend */ + start_date?: string | null; + /** @description Time till which to view spend */ + end_date?: string | null; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["LiteLLM_SpendLogs"][]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + global_view_all_end_users_global_all_end_users_get: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + }; + }; + global_spend_global_spend_get: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + }; + }; + global_get_all_tag_names_global_spend_all_tag_names_get: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["LiteLLM_SpendLogs"][]; + }; + }; + }; + }; + global_spend_end_users_global_spend_end_users_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: { + content: { + "application/json": components["schemas"]["GlobalEndUsersSpend"] | null; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + global_spend_keys_global_spend_keys_get: { + parameters: { + query?: { + /** @description Number of keys to get. Will return Top 'n' keys. */ + limit?: number; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + global_spend_logs_global_spend_logs_get: { + parameters: { + query?: { + /** @description API Key to get global spend (spend per day for last 30d). Admin-only endpoint */ + api_key?: string | null; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + global_spend_models_global_spend_models_get: { + parameters: { + query?: { + /** @description Number of models to get. Will return Top 'n' models. */ + limit?: number; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + get_global_spend_provider_global_spend_provider_get: { + parameters: { + query?: { + /** @description Time from which to start viewing spend */ + start_date?: string | null; + /** @description Time till which to view spend */ + end_date?: string | null; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["LiteLLM_SpendLogs"][]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + global_spend_refresh_global_spend_refresh_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + }; + }; get_global_spend_report_global_spend_report_get: { parameters: { query?: { @@ -35447,6 +39387,26 @@ export interface operations { }; }; }; + global_spend_per_team_global_spend_teams_get: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + }; + }; create_guardrail_guardrails_post: { parameters: { query?: never; @@ -36695,6 +40655,136 @@ export interface operations { }; }; }; + invitation_delete_invitation_delete_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["InvitationDelete"]; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["InvitationModel"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + invitation_info_invitation_info_get: { + parameters: { + query: { + invitation_id: string; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["InvitationModel"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + new_invitation_invitation_new_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["InvitationNew"]; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["InvitationModel"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + invitation_update_invitation_update_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["InvitationUpdate"]; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["InvitationModel"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; delete_jwt_key_mapping_jwt_key_mapping_delete_post: { parameters: { query?: never; @@ -37536,6 +41626,37 @@ export interface operations { }; }; }; + warm_lazy_warm__name__post: { + parameters: { + query?: never; + header?: never; + path: { + name: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; get_ui_config_litellm__well_known_litellm_ui_config_get: { parameters: { query?: never; @@ -37556,6 +41677,26 @@ export interface operations { }; }; }; + login_login_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + }; + }; test_connection_mcp_rest_test_connection_post: { parameters: { query?: never; @@ -37676,6 +41817,46 @@ export interface operations { }; }; }; + memory_usage_in_mem_cache_memory_usage_in_mem_cache_get: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + }; + }; + memory_usage_in_mem_cache_items_memory_usage_in_mem_cache_items_get: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + }; + }; milvus_proxy_route_milvus__endpoint__get: { parameters: { query?: never; @@ -37986,6 +42167,26 @@ export interface operations { }; }; }; + get_model_cost_map_source_model_cost_map_source_get: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + }; + }; delete_model_model_delete_post: { parameters: { query?: never; @@ -38050,6 +42251,111 @@ export interface operations { }; }; }; + model_metrics_model_metrics_get: { + parameters: { + query?: { + _selected_model_group?: string | null; + startTime?: string | null; + endTime?: string | null; + api_key?: string | null; + customer?: string | null; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + model_metrics_exceptions_model_metrics_exceptions_get: { + parameters: { + query?: { + _selected_model_group?: string | null; + startTime?: string | null; + endTime?: string | null; + api_key?: string | null; + customer?: string | null; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + model_metrics_slow_responses_model_metrics_slow_responses_get: { + parameters: { + query?: { + _selected_model_group?: string | null; + startTime?: string | null; + endTime?: string | null; + api_key?: string | null; + customer?: string | null; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; add_new_model_model_new_post: { parameters: { query?: never; @@ -38083,6 +42389,59 @@ export interface operations { }; }; }; + model_settings_model_settings_get: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + }; + }; + model_streaming_metrics_model_streaming_metrics_get: { + parameters: { + query?: { + _selected_model_group?: string | null; + startTime?: string | null; + endTime?: string | null; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; update_model_model_update_post: { parameters: { query?: never; @@ -38449,6 +42808,70 @@ export interface operations { }; }; }; + claim_onboarding_link_onboarding_claim_token_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["InvitationClaim"]; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + onboarding_onboarding_get_token_get: { + parameters: { + query: { + invite_link: string; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; chat_completion_openai_deployments__model__chat_completions_post: { parameters: { query?: never; @@ -38696,13 +43119,13 @@ export interface operations { /** * @description Unified rate-limit error. * - * Every rate-limit condition surfaced by litellm — whether it originated from - * an upstream LLM provider, a vendor batch endpoint, or one of litellm's own - * proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget, - * max-iterations, etc.) — is raised as an instance of this class. + * Every rate-limit condition surfaced by litellm — whether it originated from + * an upstream LLM provider, a vendor batch endpoint, or one of litellm's own + * proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget, + * max-iterations, etc.) — is raised as an instance of this class. * - * The :attr:`category` attribute lets callers distinguish the source. See - * :class:`RateLimitErrorCategory` for the available values. + * The :attr:`category` attribute lets callers distinguish the source. See + * :class:`RateLimitErrorCategory` for the available values. */ 429: { headers: { @@ -39784,6 +44207,26 @@ export interface operations { }; }; }; + get_otel_spans_otel_spans_get: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + }; + }; create_policy_policies_post: { parameters: { query?: never; @@ -41355,6 +45798,37 @@ export interface operations { }; }; }; + async_queue_request_queue_chat_completions_post: { + parameters: { + query?: { + model?: string | null; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; rag_ingest_rag_ingest_post: { parameters: { query?: never; @@ -41453,6 +45927,46 @@ export interface operations { }; }; }; + reload_anthropic_beta_headers_reload_anthropic_beta_headers_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + }; + }; + reload_model_cost_map_reload_model_cost_map_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + }; + }; rerank_rerank_post: { parameters: { query?: never; @@ -41735,6 +46249,148 @@ export interface operations { }; }; }; + schedule_anthropic_beta_headers_reload_schedule_anthropic_beta_headers_reload_post: { + parameters: { + query: { + hours: number; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + cancel_anthropic_beta_headers_reload_schedule_anthropic_beta_headers_reload_delete: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + }; + }; + get_anthropic_beta_headers_reload_status_schedule_anthropic_beta_headers_reload_status_get: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + }; + }; + schedule_model_cost_map_reload_schedule_model_cost_map_reload_post: { + parameters: { + query: { + hours: number; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + cancel_model_cost_map_reload_schedule_model_cost_map_reload_delete: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + }; + }; + get_model_cost_map_reload_status_schedule_model_cost_map_reload_status_get: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + }; + }; get_scim_base_scim_v2_get: { parameters: { query?: { @@ -42683,6 +47339,26 @@ export interface operations { }; }; }; + spend_key_fn_spend_keys_get: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + }; + }; view_spend_logs_spend_logs_get: { parameters: { query?: { @@ -42725,6 +47401,148 @@ export interface operations { }; }; }; + ui_view_session_spend_logs_spend_logs_session_ui_get: { + parameters: { + query: { + /** @description Get all spend logs for a particular session */ + session_id: string; + /** @description Page number for pagination */ + page?: number; + /** @description Number of items per page */ + page_size?: number; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["LiteLLM_SpendLogs"][]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + ui_view_spend_logs_spend_logs_ui_get: { + parameters: { + query?: { + /** @description Get spend logs based on api key */ + api_key?: string | null; + /** @description Get spend logs based on user_id */ + user_id?: string | null; + /** @description request_id to get spend logs for specific request_id */ + request_id?: string | null; + /** @description Filter spend logs by team_id */ + team_id?: string | null; + /** @description Filter logs with spend greater than or equal to this value */ + min_spend?: number | null; + /** @description Filter logs with spend less than or equal to this value */ + max_spend?: number | null; + /** @description Time from which to start viewing key spend */ + start_date?: string | null; + /** @description Time till which to view key spend */ + end_date?: string | null; + /** @description Page number for pagination */ + page?: number; + /** @description Number of items per page */ + page_size?: number; + /** @description Filter logs by status (e.g., success, failure) */ + status_filter?: string | null; + /** @description Filter logs by model */ + model?: string | null; + /** @description Filter logs by model ID (litellm model deployment id) */ + model_id?: string | null; + /** @description Filter logs by model group */ + model_group?: string | null; + /** @description Filter logs by key alias */ + key_alias?: string | null; + /** @description Filter logs by end user */ + end_user?: string | null; + /** @description Filter logs by error code (e.g., '404', '500') */ + error_code?: string | null; + /** @description Filter logs by error message (partial string match) */ + error_message?: string | null; + /** @description Sort logs by field: spend, total_tokens, startTime, endTime, request_duration_ms, model, or ttft_ms */ + sort_by?: string; + /** @description Sort order: asc or desc */ + sort_order?: string | null; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["LiteLLM_SpendLogs"][]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + ui_view_request_response_for_request_id_spend_logs_ui__request_id__get: { + parameters: { + query?: { + /** @description Time from which to start viewing key spend */ + start_date?: string | null; + /** @description Time till which to view key spend */ + end_date?: string | null; + }; + header?: never; + path: { + request_id: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; ui_view_spend_logs_spend_logs_v2_get: { parameters: { query?: { @@ -42831,6 +47649,249 @@ export interface operations { }; }; }; + spend_user_fn_spend_users_get: { + parameters: { + query?: { + /** @description Get User Table row for user_id */ + user_id?: string | null; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + auth_callback_sso_callback_get: { + parameters: { + query?: { + state?: string | null; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + cli_sso_complete_sso_cli_complete__login_id__post: { + parameters: { + query?: never; + header?: never; + path: { + login_id: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + cli_poll_key_sso_cli_poll__key_id__get: { + parameters: { + query?: { + team_id?: string | null; + }; + header?: { + "x-litellm-cli-poll-secret"?: string | null; + }; + path: { + key_id: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + cli_sso_start_sso_cli_start_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + }; + }; + debug_sso_callback_sso_debug_callback_get: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + }; + }; + debug_sso_login_sso_debug_login_get: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + }; + }; + get_ui_settings_sso_get_ui_settings_get: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + }; + }; + google_login_sso_key_generate_get: { + parameters: { + query?: { + source?: string | null; + key?: string | null; + existing_key?: string | null; + return_to?: string | null; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; sso_readiness_sso_readiness_get: { parameters: { query?: never; @@ -43423,6 +48484,44 @@ export interface operations { }; }; }; + ui_view_teams_team_filter_ui_get: { + parameters: { + query?: { + /** @description Team ID in the request parameters */ + team_id?: string | null; + /** @description Team alias in the request parameters */ + team_alias?: string | null; + /** @description Page number for pagination */ + page?: number; + /** @description Number of items per page */ + page_size?: number; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["LiteLLM_TeamTable"][]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; team_info_team_info_get: { parameters: { query?: { @@ -44673,6 +49772,26 @@ export interface operations { }; }; }; + ui_get_available_role_user_available_roles_get: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + }; + }; available_enterprise_users_user_available_users_get: { parameters: { query?: never; @@ -44853,6 +49972,46 @@ export interface operations { }; }; }; + ui_view_users_user_filter_ui_get: { + parameters: { + query?: { + /** @description User ID in the request parameters */ + user_id?: string | null; + /** @description User email in the request parameters */ + user_email?: string | null; + /** @description Team ID — used when a team admin searches for users to add to their team */ + team_id?: string | null; + /** @description Page number for pagination */ + page?: number; + /** @description Number of items per page */ + page_size?: number; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["LiteLLM_UserTableFiltered"][]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; user_info_user_info_get: { parameters: { query?: { @@ -50595,6 +55754,110 @@ export interface operations { }; }; }; + info_key_fn_v2_v2_key_info_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: { + content: { + "application/json": components["schemas"]["KeyRequest"] | null; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + login_v2_v2_login_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + }; + }; + model_info_v2_v2_model_info_get: { + parameters: { + query?: { + /** @description Specify the model name (optional) */ + model?: string | null; + /** @description Only return models added by this user */ + user_models_only?: boolean | null; + /** @description Return all models across all teams user is in. */ + include_team_models?: boolean | null; + debug?: boolean | null; + /** @description Page number */ + page?: number; + /** @description Page size */ + size?: number; + /** @description Search model names (case-insensitive partial match) */ + search?: string | null; + /** @description Search for a specific model by its unique ID */ + modelId?: string | null; + /** @description Filter models by team ID. Returns models with direct_access=True or teamId in access_via_team_ids */ + teamId?: string | null; + /** @description Field to sort by. Options: model_name, created_at, updated_at, costs, status */ + sortBy?: string | null; + /** @description Sort order. Options: asc, desc */ + sortOrder?: string | null; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; rerank_v2_rerank_post: { parameters: { query?: never; @@ -50697,6 +55960,46 @@ export interface operations { }; }; }; + login_v3_v3_login_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + }; + }; + login_v3_exchange_v3_login_exchange_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + }; + }; delete_vantage_settings_vantage_delete_delete: { parameters: { query?: never; @@ -51401,6 +56704,161 @@ export interface operations { }; }; }; + vertex_proxy_route_vertex_ai__endpoint__get_2: { + parameters: { + query?: never; + header?: never; + path: { + endpoint: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + vertex_proxy_route_vertex_ai__endpoint__put_2: { + parameters: { + query?: never; + header?: never; + path: { + endpoint: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + vertex_proxy_route_vertex_ai__endpoint__post_2: { + parameters: { + query?: never; + header?: never; + path: { + endpoint: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + vertex_proxy_route_vertex_ai__endpoint__delete_2: { + parameters: { + query?: never; + header?: never; + path: { + endpoint: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + vertex_proxy_route_vertex_ai__endpoint__patch_2: { + parameters: { + query?: never; + header?: never; + path: { + endpoint: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; vertex_discovery_proxy_route_vertex_ai_discovery__endpoint__get: { parameters: { query?: never; From f5b11b72a6dcc8f4e7a16f00a61bdd124cc171d2 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Mon, 8 Jun 2026 22:03:35 +0530 Subject: [PATCH 011/185] feat(proxy): publish /v2/model/info in Swagger OpenAPI spec (#29900) * feat(proxy): publish /v2/model/info in Swagger OpenAPI spec Expose the v2 model info endpoint in /docs by removing include_in_schema=False and documenting query parameters used by the admin UI and proxy CLI consumers. Co-authored-by: Cursor * chore(ui): regenerate schema.d.ts for /v2/model/info OpenAPI docs Co-authored-by: Cursor --------- Co-authored-by: Cursor --- litellm/proxy/proxy_server.py | 46 ++++++++++++++- .../proxy_server/test_routes_model_info.py | 9 +++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 56 ++++++++++++++++--- 3 files changed, 101 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 2c21e19dcec..213f682b8f7 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11742,10 +11742,8 @@ async def _find_model_by_id( @router.get( "/v2/model/info", - description="v2 - returns models available to the user based on their API key permissions. Shows model info from config.yaml (except api key and api base). Filter to just user-added models with ?user_models_only=true", tags=["model management"], dependencies=[Depends(user_api_key_auth)], - include_in_schema=False, ) async def model_info_v2( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), @@ -11781,7 +11779,49 @@ async def model_info_v2( ), ): """ - BETA ENDPOINT. Might change unexpectedly. Use `/v1/model/info` for now. + Paginated model metadata for proxy deployments (pricing, provider, team access). + + Returns configured router deployments with enriched `model_info` (costs, provider, + context window, etc.). Sensitive fields such as API keys and api_base are omitted. + + Query parameters: + model: Filter to a single public `model_name`. + user_models_only: When true, only return models created by the calling user. + include_team_models: When true, populate `access_via_team_ids` and `direct_access` + on each model and filter to deployments the caller can use. + page / size: Pagination controls (defaults: page=1, size=50). + search: Case-insensitive partial match on model name or team public name. + modelId: Return a single deployment by LiteLLM model id. + teamId: Filter to models with direct access or team membership for this team id. + sortBy / sortOrder: Sort by model_name, created_at, updated_at, costs, or status. + + Example request: + ``` + curl -X GET 'http://localhost:4000/v2/model/info?include_team_models=true&page=1&size=50' \\ + --header 'Authorization: Bearer sk-1234' + ``` + + Example response: + ```json + { + "data": [ + { + "model_name": "gpt-4", + "litellm_params": {"model": "openai/gpt-4.1"}, + "model_info": { + "id": "abc123", + "litellm_provider": "openai", + "access_via_team_ids": ["team-1"], + "direct_access": true + } + } + ], + "total_count": 1, + "current_page": 1, + "total_pages": 1, + "size": 50 + } + ``` """ global llm_model_list, general_settings, user_config_file_path, proxy_config, llm_router diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py b/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py index 98259824378..017f4bd4368 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py @@ -60,6 +60,15 @@ def test_v2_model_info_invalid_page_returns_422(client, auth_as, empty_router): assert "detail" in response.json() +def test_v2_model_info_in_openapi_schema(): + """``GET /v2/model/info`` is published in the proxy OpenAPI/Swagger spec.""" + from litellm.proxy.proxy_server import get_openapi_schema + + schema = get_openapi_schema() + assert "/v2/model/info" in schema["paths"] + assert "get" in schema["paths"]["/v2/model/info"] + + # --------------------------------------------------------------------------- # GET /v1/model/info, GET /model/info # --------------------------------------------------------------------------- diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index c8ec9d3c727..75a14e09852 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -18441,7 +18441,49 @@ export interface paths { }; /** * Model Info V2 - * @description v2 - returns models available to the user based on their API key permissions. Shows model info from config.yaml (except api key and api base). Filter to just user-added models with ?user_models_only=true + * @description Paginated model metadata for proxy deployments (pricing, provider, team access). + * + * Returns configured router deployments with enriched `model_info` (costs, provider, + * context window, etc.). Sensitive fields such as API keys and api_base are omitted. + * + * Query parameters: + * model: Filter to a single public `model_name`. + * user_models_only: When true, only return models created by the calling user. + * include_team_models: When true, populate `access_via_team_ids` and `direct_access` + * on each model and filter to deployments the caller can use. + * page / size: Pagination controls (defaults: page=1, size=50). + * search: Case-insensitive partial match on model name or team public name. + * modelId: Return a single deployment by LiteLLM model id. + * teamId: Filter to models with direct access or team membership for this team id. + * sortBy / sortOrder: Sort by model_name, created_at, updated_at, costs, or status. + * + * Example request: + * ``` + * curl -X GET 'http://localhost:4000/v2/model/info?include_team_models=true&page=1&size=50' \ + * --header 'Authorization: Bearer sk-1234' + * ``` + * + * Example response: + * ```json + * { + * "data": [ + * { + * "model_name": "gpt-4", + * "litellm_params": {"model": "openai/gpt-4.1"}, + * "model_info": { + * "id": "abc123", + * "litellm_provider": "openai", + * "access_via_team_ids": ["team-1"], + * "direct_access": true + * } + * } + * ], + * "total_count": 1, + * "current_page": 1, + * "total_pages": 1, + * "size": 50 + * } + * ``` */ get: operations["model_info_v2_v2_model_info_get"]; put?: never; @@ -43119,13 +43161,13 @@ export interface operations { /** * @description Unified rate-limit error. * - * Every rate-limit condition surfaced by litellm — whether it originated from - * an upstream LLM provider, a vendor batch endpoint, or one of litellm's own - * proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget, - * max-iterations, etc.) — is raised as an instance of this class. + * Every rate-limit condition surfaced by litellm — whether it originated from + * an upstream LLM provider, a vendor batch endpoint, or one of litellm's own + * proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget, + * max-iterations, etc.) — is raised as an instance of this class. * - * The :attr:`category` attribute lets callers distinguish the source. See - * :class:`RateLimitErrorCategory` for the available values. + * The :attr:`category` attribute lets callers distinguish the source. See + * :class:`RateLimitErrorCategory` for the available values. */ 429: { headers: { From ff6cea4833df88690cd282e51654411f0bca0d84 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Mon, 8 Jun 2026 11:25:50 -0700 Subject: [PATCH 012/185] refactor(ui): single source of truth for migrated-page routing (#29949) Consolidate the three hand-synced copies of the migrated-pages map (LEGACY_REDIRECTS in app/page.tsx, MIGRATED_PAGES in the dashboard layout, and MIGRATED_PAGES in leftnav) into one shared module, src/utils/migratedPages.ts, which also owns the migratedHref helper. Delete the unused, incomplete Sidebar2 prototype. No runtime behavior change: the map is still empty and Sidebar2 had no importers, so this is pure deduplication ahead of the per-page App Router migration. Follow-up work will unify the remaining base-URL builders (layout's withBase and page.tsx's redirect) onto migratedHref. --- .../app/(dashboard)/components/Sidebar2.tsx | 472 ------------------ .../src/app/(dashboard)/layout.tsx | 10 +- ui/litellm-dashboard/src/app/page.tsx | 12 +- .../src/components/leftnav.tsx | 24 +- .../src/utils/migratedPages.test.ts | 43 ++ .../src/utils/migratedPages.ts | 25 + 6 files changed, 73 insertions(+), 513 deletions(-) delete mode 100644 ui/litellm-dashboard/src/app/(dashboard)/components/Sidebar2.tsx create mode 100644 ui/litellm-dashboard/src/utils/migratedPages.test.ts create mode 100644 ui/litellm-dashboard/src/utils/migratedPages.ts diff --git a/ui/litellm-dashboard/src/app/(dashboard)/components/Sidebar2.tsx b/ui/litellm-dashboard/src/app/(dashboard)/components/Sidebar2.tsx deleted file mode 100644 index 90f498912a8..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/components/Sidebar2.tsx +++ /dev/null @@ -1,472 +0,0 @@ -"use client"; - -import { Layout, Menu, ConfigProvider } from "antd"; -import { - KeyOutlined, - PlayCircleOutlined, - BlockOutlined, - BarChartOutlined, - TeamOutlined, - BankOutlined, - UserOutlined, - SettingOutlined, - ApiOutlined, - AppstoreOutlined, - DatabaseOutlined, - FileTextOutlined, - LineChartOutlined, - SafetyOutlined, - ExperimentOutlined, - ToolOutlined, - TagsOutlined, - AuditOutlined, -} from "@ant-design/icons"; -// import { -// all_admin_roles, -// rolesWithWriteAccess, -// internalUserRoles, -// isAdminRole, -// } from "../utils/roles"; -// import UsageIndicator from "./usage_indicator"; -import * as React from "react"; -import { useRouter, usePathname } from "next/navigation"; -import { all_admin_roles, internalUserRoles, isAdminRole, rolesWithWriteAccess } from "@/utils/roles"; -import UsageIndicator from "@/components/UsageIndicator"; -import { serverRootPath } from "@/components/networking"; - -const { Sider } = Layout; - -// -------- Types -------- -interface SidebarProps { - accessToken: string | null; - userRole: string; - /** Fallback selection id (legacy), used if path can't be matched */ - defaultSelectedKey: string; - collapsed?: boolean; -} - -interface MenuItemCfg { - key: string; - newTab?: boolean; - page: string; // legacy id; we map this to a path below - label: string; - roles?: string[]; - children?: MenuItemCfg[]; - icon?: React.ReactNode; -} - -/** ---------- Base URL helpers ---------- */ -/** - * Normalizes NEXT_PUBLIC_BASE_URL to either "/" or "/ui/" (always with a trailing slash). - * Supported env values: "" or "ui/". - * Also considers the serverRootPath from the proxy config (e.g., "/my-custom-path"). - */ -const getBasePath = () => { - const raw = process.env.NEXT_PUBLIC_BASE_URL ?? ""; - const trimmed = raw.replace(/^\/+|\/+$/g, ""); // strip leading/trailing slashes - const uiPath = trimmed ? `/${trimmed}/` : "/"; - - // If serverRootPath is set and not "/", prepend it to the UI path - if (serverRootPath && serverRootPath !== "/") { - // Remove trailing slash from serverRootPath and ensure uiPath has no leading slash for proper joining - const cleanServerRoot = serverRootPath.replace(/\/+$/, ""); - const cleanUiPath = uiPath.replace(/^\/+/, ""); - return `${cleanServerRoot}/${cleanUiPath}`; - } - - return uiPath; -}; - -/** Map legacy `page` ids to real app routes (relative, no leading slash). */ -const routeFor = (slug: string): string => { - switch (slug) { - // top level - case "api-keys": - return "virtual-keys"; - case "llm-playground": - return "test-key"; - case "models": - return "models-and-endpoints"; - case "new_usage": - return "usage"; - case "teams": - return "teams"; - case "organizations": - return "organizations"; - case "users": - return "users"; - case "api_ref": - return "api-reference"; - case "model-hub-table": - // If you intend the newer in-dashboard page, use "model-hub". - return "model-hub"; - case "logs": - return "logs"; - case "guardrails": - return "guardrails"; - case "policies": - return "policies"; - case "chat": - return "chat"; - - // tools - case "mcp-servers": - return "tools/mcp-servers"; - case "vector-stores": - return "tools/vector-stores"; - case "byok-demo": - return "tools/byok-demo"; - - // experimental - case "caching": - return "experimental/caching"; - case "prompts": - return "experimental/prompts"; - case "budgets": - return "experimental/budgets"; - case "transform-request": - return "experimental/api-playground"; - case "tag-management": - return "experimental/tag-management"; - case "claude-code-plugins": - return "experimental/claude-code-plugins"; - case "usage": // "Old Usage" - return "experimental/old-usage"; - - // settings - case "general-settings": - return "settings/router-settings"; - case "settings": // "Logging & Alerts" - return "settings/logging-and-alerts"; - case "admin-panel": - return "settings/admin-settings"; - case "ui-theme": - return "settings/ui-theme"; - - default: - // treat as already a relative path - return slug.replace(/^\/+/, ""); - } -}; - -/** Prefix base path ("/" or "/ui/") */ -const toHref = (slugOrPath: string) => { - const base = getBasePath(); // "/" or "/ui/" - const rel = routeFor(slugOrPath).replace(/^\/+|\/+$/g, ""); - return `${base}${rel}`; -}; - -// ----- Menu config (unchanged labels/icons; same appearance) ----- -const menuItems: MenuItemCfg[] = [ - { key: "1", page: "api-keys", label: "Virtual Keys", icon: }, - { - key: "3", - page: "llm-playground", - label: "Test Key", - icon: , - roles: rolesWithWriteAccess, - }, - { - key: "2", - page: "models", - label: "Models + Endpoints", - icon: , - roles: rolesWithWriteAccess, - }, - { - key: "12", - page: "new_usage", - label: "Usage", - icon: , - roles: [...all_admin_roles, ...internalUserRoles], - }, - { key: "6", page: "teams", label: "Teams", icon: }, - { - key: "17", - page: "organizations", - label: "Organizations", - icon: , - roles: all_admin_roles, - }, - { - key: "5", - page: "users", - label: "Internal Users", - icon: , - roles: all_admin_roles, - }, - { key: "14", page: "api-reference", label: "API Reference", icon: }, - { - key: "16", - page: "model-hub-table", - label: "Model Hub", - icon: , - }, - { key: "15", page: "logs", label: "Logs", icon: }, - { - key: "11", - page: "guardrails", - label: "Guardrails", - icon: , - roles: all_admin_roles, - }, - { - key: "28", - page: "policies", - label: "Policies", - icon: , - roles: all_admin_roles, - }, - { - key: "26", - page: "tools", - label: "Tools", - icon: , - children: [ - { key: "18", page: "mcp-servers", label: "MCP Servers", icon: }, - { - key: "21", - page: "vector-stores", - label: "Vector Stores", - icon: , - roles: all_admin_roles, - }, - ], - }, - { - key: "experimental", - page: "experimental", - label: "Experimental", - icon: , - children: [ - { - key: "9", - page: "caching", - label: "Caching", - icon: , - roles: all_admin_roles, - }, - { - key: "25", - page: "prompts", - label: "Prompts", - icon: , - roles: all_admin_roles, - }, - { - key: "10", - page: "budgets", - label: "Budgets", - icon: , - roles: all_admin_roles, - }, - { - key: "20", - page: "transform-request", - label: "API Playground", - icon: , - roles: [...all_admin_roles, ...internalUserRoles], - }, - { - key: "19", - page: "tag-management", - label: "Tag Management", - icon: , - roles: all_admin_roles, - }, - { - key: "27", - page: "claude-code-plugins", - label: "Claude Code Plugins", - icon: , - roles: all_admin_roles, - }, - { key: "4", page: "usage", label: "Old Usage", icon: }, - ], - }, - { - key: "settings", - page: "settings", - label: "Settings", - icon: , - roles: all_admin_roles, - children: [ - { - key: "11", - page: "general-settings", - label: "Router Settings", - icon: , - roles: all_admin_roles, - }, - { - key: "8", - page: "settings", - label: "Logging & Alerts", - icon: , - roles: all_admin_roles, - }, - { - key: "13", - page: "admin-panel", - label: "Admin Settings", - icon: , - roles: all_admin_roles, - }, - { - key: "14", - page: "ui-theme", - label: "UI Theme", - icon: , - roles: all_admin_roles, - }, - ], - }, -]; - -const Sidebar2: React.FC = ({ accessToken, userRole, defaultSelectedKey, collapsed = false }) => { - const router = useRouter(); - const pathname = usePathname() || "/"; - - // ----- Filter by role without mutating originals ----- - const filteredMenuItems = React.useMemo(() => { - return menuItems - .filter((item) => !item.roles || item.roles.includes(userRole)) - .map((item) => ({ - ...item, - children: item.children ? item.children.filter((c) => !c.roles || c.roles.includes(userRole)) : undefined, - })); - }, [userRole]); - - // ----- Compute selected key from current path ----- - const selectedMenuKey = React.useMemo(() => { - const base = getBasePath(); - // strip base prefix and leading slash -> "virtual-keys", "tools/mcp-servers", etc. - const rel = pathname.startsWith(base) ? pathname.slice(base.length) : pathname.replace(/^\/+/, ""); - const relLower = rel.toLowerCase(); - - const matchesPath = (slug: string) => { - const route = routeFor(slug).toLowerCase(); - return relLower === route || relLower.startsWith(`${route}/`); - }; - - // search top-level - for (const item of filteredMenuItems) { - if (!item.children && matchesPath(item.page)) return item.key; - if (item.children) { - for (const child of item.children) { - if (matchesPath(child.page)) return child.key; - } - } - } - - // fallback to legacy defaultSelectedKey mapping - const fallback = filteredMenuItems.find((i) => i.page === defaultSelectedKey)?.key; - if (fallback) return fallback; - - for (const item of filteredMenuItems) { - if (item.children?.some((c) => c.page === defaultSelectedKey)) { - const child = item.children.find((c) => c.page === defaultSelectedKey)!; - return child.key; - } - } - - return "1"; - }, [pathname, filteredMenuItems, defaultSelectedKey]); - - // ----- Navigation ----- - const goTo = (slug: string, newTab?: boolean) => { - const href = toHref(slug); - if (newTab) { - window.open(href, "_blank"); - } else { - router.push(href); - } - }; - - // Wrap label in so every nav item supports right-click → "Open in new tab" - // and Ctrl/Cmd+click to open in a new tab, while preserving SPA navigation for normal clicks. - const renderNavLink = (label: string, page: string, newTab?: boolean): React.ReactNode => { - const href = toHref(page); - return ( - { - if (newTab) { - e.stopPropagation(); - return; - } - if (e.metaKey || e.ctrlKey || e.shiftKey || e.button === 1) { - e.stopPropagation(); - return; - } - e.preventDefault(); - }} - style={{ color: "inherit", textDecoration: "none" }} - > - {label} - - ); - }; - - return ( - - - - ({ - key: item.key, - icon: item.icon, - label: renderNavLink(item.label, item.page, item.newTab), - children: item.children?.map((child) => ({ - key: child.key, - icon: child.icon, - label: renderNavLink(child.label, child.page, child.newTab), - onClick: () => goTo(child.page, child.newTab), - })), - onClick: !item.children ? () => goTo(item.page, item.newTab) : undefined, - }))} - /> - - {isAdminRole(userRole) && !collapsed && } - - - ); -}; - -export default Sidebar2; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx index a611d619cc1..7dc2b2ad615 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx @@ -7,6 +7,7 @@ import SidebarProvider from "@/app/(dashboard)/components/SidebarProvider"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { useRouter, useSearchParams } from "next/navigation"; import { DebugWarningBanner } from "@/components/DebugWarningBanner"; +import { MIGRATED_PAGES } from "@/utils/migratedPages"; /** ---- BASE URL HELPERS ---- */ function normalizeBasePrefix(raw: string | undefined | null): string { @@ -23,15 +24,6 @@ function withBase(path: string): string { } /** -------------------------------- */ -/** - * Pages that have been migrated to path-based routing under (dashboard)/. - * When the leftnav triggers one of these, navigate to the path route instead - * of the legacy query-param root page. - * - * Key = legacy page id used in leftnav, Value = route segment under (dashboard)/ - */ -const MIGRATED_PAGES: Record = {}; - function LayoutContent({ children }: { children: React.ReactNode }) { const router = useRouter(); const searchParams = useSearchParams(); diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx index da6a0d5a76f..f329a00d60d 100644 --- a/ui/litellm-dashboard/src/app/page.tsx +++ b/ui/litellm-dashboard/src/app/page.tsx @@ -54,6 +54,7 @@ import { storeReturnUrl, } from "@/utils/returnUrlUtils"; import { isAdminRole } from "@/utils/roles"; +import { MIGRATED_PAGES } from "@/utils/migratedPages"; import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; import { useRouter, useSearchParams } from "next/navigation"; import { Suspense, useEffect, useMemo, useRef, useState } from "react"; @@ -65,13 +66,6 @@ interface ProxySettings { LITELLM_UI_API_DOC_BASE_URL?: string | null; } -/** - * Map of legacy query-param page keys → new path-based route segments. - * When a user visits ?page=, they are redirected to /ui/. - * Add entries here as pages are migrated from the if/else chain to path-based routes. - */ -const LEGACY_REDIRECTS: Record = {}; - function CreateKeyPageContent() { const { authLoading, token, userID, userRole, userEmail, accessToken, premiumUser, setUserRole, setUserEmail } = useAuth(); @@ -202,11 +196,11 @@ function CreateKeyPageContent() { }, [redirectToLogin]); // Redirect legacy query-param pages to their new path-based routes - const isLegacyRedirect = page in LEGACY_REDIRECTS; + const isLegacyRedirect = page in MIGRATED_PAGES; useEffect(() => { if (!authLoading && isLegacyRedirect) { const base = (proxyBaseUrl || "") + "/ui"; - router.replace(`${base}/${LEGACY_REDIRECTS[page]}`); + router.replace(`${base}/${MIGRATED_PAGES[page]}`); } }, [authLoading, isLegacyRedirect, page, router]); diff --git a/ui/litellm-dashboard/src/components/leftnav.tsx b/ui/litellm-dashboard/src/components/leftnav.tsx index 43a28d94046..ff0dac703d3 100644 --- a/ui/litellm-dashboard/src/components/leftnav.tsx +++ b/ui/litellm-dashboard/src/components/leftnav.tsx @@ -43,31 +43,9 @@ import { import NewBadge from "./common_components/NewBadge"; import type { Organization } from "./networking"; import UsageIndicator from "./UsageIndicator"; -import { serverRootPath } from "./networking"; +import { MIGRATED_PAGES, migratedHref } from "@/utils/migratedPages"; const { Sider } = Layout; -/** - * Pages migrated to path-based routing under (dashboard)/. - * Key = legacy page id, Value = route segment. - * Keep in sync with MIGRATED_PAGES in (dashboard)/layout.tsx. - */ -const MIGRATED_PAGES: Record = {}; - -/** Build an absolute href for a migrated page, respecting base URL + serverRootPath. */ -function migratedHref(routeSegment: string): string { - const raw = process.env.NEXT_PUBLIC_BASE_URL ?? ""; - const trimmed = raw.replace(/^\/+|\/+$/g, ""); - let base = trimmed ? `/${trimmed}/` : "/"; - - if (serverRootPath && serverRootPath !== "/") { - const cleanRoot = serverRootPath.replace(/\/+$/, ""); - const cleanBase = base.replace(/^\/+/, ""); - base = `${cleanRoot}/${cleanBase}`; - } - - return `${base}${routeSegment}`; -} - // Define the props type interface SidebarProps { setPage: (page: string) => void; diff --git a/ui/litellm-dashboard/src/utils/migratedPages.test.ts b/ui/litellm-dashboard/src/utils/migratedPages.test.ts new file mode 100644 index 00000000000..2f9ed11b3a0 --- /dev/null +++ b/ui/litellm-dashboard/src/utils/migratedPages.test.ts @@ -0,0 +1,43 @@ +import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; + +describe("migratedHref", () => { + beforeEach(() => { + vi.resetModules(); + vi.unstubAllEnvs(); + }); + + afterEach(() => { + vi.unstubAllEnvs(); + }); + + it("returns an absolute path rooted at / when serverRootPath is /", async () => { + vi.doMock("@/components/networking", () => ({ serverRootPath: "/" })); + const { migratedHref } = await import("./migratedPages"); + + expect(migratedHref("api-reference")).toBe("/api-reference"); + expect(migratedHref("virtual-keys")).toBe("/virtual-keys"); + }); + + it("prefixes a non-root serverRootPath without duplicating slashes", async () => { + vi.doMock("@/components/networking", () => ({ serverRootPath: "/team-x/" })); + const { migratedHref } = await import("./migratedPages"); + + expect(migratedHref("api-reference")).toBe("/team-x/api-reference"); + }); + + it("honors NEXT_PUBLIC_BASE_URL as a base path segment", async () => { + vi.stubEnv("NEXT_PUBLIC_BASE_URL", "ui"); + vi.doMock("@/components/networking", () => ({ serverRootPath: "/" })); + const { migratedHref } = await import("./migratedPages"); + + expect(migratedHref("api-reference")).toBe("/ui/api-reference"); + }); + + it("combines NEXT_PUBLIC_BASE_URL with a non-root serverRootPath", async () => { + vi.stubEnv("NEXT_PUBLIC_BASE_URL", "ui"); + vi.doMock("@/components/networking", () => ({ serverRootPath: "/team-x" })); + const { migratedHref } = await import("./migratedPages"); + + expect(migratedHref("api-reference")).toBe("/team-x/ui/api-reference"); + }); +}); diff --git a/ui/litellm-dashboard/src/utils/migratedPages.ts b/ui/litellm-dashboard/src/utils/migratedPages.ts new file mode 100644 index 00000000000..5421f00bf05 --- /dev/null +++ b/ui/litellm-dashboard/src/utils/migratedPages.ts @@ -0,0 +1,25 @@ +import { serverRootPath } from "@/components/networking"; + +/** + * Single source of truth for pages cut over from the legacy `?page=` switch in + * app/page.tsx to path-based routes under app/(dashboard)/. + * + * Key = legacy page id emitted by the sidebar. Value = route segment under (dashboard)/. + * Add an entry to route the sidebar and deep links to the new path and redirect the + * legacy `?page=` URL; remove it to roll back. Empty until a page is migrated. + */ +export const MIGRATED_PAGES: Record = {}; + +export function migratedHref(routeSegment: string): string { + const raw = process.env.NEXT_PUBLIC_BASE_URL ?? ""; + const trimmed = raw.replace(/^\/+|\/+$/g, ""); + let base = trimmed ? `/${trimmed}/` : "/"; + + if (serverRootPath && serverRootPath !== "/") { + const cleanRoot = serverRootPath.replace(/\/+$/, ""); + const cleanBase = base.replace(/^\/+/, ""); + base = `${cleanRoot}/${cleanBase}`; + } + + return `${base}${routeSegment}`; +} From 26fe26a5c041d0aa03f9e669b255d89b2ab67859 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 8 Jun 2026 12:12:07 -0700 Subject: [PATCH 013/185] fix(ui/model-hub): render provider icons on the public model hub (#29958) The provider logo base path was relative ("../ui/assets/logos/"). With trailingSlash enabled, the public model hub is served at /ui/model_hub_table/, so the browser resolved the base to /ui/ui/assets/logos/ (a doubled /ui/), which 404s every icon. The authenticated hub renders inside the single-level /ui/ SPA route where the relative path resolves correctly, so only the public hub broke. Make the base root-absolute so it resolves at any route depth. --- .../components/provider_info_helpers.test.tsx | 21 +++++++++++++++++++ .../src/components/provider_info_helpers.tsx | 2 +- 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx index ab22b0ef49a..a73cf699ac8 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx @@ -97,6 +97,27 @@ describe("provider_info_helpers", () => { }); }); + describe("provider logo asset paths", () => { + // Regression: a relative "../ui/assets/logos/" base resolved to + // "/ui/ui/assets/logos/..." (404) on the public model hub at + // /ui/model_hub_table/, which sits a level below the /ui/ SPA. Root-absolute + // paths resolve correctly at any route depth. + it("should expose every provider logo as a root-absolute /ui path", () => { + const logos = Object.values(providerLogoMap); + expect(logos.length).toBeGreaterThan(0); + logos.forEach((logo) => { + expect(logo.startsWith("/ui/assets/logos/")).toBe(true); + expect(logo).not.toContain("../"); + }); + }); + + it("should resolve a provider logo to a root-absolute path via getProviderLogoAndName", () => { + const { logo } = getProviderLogoAndName("openai"); + expect(logo.startsWith("/ui/assets/logos/")).toBe(true); + expect(logo).not.toContain("../"); + }); + }); + describe("getPlaceholder", () => { it("should return aiml placeholder for AIML provider", () => { expect(getPlaceholder(Providers.AIML)).toBe("aiml/flux-pro/v1.1"); diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx index 179ce63f457..727eebfe951 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx @@ -216,7 +216,7 @@ export const provider_map: Record = { ZAI: "zai", }; -const asset_logos_folder = "../ui/assets/logos/"; +const asset_logos_folder = "/ui/assets/logos/"; export const providerLogoMap: Record = { [Providers.A2A_Agent]: `${asset_logos_folder}a2a_agent.png`, From 47b383dbbfc9eb3336ee1b7e57a0e7d4d380c46f Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 8 Jun 2026 12:25:42 -0700 Subject: [PATCH 014/185] fix(ui): keep create guardrail modal open on outside click (#29871) The create guardrail modal used antd's default maskClosable, so clicking outside it dismissed the modal and reset every field the user had entered. Setting maskClosable={false} keeps the modal open; it now closes only via the explicit close button or Cancel, matching the other form modals in the dashboard --- .../guardrails/add_guardrail_form.test.tsx | 43 +++++++++++++++++++ .../guardrails/add_guardrail_form.tsx | 1 + 2 files changed, 44 insertions(+) create mode 100644 ui/litellm-dashboard/src/components/guardrails/add_guardrail_form.test.tsx diff --git a/ui/litellm-dashboard/src/components/guardrails/add_guardrail_form.test.tsx b/ui/litellm-dashboard/src/components/guardrails/add_guardrail_form.test.tsx new file mode 100644 index 00000000000..a13f12e4b42 --- /dev/null +++ b/ui/litellm-dashboard/src/components/guardrails/add_guardrail_form.test.tsx @@ -0,0 +1,43 @@ +import React from "react"; +import { fireEvent, screen } from "@testing-library/react"; +import { renderWithProviders } from "../../../tests/test-utils"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import AddGuardrailForm from "./add_guardrail_form"; + +vi.mock("@/components/networking", () => ({ + createGuardrailCall: vi.fn(), + getGuardrailProviderSpecificParams: vi.fn().mockResolvedValue({}), + getGuardrailUISettings: vi.fn().mockResolvedValue({}), + modelAvailableCall: vi.fn().mockResolvedValue({ data: [] }), +})); + +const renderForm = () => { + const onClose = vi.fn(); + renderWithProviders(); + return { onClose }; +}; + +describe("AddGuardrailForm close behavior", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("does not close when the user clicks outside the modal on the mask", () => { + const { onClose } = renderForm(); + expect(screen.getByText("Create guardrail")).toBeInTheDocument(); + + const wrap = document.querySelector(".ant-modal-wrap") as HTMLElement; + expect(wrap).toBeTruthy(); + fireEvent.mouseDown(wrap); + fireEvent.mouseUp(wrap); + fireEvent.click(wrap); + + expect(onClose).not.toHaveBeenCalled(); + }); + + it("closes when the user clicks the explicit close button", () => { + const { onClose } = renderForm(); + fireEvent.click(screen.getByRole("button", { name: "✕" })); + expect(onClose).toHaveBeenCalledTimes(1); + }); +}); diff --git a/ui/litellm-dashboard/src/components/guardrails/add_guardrail_form.tsx b/ui/litellm-dashboard/src/components/guardrails/add_guardrail_form.tsx index 05a65d1c9e5..41a2dd7f67c 100644 --- a/ui/litellm-dashboard/src/components/guardrails/add_guardrail_form.tsx +++ b/ui/litellm-dashboard/src/components/guardrails/add_guardrail_form.tsx @@ -1144,6 +1144,7 @@ const AddGuardrailForm: React.FC = ({ visible, onClose, a title={null} open={visible} onCancel={handleClose} + maskClosable={false} footer={null} width={1000} closable={false} From 728f057c5ec3334604fc65c57c41bac873d4555f Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 8 Jun 2026 12:25:52 -0700 Subject: [PATCH 015/185] fix(ui): label default key type as "Full Access" on key edit page (#29870) The key edit page showed the default key type (no allowed_routes restriction) as "Default", while the key creation form already labels the same value "Full Access". Align the edit page to the create form so the two surfaces agree on both the label and its description. --- .../src/components/templates/key_edit_view.tsx | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx index d5e410029a7..fe0fa1ab0c2 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx @@ -394,11 +394,11 @@ export function KeyEditView({ } }} > - +
-
Default
+
Full Access
- Can call AI APIs + Management routes + Can call all routes (AI APIs, Management, and read-only)
From 1afc41cb295e78a25a5d67389ac5a831e290ee8d Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Mon, 8 Jun 2026 13:05:12 -0700 Subject: [PATCH 016/185] fix(ui): unify migrated-route URLs and migrate the API Reference page (#29953) * fix(ui): unify migrated-route URLs and cut the API Reference page over to path routing Route all migrated-page navigation through one /ui-prefixed, serverRootPath-aware builder in migratedPages.ts (migratedHref/legacyPageHref/legacyKeyForPathname), replacing the three divergent base-URL constructions that lived in the dashboard layout's withBase, leftnav, and the page.tsx redirect. The previous migratedHref read NEXT_PUBLIC_BASE_URL, which no build sets, so it produced URLs without the /ui prefix the app is served under; every other internal link hardcodes /ui and this now matches that convention. Remove the sidebar's own pushState navigation so the parent (legacy root page or dashboard layout) is the single owner of navigation, fixing the double-navigate that fired when moving between a path route and a legacy ?page= route. Cut API Reference over to its path route: add api_ref -> api-reference to MIGRATED_PAGES and delete its arm from the legacy switch. Visiting /ui/?page=api_ref redirects to /ui/api-reference, the sidebar links to and highlights it, and navigating away returns to the legacy switch. * fix(ui): address review on migrated-page routing Keep the legacy hyphenated ?page=api-reference form working by mapping it to the api-reference route alongside api_ref; the old switch matched both, so a bookmark using the hyphen would otherwise fall through to the Usage default. Add legacyKeyForPathname coverage: a migrated path (with and without trailing slash) resolves to the api_ref sidebar key rather than the alias, a non-migrated path returns null, and a non-root serverRootPath prefix is stripped before matching. * fix(ui): populate serverRootPath from getUiConfig so migrated nav links keep the root path getUiConfig updated proxyBaseUrl but never called updateServerRootPath, so the module-level serverRootPath stayed at its "/" default. Under a custom server_root_path the unified migratedHref/legacyPageHref builders then dropped the prefix and the sidebar produced /ui/api-reference (404) instead of //ui/api-reference. Adds the missing updateServerRootPath call plus a regression test asserting getUiConfig sets serverRootPath and that migratedHref carries the prefix --- .../src/app/(dashboard)/layout.tsx | 36 ++-------- ui/litellm-dashboard/src/app/page.tsx | 19 +++--- .../src/components/leftnav.tsx | 26 ++----- .../src/components/networking.test.ts | 15 ++++ .../src/components/networking.tsx | 1 + .../src/utils/migratedPages.test.ts | 68 +++++++++++++------ .../src/utils/migratedPages.ts | 43 ++++++++---- 7 files changed, 112 insertions(+), 96 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx index 7dc2b2ad615..5f5c240d025 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx @@ -5,51 +5,29 @@ import Navbar from "@/components/navbar"; import { ThemeProvider } from "@/contexts/ThemeContext"; import SidebarProvider from "@/app/(dashboard)/components/SidebarProvider"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; -import { useRouter, useSearchParams } from "next/navigation"; +import { useRouter, useSearchParams, usePathname } from "next/navigation"; import { DebugWarningBanner } from "@/components/DebugWarningBanner"; -import { MIGRATED_PAGES } from "@/utils/migratedPages"; - -/** ---- BASE URL HELPERS ---- */ -function normalizeBasePrefix(raw: string | undefined | null): string { - const trimmed = (raw ?? "").trim(); - if (!trimmed) return ""; - const core = trimmed.replace(/^\/+/, "").replace(/\/+$/, ""); - return core ? `/${core}/` : "/"; -} -const BASE_PREFIX = normalizeBasePrefix(process.env.NEXT_PUBLIC_BASE_URL); -function withBase(path: string): string { - const body = path.startsWith("/") ? path.slice(1) : path; - const combined = `${BASE_PREFIX}${body}`; - return combined.startsWith("/") ? combined : `/${combined}`; -} -/** -------------------------------- */ +import { MIGRATED_PAGES, migratedHref, legacyPageHref, legacyKeyForPathname } from "@/utils/migratedPages"; function LayoutContent({ children }: { children: React.ReactNode }) { const router = useRouter(); const searchParams = useSearchParams(); + const pathname = usePathname(); const { accessToken } = useAuthorized(); const [sidebarCollapsed, setSidebarCollapsed] = React.useState(false); const [page, setPage] = useState(() => { - return searchParams.get("page") || "api-keys"; + return legacyKeyForPathname(pathname) || searchParams.get("page") || "api-keys"; }); const handleSetPage = (newPage: string) => { - // If the page has been migrated to path routing, navigate there const migratedRoute = MIGRATED_PAGES[newPage]; - if (migratedRoute) { - router.push(withBase(migratedRoute)); - setPage(newPage); - return; - } - - // Otherwise, navigate back to the legacy root page with query params - router.push(withBase(`?page=${newPage}`)); + router.push(migratedRoute ? migratedHref(migratedRoute) : legacyPageHref(newPage)); setPage(newPage); }; useEffect(() => { - setPage(searchParams.get("page") || "api-keys"); - }, [searchParams]); + setPage(legacyKeyForPathname(pathname) || searchParams.get("page") || "api-keys"); + }, [pathname, searchParams]); const toggleSidebar = () => setSidebarCollapsed((v) => !v); diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx index f329a00d60d..12dd39a1c21 100644 --- a/ui/litellm-dashboard/src/app/page.tsx +++ b/ui/litellm-dashboard/src/app/page.tsx @@ -1,6 +1,5 @@ "use client"; -import APIReferenceView from "@/app/(dashboard)/api-reference/APIReferenceView"; import SidebarProvider from "@/app/(dashboard)/components/SidebarProvider"; import OldModelDashboard from "@/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView"; import PlaygroundPage from "@/app/(dashboard)/playground/page"; @@ -54,7 +53,7 @@ import { storeReturnUrl, } from "@/utils/returnUrlUtils"; import { isAdminRole } from "@/utils/roles"; -import { MIGRATED_PAGES } from "@/utils/migratedPages"; +import { MIGRATED_PAGES, migratedHref } from "@/utils/migratedPages"; import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; import { useRouter, useSearchParams } from "next/navigation"; import { Suspense, useEffect, useMemo, useRef, useState } from "react"; @@ -156,15 +155,16 @@ function CreateKeyPageContent() { return searchParams.get("page") || "api-keys"; }); - // Custom setPage function that updates URL const updatePage = (newPage: string) => { - // Update URL without full page reload + const migratedRoute = MIGRATED_PAGES[newPage]; + if (migratedRoute) { + router.push(migratedHref(migratedRoute)); + setPage(newPage); + return; + } const newSearchParams = new URLSearchParams(searchParams); newSearchParams.set("page", newPage); - - // Use Next.js router to update URL window.history.pushState(null, "", `?${newSearchParams.toString()}`); - setPage(newPage); }; @@ -199,8 +199,7 @@ function CreateKeyPageContent() { const isLegacyRedirect = page in MIGRATED_PAGES; useEffect(() => { if (!authLoading && isLegacyRedirect) { - const base = (proxyBaseUrl || "") + "/ui"; - router.replace(`${base}/${MIGRATED_PAGES[page]}`); + router.replace(migratedHref(MIGRATED_PAGES[page])); } }, [authLoading, isLegacyRedirect, page, router]); @@ -441,8 +440,6 @@ function CreateKeyPageContent() { /> ) : page == "admin-panel" ? ( - ) : page == "api_ref" || page == "api-reference" ? ( - ) : page == "logging-and-alerts" ? ( ) : page == "budgets" ? ( diff --git a/ui/litellm-dashboard/src/components/leftnav.tsx b/ui/litellm-dashboard/src/components/leftnav.tsx index ff0dac703d3..f4946b81d68 100644 --- a/ui/litellm-dashboard/src/components/leftnav.tsx +++ b/ui/litellm-dashboard/src/components/leftnav.tsx @@ -43,7 +43,7 @@ import { import NewBadge from "./common_components/NewBadge"; import type { Organization } from "./networking"; import UsageIndicator from "./UsageIndicator"; -import { MIGRATED_PAGES, migratedHref } from "@/utils/migratedPages"; +import { MIGRATED_PAGES, migratedHref, legacyPageHref } from "@/utils/migratedPages"; const { Sider } = Layout; // Define the props type @@ -418,18 +418,9 @@ const Sidebar: React.FC = ({ // Check if user is a team admin for any team const isTeamAdmin = useMemo(() => isUserTeamAdminForAnyTeam(teams ?? null, userId ?? ""), [teams, userId]); - // Navigate to page helper - const navigateToPage = (page: string) => { - // For migrated pages, just call setPage — the parent layout handles routing - if (MIGRATED_PAGES[page]) { - setPage(page); - return; - } - const newSearchParams = new URLSearchParams(window.location.search); - newSearchParams.set("page", page); - window.history.pushState(null, "", `?${newSearchParams.toString()}`); - setPage(page); - }; + // The parent (legacy root page or dashboard layout) owns navigation for both + // migrated and legacy pages; the sidebar only reports the selected page. + const navigateToPage = (page: string) => setPage(page); // Wrap label in so every nav item supports right-click → "Open in new tab" // and Ctrl/Cmd+click to open in a new tab, while preserving SPA navigation for normal clicks. @@ -447,15 +438,8 @@ const Sidebar: React.FC = ({ ); } - // For migrated pages, generate a path-based href for right-click "Open in new tab" const migratedRoute = MIGRATED_PAGES[page]; - const href = migratedRoute - ? migratedHref(migratedRoute) - : (() => { - const params = new URLSearchParams(window.location.search); - params.set("page", page); - return `?${params.toString()}`; - })(); + const href = migratedRoute ? migratedHref(migratedRoute) : legacyPageHref(page); return ( ({ clearTokenCookies: vi.fn(), @@ -349,6 +350,20 @@ describe("UI config and public endpoints", () => { ); expect(configCall).toBeDefined(); }); + + it("updates serverRootPath so path-based nav links carry the root path", async () => { + const uiConfig = { + server_root_path: "/litellm", + proxy_base_url: "https://example.com", + }; + + setupMockFetch([{ url: "/litellm/.well-known/litellm-ui-config", data: uiConfig }]); + + await Networking.getUiConfig(); + + expect(Networking.serverRootPath).toBe("/litellm"); + expect(migratedHref("api-reference")).toBe("/litellm/ui/api-reference"); + }); }); describe("individualModelHealthCheckCall", () => { diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index d303734dffd..54a7a19362b 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -404,6 +404,7 @@ export const getUiConfig = async () => { * Update the proxy base url and server root path */ console.log("jsonData in getUiConfig:", jsonData); + updateServerRootPath(jsonData.server_root_path); updateProxyBaseUrl(jsonData.server_root_path, jsonData.proxy_base_url); return jsonData; }; diff --git a/ui/litellm-dashboard/src/utils/migratedPages.test.ts b/ui/litellm-dashboard/src/utils/migratedPages.test.ts index 2f9ed11b3a0..e9aceea8148 100644 --- a/ui/litellm-dashboard/src/utils/migratedPages.test.ts +++ b/ui/litellm-dashboard/src/utils/migratedPages.test.ts @@ -1,43 +1,69 @@ -import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; +import { describe, it, expect, vi, beforeEach } from "vitest"; -describe("migratedHref", () => { +describe("migratedHref / legacyPageHref", () => { beforeEach(() => { vi.resetModules(); - vi.unstubAllEnvs(); }); - afterEach(() => { - vi.unstubAllEnvs(); - }); - - it("returns an absolute path rooted at / when serverRootPath is /", async () => { + it("builds a /ui-rooted path when serverRootPath is /", async () => { vi.doMock("@/components/networking", () => ({ serverRootPath: "/" })); - const { migratedHref } = await import("./migratedPages"); + const { migratedHref, legacyPageHref } = await import("./migratedPages"); - expect(migratedHref("api-reference")).toBe("/api-reference"); - expect(migratedHref("virtual-keys")).toBe("/virtual-keys"); + expect(migratedHref("api-reference")).toBe("/ui/api-reference"); + expect(legacyPageHref("models")).toBe("/ui/?page=models"); }); it("prefixes a non-root serverRootPath without duplicating slashes", async () => { vi.doMock("@/components/networking", () => ({ serverRootPath: "/team-x/" })); - const { migratedHref } = await import("./migratedPages"); + const { migratedHref, legacyPageHref } = await import("./migratedPages"); - expect(migratedHref("api-reference")).toBe("/team-x/api-reference"); + expect(migratedHref("api-reference")).toBe("/team-x/ui/api-reference"); + expect(legacyPageHref("models")).toBe("/team-x/ui/?page=models"); }); - it("honors NEXT_PUBLIC_BASE_URL as a base path segment", async () => { - vi.stubEnv("NEXT_PUBLIC_BASE_URL", "ui"); + it("tolerates a leading slash in the route segment", async () => { vi.doMock("@/components/networking", () => ({ serverRootPath: "/" })); const { migratedHref } = await import("./migratedPages"); - expect(migratedHref("api-reference")).toBe("/ui/api-reference"); + expect(migratedHref("/api-reference")).toBe("/ui/api-reference"); }); - it("combines NEXT_PUBLIC_BASE_URL with a non-root serverRootPath", async () => { - vi.stubEnv("NEXT_PUBLIC_BASE_URL", "ui"); - vi.doMock("@/components/networking", () => ({ serverRootPath: "/team-x" })); - const { migratedHref } = await import("./migratedPages"); + it("maps both the api_ref id and the hyphenated alias to the api-reference route", async () => { + vi.doMock("@/components/networking", () => ({ serverRootPath: "/" })); + const { MIGRATED_PAGES } = await import("./migratedPages"); - expect(migratedHref("api-reference")).toBe("/team-x/ui/api-reference"); + expect(MIGRATED_PAGES.api_ref).toBe("api-reference"); + expect(MIGRATED_PAGES["api-reference"]).toBe("api-reference"); + }); +}); + +describe("legacyKeyForPathname", () => { + beforeEach(() => { + vi.resetModules(); + }); + + it("maps a migrated path back to its legacy sidebar key (including trailing slash)", async () => { + vi.doMock("@/components/networking", () => ({ serverRootPath: "/" })); + const { legacyKeyForPathname } = await import("./migratedPages"); + + // Resolves to the sidebar key api_ref, not the hyphenated alias, so highlighting works. + expect(legacyKeyForPathname("/ui/api-reference")).toBe("api_ref"); + expect(legacyKeyForPathname("/ui/api-reference/")).toBe("api_ref"); + }); + + it("returns null for a not-yet-migrated path", async () => { + vi.doMock("@/components/networking", () => ({ serverRootPath: "/" })); + const { legacyKeyForPathname } = await import("./migratedPages"); + + expect(legacyKeyForPathname("/ui/")).toBeNull(); + expect(legacyKeyForPathname("/ui/some-legacy-page")).toBeNull(); + }); + + it("strips a non-root serverRootPath prefix before matching", async () => { + vi.doMock("@/components/networking", () => ({ serverRootPath: "/team-x/" })); + const { legacyKeyForPathname } = await import("./migratedPages"); + + expect(legacyKeyForPathname("/team-x/ui/api-reference")).toBe("api_ref"); + expect(legacyKeyForPathname("/ui/api-reference")).toBeNull(); }); }); diff --git a/ui/litellm-dashboard/src/utils/migratedPages.ts b/ui/litellm-dashboard/src/utils/migratedPages.ts index 5421f00bf05..2c27e4fee64 100644 --- a/ui/litellm-dashboard/src/utils/migratedPages.ts +++ b/ui/litellm-dashboard/src/utils/migratedPages.ts @@ -6,20 +6,35 @@ import { serverRootPath } from "@/components/networking"; * * Key = legacy page id emitted by the sidebar. Value = route segment under (dashboard)/. * Add an entry to route the sidebar and deep links to the new path and redirect the - * legacy `?page=` URL; remove it to roll back. Empty until a page is migrated. + * legacy `?page=` URL; remove it to roll back. */ -export const MIGRATED_PAGES: Record = {}; +export const MIGRATED_PAGES: Record = { + api_ref: "api-reference", + // Legacy alias: older bookmarks used the hyphenated ?page=api-reference form. + "api-reference": "api-reference", +}; -export function migratedHref(routeSegment: string): string { - const raw = process.env.NEXT_PUBLIC_BASE_URL ?? ""; - const trimmed = raw.replace(/^\/+|\/+$/g, ""); - let base = trimmed ? `/${trimmed}/` : "/"; - - if (serverRootPath && serverRootPath !== "/") { - const cleanRoot = serverRootPath.replace(/\/+$/, ""); - const cleanBase = base.replace(/^\/+/, ""); - base = `${cleanRoot}/${cleanBase}`; - } - - return `${base}${routeSegment}`; +function uiBase(): string { + const root = serverRootPath && serverRootPath !== "/" ? `/${serverRootPath.replace(/^\/+|\/+$/g, "")}` : ""; + return `${root}/ui`; +} + +/** Absolute (same-origin) href for a migrated route segment, e.g. "api-reference" -> "/ui/api-reference". */ +export function migratedHref(routeSegment: string): string { + return `${uiBase()}/${routeSegment.replace(/^\/+/, "")}`; +} + +/** Href for a not-yet-migrated page, served by the legacy `?page=` switch at the UI root. */ +export function legacyPageHref(pageKey: string): string { + return `${uiBase()}/?page=${pageKey}`; +} + +/** Reverse-maps a path-routed location back to its legacy page id, e.g. "/ui/api-reference" -> "api_ref". */ +export function legacyKeyForPathname(pathname: string): string | null { + const base = uiBase(); + const rel = (pathname.startsWith(base) ? pathname.slice(base.length) : pathname).replace(/^\/+|\/+$/g, ""); + for (const [key, segment] of Object.entries(MIGRATED_PAGES)) { + if (rel === segment) return key; + } + return null; } From 1528f43d4c5f6c9981d065761834957cf559f7b5 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Mon, 8 Jun 2026 13:35:49 -0700 Subject: [PATCH 017/185] fix(mcp): let non-creator users OAuth into OBO-mode MCP servers from the Tools page (#29867) * fix(ui): let non-creator users OAuth into OBO-mode MCP servers from the Tools page * fix(ui): clear OBO Tools-tab one-shot on navigate-back and gate on credential-status errors --- .../components/mcp_tools/mcp_server_view.tsx | 4 +- .../src/components/mcp_tools/mcp_servers.tsx | 41 +++++- .../components/mcp_tools/mcp_tools.test.tsx | 55 +++++++- .../src/components/mcp_tools/mcp_tools.tsx | 119 +++++++++++++++--- .../src/hooks/mcpOAuthUtils.ts | 9 ++ 5 files changed, 209 insertions(+), 19 deletions(-) diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_view.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_view.tsx index cafecdafd45..fc1fe7cd405 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_view.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_view.tsx @@ -21,6 +21,7 @@ interface MCPServerViewProps { userRole: string | null; userID: string | null; availableAccessGroups: string[]; + initialTabIndex?: number; } export const MCPServerView: React.FC = ({ @@ -32,11 +33,12 @@ export const MCPServerView: React.FC = ({ userRole, userID, availableAccessGroups, + initialTabIndex = 0, }) => { const [editing, setEditing] = useState(isEditing); const [showFullUrl, setShowFullUrl] = useState(false); const [copiedStates, setCopiedStates] = useState>({}); - const [selectedTabIndex, setSelectedTabIndex] = useState(0); + const [selectedTabIndex, setSelectedTabIndex] = useState(initialTabIndex); const handleSuccess = (updated: MCPServer) => { setEditing(false); diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx index 4b383f701cb..0a445265e2d 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx @@ -21,6 +21,7 @@ import MCPNetworkSettings from "./MCPNetworkSettings"; import MCPDiscovery from "./mcp_discovery"; import { ByokCredentialModal } from "./ByokCredentialModal"; import { getSecureItem } from "@/utils/secureStorage"; +import { TOOLS_OAUTH_UI_STATE_KEY } from "@/hooks/mcpOAuthUtils"; import UserEnvVarsModal from "./UserEnvVarsModal"; import { listMCPUserEnvVarStatus } from "../networking"; @@ -71,6 +72,23 @@ const compareServers = (a: MCPServer, b: MCPServer, sort: SortKey): number => { const { Text: AntdText, Title: AntdTitle } = Typography; const EDIT_OAUTH_UI_STATE_KEY = "litellm-mcp-oauth-edit-state"; +// Server id stashed by the Tools tab before an OBO OAuth redirect, read once at +// mount so the redirect returns straight to that server's Tools tab. +const readToolsOAuthServerId = (): string | null => { + if (typeof window === "undefined") { + return null; + } + try { + const stored = getSecureItem(TOOLS_OAUTH_UI_STATE_KEY); + if (!stored) { + return null; + } + return JSON.parse(stored)?.serverId ?? null; + } catch { + return null; + } +}; + const { Option } = Select; const MCPServers: React.FC = ({ accessToken, userRole, userID }) => { @@ -103,7 +121,12 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) // state const [serverIdToDelete, setServerToDelete] = useState(null); const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false); - const [selectedServerId, setSelectedServerId] = useState(null); + // Server whose Tools tab should be reopened after an OBO OAuth redirect; read + // once from sessionStorage so the restored server selection is correct on the + // first render. Cleared when the user navigates back to the list (handleBack) + // so a later visit to the same server defaults to Overview, not the Tools tab. + const [toolsTabServerId, setToolsTabServerId] = useState(readToolsOAuthServerId); + const [selectedServerId, setSelectedServerId] = useState(toolsTabServerId); const [editServer, setEditServer] = useState(false); const [selectedTeam, setSelectedTeam] = useState("all"); const [selectedMcpAccessGroup, setSelectedMcpAccessGroup] = useState("all"); @@ -178,6 +201,19 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) } }, []); + // The restored server id was consumed by the initializer above; remove the + // one-shot sessionStorage key so a full page reload doesn't reopen the Tools + // tab (removeItem only, no setState). + useEffect(() => { + if (typeof window !== "undefined") { + try { + window.sessionStorage.removeItem(TOOLS_OAUTH_UI_STATE_KEY); + } catch { + // ignore storage errors + } + } + }, []); + // Get unique teams from all servers const uniqueTeams = React.useMemo(() => { if (!serversWithHealth) return []; @@ -338,6 +374,8 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) const handleBack = React.useCallback(() => { setEditServer(false); setSelectedServerId(null); + // Drop the post-redirect one-shot so re-selecting that server opens Overview. + setToolsTabServerId(null); refetch(); }, [refetch]); @@ -483,6 +521,7 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) userID={userID} userRole={userRole} availableAccessGroups={uniqueMcpAccessGroups} + initialTabIndex={selectedServerId === toolsTabServerId ? 1 : 0} /> ) : (
diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_tools.test.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_tools.test.tsx index 4917f1fcf1d..72ed3984244 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_tools.test.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_tools.test.tsx @@ -2,12 +2,13 @@ import { render, screen, waitFor } from "@testing-library/react"; import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; import { describe, expect, it, vi, beforeEach } from "vitest"; import MCPToolsViewer from "./mcp_tools"; -import { listMCPTools } from "../networking"; +import { listMCPTools, getMCPOAuthUserCredentialStatus } from "../networking"; import { isTokenValid, getToken } from "@/utils/mcpTokenStore"; vi.mock("../networking", () => ({ listMCPTools: vi.fn(), callMCPTool: vi.fn(), + getMCPOAuthUserCredentialStatus: vi.fn(), })); vi.mock("@/utils/mcpTokenStore", () => ({ @@ -20,6 +21,10 @@ vi.mock("@/hooks/useToolsOAuthFlow", () => ({ useToolsOAuthFlow: () => ({ startOAuthFlow: vi.fn(), status: "idle", error: null }), })); +vi.mock("@/hooks/useUserMcpOAuthFlow", () => ({ + useUserMcpOAuthFlow: () => ({ startOAuthFlow: vi.fn(), status: "idle", error: null }), +})); + const GATE_TEXT = "Authentication required"; // Realistic interactive servers carry a token endpoint; the old heuristic // (`oauth2 && !tokenUrl`) mislabeled exactly these as M2M. Setting it here is @@ -42,6 +47,13 @@ const renderViewer = (props: Record) => , ); +const credStatus = (overrides: Record = {}) => ({ + server_id: "srv-1", + has_credential: true, + is_expired: false, + ...overrides, +}); + describe("MCPToolsViewer auth gate routing", () => { beforeEach(() => { vi.mocked(listMCPTools).mockReset().mockResolvedValue({ tools: [], error: null }); @@ -49,6 +61,8 @@ describe("MCPToolsViewer auth gate routing", () => { vi.mocked(getToken) .mockReset() .mockReturnValue(undefined as any); + // Default: the OBO credential exists and is valid, so OBO servers list tools. + vi.mocked(getMCPOAuthUserCredentialStatus).mockReset().mockResolvedValue(credStatus()); }); it("shows the Authorize gate for a passthrough server with a token endpoint and does not list tools", async () => { @@ -57,6 +71,8 @@ describe("MCPToolsViewer auth gate routing", () => { expect(await screen.findByText(GATE_TEXT)).toBeInTheDocument(); expect(screen.getByRole("button", { name: "Authorize" })).toBeInTheDocument(); expect(vi.mocked(listMCPTools)).not.toHaveBeenCalled(); + // Passthrough must not consult the per-user DB credential. + expect(vi.mocked(getMCPOAuthUserCredentialStatus)).not.toHaveBeenCalled(); }); it("forwards the session token via the x-mcp header for a passthrough server that has one", async () => { @@ -75,7 +91,40 @@ describe("MCPToolsViewer auth gate routing", () => { expect(screen.queryByText(GATE_TEXT)).not.toBeInTheDocument(); }); - it("does not gate an OBO server with a token endpoint; lists with the LiteLLM key and no x-mcp header", async () => { + it("lists tools for an OBO server when the user has a DB credential, with no x-mcp header", async () => { + renderViewer({ oauth2_flow: null, delegate_auth_to_upstream: false }); + + await waitFor(() => expect(vi.mocked(listMCPTools)).toHaveBeenCalledWith("litellm-key", "srv-1", undefined)); + expect(screen.queryByText(GATE_TEXT)).not.toBeInTheDocument(); + }); + + it("shows the Authorize gate for an OBO server when the user has no DB credential and does not list tools", async () => { + vi.mocked(getMCPOAuthUserCredentialStatus).mockResolvedValue(credStatus({ has_credential: false })); + + renderViewer({ oauth2_flow: null, delegate_auth_to_upstream: false }); + + expect(await screen.findByText(GATE_TEXT)).toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Authorize" })).toBeInTheDocument(); + expect(vi.mocked(listMCPTools)).not.toHaveBeenCalled(); + }); + + it("shows the Authorize gate for an OBO server when the credential-status check fails", async () => { + vi.mocked(getMCPOAuthUserCredentialStatus).mockRejectedValue(new Error("network down")); + + renderViewer({ oauth2_flow: null, delegate_auth_to_upstream: false }); + + expect(await screen.findByText(GATE_TEXT)).toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Authorize" })).toBeInTheDocument(); + expect(vi.mocked(listMCPTools)).not.toHaveBeenCalled(); + }); + + it("does not gate an OBO server whose stored token is expired; the list call refreshes it server-side", async () => { + // has_credential=true with is_expired=true must NOT gate: resolve_valid_user_oauth_token + // refreshes from the stored refresh_token on the list call, so the user never reauthorizes. + vi.mocked(getMCPOAuthUserCredentialStatus).mockResolvedValue( + credStatus({ has_credential: true, is_expired: true }), + ); + renderViewer({ oauth2_flow: null, delegate_auth_to_upstream: false }); await waitFor(() => expect(vi.mocked(listMCPTools)).toHaveBeenCalledWith("litellm-key", "srv-1", undefined)); @@ -87,5 +136,7 @@ describe("MCPToolsViewer auth gate routing", () => { await waitFor(() => expect(vi.mocked(listMCPTools)).toHaveBeenCalledWith("litellm-key", "srv-1", undefined)); expect(screen.queryByText(GATE_TEXT)).not.toBeInTheDocument(); + // M2M uses the backend service token, not a per-user DB credential. + expect(vi.mocked(getMCPOAuthUserCredentialStatus)).not.toHaveBeenCalled(); }); }); diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_tools.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_tools.tsx index 5ec9a3d683a..8957c7b30b8 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_tools.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_tools.tsx @@ -1,11 +1,14 @@ -import React, { useEffect, useState } from "react"; +import React, { useCallback, useEffect, useState } from "react"; import { useQuery, useMutation } from "@tanstack/react-query"; import { ToolTestPanel } from "./ToolTestPanel"; import { MCPTool, MCPToolsViewerProps, MCPContent, CallMCPToolResponse, getMcpOAuthMode } from "./types"; -import { listMCPTools, callMCPTool } from "../networking"; +import { listMCPTools, callMCPTool, getMCPOAuthUserCredentialStatus } from "../networking"; import { isTokenValid, getToken, removeToken } from "@/utils/mcpTokenStore"; import { sanitizeMcpAliasForHeader } from "@/utils/mcpHeaderUtils"; import { useToolsOAuthFlow } from "@/hooks/useToolsOAuthFlow"; +import { useUserMcpOAuthFlow } from "@/hooks/useUserMcpOAuthFlow"; +import { TOOLS_OAUTH_UI_STATE_KEY } from "@/hooks/mcpOAuthUtils"; +import { setSecureItem } from "@/utils/secureStorage"; import { Card, Title, Text } from "@tremor/react"; import { RobotOutlined, ToolOutlined, SearchOutlined, KeyOutlined, LockOutlined } from "@ant-design/icons"; @@ -31,11 +34,14 @@ const MCPToolsViewer = ({ const [passthroughHeaders, setPassthroughHeaders] = useState>({}); const [showHeaderInput, setShowHeaderInput] = useState(false); - // Only PKCE passthrough uses a browser-held session token (sessionStorage, - // cleared on tab/browser close) and a user-facing auth gate. OBO uses the - // backend-stored per-user token and M2M uses the backend's own service token, - // so neither needs a gate — they list tools with just the LiteLLM key. - const isPassthrough = getMcpOAuthMode({ auth_type, oauth2_flow, delegate_auth_to_upstream }) === "passthrough"; + // PKCE passthrough holds a browser-side session token (sessionStorage) and + // gates tool listing behind it. OBO uses a backend-stored per-user token that + // the user must establish once via an interactive login; we gate on whether + // that DB credential exists. M2M uses the backend's own service token and + // needs no gate. + const oauthMode = getMcpOAuthMode({ auth_type, oauth2_flow, delegate_auth_to_upstream }); + const isPassthrough = oauthMode === "passthrough"; + const isObo = oauthMode === "obo"; const [oauthToken, setOauthToken] = useState(() => isPassthrough && isTokenValid(serverId, userID) ? getToken(serverId, userID)?.access_token ?? null : null, ); @@ -61,6 +67,31 @@ const MCPToolsViewer = ({ onSuccess: setOauthToken, }); + // OBO servers list tools using a per-user token the backend stores in the DB; + // check whether the current user has a valid one so we can prompt them to + // authorize when they don't (otherwise the backend silently returns no tools). + const { + data: oboCredStatus, + isLoading: isLoadingOboCred, + isError: isOboCredError, + refetch: refetchOboCred, + } = useQuery({ + queryKey: ["mcpOauthUserCredStatus", serverId, userID], + queryFn: () => getMCPOAuthUserCredentialStatus(accessToken ?? "", serverId), + enabled: !!accessToken && isObo, + staleTime: 30000, + }); + + // A stored credential is sufficient: the backend proactively refreshes an + // expired or near-expiry token from the stored refresh_token on the next list + // call, so the user only needs to authorize when no credential row exists. If + // the status check itself fails we can't confirm a credential, so surface the + // Authorize gate rather than a silent empty tool list; re-authorizing only + // overwrites the user's own row, so it is safe when a credential did exist. + const hasOboCred = !!oboCredStatus?.has_credential; + const oboNeedsAuth = isObo && !isLoadingOboCred && (isOboCredError || (!!oboCredStatus && !hasOboCred)); + const oboStatusLoading = isObo && isLoadingOboCred; + // Check if this server has extra headers configured const hasExtraHeaders = extraHeaders && extraHeaders.length > 0; @@ -135,8 +166,9 @@ const MCPToolsViewer = ({ } return result; }, - // For OAuth servers, block the query until a session token is available - enabled: !!accessToken && (!isPassthrough || oauthToken !== null), + // Passthrough blocks until a browser session token exists; OBO blocks until + // the user has a valid DB credential (else the backend returns no tools). + enabled: !!accessToken && (isPassthrough ? oauthToken !== null : isObo ? hasOboCred : true), staleTime: 30000, // Consider data fresh for 30 seconds retry: (failureCount, error: any) => { // Don't retry on 401 — token is invalid, user must re-authenticate @@ -145,6 +177,33 @@ const MCPToolsViewer = ({ }, }); + // OBO authorize: same redirect+exchange flow as the admin "Authorize & Fetch" + // and the chat "Connect" button, but persists the token to the per-user DB. + const onOboAuthSuccess = useCallback(() => { + refetchOboCred(); + refetchTools(); + }, [refetchOboCred, refetchTools]); + + const { + startOAuthFlow: startDbOAuthFlow, + status: dbOAuthStatus, + error: dbOAuthError, + } = useUserMcpOAuthFlow({ + accessToken: accessToken ?? "", + serverId, + serverAlias, + onSuccess: onOboAuthSuccess, + }); + + // Stash which server started the redirect so the MCP Servers page can reopen + // this Tools tab on return and let the flow resume to persist the credential. + const startOboAuthorize = useCallback(() => { + try { + setSecureItem(TOOLS_OAUTH_UI_STATE_KEY, JSON.stringify({ serverId })); + } catch (_) {} + startDbOAuthFlow(); + }, [serverId, startDbOAuthFlow]); + // If the tools query fails with 401, the cached OAuth token is invalid — // clear it so the auth gate is shown again and the user can re-authenticate. useEffect(() => { @@ -187,6 +246,13 @@ const MCPToolsViewer = ({ const toolsData = mcpToolsResponse?.tools || []; + // An auth gate replaces the tool list when the user must authenticate first: + // passthrough needs a browser token, OBO needs a stored DB credential. + const authGateActive = (isPassthrough && !oauthToken) || oboNeedsAuth; + // Treat OBO credential-status loading as "tools loading" so the empty state + // doesn't flash before we know whether the user needs to authorize. + const toolsAreaLoading = isLoadingTools || oboStatusLoading; + // Filter tools based on search term const filteredTools = toolsData.filter((tool: MCPTool) => { const searchLower = toolSearchTerm.toLowerCase(); @@ -287,7 +353,7 @@ const MCPToolsViewer = ({ )} - {/* OAuth Auth Gate — shown when token is absent for OAuth servers */} + {/* Passthrough auth gate — browser session token absent */} {isPassthrough && !oauthToken && (
@@ -306,8 +372,31 @@ const MCPToolsViewer = ({
)} + {/* OBO auth gate — only when no credential row exists for this user. + An existing-but-expired token is refreshed server-side on the + list call, so the gate never appears for a stored credential. */} + {oboNeedsAuth && ( +
+ +

Authentication required

+

+ Authenticate with the upstream provider to view available tools +

+ + Authorize + + {dbOAuthError &&

{dbOAuthError}

} +
+ )} + {/* Search Bar — only shown when tools are loaded */} - {!isPassthrough || oauthToken ? ( + {!authGateActive ? ( <> {toolsData.length > 0 && (
@@ -324,7 +413,7 @@ const MCPToolsViewer = ({ )} {/* Loading State */} - {isLoadingTools && ( + {toolsAreaLoading && (
@@ -335,7 +424,7 @@ const MCPToolsViewer = ({ )} {/* Error State */} - {(mcpToolsResponse?.error || mcpToolsError) && !isLoadingTools && !toolsData.length && ( + {(mcpToolsResponse?.error || mcpToolsError) && !toolsAreaLoading && !toolsData.length && (

Error: {mcpToolsResponse?.message || (mcpToolsError as Error)?.message} @@ -344,7 +433,7 @@ const MCPToolsViewer = ({ )} {/* No Tools State */} - {!isLoadingTools && + {!toolsAreaLoading && !mcpToolsResponse?.error && !mcpToolsError && (!toolsData || toolsData.length === 0) && ( @@ -370,7 +459,7 @@ const MCPToolsViewer = ({ )} {/* Tools List */} - {!isLoadingTools && !mcpToolsResponse?.error && toolsData.length > 0 && ( + {!toolsAreaLoading && !mcpToolsResponse?.error && toolsData.length > 0 && ( <> {filteredTools.length === 0 ? (

diff --git a/ui/litellm-dashboard/src/hooks/mcpOAuthUtils.ts b/ui/litellm-dashboard/src/hooks/mcpOAuthUtils.ts index 3aff8af6eef..83b2cb21bed 100644 --- a/ui/litellm-dashboard/src/hooks/mcpOAuthUtils.ts +++ b/ui/litellm-dashboard/src/hooks/mcpOAuthUtils.ts @@ -7,6 +7,15 @@ import { getProxyBaseUrl, serverRootPath } from "@/components/networking"; +/** + * sessionStorage key used to restore the MCP server detail view on the Tools + * tab after a full-page OAuth redirect. The OBO authorize flow redirects to the + * IdP and back to the MCP Servers page; without this the user lands on the + * server list and useUserMcpOAuthFlow never re-mounts to persist the credential. + * Mirrors the admin edit flow's EDIT_OAUTH_UI_STATE_KEY. + */ +export const TOOLS_OAUTH_UI_STATE_KEY = "litellm-mcp-oauth-tools-state"; + /** * Build the OAuth callback URL for the current UI deployment. * From 32c88ca74f29143aadc4cfaa708f89ce1086386c Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 9 Jun 2026 02:19:52 +0530 Subject: [PATCH 018/185] Litellm oss staging 080626 (#29932) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(bedrock_mantle): add SigV4/IAM auth to Responses API route (fixes #29665) (#29788) * feat(responses): add default no-op sign_request to BaseResponsesAPIConfig * feat(responses): call sign_request after body is final, send signed bytes when signed * feat(bedrock_mantle): add SigV4 sign_request via composed BaseAWSLLM (bearer path) * test(bedrock_mantle): cover SigV4 access-key, AssumeRole, body bytes, region/auth consistency * feat(bedrock_mantle): defer auth to sign_request; validate_environment no longer requires bearer * docs(bedrock_mantle): document SigV4 + Bearer auth on Responses route * test(responses): cover fake-stream signing order and mantle bearer arg/env precedence * fix(bedrock_mantle): wrap all botocore credential errors with both-paths guidance * fix(bedrock_mantle): catch specific credential errors, not all BotoCoreError, so STS transport failures are not masked * fix(bedrock_mantle): sign the compact Responses route too, not just create * fix(github-copilot): route per-model on /v1/responses based on model info (#29747) * feat(focus): add GCS destination for FOCUS export (#29751) * test: add failing tests for FocusGCSDestination * feat: add FocusGCSDestination reusing GCSBucketBase auth * feat: register FocusGCSDestination in factory; export from __init__ * fix(focus): preserve GCS_PATH_SERVICE_ACCOUNT when service_account_json not in config * style: apply Black formatting to gcs_destination and tests * style: apply Black formatting to factory.py * fix(bedrock): omit empty additionalModelRequestFields and system from Converse API payload (#29565) Amazon Nova Pro (and other strict Bedrock models) return 400 Malformed input request when additionalModelRequestFields: {} or system: [] are present in the payload. Both fields are optional in CommonRequestObject (total=False) and must be omitted rather than sent as empty structures. Co-authored-by: shin-berri Co-authored-by: yuneng-jiang Co-authored-by: Claude Sonnet 4.6 * fix(proxy): recognize *.cognitiveservices.azure.com as OpenAI-compatible in pass-through cost tracking (#29730) * fix(proxy): recognize *.cognitiveservices.azure.com as OpenAI-compatible Azure OpenAI resources created via the newer "Azure AI Foundry" / Cognitive Services pathway live on `*.cognitiveservices.azure.com` subdomains, not the older `openai.azure.com`. Both are valid Azure OpenAI surfaces in production today. The OpenAI pass-through cost-tracking handler hard-codes only the older hostname in five places (four `is_openai_*_route` methods on OpenAIPassthroughLoggingHandler, plus is_openai_route on PassThroughEndpointLogging). As a result, calls from newer Azure deployments are silently classified as "not an OpenAI route", the dispatch into the cost-tracking handler is skipped, and tokens/cost never get extracted into LiteLLM_SpendLogs — the row gets written with prompt_tokens=0, completion_tokens=0, spend=0, model='unknown'. Reproduced 2026-06-04 against a real Azure OpenAI deployment on `*.cognitiveservices.azure.com` proxied through LiteLLM v1.88.0. Fix: factor the hostname check into a single helper `_is_openai_compatible_host` listing all three recognized surfaces (api.openai.com, openai.azure.com, cognitiveservices.azure.com), and have all five call sites delegate to it. Purely additive — never weakens recognition for the originally-supported hostnames. Adds a test `test_is_openai_route_recognizes_cognitiveservices_azure_com` that exercises all four `is_openai_*_route` static methods against `*.cognitiveservices.azure.com` URLs (positive cases per route + a small cross-route negative to confirm route-specific path matching still works on the new hostname). Out of scope for this PR (separate followup): - `openai_passthrough_handler` calls chat/completions `transform_response` on Responses API payloads (`output:` not `choices:`), which throws inside the dispatch and drops the SpendLogs row entirely. Recognized + tracked separately. * ci: trigger fresh run Empty commit to re-run checks. The previous auth-and-jwt failure was a transient HuggingFace Hub 429 rate-limit hitting tokenizer downloads in tests/proxy_unit_tests/test_custom_tokenizer_bug.py — unrelated to this PR's scope (hostname recognition in pass-through cost tracking). No code change. --------- Co-authored-by: shin-berri Co-authored-by: yuneng-jiang * fix(responses): preserve forced-function tool_choice name in Responses to Chat transform (#29812) The Responses API forces a specific function with a top-level name ({"type": "function", "name": "X"}), but _transform_tool_choice only handled the nested Chat Completions shape and fell through to returning "required" for the flat form, silently dropping the function name and degrading a forced function call to force-any-tool. Map the flat Responses shape to the nested Chat shape, keeping the "required" fallback when no name is present. * Preserve x-anthropic-billing-header system blocks for first-party Anthropic (#29584) * Preserve x-anthropic-billing-header system blocks for first-party Anthropic PR #20951 strips system blocks beginning with "x-anthropic-billing-header:" for every Anthropic target. That block is how the first-party Anthropic API recognizes Claude Code subscription (OAuth) traffic, so dropping it makes requests that carry only that block, such as the auto-mode tool-safety classifier, fail with a misleading 429 rate_limit_error; normal turns still work because they also carry the "You are Claude Code" identity block. Gate the strip behind should_strip_billing_metadata(), defaulting to False on the first-party AnthropicConfig and AnthropicMessagesConfig so the block is kept, and overridden to True on the providers that reach these transforms and reject the block (Bedrock platform, Vertex, Azure for the chat path; Minimax, Azure, DeepSeek for the messages path). Behavior for those providers is unchanged. * Strip billing header on Bedrock invoke and Vertex messages pass-through Two more subclasses reach the gated strip but inherited keep-by-default. AmazonAnthropicClaudeConfig (Bedrock invoke) calls AnthropicConfig.transform_request, which calls translate_system_message, and VertexAIPartnerModelsAnthropicMessagesConfig (Vertex messages pass-through) calls super().transform_anthropic_messages_request. Override should_strip_billing_metadata() to True on both. Add a parametrized test asserting the flag for every first-party base (False) and provider subclass (True), covering all overrides, plus a translate_system_message regression test for the Bedrock invoke path. * fix(cache): log hashed cache keys (#29890) * fix(ui): save routing groups as list (#29889) * Revert "fix(ui): save routing groups as list (#29889)" (#29928) This reverts commit 9b1f78ffa7a309cabe5e9a7ab5f94d1224d192c9. * feat(parasail): add Parasail as a JSON-configured OpenAI-compatible provider (#29842) * feat(parasail): add Parasail as a JSON-configured OpenAI-compatible provider Registers parasail in the openai_like JSON provider loader with both /v1/chat/completions and /v1/responses support. Parasail's Responses API rejects store:true and any request that omits store, so the loader gains a force_store_false special_handling flag; the parasail entry sets it and the generated Responses config overrides store=false on every call. This keeps callers from hitting "State storage not supported" and matches what Parasail's docs require. Adds the PARASAIL enum value, listing under openai_compatible_providers, provider documentation at docs/my-website/docs/providers/parasail.md, and a focused unit test file under tests/test_litellm/llms/parasail/ that covers JSON registration, chat URL construction, Responses URL construction with PARASAIL_API_BASE override, and the force_store_false regression in both the caller-sent-store=true and caller-omitted cases. * fix(parasail): register in provider_endpoints_support, drop in-repo docs Greptile review feedback. The provider doc belongs in the litellm-docs repo, not this one's docs/my-website tree; removing it here. Adds the parasail entry to provider_endpoints_support.json so the check_provider_folders_documented.py CI check passes (chat_completions and responses true; others false). * fix: normalize Anthropic passthrough server tool usage (#29827) * test(anthropic): cover server_tool_use dict cost tracking * fix: normalize Anthropic server tool usage (cherry picked from commit 982f726bed7d3ec05e463c5dd3d090bebae91d19) * fix: keep server tool usage subscriptable (cherry picked from commit 70280b9b272455b2f974d08bc697f67f929755bf) --------- Co-authored-by: Genmin * fix(proxy): fix typo generic_role_mappoings -> generic_role_mappings in ui_sso.py (#29753) Co-authored-by: shin-berri Co-authored-by: yuneng-jiang * feat(proxy): add disable_budget_reservation general setting (#27639) (#29493) * feat(proxy): add disable_budget_reservation general setting (#27639) * feat(proxy): register disable_budget_reservation in ConfigGeneralSettings (#27639) * docs(proxy): document disable_budget_reservation concurrency tradeoff (#27639) * ci: re-trigger flaky docker build (prisma generate ECONNRESET) * fix(proxy): warn and document budget enforcement tradeoff when disable_budget_reservation is set (#27639) * feat(gemini_tts): adding support to Gemini TTS languageCode parameters (#29623) * Adding support to Gemini TTS Language Code parameters * Mapping Gemini TTS languageCode param in Docstring * Use snake_case for language_code input keyMapping Gemini TTS languageCode param in Docstring * Restoring files modified under enterprise/litellm_enterprise due to lint/formatting checks --------- Co-authored-by: João Garrido * feat(guardrails): capture user and model metadata in CrowdStrike AIDR (#29517) * fix(proxy): require OpenAI path segment for shared Azure Cognitive Services domains Address Greptile review: the `*.cognitiveservices.azure.com` / `*.openai.azure.com` domains are shared by every Azure Cognitive Service (Speech, Vision, Language, ...), so a hostname-only substring match misclassified non-OpenAI Azure traffic as OpenAI routes. - Replace the substring host test with suffix matching (rejects look-alike domains like cognitiveservices.azure.com.attacker.example). - Add `_is_openai_compatible_url` that requires an OpenAI-style path marker (`/openai/` or `/v1/`) on the shared Azure domains, and use it in PassThroughEndpointLogging.is_openai_route (previously hostname-only). - Add negative tests for Azure Speech/Vision paths and look-alike domains. Co-Authored-By: Claude Opus 4.8 * fix: support Responses input in Redis semantic cache (#29581) * fix: support responses input in redis semantic cache * test: cover redis semantic prompt extraction * test: handle blank redis semantic text fallbacks * chore: remove async cache dead statement * test: cover redis semantic cache miss paths * fix: filter sensitive cache lookup kwargs * chore: rerun ci after huggingface rate limit * chore(ui): regenerate dashboard API types (npm run gen:api) Sync src/lib/http/schema.d.ts with the proxy OpenAPI spec: adds the disable_budget_reservation general-settings field and picks up the RateLimitError docstring reindent. Fixes the gen:api CI drift check. Co-Authored-By: Claude Opus 4.8 * test(bedrock): assert empty additionalModelRequestFields is omitted The Converse transformer now drops an empty additionalModelRequestFields block instead of sending it as `{}`. Update test_bedrock_top_k_param so models without top_k support (llama3) assert the key is absent rather than equal to an empty dict. Co-Authored-By: Claude Opus 4.8 --------- Co-authored-by: Kent <72616338+kingdoooo@users.noreply.github.com> Co-authored-by: codgician <15964984+codgician@users.noreply.github.com> Co-authored-by: Praveen Ghuge <95286176+pghuge-cloudwiz@users.noreply.github.com> Co-authored-by: Roi Co-authored-by: shin-berri Co-authored-by: yuneng-jiang Co-authored-by: Claude Sonnet 4.6 Co-authored-by: Liam Scott Co-authored-by: abhay23-AI Co-authored-by: Ceder Dens Co-authored-by: 冯基魁 <56265583+fengjikui@users.noreply.github.com> Co-authored-by: Kai Huang Co-authored-by: rinto <54238243+ririnto@users.noreply.github.com> Co-authored-by: Genmin Co-authored-by: Arnav Bhilwariya Co-authored-by: Armaan Sandhu <74664101+Ar-maan05@users.noreply.github.com> Co-authored-by: João Garrido <48538534+johngarrido@users.noreply.github.com> Co-authored-by: João Garrido Co-authored-by: Kenan Yildirim Co-authored-by: Dávid Balatoni --- litellm/caching/caching.py | 47 +- litellm/caching/redis_semantic_cache.py | 105 +++- litellm/constants.py | 1 + .../focus/destinations/__init__.py | 2 + .../focus/destinations/factory.py | 15 + .../focus/destinations/gcs_destination.py | 74 +++ litellm/llms/anthropic/chat/transformation.py | 22 +- .../messages/transformation.py | 13 +- .../anthropic/messages_transformation.py | 3 + .../llms/azure_ai/anthropic/transformation.py | 3 + .../llms/base_llm/responses/transformation.py | 20 + .../bedrock/chat/converse_transformation.py | 6 +- .../anthropic_claude3_transformation.py | 3 + .../bedrock/claude_platform/transformation.py | 3 + .../responses/transformation.py | 120 ++++- litellm/llms/custom_httpx/llm_http_handler.py | 105 +++- .../llms/deepseek/messages/transformation.py | 3 + .../responses/transformation.py | 45 +- .../llms/minimax/messages/transformation.py | 3 + litellm/llms/openai_like/dynamic_config.py | 19 + litellm/llms/openai_like/providers.json | 9 + .../vertex_and_google_ai_studio_gemini.py | 7 +- .../transformation.py | 3 + .../anthropic/transformation.py | 3 + litellm/proxy/_types.py | 18 + litellm/proxy/auth/user_api_key_auth.py | 12 + .../crowdstrike_aidr/crowdstrike_aidr.py | 18 +- litellm/proxy/management_endpoints/ui_sso.py | 4 +- .../openai_passthrough_logging_handler.py | 98 +++- .../pass_through_endpoints/success_handler.py | 15 +- .../transformation.py | 4 +- litellm/types/llms/vertex_ai.py | 1 + litellm/types/utils.py | 11 +- litellm/utils.py | 8 +- provider_endpoints_support.json | 17 + .../test_bedrock_completion.py | 6 +- tests/test_litellm/caching/test_caching.py | 48 ++ .../caching/test_redis_semantic_cache.py | 465 ++++++++++++++++++ .../focus/test_focus_gcs_destination.py | 180 +++++++ .../test_tool_call_cost_tracking.py | 21 +- .../test_anthropic_chat_transformation.py | 169 +++++++ .../chat/test_converse_transformation.py | 7 +- .../test_bedrock_files_transformation.py | 64 +++ ...bedrock_mantle_responses_transformation.py | 398 ++++++++++++++- .../custom_httpx/test_llm_http_handler.py | 238 +++++++++ .../llms/gemini/test_gemini_tts.py | 92 ++++ ...github_copilot_responses_transformation.py | 216 +++++++- .../llms/parasail/test_parasail.py | 172 +++++++ .../proxy/auth/test_user_api_key_auth.py | 60 +++ .../guardrail_hooks/test_crowdstrike_aidr.py | 121 ++++- ...t_anthropic_passthrough_logging_handler.py | 38 ++ ...test_openai_passthrough_logging_handler.py | 88 ++++ .../test_litellm_completion_responses.py | 22 + tests/test_litellm/types/test_types_utils.py | 46 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 + 55 files changed, 3154 insertions(+), 142 deletions(-) create mode 100644 litellm/integrations/focus/destinations/gcs_destination.py create mode 100644 tests/test_litellm/caching/test_caching.py create mode 100644 tests/test_litellm/integrations/focus/test_focus_gcs_destination.py create mode 100644 tests/test_litellm/llms/parasail/test_parasail.py diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 11733ce4cee..c1afde16250 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -309,9 +309,13 @@ class Cache: param_value = kwargs[param] cache_key += f"{str(param)}: {str(param_value)}" - verbose_logger.debug("\nCreated cache key: %s", cache_key) hashed_cache_key = Cache._get_hashed_cache_key(cache_key) hashed_cache_key = self._add_namespace_to_cache_key(hashed_cache_key, **kwargs) + verbose_logger.debug( + "\nCreated cache key: %s (source material length: %d)", + hashed_cache_key, + len(cache_key), + ) # Remove preset_cache_key from kwargs to avoid "got multiple values" TypeError # when kwargs already contains preset_cache_key from upstream callers kwargs_for_preset = {k: v for k, v in kwargs.items() if k != "preset_cache_key"} @@ -497,6 +501,34 @@ class Cache: return cached_response return cached_result + @staticmethod + def _get_safe_cache_lookup_kwargs(kwargs: Dict[str, Any]) -> Dict[str, Any]: + cache_lookup_kwargs: Dict[str, Any] = {} + for prompt_kwarg in ("messages", "input"): + if prompt_kwarg in kwargs: + cache_lookup_kwargs[prompt_kwarg] = kwargs[prompt_kwarg] + + if isinstance(kwargs.get("metadata"), dict): + cache_lookup_kwargs["metadata"] = {} + + return cache_lookup_kwargs + + @staticmethod + def _update_metadata_from_cache_lookup_kwargs( + original_kwargs: Dict[str, Any], cache_lookup_kwargs: Dict[str, Any] + ) -> None: + original_metadata = original_kwargs.get("metadata") + cache_lookup_metadata = cache_lookup_kwargs.get("metadata") + if not isinstance(original_metadata, dict) or not isinstance( + cache_lookup_metadata, dict + ): + return + + if "semantic-similarity" in cache_lookup_metadata: + original_metadata["semantic-similarity"] = cache_lookup_metadata[ + "semantic-similarity" + ] + def get_cache(self, dynamic_cache_object: Optional[BaseCache] = None, **kwargs): """ Retrieves the cached result for the given arguments. @@ -511,7 +543,6 @@ class Cache: try: # never block execution if self.should_use_cache(**kwargs) is not True: return - messages = kwargs.get("messages", []) if "cache_key" in kwargs: cache_key = kwargs["cache_key"] else: @@ -523,12 +554,19 @@ class Cache: or cache_control_args.get("s-max-age") or float("inf") ) + cache_lookup_kwargs = self._get_safe_cache_lookup_kwargs(kwargs) if dynamic_cache_object is not None: cached_result = dynamic_cache_object.get_cache( - cache_key, messages=messages + cache_key, **cache_lookup_kwargs ) else: - cached_result = self.cache.get_cache(cache_key, messages=messages) + cached_result = self.cache.get_cache( + cache_key, **cache_lookup_kwargs + ) + self._update_metadata_from_cache_lookup_kwargs( + original_kwargs=kwargs, + cache_lookup_kwargs=cache_lookup_kwargs, + ) return self._get_cache_logic( cached_result=cached_result, max_age=max_age ) @@ -549,7 +587,6 @@ class Cache: if self.should_use_cache(**kwargs) is not True: return - kwargs.get("messages", []) if "cache_key" in kwargs: cache_key = kwargs["cache_key"] else: diff --git a/litellm/caching/redis_semantic_cache.py b/litellm/caching/redis_semantic_cache.py index da9e7b1e587..cce4b75795f 100644 --- a/litellm/caching/redis_semantic_cache.py +++ b/litellm/caching/redis_semantic_cache.py @@ -213,6 +213,78 @@ class RedisSemanticCache(BaseCache): ttl = int(ttl) return ttl + @classmethod + def _get_prompt_from_kwargs(cls, **kwargs) -> Optional[str]: + """ + Extract a semantic-cache prompt from chat or Responses API request kwargs. + """ + messages = kwargs.get("messages") + if messages: + return get_str_from_messages(messages) + + if "input" not in kwargs: + return None + + prompt_parts: List[str] = [] + cls._collect_responses_input_text(kwargs.get("input"), prompt_parts) + prompt = "\n".join(prompt_parts).strip() + return prompt or None + + @classmethod + def _collect_responses_input_text(cls, value: Any, prompt_parts: List[str]) -> None: + value = cls._coerce_response_input_value(value) + if value is None: + return + + if isinstance(value, str): + stripped_value = value.strip() + if stripped_value: + prompt_parts.append(stripped_value) + return + + if isinstance(value, (list, tuple)): + for item in value: + cls._collect_responses_input_text(item, prompt_parts) + return + + if isinstance(value, dict): + content = value.get("content") + if content is not None: + cls._collect_responses_input_text(content, prompt_parts) + return + + for text_key in ("text", "output", "input_text", "output_text"): + text_value = value.get(text_key) + if isinstance(text_value, str): + stripped_text = text_value.strip() + if stripped_text: + prompt_parts.append(stripped_text) + return + return + + content = getattr(value, "content", None) + if content is not None: + cls._collect_responses_input_text(content, prompt_parts) + return + + for text_key in ("text", "output", "input_text", "output_text"): + text_value = getattr(value, text_key, None) + if isinstance(text_value, str): + stripped_text = text_value.strip() + if stripped_text: + prompt_parts.append(stripped_text) + return + + @staticmethod + def _coerce_response_input_value(value: Any) -> Any: + model_dump = getattr(value, "model_dump", None) + if callable(model_dump): + return model_dump() + dict_method = getattr(value, "dict", None) + if callable(dict_method): + return dict_method() + return value + def _get_embedding(self, prompt: str) -> List[float]: """ Generate an embedding vector for the given prompt using the configured embedding model. @@ -278,13 +350,11 @@ class RedisSemanticCache(BaseCache): value_str: Optional[str] = None try: - # Extract the prompt from messages - messages = kwargs.get("messages", []) - if not messages: - print_verbose("No messages provided for semantic caching") + prompt = self._get_prompt_from_kwargs(**kwargs) + if prompt is None: + print_verbose("No prompt provided for semantic caching") return - prompt = get_str_from_messages(messages) value_str = str(value) store_kwargs: Dict[str, Any] = { @@ -315,14 +385,12 @@ class RedisSemanticCache(BaseCache): print_verbose(f"Redis semantic-cache get_cache, kwargs: {kwargs}") try: - # Extract the prompt from messages - messages = kwargs.get("messages", []) - if not messages: - print_verbose("No messages provided for semantic cache lookup") + prompt = self._get_prompt_from_kwargs(**kwargs) + if prompt is None: + print_verbose("No prompt provided for semantic cache lookup") kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 return None - prompt = get_str_from_messages(messages) # Check the cache for semantically similar prompts in this exact # LiteLLM cache-key scope. check_kwargs: Dict[str, Any] = { @@ -428,13 +496,11 @@ class RedisSemanticCache(BaseCache): print_verbose(f"Async Redis semantic-cache set_cache, kwargs: {kwargs}") try: - # Extract the prompt from messages - messages = kwargs.get("messages", []) - if not messages: - print_verbose("No messages provided for semantic caching") + prompt = self._get_prompt_from_kwargs(**kwargs) + if prompt is None: + print_verbose("No prompt provided for semantic caching") return - prompt = get_str_from_messages(messages) value_str = str(value) # Generate embedding for the value (response) to cache @@ -471,15 +537,12 @@ class RedisSemanticCache(BaseCache): print_verbose(f"Async Redis semantic-cache get_cache, kwargs: {kwargs}") try: - # Extract the prompt from messages - messages = kwargs.get("messages", []) - if not messages: - print_verbose("No messages provided for semantic cache lookup") + prompt = self._get_prompt_from_kwargs(**kwargs) + if prompt is None: + print_verbose("No prompt provided for semantic cache lookup") kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 return None - prompt = get_str_from_messages(messages) - # Generate embedding for the prompt prompt_embedding = await self._get_async_embedding(prompt, **kwargs) diff --git a/litellm/constants.py b/litellm/constants.py index 36e578bd323..f10cec034f0 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -831,6 +831,7 @@ openai_compatible_providers: List = [ "nano-gpt", # Nano-GPT - JSON-configured provider "poe", # Poe - JSON-configured provider "chutes", # Chutes - JSON-configured provider + "parasail", # Parasail - JSON-configured provider "featherless_ai", "nscale", "nebius", diff --git a/litellm/integrations/focus/destinations/__init__.py b/litellm/integrations/focus/destinations/__init__.py index 775d3a259d2..e0cd90c1d61 100644 --- a/litellm/integrations/focus/destinations/__init__.py +++ b/litellm/integrations/focus/destinations/__init__.py @@ -2,12 +2,14 @@ from .base import FocusDestination, FocusTimeWindow from .factory import FocusDestinationFactory +from .gcs_destination import FocusGCSDestination from .s3_destination import FocusS3Destination from .vantage_destination import FocusVantageDestination __all__ = [ "FocusDestination", "FocusDestinationFactory", + "FocusGCSDestination", "FocusTimeWindow", "FocusS3Destination", "FocusVantageDestination", diff --git a/litellm/integrations/focus/destinations/factory.py b/litellm/integrations/focus/destinations/factory.py index 706e10624ce..7ce21d4040a 100644 --- a/litellm/integrations/focus/destinations/factory.py +++ b/litellm/integrations/focus/destinations/factory.py @@ -6,6 +6,7 @@ import os from typing import Any, Dict, Optional from .base import FocusDestination +from .gcs_destination import FocusGCSDestination from .s3_destination import FocusS3Destination from .vantage_destination import FocusVantageDestination @@ -29,6 +30,8 @@ class FocusDestinationFactory: return FocusS3Destination(prefix=prefix, config=normalized_config) if provider_lower == "vantage": return FocusVantageDestination(prefix=prefix, config=normalized_config) + if provider_lower == "gcs": + return FocusGCSDestination(prefix=prefix, config=normalized_config) raise NotImplementedError( f"Provider '{provider}' not supported for Focus export" ) @@ -72,6 +75,18 @@ class FocusDestinationFactory: "VANTAGE_INTEGRATION_TOKEN must be provided for Vantage exports" ) return {k: v for k, v in resolved.items() if v is not None} + if provider == "gcs": + resolved = { + "bucket_name": overrides.get("bucket_name") + or os.getenv("FOCUS_GCS_BUCKET_NAME"), + "service_account_json": overrides.get("service_account_json") + or os.getenv("FOCUS_GCS_PATH_SERVICE_ACCOUNT"), + } + if not resolved.get("bucket_name"): + raise ValueError( + "FOCUS_GCS_BUCKET_NAME must be provided for GCS exports" + ) + return {k: v for k, v in resolved.items() if v is not None} raise NotImplementedError( f"Provider '{provider}' not supported for Focus export configuration" ) diff --git a/litellm/integrations/focus/destinations/gcs_destination.py b/litellm/integrations/focus/destinations/gcs_destination.py new file mode 100644 index 00000000000..b04c16c9d32 --- /dev/null +++ b/litellm/integrations/focus/destinations/gcs_destination.py @@ -0,0 +1,74 @@ +"""GCS destination for Focus export — reuses GCSBucketBase auth and httpx client.""" + +from __future__ import annotations + +from datetime import timezone +from typing import Any, Optional + +from litellm._logging import verbose_logger +from litellm.integrations.gcs_bucket.gcs_bucket_base import GCSBucketBase +from litellm.litellm_core_utils.cloud_storage_security import ( + encode_gcs_object_name_for_url, +) + +from .base import FocusDestination, FocusTimeWindow + + +class FocusGCSDestination(GCSBucketBase, FocusDestination): + """Upload serialized Focus exports to GCS using the GCS JSON API.""" + + def __init__( + self, + *, + prefix: str, + config: Optional[dict[str, Any]] = None, + ) -> None: + config = config or {} + bucket_name = config.get("bucket_name") + if not bucket_name: + raise ValueError("bucket_name must be provided for GCS destination") + super().__init__(bucket_name=bucket_name) + service_account_json = config.get("service_account_json") + if service_account_json is not None: + self.path_service_account_json = service_account_json + self.prefix = prefix.rstrip("/") + + async def deliver( + self, + *, + content: bytes, + time_window: FocusTimeWindow, + filename: str, + ) -> None: + object_name = self._build_object_key(time_window=time_window, filename=filename) + headers = await self.construct_request_headers( + service_account_json=self.path_service_account_json + ) + headers["Content-Type"] = "application/octet-stream" + encoded_name = encode_gcs_object_name_for_url(object_name) + url = ( + f"https://storage.googleapis.com/upload/storage/v1/b/" + f"{self.BUCKET_NAME}/o?uploadType=media&name={encoded_name}" + ) + response = await self.async_httpx_client.post( + url=url, headers=headers, data=content + ) + if response.status_code != 200: + raise RuntimeError( + f"GCS upload failed: status={response.status_code} body={response.text}" + ) + verbose_logger.debug( + "Focus GCS: uploaded %d bytes to gs://%s/%s", + len(content), + self.BUCKET_NAME, + object_name, + ) + + def _build_object_key(self, *, time_window: FocusTimeWindow, filename: str) -> str: + start_utc = time_window.start_time.astimezone(timezone.utc) + date_component = f"date={start_utc.strftime('%Y-%m-%d')}" + parts = [self.prefix, date_component] + if time_window.frequency == "hourly": + parts.append(f"hour={start_utc.strftime('%H')}") + key_prefix = "/".join(filter(None, parts)) + return f"{key_prefix}/{filename}" if key_prefix else filename diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 7949e150c23..3f30d5d6807 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -1607,6 +1607,15 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ) return _tool + def should_strip_billing_metadata(self) -> bool: + """ + Whether to drop x-anthropic-billing-header system blocks before sending upstream. + + The first-party Anthropic API uses these blocks for Claude Code attribution, so the + base config keeps them. Providers that reject them (e.g. Bedrock) override this to True. + """ + return False + def translate_system_message( self, messages: List[AllMessageValues] ) -> List[AnthropicSystemMessageContent]: @@ -1614,7 +1623,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): Translate system message to anthropic format. Removes system message from the original list and returns a new list of anthropic system message content. - Filters out system messages containing x-anthropic-billing-header metadata. + When should_strip_billing_metadata() is True, x-anthropic-billing-header system blocks are dropped. """ system_prompt_indices = [] anthropic_system_message_list: List[AnthropicSystemMessageContent] = [] @@ -1626,10 +1635,9 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): # Skip empty text blocks - Anthropic API raises errors for empty text if not system_message_block["content"]: continue - # Skip system messages containing x-anthropic-billing-header metadata - if system_message_block["content"].startswith( - "x-anthropic-billing-header:" - ): + if self.should_strip_billing_metadata() and system_message_block[ + "content" + ].startswith("x-anthropic-billing-header:"): continue anthropic_system_message_content = AnthropicSystemMessageContent( type="text", @@ -1648,9 +1656,9 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): text_value = _content.get("text") if _content.get("type") == "text" and not text_value: continue - # Skip system messages containing x-anthropic-billing-header metadata if ( - _content.get("type") == "text" + self.should_strip_billing_metadata() + and _content.get("type") == "text" and text_value and text_value.startswith("x-anthropic-billing-header:") ): diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py index 3a2c09f2183..07e8270b496 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py @@ -84,6 +84,15 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): if isinstance(content, list): _process_content_list(content) + def should_strip_billing_metadata(self) -> bool: + """ + Whether to drop x-anthropic-billing-header system blocks before sending upstream. + + The first-party Anthropic API uses these blocks for Claude Code attribution, so the + base config keeps them. Providers that reject them override this to True. + """ + return False + @staticmethod def _filter_billing_headers_from_system(system_param): """ @@ -286,14 +295,12 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): optional_params=anthropic_messages_optional_request_params, ) - # Filter out x-anthropic-billing-header from system messages system_param = anthropic_messages_optional_request_params.get("system") - if system_param is not None: + if self.should_strip_billing_metadata() and system_param is not None: filtered_system = self._filter_billing_headers_from_system(system_param) if filtered_system is not None and len(filtered_system) > 0: anthropic_messages_optional_request_params["system"] = filtered_system else: - # Remove system parameter if all content was filtered out anthropic_messages_optional_request_params.pop("system", None) # Transform context_management from OpenAI format to Anthropic format if needed diff --git a/litellm/llms/azure_ai/anthropic/messages_transformation.py b/litellm/llms/azure_ai/anthropic/messages_transformation.py index a81218ab76a..59b6ee2b424 100644 --- a/litellm/llms/azure_ai/anthropic/messages_transformation.py +++ b/litellm/llms/azure_ai/anthropic/messages_transformation.py @@ -21,6 +21,9 @@ class AzureAnthropicMessagesConfig(AnthropicMessagesConfig): and Azure endpoint format. """ + def should_strip_billing_metadata(self) -> bool: + return True + def validate_anthropic_messages_environment( self, headers: dict, diff --git a/litellm/llms/azure_ai/anthropic/transformation.py b/litellm/llms/azure_ai/anthropic/transformation.py index e176a4d860e..367ca75c196 100644 --- a/litellm/llms/azure_ai/anthropic/transformation.py +++ b/litellm/llms/azure_ai/anthropic/transformation.py @@ -40,6 +40,9 @@ class AzureAnthropicConfig(AnthropicConfig): def custom_llm_provider(self) -> Optional[str]: return "azure_ai" + def should_strip_billing_metadata(self) -> bool: + return True + def validate_environment( self, headers: dict, diff --git a/litellm/llms/base_llm/responses/transformation.py b/litellm/llms/base_llm/responses/transformation.py index 853eb282758..407d5ad8146 100644 --- a/litellm/llms/base_llm/responses/transformation.py +++ b/litellm/llms/base_llm/responses/transformation.py @@ -62,6 +62,26 @@ class BaseResponsesAPIConfig(ABC): """ return False + def sign_request( + self, + headers: dict, + optional_params: dict, + request_data: dict, + api_base: str, + api_key: Optional[str] = None, + model: Optional[str] = None, + stream: Optional[bool] = None, + fake_stream: Optional[bool] = None, + ) -> Tuple[dict, Optional[bytes]]: + """Sign the request after the body is finalized. + + Default is a no-op (returns headers unchanged, no signed body). Providers + whose endpoint requires request signing (e.g. Bedrock Mantle SigV4) + override this and return the signed body bytes so the handler sends those + exact bytes. + """ + return headers, None + @abstractmethod def get_supported_openai_params(self, model: str) -> list: pass diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 90dfa13e938..ea0326dffd1 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -1649,12 +1649,14 @@ class AmazonConverseConfig(BaseConfig): bedrock_tool_config["toolChoice"] = tool_choice_values data: CommonRequestObject = { - "additionalModelRequestFields": additional_request_params, - "system": system_content_blocks, "inferenceConfig": self._transform_inference_params( inference_params=inference_params ), } + if additional_request_params: + data["additionalModelRequestFields"] = additional_request_params + if system_content_blocks: + data["system"] = system_content_blocks # Handle all config blocks for config_name, config_class in self.get_config_blocks().items(): diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py index a13336b6c88..4887cbd23be 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py @@ -60,6 +60,9 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): def custom_llm_provider(self) -> Optional[str]: return "bedrock" + def should_strip_billing_metadata(self) -> bool: + return True + def get_supported_openai_params(self, model: str) -> List[str]: return AnthropicConfig.get_supported_openai_params(self, model) diff --git a/litellm/llms/bedrock/claude_platform/transformation.py b/litellm/llms/bedrock/claude_platform/transformation.py index 0167c457c96..c20dc63444f 100644 --- a/litellm/llms/bedrock/claude_platform/transformation.py +++ b/litellm/llms/bedrock/claude_platform/transformation.py @@ -17,6 +17,9 @@ class BedrockClaudePlatformConfig(BedrockClaudePlatformMixin, AnthropicConfig): def custom_llm_provider(self) -> Optional[str]: return "bedrock" + def should_strip_billing_metadata(self) -> bool: + return True + def validate_environment( self, headers: dict, diff --git a/litellm/llms/bedrock_mantle/responses/transformation.py b/litellm/llms/bedrock_mantle/responses/transformation.py index b63fd0ecdb1..df219091074 100644 --- a/litellm/llms/bedrock_mantle/responses/transformation.py +++ b/litellm/llms/bedrock_mantle/responses/transformation.py @@ -4,14 +4,26 @@ Amazon Bedrock Mantle - Responses API backend. gpt-5.5 / gpt-5.4 on Mantle are exposed ONLY on the `/openai/v1/responses` path (not the standard `/v1/responses`). Payloads and SSE follow the OpenAI Responses spec, so this config inherits OpenAIResponsesAPIConfig and overrides -only the endpoint URL and Bearer authentication. +only the endpoint URL and authentication. -Auth: AWS Bedrock API key as Bearer token (BEDROCK_MANTLE_API_KEY or the -standard AWS_BEARER_TOKEN_BEDROCK), NOT SigV4. +Auth: Bearer token (BEDROCK_MANTLE_API_KEY or the standard +AWS_BEARER_TOKEN_BEDROCK, or litellm_params.api_key) when present; otherwise +AWS SigV4 (service name "bedrock") using the standard credential chain (IAM +role / access key / profile / web identity), signed via the shared +BaseAWSLLM._sign_request after the request body is finalized. """ -from typing import Optional +import re +from typing import Optional, Tuple +from botocore.exceptions import ( + CredentialRetrievalError, + NoCredentialsError, + PartialCredentialsError, + ProfileNotFound, +) + +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.secret_managers.main import get_secret_str from litellm.types.router import GenericLiteLLMParams @@ -29,22 +41,44 @@ _BASE_SUFFIXES_TO_STRIP = ( "/v1", ) +# Standard Mantle host: https://bedrock-mantle..api.aws (group 1 = region). +_MANTLE_HOST_RE = re.compile( + r"^https?://bedrock-mantle\.([^/.]+)\.api\.aws", re.IGNORECASE +) + class BedrockMantleResponsesAPIConfig(OpenAIResponsesAPIConfig): + def __init__(self, aws_signer: Optional[BaseAWSLLM] = None): + super().__init__() + self._aws_signer = aws_signer or BaseAWSLLM() + @property def custom_llm_provider(self) -> LlmProviders: return LlmProviders.BEDROCK_MANTLE + @staticmethod + def _resolve_region(params: dict) -> str: + region = params.get("aws_region_name") + if region: + return region + base = params.get("api_base") or get_secret_str("BEDROCK_MANTLE_API_BASE") + if base: + match = _MANTLE_HOST_RE.match(base.rstrip("/")) + if match: + return match.group(1) + return ( + get_secret_str("BEDROCK_MANTLE_REGION") + or get_secret_str("AWS_REGION_NAME") + or get_secret_str("AWS_REGION") + or BEDROCK_MANTLE_DEFAULT_REGION + ) + def get_complete_url( self, api_base: Optional[str], litellm_params: dict, ) -> str: - region = ( - get_secret_str("BEDROCK_MANTLE_REGION") - or get_secret_str("AWS_REGION") - or BEDROCK_MANTLE_DEFAULT_REGION - ) + region = self._resolve_region({**litellm_params, "api_base": api_base}) base = ( api_base or get_secret_str("BEDROCK_MANTLE_API_BASE") @@ -55,6 +89,11 @@ class BedrockMantleResponsesAPIConfig(OpenAIResponsesAPIConfig): if base.endswith(suffix): base = base[: -len(suffix)] break + # For the standard Mantle host (including the default-region base that + # responses/main.py auto-injects into litellm_params.api_base), pin to the + # single resolved region so aws_region_name wins; preserve custom proxy hosts. + if _MANTLE_HOST_RE.match(base): + base = f"https://bedrock-mantle.{region}.api.aws" return f"{base}/openai/v1/responses" def validate_environment( @@ -66,12 +105,8 @@ class BedrockMantleResponsesAPIConfig(OpenAIResponsesAPIConfig): or get_secret_str("BEDROCK_MANTLE_API_KEY") or get_secret_str("AWS_BEARER_TOKEN_BEDROCK") ) - if not api_key: - raise ValueError( - "Bedrock Mantle API key is required. Set BEDROCK_MANTLE_API_KEY " - "(or AWS_BEARER_TOKEN_BEDROCK) or pass api_key." - ) - headers["Authorization"] = f"Bearer {api_key}" + if api_key: + headers["Authorization"] = f"Bearer {api_key}" return headers def supports_native_file_search(self) -> bool: @@ -79,3 +114,58 @@ class BedrockMantleResponsesAPIConfig(OpenAIResponsesAPIConfig): def supports_native_websocket(self) -> bool: return False + + def sign_request( + self, + headers: dict, + optional_params: dict, + request_data: dict, + api_base: str, + api_key: Optional[str] = None, + model: Optional[str] = None, + stream: Optional[bool] = None, + fake_stream: Optional[bool] = None, + ) -> Tuple[dict, Optional[bytes]]: + bearer = ( + api_key + or get_secret_str("BEDROCK_MANTLE_API_KEY") + or get_secret_str("AWS_BEARER_TOKEN_BEDROCK") + ) + if not bearer: + # SigV4 path. Pin the credential-scope region to the region of the actual + # signing URL (api_base, already region-resolved by get_complete_url) so the + # SigV4 scope and the URL host can never disagree. Resolve from api_base first, + # then fall back to the regular precedence. Also drop any caller Authorization + # so _sign_request's restore-original-Authorization step cannot override the + # SigV4 header. + optional_params = { + **optional_params, + "aws_region_name": self._resolve_region( + {**optional_params, "api_base": api_base} + ), + } + headers = {k: v for k, v in headers.items() if k.lower() != "authorization"} + try: + return self._aws_signer._sign_request( + service_name="bedrock", + headers=headers, + optional_params=optional_params, + request_data=request_data, + api_base=api_base, + api_key=bearer, + model=model, + stream=stream, + fake_stream=fake_stream, + ) + except ( + NoCredentialsError, + PartialCredentialsError, + ProfileNotFound, + CredentialRetrievalError, + ) as e: + raise ValueError( + "Bedrock Mantle auth failed: no Bearer token and no usable AWS " + "credentials. Set BEDROCK_MANTLE_API_KEY (or AWS_BEARER_TOKEN_BEDROCK) " + "or pass api_key for Bearer auth, or provide AWS credentials " + "(IAM role / access key / profile / web identity) for SigV4." + ) from e diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 31c772510ba..25424feaeb4 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2318,6 +2318,31 @@ class BaseLLMHTTPHandler: # but never included in the outbound provider payload. request_context["litellm_params"] = dict(litellm_params) + is_stream_request = bool(stream) + if is_stream_request and fake_stream is True: + stream, data = self._prepare_fake_stream_request( + stream=stream, + data=data, + fake_stream=fake_stream, + ) + + # Sign after the body is final (post-transform/normalize/extra_body and post + # fake-stream prep) so signed bytes match what we send. No-op for providers + # that inherit the default sign_request. + headers, signed_body = responses_api_provider_config.sign_request( + headers=headers, + optional_params=dict(litellm_params), + request_data=data, + api_base=api_base, + api_key=litellm_params.api_key, + model=model, + stream=stream, + fake_stream=fake_stream, + ) + body_kwargs: Dict[str, Any] = ( + {"data": signed_body} if signed_body is not None else {"json": data} + ) + ## LOGGING logging_obj.pre_call( input=input, @@ -2330,22 +2355,14 @@ class BaseLLMHTTPHandler: ) try: - if stream: - # For streaming, use stream=True in the request - if fake_stream is True: - stream, data = self._prepare_fake_stream_request( - stream=stream, - data=data, - fake_stream=fake_stream, - ) - + if is_stream_request: response = sync_httpx_client.post( url=api_base, headers=headers, - json=data, timeout=timeout or float(response_api_optional_request_params.get("timeout", 0)), stream=stream, + **body_kwargs, ) if fake_stream is True: return MockResponsesAPIStreamingIterator( @@ -2370,13 +2387,12 @@ class BaseLLMHTTPHandler: call_type=CallTypes.responses.value, ) else: - # For non-streaming requests response = sync_httpx_client.post( url=api_base, headers=headers, - json=data, timeout=timeout or float(response_api_optional_request_params.get("timeout", 0)), + **body_kwargs, ) except Exception as e: raise self._handle_error( @@ -2464,6 +2480,28 @@ class BaseLLMHTTPHandler: # but never included in the outbound provider payload. request_context["litellm_params"] = dict(litellm_params) + is_stream_request = bool(stream) + if is_stream_request and fake_stream is True: + stream, data = self._prepare_fake_stream_request( + stream=stream, + data=data, + fake_stream=fake_stream, + ) + + headers, signed_body = responses_api_provider_config.sign_request( + headers=headers, + optional_params=dict(litellm_params), + request_data=data, + api_base=api_base, + api_key=litellm_params.api_key, + model=model, + stream=stream, + fake_stream=fake_stream, + ) + body_kwargs: Dict[str, Any] = ( + {"data": signed_body} if signed_body is not None else {"json": data} + ) + ## LOGGING logging_obj.pre_call( input=input, @@ -2476,22 +2514,14 @@ class BaseLLMHTTPHandler: ) try: - if stream: - # For streaming, we need to use stream=True in the request - if fake_stream is True: - stream, data = self._prepare_fake_stream_request( - stream=stream, - data=data, - fake_stream=fake_stream, - ) - + if is_stream_request: response = await async_httpx_client.post( url=api_base, headers=headers, - json=data, timeout=timeout or float(response_api_optional_request_params.get("timeout", 0)), stream=stream, + **body_kwargs, ) if fake_stream is True: @@ -2518,13 +2548,12 @@ class BaseLLMHTTPHandler: call_type=CallTypes.responses.value, ) else: - # For non-streaming, proceed as before response = await async_httpx_client.post( url=api_base, headers=headers, - json=data, timeout=timeout or float(response_api_optional_request_params.get("timeout", 0)), + **body_kwargs, ) except Exception as e: @@ -4005,6 +4034,18 @@ class BaseLLMHTTPHandler: ) data = BaseResponsesAPIConfig.normalize_responses_api_request_dict(data) + headers, signed_body = responses_api_provider_config.sign_request( + headers=headers, + optional_params=dict(litellm_params), + request_data=data, + api_base=url, + api_key=litellm_params.api_key, + model=model, + ) + body_kwargs: Dict[str, Any] = ( + {"data": signed_body} if signed_body is not None else {"json": data} + ) + ## LOGGING logging_obj.pre_call( input=input, @@ -4018,7 +4059,7 @@ class BaseLLMHTTPHandler: try: response = sync_httpx_client.post( - url=url, headers=headers, json=data, timeout=timeout + url=url, headers=headers, timeout=timeout, **body_kwargs ) except Exception as e: @@ -4088,6 +4129,18 @@ class BaseLLMHTTPHandler: ) data = BaseResponsesAPIConfig.normalize_responses_api_request_dict(data) + headers, signed_body = responses_api_provider_config.sign_request( + headers=headers, + optional_params=dict(litellm_params), + request_data=data, + api_base=url, + api_key=litellm_params.api_key, + model=model, + ) + body_kwargs: Dict[str, Any] = ( + {"data": signed_body} if signed_body is not None else {"json": data} + ) + ## LOGGING logging_obj.pre_call( input=input, @@ -4101,7 +4154,7 @@ class BaseLLMHTTPHandler: try: response = await async_httpx_client.post( - url=url, headers=headers, json=data, timeout=timeout + url=url, headers=headers, timeout=timeout, **body_kwargs ) except Exception as e: diff --git a/litellm/llms/deepseek/messages/transformation.py b/litellm/llms/deepseek/messages/transformation.py index ad60478960e..63b736ffd1d 100644 --- a/litellm/llms/deepseek/messages/transformation.py +++ b/litellm/llms/deepseek/messages/transformation.py @@ -26,6 +26,9 @@ class DeepSeekAnthropicMessagesConfig(AnthropicMessagesConfig): def custom_llm_provider(self) -> Optional[str]: return "deepseek" + def should_strip_billing_metadata(self) -> bool: + return True + @staticmethod def get_api_key(api_key: Optional[str] = None) -> Optional[str]: return api_key or get_secret_str("DEEPSEEK_API_KEY") or litellm.api_key diff --git a/litellm/llms/github_copilot/responses/transformation.py b/litellm/llms/github_copilot/responses/transformation.py index 0929f95cf43..3406538c774 100644 --- a/litellm/llms/github_copilot/responses/transformation.py +++ b/litellm/llms/github_copilot/responses/transformation.py @@ -2,7 +2,7 @@ GitHub Copilot Responses API Configuration. This module provides the configuration for GitHub Copilot's Responses API, -which is required for models like gpt-5.1-codex that only support the /responses endpoint. +which is required for models like gpt-5.3-codex that only support the /responses endpoint. Implementation based on analysis of the copilot-api project by caozhiyuan: https://github.com/caozhiyuan/copilot-api @@ -12,6 +12,7 @@ from typing import TYPE_CHECKING, Any, Dict, Optional, Union import os +import litellm from litellm._logging import verbose_logger from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH from litellm.exceptions import AuthenticationError @@ -22,6 +23,7 @@ from litellm.types.llms.openai import ( ) from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import LlmProviders +from litellm.utils import _cached_get_model_info_helper from ..authenticator import Authenticator from ..common_utils import ( @@ -38,6 +40,47 @@ else: LiteLLMLoggingObj = Any +def github_copilot_supports_responses_api(model: str) -> bool: + """ + Gate native /v1/responses dispatch per github_copilot model. + + Resolution (first match wins): mode "responses" -> True; mode "chat" -> + False (opt-out wins for dual-endpoint models); "/v1/responses" in + supported_endpoints -> True; else False. Unknown model -> False (the bridge + always works since every Copilot model supports /chat/completions). + + Reads merged model info (per-deployment model_info applied via the router's + register_model, which also clears the cache used here). + """ + try: + info = _cached_get_model_info_helper( + model=model, custom_llm_provider="github_copilot" + ) + except Exception as e: + verbose_logger.debug( + "github_copilot_supports_responses_api: get_model_info failed " + "for %s: %s", + model, + e, + ) + return False + + mode = info.get("mode") + if mode == "responses": + return True + if mode == "chat": + return False + + # supported_endpoints is dropped by ModelInfoBase; read it from the raw + # model_cost entry via the resolved key. + key = info.get("key") + raw_info = litellm.model_cost.get(key) if isinstance(key, str) else None + endpoints = ( + raw_info.get("supported_endpoints") if isinstance(raw_info, dict) else None + ) + return isinstance(endpoints, list) and "/v1/responses" in endpoints + + class GithubCopilotResponsesAPIConfig(OpenAIResponsesAPIConfig): """ Configuration for GitHub Copilot's Responses API. diff --git a/litellm/llms/minimax/messages/transformation.py b/litellm/llms/minimax/messages/transformation.py index 3190a5f5412..57cfcbf0621 100644 --- a/litellm/llms/minimax/messages/transformation.py +++ b/litellm/llms/minimax/messages/transformation.py @@ -28,6 +28,9 @@ class MinimaxMessagesConfig(AnthropicMessagesConfig): def custom_llm_provider(self) -> Optional[str]: return "minimax" + def should_strip_billing_metadata(self) -> bool: + return True + @staticmethod def get_api_key(api_key: Optional[str] = None) -> Optional[str]: """ diff --git a/litellm/llms/openai_like/dynamic_config.py b/litellm/llms/openai_like/dynamic_config.py index fac453447fa..9ed9734edae 100644 --- a/litellm/llms/openai_like/dynamic_config.py +++ b/litellm/llms/openai_like/dynamic_config.py @@ -187,6 +187,7 @@ def create_responses_config_class(provider: SimpleProviderConfig): from litellm.llms.openai_like.responses.transformation import ( OpenAILikeResponsesConfig, ) + from litellm.types.llms.openai import ResponseInputParam from litellm.types.router import GenericLiteLLMParams class JSONProviderResponsesConfig(OpenAILikeResponsesConfig): @@ -223,5 +224,23 @@ def create_responses_config_class(provider: SimpleProviderConfig): api_base = api_base.rstrip("/") return f"{api_base}/responses" + def transform_responses_api_request( + self, + model: str, + input: Union[str, ResponseInputParam], + response_api_optional_request_params: dict, + litellm_params: GenericLiteLLMParams, + headers: dict, + ) -> dict: + if provider.special_handling.get("force_store_false"): + response_api_optional_request_params["store"] = False + return super().transform_responses_api_request( + model=model, + input=input, + response_api_optional_request_params=response_api_optional_request_params, + litellm_params=litellm_params, + headers=headers, + ) + _responses_config_cache[provider.slug] = JSONProviderResponsesConfig return JSONProviderResponsesConfig diff --git a/litellm/llms/openai_like/providers.json b/litellm/llms/openai_like/providers.json index 49b3801c82f..13d22488838 100644 --- a/litellm/llms/openai_like/providers.json +++ b/litellm/llms/openai_like/providers.json @@ -132,5 +132,14 @@ "param_mappings": { "max_completion_tokens": "max_tokens" } + }, + "parasail": { + "base_url": "https://api.parasail.io/v1", + "api_key_env": "PARASAIL_API_KEY", + "api_base_env": "PARASAIL_API_BASE", + "supported_endpoints": ["/v1/chat/completions", "/v1/responses"], + "special_handling": { + "force_store_false": true + } } } diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 5cd02293f14..7d355a2e908 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -1111,6 +1111,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): { "voice": "alloy", "format": "mp3", + "language_code": "en-US", } Expected output: @@ -1119,7 +1120,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): prebuiltVoiceConfig: { voiceName: "alloy", } - } + }, + languageCode: "en-US", } """ from litellm.types.llms.vertex_ai import ( @@ -1145,6 +1147,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): voice_config: VoiceConfig = {"prebuiltVoiceConfig": prebuilt_voice_config} speech_config["voiceConfig"] = voice_config + if "language_code" in value: + speech_config["languageCode"] = value["language_code"] + return cast(dict, speech_config) @staticmethod diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py index 1e92754857b..8a92e7ec4a5 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py @@ -17,6 +17,9 @@ from ..output_params_utils import sanitize_vertex_anthropic_output_params class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, VertexBase): + def should_strip_billing_metadata(self) -> bool: + return True + def validate_anthropic_messages_environment( self, headers: dict, diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py index c852909d475..ae8bdc55443 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py @@ -52,6 +52,9 @@ class VertexAIAnthropicConfig(AnthropicConfig): def custom_llm_provider(self) -> Optional[str]: return "vertex_ai" + def should_strip_billing_metadata(self) -> bool: + return True + def _add_context_management_beta_headers( self, beta_set: set, context_management: dict ) -> None: diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 57a7d860baa..fba968dd95d 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2312,6 +2312,24 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="List of MCP server fields that must be filled in for a submission to pass standards checks (e.g. ['description', 'source_url', 'alias']).", ) + disable_budget_reservation: Optional[bool] = Field( + None, + description=( + "If True, disables the optimistic per-request budget reservation " + "introduced in v1.84.0. " + "WARNING: This weakens hard budget enforcement. Without the reservation, " + "a burst of concurrent requests from a single key can each pass the " + "read-time spend check before any of them is charged, allowing a " + "configured budget to be exceeded under high concurrency. " + "Budgets are still evaluated on every request at read time, so " + "an already-exhausted budget is still rejected. " + "Enable only if your deployment is experiencing phantom " + "BudgetExceededError responses caused by leaked reservations " + "(see GitHub issue #27639). " + "A proxy-level WARNING is logged on every request while this flag " + "is active as a reminder that hard enforcement is relaxed." + ), + ) class ConfigYAML(LiteLLMPydanticObjectBase): diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index a970e0ddee8..666c01562b5 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -2425,6 +2425,7 @@ async def _run_centralized_common_checks( # noqa: PLR0915 user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, skip_budget_checks=skip_budget_checks, + general_settings=general_settings, ) @@ -2445,12 +2446,23 @@ async def _reserve_budget_after_common_checks( user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, skip_budget_checks: bool, + general_settings: dict, end_user_id: Optional[str] = None, end_user_object: Optional[LiteLLM_EndUserTable] = None, ) -> None: user_api_key_auth_obj.budget_reservation = None if skip_budget_checks: return + if general_settings.get("disable_budget_reservation") is True: + verbose_proxy_logger.warning( + "disable_budget_reservation is enabled: skipping optimistic budget " + "reservation. Budget enforcement is read-time only — concurrent " + "requests can each pass the spend check before their cost is recorded, " + "so a configured budget may be briefly exceeded under high concurrency. " + "Set disable_budget_reservation to False or remove it to restore " + "hard per-request budget enforcement." + ) + return from litellm.proxy.spend_tracking.budget_reservation import ( reserve_budget_for_request, diff --git a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py index 14d950ecdf4..d1ef165b46e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py +++ b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py @@ -312,11 +312,27 @@ class CrowdStrikeAIDRHandler(CustomGuardrail): event_type = "output" hook_name = "apply_guardrail (response)" - ai_guard_payload = { + ai_guard_payload: dict[str, Any] = { "guard_input": guard_input.model_dump(mode="json"), "event_type": event_type, } + model = inputs.get("model") + if model: + ai_guard_payload["model"] = model + + metadata = request_data.get("litellm_metadata", request_data.get("metadata")) + if isinstance(metadata, Mapping): + user_id = metadata.get("user_api_key_user_id") + if user_id: + ai_guard_payload["user_id"] = user_id + + extra_info: dict[str, str] = {} + user_email = metadata.get("user_api_key_user_email") + if user_email: + extra_info["user_name"] = user_email + ai_guard_payload["extra_info"] = extra_info + ai_guard_response = await self._call_crowdstrike_aidr_guard( ai_guard_payload, hook_name ) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index a9570616850..4812bed2f21 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -1220,7 +1220,7 @@ async def _setup_role_mappings() -> Optional["RoleMappings"]: generic_role_mappings_group_claim = os.getenv( "GENERIC_ROLE_MAPPINGS_GROUP_CLAIM", None ) - generic_role_mappoings_default_role = os.getenv( + generic_role_mappings_default_role = os.getenv( "GENERIC_ROLE_MAPPINGS_DEFAULT_ROLE", None ) if generic_role_mappings is not None: @@ -1239,7 +1239,7 @@ async def _setup_role_mappings() -> Optional["RoleMappings"]: role_mappings_data = { "provider": "generic", "group_claim": generic_role_mappings_group_claim, - "default_role": generic_role_mappoings_default_role, + "default_role": generic_role_mappings_default_role, "roles": generic_user_role_mappings_data, } diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py index 29bbb37501f..6dd1f8548eb 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py @@ -32,6 +32,71 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import ( from litellm.types.utils import ImageResponse, LlmProviders, PassthroughCallTypes from litellm.utils import ModelResponse, TextCompletionResponse +# Hostnames that route to OpenAI-compatible APIs. +# +# `api.openai.com` is OpenAI proper. The two Azure domains below are *shared by +# every Azure Cognitive Service* (Speech, Vision, Language, ...), not just Azure +# OpenAI: `openai.azure.com` is the classic Azure OpenAI domain, while +# `cognitiveservices.azure.com` is used by newer "Azure AI Foundry" / +# Cognitive Services-hosted Azure OpenAI deployments. Because the hostname alone +# cannot tell Azure OpenAI apart from the other Cognitive Services on those +# domains, requests there must additionally carry an OpenAI-style path segment. +_OPENAI_HOSTNAMES = ("api.openai.com",) +_AZURE_OPENAI_HOSTNAMES = ("openai.azure.com", "cognitiveservices.azure.com") +# Path markers that identify an Azure request as Azure OpenAI rather than Speech +# / Vision / Language / ... `/openai/` is the native Azure OpenAI path prefix; +# `/v1/` is the OpenAI-v1 surface used by LiteLLM's pass-through routing. Other +# Cognitive Services use service-named prefixes and versions like `/v3.1/`, +# `/v1.0/`, so they do not collide with these markers. +_AZURE_OPENAI_PATH_MARKERS = ("/openai/", "/v1/") + + +def _hostname_matches(hostname: str, suffixes: tuple) -> bool: + """True if hostname equals one of `suffixes` or is a subdomain of it. + + Uses suffix matching (not a bare substring test) so look-alikes such as + `cognitiveservices.azure.com.attacker.example` are not accepted. + """ + return any( + hostname == suffix or hostname.endswith("." + suffix) for suffix in suffixes + ) + + +def _is_openai_compatible_host(hostname: Optional[str]) -> bool: + """True if the hostname is OpenAI proper or one of the Azure OpenAI domains. + + Hostname-only check, kept for the route-level helpers that additionally + require a specific OpenAI path (e.g. `/v1/chat/completions`). When only the + hostname would otherwise gate dispatch, use `_is_openai_compatible_url` so + non-OpenAI Azure Cognitive Services on the shared domains are excluded. + """ + if not hostname: + return False + return _hostname_matches(hostname, _OPENAI_HOSTNAMES) or _hostname_matches( + hostname, _AZURE_OPENAI_HOSTNAMES + ) + + +def _is_openai_compatible_url(url_route: Optional[str]) -> bool: + """True if the URL targets an OpenAI-compatible API surface. + + For the shared Azure Cognitive Services domains we additionally require an + OpenAI-style path segment (`/openai/` or `/v1/`) so non-OpenAI Azure services + (Speech, Vision, Language, ...) on the same domain are not misclassified as + OpenAI routes. + """ + if not url_route: + return False + parsed_url = urlparse(url_route) + hostname = parsed_url.hostname + if not hostname: + return False + if _hostname_matches(hostname, _OPENAI_HOSTNAMES): + return True + if _hostname_matches(hostname, _AZURE_OPENAI_HOSTNAMES): + return any(marker in parsed_url.path for marker in _AZURE_OPENAI_PATH_MARKERS) + return False + class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): """ @@ -52,12 +117,8 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): if not url_route: return False parsed_url = urlparse(url_route) - return bool( - parsed_url.hostname - and ( - "api.openai.com" in parsed_url.hostname - or "openai.azure.com" in parsed_url.hostname - ) + return ( + _is_openai_compatible_host(parsed_url.hostname) and "/v1/chat/completions" in parsed_url.path ) @@ -67,12 +128,8 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): if not url_route: return False parsed_url = urlparse(url_route) - return bool( - parsed_url.hostname - and ( - "api.openai.com" in parsed_url.hostname - or "openai.azure.com" in parsed_url.hostname - ) + return ( + _is_openai_compatible_host(parsed_url.hostname) and "/v1/images/generations" in parsed_url.path ) @@ -82,12 +139,8 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): if not url_route: return False parsed_url = urlparse(url_route) - return bool( - parsed_url.hostname - and ( - "api.openai.com" in parsed_url.hostname - or "openai.azure.com" in parsed_url.hostname - ) + return ( + _is_openai_compatible_host(parsed_url.hostname) and "/v1/images/edits" in parsed_url.path ) @@ -97,13 +150,8 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): if not url_route: return False parsed_url = urlparse(url_route) - return bool( - parsed_url.hostname - and ( - "api.openai.com" in parsed_url.hostname - or "openai.azure.com" in parsed_url.hostname - ) - and ("/v1/responses" in parsed_url.path or "/responses" in parsed_url.path) + return _is_openai_compatible_host(parsed_url.hostname) and ( + "/v1/responses" in parsed_url.path or "/responses" in parsed_url.path ) def _get_user_from_metadata( diff --git a/litellm/proxy/pass_through_endpoints/success_handler.py b/litellm/proxy/pass_through_endpoints/success_handler.py index 292871bae67..af1d39da020 100644 --- a/litellm/proxy/pass_through_endpoints/success_handler.py +++ b/litellm/proxy/pass_through_endpoints/success_handler.py @@ -434,15 +434,20 @@ class PassThroughEndpointLogging: return False def is_openai_route(self, url_route: str): - """Check if the URL route is an OpenAI API route.""" + """Check if the URL route is an OpenAI API route. + + Uses the URL-aware helper so that non-OpenAI Azure Cognitive Services + (Speech, Vision, Language, ...) sharing the `*.cognitiveservices.azure.com` + / `*.openai.azure.com` domains are not misclassified as OpenAI routes. + """ if not url_route: return False - parsed_url = urlparse(url_route) - return parsed_url.hostname and ( - "api.openai.com" in parsed_url.hostname - or "openai.azure.com" in parsed_url.hostname + from .llm_provider_handlers.openai_passthrough_logging_handler import ( + _is_openai_compatible_url, ) + return _is_openai_compatible_url(url_route) + def is_gemini_route( self, url_route: str, custom_llm_provider: Optional[str] = None ): diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index e2ba8353591..d3d30642216 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -148,7 +148,9 @@ class LiteLLMCompletionResponsesConfig: # which is equivalent to "required" in OpenAI format return "required" elif tool_choice_type == "function": - # function type without name - fall back to required + function_name = tool_choice.get("name") + if function_name: + return {"type": "function", "function": {"name": function_name}} return "required" # Return as-is for unknown formats diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index 51429d0769e..00db7b199b6 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -232,6 +232,7 @@ class VoiceConfig(TypedDict): class SpeechConfig(TypedDict, total=False): voiceConfig: VoiceConfig + languageCode: str class GenerationConfig(TypedDict, total=False): diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 9633cecf96c..c3ea605dd9e 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1540,6 +1540,11 @@ class ServerToolUse(BaseModel): web_search_requests: Optional[int] = None tool_search_requests: Optional[int] = None + def __getitem__(self, key: str) -> Optional[int]: + if key not in self.__class__.model_fields: + raise KeyError(key) + return getattr(self, key) + class Usage(SafeAttributeModel, CompletionUsage): _cache_creation_input_tokens: int = PrivateAttr( @@ -1570,7 +1575,7 @@ class Usage(SafeAttributeModel, CompletionUsage): completion_tokens_details: Optional[ Union[CompletionTokensDetailsWrapper, dict] ] = None, - server_tool_use: Optional[ServerToolUse] = None, + server_tool_use: Optional[Union[ServerToolUse, dict]] = None, cost: Optional[float] = None, **params, ): @@ -1671,6 +1676,9 @@ class Usage(SafeAttributeModel, CompletionUsage): prompt_tokens_details=_prompt_tokens_details or None, ) + if isinstance(server_tool_use, dict): + server_tool_use = ServerToolUse(**server_tool_use) + if server_tool_use is not None: self.server_tool_use = server_tool_use else: # maintain openai compatibility in usage object if possible @@ -3392,6 +3400,7 @@ class LlmProviders(str, Enum): POE = "poe" CHUTES = "chutes" NEOSANTARA = "neosantara" + PARASAIL = "parasail" XIAOMI_MIMO = "xiaomi_mimo" TENSORMESH = "tensormesh" LITELLM_AGENT = "litellm_agent" diff --git a/litellm/utils.py b/litellm/utils.py index 7312e71bbd1..8d9d0a409c6 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8895,7 +8895,13 @@ class ProviderConfigManager: elif litellm.LlmProviders.XAI == provider: return litellm.XAIResponsesAPIConfig() elif litellm.LlmProviders.GITHUB_COPILOT == provider: - return litellm.GithubCopilotResponsesAPIConfig() + from litellm.llms.github_copilot.responses.transformation import ( + github_copilot_supports_responses_api, + ) + + if model is None or github_copilot_supports_responses_api(model=model): + return litellm.GithubCopilotResponsesAPIConfig() + return None elif litellm.LlmProviders.CHATGPT == provider: return litellm.ChatGPTResponsesAPIConfig() elif litellm.LlmProviders.LITELLM_PROXY == provider: diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index a1ad20fffd1..6caab585ac9 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -1834,6 +1834,23 @@ "search": true } }, + "parasail": { + "display_name": "Parasail (`parasail`)", + "url": "https://docs.litellm.ai/docs/providers/parasail", + "endpoints": { + "chat_completions": true, + "messages": false, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": false + } + }, "perplexity": { "display_name": "Perplexity AI (`perplexity`)", "url": "https://docs.litellm.ai/docs/providers/perplexity", diff --git a/tests/llm_translation/test_bedrock_completion.py b/tests/llm_translation/test_bedrock_completion.py index 9cf253c379d..fa22ff6b392 100644 --- a/tests/llm_translation/test_bedrock_completion.py +++ b/tests/llm_translation/test_bedrock_completion.py @@ -2712,6 +2712,10 @@ def test_bedrock_top_k_param(model, expected_params): data = json.loads(mock_post.call_args.kwargs["data"]) if "mistral" in model: assert data["top_k"] == 2 + elif expected_params == {}: + # Models that don't support top_k produce no additionalModelRequestFields; + # the empty block is now omitted entirely rather than sent as `{}`. + assert "additionalModelRequestFields" not in data else: assert data["additionalModelRequestFields"] == expected_params @@ -3059,8 +3063,6 @@ async def test_bedrock_max_completion_tokens(model: str): assert request_body == { "messages": [{"role": "user", "content": [{"text": "Hello!"}]}], - "additionalModelRequestFields": {}, - "system": [], "inferenceConfig": {"maxTokens": 10}, } diff --git a/tests/test_litellm/caching/test_caching.py b/tests/test_litellm/caching/test_caching.py new file mode 100644 index 00000000000..02d62a19152 --- /dev/null +++ b/tests/test_litellm/caching/test_caching.py @@ -0,0 +1,48 @@ +import logging +import re + +from litellm.caching.caching import Cache +from litellm.types.caching import LiteLLMCacheType + + +def test_cache_key_debug_log_does_not_include_prompt_material(caplog): + cache = Cache(type=LiteLLMCacheType.LOCAL) + prompt_marker = "secret prompt material " + + with caplog.at_level(logging.DEBUG, logger="LiteLLM"): + cache_key = cache.get_cache_key( + model="gpt-4.1-mini", + messages=[ + {"role": "system", "content": prompt_marker * 100}, + {"role": "user", "content": "hello"}, + ], + tools=[ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": { + "type": "object", + "properties": {"query": {"type": "string"}}, + }, + }, + } + ], + response_format={ + "type": "json_schema", + "json_schema": { + "name": "lookup_response", + "schema": {"type": "object"}, + }, + }, + stream=True, + ) + + assert re.fullmatch(r"[0-9a-f]{64}", cache_key) + + created_cache_key_logs = [ + record.getMessage() for record in caplog.records if "Created cache key:" in record.getMessage() + ] + assert created_cache_key_logs + assert all(prompt_marker not in message for message in created_cache_key_logs) + assert any(cache_key in message for message in created_cache_key_logs) diff --git a/tests/test_litellm/caching/test_redis_semantic_cache.py b/tests/test_litellm/caching/test_redis_semantic_cache.py index b50a35ef50e..13f9d00136d 100644 --- a/tests/test_litellm/caching/test_redis_semantic_cache.py +++ b/tests/test_litellm/caching/test_redis_semantic_cache.py @@ -523,3 +523,468 @@ async def test_redis_semantic_cache_async_set_cache_stores_cache_key_filter( filters={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"}, ttl=60, ) + + +def test_redis_semantic_cache_set_cache_uses_responses_string_input(): + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache) + redis_semantic_cache.llmcache = MagicMock() + redis_semantic_cache._get_cache_filters = MagicMock( + return_value={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"} + ) + redis_semantic_cache._get_ttl = MagicMock(return_value=None) + + redis_semantic_cache.set_cache( + key="test_key", + value={"content": "Paris"}, + input="What is the capital of France?", + ) + + redis_semantic_cache.llmcache.store.assert_called_once_with( + "What is the capital of France?", + "{'content': 'Paris'}", + filters={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"}, + ) + + +def test_redis_semantic_cache_get_cache_uses_responses_string_input(): + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache) + redis_semantic_cache.similarity_threshold = 0.8 + redis_semantic_cache.llmcache = MagicMock() + redis_semantic_cache.llmcache.check = MagicMock( + return_value=[ + { + "prompt": "What is the capital of France?", + "response": '{"content": "Paris"}', + "vector_distance": 0.1, + RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key", + } + ] + ) + + with patch.object( + redis_semantic_cache, + "_get_cache_key_filter_expression", + return_value="cache-key-filter", + ): + metadata = {} + result = redis_semantic_cache.get_cache( + key="test_key", + input="What is the capital of France?", + metadata=metadata, + ) + + assert result == {"content": "Paris"} + assert metadata["semantic-similarity"] == pytest.approx(0.9) + redis_semantic_cache.llmcache.check.assert_called_once_with( + prompt="What is the capital of France?", + filter_expression="cache-key-filter", + ) + + +def test_redis_semantic_cache_set_cache_flattens_structured_responses_input(): + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache) + redis_semantic_cache.llmcache = MagicMock() + redis_semantic_cache._get_cache_filters = MagicMock( + return_value={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"} + ) + redis_semantic_cache._get_ttl = MagicMock(return_value=None) + + redis_semantic_cache.set_cache( + key="test_key", + value={"content": "Paris"}, + input=[ + { + "role": "user", + "content": [ + {"type": "input_text", "text": "What is the capital of France?"}, + {"type": "input_text", "text": "Answer briefly."}, + { + "type": "input_image", + "image_url": "https://example.com/paris.png", + }, + ], + } + ], + ) + + redis_semantic_cache.llmcache.store.assert_called_once_with( + "What is the capital of France?\nAnswer briefly.", + "{'content': 'Paris'}", + filters={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"}, + ) + + +def test_redis_semantic_cache_prompt_extraction_prefers_messages(): + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + prompt = RedisSemanticCache._get_prompt_from_kwargs( + messages=[{"content": "message prompt"}], + input="responses prompt", + ) + + assert prompt == "message prompt" + + +def test_redis_semantic_cache_prompt_extraction_handles_model_objects(): + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + class ModelDumpInput: + def model_dump(self): + return {"content": [{"text": "model dump prompt"}]} + + class DictInput: + def dict(self): + return {"content": [{"output_text": "dict prompt"}]} + + prompt = RedisSemanticCache._get_prompt_from_kwargs( + input=[ + ModelDumpInput(), + DictInput(), + {"content": [{"input_text": "inline prompt"}]}, + {"content": [{"type": "input_image", "image_url": "https://example.com"}]}, + ] + ) + + assert prompt == "model dump prompt\ndict prompt\ninline prompt" + + +def test_redis_semantic_cache_prompt_extraction_returns_none_without_text(): + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + assert RedisSemanticCache._get_prompt_from_kwargs() is None + assert RedisSemanticCache._get_prompt_from_kwargs(input=None) is None + assert RedisSemanticCache._get_prompt_from_kwargs(input=" ") is None + assert ( + RedisSemanticCache._get_prompt_from_kwargs( + input=[{"type": "input_image", "image_url": "https://example.com"}] + ) + is None + ) + + +def test_redis_semantic_cache_prompt_extraction_skips_blank_dict_text_keys(): + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + prompt = RedisSemanticCache._get_prompt_from_kwargs( + input={"text": " ", "input_text": "fallback prompt"} + ) + + assert prompt == "fallback prompt" + + +def test_redis_semantic_cache_prompt_extraction_skips_blank_object_text_keys(): + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + class ResponseInput: + text = " " + input_text = "fallback prompt" + + prompt = RedisSemanticCache._get_prompt_from_kwargs(input=ResponseInput()) + + assert prompt == "fallback prompt" + + +def test_redis_semantic_cache_prompt_extraction_handles_object_content(): + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + class ResponseInput: + content = [{"text": "object content prompt"}] + + prompt = RedisSemanticCache._get_prompt_from_kwargs(input=ResponseInput()) + + assert prompt == "object content prompt" + + +def test_redis_semantic_cache_set_cache_skips_blank_responses_input(): + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache) + redis_semantic_cache.llmcache = MagicMock() + + redis_semantic_cache.set_cache( + key="test_key", + value={"content": "Paris"}, + input=" ", + ) + + redis_semantic_cache.llmcache.store.assert_not_called() + + +def test_redis_semantic_cache_get_cache_sets_similarity_on_blank_responses_input(): + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache) + redis_semantic_cache.llmcache = MagicMock() + metadata = {} + + result = redis_semantic_cache.get_cache( + key="test_key", + input=" ", + metadata=metadata, + ) + + assert result is None + assert metadata["semantic-similarity"] == 0.0 + redis_semantic_cache.llmcache.check.assert_not_called() + + +def test_redis_semantic_cache_get_cache_sets_similarity_when_no_results(): + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache) + redis_semantic_cache.llmcache = MagicMock() + redis_semantic_cache.llmcache.check = MagicMock(return_value=[]) + + with patch.object( + redis_semantic_cache, + "_get_cache_key_filter_expression", + return_value="cache-key-filter", + ): + metadata = {} + result = redis_semantic_cache.get_cache( + key="test_key", + input="What is the capital of France?", + metadata=metadata, + ) + + assert result is None + assert metadata["semantic-similarity"] == 0.0 + redis_semantic_cache.llmcache.check.assert_called_once_with( + prompt="What is the capital of France?", + filter_expression="cache-key-filter", + ) + + +@pytest.mark.asyncio +async def test_redis_semantic_cache_async_paths_use_responses_string_input(): + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache) + redis_semantic_cache.similarity_threshold = 0.8 + redis_semantic_cache.llmcache = MagicMock() + redis_semantic_cache.llmcache.astore = AsyncMock() + redis_semantic_cache.llmcache.acheck = AsyncMock( + return_value=[ + { + "prompt": "What is the capital of France?", + "response": '{"content": "Paris"}', + "vector_distance": 0.1, + RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key", + } + ] + ) + redis_semantic_cache._get_cache_filters = MagicMock( + return_value={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"} + ) + redis_semantic_cache._get_ttl = MagicMock(return_value=None) + redis_semantic_cache._get_async_embedding = AsyncMock(return_value=[0.1, 0.2, 0.3]) + + await redis_semantic_cache.async_set_cache( + key="test_key", + value={"content": "Paris"}, + input="What is the capital of France?", + ) + + with patch.object( + redis_semantic_cache, + "_get_cache_key_filter_expression", + return_value="cache-key-filter", + ): + metadata = {} + result = await redis_semantic_cache.async_get_cache( + key="test_key", + input="What is the capital of France?", + metadata=metadata, + ) + + redis_semantic_cache.llmcache.astore.assert_called_once_with( + "What is the capital of France?", + "{'content': 'Paris'}", + vector=[0.1, 0.2, 0.3], + filters={RedisSemanticCache.CACHE_KEY_FIELD_NAME: "test_key"}, + ) + assert result == {"content": "Paris"} + assert metadata["semantic-similarity"] == pytest.approx(0.9) + redis_semantic_cache.llmcache.acheck.assert_called_once_with( + prompt="What is the capital of France?", + vector=[0.1, 0.2, 0.3], + filter_expression="cache-key-filter", + ) + + +@pytest.mark.asyncio +async def test_redis_semantic_cache_async_paths_set_similarity_on_misses(): + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache) + redis_semantic_cache.llmcache = MagicMock() + redis_semantic_cache.llmcache.astore = AsyncMock() + redis_semantic_cache.llmcache.acheck = AsyncMock(return_value=[]) + redis_semantic_cache._get_async_embedding = AsyncMock(return_value=[0.1, 0.2, 0.3]) + + await redis_semantic_cache.async_set_cache( + key="test_key", + value={"content": "Paris"}, + input=" ", + ) + + redis_semantic_cache.llmcache.astore.assert_not_called() + redis_semantic_cache._get_async_embedding.assert_not_called() + + blank_metadata = {} + blank_result = await redis_semantic_cache.async_get_cache( + key="test_key", + input=" ", + metadata=blank_metadata, + ) + + assert blank_result is None + assert blank_metadata["semantic-similarity"] == 0.0 + redis_semantic_cache.llmcache.acheck.assert_not_called() + redis_semantic_cache._get_async_embedding.assert_not_called() + + with patch.object( + redis_semantic_cache, + "_get_cache_key_filter_expression", + return_value="cache-key-filter", + ): + miss_metadata = {} + miss_result = await redis_semantic_cache.async_get_cache( + key="test_key", + input="What is the capital of France?", + metadata=miss_metadata, + ) + + assert miss_result is None + assert miss_metadata["semantic-similarity"] == 0.0 + redis_semantic_cache.llmcache.acheck.assert_called_once_with( + prompt="What is the capital of France?", + vector=[0.1, 0.2, 0.3], + filter_expression="cache-key-filter", + ) + + +def test_cache_get_cache_passes_responses_input_to_backend_cache(): + from litellm.caching.caching import Cache + + cache = Cache.__new__(Cache) + cache.cache = MagicMock() + cache.cache.get_cache = MagicMock(return_value=None) + cache.should_use_cache = MagicMock(return_value=True) + cache.get_cache_key = MagicMock(return_value="test_key") + + metadata = {} + cache.get_cache( + input="What is the capital of France?", + metadata=metadata, + cache={}, + ) + + cache.cache.get_cache.assert_called_once_with( + "test_key", + input="What is the capital of France?", + metadata=metadata, + ) + + +def test_cache_get_cache_filters_sensitive_kwargs_from_backend_cache(): + from litellm.caching.caching import Cache + + cache = Cache.__new__(Cache) + cache.cache = MagicMock() + cache.should_use_cache = MagicMock(return_value=True) + cache.get_cache_key = MagicMock(return_value="test_key") + cache._get_cache_logic = MagicMock(return_value={"content": "Paris"}) + + def _cache_hit(_cache_key, **cache_kwargs): + cache_kwargs["metadata"]["semantic-similarity"] = 0.7 + return {"content": "Paris"} + + cache.cache.get_cache = MagicMock(side_effect=_cache_hit) + + metadata = {"user_api_key": "sk-secret", "trace_id": "trace-id"} + result = cache.get_cache( + input="What is the capital of France?", + metadata=metadata, + cache={"s-maxage": 10}, + api_key="sk-secret", + headers={"authorization": "Bearer sk-secret"}, + ) + + assert result == {"content": "Paris"} + assert metadata == { + "user_api_key": "sk-secret", + "trace_id": "trace-id", + "semantic-similarity": 0.7, + } + + forwarded_kwargs = cache.cache.get_cache.call_args.kwargs + assert forwarded_kwargs == { + "input": "What is the capital of France?", + "metadata": {"semantic-similarity": 0.7}, + } + assert forwarded_kwargs["metadata"] is not metadata + cache._get_cache_logic.assert_called_once_with( + cached_result={"content": "Paris"}, + max_age=10, + ) + + +def test_cache_get_cache_filters_sensitive_kwargs_without_metadata(): + from litellm.caching.caching import Cache + + cache = Cache.__new__(Cache) + cache.cache = MagicMock() + cache.cache.get_cache = MagicMock(return_value={"content": "Paris"}) + cache.should_use_cache = MagicMock(return_value=True) + cache.get_cache_key = MagicMock(return_value="test_key") + cache._get_cache_logic = MagicMock(return_value={"content": "Paris"}) + + result = cache.get_cache( + input="What is the capital of France?", + cache={"s-maxage": 10}, + api_key="sk-secret", + headers={"authorization": "Bearer sk-secret"}, + ) + + assert result == {"content": "Paris"} + cache.cache.get_cache.assert_called_once_with( + "test_key", + input="What is the capital of France?", + ) + + +def test_cache_get_cache_passes_responses_input_to_dynamic_cache(): + from litellm.caching.caching import Cache + + cache = Cache.__new__(Cache) + cache.should_use_cache = MagicMock(return_value=True) + cache.get_cache_key = MagicMock(return_value="test_key") + cache._get_cache_logic = MagicMock(return_value={"content": "Paris"}) + dynamic_cache_object = MagicMock() + dynamic_cache_object.get_cache = MagicMock(return_value={"content": "Paris"}) + + metadata = {} + result = cache.get_cache( + dynamic_cache_object=dynamic_cache_object, + input="What is the capital of France?", + metadata=metadata, + cache={}, + ) + + assert result == {"content": "Paris"} + dynamic_cache_object.get_cache.assert_called_once_with( + "test_key", + input="What is the capital of France?", + metadata=metadata, + ) + cache._get_cache_logic.assert_called_once_with( + cached_result={"content": "Paris"}, + max_age=float("inf"), + ) diff --git a/tests/test_litellm/integrations/focus/test_focus_gcs_destination.py b/tests/test_litellm/integrations/focus/test_focus_gcs_destination.py new file mode 100644 index 00000000000..35cdb18326a --- /dev/null +++ b/tests/test_litellm/integrations/focus/test_focus_gcs_destination.py @@ -0,0 +1,180 @@ +"""Tests for FocusGCSDestination.""" + +from __future__ import annotations + +from datetime import datetime, timezone +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from litellm.integrations.focus.destinations.base import FocusTimeWindow + + +def _make_window(frequency: str = "hourly") -> FocusTimeWindow: + return FocusTimeWindow( + start_time=datetime(2026, 1, 1, 10, 0, 0, tzinfo=timezone.utc), + end_time=datetime(2026, 1, 1, 11, 0, 0, tzinfo=timezone.utc), + frequency=frequency, + ) + + +@pytest.mark.asyncio +async def test_deliver_posts_to_gcs_upload_endpoint(): + """deliver() must POST raw bytes to the GCS upload endpoint.""" + from litellm.integrations.focus.destinations.gcs_destination import ( + FocusGCSDestination, + ) + + dest = FocusGCSDestination( + prefix="focus_exports", + config={"bucket_name": "my-bucket", "service_account_json": None}, + ) + + mock_response = MagicMock() + mock_response.status_code = 200 + + mock_client = MagicMock() + mock_client.post = AsyncMock(return_value=mock_response) + dest.async_httpx_client = mock_client + + with patch.object( + dest, + "construct_request_headers", + new=AsyncMock(return_value={"Authorization": "Bearer tok-123"}), + ): + await dest.deliver( + content=b"col1,col2\nval1,val2\n", + time_window=_make_window(), + filename="usage_20260101T100000Z_20260101T110000Z.csv", + ) + + mock_client.post.assert_called_once() + call_kwargs = mock_client.post.call_args + url = call_kwargs.kwargs.get("url") or call_kwargs.args[0] + assert "my-bucket" in url + assert "uploadType=media" in url + headers = call_kwargs.kwargs["headers"] + assert headers["Authorization"] == "Bearer tok-123" + + +@pytest.mark.asyncio +async def test_deliver_raises_on_gcs_error(): + """deliver() must raise RuntimeError when GCS returns non-200.""" + from litellm.integrations.focus.destinations.gcs_destination import ( + FocusGCSDestination, + ) + + dest = FocusGCSDestination( + prefix="focus_exports", + config={"bucket_name": "my-bucket"}, + ) + + mock_response = MagicMock() + mock_response.status_code = 403 + mock_response.text = "Permission denied" + + mock_client = MagicMock() + mock_client.post = AsyncMock(return_value=mock_response) + dest.async_httpx_client = mock_client + + with patch.object( + dest, + "construct_request_headers", + new=AsyncMock(return_value={"Authorization": "Bearer tok-bad"}), + ): + with pytest.raises(RuntimeError, match="GCS upload failed"): + await dest.deliver( + content=b"data", + time_window=_make_window(), + filename="usage.csv", + ) + + +def test_build_object_key_hourly(): + """Hourly key must include date= and hour= components.""" + from litellm.integrations.focus.destinations.gcs_destination import ( + FocusGCSDestination, + ) + + dest = FocusGCSDestination(prefix="focus_exports", config={"bucket_name": "b"}) + key = dest._build_object_key( + time_window=_make_window("hourly"), filename="usage.parquet" + ) + + assert key == "focus_exports/date=2026-01-01/hour=10/usage.parquet" + + +def test_build_object_key_daily(): + """Daily key must include date= but not hour=.""" + from litellm.integrations.focus.destinations.gcs_destination import ( + FocusGCSDestination, + ) + + dest = FocusGCSDestination(prefix="focus_exports", config={"bucket_name": "b"}) + window = FocusTimeWindow( + start_time=datetime(2026, 1, 1, 0, 0, 0, tzinfo=timezone.utc), + end_time=datetime(2026, 1, 2, 0, 0, 0, tzinfo=timezone.utc), + frequency="daily", + ) + key = dest._build_object_key(time_window=window, filename="usage.parquet") + + assert key == "focus_exports/date=2026-01-01/usage.parquet" + + +def test_missing_bucket_name_raises(): + """Constructing without bucket_name must raise ValueError.""" + from litellm.integrations.focus.destinations.gcs_destination import ( + FocusGCSDestination, + ) + + with pytest.raises(ValueError, match="bucket_name"): + FocusGCSDestination(prefix="focus_exports", config={}) + + +def test_global_gcs_service_account_not_overwritten_when_absent(monkeypatch): + """service_account_json absent from config must not overwrite GCS_PATH_SERVICE_ACCOUNT. + + GCSBucketBase sets self.path_service_account_json from GCS_PATH_SERVICE_ACCOUNT. + If config has no service_account_json key, we must leave the parent value intact + so deployments using the global credential don't silently fall back to ADC. + """ + monkeypatch.setenv("GCS_PATH_SERVICE_ACCOUNT", "/global/sa.json") + + from litellm.integrations.focus.destinations.gcs_destination import ( + FocusGCSDestination, + ) + + dest = FocusGCSDestination(prefix="focus_exports", config={"bucket_name": "b"}) + + assert dest.path_service_account_json == "/global/sa.json" + + +def test_explicit_service_account_overrides_global(monkeypatch): + """Explicit service_account_json in config must take precedence over GCS_PATH_SERVICE_ACCOUNT.""" + monkeypatch.setenv("GCS_PATH_SERVICE_ACCOUNT", "/global/sa.json") + + from litellm.integrations.focus.destinations.gcs_destination import ( + FocusGCSDestination, + ) + + dest = FocusGCSDestination( + prefix="focus_exports", + config={"bucket_name": "b", "service_account_json": "/focus/sa.json"}, + ) + + assert dest.path_service_account_json == "/focus/sa.json" + + +def test_factory_creates_gcs_destination(monkeypatch): + """FocusDestinationFactory.create(provider='gcs') must return FocusGCSDestination.""" + monkeypatch.setenv("FOCUS_GCS_BUCKET_NAME", "env-bucket") + + from litellm.integrations.focus.destinations.factory import FocusDestinationFactory + from litellm.integrations.focus.destinations.gcs_destination import ( + FocusGCSDestination, + ) + + dest = FocusDestinationFactory.create(provider="gcs", prefix="focus_exports") + + assert isinstance(dest, FocusGCSDestination) + assert dest.BUCKET_NAME == "env-bucket" diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py index a04f6407e4b..c43291566b6 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py @@ -1,17 +1,14 @@ -import json import os import sys -from unittest.mock import MagicMock import pytest -from fastapi.testclient import TestClient import litellm from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( StandardBuiltInToolCostTracking, ) from litellm.types.llms.openai import FileSearchTool, WebSearchOptions -from litellm.types.utils import ModelInfo, ModelResponse, StandardBuiltInToolsParams +from litellm.types.utils import ModelResponse, StandardBuiltInToolsParams sys.path.insert( 0, os.path.abspath("../../..") @@ -139,6 +136,22 @@ def test_get_cost_for_anthropic_web_search(): assert cost > 0.0 +def test_get_cost_for_anthropic_web_search_with_server_tool_use_dict(): + """ + Anthropic-compatible passthrough responses can construct Usage from a raw + usage payload. Ensure dict server_tool_use values are normalized before + built-in tool cost tracking reads server_tool_use.web_search_requests. + """ + from litellm.types.utils import ServerToolUse, Usage + + usage = Usage(server_tool_use={"web_search_requests": 1}) + + assert isinstance(usage.server_tool_use, ServerToolUse) + assert StandardBuiltInToolCostTracking.response_object_includes_web_search_call( + response_object=None, usage=usage + ) + + @pytest.mark.parametrize( "model", ["gemini/gemini-2.0-flash-001", "gemini-2.0-flash-001"] ) diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py index 4c330312930..75038574c63 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -5092,6 +5092,175 @@ def test_map_tool_helper_collision_prefers_definitions_over_components_schemas() assert transformed["input_schema"]["properties"]["from_components"] == expected +BILLING_HEADER_BLOCK = { + "type": "text", + "text": "x-anthropic-billing-header: cc_version=1.0.abc; cc_entrypoint=cli; cch=00000;", +} + + +def _system_with_billing_header(real_text: str) -> list: + return [ + { + "role": "system", + "content": [BILLING_HEADER_BLOCK, {"type": "text", "text": real_text}], + } + ] + + +def test_translate_system_message_keeps_billing_header_for_first_party_anthropic(): + config = AnthropicConfig() + assert config.should_strip_billing_metadata() is False + + result = config.translate_system_message( + messages=_system_with_billing_header( + "You are Claude Code, Anthropic's official CLI for Claude." + ) + ) + + texts = [block["text"] for block in result] + assert any(t.startswith("x-anthropic-billing-header:") for t in texts) + assert "You are Claude Code, Anthropic's official CLI for Claude." in texts + + +def test_translate_system_message_strips_billing_header_for_bedrock(): + from litellm.llms.bedrock.claude_platform.transformation import ( + BedrockClaudePlatformConfig, + ) + + config = BedrockClaudePlatformConfig() + assert config.should_strip_billing_metadata() is True + + result = config.translate_system_message( + messages=_system_with_billing_header("real system prompt") + ) + + texts = [block["text"] for block in result] + assert all(not t.startswith("x-anthropic-billing-header:") for t in texts) + assert "real system prompt" in texts + + +def test_anthropic_messages_request_keeps_billing_header_for_first_party(): + from litellm.types.router import GenericLiteLLMParams + + config = AnthropicMessagesConfig() + assert config.should_strip_billing_metadata() is False + + optional_params = { + "max_tokens": 16, + "system": [ + BILLING_HEADER_BLOCK, + {"type": "text", "text": "real system prompt"}, + ], + } + result = config.transform_anthropic_messages_request( + model="claude-3-5-sonnet-latest", + messages=[{"role": "user", "content": "hi"}], + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + texts = [block["text"] for block in result["system"]] + assert any(t.startswith("x-anthropic-billing-header:") for t in texts) + + +def test_anthropic_messages_request_strips_billing_header_for_minimax(): + from litellm.llms.minimax.messages.transformation import MinimaxMessagesConfig + from litellm.types.router import GenericLiteLLMParams + + config = MinimaxMessagesConfig() + assert config.should_strip_billing_metadata() is True + + optional_params = { + "max_tokens": 16, + "system": [ + BILLING_HEADER_BLOCK, + {"type": "text", "text": "real system prompt"}, + ], + } + result = config.transform_anthropic_messages_request( + model="MiniMax-M2", + messages=[{"role": "user", "content": "hi"}], + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + texts = [block["text"] for block in result.get("system", [])] + assert all(not t.startswith("x-anthropic-billing-header:") for t in texts) + + +def test_translate_system_message_strips_billing_header_for_bedrock_invoke(): + from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import ( + AmazonAnthropicClaudeConfig, + ) + + config = AmazonAnthropicClaudeConfig() + assert config.should_strip_billing_metadata() is True + + result = config.translate_system_message( + messages=_system_with_billing_header("real system prompt") + ) + + texts = [block["text"] for block in result] + assert all(not t.startswith("x-anthropic-billing-header:") for t in texts) + assert "real system prompt" in texts + + +@pytest.mark.parametrize( + "module_path, class_name, expected_strip", + [ + ("litellm.llms.anthropic.chat.transformation", "AnthropicConfig", False), + ( + "litellm.llms.anthropic.experimental_pass_through.messages.transformation", + "AnthropicMessagesConfig", + False, + ), + ( + "litellm.llms.bedrock.claude_platform.transformation", + "BedrockClaudePlatformConfig", + True, + ), + ( + "litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation", + "AmazonAnthropicClaudeConfig", + True, + ), + ( + "litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation", + "VertexAIAnthropicConfig", + True, + ), + ( + "litellm.llms.azure_ai.anthropic.transformation", + "AzureAnthropicConfig", + True, + ), + ("litellm.llms.minimax.messages.transformation", "MinimaxMessagesConfig", True), + ( + "litellm.llms.azure_ai.anthropic.messages_transformation", + "AzureAnthropicMessagesConfig", + True, + ), + ( + "litellm.llms.deepseek.messages.transformation", + "DeepSeekAnthropicMessagesConfig", + True, + ), + ( + "litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.experimental_pass_through.transformation", + "VertexAIPartnerModelsAnthropicMessagesConfig", + True, + ), + ], +) +def test_should_strip_billing_metadata_by_provider( + module_path, class_name, expected_strip +): + import importlib + + config_cls = getattr(importlib.import_module(module_path), class_name) + assert config_cls().should_strip_billing_metadata() is expected_strip def test_namespace_tool_flat_nested_tools_are_extracted(): """Codex sends nested tools in flat format {type, name, description, parameters} with no 'function' wrapper. These must be normalized and mapped without raising KeyError: 'function'.""" diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index a6aa35ee6d1..ed978113b8b 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -1467,11 +1467,10 @@ def test_transform_request_with_function_tool(): ) # Verify the structure - assert "additionalModelRequestFields" in request_data - additional_fields = request_data["additionalModelRequestFields"] + # Function tools are not computer use tools, so they don't get anthropic_beta — + # additionalModelRequestFields should be absent (not serialized as empty {}) + assert "additionalModelRequestFields" not in request_data - # Function tools are not computer use tools, so they don't get anthropic_beta - # They are processed through the regular tool config assert "toolConfig" in request_data assert "tools" in request_data["toolConfig"] assert len(request_data["toolConfig"]["tools"]) == 1 diff --git a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py index ba41fc47e8b..4731be13e78 100644 --- a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py @@ -128,6 +128,70 @@ class TestBedrockFilesTransformation: # Must have messages assert "messages" in model_input + # Nova Pro rejects empty additionalModelRequestFields / system — they must be absent + assert ( + "additionalModelRequestFields" not in model_input + ), "Nova: empty additionalModelRequestFields must be omitted, not serialized as {}" + assert ( + "system" not in model_input + ), "Nova: empty system must be omitted, not serialized as []" + + def test_nova_batch_jsonl_omits_empty_converse_fields(self): + """ + Regression test: Amazon Nova Pro returns 400 Malformed input request when + additionalModelRequestFields or system are present but empty in the Converse + API payload. The proxy must strip these keys when they carry no data. + """ + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + config = BedrockFilesConfig() + + openai_jsonl_content = [ + { + "custom_id": "req-0", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": "us.amazon.nova-pro-v1:0", + "messages": [ + { + "role": "user", + "content": "What is 1 + 1? Answer with just the number.", + } + ], + "max_tokens": 16, + }, + } + ] + + result = config._transform_openai_jsonl_content_to_bedrock_jsonl_content( + openai_jsonl_content + ) + + assert len(result) == 1 + model_input = result[0]["modelInput"] + + assert ( + "additionalModelRequestFields" not in model_input + or model_input["additionalModelRequestFields"] + ), "additionalModelRequestFields must be absent or non-empty — Nova rejects {}" + assert ( + "system" not in model_input or model_input["system"] + ), "system must be absent or non-empty — Nova rejects []" + + # Validate the exact shape AWS accepts + assert model_input == { + "messages": [ + { + "role": "user", + "content": [ + {"text": "What is 1 + 1? Answer with just the number."} + ], + } + ], + "inferenceConfig": {"maxTokens": 16}, + } + def test_nova_image_content_uses_converse_image_blocks(self): """ Test that image_url content blocks are converted to Bedrock Converse diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py index e2133d56f89..92b5ca7b10b 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py @@ -12,6 +12,11 @@ import sys sys.path.insert(0, os.path.abspath("../../../../..")) import pytest +from botocore.exceptions import ( + ConnectTimeoutError, + PartialCredentialsError, + ProfileNotFound, +) import litellm from litellm.llms.bedrock_mantle.responses.transformation import ( @@ -114,16 +119,15 @@ class TestBedrockMantleResponsesAuth: ) assert headers["Authorization"] == "Bearer bearer-key" - def test_missing_key_raises(self, monkeypatch): + def test_missing_bearer_does_not_raise_in_validate_environment(self, monkeypatch): + # SigV4 may still apply, so validate_environment must defer instead of raising. monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) cfg = BedrockMantleResponsesAPIConfig() - with pytest.raises(ValueError, match="Bedrock Mantle API key"): - cfg.validate_environment( - headers={}, - model="openai.gpt-5.5", - litellm_params=GenericLiteLLMParams(), - ) + headers = cfg.validate_environment( + headers={}, model="openai.gpt-5.5", litellm_params=GenericLiteLLMParams() + ) + assert "Authorization" not in headers def test_custom_llm_provider(self): cfg = BedrockMantleResponsesAPIConfig() @@ -261,6 +265,386 @@ def local_cost_map(monkeypatch): litellm.get_model_info.cache_clear() +class TestBedrockMantleResponsesSigV4: + def test_bearer_short_circuits_without_credentials(self, monkeypatch): + from unittest.mock import MagicMock + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + + signer = BaseAWSLLM() + signer.get_credentials = MagicMock( + side_effect=AssertionError("get_credentials must not run for bearer auth") + ) + cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer) + + headers, signed_body = cfg.sign_request( + headers={}, + optional_params={}, + request_data={"input": "hi"}, + api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses", + api_key="bearer-from-config", + ) + assert headers["Authorization"] == "Bearer bearer-from-config" + assert signed_body == b'{"input": "hi"}' + signer.get_credentials.assert_not_called() + + def test_bearer_resolved_from_mantle_env_key(self, monkeypatch): + from unittest.mock import MagicMock + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.setenv("BEDROCK_MANTLE_API_KEY", "env-bearer") + + signer = BaseAWSLLM() + signer.get_credentials = MagicMock( + side_effect=AssertionError("get_credentials must not run for bearer auth") + ) + cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer) + + headers, _ = cfg.sign_request( + headers={}, + optional_params={}, + request_data={"input": "hi"}, + api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses", + api_key=None, + ) + assert headers["Authorization"] == "Bearer env-bearer" + + def test_bearer_arg_takes_priority_over_mantle_env_key(self, monkeypatch): + # The passed api_key (e.g. litellm_params.api_key) must win over the env + # bearer; a reordered precedence chain would silently use the wrong token. + from unittest.mock import MagicMock + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.setenv("BEDROCK_MANTLE_API_KEY", "env-bearer") + + signer = BaseAWSLLM() + signer.get_credentials = MagicMock( + side_effect=AssertionError("get_credentials must not run for bearer auth") + ) + cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer) + + headers, _ = cfg.sign_request( + headers={}, + optional_params={}, + request_data={"input": "hi"}, + api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses", + api_key="arg-bearer", + ) + assert headers["Authorization"] == "Bearer arg-bearer" + signer.get_credentials.assert_not_called() + + def test_access_key_produces_sigv4_headers(self, monkeypatch): + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + + cfg = BedrockMantleResponsesAPIConfig(aws_signer=BaseAWSLLM()) + headers, signed_body = cfg.sign_request( + headers={}, + optional_params={ + "aws_access_key_id": "AKIAEXAMPLE", + "aws_secret_access_key": "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0", + "aws_session_token": "session-token-test", + "aws_region_name": "us-east-2", + }, + request_data={"input": "hi"}, + api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses", + api_key=None, + ) + assert headers["Authorization"].startswith("AWS4-HMAC-SHA256") + assert "Credential=AKIAEXAMPLE/" in headers["Authorization"] + assert "/us-east-2/bedrock/aws4_request" in headers["Authorization"] + assert "X-Amz-Date" in headers + assert headers["X-Amz-Security-Token"] == "session-token-test" + assert signed_body == b'{"input": "hi"}' + + def test_assume_role_path_produces_sigv4_headers(self, monkeypatch): + from unittest.mock import MagicMock + from botocore.credentials import Credentials + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + + signer = BaseAWSLLM() + signer.get_credentials = MagicMock( + return_value=Credentials( + access_key="ASIAEXAMPLE", + secret_key="YXNzdW1lZC1yb2xlLXNlY3JldC1hc3N1bWVk", + token="assumed-session-token", + ) + ) + cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer) + + headers, _ = cfg.sign_request( + headers={}, + optional_params={ + "aws_role_name": "arn:aws:iam::000000000000:role/test-role", + "aws_session_name": "litellm-test", + "aws_region_name": "us-east-2", + }, + request_data={"input": "hi"}, + api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses", + api_key=None, + ) + signer.get_credentials.assert_called_once() + call = signer.get_credentials.call_args.kwargs + assert call["aws_role_name"] == "arn:aws:iam::000000000000:role/test-role" + assert call["aws_session_name"] == "litellm-test" + assert headers["Authorization"].startswith("AWS4-HMAC-SHA256") + assert "/us-east-2/bedrock/aws4_request" in headers["Authorization"] + + def test_signed_body_matches_final_data_after_normalize(self, monkeypatch): + """Core regression: the signed bytes must equal the bytes actually sent. + + Sign the *final* data dict and assert the returned signed_body decodes to + exactly that dict, so a later change to the data would break the SigV4 hash. + """ + import json + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + + final_data = {"model": "openai.gpt-5.5", "input": "hi", "max_output_tokens": 16} + cfg = BedrockMantleResponsesAPIConfig(aws_signer=BaseAWSLLM()) + _, signed_body = cfg.sign_request( + headers={}, + optional_params={ + "aws_access_key_id": "AKIAEXAMPLE", + "aws_secret_access_key": "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0", + "aws_region_name": "us-east-2", + }, + request_data=final_data, + api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses", + api_key=None, + ) + assert signed_body is not None + assert json.loads(signed_body) == final_data + + def test_region_comes_from_optional_params(self, monkeypatch): + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + monkeypatch.delenv("AWS_REGION", raising=False) + monkeypatch.delenv("AWS_REGION_NAME", raising=False) + + cfg = BedrockMantleResponsesAPIConfig(aws_signer=BaseAWSLLM()) + headers, _ = cfg.sign_request( + headers={}, + optional_params={ + "aws_access_key_id": "AKIAEXAMPLE", + "aws_secret_access_key": "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0", + "aws_region_name": "eu-west-1", + }, + request_data={"input": "hi"}, + api_base="https://bedrock-mantle.eu-west-1.api.aws/openai/v1/responses", + api_key=None, + ) + assert "/eu-west-1/bedrock/aws4_request" in headers["Authorization"] + + def test_url_region_and_sigv4_region_agree_from_litellm_params(self, monkeypatch): + """Adversarial-review regression: a caller-supplied aws_region_name (no region + env set) must shape BOTH the URL host and the SigV4 credential scope, or the + request is signed for one region and sent to another -> 401. + """ + monkeypatch.delenv("BEDROCK_MANTLE_REGION", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) + monkeypatch.delenv("AWS_REGION", raising=False) + monkeypatch.delenv("AWS_REGION_NAME", raising=False) + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + params = { + "aws_region_name": "ap-southeast-2", + "aws_access_key_id": "AKIAEXAMPLE", + "aws_secret_access_key": "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0", + } + cfg = BedrockMantleResponsesAPIConfig(aws_signer=BaseAWSLLM()) + url = cfg.get_complete_url(api_base=None, litellm_params=params) + assert ( + url == "https://bedrock-mantle.ap-southeast-2.api.aws/openai/v1/responses" + ) + + headers, _ = cfg.sign_request( + headers={}, + optional_params=params, + request_data={"input": "hi"}, + api_base=url, + api_key=None, + ) + assert "/ap-southeast-2/bedrock/aws4_request" in headers["Authorization"] + + def test_injected_default_region_base_does_not_override_aws_region_name( + self, monkeypatch + ): + """2nd-round adversarial regression: responses/main.py auto-injects + litellm_params.api_base = https://bedrock-mantle..api.aws/v1 (default + region, ignoring aws_region_name). The config must still pin BOTH the URL host + and the SigV4 scope to aws_region_name, or the IAM deployment 401s. A naive + 'resolve region only when api_base is None' fix would fail this test. + """ + monkeypatch.delenv("BEDROCK_MANTLE_REGION", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) + monkeypatch.delenv("AWS_REGION", raising=False) + monkeypatch.delenv("AWS_REGION_NAME", raising=False) + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + injected_base = "https://bedrock-mantle.us-east-1.api.aws/v1" # default region + params = { + "aws_region_name": "us-east-2", # what the caller actually wants + "api_base": injected_base, + "aws_access_key_id": "AKIAEXAMPLE", + "aws_secret_access_key": "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0", + } + cfg = BedrockMantleResponsesAPIConfig(aws_signer=BaseAWSLLM()) + url = cfg.get_complete_url(api_base=injected_base, litellm_params=params) + assert url == "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses" + + headers, _ = cfg.sign_request( + headers={}, + optional_params=params, + request_data={"input": "hi"}, + api_base=url, + api_key=None, + ) + assert "/us-east-2/bedrock/aws4_request" in headers["Authorization"] + assert "us-east-1" not in headers["Authorization"] + + def test_custom_proxy_host_is_preserved(self, monkeypatch): + """A genuinely custom (non-Mantle) api_base host must be preserved, not rewritten + to a bedrock-mantle host. Only standard Mantle hosts are region-pinned. + """ + monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) + cfg = BedrockMantleResponsesAPIConfig() + url = cfg.get_complete_url( + api_base="https://mantle-proxy.internal.example/openai/v1", + litellm_params={"aws_region_name": "us-east-2"}, + ) + assert url == "https://mantle-proxy.internal.example/openai/v1/responses" + + def test_caller_authorization_does_not_override_sigv4(self, monkeypatch): + """Adversarial-review regression: a caller-supplied Authorization header (e.g. + from extra_headers, surviving the relaxed validate_environment) must not clobber + the SigV4 Authorization that _sign_request would otherwise restore. + """ + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + cfg = BedrockMantleResponsesAPIConfig(aws_signer=BaseAWSLLM()) + headers, _ = cfg.sign_request( + headers={"Authorization": "Bearer stale-caller-token"}, + optional_params={ + "aws_access_key_id": "AKIAEXAMPLE", + "aws_secret_access_key": "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0", + "aws_region_name": "us-east-2", + }, + request_data={"input": "hi"}, + api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses", + api_key=None, + ) + assert headers["Authorization"].startswith("AWS4-HMAC-SHA256") + assert "Bearer stale-caller-token" not in headers["Authorization"] + + def test_no_bearer_and_no_credentials_raises_both_paths(self, monkeypatch): + from unittest.mock import MagicMock + from botocore.exceptions import NoCredentialsError + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + + signer = BaseAWSLLM() + signer.get_credentials = MagicMock(side_effect=NoCredentialsError()) + cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer) + + with pytest.raises(ValueError) as exc: + cfg.sign_request( + headers={}, + optional_params={"aws_region_name": "us-east-2"}, + request_data={"input": "hi"}, + api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses", + api_key=None, + ) + msg = str(exc.value) + assert "Bearer" in msg + assert "SigV4" in msg or "IAM" in msg + + @pytest.mark.parametrize( + "cred_error", + [ + PartialCredentialsError(provider="env", cred_var="aws_secret_access_key"), + ProfileNotFound(profile="missing-profile"), + ], + ) + def test_partial_credentials_raises_both_paths(self, monkeypatch, cred_error): + from unittest.mock import MagicMock + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + + signer = BaseAWSLLM() + signer.get_credentials = MagicMock(side_effect=cred_error) + cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer) + + with pytest.raises(ValueError) as exc: + cfg.sign_request( + headers={}, + optional_params={"aws_region_name": "us-east-2"}, + request_data={"input": "hi"}, + api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses", + api_key=None, + ) + msg = str(exc.value) + assert "Bearer" in msg + assert "SigV4" in msg or "IAM" in msg + + def test_sts_transport_error_is_not_masked_as_credentials(self, monkeypatch): + # An AssumeRole / web-identity flow hits STS over the network, so a transient + # connection error must surface as itself, not be rewritten into the + # "no usable AWS credentials" message that would send the user to fix the + # wrong thing. + from unittest.mock import MagicMock + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + monkeypatch.delenv("BEDROCK_MANTLE_API_KEY", raising=False) + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + + signer = BaseAWSLLM() + signer.get_credentials = MagicMock( + side_effect=ConnectTimeoutError( + endpoint_url="https://sts.us-east-2.amazonaws.com" + ) + ) + cfg = BedrockMantleResponsesAPIConfig(aws_signer=signer) + + with pytest.raises(ConnectTimeoutError): + cfg.sign_request( + headers={}, + optional_params={ + "aws_role_name": "arn:aws:iam::000000000000:role/test-role", + "aws_region_name": "us-east-2", + }, + request_data={"input": "hi"}, + api_base="https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses", + api_key=None, + ) + + class TestBedrockMantleResponsesPricing: def test_gpt_5_5_pricing_and_mode(self, local_cost_map): info = litellm.get_model_info("bedrock_mantle/openai.gpt-5.5") diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index 279e9730e69..7321abcee46 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -742,3 +742,241 @@ async def test_anthropic_post_retry_reserializes_mutated_body(): assert first_sent == prebuilt # attempt 0 used prebuilt assert second_sent == _json.dumps(request_body) # attempt 1 re-serialized assert "MUTATED" in second_sent # ... the mutated body + + +def test_base_responses_config_sign_request_is_noop_by_default(): + """Default responses sign_request must be a no-op: unchanged headers, no signed body. + + Guards the 15 existing responses providers from accidental signing when the + handler starts calling sign_request. + """ + from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig + + cfg = OpenAIResponsesAPIConfig() + headers = {"Authorization": "Bearer sk-existing"} + out_headers, signed_body = cfg.sign_request( + headers=headers, + optional_params={}, + request_data={"input": "hi"}, + api_base="https://api.openai.com/v1/responses", + ) + assert out_headers == {"Authorization": "Bearer sk-existing"} + assert signed_body is None + + +def _make_responses_handler_call(signed_body): + """Drive BaseLLMHTTPHandler.response_api_handler with a fully mocked provider + config + sync client, returning the kwargs the client.post was called with. + + signed_body=None simulates a no-op (non-signing) provider; bytes simulates a + signing provider (e.g. Bedrock Mantle). + """ + from unittest.mock import MagicMock + from litellm.llms.custom_httpx.http_handler import HTTPHandler + from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler + from litellm.types.router import GenericLiteLLMParams + + provider_config = MagicMock() + provider_config.validate_environment.return_value = {} + provider_config.get_complete_url.return_value = ( + "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses" + ) + provider_config.transform_responses_api_request.return_value = {"input": "hi"} + provider_config.should_fake_stream.return_value = False + provider_config.sign_request.return_value = ({"X-Signed": "1"}, signed_body) + + mock_client = MagicMock(spec=HTTPHandler) + mock_client.post.return_value = MagicMock() + + handler = BaseLLMHTTPHandler() + handler.response_api_handler( + model="openai.gpt-5.5", + input="hi", + responses_api_provider_config=provider_config, + response_api_optional_request_params={}, + custom_llm_provider="bedrock_mantle", + litellm_params=GenericLiteLLMParams(aws_region_name="us-east-2"), + logging_obj=MagicMock(), + client=mock_client, + _is_async=False, + ) + return mock_client.post.call_args.kwargs + + +def test_responses_handler_sends_json_when_not_signed(): + """No-op provider (signed_body is None) -> handler posts json=data, no data= bytes.""" + kwargs = _make_responses_handler_call(signed_body=None) + assert kwargs.get("json") == {"input": "hi"} + assert "data" not in kwargs + + +def test_responses_handler_sends_signed_bytes_when_signed(): + """Signing provider -> handler posts the exact signed bytes via data=, not json=.""" + kwargs = _make_responses_handler_call(signed_body=b'{"input": "hi"}') + assert kwargs.get("data") == b'{"input": "hi"}' + assert "json" not in kwargs + assert kwargs["headers"] == {"X-Signed": "1"} + + +def test_responses_handler_signs_after_fake_stream_prep_strips_stream(): + """Fake-stream signing-order invariant: the bytes SIGNED must equal the bytes SENT. + + In the streaming + fake-stream path the handler first runs + _prepare_fake_stream_request, which pops "stream" out of the body, and only + then calls sign_request. If signing ran before that pop, the signed body + would still carry "stream" while the body sent over the wire would not, + producing a SigV4 payload-hash mismatch (401) for a real Mantle deployment. + We snapshot request_data at sign time and assert "stream" is already gone. + """ + from unittest.mock import MagicMock + from litellm.llms.custom_httpx.http_handler import HTTPHandler + from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler + from litellm.types.llms.openai import ResponsesAPIResponse + from litellm.types.router import GenericLiteLLMParams + + provider_config = MagicMock() + provider_config.validate_environment.return_value = {} + provider_config.get_complete_url.return_value = ( + "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses" + ) + provider_config.transform_responses_api_request.return_value = { + "input": "hi", + "stream": True, + } + provider_config.should_fake_stream.return_value = True + provider_config.transform_response_api_response.return_value = ResponsesAPIResponse( + id="resp_1", + created_at=0, + output=[], + status="completed", + model="openai.gpt-5.5", + ) + + captured = {} + + def _capture_sign(**kwargs): + captured["request_data"] = dict(kwargs["request_data"]) + return ({"X-Signed": "1"}, b'{"input": "hi"}') + + provider_config.sign_request.side_effect = _capture_sign + + mock_client = MagicMock(spec=HTTPHandler) + mock_client.post.return_value = MagicMock() + + handler = BaseLLMHTTPHandler() + handler.response_api_handler( + model="openai.gpt-5.5", + input="hi", + responses_api_provider_config=provider_config, + response_api_optional_request_params={"stream": True}, + custom_llm_provider="bedrock_mantle", + litellm_params=GenericLiteLLMParams(aws_region_name="us-east-2"), + logging_obj=MagicMock(), + client=mock_client, + _is_async=False, + fake_stream=True, + ) + + assert "stream" not in captured["request_data"] + assert "input" in captured["request_data"] + + post_kwargs = mock_client.post.call_args.kwargs + assert post_kwargs.get("data") == b'{"input": "hi"}' + assert "json" not in post_kwargs + assert "stream" in post_kwargs + + +def _make_compact_handler_call(signed_body, is_async): + """Drive (async_)compact_response_api_handler with a fully mocked provider config + + client, returning the kwargs the client.post was called with. + + signed_body=None simulates a no-op (non-signing) provider; bytes simulates a + signing provider (e.g. Bedrock Mantle SigV4 / bearer). + """ + from unittest.mock import MagicMock + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler + from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler + from litellm.types.router import GenericLiteLLMParams + + compact_url = "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses/compact" + provider_config = MagicMock() + provider_config.validate_environment.return_value = {} + provider_config.get_complete_url.return_value = ( + "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses" + ) + provider_config.transform_compact_response_api_request.return_value = ( + compact_url, + {"model": "openai.gpt-5.5", "input": "hi"}, + ) + provider_config.sign_request.return_value = ({"X-Signed": "1"}, signed_body) + provider_config.transform_compact_response_api_response.return_value = "ok" + + spec = AsyncHTTPHandler if is_async else HTTPHandler + mock_client = MagicMock(spec=spec) + if is_async: + mock_client.post = AsyncMock(return_value=MagicMock()) + else: + mock_client.post.return_value = MagicMock() + + handler = BaseLLMHTTPHandler() + result = handler.compact_response_api_handler( + model="openai.gpt-5.5", + input="hi", + responses_api_provider_config=provider_config, + response_api_optional_request_params={}, + custom_llm_provider="bedrock_mantle", + litellm_params=GenericLiteLLMParams(aws_region_name="us-east-2"), + logging_obj=MagicMock(), + client=mock_client, + _is_async=is_async, + ) + if is_async: + asyncio.run(result) + return provider_config, mock_client.post.call_args.kwargs + + +def test_compact_handler_sends_json_when_not_signed(): + """No-op provider on compact (signed_body is None) -> posts json=data, no data= bytes.""" + provider_config, kwargs = _make_compact_handler_call( + signed_body=None, is_async=False + ) + provider_config.sign_request.assert_called_once() + assert kwargs.get("json") == {"model": "openai.gpt-5.5", "input": "hi"} + assert "data" not in kwargs + + +def test_compact_handler_sends_signed_bytes_when_signed(): + """Signing provider on compact -> posts the signed bytes via data=, not json=. + + Regression for the adversarial-review finding that /responses/compact bypassed + the SigV4 signing hook, so IAM-only Mantle callers sent unsigned bodies. + """ + provider_config, kwargs = _make_compact_handler_call( + signed_body=b'{"model": "openai.gpt-5.5", "input": "hi"}', is_async=False + ) + assert kwargs.get("data") == b'{"model": "openai.gpt-5.5", "input": "hi"}' + assert "json" not in kwargs + assert kwargs["headers"] == {"X-Signed": "1"} + # signing must use the compact endpoint as api_base, not the create URL + assert provider_config.sign_request.call_args.kwargs["api_base"].endswith( + "/openai/v1/responses/compact" + ) + + +def test_async_compact_handler_sends_signed_bytes_when_signed(): + """Async compact must sign identically to sync (same omission in the async twin).""" + provider_config, kwargs = _make_compact_handler_call( + signed_body=b'{"model": "openai.gpt-5.5", "input": "hi"}', is_async=True + ) + assert kwargs.get("data") == b'{"model": "openai.gpt-5.5", "input": "hi"}' + assert "json" not in kwargs + assert kwargs["headers"] == {"X-Signed": "1"} + + +def test_async_compact_handler_sends_json_when_not_signed(): + """Async no-op provider on compact -> posts json=data, no data= bytes.""" + _provider_config, kwargs = _make_compact_handler_call( + signed_body=None, is_async=True + ) + assert kwargs.get("json") == {"model": "openai.gpt-5.5", "input": "hi"} + assert "data" not in kwargs diff --git a/tests/test_litellm/llms/gemini/test_gemini_tts.py b/tests/test_litellm/llms/gemini/test_gemini_tts.py index 65eefca5af1..98f3ac0f4e5 100644 --- a/tests/test_litellm/llms/gemini/test_gemini_tts.py +++ b/tests/test_litellm/llms/gemini/test_gemini_tts.py @@ -80,6 +80,46 @@ class TestGeminiTTSTransformation: assert "responseModalities" in result assert "AUDIO" in result["responseModalities"] + def test_gemini_tts_audio_parameter_mapping_with_language_code(self): + config = GoogleAIStudioGeminiConfig() + + non_default_params = { + "audio": {"voice": "Kore", "format": "pcm16", "language_code": "en-US"} + } + optional_params = {} + + result = config.map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model="gemini-2.5-flash-preview-tts", + drop_params=False, + ) + + assert "speechConfig" in result + assert result["speechConfig"]["languageCode"] == "en-US" + assert ( + result["speechConfig"]["voiceConfig"]["prebuiltVoiceConfig"]["voiceName"] + == "Kore" + ) + + def test_map_audio_params_language_code(self): + config = GoogleAIStudioGeminiConfig() + + result = config._map_audio_params( + {"voice": "Kore", "format": "pcm16", "language_code": "de-DE"} + ) + + assert result["languageCode"] == "de-DE" + assert result["voiceConfig"]["prebuiltVoiceConfig"]["voiceName"] == "Kore" + + def test_map_audio_params_no_language_code(self): + config = GoogleAIStudioGeminiConfig() + + result = config._map_audio_params({"voice": "Kore", "format": "pcm16"}) + + assert "languageCode" not in result + assert result["voiceConfig"]["prebuiltVoiceConfig"]["voiceName"] == "Kore" + def test_gemini_tts_audio_parameter_with_existing_modalities(self): """Test audio parameter mapping when modalities already exist""" config = GoogleAIStudioGeminiConfig() @@ -328,5 +368,57 @@ class TestGeminiTTSSpeechConfigInRequestBody: assert "AUDIO" in generation_config["responseModalities"] + @pytest.mark.parametrize( + "model,custom_llm_provider", + [ + ("gemini-2.5-flash-tts", "vertex_ai"), + ("gemini-2.5-flash-tts", "gemini"), + ("gemini-2.5-flash-preview-tts", "vertex_ai"), + ], + ) + def test_language_code_end_to_end_mapping(self, model, custom_llm_provider): + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexGeminiConfig, + ) + from litellm.llms.vertex_ai.gemini.transformation import ( + _transform_request_body, + ) + + config = VertexGeminiConfig() + + non_default_params = { + "audio": {"voice": "Puck", "format": "pcm16", "language_code": "pt-BR"} + } + optional_params = {} + + mapped_params = config.map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model=model, + drop_params=False, + ) + + assert mapped_params["speechConfig"]["languageCode"] == "pt-BR" + + request_body = _transform_request_body( + messages=[{"role": "user", "content": "Hello world"}], + model=model, + optional_params=mapped_params, + custom_llm_provider=custom_llm_provider, + litellm_params={}, + cached_content=None, + ) + + generation_config = request_body["generationConfig"] + assert generation_config["speechConfig"]["languageCode"] == "pt-BR" + assert ( + generation_config["speechConfig"]["voiceConfig"]["prebuiltVoiceConfig"][ + "voiceName" + ] + == "Puck" + ) + assert "AUDIO" in generation_config["responseModalities"] + + if __name__ == "__main__": pytest.main([__file__]) diff --git a/tests/test_litellm/llms/github_copilot/responses/test_github_copilot_responses_transformation.py b/tests/test_litellm/llms/github_copilot/responses/test_github_copilot_responses_transformation.py index 54e7170bb20..17373f24a97 100644 --- a/tests/test_litellm/llms/github_copilot/responses/test_github_copilot_responses_transformation.py +++ b/tests/test_litellm/llms/github_copilot/responses/test_github_copilot_responses_transformation.py @@ -14,6 +14,8 @@ from unittest.mock import patch, MagicMock sys.path.insert(0, os.path.abspath("../../../../..")) import pytest +import litellm +from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map from litellm.types.utils import LlmProviders from litellm.utils import ProviderConfigManager from litellm.llms.github_copilot.responses.transformation import ( @@ -22,13 +24,26 @@ from litellm.llms.github_copilot.responses.transformation import ( from litellm.types.llms.openai import ResponsesAPIOptionalRequestParams +@pytest.fixture(autouse=True) +def use_local_model_cost_map(monkeypatch: pytest.MonkeyPatch): + """Pin litellm.model_cost to the bundled local backup so tests don't depend + on remote catalog fetches (and don't change behavior across remote refreshes).""" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr( + litellm, "model_cost", get_model_cost_map(url=litellm.model_cost_map_url) + ) + litellm.add_known_models(model_cost_map=litellm.model_cost) + + class TestGithubCopilotResponsesAPITransformation: """Test GitHub Copilot Responses API configuration and transformations""" def test_github_copilot_provider_config_registration(self): - """Test that GitHub Copilot provider returns GithubCopilotResponsesAPIConfig""" + """Test that GitHub Copilot provider returns the native Responses API + config for a Responses-capable catalog model. Exercises the full stack: + catalog lookup -> github_copilot_supports_responses_api -> native config.""" config = ProviderConfigManager.get_provider_responses_api_config( - model="github_copilot/gpt-5.1-codex", + model="github_copilot/gpt-5.3-codex", provider=LlmProviders.GITHUB_COPILOT, ) @@ -373,3 +388,200 @@ class TestGithubCopilotResponsesAPITransformation: # Non-reasoning items should pass through unchanged assert result == message_item + + +class TestGithubCopilotResponsesAPIRouting: + """``ProviderConfigManager.get_provider_responses_api_config`` for github_copilot + returns the native Responses config only when the model has ``mode=responses`` + in the (already-merged) model info; otherwise returns None so the dispatcher + routes through the chat-completions translation bridge.""" + + @patch( + "litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper" + ) + def test_returns_config_when_mode_is_responses(self, mock_get_info): + """``mode=responses`` returns native config.""" + mock_get_info.return_value = {"mode": "responses"} + config = ProviderConfigManager.get_provider_responses_api_config( + model="github_copilot/some-responses-model", + provider=LlmProviders.GITHUB_COPILOT, + ) + assert isinstance(config, GithubCopilotResponsesAPIConfig) + + @patch( + "litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper" + ) + def test_returns_none_when_mode_is_chat(self, mock_get_info): + """``mode=chat`` returns None so dispatcher uses bridge.""" + mock_get_info.return_value = {"mode": "chat"} + config = ProviderConfigManager.get_provider_responses_api_config( + model="github_copilot/some-chat-only-model", + provider=LlmProviders.GITHUB_COPILOT, + ) + assert config is None + + @patch( + "litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper" + ) + def test_returns_none_when_mode_is_unset_and_no_endpoints(self, mock_get_info): + """Entry without ``mode`` and without ``supported_endpoints`` returns None + (conservative default).""" + mock_get_info.return_value = {} + config = ProviderConfigManager.get_provider_responses_api_config( + model="github_copilot/some-model", + provider=LlmProviders.GITHUB_COPILOT, + ) + assert config is None + + def test_returns_config_when_mode_unset_but_endpoints_have_responses(self): + """``mode`` unset but ``supported_endpoints`` declaring /v1/responses + returns native config (endpoint-list fallback for stale-but-correct + catalog entries that lack ``mode``). + + Exercises the real ``_cached_get_model_info_helper`` plumbing via + ``register_model`` (no mock). ``supported_endpoints`` is not carried on + the normalized ``ModelInfoBase`` the helper returns, so the gate must + read it from the raw ``litellm.model_cost`` entry; a mock-based test + would mask that. + """ + litellm.register_model( + { + "github_copilot/test-endpoints-only-model": { + "litellm_provider": "github_copilot", + "max_tokens": 1, + "input_cost_per_token": 0, + "output_cost_per_token": 0, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + ], + } + } + ) + config = ProviderConfigManager.get_provider_responses_api_config( + model="github_copilot/test-endpoints-only-model", + provider=LlmProviders.GITHUB_COPILOT, + ) + assert isinstance(config, GithubCopilotResponsesAPIConfig) + + def test_mode_chat_overrides_endpoints_with_responses(self): + """``mode=chat`` is a hard opt-out: forces bridge even when + ``supported_endpoints`` includes /v1/responses. Lets users force the + bridge for dual-endpoint models without clearing endpoint metadata. + + Exercises the real ``_cached_get_model_info_helper`` plumbing via + ``register_model`` (no mock) so the ``mode``-over-endpoints precedence + is verified against the actual model-info resolution. + """ + litellm.register_model( + { + "github_copilot/test-chat-override-model": { + "litellm_provider": "github_copilot", + "max_tokens": 1, + "input_cost_per_token": 0, + "output_cost_per_token": 0, + "mode": "chat", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + ], + } + } + ) + config = ProviderConfigManager.get_provider_responses_api_config( + model="github_copilot/test-chat-override-model", + provider=LlmProviders.GITHUB_COPILOT, + ) + assert config is None + + def test_returns_config_when_model_is_none(self): + """Follow-up GET/DELETE operations pass model=None and keep the native + config path (no per-model lookup is possible).""" + config = ProviderConfigManager.get_provider_responses_api_config( + model=None, + provider=LlmProviders.GITHUB_COPILOT, + ) + assert isinstance(config, GithubCopilotResponsesAPIConfig) + + @patch( + "litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper" + ) + def test_returns_none_when_get_model_info_raises(self, mock_get_info): + """Catalog lookup failure (model not registered) returns None + (conservative default; bridge handles unknown models safely).""" + mock_get_info.side_effect = Exception("model not in catalog") + config = ProviderConfigManager.get_provider_responses_api_config( + model="github_copilot/never-seen-model", + provider=LlmProviders.GITHUB_COPILOT, + ) + assert config is None + + @patch( + "litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper" + ) + def test_user_override_via_register_model(self, mock_get_info): + """User-supplied per-deployment ``model_info`` flows through + ``litellm.register_model`` (called by the router) into the merged + catalog read by ``_cached_get_model_info_helper``. Setting ``mode=responses`` + for a model whose catalog entry says ``mode=chat`` therefore opts in + to native dispatch without any per-call argument plumbing.""" + mock_get_info.return_value = {"mode": "responses"} + config = ProviderConfigManager.get_provider_responses_api_config( + model="github_copilot/some-chat-only-model", + provider=LlmProviders.GITHUB_COPILOT, + ) + assert isinstance(config, GithubCopilotResponsesAPIConfig) + + @patch( + "litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper" + ) + def test_realistic_chat_only_entry_returns_none(self, mock_get_info): + """Realistic ``model_prices_and_context_window.json`` shape for a + chat-only Copilot model (e.g. github_copilot/gemini-3.1-pro-preview) + returns None so /v1/responses calls fall back to the bridge.""" + mock_get_info.return_value = { + "litellm_provider": "github_copilot", + "max_input_tokens": 136000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "supported_endpoints": ["/v1/chat/completions"], + "supports_function_calling": True, + "supports_tool_choice": True, + "supports_parallel_function_calling": True, + "supports_vision": True, + "supports_reasoning": True, + } + config = ProviderConfigManager.get_provider_responses_api_config( + model="github_copilot/some-chat-only-model", + provider=LlmProviders.GITHUB_COPILOT, + ) + assert config is None + + @patch( + "litellm.llms.github_copilot.responses.transformation._cached_get_model_info_helper" + ) + def test_realistic_responses_only_entry_returns_config(self, mock_get_info): + """Realistic catalog entry for a Responses-only Copilot model + (e.g. github_copilot/gpt-5.5) returns the native config.""" + mock_get_info.return_value = { + "litellm_provider": "github_copilot", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "responses", + "supported_endpoints": ["/v1/responses"], + "supports_function_calling": True, + "supports_tool_choice": True, + "supports_parallel_function_calling": True, + "supports_response_schema": True, + "supports_vision": True, + "supports_reasoning": True, + "supports_none_reasoning_effort": True, + "supports_xhigh_reasoning_effort": True, + } + config = ProviderConfigManager.get_provider_responses_api_config( + model="github_copilot/some-responses-only-model", + provider=LlmProviders.GITHUB_COPILOT, + ) + assert isinstance(config, GithubCopilotResponsesAPIConfig) diff --git a/tests/test_litellm/llms/parasail/test_parasail.py b/tests/test_litellm/llms/parasail/test_parasail.py new file mode 100644 index 00000000000..8fb9b22b5f6 --- /dev/null +++ b/tests/test_litellm/llms/parasail/test_parasail.py @@ -0,0 +1,172 @@ +import os +from unittest.mock import patch + +PARASAIL_API_BASE = "https://api.parasail.io/v1" +PARASAIL_RESPONSES_GATEWAY = "https://api-webflux.saas.parasail.io/v1" + + +def test_parasail_json_registry(): + import litellm + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + assert litellm.LlmProviders.PARASAIL.value == "parasail" + assert litellm.LlmProviders("parasail") == litellm.LlmProviders.PARASAIL + assert JSONProviderRegistry.exists("parasail") + config = JSONProviderRegistry.get("parasail") + assert config is not None + assert config.base_url == PARASAIL_API_BASE + assert config.api_key_env == "PARASAIL_API_KEY" + assert config.api_base_env == "PARASAIL_API_BASE" + assert "/v1/chat/completions" in config.supported_endpoints + assert "/v1/responses" in config.supported_endpoints + assert config.special_handling.get("force_store_false") is True + + +def test_parasail_listed_in_openai_compatible_providers(): + from litellm.constants import openai_compatible_providers + + assert "parasail" in openai_compatible_providers + + +def test_parasail_dynamic_config_env_vars(): + from litellm.llms.openai_like.dynamic_config import create_config_class + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + config = create_config_class(JSONProviderRegistry.get("parasail"))() + + with patch.dict( + os.environ, + { + "PARASAIL_API_KEY": "test-key", + "PARASAIL_API_BASE": PARASAIL_RESPONSES_GATEWAY, + }, + ): + api_base, api_key = config._get_openai_compatible_provider_info(None, None) + + assert api_base == PARASAIL_RESPONSES_GATEWAY + assert api_key == "test-key" + + +def test_parasail_provider_detection_by_prefix(): + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + model, provider, _, api_base = get_llm_provider( + "parasail/parasail-llama-33-70b-fp8" + ) + + assert model == "parasail-llama-33-70b-fp8" + assert provider == "parasail" + assert api_base == PARASAIL_API_BASE + + +def test_parasail_chat_complete_url(): + from litellm.llms.openai_like.dynamic_config import create_config_class + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + config = create_config_class(JSONProviderRegistry.get("parasail"))() + + assert ( + config.get_complete_url( + api_base=None, + api_key=None, + model="parasail-llama-33-70b-fp8", + optional_params={}, + litellm_params={}, + ) + == f"{PARASAIL_API_BASE}/chat/completions" + ) + + +def test_parasail_responses_api_config(): + from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig + from litellm.utils import ProviderConfigManager + + config = ProviderConfigManager.get_provider_responses_api_config( + provider="parasail", + model="parasail-kimi-k25-elicit", + ) + + assert isinstance(config, OpenAIResponsesAPIConfig) + assert config.custom_llm_provider == "parasail" + assert ( + config.get_complete_url(api_base=None, litellm_params={}) + == f"{PARASAIL_API_BASE}/responses" + ) + + +def test_parasail_responses_api_honors_api_base_override(): + from litellm.utils import ProviderConfigManager + + config = ProviderConfigManager.get_provider_responses_api_config( + provider="parasail", + model="parasail-kimi-k25-elicit", + ) + + with patch.dict( + os.environ, + {"PARASAIL_API_BASE": PARASAIL_RESPONSES_GATEWAY}, + ): + url = config.get_complete_url(api_base=None, litellm_params={}) + + assert url == f"{PARASAIL_RESPONSES_GATEWAY}/responses" + + +def test_parasail_responses_api_forces_store_false_when_caller_sets_true(): + from litellm.types.router import GenericLiteLLMParams + from litellm.utils import ProviderConfigManager + + config = ProviderConfigManager.get_provider_responses_api_config( + provider="parasail", + model="parasail-kimi-k25-elicit", + ) + + request_params: dict = {"store": True, "temperature": 0.2} + transformed = config.transform_responses_api_request( + model="parasail-kimi-k25-elicit", + input="hello", + response_api_optional_request_params=request_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert transformed["store"] is False + assert transformed["temperature"] == 0.2 + + +def test_parasail_responses_api_forces_store_false_when_caller_omits_store(): + from litellm.types.router import GenericLiteLLMParams + from litellm.utils import ProviderConfigManager + + config = ProviderConfigManager.get_provider_responses_api_config( + provider="parasail", + model="parasail-kimi-k25-elicit", + ) + + transformed = config.transform_responses_api_request( + model="parasail-kimi-k25-elicit", + input="hello", + response_api_optional_request_params={}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert transformed["store"] is False + + +def test_parasail_responses_api_validate_environment_sets_bearer_token(): + from litellm.types.router import GenericLiteLLMParams + from litellm.utils import ProviderConfigManager + + config = ProviderConfigManager.get_provider_responses_api_config( + provider="parasail", + model="parasail-kimi-k25-elicit", + ) + + with patch.dict(os.environ, {"PARASAIL_API_KEY": "secret-from-env"}): + headers = config.validate_environment( + headers={}, + model="parasail-kimi-k25-elicit", + litellm_params=GenericLiteLLMParams(), + ) + + assert headers["Authorization"] == "Bearer secret-from-env" diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index fa6cc8bed1b..0236646c796 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -112,11 +112,71 @@ async def test_should_clear_stale_budget_reservation_when_budget_checks_skip(): user_api_key_cache=MagicMock(), proxy_logging_obj=MagicMock(), skip_budget_checks=True, + general_settings={}, ) assert user_api_key_auth_obj.budget_reservation is None +@pytest.mark.asyncio +async def test_disable_budget_reservation_skips_reservation(): + """#27639: general_settings.disable_budget_reservation turns off the optimistic Redis + reservation so operators hit by phantom BudgetExceededError can opt out of it.""" + user_api_key_auth_obj = UserAPIKeyAuth(token="test_token") + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.reserve_budget_for_request", + new=AsyncMock(return_value={"reserved_cost": 0.5, "entries": []}), + ) as mock_reserve: + await _reserve_budget_after_common_checks( + user_api_key_auth_obj=user_api_key_auth_obj, + request_data={"model": "gpt-4o"}, + route="/v1/chat/completions", + llm_router=None, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + skip_budget_checks=False, + general_settings={"disable_budget_reservation": True}, + ) + + mock_reserve.assert_not_called() + assert user_api_key_auth_obj.budget_reservation is None + + +@pytest.mark.asyncio +async def test_budget_reservation_runs_when_not_disabled(): + """Control for #27639: with the flag absent, the reservation still runs and is stored.""" + user_api_key_auth_obj = UserAPIKeyAuth(token="test_token") + reservation = { + "reserved_cost": 0.5, + "entries": [{"counter_key": "spend:key:test_token"}], + } + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.reserve_budget_for_request", + new=AsyncMock(return_value=reservation), + ) as mock_reserve: + await _reserve_budget_after_common_checks( + user_api_key_auth_obj=user_api_key_auth_obj, + request_data={"model": "gpt-4o"}, + route="/v1/chat/completions", + llm_router=None, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + skip_budget_checks=False, + general_settings={}, + ) + + mock_reserve.assert_awaited_once() + assert user_api_key_auth_obj.budget_reservation == reservation + + @pytest.mark.asyncio async def test_should_not_reuse_cached_key_object_for_request_state(): key_cache = DualCache() diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py index c58c94cbbc7..e7f72ff7a3f 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py @@ -41,7 +41,8 @@ def test_crowdstrike_aidr_guardrail_config() -> None: ) -def test_crowdstrike_aidr_guardrail_config_no_api_key() -> None: +def test_crowdstrike_aidr_guardrail_config_no_api_key(monkeypatch) -> None: + monkeypatch.delenv("CS_AIDR_TOKEN", raising=False) with pytest.raises(CrowdStrikeAIDRGuardrailMissingSecrets): init_guardrails_v2( all_guardrails=[ @@ -59,7 +60,8 @@ def test_crowdstrike_aidr_guardrail_config_no_api_key() -> None: ) -def test_crowdstrike_aidr_guardrail_config_no_api_base() -> None: +def test_crowdstrike_aidr_guardrail_config_no_api_base(monkeypatch) -> None: + monkeypatch.delenv("CS_AIDR_BASE_URL", raising=False) with pytest.raises(CrowdStrikeAIDRGuardrailMissingSecrets): init_guardrails_v2( all_guardrails=[ @@ -412,6 +414,121 @@ async def test_apply_guardrail_response_ok( assert result["texts"] == inputs["texts"] +@pytest.mark.asyncio +async def test_apply_guardrail_sends_user_id_model_and_extra_info( + crowdstrike_aidr_guardrail: CrowdStrikeAIDRHandler, +) -> None: + inputs: GenericGuardrailAPIInputs = { + "texts": ["Hello"], + "structured_messages": [{"role": "user", "content": "Hello"}], + "model": "gpt-4o", + } + request_data = { + "messages": inputs["structured_messages"], + "model": "gpt-4o", + "litellm_metadata": { + "user_api_key_user_id": "uid-abc", + "user_api_key_user_email": "alice@example.com", + }, + } + guardrail_endpoint = ( + f"{crowdstrike_aidr_guardrail.api_base}/v1/guard_chat_completions" + ) + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + return_value=httpx.Response( + status_code=200, + json={"result": {"blocked": False, "transformed": False}}, + request=httpx.Request(method="POST", url=guardrail_endpoint), + ), + ) as mock_method: + await crowdstrike_aidr_guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + ) + + payload = mock_method.call_args.kwargs["json"] + assert payload["user_id"] == "uid-abc" + assert payload["model"] == "gpt-4o" + assert payload["extra_info"] == {"user_name": "alice@example.com"} + + +@pytest.mark.asyncio +async def test_apply_guardrail_empty_extra_info_when_no_email( + crowdstrike_aidr_guardrail: CrowdStrikeAIDRHandler, +) -> None: + inputs: GenericGuardrailAPIInputs = { + "texts": ["Hello"], + "structured_messages": [{"role": "user", "content": "Hello"}], + "model": "gemini-flash", + } + request_data = { + "messages": inputs["structured_messages"], + "model": "gemini-flash", + "litellm_metadata": { + "user_api_key_user_id": "uid-no-email", + "user_api_key_user_email": None, + }, + } + guardrail_endpoint = ( + f"{crowdstrike_aidr_guardrail.api_base}/v1/guard_chat_completions" + ) + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + return_value=httpx.Response( + status_code=200, + json={"result": {"blocked": False, "transformed": False}}, + request=httpx.Request(method="POST", url=guardrail_endpoint), + ), + ) as mock_method: + await crowdstrike_aidr_guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + ) + + payload = mock_method.call_args.kwargs["json"] + assert payload["user_id"] == "uid-no-email" + assert payload["model"] == "gemini-flash" + assert payload["extra_info"] == {} + + +@pytest.mark.asyncio +async def test_apply_guardrail_no_metadata_skips_user_fields( + crowdstrike_aidr_guardrail: CrowdStrikeAIDRHandler, +) -> None: + inputs: GenericGuardrailAPIInputs = { + "texts": ["Hello"], + "structured_messages": [{"role": "user", "content": "Hello"}], + } + request_data = {"messages": inputs["structured_messages"]} + guardrail_endpoint = ( + f"{crowdstrike_aidr_guardrail.api_base}/v1/guard_chat_completions" + ) + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + return_value=httpx.Response( + status_code=200, + json={"result": {"blocked": False, "transformed": False}}, + request=httpx.Request(method="POST", url=guardrail_endpoint), + ), + ) as mock_method: + await crowdstrike_aidr_guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + ) + + payload = mock_method.call_args.kwargs["json"] + assert "user_id" not in payload + assert "model" not in payload + assert "extra_info" not in payload + + @pytest.mark.asyncio async def test_apply_guardrail_request_skipped_messages_stay_aligned( crowdstrike_aidr_guardrail: CrowdStrikeAIDRHandler, diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py index f8b6fbde3dc..1114b3df0c2 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py @@ -321,6 +321,44 @@ class TestAzureAnthropicCostCalculation: assert call_kwargs["model"] == "azure_ai/claude-sonnet-4-5_gb_20250929" assert call_kwargs["custom_llm_provider"] == "azure_ai" + def test_passthrough_logging_sets_response_cost_with_server_tool_use_dict(self): + from litellm.types.utils import Choices, Message, ModelResponse + + logging_obj = self._create_mock_logging_obj(model="claude-3-7-sonnet-20250219") + logging_obj.get_router_model_id.return_value = None + logging_obj.litellm_params = {} + + response = ModelResponse( + id="test-id", + choices=[ + Choices( + finish_reason="stop", + index=0, + message=Message(content="test", role="assistant"), + ) + ], + created=1234567890, + model="claude-3-7-sonnet-20250219", + usage={ + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15, + "server_tool_use": {"web_search_requests": 1}, + }, + ) + + kwargs = AnthropicPassthroughLoggingHandler._create_anthropic_response_logging_payload( + litellm_model_response=response, + model="claude-3-7-sonnet-20250219", + kwargs={}, + start_time=datetime.now(), + end_time=datetime.now(), + logging_obj=logging_obj, + ) + + assert "response_cost" in kwargs + assert kwargs["response_cost"] > 0 + class TestAnthropicBatchPassthroughCostTracking: """Test cases for Anthropic batch passthrough cost tracking functionality""" diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py index bfcaaafd335..3c6af3e528a 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py @@ -257,6 +257,64 @@ class TestOpenAIPassthroughLoggingHandler: ) assert OpenAIPassthroughLoggingHandler.is_openai_responses_route("") == False + def test_is_openai_route_recognizes_cognitiveservices_azure_com(self): + """Azure OpenAI resources created via the newer "Azure AI Foundry" / + Cognitive Services pathway live on `*.cognitiveservices.azure.com` + subdomains rather than the older `openai.azure.com`. All four + is_openai_*_route methods must recognize both Azure subdomains so + cost tracking applies regardless of which Azure naming the user's + resource happens to be on. + """ + cognitive_chat = ( + "https://my-resource.cognitiveservices.azure.com/v1/chat/completions" + ) + cognitive_images_gen = ( + "https://my-resource.cognitiveservices.azure.com/v1/images/generations" + ) + cognitive_images_edit = ( + "https://my-resource.cognitiveservices.azure.com/v1/images/edits" + ) + cognitive_responses = ( + "https://my-resource.cognitiveservices.azure.com/v1/responses" + ) + + assert ( + OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route( + cognitive_chat + ) + is True + ) + assert ( + OpenAIPassthroughLoggingHandler.is_openai_image_generation_route( + cognitive_images_gen + ) + is True + ) + assert ( + OpenAIPassthroughLoggingHandler.is_openai_image_editing_route( + cognitive_images_edit + ) + is True + ) + assert ( + OpenAIPassthroughLoggingHandler.is_openai_responses_route( + cognitive_responses + ) + is True + ) + + # Cross-route negatives still hold for cognitiveservices hosts. + assert ( + OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route( + cognitive_responses + ) + is False + ) + assert ( + OpenAIPassthroughLoggingHandler.is_openai_responses_route(cognitive_chat) + is False + ) + @patch("litellm.completion_cost") @patch( "litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload" @@ -766,6 +824,14 @@ class TestOpenAIPassthroughIntegration: == True ) assert self.handler.is_openai_route("https://api.openai.com/v1/models") == True + # Azure OpenAI on the shared Cognitive Services domain, identified by an + # OpenAI-style path segment. + assert ( + self.handler.is_openai_route( + "https://my-resource.cognitiveservices.azure.com/v1/chat/completions" + ) + == True + ) # Negative cases assert ( @@ -782,6 +848,28 @@ class TestOpenAIPassthroughIntegration: self.handler.is_openai_route("https://api.assemblyai.com/v2/transcript") == False ) + # Non-OpenAI Azure Cognitive Services share the `cognitiveservices.azure.com` + # domain but must NOT be classified as OpenAI routes (no OpenAI path segment). + assert ( + self.handler.is_openai_route( + "https://my-resource.cognitiveservices.azure.com/speechtotext/v3.1/recognize" + ) + == False + ) + assert ( + self.handler.is_openai_route( + "https://my-resource.cognitiveservices.azure.com/vision/v3.2/analyze" + ) + == False + ) + # A look-alike domain that merely contains an OpenAI host as a substring + # must be rejected by the suffix-based hostname match. + assert ( + self.handler.is_openai_route( + "https://cognitiveservices.azure.com.attacker.example/v1/chat/completions" + ) + == False + ) assert self.handler.is_openai_route("") == False @patch( diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 503a610e016..960fca205ce 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -949,6 +949,28 @@ class TestToolChoiceTransformation: result = LiteLLMCompletionResponsesConfig._transform_tool_choice(tool_choice) assert result == tool_choice + def test_transform_tool_choice_responses_flat_function_name(self): + """Responses-API forced-function with a top-level name maps to the nested Chat + Completions shape instead of degrading to required and dropping the name""" + result = LiteLLMCompletionResponsesConfig._transform_tool_choice( + {"type": "function", "name": "get_weather"} + ) + assert result == {"type": "function", "function": {"name": "get_weather"}} + + def test_transform_tool_choice_function_without_name_falls_back_to_required(self): + """A function-type dict with no name still falls back to required""" + result = LiteLLMCompletionResponsesConfig._transform_tool_choice( + {"type": "function"} + ) + assert result == "required" + + def test_transform_tool_choice_function_empty_name_falls_back_to_required(self): + """An empty top-level name is falsy and must not produce an empty function name""" + result = LiteLLMCompletionResponsesConfig._transform_tool_choice( + {"type": "function", "name": ""} + ) + assert result == "required" + class TestContentTypeTransformation: """Test content type transformation from Responses API to Chat Completion format""" diff --git a/tests/test_litellm/types/test_types_utils.py b/tests/test_litellm/types/test_types_utils.py index c146847f391..a4074ccdaaa 100644 --- a/tests/test_litellm/types/test_types_utils.py +++ b/tests/test_litellm/types/test_types_utils.py @@ -1,13 +1,9 @@ -import asyncio import os import sys -from typing import Optional -from unittest.mock import AsyncMock, patch import pytest sys.path.insert(0, os.path.abspath("../..")) -import json from litellm.types.utils import HiddenParams @@ -75,6 +71,48 @@ def test_usage_dump(): assert new_usage.prompt_tokens_details.web_search_requests == 1 +def test_usage_server_tool_use_dict_is_coerced_and_round_trips(): + from litellm.types.utils import ServerToolUse, Usage + + current_usage = Usage( + completion_tokens=1, + prompt_tokens=1, + total_tokens=2, + server_tool_use={"web_search_requests": 1}, + ) + + assert isinstance(current_usage.server_tool_use, ServerToolUse) + assert current_usage.server_tool_use.web_search_requests == 1 + + new_usage = Usage(**current_usage.model_dump()) + assert isinstance(new_usage.server_tool_use, ServerToolUse) + assert new_usage.server_tool_use.web_search_requests == 1 + + +def test_usage_converts_server_tool_use_dict(): + from litellm.types.utils import ServerToolUse, Usage + + usage = Usage( + completion_tokens=2, + prompt_tokens=1, + total_tokens=3, + server_tool_use={"web_search_requests": 4, "tool_search_requests": 1}, + ) + + assert isinstance(usage.server_tool_use, ServerToolUse) + assert usage.server_tool_use.web_search_requests == 4 + assert usage.server_tool_use["web_search_requests"] == 4 + assert usage.server_tool_use.tool_search_requests == 1 + with pytest.raises(KeyError): + usage.server_tool_use["unknown_metric"] + + round_trip = Usage(**usage.model_dump()) + assert isinstance(round_trip.server_tool_use, ServerToolUse) + assert round_trip.server_tool_use.web_search_requests == 4 + assert round_trip.server_tool_use["web_search_requests"] == 4 + assert round_trip.server_tool_use.tool_search_requests == 1 + + def test_usage_completion_tokens_details_text_tokens(): from litellm.types.utils import Usage diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 75a14e09852..d9992680c8a 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -22034,6 +22034,11 @@ export interface components { * @description connect to a postgres db - needed for generating temporary keys + tracking spend / key */ database_url?: string | null; + /** + * Disable Budget Reservation + * @description If True, disables the optimistic per-request budget reservation introduced in v1.84.0. WARNING: This weakens hard budget enforcement. Without the reservation, a burst of concurrent requests from a single key can each pass the read-time spend check before any of them is charged, allowing a configured budget to be exceeded under high concurrency. Budgets are still evaluated on every request at read time, so an already-exhausted budget is still rejected. Enable only if your deployment is experiencing phantom BudgetExceededError responses caused by leaked reservations (see GitHub issue #27639). A proxy-level WARNING is logged on every request while this flag is active as a reminder that hard enforcement is relaxed. + */ + disable_budget_reservation?: boolean | null; /** * Enable Public Model Hub * @description Public model hub for users to see what models they have access to, supported openai params, etc. From dfb68a23de435afc4c2503c1caf552f1a116f154 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 9 Jun 2026 02:27:03 +0530 Subject: [PATCH 019/185] feat(galileo): add health check support for UI callback test (#29908) * feat(galileo): add health check support for UI callback test Register galileo in /health/services so the proxy UI callback connection test works. Co-authored-by: Cursor * feat(galileo): verify API key via /current_user health check Call Galileo's current_user endpoint so the UI callback test validates credentials against the provider. Co-authored-by: Cursor * chore(ui): regenerate schema.d.ts for galileo health service Co-authored-by: Cursor * fix(galileo): return IntegrationHealthCheckStatus from async_health_check Fixes mypy assignment error in health_services_endpoint where response was narrowed to IntegrationHealthCheckStatus from earlier branches. Co-authored-by: Cursor * Fix Galileo logging to match Langfuse across all endpoint types. Stop skipping ingest when output is empty and log embeddings with a placeholder so embedding, speech, and other non-text responses are recorded like Langfuse. Co-authored-by: Cursor * fix(galileo): remove unreachable health-check guard and None output sentinel The use_v2_api flag is derived from bool(api_key), so the inner GALILEO_API_KEY check inside the v2 branch could never run; collapse the credential validation into the username/password path with a combined message. _serialize_galileo_output now returns an empty string for None, so _get_galileo_input_output_content always yields a str and the post-call None coalescing guard is no longer needed. * test(galileo): cover async_health_check failure paths and empty model response Add regression tests for the Galileo health check unhealthy branches (missing project id, missing base url, missing credentials, auth failure, and request exception) and for logging a model response with no choices, which now queues an empty output instead of being skipped. --------- Co-authored-by: Cursor Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/integrations/galileo.py | 71 ++++-- .../health_endpoints/_health_endpoints.py | 15 ++ .../test_litellm/integrations/test_galileo.py | 206 +++++++++++++++++- .../health_endpoints/test_health_endpoints.py | 27 +++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 +- 5 files changed, 303 insertions(+), 18 deletions(-) diff --git a/litellm/integrations/galileo.py b/litellm/integrations/galileo.py index 8fef90c24e0..f9ff7e8c7a1 100644 --- a/litellm/integrations/galileo.py +++ b/litellm/integrations/galileo.py @@ -26,6 +26,7 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, ) +from litellm.types.integrations.base_health_check import IntegrationHealthCheckStatus GALILEO_CLOUD_API_BASE_URL = "https://api.galileo.ai" # Cap the in-memory buffer so persistent flush failures (e.g. Galileo @@ -89,6 +90,52 @@ class GalileoObserve(CustomLogger): return bool(self.api_key) return bool(self.username and self.password) + async def async_health_check(self) -> IntegrationHealthCheckStatus: + try: + if not self.project_id: + return IntegrationHealthCheckStatus( + status="unhealthy", + error_message="GALILEO_PROJECT_ID environment variable not set", + ) + + if not self.base_url: + return IntegrationHealthCheckStatus( + status="unhealthy", + error_message="GALILEO_BASE_URL environment variable not set", + ) + + if not self.use_v2_api and (not self.username or not self.password): + return IntegrationHealthCheckStatus( + status="unhealthy", + error_message=( + "GALILEO_API_KEY or GALILEO_USERNAME and GALILEO_PASSWORD " + "environment variables must be set" + ), + ) + + if not await self._ensure_headers(): + return IntegrationHealthCheckStatus( + status="unhealthy", + error_message="Galileo authentication failed", + ) + + response = await self.async_httpx_handler.get( + url=f"{self.base_url}/current_user", + headers=self.headers, + ) + if response.status_code >= 400: + return IntegrationHealthCheckStatus( + status="unhealthy", + error_message=(f"Galileo API returned HTTP {response.status_code}"), + ) + + return IntegrationHealthCheckStatus(status="healthy", error_message=None) + except Exception as e: + return IntegrationHealthCheckStatus( + status="unhealthy", + error_message=f"Galileo health check failed: {str(e)}", + ) + async def async_set_galileo_headers(self) -> None: galileo_login_response = await self.async_httpx_handler.post( url=f"{self.base_url}/login", @@ -399,9 +446,9 @@ class GalileoObserve(CustomLogger): return prompt @staticmethod - def _serialize_galileo_output(value: Any) -> Optional[str]: + def _serialize_galileo_output(value: Any) -> str: if value is None: - return None + return "" if isinstance(value, str): return value @@ -460,11 +507,11 @@ class GalileoObserve(CustomLogger): response_obj: Any, level: str = "DEFAULT", status_message: Optional[str] = None, - ) -> Tuple[str, Optional[str], Any]: + ) -> Tuple[str, str, Any]: """ Mirror Langfuse _get_langfuse_input_output_content for Galileo ingest. - Returns (input_text, output_text, messages_for_span). output_text None skips ingest. + Returns (input_text, output_text, messages_for_span). """ call_type = kwargs.get("call_type") prompt = self._build_prompt(kwargs) @@ -477,10 +524,11 @@ class GalileoObserve(CustomLogger): return self._prompt_to_input_text(prompt), status_message, prompt if response_obj is not None and ( - call_type == "embedding" + call_type in ("embedding", "aembedding") or isinstance(response_obj, litellm.EmbeddingResponse) ): - return self._prompt_to_input_text(prompt), None, prompt + # Match Langfuse OTEL: log embeddings without serializing vectors. + return self._prompt_to_input_text(prompt), "embedding-output", prompt if response_obj is not None and isinstance(response_obj, litellm.ModelResponse): output = self._get_chat_content_for_galileo(response_obj) @@ -549,7 +597,7 @@ class GalileoObserve(CustomLogger): ): input_val = kwargs.get("input") return ( - self._serialize_galileo_output(input_val) or "", + self._serialize_galileo_output(input_val), self._serialize_galileo_output(response_obj), input_val, ) @@ -574,11 +622,11 @@ class GalileoObserve(CustomLogger): kwargs.get("messages") or [], ) - return self._prompt_to_input_text(prompt), None, kwargs.get("messages") or [] + return self._prompt_to_input_text(prompt), "", kwargs.get("messages") or [] def get_output_str_from_response( self, response_obj: Any, kwargs: Dict[str, Any] - ) -> Optional[str]: + ) -> str: _, output_text, _ = self._get_galileo_input_output_content( kwargs=kwargs, response_obj=response_obj ) @@ -659,11 +707,6 @@ class GalileoObserve(CustomLogger): input_text, output_text, messages = self._get_galileo_input_output_content( kwargs=kwargs, response_obj=response_obj ) - if output_text is None: - verbose_logger.debug( - "Galileo Logger: skipping %s — no text output to log", _call_type - ) - return raw_start = slo.get("startTime") raw_end = slo.get("endTime") diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index c109f374993..6ef8bbc4006 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -129,6 +129,7 @@ services = Union[ "datadog_llm_observability", "generic_api", "arize", + "galileo", "sqs", ], str, @@ -206,6 +207,7 @@ async def health_services_endpoint( # noqa: PLR0915 "datadog_llm_observability", "generic_api", "arize", + "galileo", "sqs", ]: raise HTTPException( @@ -295,6 +297,19 @@ async def health_services_endpoint( # noqa: PLR0915 else "Arize is healthy" ), } + elif service == "galileo": + from litellm.integrations.galileo import GalileoObserve + + galileo_logger = GalileoObserve() + response = await galileo_logger.async_health_check() + return { + "status": response["status"], + "message": ( + response["error_message"] + if response["status"] == "unhealthy" + else "Galileo is healthy" + ), + } elif service == "langfuse": from litellm.integrations.langfuse.langfuse import LangFuseLogger diff --git a/tests/test_litellm/integrations/test_galileo.py b/tests/test_litellm/integrations/test_galileo.py index 8ce5eb776f9..0533b7ca7d1 100644 --- a/tests/test_litellm/integrations/test_galileo.py +++ b/tests/test_litellm/integrations/test_galileo.py @@ -357,12 +357,18 @@ def test_galileo_record_to_v2_span_with_tags_and_offset(): def test_galileo_get_output_str_variants(galileo_v2_env): logger = GalileoObserve() - assert logger.get_output_str_from_response(None, {}) is None + assert logger.get_output_str_from_response(None, {}) == "" assert ( logger.get_output_str_from_response( EmbeddingResponse(), {"call_type": "embedding"} ) - is None + == "embedding-output" + ) + assert ( + logger.get_output_str_from_response( + EmbeddingResponse(), {"call_type": "aembedding"} + ) + == "embedding-output" ) text_resp = TextCompletionResponse() @@ -414,7 +420,7 @@ def test_galileo_get_output_str_variants(galileo_v2_env): {"call_type": "acompletion", "messages": [{"role": "user", "content": "hi"}]}, ) - assert logger.get_output_str_from_response("not-a-supported-type", {}) is None + assert logger.get_output_str_from_response("not-a-supported-type", {}) == "" def test_galileo_get_input_output_error_status_message(galileo_v2_env): @@ -445,6 +451,48 @@ def test_galileo_get_output_str_rerank_response(galileo_v2_env): assert '"relevance_score": 0.98' in output +@pytest.mark.asyncio +async def test_galileo_async_log_success_embedding(galileo_v2_env): + import datetime + + logger = GalileoObserve() + embedding_response = EmbeddingResponse( + data=[{"object": "embedding", "embedding": [0.1, 0.2, 0.3], "index": 0}] + ) + + mock_response = MagicMock() + mock_response.is_success = True + mock_response.status_code = 201 + + with patch.object(logger.async_httpx_handler, "post", return_value=mock_response): + await logger.async_log_success_event( + kwargs={ + "call_type": "aembedding", + "model": "text-embedding-3-small", + "input": "hello world", + "standard_logging_object": { + "call_type": "aembedding", + "model": "text-embedding-3-small", + "prompt_tokens": 2, + "completion_tokens": 0, + "total_tokens": 2, + "response_cost": 0.0, + "startTime": datetime.datetime( + 2026, 5, 25, 12, 0, 0, tzinfo=datetime.timezone.utc + ).timestamp(), + "endTime": datetime.datetime( + 2026, 5, 25, 12, 0, 1, tzinfo=datetime.timezone.utc + ).timestamp(), + }, + }, + response_obj=embedding_response, + start_time=datetime.datetime(2026, 5, 25, 12, 0, 0), + end_time=datetime.datetime(2026, 5, 25, 12, 0, 1), + ) + + assert logger.in_memory_records == [] + + @pytest.mark.asyncio async def test_galileo_async_log_success_rerank(galileo_v2_env): import datetime @@ -524,6 +572,158 @@ def test_galileo_get_ingest_request_legacy(monkeypatch): assert payload["traces"][0]["input"] == "hi" +@pytest.mark.asyncio +async def test_galileo_async_health_check_success(galileo_v2_env): + logger = GalileoObserve() + current_user_resp = MagicMock() + current_user_resp.status_code = 200 + + with patch.object( + logger.async_httpx_handler, "get", new_callable=AsyncMock + ) as mock_get: + mock_get.return_value = current_user_resp + result = await logger.async_health_check() + + assert result["status"] == "healthy" + mock_get.assert_awaited_once_with( + url="https://api.galileo.ai/current_user", + headers={ + "accept": "application/json", + "Content-Type": "application/json", + "Galileo-API-Key": "test-api-key", + }, + ) + + +@pytest.mark.asyncio +async def test_galileo_async_health_check_api_error(galileo_v2_env): + logger = GalileoObserve() + current_user_resp = MagicMock() + current_user_resp.status_code = 401 + + with patch.object( + logger.async_httpx_handler, "get", new_callable=AsyncMock + ) as mock_get: + mock_get.return_value = current_user_resp + result = await logger.async_health_check() + + assert result["status"] == "unhealthy" + assert "HTTP 401" in result["error_message"] + + +@pytest.mark.asyncio +async def test_galileo_async_health_check_missing_project_id(monkeypatch): + monkeypatch.setenv("GALILEO_API_KEY", "test-api-key") + monkeypatch.setenv("GALILEO_BASE_URL", "https://api.galileo.ai") + monkeypatch.delenv("GALILEO_PROJECT_ID", raising=False) + logger = GalileoObserve() + + result = await logger.async_health_check() + + assert result["status"] == "unhealthy" + assert "GALILEO_PROJECT_ID" in result["error_message"] + + +@pytest.mark.asyncio +async def test_galileo_async_health_check_missing_base_url(monkeypatch): + monkeypatch.delenv("GALILEO_API_KEY", raising=False) + monkeypatch.delenv("GALILEO_BASE_URL", raising=False) + monkeypatch.setenv("GALILEO_PROJECT_ID", "p") + monkeypatch.setenv("GALILEO_USERNAME", "u") + monkeypatch.setenv("GALILEO_PASSWORD", "pw") + logger = GalileoObserve() + + result = await logger.async_health_check() + + assert result["status"] == "unhealthy" + assert "GALILEO_BASE_URL" in result["error_message"] + + +@pytest.mark.asyncio +async def test_galileo_async_health_check_missing_credentials(monkeypatch): + monkeypatch.delenv("GALILEO_API_KEY", raising=False) + monkeypatch.delenv("GALILEO_USERNAME", raising=False) + monkeypatch.delenv("GALILEO_PASSWORD", raising=False) + monkeypatch.setenv("GALILEO_PROJECT_ID", "p") + monkeypatch.setenv("GALILEO_BASE_URL", "https://galileo.example") + logger = GalileoObserve() + + result = await logger.async_health_check() + + assert result["status"] == "unhealthy" + assert "GALILEO_USERNAME" in result["error_message"] + + +@pytest.mark.asyncio +async def test_galileo_async_health_check_auth_failed(monkeypatch): + monkeypatch.delenv("GALILEO_API_KEY", raising=False) + monkeypatch.setenv("GALILEO_PROJECT_ID", "p") + monkeypatch.setenv("GALILEO_BASE_URL", "https://galileo.example") + monkeypatch.setenv("GALILEO_USERNAME", "u") + monkeypatch.setenv("GALILEO_PASSWORD", "pw") + logger = GalileoObserve() + + with patch.object( + logger.async_httpx_handler, "post", new_callable=AsyncMock + ) as mock_post: + mock_post.side_effect = Exception("login failed") + result = await logger.async_health_check() + + assert result["status"] == "unhealthy" + assert result["error_message"] == "Galileo authentication failed" + + +@pytest.mark.asyncio +async def test_galileo_async_health_check_request_exception(galileo_v2_env): + logger = GalileoObserve() + + with patch.object( + logger.async_httpx_handler, "get", new_callable=AsyncMock + ) as mock_get: + mock_get.side_effect = Exception("connection refused") + result = await logger.async_health_check() + + assert result["status"] == "unhealthy" + assert "connection refused" in result["error_message"] + + +@pytest.mark.asyncio +async def test_galileo_async_log_success_empty_model_response(galileo_v2_env): + import datetime + + logger = GalileoObserve() + logger.batch_size = 2 + empty_response = ModelResponse(choices=[]) + + await logger.async_log_success_event( + kwargs={ + "call_type": "acompletion", + "model": "gpt-5.2", + "messages": [{"role": "user", "content": "hi"}], + "standard_logging_object": { + "call_type": "acompletion", + "model": "gpt-5.2", + "prompt_tokens": 1, + "completion_tokens": 0, + "total_tokens": 1, + "response_cost": 0.0, + "startTime": datetime.datetime( + 2026, 5, 25, 12, 0, 0, tzinfo=datetime.timezone.utc + ).timestamp(), + "endTime": datetime.datetime( + 2026, 5, 25, 12, 0, 1, tzinfo=datetime.timezone.utc + ).timestamp(), + }, + }, + response_obj=empty_response, + start_time=datetime.datetime(2026, 5, 25, 12, 0, 0), + end_time=datetime.datetime(2026, 5, 25, 12, 0, 1), + ) + + assert len(logger.in_memory_records) == 1 + assert logger.in_memory_records[0]["output_text"] == "" + + @pytest.mark.asyncio async def test_galileo_ensure_headers_v2_missing_key(monkeypatch): monkeypatch.delenv("GALILEO_API_KEY", raising=False) diff --git a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py index 80a4804956c..64c57ab90e3 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -696,6 +696,33 @@ async def test_test_model_connection_falls_back_to_deployments_zero_without_id() assert model_params.get("api_key") == "fake-key-A" +@pytest.mark.asyncio +@pytest.mark.parametrize( + "status,error_message", + [ + ("healthy", ""), + ("unhealthy", "Galileo authentication failed"), + ], +) +async def test_health_services_endpoint_galileo(status, error_message): + with patch("litellm.integrations.galileo.GalileoObserve") as MockGalileoObserve: + mock_instance = MagicMock() + mock_instance.async_health_check = AsyncMock( + return_value={"status": status, "error_message": error_message} + ) + MockGalileoObserve.return_value = mock_instance + + result = await health_services_endpoint(service="galileo") + + if status == "healthy": + assert result["status"] == "healthy" + assert result["message"] == "Galileo is healthy" + else: + assert result["status"] == "unhealthy" + assert result["message"] == error_message + mock_instance.async_health_check.assert_awaited_once() + + @pytest.mark.asyncio async def test_health_services_endpoint_datadog_llm_observability(): """ diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index d9992680c8a..8379d0536a6 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -40422,7 +40422,7 @@ export interface operations { parameters: { query: { /** @description Specify the service being hit. */ - service: ("slack_budget_alerts" | "langfuse" | "langfuse_otel" | "slack" | "openmeter" | "webhook" | "email" | "braintrust" | "datadog" | "datadog_llm_observability" | "generic_api" | "arize" | "sqs") | string; + service: ("slack_budget_alerts" | "langfuse" | "langfuse_otel" | "slack" | "openmeter" | "webhook" | "email" | "braintrust" | "datadog" | "datadog_llm_observability" | "generic_api" | "arize" | "galileo" | "sqs") | string; }; header?: never; path?: never; From 69a7bdb24728eb13280bba5e4d7c3427e8ac0768 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 8 Jun 2026 14:28:39 -0700 Subject: [PATCH 020/185] fix(model-management): allow deleting a BYOK model after its team is deleted (#29875) * fix(model-management): allow deleting a BYOK model after its team is deleted A team BYOK model (model_info.team_id set) became undeletable once its team was deleted: POST /model/delete ran can_user_make_model_call, which looked the team up and raised 400 "Team id=... does not exist in db" before the delete could run, so the model lingered on the Models + Endpoints page with no way to remove it. Drop the team-existence prerequisite from the delete path. When the model's team still exists the normal auth check runs unchanged; when it is gone a proxy admin may delete the orphan and any other caller gets a 403. The check is fail-closed, so a missing or errored team lookup can only block the delete or require an admin, never grant a non-admin access. Add/update/health keep their team-existence validation. * refactor(model-management): drop redundant team lookup on model delete Move the orphaned-team handling into can_user_make_model_call behind an allow_missing_team flag instead of pre-checking team existence in delete_model. The endpoint no longer issues its own litellm_teamtable lookup, so deleting a model whose team still exists hits the team table once instead of twice. The auth behavior is unchanged: a proxy admin can delete a model whose team was deleted, any other caller gets a 403, and add/update/health keep the strict "team must exist" validation. --- .../model_management_endpoints.py | 14 ++ .../test_model_management_endpoints.py | 184 ++++++++++++++++++ 2 files changed, 198 insertions(+) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 566ef845333..0cbccfc18ad 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -860,6 +860,7 @@ class ModelManagementAuthChecks: user_api_key_dict: UserAPIKeyAuth, prisma_client: PrismaClient, premium_user: bool, + allow_missing_team: bool = False, ) -> Literal[True]: ## Check team model auth if ( @@ -870,6 +871,18 @@ class ModelManagementAuthChecks: where={"team_id": model_params.model_info.team_id} ) if team_obj_row is None: + # The team was deleted. Callers that opt in (e.g. model deletion) may + # act on the orphaned model, but only as a proxy admin -- without the + # team there is no team-admin membership left to verify. + if allow_missing_team: + if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: + return True + raise HTTPException( + status_code=403, + detail={ + "error": "Only a proxy admin can delete a model whose team has been deleted." + }, + ) raise HTTPException( status_code=400, detail={ @@ -955,6 +968,7 @@ async def delete_model( user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, premium_user=premium_user, + allow_missing_team=True, ) # update DB diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index a9bb2b09a15..2ba604e5da6 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -1915,6 +1915,190 @@ class TestDeleteTeamBYOKModelGhost: mock_refresh.assert_not_awaited() +class TestDeleteModelTeamAuth: + """Team auth on the /model/delete path. + + A model added via /model/new with model_info.team_id is orphaned once its + team is deleted: can_user_make_model_call looked the team up and raised + 'Team id=... does not exist in db' before the delete could run, so the model + was undeletable from the Models + Endpoints page. Without the team, team-admin + membership can't be verified, so a proxy admin (and only a proxy admin) may + delete the orphan; a missing team must never let a non-admin through. The team + is also looked up exactly once -- the auth check must not add a second query. + """ + + def _orphaned_model_mocks(self, team_id, model_id): + db_row = LiteLLM_ProxyModelTable( + model_id=model_id, + model_name=f"model_name_{team_id}_abc-uuid", + litellm_params={"model": "openai/gpt-4.1-nano"}, + model_info={ + "id": model_id, + "team_id": team_id, + "team_public_model_name": "orphaned-gpt", + }, + created_by="admin", + updated_by="admin", + ) + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.litellm_proxymodeltable = AsyncMock() + mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock( + return_value=db_row + ) + mock_prisma.db.litellm_proxymodeltable.delete = AsyncMock(return_value=db_row) + mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + # The team is gone -> every team lookup returns None. + mock_prisma.db.litellm_teamtable = AsyncMock() + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=None) + mock_prisma.db.litellm_teamtable.update = AsyncMock() + mock_prisma.db.litellm_modeltable = AsyncMock() + mock_prisma.db.litellm_modeltable.find_many = AsyncMock(return_value=[]) + return mock_prisma + + @pytest.mark.asyncio + async def test_proxy_admin_can_delete_model_when_team_deleted(self): + from litellm.proxy.management_endpoints.model_management_endpoints import ( + ModelInfoDelete, + delete_model as delete_model_endpoint, + ) + + team_id = "deleted-team-xyz" + model_id = "orphaned-byok-1" + mock_prisma = self._orphaned_model_mocks(team_id, model_id) + + admin_user = UserAPIKeyAuth( + user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + _PS = "litellm.proxy.proxy_server" + _MOD = "litellm.proxy.management_endpoints.model_management_endpoints" + with ( + patch(f"{_PS}.prisma_client", mock_prisma), + patch(f"{_PS}.store_model_in_db", True), + patch(f"{_PS}.premium_user", True), + patch(f"{_PS}.llm_router", MagicMock()), + patch(f"{_PS}.proxy_logging_obj", MagicMock()), + patch(f"{_PS}.user_api_key_cache", MagicMock()), + patch(f"{_MOD}._refresh_cached_team", new=AsyncMock()), + ): + result = await delete_model_endpoint( + model_info=ModelInfoDelete(id=model_id), + user_api_key_dict=admin_user, + ) + + assert "deleted successfully" in result["message"] + mock_prisma.db.litellm_proxymodeltable.delete.assert_awaited_once() + # Team is gone -> no team.models cleanup to do. + mock_prisma.db.litellm_teamtable.update.assert_not_awaited() + + @pytest.mark.asyncio + async def test_non_admin_cannot_delete_model_when_team_deleted(self): + """A missing team must never let a non-admin delete the orphan (no fail-open).""" + from litellm.proxy.management_endpoints.model_management_endpoints import ( + ModelInfoDelete, + delete_model as delete_model_endpoint, + ) + from litellm.proxy.proxy_server import ProxyException + + team_id = "deleted-team-abc" + model_id = "orphaned-byok-2" + mock_prisma = self._orphaned_model_mocks(team_id, model_id) + + non_admin = UserAPIKeyAuth( + user_id="someone", user_role=LitellmUserRoles.INTERNAL_USER + ) + + _PS = "litellm.proxy.proxy_server" + _MOD = "litellm.proxy.management_endpoints.model_management_endpoints" + with ( + patch(f"{_PS}.prisma_client", mock_prisma), + patch(f"{_PS}.store_model_in_db", True), + patch(f"{_PS}.premium_user", True), + patch(f"{_PS}.llm_router", MagicMock()), + patch(f"{_PS}.proxy_logging_obj", MagicMock()), + patch(f"{_PS}.user_api_key_cache", MagicMock()), + patch(f"{_MOD}._refresh_cached_team", new=AsyncMock()), + ): + with pytest.raises(ProxyException) as exc_info: + await delete_model_endpoint( + model_info=ModelInfoDelete(id=model_id), + user_api_key_dict=non_admin, + ) + + assert str(exc_info.value.code) == "403" + mock_prisma.db.litellm_proxymodeltable.delete.assert_not_awaited() + + @pytest.mark.asyncio + async def test_live_team_delete_looks_up_team_once(self): + """The auth check must not add a redundant team query on the live-team path.""" + from litellm.proxy.management_endpoints.model_management_endpoints import ( + ModelInfoDelete, + delete_model as delete_model_endpoint, + ) + from litellm.proxy.proxy_server import ProxyException + + team_id = "live-team-1" + model_id = "live-byok-1" + db_row = LiteLLM_ProxyModelTable( + model_id=model_id, + model_name=f"model_name_{team_id}_abc-uuid", + litellm_params={"model": "openai/gpt-4.1-nano"}, + model_info={ + "id": model_id, + "team_id": team_id, + "team_public_model_name": "live-gpt", + }, + created_by="admin", + updated_by="admin", + ) + team_row = LiteLLM_TeamTable( + team_id=team_id, + team_alias="live-team", + members_with_roles=[Member(user_id="admin", role="admin")], + models=["live-gpt"], + ) + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.litellm_proxymodeltable = AsyncMock() + mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock( + return_value=db_row + ) + mock_prisma.db.litellm_proxymodeltable.delete = AsyncMock(return_value=db_row) + mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_teamtable = AsyncMock() + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + mock_prisma.db.litellm_modeltable = AsyncMock() + mock_prisma.db.litellm_modeltable.find_many = AsyncMock(return_value=[]) + + # A team member who is not the team admin: rejected before the delete runs, + # so the only team lookup is the single one inside the auth check. + non_admin = UserAPIKeyAuth( + user_id="someone", user_role=LitellmUserRoles.INTERNAL_USER + ) + + _PS = "litellm.proxy.proxy_server" + _MOD = "litellm.proxy.management_endpoints.model_management_endpoints" + with ( + patch(f"{_PS}.prisma_client", mock_prisma), + patch(f"{_PS}.store_model_in_db", True), + patch(f"{_PS}.premium_user", True), + patch(f"{_PS}.llm_router", MagicMock()), + patch(f"{_PS}.proxy_logging_obj", MagicMock()), + patch(f"{_PS}.user_api_key_cache", MagicMock()), + patch(f"{_MOD}._refresh_cached_team", new=AsyncMock()), + ): + with pytest.raises(ProxyException) as exc_info: + await delete_model_endpoint( + model_info=ModelInfoDelete(id=model_id), + user_api_key_dict=non_admin, + ) + + assert str(exc_info.value.code) == "403" + assert mock_prisma.db.litellm_teamtable.find_unique.await_count == 1 + mock_prisma.db.litellm_proxymodeltable.delete.assert_not_awaited() + + class TestGetTeamDeployments: """Tests for _get_team_deployments which filters by model_name prefix + Python-side team_id check.""" From a7ecf6b5b1452410cb83f58054d939561ef322a6 Mon Sep 17 00:00:00 2001 From: milan-berri Date: Tue, 9 Jun 2026 01:09:03 +0300 Subject: [PATCH 021/185] feat(jwt-auth): opt-in fallback to DB team on unresolved JWT claim (#28913) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(jwt-auth): defer to single-team DB fallback on claim mismatch Extends the single-team DB fallback introduced in #26418 to two more cases where it previously could not run: * `find_and_validate_specific_team_id`: when `team_id_jwt_field` is configured and a claim value is present in the token but the team does not exist in the LiteLLM DB (HTTPException 404 from `get_team_object`), return `(None, None)` instead of raising — the auth_builder fallback then attributes the request to the user's single DB team. Only HTTPException is caught; other errors (e.g. "No DB Connected") still propagate. * `find_team_with_model_access`: when none of the `team_ids_jwt_field` groups resolve to a real LiteLLM team, return `(None, None)` instead of raising 403 so the same fallback path runs. If at least one group DID resolve to a team but none granted the requested model, the original 403 is preserved (legitimate access denial — not a claim mismatch). Tracked via the new `any_claim_team_resolved` flag. The strict `is_required_team_id` raise and `enforce_team_based_model_access` raise remain unchanged. Unit tests cover both new soft-fail paths and guard each preserved path (strict required, enforce_team_based, the preserved 403, and the non-HTTPException propagation). Co-authored-by: Cursor * fix(jwt-auth): narrow HTTPException catch to 404 (greptile review) Address Greptile review comments on #28913: * `find_and_validate_specific_team_id`: re-raise HTTPException when `status_code != 404`, pinning the catch to the "team doesn't exist in db" path documented for `get_team_object`. A future change that introduces a different status code (e.g. 403 for a blocked team) will now propagate instead of silently falling through to the single-team DB fallback. * Add `test_find_and_validate_specific_team_id_non_404_http_exception_propagates` parametrised over 400 / 403 / 500 to lock in the contract. Co-authored-by: Cursor * fix(jwt-auth): gate claim-mismatch fallback behind opt-in flag The unresolved-team-claim fallback added in the previous commit weakened the strict claim-based authorization contract by default — an authenticated user whose JWT carries a stale or invalid team claim could still consume their single DB team's models/quota via the fallback. Gate both soft-fail paths in `find_and_validate_specific_team_id` and `find_team_with_model_access` behind a new opt-in flag `team_claim_fallback` on `LiteLLM_JWTAuth` (default False). Default-off preserves the pre-existing strict behavior. Operators who intentionally treat IdP team claims as advisory (e.g. machine tokens whose group claims live in a separate namespace from LiteLLM team_ids) opt in via config. Adds two regression tests guarding the default-off behavior. Co-authored-by: Cursor --------- Co-authored-by: Cursor --- litellm/proxy/_types.py | 10 + litellm/proxy/auth/handle_jwt.py | 43 ++- .../proxy/auth/test_handle_jwt.py | 304 ++++++++++++++++++ 3 files changed, 347 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index fba968dd95d..f1443edf455 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -4195,6 +4195,16 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase): default=None, description="Optional claim-based routing overrides for JWT-shaped tokens. Matching rules route requests to oauth2 before default JWT flow.", ) + team_claim_fallback: bool = Field( + default=False, + description=( + "If True, when a configured team_id_jwt_field / team_ids_jwt_field " + "claim is present but does not resolve to any known team, defer to " + "the single-team DB fallback (caller's only team membership) " + "instead of raising. Default False preserves strict claim-based " + "authorization." + ), + ) issuers: Optional[List[JWTIssuerConfig]] = Field( default=None, description="Optional issuer-bound JWT validation rules. When a token's `iss` matches a configured issuer, validation uses that issuer's JWKS, audience, and claim mappings. Tokens with an unlisted `iss` fall back to the global JWT_AUDIENCE/JWT_ISSUER validation path — this is additive routing, not an allow-list.", diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 536d2867855..fd6ff2ada7f 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -1299,15 +1299,29 @@ class JWTAuthManager: # First try to get team by team_id if individual_team_id: - team_object = await get_team_object( - team_id=individual_team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - team_id_upsert=jwt_handler.litellm_jwtauth.team_id_upsert, - ) - return individual_team_id, team_object + try: + team_object = await get_team_object( + team_id=individual_team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + team_id_upsert=jwt_handler.litellm_jwtauth.team_id_upsert, + ) + return individual_team_id, team_object + except HTTPException as e: + if ( + e.status_code != 404 + or not jwt_handler.litellm_jwtauth.team_claim_fallback + ): + raise + # Claim doesn't map to a known team — defer to fallback. + verbose_proxy_logger.debug( + "JWT team_id claim '%s' did not resolve to a team: %s", + individual_team_id, + e.detail, + ) + return None, None # If no team_id found, try to resolve via team_alias_jwt_field team_alias = jwt_handler.get_team_alias( @@ -1431,6 +1445,7 @@ class JWTAuthManager: ) return None, None + any_claim_team_resolved = False for team_id in team_ids: try: team_object = await get_team_object( @@ -1441,6 +1456,9 @@ class JWTAuthManager: proxy_logging_obj=proxy_logging_obj, ) + if team_object is not None: + any_claim_team_resolved = True + if team_object and team_object.models is not None: team_models = team_object.models if isinstance(team_models, list) and ( @@ -1478,12 +1496,17 @@ class JWTAuthManager: if denied_auth_enforced_pass_through_route: JWTAuthManager._raise_team_passthrough_route_denial(route=route) - if requested_model: + if requested_model and ( + any_claim_team_resolved + or not jwt_handler.litellm_jwtauth.team_claim_fallback + ): + # Claim resolved but no model access, or fallback disabled — deny. raise HTTPException( status_code=403, detail=f"No team has access to the requested model: {requested_model}. Checked teams={team_ids}. Check `/models` to see all available models.", ) + # No claim team resolved and fallback enabled — defer to fallback. return None, None @staticmethod diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py index 7dadca24504..63510086f95 100644 --- a/tests/test_litellm/proxy/auth/test_handle_jwt.py +++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py @@ -3182,6 +3182,310 @@ def test_build_decode_kwargs_no_warning_when_scoped( assert matching == [] +# --------------------------------------------------------------------------- +# Defer to single-team DB fallback (PR #26418) when JWT claims are present +# but do not resolve to a LiteLLM team. +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_find_and_validate_specific_team_id_unresolved_claim_returns_none(): + """With `team_claim_fallback=True`: team_id claim is present in the JWT + but the team is missing in the DB — return (None, None) so the + auth_builder single-team fallback can run, instead of raising and + failing auth.""" + from fastapi import HTTPException + + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( + team_id_jwt_field="team_id", + team_claim_fallback=True, + ) + token = {"sub": "user-1", "team_id": "claim-team-not-in-db"} + + with patch( + "litellm.proxy.auth.handle_jwt.get_team_object", + new_callable=AsyncMock, + ) as mock_get_team: + mock_get_team.side_effect = HTTPException(status_code=404, detail="missing") + + team_id, team_object = await JWTAuthManager.find_and_validate_specific_team_id( + jwt_handler=jwt_handler, + jwt_valid_token=token, + prisma_client=None, + user_api_key_cache=None, + parent_otel_span=None, + proxy_logging_obj=None, + ) + + assert team_id is None + assert team_object is None + + +@pytest.mark.asyncio +async def test_find_team_with_model_access_unresolved_group_claim_returns_none( + monkeypatch, +): + """With `team_claim_fallback=True`: group claim resolves to team_ids that + don't exist in the DB — return (None, None) instead of raising 403, so + the single-team fallback can run.""" + import sys + import types + + from fastapi import HTTPException + + from litellm.router import Router + + router = Router( + model_list=[ + {"model_name": "gpt-4o-mini", "litellm_params": {"model": "gpt-4o-mini"}} + ] + ) + proxy_server_module = types.ModuleType("proxy_server") + proxy_server_module.llm_router = router + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_module) + + async def raise_404(*_args, **_kwargs): + raise HTTPException(status_code=404, detail="missing") + + monkeypatch.setattr("litellm.proxy.auth.handle_jwt.get_team_object", raise_404) + + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_claim_fallback=True) + + team_id, team_object = await JWTAuthManager.find_team_with_model_access( + team_ids={"idp-group-a", "idp-group-b"}, + requested_model="gpt-4o-mini", + route="/chat/completions", + jwt_handler=jwt_handler, + prisma_client=None, + user_api_key_cache=None, + parent_otel_span=None, + proxy_logging_obj=None, + ) + + assert team_id is None + assert team_object is None + + +@pytest.mark.asyncio +async def test_find_and_validate_specific_team_id_non_http_exception_still_propagates(): + """Regression guard: only the 404 HTTPException raised by + `get_team_object` ("team doesn't exist in db") is softened. Other + errors — e.g. "No DB Connected" — must still propagate so operator-side + problems are loud.""" + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_id_jwt_field="team_id") + token = {"sub": "user-1", "team_id": "some-claim-team"} + + with patch( + "litellm.proxy.auth.handle_jwt.get_team_object", + new_callable=AsyncMock, + ) as mock_get_team: + mock_get_team.side_effect = RuntimeError("simulated infrastructure error") + + with pytest.raises(RuntimeError, match="simulated infrastructure error"): + await JWTAuthManager.find_and_validate_specific_team_id( + jwt_handler=jwt_handler, + jwt_valid_token=token, + prisma_client=None, + user_api_key_cache=None, + parent_otel_span=None, + proxy_logging_obj=None, + ) + + +@pytest.mark.asyncio +async def test_find_and_validate_specific_team_id_non_404_http_exception_propagates(): + """Regression guard: only 404 HTTPException is softened. If + `get_team_object` is ever updated to raise a different HTTP status code + (e.g. 403 for a blocked team), that error must still propagate rather + than silently fall through to the single-team DB fallback.""" + from fastapi import HTTPException + + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_id_jwt_field="team_id") + token = {"sub": "user-1", "team_id": "some-claim-team"} + + for status_code in (400, 403, 500): + with patch( + "litellm.proxy.auth.handle_jwt.get_team_object", + new_callable=AsyncMock, + ) as mock_get_team: + mock_get_team.side_effect = HTTPException( + status_code=status_code, detail="non-404 failure" + ) + + with pytest.raises(HTTPException) as exc_info: + await JWTAuthManager.find_and_validate_specific_team_id( + jwt_handler=jwt_handler, + jwt_valid_token=token, + prisma_client=None, + user_api_key_cache=None, + parent_otel_span=None, + proxy_logging_obj=None, + ) + assert exc_info.value.status_code == status_code + + +@pytest.mark.asyncio +async def test_find_team_with_model_access_enforce_team_based_access_still_raises(): + """Regression guard: when no group claims are present and + `enforce_team_based_model_access` is on, the original 403 still fires — + the new soft-fail only applies to the unresolved-claim path inside the + loop, not to the no-team-claims-at-all path at the top.""" + from fastapi import HTTPException + + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(enforce_team_based_model_access=True) + + with pytest.raises(HTTPException) as exc_info: + await JWTAuthManager.find_team_with_model_access( + team_ids=set(), + requested_model="gpt-4o-mini", + route="/chat/completions", + jwt_handler=jwt_handler, + prisma_client=None, + user_api_key_cache=None, + parent_otel_span=None, + proxy_logging_obj=None, + ) + + assert exc_info.value.status_code == 403 + assert "enforce_team_based_model_access" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_find_team_with_model_access_resolved_team_without_model_still_raises_403( + monkeypatch, +): + """Regression guard: when the JWT group claim DOES resolve to a real + LiteLLM team but that team does not grant the requested model, keep the + original 403. Only the unresolved-claim case is softened.""" + import sys + import types + + from fastapi import HTTPException + + from litellm.router import Router + + router = Router( + model_list=[ + {"model_name": "gpt-4o-mini", "litellm_params": {"model": "gpt-4o-mini"}}, + { + "model_name": "gpt-3.5-turbo", + "litellm_params": {"model": "gpt-3.5-turbo"}, + }, + ] + ) + proxy_server_module = types.ModuleType("proxy_server") + proxy_server_module.llm_router = router + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_module) + + team = LiteLLM_TeamTable(team_id="real-team", models=["gpt-3.5-turbo"]) + + async def mock_get_team_object(*_args, **_kwargs): + return team + + monkeypatch.setattr( + "litellm.proxy.auth.handle_jwt.get_team_object", mock_get_team_object + ) + + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth() + + with pytest.raises(HTTPException) as exc_info: + await JWTAuthManager.find_team_with_model_access( + team_ids={"real-team"}, + requested_model="gpt-4o-mini", + route="/chat/completions", + jwt_handler=jwt_handler, + prisma_client=None, + user_api_key_cache=None, + parent_otel_span=None, + proxy_logging_obj=None, + ) + + assert exc_info.value.status_code == 403 + assert "No team has access to the requested model" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_find_and_validate_specific_team_id_unresolved_claim_default_raises(): + """Default `team_claim_fallback=False`: unresolved team_id claim must + still raise — preserves the strict claim-based authorization boundary + when the operator has not opted in to the fallback.""" + from fastapi import HTTPException + + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_id_jwt_field="team_id") + token = {"sub": "user-1", "team_id": "claim-team-not-in-db"} + + with patch( + "litellm.proxy.auth.handle_jwt.get_team_object", + new_callable=AsyncMock, + ) as mock_get_team: + mock_get_team.side_effect = HTTPException(status_code=404, detail="missing") + + with pytest.raises(HTTPException) as exc_info: + await JWTAuthManager.find_and_validate_specific_team_id( + jwt_handler=jwt_handler, + jwt_valid_token=token, + prisma_client=None, + user_api_key_cache=None, + parent_otel_span=None, + proxy_logging_obj=None, + ) + + assert exc_info.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_find_team_with_model_access_unresolved_group_claim_default_raises( + monkeypatch, +): + """Default `team_claim_fallback=False`: group claims that don't resolve + to any LiteLLM team must still raise 403 — preserves the strict + claim-based authorization boundary.""" + import sys + import types + + from fastapi import HTTPException + + from litellm.router import Router + + router = Router( + model_list=[ + {"model_name": "gpt-4o-mini", "litellm_params": {"model": "gpt-4o-mini"}} + ] + ) + proxy_server_module = types.ModuleType("proxy_server") + proxy_server_module.llm_router = router + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_module) + + async def raise_404(*_args, **_kwargs): + raise HTTPException(status_code=404, detail="missing") + + monkeypatch.setattr("litellm.proxy.auth.handle_jwt.get_team_object", raise_404) + + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth() + + with pytest.raises(HTTPException) as exc_info: + await JWTAuthManager.find_team_with_model_access( + team_ids={"idp-group-a", "idp-group-b"}, + requested_model="gpt-4o-mini", + route="/chat/completions", + jwt_handler=jwt_handler, + prisma_client=None, + user_api_key_cache=None, + parent_otel_span=None, + proxy_logging_obj=None, + ) + + assert exc_info.value.status_code == 403 + + # GH #26789: JWT claim user_id must rebind to legacy DB row after fuzzy match. From 9ccda119194415313b3a40acd0dfedf0136b2826 Mon Sep 17 00:00:00 2001 From: milan-berri Date: Tue, 9 Jun 2026 01:14:24 +0300 Subject: [PATCH 022/185] fix(team_endpoints): don't block /team/update on unchanged team budget (#29525) On /team/update for a standalone (no-org) team, _check_user_team_limits() compared the request max_budget against the caller's personal max_budget whenever max_budget was present in the payload. A team admin whose personal budget is lower than the team's budget could not edit any field (tpm_limit, team name, etc.) because the UI re-sends the unchanged max_budget on every update, tripping the personal-budget check. Pass the team's current max_budget into _check_user_team_limits() and skip the personal-budget comparison when the incoming value is unchanged or lower than the team's current budget. Only genuine increases above the team's current budget are still validated against the caller's personal limit, so no over-relaxation. Proxy admins and the org-scoped path are unaffected. Adds two regression tests for the standalone update path (unchanged budget + tpm_limit change, and lowering the budget), both for a caller whose personal budget is below the team budget. Co-authored-by: Cursor --- .../management_endpoints/team_endpoints.py | 48 +++-- .../test_team_endpoints.py | 193 ++++++++++++++++++ 2 files changed, 226 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 7a784ee4622..f2eafcbf839 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -779,6 +779,7 @@ async def _check_user_team_limits( user_api_key_dict: UserAPIKeyAuth, prisma_client: PrismaClient, user_api_key_cache: Any, + existing_team_max_budget: Optional[float] = None, ) -> None: """ Check user team limits for standalone teams (not org-scoped). @@ -789,28 +790,44 @@ async def _check_user_team_limits( Should only be called for standalone teams (when organization_id is None). For org-scoped teams, use _check_org_team_limits() instead. + + `existing_team_max_budget` is the team's current `max_budget` on the + /team/update path. When the incoming `max_budget` is unchanged or lower + than the team's current budget, the personal-budget comparison is skipped + so a team admin can edit other fields (e.g. tpm_limit, team name) without + being blocked by a budget the team already has. The UI sends the full team + object on every update, so the unchanged `max_budget` would otherwise fail. """ # Validate team budget against user's max_budget if data.max_budget is not None and user_api_key_dict.user_id is not None: - user_obj = await get_user_object( - user_id=user_api_key_dict.user_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - user_id_upsert=False, + # On /team/update, allow unchanged or lower budgets without checking + # the caller's personal max_budget. Only increases above the team's + # current budget are validated against the user's personal limit. + budget_unchanged_or_lower = ( + existing_team_max_budget is not None + and data.max_budget <= existing_team_max_budget ) - if ( - user_obj is not None - and user_obj.max_budget is not None - and data.max_budget > user_obj.max_budget - ): - raise HTTPException( - status_code=400, - detail={ - "error": f"max budget higher than user max. User max budget={user_obj.max_budget}. User role={user_api_key_dict.user_role}" - }, + if not budget_unchanged_or_lower: + user_obj = await get_user_object( + user_id=user_api_key_dict.user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, ) + if ( + user_obj is not None + and user_obj.max_budget is not None + and data.max_budget > user_obj.max_budget + ): + raise HTTPException( + status_code=400, + detail={ + "error": f"max budget higher than user max. User max budget={user_obj.max_budget}. User role={user_api_key_dict.user_role}" + }, + ) + # Validate team models against user's allowed models if data.models is not None and len(user_api_key_dict.models) > 0: for m in data.models: @@ -1824,6 +1841,7 @@ async def update_team( # noqa: PLR0915 user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, + existing_team_max_budget=existing_team_row.max_budget, ) updated_kv = data.json(exclude_unset=True) diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index d580f1f7703..400ed287802 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -4531,6 +4531,199 @@ async def test_update_team_standalone_budget_exceeds_user_limit(): assert "budget" in str(exc_info.value.message).lower() +@pytest.mark.asyncio +async def test_update_team_standalone_unchanged_budget_allowed(): + """ + Test that /team/update for a standalone team does NOT compare against the + caller's personal max_budget when the budget is unchanged. + + This is the LiteLLM UI scenario: the UI sends the full team object on every + update (including the unchanged max_budget). A team admin only changing + tpm_limit should not be blocked by a budget the team already has. + + Scenario: + - User (team admin) has personal max_budget=$100 + - Standalone team exists with current budget=$500 + - User updates tpm_limit and re-sends the unchanged max_budget=$500 + - Expected: Should succeed (budget unchanged, not an increase) + """ + from fastapi import Request + + from litellm.proxy._types import ( + LiteLLM_UserTable, + UpdateTeamRequest, + UserAPIKeyAuth, + ) + from litellm.proxy.management_endpoints.team_endpoints import update_team + + team_admin_user = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="standalone-unchanged-budget-admin", + models=[], + ) + + # UI re-sends the unchanged max_budget alongside the tpm_limit change. + update_request = UpdateTeamRequest( + team_id="standalone-unchanged-budget-123", + max_budget=500.0, # Unchanged from the team's current budget + tpm_limit=50000, + ) + + dummy_request = MagicMock(spec=Request) + + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch( + "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() + ) as mock_audit, + ): + # Mock existing standalone team (no organization_id) with budget=$500 + mock_existing_team = MagicMock() + mock_existing_team.team_id = "standalone-unchanged-budget-123" + mock_existing_team.organization_id = None + mock_existing_team.max_budget = 500.0 + mock_existing_team.model_id = None + mock_existing_team.model_dump.return_value = { + "team_id": "standalone-unchanged-budget-123", + "organization_id": None, + "max_budget": 500.0, + "members_with_roles": [ + {"user_id": "standalone-unchanged-budget-admin", "role": "admin"} + ], + } + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock( + return_value=mock_existing_team + ) + mock_prisma.jsonify_team_object = lambda db_data: db_data + + # User has a restrictive personal budget that is lower than the team's. + mock_user_obj = LiteLLM_UserTable( + user_id="standalone-unchanged-budget-admin", + max_budget=100.0, + ) + mock_cache.async_get_cache = AsyncMock(return_value=mock_user_obj) + mock_cache.async_set_cache = AsyncMock() + + mock_updated_team = MagicMock() + mock_updated_team.team_id = "standalone-unchanged-budget-123" + mock_updated_team.organization_id = None + mock_updated_team.max_budget = 500.0 + mock_updated_team.litellm_model_table = None + mock_updated_team.model_dump.return_value = { + "team_id": "standalone-unchanged-budget-123", + "organization_id": None, + "max_budget": 500.0, + "tpm_limit": 50000, + } + mock_prisma.db.litellm_teamtable.update = AsyncMock( + return_value=mock_updated_team + ) + + # Should NOT raise - unchanged budget skips the personal-budget check. + result = await update_team( + data=update_request, + http_request=dummy_request, + user_api_key_dict=team_admin_user, + ) + + assert result is not None + assert result["data"].max_budget == 500.0 + + +@pytest.mark.asyncio +async def test_update_team_standalone_lower_budget_allowed(): + """ + Test that /team/update for a standalone team allows lowering the budget + below the team's current value even when the new value still exceeds the + caller's personal max_budget. + + Scenario: + - User (team admin) has personal max_budget=$100 + - Standalone team exists with current budget=$500 + - User lowers team budget to $300 (a decrease, still above user's $100) + - Expected: Should succeed (decrease is not an increase above team budget) + """ + from fastapi import Request + + from litellm.proxy._types import ( + LiteLLM_UserTable, + UpdateTeamRequest, + UserAPIKeyAuth, + ) + from litellm.proxy.management_endpoints.team_endpoints import update_team + + team_admin_user = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="standalone-lower-budget-admin", + models=[], + ) + + update_request = UpdateTeamRequest( + team_id="standalone-lower-budget-123", + max_budget=300.0, # Lower than current $500, still above user's $100 + ) + + dummy_request = MagicMock(spec=Request) + + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch( + "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() + ) as mock_audit, + ): + mock_existing_team = MagicMock() + mock_existing_team.team_id = "standalone-lower-budget-123" + mock_existing_team.organization_id = None + mock_existing_team.max_budget = 500.0 + mock_existing_team.model_id = None + mock_existing_team.model_dump.return_value = { + "team_id": "standalone-lower-budget-123", + "organization_id": None, + "max_budget": 500.0, + "members_with_roles": [ + {"user_id": "standalone-lower-budget-admin", "role": "admin"} + ], + } + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock( + return_value=mock_existing_team + ) + mock_prisma.jsonify_team_object = lambda db_data: db_data + + mock_user_obj = LiteLLM_UserTable( + user_id="standalone-lower-budget-admin", + max_budget=100.0, + ) + mock_cache.async_get_cache = AsyncMock(return_value=mock_user_obj) + mock_cache.async_set_cache = AsyncMock() + + mock_updated_team = MagicMock() + mock_updated_team.team_id = "standalone-lower-budget-123" + mock_updated_team.organization_id = None + mock_updated_team.max_budget = 300.0 + mock_updated_team.litellm_model_table = None + mock_updated_team.model_dump.return_value = { + "team_id": "standalone-lower-budget-123", + "organization_id": None, + "max_budget": 300.0, + } + mock_prisma.db.litellm_teamtable.update = AsyncMock( + return_value=mock_updated_team + ) + + result = await update_team( + data=update_request, + http_request=dummy_request, + user_api_key_dict=team_admin_user, + ) + + assert result is not None + assert result["data"].max_budget == 300.0 + + @pytest.mark.asyncio async def test_update_team_org_scoped_budget_exceeds_org_limit(): """ From 1c881eee5da8467ca2d07ea5caedf8bdf5b69069 Mon Sep 17 00:00:00 2001 From: milan-berri Date: Tue, 9 Jun 2026 01:54:19 +0300 Subject: [PATCH 023/185] fix(fireworks): enable tool calling for glm-5p1 in model cost map (#29697) glm-5p1 supports native tools on Fireworks; explicit false flags caused drop_params to strip tools and tool_choice before the provider request. Co-authored-by: Cursor --- litellm/model_prices_and_context_window_backup.json | 12 ++++++------ model_prices_and_context_window.json | 12 ++++++------ .../chat/test_fireworks_ai_chat_transformation.py | 10 ++++++---- 3 files changed, 18 insertions(+), 16 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 397f96fdb1e..e765512175b 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -14286,10 +14286,10 @@ "mode": "chat", "output_cost_per_token": 4.4e-06, "source": "https://fireworks.ai/models/fireworks/glm-5p1", - "supports_function_calling": false, + "supports_function_calling": true, "supports_reasoning": true, - "supports_response_schema": false, - "supports_tool_choice": false + "supports_response_schema": true, + "supports_tool_choice": true }, "fireworks_ai/accounts/fireworks/models/gpt-oss-120b": { "input_cost_per_token": 1.5e-07, @@ -14567,10 +14567,10 @@ "mode": "chat", "output_cost_per_token": 4.4e-06, "source": "https://fireworks.ai/models/fireworks/glm-5p1", - "supports_function_calling": false, + "supports_function_calling": true, "supports_reasoning": true, - "supports_response_schema": false, - "supports_tool_choice": false + "supports_response_schema": true, + "supports_tool_choice": true }, "fireworks_ai/kimi-k2p5": { "cache_read_input_token_cost": 1e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index b2836a096b7..c1c05b982f6 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -14286,10 +14286,10 @@ "mode": "chat", "output_cost_per_token": 4.4e-06, "source": "https://fireworks.ai/models/fireworks/glm-5p1", - "supports_function_calling": false, + "supports_function_calling": true, "supports_reasoning": true, - "supports_response_schema": false, - "supports_tool_choice": false + "supports_response_schema": true, + "supports_tool_choice": true }, "fireworks_ai/accounts/fireworks/models/gpt-oss-120b": { "input_cost_per_token": 1.5e-07, @@ -14567,10 +14567,10 @@ "mode": "chat", "output_cost_per_token": 4.4e-06, "source": "https://fireworks.ai/models/fireworks/glm-5p1", - "supports_function_calling": false, + "supports_function_calling": true, "supports_reasoning": true, - "supports_response_schema": false, - "supports_tool_choice": false + "supports_response_schema": true, + "supports_tool_choice": true }, "fireworks_ai/kimi-k2p5": { "cache_read_input_token_cost": 1e-07, diff --git a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index ca340b5f275..0221db1b23d 100644 --- a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py @@ -127,12 +127,14 @@ def test_get_supported_openai_params_parallel_tool_calls(): config = FireworksAIConfig() supported_params = config.get_supported_openai_params( - "fireworks_ai/accounts/fireworks/models/glm-4p6" + "fireworks_ai/accounts/fireworks/models/glm-5p1" ) assert "parallel_tool_calls" in supported_params + assert "tools" in supported_params + assert "tool_choice" in supported_params unsupported_params = config.get_supported_openai_params( - "fireworks_ai/accounts/fireworks/models/glm-5p1" + "fireworks_ai/accounts/fireworks/models/llama-v3p1-8b-instruct" ) assert "parallel_tool_calls" not in unsupported_params @@ -163,9 +165,9 @@ def test_get_model_info_respects_explicit_fireworks_capabilities(): """Test that get_model_info preserves explicit capability flags from the model map.""" model_info = get_model_info("fireworks_ai/accounts/fireworks/models/glm-5p1") - assert model_info["supports_function_calling"] is False + assert model_info["supports_function_calling"] is True assert model_info["supports_reasoning"] is True - assert model_info["supports_tool_choice"] is False + assert model_info["supports_tool_choice"] is True def test_get_provider_info_omits_false_supports_reasoning(monkeypatch): From dfd6cbc514f7af84179ab126d15cf68833ca410e Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 9 Jun 2026 04:44:30 +0530 Subject: [PATCH 024/185] fix(vertex): propagate Vertex AI metadata in streaming success callbacks (#29899) * fix(vertex): propagate Vertex AI metadata in streaming success callbacks Streaming calls assembled via stream_chunk_builder were missing vertex_ai_grounding_metadata and vertex_ai_url_context_metadata in standard_logging_object.response. Merge metadata from chunks into the assembled response and mirror non-streaming hidden_params on Gemini chunks. Co-authored-by: Cursor * refactor(vertex): move streaming metadata merge into provider config hook Address review feedback by delegating assembled-stream metadata propagation to VertexGeminiConfig via BaseConfig.apply_assembled_streaming_response_metadata, and only write chunk hidden_params when metadata is non-empty. Co-authored-by: Cursor * fix(redaction): scrub Vertex provider metadata when message logging is off Clear vertex_ai_grounding_metadata and related fields from standard logging responses and assembled streaming ModelResponse objects so turn_off_message_logging cannot leak prompt-derived web search queries. Co-authored-by: Cursor * Use assembled model for streaming metadata hook * Fix Vertex metadata redaction bypass in logging callbacks. Scrub Vertex provider fields from litellm_params.metadata.hidden_params during perform_redaction so streaming success_handler merges do not leak prompt-derived metadata when message logging is disabled. Co-authored-by: Cursor * Fix Vertex streaming metadata from hidden params * fix(vertex): mirror vertex_ai_safety_results on assembled streaming responses The non-streaming transform_response stores safety data under vertex_ai_safety_results, but the streaming path only wrote vertex_ai_safety_ratings. Assembled streaming responses therefore never carried vertex_ai_safety_results, so any consumer reading that field saw a silent difference between streaming and non-streaming calls. Set vertex_ai_safety_results alongside vertex_ai_safety_ratings in the shared stream metadata setter and add it to the assembled metadata field list so it propagates through stream_chunk_builder. * fix(streaming): log provider streaming metadata hook failures instead of swallowing them * refactor(vertex): share single Vertex metadata field tuple across redaction and streaming * refactor(vertex): move Vertex metadata redaction helpers into llms/vertex_ai --------- Co-authored-by: Cursor Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/litellm_core_utils/redact_messages.py | 10 ++ .../streaming_chunk_builder_utils.py | 49 ++++++ litellm/llms/base_llm/chat/transformation.py | 8 + litellm/llms/vertex_ai/common_utils.py | 47 +++++- .../vertex_and_google_ai_studio_gemini.py | 77 ++++++++- litellm/main.py | 6 + litellm/types/llms/vertex_ai.py | 9 ++ .../test_litellm_logging.py | 35 ++++ .../test_redact_messages.py | 93 +++++++++++ .../test_streaming_chunk_builder_utils.py | 150 ++++++++++++++++++ ...test_vertex_and_google_ai_studio_gemini.py | 20 +++ 11 files changed, 499 insertions(+), 5 deletions(-) diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index dbc9cabdc7a..763596336a0 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -17,6 +17,10 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import ( get_metadata_variable_name_from_kwargs, ) +from litellm.llms.vertex_ai.common_utils import ( + redact_vertex_ai_metadata_from_litellm_params, + redact_vertex_ai_metadata_from_logged_object, +) from litellm.secret_managers.main import str_to_bool from litellm.types.utils import StandardCallbackDynamicParams @@ -119,10 +123,12 @@ def _redact_standard_logging_object(model_call_details: dict): # ResponsesAPIResponse format - redact content in output items if isinstance(response.get("output"), list): _redact_responses_api_output_dict(response["output"], redacted_str) + redact_vertex_ai_metadata_from_logged_object(response) elif isinstance(response, dict) and "choices" in response: # ModelResponse dict format - redact content in choices if isinstance(response.get("choices"), list): _redact_model_response_dict_choices(response["choices"], redacted_str) + redact_vertex_ai_metadata_from_logged_object(response) elif isinstance(response, str): standard_logging_object["response"] = redacted_str else: @@ -164,6 +170,7 @@ def perform_redaction(model_call_details: dict, result): model_call_details["prompt"] = "" model_call_details["input"] = "" _redact_standard_logging_object(model_call_details) + redact_vertex_ai_metadata_from_litellm_params(model_call_details) # Redact streaming response if ( @@ -174,6 +181,7 @@ def perform_redaction(model_call_details: dict, result): if hasattr(_streaming_response, "choices"): for choice in _streaming_response.choices: _redact_choice_content(choice) + redact_vertex_ai_metadata_from_logged_object(_streaming_response) elif hasattr(_streaming_response, "output"): _redact_responses_api_output(_streaming_response.output) # Redact reasoning field in ResponsesAPIResponse @@ -200,12 +208,14 @@ def perform_redaction(model_call_details: dict, result): if hasattr(_result, "choices") and _result.choices is not None: for choice in _result.choices: _redact_choice_content(choice) + redact_vertex_ai_metadata_from_logged_object(_result) elif isinstance(_result, dict) and "choices" in _result: # Handle dict representation of ModelResponse (e.g., from model_dump()) if _result.get("choices") is not None: _redact_model_response_dict_choices( _result["choices"], "redacted-by-litellm" ) + redact_vertex_ai_metadata_from_logged_object(_result) elif isinstance(_result, dict) and "output" in _result: if isinstance(_result.get("output"), list): _redact_responses_api_output_dict( diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index fe7c62c3842..6257cce9aec 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -20,6 +20,7 @@ from litellm.types.utils import ( ServerToolUse, Usage, ) +from litellm._logging import verbose_logger from litellm.utils import print_verbose, token_counter if TYPE_CHECKING: @@ -79,6 +80,54 @@ class ChunkProcessor: model_response._hidden_params = chunk.get("_hidden_params", {}) return model_response + @staticmethod + def apply_provider_assembled_streaming_metadata( + response: ModelResponse, + chunks: List[Any], + logging_obj: Optional[Any] = None, + ) -> None: + if not chunks: + return + + model = getattr(response, "model", None) + if not model: + return + + custom_llm_provider = None + if logging_obj is not None: + custom_llm_provider = logging_obj.model_call_details.get( + "custom_llm_provider" + ) + + try: + from litellm.litellm_core_utils.get_llm_provider_logic import ( + get_llm_provider, + ) + from litellm.types.utils import LlmProviders + from litellm.utils import ProviderConfigManager + + if custom_llm_provider: + provider = LlmProviders(custom_llm_provider) + else: + _, provider_str, _, _ = get_llm_provider(model) + provider = LlmProviders(provider_str) + + provider_config = ProviderConfigManager.get_provider_chat_config( + model=model, + provider=provider, + ) + if provider_config is not None: + provider_config.apply_assembled_streaming_response_metadata( + response=response, + chunks=chunks, + ) + except Exception as e: + verbose_logger.debug( + "apply_provider_assembled_streaming_metadata failed for model=%s: %s", + model, + e, + ) + @staticmethod def _get_chunk_id(chunks: List[Dict[str, Any]]) -> str: """ diff --git a/litellm/llms/base_llm/chat/transformation.py b/litellm/llms/base_llm/chat/transformation.py index 5f35a58ce1f..8f9d5cad7c4 100644 --- a/litellm/llms/base_llm/chat/transformation.py +++ b/litellm/llms/base_llm/chat/transformation.py @@ -442,6 +442,14 @@ class BaseConfig(ABC): """Hook for providers to post-process streaming responses. Default: pass-through.""" return stream + def apply_assembled_streaming_response_metadata( + self, + response: "ModelResponse", + chunks: List[Any], + ) -> None: + """Hook for providers to merge chunk metadata into assembled streaming responses.""" + return None + def calculate_additional_costs( self, model: str, prompt_tokens: int, completion_tokens: int ) -> Optional[dict]: diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index e6e39651109..85c23d8603c 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -12,7 +12,11 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import unpack_defs from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.llms.openai import AllMessageValues -from litellm.types.llms.vertex_ai import PartType, Schema +from litellm.types.llms.vertex_ai import ( + VERTEX_AI_PROVIDER_METADATA_FIELDS, + PartType, + Schema, +) from litellm.types.utils import TokenCountResponse from litellm.utils import supports_response_schema, supports_system_messages @@ -27,6 +31,47 @@ class VertexAIError(BaseLLMException): super().__init__(message=message, status_code=status_code, headers=headers) +def redact_vertex_ai_metadata_from_logged_object(obj: Any) -> None: + if isinstance(obj, dict): + for field in VERTEX_AI_PROVIDER_METADATA_FIELDS: + if field in obj: + obj[field] = [] + hidden_params = obj.get("_hidden_params") + if isinstance(hidden_params, dict): + for field in VERTEX_AI_PROVIDER_METADATA_FIELDS: + hidden_params.pop(field, None) + return + + for field in VERTEX_AI_PROVIDER_METADATA_FIELDS: + if hasattr(obj, field): + setattr(obj, field, []) + hidden_params = getattr(obj, "_hidden_params", None) + if isinstance(hidden_params, dict): + for field in VERTEX_AI_PROVIDER_METADATA_FIELDS: + hidden_params.pop(field, None) + + +def redact_vertex_ai_metadata_from_litellm_params(model_call_details: dict) -> None: + """ + success_handler() merges response._hidden_params into + litellm_params.metadata['hidden_params'] before redaction runs, so the Vertex + metadata must be scrubbed from that copy too. + """ + litellm_params = model_call_details.get("litellm_params") + if not isinstance(litellm_params, dict): + return + + for metadata_key in ("metadata", "litellm_metadata"): + metadata = litellm_params.get(metadata_key) + if not isinstance(metadata, dict): + continue + hidden_params = metadata.get("hidden_params") + if not isinstance(hidden_params, dict): + continue + for field in VERTEX_AI_PROVIDER_METADATA_FIELDS: + hidden_params.pop(field, None) + + def vertex_request_labels_from_litellm_params( litellm_params: Optional[dict], ) -> Optional[Dict[str, str]]: diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 7d355a2e908..430a789d2a0 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -63,6 +63,7 @@ from litellm.types.llms.openai import ( OpenAIChatCompletionFinishReason, ) from litellm.types.llms.vertex_ai import ( + VERTEX_AI_PROVIDER_METADATA_FIELDS, VERTEX_CREDENTIALS_TYPES, Candidates, ContentType, @@ -2258,6 +2259,71 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): citation_metadata, ) + @staticmethod + def _get_stream_chunk_attr(chunk: Any, field_name: str) -> Any: + if isinstance(chunk, dict): + value = chunk.get(field_name) + if value is not None: + return value + model_extra = chunk.get("model_extra") + if isinstance(model_extra, dict): + value = model_extra.get(field_name) + if value is not None: + return value + hidden_params = chunk.get("_hidden_params") + if isinstance(hidden_params, dict): + return hidden_params.get(field_name) + return None + return getattr(chunk, field_name, None) + + @staticmethod + def _set_stream_metadata_on_response( + model_response: Any, + grounding_metadata: List[dict], + url_context_metadata: List[dict], + safety_ratings: List[dict], + citation_metadata: List[dict], + ) -> None: + setattr(model_response, "vertex_ai_grounding_metadata", grounding_metadata) # type: ignore + if grounding_metadata: + model_response._hidden_params["vertex_ai_grounding_metadata"] = ( + grounding_metadata + ) + setattr(model_response, "vertex_ai_url_context_metadata", url_context_metadata) # type: ignore + if url_context_metadata: + model_response._hidden_params["vertex_ai_url_context_metadata"] = ( + url_context_metadata + ) + setattr(model_response, "vertex_ai_safety_ratings", safety_ratings) # type: ignore + setattr(model_response, "vertex_ai_safety_results", safety_ratings) # type: ignore + if safety_ratings: + model_response._hidden_params["vertex_ai_safety_ratings"] = safety_ratings + model_response._hidden_params["vertex_ai_safety_results"] = safety_ratings + setattr(model_response, "vertex_ai_citation_metadata", citation_metadata) # type: ignore + if citation_metadata: + model_response._hidden_params["vertex_ai_citation_metadata"] = ( + citation_metadata + ) + + def apply_assembled_streaming_response_metadata( + self, + response: ModelResponse, + chunks: List[Any], + ) -> None: + for field_name in VERTEX_AI_PROVIDER_METADATA_FIELDS: + merged: List[Any] = [] + for chunk in chunks: + value = VertexGeminiConfig._get_stream_chunk_attr(chunk, field_name) + if not value: + continue + if isinstance(value, list): + merged.extend(value) + else: + merged.append(value) + if merged: + setattr(response, field_name, merged) + response._hidden_params[field_name] = merged + @staticmethod def _convert_grounding_metadata_to_annotations( grounding_metadata: List[dict], @@ -3390,10 +3456,13 @@ class ModelResponseIterator: if choice.finish_reason == "stop": choice.finish_reason = "tool_calls" - setattr(model_response, "vertex_ai_grounding_metadata", grounding_metadata) # type: ignore - setattr(model_response, "vertex_ai_url_context_metadata", url_context_metadata) # type: ignore - setattr(model_response, "vertex_ai_safety_ratings", safety_ratings) # type: ignore - setattr(model_response, "vertex_ai_citation_metadata", citation_metadata) # type: ignore + VertexGeminiConfig._set_stream_metadata_on_response( + model_response, + grounding_metadata, + url_context_metadata, + safety_ratings, + citation_metadata, + ) return ( grounding_metadata, diff --git a/litellm/main.py b/litellm/main.py index 64891e2def9..1a0d0312d73 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -7761,6 +7761,9 @@ def stream_chunk_builder( # noqa: PLR0915 "cost", logging_obj._response_cost_calculator(result=response), ) + processor.apply_provider_assembled_streaming_metadata( + response, chunks, logging_obj + ) return response tool_call_chunks = [ @@ -7940,6 +7943,9 @@ def stream_chunk_builder( # noqa: PLR0915 usage, "cost", logging_obj._response_cost_calculator(result=response) ) + processor.apply_provider_assembled_streaming_metadata( + response, chunks, logging_obj + ) return response except Exception as e: verbose_logger.exception( diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index 00db7b199b6..b28fee51284 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -758,3 +758,12 @@ class VertexPartnerProvider(str, Enum): llama = "llama" ai21 = "ai21" claude = "claude" + + +VERTEX_AI_PROVIDER_METADATA_FIELDS = ( + "vertex_ai_grounding_metadata", + "vertex_ai_url_context_metadata", + "vertex_ai_safety_ratings", + "vertex_ai_safety_results", + "vertex_ai_citation_metadata", +) diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index d57d8dafdbd..34edd6eccf3 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -2165,6 +2165,41 @@ def test_get_assembled_streaming_response_returns_result_for_streaming(): assert assembled is result +def test_streaming_success_handler_includes_vertex_ai_metadata_in_standard_logging(): + """Assembled streaming responses should include Vertex AI metadata in logging payload.""" + import datetime + + from litellm.types.utils import Choices, Message + + logging_obj = _make_logging_obj(stream=True) + grounding_metadata = [{"webSearchQueries": ["weather in SF"]}] + url_context_metadata = [{"urlMetadata": [{"retrievedUrl": "https://example.com"}]}] + result = ModelResponse( + id="resp-1", + choices=[ + Choices( + index=0, + message=Message(role="assistant", content="hello"), + finish_reason="stop", + ) + ], + model="gemini-2.5-flash", + ) + setattr(result, "vertex_ai_grounding_metadata", grounding_metadata) + setattr(result, "vertex_ai_url_context_metadata", url_context_metadata) + result._hidden_params["vertex_ai_grounding_metadata"] = grounding_metadata + result._hidden_params["vertex_ai_url_context_metadata"] = url_context_metadata + + start = datetime.datetime.now() + end = datetime.datetime.now() + logging_obj.success_handler(result=result, start_time=start, end_time=end) + + payload = logging_obj.model_call_details.get("standard_logging_object") + assert payload is not None + assert payload["response"]["vertex_ai_grounding_metadata"] == grounding_metadata + assert payload["response"]["vertex_ai_url_context_metadata"] == url_context_metadata + + def test_get_assembled_streaming_response_returns_none_for_non_streaming_text_completion(): """Non-streaming TextCompletionResponse should also return None.""" import datetime diff --git a/tests/test_litellm/litellm_core_utils/test_redact_messages.py b/tests/test_litellm/litellm_core_utils/test_redact_messages.py index 60cfff6e4a0..36f220f9a2c 100644 --- a/tests/test_litellm/litellm_core_utils/test_redact_messages.py +++ b/tests/test_litellm/litellm_core_utils/test_redact_messages.py @@ -349,3 +349,96 @@ class TestPerformRedaction: assert redacted.output[0].content[0].text == "redacted-by-litellm" assert response.output[0].content[0].text == "sensitive output" + + def test_redacts_vertex_provider_metadata_in_standard_logging_response(self): + details = { + "standard_logging_object": { + "messages": [{"role": "user", "content": "sensitive prompt"}], + "response": { + "choices": [ + { + "message": { + "content": "sensitive answer", + "role": "assistant", + } + } + ], + "vertex_ai_grounding_metadata": [ + {"webSearchQueries": ["sensitive search term"]} + ], + "vertex_ai_url_context_metadata": [ + {"urlMetadata": [{"retrievedUrl": "https://example.com"}]} + ], + }, + } + } + + perform_redaction(details, None) + + response = details["standard_logging_object"]["response"] + assert response["choices"][0]["message"]["content"] == "redacted-by-litellm" + assert response["vertex_ai_grounding_metadata"] == [] + assert response["vertex_ai_url_context_metadata"] == [] + + def test_redacts_vertex_provider_metadata_on_streaming_model_response(self): + response = litellm.ModelResponse( + id="resp-1", + choices=[ + litellm.Choices( + message=litellm.Message( + content="sensitive answer", + role="assistant", + ) + ) + ], + model="gemini-2.5-flash", + ) + setattr( + response, + "vertex_ai_grounding_metadata", + [{"webSearchQueries": ["sensitive search term"]}], + ) + response._hidden_params["vertex_ai_grounding_metadata"] = [ + {"webSearchQueries": ["sensitive search term"]} + ] + + details = { + "stream": True, + "complete_streaming_response": response, + } + + perform_redaction(details, response) + + assert response.choices[0].message.content == "redacted-by-litellm" + assert getattr(response, "vertex_ai_grounding_metadata") == [] + assert "vertex_ai_grounding_metadata" not in response._hidden_params + + def test_redacts_vertex_provider_metadata_from_metadata_hidden_params(self): + """Streaming success_handler copies _hidden_params into metadata before redaction.""" + details = { + "stream": True, + "litellm_params": { + "metadata": { + "hidden_params": { + "response_cost": 0.01, + "vertex_ai_grounding_metadata": [ + {"webSearchQueries": ["sensitive search term"]} + ], + "vertex_ai_url_context_metadata": [ + {"urlMetadata": [{"retrievedUrl": "https://example.com"}]} + ], + "vertex_ai_safety_ratings": [{"category": "HARM"}], + "vertex_ai_citation_metadata": [{"citations": ["source"]}], + } + } + }, + } + + perform_redaction(details, None) + + hidden_params = details["litellm_params"]["metadata"]["hidden_params"] + assert hidden_params["response_cost"] == 0.01 + assert "vertex_ai_grounding_metadata" not in hidden_params + assert "vertex_ai_url_context_metadata" not in hidden_params + assert "vertex_ai_safety_ratings" not in hidden_params + assert "vertex_ai_citation_metadata" not in hidden_params diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py index e40a0817fd9..77765340c61 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -613,3 +613,153 @@ def test_stream_chunk_builder_dict_snapshot_preserves_hidden_provider_fields(): assert ( response._hidden_params["provider_specific_fields"]["traffic_type"] == "default" ) + + +def test_stream_chunk_builder_propagates_vertex_ai_metadata_from_chunks(): + """Vertex AI metadata on streaming chunks must appear on assembled response.""" + grounding_metadata = [{"webSearchQueries": ["weather in SF"]}] + url_context_metadata = [{"urlMetadata": [{"retrievedUrl": "https://example.com"}]}] + + chunk1 = ModelResponseStream( + id="chatcmpl-vertex-1", + created=1, + model="gemini-2.5-flash", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta(content="The weather", role="assistant"), + ) + ], + ) + setattr(chunk1, "vertex_ai_grounding_metadata", grounding_metadata) + chunk1._hidden_params["vertex_ai_grounding_metadata"] = grounding_metadata + + chunk2 = ModelResponseStream( + id="chatcmpl-vertex-1", + created=1, + model="gemini-2.5-flash", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(content=" is sunny.", role="assistant"), + ) + ], + ) + setattr(chunk2, "vertex_ai_url_context_metadata", url_context_metadata) + chunk2._hidden_params["vertex_ai_url_context_metadata"] = url_context_metadata + + response = stream_chunk_builder(chunks=[chunk1, chunk2]) + assert response is not None + assert getattr(response, "vertex_ai_grounding_metadata") == grounding_metadata + assert getattr(response, "vertex_ai_url_context_metadata") == url_context_metadata + assert response._hidden_params["vertex_ai_grounding_metadata"] == grounding_metadata + assert ( + response._hidden_params["vertex_ai_url_context_metadata"] + == url_context_metadata + ) + + dumped = response.model_dump() + assert dumped["vertex_ai_grounding_metadata"] == grounding_metadata + assert dumped["vertex_ai_url_context_metadata"] == url_context_metadata + + +def test_stream_chunk_builder_uses_assembled_model_for_provider_metadata(): + grounding_metadata = [{"webSearchQueries": ["weather in SF"]}] + + chunk1 = ModelResponseStream( + id="chatcmpl-vertex-router", + created=1, + model="gpt-4o", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta(content="The weather", role="assistant"), + ) + ], + ) + chunk2 = ModelResponseStream( + id="chatcmpl-vertex-router", + created=1, + model="gemini-2.5-flash", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(content=" is sunny.", role=None), + ) + ], + ) + setattr(chunk2, "vertex_ai_grounding_metadata", grounding_metadata) + chunk2._hidden_params["vertex_ai_grounding_metadata"] = grounding_metadata + + response = stream_chunk_builder(chunks=[chunk1, chunk2]) + assert response is not None + assert response.model == "gemini-2.5-flash" + assert getattr(response, "vertex_ai_grounding_metadata") == grounding_metadata + + +def test_stream_chunk_builder_propagates_vertex_ai_safety_results(): + """Assembled response must expose safety data under the non-streaming field name.""" + safety_ratings = [ + [{"category": "HARM_CATEGORY_HATE_SPEECH", "probability": "NEGLIGIBLE"}] + ] + + chunk = ModelResponseStream( + id="chatcmpl-vertex-safety", + created=1, + model="gemini-2.5-flash", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(content="hello", role="assistant"), + ) + ], + ) + setattr(chunk, "vertex_ai_safety_ratings", safety_ratings) + setattr(chunk, "vertex_ai_safety_results", safety_ratings) + chunk._hidden_params["vertex_ai_safety_ratings"] = safety_ratings + chunk._hidden_params["vertex_ai_safety_results"] = safety_ratings + + response = stream_chunk_builder(chunks=[chunk]) + assert response is not None + assert getattr(response, "vertex_ai_safety_results") == safety_ratings + assert response._hidden_params["vertex_ai_safety_results"] == safety_ratings + assert response.model_dump()["vertex_ai_safety_results"] == safety_ratings + + +def test_stream_chunk_builder_propagates_vertex_ai_metadata_from_dict_chunks(): + """Dict snapshot chunks (model_dump) should also propagate Vertex AI metadata.""" + chunk_dict = ModelResponseStream( + id="chatcmpl-vertex-2", + created=1, + model="gemini-2.5-flash", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(content="hello", role="assistant"), + ) + ], + ).model_dump() + chunk_dict["_hidden_params"] = { + "vertex_ai_grounding_metadata": [{"webSearchQueries": ["test query"]}] + } + + response = stream_chunk_builder(chunks=[chunk_dict]) + assert response is not None + assert getattr(response, "vertex_ai_grounding_metadata") == [ + {"webSearchQueries": ["test query"]} + ] + assert response.model_dump()["vertex_ai_grounding_metadata"] == [ + {"webSearchQueries": ["test query"]} + ] diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 0d02521433a..671d7355e8f 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -1459,6 +1459,26 @@ def test_vertex_ai_process_candidates_with_grounding_metadata(): assert len(result[0]) == 1 +def test_set_stream_metadata_mirrors_non_streaming_safety_field_names(): + safety_ratings = [ + [{"category": "HARM_CATEGORY_HATE_SPEECH", "probability": "NEGLIGIBLE"}] + ] + + model_response = ModelResponse() + VertexGeminiConfig._set_stream_metadata_on_response( + model_response=model_response, + grounding_metadata=[], + url_context_metadata=[], + safety_ratings=safety_ratings, + citation_metadata=[], + ) + + assert getattr(model_response, "vertex_ai_safety_ratings") == safety_ratings + assert getattr(model_response, "vertex_ai_safety_results") == safety_ratings + assert model_response._hidden_params["vertex_ai_safety_ratings"] == safety_ratings + assert model_response._hidden_params["vertex_ai_safety_results"] == safety_ratings + + def test_vertex_ai_tool_call_id_format(): """ Test that tool call IDs have the correct format and length. From f59e4ebc9e98e26b434aa3ee338f41b3f3cf3fa6 Mon Sep 17 00:00:00 2001 From: milan-berri Date: Tue, 9 Jun 2026 02:27:35 +0300 Subject: [PATCH 025/185] fix(ui): show team projects to internal users (#28855) Allow internal users to fetch their backend-scoped project list so the key creation project dropdown can populate for selected teams. --- .../app/(dashboard)/hooks/projects/useProjects.test.ts | 10 +++++++++- .../src/app/(dashboard)/hooks/projects/useProjects.ts | 6 ++++-- 2 files changed, 13 insertions(+), 3 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useProjects.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useProjects.test.ts index 13b9107bdc1..39d1b28303d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useProjects.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useProjects.test.ts @@ -115,8 +115,16 @@ describe("useProjects", () => { expect(global.fetch).not.toHaveBeenCalled(); }); - it("should not fetch when userRole is not an admin role", () => { + it("should fetch when userRole is an internal user role", async () => { mockUseAuthorized.mockReturnValue({ accessToken: "test-token", userRole: "Internal User" }); + (global.fetch as any).mockResolvedValue({ ok: true, json: async () => mockProjects }); + const { result } = renderHook(() => useProjects(), { wrapper: makeWrapper(queryClient) }); + await waitFor(() => expect(result.current.isSuccess).toBe(true)); + expect(global.fetch).toHaveBeenCalled(); + }); + + it("should not fetch when userRole cannot read projects", () => { + mockUseAuthorized.mockReturnValue({ accessToken: "test-token", userRole: "regular_user" }); const { result } = renderHook(() => useProjects(), { wrapper: makeWrapper(queryClient) }); expect(result.current.isFetched).toBe(false); expect(global.fetch).not.toHaveBeenCalled(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useProjects.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useProjects.ts index 7bdc8a4fe6d..c240dbb0170 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useProjects.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useProjects.ts @@ -2,7 +2,7 @@ import { useQuery } from "@tanstack/react-query"; import { createQueryKeys } from "../common/queryKeysFactory"; import { getProxyBaseUrl, getGlobalLitellmHeaderName, deriveErrorMessage, handleError } from "@/components/networking"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; -import { all_admin_roles } from "@/utils/roles"; +import { all_admin_roles, internalUserRoles } from "@/utils/roles"; // ── Types ──────────────────────────────────────────────────────────────────── @@ -42,6 +42,8 @@ export interface ProjectResponse { export const projectKeys = createQueryKeys("projects"); +const projectReaderRoles = [...all_admin_roles, ...internalUserRoles]; + // ── Fetch function ─────────────────────────────────────────────────────────── const fetchProjects = async (accessToken: string): Promise => { @@ -74,6 +76,6 @@ export const useProjects = () => { return useQuery({ queryKey: projectKeys.list({}), queryFn: async () => fetchProjects(accessToken!), - enabled: Boolean(accessToken) && all_admin_roles.includes(userRole!), + enabled: Boolean(accessToken) && projectReaderRoles.includes(userRole!), }); }; From bac2590b39b3a32a25ae4d6a932a342349def497 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 8 Jun 2026 16:39:21 -0700 Subject: [PATCH 026/185] build(deps): bump pyjwt to 2.13.0 and ws override to 8.20.1 (#29982) Raise the PyJWT floor in pyproject (>=2.13.0,<3.0) and re-resolve uv.lock so the proxy installs 2.13.0 instead of 2.12.0. Bump the ws transitive-version override in the dashboard from 8.19.0 to 8.20.1 and regenerate package-lock; jsdom and openai both dedupe onto the single 8.20.1 copy. Both are routine dependency maintenance bumps to keep pinned versions current. --- pyproject.toml | 2 +- ui/litellm-dashboard/package-lock.json | 6 +++--- ui/litellm-dashboard/package.json | 2 +- uv.lock | 13 ++++++++----- 4 files changed, 13 insertions(+), 10 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 577e800d79b..28e6f48dc4c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -53,7 +53,7 @@ proxy = [ "orjson>=3.11.6,<4.0", "apscheduler>=3.11.2,<4.0", "fastapi-sso>=0.19.0,<1.0", - "PyJWT>=2.12.0,<3.0", + "PyJWT>=2.13.0,<3.0", "python-multipart>=0.0.27,<1.0", "cryptography>=46.0.7,<47.0", "pynacl>=1.6.2,<2.0", diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json index 7000efd63b4..b7dd2a6f59b 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -13770,9 +13770,9 @@ } }, "node_modules/ws": { - "version": "8.19.0", - "resolved": "https://registry.npmjs.org/ws/-/ws-8.19.0.tgz", - "integrity": "sha512-blAT2mjOEIi0ZzruJfIhb3nps74PRWTCz1IjglWEEpQl5XS/UNama6u2/rjFkDDouqr4L67ry+1aGIALViWjDg==", + "version": "8.20.1", + "resolved": "https://registry.npmjs.org/ws/-/ws-8.20.1.tgz", + "integrity": "sha512-It4dO0K5v//JtTXuPkfEOaI3uUN87iYPnqo/ZzqCoG3g8uhA66QUMs/SrM0YK7/NAu+r4LMh/9dq2A7k+rHs+w==", "devOptional": true, "license": "MIT", "engines": { diff --git a/ui/litellm-dashboard/package.json b/ui/litellm-dashboard/package.json index ca753e59dc5..623389e76ed 100644 --- a/ui/litellm-dashboard/package.json +++ b/ui/litellm-dashboard/package.json @@ -90,7 +90,7 @@ "glob": "13.0.0", "minimatch": "10.2.4", "lodash": "4.18.1", - "ws": "8.19.0", + "ws": "8.20.1", "braces": "3.0.3", "axios": "1.13.6", "postcss": "8.5.13" diff --git a/uv.lock b/uv.lock index 4386fce1b4d..2403a7fbf03 100644 --- a/uv.lock +++ b/uv.lock @@ -9,7 +9,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-06-03T21:40:52.018333Z" +exclude-newer = "2026-06-05T23:18:37.734017Z" exclude-newer-span = "P3D" [manifest] @@ -3522,7 +3522,7 @@ requires-dist = [ { name = "prometheus-client", marker = "extra == 'proxy-runtime'", specifier = ">=0.20.0,<1.0" }, { name = "pydantic", specifier = ">=2.10.0,<3.0.0" }, { name = "pydantic-settings", marker = "extra == 'proxy'", specifier = ">=2.14.1,<3.0" }, - { name = "pyjwt", marker = "extra == 'proxy'", specifier = ">=2.12.0,<3.0" }, + { name = "pyjwt", marker = "extra == 'proxy'", specifier = ">=2.13.0,<3.0" }, { name = "pynacl", marker = "extra == 'proxy'", specifier = ">=1.6.2,<2.0" }, { name = "pypdf", marker = "python_full_version < '3.14' and extra == 'proxy-runtime'", specifier = ">=6.10.2,<7.0" }, { name = "pyroscope-io", marker = "sys_platform != 'win32' and extra == 'proxy'", specifier = ">=0.8.16,<1.0" }, @@ -5983,11 +5983,14 @@ wheels = [ [[package]] name = "pyjwt" -version = "2.12.0" +version = "2.13.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/a8/10/e8192be5f38f3e8e7e046716de4cae33d56fd5ae08927a823bb916be36c1/pyjwt-2.12.0.tar.gz", hash = "sha256:2f62390b667cd8257de560b850bb5a883102a388829274147f1d724453f8fb02", size = 102511, upload-time = "2026-03-12T17:15:30.831Z" } +dependencies = [ + { name = "typing-extensions", marker = "python_full_version < '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/3b/81/58d0ac84e1ef3a3843791d6954d94c0b33d526c75eeb1efbce9d0a4c4077/pyjwt-2.13.0.tar.gz", hash = "sha256:41571c89ca91598c79e8ef18a2d07367d4810fbbd6f637794879baf1b7703423", size = 107515, upload-time = "2026-05-21T19:54:36.618Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/15/70/70f895f404d363d291dcf62c12c85fdd47619ad9674ac0f53364d035925a/pyjwt-2.12.0-py3-none-any.whl", hash = "sha256:9bb459d1bdd0387967d287f5656bf7ec2b9a26645d1961628cda1764e087fd6e", size = 29700, upload-time = "2026-03-12T17:15:29.257Z" }, + { url = "https://files.pythonhosted.org/packages/a3/5e/ecf12fdb62546d64385c158514e9b2b671f7832108ef2ecd2020ce0af2d1/pyjwt-2.13.0-py3-none-any.whl", hash = "sha256:66adcc2aff09b3f1bbd95fc1e1577df8ac8723c978552fd43304c8a290ac5728", size = 31274, upload-time = "2026-05-21T19:54:35.362Z" }, ] [package.optional-dependencies] From c24a3603d971c8af8a29d4f6ef2fa16cfd097100 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 8 Jun 2026 16:55:35 -0700 Subject: [PATCH 027/185] fix(team-management): delete a team's BYOK models when the team is deleted (#29977) A team's BYOK models (rows in LiteLLM_ProxyModelTable with model_info.team_id set) were left orphaned when the team was deleted; they lingered in the database and kept showing on the Models + Endpoints page. delete_team now removes them via a new delete_team_models helper that deletes the rows in one transaction and syncs the in-memory router only after that transaction commits, run before the team rows are deleted so a mid-flight failure never leaves the team gone with its models orphaned --- .../model_management_endpoints.py | 44 +++++- .../management_endpoints/team_endpoints.py | 14 ++ .../test_model_management_endpoints.py | 143 ++++++++++++++++++ .../test_team_endpoints.py | 10 ++ 4 files changed, 209 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 0cbccfc18ad..0be476469a6 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -561,7 +561,7 @@ async def _setup_new_team_model_assignment( async def _get_team_deployments( - team_id: str, prisma_client: PrismaClient + team_id: str, prisma_client: PrismaClient, table: Optional[Any] = None ) -> List[LiteLLM_ProxyModelTable]: """ Fetch all deployments for a given team_id from the database. @@ -572,9 +572,13 @@ async def _get_team_deployments( Note: prisma-client-py 0.11.0 does not support JSON path filtering, so we filter by the model_name prefix (team models use "model_name_{team_id}_*") and confirm team_id in model_info with Python-side filtering. + + Pass ``table`` (a transaction's proxy-model table) to run the read inside an + existing transaction. """ prefix = f"model_name_{team_id}_" - response = await ModelRepository(prisma_client).table.find_many( + table = table or ModelRepository(prisma_client).table + response = await table.find_many( where={ "model_name": {"startswith": prefix}, } @@ -596,6 +600,42 @@ async def _get_team_deployments( return result +async def delete_team_models( + team_ids: List[str], + prisma_client: PrismaClient, + llm_router: Optional[Any], +) -> List[str]: + """ + Delete every BYOK model owned by the given teams, from the DB and the router. + + The DB rows are removed inside a single transaction, so deletion is atomic + across all team_ids. Each team's rows are deleted by the exact model_ids read + in the same transaction, which keeps the deleted set identical to the set + handed to the router. The router is synced only after the transaction commits, + so a rollback can never leave a deployment live in the router without its row. + + Returns the model_ids that were deleted. + """ + deleted_model_ids: List[str] = [] + async with prisma_client.db.tx() as tx: + for team_id in team_ids: + rows = await _get_team_deployments( + team_id, prisma_client, table=tx.litellm_proxymodeltable + ) + model_ids = [row.model_id for row in rows] + if model_ids: + await tx.litellm_proxymodeltable.delete_many( + where={"model_id": {"in": model_ids}} + ) + deleted_model_ids.extend(model_ids) + + if llm_router is not None: + for model_id in deleted_model_ids: + llm_router.delete_deployment(id=model_id) + + return deleted_model_ids + + async def _get_team_public_model_names( team_id: str, prisma_client: PrismaClient, diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index f2eafcbf839..0e69de87ce2 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -3289,6 +3289,20 @@ async def delete_team( await prisma_client.delete_data(team_id_list=data.team_ids, table_name="key") + ## DELETE ASSOCIATED BYOK MODELS + # Runs before the team rows are deleted so a mid-flight failure never leaves + # the team gone with its models orphaned. + from litellm.proxy.management_endpoints.model_management_endpoints import ( + delete_team_models, + ) + from litellm.proxy.proxy_server import llm_router + + await delete_team_models( + team_ids=data.team_ids, + prisma_client=prisma_client, + llm_router=llm_router, + ) + # ## DELETE TEAM MEMBERSHIPS for team_row in team_rows: ### get all team members diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index 2ba604e5da6..6a81b1b613b 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -24,6 +24,7 @@ from litellm.proxy.management_endpoints.model_management_endpoints import ( ModelManagementAuthChecks, _get_team_deployments, clear_cache, + delete_team_models, ) from litellm.proxy.utils import PrismaClient from litellm.types.router import Deployment, LiteLLM_Params, updateDeployment @@ -2184,6 +2185,148 @@ class TestGetTeamDeployments: assert result[0] is dep1 +def _model_row(model_id: str, team_id: str): + row = MagicMock() + row.model_id = model_id + row.model_name = f"model_name_{team_id}_{model_id}" + row.model_info = {"team_id": team_id} + return row + + +class _TxProxyModelTable: + """Transactional proxy-model table that records the order of DB writes.""" + + def __init__(self, rows, events): + self._rows = list(rows) + self.events = events + + async def find_many(self, where): + prefix = where["model_name"]["startswith"] + return [r for r in self._rows if r.model_name.startswith(prefix)] + + async def delete_many(self, where): + ids = list(where["model_id"]["in"]) + self.events.append(("delete_many", tuple(ids))) + self._rows = [r for r in self._rows if r.model_id not in ids] + return len(ids) + + +class _TxPrismaClient: + """Minimal prisma stub whose ``db.tx()`` yields a transaction and records commit.""" + + def __init__(self, rows): + self.events: list = [] + self._table = _TxProxyModelTable(rows, self.events) + tx = MagicMock() + tx.litellm_proxymodeltable = self._table + outer = self + + class _TxCM: + async def __aenter__(self): + return tx + + async def __aexit__(self, *exc): + outer.events.append(("commit",)) + return False + + self.db = MagicMock() + self.db.tx = MagicMock(return_value=_TxCM()) + + +class _RecordingRouter: + def __init__(self, events): + self.events = events + self.deleted: list = [] + + def delete_deployment(self, id): # noqa: A002 - matches router signature + self.events.append(("router", id)) + self.deleted.append(id) + + +class TestDeleteTeamModels: + """delete_team_models must remove every team's BYOK models in one transaction + and sync the in-memory router only after that transaction commits.""" + + @pytest.mark.asyncio + async def test_deletes_all_teams_models_and_syncs_router(self): + rows = [_model_row("a1", "team_a"), _model_row("b1", "team_b")] + prisma = _TxPrismaClient(rows) + router = _RecordingRouter(prisma.events) + + deleted = await delete_team_models( + team_ids=["team_a", "team_b"], + prisma_client=prisma, + llm_router=router, + ) + + assert sorted(deleted) == ["a1", "b1"] + assert sorted(router.deleted) == ["a1", "b1"] + + @pytest.mark.asyncio + async def test_router_sync_happens_after_commit(self): + """Race-safety: the router is touched only once the DB transaction has + committed, so a rollback can never leave a deployment without its row.""" + rows = [_model_row("a1", "team_a"), _model_row("b1", "team_b")] + prisma = _TxPrismaClient(rows) + router = _RecordingRouter(prisma.events) + + await delete_team_models( + team_ids=["team_a", "team_b"], prisma_client=prisma, llm_router=router + ) + + commit_idx = prisma.events.index(("commit",)) + router_indices = [i for i, e in enumerate(prisma.events) if e[0] == "router"] + delete_indices = [ + i for i, e in enumerate(prisma.events) if e[0] == "delete_many" + ] + assert router_indices, "router was never synced" + assert all(i > commit_idx for i in router_indices) + assert all(i < commit_idx for i in delete_indices) + + @pytest.mark.asyncio + async def test_only_owning_team_models_deleted(self): + """A row sharing the prefix but a different model_info.team_id is left alone.""" + mine = _model_row("a1", "team_a") + intruder = MagicMock() + intruder.model_id = "x9" + intruder.model_name = "model_name_team_a_x9" + intruder.model_info = {"team_id": "someone_else"} + prisma = _TxPrismaClient([mine, intruder]) + router = _RecordingRouter(prisma.events) + + deleted = await delete_team_models( + team_ids=["team_a"], prisma_client=prisma, llm_router=router + ) + + assert deleted == ["a1"] + assert router.deleted == ["a1"] + + @pytest.mark.asyncio + async def test_no_models_no_writes(self): + prisma = _TxPrismaClient([]) + router = _RecordingRouter(prisma.events) + + deleted = await delete_team_models( + team_ids=["team_a"], prisma_client=prisma, llm_router=router + ) + + assert deleted == [] + assert router.deleted == [] + assert not any(e[0] == "delete_many" for e in prisma.events) + + @pytest.mark.asyncio + async def test_missing_router_is_safe(self): + rows = [_model_row("a1", "team_a")] + prisma = _TxPrismaClient(rows) + + deleted = await delete_team_models( + team_ids=["team_a"], prisma_client=prisma, llm_router=None + ) + + assert deleted == ["a1"] + assert any(e[0] == "delete_many" for e in prisma.events) + + def _build_db_model_for_blocked_test(): from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 400ed287802..06adcf80707 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -6350,6 +6350,14 @@ async def test_delete_team_persists_deleted_teams(monkeypatch): mock_find_many_keys = AsyncMock(return_value=[]) mock_prisma_client.db.litellm_verificationtoken.find_many = mock_find_many_keys + # delete_team now deletes team BYOK models inside a transaction; this team has none. + mock_tx = AsyncMock() + mock_tx.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + mock_tx_cm = MagicMock() + mock_tx_cm.__aenter__ = AsyncMock(return_value=mock_tx) + mock_tx_cm.__aexit__ = AsyncMock(return_value=False) + mock_prisma_client.db.tx = MagicMock(return_value=mock_tx_cm) + monkeypatch.setattr( "litellm.proxy.proxy_server.prisma_client", mock_prisma_client, @@ -8499,6 +8507,8 @@ async def test_new_team_encrypts_callback_vars( assert cv["langfuse_secret_key"] != "sk-real" recovered = decrypt_callback_vars(metadata)["logging"][0]["callback_vars"] assert recovered["langfuse_secret_key"] == "sk-real" + + def _non_admin_auth(): return UserAPIKeyAuth( user_id="u-team-admin", user_role=LitellmUserRoles.INTERNAL_USER From 411bd3da5ba6abb9b2aae173d3274317e5805364 Mon Sep 17 00:00:00 2001 From: milan-berri Date: Tue, 9 Jun 2026 02:59:21 +0300 Subject: [PATCH 028/185] feat(vantage): include organization metadata in FOCUS Tags export (#28184) * feat(vantage): include organization metadata in FOCUS Tags export Join LiteLLM_OrganizationTable when building Vantage/FOCUS export rows so organization_id and organization_alias appear in Tags for org-level filtering. Co-authored-by: Cursor * test(focus): include api_requests in organization Tags tests FocusTransformer now requires api_requests after staging merge; add the column to test fixtures so integrations CI can run the Tags assertions. Co-authored-by: Cursor --------- Co-authored-by: Cursor --- litellm/integrations/focus/database.py | 6 +- litellm/integrations/focus/transformer.py | 2 + .../integrations/focus/test_focus_database.py | 15 +++++ .../integrations/focus/test_transformer.py | 62 +++++++++++++++++++ 4 files changed, 84 insertions(+), 1 deletion(-) create mode 100644 tests/test_litellm/integrations/focus/test_transformer.py diff --git a/litellm/integrations/focus/database.py b/litellm/integrations/focus/database.py index 298254670eb..3ae3f6b53ac 100644 --- a/litellm/integrations/focus/database.py +++ b/litellm/integrations/focus/database.py @@ -80,11 +80,15 @@ class FocusLiteLLMDatabase: vt.team_id, vt.key_alias as api_key_alias, tt.team_alias, - ut.user_email as user_email + ut.user_email as user_email, + COALESCE(vt.organization_id, tt.organization_id) as organization_id, + ot.organization_alias as organization_alias FROM "LiteLLM_DailyUserSpend" dus LEFT JOIN "LiteLLM_VerificationToken" vt ON dus.api_key = vt.token LEFT JOIN "LiteLLM_TeamTable" tt ON vt.team_id = tt.team_id LEFT JOIN "LiteLLM_UserTable" ut ON dus.user_id = ut.user_id + LEFT JOIN "LiteLLM_OrganizationTable" ot + ON ot.organization_id = COALESCE(vt.organization_id, tt.organization_id) {where_clause} ORDER BY dus.date DESC, dus.created_at DESC {limit_clause} diff --git a/litellm/integrations/focus/transformer.py b/litellm/integrations/focus/transformer.py index 8496b7ec159..a17df29b912 100644 --- a/litellm/integrations/focus/transformer.py +++ b/litellm/integrations/focus/transformer.py @@ -12,6 +12,8 @@ from .schema import FOCUS_NORMALIZED_SCHEMA _TAG_KEYS = ( "team_id", "team_alias", + "organization_id", + "organization_alias", "user_id", "user_email", "api_key_alias", diff --git a/tests/test_litellm/integrations/focus/test_focus_database.py b/tests/test_litellm/integrations/focus/test_focus_database.py index 5ee98cc9dd0..d77af2dd170 100644 --- a/tests/test_litellm/integrations/focus/test_focus_database.py +++ b/tests/test_litellm/integrations/focus/test_focus_database.py @@ -72,3 +72,18 @@ async def test_should_reject_invalid_limit(monkeypatch: pytest.MonkeyPatch): await db.get_usage_data(limit="invalid") assert query_mock.await_count == 0 + + +@pytest.mark.asyncio +async def test_should_join_organization_table(monkeypatch: pytest.MonkeyPatch): + db, query_mock = _setup_db(monkeypatch, []) + + await db.get_usage_data() + + query_text, *_ = query_mock.await_args.args + assert ( + "COALESCE(vt.organization_id, tt.organization_id) as organization_id" + in query_text + ) + assert "ot.organization_alias as organization_alias" in query_text + assert 'LEFT JOIN "LiteLLM_OrganizationTable" ot' in query_text diff --git a/tests/test_litellm/integrations/focus/test_transformer.py b/tests/test_litellm/integrations/focus/test_transformer.py new file mode 100644 index 00000000000..4461d19efde --- /dev/null +++ b/tests/test_litellm/integrations/focus/test_transformer.py @@ -0,0 +1,62 @@ +"""Tests for FocusTransformer organization metadata in Tags.""" + +from __future__ import annotations + +import json +from datetime import date + +import polars as pl + +from litellm.integrations.focus.transformer import FocusTransformer + + +def test_should_include_organization_fields_in_tags(): + frame = pl.DataFrame( + { + "date": [date(2024, 1, 2)], + "spend": [1.25], + "api_requests": [1], + "api_key": ["hashed-key"], + "api_key_alias": ["prod-key"], + "model": ["gpt-4o"], + "model_group": ["gpt-4o"], + "custom_llm_provider": ["openai"], + "team_id": ["team-1"], + "team_alias": ["Platform"], + "organization_id": ["org-123"], + "organization_alias": ["Acme Corp"], + "user_id": ["user-1"], + "user_email": ["user@example.com"], + } + ) + + normalized = FocusTransformer().transform(frame) + + tags = json.loads(normalized["Tags"][0]) + assert tags["organization_id"] == "org-123" + assert tags["organization_alias"] == "Acme Corp" + assert tags["team_id"] == "team-1" + + +def test_should_omit_missing_organization_fields_from_tags(): + frame = pl.DataFrame( + { + "date": [date(2024, 1, 2)], + "spend": [0.5], + "api_requests": [1], + "api_key": ["hashed-key"], + "api_key_alias": ["prod-key"], + "model": ["gpt-4o-mini"], + "model_group": ["gpt-4o-mini"], + "custom_llm_provider": ["openai"], + "team_id": ["team-1"], + "team_alias": ["Platform"], + } + ) + + normalized = FocusTransformer().transform(frame) + + tags = json.loads(normalized["Tags"][0]) + assert "organization_id" not in tags + assert "organization_alias" not in tags + assert tags["team_id"] == "team-1" From 1bbaf1c39dda367f5a2b4b6b9ab4cac46d71ab14 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Mon, 8 Jun 2026 17:46:28 -0700 Subject: [PATCH 029/185] fix(guardrails): read CrowdStrike AIDR identity from both metadata bags (#29991) Capture user_id and extra_info from metadata or litellm_metadata. The single-bag read dropped identity whenever a request carried a present litellm_metadata field (null or a user-supplied dict), since /chat/completions routes the authenticated identity into metadata while the guardrail read litellm_metadata first --- .../crowdstrike_aidr/crowdstrike_aidr.py | 14 +++++- .../guardrail_hooks/test_crowdstrike_aidr.py | 50 +++++++++++++++++++ 2 files changed, 62 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py index d1ef165b46e..248202b644c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py +++ b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py @@ -105,6 +105,16 @@ def _extract_text_from_content(content: object) -> str: return "" +def _merge_metadata_bags(request_data: Mapping[str, Any]) -> Optional[dict[str, Any]]: + merged: dict[str, Any] = {} + present = False + for bag in (request_data.get("metadata"), request_data.get("litellm_metadata")): + if isinstance(bag, Mapping): + present = True + merged.update(bag) + return merged if present else None + + class CrowdStrikeAIDRHandler(CustomGuardrail): """ CrowdStrike AIDR AI Guardrail handler to interact with the CrowdStrike AIDR @@ -321,8 +331,8 @@ class CrowdStrikeAIDRHandler(CustomGuardrail): if model: ai_guard_payload["model"] = model - metadata = request_data.get("litellm_metadata", request_data.get("metadata")) - if isinstance(metadata, Mapping): + metadata = _merge_metadata_bags(request_data) + if metadata is not None: user_id = metadata.get("user_api_key_user_id") if user_id: ai_guard_payload["user_id"] = user_id diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py index e7f72ff7a3f..f8fd9a0a185 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py @@ -529,6 +529,56 @@ async def test_apply_guardrail_no_metadata_skips_user_fields( assert "extra_info" not in payload +@pytest.mark.asyncio +@pytest.mark.parametrize( + "litellm_metadata, metadata", + [ + (None, {"user_api_key_user_id": "uid-abc", "user_api_key_user_email": "alice@example.com"}), + ({"trace_id": "t1"}, {"user_api_key_user_id": "uid-abc", "user_api_key_user_email": "alice@example.com"}), + (["unexpected"], {"user_api_key_user_id": "uid-abc", "user_api_key_user_email": "alice@example.com"}), + ({"user_api_key_user_id": "uid-abc", "user_api_key_user_email": "alice@example.com"}, {"trace_id": "t1"}), + ], + ids=["identity_in_metadata_llm_none", "identity_in_metadata_llm_user_dict", "identity_in_metadata_llm_non_mapping", "identity_in_litellm_metadata"], +) +async def test_apply_guardrail_reads_identity_from_either_metadata_bag( + crowdstrike_aidr_guardrail: CrowdStrikeAIDRHandler, + litellm_metadata, + metadata, +) -> None: + inputs: GenericGuardrailAPIInputs = { + "texts": ["Hello"], + "structured_messages": [{"role": "user", "content": "Hello"}], + "model": "gpt-4o", + } + request_data = { + "messages": inputs["structured_messages"], + "model": "gpt-4o", + "litellm_metadata": litellm_metadata, + "metadata": metadata, + } + guardrail_endpoint = ( + f"{crowdstrike_aidr_guardrail.api_base}/v1/guard_chat_completions" + ) + + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + return_value=httpx.Response( + status_code=200, + json={"result": {"blocked": False, "transformed": False}}, + request=httpx.Request(method="POST", url=guardrail_endpoint), + ), + ) as mock_method: + await crowdstrike_aidr_guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + ) + + payload = mock_method.call_args.kwargs["json"] + assert payload["user_id"] == "uid-abc" + assert payload["extra_info"] == {"user_name": "alice@example.com"} + + @pytest.mark.asyncio async def test_apply_guardrail_request_skipped_messages_stay_aligned( crowdstrike_aidr_guardrail: CrowdStrikeAIDRHandler, From 92817cb65bc351b71543e6f980850953f7c1511e Mon Sep 17 00:00:00 2001 From: tin-berri Date: Mon, 8 Jun 2026 18:13:06 -0700 Subject: [PATCH 030/185] changing expires_in default to use actual slack return details (#29951) --- .../mcp_server/discoverable_endpoints.py | 7 +- .../mcp_server/test_discoverable_endpoints.py | 72 +++++++++++++++++++ 2 files changed, 76 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index ed374635fea..3beddd2c435 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -512,12 +512,13 @@ async def exchange_token_with_server( result = { "access_token": access_token, "token_type": token_response.get("token_type", "Bearer"), - "expires_in": token_response.get("expires_in", 3600), } - if "refresh_token" in token_response and token_response["refresh_token"]: + if token_response.get("expires_in") is not None: + result["expires_in"] = token_response["expires_in"] + if token_response.get("refresh_token"): result["refresh_token"] = token_response["refresh_token"] - if "scope" in token_response and token_response["scope"]: + if token_response.get("scope"): result["scope"] = token_response["scope"] # RFC 6749 §5.1: token responses must not be cached. diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index da66d60aed8..6fd935e3364 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -1,5 +1,6 @@ """Tests for MCP OAuth discoverable endpoints""" +import json from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -2661,3 +2662,74 @@ async def test_token_endpoint_sets_no_store_cache_control(): assert response.headers["cache-control"] == "no-store" assert response.headers["pragma"] == "no-cache" + + +async def _exchange_with_upstream_token_response(upstream_body): + from fastapi import Request + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + exchange_token_with_server, + ) + from litellm.proxy._types import MCPTransport + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="t", + name="t", + server_name="t", + alias="t", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + client_id="cid", + client_secret="cs", + authorization_url="https://provider.com/oauth/authorize", + token_url="https://provider.com/oauth/token", + ) + mock_request = MagicMock(spec=Request) + mock_request.base_url = "https://litellm.example.com/" + mock_request.headers = {} + + fake_http_response = MagicMock() + fake_http_response.json.return_value = upstream_body + fake_http_response.raise_for_status = MagicMock() + fake_http_client = MagicMock() + fake_http_client.post = AsyncMock(return_value=fake_http_response) + + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", + return_value=fake_http_client, + ): + response = await exchange_token_with_server( + request=mock_request, + mcp_server=server, + grant_type="authorization_code", + code="c", + redirect_uri="http://127.0.0.1:3000/cb", + client_id="cid", + client_secret=None, + code_verifier=None, + ) + return json.loads(response.body) + + +@pytest.mark.asyncio +async def test_token_exchange_omits_expires_in_when_upstream_omits_it(): + """A provider that issues a non-expiring token (e.g. Slack without token + rotation) returns no ``expires_in``. The exchange must mirror that and omit + ``expires_in`` rather than fabricate a 1-hour TTL, so the stored credential + is treated as non-expiring instead of dying after an hour.""" + body = await _exchange_with_upstream_token_response( + {"access_token": "tok", "token_type": "Bearer"} + ) + assert "expires_in" not in body + + +@pytest.mark.asyncio +async def test_token_exchange_passes_through_upstream_expires_in(): + """When the provider does send ``expires_in`` (e.g. Slack with token + rotation), the exchange forwards the real value unchanged.""" + body = await _exchange_with_upstream_token_response( + {"access_token": "tok", "token_type": "Bearer", "expires_in": 43200} + ) + assert body["expires_in"] == 43200 From 424db6a9807d6e830492825598332f9bb7f234c6 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 9 Jun 2026 06:57:04 +0530 Subject: [PATCH 031/185] feat(azure_ai): add MAI-Image-2.5 image generation support (#29688) * feat(azure_ai): add MAI-Image-2.5 image generation support Route azure_ai MAI models to /mai/v1/images/generations and map OpenAI size to width/height for the serverless API. Co-authored-by: Cursor * fix(azure_ai): address MAI image generation review feedback Validate unsupported size values, default width/height independently, add MAI-Image-2.5 pricing, and expand test coverage. @greptileai Co-authored-by: Cursor * feat(azure_ai): add MAI image edit and expand model cost map Add MAI image edit support with usage normalization for Azure response format, and register MAI-Image-2.5-Flash and MAI-Image-2e pricing in the model map. Co-authored-by: Cursor * fix(azure_ai): validate MAI edit size by consuming map iterator Greptile: lazy map() never evaluated int() so values like 1024xabc passed through. Co-authored-by: Cursor * fix(azure_ai): normalize MAI usage in generation response handler Apply normalize_mai_image_usage before building ImageResponse so token-based cost calculation works when Azure returns num_output_tokens fields. Co-authored-by: Cursor * fix(azure_ai): narrow MAI edit size param type for mypy Co-authored-by: Cursor * Fix Azure MAI image response handling * Fix MAI image generation base model routing * fix(azure_ai): preserve zero num_output_tokens in MAI usage normalization * fix(azure_ai): wrap MAI generation response JSON parsing in error handling * fix(azure_ai): build MAI image edit URL correctly for /mai/ root bases * fix(azure_ai): build MAI image generation URL correctly for /mai/ root bases --------- Co-authored-by: Cursor Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/llms/azure/azure.py | 39 +- .../llms/azure/image_generation/__init__.py | 3 + litellm/llms/azure_ai/image_edit/__init__.py | 14 +- .../azure_ai/image_edit/mai_transformation.py | 199 +++++++++ .../azure_ai/image_generation/__init__.py | 4 + .../image_generation/cost_calculator.py | 26 +- .../image_generation/mai_transformation.py | 236 +++++++++++ ...odel_prices_and_context_window_backup.json | 37 ++ model_prices_and_context_window.json | 37 ++ .../test_azure_image_generation_init.py | 16 + .../test_mai_image_edit_transformation.py | 171 ++++++++ .../test_mai_image_generation.py | 380 ++++++++++++++++++ 12 files changed, 1149 insertions(+), 13 deletions(-) create mode 100644 litellm/llms/azure_ai/image_edit/mai_transformation.py create mode 100644 litellm/llms/azure_ai/image_generation/mai_transformation.py create mode 100644 tests/test_litellm/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py create mode 100644 tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index 734b8ecef16..56cf035d0f7 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -43,7 +43,10 @@ from .common_utils import ( process_azure_headers, select_azure_base_url_or_endpoint, ) -from .image_generation import get_azure_image_generation_config +from .image_generation import ( + AzureFoundryMAIImageGenerationConfig, + get_azure_image_generation_config, +) from .image_generation.http_utils import azure_deployment_image_generation_json_body @@ -1097,10 +1100,14 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): ) def create_azure_base_url( - self, azure_client_params: dict, model: Optional[str] + self, + azure_client_params: dict, + model: Optional[str], + base_model: Optional[str] = None, ) -> str: from litellm.llms.azure_ai.image_generation import ( AzureFoundryFluxImageGenerationConfig, + AzureFoundryMAIImageGenerationConfig, ) api_base: str = azure_client_params.get( @@ -1112,6 +1119,12 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): if model is None: model = "" + if AzureFoundryMAIImageGenerationConfig.is_mai_model(base_model or model): + return AzureFoundryMAIImageGenerationConfig.get_mai_image_generation_url( + api_base=api_base, + api_version=api_version, + ) + # Handle FLUX 2 models on Azure AI which use a different URL pattern # e.g., /providers/blackforestlabs/v1/flux-2-pro instead of /openai/deployments/{model}/images/generations if AzureFoundryFluxImageGenerationConfig.is_flux2_model(model): @@ -1153,10 +1166,10 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): if api_base.endswith("/"): api_base = api_base.rstrip("/") api_version: str = azure_client_params.get("api_version", "") - # Use the deployment name (model) for URL construction, not the base_model from data img_gen_api_base = self.create_azure_base_url( azure_client_params=azure_client_params, model=model or data.get("model", ""), + base_model=data.get("model", ""), ) ## LOGGING @@ -1285,9 +1298,10 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): if aimg_generation is True: return self.aimage_generation(data=data, input=input, logging_obj=logging_obj, model_response=model_response, api_key=api_key, client=client, azure_client_params=azure_client_params, timeout=timeout, headers=headers, model=model) # type: ignore - # Use the deployment name (model) for URL construction, not the base_model from data img_gen_api_base = self.create_azure_base_url( - azure_client_params=azure_client_params, model=model + azure_client_params=azure_client_params, + model=model, + base_model=base_model, ) ## LOGGING @@ -1309,6 +1323,21 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): data=data, headers=headers, ) + provider_config = get_azure_image_generation_config( + data.get("model", "dall-e-2") + ) + if isinstance(provider_config, AzureFoundryMAIImageGenerationConfig): + return provider_config.transform_image_generation_response( + model=data.get("model", "dall-e-2"), + raw_response=httpx_response, + model_response=model_response or ImageResponse(), + logging_obj=logging_obj, + request_data=data, + optional_params=data, + litellm_params=data, + encoding=litellm.encoding, + ) + response = httpx_response.json() ## LOGGING diff --git a/litellm/llms/azure/image_generation/__init__.py b/litellm/llms/azure/image_generation/__init__.py index f60e446f0c4..64636bc689d 100644 --- a/litellm/llms/azure/image_generation/__init__.py +++ b/litellm/llms/azure/image_generation/__init__.py @@ -1,4 +1,5 @@ from litellm._logging import verbose_logger +from litellm.llms.azure_ai.image_generation import AzureFoundryMAIImageGenerationConfig from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, ) @@ -24,6 +25,8 @@ def get_azure_image_generation_config(model: str) -> BaseImageGenerationConfig: return AzureDallE2ImageGenerationConfig() elif "dalle3" in model: return AzureDallE3ImageGenerationConfig() + elif AzureFoundryMAIImageGenerationConfig.is_mai_model(model): + return AzureFoundryMAIImageGenerationConfig() else: verbose_logger.debug( f"Using AzureGPTImageGenerationConfig for model: {model}. This follows the gpt-image model format." diff --git a/litellm/llms/azure_ai/image_edit/__init__.py b/litellm/llms/azure_ai/image_edit/__init__.py index e3acd610446..42ece6d19ec 100644 --- a/litellm/llms/azure_ai/image_edit/__init__.py +++ b/litellm/llms/azure_ai/image_edit/__init__.py @@ -1,21 +1,33 @@ from litellm.llms.azure_ai.image_generation.flux_transformation import ( AzureFoundryFluxImageGenerationConfig, ) +from litellm.llms.azure_ai.image_generation.mai_transformation import ( + AzureFoundryMAIImageGenerationConfig, +) from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from .flux2_transformation import AzureFoundryFlux2ImageEditConfig +from .mai_transformation import AzureFoundryMAIImageEditConfig from .transformation import AzureFoundryFluxImageEditConfig -__all__ = ["AzureFoundryFluxImageEditConfig", "AzureFoundryFlux2ImageEditConfig"] +__all__ = [ + "AzureFoundryFluxImageEditConfig", + "AzureFoundryFlux2ImageEditConfig", + "AzureFoundryMAIImageEditConfig", +] def get_azure_ai_image_edit_config(model: str) -> BaseImageEditConfig: """ Get the appropriate image edit config for an Azure AI model. + - MAI models use /mai/v1/images/edits with multipart form data and size - FLUX 2 models use JSON with base64 image - FLUX 1 models use multipart/form-data """ + if AzureFoundryMAIImageGenerationConfig.is_mai_model(model): + return AzureFoundryMAIImageEditConfig() + # Check if it's a FLUX 2 model if AzureFoundryFluxImageGenerationConfig.is_flux2_model(model): return AzureFoundryFlux2ImageEditConfig() diff --git a/litellm/llms/azure_ai/image_edit/mai_transformation.py b/litellm/llms/azure_ai/image_edit/mai_transformation.py new file mode 100644 index 00000000000..75bfc913a8f --- /dev/null +++ b/litellm/llms/azure_ai/image_edit/mai_transformation.py @@ -0,0 +1,199 @@ +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, cast + +import httpx +from httpx._types import RequestFiles + +from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo +from litellm.llms.azure_ai.image_generation.mai_transformation import ( + AzureFoundryMAIImageGenerationConfig, +) +from litellm.llms.openai.common_utils import OpenAIError +from litellm.llms.openai.image_edit.transformation import OpenAIImageEditConfig +from litellm.secret_managers.main import get_secret_str +from litellm.types.images.main import ImageEditOptionalRequestParams +from litellm.types.llms.openai import FileTypes +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import ImageResponse +from litellm.utils import convert_to_model_response_object + +if TYPE_CHECKING: + from litellm.litellm_core_utils.logging import Logging as LiteLLMLoggingObj + + +class AzureFoundryMAIImageEditConfig(OpenAIImageEditConfig): + """Azure AI Foundry MAI image editing (e.g. MAI-Image-2.5).""" + + DEFAULT_SIZE = "1024x1024" + + def get_supported_openai_params(self, model: str) -> list: + return ["prompt", "image", "model", "n", "size"] + + def map_openai_params( + self, + image_edit_optional_params: ImageEditOptionalRequestParams, + model: str, + drop_params: bool, + ) -> Dict: + optional_params: Dict[str, Any] = {} + supported_params = self.get_supported_openai_params(model) + + for key, value in dict(image_edit_optional_params).items(): + if value is None or key in optional_params: + continue + + if key in supported_params: + if key == "size" and value: + size_param = cast(str, value) + self._validate_size_param(size_param) + optional_params[key] = size_param + else: + optional_params[key] = value + elif not drop_params: + raise ValueError( + f"Parameter {key} is not supported for model {model}. " + f"Supported parameters are {supported_params}. " + f"Set drop_params=True to drop unsupported parameters." + ) + + if "size" not in optional_params: + optional_params["size"] = self.DEFAULT_SIZE + + return optional_params + + def _validate_size_param(self, size: str) -> None: + known_sizes = { + "1024x1024", + "1792x1024", + "1024x1792", + "512x512", + "256x256", + } + + if size in known_sizes: + return + + if "x" in size: + try: + tuple(map(int, size.lower().split("x", 1))) + return + except ValueError: + raise ValueError( + f"Invalid size format: '{size}'. Expected format 'WIDTHxHEIGHT' (e.g., '1024x1024')." + ) + + raise ValueError( + f"Unsupported size value: '{size}'. " + f"Use a known size (e.g., '1024x1024') or a custom 'WIDTHxHEIGHT' string." + ) + + def validate_environment( + self, + headers: dict, + model: str, + api_key: Optional[str] = None, + litellm_params: Optional[dict] = None, + api_base: Optional[str] = None, + ) -> dict: + api_key = AzureFoundryModelInfo.get_api_key(api_key) + + if not api_key: + raise ValueError( + f"Azure AI API key is required for model {model}. " + "Set AZURE_AI_API_KEY environment variable or pass api_key parameter." + ) + + headers.update({"api-key": api_key}) + return headers + + def get_complete_url( + self, + model: str, + api_base: Optional[str], + litellm_params: dict, + ) -> str: + api_base = AzureFoundryModelInfo.get_api_base(api_base) + + if api_base is None: + raise ValueError( + "Azure AI API base is required. Set AZURE_AI_API_BASE environment variable or pass api_base parameter." + ) + + api_version = ( + litellm_params.get("api_version") + or get_secret_str("AZURE_AI_API_VERSION") + or "preview" + ) + + return AzureFoundryMAIImageGenerationConfig.get_mai_image_edit_url( + api_base=api_base, + api_version=api_version, + ) + + def transform_image_edit_request( + self, + model: str, + prompt: Optional[str], + image: Optional[FileTypes], + image_edit_optional_request_params: Dict, + litellm_params: GenericLiteLLMParams, + headers: dict, + ) -> Tuple[Dict, RequestFiles]: + request_params = { + "model": model, + **image_edit_optional_request_params, + } + if prompt is not None: + request_params["prompt"] = prompt + + data_without_files = { + key: value + for key, value in request_params.items() + if key not in ["image", "mask"] + } + files_list: List[Tuple[str, Any]] = [] + + if image is not None: + image_list = [image] if not isinstance(image, list) else image + for _image in image_list: + if _image is not None: + self._add_image_to_files( + files_list=files_list, + image=_image, + field_name="image", + ) + break + + return data_without_files, files_list + + def transform_image_edit_response( + self, + model: str, + raw_response: httpx.Response, + logging_obj: "LiteLLMLoggingObj", + ) -> ImageResponse: + try: + response = raw_response.json() + except Exception: + raise OpenAIError( + message=raw_response.text, status_code=raw_response.status_code + ) + + if "usage" in response: + response["usage"] = ( + AzureFoundryMAIImageGenerationConfig.normalize_mai_image_usage( + response.get("usage") + ) + ) + + logging_obj.post_call( + input="", + api_key="", + additional_args={"complete_input_dict": {}}, + original_response=response, + ) + + return convert_to_model_response_object( + response_object=response, + model_response_object=ImageResponse(), + response_type="image_generation", + ) diff --git a/litellm/llms/azure_ai/image_generation/__init__.py b/litellm/llms/azure_ai/image_generation/__init__.py index cebab3de16e..70821d5d764 100644 --- a/litellm/llms/azure_ai/image_generation/__init__.py +++ b/litellm/llms/azure_ai/image_generation/__init__.py @@ -7,12 +7,14 @@ from .dall_e_2_transformation import AzureFoundryDallE2ImageGenerationConfig from .dall_e_3_transformation import AzureFoundryDallE3ImageGenerationConfig from .flux_transformation import AzureFoundryFluxImageGenerationConfig from .gpt_transformation import AzureFoundryGPTImageGenerationConfig +from .mai_transformation import AzureFoundryMAIImageGenerationConfig __all__ = [ "AzureFoundryFluxImageGenerationConfig", "AzureFoundryGPTImageGenerationConfig", "AzureFoundryDallE2ImageGenerationConfig", "AzureFoundryDallE3ImageGenerationConfig", + "AzureFoundryMAIImageGenerationConfig", ] @@ -24,6 +26,8 @@ def get_azure_ai_image_generation_config(model: str) -> BaseImageGenerationConfi return AzureFoundryDallE2ImageGenerationConfig() elif "dalle3" in model: return AzureFoundryDallE3ImageGenerationConfig() + elif AzureFoundryMAIImageGenerationConfig.is_mai_model(model): + return AzureFoundryMAIImageGenerationConfig() elif "flux" in model: return AzureFoundryFluxImageGenerationConfig() else: diff --git a/litellm/llms/azure_ai/image_generation/cost_calculator.py b/litellm/llms/azure_ai/image_generation/cost_calculator.py index b67de9cb70d..f8c876bb5be 100644 --- a/litellm/llms/azure_ai/image_generation/cost_calculator.py +++ b/litellm/llms/azure_ai/image_generation/cost_calculator.py @@ -1,6 +1,9 @@ from typing import Any import litellm +from litellm.litellm_core_utils.llm_cost_calc.utils import ( + calculate_image_response_cost_from_usage, +) from litellm.types.utils import ImageResponse @@ -9,19 +12,28 @@ def cost_calculator( image_response: Any, ) -> float: """ - Recraft image generation cost calculator + Azure AI image generation cost calculator """ _model_info = litellm.get_model_info( model=model, custom_llm_provider=litellm.LlmProviders.AZURE_AI.value, ) - output_cost_per_image: float = _model_info.get("output_cost_per_image") or 0.0 - num_images: int = 0 + if isinstance(image_response, ImageResponse): + token_based_cost = calculate_image_response_cost_from_usage( + model=model, + image_response=image_response, + custom_llm_provider=litellm.LlmProviders.AZURE_AI.value, + ) + if token_based_cost is not None: + return token_based_cost + + output_cost_per_image: float = _model_info.get("output_cost_per_image") or 0.0 + num_images: int = 0 if image_response.data: num_images = len(image_response.data) return output_cost_per_image * num_images - else: - raise ValueError( - f"image_response must be of type ImageResponse got type={type(image_response)}" - ) + + raise ValueError( + f"image_response must be of type ImageResponse got type={type(image_response)}" + ) diff --git a/litellm/llms/azure_ai/image_generation/mai_transformation.py b/litellm/llms/azure_ai/image_generation/mai_transformation.py new file mode 100644 index 00000000000..071ca9d9895 --- /dev/null +++ b/litellm/llms/azure_ai/image_generation/mai_transformation.py @@ -0,0 +1,236 @@ +from typing import TYPE_CHECKING, Any, Dict, List, Optional + +import httpx + +from litellm.llms.base_llm.image_generation.transformation import ( + BaseImageGenerationConfig, +) +from litellm.llms.openai.common_utils import OpenAIError +from litellm.types.llms.openai import OpenAIImageGenerationOptionalParams +from litellm.types.utils import ImageResponse +from litellm.utils import convert_to_model_response_object + +if TYPE_CHECKING: + from litellm.litellm_core_utils.logging import Logging as LiteLLMLoggingObj + + +class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig): + """Azure AI Foundry MAI image generation (e.g. MAI-Image-2.5).""" + + DEFAULT_WIDTH = 1024 + DEFAULT_HEIGHT = 1024 + + @staticmethod + def get_mai_image_generation_url( + api_base: Optional[str], + api_version: Optional[str], + ) -> str: + if api_base is None: + raise ValueError("api_base is required for Azure AI MAI image generation") + + api_version = api_version or "preview" + path, separator, query = api_base.partition("?") + path = path.rstrip("/") + + if "/mai/" in path: + prefix, _, _ = path.partition("/images/") + path = f"{prefix}/images/generations" + else: + path = f"{path}/mai/v1/images/generations" + + if separator: + return f"{path}?{query}" + return f"{path}?api-version={api_version}" + + @staticmethod + def get_mai_image_edit_url( + api_base: Optional[str], + api_version: Optional[str], + ) -> str: + if api_base is None: + raise ValueError("api_base is required for Azure AI MAI image editing") + + api_version = api_version or "preview" + path, separator, query = api_base.partition("?") + path = path.rstrip("/") + + if "/mai/" in path: + prefix, _, _ = path.partition("/images/") + path = f"{prefix}/images/edits" + else: + path = f"{path}/mai/v1/images/edits" + + if separator: + return f"{path}?{query}" + return f"{path}?api-version={api_version}" + + @staticmethod + def is_mai_model(model: str) -> bool: + model_normalized = model.lower().replace("-", "").replace("_", "") + return "maiimage" in model_normalized + + @staticmethod + def normalize_mai_image_usage(usage: Optional[Dict[str, Any]]) -> Dict[str, Any]: + """Map Azure MAI usage fields to OpenAI ImageUsage schema.""" + if usage is None: + return { + "input_tokens": 0, + "input_tokens_details": {"image_tokens": 0, "text_tokens": 0}, + "output_tokens": 0, + "total_tokens": 0, + } + + normalized_usage = dict(usage) + input_tokens_details = normalized_usage.get("input_tokens_details") + if not isinstance(input_tokens_details, dict): + input_tokens_details = {} + + text_tokens = normalized_usage.get("num_input_text_tokens") + if text_tokens is None: + text_tokens = input_tokens_details.get("text_tokens") + if text_tokens is None: + text_tokens = normalized_usage.get("input_tokens", 0) or 0 + + image_tokens = normalized_usage.get("num_input_image_tokens") + if image_tokens is None: + image_tokens = input_tokens_details.get("image_tokens") + if image_tokens is None: + image_tokens = 0 + + output_tokens = normalized_usage.get("output_tokens") + if output_tokens is None: + output_tokens = normalized_usage.get("num_output_tokens") + if output_tokens is None: + output_tokens = normalized_usage.get("output_image_tokens") + if output_tokens is None: + output_tokens = 0 + + input_tokens = normalized_usage.get("input_tokens") + if input_tokens is None: + input_tokens = text_tokens + image_tokens + + total_tokens = normalized_usage.get("total_tokens") + if total_tokens is None: + total_tokens = input_tokens + output_tokens + + normalized_usage.update( + { + "input_tokens": input_tokens, + "input_tokens_details": { + "image_tokens": image_tokens, + "text_tokens": text_tokens, + }, + "output_tokens": output_tokens, + "total_tokens": total_tokens, + } + ) + return normalized_usage + + def get_supported_openai_params( + self, model: str + ) -> List[OpenAIImageGenerationOptionalParams]: + return ["n", "size"] + + def map_openai_params( + self, + non_default_params: dict, + optional_params: dict, + model: str, + drop_params: bool, + ) -> dict: + supported_params = self.get_supported_openai_params(model) + + for k, v in non_default_params.items(): + if k in optional_params: + continue + + if k in supported_params: + if k == "size" and v: + self._map_size_param(v, optional_params) + else: + optional_params[k] = v + elif k in ("width", "height"): + optional_params[k] = v + elif not drop_params: + raise ValueError( + f"Parameter {k} is not supported for model {model}. " + f"Supported parameters are {supported_params} and width/height. " + f"Set drop_params=True to drop unsupported parameters." + ) + + if "width" not in optional_params: + optional_params["width"] = self.DEFAULT_WIDTH + if "height" not in optional_params: + optional_params["height"] = self.DEFAULT_HEIGHT + + optional_params.pop("size", None) + return optional_params + + def _map_size_param(self, size: str, optional_params: dict) -> None: + size_mapping = { + "1024x1024": (1024, 1024), + "1792x1024": (1792, 1024), + "1024x1792": (1024, 1792), + "512x512": (512, 512), + "256x256": (256, 256), + } + + if size in size_mapping: + width, height = size_mapping[size] + optional_params["width"] = width + optional_params["height"] = height + elif "x" in size: + try: + width, height = map(int, size.lower().split("x")) + optional_params["width"] = width + optional_params["height"] = height + except ValueError: + raise ValueError( + f"Invalid size format: '{size}'. Expected format 'WIDTHxHEIGHT' (e.g., '1024x1024')." + ) + else: + raise ValueError( + f"Unsupported size value: '{size}'. " + f"Use a known size (e.g., '1024x1024') or a custom 'WIDTHxHEIGHT' string." + ) + + def transform_image_generation_response( + self, + model: str, + raw_response: httpx.Response, + model_response: ImageResponse, + logging_obj: "LiteLLMLoggingObj", + request_data: dict, + optional_params: dict, + litellm_params: dict, + encoding: Any, + api_key: Optional[str] = None, + json_mode: Optional[bool] = None, + ) -> ImageResponse: + try: + response = raw_response.json() + except Exception: + raise OpenAIError( + message=raw_response.text, status_code=raw_response.status_code + ) + + if "usage" in response: + response["usage"] = self.normalize_mai_image_usage(response.get("usage")) + + logging_obj.post_call( + input=request_data.get("prompt", ""), + api_key=api_key, + additional_args={"complete_input_dict": request_data}, + original_response=response, + ) + + image_response: ImageResponse = convert_to_model_response_object( + response_object=response, + model_response_object=model_response, + response_type="image_generation", + ) + + width = optional_params.get("width", self.DEFAULT_WIDTH) + height = optional_params.get("height", self.DEFAULT_HEIGHT) + image_response.size = f"{width}x{height}" # type: ignore[assignment] + return image_response diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index e765512175b..282a292ab17 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -6889,6 +6889,43 @@ "/v1/images/generations" ] }, + "azure_ai/MAI-Image-2.5": { + "input_cost_per_image_token": 8e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "azure_ai", + "mode": "image_generation", + "output_cost_per_image": 0.05, + "output_cost_per_image_token": 4.7e-05, + "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/new-mai-models-in-microsoft-foundry-across-text-image-voice-and-speech/4524632", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ] + }, + "azure_ai/MAI-Image-2.5-Flash": { + "input_cost_per_image_token": 1.75e-06, + "input_cost_per_token": 1.75e-06, + "litellm_provider": "azure_ai", + "mode": "image_generation", + "output_cost_per_image": 0.0338, + "output_cost_per_image_token": 3.3e-05, + "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/new-mai-models-in-microsoft-foundry-across-text-image-voice-and-speech/4524632", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ] + }, + "azure_ai/MAI-Image-2e": { + "input_cost_per_token": 5e-06, + "litellm_provider": "azure_ai", + "mode": "image_generation", + "output_cost_per_image": 0.02, + "output_cost_per_image_token": 1.95e-05, + "source": "https://aka.ms/mai-image-2e-foundryblog", + "supported_endpoints": [ + "/v1/images/generations" + ] + }, "azure_ai/Llama-3.2-11B-Vision-Instruct": { "input_cost_per_token": 3.7e-07, "litellm_provider": "azure_ai", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index c1c05b982f6..b0ffc66d03b 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -6889,6 +6889,43 @@ "/v1/images/generations" ] }, + "azure_ai/MAI-Image-2.5": { + "input_cost_per_image_token": 8e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "azure_ai", + "mode": "image_generation", + "output_cost_per_image": 0.05, + "output_cost_per_image_token": 4.7e-05, + "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/new-mai-models-in-microsoft-foundry-across-text-image-voice-and-speech/4524632", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ] + }, + "azure_ai/MAI-Image-2.5-Flash": { + "input_cost_per_image_token": 1.75e-06, + "input_cost_per_token": 1.75e-06, + "litellm_provider": "azure_ai", + "mode": "image_generation", + "output_cost_per_image": 0.0338, + "output_cost_per_image_token": 3.3e-05, + "source": "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/new-mai-models-in-microsoft-foundry-across-text-image-voice-and-speech/4524632", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ] + }, + "azure_ai/MAI-Image-2e": { + "input_cost_per_token": 5e-06, + "litellm_provider": "azure_ai", + "mode": "image_generation", + "output_cost_per_image": 0.02, + "output_cost_per_image_token": 1.95e-05, + "source": "https://aka.ms/mai-image-2e-foundryblog", + "supported_endpoints": [ + "/v1/images/generations" + ] + }, "azure_ai/Llama-3.2-11B-Vision-Instruct": { "input_cost_per_token": 3.7e-07, "litellm_provider": "azure_ai", diff --git a/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py b/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py index f49b3f09d6d..a211a69b9c7 100644 --- a/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py +++ b/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py @@ -57,6 +57,22 @@ def test_azure_providers_image_generation_json_body_keeps_model(): assert out == data +def test_azure_image_generation_mai_base_model_uses_mai_url(): + azure_chat = AzureChatCompletion() + url = azure_chat.create_azure_base_url( + azure_client_params={ + "azure_endpoint": "https://my-resource.services.ai.azure.com", + "api_version": "preview", + }, + model="image-deployment-alias", + base_model="MAI-Image-2.5", + ) + assert ( + url + == "https://my-resource.services.ai.azure.com/mai/v1/images/generations?api-version=preview" + ) + + def test_azure_image_generation_flattens_extra_body(): """ Test that Azure image generation correctly flattens extra_body parameters. diff --git a/tests/test_litellm/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py b/tests/test_litellm/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py new file mode 100644 index 00000000000..d5256be02d7 --- /dev/null +++ b/tests/test_litellm/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py @@ -0,0 +1,171 @@ +import io +import os +import sys +from unittest.mock import MagicMock + +import httpx +import pytest + +sys.path.insert(0, os.path.abspath("../../../../../..")) + +from litellm.llms.azure_ai.image_edit import ( + AzureFoundryMAIImageEditConfig, + get_azure_ai_image_edit_config, +) +from litellm.llms.azure_ai.image_generation.mai_transformation import ( + AzureFoundryMAIImageGenerationConfig, +) + + +class TestAzureMAIImageEdit: + def test_get_mai_image_edit_url(self): + url = AzureFoundryMAIImageGenerationConfig.get_mai_image_edit_url( + api_base="https://my-resource.services.ai.azure.com", + api_version="preview", + ) + assert ( + url + == "https://my-resource.services.ai.azure.com/mai/v1/images/edits?api-version=preview" + ) + + def test_get_mai_image_edit_url_rewrites_generation_url(self): + url = AzureFoundryMAIImageGenerationConfig.get_mai_image_edit_url( + api_base=( + "https://my-resource.services.ai.azure.com/mai/v1/images/generations" + "?api-version=preview" + ), + api_version="preview", + ) + assert ( + url + == "https://my-resource.services.ai.azure.com/mai/v1/images/edits?api-version=preview" + ) + + def test_get_mai_image_edit_url_appends_edits_to_mai_root(self): + url = AzureFoundryMAIImageGenerationConfig.get_mai_image_edit_url( + api_base="https://my-resource.services.ai.azure.com/mai/v1", + api_version="preview", + ) + assert ( + url + == "https://my-resource.services.ai.azure.com/mai/v1/images/edits?api-version=preview" + ) + + def test_get_azure_ai_image_edit_config_returns_mai(self): + config = get_azure_ai_image_edit_config("MAI-Image-2.5") + assert isinstance(config, AzureFoundryMAIImageEditConfig) + + def test_validate_environment_uses_api_key_header(self): + config = AzureFoundryMAIImageEditConfig() + headers: dict = {} + config.validate_environment(headers, "MAI-Image-2.5", api_key="test-key") + assert headers["api-key"] == "test-key" + assert "Api-Key" not in headers + + def test_get_complete_url(self): + config = AzureFoundryMAIImageEditConfig() + url = config.get_complete_url( + model="MAI-Image-2.5", + api_base="https://my-resource.services.ai.azure.com", + litellm_params={"api_version": "preview"}, + ) + assert "/mai/v1/images/edits" in url + assert "api-version=preview" in url + + def test_map_openai_params_keeps_size(self): + config = AzureFoundryMAIImageEditConfig() + optional_params = config.map_openai_params( + image_edit_optional_params={"size": "1792x1024", "n": 1}, + model="MAI-Image-2.5", + drop_params=True, + ) + assert optional_params["size"] == "1792x1024" + assert optional_params["n"] == 1 + assert "width" not in optional_params + assert "height" not in optional_params + + def test_map_openai_params_defaults_size(self): + config = AzureFoundryMAIImageEditConfig() + optional_params = config.map_openai_params( + image_edit_optional_params={}, + model="MAI-Image-2.5", + drop_params=True, + ) + assert optional_params["size"] == "1024x1024" + + def test_map_openai_params_unsupported_size_raises(self): + config = AzureFoundryMAIImageEditConfig() + with pytest.raises(ValueError, match="Unsupported size value: 'auto'"): + config.map_openai_params( + image_edit_optional_params={"size": "auto"}, + model="MAI-Image-2.5", + drop_params=True, + ) + + def test_map_openai_params_invalid_size_format_raises(self): + config = AzureFoundryMAIImageEditConfig() + with pytest.raises(ValueError, match="Invalid size format: '1024xabc'"): + config.map_openai_params( + image_edit_optional_params={"size": "1024xabc"}, + model="MAI-Image-2.5", + drop_params=True, + ) + + def test_transform_image_edit_request_uses_image_field(self): + config = AzureFoundryMAIImageEditConfig() + image_bytes = io.BytesIO(b"fake-image-bytes") + + data, files = config.transform_image_edit_request( + model="MAI-Image-2.5", + prompt="Turn this into a studio product shot", + image=image_bytes, + image_edit_optional_request_params={"size": "1024x1024", "n": 1}, + litellm_params={}, + headers={}, + ) + + assert data["model"] == "MAI-Image-2.5" + assert data["prompt"] == "Turn this into a studio product shot" + assert data["size"] == "1024x1024" + assert data["n"] == 1 + assert len(files) == 1 + assert files[0][0] == "image" + assert files[0][0] != "image[]" + + def test_normalize_mai_image_usage_maps_edit_response_fields(self): + usage = AzureFoundryMAIImageGenerationConfig.normalize_mai_image_usage( + { + "num_output_tokens": 1024, + "output_image_tokens": 1024, + } + ) + assert usage["output_tokens"] == 1024 + assert usage["input_tokens"] == 0 + assert usage["total_tokens"] == 1024 + assert usage["input_tokens_details"]["text_tokens"] == 0 + assert usage["input_tokens_details"]["image_tokens"] == 0 + + def test_transform_image_edit_response_parses_mai_usage(self): + config = AzureFoundryMAIImageEditConfig() + raw_response = MagicMock(spec=httpx.Response) + raw_response.status_code = 200 + raw_response.text = "" + raw_response.json.return_value = { + "created": 1780897477, + "data": [{"b64_json": "abc123"}], + "usage": { + "num_output_tokens": 1024, + "output_image_tokens": 1024, + }, + } + + logging_obj = MagicMock() + image_response = config.transform_image_edit_response( + model="MAI-Image-2.5", + raw_response=raw_response, + logging_obj=logging_obj, + ) + + assert image_response.data[0].b64_json == "abc123" + assert image_response.usage.output_tokens == 1024 + assert image_response.usage.total_tokens == 1024 diff --git a/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py b/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py new file mode 100644 index 00000000000..f7ad333293c --- /dev/null +++ b/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py @@ -0,0 +1,380 @@ +import os +import sys +from unittest.mock import MagicMock + +import httpx +import pytest + +sys.path.insert(0, os.path.abspath("../../../../../..")) + +import litellm +from litellm.llms.azure.azure import AzureChatCompletion +from litellm.llms.azure.image_generation import get_azure_image_generation_config +from litellm.llms.azure.image_generation.http_utils import ( + azure_deployment_image_generation_json_body, +) +from litellm.llms.azure_ai.image_generation import ( + AzureFoundryMAIImageGenerationConfig, + get_azure_ai_image_generation_config, +) +from litellm.llms.azure_ai.image_generation.cost_calculator import ( + cost_calculator as azure_ai_image_cost_calculator, +) +from litellm.types.utils import ( + ImageObject, + ImageResponse, + ImageUsage, + ImageUsageInputTokensDetails, +) +from litellm.utils import get_optional_params_image_gen + + +class TestAzureMAIImageGeneration: + def test_is_mai_model(self): + assert AzureFoundryMAIImageGenerationConfig.is_mai_model("MAI-Image-2.5") + assert AzureFoundryMAIImageGenerationConfig.is_mai_model( + "azure_ai/MAI-Image-2.5" + ) + assert AzureFoundryMAIImageGenerationConfig.is_mai_model("MAI-Image-2.5-Flash") + assert AzureFoundryMAIImageGenerationConfig.is_mai_model("MAI-Image-2e") + assert not AzureFoundryMAIImageGenerationConfig.is_mai_model("flux.2-pro") + assert not AzureFoundryMAIImageGenerationConfig.is_mai_model("MAI-DS-R1") + + def test_mai_flash_and_2e_model_pricing_in_cost_map(self): + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + flash_info = litellm.get_model_info( + model="azure_ai/MAI-Image-2.5-Flash", + custom_llm_provider="azure_ai", + ) + assert flash_info["input_cost_per_token"] == 1.75e-06 + assert flash_info["input_cost_per_image_token"] == 1.75e-06 + assert flash_info["output_cost_per_image_token"] == 3.3e-05 + + image_2e_info = litellm.get_model_info( + model="azure_ai/MAI-Image-2e", + custom_llm_provider="azure_ai", + ) + assert image_2e_info["input_cost_per_token"] == 5e-06 + assert image_2e_info["output_cost_per_image_token"] == 1.95e-05 + + def test_get_mai_image_generation_url(self): + url = AzureFoundryMAIImageGenerationConfig.get_mai_image_generation_url( + api_base="https://my-resource.services.ai.azure.com", + api_version="preview", + ) + assert ( + url + == "https://my-resource.services.ai.azure.com/mai/v1/images/generations?api-version=preview" + ) + + def test_get_mai_image_generation_url_preserves_full_path(self): + api = ( + "https://my-resource.services.ai.azure.com/mai/v1/images/generations" + "?api-version=preview" + ) + url = AzureFoundryMAIImageGenerationConfig.get_mai_image_generation_url( + api_base=api, + api_version="preview", + ) + assert url == api + + def test_get_mai_image_generation_url_appends_generations_to_mai_root(self): + url = AzureFoundryMAIImageGenerationConfig.get_mai_image_generation_url( + api_base="https://my-resource.services.ai.azure.com/mai/v1", + api_version="preview", + ) + assert ( + url + == "https://my-resource.services.ai.azure.com/mai/v1/images/generations?api-version=preview" + ) + + def test_get_azure_ai_image_generation_config_returns_mai(self): + config = get_azure_ai_image_generation_config("MAI-Image-2.5") + assert isinstance(config, AzureFoundryMAIImageGenerationConfig) + + def test_azure_image_generation_config_returns_mai(self): + config = get_azure_image_generation_config("MAI-Image-2.5") + assert isinstance(config, AzureFoundryMAIImageGenerationConfig) + + def test_map_openai_params_size_to_width_height(self): + config = AzureFoundryMAIImageGenerationConfig() + optional_params = config.map_openai_params( + non_default_params={"size": "1024x1024", "n": 1}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=True, + ) + assert optional_params["width"] == 1024 + assert optional_params["height"] == 1024 + assert optional_params["n"] == 1 + assert "size" not in optional_params + + def test_map_openai_params_defaults(self): + config = AzureFoundryMAIImageGenerationConfig() + optional_params = config.map_openai_params( + non_default_params={}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=True, + ) + assert optional_params["width"] == 1024 + assert optional_params["height"] == 1024 + + def test_get_optional_params_image_gen_mai(self): + config = AzureFoundryMAIImageGenerationConfig() + optional_params = get_optional_params_image_gen( + model="MAI-Image-2.5", + size="1792x1024", + n=1, + custom_llm_provider="azure_ai", + provider_config=config, + drop_params=True, + ) + assert optional_params["width"] == 1792 + assert optional_params["height"] == 1024 + assert "size" not in optional_params + + def test_azure_create_azure_base_url_mai(self): + azure_chat = AzureChatCompletion() + url = azure_chat.create_azure_base_url( + azure_client_params={ + "azure_endpoint": "https://my-resource.services.ai.azure.com", + "api_version": "preview", + }, + model="MAI-Image-2.5", + ) + assert "/mai/v1/images/generations" in url + assert "api-version=preview" in url + + def test_mai_json_body_keeps_model(self): + api = ( + "https://my-resource.services.ai.azure.com/mai/v1/images/generations" + "?api-version=preview" + ) + data = { + "model": "MAI-Image-2.5", + "prompt": "A photograph of a red fox", + "width": 1024, + "height": 1024, + "n": 1, + } + out = azure_deployment_image_generation_json_body(api, data) + assert out == data + + def test_map_openai_params_custom_size(self): + config = AzureFoundryMAIImageGenerationConfig() + optional_params = config.map_openai_params( + non_default_params={"size": "768x768"}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=True, + ) + assert optional_params["width"] == 768 + assert optional_params["height"] == 768 + + def test_map_openai_params_width_only_gets_height_default(self): + config = AzureFoundryMAIImageGenerationConfig() + optional_params = config.map_openai_params( + non_default_params={"width": 1792}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=True, + ) + assert optional_params["width"] == 1792 + assert optional_params["height"] == config.DEFAULT_HEIGHT + + def test_map_openai_params_height_only_gets_width_default(self): + config = AzureFoundryMAIImageGenerationConfig() + optional_params = config.map_openai_params( + non_default_params={"height": 1792}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=True, + ) + assert optional_params["width"] == config.DEFAULT_WIDTH + assert optional_params["height"] == 1792 + + def test_map_openai_params_unsupported_size_raises(self): + config = AzureFoundryMAIImageGenerationConfig() + with pytest.raises(ValueError, match="Unsupported size value: 'auto'"): + config.map_openai_params( + non_default_params={"size": "auto"}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=True, + ) + + def test_map_openai_params_invalid_custom_size_raises(self): + config = AzureFoundryMAIImageGenerationConfig() + with pytest.raises(ValueError, match="Invalid size format: '1024xabc'"): + config.map_openai_params( + non_default_params={"size": "1024xabc"}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=True, + ) + + def test_map_openai_params_unsupported_param_raises(self): + config = AzureFoundryMAIImageGenerationConfig() + with pytest.raises(ValueError, match="Parameter quality is not supported"): + config.map_openai_params( + non_default_params={"quality": "hd"}, + optional_params={}, + model="MAI-Image-2.5", + drop_params=False, + ) + + def test_transform_image_generation_response_normalizes_mai_usage(self): + config = AzureFoundryMAIImageGenerationConfig() + raw_response = MagicMock(spec=httpx.Response) + raw_response.json.return_value = { + "created": 1780897477, + "data": [{"b64_json": "abc123"}], + "usage": { + "num_output_tokens": 1024, + "num_input_text_tokens": 22, + "output_image_tokens": 1024, + }, + } + + logging_obj = MagicMock() + image_response = config.transform_image_generation_response( + model="MAI-Image-2.5", + raw_response=raw_response, + model_response=ImageResponse(), + logging_obj=logging_obj, + request_data={"prompt": "A red fox"}, + optional_params={"width": 1024, "height": 1024}, + litellm_params={}, + encoding=None, + ) + + assert image_response.data[0].b64_json == "abc123" + assert image_response.usage.output_tokens == 1024 + assert image_response.usage.input_tokens == 22 + assert image_response.usage.total_tokens == 1046 + + def test_transform_image_generation_response_non_json_raises_openai_error(self): + from litellm.llms.openai.common_utils import OpenAIError + + config = AzureFoundryMAIImageGenerationConfig() + raw_response = MagicMock(spec=httpx.Response) + raw_response.json.side_effect = ValueError("not json") + raw_response.text = "upstream gateway error" + raw_response.status_code = 502 + + with pytest.raises(OpenAIError) as exc_info: + config.transform_image_generation_response( + model="MAI-Image-2.5", + raw_response=raw_response, + model_response=ImageResponse(), + logging_obj=MagicMock(), + request_data={"prompt": "A red fox"}, + optional_params={"width": 1024, "height": 1024}, + litellm_params={}, + encoding=None, + ) + + assert exc_info.value.status_code == 502 + assert exc_info.value.message == "upstream gateway error" + + def test_normalize_mai_usage_preserves_zero_output_tokens(self): + config = AzureFoundryMAIImageGenerationConfig() + normalized = config.normalize_mai_image_usage( + { + "num_output_tokens": 0, + "output_image_tokens": 1024, + "num_input_text_tokens": 22, + } + ) + assert normalized["output_tokens"] == 0 + assert normalized["input_tokens"] == 22 + assert normalized["total_tokens"] == 22 + + def test_azure_sync_image_generation_uses_mai_response_transform(self): + raw_response = MagicMock(spec=httpx.Response) + raw_response.json.return_value = { + "created": 1780897477, + "data": [{"b64_json": "abc123"}], + "usage": { + "num_output_tokens": 1024, + "num_input_text_tokens": 22, + }, + } + + class MAIImageGenerationAzureChatCompletion(AzureChatCompletion): + def make_sync_azure_httpx_request(self, **kwargs): + return raw_response + + logging_obj = MagicMock() + image_response = MAIImageGenerationAzureChatCompletion().image_generation( + prompt="A red fox", + timeout=60.0, + optional_params={"width": 1792, "height": 1024}, + logging_obj=logging_obj, + headers={}, + model="MAI-Image-2.5", + api_key="test-key", + api_base="https://my-resource.services.ai.azure.com", + api_version="preview", + litellm_params={}, + ) + + assert image_response.data[0].b64_json == "abc123" + assert image_response.usage.output_tokens == 1024 + assert image_response.usage.input_tokens == 22 + assert image_response.usage.total_tokens == 1046 + assert image_response.size == "1792x1024" + + def test_mai_image_cost_calculator_token_based(self): + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + model = "azure_ai/MAI-Image-2.5" + model_info = litellm.get_model_info(model=model, custom_llm_provider="azure_ai") + input_text_tokens = 100 + output_image_tokens = 1024 + + image_response = ImageResponse( + data=[ImageObject(b64_json="img1")], + usage=ImageUsage( + input_tokens=input_text_tokens, + input_tokens_details=ImageUsageInputTokensDetails( + text_tokens=input_text_tokens, + image_tokens=0, + ), + output_tokens=output_image_tokens, + total_tokens=input_text_tokens + output_image_tokens, + ), + ) + + cost = azure_ai_image_cost_calculator( + model=model, + image_response=image_response, + ) + + expected_cost = ( + input_text_tokens * model_info["input_cost_per_token"] + + output_image_tokens * model_info["output_cost_per_image_token"] + ) + assert round(cost, 10) == round(expected_cost, 10) + + def test_mai_image_cost_calculator_falls_back_to_flat_image_pricing(self): + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + model = "azure_ai/MAI-Image-2.5" + model_info = litellm.get_model_info(model=model, custom_llm_provider="azure_ai") + image_response = ImageResponse( + data=[ImageObject(b64_json="img1"), ImageObject(b64_json="img2")] + ) + + cost = azure_ai_image_cost_calculator( + model=model, + image_response=image_response, + ) + + assert ( + cost == len(image_response.data or []) * model_info["output_cost_per_image"] + ) + assert cost > 0 From 51ba6e39cd23576b9c2110361f1045782762f3e4 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Mon, 8 Jun 2026 19:58:51 -0700 Subject: [PATCH 032/185] fix(mcp): load MCP tool configuration tools via the OBO/passthrough-aware GET path (#29960) * fix(ui): load MCP tool configuration tools via the OBO/passthrough-aware GET path * fix(mcp): admin-only include_disabled_tools so the settings UI shows toggled-off tools * fix(ui): repopulate MCP server edit form when server data loads after mount (OAuth return) * fix(ui): persist MCP OAuth token on save and return to the Settings tab after authorize * fix(ui): scope MCP OAuth callback to the initiating form so create and edit flows don't cross-talk * fix(ui): derive OAuth-return Settings tab via lazy state init instead of setState-in-effect * Fix MCP OAuth edit token handling --------- Co-authored-by: Cursor Agent --- .../mcp_server/rest_endpoints.py | 31 +- .../mcp_server/test_rest_endpoints.py | 125 +++++++ ui/litellm-dashboard/eslint-suppressions.json | 10 - .../mcp_tools/create_mcp_server.tsx | 2 +- .../mcp_tools/mcp_server_edit.test.tsx | 318 +++++++++++++++++- .../components/mcp_tools/mcp_server_edit.tsx | 123 ++++++- .../components/mcp_tools/mcp_server_view.tsx | 29 +- .../mcp_tools/mcp_tool_configuration.tsx | 23 +- .../src/components/mcp_tools/mcp_tools.tsx | 13 +- .../src/components/networking.tsx | 15 +- .../src/hooks/useMcpOAuthFlow.tsx | 18 + .../src/utils/mcpHeaderUtils.ts | 15 + 12 files changed, 658 insertions(+), 64 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 725f7a335bc..2149f079a3d 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -386,8 +386,15 @@ if MCP_AVAILABLE: raw_headers: Optional[Dict[str, str]] = None, user_api_key_auth: Optional[UserAPIKeyAuth] = None, extra_headers: Optional[Dict[str, str]] = None, + apply_tool_filters: bool = True, ): - """Helper function to get tools for a single server.""" + """Helper function to get tools for a single server. + + When ``apply_tool_filters`` is False the raw server catalog is returned + without the allowed_tools/disallowed_tools gate or the per-key tool + permissions. This is the admin-only configuration view; every runtime + path keeps the default True so callable tools stay filtered. + """ tools = await global_mcp_server_manager._get_tools_from_server( server=server, mcp_auth_header=server_auth_header, @@ -397,6 +404,9 @@ if MCP_AVAILABLE: user_api_key_auth=user_api_key_auth, ) + if not apply_tool_filters: + return _create_tool_response_objects(tools, server.mcp_info) + # Always apply allowed_tools/disallowed_tools so the blacklist is # enforced even when no allowlist is set (matches the SSE/HTTP path). tools = filter_tools_by_allowed_tools(tools, server) @@ -463,6 +473,7 @@ if MCP_AVAILABLE: mcp_auth_header: Optional[str], raw_headers_from_request: dict, user_api_key_dict: UserAPIKeyAuth, + apply_tool_filters: bool = True, ) -> dict: """Handle tool listing for a single server_id request.""" # Resolve a server name to its UUID if needed @@ -527,6 +538,7 @@ if MCP_AVAILABLE: raw_headers_from_request, user_api_key_dict, extra_headers=user_oauth_extra_headers, + apply_tool_filters=apply_tool_filters, ) except MCPUpstreamAuthError: # Surface the upstream 401/403 to the caller so it can emit the @@ -552,6 +564,14 @@ if MCP_AVAILABLE: server_id: Optional[str] = Query( None, description="The server id to list tools for" ), + include_disabled_tools: bool = Query( + False, + description=( + "Admin only. Return the full server tool catalog without the " + "allowed_tools filter or per-key tool permissions, so the MCP " + "settings UI can configure the allowlist. Ignored for non-admins." + ), + ), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ) -> dict: """ @@ -579,6 +599,13 @@ if MCP_AVAILABLE: ) try: + # The full catalog (allowlist filter skipped) is admin-only so the + # REST endpoint can't be used to enumerate deliberately-disabled tools. + apply_tool_filters = not ( + include_disabled_tools + and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN + ) + # Extract auth headers from request headers = request.headers raw_headers_from_request = dict(headers) @@ -620,6 +647,7 @@ if MCP_AVAILABLE: mcp_auth_header=mcp_auth_header, raw_headers_from_request=raw_headers_from_request, user_api_key_dict=user_api_key_dict, + apply_tool_filters=apply_tool_filters, ) else: if not allowed_server_ids: @@ -677,6 +705,7 @@ if MCP_AVAILABLE: raw_headers_from_request, user_api_key_dict, extra_headers=user_oauth_extra_headers, + apply_tool_filters=apply_tool_filters, ) list_tools_result.extend(tools_result) except Exception as e: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index caff9ea2d28..47c9396f121 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -501,6 +501,7 @@ class TestListToolsRestAPI: raw_headers=None, user_api_key_auth=None, extra_headers=None, + apply_tool_filters=True, ): captured["called"] = True captured["server"] = server @@ -545,6 +546,78 @@ class TestListToolsRestAPI: assert result["error"] is None assert result["message"] == "Successfully retrieved tools" + async def test_include_disabled_tools_is_admin_only(self, monkeypatch): + """include_disabled_tools skips the allowlist filter only for PROXY_ADMIN; + a non-admin passing it stays filtered so the REST endpoint can't be used + to enumerate deliberately-disabled tools.""" + from litellm.proxy._types import LitellmUserRoles + + async def fake_contexts(user_api_key_auth): + return [user_api_key_auth] + + async def fake_get_allowed_mcp_servers(*args, **kwargs): + return ["server-1"] + + class StubServer: + alias = "server-1" + server_name = "server-1" + name = "stub" + allowed_tools = ["tool1"] + mcp_info = {"server_name": "stub"} + available_on_public_internet = True + + stub_server = StubServer() + captured = {} + + async def fake_get_tools( + server, server_auth_header, *args, apply_tool_filters=True, **kwargs + ): + captured["apply_tool_filters"] = apply_tool_filters + return ["tool-1"] + + monkeypatch.setattr( + rest_endpoints, + "build_effective_auth_contexts", + fake_contexts, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "get_allowed_mcp_servers", + fake_get_allowed_mcp_servers, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "get_mcp_server_by_id", + lambda server_id: stub_server if server_id == "server-1" else None, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints, + "_get_tools_for_single_server", + fake_get_tools, + raising=False, + ) + + request = _build_request(path="/mcp-rest/tools/list", method="GET") + + await rest_endpoints.list_tool_rest_api( + request, + server_id="server-1", + include_disabled_tools=True, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + assert captured["apply_tool_filters"] is False + + await rest_endpoints.list_tool_rest_api( + request, + server_id="server-1", + include_disabled_tools=True, + user_api_key_dict=UserAPIKeyAuth(), + ) + assert captured["apply_tool_filters"] is True + @pytest.mark.parametrize("upstream_status", [401, 403]) async def test_upstream_auth_failure_surfaces_status_and_challenge( self, monkeypatch, upstream_status @@ -649,6 +722,7 @@ class TestListToolsRestAPI: raw_headers=None, user_api_key_auth=None, extra_headers=None, + apply_tool_filters=True, ): captured["called"] = True captured["server_arg"] = server @@ -792,6 +866,7 @@ class TestListToolsRestAPI: raw_headers=None, user_api_key_auth=None, extra_headers=None, + apply_tool_filters=True, ): captured["server"] = server captured["auth_header"] = server_auth_header @@ -1284,6 +1359,56 @@ class TestGetToolsForSingleServer: assert "tool1" not in tool_names assert "tool4" not in tool_names + async def test_apply_tool_filters_false_returns_full_catalog(self, monkeypatch): + """apply_tool_filters=False returns the raw catalog without the server + allowed_tools gate, so the config UI can render disabled tools as off.""" + from litellm.proxy._experimental.mcp_server.server import MCPServer + from litellm.types.mcp import MCPTransport + + class MockTool: + def __init__(self, name): + self.name = name + self.description = name + self.inputSchema = {} + + mock_tools = [MockTool("tool1"), MockTool("tool2"), MockTool("tool3")] + + async def fake_get_tools_from_server(**kwargs): + return mock_tools + + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "_get_tools_from_server", + fake_get_tools_from_server, + raising=False, + ) + + # Server enforces an allowlist of just tool1. + server = MCPServer( + server_id="test-server-id", + name="test-server", + transport=MCPTransport.sse, + allowed_tools=["tool1"], + ) + user_api_key_dict = UserAPIKeyAuth(api_key="test-key", object_permission=None) + + # Runtime default: only the allowed tool comes back. + filtered = await rest_endpoints._get_tools_for_single_server( + server=server, + server_auth_header=None, + user_api_key_auth=user_api_key_dict, + ) + assert [t.name for t in filtered] == ["tool1"] + + # Config view: full catalog, including the disabled tools. + full = await rest_endpoints._get_tools_for_single_server( + server=server, + server_auth_header=None, + user_api_key_auth=user_api_key_dict, + apply_tool_filters=False, + ) + assert {t.name for t in full} == {"tool1", "tool2", "tool3"} + class TestStdioCommandAllowlist: """Tests for MCP stdio command allowlist validation.""" diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index bbab73b07f1..233741652a9 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -1351,11 +1351,6 @@ "count": 1 } }, - "src/components/mcp_tools/mcp_server_edit.test.tsx": { - "unused-imports/no-unused-imports": { - "count": 1 - } - }, "src/components/mcp_tools/mcp_server_edit.tsx": { "no-restricted-imports": { "count": 1 @@ -1517,11 +1512,6 @@ "count": 1 } }, - "src/components/organisms/RegenerateKeyModal.tsx": { - "react-hooks/set-state-in-effect": { - "count": 1 - } - }, "src/components/organisms/create_key_button.test.tsx": { "@typescript-eslint/no-require-imports": { "count": 2 diff --git a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx index 6a8dc353f21..28c38c459aa 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx @@ -188,6 +188,7 @@ const CreateMCPServer: React.FC = ({ } }, onBeforeRedirect: persistCreateUiState, + flowSource: "create", }); React.useEffect(() => { @@ -1088,7 +1089,6 @@ const CreateMCPServer: React.FC = ({
({ updateMCPServer: vi.fn(), listMCPTools: vi.fn().mockResolvedValue({ tools: [], error: null }), + storeMCPOAuthUserCredential: vi.fn().mockResolvedValue({}), })); vi.mock("../molecules/notifications_manager", () => ({ @@ -17,12 +18,13 @@ vi.mock("../molecules/notifications_manager", () => ({ }, })); +const mockOauth: { tokenResponse: any } = { tokenResponse: null }; vi.mock("@/hooks/useMcpOAuthFlow", () => ({ useMcpOAuthFlow: () => ({ startOAuthFlow: vi.fn(), status: "idle", error: null, - tokenResponse: null, + tokenResponse: mockOauth.tokenResponse, }), })); @@ -37,12 +39,19 @@ vi.mock("./MCPPermissionManagement", () => ({ vi.mock("./mcp_tool_configuration", () => ({ default: ({ existingAllowedTools, + externalTools, + externalError, onAllowedToolsChange, onToolAllowlistInteraction, onToolNameToDisplayNameChange, onToolNameToDescriptionChange, }: any) => ( -
+
- ); - }, - }, - ]; - - return allColumns; -}; diff --git a/ui/litellm-dashboard/src/components/Projects/types.ts b/ui/litellm-dashboard/src/components/Projects/types.ts deleted file mode 100644 index 51429902dff..00000000000 --- a/ui/litellm-dashboard/src/components/Projects/types.ts +++ /dev/null @@ -1,14 +0,0 @@ -export interface Project { - id: string; - name: string; - description: string; - teamId: string; - teamAlias: string; - models: string[]; - status: "active" | "blocked"; - spend: number; - createdAt: string; - createdBy: string; - updatedAt: string; - updatedBy: string; -} diff --git a/ui/litellm-dashboard/src/components/agents/agent_table.tsx b/ui/litellm-dashboard/src/components/agents/agent_table.tsx deleted file mode 100644 index cc170ca4f26..00000000000 --- a/ui/litellm-dashboard/src/components/agents/agent_table.tsx +++ /dev/null @@ -1,211 +0,0 @@ -import React, { useState } from "react"; -import { Table, TableBody, TableCell, TableHead, TableHeaderCell, TableRow, Button } from "@tremor/react"; -import { SwitchVerticalIcon, ChevronUpIcon, ChevronDownIcon, TrashIcon } from "@heroicons/react/outline"; -import { Tooltip } from "antd"; -import { CopyOutlined } from "@ant-design/icons"; -import { Agent } from "./types"; -import { - ColumnDef, - flexRender, - getCoreRowModel, - getSortedRowModel, - SortingState, - useReactTable, -} from "@tanstack/react-table"; - -interface AgentTableProps { - agentsList: Agent[]; - isLoading: boolean; - onDeleteClick: (agentId: string, agentName: string) => void; - accessToken: string | null; - onAgentUpdated: () => void; - isAdmin: boolean; - onAgentClick: (agentId: string) => void; -} - -const AgentTable: React.FC = ({ - agentsList, - isLoading, - onDeleteClick, - accessToken, - onAgentUpdated, - isAdmin, - onAgentClick, -}) => { - const [sorting, setSorting] = useState([{ id: "created_at", desc: true }]); - - const formatDate = (dateString?: string) => { - if (!dateString) return "-"; - const date = new Date(dateString); - return date.toLocaleString(); - }; - - const copyToClipboard = (text: string) => { - navigator.clipboard.writeText(text); - }; - - const columns: ColumnDef[] = [ - { - header: "Agent Name", - accessorKey: "agent_name", - cell: ({ row }) => { - const agent = row.original; - const name = agent.agent_name || ""; - return ( -
- - - - - { - e.stopPropagation(); - copyToClipboard(agent.agent_id); - }} - className="cursor-pointer text-gray-500 hover:text-blue-500 text-xs" - /> - -
- ); - }, - }, - { - header: "Description", - accessorKey: "agent_card_params.description", - cell: ({ row }) => { - const description = row.original.agent_card_params?.description || "No description"; - return {description}; - }, - }, - { - header: "Created At", - accessorKey: "created_at", - cell: ({ row }) => { - const agent = row.original; - return ( - - {formatDate(agent.created_at)} - - ); - }, - }, - ...(isAdmin - ? [ - { - header: "Actions", - id: "actions", - enableSorting: false, - cell: ({ row }: any) => { - const agent = row.original; - - return ( -
- -
- ); - }, - }, - ] - : []), - ]; - - const table = useReactTable({ - data: agentsList, - columns, - state: { - sorting, - }, - onSortingChange: setSorting, - getCoreRowModel: getCoreRowModel(), - getSortedRowModel: getSortedRowModel(), - enableSorting: true, - }); - - return ( -
-
- - - {table.getHeaderGroups().map((headerGroup) => ( - - {headerGroup.headers.map((header) => ( - -
-
- {header.isPlaceholder ? null : flexRender(header.column.columnDef.header, header.getContext())} -
-
- {header.column.getIsSorted() ? ( - { - asc: , - desc: , - }[header.column.getIsSorted() as string] - ) : ( - - )} -
-
- - ))} - - ))} - - - {isLoading ? ( - - -
-

Loading...

-
- - - ) : agentsList && agentsList.length > 0 ? ( - table.getRowModel().rows.map((row) => ( - - {row.getVisibleCells().map((cell) => ( - - {flexRender(cell.column.columnDef.cell, cell.getContext())} - - ))} - - )) - ) : ( - - -
-

No agents found. Create one to get started.

-
- - - )} - -
-
-
- ); -}; - -export default AgentTable; diff --git a/ui/litellm-dashboard/src/components/claude_code_plugins/plugin_info.tsx b/ui/litellm-dashboard/src/components/claude_code_plugins/plugin_info.tsx deleted file mode 100644 index 55347025201..00000000000 --- a/ui/litellm-dashboard/src/components/claude_code_plugins/plugin_info.tsx +++ /dev/null @@ -1,313 +0,0 @@ -import { CopyOutlined } from "@ant-design/icons"; -import { ArrowLeftIcon, ExternalLinkIcon } from "@heroicons/react/outline"; -import { Badge, Button, Card, Grid, Text, Title } from "@tremor/react"; -import { Spin, Switch, Tooltip } from "antd"; -import React, { useEffect, useState } from "react"; -import NotificationsManager from "../molecules/notifications_manager"; -import { disableClaudeCodePlugin, enableClaudeCodePlugin, getClaudeCodePluginDetails } from "../networking"; -import { - formatDateString, - formatInstallCommand, - getCategoryBadgeColor, - getSourceDisplayText, - getSourceLink, -} from "./helpers"; -import { Plugin } from "./types"; - -interface PluginInfoViewProps { - pluginId: string; - onClose: () => void; - accessToken: string | null; - isAdmin: boolean; - onPluginUpdated: () => void; -} - -const PluginInfoView: React.FC = ({ - pluginId, - onClose, - accessToken, - isAdmin, - onPluginUpdated, -}) => { - const [plugin, setPlugin] = useState(null); - const [isLoading, setIsLoading] = useState(true); - const [isToggling, setIsToggling] = useState(false); - - useEffect(() => { - fetchPluginInfo(); - }, [pluginId, accessToken]); - - const fetchPluginInfo = async () => { - if (!accessToken) return; - - setIsLoading(true); - try { - // The backend expects plugin name, not ID - // We'll need to find the plugin by ID from the list - // For now, assume pluginId is actually the plugin name - const data = await getClaudeCodePluginDetails(accessToken, pluginId as string); - setPlugin(data.plugin); - } catch (error) { - console.error("Error fetching plugin info:", error); - NotificationsManager.error("Failed to load plugin information"); - } finally { - setIsLoading(false); - } - }; - - const handleToggleEnabled = async () => { - if (!accessToken || !plugin) return; - - setIsToggling(true); - try { - if (plugin.enabled) { - await disableClaudeCodePlugin(accessToken, plugin.name); - NotificationsManager.success(`Plugin "${plugin.name}" disabled`); - } else { - await enableClaudeCodePlugin(accessToken, plugin.name); - NotificationsManager.success(`Plugin "${plugin.name}" enabled`); - } - onPluginUpdated(); - fetchPluginInfo(); - } catch (error) { - NotificationsManager.error("Failed to toggle plugin status"); - } finally { - setIsToggling(false); - } - }; - - const copyToClipboard = (text: string) => { - navigator.clipboard.writeText(text); - NotificationsManager.success("Copied to clipboard!"); - }; - - if (isLoading) { - return ( -
- -
- ); - } - - if (!plugin) { - return ( -
-

Plugin not found

- -
- ); - } - - const installCommand = formatInstallCommand(plugin); - const sourceLink = getSourceLink(plugin.source); - const categoryBadgeColor = getCategoryBadgeColor(plugin.category); - - return ( -
- {/* Header with Back Button */} -
- -

{plugin.name}

- {plugin.version && ( - - v{plugin.version} - - )} - {plugin.category && ( - - {plugin.category} - - )} - - {plugin.enabled ? "Enabled" : "Disabled"} - -
- - {/* Install Command */} - -
-
- Install Command -
{installCommand}
-
- - - -
- - - {/* Plugin Details */} - - Plugin Details - - {/* Plugin ID */} -
- Plugin ID -
- {plugin.id} - copyToClipboard(plugin.id)} - /> -
-
- - {/* Name */} -
- Name - {plugin.name} -
- - {/* Version */} -
- Version - {plugin.version || "N/A"} -
- - {/* Source */} - - - {/* Category */} -
- Category -
- {plugin.category ? ( - - {plugin.category} - - ) : ( - Uncategorized - )} -
-
- - {/* Enabled Status */} - {isAdmin && ( -
- Status -
- - - {plugin.enabled - ? "Plugin is enabled and visible in marketplace" - : "Plugin is disabled and hidden from marketplace"} - -
-
- )} - - - - {/* Description */} - {plugin.description && ( - - Description - {plugin.description} - - )} - - {/* Keywords */} - {plugin.keywords && plugin.keywords.length > 0 && ( - - Keywords -
- {plugin.keywords.map((keyword, index) => ( - - {keyword} - - ))} -
- - )} - - {/* Author Information */} - {plugin.author && ( - - Author Information - - {plugin.author.name && ( -
- Name - {plugin.author.name} -
- )} - {plugin.author.email && ( - - )} - - - )} - - {/* Additional Links */} - {plugin.homepage && ( - - Homepage - - {plugin.homepage} - - - - )} - - {/* Timestamps */} - - Metadata - -
- Created At - {formatDateString(plugin.created_at)} -
-
- Updated At - {formatDateString(plugin.updated_at)} -
- {plugin.created_by && ( -
- Created By - {plugin.created_by} -
- )} - - -
- ); -}; - -export default PluginInfoView; diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_columns.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_columns.tsx deleted file mode 100644 index bc13ed72a8e..00000000000 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_columns.tsx +++ /dev/null @@ -1,321 +0,0 @@ -import { useState } from "react"; -import { ColumnDef } from "@tanstack/react-table"; -import { MCPServer } from "./types"; -import { Icon } from "@tremor/react"; -import { PencilAltIcon, TrashIcon } from "@heroicons/react/outline"; -import { getMaskedAndFullUrl } from "./utils"; -import { Tooltip } from "antd"; -import { CheckOutlined } from "@ant-design/icons"; - -const HealthStatusBadge: React.FC<{ - server: MCPServer; - isLoadingHealth?: boolean; - isRechecking?: boolean; - onRecheck?: (serverId: string) => void; -}> = ({ server, isLoadingHealth, isRechecking, onRecheck }) => { - const [isHovered, setIsHovered] = useState(false); - const status = server.status || "unknown"; - const lastCheck = server.last_health_check; - const error = server.health_check_error; - - if (isLoadingHealth || isRechecking) { - return ( - - - Checking - - ); - } - - const getStatusColor = (status: string) => { - switch (status) { - case "healthy": - return "text-green-700 bg-green-50 border border-green-200"; - case "unhealthy": - return "text-red-700 bg-red-50 border border-red-200"; - default: - return "text-gray-600 bg-gray-50 border border-gray-200"; - } - }; - - const getStatusIcon = (status: string) => { - switch (status) { - case "healthy": - return "✓"; - case "unhealthy": - return "✗"; - default: - return "?"; - } - }; - - const isClickable = !!onRecheck; - - const tooltipContent = ( -
-
Health Status: {status}
- {lastCheck &&
Last Check: {new Date(lastCheck).toLocaleString()}
} - {error && ( -
-
Error:
-
{error}
-
- )} - {!lastCheck && !error &&
No health check data available
} - {isClickable &&
Click to recheck
} -
- ); - - return ( - - setIsHovered(true)} - onMouseLeave={() => setIsHovered(false)} - onClick={isClickable ? () => onRecheck(server.server_id) : undefined} - > - {isHovered && isClickable ? "↻" : getStatusIcon(status)} - {isHovered && isClickable ? "Recheck" : status.charAt(0).toUpperCase() + status.slice(1)} - - - ); -}; - -export const mcpServerColumns = ( - userRole: string, - onView: (serverId: string) => void, - onEdit: (serverId: string) => void, - onDelete: (serverId: string) => void, - isLoadingHealth?: boolean, - onByokConnect?: (server: MCPServer) => void, - onRecheckHealth?: (serverId: string) => void, - recheckingServerIds?: Set, -): ColumnDef[] => [ - { - accessorKey: "server_id", - header: "Server ID", - enableSorting: true, - cell: ({ row }) => ( - - ), - }, - { - accessorKey: "server_name", - header: "Name", - enableSorting: true, - cell: ({ row }) => { - const logoUrl = row.original.mcp_info?.logo_url; - const name = row.original.server_name; - return ( -
- {logoUrl ? ( - {`${name { - (e.target as HTMLImageElement).style.display = "none"; - }} - /> - ) : null} - {name} -
- ); - }, - }, - { - accessorKey: "alias", - header: "Alias", - enableSorting: true, - }, - { - id: "url", - header: "URL", - cell: ({ row }) => { - const url = row.original.url; - if (!url) { - return —; - } - const { maskedUrl } = getMaskedAndFullUrl(url); - return {maskedUrl}; - }, - }, - { - accessorKey: "transport", - header: "Transport", - enableSorting: true, - cell: ({ row }) => { - const transport = row.original.transport || "http"; - const specPath = row.original.spec_path; - const displayTransport = specPath && transport !== "stdio" ? "OPENAPI" : transport; - const label = displayTransport.toUpperCase(); - return ( - - {label} - - ); - }, - }, - { - accessorKey: "auth_type", - header: "Auth Type", - enableSorting: true, - cell: ({ getValue }) => { - const authType = (getValue() as string) || "none"; - return ( - - {authType} - - ); - }, - }, - { - id: "health_status", - header: "Health Status", - cell: ({ row }) => ( - - ), - }, - { - id: "mcp_access_groups", - header: "Access Groups", - cell: ({ row }) => { - const groups = row.original.mcp_access_groups; - if (Array.isArray(groups) && groups.length > 0) { - if (typeof groups[0] === "string") { - const joined = groups.join(", "); - return ( - -
- - {groups[0]} - - {groups.length > 1 && +{groups.length - 1}} -
- - ); - } - } - return —; - }, - }, - { - id: "available_on_public_internet", - header: "Network Access", - cell: ({ row }) => { - const isPublic = row.original.available_on_public_internet; - return isPublic ? ( - - - Public - - ) : ( - - - Internal - - ); - }, - }, - { - header: "Created", - accessorKey: "created_at", - enableSorting: true, - sortingFn: "datetime", - cell: ({ row }) => { - const server = row.original; - if (!server.created_at) return —; - const date = new Date(server.created_at); - return ( - - {date.toLocaleDateString()} - - ); - }, - }, - { - header: "Updated", - accessorKey: "updated_at", - enableSorting: true, - sortingFn: "datetime", - cell: ({ row }) => { - const server = row.original; - if (!server.updated_at) return —; - const date = new Date(server.updated_at); - return ( - - {date.toLocaleDateString()} - - ); - }, - }, - { - id: "byok_credential", - header: "Credential", - cell: ({ row }) => { - const server = row.original; - if (!server.is_byok) { - return —; - } - if (server.has_user_credential) { - return ( -
- - Connected - - {onByokConnect && ( - - )} -
- ); - } - return onByokConnect ? ( - - ) : null; - }, - }, - { - id: "actions", - header: "Actions", - cell: ({ row }) => ( -
- - - - - - -
- ), - }, -]; From 2cd7e874859eade595928b97b73fab5e10620629 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 10 Jun 2026 08:22:15 +0530 Subject: [PATCH 042/185] fix(proxy): authorize batch files using upload target_model_names (LIT-3593) (#30009) * fix(proxy): authorize batch files using upload target_model_names (LIT-3593) After replace_model_in_jsonl, body.model is a stripped provider id. Reverse-mapping it via resolve_model_name_from_model_id is first-match on model_list and caused false 403s when multiple deployments share the same stripped name. Use target_model_names from the unified file id instead. Co-authored-by: Cursor * fix(proxy): restore resolve_model_name_from_model_id for JSONL fallback path (LIT-3593) Restores the reverse-lookup for the JSONL body.model fallback path so that legacy/pre-target_model_names managed files still map stripped provider IDs back to proxy aliases before auth. Also cleans up redundant `or None`. Co-Authored-By: Claude Sonnet 4.6 * Revert "fix(proxy): restore resolve_model_name_from_model_id for JSONL fallback path (LIT-3593)" This reverts commit 30d2e96f77ef521ccaaf2193fe554980380eb669. --------- Co-authored-by: Cursor Co-authored-by: Claude Sonnet 4.6 --- litellm/proxy/hooks/batch_rate_limiter.py | 29 ++++-- .../proxy/hooks/test_batch_file_validation.py | 99 ++++++++++++++++++- 2 files changed, 116 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index 3957e3a7fbb..5b691beccbf 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -531,11 +531,17 @@ class _PROXY_BatchRateLimiter(CustomLogger): # Check if this is a managed file (base64 encoded unified file ID) from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, + get_models_from_unified_file_id, ) # Managed files require bypassing the HTTP endpoint (which runs access-check hooks) # and calling the managed files hook directly with the user's credentials. is_managed_file = _is_base64_encoded_unified_file_id(file_id) + target_model_names = ( + get_models_from_unified_file_id(is_managed_file) + if is_managed_file + else [] + ) if is_managed_file and user_api_key_dict is not None: file_content = await self._fetch_managed_file_content( file_id=file_id, @@ -573,6 +579,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): await self._enforce_batch_file_model_access( user_api_key_dict=user_api_key_dict, file_content_as_dict=file_content_as_dict, + target_model_names=target_model_names or None, ) input_file_usage = _get_batch_job_input_file_usage( @@ -608,9 +615,13 @@ class _PROXY_BatchRateLimiter(CustomLogger): self, user_api_key_dict: UserAPIKeyAuth, file_content_as_dict: List[dict], + target_model_names: Optional[List[str]] = None, ) -> None: - """Reject the batch if the caller is not authorized for every - ``body.model`` named inside the JSONL. + """Reject the batch if the caller is not authorized for the upload target. + + For managed files, ``target_model_names`` (from the unified file id) is + the proxy alias the file was uploaded for and is used directly for auth. + For legacy/non-managed files, falls back to ``body.model`` values in the JSONL. Reuses standard auth helpers so the same model access rules the proxy enforces on `/chat/completions` apply here. @@ -627,9 +638,12 @@ class _PROXY_BatchRateLimiter(CustomLogger): from litellm.proxy.proxy_server import proxy_logging_obj from litellm.proxy.proxy_server import user_api_key_cache - models = _get_models_from_batch_input_file_content(file_content_as_dict) - if not models: - return + if target_model_names: + models = target_model_names + else: + models = _get_models_from_batch_input_file_content(file_content_as_dict) + if not models: + return team_object = None if ( @@ -660,12 +674,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): llm_model_list = llm_router.model_list if llm_router is not None else None for model in models: - # body.model may be the provider id after replace_model_in_jsonl; map to proxy model_name for auth. model_to_check = model - if llm_router is not None: - proxy_model_name = llm_router.resolve_model_name_from_model_id(model) - if proxy_model_name is not None: - model_to_check = proxy_model_name try: if team_object is not None: try: diff --git a/tests/test_litellm/proxy/hooks/test_batch_file_validation.py b/tests/test_litellm/proxy/hooks/test_batch_file_validation.py index ae71d10b378..a6f6e651487 100644 --- a/tests/test_litellm/proxy/hooks/test_batch_file_validation.py +++ b/tests/test_litellm/proxy/hooks/test_batch_file_validation.py @@ -713,7 +713,8 @@ async def test_count_input_file_usage_decodes_model_embedded_file_id(): @pytest.mark.asyncio async def test_pre_call_allows_stripped_provider_model_when_key_has_proxy_alias(): """After replace_model_in_jsonl, body.model is the provider id (e.g. gpt-5.5). - Auth must check the proxy model_name the key was granted, not the stripped id.""" + Auth must check target_model_names from the unified file id, not reverse-map + the stripped id.""" from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter rate_limiter = _PROXY_BatchRateLimiter( @@ -732,7 +733,6 @@ async def test_pre_call_allows_stripped_provider_model_when_key_has_proxy_alias( ) mock_router = MagicMock() mock_router.model_list = [] - mock_router.resolve_model_name_from_model_id.return_value = proxy_alias can_key_call_model = AsyncMock(return_value=True) with ( @@ -745,10 +745,105 @@ async def test_pre_call_allows_stripped_provider_model_when_key_has_proxy_alias( await rate_limiter._enforce_batch_file_model_access( user_api_key_dict=user, file_content_as_dict=file_dict, + target_model_names=[proxy_alias], ) can_key_call_model.assert_awaited_once() assert can_key_call_model.await_args.kwargs["model"] == proxy_alias + mock_router.resolve_model_name_from_model_id.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model_list_order", + [ + [ + "openai/openai/gpt-5.5", + "openai/openai/gpt-5.5-batch", + "us/azure/openai/gpt-5.5", + ], + [ + "us/azure/openai/gpt-5.5", + "openai/openai/gpt-5.5", + "openai/openai/gpt-5.5-batch", + ], + [ + "openai/openai/gpt-5.5-batch", + "us/azure/openai/gpt-5.5", + "openai/openai/gpt-5.5", + ], + ], +) +async def test_pre_call_uses_target_model_names_not_stripped_reverse_lookup( + model_list_order, +): + """LIT-3593: three deployments strip to gpt-5.5; auth must use the upload + target alias from target_model_names, not first-match reverse lookup.""" + from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter + + rate_limiter = _PROXY_BatchRateLimiter( + internal_usage_cache=MagicMock(), + parallel_request_limiter=MagicMock(), + ) + batch_alias = "openai/openai/gpt-5.5-batch" + deployment_templates = { + "openai/openai/gpt-5.5": { + "model_name": "openai/openai/gpt-5.5", + "litellm_params": {"model": "openai/gpt-5.5"}, + "model_info": {"id": "openai/openai/gpt-5.5", "mode": "chat"}, + }, + "openai/openai/gpt-5.5-batch": { + "model_name": "openai/openai/gpt-5.5-batch", + "litellm_params": {"model": "openai/gpt-5.5"}, + "model_info": {"id": "openai/openai/gpt-5.5-batch", "mode": "batch"}, + }, + "us/azure/openai/gpt-5.5": { + "model_name": "us/azure/openai/gpt-5.5", + "litellm_params": {"model": "azure/gpt-5.5"}, + "model_info": {"id": "openai/openai/gpt-5.5", "mode": "chat"}, + }, + } + mock_router = MagicMock() + mock_router.model_list = [deployment_templates[name] for name in model_list_order] + + def _resolve(model_id): + for deployment in mock_router.model_list: + actual_model = deployment.get("litellm_params", {}).get("model") + if actual_model == model_id or ( + actual_model and actual_model.endswith(f"/{model_id}") + ): + return deployment.get("model_name") + return None + + mock_router.resolve_model_name_from_model_id.side_effect = _resolve + + file_dict = [ + {"body": {"model": "gpt-5.5", "messages": [{"role": "user", "content": "x"}]}} + ] + user = UserAPIKeyAuth( + api_key="sk-ok", + user_id="alice", + models=[batch_alias], + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + can_key_call_model = AsyncMock(return_value=True) + + with ( + patch( + "litellm.proxy.auth.auth_checks.can_key_call_model", + new=can_key_call_model, + ), + patch("litellm.proxy.proxy_server.llm_router", mock_router), + ): + await rate_limiter._enforce_batch_file_model_access( + user_api_key_dict=user, + file_content_as_dict=file_dict, + target_model_names=[batch_alias], + ) + + can_key_call_model.assert_awaited_once() + assert can_key_call_model.await_args.kwargs["model"] == batch_alias + mock_router.resolve_model_name_from_model_id.assert_not_called() @pytest.mark.asyncio From e15b37a18eac240c690763c60ca409d13c7be2e4 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 9 Jun 2026 20:20:15 -0700 Subject: [PATCH 043/185] Add Claude Fable 5 across Anthropic, Bedrock, Vertex AI, and Azure AI (#30064) * Add Claude Fable 5 across Anthropic, Bedrock, Vertex AI, and Azure AI Adds cost map entries for claude-fable-5 ($10/$50 per MTok, 1M context, 128K output, adaptive thinking only) on the Anthropic API, Bedrock converse (base, global, and us/eu geo inference profiles at the 10% regional premium), Vertex AI, and Azure AI (Microsoft Foundry, which serves Fable 5 with the full 1M context window unlike Opus 4.8). Registers anthropic.claude-fable-5 in BEDROCK_CONVERSE_MODELS, lists the model in the setup wizard, and extends the reasoning effort e2e grid. The Bedrock, Vertex, and Azure grid cells carry fail_reason markers until the CI accounts are provisioned: Bedrock needs the provider data sharing opt-in Fable 5 requires, and the Foundry resource needs a claude-fable-5 deployment. The first-party entry carries provider_specific_entry {us: 1.1} for the inference_geo premium and deliberately no fast multiplier since Fable 5 has no fast mode. https://claude.ai/code/session_01MZarYYT3aS7DxaNjoax6Gm * Drop removed sampling params for Claude 4.7+ when drop_params is set Fable 5, Opus 4.7, and Opus 4.8 removed sampling params: the API rejects top_p, top_k, and any temperature other than 1 with a 400. LiteLLM was forwarding them even with drop_params enabled because the Anthropic and Bedrock converse transformations passed temperature/top_p through unconditionally. Mirror the GPT-5/o-series handling: temperature=1 still passes through, other values and any top_p are dropped when drop_params is set, and without drop_params a clean client-side UnsupportedParamsError tells the caller how to opt in, instead of surfacing the raw provider error. https://claude.ai/code/session_01MZarYYT3aS7DxaNjoax6Gm * Drive sampling param gating from the cost map and cover top_k Greptile review follow-ups on the sampling param fix: the restriction for Fable 5 / Opus 4.7 / 4.8 is now declared as supports_sampling_params: false on every affected cost map entry (perplexity excluded; that route is OpenAI-compatible and maps sampling params upstream) and read back through a tri-state map lookup, keeping the name check only as a fallback for provider-routed ids whose hosted map entries predate the flag, the same layering supports_adaptive_thinking uses. top_k bypasses map_openai_params as a provider-specific kwarg, so it is gated at the shared AnthropicConfig.transform_request boundary (direct, Bedrock invoke, Vertex, Azure) and in the Bedrock converse _handle_top_k_value path, with drop_params threaded through the converse transform helpers. Also updates the reasoning effort grid cell count assertion for the four Fable 5 rows added on this branch (29 x 11 cells). https://claude.ai/code/session_01MZarYYT3aS7DxaNjoax6Gm * Declare supports_sampling_params in the cost map schema The model map validation schema uses additionalProperties: false, so the new flag must be declared for the 28 entries that carry it; this was the one failing job (misc / Run tests) on the previous commit. https://claude.ai/code/session_01MZarYYT3aS7DxaNjoax6Gm * fix(bedrock): gate top_k=0 on converse to match Anthropic boundary Truthiness check let top_k=0 silently disappear on models that removed sampling params, while AnthropicConfig.transform_request treats 0 as present and raises UnsupportedParamsError (or drops when drop_params is set). Switch to 'is not None' so converse, direct Anthropic, invoke, Vertex, and Azure all behave the same for top_k=0. --------- Co-authored-by: Cursor Agent --- litellm/constants.py | 1 + litellm/llms/anthropic/chat/transformation.py | 27 +- litellm/llms/anthropic/common_utils.py | 112 +++++-- .../bedrock/chat/converse_transformation.py | 39 ++- ...odel_prices_and_context_window_backup.json | 276 ++++++++++++++++++ litellm/setup_wizard.py | 3 +- model_prices_and_context_window.json | 276 ++++++++++++++++++ .../reasoning_effort_grid/grid_spec.py | 52 +++- .../test_reasoning_effort_grid.py | 5 +- .../test_anthropic_chat_transformation.py | 139 +++++++++ .../chat/test_converse_transformation.py | 119 ++++++++ .../test_claude_fable_5_config.py | 230 +++++++++++++++ tests/test_litellm/test_utils.py | 1 + 13 files changed, 1240 insertions(+), 40 deletions(-) create mode 100644 tests/test_litellm/test_claude_fable_5_config.py diff --git a/litellm/constants.py b/litellm/constants.py index f10cec034f0..57f55e6c177 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1158,6 +1158,7 @@ BEDROCK_CONVERSE_MODELS = [ "openai.gpt-oss-120b-1:0", "anthropic.claude-haiku-4-5-20251001-v1:0", "anthropic.claude-sonnet-4-5-20250929-v1:0", + "anthropic.claude-fable-5", "anthropic.claude-opus-4-8", "anthropic.claude-opus-4-7", "anthropic.claude-opus-4-6-v1:0", diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 3f30d5d6807..9ecd0df0cb8 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -1455,10 +1455,15 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): _value = self._map_stop_sequences(value) if _value is not None: optional_params["stop_sequences"] = _value - elif param == "temperature": - optional_params["temperature"] = value - elif param == "top_p": - optional_params["top_p"] = value + elif param == "temperature" or param == "top_p": + AnthropicConfig._apply_sampling_param( + optional_params=optional_params, + model=model, + param=param, + value=value, + drop_params=drop_params, + output_key=param, + ) elif param == "response_format" and isinstance(value, dict): if any( substring in model @@ -1975,6 +1980,20 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): optional_params.pop("is_vertex_request", None) optional_params.pop("client_metadata", None) + # ``top_k`` is a provider-specific kwarg that bypasses + # ``map_openai_params``; gate it here, the single boundary shared by + # the direct Anthropic, Bedrock invoke, Vertex, and Azure paths. + top_k = optional_params.pop("top_k", None) + if top_k is not None: + AnthropicConfig._apply_sampling_param( + optional_params=optional_params, + model=model, + param="top_k", + value=top_k, + drop_params=litellm_params.get("drop_params") is True, + output_key="top_k", + ) + data = { "model": model, "messages": anthropic_messages, diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 3f002d73cbc..5741513903c 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -272,23 +272,68 @@ class AnthropicModelInfo(BaseLLMModelInfo): ) @staticmethod - def _supports_model_capability(model: str, key: str) -> bool: - """Check a boolean capability ``key`` in the model map. + def _supports_sampling_params(model: str) -> bool: + """Claude 4.7+ (Opus 4.7/4.8, Fable 5) removed sampling params: the API + rejects ``top_p``, ``top_k``, and any ``temperature`` other than 1 with + a 400 ("`temperature` is deprecated for this model"). - Strips bedrock/vertex prefixes so a provider-routed Claude still - resolves to the Anthropic model-map entry. - """ - from litellm.utils import _supports_factory + Driven by the ``supports_sampling_params`` flag in the model map; the + name check remains only as a fallback for provider-routed ids whose + map entries predate the flag.""" + flag = AnthropicModelInfo._get_model_capability( + model, "supports_sampling_params" + ) + if flag is not None: + return flag + model_lower = model.lower() + return not any( + v in model_lower + for v in ( + "fable", + "opus-4-7", + "opus_4_7", + "opus-4.7", + "opus_4.7", + "opus-4-8", + "opus_4_8", + "opus-4.8", + "opus_4.8", + ) + ) - try: - if _supports_factory( - model=model, - custom_llm_provider="anthropic", - key=key, - ): - return True - except Exception: - pass + @staticmethod + def _apply_sampling_param( + optional_params: dict, + model: str, + param: str, + value: Any, + drop_params: bool, + output_key: str, + ) -> None: + """Forward ``temperature``/``top_p``/``top_k`` to + ``optional_params[output_key]`` unless the model removed sampling + params, in which case drop the param (with drop_params) or raise a + clean client-side 400.""" + if AnthropicModelInfo._supports_sampling_params(model) or ( + param == "temperature" and value == 1 + ): + optional_params[output_key] = value + elif not (litellm.drop_params or drop_params): + supported_hint = ( + "Only temperature=1 is supported. " if param == "temperature" else "" + ) + raise litellm.utils.UnsupportedParamsError( + message=( + f"{model} does not support {param}={value}. {supported_hint}" + "To drop unsupported params, set `litellm.drop_params = True`." + ), + status_code=400, + ) + + @staticmethod + def _model_map_lookup_candidates(model: str) -> List[str]: + """Model-map keys to try for ``model``, stripping bedrock/vertex + prefixes so a provider-routed Claude still resolves to its entry.""" candidates = [model] for prefix in ( "bedrock/converse/", @@ -307,15 +352,40 @@ class AnthropicModelInfo(BaseLLMModelInfo): candidates.append(f"bedrock/{base}") except Exception: pass + return candidates + + @staticmethod + def _get_model_capability(model: str, key: str) -> Optional[bool]: + """Read boolean capability ``key`` from the model map, or None when + no entry declares it.""" try: - for cand in candidates: - if cand in litellm.model_cost and ( - litellm.model_cost[cand].get(key) is True - ): - return True + for cand in AnthropicModelInfo._model_map_lookup_candidates(model): + value = litellm.model_cost.get(cand, {}).get(key) + if isinstance(value, bool): + return value except Exception: pass - return False + return None + + @staticmethod + def _supports_model_capability(model: str, key: str) -> bool: + """Check a boolean capability ``key`` in the model map. + + Strips bedrock/vertex prefixes so a provider-routed Claude still + resolves to the Anthropic model-map entry. + """ + from litellm.utils import _supports_factory + + try: + if _supports_factory( + model=model, + custom_llm_provider="anthropic", + key=key, + ): + return True + except Exception: + pass + return AnthropicModelInfo._get_model_capability(model, key) is True @staticmethod def _is_adaptive_thinking_model(model: str) -> bool: diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index ea0326dffd1..b5e5e4de6fc 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -920,10 +920,15 @@ class AmazonConverseConfig(BaseConfig): continue value = [value] optional_params["stopSequences"] = value - if param == "temperature": - optional_params["temperature"] = value - if param == "top_p": - optional_params["topP"] = value + if param == "temperature" or param == "top_p": + AnthropicConfig._apply_sampling_param( + optional_params=optional_params, + model=model, + param=param, + value=value, + drop_params=drop_params, + output_key="topP" if param == "top_p" else param, + ) if param == "tools" and isinstance(value, list): self._apply_tool_call_transformation( tools=cast(List[OpenAIChatCompletionToolParam], value), @@ -1221,7 +1226,9 @@ class AmazonConverseConfig(BaseConfig): inference_params["topK"] = inference_params.pop("top_k") return InferenceConfig(**inference_params) - def _handle_top_k_value(self, model: str, inference_params: dict) -> dict: + def _handle_top_k_value( + self, model: str, inference_params: dict, drop_params: bool = False + ) -> dict: base_model = BedrockModelInfo.get_base_model(model) val_top_k = None @@ -1230,16 +1237,25 @@ class AmazonConverseConfig(BaseConfig): elif "top_k" in inference_params: val_top_k = inference_params.pop("top_k") - if val_top_k: + if val_top_k is not None: if base_model.startswith("anthropic"): - return {"top_k": val_top_k} + top_k_params: dict = {} + AnthropicConfig._apply_sampling_param( + optional_params=top_k_params, + model=model, + param="top_k", + value=val_top_k, + drop_params=drop_params, + output_key="top_k", + ) + return top_k_params if base_model.startswith("amazon.nova"): return {"inferenceConfig": {"topK": val_top_k}} return {} def _prepare_request_params( - self, optional_params: dict, model: str + self, optional_params: dict, model: str, drop_params: bool = False ) -> Tuple[dict, dict, dict, Optional[OutputConfigBlock]]: """Prepare and separate request parameters.""" # Consume the internal ``_output_config_normalized`` marker set by @@ -1338,7 +1354,7 @@ class AmazonConverseConfig(BaseConfig): # Only set the topK value in for models that support it additional_request_params.update( - self._handle_top_k_value(model, inference_params) + self._handle_top_k_value(model, inference_params, drop_params) ) # Filter out internal/MCP-related parameters that shouldn't be sent to the API @@ -1572,6 +1588,7 @@ class AmazonConverseConfig(BaseConfig): optional_params: dict, messages: Optional[List[AllMessageValues]] = None, headers: Optional[dict] = None, + drop_params: bool = False, ) -> CommonRequestObject: ## VALIDATE REQUEST """ @@ -1618,7 +1635,7 @@ class AmazonConverseConfig(BaseConfig): additional_request_params, request_metadata, output_config, - ) = self._prepare_request_params(optional_params, model) + ) = self._prepare_request_params(optional_params, model, drop_params) original_tools = inference_params.pop("tools", []) @@ -1701,6 +1718,7 @@ class AmazonConverseConfig(BaseConfig): optional_params=optional_params, messages=messages, headers=headers, + drop_params=litellm_params.get("drop_params") is True, ) bedrock_messages = ( @@ -1758,6 +1776,7 @@ class AmazonConverseConfig(BaseConfig): optional_params=optional_params, messages=messages, headers=headers, + drop_params=litellm_params.get("drop_params") is True, ) ## TRANSFORMATION ## diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 282a292ab17..3782da1350f 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -1156,6 +1156,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1202,6 +1203,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1233,6 +1235,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1264,6 +1267,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1295,6 +1299,139 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh" + }, + "anthropic.claude-fable-5": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 1e-06, + "input_cost_per_token": 1e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh" + }, + "global.anthropic.claude-fable-5": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 1e-06, + "input_cost_per_token": 1e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh" + }, + "us.anthropic.claude-fable-5": { + "cache_creation_input_token_cost": 1.375e-05, + "cache_creation_input_token_cost_above_1hr": 2.2e-05, + "cache_read_input_token_cost": 1.1e-06, + "input_cost_per_token": 1.1e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh" + }, + "eu.anthropic.claude-fable-5": { + "cache_creation_input_token_cost": 1.375e-05, + "cache_creation_input_token_cost_above_1hr": 2.2e-05, + "cache_read_input_token_cost": 1.1e-06, + "input_cost_per_token": 1.1e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1327,6 +1464,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1359,6 +1497,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1391,6 +1530,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1423,6 +1563,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1455,6 +1596,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1485,6 +1627,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -2208,6 +2351,37 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true + }, + "azure_ai/claude-fable-5": { + "input_cost_per_token": 1e-05, + "output_cost_per_token": 5e-05, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 1e-06, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -2237,6 +2411,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -10170,6 +10345,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -10204,6 +10380,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -10214,6 +10391,40 @@ }, "supports_output_config": true }, + "claude-fable-5": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 1e-06, + "input_cost_per_token": 1e-05, + "litellm_provider": "anthropic", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true + }, "claude-opus-4-8": { "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, @@ -10238,6 +10449,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -34004,6 +34216,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -34032,6 +34245,67 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true + }, + "vertex_ai/claude-fable-5": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 1e-06, + "input_cost_per_token": 1e-05, + "litellm_provider": "vertex_ai-anthropic_models", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true + }, + "vertex_ai/claude-fable-5@default": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 1e-06, + "input_cost_per_token": 1e-05, + "litellm_provider": "vertex_ai-anthropic_models", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -34061,6 +34335,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -34090,6 +34365,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, diff --git a/litellm/setup_wizard.py b/litellm/setup_wizard.py index 862ca13e7ba..2f0cb1233ae 100644 --- a/litellm/setup_wizard.py +++ b/litellm/setup_wizard.py @@ -52,11 +52,12 @@ PROVIDERS: List[Dict] = [ { "id": "anthropic", "name": "Anthropic", - "description": "Claude Opus 4.8, Opus 4.7, Opus 4.6, Sonnet 4.6, Haiku 4.5", + "description": "Claude Fable 5, Opus 4.8, Opus 4.7, Opus 4.6, Sonnet 4.6, Haiku 4.5", "env_key": "ANTHROPIC_API_KEY", "key_hint": "sk-ant-...", "test_model": "claude-haiku-4-5-20251001", "models": [ + "claude-fable-5", "claude-opus-4-8", "claude-opus-4-7", "claude-opus-4-6", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index b0ffc66d03b..85cb06b7f19 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -1156,6 +1156,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1202,6 +1203,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1233,6 +1235,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1264,6 +1267,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1295,6 +1299,139 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh" + }, + "anthropic.claude-fable-5": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 1e-06, + "input_cost_per_token": 1e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh" + }, + "global.anthropic.claude-fable-5": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 1e-06, + "input_cost_per_token": 1e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh" + }, + "us.anthropic.claude-fable-5": { + "cache_creation_input_token_cost": 1.375e-05, + "cache_creation_input_token_cost_above_1hr": 2.2e-05, + "cache_read_input_token_cost": 1.1e-06, + "input_cost_per_token": 1.1e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh" + }, + "eu.anthropic.claude-fable-5": { + "cache_creation_input_token_cost": 1.375e-05, + "cache_creation_input_token_cost_above_1hr": 2.2e-05, + "cache_read_input_token_cost": 1.1e-06, + "input_cost_per_token": 1.1e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1327,6 +1464,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1359,6 +1497,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1391,6 +1530,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1423,6 +1563,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1455,6 +1596,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -1485,6 +1627,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -2208,6 +2351,37 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true + }, + "azure_ai/claude-fable-5": { + "input_cost_per_token": 1e-05, + "output_cost_per_token": 5e-05, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 1e-06, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -2237,6 +2411,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -10170,6 +10345,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -10204,6 +10380,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -10214,6 +10391,40 @@ }, "supports_output_config": true }, + "claude-fable-5": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 1e-06, + "input_cost_per_token": 1e-05, + "litellm_provider": "anthropic", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true + }, "claude-opus-4-8": { "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, @@ -10238,6 +10449,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -34044,6 +34256,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -34072,6 +34285,67 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true + }, + "vertex_ai/claude-fable-5": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 1e-06, + "input_cost_per_token": 1e-05, + "litellm_provider": "vertex_ai-anthropic_models", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true + }, + "vertex_ai/claude-fable-5@default": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 1e-06, + "input_cost_per_token": 1e-05, + "litellm_provider": "vertex_ai-anthropic_models", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -34101,6 +34375,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, @@ -34130,6 +34405,7 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, + "supports_sampling_params": false, "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, diff --git a/tests/llm_translation/reasoning_effort_grid/grid_spec.py b/tests/llm_translation/reasoning_effort_grid/grid_spec.py index a08013cd439..83a2c286d64 100644 --- a/tests/llm_translation/reasoning_effort_grid/grid_spec.py +++ b/tests/llm_translation/reasoning_effort_grid/grid_spec.py @@ -1,7 +1,6 @@ from dataclasses import dataclass, field from typing import Dict, FrozenSet, List, Optional, Tuple - OMIT = object() @@ -136,6 +135,13 @@ _CAPS_NONE: FrozenSet[str] = frozenset() ANTHROPIC_DIRECT_MODELS: Tuple[ModelEntry, ...] = ( + ModelEntry( + alias="claude-fable-5", + model="anthropic/claude-fable-5", + mode="adaptive", + required_env=_ANTHROPIC_REQ, + caps=_CAPS_XHIGH_MAX, + ), ModelEntry( alias="claude-opus-4-8", model="anthropic/claude-opus-4-8", @@ -168,6 +174,19 @@ ANTHROPIC_DIRECT_MODELS: Tuple[ModelEntry, ...] = ( AZURE_AI_MODELS: Tuple[ModelEntry, ...] = ( + ModelEntry( + alias="azure-claude-fable-5", + model="azure_ai/claude-fable-5", + mode="adaptive", + required_env=_AZURE_FOUNDRY_REQ, + caps=_CAPS_XHIGH_MAX, + fail_reason=( + "claude-fable-5 has no deployment on the CI Microsoft Foundry " + "resource yet; Foundry returns DeploymentNotFound until someone " + "creates the fable-5 deployment, so this cell stays loud in CI. " + "Remove this fail_reason once the deployment exists." + ), + ), ModelEntry( alias="azure-claude-opus-4-8", model="azure_ai/claude-opus-4-8", @@ -213,6 +232,20 @@ AZURE_AI_MODELS: Tuple[ModelEntry, ...] = ( VERTEX_AI_MODELS: Tuple[ModelEntry, ...] = ( + ModelEntry( + alias="vertex-claude-fable-5", + model="vertex_ai/claude-fable-5", + mode="adaptive", + extra_params=(("vertex_location", "global"),), + required_env=_VERTEX_REQ, + caps=_CAPS_XHIGH_MAX, + fail_reason=( + "claude-fable-5 availability on the CI Vertex project is not yet " + "confirmed for this brand-new release, so this cell stays loud in " + "CI until verified. Remove this fail_reason once the model is " + "confirmed available on the global Vertex endpoint." + ), + ), ModelEntry( alias="vertex-claude-opus-4-8", model="vertex_ai/claude-opus-4-8", @@ -263,6 +296,23 @@ VERTEX_AI_MODELS: Tuple[ModelEntry, ...] = ( BEDROCK_CONVERSE_MODELS: Tuple[ModelEntry, ...] = ( + ModelEntry( + alias="bedrock-claude-fable-5", + model="bedrock/converse/us.anthropic.claude-fable-5", + mode="adaptive", + extra_params=(("aws_region_name", "us-east-1"),), + required_env=_BEDROCK_REQ, + caps=_CAPS_XHIGH_MAX, + bedrock_effort_ceiling="xhigh", + unavailable_error="is not available for this account", + fail_reason=( + "claude-fable-5 on Bedrock requires the account to opt in to " + "provider data sharing (data retention mode " + "'provider_data_sharing' via the Data Retention API); the CI " + "account has not opted in yet, so this cell stays loud in CI. " + "Remove this fail_reason once the opt-in is done." + ), + ), ModelEntry( alias="bedrock-claude-opus-4-8", model="bedrock/converse/us.anthropic.claude-opus-4-8", diff --git a/tests/llm_translation/reasoning_effort_grid/test_reasoning_effort_grid.py b/tests/llm_translation/reasoning_effort_grid/test_reasoning_effort_grid.py index 551ab8459d1..a5f16f928e5 100644 --- a/tests/llm_translation/reasoning_effort_grid/test_reasoning_effort_grid.py +++ b/tests/llm_translation/reasoning_effort_grid/test_reasoning_effort_grid.py @@ -15,7 +15,6 @@ from .grid_spec import ( all_cells, ) - _PROMPT_MESSAGES: List[Dict[str, str]] = [ {"role": "user", "content": "Step by step, calculate 47 * 53. Show your work."} ] @@ -201,8 +200,8 @@ async def test_reasoning_effort_grid( def test_grid_cell_count() -> None: - assert len(_PARAMS) == 25 * 11, ( - f"expected 275 cells (25 provider x model combos x 11 efforts), " + assert len(_PARAMS) == 29 * 11, ( + f"expected 319 cells (29 provider x model combos x 11 efforts), " f"got {len(_PARAMS)}" ) diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py index 75038574c63..abb162e9ddb 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -5261,6 +5261,8 @@ def test_should_strip_billing_metadata_by_provider( config_cls = getattr(importlib.import_module(module_path), class_name) assert config_cls().should_strip_billing_metadata() is expected_strip + + def test_namespace_tool_flat_nested_tools_are_extracted(): """Codex sends nested tools in flat format {type, name, description, parameters} with no 'function' wrapper. These must be normalized and mapped without raising KeyError: 'function'.""" @@ -5357,3 +5359,140 @@ def test_client_metadata_stripped_from_anthropic_request(): headers={}, ) assert "client_metadata" not in result + + +@pytest.mark.parametrize( + "model", + ["claude-fable-5", "claude-opus-4-7", "claude-opus-4-8-20260120"], +) +def test_sampling_params_dropped_for_models_that_removed_them(model): + """Fable 5 / Opus 4.7 / 4.8 reject temperature != 1 and any top_p with a + 400; with drop_params set they must be dropped, not forwarded (#30064).""" + config = AnthropicConfig() + + result = config.map_openai_params( + non_default_params={"temperature": 0.5, "top_p": 0.9}, + optional_params={}, + model=model, + drop_params=True, + ) + + assert "temperature" not in result + assert "top_p" not in result + + +@pytest.mark.parametrize("params", [{"temperature": 0.5}, {"top_p": 0.9}, {"top_p": 1}]) +def test_sampling_params_raise_clean_error_without_drop_params(params, monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + config = AnthropicConfig() + + with pytest.raises(litellm.utils.UnsupportedParamsError, match="drop_params"): + config.map_openai_params( + non_default_params=params, + optional_params={}, + model="claude-fable-5", + drop_params=False, + ) + + +def test_temperature_1_forwarded_on_models_that_removed_sampling_params(): + """temperature=1 (the API default) is still accepted and must pass through.""" + config = AnthropicConfig() + + result = config.map_openai_params( + non_default_params={"temperature": 1}, + optional_params={}, + model="claude-fable-5", + drop_params=False, + ) + + assert result["temperature"] == 1 + + +@pytest.mark.parametrize("model", ["claude-opus-4-6", "claude-sonnet-4-6"]) +def test_sampling_params_forwarded_on_models_that_accept_them(model): + config = AnthropicConfig() + + result = config.map_openai_params( + non_default_params={"temperature": 0.5, "top_p": 0.9}, + optional_params={}, + model=model, + drop_params=True, + ) + + assert result["temperature"] == 0.5 + assert result["top_p"] == 0.9 + + +def test_sampling_param_gating_driven_by_model_map_flag(monkeypatch): + """The drop/raise decision must come from ``supports_sampling_params`` in + the model map, not just name matching: a flagged entry gates a model whose + name says nothing, and an explicit ``true`` overrides the name fallback.""" + monkeypatch.setitem( + litellm.model_cost, "claude-zeta-9", {"supports_sampling_params": False} + ) + monkeypatch.setitem( + litellm.model_cost, "claude-fable-5-test", {"supports_sampling_params": True} + ) + config = AnthropicConfig() + + flagged_off = config.map_openai_params( + non_default_params={"top_p": 0.9}, + optional_params={}, + model="claude-zeta-9", + drop_params=True, + ) + assert "top_p" not in flagged_off + + flagged_on = config.map_openai_params( + non_default_params={"top_p": 0.9}, + optional_params={}, + model="claude-fable-5-test", + drop_params=True, + ) + assert flagged_on["top_p"] == 0.9 + + +def test_top_k_dropped_at_transform_for_models_that_removed_it(): + """``top_k`` is a provider-specific kwarg that bypasses + ``map_openai_params``, so it must be stripped at the transform_request + boundary shared by the direct, invoke, Vertex, and Azure paths (#30064).""" + config = AnthropicConfig() + + result = config.transform_request( + model="claude-fable-5", + messages=[{"role": "user", "content": "hello"}], + optional_params={"max_tokens": 10, "top_k": 40}, + litellm_params={"drop_params": True}, + headers={}, + ) + + assert "top_k" not in result + + +def test_top_k_raises_at_transform_without_drop_params(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + config = AnthropicConfig() + + with pytest.raises(litellm.utils.UnsupportedParamsError, match="drop_params"): + config.transform_request( + model="claude-fable-5", + messages=[{"role": "user", "content": "hello"}], + optional_params={"max_tokens": 10, "top_k": 40}, + litellm_params={}, + headers={}, + ) + + +def test_top_k_forwarded_at_transform_on_models_that_accept_it(): + config = AnthropicConfig() + + result = config.transform_request( + model="claude-sonnet-4-6", + messages=[{"role": "user", "content": "hello"}], + optional_params={"max_tokens": 10, "top_k": 40}, + litellm_params={"drop_params": True}, + headers={}, + ) + + assert result["top_k"] == 40 diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index ed978113b8b..5c83f8b34f7 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -5267,3 +5267,122 @@ def test_transform_response_does_not_leak_body_on_parse_failure(): msg = str(exc_info.value) assert "secret content" not in msg assert "Error converting to valid response block" in msg + + +def test_converse_drops_sampling_params_for_models_that_removed_them(): + """Fable 5 / Opus 4.7 / 4.8 reject temperature != 1 and any top_p; with + drop_params set, converse must drop them instead of forwarding (#30064).""" + config = AmazonConverseConfig() + + result = config.map_openai_params( + non_default_params={"temperature": 0.5, "top_p": 0.9}, + optional_params={}, + model="us.anthropic.claude-fable-5", + drop_params=True, + ) + + assert "temperature" not in result + assert "topP" not in result + + +def test_converse_sampling_params_raise_without_drop_params(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + config = AmazonConverseConfig() + + with pytest.raises(litellm.utils.UnsupportedParamsError, match="drop_params"): + config.map_openai_params( + non_default_params={"temperature": 0.5}, + optional_params={}, + model="global.anthropic.claude-opus-4-8-v1:0", + drop_params=False, + ) + + +def test_converse_sampling_params_forwarded_on_models_that_accept_them(): + config = AmazonConverseConfig() + + result = config.map_openai_params( + non_default_params={"temperature": 0.5, "top_p": 0.9}, + optional_params={}, + model="us.anthropic.claude-sonnet-4-6", + drop_params=True, + ) + + assert result["temperature"] == 0.5 + assert result["topP"] == 0.9 + + +def test_converse_top_k_dropped_for_models_that_removed_it(): + """``top_k`` reaches converse as a provider-specific kwarg destined for + ``additionalModelRequestFields``, bypassing ``map_openai_params``; the + transform must strip it for models that removed sampling params (#30064).""" + config = AmazonConverseConfig() + + result = config.transform_request( + model="us.anthropic.claude-fable-5", + messages=[{"role": "user", "content": "hello"}], + optional_params={"top_k": 40}, + litellm_params={"drop_params": True}, + headers={}, + ) + + assert "top_k" not in result.get("additionalModelRequestFields", {}) + + +def test_converse_top_k_raises_without_drop_params(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + config = AmazonConverseConfig() + + with pytest.raises(litellm.utils.UnsupportedParamsError, match="drop_params"): + config.transform_request( + model="us.anthropic.claude-fable-5", + messages=[{"role": "user", "content": "hello"}], + optional_params={"top_k": 40}, + litellm_params={}, + headers={}, + ) + + +def test_converse_top_k_forwarded_on_models_that_accept_it(): + config = AmazonConverseConfig() + + result = config.transform_request( + model="us.anthropic.claude-sonnet-4-6", + messages=[{"role": "user", "content": "hello"}], + optional_params={"top_k": 40}, + litellm_params={"drop_params": True}, + headers={}, + ) + + assert result["additionalModelRequestFields"]["top_k"] == 40 + + +def test_converse_top_k_zero_raises_without_drop_params(monkeypatch): + """``top_k=0`` must hit the same gating as any other value; previously the + truthiness check let it silently disappear on models that removed sampling + params, diverging from the Anthropic boundary that treats ``0`` as present.""" + monkeypatch.setattr(litellm, "drop_params", False) + config = AmazonConverseConfig() + + with pytest.raises(litellm.utils.UnsupportedParamsError, match="drop_params"): + config.transform_request( + model="us.anthropic.claude-fable-5", + messages=[{"role": "user", "content": "hello"}], + optional_params={"top_k": 0}, + litellm_params={}, + headers={}, + ) + + +def test_converse_top_k_zero_forwarded_on_models_that_accept_it(): + config = AmazonConverseConfig() + + result = config.transform_request( + model="us.anthropic.claude-sonnet-4-6", + messages=[{"role": "user", "content": "hello"}], + optional_params={"top_k": 0}, + litellm_params={"drop_params": True}, + headers={}, + ) + + assert result["additionalModelRequestFields"]["top_k"] == 0 diff --git a/tests/test_litellm/test_claude_fable_5_config.py b/tests/test_litellm/test_claude_fable_5_config.py new file mode 100644 index 00000000000..d8d95fba0da --- /dev/null +++ b/tests/test_litellm/test_claude_fable_5_config.py @@ -0,0 +1,230 @@ +""" +Validate Claude Fable 5 model configuration entries. + +Fable 5 is a new tier above Opus ($10/$50 per MTok) with the same adaptive-only +API surface as Opus 4.7/4.8. The cost-map entries below are what make the model +resolvable across Anthropic, Bedrock, Vertex AI, and Azure AI (Microsoft +Foundry), and the ``supports_adaptive_thinking`` flag is what makes LiteLLM send +``thinking.type='adaptive'`` instead of the legacy ``enabled``/``budget_tokens`` +shape, which Fable 5 rejects with a 400. +""" + +import json +import os + +import pytest + +import litellm +from litellm.constants import BEDROCK_CONVERSE_MODELS +from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap + +REPO_ROOT = os.path.join(os.path.dirname(__file__), "../..") + + +def _load_root_cost_map() -> dict: + json_path = os.path.join(REPO_ROOT, "model_prices_and_context_window.json") + with open(json_path) as f: + return json.load(f) + + +@pytest.fixture +def local_model_cost_map(monkeypatch): + """Force the bundled backup cost map so assertions don't depend on the + network-fetched ``main`` copy (which lags this branch until merge).""" + original_model_cost = litellm.model_cost + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + litellm.get_model_info.cache_clear() + try: + yield + finally: + litellm.model_cost = original_model_cost + litellm.get_model_info.cache_clear() + + +def test_fable_5_model_pricing_and_capabilities(): + model_data = _load_root_cost_map() + + expected_models = [ + ("claude-fable-5", "anthropic"), + ("anthropic.claude-fable-5", "bedrock_converse"), + ("vertex_ai/claude-fable-5", "vertex_ai-anthropic_models"), + # Unlike Opus 4.8 (200k on Foundry), Fable 5 has the full 1M context + # window on Microsoft Foundry. + ("azure_ai/claude-fable-5", "azure_ai"), + ] + + for model_name, provider in expected_models: + assert model_name in model_data, f"Missing model entry: {model_name}" + info = model_data[model_name] + + assert info["litellm_provider"] == provider + assert info["mode"] == "chat" + assert info["max_input_tokens"] == 1000000 + assert info["max_output_tokens"] == 128000 + assert info["max_tokens"] == 128000 + + # $10 / $50 per MTok (2x Opus 4.8), with the standard 1.25x 5m + # cache-write, 2x 1h cache-write, and 0.1x cache-read multipliers. + assert info["input_cost_per_token"] == 1e-05 + assert info["output_cost_per_token"] == 5e-05 + assert info["cache_creation_input_token_cost"] == 1.25e-05 + assert info["cache_creation_input_token_cost_above_1hr"] == 2e-05 + assert info["cache_read_input_token_cost"] == 1e-06 + + # Flat-rate across the full 1M context window. + assert "input_cost_per_token_above_200k_tokens" not in info + assert "output_cost_per_token_above_200k_tokens" not in info + + assert info["supports_assistant_prefill"] is False + assert info["supports_function_calling"] is True + assert info["supports_prompt_caching"] is True + assert info["supports_reasoning"] is True + assert info["supports_tool_choice"] is True + assert info["supports_vision"] is True + assert info["supports_xhigh_reasoning_effort"] is True + assert info["supports_max_reasoning_effort"] is True + + +def test_fable_5_bedrock_regional_model_pricing(): + model_data = _load_root_cost_map() + + # Fable 5 launched with us/eu geo inference profiles plus a global profile + # (no au/apac/jp). Global uses base pricing; geo profiles carry the + # standard 10% regional premium. + expected_models = { + "global.anthropic.claude-fable-5": { + "input_cost_per_token": 1e-05, + "output_cost_per_token": 5e-05, + "cache_creation_input_token_cost": 1.25e-05, + "cache_read_input_token_cost": 1e-06, + }, + "us.anthropic.claude-fable-5": { + "input_cost_per_token": 1.1e-05, + "output_cost_per_token": 5.5e-05, + "cache_creation_input_token_cost": 1.375e-05, + "cache_read_input_token_cost": 1.1e-06, + }, + "eu.anthropic.claude-fable-5": { + "input_cost_per_token": 1.1e-05, + "output_cost_per_token": 5.5e-05, + "cache_creation_input_token_cost": 1.375e-05, + "cache_read_input_token_cost": 1.1e-06, + }, + } + + for model_name, expected in expected_models.items(): + assert model_name in model_data, f"Missing model entry: {model_name}" + info = model_data[model_name] + assert info["litellm_provider"] == "bedrock_converse" + assert info["max_input_tokens"] == 1000000 + assert info["max_output_tokens"] == 128000 + assert info["bedrock_output_config_effort_ceiling"] == "xhigh" + for key, value in expected.items(): + assert info[key] == value + + +def test_fable_5_geo_multiplier_without_fast_mode(): + """First-party ``inference_geo='us'`` carries the 1.1x premium, but unlike + the Opus line there is no fast-mode variant for Fable 5; a ``fast`` key + here would silently misprice ``speed='fast'`` requests.""" + model_data = _load_root_cost_map() + entry = model_data["claude-fable-5"]["provider_specific_entry"] + assert entry == {"us": 1.1} + + +def test_fable_5_present_in_bundled_backup(): + """The bundled backup is the runtime fallback (and what tests load with + ``LITELLM_LOCAL_MODEL_COST_MAP=True``) — it must carry the same entries as + the root cost map, otherwise the model resolves on one path but not the + other.""" + backup = GetModelCostMap.load_local_model_cost_map() + root = _load_root_cost_map() + for model_name in ( + "claude-fable-5", + "anthropic.claude-fable-5", + "global.anthropic.claude-fable-5", + "us.anthropic.claude-fable-5", + "eu.anthropic.claude-fable-5", + "vertex_ai/claude-fable-5", + "vertex_ai/claude-fable-5@default", + "azure_ai/claude-fable-5", + ): + assert model_name in backup, f"Missing from backup cost map: {model_name}" + assert backup[model_name] == root[model_name], model_name + + +def test_fable_5_registered_for_bedrock_converse(): + assert "anthropic.claude-fable-5" in BEDROCK_CONVERSE_MODELS + + +def test_fable_5_provider_resolves_via_model_info(local_model_cost_map): + info = litellm.get_model_info(model="claude-fable-5") + assert info["litellm_provider"] == "anthropic" + assert info["max_input_tokens"] == 1000000 + assert info["max_output_tokens"] == 128000 + + +@pytest.mark.parametrize( + "cost_map", + [_load_root_cost_map(), GetModelCostMap.load_local_model_cost_map()], + ids=["root", "bundled_backup"], +) +def test_fable_5_all_variants_carry_adaptive_thinking_flag(cost_map): + """Every Fable 5 entry must advertise ``supports_adaptive_thinking``. + + Adaptive-thinking detection is cost-map driven, so a single variant missing + the flag silently sends the legacy ``thinking.type='enabled'`` shape and the + provider 400s (issue #29188 for the Opus 4.8 equivalent). Fable 5 is even + stricter than Opus 4.8: an explicit ``thinking.type='disabled'`` also 400s, + so adaptive is the only valid thinking shape LiteLLM can emit for it.""" + variants = [k for k in cost_map if "claude-fable-5" in k] + assert variants, "no claude-fable-5 entries found in cost map" + missing = [ + k for k in variants if cost_map[k].get("supports_adaptive_thinking") is not True + ] + assert not missing, f"missing supports_adaptive_thinking: {missing}" + + +@pytest.mark.parametrize( + "model", + [ + "claude-fable-5", + "anthropic/claude-fable-5", + "anthropic.claude-fable-5", + "bedrock/us.anthropic.claude-fable-5", + "bedrock/invoke/eu.anthropic.claude-fable-5", + "bedrock/global.anthropic.claude-fable-5", + "vertex_ai/claude-fable-5", + "azure_ai/claude-fable-5", + ], +) +def test_adaptive_thinking_detected_for_fable_5(local_model_cost_map, model): + """Provider-routed ids must resolve to a flagged entry so ``reasoning_effort`` + maps to ``thinking.type='adaptive'`` + ``output_config.effort``.""" + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + assert AnthropicModelInfo._is_adaptive_thinking_model(model) is True + + +@pytest.mark.parametrize( + "cost_map", + [_load_root_cost_map(), GetModelCostMap.load_local_model_cost_map()], + ids=["root", "bundled_backup"], +) +def test_sampling_params_flag_on_all_models_that_removed_them(cost_map): + """Fable 5 and Opus 4.7/4.8 reject ``top_p``/``top_k``/``temperature != 1``; + the drop/raise gating is cost-map driven, so every variant must carry an + explicit ``supports_sampling_params: false``. The perplexity route is + exempt: it is OpenAI-compatible and maps sampling params upstream.""" + variants = [ + k + for k in cost_map + if any(v in k for v in ("claude-fable-5", "claude-opus-4-7", "claude-opus-4-8")) + and not k.startswith("perplexity/") + ] + assert variants, "no matching entries found in cost map" + missing = [ + k for k in variants if cost_map[k].get("supports_sampling_params") is not False + ] + assert not missing, f"missing supports_sampling_params=false: {missing}" diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index f179e9c8f93..4c4d9e1133b 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -858,6 +858,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "supports_xhigh_reasoning_effort": {"type": "boolean"}, "supports_max_reasoning_effort": {"type": "boolean"}, "supports_adaptive_thinking": {"type": "boolean"}, + "supports_sampling_params": {"type": "boolean"}, "supports_service_tier": {"type": "boolean"}, "supports_preset": {"type": "boolean"}, "supports_output_config": {"type": "boolean"}, From 2fe9feda71d7e3d397579b272a45d8c26902e9a3 Mon Sep 17 00:00:00 2001 From: michelligabriele Date: Wed, 10 Jun 2026 12:19:20 +0200 Subject: [PATCH 044/185] fix(caching): restore stored prompt_tokens on embedding cache hits instead of recomputing (#30046) --- litellm/caching/caching.py | 34 ++++- litellm/caching/caching_handler.py | 7 +- litellm/types/caching.py | 1 + tests/test_litellm/caching/test_caching.py | 32 ++++- .../caching/test_caching_handler.py | 120 ++++++++++++++++++ 5 files changed, 190 insertions(+), 4 deletions(-) diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index c1afde16250..b6cfc8e7907 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -691,6 +691,7 @@ class Cache: self, embedding_response: Any, model: Optional[str], + prompt_tokens: Optional[int] = None, prompt_tokens_details: Optional[dict] = None, ) -> CachedEmbedding: """ @@ -703,6 +704,7 @@ class Cache: "index": embedding_response.get("index"), "object": embedding_response.get("object"), "model": model, + "prompt_tokens": prompt_tokens, "prompt_tokens_details": prompt_tokens_details, } elif hasattr(embedding_response, "model_dump"): @@ -712,6 +714,7 @@ class Cache: "index": data.get("index"), "object": data.get("object"), "model": model, + "prompt_tokens": prompt_tokens, "prompt_tokens_details": prompt_tokens_details, } else: @@ -721,6 +724,7 @@ class Cache: "index": data.get("index"), "object": data.get("object"), "model": model, + "prompt_tokens": prompt_tokens, "prompt_tokens_details": prompt_tokens_details, } except KeyError as e: @@ -769,6 +773,29 @@ class Cache: per_item[key] = value return per_item if per_item else None + def _get_per_item_prompt_tokens( + self, + result: EmbeddingResponse, + idx_in_result_data: int, + ) -> Optional[int]: + """ + Extract the per-item prompt_tokens from a response for caching. + + Single-item responses store the full usage.prompt_tokens. Multi-item + responses distribute it evenly (with remainder) so that summing all + per-item values on retrieval reconstructs the original total. + """ + if result.usage is None or result.usage.prompt_tokens is None: + return None + + total = result.usage.prompt_tokens + num_items = len(result.data) + if num_items <= 1: + return total + + quotient, remainder = divmod(total, num_items) + return quotient + (1 if idx_in_result_data < remainder else 0) + def add_embedding_response_to_cache( self, result: EmbeddingResponse, @@ -780,7 +807,11 @@ class Cache: kwargs["cache_key"] = preset_cache_key embedding_response = result.data[idx_in_result_data] - # Extract per-item prompt_tokens_details from response usage + # Extract per-item prompt_tokens + details from response usage + prompt_tokens = self._get_per_item_prompt_tokens( + result=result, + idx_in_result_data=idx_in_result_data, + ) prompt_tokens_details = self._get_per_item_prompt_tokens_details( result=result, idx_in_result_data=idx_in_result_data, @@ -791,6 +822,7 @@ class Cache: embedding_dict: CachedEmbedding = self._convert_to_cached_embedding( embedding_response, model_name, + prompt_tokens=prompt_tokens, prompt_tokens_details=prompt_tokens_details, ) diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index 3f4e54382c9..48691335b40 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -394,7 +394,7 @@ class LLMCachingHandler: return cr["model"] return None - def _process_async_embedding_cached_response( + def _process_async_embedding_cached_response( # noqa: PLR0915 self, final_embedding_cached_response: Optional[EmbeddingResponse], cached_result: List[Optional[CachedEmbedding]], @@ -456,7 +456,10 @@ class LLMCachingHandler: index=idx, object="embedding", ) - if isinstance(kwargs_input_as_list[idx], str): + cached_prompt_tokens = cr.get("prompt_tokens") + if cached_prompt_tokens is not None: + prompt_tokens += cached_prompt_tokens + elif isinstance(kwargs_input_as_list[idx], str): from litellm.utils import token_counter prompt_tokens += token_counter( diff --git a/litellm/types/caching.py b/litellm/types/caching.py index f8050b292c7..10453c74a15 100644 --- a/litellm/types/caching.py +++ b/litellm/types/caching.py @@ -118,4 +118,5 @@ class CachedEmbedding(TypedDict): index: Optional[int] object: Optional[str] model: Optional[str] + prompt_tokens: Optional[int] prompt_tokens_details: Optional[dict] diff --git a/tests/test_litellm/caching/test_caching.py b/tests/test_litellm/caching/test_caching.py index 02d62a19152..20614103ed2 100644 --- a/tests/test_litellm/caching/test_caching.py +++ b/tests/test_litellm/caching/test_caching.py @@ -3,6 +3,7 @@ import re from litellm.caching.caching import Cache from litellm.types.caching import LiteLLMCacheType +from litellm.types.utils import Embedding, EmbeddingResponse, Usage def test_cache_key_debug_log_does_not_include_prompt_material(caplog): @@ -41,8 +42,37 @@ def test_cache_key_debug_log_does_not_include_prompt_material(caplog): assert re.fullmatch(r"[0-9a-f]{64}", cache_key) created_cache_key_logs = [ - record.getMessage() for record in caplog.records if "Created cache key:" in record.getMessage() + record.getMessage() + for record in caplog.records + if "Created cache key:" in record.getMessage() ] assert created_cache_key_logs assert all(prompt_marker not in message for message in created_cache_key_logs) assert any(cache_key in message for message in created_cache_key_logs) + + +def _embedding_response(prompt_tokens, num_items): + return EmbeddingResponse( + model="amazon.titan-embed-image-v1", + data=[ + Embedding(embedding=[0.0], index=i, object="embedding") + for i in range(num_items) + ], + usage=Usage( + prompt_tokens=prompt_tokens, completion_tokens=0, total_tokens=prompt_tokens + ), + ) + + +def test_get_per_item_prompt_tokens_single_item_returns_full_value(): + cache = Cache(type=LiteLLMCacheType.LOCAL) + result = _embedding_response(prompt_tokens=0, num_items=1) + assert cache._get_per_item_prompt_tokens(result, 0) == 0 + + +def test_get_per_item_prompt_tokens_distributes_with_remainder(): + cache = Cache(type=LiteLLMCacheType.LOCAL) + result = _embedding_response(prompt_tokens=10, num_items=3) + per_item = [cache._get_per_item_prompt_tokens(result, i) for i in range(3)] + assert sum(per_item) == 10 # 4 + 3 + 3 + assert per_item == [4, 3, 3] diff --git a/tests/test_litellm/caching/test_caching_handler.py b/tests/test_litellm/caching/test_caching_handler.py index 3eb949d7f29..01327529410 100644 --- a/tests/test_litellm/caching/test_caching_handler.py +++ b/tests/test_litellm/caching/test_caching_handler.py @@ -436,3 +436,123 @@ def test_convert_cached_responses_legacy_stream_path(): ) assert isinstance(result, CachedResponsesAPIStreamingIterator) + + +@pytest.mark.asyncio +async def test_embedding_cache_restores_stored_prompt_tokens_for_image_input(): + """Image-embedding cache hit restores prompt_tokens=0 from the stored value + instead of recomputing a bogus count by tokenizing the base64 input.""" + llm_caching_handler = LLMCachingHandler( + original_function=MagicMock(), + request_kwargs={}, + start_time=datetime.now(), + ) + + # base64-like blob — token_counter over this would return a large nonzero count + image_input = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk" * 50 + + cached_result = [ + { + "embedding": [-0.025, -0.019], + "index": 0, + "object": "embedding", + "model": "amazon.titan-embed-image-v1", + "prompt_tokens": 0, + "prompt_tokens_details": {"image_count": 1}, + } + ] + + mock_logging_obj = MagicMock() + mock_logging_obj.async_success_handler = AsyncMock() + response, cache_hit = llm_caching_handler._process_async_embedding_cached_response( + final_embedding_cached_response=None, + cached_result=cached_result, + kwargs={"model": "amazon.titan-embed-image-v1", "input": image_input}, + logging_obj=mock_logging_obj, + start_time=datetime.now(), + model="amazon.titan-embed-image-v1", + ) + + assert cache_hit + assert response.usage is not None + assert response.usage.prompt_tokens == 0 + assert response.usage.total_tokens == 0 + assert response.usage.prompt_tokens_details.image_count == 1 + + +@pytest.mark.asyncio +async def test_embedding_cache_sums_stored_prompt_tokens_across_items(): + """A multi-item cache hit sums the stored per-item prompt_tokens back to the total.""" + llm_caching_handler = LLMCachingHandler( + original_function=MagicMock(), + request_kwargs={}, + start_time=datetime.now(), + ) + + cached_result = [ + { + "embedding": [-0.01], + "index": 0, + "object": "embedding", + "model": "text-embedding-3-small", + "prompt_tokens": 5, + }, + { + "embedding": [-0.02], + "index": 1, + "object": "embedding", + "model": "text-embedding-3-small", + "prompt_tokens": 4, + }, + ] + + mock_logging_obj = MagicMock() + mock_logging_obj.async_success_handler = AsyncMock() + response, cache_hit = llm_caching_handler._process_async_embedding_cached_response( + final_embedding_cached_response=None, + cached_result=cached_result, + kwargs={"model": "text-embedding-3-small", "input": ["hello world", "foo bar"]}, + logging_obj=mock_logging_obj, + start_time=datetime.now(), + model="text-embedding-3-small", + ) + + assert cache_hit + assert response.usage.prompt_tokens == 9 + assert response.usage.total_tokens == 9 + + +@pytest.mark.asyncio +async def test_embedding_cache_falls_back_to_token_counter_for_legacy_entries(): + """Legacy cache entries with no stored prompt_tokens still recompute via token_counter + for str inputs (backward compatibility).""" + llm_caching_handler = LLMCachingHandler( + original_function=MagicMock(), + request_kwargs={}, + start_time=datetime.now(), + ) + + # No prompt_tokens key — pre-fix entry + cached_result = [ + { + "embedding": [-0.025, -0.019], + "index": 0, + "object": "embedding", + "model": "text-embedding-ada-002", + }, + ] + + mock_logging_obj = MagicMock() + mock_logging_obj.async_success_handler = AsyncMock() + response, cache_hit = llm_caching_handler._process_async_embedding_cached_response( + final_embedding_cached_response=None, + cached_result=cached_result, + kwargs={"model": "text-embedding-ada-002", "input": "hello world"}, + logging_obj=mock_logging_obj, + start_time=datetime.now(), + model="text-embedding-ada-002", + ) + + assert cache_hit + # token_counter over "hello world" yields a nonzero count — fallback path still runs + assert response.usage.prompt_tokens > 0 From 3b40ac987fb4fe08061b67dda91b286dc41bee28 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 10 Jun 2026 23:04:07 +0530 Subject: [PATCH 045/185] Litellm oss 090626 (#30021) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(mcp): report scoped server name during initialize (#29865) * fix mcp scoped server name * Update litellm/proxy/_experimental/mcp_server/mcp_context.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * test(mcp): cover scoped server name in the SSE initialize handler --------- Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * fix(ui): show all session logs in the drawer, not just the first 50 (#29795) * fix(ui): show newest session logs first * test(ui): keep session log pagination coverage * fix(ui): show all session logs in the drawer, not just the first page The session detail drawer fetched session logs via sessionSpendLogsCall without page/page_size, so it only ever received the backend default of one page (50 rows). Sessions with more than 50 calls had the rest unreachable in the UI (#29153). sessionSpendLogsCall now takes page/page_size, and the drawer fetches the first page, reads total_pages, then fetches the remaining pages and accumulates them before the existing client-side sort. This keeps the single continuous list (and the selected-log lookup and keyboard navigation, which all assume the full session) correct. Fetching is bounded by a page cap, and the sidebar shows a "showing most recent N" note if a session exceeds it. The rows are lightweight metadata (the endpoint excludes messages/response), so the full set is small; request/response bodies are still loaded per log on demand. * fix(ui): default session drawer to most recent log, newest first Open a session with its most recent log selected, and order the sidebar newest-first to match the all-sessions logs overview. MCP calls stay grouped last. The latest log by time is computed explicitly, since the MCP grouping means it is not always the first row. * Apply fetching pages in batches suggestion from @greptile-apps[bot] Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * fix(ui): derive session total from accumulated rows when backend omits it Compute the session total after all pages are fetched, falling back to the accumulated row count rather than the first page's. Guards the truncation note against a backend response that omits total but spans multiple pages. --------- Co-authored-by: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * fix(proxy): handle Mistral multipart passthrough (#29927) * fix(proxy): handle Mistral multipart passthrough * chore: satisfy passthrough ci formatting * test(proxy): cover Mistral passthrough in CI shard * fix(vertex_ai): use REP host for context caching on eu/us multi-region endpoints (#29573) Context caching built the cachedContents URL as https://{location}-aiplatform.googleapis.com, which is an invalid host for the eu/us multi-region endpoints and returns 404. The inference path already resolves these to the REP host (https://aiplatform.{geo}.rep.googleapis.com) via get_vertex_base_url(); reuse that helper in _get_token_and_url_context_caching so caching uses the same host as inference. Adds tests covering the eu/us multi-region cachedContents URLs (v1 and v1beta1). Fixes #29571 * Support per-model encrypted content affinity config (#29760) Co-authored-by: shin-berri Co-authored-by: yuneng-jiang * fix: propagate upstream status code in proxy API exception handler (#29402) * fix: propagate upstream status code in proxy API exception handler When Google GenAI / Vertex returns a 404 for deprecated or missing models via streamGenerateContent, the exception was falling through to a generic handler that defaulted to 500. Now provider exceptions carrying a valid HTTP status_code correctly propagate it through to the ProxyException. * fix: apply black formatting to common_request_processing.py * fix: tighten status code range to 400-599 and deduplicate ProxyException raise * fix(tests): use valid vertex_location in context caching tests Replace "test_location" (contains underscore) with "us-central1" so tests pass the regex validation added in get_vertex_base_url(). Co-Authored-By: Claude Sonnet 4.6 * feat(sdk): add xAI OAuth provider (#29866) * Add xAI OAuth provider * Update oauth.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * Fix xAI OAuth CI failures * Add xAI OAuth coverage tests * Move xAI OAuth coverage tests to core utils * Address xAI OAuth review comments * Prevent xAI OAuth api_base token exfiltration * Treat blank xAI OAuth api keys as absent * Wrap invalid xAI OAuth JSON responses * Use xAI OAuth behind explicit flag --------- Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * fix(proxy) #27734 allow clearing budget_duration and team_member fields by sending null on /key/update and /team/update (#27751) * fix(proxy): allow clearing budget_duration and team_member fields by sending null on /key/update and /team/update Fixes #27734 Sending null for budget_duration, team_member_budget, team_member_budget_duration, team_member_rpm_limit, or team_member_tpm_limit via /key/update or /team/update returned 200 OK but silently ignored the null value. The fields remained unchanged in the database. Root causes: - /key/update: prepare_key_update_data() popped budget_duration from the update dict but never re-added it (or budget_reset_at) when the value was None. - /team/update: _set_budget_reset_at() only acted when budget_duration was non-None, leaving a stale budget_reset_at in the DB. - /team/update: team_member_* null values bypassed the budget table update entirely because should_create_budget() requires at least one non-None field. * test(proxy): cover no-budget-row path in clear_team_member_budget_fields * fix(presidio): unmask PII tokens in Anthropic native SSE streaming bytes (#30028) * fix(presidio): unmask PII tokens in Anthropic native SSE streaming bytes When output_parse_pii=true on the Anthropic native path (anthropic/claude-*), response chunks arrive as raw bytes in SSE format. _stream_pii_unmasking was yielding those bytes unchanged, so tokens were never replaced with the original values before reaching the caller. Add _unmask_sse_bytes_chunk to parse each data: line, find content_block_delta / text_delta events, and apply _unmask_pii_text before re-encoding. Wire it into _stream_pii_unmasking so bytes chunks are unmasked when pii_tokens exist. * fix(presidio): handle CRLF line endings and non-ASCII PII in SSE unmask Strip trailing \r before the [DONE] guard so CRLF-terminated SSE chunks don't bypass it and silently swallow a JSONDecodeError. Add ensure_ascii=False to json.dumps so non-ASCII replacement values like accented names are preserved as UTF-8 on the wire rather than being \uXXXX-escaped. Add regression tests for both cases. * feat(bedrock_mantle): path-aware Responses routing (/v1/responses vs /openai/v1/responses) (#29925) * feat(bedrock_mantle): path-aware Responses routing (/v1/responses vs /openai/v1/responses) Bedrock Mantle serves the Responses API on two upstream paths: - gpt frontier models (gpt-5.5 / gpt-5.4) on /openai/v1/responses - every other Responses-capable model (e.g. gpt-oss) on the standard /v1/responses BedrockMantleResponsesAPIConfig gains a `use_openai_path` flag; the provider gate in utils.py picks the path per model: openai.gpt-* (non gpt-oss) -> /openai/v1/responses; any model declared mode=responses (price-map entry or user model_info) -> /v1/responses; everything else returns None and keeps the existing chat-completions emulation. Adds gpt-5.5 / gpt-5.4 price-map entries, registry wiring, and the routing-matrix tests. * feat(bedrock_mantle): data-driven frontier routing via use_openai_responses_path Addresses the Greptile review point that frontier detection should be a price-map field rather than a hardcoded name match. The gate now routes a model to /openai/v1/responses when its price-map entry declares use_openai_responses_path, so a frontier model whose name does not follow the openai.gpt- convention can be onboarded by JSON alone. The name-convention check is kept as a fallback that needs no price-map entry, which preserves zero-change routing for a future gpt-6 before its entry loads. gpt-5.5 / gpt-5.4 get the flag in both price maps. Adds tests for the data-driven flag path and for the flag presence on the gpt-5.x entries; both branches are mutation-tested. * test(model_prices): allow use_openai_responses_path in price-map schema The model_prices_and_context_window.json schema validator (test_aaamodel_prices_and_context_window_json_is_valid) enforces additionalProperties: false, so the new use_openai_responses_path flag on the gpt-5.5 / gpt-5.4 entries failed validation. Add it to the schema as a boolean, alongside the other supports_* / capability flags. * Add Tensormesh serverless models to the model cost map (#30037) * Add Tensormesh serverless models to the model cost map * Flag reasoning support on the Tensormesh models that expose thinking mode * fix(proxy): invalidate stale key spend counter after budget reset or manual spend update (#30001) * fix(proxy): reconcile stale key spend counter after budget reset * fix(proxy): invalidate stale key spend counter after budget reset or manual spend update * fix(proxy): remove read-time stale counter reconciliation to prevent budget bypass * revert: undo unrelated formatting changes in enterprise directory * test(proxy): add unit test for key spend update invalidating counter * test(proxy): fix mocked update_data and hash token expectations in unit test * fix(proxy): use Responses-API transformer in pass-through cost tracking (#29728) The `elif is_responses:` branch of `openai_passthrough_handler` was calling the chat-completions `transform_response` on a Responses API payload. The chat-completions transformer expects `choices: [...]` in the raw response; the Responses API uses `output: [...]` and `usage.input_tokens` / `usage.output_tokens` (not `prompt_tokens` / `completion_tokens`). The result was a KeyError 'choices' deep inside `convert_to_model_response_object`, swallowed by the surrounding `except Exception` in the handler, and the SpendLogs row was written by the fallback path with zeroed-out tokens, spend, and model. This bug silently undercounts cost for every successful pass-through call to either OpenAI's `/v1/responses` or Azure's `/openai/v1/responses` (deployments configured for the Responses API). Reproduced 2026-06-04 against a real Azure OpenAI Responses API deployment proxied through LiteLLM v1.88.0. Fix: use the dedicated `OpenAIResponsesAPIConfig.transform_response_api_response` for the Responses branch. This transformer already exists in LiteLLM (`litellm/llms/openai/responses/transformation.py`) and knows the Responses-API on-the-wire shape. `litellm.completion_cost` already handles `ResponsesAPIResponse` natively with `call_type="responses"`, so no downstream changes are needed. Tests: test_responses_api_uses_responses_transformer_not_chat_completions NEW. Real regression test — exercises the openai_passthrough_handler with a real-shaped Responses payload (no `choices`, has `output` and Responses-API `usage` keys) and NO mocked `get_provider_config`. Pre-fix: raises KeyError 'choices' inside the chat-completions transformer (the bug). Post-fix: returns a ResponsesAPIResponse, completion_cost is called with call_type="responses" and a ResponsesAPIResponse instance (asserted). Verified to fail on un-fixed handler + pass on fixed handler before commit. test_responses_api_cost_tracking UPDATED. Old test mocked `get_provider_config` (no longer called in the responses branch post-fix). Now mocks the Responses transformer directly (`OpenAIResponsesAPIConfig.transform_response_api_response`) to test the downstream cost-calc contract. Out of scope for this PR (separate followup): - Recognizing *.cognitiveservices.azure.com (the newer Azure OpenAI hostname) in the is_openai_*_route checks. Separate PR. Co-authored-by: shin-berri Co-authored-by: yuneng-jiang * fix(skills): execute DB skills by matching the litellm_skill_ tool name prefix (#30116) Skill IDs are generated as litellm_skill_ and the model-facing tool name is the sanitized skill ID, but the post-call execution gates in SkillsInjectionHook only ran tools whose name starts with "skill_", so DB skills were silently returned to the client as raw tool calls. Fixes #28122. Co-authored-by: Cursor * fix(anthropic): synthesize content_block_start when Responses stream omits output_item.added (#30115) * fix(team): reserve team budget raises for proxy admins on /team/update (#30030) The caller's PERSONAL max_budget was the wrong yardstick for /team/update: a team's spend ceiling has nothing to do with the admin's own key budget. That comparison was an unintended side effect of reusing _check_user_team_limits() (which exists for the /team/new path) and broke the UI, which re-sends the unchanged budget on every save. New behavior on /team/update for standalone teams: - A team admin (already authorized via _verify_team_access) may freely KEEP or LOWER the team budget, and change models/tpm/rpm, without being gated by their personal limits. - GROWING a team's spend ceiling is a budget-authority action reserved for proxy admins -> 403 for team admins. "Growing" covers both raising max_budget above the team's current finite value and removing the cap entirely (max_budget=null, detected via model_fields_set so an explicit null is distinguished from an omitted field). For a team that currently has no cap, setting a finite value is a restriction and is allowed. - Org-scoped teams remain governed by _check_org_team_limits() (capped by the org budget). Also reverts the #29525 existing_team_max_budget workaround in _check_user_team_limits() back to the create-only form; /team/new still enforces the creator's personal caps. docs(access_control): resolve the contradiction in the team-admin section — team admins can keep/lower the budget and manage rate limits/models, but cannot raise the team budget (proxy-admin only). tests: unit + behavior coverage for raise-blocked, cap-removal-blocked (team admin), raise/removal allowed (proxy admin), uncapped-team restriction allowed, keep/lower/resend allowed, and unchanged create-path guards. Co-authored-by: Cursor * test(ui): data-driven App Router migration E2E smoke (default + server-root-path) (#29974) * test(ui): add a data-driven App Router migration E2E smoke Add a growing Playwright smoke for migrated pages: for each segment it deep-links to the path route, asserts the URL and that the dashboard shell rendered, then clicks off to a legacy page and asserts navigation still works. Driven by e2e_tests/fixtures/migratedPages.ts, so adding a page is one line. Runs in two situations against the same proxy: the default mount (npm run e2e:migration) and a non-root SERVER_ROOT_PATH mount (npm run e2e:migration:root). globalSetup now logs in at `${SERVER_ROOT_PATH}/ui/login` so the admin storage state is valid under a prefix. Seeded with api-reference; append the rest as their migrations merge. * test(ui): support headed slow-motion + watch pauses in the migration smoke Honor SLOWMO in the server-root-path config (the default config already did), and add an env-gated E2E_WATCH_MS pause so a headed run lingers on each state. Both are no-ops by default, so CI behavior is unchanged. * test(ui): make the migration smoke a sidebar-click user journey Rework the smoke from deep-linking to a real navigation journey: start at the landing page, click the migrated page in the sidebar (expanding submenus for nested items), assert the path route rendered, reload it (the check a wrong server_root_path breaks), bounce to a legacy page and back, and — once two pages are migrated — navigate directly between two migrated pages. Verifies via URL + shell render, driven by the same fixture list. * test(ui): address review on the migration smoke Escape ROOT and segment before interpolating them into RegExp URL matchers so a future segment containing regex metacharacters can't silently widen the match. Make the server-root-path config fail fast when SERVER_ROOT_PATH is unset instead of silently re-running the default mount and passing without exercising the prefix. * test(ui): drop unused watch helper and fix stale smoke README * test(ui): run the migration smoke under a server root path in CI * test(ui): harden + instrument the server-root-path proxy reboot in CI * test(ui): run the server-root-path migration smoke as its own CI job Replace the in-place proxy reboot in e2e_ui_testing with a dedicated e2e_ui_testing_server_root_path job that boots the proxy once with SERVER_ROOT_PATH=/litellm, matching how every other proxy variant in the config gets its own job rather than killing and relaunching the live proxy. The reboot was failing deterministically: after pkill -9 and relaunch the prefixed proxy never came back up on :4000 (connection refused), so the smoke never ran. The readiness step that was supposed to surface the cause could never reach its boot-log tail because CircleCI runs steps under bash -eo pipefail and the preceding `curl -sv ... | tail` aborted the step with curl's exit 7. Booting the proxy as the job's own background step lets any boot crash land in that step's log instead of being swallowed. The default e2e_ui_testing job is unchanged aside from dropping the reboot, prefixed-readiness, and prefixed-smoke steps; the migration smoke still runs at the root mount there via the default Playwright config. * fix(proxy): extend response headers hook to streaming, TTS, image gen, and pass-through (#24232) * fix(proxy): extend response headers hook to streaming, TTS, image gen, and pass-through * test: mock post_call_response_headers_hook in audio speech route tests * chore(ui): remove dead App Router route stubs under (dashboard) (#30045) models-and-endpoints, organizations, and virtual-keys each had a page.tsx route under (dashboard)/ that is not in MIGRATED_PAGES, so the sidebar and deep links never resolve to it and the route is unreachable. Each was a thin wrapper that handed the shared view empty or no-op props (empty modelData with a no-op setModelData, hardcoded empty organizations, no-op setUserRole/setUserEmail), so reaching one would render a degraded page in any case. The real wrapper belongs in the PR that flips each page into MIGRATED_PAGES, written with eyes on it and a test This continues the dead-scaffolding cleanup from #28891. The shared components these wrappers rendered (ModelsAndEndpointsView, OrganizationFilters) stay, since the legacy ?page= switch in app/page.tsx and src/components still import them * fix(ui/mcp): reset OAuth state on create-server modal close so a prior server's token no longer leaks into the next add-server session (#30000) * fix(ui/mcp): reset OAuth hook state on modal close so a prior server's token no longer leaks into the next add-server session * fix(ui/mcp): clear in-flight OAuth guard on reset and reset form/tools on modal close so nothing leaks on a parent-driven dismiss * fix(mcp): allow team access-group grants in OAuth authorize/token access check (#30041) * fix(mcp): honor team access-group grants in OAuth authorize/token access check * test(mcp): mock build_effective_auth_contexts in non-admin authorize tests for isolation * docs(security): require a reproduction video for vulnerability reports (#30048) (#30063) With AI models capable of automated vulnerability discovery now publicly available, we expect a large increase in report volume, much of it unverified. Requiring a video of the exploit running against a live instance raises the bar for submissions and keeps triage focused on reproducible issues. Reports without a video will be closed and reopened if one is added later. Co-authored-by: stuxf <70670632+stuxf@users.noreply.github.com> * feat(ui): add admin flag to disable in-product UI nudges for everyone (#29796) * feat(ui): add admin flag to disable in-product UI nudges for everyone Admins can now suppress the survey and Claude Code feedback popups for all users via a single disable_ui_nudges UI setting, instead of relying on each user dismissing them individually. * fix(ui): suppress nudges while ui settings are loading Gate nudgesDisabled on the ui-settings loading state so an admin with disable_ui_nudges on doesn't see the survey prompt flash, and the getInProductNudgesCall fetch doesn't fire, on a cold page load before the flag resolves. Falls back to showing nudges if the fetch errors. * test(ui): wrap CreateKeyPage test in QueryClientProvider page.tsx now calls useUISettings (react-query), which needs a QueryClient that layout.tsx supplies in production but the test did not. Add the provider and mock getUiSettings so the query resolves. * chore(ui): remove dead dashboard files and unused dependencies (#30047) * chore(ui): remove dead dashboard files and unused dependencies knip flagged seven orphaned source/config files with no importers and five declared dependencies that nothing in the tree uses. Removing them shrinks the dashboard bundle's source surface and keeps the manifest honest; vite stays installed transitively via vitest, so test tooling is unaffected. * fix(ci): restore serverRootPath.config.ts referenced by SERVER_ROOT_PATH workflow The dead-code sweep removed e2e_tests/serverRootPath.config.ts, but its spec (tests/login/serverRootPathRedirect.spec.ts) and the test_server_root_path.yml workflow step still depend on it, so the redirect e2e job failed to load a config that no longer existed. * fix(proxy): authorize batch files using upload target_model_names (LIT-3593) (#30009) * fix(proxy): authorize batch files using upload target_model_names (LIT-3593) After replace_model_in_jsonl, body.model is a stripped provider id. Reverse-mapping it via resolve_model_name_from_model_id is first-match on model_list and caused false 403s when multiple deployments share the same stripped name. Use target_model_names from the unified file id instead. Co-authored-by: Cursor * fix(proxy): restore resolve_model_name_from_model_id for JSONL fallback path (LIT-3593) Restores the reverse-lookup for the JSONL body.model fallback path so that legacy/pre-target_model_names managed files still map stripped provider IDs back to proxy aliases before auth. Also cleans up redundant `or None`. Co-Authored-By: Claude Sonnet 4.6 * Revert "fix(proxy): restore resolve_model_name_from_model_id for JSONL fallback path (LIT-3593)" This reverts commit 30d2e96f77ef521ccaaf2193fe554980380eb669. --------- Co-authored-by: Cursor Co-authored-by: Claude Sonnet 4.6 * Add Claude Fable 5 across Anthropic, Bedrock, Vertex AI, and Azure AI (#30064) * Add Claude Fable 5 across Anthropic, Bedrock, Vertex AI, and Azure AI Adds cost map entries for claude-fable-5 ($10/$50 per MTok, 1M context, 128K output, adaptive thinking only) on the Anthropic API, Bedrock converse (base, global, and us/eu geo inference profiles at the 10% regional premium), Vertex AI, and Azure AI (Microsoft Foundry, which serves Fable 5 with the full 1M context window unlike Opus 4.8). Registers anthropic.claude-fable-5 in BEDROCK_CONVERSE_MODELS, lists the model in the setup wizard, and extends the reasoning effort e2e grid. The Bedrock, Vertex, and Azure grid cells carry fail_reason markers until the CI accounts are provisioned: Bedrock needs the provider data sharing opt-in Fable 5 requires, and the Foundry resource needs a claude-fable-5 deployment. The first-party entry carries provider_specific_entry {us: 1.1} for the inference_geo premium and deliberately no fast multiplier since Fable 5 has no fast mode. https://claude.ai/code/session_01MZarYYT3aS7DxaNjoax6Gm * Drop removed sampling params for Claude 4.7+ when drop_params is set Fable 5, Opus 4.7, and Opus 4.8 removed sampling params: the API rejects top_p, top_k, and any temperature other than 1 with a 400. LiteLLM was forwarding them even with drop_params enabled because the Anthropic and Bedrock converse transformations passed temperature/top_p through unconditionally. Mirror the GPT-5/o-series handling: temperature=1 still passes through, other values and any top_p are dropped when drop_params is set, and without drop_params a clean client-side UnsupportedParamsError tells the caller how to opt in, instead of surfacing the raw provider error. https://claude.ai/code/session_01MZarYYT3aS7DxaNjoax6Gm * Drive sampling param gating from the cost map and cover top_k Greptile review follow-ups on the sampling param fix: the restriction for Fable 5 / Opus 4.7 / 4.8 is now declared as supports_sampling_params: false on every affected cost map entry (perplexity excluded; that route is OpenAI-compatible and maps sampling params upstream) and read back through a tri-state map lookup, keeping the name check only as a fallback for provider-routed ids whose hosted map entries predate the flag, the same layering supports_adaptive_thinking uses. top_k bypasses map_openai_params as a provider-specific kwarg, so it is gated at the shared AnthropicConfig.transform_request boundary (direct, Bedrock invoke, Vertex, Azure) and in the Bedrock converse _handle_top_k_value path, with drop_params threaded through the converse transform helpers. Also updates the reasoning effort grid cell count assertion for the four Fable 5 rows added on this branch (29 x 11 cells). https://claude.ai/code/session_01MZarYYT3aS7DxaNjoax6Gm * Declare supports_sampling_params in the cost map schema The model map validation schema uses additionalProperties: false, so the new flag must be declared for the 28 entries that carry it; this was the one failing job (misc / Run tests) on the previous commit. https://claude.ai/code/session_01MZarYYT3aS7DxaNjoax6Gm * fix(bedrock): gate top_k=0 on converse to match Anthropic boundary Truthiness check let top_k=0 silently disappear on models that removed sampling params, while AnthropicConfig.transform_request treats 0 as present and raises UnsupportedParamsError (or drops when drop_params is set). Switch to 'is not None' so converse, direct Anthropic, invoke, Vertex, and Azure all behave the same for top_k=0. --------- Co-authored-by: Cursor Agent * fix(anthropic): avoid index -1 content_block_delta in messages stream When a /v1/messages request is routed through the Responses API adapter, AnthropicResponsesStreamWrapper only emits content_block_start on response.output_item.added. Some upstreams (LMStudio for example) never send that event, so the text delta handler fell back to _current_block_index, which starts at -1, and clients received content_block_delta events with index -1 and no preceding content_block_start. Anthropic SDKs then fail with "text part -1 not found" The text delta handler now synthesizes a content_block_start with a fresh block index whenever the delta references an unregistered item_id or no block is open yet, and registers the item_id so follow-up deltas reuse the same index Addresses the /v1/messages defect in #27442 * Make test sys.path shim resolve relative to the file, not the CWD os.path.abspath("../../../../../../..") depends on where pytest is invoked from; anchoring on os.path.dirname(__file__) makes the import work from any working directory. Also corrects the depth: the repo root is six levels above this file, not seven. --------- Co-authored-by: milan-berri Co-authored-by: Cursor Co-authored-by: ryan-crabbe-berri Co-authored-by: michelligabriele Co-authored-by: tin-berri Co-authored-by: yuneng-jiang Co-authored-by: stuxf <70670632+stuxf@users.noreply.github.com> Co-authored-by: Sameer Kankute Co-authored-by: Claude Sonnet 4.6 Co-authored-by: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> * fix: enable compact-2026-01-12 beta header for vertex_ai provider (#30114) * fix(team): reserve team budget raises for proxy admins on /team/update (#30030) The caller's PERSONAL max_budget was the wrong yardstick for /team/update: a team's spend ceiling has nothing to do with the admin's own key budget. That comparison was an unintended side effect of reusing _check_user_team_limits() (which exists for the /team/new path) and broke the UI, which re-sends the unchanged budget on every save. New behavior on /team/update for standalone teams: - A team admin (already authorized via _verify_team_access) may freely KEEP or LOWER the team budget, and change models/tpm/rpm, without being gated by their personal limits. - GROWING a team's spend ceiling is a budget-authority action reserved for proxy admins -> 403 for team admins. "Growing" covers both raising max_budget above the team's current finite value and removing the cap entirely (max_budget=null, detected via model_fields_set so an explicit null is distinguished from an omitted field). For a team that currently has no cap, setting a finite value is a restriction and is allowed. - Org-scoped teams remain governed by _check_org_team_limits() (capped by the org budget). Also reverts the #29525 existing_team_max_budget workaround in _check_user_team_limits() back to the create-only form; /team/new still enforces the creator's personal caps. docs(access_control): resolve the contradiction in the team-admin section — team admins can keep/lower the budget and manage rate limits/models, but cannot raise the team budget (proxy-admin only). tests: unit + behavior coverage for raise-blocked, cap-removal-blocked (team admin), raise/removal allowed (proxy admin), uncapped-team restriction allowed, keep/lower/resend allowed, and unchanged create-path guards. Co-authored-by: Cursor * test(ui): data-driven App Router migration E2E smoke (default + server-root-path) (#29974) * test(ui): add a data-driven App Router migration E2E smoke Add a growing Playwright smoke for migrated pages: for each segment it deep-links to the path route, asserts the URL and that the dashboard shell rendered, then clicks off to a legacy page and asserts navigation still works. Driven by e2e_tests/fixtures/migratedPages.ts, so adding a page is one line. Runs in two situations against the same proxy: the default mount (npm run e2e:migration) and a non-root SERVER_ROOT_PATH mount (npm run e2e:migration:root). globalSetup now logs in at `${SERVER_ROOT_PATH}/ui/login` so the admin storage state is valid under a prefix. Seeded with api-reference; append the rest as their migrations merge. * test(ui): support headed slow-motion + watch pauses in the migration smoke Honor SLOWMO in the server-root-path config (the default config already did), and add an env-gated E2E_WATCH_MS pause so a headed run lingers on each state. Both are no-ops by default, so CI behavior is unchanged. * test(ui): make the migration smoke a sidebar-click user journey Rework the smoke from deep-linking to a real navigation journey: start at the landing page, click the migrated page in the sidebar (expanding submenus for nested items), assert the path route rendered, reload it (the check a wrong server_root_path breaks), bounce to a legacy page and back, and — once two pages are migrated — navigate directly between two migrated pages. Verifies via URL + shell render, driven by the same fixture list. * test(ui): address review on the migration smoke Escape ROOT and segment before interpolating them into RegExp URL matchers so a future segment containing regex metacharacters can't silently widen the match. Make the server-root-path config fail fast when SERVER_ROOT_PATH is unset instead of silently re-running the default mount and passing without exercising the prefix. * test(ui): drop unused watch helper and fix stale smoke README * test(ui): run the migration smoke under a server root path in CI * test(ui): harden + instrument the server-root-path proxy reboot in CI * test(ui): run the server-root-path migration smoke as its own CI job Replace the in-place proxy reboot in e2e_ui_testing with a dedicated e2e_ui_testing_server_root_path job that boots the proxy once with SERVER_ROOT_PATH=/litellm, matching how every other proxy variant in the config gets its own job rather than killing and relaunching the live proxy. The reboot was failing deterministically: after pkill -9 and relaunch the prefixed proxy never came back up on :4000 (connection refused), so the smoke never ran. The readiness step that was supposed to surface the cause could never reach its boot-log tail because CircleCI runs steps under bash -eo pipefail and the preceding `curl -sv ... | tail` aborted the step with curl's exit 7. Booting the proxy as the job's own background step lets any boot crash land in that step's log instead of being swallowed. The default e2e_ui_testing job is unchanged aside from dropping the reboot, prefixed-readiness, and prefixed-smoke steps; the migration smoke still runs at the root mount there via the default Playwright config. * fix(proxy): extend response headers hook to streaming, TTS, image gen, and pass-through (#24232) * fix(proxy): extend response headers hook to streaming, TTS, image gen, and pass-through * test: mock post_call_response_headers_hook in audio speech route tests * chore(ui): remove dead App Router route stubs under (dashboard) (#30045) models-and-endpoints, organizations, and virtual-keys each had a page.tsx route under (dashboard)/ that is not in MIGRATED_PAGES, so the sidebar and deep links never resolve to it and the route is unreachable. Each was a thin wrapper that handed the shared view empty or no-op props (empty modelData with a no-op setModelData, hardcoded empty organizations, no-op setUserRole/setUserEmail), so reaching one would render a degraded page in any case. The real wrapper belongs in the PR that flips each page into MIGRATED_PAGES, written with eyes on it and a test This continues the dead-scaffolding cleanup from #28891. The shared components these wrappers rendered (ModelsAndEndpointsView, OrganizationFilters) stay, since the legacy ?page= switch in app/page.tsx and src/components still import them * fix(ui/mcp): reset OAuth state on create-server modal close so a prior server's token no longer leaks into the next add-server session (#30000) * fix(ui/mcp): reset OAuth hook state on modal close so a prior server's token no longer leaks into the next add-server session * fix(ui/mcp): clear in-flight OAuth guard on reset and reset form/tools on modal close so nothing leaks on a parent-driven dismiss * fix(mcp): allow team access-group grants in OAuth authorize/token access check (#30041) * fix(mcp): honor team access-group grants in OAuth authorize/token access check * test(mcp): mock build_effective_auth_contexts in non-admin authorize tests for isolation * docs(security): require a reproduction video for vulnerability reports (#30048) (#30063) With AI models capable of automated vulnerability discovery now publicly available, we expect a large increase in report volume, much of it unverified. Requiring a video of the exploit running against a live instance raises the bar for submissions and keeps triage focused on reproducible issues. Reports without a video will be closed and reopened if one is added later. Co-authored-by: stuxf <70670632+stuxf@users.noreply.github.com> * feat(ui): add admin flag to disable in-product UI nudges for everyone (#29796) * feat(ui): add admin flag to disable in-product UI nudges for everyone Admins can now suppress the survey and Claude Code feedback popups for all users via a single disable_ui_nudges UI setting, instead of relying on each user dismissing them individually. * fix(ui): suppress nudges while ui settings are loading Gate nudgesDisabled on the ui-settings loading state so an admin with disable_ui_nudges on doesn't see the survey prompt flash, and the getInProductNudgesCall fetch doesn't fire, on a cold page load before the flag resolves. Falls back to showing nudges if the fetch errors. * test(ui): wrap CreateKeyPage test in QueryClientProvider page.tsx now calls useUISettings (react-query), which needs a QueryClient that layout.tsx supplies in production but the test did not. Add the provider and mock getUiSettings so the query resolves. * chore(ui): remove dead dashboard files and unused dependencies (#30047) * chore(ui): remove dead dashboard files and unused dependencies knip flagged seven orphaned source/config files with no importers and five declared dependencies that nothing in the tree uses. Removing them shrinks the dashboard bundle's source surface and keeps the manifest honest; vite stays installed transitively via vitest, so test tooling is unaffected. * fix(ci): restore serverRootPath.config.ts referenced by SERVER_ROOT_PATH workflow The dead-code sweep removed e2e_tests/serverRootPath.config.ts, but its spec (tests/login/serverRootPathRedirect.spec.ts) and the test_server_root_path.yml workflow step still depend on it, so the redirect e2e job failed to load a config that no longer existed. * fix(proxy): authorize batch files using upload target_model_names (LIT-3593) (#30009) * fix(proxy): authorize batch files using upload target_model_names (LIT-3593) After replace_model_in_jsonl, body.model is a stripped provider id. Reverse-mapping it via resolve_model_name_from_model_id is first-match on model_list and caused false 403s when multiple deployments share the same stripped name. Use target_model_names from the unified file id instead. Co-authored-by: Cursor * fix(proxy): restore resolve_model_name_from_model_id for JSONL fallback path (LIT-3593) Restores the reverse-lookup for the JSONL body.model fallback path so that legacy/pre-target_model_names managed files still map stripped provider IDs back to proxy aliases before auth. Also cleans up redundant `or None`. Co-Authored-By: Claude Sonnet 4.6 * Revert "fix(proxy): restore resolve_model_name_from_model_id for JSONL fallback path (LIT-3593)" This reverts commit 30d2e96f77ef521ccaaf2193fe554980380eb669. --------- Co-authored-by: Cursor Co-authored-by: Claude Sonnet 4.6 * Add Claude Fable 5 across Anthropic, Bedrock, Vertex AI, and Azure AI (#30064) * Add Claude Fable 5 across Anthropic, Bedrock, Vertex AI, and Azure AI Adds cost map entries for claude-fable-5 ($10/$50 per MTok, 1M context, 128K output, adaptive thinking only) on the Anthropic API, Bedrock converse (base, global, and us/eu geo inference profiles at the 10% regional premium), Vertex AI, and Azure AI (Microsoft Foundry, which serves Fable 5 with the full 1M context window unlike Opus 4.8). Registers anthropic.claude-fable-5 in BEDROCK_CONVERSE_MODELS, lists the model in the setup wizard, and extends the reasoning effort e2e grid. The Bedrock, Vertex, and Azure grid cells carry fail_reason markers until the CI accounts are provisioned: Bedrock needs the provider data sharing opt-in Fable 5 requires, and the Foundry resource needs a claude-fable-5 deployment. The first-party entry carries provider_specific_entry {us: 1.1} for the inference_geo premium and deliberately no fast multiplier since Fable 5 has no fast mode. https://claude.ai/code/session_01MZarYYT3aS7DxaNjoax6Gm * Drop removed sampling params for Claude 4.7+ when drop_params is set Fable 5, Opus 4.7, and Opus 4.8 removed sampling params: the API rejects top_p, top_k, and any temperature other than 1 with a 400. LiteLLM was forwarding them even with drop_params enabled because the Anthropic and Bedrock converse transformations passed temperature/top_p through unconditionally. Mirror the GPT-5/o-series handling: temperature=1 still passes through, other values and any top_p are dropped when drop_params is set, and without drop_params a clean client-side UnsupportedParamsError tells the caller how to opt in, instead of surfacing the raw provider error. https://claude.ai/code/session_01MZarYYT3aS7DxaNjoax6Gm * Drive sampling param gating from the cost map and cover top_k Greptile review follow-ups on the sampling param fix: the restriction for Fable 5 / Opus 4.7 / 4.8 is now declared as supports_sampling_params: false on every affected cost map entry (perplexity excluded; that route is OpenAI-compatible and maps sampling params upstream) and read back through a tri-state map lookup, keeping the name check only as a fallback for provider-routed ids whose hosted map entries predate the flag, the same layering supports_adaptive_thinking uses. top_k bypasses map_openai_params as a provider-specific kwarg, so it is gated at the shared AnthropicConfig.transform_request boundary (direct, Bedrock invoke, Vertex, Azure) and in the Bedrock converse _handle_top_k_value path, with drop_params threaded through the converse transform helpers. Also updates the reasoning effort grid cell count assertion for the four Fable 5 rows added on this branch (29 x 11 cells). https://claude.ai/code/session_01MZarYYT3aS7DxaNjoax6Gm * Declare supports_sampling_params in the cost map schema The model map validation schema uses additionalProperties: false, so the new flag must be declared for the 28 entries that carry it; this was the one failing job (misc / Run tests) on the previous commit. https://claude.ai/code/session_01MZarYYT3aS7DxaNjoax6Gm * fix(bedrock): gate top_k=0 on converse to match Anthropic boundary Truthiness check let top_k=0 silently disappear on models that removed sampling params, while AnthropicConfig.transform_request treats 0 as present and raises UnsupportedParamsError (or drops when drop_params is set). Switch to 'is not None' so converse, direct Anthropic, invoke, Vertex, and Azure all behave the same for top_k=0. --------- Co-authored-by: Cursor Agent * fix: enable compact-2026-01-12 beta header for vertex_ai provider The vertex_ai block in anthropic_beta_headers_config.json mapped compact-2026-01-12 to null, so update_headers_with_filtered_beta stripped the header before the request reached Vertex while the compact_20260112 context edit stayed in the body, and Vertex rejected the request with HTTP 400. Vertex rawPredict accepts the header, and the bedrock and databricks blocks already forward it. Mirrors #21867, which enabled context-1m-2025-08-07 for vertex_ai the same way. Fixes #27290. --------- Co-authored-by: milan-berri Co-authored-by: Cursor Co-authored-by: ryan-crabbe-berri Co-authored-by: michelligabriele Co-authored-by: tin-berri Co-authored-by: yuneng-jiang Co-authored-by: stuxf <70670632+stuxf@users.noreply.github.com> Co-authored-by: Sameer Kankute Co-authored-by: Claude Sonnet 4.6 Co-authored-by: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> * fix(proxy): coerce litellm_settings.max_budget env var to float (#30113) * fix(team): reserve team budget raises for proxy admins on /team/update (#30030) The caller's PERSONAL max_budget was the wrong yardstick for /team/update: a team's spend ceiling has nothing to do with the admin's own key budget. That comparison was an unintended side effect of reusing _check_user_team_limits() (which exists for the /team/new path) and broke the UI, which re-sends the unchanged budget on every save. New behavior on /team/update for standalone teams: - A team admin (already authorized via _verify_team_access) may freely KEEP or LOWER the team budget, and change models/tpm/rpm, without being gated by their personal limits. - GROWING a team's spend ceiling is a budget-authority action reserved for proxy admins -> 403 for team admins. "Growing" covers both raising max_budget above the team's current finite value and removing the cap entirely (max_budget=null, detected via model_fields_set so an explicit null is distinguished from an omitted field). For a team that currently has no cap, setting a finite value is a restriction and is allowed. - Org-scoped teams remain governed by _check_org_team_limits() (capped by the org budget). Also reverts the #29525 existing_team_max_budget workaround in _check_user_team_limits() back to the create-only form; /team/new still enforces the creator's personal caps. docs(access_control): resolve the contradiction in the team-admin section — team admins can keep/lower the budget and manage rate limits/models, but cannot raise the team budget (proxy-admin only). tests: unit + behavior coverage for raise-blocked, cap-removal-blocked (team admin), raise/removal allowed (proxy admin), uncapped-team restriction allowed, keep/lower/resend allowed, and unchanged create-path guards. Co-authored-by: Cursor * test(ui): data-driven App Router migration E2E smoke (default + server-root-path) (#29974) * test(ui): add a data-driven App Router migration E2E smoke Add a growing Playwright smoke for migrated pages: for each segment it deep-links to the path route, asserts the URL and that the dashboard shell rendered, then clicks off to a legacy page and asserts navigation still works. Driven by e2e_tests/fixtures/migratedPages.ts, so adding a page is one line. Runs in two situations against the same proxy: the default mount (npm run e2e:migration) and a non-root SERVER_ROOT_PATH mount (npm run e2e:migration:root). globalSetup now logs in at `${SERVER_ROOT_PATH}/ui/login` so the admin storage state is valid under a prefix. Seeded with api-reference; append the rest as their migrations merge. * test(ui): support headed slow-motion + watch pauses in the migration smoke Honor SLOWMO in the server-root-path config (the default config already did), and add an env-gated E2E_WATCH_MS pause so a headed run lingers on each state. Both are no-ops by default, so CI behavior is unchanged. * test(ui): make the migration smoke a sidebar-click user journey Rework the smoke from deep-linking to a real navigation journey: start at the landing page, click the migrated page in the sidebar (expanding submenus for nested items), assert the path route rendered, reload it (the check a wrong server_root_path breaks), bounce to a legacy page and back, and — once two pages are migrated — navigate directly between two migrated pages. Verifies via URL + shell render, driven by the same fixture list. * test(ui): address review on the migration smoke Escape ROOT and segment before interpolating them into RegExp URL matchers so a future segment containing regex metacharacters can't silently widen the match. Make the server-root-path config fail fast when SERVER_ROOT_PATH is unset instead of silently re-running the default mount and passing without exercising the prefix. * test(ui): drop unused watch helper and fix stale smoke README * test(ui): run the migration smoke under a server root path in CI * test(ui): harden + instrument the server-root-path proxy reboot in CI * test(ui): run the server-root-path migration smoke as its own CI job Replace the in-place proxy reboot in e2e_ui_testing with a dedicated e2e_ui_testing_server_root_path job that boots the proxy once with SERVER_ROOT_PATH=/litellm, matching how every other proxy variant in the config gets its own job rather than killing and relaunching the live proxy. The reboot was failing deterministically: after pkill -9 and relaunch the prefixed proxy never came back up on :4000 (connection refused), so the smoke never ran. The readiness step that was supposed to surface the cause could never reach its boot-log tail because CircleCI runs steps under bash -eo pipefail and the preceding `curl -sv ... | tail` aborted the step with curl's exit 7. Booting the proxy as the job's own background step lets any boot crash land in that step's log instead of being swallowed. The default e2e_ui_testing job is unchanged aside from dropping the reboot, prefixed-readiness, and prefixed-smoke steps; the migration smoke still runs at the root mount there via the default Playwright config. * fix(proxy): extend response headers hook to streaming, TTS, image gen, and pass-through (#24232) * fix(proxy): extend response headers hook to streaming, TTS, image gen, and pass-through * test: mock post_call_response_headers_hook in audio speech route tests * chore(ui): remove dead App Router route stubs under (dashboard) (#30045) models-and-endpoints, organizations, and virtual-keys each had a page.tsx route under (dashboard)/ that is not in MIGRATED_PAGES, so the sidebar and deep links never resolve to it and the route is unreachable. Each was a thin wrapper that handed the shared view empty or no-op props (empty modelData with a no-op setModelData, hardcoded empty organizations, no-op setUserRole/setUserEmail), so reaching one would render a degraded page in any case. The real wrapper belongs in the PR that flips each page into MIGRATED_PAGES, written with eyes on it and a test This continues the dead-scaffolding cleanup from #28891. The shared components these wrappers rendered (ModelsAndEndpointsView, OrganizationFilters) stay, since the legacy ?page= switch in app/page.tsx and src/components still import them * fix(ui/mcp): reset OAuth state on create-server modal close so a prior server's token no longer leaks into the next add-server session (#30000) * fix(ui/mcp): reset OAuth hook state on modal close so a prior server's token no longer leaks into the next add-server session * fix(ui/mcp): clear in-flight OAuth guard on reset and reset form/tools on modal close so nothing leaks on a parent-driven dismiss * fix(mcp): allow team access-group grants in OAuth authorize/token access check (#30041) * fix(mcp): honor team access-group grants in OAuth authorize/token access check * test(mcp): mock build_effective_auth_contexts in non-admin authorize tests for isolation * docs(security): require a reproduction video for vulnerability reports (#30048) (#30063) With AI models capable of automated vulnerability discovery now publicly available, we expect a large increase in report volume, much of it unverified. Requiring a video of the exploit running against a live instance raises the bar for submissions and keeps triage focused on reproducible issues. Reports without a video will be closed and reopened if one is added later. Co-authored-by: stuxf <70670632+stuxf@users.noreply.github.com> * feat(ui): add admin flag to disable in-product UI nudges for everyone (#29796) * feat(ui): add admin flag to disable in-product UI nudges for everyone Admins can now suppress the survey and Claude Code feedback popups for all users via a single disable_ui_nudges UI setting, instead of relying on each user dismissing them individually. * fix(ui): suppress nudges while ui settings are loading Gate nudgesDisabled on the ui-settings loading state so an admin with disable_ui_nudges on doesn't see the survey prompt flash, and the getInProductNudgesCall fetch doesn't fire, on a cold page load before the flag resolves. Falls back to showing nudges if the fetch errors. * test(ui): wrap CreateKeyPage test in QueryClientProvider page.tsx now calls useUISettings (react-query), which needs a QueryClient that layout.tsx supplies in production but the test did not. Add the provider and mock getUiSettings so the query resolves. * chore(ui): remove dead dashboard files and unused dependencies (#30047) * chore(ui): remove dead dashboard files and unused dependencies knip flagged seven orphaned source/config files with no importers and five declared dependencies that nothing in the tree uses. Removing them shrinks the dashboard bundle's source surface and keeps the manifest honest; vite stays installed transitively via vitest, so test tooling is unaffected. * fix(ci): restore serverRootPath.config.ts referenced by SERVER_ROOT_PATH workflow The dead-code sweep removed e2e_tests/serverRootPath.config.ts, but its spec (tests/login/serverRootPathRedirect.spec.ts) and the test_server_root_path.yml workflow step still depend on it, so the redirect e2e job failed to load a config that no longer existed. * fix(proxy): authorize batch files using upload target_model_names (LIT-3593) (#30009) * fix(proxy): authorize batch files using upload target_model_names (LIT-3593) After replace_model_in_jsonl, body.model is a stripped provider id. Reverse-mapping it via resolve_model_name_from_model_id is first-match on model_list and caused false 403s when multiple deployments share the same stripped name. Use target_model_names from the unified file id instead. Co-authored-by: Cursor * fix(proxy): restore resolve_model_name_from_model_id for JSONL fallback path (LIT-3593) Restores the reverse-lookup for the JSONL body.model fallback path so that legacy/pre-target_model_names managed files still map stripped provider IDs back to proxy aliases before auth. Also cleans up redundant `or None`. Co-Authored-By: Claude Sonnet 4.6 * Revert "fix(proxy): restore resolve_model_name_from_model_id for JSONL fallback path (LIT-3593)" This reverts commit 30d2e96f77ef521ccaaf2193fe554980380eb669. --------- Co-authored-by: Cursor Co-authored-by: Claude Sonnet 4.6 * Add Claude Fable 5 across Anthropic, Bedrock, Vertex AI, and Azure AI (#30064) * Add Claude Fable 5 across Anthropic, Bedrock, Vertex AI, and Azure AI Adds cost map entries for claude-fable-5 ($10/$50 per MTok, 1M context, 128K output, adaptive thinking only) on the Anthropic API, Bedrock converse (base, global, and us/eu geo inference profiles at the 10% regional premium), Vertex AI, and Azure AI (Microsoft Foundry, which serves Fable 5 with the full 1M context window unlike Opus 4.8). Registers anthropic.claude-fable-5 in BEDROCK_CONVERSE_MODELS, lists the model in the setup wizard, and extends the reasoning effort e2e grid. The Bedrock, Vertex, and Azure grid cells carry fail_reason markers until the CI accounts are provisioned: Bedrock needs the provider data sharing opt-in Fable 5 requires, and the Foundry resource needs a claude-fable-5 deployment. The first-party entry carries provider_specific_entry {us: 1.1} for the inference_geo premium and deliberately no fast multiplier since Fable 5 has no fast mode. https://claude.ai/code/session_01MZarYYT3aS7DxaNjoax6Gm * Drop removed sampling params for Claude 4.7+ when drop_params is set Fable 5, Opus 4.7, and Opus 4.8 removed sampling params: the API rejects top_p, top_k, and any temperature other than 1 with a 400. LiteLLM was forwarding them even with drop_params enabled because the Anthropic and Bedrock converse transformations passed temperature/top_p through unconditionally. Mirror the GPT-5/o-series handling: temperature=1 still passes through, other values and any top_p are dropped when drop_params is set, and without drop_params a clean client-side UnsupportedParamsError tells the caller how to opt in, instead of surfacing the raw provider error. https://claude.ai/code/session_01MZarYYT3aS7DxaNjoax6Gm * Drive sampling param gating from the cost map and cover top_k Greptile review follow-ups on the sampling param fix: the restriction for Fable 5 / Opus 4.7 / 4.8 is now declared as supports_sampling_params: false on every affected cost map entry (perplexity excluded; that route is OpenAI-compatible and maps sampling params upstream) and read back through a tri-state map lookup, keeping the name check only as a fallback for provider-routed ids whose hosted map entries predate the flag, the same layering supports_adaptive_thinking uses. top_k bypasses map_openai_params as a provider-specific kwarg, so it is gated at the shared AnthropicConfig.transform_request boundary (direct, Bedrock invoke, Vertex, Azure) and in the Bedrock converse _handle_top_k_value path, with drop_params threaded through the converse transform helpers. Also updates the reasoning effort grid cell count assertion for the four Fable 5 rows added on this branch (29 x 11 cells). https://claude.ai/code/session_01MZarYYT3aS7DxaNjoax6Gm * Declare supports_sampling_params in the cost map schema The model map validation schema uses additionalProperties: false, so the new flag must be declared for the 28 entries that carry it; this was the one failing job (misc / Run tests) on the previous commit. https://claude.ai/code/session_01MZarYYT3aS7DxaNjoax6Gm * fix(bedrock): gate top_k=0 on converse to match Anthropic boundary Truthiness check let top_k=0 silently disappear on models that removed sampling params, while AnthropicConfig.transform_request treats 0 as present and raises UnsupportedParamsError (or drops when drop_params is set). Switch to 'is not None' so converse, direct Anthropic, invoke, Vertex, and Azure all behave the same for top_k=0. --------- Co-authored-by: Cursor Agent * fix(proxy): coerce litellm_settings.max_budget env var to float When max_budget is set in litellm_settings via os.environ/MAX_BUDGET, the env var resolves to a string and the generic setattr branch in ProxyConfig.load_config stored it as-is, so the startup check litellm.max_budget > 0 raised TypeError. The earlier fix (#23855) only covered the CLI initialize() path. Coerce the value to float in the settings loop, matching the existing max_internal_user_budget handling. Fixes #26696. --------- Co-authored-by: milan-berri Co-authored-by: Cursor Co-authored-by: ryan-crabbe-berri Co-authored-by: michelligabriele Co-authored-by: tin-berri Co-authored-by: yuneng-jiang Co-authored-by: stuxf <70670632+stuxf@users.noreply.github.com> Co-authored-by: Sameer Kankute Co-authored-by: Claude Sonnet 4.6 Co-authored-by: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> * fix(router): don't drop bedrock pass-through deployments using IAM credentials (#30111) * Fix Bedrock passthrough deployment dropped when using IAM credentials Bedrock deployments with use_in_pass_through enabled and IAM/OIDC auth (aws_role_name, no api_key) hit the generic pass-through branch in Router._initialize_deployment_for_pass_through, which calls set_pass_through_credentials and raises "api_key is required". The exception drops the deployment from the router entirely, breaking both passthrough and normal routing for that model. Skip the credential store write when no api_key is set; the bedrock passthrough route resolves AWS credentials at request time via BedrockConverseLLM.get_credentials(), not the passthrough credential store, so there is nothing to register here. Fixes #27728. * Reset passthrough credentials singleton before api_key credential test The test reads the module-level passthrough_endpoint_router singleton, so a stale "openai" entry written by an earlier test in the same process could make the assertion pass without exercising the code path. Clearing the credentials dict up front makes the test order-independent. * fix(sdk): stop mirroring reasoning_content in provider_specific_fields (#30110) The dict-to-response conversion path mirrored reasoning_content into provider_specific_fields, while live provider transforms (Anthropic's _build_provider_specific_fields) only set it top-level on the Message. Cache-replayed messages therefore serialized differently from live ones, breaking disk cache key stability for multi-turn conversations with extended thinking. The mirror was added for DeepSeek before Message.reasoning_content existed as a top-level attribute. The top-level field is still set by the converter, so DeepSeek's request-side promotion is unaffected. Fixes #27337. * fix(mcp): coerce mcp_server_cost_info values to float at ingest (#30109) * fix(mcp): coerce mcp_server_cost_info values to float at ingest YAML 1.1 parses scientific notation without a decimal point (e.g. 7e-05) as a string, and MCPServerCostInfo is a TypedDict with no runtime validation, so a string-typed default_cost_per_query from config.yaml flowed through the proxy untouched and crashed the MCP server settings page with '.toFixed is not a function'. Normalize mcp_server_cost_info on both the config and DB load paths, dropping non-numeric values with a warning instead of failing the server load. Fixes #27097. * fix(mcp): drop non-numeric default_cost_per_query instead of nulling it Keeping the key with a None value still exposes a null to the UI, which can crash .toFixed formatting when the consumer checks key existence rather than truthiness. Delete the key on coercion failure, matching how non-numeric per-tool cost entries are already omitted. * fix(proxy): count embedding and text completion tokens toward TPM limits (#30105) * fix(proxy): count embedding and text completion tokens toward TPM limits The parallel request limiters only read token usage off ModelResponse, so EmbeddingResponse and TextCompletionResponse objects left total_tokens at 0 and the per key, user, team, and end user TPM counters never incremented. Requests to /v1/embeddings and /v1/completions were effectively free against any tpm_limit. In the v3 limiter this was worse: the post-call reconciliation computed actual usage as 0 and refunded the pre-call reservation made at request time. Broaden the isinstance checks to accept EmbeddingResponse and TextCompletionResponse, which both expose a Usage object, at the four per-scope sites in parallel_request_limiter.py and at the usage extraction in parallel_request_limiter_v3.py. ResponsesAPIResponse was already covered in v3 via BaseLiteLLMOpenAIResponseObject. Fixes #27738. * test(proxy): cover v1 limiter TPM counting for embedding and text completion responses Exercise the broadened isinstance sites in parallel_request_limiter.py by asserting that async_log_success_event adds total_tokens to the per key, user, team, and end user TPM counters for EmbeddingResponse and TextCompletionResponse objects. The counters are pre-seeded at zero so the assertion is exactly the increment; on the pre-fix code these responses left total_tokens at 0 and the test fails. * fix(openai): forward client headers on the text completion path (#30103) * fix(openai): forward client headers on the text completion path litellm.completion() merges caller headers with extra_headers, but the text-completion-openai branch never passed the merged dict to openai_text_completions.completion(), and the handler only used its headers argument for logging. Pass the merged headers through the call site and set them as extra_headers on the outgoing request, mirroring the chat completion handler, so x-* client headers forwarded by the proxy reach the provider on /v1/completions. Fixes #27410. * Drop redundant extra_headers assignment and fix test module collision completion() merges extra_headers into headers before the text-completion-openai branch, and the handler now sets the merged headers as extra_headers on the request, so the branch-local optional_params["extra_headers"] assignment was a dead duplicate. Removing it keeps the assignment in one place while both entry paths (litellm.text_completion and direct handler callers) still forward headers; a new regression test pins the extra_headers kwarg path. Also rename the test module to test_completion_handler.py since its basename collided with tests/test_litellm/llms/bedrock/batches/ test_handler.py and broke pytest collection. * fix(bedrock): route Anthropic-shape count_tokens to InvokeModel and base64-encode the body (#30102) * fix(bedrock): route Anthropic-shape count_tokens to InvokeModel POST /v1/messages/count_tokens with Anthropic content blocks ({"type": "text"|"tool_use"|...}) was routed to the Converse input of the Bedrock CountTokens API. The Converse transform copies list content through verbatim, so Bedrock rejected the request with a 400 and the caller silently fell back to the local tokenizer, returning counts that can be off by ~50% on tool-heavy payloads. _detect_input_type now routes messages whose content blocks carry a "type" key (Anthropic shape) to the invokeModel input, which forwards the body verbatim. The invokeModel body is now base64-encoded as the CountTokens API requires (InvokeModelTokensRequest.body is a base64-encoded blob), and Anthropic Messages bodies get the anthropic_version and max_tokens fields Bedrock validates against. Fixes #27632. * refactor(bedrock): name the CountTokens max_tokens placeholder Replace the magic 1024 with a module-level DEFAULT_ANTHROPIC_INVOKE_MODEL_MAX_TOKENS constant so the intent is explicit and there is a single place to update if Bedrock's InvokeModel schema ever changes. Module-local rather than litellm/constants.py because the value is only a schema-validation placeholder for token counting, not a user-tunable generation default. * Add above-512k pricing tier for MiniMax-M3 and correct its base rates (#30095) * Add above-512k pricing tier support for MiniMax-M3 MiniMax-M3 doubles its per-token rates once a prompt exceeds 512k input tokens. The tiered cost parser already handles arbitrary thresholds, but get_model_info only copies whitelisted keys from ModelInfoBase, which had no 512k variants, so above_512k keys were silently dropped and long-context requests were priced at the flat rate. Add the input, output, and cache-read above_512k_tokens fields to ModelInfoBase and pass them through in get_model_info. Update the minimax/MiniMax-M3 entry with the tiered rates and correct the base rates, which matched the above-512k tier instead of the published base tier (https://platform.minimax.io/docs/guides/pricing-paygo). Fixes #29663. * Add above-512k keys to pricing schema, set MiniMax-M3 context to 1M Register the three new above_512k_tokens cost keys in the INTENDED_SCHEMA of test_aaamodel_prices_and_context_window_json_is_valid, declared the same way as the existing above_200k/above_272k tier keys, so the schema check accepts the MiniMax-M3 tiered pricing entry. Also raise MiniMax-M3 max_input_tokens from 512000 to 1000000 in both pricing JSONs. The MiniMax API docs (https://platform.minimax.io/docs/guides/text-generation) state the model supports a 1,000,000-token context window, and the pay-as-you-go pricing page (https://platform.minimax.io/docs/guides/pricing-paygo) prices input above 512k tokens, which only makes sense if inputs beyond 512k are accepted. This makes the above-512k pricing tier reachable. * fix(bedrock): make document names unique across conversation turns (#30093) * fix(bedrock): make document names unique across conversation turns PR #16275 derived Bedrock document names purely from a content hash so that names stay deterministic for prompt caching. When the same PDF or document appears in more than one conversation turn, every occurrence gets the identical name and Bedrock rejects the request with "Messages can not contain duplicate document names". Add _rename_duplicate_bedrock_document_names, a post-pass over the assembled message blocks that keeps the first occurrence's hash-based name and appends a positional suffix (_2, _3, ...) to later occurrences. Apply it in both _bedrock_converse_messages_pt and _bedrock_converse_messages_pt_async. Names remain deterministic across requests and the first occurrence is unchanged, so prompt cache prefixes stay stable. Fixes #29418. * fix(bedrock): avoid suffix collisions with organic document names A renamed duplicate could collide with a document whose hash-derived name already ends in the same positional suffix (e.g. an organic report_2 next to two documents named report). Collect every document name up front and bump the suffix until the candidate is unused, so renames can collide neither with organic names nor with each other. * fix(_types): remove ResponsesAPIResponse from PassThroughEndpointLoggingResultValues The import of ResponsesAPIResponse was removed from the file but a usage was left in the Union type, causing a NameError on import and breaking all CI tests. Remove the stale reference to match the cleanup intent. Co-Authored-By: Claude Sonnet 4.6 * fix(_types): restore ResponsesAPIResponse import and add use_xai_oauth to filter list Two related fixes: 1. Re-add ResponsesAPIResponse import in _types.py — it was removed but still needed in PassThroughEndpointLoggingResultValues (used in openai_passthrough_logging_handler.py). 2. Add use_xai_oauth to all_litellm_params so it is filtered before forwarding kwargs to providers like OpenAI that do not recognize it. Co-Authored-By: Claude Sonnet 4.6 --------- Co-authored-by: Hari Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Co-authored-by: Ceder Dens Co-authored-by: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Co-authored-by: 冯基魁 <56265583+fengjikui@users.noreply.github.com> Co-authored-by: victoruce <161634297+victoruce@users.noreply.github.com> Co-authored-by: kejunleng <33445544+silencedoctor@users.noreply.github.com> Co-authored-by: shin-berri Co-authored-by: yuneng-jiang Co-authored-by: Tyson Cung <45380903+tysoncung@users.noreply.github.com> Co-authored-by: Claude Sonnet 4.6 Co-authored-by: Jeremy Chapeau <113923302+jychp@users.noreply.github.com> Co-authored-by: Daan <255322319+daanhendrio@users.noreply.github.com> Co-authored-by: Avani Prajapati <143805019+Avani-prajapati@users.noreply.github.com> Co-authored-by: Kent <72616338+kingdoooo@users.noreply.github.com> Co-authored-by: daitran-tensormesh Co-authored-by: Dimitris Spachos Co-authored-by: Liam Scott Co-authored-by: Cursor Co-authored-by: Filippo Menghi <113345637+Cyberfilo@users.noreply.github.com> Co-authored-by: milan-berri Co-authored-by: ryan-crabbe-berri Co-authored-by: michelligabriele Co-authored-by: tin-berri Co-authored-by: stuxf <70670632+stuxf@users.noreply.github.com> Co-authored-by: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> --- litellm/anthropic_beta_headers_config.json | 2 +- .../litellm_core_utils/get_litellm_params.py | 1 + .../convert_dict_to_response.py | 5 - .../prompt_templates/factory.py | 47 +- .../responses_adapters/streaming_iterator.py | 16 +- .../bedrock/count_tokens/transformation.py | 36 +- .../responses/transformation.py | 16 +- litellm/llms/litellm_proxy/skills/README.md | 16 +- .../llms/litellm_proxy/skills/constants.py | 4 + litellm/llms/litellm_proxy/skills/handler.py | 3 +- litellm/llms/openai/completion/handler.py | 2 + .../vertex_ai_context_caching.py | 14 +- litellm/llms/xai/chat/transformation.py | 67 ++ litellm/llms/xai/oauth.py | 421 +++++++++ litellm/llms/xai/responses/transformation.py | 38 +- litellm/main.py | 5 +- ...odel_prices_and_context_window_backup.json | 172 +++- .../_experimental/mcp_server/mcp_context.py | 6 + .../mcp_server/mcp_server_manager.py | 48 ++ .../proxy/_experimental/mcp_server/server.py | 29 +- litellm/proxy/_types.py | 2 + litellm/proxy/auth/auth_checks.py | 7 +- litellm/proxy/common_request_processing.py | 16 +- .../guardrails/guardrail_hooks/presidio.py | 41 +- litellm/proxy/hooks/litellm_skills/main.py | 17 +- .../proxy/hooks/parallel_request_limiter.py | 21 +- .../hooks/parallel_request_limiter_v3.py | 19 +- .../key_management_endpoints.py | 29 +- .../management_endpoints/team_endpoints.py | 69 +- .../llm_passthrough_endpoints.py | 7 +- .../openai_passthrough_logging_handler.py | 76 +- .../pass_through_endpoints/success_handler.py | 12 +- litellm/proxy/proxy_cli.py | 16 + litellm/proxy/proxy_server.py | 2 + .../spend_management_endpoints.py | 2 +- litellm/router.py | 87 +- .../deployment_affinity_check.py | 7 +- .../encrypted_content_affinity_check.py | 36 +- litellm/types/router.py | 4 + litellm/types/utils.py | 8 + litellm/utils.py | 61 +- model_prices_and_context_window.json | 172 +++- .../test_convert_dict_to_chat_completion.py | 36 + .../test_router_endpoints.py | 57 ++ .../llm_cost_calc/test_llm_cost_calc_utils.py | 35 + ...llm_core_utils_prompt_templates_factory.py | 88 ++ .../test_xai_oauth_routing.py | 81 ++ ...t_responses_adapters_streaming_iterator.py | 79 ++ ...est_bedrock_count_tokens_transformation.py | 72 +- ...bedrock_mantle_responses_transformation.py | 248 +++++- .../completion/test_completion_handler.py | 93 ++ .../openai_like/test_tensormesh_provider.py | 72 ++ .../test_vertex_ai_context_caching.py | 52 +- tests/test_litellm/llms/xai/test_xai_oauth.py | 801 ++++++++++++++++++ .../mcp_server/test_mcp_server.py | 130 ++- .../mcp_server/test_mcp_server_manager.py | 64 ++ .../proxy/auth/test_auth_checks.py | 2 + .../guardrail_hooks/test_presidio.py | 162 +++- .../proxy/hooks/litellm_skills/test_main.py | 67 ++ .../hooks/test_parallel_request_limiter.py | 86 ++ .../hooks/test_parallel_request_limiter_v3.py | 69 +- .../test_key_management_endpoints.py | 154 ++++ .../test_team_endpoints.py | 326 +++++++ ...test_openai_passthrough_logging_handler.py | 268 +++++- .../test_llm_pass_through_endpoints.py | 80 +- .../test_spend_management_endpoints.py | 5 +- .../proxy/test_common_request_processing.py | 30 + tests/test_litellm/proxy/test_proxy_server.py | 30 + .../test_encrypted_content_affinity_check.py | 188 ++++ .../test_anthropic_beta_headers_filtering.py | 10 + tests/test_litellm/test_router.py | 292 +++++++ tests/test_litellm/test_utils.py | 4 + .../src/components/networking.test.ts | 47 + .../src/components/networking.tsx | 22 +- .../LogDetailsDrawer/LogDetailsDrawer.tsx | 88 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 12 + 76 files changed, 5277 insertions(+), 232 deletions(-) create mode 100644 litellm/llms/xai/oauth.py create mode 100644 tests/test_litellm/litellm_core_utils/test_xai_oauth_routing.py create mode 100644 tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py create mode 100644 tests/test_litellm/llms/openai/completion/test_completion_handler.py create mode 100644 tests/test_litellm/llms/xai/test_xai_oauth.py create mode 100644 tests/test_litellm/proxy/hooks/litellm_skills/test_main.py create mode 100644 tests/test_litellm/proxy/hooks/test_parallel_request_limiter.py diff --git a/litellm/anthropic_beta_headers_config.json b/litellm/anthropic_beta_headers_config.json index d02afe37569..a0d63f5043c 100644 --- a/litellm/anthropic_beta_headers_config.json +++ b/litellm/anthropic_beta_headers_config.json @@ -129,7 +129,7 @@ "bash_20241022": null, "bash_20250124": null, "code-execution-2025-08-25": null, - "compact-2026-01-12": null, + "compact-2026-01-12": "compact-2026-01-12", "computer-use-2025-01-24": "computer-use-2025-01-24", "computer-use-2025-11-24": "computer-use-2025-11-24", "context-1m-2025-08-07": "context-1m-2025-08-07", diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index b32803b5dfc..f80cb41dc3f 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -34,6 +34,7 @@ _OPTIONAL_KWARGS_KEYS = frozenset( "aws_bedrock_runtime_endpoint", "tpm", "rpm", + "use_xai_oauth", } ) diff --git a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py index 2547fd4d8c6..4e5b53a13d7 100644 --- a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py +++ b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py @@ -633,11 +633,6 @@ def convert_to_model_response_object( # noqa: PLR0915 thinking_blocks = choice["message"]["thinking_blocks"] provider_specific_fields["thinking_blocks"] = thinking_blocks - if reasoning_content: - provider_specific_fields["reasoning_content"] = ( - reasoning_content - ) - message = Message( content=content, role=choice["message"]["role"] or "assistant", diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 81a4c8b14b6..b09f2bb130e 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -4290,6 +4290,49 @@ def _deduplicate_bedrock_tool_content( return _deduplicate_bedrock_content_blocks(tool_content, "toolResult") +def _rename_duplicate_bedrock_document_names( + contents: List[BedrockMessageBlock], +) -> List[BedrockMessageBlock]: + """ + Rename duplicate document names across all messages in a Bedrock request. + + Document names are derived from a content hash, so the same file appearing + in multiple conversation turns produces identical names and Bedrock rejects + the request with "Messages can not contain duplicate document names". The + first occurrence keeps its original name so prompt-cache prefixes stay + stable; later occurrences get a deterministic positional suffix + (``_2``, ``_3``, ...), bumped further if the suffixed name already + belongs to another document (e.g. an organic name ending in ``_2``). + """ + used_names: Set[str] = set() + for message in contents: + for block in message.get("content") or []: + document = block.get("document") + if isinstance(document, dict) and document.get("name"): + used_names.add(document["name"]) + + name_counts: Dict[str, int] = {} + for message in contents: + for block in message.get("content") or []: + document = block.get("document") + if not isinstance(document, dict): + continue + name = document.get("name") + if not name: + continue + count = name_counts.get(name, 0) + 1 + name_counts[name] = count + if count > 1: + suffix = count + new_name = f"{name}_{suffix}" + while new_name in used_names: + suffix += 1 + new_name = f"{name}_{suffix}" + used_names.add(new_name) + document["name"] = new_name + return contents + + def _sort_bedrock_assistant_content_blocks( blocks: List[BedrockContentBlock], ) -> List[BedrockContentBlock]: @@ -4938,7 +4981,7 @@ class BedrockConverseMessagesProcessor: llm_provider=llm_provider, ) - return contents + return _rename_duplicate_bedrock_document_names(contents) @staticmethod def translate_thinking_blocks_to_reasoning_content_blocks( @@ -5360,7 +5403,7 @@ def _bedrock_converse_messages_pt( # noqa: PLR0915 llm_provider=llm_provider, ) - return contents + return _rename_duplicate_bedrock_document_names(contents) def make_valid_bedrock_tool_name(input_tool_name: str) -> str: diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py index 94c5200be64..5f1362e259f 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py @@ -155,10 +155,24 @@ class AnthropicResponsesStreamWrapper: event.get("delta", "") if isinstance(event, dict) else "" ) block_idx = ( - self._item_id_to_block_index.get(item_id, self._current_block_index) + self._item_id_to_block_index.get(item_id, -1) if item_id else self._current_block_index ) + if block_idx < 0: + # Some providers (e.g. LMStudio) skip response.output_item.added, + # so no text block is open yet; synthesize content_block_start + # instead of emitting a delta with index -1 + block_idx = self._next_block_index() + if item_id: + self._item_id_to_block_index[item_id] = block_idx + self._chunk_queue.append( + { + "type": "content_block_start", + "index": block_idx, + "content_block": {"type": "text", "text": ""}, + } + ) self._chunk_queue.append( { "type": "content_block_delta", diff --git a/litellm/llms/bedrock/count_tokens/transformation.py b/litellm/llms/bedrock/count_tokens/transformation.py index c967fd334bc..bdef3349e00 100644 --- a/litellm/llms/bedrock/count_tokens/transformation.py +++ b/litellm/llms/bedrock/count_tokens/transformation.py @@ -11,6 +11,11 @@ from typing import Any, Dict, List, Optional from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.bedrock.common_utils import get_bedrock_base_model +# Placeholder satisfying the Anthropic InvokeModel schema's required +# max_tokens field; CountTokens only counts input, so it has no effect +# on any generation. +DEFAULT_ANTHROPIC_INVOKE_MODEL_MAX_TOKENS = 1024 + class BedrockCountTokensConfig(BaseAWSLLM): """ @@ -32,8 +37,20 @@ class BedrockCountTokensConfig(BaseAWSLLM): Returns: 'converse' or 'invokeModel' """ - # If the request has messages in the expected Anthropic format, use converse - if "messages" in request_data and isinstance(request_data["messages"], list): + messages = request_data.get("messages") + if isinstance(messages, list): + # Anthropic content blocks carry a "type" key ({"type": "text", ...}); + # Converse blocks don't ({"text": ...}, {"toolUse": ...}). Converse + # rejects Anthropic-shape blocks, so route those to invokeModel, + # which forwards the body verbatim. + for message in messages: + if not isinstance(message, dict): + continue + content = message.get("content") + if isinstance(content, list) and any( + isinstance(block, dict) and "type" in block for block in content + ): + return "invokeModel" return "converse" # For raw text or other formats, use invokeModel @@ -68,7 +85,7 @@ class BedrockCountTokensConfig(BaseAWSLLM): { "input": { "invokeModel": { - "body": "{...raw model input...}" + "body": "" } } } @@ -168,13 +185,24 @@ class BedrockCountTokensConfig(BaseAWSLLM): self, request_data: Dict[str, Any] ) -> Dict[str, Any]: """Transform to InvokeModel input format.""" + import base64 import json # For InvokeModel, we need to provide the raw body that would be sent to the model # Remove the 'model' field from the body as it's not part of the model input body_data = {k: v for k, v in request_data.items() if k != "model"} - return {"input": {"invokeModel": {"body": json.dumps(body_data)}}} + if "messages" in body_data: + # Bedrock validates the body against the model's InvokeModel schema; + # Anthropic Messages bodies require these fields. + body_data.setdefault("anthropic_version", "bedrock-2023-05-31") + body_data.setdefault( + "max_tokens", DEFAULT_ANTHROPIC_INVOKE_MODEL_MAX_TOKENS + ) + + # The CountTokens API expects invokeModel.body as a base64-encoded blob + encoded_body = base64.b64encode(json.dumps(body_data).encode()).decode() + return {"input": {"invokeModel": {"body": encoded_body}}} def get_bedrock_count_tokens_endpoint( self, diff --git a/litellm/llms/bedrock_mantle/responses/transformation.py b/litellm/llms/bedrock_mantle/responses/transformation.py index df219091074..dfa108833ac 100644 --- a/litellm/llms/bedrock_mantle/responses/transformation.py +++ b/litellm/llms/bedrock_mantle/responses/transformation.py @@ -1,8 +1,10 @@ """ Amazon Bedrock Mantle - Responses API backend. -gpt-5.5 / gpt-5.4 on Mantle are exposed ONLY on the `/openai/v1/responses` -path (not the standard `/v1/responses`). Payloads and SSE follow the OpenAI +Mantle serves Responses on two upstream paths: gpt frontier models (gpt-5.5 / +gpt-5.4) on `/openai/v1/responses`, and everything else that supports Responses +(e.g. gpt-oss) on the standard `/v1/responses`. The gate picks the path per +model and injects it via `use_openai_path`. Payloads and SSE follow the OpenAI Responses spec, so this config inherits OpenAIResponsesAPIConfig and overrides only the endpoint URL and authentication. @@ -48,9 +50,14 @@ _MANTLE_HOST_RE = re.compile( class BedrockMantleResponsesAPIConfig(OpenAIResponsesAPIConfig): - def __init__(self, aws_signer: Optional[BaseAWSLLM] = None): + def __init__( + self, + aws_signer: Optional[BaseAWSLLM] = None, + use_openai_path: bool = True, + ): super().__init__() self._aws_signer = aws_signer or BaseAWSLLM() + self.use_openai_path = use_openai_path @property def custom_llm_provider(self) -> LlmProviders: @@ -94,7 +101,8 @@ class BedrockMantleResponsesAPIConfig(OpenAIResponsesAPIConfig): # single resolved region so aws_region_name wins; preserve custom proxy hosts. if _MANTLE_HOST_RE.match(base): base = f"https://bedrock-mantle.{region}.api.aws" - return f"{base}/openai/v1/responses" + path = "/openai/v1/responses" if self.use_openai_path else "/v1/responses" + return f"{base}{path}" def validate_environment( self, headers: dict, model: str, litellm_params: Optional[GenericLiteLLMParams] diff --git a/litellm/llms/litellm_proxy/skills/README.md b/litellm/llms/litellm_proxy/skills/README.md index 1dfeff1a42c..a896aa1166e 100644 --- a/litellm/llms/litellm_proxy/skills/README.md +++ b/litellm/llms/litellm_proxy/skills/README.md @@ -18,7 +18,7 @@ flowchart TB F[Request with container.skills] --> G[SkillsInjectionHook] G --> H{skill_id prefix?} - H -->|"litellm:skill_abc"| I[Fetch from LiteLLM DB] + H -->|"litellm_skill_abc"| I[Fetch from LiteLLM DB] H -->|"skill_xyz" no prefix| J[Pass to Anthropic as native skill] I --> K{Model provider?} @@ -57,7 +57,7 @@ sequenceDiagram Note over LiteLLM,PreHook: PRE-CALL HOOK LiteLLM->>PreHook: Intercept request - PreHook->>PreHook: Fetch skill from DB (litellm:skill_id) + PreHook->>PreHook: Fetch skill from DB (litellm_skill_id) PreHook->>PreHook: Extract SKILL.md from ZIP PreHook->>PreHook: Inject SKILL.md into system prompt PreHook->>PreHook: Add litellm_code_execution tool @@ -105,7 +105,7 @@ response = await litellm.acompletion( model="gpt-4o-mini", messages=[{"role": "user", "content": "Create a bouncing ball GIF"}], container={ - "skills": [{"type": "custom", "skill_id": "litellm:skill_abc123"}] + "skills": [{"type": "custom", "skill_id": "litellm_skill_abc123"}] }, ) @@ -261,7 +261,7 @@ response = litellm.completion( messages=[{"role": "user", "content": "Analyze this data..."}], container={ "skills": [ - {"type": "custom", "skill_id": "litellm:skill_abc123"} # litellm: prefix + {"type": "custom", "skill_id": "litellm_skill_abc123"} # litellm_skill_ prefix ] } ) @@ -277,7 +277,7 @@ response = litellm.completion( "messages": [{"role": "user", "content": "Help me analyze data"}], "container": { "skills": [ - {"type": "custom", "skill_id": "litellm:skill_abc123"} + {"type": "custom", "skill_id": "litellm_skill_abc123"} ] } } @@ -287,7 +287,7 @@ response = litellm.completion( The hook (`litellm/proxy/hooks/litellm_skills/main.py`) intercepts the request: -1. **Detects `litellm:` prefix** → Fetches skill from database +1. **Detects `litellm_skill_` prefix** → Fetches skill from database 2. **Checks model provider** → Bedrock is not Anthropic 3. **Extracts SKILL.md** from stored ZIP file 4. **Converts skill to tool** + **Injects content into system prompt** @@ -361,8 +361,8 @@ model LiteLLM_SkillsTable { | Create skill on Anthropic | `anthropic` | N/A | Forward to Anthropic API | | Create skill in LiteLLM DB | `litellm_proxy` | N/A | Store in database | | Use Anthropic native skill | N/A | `skill_xyz` | Pass to Anthropic container.skills | -| Use LiteLLM skill on Anthropic | N/A | `litellm:skill_abc` | Convert to tools | -| Use LiteLLM skill on Bedrock/OpenAI | N/A | `litellm:skill_abc` | Convert to tools + inject SKILL.md | +| Use LiteLLM skill on Anthropic | N/A | `litellm_skill_abc` | Convert to tools | +| Use LiteLLM skill on Bedrock/OpenAI | N/A | `litellm_skill_abc` | Convert to tools + inject SKILL.md | ## Testing diff --git a/litellm/llms/litellm_proxy/skills/constants.py b/litellm/llms/litellm_proxy/skills/constants.py index a8c2697fcee..0c60a60842a 100644 --- a/litellm/llms/litellm_proxy/skills/constants.py +++ b/litellm/llms/litellm_proxy/skills/constants.py @@ -4,6 +4,10 @@ Constants for LiteLLM Skills Centralized constants for skills processing, code execution, and sandbox configuration. """ +LITELLM_SKILL_ID_PREFIX: str = "litellm_skill_" +"""Prefix for DB-backed skill IDs. The model-facing tool name is the skill ID +with hyphens/spaces replaced by underscores, which leaves this prefix intact.""" + # Code execution loop settings DEFAULT_MAX_ITERATIONS: int = 10 """Maximum number of iterations for the automatic code execution loop.""" diff --git a/litellm/llms/litellm_proxy/skills/handler.py b/litellm/llms/litellm_proxy/skills/handler.py index 7b259c1ed66..9138b9a712f 100644 --- a/litellm/llms/litellm_proxy/skills/handler.py +++ b/litellm/llms/litellm_proxy/skills/handler.py @@ -10,6 +10,7 @@ from typing import Any, Dict, List, Optional from litellm._logging import verbose_logger from litellm.caching.in_memory_cache import InMemoryCache +from litellm.llms.litellm_proxy.skills.constants import LITELLM_SKILL_ID_PREFIX from litellm.proxy._types import LiteLLM_SkillsTable, NewSkillRequest, UserAPIKeyAuth from litellm.proxy.common_utils.resource_ownership import ( get_primary_resource_owner_scope, @@ -68,7 +69,7 @@ class LiteLLMSkillsHandler: ) -> LiteLLM_SkillsTable: prisma_client = await LiteLLMSkillsHandler._get_prisma_client() - skill_id = f"litellm_skill_{uuid.uuid4()}" + skill_id = f"{LITELLM_SKILL_ID_PREFIX}{uuid.uuid4()}" owner = get_primary_resource_owner_scope(user_api_key_dict) or user_id if owner is None: # Identity-less callers (no user_id / team_id / org_id / diff --git a/litellm/llms/openai/completion/handler.py b/litellm/llms/openai/completion/handler.py index 1641615126e..63d39151254 100644 --- a/litellm/llms/openai/completion/handler.py +++ b/litellm/llms/openai/completion/handler.py @@ -49,6 +49,8 @@ class OpenAITextCompletion(BaseLLM): headers: Optional[dict] = None, ): try: + if headers: + optional_params = {**optional_params, "extra_headers": headers} if headers is None: headers = self.validate_environment(api_key=api_key) if model is None or messages is None: diff --git a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py index e9f08f403f9..103801a1e8d 100644 --- a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py +++ b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py @@ -19,7 +19,7 @@ from litellm.types.llms.vertex_ai import ( VertexAICachedContentResponseObject, ) -from ..common_utils import VertexAIError +from ..common_utils import VertexAIError, get_vertex_base_url from ..vertex_llm_base import VertexBase from .transformation import ( separate_cached_messages, @@ -69,17 +69,13 @@ class ContextCachingEndpoints(VertexBase): elif custom_llm_provider == "vertex_ai": auth_header = vertex_auth_header endpoint = "cachedContents" - if vertex_location == "global": - url = f"https://aiplatform.googleapis.com/v1/projects/{vertex_project}/locations/{vertex_location}/{endpoint}" - else: - url = f"https://{vertex_location}-aiplatform.googleapis.com/v1/projects/{vertex_project}/locations/{vertex_location}/{endpoint}" + base_url = get_vertex_base_url(vertex_location) + url = f"{base_url}/v1/projects/{vertex_project}/locations/{vertex_location}/{endpoint}" else: auth_header = vertex_auth_header endpoint = "cachedContents" - if vertex_location == "global": - url = f"https://aiplatform.googleapis.com/v1beta1/projects/{vertex_project}/locations/{vertex_location}/{endpoint}" - else: - url = f"https://{vertex_location}-aiplatform.googleapis.com/v1beta1/projects/{vertex_project}/locations/{vertex_location}/{endpoint}" + base_url = get_vertex_base_url(vertex_location) + url = f"{base_url}/v1beta1/projects/{vertex_project}/locations/{vertex_location}/{endpoint}" return self._check_custom_proxy( api_base=api_base, diff --git a/litellm/llms/xai/chat/transformation.py b/litellm/llms/xai/chat/transformation.py index c06928516ef..8019bb67991 100644 --- a/litellm/llms/xai/chat/transformation.py +++ b/litellm/llms/xai/chat/transformation.py @@ -5,6 +5,7 @@ import httpx import litellm from litellm._logging import verbose_logger from litellm.constants import XAI_API_BASE +from litellm.exceptions import AuthenticationError from litellm.litellm_core_utils.prompt_templates.common_utils import ( filter_value_from_dict, strip_name_from_messages, @@ -39,6 +40,72 @@ class XAIChatConfig(OpenAIGPTConfig): dynamic_api_key = XAIModelInfo.get_api_key(api_key) return api_base, dynamic_api_key + def validate_environment( + self, + headers: dict, + model: str, + messages: List[AllMessageValues], + optional_params: dict, + litellm_params: dict, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + ) -> dict: + from litellm.llms.xai.oauth import ( + XAIOAuthAuthenticator, + XAIOAuthError, + should_use_xai_oauth, + ) + + dynamic_api_key = XAIModelInfo.get_api_key(api_key) + if should_use_xai_oauth(litellm_params) and not dynamic_api_key: + try: + headers["Authorization"] = ( + f"Bearer {XAIOAuthAuthenticator().get_access_token()}" + ) + except XAIOAuthError as exc: + raise AuthenticationError( + model=model, + llm_provider=self.custom_llm_provider or "xai", + message=str(exc), + ) from exc + if "content-type" not in headers and "Content-Type" not in headers: + headers["Content-Type"] = "application/json" + return headers + + return super().validate_environment( + headers=headers, + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + api_key=dynamic_api_key, + api_base=api_base, + ) + + def get_complete_url( + self, + api_base: Optional[str], + api_key: Optional[str], + model: str, + optional_params: dict, + litellm_params: dict, + stream: Optional[bool] = None, + ) -> str: + from litellm.llms.xai.oauth import XAIOAuthAuthenticator, should_use_xai_oauth + + dynamic_api_key = XAIModelInfo.get_api_key(api_key) + if should_use_xai_oauth(litellm_params) and not dynamic_api_key: + api_base = XAIOAuthAuthenticator().get_api_base() + + return super().get_complete_url( + api_base=api_base, + api_key=dynamic_api_key, + model=model, + optional_params=optional_params, + litellm_params=litellm_params, + stream=stream, + ) + def get_supported_openai_params(self, model: str) -> list: base_openai_params = [ "logit_bias", diff --git a/litellm/llms/xai/oauth.py b/litellm/llms/xai/oauth.py new file mode 100644 index 00000000000..30c717b7ca0 --- /dev/null +++ b/litellm/llms/xai/oauth.py @@ -0,0 +1,421 @@ +import base64 +import hashlib +import json +import os +import secrets +import sys +import threading +import time +import uuid +import webbrowser +from http.server import BaseHTTPRequestHandler, HTTPServer +from typing import Any, Dict, Optional, Tuple, Union +from urllib.parse import parse_qs, urlencode, urlparse + +import httpx + +from litellm._logging import verbose_logger +from litellm.constants import XAI_API_BASE +from litellm.llms.custom_httpx.http_handler import HTTPHandler, _get_httpx_client +from litellm.secret_managers.main import get_secret_str + +XAI_OAUTH_ISSUER = "https://auth.x.ai" +XAI_OAUTH_DISCOVERY_URL = f"{XAI_OAUTH_ISSUER}/.well-known/openid-configuration" +XAI_OAUTH_CLIENT_ID = "b1a00492-073a-47ea-816f-4c329264a828" +XAI_OAUTH_SCOPE = "openid profile email offline_access grok-cli:access api:access" +XAI_OAUTH_REDIRECT_HOST = "127.0.0.1" +XAI_OAUTH_REDIRECT_PORT = 56121 +XAI_OAUTH_REDIRECT_PATH = "/callback" +XAI_OAUTH_EXPIRY_SKEW_SECONDS = 120 +XAI_OAUTH_CALLBACK_TIMEOUT_SECONDS = 180 +_XAI_OAUTH_REFRESH_LOCK = threading.Lock() + + +class XAIOAuthError(Exception): + pass + + +class XAIOAuthLoginRequiredError(XAIOAuthError): + pass + + +class _CallbackHandler(BaseHTTPRequestHandler): + server: "_CallbackServer" + + def do_GET(self) -> None: + parsed = urlparse(self.path) + if parsed.path != XAI_OAUTH_REDIRECT_PATH: + self.send_response(404) + self.end_headers() + return + + params = parse_qs(parsed.query) + result = { + "code": params.get("code", [None])[0], + "state": params.get("state", [None])[0], + "error": params.get("error", [None])[0], + "error_description": params.get("error_description", [None])[0], + } + self.server.callback_result = result + + if result["state"] != self.server.expected_state: + self.send_response(400) + self.send_header("Content-Type", "text/html; charset=utf-8") + self.end_headers() + self.wfile.write( + b"

xAI authorization state mismatch.

" + ) + return + + self.send_response(200) + self.send_header("Content-Type", "text/html; charset=utf-8") + self.end_headers() + body = ( + b"

xAI authorization failed.

You can close this tab." + if result["error"] + else b"

xAI authorization received.

You can close this tab." + ) + self.wfile.write(body) + + def log_message(self, format: str, *args: Any) -> None: + return + + +class _CallbackServer(HTTPServer): + expected_state: str + callback_result: Optional[Dict[str, Optional[str]]] + + +class XAIOAuthAuthenticator: + def __init__( + self, http_client: Optional[Union[httpx.Client, HTTPHandler]] = None + ) -> None: + self.token_dir = get_secret_str("XAI_OAUTH_TOKEN_DIR") or os.path.expanduser( + "~/.config/litellm/xai_oauth" + ) + self.auth_file = os.path.join( + self.token_dir, get_secret_str("XAI_OAUTH_AUTH_FILE") or "auth.json" + ) + self.http_client = http_client + + def get_api_base(self) -> str: + return ( + get_secret_str("XAI_OAUTH_API_BASE") + or get_secret_str("XAI_API_BASE") + or XAI_API_BASE + ) + + def get_access_token(self) -> str: + auth_data = self._read_auth_file() + if not auth_data: + raise XAIOAuthLoginRequiredError( + "xAI OAuth login required. Run `litellm xai-oauth login`." + ) + + access_token = auth_data.get("access_token") + if access_token and not self._is_expired(auth_data): + return access_token + + refresh_token = auth_data.get("refresh_token") + if not refresh_token: + raise XAIOAuthLoginRequiredError( + "xAI OAuth refresh token missing. Run `litellm xai-oauth login`." + ) + + with _XAI_OAUTH_REFRESH_LOCK: + locked_auth_data = self._read_auth_file() or auth_data + access_token = locked_auth_data.get("access_token") + if access_token and not self._is_expired(locked_auth_data): + return access_token + + refreshed = self._refresh_tokens(locked_auth_data) + return refreshed["access_token"] + + def login(self, force: bool = False, no_browser: bool = False) -> Dict[str, Any]: + existing = self._read_auth_file() + if existing and not force and existing.get("access_token"): + if not self._is_expired(existing): + return existing + if existing.get("refresh_token"): + try: + return self._refresh_tokens(existing) + except XAIOAuthError: + pass + + discovery = self._discover() + verifier, challenge = self._pkce_pair() + state = uuid.uuid4().hex + nonce = uuid.uuid4().hex + server, redirect_uri = self._start_callback_server(state) + authorize_url = self._build_authorize_url( + authorization_endpoint=discovery["authorization_endpoint"], + redirect_uri=redirect_uri, + challenge=challenge, + state=state, + nonce=nonce, + ) + + if no_browser or not webbrowser.open(authorize_url): + sys.stdout.write( + f"Open this URL to authenticate with xAI:\n{authorize_url}\n" + ) + sys.stdout.flush() + + result = self._wait_for_callback(server) + if result.get("state") != state: + raise XAIOAuthError("xAI OAuth state mismatch") + if result.get("error"): + description = result.get("error_description") or result["error"] + raise XAIOAuthError(f"xAI authorization failed: {description}") + code = result.get("code") + if not code: + raise XAIOAuthError("xAI authorization failed: no code returned") + + token_payload = self._exchange_token( + discovery["token_endpoint"], + { + "grant_type": "authorization_code", + "code": code, + "redirect_uri": redirect_uri, + "client_id": XAI_OAUTH_CLIENT_ID, + "code_verifier": verifier, + }, + ) + auth_data = self._build_auth_record(token_payload, discovery["token_endpoint"]) + self._write_auth_file(auth_data) + return auth_data + + def _client(self) -> Union[httpx.Client, HTTPHandler]: + return self.http_client or _get_httpx_client() + + def _ensure_token_dir(self) -> None: + os.makedirs(self.token_dir, mode=0o700, exist_ok=True) + try: + os.chmod(self.token_dir, 0o700) + except OSError: + verbose_logger.debug("Could not chmod xAI OAuth token directory") + + def _read_auth_file(self) -> Optional[Dict[str, Any]]: + try: + with open(self.auth_file, "r") as f: + data = json.load(f) + return data if isinstance(data, dict) else None + except (IOError, json.JSONDecodeError): + return None + + def _write_auth_file(self, data: Dict[str, Any]) -> None: + self._ensure_token_dir() + tmp_file = os.path.join( + self.token_dir, + f".{os.path.basename(self.auth_file)}.{uuid.uuid4().hex}.tmp", + ) + flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL + if hasattr(os, "O_NOFOLLOW"): + flags |= os.O_NOFOLLOW + fd = os.open(tmp_file, flags, 0o600) + try: + with os.fdopen(fd, "w") as f: + json.dump(data, f) + f.flush() + os.fsync(f.fileno()) + os.replace(tmp_file, self.auth_file) + try: + os.chmod(self.auth_file, 0o600) + except OSError: + verbose_logger.debug("Could not chmod xAI OAuth auth file") + except Exception: + try: + os.close(fd) + except OSError: + pass + try: + os.unlink(tmp_file) + except OSError: + pass + raise + + def _is_expired(self, auth_data: Dict[str, Any]) -> bool: + expires_at = auth_data.get("expires_at") + if expires_at is None: + return True + try: + return time.time() >= float(expires_at) - XAI_OAUTH_EXPIRY_SKEW_SECONDS + except (TypeError, ValueError): + return True + + def _discover(self) -> Dict[str, str]: + try: + response = self._client().get( + XAI_OAUTH_DISCOVERY_URL, headers={"Accept": "application/json"} + ) + response.raise_for_status() + except httpx.HTTPStatusError as exc: + raise XAIOAuthError( + f"xAI OAuth discovery request failed: {exc.response.status_code} {exc.response.text}" + ) from exc + try: + data = response.json() + except ValueError as exc: + raise XAIOAuthError( + "xAI OAuth discovery response was not valid JSON" + ) from exc + authorization_endpoint = data.get("authorization_endpoint") + token_endpoint = data.get("token_endpoint") + if not authorization_endpoint or not token_endpoint: + raise XAIOAuthError("xAI OAuth discovery missing endpoints") + return { + "authorization_endpoint": self._validate_xai_endpoint( + authorization_endpoint + ), + "token_endpoint": self._validate_xai_endpoint(token_endpoint), + } + + def _validate_xai_endpoint(self, url: str) -> str: + parsed = urlparse(url) + host = (parsed.hostname or "").lower() + if parsed.scheme != "https" or (host != "x.ai" and not host.endswith(".x.ai")): + raise XAIOAuthError( + f"xAI OAuth discovery returned unexpected endpoint: {url}" + ) + return url + + def _pkce_pair(self) -> Tuple[str, str]: + verifier = ( + base64.urlsafe_b64encode(secrets.token_bytes(32)).rstrip(b"=").decode() + ) + challenge = ( + base64.urlsafe_b64encode(hashlib.sha256(verifier.encode()).digest()) + .rstrip(b"=") + .decode() + ) + return verifier, challenge + + def _start_callback_server(self, state: str) -> Tuple[_CallbackServer, str]: + last_error: Optional[OSError] = None + for port in (XAI_OAUTH_REDIRECT_PORT, 0): + try: + server = _CallbackServer( + (XAI_OAUTH_REDIRECT_HOST, port), _CallbackHandler + ) + server.expected_state = state + server.callback_result = None + actual_port = server.server_address[1] + redirect_uri = f"http://{XAI_OAUTH_REDIRECT_HOST}:{actual_port}{XAI_OAUTH_REDIRECT_PATH}" + return server, redirect_uri + except OSError as exc: + last_error = exc + raise XAIOAuthError(f"Could not start xAI OAuth callback server: {last_error}") + + def _build_authorize_url( + self, + authorization_endpoint: str, + redirect_uri: str, + challenge: str, + state: str, + nonce: str, + ) -> str: + params = { + "response_type": "code", + "client_id": XAI_OAUTH_CLIENT_ID, + "redirect_uri": redirect_uri, + "scope": XAI_OAUTH_SCOPE, + "code_challenge": challenge, + "code_challenge_method": "S256", + "state": state, + "nonce": nonce, + } + return f"{authorization_endpoint}?{urlencode(params)}" + + def _wait_for_callback(self, server: _CallbackServer) -> Dict[str, Optional[str]]: + server.timeout = 1 + deadline = time.time() + XAI_OAUTH_CALLBACK_TIMEOUT_SECONDS + try: + while time.time() < deadline: + server.handle_request() + if server.callback_result is not None: + return server.callback_result + finally: + server.server_close() + raise XAIOAuthError("Timed out waiting for xAI OAuth callback") + + def _exchange_token( + self, token_endpoint: str, data: Dict[str, str] + ) -> Dict[str, Any]: + try: + response = self._client().post( + token_endpoint, + headers={ + "Accept": "application/json", + "Content-Type": "application/x-www-form-urlencoded", + }, + data=data, + ) + response.raise_for_status() + except httpx.HTTPStatusError as exc: + raise XAIOAuthError( + f"xAI OAuth token request failed: {exc.response.status_code} {exc.response.text}" + ) from exc + try: + body = response.json() + except ValueError as exc: + raise XAIOAuthError("xAI OAuth token response was not valid JSON") from exc + if not isinstance(body, dict): + raise XAIOAuthError("xAI OAuth token response was not an object") + return body + + def _build_auth_record( + self, + token_payload: Dict[str, Any], + token_endpoint: str, + fallback_refresh_token: Optional[str] = None, + ) -> Dict[str, Any]: + access_token = token_payload.get("access_token") + refresh_token = token_payload.get("refresh_token") or fallback_refresh_token + if not access_token: + raise XAIOAuthError("xAI OAuth token response missing access_token") + if not refresh_token: + raise XAIOAuthError("xAI OAuth token response missing refresh_token") + expires_in = token_payload.get("expires_in") or 3600 + try: + expires_at = int(time.time() + int(expires_in)) + except (TypeError, ValueError): + expires_at = int(time.time() + 3600) + return { + "access_token": access_token, + "refresh_token": refresh_token, + "id_token": token_payload.get("id_token"), + "token_type": token_payload.get("token_type") or "Bearer", + "token_endpoint": token_endpoint, + "expires_at": expires_at, + } + + def _refresh_tokens(self, auth_data: Dict[str, Any]) -> Dict[str, Any]: + token_endpoint = auth_data.get("token_endpoint") + if not token_endpoint: + token_endpoint = self._discover()["token_endpoint"] + token_endpoint = self._validate_xai_endpoint(token_endpoint) + refresh_token = auth_data.get("refresh_token") + if not refresh_token: + raise XAIOAuthLoginRequiredError( + "xAI OAuth refresh token missing. Run `litellm xai-oauth login`." + ) + + token_payload = self._exchange_token( + token_endpoint, + { + "grant_type": "refresh_token", + "refresh_token": refresh_token, + "client_id": XAI_OAUTH_CLIENT_ID, + }, + ) + refreshed = self._build_auth_record( + token_payload, + token_endpoint, + fallback_refresh_token=refresh_token, + ) + self._write_auth_file(refreshed) + return refreshed + + +def should_use_xai_oauth(litellm_params: Optional[Dict[str, Any]]) -> bool: + return bool((litellm_params or {}).get("use_xai_oauth")) diff --git a/litellm/llms/xai/responses/transformation.py b/litellm/llms/xai/responses/transformation.py index 55805ddaede..f81e860a8ce 100644 --- a/litellm/llms/xai/responses/transformation.py +++ b/litellm/llms/xai/responses/transformation.py @@ -3,6 +3,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union import litellm from litellm._logging import verbose_logger from litellm.constants import XAI_API_BASE +from litellm.exceptions import AuthenticationError from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.llms.xai.common_utils import XAIModelInfo from litellm.secret_managers.main import get_secret_str @@ -220,10 +221,27 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): litellm_params.api_key, legacy_generic_before_env=True ) + if not api_key: + from litellm.llms.xai.oauth import ( + XAIOAuthAuthenticator, + XAIOAuthError, + should_use_xai_oauth, + ) + + if should_use_xai_oauth(litellm_params.model_dump()): + try: + api_key = XAIOAuthAuthenticator().get_access_token() + except XAIOAuthError as exc: + raise AuthenticationError( + model=model, + llm_provider=self.custom_llm_provider.value, + message=str(exc), + ) from exc + if not api_key: raise ValueError( "XAI API key is required. Set api_key, litellm.xai_key, " - "litellm.api_key, or XAI_API_KEY." + "litellm.api_key, XAI_API_KEY, or use_xai_oauth=True." ) headers.update( @@ -244,12 +262,20 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): Returns: str: The full URL for the XAI /responses endpoint """ - api_base = ( - api_base - or litellm.api_base - or get_secret_str("XAI_API_BASE") - or XAI_API_BASE + from litellm.llms.xai.oauth import XAIOAuthAuthenticator, should_use_xai_oauth + + api_key = XAIModelInfo.get_api_key( + litellm_params.get("api_key"), legacy_generic_before_env=True ) + if should_use_xai_oauth(litellm_params) and not api_key: + api_base = XAIOAuthAuthenticator().get_api_base() + else: + api_base = ( + api_base + or litellm.api_base + or get_secret_str("XAI_API_BASE") + or XAI_API_BASE + ) # Remove trailing slashes api_base = api_base.rstrip("/") diff --git a/litellm/main.py b/litellm/main.py index 1a0d0312d73..2c416a595c4 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -1638,6 +1638,7 @@ def completion( # type: ignore # noqa: PLR0915 litellm_request_debug=kwargs.get("litellm_request_debug", False), tpm=kwargs.get("tpm"), rpm=kwargs.get("rpm"), + use_xai_oauth=kwargs.get("use_xai_oauth", False), ) cast(LiteLLMLoggingObj, logging).update_environment_variables( model=model, @@ -2134,9 +2135,6 @@ def completion( # type: ignore # noqa: PLR0915 headers = headers or litellm.headers - if extra_headers is not None: - optional_params["extra_headers"] = extra_headers - ## LOAD CONFIG - if set config = litellm.OpenAITextCompletionConfig.get_config() for k, v in config.items(): @@ -2162,6 +2160,7 @@ def completion( # type: ignore # noqa: PLR0915 _response = openai_text_completions.completion( model=model, messages=messages, + headers=headers, model_response=model_response, print_verbose=print_verbose, api_key=api_key, diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 3782da1350f..aab0e4264d0 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -24392,9 +24392,12 @@ "max_output_tokens": 8192 }, "minimax/MiniMax-M3": { - "input_cost_per_token": 6e-07, - "output_cost_per_token": 2.4e-06, - "cache_read_input_token_cost": 1.2e-07, + "input_cost_per_token": 3e-07, + "input_cost_per_token_above_512k_tokens": 6e-07, + "output_cost_per_token": 1.2e-06, + "output_cost_per_token_above_512k_tokens": 2.4e-06, + "cache_read_input_token_cost": 6e-08, + "cache_read_input_token_cost_above_512k_tokens": 1.2e-07, "litellm_provider": "minimax", "mode": "chat", "supports_function_calling": true, @@ -24403,7 +24406,7 @@ "supports_reasoning": true, "supports_system_messages": true, "supports_vision": true, - "max_input_tokens": 512000, + "max_input_tokens": 1000000, "max_output_tokens": 128000 }, "mistral.devstral-2-123b": { @@ -41646,6 +41649,7 @@ "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", + "use_openai_responses_path": true, "supported_endpoints": ["/v1/responses"], "supported_modalities": ["text", "image"], "supported_output_modalities": ["text"], @@ -41665,6 +41669,7 @@ "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", + "use_openai_responses_path": true, "supported_endpoints": ["/v1/responses"], "supported_modalities": ["text", "image"], "supported_output_modalities": ["text"], @@ -41966,5 +41971,164 @@ "/v1/audio/transcriptions" ], "supports_audio_input": true + }, + "tensormesh/Qwen/Qwen3.5-397B-A17B-FP8": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 6e-07, + "output_cost_per_token": 3.6e-06, + "cache_read_input_token_cost": 0, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/Qwen/Qwen3-Coder-480B-A35B-Instruct-FP8": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 4.5e-07, + "output_cost_per_token": 1.8e-06, + "cache_read_input_token_cost": 0, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/Qwen/Qwen3.6-27B-FP8": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 3.2e-07, + "output_cost_per_token": 3.2e-06, + "cache_read_input_token_cost": 0, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/lukealonso/GLM-5.1-NVFP4-MTP": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 1.4e-06, + "output_cost_per_token": 4.4e-06, + "cache_read_input_token_cost": 0, + "max_input_tokens": 202752, + "max_output_tokens": 202752, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/deepseek-ai/DeepSeek-V4-Flash": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 2.8e-07, + "cache_read_input_token_cost": 0, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/moonshotai/Kimi-K2.6": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 9.6e-07, + "output_cost_per_token": 4e-06, + "cache_read_input_token_cost": 0, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/MiniMaxAI/MiniMax-M2.5": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.2e-06, + "cache_read_input_token_cost": 0, + "max_input_tokens": 196608, + "max_output_tokens": 196608, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/google/gemma-4-31B-it": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 5.6e-07, + "cache_read_input_token_cost": 0, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/openai/gpt-oss-120b": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 6e-07, + "cache_read_input_token_cost": 0, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/openai/gpt-oss-20b": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 7e-08, + "output_cost_per_token": 2.8e-07, + "cache_read_input_token_cost": 0, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" } } \ No newline at end of file diff --git a/litellm/proxy/_experimental/mcp_server/mcp_context.py b/litellm/proxy/_experimental/mcp_server/mcp_context.py index a60138dd340..51918509441 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_context.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_context.py @@ -19,3 +19,9 @@ _mcp_active_toolset_id: ContextVar[Optional[str]] = ContextVar( _mcp_gateway_initialize_instructions: ContextVar[Optional[str]] = ContextVar( "_mcp_gateway_initialize_instructions", default=None ) + +# Per-request scoped server name; set in MCP HTTP/SSE handlers when the path +# identifies exactly one upstream server. Never populated from client-supplied headers. +_mcp_gateway_server_name: ContextVar[Optional[str]] = ContextVar( + "_mcp_gateway_server_name", default=None +) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 85ac6b399f4..73935beeb3a 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -354,6 +354,52 @@ def _deserialize_json_list(data: Any) -> Optional[List[Dict[str, Any]]]: ] +def _normalize_mcp_server_cost_info(mcp_info: MCPInfo) -> None: + """Coerce ``mcp_server_cost_info`` numeric fields to ``float`` at ingest. + + YAML 1.1 parses scientific notation without a decimal point (e.g. + ``7e-05``) as a string, and ``MCPServerCostInfo`` is a TypedDict with no + runtime validation, so string-typed costs flow through to the UI and + crash its ``.toFixed`` formatting. Values that cannot be coerced are + dropped with a warning instead of failing the server load. + """ + cost_info = mcp_info.get("mcp_server_cost_info") + if not isinstance(cost_info, dict): + return + + server_name = mcp_info.get("server_name") + normalized = dict(cost_info) + + default_cost = normalized.get("default_cost_per_query") + if default_cost is not None: + try: + normalized["default_cost_per_query"] = float(default_cost) + except (TypeError, ValueError): + verbose_logger.warning( + "MCP server '%s' has non-numeric default_cost_per_query %r; ignoring it", + server_name, + default_cost, + ) + del normalized["default_cost_per_query"] + + tool_costs = normalized.get("tool_name_to_cost_per_query") + if isinstance(tool_costs, dict): + normalized_tool_costs = {} + for tool_name, cost in tool_costs.items(): + try: + normalized_tool_costs[tool_name] = float(cost) + except (TypeError, ValueError): + verbose_logger.warning( + "MCP server '%s' has non-numeric cost %r for tool '%s'; ignoring it", + server_name, + cost, + tool_name, + ) + normalized["tool_name_to_cost_per_query"] = normalized_tool_costs + + mcp_info["mcp_server_cost_info"] = normalized + + def _create_sampling_callback(user_api_key_auth: Optional[Any] = None): """ Create a sampling callback for MCP ClientSession. @@ -621,6 +667,7 @@ class MCPServerManager: mcp_info["server_name"] = server_name if "description" not in mcp_info and server_config.get("description"): mcp_info["description"] = server_config.get("description") + _normalize_mcp_server_cost_info(mcp_info) # Use alias for name if present, else server_name alias = server_config.get("alias", None) @@ -1091,6 +1138,7 @@ class MCPServerManager: mcp_info["server_name"] = mcp_server.server_name or mcp_server.server_id if "description" not in mcp_info and mcp_server.description: mcp_info["description"] = mcp_server.description + _normalize_mcp_server_cost_info(mcp_info) auth_type = cast(MCPAuthType, mcp_server.auth_type) server_url = mcp_server.url diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 0477a5d3244..731493b1337 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -47,6 +47,7 @@ from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( from litellm.proxy._experimental.mcp_server.mcp_context import ( _mcp_active_toolset_id, _mcp_gateway_initialize_instructions, + _mcp_gateway_server_name, ) from litellm.proxy._experimental.mcp_server.mcp_debug import MCPDebug from litellm.proxy._experimental.mcp_server.utils import ( @@ -323,10 +324,14 @@ if MCP_AVAILABLE: notification_options=notification_options, experimental_capabilities=experimental_capabilities or {}, ) + updates: Dict[str, Any] = {} merged = _mcp_gateway_initialize_instructions.get() if merged is not None: - return opts.model_copy(update={"instructions": merged}) - return opts + updates["instructions"] = merged + scoped_server_name = _mcp_gateway_server_name.get() + if scoped_server_name is not None: + updates["server_name"] = scoped_server_name + return opts.model_copy(update=updates) if updates else opts ######################################################## ############ Initialize the MCP Server ################# @@ -1544,6 +1549,7 @@ if MCP_AVAILABLE: user_api_key_auth: Optional[UserAPIKeyAuth], mcp_servers: Optional[List[str]], client_ip: Optional[str], + scoped_server_endpoint: bool = False, ) -> AsyncIterator[None]: allowed = await _get_allowed_mcp_servers( user_api_key_auth=user_api_key_auth, @@ -1565,11 +1571,22 @@ if MCP_AVAILABLE: return_exceptions=True, ) merged = _merge_gateway_initialize_instructions(allowed_mcp_servers=allowed) - tok = _mcp_gateway_initialize_instructions.set(merged) + scoped_server_name = None + if scoped_server_endpoint and len(allowed) == 1: + scoped_server = allowed[0] + scoped_server_name = ( + scoped_server.alias + or scoped_server.server_name + or scoped_server.name + or scoped_server.server_id + ) + instructions_token = _mcp_gateway_initialize_instructions.set(merged) + server_name_token = _mcp_gateway_server_name.set(scoped_server_name) try: yield finally: - _mcp_gateway_initialize_instructions.reset(tok) + _mcp_gateway_initialize_instructions.reset(instructions_token) + _mcp_gateway_server_name.reset(server_name_token) async def _get_tools_from_mcp_servers( # noqa: PLR0915 user_api_key_auth: Optional[UserAPIKeyAuth], @@ -3620,6 +3637,7 @@ if MCP_AVAILABLE: oauth2_headers, raw_headers, ) = await extract_mcp_auth_context(scope, path) + scoped_server_endpoint = len(_get_mcp_servers_in_path(path) or []) == 1 # Extract client IP for MCP access control _client_ip = IPAddressUtils.get_mcp_client_ip(StarletteRequest(scope)) @@ -3896,6 +3914,7 @@ if MCP_AVAILABLE: user_api_key_auth, mcp_servers, _client_ip, + scoped_server_endpoint=scoped_server_endpoint, ): await target_manager.handle_request(scope, receive, local_send) if use_stateful and session_id and scope.get("method") == "DELETE": @@ -3980,6 +3999,7 @@ if MCP_AVAILABLE: oauth2_headers, raw_headers, ) = await extract_mcp_auth_context(scope, path) + scoped_server_endpoint = len(_get_mcp_servers_in_path(path) or []) == 1 # Extract client IP for MCP access control _sse_client_ip = IPAddressUtils.get_mcp_client_ip(StarletteRequest(scope)) @@ -4052,6 +4072,7 @@ if MCP_AVAILABLE: user_api_key_auth, mcp_servers, _sse_client_ip, + scoped_server_endpoint=scoped_server_endpoint, ): await sse_session_manager.handle_request(scope, receive, send) except MCPUpstreamAuthError as e: diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index f1443edf455..33a1e4179fa 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -23,6 +23,7 @@ from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( from litellm.types.integrations.slack_alerting import AlertType from litellm.types.llms.openai import ( AllMessageValues, + ResponsesAPIResponse, ) from litellm.types.mcp import ( MCPAuthType, @@ -3834,6 +3835,7 @@ PassThroughEndpointLoggingResultValues = Union[ EmbeddingResponse, VideoObject, StandardPassThroughResponseObject, + ResponsesAPIResponse, ] diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 45007861d55..6eae9d0d475 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -3516,10 +3516,13 @@ async def _virtual_key_max_budget_check( if valid_token.max_budget is not None: from litellm.proxy.proxy_server import get_current_spend + fallback_spend = valid_token.spend or 0.0 + counter_key = f"spend:key:{valid_token.token}" + # Read spend from cross-pod counter (Redis-first) or cached object (fallback) spend = await get_current_spend( - counter_key=f"spend:key:{valid_token.token}", - fallback_spend=valid_token.spend or 0.0, + counter_key=counter_key, + fallback_spend=fallback_spend, ) #################################### diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 6558543370d..b9a9f3cebb7 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1949,12 +1949,26 @@ class ProxyBaseLLMRequestProcessing: code=status.HTTP_400_BAD_REQUEST, headers=headers, ) + # Extract status_code from the exception if it carries one. + # Provider exceptions (NotFoundError, BadRequestError, GeminiError, + # VertexAIError, etc.) all have a status_code attribute reflecting + # the upstream API response. Use it to return the correct HTTP code + # instead of defaulting to 500. + _exc_status_code = getattr(e, "status_code", None) + if ( + _exc_status_code is not None + and isinstance(_exc_status_code, int) + and 400 <= _exc_status_code <= 599 + ): + _code = _exc_status_code + else: + _code = status.HTTP_500_INTERNAL_SERVER_ERROR raise ProxyException( message=getattr(e, "message", error_msg), type=getattr(e, "type", "None"), param=getattr(e, "param", "None"), openai_code=getattr(e, "code", None), - code=getattr(e, "status_code", 500), + code=_code, provider_specific_fields=getattr(e, "provider_specific_fields", None), headers=headers, ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index fc414ab7b54..033e1d0b8e7 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -1225,6 +1225,39 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): for chunk in all_chunks: yield chunk + @staticmethod + def _unmask_sse_bytes_chunk(chunk: bytes, pii_tokens: Dict[str, str]) -> bytes: + try: + text = chunk.decode("utf-8") + except UnicodeDecodeError: + return chunk + + result_lines: List[str] = [] + for line in text.split("\n"): + line = line.rstrip("\r") + if line.startswith("data: ") and line != "data: [DONE]": + raw_json = line[6:] + try: + event = json.loads(raw_json) + delta = event.get("delta") if isinstance(event, dict) else None + if ( + isinstance(delta, dict) + and event.get("type") == "content_block_delta" + and delta.get("type") == "text_delta" + and isinstance(delta.get("text"), str) + ): + unmasked = _OPTIONAL_PresidioPIIMasking._unmask_pii_text( + delta["text"], pii_tokens + ) + if unmasked != delta["text"]: + event["delta"]["text"] = unmasked + line = "data: " + json.dumps(event, ensure_ascii=False) + except (json.JSONDecodeError, KeyError, TypeError): + pass + result_lines.append(line) + + return "\n".join(result_lines).encode("utf-8") + async def _stream_pii_unmasking( self, response: Any, @@ -1237,13 +1270,19 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): from litellm.main import stream_chunk_builder from litellm.types.utils import ModelResponse + metadata = (request_data.get("metadata") or {}) if request_data else {} + pii_tokens: Dict[str, str] = metadata.get("pii_tokens", {}) + remaining_chunks: List[ModelResponseStream] = [] try: async for chunk in response: if isinstance(chunk, ModelResponseStream): remaining_chunks.append(chunk) elif isinstance(chunk, bytes): - yield chunk # type: ignore[misc] + if pii_tokens: + yield self._unmask_sse_bytes_chunk(chunk, pii_tokens) # type: ignore[misc] + else: + yield chunk # type: ignore[misc] continue if not remaining_chunks: diff --git a/litellm/proxy/hooks/litellm_skills/main.py b/litellm/proxy/hooks/litellm_skills/main.py index 21e8bbbd308..77ed3493a0c 100644 --- a/litellm/proxy/hooks/litellm_skills/main.py +++ b/litellm/proxy/hooks/litellm_skills/main.py @@ -19,7 +19,7 @@ Usage: response = await litellm.acompletion( model="gpt-4o-mini", messages=[{"role": "user", "content": "Create a bouncing ball GIF"}], - container={"skills": [{"skill_id": "litellm:skill_abc123"}]}, + container={"skills": [{"skill_id": "litellm_skill_abc123"}]}, ) # Response includes file_ids for generated files """ @@ -31,6 +31,7 @@ from typing import Any, Dict, List, Optional, Union from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger +from litellm.llms.litellm_proxy.skills.constants import LITELLM_SKILL_ID_PREFIX from litellm.llms.litellm_proxy.skills.prompt_injection import ( SkillPromptInjectionHandler, ) @@ -43,7 +44,7 @@ class SkillsInjectionHook(CustomLogger): Pre/Post-call hook that processes skills from container.skills parameter. Pre-call (async_pre_call_hook): - - Skills with 'litellm:' prefix are fetched from LiteLLM DB + - Skills with 'litellm_skill_' prefix are fetched from LiteLLM DB - For Anthropic models: native skills pass through, LiteLLM skills converted to tools - For non-Anthropic models: LiteLLM skills are converted to tools + execute_code tool @@ -78,7 +79,7 @@ class SkillsInjectionHook(CustomLogger): Process skills from container.skills before the LLM call. 1. Check if container.skills exists in request - 2. Separate skills by prefix (litellm: vs native) + 2. Separate skills by prefix (litellm_skill_ vs native) 3. Fetch LiteLLM skills from database 4. For Anthropic: keep native skills in container 5. For non-Anthropic: convert LiteLLM skills to tools, inject content, add execute_code @@ -108,7 +109,7 @@ class SkillsInjectionHook(CustomLogger): continue skill_id = skill.get("skill_id", "") - if skill_id.startswith("litellm_"): + if skill_id.startswith(LITELLM_SKILL_ID_PREFIX): # Fetch from LiteLLM DB db_skill = await self._fetch_skill_from_db( skill_id, @@ -287,7 +288,7 @@ class SkillsInjectionHook(CustomLogger): Fetch a skill from the LiteLLM database. Args: - skill_id: The skill ID (without 'litellm:' prefix) + skill_id: The skill ID (including the 'litellm_skill_' prefix) Returns: LiteLLM_SkillsTable or None if not found @@ -382,10 +383,10 @@ class SkillsInjectionHook(CustomLogger): has_executable_tool = False for tc in tool_calls: tool_name = tc.get("name", "") - # Execute if it's litellm_code_execution OR a skill tool (skill_xxx) + # Execute if it's litellm_code_execution OR a skill tool (litellm_skill_xxx) if ( tool_name == LiteLLMInternalTools.CODE_EXECUTION.value - or tool_name.startswith("skill_") + or tool_name.startswith(LITELLM_SKILL_ID_PREFIX) ): has_executable_tool = True break @@ -543,7 +544,7 @@ class SkillsInjectionHook(CustomLogger): result = await self._execute_code( code, skill_files, executor, generated_files ) - elif tool_name.startswith("skill_"): + elif tool_name.startswith(LITELLM_SKILL_ID_PREFIX): # Skill tool - execute the skill's code result = await self._execute_skill_tool( tool_name, tool_input, skill_files, executor, generated_files diff --git a/litellm/proxy/hooks/parallel_request_limiter.py b/litellm/proxy/hooks/parallel_request_limiter.py index b622241dfa5..874e5aa1939 100644 --- a/litellm/proxy/hooks/parallel_request_limiter.py +++ b/litellm/proxy/hooks/parallel_request_limiter.py @@ -7,7 +7,7 @@ from pydantic import BaseModel from typing_extensions import TypedDict import litellm -from litellm import DualCache, ModelResponse +from litellm import DualCache, EmbeddingResponse, ModelResponse, TextCompletionResponse from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import _get_parent_otel_span_from_kwargs @@ -570,7 +570,9 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): total_tokens = 0 - if isinstance(response_obj, ModelResponse): + if isinstance( + response_obj, (ModelResponse, EmbeddingResponse, TextCompletionResponse) + ): total_tokens = response_obj.usage.total_tokens # type: ignore # ------------ @@ -659,7 +661,10 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): if user_api_key_user_id is not None: total_tokens = 0 - if isinstance(response_obj, ModelResponse): + if isinstance( + response_obj, + (ModelResponse, EmbeddingResponse, TextCompletionResponse), + ): total_tokens = response_obj.usage.total_tokens # type: ignore request_count_api_key = ( @@ -692,7 +697,10 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): if user_api_key_team_id is not None: total_tokens = 0 - if isinstance(response_obj, ModelResponse): + if isinstance( + response_obj, + (ModelResponse, EmbeddingResponse, TextCompletionResponse), + ): total_tokens = response_obj.usage.total_tokens # type: ignore request_count_api_key = ( @@ -725,7 +733,10 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): if user_api_key_end_user_id is not None: total_tokens = 0 - if isinstance(response_obj, ModelResponse): + if isinstance( + response_obj, + (ModelResponse, EmbeddingResponse, TextCompletionResponse), + ): total_tokens = response_obj.usage.total_tokens # type: ignore request_count_api_key = ( diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 62751fb68a4..6b70cea65a3 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -39,7 +39,13 @@ from litellm.proxy.common_utils.proxy_rate_limit_error import ( from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.llms.openai import BaseLiteLLMOpenAIResponseObject -from litellm.types.utils import CallTypes, ModelResponse, Usage +from litellm.types.utils import ( + CallTypes, + EmbeddingResponse, + ModelResponse, + TextCompletionResponse, + Usage, +) if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -2736,9 +2742,14 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # Get total tokens from response total_tokens = 0 - # spot fix for /responses api - if isinstance(response_obj, ModelResponse) or isinstance( - response_obj, BaseLiteLLMOpenAIResponseObject + if isinstance( + response_obj, + ( + ModelResponse, + EmbeddingResponse, + TextCompletionResponse, + BaseLiteLLMOpenAIResponseObject, + ), ): _usage = getattr(response_obj, "usage", None) total_tokens = self._get_total_tokens_from_usage( diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 8f606fdf90d..eba16c077b0 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1850,11 +1850,10 @@ async def prepare_key_update_data( if "budget_duration" in non_default_values: budget_duration = non_default_values.pop("budget_duration") - if ( - budget_duration - and (isinstance(budget_duration, str)) - and len(budget_duration) > 0 - ): + if budget_duration is None: + non_default_values["budget_duration"] = None + non_default_values["budget_reset_at"] = None + elif isinstance(budget_duration, str) and len(budget_duration) > 0: from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time key_reset_at = get_budget_reset_time(budget_duration=budget_duration) @@ -2518,7 +2517,7 @@ async def update_key_fn( # noqa: PLR0915 }, ) - data_json: dict = data.model_dump(exclude_unset=True, exclude_none=True) + data_json: dict = data.model_dump(exclude_unset=True) key = data_json.pop("key") # get the row from db @@ -2588,6 +2587,17 @@ async def update_key_fn( # noqa: PLR0915 proxy_logging_obj=proxy_logging_obj, ) + if data.spend is not None: + try: + from litellm.proxy.proxy_server import _invalidate_spend_counter + + token_to_invalidate = _hash_token_if_needed(key) + await _invalidate_spend_counter( + counter_key=f"spend:key:{token_to_invalidate}" + ) + except Exception: + pass + asyncio.create_task( KeyManagementEventHooks.async_key_updated_hook( data=data, @@ -4775,6 +4785,13 @@ async def reset_key_spend_fn( proxy_logging_obj=proxy_logging_obj, ) + try: + from litellm.proxy.proxy_server import _invalidate_spend_counter + + await _invalidate_spend_counter(counter_key=f"spend:key:{hashed_api_key}") + except Exception: + pass + max_budget = updated_key.max_budget budget_reset_at = updated_key.budget_reset_at diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 435a8cae379..c894813ada4 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -357,6 +357,42 @@ class TeamMemberBudgetHandler: data_dict.pop("team_member_rpm_limit", None) data_dict.pop("team_member_tpm_limit", None) + @staticmethod + async def clear_team_member_budget_fields( + team_table: LiteLLM_TeamTable, + user_api_key_dict: "UserAPIKeyAuth", + updated_kv: dict, + explicitly_set_fields: set, + ) -> dict: + """Clear explicitly-nulled fields on the team member budget row.""" + from litellm.proxy._types import BudgetNewRequest + from litellm.proxy.management_endpoints.budget_management_endpoints import ( + update_budget, + ) + + if team_table.metadata is None: + team_table.metadata = {} + + team_member_budget_id = team_table.metadata.get("team_member_budget_id") + if team_member_budget_id is not None and isinstance(team_member_budget_id, str): + budget_request = BudgetNewRequest(budget_id=team_member_budget_id) + if "team_member_budget" in explicitly_set_fields: + budget_request.max_budget = None + if "team_member_budget_duration" in explicitly_set_fields: + budget_request.budget_duration = None + budget_request.budget_reset_at = None + if "team_member_rpm_limit" in explicitly_set_fields: + budget_request.rpm_limit = None + if "team_member_tpm_limit" in explicitly_set_fields: + budget_request.tpm_limit = None + await update_budget( + budget_obj=budget_request, + user_api_key_dict=user_api_key_dict, + ) + + TeamMemberBudgetHandler._clean_team_member_fields(updated_kv) + return updated_kv + @staticmethod async def backfill_team_member_budget_entries( team_id: str, @@ -1872,11 +1908,25 @@ async def update_team( # noqa: PLR0915 # Check budget_duration and budget_reset_at _set_budget_reset_at(data, updated_kv) - if TeamMemberBudgetHandler.should_create_budget( - team_member_budget=data.team_member_budget, - team_member_rpm_limit=data.team_member_rpm_limit, - team_member_tpm_limit=data.team_member_tpm_limit, - team_member_budget_duration=data.team_member_budget_duration, + _team_member_fields_in_request = { + field + for field in [ + "team_member_budget", + "team_member_rpm_limit", + "team_member_tpm_limit", + "team_member_budget_duration", + ] + if field in updated_kv + } + + if ( + _team_member_fields_in_request + and TeamMemberBudgetHandler.should_create_budget( + team_member_budget=data.team_member_budget, + team_member_rpm_limit=data.team_member_rpm_limit, + team_member_tpm_limit=data.team_member_tpm_limit, + team_member_budget_duration=data.team_member_budget_duration, + ) ): updated_kv = await TeamMemberBudgetHandler.upsert_team_member_budget_table( team_table=existing_team_row, @@ -1899,6 +1949,13 @@ async def update_team( # noqa: PLR0915 team_member_budget_id=_backfill_budget_id, prisma_client=prisma_client, ) + elif _team_member_fields_in_request: + updated_kv = await TeamMemberBudgetHandler.clear_team_member_budget_fields( + team_table=existing_team_row, + user_api_key_dict=user_api_key_dict, + updated_kv=updated_kv, + explicitly_set_fields=_team_member_fields_in_request, + ) else: TeamMemberBudgetHandler._clean_team_member_fields(updated_kv) @@ -1987,6 +2044,8 @@ def _set_budget_reset_at(data: UpdateTeamRequest, updated_kv: dict) -> None: reset_at = get_budget_reset_time(budget_duration=data.budget_duration) updated_kv["budget_reset_at"] = reset_at + elif "budget_duration" in updated_kv and updated_kv["budget_duration"] is None: + updated_kv["budget_reset_at"] = None if data.budget_limits is not None and len(data.budget_limits) > 0: from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 7c3a6f19013..c7db818a07e 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -426,12 +426,7 @@ async def mistral_proxy_route( ) ## check for streaming - is_streaming_request = False - # anthropic is streaming when 'stream' = True is in the body - if request.method == "POST": - _request_body = await request.json() - if _request_body.get("stream"): - is_streaming_request = True + is_streaming_request = await is_streaming_request_fn(request) ## CREATE PASS-THROUGH endpoint_func = create_pass_through_route( diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py index 6dd1f8548eb..9f353226dd0 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py @@ -5,7 +5,7 @@ Handles cost tracking and logging for OpenAI passthrough endpoints, specifically """ from datetime import datetime -from typing import List, Optional, Union +from typing import List, Optional, Tuple, Union from urllib.parse import urlparse import httpx @@ -18,6 +18,7 @@ from litellm.litellm_core_utils.litellm_logging import ( ) from litellm.llms.openai.openai import OpenAIConfig from litellm.llms.openai.openai import OpenAIConfig as OpenAIConfigType +from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.proxy._types import PassThroughEndpointLoggingTypedDict from litellm.proxy.pass_through_endpoints.llm_provider_handlers.base_passthrough_logging_handler import ( BasePassthroughLoggingHandler, @@ -29,6 +30,7 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import ( EndpointType, PassthroughStandardLoggingPayload, ) +from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.utils import ImageResponse, LlmProviders, PassthroughCallTypes from litellm.utils import ModelResponse, TextCompletionResponse @@ -236,6 +238,42 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): ) return 0.0 + @staticmethod + def _build_responses_api_response_and_cost( + model: str, + httpx_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + custom_llm_provider: str, + ) -> Tuple[ResponsesAPIResponse, float]: + """Transform a Responses API raw response into a ResponsesAPIResponse + and compute its cost. + + The Responses API has a different on-the-wire shape from chat + completions (`output: [...]` instead of `choices: [...]`), so the + chat-completions `transform_response` raises KeyError 'choices' on + a Responses payload. Use the dedicated Responses-API transformer + (`OpenAIResponsesAPIConfig.transform_response_api_response`) here. + + Returns (litellm_model_response, response_cost) — symmetric with the + chat-completions branch which produces the same two values inline, + and analogous to the image branches' `_calculate_image_*_cost` helpers + (which return cost only because the image-response object is trivial + to build inline; the Responses payload needs a real transformer). + """ + responses_config = OpenAIResponsesAPIConfig() + litellm_model_response = responses_config.transform_response_api_response( + model=model, + raw_response=httpx_response, + logging_obj=logging_obj, + ) + response_cost = litellm.completion_cost( + completion_response=litellm_model_response, + model=model, + custom_llm_provider=custom_llm_provider, + call_type="responses", + ) + return litellm_model_response, response_cost + @staticmethod def openai_passthrough_handler( # noqa: PLR0915 httpx_response: httpx.Response, @@ -301,7 +339,12 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): try: response_cost = 0.0 litellm_model_response: Optional[ - Union[ModelResponse, TextCompletionResponse, ImageResponse] + Union[ + ModelResponse, + TextCompletionResponse, + ImageResponse, + ResponsesAPIResponse, + ] ] = None handler_instance = OpenAIPassthroughLoggingHandler() @@ -384,29 +427,18 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): litellm_model_response._hidden_params = {} litellm_model_response._hidden_params["response_cost"] = response_cost elif is_responses: - # Handle responses API cost calculation - provider_config = handler_instance.get_provider_config(model=model) - existing_litellm_params = kwargs.get("litellm_params", {}) or {} - litellm_model_response = provider_config.transform_response( - raw_response=httpx_response, - model_response=litellm.ModelResponse(), + # Responses-API cost tracking — see + # `_build_responses_api_response_and_cost` for why this needs + # a dedicated transformer (the chat-completions transform + # crashes on the Responses payload shape). + ( + litellm_model_response, + response_cost, + ) = OpenAIPassthroughLoggingHandler._build_responses_api_response_and_cost( model=model, - messages=request_body.get("messages", []), + httpx_response=httpx_response, logging_obj=logging_obj, - optional_params=request_body.get("optional_params", {}), - api_key="", - request_data=request_body, - encoding=litellm.encoding, - json_mode=False, - litellm_params=existing_litellm_params, - ) - - # Calculate cost using LiteLLM's cost calculator with responses call type - response_cost = litellm.completion_cost( - completion_response=litellm_model_response, - model=model, custom_llm_provider=custom_llm_provider, - call_type="responses", ) # Update kwargs with cost information diff --git a/litellm/proxy/pass_through_endpoints/success_handler.py b/litellm/proxy/pass_through_endpoints/success_handler.py index af1d39da020..46043d10a06 100644 --- a/litellm/proxy/pass_through_endpoints/success_handler.py +++ b/litellm/proxy/pass_through_endpoints/success_handler.py @@ -458,7 +458,16 @@ class PassThroughEndpointLogging: return False def _is_supported_openai_endpoint(self, url_route: str) -> bool: - """Check if the OpenAI endpoint is supported by the passthrough logging handler.""" + """Check if the OpenAI endpoint is supported by the passthrough logging handler. + + The Responses API route is included because + `openai_passthrough_handler` has a dedicated `elif is_responses:` + branch that knows how to extract usage + cost from the + Responses-API on-the-wire shape. Without including it here, the + outer dispatch filters Responses calls out before reaching the + handler — the inner branch is then unreachable and Responses + calls land in `LiteLLM_SpendLogs` with zero tokens / zero spend. + """ from .llm_provider_handlers.openai_passthrough_logging_handler import ( OpenAIPassthroughLoggingHandler, ) @@ -469,6 +478,7 @@ class PassThroughEndpointLogging: url_route ) or OpenAIPassthroughLoggingHandler.is_openai_image_editing_route(url_route) + or OpenAIPassthroughLoggingHandler.is_openai_responses_route(url_route) ) def _set_cost_per_request( diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index e4567b9f494..ae831ef1b53 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -555,6 +555,7 @@ class ProxyInitializationHelpers: @click.command() +@click.argument("cli_args", nargs=-1) @click.option( "--host", default="0.0.0.0", help="Host for the server to listen on.", envvar="HOST" ) @@ -808,6 +809,7 @@ class ProxyInitializationHelpers: help="Enable uvicorn hot reload (dev only). Also reloads when the --config YAML file changes. Incompatible with --num_workers>1, --run_gunicorn, and --run_hypercorn.", ) def run_server( # noqa: PLR0915 + cli_args, host, port, api_base, @@ -854,6 +856,20 @@ def run_server( # noqa: PLR0915 use_v2_migration_resolver: bool, reload: bool, ): + if cli_args: + if cli_args == ("xai-oauth", "login"): + from litellm.llms.xai.oauth import XAIOAuthAuthenticator + + authenticator = XAIOAuthAuthenticator() + auth_data = authenticator.login() + click.echo( + f"xAI OAuth login successful. Credentials saved to {authenticator.auth_file}." + ) + if auth_data.get("expires_at"): + click.echo(f"Access token expires at {auth_data['expires_at']}.") + return + raise click.UsageError(f"Unknown command: {' '.join(cli_args)}") + if setup: from litellm.setup_wizard import run_setup_wizard diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index e50e58838e6..37a0285b196 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -4089,6 +4089,8 @@ class ProxyConfig: verbose_proxy_logger.debug( f"litellm.post_call_rules: {litellm.post_call_rules}" ) + elif key == "max_budget": + litellm.max_budget = float(value) elif key == "max_internal_user_budget": litellm.max_internal_user_budget = float(value) # type: ignore elif key == "default_max_internal_user_budget": diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index f651e6e5f7b..aa85be6671a 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -3405,7 +3405,7 @@ async def ui_view_session_spend_logs( session_id, status, mcp_namespaced_tool_name, agent_id FROM "LiteLLM_SpendLogs" WHERE session_id = $1 - ORDER BY "startTime" ASC + ORDER BY "startTime" DESC LIMIT $2 OFFSET $3 """ result = await prisma_client.db.query_raw( diff --git a/litellm/router.py b/litellm/router.py index d0f4e5ff44d..8966a2fc191 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1658,6 +1658,67 @@ class Router: f"Dictionary '{fallback_dict}' must have exactly one key, but has {len(fallback_dict)} keys." ) + def _add_encrypted_content_affinity_check( + self, enable_global_affinity: bool + ) -> None: + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + def _move_before_deployment_affinity( + callback_list: List[Any], + callback_to_move: EncryptedContentAffinityCheck, + ) -> None: + if callback_to_move not in callback_list: + return + callback_list.remove(callback_to_move) + insert_index = next( + ( + idx + for idx, callback in enumerate(callback_list) + if isinstance(callback, DeploymentAffinityCheck) + ), + len(callback_list), + ) + callback_list.insert(insert_index, callback_to_move) + + if ( + enable_global_affinity + or EncryptedContentAffinityCheck.has_model_group_affinity_enabled( + self.model_group_affinity_config + ) + ): + if self.optional_callbacks is None: + self.optional_callbacks = [] + + existing_ec_callback: Optional[EncryptedContentAffinityCheck] = None + for cb in self.optional_callbacks: + if isinstance(cb, EncryptedContentAffinityCheck): + existing_ec_callback = cb + break + + if existing_ec_callback is not None: + existing_ec_callback.router = self + existing_ec_callback.enable_global_affinity = ( + existing_ec_callback.enable_global_affinity + or enable_global_affinity + ) + existing_ec_callback.model_group_affinity_config = ( + self.model_group_affinity_config or {} + ) + ec_callback = existing_ec_callback + else: + ec_callback = EncryptedContentAffinityCheck( + router=self, + enable_global_affinity=enable_global_affinity, + model_group_affinity_config=self.model_group_affinity_config, + ) + self.optional_callbacks.append(ec_callback) + litellm.logging_callback_manager.add_litellm_callback(ec_callback) + + _move_before_deployment_affinity(self.optional_callbacks, ec_callback) + _move_before_deployment_affinity(litellm.callbacks, ec_callback) + def add_optional_pre_call_checks( self, optional_pre_call_checks: Optional[OptionalPreCallChecks] ): @@ -1721,22 +1782,11 @@ class Router: # --------------------------------------------------------------------- # Encrypted content affinity # --------------------------------------------------------------------- - if "encrypted_content_affinity" in optional_pre_call_checks: - from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( - EncryptedContentAffinityCheck, + self._add_encrypted_content_affinity_check( + enable_global_affinity=( + "encrypted_content_affinity" in optional_pre_call_checks ) - - if self.optional_callbacks is None: - self.optional_callbacks = [] - - already_registered = any( - isinstance(cb, EncryptedContentAffinityCheck) - for cb in self.optional_callbacks - ) - if not already_registered: - ec_callback = EncryptedContentAffinityCheck(router=self) - self.optional_callbacks.append(ec_callback) - litellm.logging_callback_manager.add_litellm_callback(ec_callback) + ) # --------------------------------------------------------------------- # Remaining optional pre-call checks @@ -8471,6 +8521,13 @@ class Router: credential_values.get("api_key") or deployment.litellm_params.api_key ) + if api_key is None: + verbose_router_logger.debug( + "Skipping pass-through credential setup for deployment model=%s, custom_llm_provider=%s; no api_key set. Providers like bedrock resolve credentials at request time.", + model, + custom_llm_provider, + ) + return passthrough_endpoint_router.set_pass_through_credentials( custom_llm_provider=custom_llm_provider, api_base=api_base, diff --git a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py index 148b7fce0ee..d3e7e2ffa34 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -39,7 +39,12 @@ class DeploymentAffinityCheck(CustomLogger): CACHE_KEY_PREFIX = "deployment_affinity:v1" VALID_FLAGS = frozenset( - {"deployment_affinity", "responses_api_deployment_check", "session_affinity"} + { + "deployment_affinity", + "responses_api_deployment_check", + "session_affinity", + "encrypted_content_affinity", + } ) def __init__( diff --git a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py index 4ed19c5cd26..5fd2be9c6dd 100644 --- a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py @@ -37,7 +37,7 @@ Safe to enable globally: """ import time -from typing import TYPE_CHECKING, Any, List, Optional, cast +from typing import TYPE_CHECKING, Any, Dict, List, Optional, cast import httpx @@ -64,17 +64,45 @@ class EncryptedContentAffinityCheck(CustomLogger): The ``model_id`` is decoded directly from the litellm-encoded item IDs – no caching or TTL management needed. - Wired via ``Router(optional_pre_call_checks=["encrypted_content_affinity"])``. + Wired via ``Router(optional_pre_call_checks=["encrypted_content_affinity"])`` or + per-model group ``model_group_affinity_config``. """ - def __init__(self, router: Optional["Router"] = None) -> None: + def __init__( + self, + router: Optional["Router"] = None, + enable_global_affinity: bool = True, + model_group_affinity_config: Optional[Dict[str, List[str]]] = None, + ) -> None: super().__init__() self.router = router + self.enable_global_affinity = enable_global_affinity + self.model_group_affinity_config: Dict[str, List[str]] = ( + model_group_affinity_config or {} + ) # ------------------------------------------------------------------ # Helpers # ------------------------------------------------------------------ + @staticmethod + def has_model_group_affinity_enabled( + model_group_affinity_config: Optional[Dict[str, List[str]]], + ) -> bool: + if not model_group_affinity_config: + return False + + return any( + "encrypted_content_affinity" in checks + for checks in model_group_affinity_config.values() + ) + + def _is_enabled_for_model_group(self, model_group: str) -> bool: + group_checks = self.model_group_affinity_config.get(model_group) + return self.enable_global_affinity or ( + group_checks is not None and "encrypted_content_affinity" in group_checks + ) + @staticmethod def _extract_model_id_from_input(request_input: Any) -> Optional[str]: """ @@ -213,6 +241,8 @@ class EncryptedContentAffinityCheck(CustomLogger): """ request_kwargs = request_kwargs or {} typed_healthy_deployments = cast(List[dict], healthy_deployments) + if not self._is_enabled_for_model_group(model): + return typed_healthy_deployments # Signal to the response post-processor that encrypted item IDs should be # encoded in the output of this request. Only set the flag when diff --git a/litellm/types/router.py b/litellm/types/router.py index ef7eb05d087..ed858557a61 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -220,6 +220,10 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): use_in_pass_through: Optional[bool] = False use_litellm_proxy: Optional[bool] = False use_chat_completions_api: Optional[bool] = None + use_xai_oauth: Optional[bool] = Field( + default=False, + description="Use stored xAI OAuth credentials when no xAI API key is configured.", + ) model_config = ConfigDict(extra="allow", arbitrary_types_allowed=True) merge_reasoning_content_in_choices: Optional[bool] = False model_info: Optional[Dict] = None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index c3ea605dd9e..21eb0c9a173 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -197,6 +197,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): ] # OpenAI priority service tier pricing cache_read_input_token_cost_above_200k_tokens: Optional[float] cache_read_input_token_cost_above_272k_tokens: Optional[float] + cache_read_input_token_cost_above_512k_tokens: Optional[float] input_cost_per_character: Optional[float] # only for vertex ai models input_cost_per_audio_token: Optional[float] input_cost_per_token_above_128k_tokens: Optional[float] # only for vertex ai models @@ -206,6 +207,9 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): input_cost_per_token_above_272k_tokens: Optional[ float ] # GPT-5.4/5.4-pro: prompts >272K priced at 2x input + input_cost_per_token_above_512k_tokens: Optional[ + float + ] # MiniMax-M3: prompts >512K priced at 2x input input_cost_per_character_above_128k_tokens: Optional[ float ] # only for vertex ai models @@ -239,6 +243,9 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): output_cost_per_token_above_272k_tokens: Optional[ float ] # GPT-5.4/5.4-pro: prompts >272K priced at 1.5x output + output_cost_per_token_above_512k_tokens: Optional[ + float + ] # MiniMax-M3: prompts >512K priced at 2x output output_cost_per_character_above_128k_tokens: Optional[ float ] # only for vertex ai models @@ -3217,6 +3224,7 @@ all_litellm_params = ( "search_tool_name", "order", "enable_json_schema_validation", + "use_xai_oauth", ] + list(StandardCallbackDynamicParams.__annotations__.keys()) + list(CustomPricingLiteLLMParams.model_fields.keys()) diff --git a/litellm/utils.py b/litellm/utils.py index 8d9d0a409c6..4f4e8d8cb9e 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5775,6 +5775,7 @@ def _get_model_info_helper( # noqa: PLR0915 ] split_model = potential_model_names["split_model"] custom_llm_provider = potential_model_names["custom_llm_provider"] + model_cost_custom_llm_provider = custom_llm_provider ######################### provider_config: Optional[BaseLLMModelInfo] = None if custom_llm_provider and custom_llm_provider in LlmProvidersSet: @@ -5840,7 +5841,8 @@ def _get_model_info_helper( # noqa: PLR0915 key = _matched_key _model_info = _get_model_info_from_model_cost(key=cast(str, key)) if not _check_provider_match( - model_info=_model_info, custom_llm_provider=custom_llm_provider + model_info=_model_info, + custom_llm_provider=model_cost_custom_llm_provider, ): _model_info = None if _model_info is None: @@ -5849,7 +5851,8 @@ def _get_model_info_helper( # noqa: PLR0915 key = _matched_key _model_info = _get_model_info_from_model_cost(key=cast(str, key)) if not _check_provider_match( - model_info=_model_info, custom_llm_provider=custom_llm_provider + model_info=_model_info, + custom_llm_provider=model_cost_custom_llm_provider, ): _model_info = None if _model_info is None: @@ -5858,7 +5861,8 @@ def _get_model_info_helper( # noqa: PLR0915 key = _matched_key _model_info = _get_model_info_from_model_cost(key=cast(str, key)) if not _check_provider_match( - model_info=_model_info, custom_llm_provider=custom_llm_provider + model_info=_model_info, + custom_llm_provider=model_cost_custom_llm_provider, ): _model_info = None if _model_info is None: @@ -5867,7 +5871,8 @@ def _get_model_info_helper( # noqa: PLR0915 key = _matched_key _model_info = _get_model_info_from_model_cost(key=cast(str, key)) if not _check_provider_match( - model_info=_model_info, custom_llm_provider=custom_llm_provider + model_info=_model_info, + custom_llm_provider=model_cost_custom_llm_provider, ): _model_info = None if _model_info is None: @@ -5876,7 +5881,8 @@ def _get_model_info_helper( # noqa: PLR0915 key = _matched_key _model_info = _get_model_info_from_model_cost(key=cast(str, key)) if not _check_provider_match( - model_info=_model_info, custom_llm_provider=custom_llm_provider + model_info=_model_info, + custom_llm_provider=model_cost_custom_llm_provider, ): _model_info = None @@ -5884,7 +5890,6 @@ def _get_model_info_helper( # noqa: PLR0915 raise ValueError( "This model isn't mapped yet. Add it here - https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json" ) - _input_cost_per_token: Optional[float] = _model_info.get( "input_cost_per_token" ) @@ -5936,6 +5941,9 @@ def _get_model_info_helper( # noqa: PLR0915 cache_read_input_token_cost_above_272k_tokens=_model_info.get( "cache_read_input_token_cost_above_272k_tokens", None ), + cache_read_input_token_cost_above_512k_tokens=_model_info.get( + "cache_read_input_token_cost_above_512k_tokens", None + ), cache_read_input_token_cost_flex=_model_info.get( "cache_read_input_token_cost_flex", None ), @@ -5957,6 +5965,9 @@ def _get_model_info_helper( # noqa: PLR0915 input_cost_per_token_above_272k_tokens=_model_info.get( "input_cost_per_token_above_272k_tokens", None ), + input_cost_per_token_above_512k_tokens=_model_info.get( + "input_cost_per_token_above_512k_tokens", None + ), input_cost_per_query=_model_info.get("input_cost_per_query", None), input_cost_per_second=_model_info.get("input_cost_per_second", None), input_cost_per_audio_token=_model_info.get( @@ -6012,6 +6023,9 @@ def _get_model_info_helper( # noqa: PLR0915 output_cost_per_token_above_272k_tokens=_model_info.get( "output_cost_per_token_above_272k_tokens", None ), + output_cost_per_token_above_512k_tokens=_model_info.get( + "output_cost_per_token_above_512k_tokens", None + ), output_cost_per_second=_model_info.get("output_cost_per_second", None), output_cost_per_second_1080p=_model_info.get( "output_cost_per_second_1080p", None @@ -8922,14 +8936,33 @@ class ProviderConfigManager: elif litellm.LlmProviders.HOSTED_VLLM == provider: return litellm.HostedVLLMResponsesAPIConfig() elif litellm.LlmProviders.BEDROCK_MANTLE == provider: - # Only OpenAI gpt frontier models (gpt-5.x, and future gpt-6 etc.) are - # served on the /openai/v1/responses path. gpt-oss and every non-OpenAI - # model on Mantle (nvidia, mistral, google, zai, ...) are chat-completions - # only and 400 on that path, so they fall through to None to keep the - # chat-completions emulation (see litellm/responses/main.py "config is None"). - model_lower = model.lower() if model else "" - if "openai.gpt-" in model_lower and "gpt-oss" not in model_lower: - return litellm.BedrockMantleResponsesAPIConfig() + # Mantle serves Responses on two upstream paths. A model takes the + # /openai/v1/responses path when its price-map entry declares + # use_openai_responses_path (data-driven, so a non-gpt-named frontier + # model can be onboarded by JSON alone), or, as a fallback needing no + # price-map entry, when its name matches the openai.gpt- frontier + # convention (minus gpt-oss) -- this keeps a future gpt-6 routing + # correctly before its entry loads. Any other model declared + # mode=responses takes the standard /v1/responses path. Everything + # else returns None and keeps the chat-completions emulation (see + # responses/main.py "config is None"). + if not model: + return None + model_lower = model.lower() + entry = litellm.model_cost.get(f"bedrock_mantle/{model}", {}) + on_openai_path = entry.get("use_openai_responses_path") is True + name_is_frontier = ( + "openai.gpt-" in model_lower and "gpt-oss" not in model_lower + ) + if on_openai_path or name_is_frontier: + return litellm.BedrockMantleResponsesAPIConfig(use_openai_path=True) + try: + if get_model_info(model, "bedrock_mantle").get("mode") == "responses": + return litellm.BedrockMantleResponsesAPIConfig( + use_openai_path=False + ) + except Exception: + pass return None return None diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 85cb06b7f19..f0b2432ddc8 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -24392,9 +24392,12 @@ "max_output_tokens": 8192 }, "minimax/MiniMax-M3": { - "input_cost_per_token": 6e-07, - "output_cost_per_token": 2.4e-06, - "cache_read_input_token_cost": 1.2e-07, + "input_cost_per_token": 3e-07, + "input_cost_per_token_above_512k_tokens": 6e-07, + "output_cost_per_token": 1.2e-06, + "output_cost_per_token_above_512k_tokens": 2.4e-06, + "cache_read_input_token_cost": 6e-08, + "cache_read_input_token_cost_above_512k_tokens": 1.2e-07, "litellm_provider": "minimax", "mode": "chat", "supports_function_calling": true, @@ -24403,7 +24406,7 @@ "supports_reasoning": true, "supports_system_messages": true, "supports_vision": true, - "max_input_tokens": 512000, + "max_input_tokens": 1000000, "max_output_tokens": 128000 }, "mistral.devstral-2-123b": { @@ -41686,6 +41689,7 @@ "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", + "use_openai_responses_path": true, "supported_endpoints": ["/v1/responses"], "supported_modalities": ["text", "image"], "supported_output_modalities": ["text"], @@ -41705,6 +41709,7 @@ "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", + "use_openai_responses_path": true, "supported_endpoints": ["/v1/responses"], "supported_modalities": ["text", "image"], "supported_output_modalities": ["text"], @@ -42178,5 +42183,164 @@ "source": "https://soniox.com/pricing", "supported_endpoints": ["/v1/audio/transcriptions"], "supports_audio_input": true + }, + "tensormesh/Qwen/Qwen3.5-397B-A17B-FP8": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 6e-07, + "output_cost_per_token": 3.6e-06, + "cache_read_input_token_cost": 0, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/Qwen/Qwen3-Coder-480B-A35B-Instruct-FP8": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 4.5e-07, + "output_cost_per_token": 1.8e-06, + "cache_read_input_token_cost": 0, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/Qwen/Qwen3.6-27B-FP8": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 3.2e-07, + "output_cost_per_token": 3.2e-06, + "cache_read_input_token_cost": 0, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/lukealonso/GLM-5.1-NVFP4-MTP": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 1.4e-06, + "output_cost_per_token": 4.4e-06, + "cache_read_input_token_cost": 0, + "max_input_tokens": 202752, + "max_output_tokens": 202752, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/deepseek-ai/DeepSeek-V4-Flash": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 2.8e-07, + "cache_read_input_token_cost": 0, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/moonshotai/Kimi-K2.6": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 9.6e-07, + "output_cost_per_token": 4e-06, + "cache_read_input_token_cost": 0, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/MiniMaxAI/MiniMax-M2.5": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.2e-06, + "cache_read_input_token_cost": 0, + "max_input_tokens": 196608, + "max_output_tokens": 196608, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/google/gemma-4-31B-it": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 5.6e-07, + "cache_read_input_token_cost": 0, + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/openai/gpt-oss-120b": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 6e-07, + "cache_read_input_token_cost": 0, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" + }, + "tensormesh/openai/gpt-oss-20b": { + "litellm_provider": "tensormesh", + "mode": "chat", + "input_cost_per_token": 7e-08, + "output_cost_per_token": 2.8e-07, + "cache_read_input_token_cost": 0, + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_prompt_caching": true, + "supports_system_messages": true, + "supports_reasoning": true, + "source": "https://serverless.tensormesh.ai/v1/models/openrouter" } } diff --git a/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py b/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py index 01bcb1a247a..9a69f513069 100644 --- a/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py +++ b/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py @@ -2425,6 +2425,42 @@ class TestConvertToModelResponseObjectCompletion: assert result.choices[0].message.content == "The answer is 4." assert result.choices[0].message.reasoning_content == "2+2=4" + def test_reasoning_content_not_mirrored_into_provider_specific_fields(self): + """Mirroring reasoning_content into provider_specific_fields made + cache-replayed messages diverge from live Anthropic messages, which + only set it top-level, breaking cache key stability (issue #27337).""" + response_object = { + "id": "chatcmpl-5", + "model": "claude-sonnet-4-5", + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "The answer is 4.", + "role": "assistant", + "reasoning_content": "2+2=4", + "thinking_blocks": [ + { + "type": "thinking", + "thinking": "2+2=4", + "signature": "sig", + } + ], + }, + } + ], + "usage": {"prompt_tokens": 5, "completion_tokens": 10, "total_tokens": 15}, + } + + result = convert_to_model_response_object( + response_object=response_object, + model_response_object=ModelResponse(), + ) + message = result.choices[0].message + assert message.reasoning_content == "2+2=4" + assert "reasoning_content" not in (message.provider_specific_fields or {}) + def test_response_none_raises(self): with pytest.raises(Exception): convert_to_model_response_object( diff --git a/tests/router_unit_tests/test_router_endpoints.py b/tests/router_unit_tests/test_router_endpoints.py index 3f0afe2a5a6..c170972d984 100644 --- a/tests/router_unit_tests/test_router_endpoints.py +++ b/tests/router_unit_tests/test_router_endpoints.py @@ -1236,3 +1236,60 @@ async def test_init_containers_api_endpoints_managed_id_without_model_id_applies assert call_kw["container_id"] == "cfile_upstream_abc" assert call_kw["file_id"] == "cfile_xyz" assert call_kw["custom_llm_provider"] == "azure" + + +def test_router_model_group_encrypted_content_affinity_callback_registration(): + from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( + DeploymentAffinityCheck, + ) + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + model_group = "openai.gpt-5.1-codex" + model_group_affinity_config = { + model_group: ["encrypted_content_affinity"], + } + router = Router( + model_list=[ + { + "model_name": model_group, + "litellm_params": { + "model": "openai/gpt-5.1-codex", + "api_key": "mock-api-key", + }, + } + ], + model_group_affinity_config=model_group_affinity_config, + num_retries=0, + ) + + try: + callbacks = router.optional_callbacks or [] + encrypted_content_callbacks = [ + cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck) + ] + deployment_callback = next( + cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck) + ) + assert len(encrypted_content_callbacks) == 1 + assert encrypted_content_callbacks[0].enable_global_affinity is False + assert ( + encrypted_content_callbacks[0].model_group_affinity_config + == model_group_affinity_config + ) + assert callbacks.index(encrypted_content_callbacks[0]) < callbacks.index( + deployment_callback + ) + + router._add_encrypted_content_affinity_check(enable_global_affinity=True) + + callbacks = router.optional_callbacks or [] + encrypted_content_callbacks = [ + cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck) + ] + assert len(encrypted_content_callbacks) == 1 + assert encrypted_content_callbacks[0].enable_global_affinity is True + assert encrypted_content_callbacks[0].router is router + finally: + router.discard() diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 2b47a232262..fe49b930c10 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -328,6 +328,41 @@ def test_generic_cost_per_token_gpt54_above_272k_tokens(): assert round(completion_cost, 10) == round(expected_completion, 10) +def test_generic_cost_per_token_minimax_m3_above_512k_tokens(): + """MiniMax-M3: prompts >512K input tokens priced at 2x input, output, and cache read.""" + model = "minimax/MiniMax-M3" + custom_llm_provider = "minimax" + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + model_cost_map = litellm.model_cost[model] + prompt_tokens = 600000 + cached_tokens = 100000 + completion_tokens = 1000 + usage = Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=cached_tokens), + ) + prompt_cost, completion_cost = generic_cost_per_token( + model=model, + usage=usage, + custom_llm_provider=custom_llm_provider, + ) + expected_prompt = ( + model_cost_map["input_cost_per_token_above_512k_tokens"] + * (prompt_tokens - cached_tokens) + + model_cost_map["cache_read_input_token_cost_above_512k_tokens"] + * cached_tokens + ) + expected_completion = ( + model_cost_map["output_cost_per_token_above_512k_tokens"] * completion_tokens + ) + assert round(prompt_cost, 10) == round(expected_prompt, 10) + assert round(completion_cost, 10) == round(expected_completion, 10) + + def test_generic_cost_per_token_gpt55(): """gpt-5.5: base pricing — $5/1M input, $30/1M output, $0.50/1M cached input.""" model = "gpt-5.5" diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index 3bf3b04bf14..ed2dfc9440e 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -11,6 +11,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( BedrockImageProcessor, _bedrock_converse_messages_pt, _bedrock_tools_pt, + _rename_duplicate_bedrock_document_names, _convert_to_bedrock_tool_call_invoke, _convert_to_bedrock_tool_call_result, anthropic_messages_pt, @@ -2809,6 +2810,93 @@ def test_bedrock_converse_messages_pt_document_deterministic_name(): assert name1 == name2 +def test_bedrock_converse_messages_pt_renames_duplicate_document_names(): + """ + The same document in multiple turns must not produce duplicate names; + Bedrock rejects requests with "Messages can not contain duplicate + document names". The first occurrence keeps its hash-based name and + later occurrences get a deterministic positional suffix. + """ + document_block = { + "type": "document", + "source": { + "type": "base64", + "media_type": "application/pdf", + "data": "dGVzdA==", + }, + } + messages = [ + { + "role": "user", + "content": [document_block, {"type": "text", "text": "summarize this"}], + }, + {"role": "assistant", "content": "It says test."}, + { + "role": "user", + "content": [document_block, {"type": "text", "text": "summarize again"}], + }, + ] + + result1 = _bedrock_converse_messages_pt( + messages, "anthropic.claude-sonnet-4-6", "bedrock" + ) + result2 = _bedrock_converse_messages_pt( + messages, "anthropic.claude-sonnet-4-6", "bedrock" + ) + + names1 = [ + block["document"]["name"] + for message in result1 + for block in message["content"] + if "document" in block + ] + names2 = [ + block["document"]["name"] + for message in result2 + for block in message["content"] + if "document" in block + ] + + assert len(names1) == 2 + assert len(set(names1)) == 2 + assert names1[1] == f"{names1[0]}_2" + assert names1 == names2 + + single_turn = _bedrock_converse_messages_pt( + [messages[0]], "anthropic.claude-sonnet-4-6", "bedrock" + ) + assert names1[0] == single_turn[0]["content"][0]["document"]["name"] + + +def test_rename_duplicate_bedrock_document_names_skips_organic_suffixes(): + """ + A renamed duplicate must not collide with a document whose organic name + already carries the would-be suffix (e.g. an existing ``report_2``), + regardless of whether that document appears before or after the rename. + """ + + def _contents(names): + return [ + { + "role": "user", + "content": [{"document": {"name": name}} for name in names], + } + ] + + def _names(contents): + return [block["document"]["name"] for block in contents[0]["content"]] + + organic_first = _rename_duplicate_bedrock_document_names( + _contents(["report", "report_2", "report"]) + ) + assert _names(organic_first) == ["report", "report_2", "report_3"] + + organic_last = _rename_duplicate_bedrock_document_names( + _contents(["report", "report", "report_2"]) + ) + assert _names(organic_last) == ["report", "report_3", "report_2"] + + def test_bedrock_converse_messages_pt_document_rejects_url_source(): """Test that a URL-type document source raises a clear error instead of KeyError.""" messages = [ diff --git a/tests/test_litellm/litellm_core_utils/test_xai_oauth_routing.py b/tests/test_litellm/litellm_core_utils/test_xai_oauth_routing.py new file mode 100644 index 00000000000..03790b220eb --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_xai_oauth_routing.py @@ -0,0 +1,81 @@ +import os +import sys + +sys.path.insert(0, os.path.abspath("../../..")) + +import litellm +from litellm import LlmProviders +from litellm.litellm_core_utils.get_litellm_params import get_litellm_params +from litellm.litellm_core_utils.get_llm_provider_logic import ( + _get_openai_compatible_provider_info, +) +from litellm.llms.xai.chat.transformation import XAIChatConfig +from litellm.llms.xai.responses.transformation import XAIResponsesAPIConfig +from litellm.types.router import GenericLiteLLMParams +from litellm.utils import ( + ProviderConfigManager, + get_optional_params, + validate_environment, +) + + +def test_xai_provider_config_routing(): + chat_config = ProviderConfigManager.get_provider_chat_config( + model="grok-3-mini", + provider=LlmProviders.XAI, + ) + responses_config = ProviderConfigManager.get_provider_responses_api_config( + model="grok-3-mini", + provider=LlmProviders.XAI, + ) + + assert isinstance(chat_config, XAIChatConfig) + assert isinstance(responses_config, XAIResponsesAPIConfig) + + +def test_xai_openai_compatible_provider_info(): + model, custom_llm_provider, dynamic_api_key, api_base = ( + _get_openai_compatible_provider_info( + model="xai/grok-3-mini", + api_base="https://api.x.ai/v1", + api_key="api-key", + dynamic_api_key=None, + ) + ) + + assert model == "grok-3-mini" + assert custom_llm_provider == "xai" + assert api_base == "https://api.x.ai/v1" + assert dynamic_api_key == "api-key" + + +def test_xai_get_model_info_uses_xai_pricing_metadata(): + model_info = litellm.get_model_info("xai/grok-3-mini") + + assert model_info["litellm_provider"] == "xai" + assert model_info["key"] == "xai/grok-3-mini" + assert model_info["mode"] == "chat" + + +def test_xai_validate_environment_reads_api_key(monkeypatch): + monkeypatch.setenv("XAI_API_KEY", "api-key") + + result = validate_environment(model="xai/grok-3-mini") + + assert result == {"keys_in_environment": True, "missing_keys": []} + + +def test_xai_oauth_flag_is_generic_litellm_param(): + litellm_params = GenericLiteLLMParams(use_xai_oauth=True) + runtime_params = get_litellm_params(use_xai_oauth=True) + result = get_optional_params( + model="grok-3-mini", + custom_llm_provider="xai", + temperature=0.2, + drop_params=True, + ) + + assert result["temperature"] == 0.2 + assert litellm_params.use_xai_oauth is True + assert runtime_params["use_xai_oauth"] is True + assert "use_xai_oauth" not in result diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py new file mode 100644 index 00000000000..450f69fb87c --- /dev/null +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py @@ -0,0 +1,79 @@ +""" +Tests for AnthropicResponsesStreamWrapper +(litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py) +""" + +import os +import sys + +sys.path.insert( + 0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../../..")) +) + +from litellm.llms.anthropic.experimental_pass_through.responses_adapters.streaming_iterator import ( + AnthropicResponsesStreamWrapper, +) + + +def _process_all(events: list) -> list: + wrapper = AnthropicResponsesStreamWrapper(responses_stream=None, model="m") + for event in events: + wrapper._process_event(event) + return list(wrapper._chunk_queue) + + +class TestProcessEventTextDeltaWithoutOutputItemAdded: + """Streams that skip response.output_item.added (e.g. LMStudio) must still + open a text block before any delta and never emit index -1.""" + + def test_process_event_synthesizes_content_block_start_before_delta(self): + chunks = _process_all( + [ + {"type": "response.output_text.delta", "item_id": "i1", "delta": "Hel"}, + {"type": "response.output_text.delta", "item_id": "i1", "delta": "lo"}, + ] + ) + assert [c["type"] for c in chunks] == [ + "content_block_start", + "content_block_delta", + "content_block_delta", + ] + assert chunks[0]["content_block"] == {"type": "text", "text": ""} + assert [c["index"] for c in chunks] == [0, 0, 0] + assert chunks[1]["delta"] == {"type": "text_delta", "text": "Hel"} + + def test_process_event_delta_without_item_id_never_yields_negative_index(self): + chunks = _process_all([{"type": "response.output_text.delta", "delta": "Hi"}]) + assert [(c["type"], c["index"]) for c in chunks] == [ + ("content_block_start", 0), + ("content_block_delta", 0), + ] + + def test_process_event_unregistered_item_id_opens_new_text_block(self): + chunks = _process_all( + [ + { + "type": "response.output_item.added", + "item": {"type": "reasoning", "id": "rs_1"}, + }, + {"type": "response.output_text.delta", "item_id": "m1", "delta": "Hi"}, + ] + ) + assert chunks[1]["type"] == "content_block_start" + assert chunks[1]["content_block"] == {"type": "text", "text": ""} + assert [c["index"] for c in chunks[1:]] == [1, 1] + + def test_process_event_registered_item_id_does_not_synthesize_start(self): + chunks = _process_all( + [ + { + "type": "response.output_item.added", + "item": {"type": "message", "id": "m1"}, + }, + {"type": "response.output_text.delta", "item_id": "m1", "delta": "Hi"}, + ] + ) + assert [(c["type"], c["index"]) for c in chunks] == [ + ("content_block_start", 0), + ("content_block_delta", 0), + ] diff --git a/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py b/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py index 0de3f833a37..6812f40829a 100644 --- a/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py +++ b/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py @@ -1,10 +1,15 @@ +import base64 +import json import os import sys sys.path.insert( 0, os.path.abspath("../../../../..") ) # Adds the parent directory to the system path -from litellm.llms.bedrock.count_tokens.transformation import BedrockCountTokensConfig +from litellm.llms.bedrock.count_tokens.transformation import ( + DEFAULT_ANTHROPIC_INVOKE_MODEL_MAX_TOKENS, + BedrockCountTokensConfig, +) def test_detect_input_type(): @@ -20,6 +25,71 @@ def test_detect_input_type(): assert config._detect_input_type(request_with_text) == "invokeModel" +def test_detect_input_type_anthropic_blocks_route_to_invoke_model(): + """Anthropic-shape content blocks must not go through the Converse path, + which Bedrock rejects with a 400 (and the caller then silently falls back + to the local tokenizer).""" + config = BedrockCountTokensConfig() + + request = { + "messages": [ + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Reading the file."}, + { + "type": "tool_use", + "id": "toolu_01", + "name": "read_file", + "input": {"path": "main.py"}, + }, + ], + }, + ], + } + assert config._detect_input_type(request) == "invokeModel" + + +def test_detect_input_type_converse_blocks_route_to_converse(): + """Converse-shape blocks (no "type" key) keep using the converse input.""" + config = BedrockCountTokensConfig() + + request = {"messages": [{"role": "user", "content": [{"text": "hi"}]}]} + assert config._detect_input_type(request) == "converse" + + +def test_transform_to_invoke_model_format_base64_encodes_body(): + """The CountTokens API expects invokeModel.body as a base64-encoded blob; + Anthropic Messages bodies additionally need anthropic_version/max_tokens + to pass Bedrock's InvokeModel schema validation.""" + config = BedrockCountTokensConfig() + + request = { + "model": "anthropic.claude-3-sonnet-20240229-v1:0", + "messages": [{"role": "user", "content": [{"type": "text", "text": "Hello"}]}], + } + + result = config.transform_anthropic_to_bedrock_count_tokens(request) + + body = json.loads(base64.b64decode(result["input"]["invokeModel"]["body"])) + assert body["messages"] == request["messages"] + assert "model" not in body + assert body["anthropic_version"] == "bedrock-2023-05-31" + assert body["max_tokens"] == DEFAULT_ANTHROPIC_INVOKE_MODEL_MAX_TOKENS + + +def test_transform_to_invoke_model_format_raw_body_unchanged(): + """Non-messages bodies (e.g. Titan inputText) must not get Anthropic fields.""" + config = BedrockCountTokensConfig() + + result = config.transform_anthropic_to_bedrock_count_tokens( + {"model": "amazon.titan-text-express-v1", "inputText": "hello"} + ) + + body = json.loads(base64.b64decode(result["input"]["invokeModel"]["body"])) + assert body == {"inputText": "hello"} + + def test_transform_anthropic_to_bedrock_request(): """Test basic request transformation""" config = BedrockCountTokensConfig() diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py index 92b5ca7b10b..e83992b6bde 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py @@ -1,11 +1,13 @@ """ Unit tests for Amazon Bedrock Mantle Responses API configuration. -Mantle's gpt-5.5 / gpt-5.4 are served ONLY on the non-standard -`/openai/v1/responses` path. These tests lock the URL construction and -Bearer auth that make that routing work. +Mantle serves Responses on two paths: gpt frontier models on +`/openai/v1/responses` and other Responses-capable models (e.g. gpt-oss) on the +standard `/v1/responses`. These tests lock the per-model path selection in the +gate, the URL construction for both paths, and the shared Bearer auth. """ +import copy import os import sys @@ -89,6 +91,42 @@ class TestBedrockMantleResponsesURL: url = cfg.get_complete_url(api_base=None, litellm_params={}) assert url == "https://bedrock-mantle.us-east-1.api.aws/openai/v1/responses" + def test_standard_path_uses_region_from_env(self, monkeypatch): + monkeypatch.setenv("BEDROCK_MANTLE_REGION", "us-east-2") + monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) + cfg = BedrockMantleResponsesAPIConfig(use_openai_path=False) + url = cfg.get_complete_url(api_base=None, litellm_params={}) + assert url == "https://bedrock-mantle.us-east-2.api.aws/v1/responses" + assert "/openai/v1/responses" not in url + + def test_standard_path_normalizes_v1_base(self, monkeypatch): + monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) + cfg = BedrockMantleResponsesAPIConfig(use_openai_path=False) + url = cfg.get_complete_url( + api_base="https://bedrock-mantle.us-east-2.api.aws/v1", + litellm_params={}, + ) + assert url == "https://bedrock-mantle.us-east-2.api.aws/v1/responses" + assert url.count("/responses") == 1 + assert "/v1/v1/responses" not in url + + def test_standard_path_full_endpoint_base_not_doubled(self, monkeypatch): + monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) + cfg = BedrockMantleResponsesAPIConfig(use_openai_path=False) + url = cfg.get_complete_url( + api_base="https://bedrock-mantle.us-east-2.api.aws/v1/responses", + litellm_params={}, + ) + assert url == "https://bedrock-mantle.us-east-2.api.aws/v1/responses" + assert url.count("/responses") == 1 + + def test_default_construction_keeps_openai_path(self, monkeypatch): + monkeypatch.setenv("BEDROCK_MANTLE_REGION", "us-east-2") + monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) + cfg = BedrockMantleResponsesAPIConfig() + url = cfg.get_complete_url(api_base=None, litellm_params={}) + assert url == "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses" + class TestBedrockMantleResponsesAuth: def test_config_api_key_takes_priority(self, monkeypatch): @@ -158,6 +196,36 @@ class TestBedrockMantleResponsesAuth: is True ) + def test_standard_path_still_uses_bearer_auth(self, monkeypatch): + monkeypatch.setenv("BEDROCK_MANTLE_API_KEY", "env-key") + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + cfg = BedrockMantleResponsesAPIConfig(use_openai_path=False) + headers = cfg.validate_environment( + headers={}, + model="openai.gpt-oss-120b", + litellm_params=GenericLiteLLMParams(), + ) + assert headers["Authorization"] == "Bearer env-key" + + def test_standard_path_opts_out_of_native_features(self): + cfg = BedrockMantleResponsesAPIConfig(use_openai_path=False) + assert cfg.supports_native_file_search() is False + assert cfg.supports_native_websocket() is False + + +class TestBedrockMantleResponsesRequestBody: + def test_standard_path_outbound_body_carries_bare_model(self): + cfg = BedrockMantleResponsesAPIConfig(use_openai_path=False) + body = cfg.transform_responses_api_request( + model="openai.gpt-oss-120b", + input="hello", + response_api_optional_request_params={}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + assert body["model"] == "openai.gpt-oss-120b" + assert "input" in body + class TestBedrockMantleResponsesRegistry: def test_registry_returns_config_for_gpt_5_5(self): @@ -168,6 +236,7 @@ class TestBedrockMantleResponsesRegistry: model="openai.gpt-5.5", ) assert isinstance(cfg, BedrockMantleResponsesAPIConfig) + assert cfg.use_openai_path is True def test_registry_returns_config_for_gpt_5_4_enum(self): from litellm.utils import ProviderConfigManager @@ -177,6 +246,7 @@ class TestBedrockMantleResponsesRegistry: model="openai.gpt-5.4", ) assert isinstance(cfg, BedrockMantleResponsesAPIConfig) + assert cfg.use_openai_path is True def test_registry_returns_none_for_gpt_oss(self): # Regression guard: gpt-oss must NOT get the native Responses config; it @@ -199,9 +269,10 @@ class TestBedrockMantleResponsesRegistry: assert cfg is None def test_registry_returns_config_for_future_frontier_model(self): - # Forward-compatibility: an unseen OpenAI gpt frontier model (e.g. gpt-6) must - # get the native Responses config without a code change. The gate allow-lists - # the openai.gpt- family (minus gpt-oss), so gpt-6 matches automatically. + # Forward-compatibility: an unseen OpenAI gpt frontier model (e.g. gpt-6), + # not yet in the price map, must get the openai-path Responses config with + # no code or JSON change. The name-convention fallback (openai.gpt- minus + # gpt-oss) catches it before any price-map entry exists. from litellm.utils import ProviderConfigManager cfg = ProviderConfigManager.get_provider_responses_api_config( @@ -209,6 +280,48 @@ class TestBedrockMantleResponsesRegistry: model="openai.gpt-6", ) assert isinstance(cfg, BedrockMantleResponsesAPIConfig) + assert cfg.use_openai_path is True + + def test_price_map_flag_routes_non_gpt_name_to_openai_path( + self, restore_model_cost + ): + # Data-driven onboarding: a frontier model whose name does NOT match the + # openai.gpt- convention can still be routed to /openai/v1/responses by + # declaring use_openai_responses_path in its price-map entry, with no code + # change. The string fallback alone could never catch this name. + from litellm.utils import ProviderConfigManager, register_model + + register_model( + { + "bedrock_mantle/somelab.frontier-x": { + "litellm_provider": "bedrock_mantle", + "mode": "responses", + "use_openai_responses_path": True, + } + } + ) + cfg = ProviderConfigManager.get_provider_responses_api_config( + provider="bedrock_mantle", + model="somelab.frontier-x", + ) + assert isinstance(cfg, BedrockMantleResponsesAPIConfig) + assert cfg.use_openai_path is True + + def test_gpt_5_5_price_map_declares_openai_responses_path(self, local_cost_map): + # The gpt-5.x entries must carry the data-driven flag so frontier routing + # does not rely on the name-string fallback alone. + assert ( + litellm.model_cost["bedrock_mantle/openai.gpt-5.5"].get( + "use_openai_responses_path" + ) + is True + ) + assert ( + litellm.model_cost["bedrock_mantle/openai.gpt-5.4"].get( + "use_openai_responses_path" + ) + is True + ) @pytest.mark.parametrize( "model", @@ -243,6 +356,129 @@ class TestBedrockMantleResponsesRegistry: ) assert cfg is None + def test_declared_responses_non_openai_routes_to_standard_path( + self, restore_model_cost + ): + # New feature: a non-OpenAI model declared mode=responses (e.g. via a + # user's proxy model_info block) must route to the STANDARD /v1/responses + # path, not the frontier /openai/v1/responses path. Fails before the + # path-aware gate exists (old gate returned None for non-gpt models). + from litellm.utils import ProviderConfigManager, register_model + + register_model( + { + "bedrock_mantle/somelab.future-model": { + "litellm_provider": "bedrock_mantle", + "mode": "responses", + } + } + ) + cfg = ProviderConfigManager.get_provider_responses_api_config( + provider="bedrock_mantle", + model="somelab.future-model", + ) + assert isinstance(cfg, BedrockMantleResponsesAPIConfig) + assert cfg.use_openai_path is False + + def test_gpt_oss_opt_in_routes_to_standard_path(self, restore_model_cost): + # When a user opts gpt-oss into native Responses via model_info mode, + # it must take the STANDARD /v1/responses path (gpt-oss Responses is on + # /v1/responses, NOT the frontier /openai/v1/responses path). + from litellm.utils import ProviderConfigManager, register_model + + register_model( + { + "bedrock_mantle/openai.gpt-oss-120b": { + "litellm_provider": "bedrock_mantle", + "mode": "responses", + } + } + ) + cfg = ProviderConfigManager.get_provider_responses_api_config( + provider="bedrock_mantle", + model="openai.gpt-oss-120b", + ) + assert isinstance(cfg, BedrockMantleResponsesAPIConfig) + assert cfg.use_openai_path is False + + def test_unmapped_model_degrades_to_none_without_crashing(self, restore_model_cost): + # A non-frontier model that is not in model_cost makes get_model_info + # raise; the gate must swallow it and return None rather than crash. + from litellm.utils import ProviderConfigManager + + litellm.model_cost.pop("bedrock_mantle/somelab.unmapped-model", None) + litellm.get_model_info.cache_clear() + cfg = ProviderConfigManager.get_provider_responses_api_config( + provider="bedrock_mantle", + model="somelab.unmapped-model", + ) + assert cfg is None + + def test_register_model_restore_undoes_existing_key_overwrite(self): + # Self-contained guard for the deepcopy requirement of restore_model_cost. + # register_model overwrites an existing key by mutating its nested dict in + # place, so the snapshot must be a deepcopy: a shallow dict() copy would + # share that nested dict and leave mode=responses after restore, making + # the final assertion fail. The in-place clear+update mirrors the fixture. + from litellm.utils import ProviderConfigManager, register_model + + snapshot = copy.deepcopy(litellm.model_cost) + litellm.get_model_info.cache_clear() + try: + register_model( + { + "bedrock_mantle/openai.gpt-oss-120b": { + "litellm_provider": "bedrock_mantle", + "mode": "responses", + } + } + ) + during = ProviderConfigManager.get_provider_responses_api_config( + provider="bedrock_mantle", model="openai.gpt-oss-120b" + ) + assert isinstance(during, BedrockMantleResponsesAPIConfig) + finally: + litellm.model_cost.clear() + litellm.model_cost.update(snapshot) + litellm.get_model_info.cache_clear() + after = ProviderConfigManager.get_provider_responses_api_config( + provider="bedrock_mantle", model="openai.gpt-oss-120b" + ) + assert after is None + + +@pytest.fixture +def restore_model_cost(): + """Snapshot litellm.model_cost so register_model edits don't leak across tests. + + register_model mutates the global litellm.model_cost, and get_model_info is + lru_cached, so without restore + cache_clear a registered model would bleed + into sibling tests in the same process. + + Two subtleties make this fixture non-obvious: + + 1. The snapshot must be a deepcopy. register_model overwrites an existing key + via `litellm.model_cost.setdefault(key, {}).update(...)`, mutating the + nested dict in place; a shallow copy would share those nested dicts and + could not capture the pre-mutation values of an existing entry. + 2. The restore must be in place (clear + update the SAME dict object), not a + reassignment. The conftest autouse `isolate_litellm_state` fixture + snapshots `litellm.model_cost` by reference and restores that reference on + its teardown, which runs after this one. Reassigning `litellm.model_cost` + to a fresh dict here is undone when conftest reinstalls its (in-place + mutated) reference, so the registered mode would leak and poison + TestBedrockMantleResponsesPricing. Mutating the original object in place + restores the contents conftest's reference points at. + """ + original_model_cost = copy.deepcopy(litellm.model_cost) + litellm.get_model_info.cache_clear() + try: + yield + finally: + litellm.model_cost.clear() + litellm.model_cost.update(original_model_cost) + litellm.get_model_info.cache_clear() + @pytest.fixture def local_cost_map(monkeypatch): diff --git a/tests/test_litellm/llms/openai/completion/test_completion_handler.py b/tests/test_litellm/llms/openai/completion/test_completion_handler.py new file mode 100644 index 00000000000..c6af96fa375 --- /dev/null +++ b/tests/test_litellm/llms/openai/completion/test_completion_handler.py @@ -0,0 +1,93 @@ +""" +Tests that client headers are forwarded to the provider on the OpenAI +text completion path. + +Regression tests for https://github.com/BerriAI/litellm/issues/27410 +""" + +import os +import sys + +import pytest +import respx +from httpx import Response + +sys.path.insert(0, os.path.abspath("../../../../..")) + +import litellm +from litellm import atext_completion, text_completion + + +@pytest.fixture(autouse=True) +def setup_env(monkeypatch): + monkeypatch.setenv("OPENAI_API_KEY", "sk-test-fake-key") + + +@pytest.fixture +def mock_completions_endpoint(): + return respx.post("https://api.openai.com/v1/completions").mock( + return_value=Response( + 200, + json={ + "id": "cmpl-test123", + "object": "text_completion", + "created": 1677652288, + "model": "gpt-3.5-turbo-instruct", + "choices": [ + { + "text": "hi", + "index": 0, + "logprobs": None, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 1, + "completion_tokens": 1, + "total_tokens": 2, + }, + }, + ) + ) + + +@respx.mock +def test_completion_forwards_client_headers_to_provider(mock_completions_endpoint): + text_completion( + model="gpt-3.5-turbo-instruct", + prompt="hello", + max_tokens=5, + headers={"x-mycorp-llmcall-id": "abc-123"}, + ) + + request_headers = mock_completions_endpoint.calls.last.request.headers + assert request_headers["x-mycorp-llmcall-id"] == "abc-123" + + +@respx.mock +def test_completion_forwards_extra_headers_to_provider(mock_completions_endpoint): + text_completion( + model="gpt-3.5-turbo-instruct", + prompt="hello", + max_tokens=5, + extra_headers={"x-mycorp-llmcall-id": "abc-123"}, + ) + + request_headers = mock_completions_endpoint.calls.last.request.headers + assert request_headers["x-mycorp-llmcall-id"] == "abc-123" + + +@respx.mock +async def test_acompletion_forwards_client_headers_to_provider( + mock_completions_endpoint, monkeypatch +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + await atext_completion( + model="gpt-3.5-turbo-instruct", + prompt="hello", + max_tokens=5, + headers={"x-mycorp-llmcall-id": "abc-123"}, + ) + + request_headers = mock_completions_endpoint.calls.last.request.headers + assert request_headers["x-mycorp-llmcall-id"] == "abc-123" diff --git a/tests/test_litellm/llms/openai_like/test_tensormesh_provider.py b/tests/test_litellm/llms/openai_like/test_tensormesh_provider.py index f81f1c00a7b..09248a779c5 100644 --- a/tests/test_litellm/llms/openai_like/test_tensormesh_provider.py +++ b/tests/test_litellm/llms/openai_like/test_tensormesh_provider.py @@ -2,8 +2,23 @@ Tests for Tensormesh provider configuration and integration. """ +import pytest + import litellm +TENSORMESH_MODELS = [ + "tensormesh/Qwen/Qwen3.5-397B-A17B-FP8", + "tensormesh/Qwen/Qwen3-Coder-480B-A35B-Instruct-FP8", + "tensormesh/Qwen/Qwen3.6-27B-FP8", + "tensormesh/lukealonso/GLM-5.1-NVFP4-MTP", + "tensormesh/deepseek-ai/DeepSeek-V4-Flash", + "tensormesh/moonshotai/Kimi-K2.6", + "tensormesh/MiniMaxAI/MiniMax-M2.5", + "tensormesh/google/gemma-4-31B-it", + "tensormesh/openai/gpt-oss-120b", + "tensormesh/openai/gpt-oss-20b", +] + class TestTensormeshProviderConfig: """Test Tensormesh provider configuration""" @@ -82,3 +97,60 @@ class TestTensormeshProviderConfig: assert len(router.model_list) == 1 assert router.model_list[0]["model_name"] == "tensormesh-chat" + + +class TestTensormeshCostMap: + """The serverless models are registered in the cost map so LiteLLM can + price requests and unblock tool-calling params on the JSON provider path.""" + + @pytest.fixture(autouse=True) + def _use_local_model_cost_map(self, monkeypatch): + original_model_cost = litellm.model_cost + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + litellm.model_cost = litellm.get_model_cost_map(url="") + litellm.get_model_info.cache_clear() + try: + yield + finally: + litellm.model_cost = original_model_cost + litellm.get_model_info.cache_clear() + + def test_models_registered_with_capabilities(self): + for model in TENSORMESH_MODELS: + info = litellm.get_model_info(model) + assert info["litellm_provider"] == "tensormesh" + assert info["mode"] == "chat" + assert litellm.supports_function_calling(model) is True, model + assert litellm.supports_response_schema(model) is True, model + assert litellm.model_cost[model]["supports_tool_choice"] is True, model + assert litellm.model_cost[model]["supports_prompt_caching"] is True, model + + def test_reasoning_flag_matches_expected_set(self): + reasoning_models = { + "tensormesh/deepseek-ai/DeepSeek-V4-Flash", + "tensormesh/Qwen/Qwen3.5-397B-A17B-FP8", + "tensormesh/Qwen/Qwen3.6-27B-FP8", + "tensormesh/lukealonso/GLM-5.1-NVFP4-MTP", + "tensormesh/MiniMaxAI/MiniMax-M2.5", + "tensormesh/moonshotai/Kimi-K2.6", + "tensormesh/openai/gpt-oss-120b", + "tensormesh/openai/gpt-oss-20b", + "tensormesh/google/gemma-4-31B-it", + } + for model in TENSORMESH_MODELS: + assert litellm.supports_reasoning(model) is (model in reasoning_models), model + + def test_cost_is_wired_and_cache_reads_are_free(self): + prompt_cost, completion_cost = litellm.cost_per_token( + model="tensormesh/openai/gpt-oss-120b", + prompt_tokens=1_000_000, + completion_tokens=1_000_000, + ) + assert prompt_cost == pytest.approx(0.15) + assert completion_cost == pytest.approx(0.60) + assert ( + litellm.model_cost["tensormesh/openai/gpt-oss-120b"][ + "cache_read_input_token_cost" + ] + == 0 + ) diff --git a/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py b/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py index 74888e6cd9e..cc8b14e5514 100644 --- a/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py +++ b/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py @@ -86,7 +86,7 @@ class TestContextCachingEndpoints: cached_content = "cached_content_123" optional_params = self.sample_optional_params.copy() test_project = "test_project" - test_location = "test_location" + test_location = "us-central1" # Execute result = self.context_caching.check_and_create_cache( @@ -129,7 +129,7 @@ class TestContextCachingEndpoints: mock_separate.return_value = ([], self.sample_messages) # No cached messages optional_params = self.sample_optional_params.copy() test_project = "test_project" - test_location = "test_location" + test_location = "us-central1" # Execute result = self.context_caching.check_and_create_cache( @@ -177,7 +177,7 @@ class TestContextCachingEndpoints: optional_params = self.sample_optional_params.copy() test_project = "test_project" - test_location = "test_location" + test_location = "us-central1" # Execute result = self.context_caching.check_and_create_cache( @@ -254,7 +254,7 @@ class TestContextCachingEndpoints: optional_params = self.sample_optional_params.copy() test_project = "test_project" - test_location = "test_location" + test_location = "us-central1" # Execute result = self.context_caching.check_and_create_cache( @@ -324,7 +324,7 @@ class TestContextCachingEndpoints: optional_params = self.sample_optional_params.copy() test_project = "test_project" - test_location = "test_location" + test_location = "us-central1" # Execute and Assert with pytest.raises(VertexAIError) as exc_info: @@ -364,7 +364,7 @@ class TestContextCachingEndpoints: cached_content = "cached_content_123" optional_params = self.sample_optional_params.copy() test_project = "test_project" - test_location = "test_location" + test_location = "us-central1" # Execute result = await self.context_caching.async_check_and_create_cache( @@ -404,7 +404,7 @@ class TestContextCachingEndpoints: mock_separate.return_value = ([], self.sample_messages) optional_params = self.sample_optional_params.copy() test_project = "test_project" - test_location = "test_location" + test_location = "us-central1" # Execute result = await self.context_caching.async_check_and_create_cache( @@ -453,7 +453,7 @@ class TestContextCachingEndpoints: optional_params = self.sample_optional_params.copy() test_project = "test_project" - test_location = "test_location" + test_location = "us-central1" # Execute result = await self.context_caching.async_check_and_create_cache( @@ -535,7 +535,7 @@ class TestContextCachingEndpoints: optional_params = self.sample_optional_params.copy() test_project = "test_project" - test_location = "test_location" + test_location = "us-central1" # Execute result = await self.context_caching.async_check_and_create_cache( @@ -606,7 +606,7 @@ class TestContextCachingEndpoints: optional_params = self.sample_optional_params.copy() test_project = "test_project" - test_location = "test_location" + test_location = "us-central1" # Execute and Assert with pytest.raises(VertexAIError) as exc_info: @@ -648,7 +648,7 @@ class TestContextCachingEndpoints: optional_params = self.sample_optional_params.copy() original_tools = optional_params["tools"].copy() test_project = "test_project" - test_location = "test_location" + test_location = "us-central1" # Mock the check_cache to return existing cache so we don't make HTTP calls with patch.object( @@ -694,7 +694,7 @@ class TestContextCachingEndpoints: optional_params = self.sample_optional_params.copy() original_tools = optional_params["tools"].copy() test_project = "test_project" - test_location = "test_location" + test_location = "us-central1" # Execute result = self.context_caching.check_and_create_cache( @@ -735,7 +735,7 @@ class TestContextCachingEndpoints: optional_params = self.sample_optional_params.copy() original_tools = optional_params["tools"].copy() test_project = "test_project" - test_location = "test_location" + test_location = "us-central1" # Execute result = await self.context_caching.async_check_and_create_cache( @@ -778,7 +778,7 @@ class TestContextCachingEndpoints: optional_params = self.sample_optional_params.copy() original_tools = optional_params["tools"].copy() test_project = "test_project" - test_location = "test_location" + test_location = "us-central1" # Mock the async_check_cache to return existing cache so we don't make HTTP calls with patch.object( @@ -837,7 +837,7 @@ class TestContextCachingEndpoints: logging_obj=self.mock_logging, custom_llm_provider=custom_llm_provider, vertex_project="test_project", - vertex_location="test_location", + vertex_location="us-central1", vertex_auth_header="vertext_test_token", ) @@ -870,7 +870,7 @@ class TestContextCachingEndpoints: logging_obj=self.mock_logging, custom_llm_provider=custom_llm_provider, vertex_project="test_project", - vertex_location="test_location", + vertex_location="us-central1", vertex_auth_header="vertext_test_token", ) @@ -908,7 +908,7 @@ class TestContextCachingEndpoints: logging_obj=self.mock_logging, custom_llm_provider=custom_llm_provider, vertex_project="test_project", - vertex_location="test_location", + vertex_location="us-central1", vertex_auth_header="vertext_test_token", ) @@ -942,7 +942,7 @@ class TestContextCachingEndpoints: logging_obj=self.mock_logging, custom_llm_provider=custom_llm_provider, vertex_project="test_project", - vertex_location="test_location", + vertex_location="us-central1", vertex_auth_header="vertext_test_token", ) @@ -1002,7 +1002,7 @@ class TestContextCachingEndpoints: logging_obj=self.mock_logging, custom_llm_provider=custom_llm_provider, vertex_project="test_project", - vertex_location="test_location", + vertex_location="us-central1", vertex_auth_header="vertext_test_token", ) @@ -1072,7 +1072,7 @@ class TestContextCachingEndpoints: logging_obj=self.mock_logging, custom_llm_provider=custom_llm_provider, vertex_project="test_project", - vertex_location="test_location", + vertex_location="us-central1", vertex_auth_header="vertext_test_token", ) @@ -1138,7 +1138,7 @@ class TestContextCachingEndpoints: logging_obj=self.mock_logging, custom_llm_provider=custom_llm_provider, vertex_project="test_project", - vertex_location="test_location", + vertex_location="us-central1", vertex_auth_header="vertext_test_token", ) @@ -1205,7 +1205,7 @@ class TestContextCachingEndpoints: logging_obj=self.mock_logging, custom_llm_provider=custom_llm_provider, vertex_project="test_project", - vertex_location="test_location", + vertex_location="us-central1", vertex_auth_header="vertext_test_token", ) @@ -1280,7 +1280,7 @@ class TestContextCachingEndpoints: logging_obj=self.mock_logging, custom_llm_provider=custom_llm_provider, vertex_project="test_project", - vertex_location="test_location", + vertex_location="us-central1", vertex_auth_header="vertext_test_token", ) @@ -1336,7 +1336,7 @@ class TestContextCachingEndpoints: logging_obj=self.mock_logging, custom_llm_provider=custom_llm_provider, vertex_project="test_project", - vertex_location="test_location", + vertex_location="us-central1", vertex_auth_header="vertext_test_token", ) @@ -1390,7 +1390,7 @@ class TestContextCachingEndpoints: cached_content=None, custom_llm_provider=custom_llm_provider, vertex_project="test_project", - vertex_location="test_location", + vertex_location="us-central1", vertex_auth_header="test_token", ) @@ -1441,7 +1441,7 @@ class TestContextCachingEndpoints: cached_content=None, custom_llm_provider=custom_llm_provider, vertex_project="test_project", - vertex_location="test_location", + vertex_location="us-central1", vertex_auth_header="test_token", ) diff --git a/tests/test_litellm/llms/xai/test_xai_oauth.py b/tests/test_litellm/llms/xai/test_xai_oauth.py new file mode 100644 index 00000000000..45fa6a405f2 --- /dev/null +++ b/tests/test_litellm/llms/xai/test_xai_oauth.py @@ -0,0 +1,801 @@ +import base64 +import hashlib +import json +import os +import threading +import time +from urllib.parse import parse_qs, urlparse +from unittest.mock import MagicMock + +import httpx +import litellm +import pytest +from click.testing import CliRunner + +import litellm.llms.xai.oauth as xai_oauth_module +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +from litellm.llms.xai.oauth import ( + XAI_OAUTH_CLIENT_ID, + XAI_OAUTH_SCOPE, + XAIOAuthError, + XAIOAuthAuthenticator, + XAIOAuthLoginRequiredError, +) +from litellm.llms.xai.chat.transformation import XAIChatConfig +from litellm.llms.xai.responses.transformation import XAIResponsesAPIConfig +from litellm.types.router import GenericLiteLLMParams +from litellm.utils import get_optional_params, validate_environment + + +def _write_auth_file(tmp_path, payload): + token_dir = tmp_path / "xai_oauth" + token_dir.mkdir() + auth_file = token_dir / "auth.json" + auth_file.write_text(json.dumps(payload)) + return token_dir, auth_file + + +def test_get_access_token_uses_fresh_local_token(tmp_path, monkeypatch): + token_dir, _ = _write_auth_file( + tmp_path, + { + "access_token": "fresh-token", + "refresh_token": "refresh-token", + "expires_at": time.time() + 3600, + }, + ) + monkeypatch.setenv("XAI_OAUTH_TOKEN_DIR", str(token_dir)) + + assert XAIOAuthAuthenticator().get_access_token() == "fresh-token" + + +def test_get_access_token_refreshes_and_preserves_refresh_token(tmp_path, monkeypatch): + token_dir, auth_file = _write_auth_file( + tmp_path, + { + "access_token": "expired-token", + "refresh_token": "refresh-token", + "token_endpoint": "https://auth.x.ai/oauth/token", + "expires_at": time.time() - 1, + }, + ) + monkeypatch.setenv("XAI_OAUTH_TOKEN_DIR", str(token_dir)) + + def handler(request: httpx.Request) -> httpx.Response: + body = dict(item.split("=") for item in request.content.decode().split("&")) + assert body["grant_type"] == "refresh_token" + assert body["refresh_token"] == "refresh-token" + assert body["client_id"] == XAI_OAUTH_CLIENT_ID + return httpx.Response( + 200, + json={ + "access_token": "new-token", + "expires_in": 3600, + "token_type": "Bearer", + }, + ) + + client = httpx.Client(transport=httpx.MockTransport(handler)) + + assert XAIOAuthAuthenticator(http_client=client).get_access_token() == "new-token" + stored = json.loads(auth_file.read_text()) + assert stored["access_token"] == "new-token" + assert stored["refresh_token"] == "refresh-token" + + +def test_get_access_token_reuses_token_refreshed_by_parallel_request(): + expired_auth_data = { + "access_token": "expired-token", + "refresh_token": "refresh-token", + "token_endpoint": "https://auth.x.ai/oauth/token", + "expires_at": time.time() - 1, + } + refreshed_auth_data = { + "access_token": "already-refreshed-token", + "refresh_token": "rotated-refresh-token", + "token_endpoint": "https://auth.x.ai/oauth/token", + "expires_at": time.time() + 3600, + } + authenticator = XAIOAuthAuthenticator() + authenticator._read_auth_file = MagicMock( + side_effect=[expired_auth_data, refreshed_auth_data] + ) + authenticator._refresh_tokens = MagicMock() + + assert authenticator.get_access_token() == "already-refreshed-token" + authenticator._refresh_tokens.assert_not_called() + + +def test_get_access_token_requires_login_without_auth_file(tmp_path, monkeypatch): + monkeypatch.setenv("XAI_OAUTH_TOKEN_DIR", str(tmp_path / "missing")) + + with pytest.raises(XAIOAuthLoginRequiredError): + XAIOAuthAuthenticator().get_access_token() + + +def test_get_access_token_ignores_invalid_auth_file(tmp_path, monkeypatch): + token_dir = tmp_path / "xai_oauth" + token_dir.mkdir() + (token_dir / "auth.json").write_text("{not-json") + monkeypatch.setenv("XAI_OAUTH_TOKEN_DIR", str(token_dir)) + + with pytest.raises(XAIOAuthLoginRequiredError): + XAIOAuthAuthenticator().get_access_token() + + +def test_refresh_failure_surfaces_oauth_error(tmp_path, monkeypatch): + token_dir, _ = _write_auth_file( + tmp_path, + { + "access_token": "expired-token", + "refresh_token": "refresh-token", + "token_endpoint": "https://auth.x.ai/oauth/token", + "expires_at": time.time() - 1, + }, + ) + monkeypatch.setenv("XAI_OAUTH_TOKEN_DIR", str(token_dir)) + + client = httpx.Client( + transport=httpx.MockTransport( + lambda request: httpx.Response(401, text="invalid_grant", request=request) + ) + ) + + with pytest.raises(XAIOAuthError) as exc_info: + XAIOAuthAuthenticator(http_client=client).get_access_token() + + assert "401 invalid_grant" in str(exc_info.value) + + +def test_build_auth_record_requires_access_and_refresh_tokens(): + authenticator = XAIOAuthAuthenticator() + + with pytest.raises(XAIOAuthError, match="access_token"): + authenticator._build_auth_record( + {"refresh_token": "refresh-token"}, + "https://auth.x.ai/oauth/token", + ) + + with pytest.raises(XAIOAuthError, match="refresh_token"): + authenticator._build_auth_record( + {"access_token": "access-token"}, + "https://auth.x.ai/oauth/token", + ) + + +def test_build_auth_record_defaults_expiry_and_token_type(): + authenticator = XAIOAuthAuthenticator() + + auth_data = authenticator._build_auth_record( + { + "access_token": "access-token", + "refresh_token": "refresh-token", + "expires_in": "not-a-number", + }, + "https://auth.x.ai/oauth/token", + ) + + assert auth_data["token_type"] == "Bearer" + assert auth_data["expires_at"] > time.time() + + +def test_is_expired_treats_missing_or_invalid_expiry_as_expired(): + authenticator = XAIOAuthAuthenticator() + + assert authenticator._is_expired({}) is True + assert authenticator._is_expired({"expires_at": "not-a-number"}) is True + + +def test_write_auth_file_creates_private_file(tmp_path, monkeypatch): + token_dir = tmp_path / "xai_oauth" + monkeypatch.setenv("XAI_OAUTH_TOKEN_DIR", str(token_dir)) + authenticator = XAIOAuthAuthenticator() + old_umask = os.umask(0o022) + replace_calls = [] + real_replace = os.replace + + def assert_private_temp_file(src, dst): + replace_calls.append((src, dst)) + assert oct(os.stat(src).st_mode & 0o777) == "0o600" + with open(src) as f: + assert json.load(f)["refresh_token"] == "refresh-token" + real_replace(src, dst) + + monkeypatch.setattr(os, "replace", assert_private_temp_file) + + try: + authenticator._write_auth_file( + { + "access_token": "access-token", + "refresh_token": "refresh-token", + "expires_at": time.time() + 3600, + } + ) + finally: + os.umask(old_umask) + + stored = json.loads((token_dir / "auth.json").read_text()) + assert stored["access_token"] == "access-token" + assert replace_calls + assert oct(os.stat(token_dir).st_mode & 0o777) == "0o700" + assert oct(os.stat(token_dir / "auth.json").st_mode & 0o777) == "0o600" + + +def test_discovery_rejects_unexpected_endpoint(): + authenticator = XAIOAuthAuthenticator() + + with pytest.raises(XAIOAuthError, match="unexpected endpoint"): + authenticator._validate_xai_endpoint("https://evil.example.com/oauth/token") + + with pytest.raises(XAIOAuthError, match="unexpected endpoint"): + authenticator._validate_xai_endpoint("http://auth.x.ai/oauth/token") + + +def test_discover_returns_validated_xai_endpoints(): + def handler(request: httpx.Request) -> httpx.Response: + assert request.url == "https://auth.x.ai/.well-known/openid-configuration" + return httpx.Response( + 200, + json={ + "authorization_endpoint": "https://auth.x.ai/oauth/authorize", + "token_endpoint": "https://auth.x.ai/oauth/token", + }, + ) + + authenticator = XAIOAuthAuthenticator( + http_client=httpx.Client(transport=httpx.MockTransport(handler)) + ) + + assert authenticator._discover() == { + "authorization_endpoint": "https://auth.x.ai/oauth/authorize", + "token_endpoint": "https://auth.x.ai/oauth/token", + } + + +def test_discover_requires_authorization_and_token_endpoints(): + authenticator = XAIOAuthAuthenticator( + http_client=httpx.Client( + transport=httpx.MockTransport(lambda request: httpx.Response(200, json={})) + ) + ) + + with pytest.raises(XAIOAuthError, match="missing endpoints"): + authenticator._discover() + + +def test_discover_wraps_http_errors(): + authenticator = XAIOAuthAuthenticator( + http_client=httpx.Client( + transport=httpx.MockTransport( + lambda request: httpx.Response( + 500, text="discovery failed", request=request + ) + ) + ) + ) + + with pytest.raises(XAIOAuthError) as exc_info: + authenticator._discover() + + assert "xAI OAuth discovery request failed: 500 discovery failed" in str( + exc_info.value + ) + + +def test_discover_wraps_invalid_json_response(): + authenticator = XAIOAuthAuthenticator( + http_client=httpx.Client( + transport=httpx.MockTransport( + lambda request: httpx.Response(200, text="not-json") + ) + ) + ) + + with pytest.raises(XAIOAuthError, match="discovery response was not valid JSON"): + authenticator._discover() + + +def test_refresh_discovers_token_endpoint_when_auth_file_is_legacy( + tmp_path, monkeypatch +): + token_dir, auth_file = _write_auth_file( + tmp_path, + { + "access_token": "expired-token", + "refresh_token": "refresh-token", + "expires_at": time.time() - 1, + }, + ) + monkeypatch.setenv("XAI_OAUTH_TOKEN_DIR", str(token_dir)) + + def handler(request: httpx.Request) -> httpx.Response: + if request.method == "GET": + return httpx.Response( + 200, + json={ + "authorization_endpoint": "https://auth.x.ai/oauth/authorize", + "token_endpoint": "https://auth.x.ai/oauth/token", + }, + ) + return httpx.Response( + 200, + json={ + "access_token": "discovered-token", + "refresh_token": "new-refresh-token", + "expires_in": 3600, + }, + ) + + authenticator = XAIOAuthAuthenticator( + http_client=httpx.Client(transport=httpx.MockTransport(handler)) + ) + + assert authenticator.get_access_token() == "discovered-token" + stored = json.loads(auth_file.read_text()) + assert stored["token_endpoint"] == "https://auth.x.ai/oauth/token" + + +def test_exchange_token_rejects_non_object_response(): + authenticator = XAIOAuthAuthenticator( + http_client=httpx.Client( + transport=httpx.MockTransport( + lambda request: httpx.Response(200, json=["not", "an", "object"]) + ) + ) + ) + + with pytest.raises(XAIOAuthError, match="was not an object"): + authenticator._exchange_token("https://auth.x.ai/oauth/token", {}) + + +def test_exchange_token_wraps_invalid_json_response(): + authenticator = XAIOAuthAuthenticator( + http_client=httpx.Client( + transport=httpx.MockTransport( + lambda request: httpx.Response(200, text="not-json") + ) + ) + ) + + with pytest.raises(XAIOAuthError, match="token response was not valid JSON"): + authenticator._exchange_token("https://auth.x.ai/oauth/token", {}) + + +def test_start_callback_server_falls_back_to_ephemeral_port(monkeypatch): + calls = [] + real_server = xai_oauth_module._CallbackServer + + class FirstPortFailsCallbackServer(real_server): + def __init__(self, server_address, handler_class): + calls.append(server_address[1]) + if server_address[1] == xai_oauth_module.XAI_OAUTH_REDIRECT_PORT: + raise OSError("port unavailable") + super().__init__(server_address, handler_class) + + monkeypatch.setattr( + xai_oauth_module, "_CallbackServer", FirstPortFailsCallbackServer + ) + + server, redirect_uri = XAIOAuthAuthenticator()._start_callback_server("state-value") + try: + assert calls == [xai_oauth_module.XAI_OAUTH_REDIRECT_PORT, 0] + assert redirect_uri.startswith("http://127.0.0.1:") + assert redirect_uri.endswith("/callback") + finally: + server.server_close() + + +def test_wait_for_callback_times_out_and_closes_server(monkeypatch): + server, _ = XAIOAuthAuthenticator()._start_callback_server("state-value") + monkeypatch.setattr(xai_oauth_module, "XAI_OAUTH_CALLBACK_TIMEOUT_SECONDS", 0) + + with pytest.raises(XAIOAuthError, match="Timed out"): + XAIOAuthAuthenticator()._wait_for_callback(server) + + +def test_callback_handler_records_success_and_rejects_state_mismatch(): + authenticator = XAIOAuthAuthenticator() + server, redirect_uri = authenticator._start_callback_server("expected-state") + thread = threading.Thread(target=server.handle_request) + thread.start() + response = httpx.get(f"{redirect_uri}?code=auth-code&state=expected-state") + thread.join(timeout=5) + + assert response.status_code == 200 + assert server.callback_result == { + "code": "auth-code", + "state": "expected-state", + "error": None, + "error_description": None, + } + + server, redirect_uri = authenticator._start_callback_server("expected-state") + thread = threading.Thread(target=server.handle_request) + thread.start() + response = httpx.get(f"{redirect_uri}?code=auth-code&state=wrong-state") + thread.join(timeout=5) + + assert response.status_code == 400 + assert server.callback_result["state"] == "wrong-state" + + +def test_login_exchanges_authorization_code_and_persists_auth_record(monkeypatch): + authenticator = XAIOAuthAuthenticator() + fake_server = MagicMock() + written_records = [] + + class FakeUUID: + def __init__(self, value): + self.hex = value + + monkeypatch.setattr( + xai_oauth_module.uuid, + "uuid4", + MagicMock(side_effect=[FakeUUID("state-value"), FakeUUID("nonce-value")]), + ) + authenticator._read_auth_file = MagicMock(return_value=None) + authenticator._discover = MagicMock( + return_value={ + "authorization_endpoint": "https://auth.x.ai/oauth/authorize", + "token_endpoint": "https://auth.x.ai/oauth/token", + } + ) + authenticator._pkce_pair = MagicMock(return_value=("verifier", "challenge")) + authenticator._start_callback_server = MagicMock( + return_value=(fake_server, "http://127.0.0.1:56121/callback") + ) + authenticator._wait_for_callback = MagicMock( + return_value={"state": "state-value", "code": "auth-code"} + ) + authenticator._exchange_token = MagicMock( + return_value={ + "access_token": "access-token", + "refresh_token": "refresh-token", + "expires_in": 3600, + } + ) + authenticator._write_auth_file = MagicMock(side_effect=written_records.append) + + auth_data = authenticator.login(no_browser=True) + + authenticator._exchange_token.assert_called_once_with( + "https://auth.x.ai/oauth/token", + { + "grant_type": "authorization_code", + "code": "auth-code", + "redirect_uri": "http://127.0.0.1:56121/callback", + "client_id": XAI_OAUTH_CLIENT_ID, + "code_verifier": "verifier", + }, + ) + assert auth_data["access_token"] == "access-token" + assert written_records == [auth_data] + + +def test_login_raises_on_callback_error_or_missing_code(monkeypatch): + authenticator = XAIOAuthAuthenticator() + + class FakeUUID: + hex = "state-value" + + monkeypatch.setattr( + xai_oauth_module.uuid, "uuid4", MagicMock(return_value=FakeUUID()) + ) + authenticator._read_auth_file = MagicMock(return_value=None) + authenticator._discover = MagicMock( + return_value={ + "authorization_endpoint": "https://auth.x.ai/oauth/authorize", + "token_endpoint": "https://auth.x.ai/oauth/token", + } + ) + authenticator._pkce_pair = MagicMock(return_value=("verifier", "challenge")) + authenticator._start_callback_server = MagicMock( + return_value=(MagicMock(), "http://127.0.0.1:56121/callback") + ) + authenticator._wait_for_callback = MagicMock( + return_value={ + "state": "state-value", + "error": "access_denied", + "error_description": "denied", + } + ) + + with pytest.raises(XAIOAuthError, match="denied"): + authenticator.login(no_browser=True) + + authenticator._wait_for_callback = MagicMock(return_value={"state": "state-value"}) + + with pytest.raises(XAIOAuthError, match="no code returned"): + authenticator.login(no_browser=True) + + +def test_pkce_pair_generates_s256_challenge(): + verifier, challenge = XAIOAuthAuthenticator()._pkce_pair() + expected = ( + base64.urlsafe_b64encode(hashlib.sha256(verifier.encode()).digest()) + .rstrip(b"=") + .decode() + ) + + assert challenge == expected + assert "=" not in verifier + assert "=" not in challenge + + +def test_build_authorize_url_contains_xai_oauth_parameters(): + authorize_url = XAIOAuthAuthenticator()._build_authorize_url( + authorization_endpoint="https://auth.x.ai/oauth/authorize", + redirect_uri="http://127.0.0.1:56121/callback", + challenge="pkce-challenge", + state="state-value", + nonce="nonce-value", + ) + parsed = urlparse(authorize_url) + params = parse_qs(parsed.query) + + assert parsed.scheme == "https" + assert parsed.netloc == "auth.x.ai" + assert params["response_type"] == ["code"] + assert params["client_id"] == [XAI_OAUTH_CLIENT_ID] + assert params["scope"] == [XAI_OAUTH_SCOPE] + assert params["code_challenge"] == ["pkce-challenge"] + assert params["code_challenge_method"] == ["S256"] + assert params["state"] == ["state-value"] + assert params["nonce"] == ["nonce-value"] + + +def test_get_llm_provider_uses_single_xai_provider(monkeypatch): + monkeypatch.setenv("XAI_API_KEY", "api-key") + + model, provider, api_key, api_base = get_llm_provider("xai/grok-4") + + assert model == "grok-4" + assert provider == "xai" + assert api_key == "api-key" + assert api_base == "https://api.x.ai/v1" + + +def test_xai_oauth_alias_is_not_a_provider(): + with pytest.raises(Exception): + get_llm_provider("xai_oauth/grok-4") + + +def test_chat_config_wraps_flagged_oauth_errors_as_authentication_error( + tmp_path, monkeypatch +): + monkeypatch.setenv("XAI_OAUTH_TOKEN_DIR", str(tmp_path / "missing")) + + with pytest.raises(litellm.AuthenticationError) as exc_info: + XAIChatConfig().validate_environment( + headers={}, + model="grok-4", + messages=[], + optional_params={}, + litellm_params={"use_xai_oauth": True}, + api_key=None, + ) + + assert exc_info.value.llm_provider == "xai" + assert "litellm xai-oauth login" in str(exc_info.value) + + +def test_chat_config_injects_flagged_oauth_token(tmp_path, monkeypatch): + token_dir, _ = _write_auth_file( + tmp_path, + { + "access_token": "chat-token", + "refresh_token": "refresh-token", + "expires_at": time.time() + 3600, + }, + ) + monkeypatch.setenv("XAI_OAUTH_TOKEN_DIR", str(token_dir)) + + headers = XAIChatConfig().validate_environment( + headers={}, + model="grok-4", + messages=[], + optional_params={}, + litellm_params={"use_xai_oauth": True}, + api_key=None, + ) + + assert headers["Authorization"] == "Bearer chat-token" + + +def test_chat_config_ignores_api_base_override_for_flagged_oauth(monkeypatch): + monkeypatch.setenv("XAI_OAUTH_API_BASE", "https://api.x.ai/v1") + + url = XAIChatConfig().get_complete_url( + api_base="https://attacker.example.com/v1", + api_key=None, + model="grok-4", + optional_params={}, + litellm_params={"use_xai_oauth": True}, + ) + + assert url == "https://api.x.ai/v1/chat/completions" + + +def test_chat_config_treats_blank_api_key_as_absent_for_flagged_oauth( + tmp_path, monkeypatch +): + token_dir, _ = _write_auth_file( + tmp_path, + { + "access_token": "stored-oauth-token", + "refresh_token": "refresh-token", + "expires_at": time.time() + 3600, + }, + ) + monkeypatch.setenv("XAI_OAUTH_TOKEN_DIR", str(token_dir)) + + headers = XAIChatConfig().validate_environment( + headers={}, + model="grok-4", + messages=[], + optional_params={}, + litellm_params={"use_xai_oauth": True}, + api_key="", + ) + + assert headers["Authorization"] == "Bearer stored-oauth-token" + + +def test_chat_config_allows_api_base_override_with_caller_api_key(): + headers = XAIChatConfig().validate_environment( + headers={}, + model="grok-4", + messages=[], + optional_params={}, + litellm_params={"use_xai_oauth": True}, + api_key="caller-api-key", + ) + url = XAIChatConfig().get_complete_url( + api_base="https://custom.example.com/v1", + api_key="caller-api-key", + model="grok-4", + optional_params={}, + litellm_params={"use_xai_oauth": True}, + ) + + assert headers["Authorization"] == "Bearer caller-api-key" + assert url == "https://custom.example.com/v1/chat/completions" + + +def test_chat_config_prioritizes_env_api_key_over_oauth_flag(monkeypatch): + monkeypatch.setenv("XAI_API_KEY", "env-api-key") + + headers = XAIChatConfig().validate_environment( + headers={}, + model="grok-4", + messages=[], + optional_params={}, + litellm_params={"use_xai_oauth": True}, + api_key=None, + ) + url = XAIChatConfig().get_complete_url( + api_base="https://custom.example.com/v1", + api_key=None, + model="grok-4", + optional_params={}, + litellm_params={"use_xai_oauth": True}, + ) + + assert headers["Authorization"] == "Bearer env-api-key" + assert url == "https://custom.example.com/v1/chat/completions" + + +def test_validate_environment_still_reports_xai_api_key(monkeypatch): + monkeypatch.setenv("XAI_API_KEY", "env-api-key") + + assert validate_environment("xai/grok-4") == { + "keys_in_environment": True, + "missing_keys": [], + } + + +def test_xai_oauth_flag_uses_xai_optional_param_mapping(): + litellm_params = GenericLiteLLMParams(use_xai_oauth=True) + optional_params = get_optional_params( + model="grok-4", + custom_llm_provider="xai", + temperature=0.2, + max_tokens=8, + ) + + assert optional_params["temperature"] == 0.2 + assert optional_params["max_tokens"] == 8 + assert litellm_params.use_xai_oauth is True + assert "use_xai_oauth" not in optional_params + + +def test_responses_config_injects_flagged_oauth_bearer_token(tmp_path, monkeypatch): + token_dir, _ = _write_auth_file( + tmp_path, + { + "access_token": "responses-token", + "refresh_token": "refresh-token", + "expires_at": time.time() + 3600, + }, + ) + monkeypatch.setenv("XAI_OAUTH_TOKEN_DIR", str(token_dir)) + + headers = XAIResponsesAPIConfig().validate_environment( + headers={}, + model="grok-4", + litellm_params=GenericLiteLLMParams(use_xai_oauth=True), + ) + + assert headers["Authorization"] == "Bearer responses-token" + + +def test_responses_config_endpoint_url_uses_oauth_authenticator(monkeypatch): + monkeypatch.setenv("XAI_OAUTH_API_BASE", "https://xai.example.com/v1/") + config = XAIResponsesAPIConfig() + + assert config.get_complete_url( + api_base=None, litellm_params={"use_xai_oauth": True} + ) == ("https://xai.example.com/v1/responses") + assert ( + config.get_complete_url( + api_base="https://custom.example.com/v1/", + litellm_params={"use_xai_oauth": True}, + ) + == "https://xai.example.com/v1/responses" + ) + assert ( + config.get_complete_url( + api_base="https://custom.example.com/v1/", + litellm_params={"api_key": "", "use_xai_oauth": True}, + ) + == "https://xai.example.com/v1/responses" + ) + assert ( + config.get_complete_url( + api_base="https://custom.example.com/v1/", + litellm_params={"api_key": "caller-api-key"}, + ) + == "https://custom.example.com/v1/responses" + ) + + +def test_responses_config_wraps_flagged_oauth_errors_as_authentication_error( + tmp_path, monkeypatch +): + monkeypatch.setenv("XAI_OAUTH_TOKEN_DIR", str(tmp_path / "missing")) + + with pytest.raises(litellm.AuthenticationError) as exc_info: + XAIResponsesAPIConfig().validate_environment( + headers={}, + model="grok-4", + litellm_params=GenericLiteLLMParams(use_xai_oauth=True), + ) + + assert XAIResponsesAPIConfig().custom_llm_provider.value == "xai" + assert exc_info.value.llm_provider == "xai" + + +def test_proxy_cli_xai_oauth_login_uses_single_authenticator(monkeypatch): + from litellm.proxy.proxy_cli import run_server + + instances = [] + + class FakeAuthenticator: + auth_file = "/tmp/xai-oauth-auth.json" + + def __init__(self): + instances.append(self) + + def login(self): + return {"expires_at": 1234567890} + + monkeypatch.setattr( + "litellm.llms.xai.oauth.XAIOAuthAuthenticator", FakeAuthenticator + ) + + result = CliRunner().invoke(run_server, ["xai-oauth", "login"]) + + assert result.exit_code == 0 + assert len(instances) == 1 + assert "Credentials saved to /tmp/xai-oauth-auth.json" in result.output + assert "Access token expires at 1234567890" in result.output diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 0b1240f8bac..b6550fee6b9 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -4970,17 +4970,143 @@ class TestGatewayCreateInitializationOptions: try: from litellm.proxy._experimental.mcp_server.mcp_context import ( _mcp_gateway_initialize_instructions, + _mcp_gateway_server_name, ) from litellm.proxy._experimental.mcp_server.server import server except ImportError: pytest.skip("MCP server not available") - tok = _mcp_gateway_initialize_instructions.set(None) + instructions_token = _mcp_gateway_initialize_instructions.set(None) + server_name_token = _mcp_gateway_server_name.set(None) try: opts = server.create_initialization_options() assert getattr(opts, "instructions", None) is None + assert opts.server_name == "litellm-mcp-server" finally: - _mcp_gateway_initialize_instructions.reset(tok) + _mcp_gateway_initialize_instructions.reset(instructions_token) + _mcp_gateway_server_name.reset(server_name_token) + + @pytest.mark.asyncio + async def test_scoped_request_uses_configured_server_alias(self): + try: + from litellm.proxy._experimental.mcp_server.server import ( + _gateway_initialize_instructions_request_scope, + global_mcp_server_manager, + server, + ) + except ImportError: + pytest.skip("MCP server not available") + + scoped_server = MCPServer( + server_id="server-123", + name="upstream-server", + alias="grafana", + transport=MCPTransport.http, + url="https://example.com/mcp", + ) + + with ( + patch( + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new_callable=AsyncMock, + return_value=[scoped_server], + ), + patch.object( + global_mcp_server_manager, + "_ensure_upstream_initialize_instructions_cached", + new_callable=AsyncMock, + ), + ): + async with _gateway_initialize_instructions_request_scope( + user_api_key_auth=None, + mcp_servers=["grafana"], + client_ip=None, + scoped_server_endpoint=True, + ): + assert server.create_initialization_options().server_name == "grafana" + + assert ( + server.create_initialization_options().server_name == "litellm-mcp-server" + ) + + @pytest.mark.asyncio + async def test_sse_handler_scopes_server_name_from_single_server_path(self): + try: + from litellm.proxy._experimental.mcp_server import server as mcp_server + from litellm.proxy._experimental.mcp_server.server import ( + global_mcp_server_manager, + handle_sse_mcp, + server, + ) + except ImportError: + pytest.skip("MCP server not available") + + scoped_server = MCPServer( + server_id="server-123", + name="upstream-server", + alias="grafana", + transport=MCPTransport.http, + url="https://example.com/mcp", + ) + captured = {} + + async def record_request(scope, receive, send): + captured["server_name"] = server.create_initialization_options().server_name + + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/grafana", + "headers": [], + } + + with ( + patch( + "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", + new_callable=AsyncMock, + return_value=( + UserAPIKeyAuth(api_key="sk-test"), + None, + ["grafana"], + None, + None, + None, + ), + ), + patch( + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new_callable=AsyncMock, + return_value=[scoped_server], + ), + patch.object( + global_mcp_server_manager, + "_ensure_upstream_initialize_instructions_cached", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy._experimental.mcp_server.server._raise_preemptive_401_for_unauthenticated_servers", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy._experimental.mcp_server.server._check_passthrough_upstream_auth", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", + True, + ), + patch.object( + mcp_server.sse_session_manager, + "handle_request", + side_effect=record_request, + ), + ): + await handle_sse_mcp(scope, AsyncMock(), AsyncMock()) + + assert captured["server_name"] == "grafana" + assert ( + server.create_initialization_options().server_name == "litellm-mcp-server" + ) def test_contextvar_set_injects_instructions(self): """When ContextVar has a value, it appears in InitializationOptions.""" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 48c09f6e456..1b815b7a1c9 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -30,6 +30,7 @@ from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( MCPServerManager, _deserialize_json_dict, _deserialize_json_list, + _normalize_mcp_server_cost_info, ) from litellm.proxy._types import ( LiteLLM_MCPServerTable, @@ -257,6 +258,69 @@ class TestMCPServerManager: assert server.alias == "friendly_alias" assert server.server_name == "validserver" + @pytest.mark.asyncio + async def test_load_servers_from_config_coerces_cost_string_to_float(self): + """YAML 1.1 parses `7e-05` as a string; ingest must coerce it to float.""" + manager = MCPServerManager() + config = { + "google_maps": { + "url": "https://example.com/mcp", + "transport": MCPTransport.http, + "mcp_info": { + "mcp_server_cost_info": { + "default_cost_per_query": "7e-05", + "tool_name_to_cost_per_query": {"geocode": "1e-3"}, + } + }, + } + } + + await manager.load_servers_from_config(config) + + server = next(iter(manager.config_mcp_servers.values())) + cost_info = server.mcp_info["mcp_server_cost_info"] + assert cost_info["default_cost_per_query"] == 7e-05 + assert isinstance(cost_info["default_cost_per_query"], float) + assert cost_info["tool_name_to_cost_per_query"]["geocode"] == 1e-3 + assert isinstance(cost_info["tool_name_to_cost_per_query"]["geocode"], float) + + def test_normalize_mcp_server_cost_info_preserves_float_values(self): + mcp_info = { + "server_name": "maps", + "mcp_server_cost_info": { + "default_cost_per_query": 0.01, + "tool_name_to_cost_per_query": {"search": 0.05}, + }, + } + + _normalize_mcp_server_cost_info(mcp_info) + + cost_info = mcp_info["mcp_server_cost_info"] + assert cost_info["default_cost_per_query"] == 0.01 + assert cost_info["tool_name_to_cost_per_query"] == {"search": 0.05} + + def test_normalize_mcp_server_cost_info_drops_non_numeric_values(self): + mcp_info = { + "server_name": "maps", + "mcp_server_cost_info": { + "default_cost_per_query": "not-a-number", + "tool_name_to_cost_per_query": {"search": "free", "geocode": "2e-4"}, + }, + } + + _normalize_mcp_server_cost_info(mcp_info) + + cost_info = mcp_info["mcp_server_cost_info"] + assert "default_cost_per_query" not in cost_info + assert cost_info["tool_name_to_cost_per_query"] == {"geocode": 2e-4} + + def test_normalize_mcp_server_cost_info_leaves_missing_cost_info_alone(self): + mcp_info = {"server_name": "maps"} + + _normalize_mcp_server_cost_info(mcp_info) + + assert "mcp_server_cost_info" not in mcp_info + def test_warns_when_custom_separator_invalid(self, monkeypatch, caplog): """Invalid MCP_TOOL_PREFIX_SEPARATOR values should log a warning.""" diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 42c76c4671d..e14ef05bd43 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -2409,6 +2409,8 @@ async def test_virtual_key_budget_check_fallback_no_counter(): assert exc_info.value.current_cost == 15.0 + + @pytest.mark.asyncio async def test_team_budget_check_reads_from_spend_counter(): """Team budget check should use get_current_spend when counter exists.""" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py index 565bf83c6a2..8a5eeeff367 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -2398,11 +2398,9 @@ async def test_anonymize_text_uses_correct_positions_no_parse_pii(): ) expected = "My name is , my email is , phone " - assert result == expected, ( - f"anonymize_text produced garbled output with PII remnants.\n" - f"Expected: {expected!r}\n" - f"Got: {result!r}" - ) + assert ( + result == expected + ), f"anonymize_text produced garbled output with PII remnants.\nExpected: {expected!r}\nGot: {result!r}" assert masked_entity_count == { "PERSON": 1, "EMAIL_ADDRESS": 1, @@ -2495,3 +2493,157 @@ async def test_anonymize_text_uses_correct_positions_with_parse_pii(): assert pii_tokens.get("") == "John Smith" assert pii_tokens.get("") == "john@example.com" assert pii_tokens.get("") == "555-867-5309" + + +def test_unmask_sse_bytes_chunk_replaces_text_delta(): + import json + + pii_tokens = {"": "Bobby"} + event = { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "Hello , how are you?"}, + } + chunk = ("data: " + json.dumps(event) + "\n\n").encode("utf-8") + + result = _OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk(chunk, pii_tokens) + + decoded = result.decode("utf-8") + parsed = json.loads(decoded.split("data: ", 1)[1].strip()) + assert parsed["delta"]["text"] == "Hello Bobby, how are you?" + + +def test_unmask_sse_bytes_chunk_ignores_non_text_delta(): + import json + + pii_tokens = {"": "Bobby"} + + # message_start event — no delta + event = {"type": "message_start", "message": {"id": "msg_01", "role": "assistant"}} + chunk = ("data: " + json.dumps(event) + "\n\n").encode("utf-8") + result = _OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk(chunk, pii_tokens) + assert result == chunk + + # input_json_delta — should not be touched + event2 = { + "type": "content_block_delta", + "index": 1, + "delta": {"type": "input_json_delta", "partial_json": '{"name": ""}'}, + } + chunk2 = ("data: " + json.dumps(event2) + "\n\n").encode("utf-8") + result2 = _OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk(chunk2, pii_tokens) + assert result2 == chunk2 + + +def test_unmask_sse_bytes_chunk_handles_malformed_json(): + chunk = b"data: {not valid json}\n\n" + result = _OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk( + chunk, {"": "Bobby"} + ) + assert result == chunk + + +def test_unmask_sse_bytes_chunk_handles_unicode_decode_error(): + chunk = b"\xff\xfe invalid utf-8" + result = _OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk( + chunk, {"": "Bobby"} + ) + assert result == chunk + + +def test_unmask_sse_bytes_chunk_non_ascii_pii_not_escaped(): + import json + + pii_tokens = {"": "José"} + event = { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "Hello !"}, + } + chunk = ("data: " + json.dumps(event) + "\n\n").encode("utf-8") + + result = _OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk(chunk, pii_tokens) + + decoded = result.decode("utf-8") + assert "Jos\\u" not in decoded + parsed = json.loads(decoded.split("data: ", 1)[1].strip()) + assert parsed["delta"]["text"] == "Hello José!" + + +def test_unmask_sse_bytes_chunk_handles_crlf_line_endings(): + import json + + pii_tokens = {"": "Bobby"} + event = { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "Hi !"}, + } + crlf_chunk = ("data: " + json.dumps(event) + "\r\ndata: [DONE]\r\n").encode("utf-8") + + result = _OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk( + crlf_chunk, pii_tokens + ) + + decoded = result.decode("utf-8") + parsed = json.loads(decoded.split("data: ", 1)[1].split("\n")[0].strip()) + assert parsed["delta"]["text"] == "Hi Bobby!" + assert "data: [DONE]" in decoded + + +@pytest.mark.asyncio +async def test_stream_pii_unmasking_unmaskes_bytes_chunks(mock_user_api_key): + import json + + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + output_parse_pii=True, + ) + + pii_tokens = {"": "Bobby"} + request_data = {"metadata": {"pii_tokens": pii_tokens}} + + def _make_sse_chunk(text: str) -> bytes: + event = { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": text}, + } + return ("data: " + json.dumps(event) + "\n\n").encode("utf-8") + + async def mock_stream(): + yield _make_sse_chunk("Hello !") + yield _make_sse_chunk(" How can I help?") + + chunks = [] + async for chunk in guardrail._stream_pii_unmasking(mock_stream(), request_data): + chunks.append(chunk) + + assert len(chunks) == 2 + first = chunks[0].decode("utf-8") + first_event = json.loads(first.split("data: ", 1)[1].strip()) + assert first_event["delta"]["text"] == "Hello Bobby!" + + second = chunks[1].decode("utf-8") + second_event = json.loads(second.split("data: ", 1)[1].strip()) + assert second_event["delta"]["text"] == " How can I help?" + + +@pytest.mark.asyncio +async def test_stream_pii_unmasking_passthrough_when_no_tokens(mock_user_api_key): + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + output_parse_pii=True, + ) + + raw_chunk = b"data: {}\n\n" + request_data: dict = {"metadata": {}} + + async def mock_stream(): + yield raw_chunk + + chunks = [] + async for chunk in guardrail._stream_pii_unmasking(mock_stream(), request_data): + chunks.append(chunk) + + assert chunks == [raw_chunk] diff --git a/tests/test_litellm/proxy/hooks/litellm_skills/test_main.py b/tests/test_litellm/proxy/hooks/litellm_skills/test_main.py new file mode 100644 index 00000000000..f716a8533d8 --- /dev/null +++ b/tests/test_litellm/proxy/hooks/litellm_skills/test_main.py @@ -0,0 +1,67 @@ +from unittest.mock import AsyncMock, patch + +import pytest + +from litellm.proxy.hooks.litellm_skills.main import SkillsInjectionHook + +SKILL_TOOL_NAME = "litellm_skill_e2b8dca8_031a_4481_b034_b9ec7d4eb7bf" + + +def _request_data(): + return { + "model": "claude-sonnet-4-5", + "messages": [{"role": "user", "content": "run the skill"}], + "litellm_metadata": { + "_litellm_code_execution_enabled": True, + "_skill_files": {SKILL_TOOL_NAME: {"main.py": b"print('hi')"}}, + }, + } + + +def _tool_use_response(tool_name): + return { + "stop_reason": "tool_use", + "content": [ + {"type": "tool_use", "id": "toolu_1", "name": tool_name, "input": {}} + ], + } + + +@pytest.mark.asyncio +async def test_post_call_success_hook_executes_litellm_skill_tool(): + """DB skill tool names carry the litellm_skill_ prefix and must trigger the execution loop.""" + hook = SkillsInjectionHook() + response = _tool_use_response(SKILL_TOOL_NAME) + + with patch.object( + hook, "_execute_code_loop_messages_api", new=AsyncMock(return_value=response) + ) as mock_loop: + result = await hook.async_post_call_success_deployment_hook( + request_data=_request_data(), response=response, call_type=None + ) + + mock_loop.assert_awaited_once() + assert result is response + + +@pytest.mark.asyncio +async def test_execute_code_loop_dispatches_litellm_skill_tool(): + """The agentic loop must route litellm_skill_ tool calls to _execute_skill_tool.""" + hook = SkillsInjectionHook() + final_response = {"stop_reason": "end_turn", "content": []} + + with ( + patch.object( + hook, "_execute_skill_tool", new=AsyncMock(return_value="skill ran") + ) as mock_exec, + patch("litellm.anthropic.acreate", new=AsyncMock(return_value=final_response)), + ): + result = await hook._execute_code_loop_messages_api( + data=_request_data(), + response=_tool_use_response(SKILL_TOOL_NAME), + skill_files={"main.py": b"print('hi')"}, + ) + + mock_exec.assert_awaited_once() + assert mock_exec.await_args.args[0] == SKILL_TOOL_NAME + assert result is final_response diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter.py new file mode 100644 index 00000000000..0e2683dcbfd --- /dev/null +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter.py @@ -0,0 +1,86 @@ +""" +Unit Tests for the max parallel request limiter v1 for the proxy +""" + +from datetime import datetime + +import pytest + +from litellm.caching.caching import DualCache +from litellm.proxy.hooks.parallel_request_limiter import ( + _PROXY_MaxParallelRequestsHandler, +) +from litellm.proxy.utils import InternalUsageCache, hash_token +from litellm.types.utils import EmbeddingResponse, TextCompletionResponse, Usage + + +@pytest.mark.parametrize( + "response_obj", + [ + EmbeddingResponse( + model="text-embedding-3-small", + usage=Usage(prompt_tokens=50, completion_tokens=0, total_tokens=50), + ), + TextCompletionResponse( + model="gpt-3.5-turbo-instruct", + usage=Usage(prompt_tokens=20, completion_tokens=30, total_tokens=50), + ), + ], +) +@pytest.mark.asyncio +async def test_async_log_success_event_counts_non_chat_response_tokens(response_obj): + """ + Embedding and text completion responses must increment the per key, user, + team, and end user TPM counters, not just chat completion ModelResponse + objects. + """ + _api_key = hash_token("sk-12345") + user_id = "ishaan" + team_id = "litellm-team" + end_user_id = "customer-1" + + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()) + ) + + current_date = datetime.now().strftime("%Y-%m-%d") + current_hour = datetime.now().strftime("%H") + current_minute = datetime.now().strftime("%M") + precise_minute = f"{current_date}-{current_hour}-{current_minute}" + + scope_ids = [_api_key, user_id, team_id, end_user_id] + for scope_id in scope_ids: + await parallel_request_handler.internal_usage_cache.async_set_cache( + key=f"{scope_id}::{precise_minute}::request_count", + value={"current_requests": 1, "current_tpm": 0, "current_rpm": 1}, + litellm_parent_otel_span=None, + ) + + kwargs = { + "litellm_params": { + "metadata": { + "user_api_key": _api_key, + "user_api_key_user_id": user_id, + "user_api_key_team_id": team_id, + "user_api_key_model_max_budget": {}, + } + }, + "user": end_user_id, + } + + await parallel_request_handler.async_log_success_event( + kwargs=kwargs, + response_obj=response_obj, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + for scope_id in scope_ids: + current = await parallel_request_handler.internal_usage_cache.async_get_cache( + key=f"{scope_id}::{precise_minute}::request_count", + litellm_parent_otel_span=None, + ) + assert current["current_tpm"] == 50, ( + f"expected 50 tokens counted for {scope_id}, " + f"got {current['current_tpm']}" + ) diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index 676f623a5dd..d10311b9f41 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -20,7 +20,12 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import ( _PROXY_MaxParallelRequestsHandler_v3 as _PROXY_MaxParallelRequestsHandler, ) from litellm.proxy.utils import InternalUsageCache, ProxyLogging, hash_token -from litellm.types.utils import ModelResponse, Usage +from litellm.types.utils import ( + EmbeddingResponse, + ModelResponse, + TextCompletionResponse, + Usage, +) class TimeController: @@ -547,6 +552,68 @@ async def test_token_rate_limit_type_respected_v3(monkeypatch, token_rate_limit_ ), f"Expected {expected_tokens[token_rate_limit_type]} tokens for type '{token_rate_limit_type}', got {tpm_operation['increment_value']}" +@pytest.mark.parametrize( + "response_obj", + [ + EmbeddingResponse( + model="text-embedding-3-small", + usage=Usage(prompt_tokens=50, completion_tokens=0, total_tokens=50), + ), + TextCompletionResponse( + model="gpt-3.5-turbo-instruct", + usage=Usage(prompt_tokens=20, completion_tokens=30, total_tokens=50), + ), + ], +) +@pytest.mark.asyncio +async def test_async_log_success_event_counts_non_chat_response_tokens( + monkeypatch, response_obj +): + """ + Embedding and text completion responses must increment the TPM counter, + not just chat completion ModelResponse objects. + """ + monkeypatch.setenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", "60") + + _api_key = hash_token("sk-12345") + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()) + ) + monkeypatch.setattr( + parallel_request_handler, "get_rate_limit_type", lambda: "total" + ) + + mock_kwargs = { + "standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}}, + "model": response_obj.model, + } + + captured_operations = [] + + async def mock_increment_pipeline(increment_list, **kwargs): + captured_operations.extend(increment_list) + return True + + monkeypatch.setattr( + parallel_request_handler.internal_usage_cache.dual_cache, + "async_increment_cache_pipeline", + mock_increment_pipeline, + ) + + await parallel_request_handler.async_log_success_event( + kwargs=mock_kwargs, + response_obj=response_obj, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + tpm_operation = next( + (op for op in captured_operations if op["key"].endswith(":tokens")), None + ) + assert tpm_operation is not None, "Should have a TPM increment operation" + assert tpm_operation["increment_value"] == 50 + + @pytest.mark.asyncio async def test_async_log_failure_event_v3(): """ diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 3c212d86e65..473d61f8a85 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -6496,6 +6496,9 @@ async def test_reset_key_spend_success(monkeypatch): patch( "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object" ) as mock_delete_cache, + patch( + "litellm.proxy.proxy_server._invalidate_spend_counter" + ) as mock_invalidate, ): mock_hash_token.return_value = hashed_key mock_check_admin.return_value = None @@ -6520,6 +6523,76 @@ async def test_reset_key_spend_success(monkeypatch): assert response["max_budget"] == 200.0 mock_prisma_client.db.litellm_verificationtoken.update.assert_called_once() mock_delete_cache.assert_awaited_once() + mock_invalidate.assert_awaited_once_with(counter_key=f"spend:key:{hashed_key}") + + +@pytest.mark.asyncio +async def test_update_key_spend_invalidates_counter(monkeypatch): + """ + Test that updating a key's spend via update_key_fn immediately invalidates the spend counter. + """ + from litellm.proxy.management_endpoints.key_management_endpoints import ( + update_key_fn, + ) + + mock_prisma_client = AsyncMock() + mock_user_api_key_cache = AsyncMock() + mock_proxy_logging_obj = MagicMock() + + hashed_key = "0d62f396c1317066f55a96086517047c737087c61eb2bf016b72e6298927b15b" + key_in_db = LiteLLM_VerificationToken( + token=hashed_key, + user_id="test-user", + spend=10.0, + max_budget=200.0, + litellm_budget_table=None, + ) + + mock_prisma_client.get_data = AsyncMock(return_value=key_in_db) + mock_prisma_client.update_data = AsyncMock(return_value={"data": {"spend": 0.0}}) + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=key_in_db + ) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr( + "litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache + ) + monkeypatch.setattr( + "litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + monkeypatch.setattr("litellm.store_audit_logs", False) + + with ( + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object" + ) as mock_delete_cache, + patch( + "litellm.proxy.proxy_server._invalidate_spend_counter" + ) as mock_invalidate, + ): + mock_delete_cache.return_value = None + + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-admin", + user_id="admin-user", + ) + + mock_request = MagicMock() + mock_request.query_params = {} + + await update_key_fn( + request=mock_request, + data=UpdateKeyRequest(key="sk-test-key", spend=0.0), + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + ) + + mock_delete_cache.assert_awaited_once() + mock_invalidate.assert_awaited_once_with(counter_key=f"spend:key:{hashed_key}") @pytest.mark.asyncio @@ -11668,3 +11741,84 @@ async def test_ghsa_q775_default_team_id_does_not_grant_session_token_exemption( msg = str(getattr(err, "detail", "")) + str(getattr(err, "message", "")) assert str(code) == "400" assert "cannot exceed" in msg.lower() + + + +@pytest.mark.asyncio +async def test_prepare_key_update_data_budget_duration_null_clears_fields(): + """ + When budget_duration is explicitly set to null, prepare_key_update_data + should produce budget_duration=None and budget_reset_at=None so Prisma + clears them in the DB. + """ + existing_key = LiteLLM_VerificationToken( + token="test-token", + key_alias="test-key", + models=[], + user_id="test-user", + team_id=None, + metadata={}, + ) + + update_request = UpdateKeyRequest(key="test-token", budget_duration=None) + + result = await prepare_key_update_data( + data=update_request, existing_key_row=existing_key + ) + + assert "budget_duration" in result + assert result["budget_duration"] is None + assert "budget_reset_at" in result + assert result["budget_reset_at"] is None + + +@pytest.mark.asyncio +async def test_prepare_key_update_data_budget_duration_not_sent_excluded(): + """ + When budget_duration is NOT sent in the request (unset), it should not + appear in the result dict at all — the existing DB value stays unchanged. + """ + existing_key = LiteLLM_VerificationToken( + token="test-token", + key_alias="test-key", + models=[], + user_id="test-user", + team_id=None, + metadata={}, + ) + + update_request = UpdateKeyRequest(key="test-token", models=["gpt-4"]) + + result = await prepare_key_update_data( + data=update_request, existing_key_row=existing_key + ) + + assert "budget_duration" not in result + assert "budget_reset_at" not in result + + +@pytest.mark.asyncio +async def test_prepare_key_update_data_budget_duration_valid_sets_reset(): + """ + When budget_duration is set to a valid duration string, both + budget_duration and budget_reset_at should be populated. + """ + existing_key = LiteLLM_VerificationToken( + token="test-token", + key_alias="test-key", + models=[], + user_id="test-user", + team_id=None, + metadata={}, + ) + + update_request = UpdateKeyRequest(key="test-token", budget_duration="30d") + + result = await prepare_key_update_data( + data=update_request, existing_key_row=existing_key + ) + + assert result["budget_duration"] == "30d" + assert result["budget_reset_at"] is not None + + diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index b750b6d022c..d4bc3841668 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -8886,3 +8886,329 @@ async def test_update_team_blocks_non_admin_passthrough_routes(mock_db_client): ) assert str(exc.value.code) == "403" assert "allowed_passthrough_routes" in str(exc.value.message) + + +def test_set_budget_reset_at_clears_when_budget_duration_null(): + """ + When budget_duration is explicitly set to null, _set_budget_reset_at + should set budget_reset_at=None in updated_kv so Prisma clears it in the DB. + """ + from litellm.proxy._types import UpdateTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import _set_budget_reset_at + + data = UpdateTeamRequest(team_id="test-team", budget_duration=None) + updated_kv = {"team_id": "test-team", "budget_duration": None} + + _set_budget_reset_at(data, updated_kv) + + assert "budget_reset_at" in updated_kv + assert updated_kv["budget_reset_at"] is None + + +def test_set_budget_reset_at_noop_when_budget_duration_not_sent(): + """ + When budget_duration is NOT sent (unset), _set_budget_reset_at should + not add budget_reset_at to updated_kv. + """ + from litellm.proxy._types import UpdateTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import _set_budget_reset_at + + data = UpdateTeamRequest(team_id="test-team") + updated_kv = {"team_id": "test-team"} + + _set_budget_reset_at(data, updated_kv) + + assert "budget_reset_at" not in updated_kv + + +def test_set_budget_reset_at_sets_value_when_budget_duration_provided(): + """ + When budget_duration is set to a valid string, _set_budget_reset_at + should compute and set budget_reset_at. + """ + from litellm.proxy._types import UpdateTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import _set_budget_reset_at + + data = UpdateTeamRequest(team_id="test-team", budget_duration="30d") + updated_kv = {"team_id": "test-team", "budget_duration": "30d"} + + _set_budget_reset_at(data, updated_kv) + + assert "budget_reset_at" in updated_kv + assert updated_kv["budget_reset_at"] is not None + + +@pytest.mark.asyncio +async def test_clear_team_member_budget_duration_calls_update_budget(): + """ + When team_member_budget_duration is explicitly null and a budget row + exists, clear_team_member_budget_fields should call update_budget + with budget_duration=None and budget_reset_at=None. + """ + from litellm.proxy.management_endpoints.team_endpoints import ( + TeamMemberBudgetHandler, + ) + + mock_user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="admin-user", + ) + + team_table = LiteLLM_TeamTable( + team_id="test-team", + metadata={"team_member_budget_id": "budget-123"}, + members_with_roles=[], + ) + + updated_kv = { + "team_id": "test-team", + "team_member_budget_duration": None, + } + + with patch( + "litellm.proxy.management_endpoints.budget_management_endpoints.update_budget", + new_callable=AsyncMock, + ) as mock_update_budget: + result = await TeamMemberBudgetHandler.clear_team_member_budget_fields( + team_table=team_table, + user_api_key_dict=mock_user_api_key_dict, + updated_kv=updated_kv, + explicitly_set_fields={"team_member_budget_duration"}, + ) + + mock_update_budget.assert_awaited_once() + budget_request = mock_update_budget.call_args.kwargs["budget_obj"] + assert budget_request.budget_id == "budget-123" + assert "budget_duration" in budget_request.model_fields_set + assert budget_request.budget_duration is None + assert "budget_reset_at" in budget_request.model_fields_set + assert budget_request.budget_reset_at is None + assert "team_member_budget_duration" not in result + + +@pytest.mark.asyncio +async def test_clear_team_member_budget_clears_max_budget(): + """ + When team_member_budget is explicitly null, clear_team_member_budget_fields + should call update_budget with max_budget=None. + """ + from litellm.proxy.management_endpoints.team_endpoints import ( + TeamMemberBudgetHandler, + ) + + mock_user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="admin-user", + ) + + team_table = LiteLLM_TeamTable( + team_id="test-team", + metadata={"team_member_budget_id": "budget-456"}, + members_with_roles=[], + ) + + updated_kv = { + "team_id": "test-team", + "team_member_budget": None, + } + + with patch( + "litellm.proxy.management_endpoints.budget_management_endpoints.update_budget", + new_callable=AsyncMock, + ) as mock_update_budget: + result = await TeamMemberBudgetHandler.clear_team_member_budget_fields( + team_table=team_table, + user_api_key_dict=mock_user_api_key_dict, + updated_kv=updated_kv, + explicitly_set_fields={"team_member_budget"}, + ) + + mock_update_budget.assert_awaited_once() + budget_request = mock_update_budget.call_args.kwargs["budget_obj"] + assert budget_request.budget_id == "budget-456" + assert "max_budget" in budget_request.model_fields_set + assert budget_request.max_budget is None + assert "team_member_budget" not in result + + +@pytest.mark.asyncio +async def test_clear_team_member_rpm_tpm_limits(): + """ + When team_member_rpm_limit and team_member_tpm_limit are explicitly null, + clear_team_member_budget_fields should clear both on the budget row. + """ + from litellm.proxy.management_endpoints.team_endpoints import ( + TeamMemberBudgetHandler, + ) + + mock_user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="admin-user", + ) + + team_table = LiteLLM_TeamTable( + team_id="test-team", + metadata={"team_member_budget_id": "budget-789"}, + members_with_roles=[], + ) + + updated_kv = { + "team_id": "test-team", + "team_member_rpm_limit": None, + "team_member_tpm_limit": None, + } + + with patch( + "litellm.proxy.management_endpoints.budget_management_endpoints.update_budget", + new_callable=AsyncMock, + ) as mock_update_budget: + result = await TeamMemberBudgetHandler.clear_team_member_budget_fields( + team_table=team_table, + user_api_key_dict=mock_user_api_key_dict, + updated_kv=updated_kv, + explicitly_set_fields={"team_member_rpm_limit", "team_member_tpm_limit"}, + ) + + mock_update_budget.assert_awaited_once() + budget_request = mock_update_budget.call_args.kwargs["budget_obj"] + assert budget_request.budget_id == "budget-789" + assert "rpm_limit" in budget_request.model_fields_set + assert budget_request.rpm_limit is None + assert "tpm_limit" in budget_request.model_fields_set + assert budget_request.tpm_limit is None + assert "team_member_rpm_limit" not in result + assert "team_member_tpm_limit" not in result + + +@pytest.mark.asyncio +async def test_clear_all_team_member_fields_at_once(): + """ + When all team_member fields are explicitly null, all corresponding + budget row fields should be cleared in a single update. + """ + from litellm.proxy.management_endpoints.team_endpoints import ( + TeamMemberBudgetHandler, + ) + + mock_user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="admin-user", + ) + + team_table = LiteLLM_TeamTable( + team_id="test-team", + metadata={"team_member_budget_id": "budget-all"}, + members_with_roles=[], + ) + + updated_kv = { + "team_id": "test-team", + "team_member_budget": None, + "team_member_budget_duration": None, + "team_member_rpm_limit": None, + "team_member_tpm_limit": None, + } + + all_fields = { + "team_member_budget", + "team_member_budget_duration", + "team_member_rpm_limit", + "team_member_tpm_limit", + } + + with patch( + "litellm.proxy.management_endpoints.budget_management_endpoints.update_budget", + new_callable=AsyncMock, + ) as mock_update_budget: + result = await TeamMemberBudgetHandler.clear_team_member_budget_fields( + team_table=team_table, + user_api_key_dict=mock_user_api_key_dict, + updated_kv=updated_kv, + explicitly_set_fields=all_fields, + ) + + mock_update_budget.assert_awaited_once() + budget_request = mock_update_budget.call_args.kwargs["budget_obj"] + assert budget_request.budget_id == "budget-all" + assert budget_request.max_budget is None + assert budget_request.budget_duration is None + assert budget_request.budget_reset_at is None + assert budget_request.rpm_limit is None + assert budget_request.tpm_limit is None + for field in all_fields: + assert field not in result + + +@pytest.mark.asyncio +async def test_team_member_budget_duration_not_sent_does_not_update(): + """ + When team_member_budget_duration is NOT sent in the request, no budget + update should occur and the field should not appear in updated_kv. + """ + from litellm.proxy.management_endpoints.team_endpoints import ( + TeamMemberBudgetHandler, + ) + + updated_kv = {"team_id": "test-team", "max_budget": 200} + + _team_member_fields_in_request = { + field + for field in [ + "team_member_budget", + "team_member_rpm_limit", + "team_member_tpm_limit", + "team_member_budget_duration", + ] + if field in updated_kv + } + + assert len(_team_member_fields_in_request) == 0 + + TeamMemberBudgetHandler._clean_team_member_fields(updated_kv) + + assert "team_member_budget_duration" not in updated_kv + assert "team_member_budget" not in updated_kv + + +@pytest.mark.asyncio +async def test_clear_team_member_budget_fields_no_budget_row_skips_update(): + from litellm.proxy.management_endpoints.team_endpoints import ( + TeamMemberBudgetHandler, + ) + + mock_user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="admin-user", + ) + + team_table = LiteLLM_TeamTable( + team_id="test-team", + metadata=None, + members_with_roles=[], + ) + + updated_kv = { + "team_id": "test-team", + "team_member_budget": None, + "team_member_rpm_limit": None, + } + + with patch( + "litellm.proxy.management_endpoints.budget_management_endpoints.update_budget", + new_callable=AsyncMock, + ) as mock_update_budget: + result = await TeamMemberBudgetHandler.clear_team_member_budget_fields( + team_table=team_table, + user_api_key_dict=mock_user_api_key_dict, + updated_kv=updated_kv, + explicitly_set_fields={"team_member_budget", "team_member_rpm_limit"}, + ) + + mock_update_budget.assert_not_awaited() + assert "team_member_budget" not in result + assert "team_member_rpm_limit" not in result diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py index 3c6af3e528a..401ea2ef589 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py @@ -683,36 +683,43 @@ class TestOpenAIPassthroughLoggingHandler: "litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload" ) @patch( - "litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler.OpenAIPassthroughLoggingHandler.get_provider_config" + "litellm.llms.openai.responses.transformation.OpenAIResponsesAPIConfig.transform_response_api_response" ) def test_responses_api_cost_tracking( - self, mock_get_provider_config, mock_get_standard_logging, mock_completion_cost + self, + mock_transform_responses, + mock_get_standard_logging, + mock_completion_cost, ): - """Test cost tracking for responses API route""" + """Test cost tracking for responses API route. + + Mocks the Responses-API transformer (the dedicated one this branch + of the handler dispatches into post-fix) so we can assert the + downstream cost-calculation contract without depending on the + real transformer's full behavior. + """ # Arrange mock_completion_cost.return_value = 0.000050 mock_get_standard_logging.return_value = {"test": "logging_payload"} - # Mock the provider config's transform_response to return a valid ModelResponse - from litellm import ModelResponse + # Mock the Responses transformer's return — a ResponsesAPIResponse + # carrying the usage fields downstream cost-calc expects. + from litellm.types.llms.openai import ResponsesAPIResponse - mock_model_response = ModelResponse( + mock_responses_api_response = ResponsesAPIResponse.model_construct( id="resp_abc123", + object="response", + created_at=1677652288, model="gpt-4o-2024-08-06", - choices=[ - { - "message": { - "role": "assistant", - "content": "Hello! How can I help you today?", - } - } - ], - usage={"prompt_tokens": 20, "completion_tokens": 15, "total_tokens": 35}, + status="completed", + output=[], + usage={ + "input_tokens": 20, + "output_tokens": 15, + "total_tokens": 35, + }, ) - - mock_provider_config = MagicMock() - mock_provider_config.transform_response.return_value = mock_model_response - mock_get_provider_config.return_value = mock_provider_config + mock_transform_responses.return_value = mock_responses_api_response # Mock responses API response mock_responses_response = { @@ -768,6 +775,109 @@ class TestOpenAIPassthroughLoggingHandler: assert mock_logging_obj.model_call_details["model"] == "gpt-4o" assert mock_logging_obj.model_call_details["custom_llm_provider"] == "openai" + @patch("litellm.completion_cost") + @patch( + "litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload" + ) + def test_responses_api_uses_responses_transformer_not_chat_completions( + self, mock_get_standard_logging, mock_completion_cost + ): + """Regression test for the Responses-API cost-tracking dispatch bug. + + BUG: the `elif is_responses:` branch in `openai_passthrough_handler` + was calling `OpenAIConfig.transform_response` (the chat-completions + transformer) on a Responses API payload. Chat-completions + transform_response expects `choices: [...]` in the raw response; + the Responses API uses `output: [...]` and `usage.input_tokens` / + `usage.output_tokens` (not `prompt_tokens` / `completion_tokens`). + The result was a KeyError 'choices' inside + `convert_to_model_response_object`, swallowed by the surrounding + try/except, and the SpendLogs row was written with zero tokens + and zero spend. + + FIX: use the dedicated `OpenAIResponsesAPIConfig.transform_response_api_response` + for the Responses branch. + + This test exercises the REAL transformer (no mocked + `get_provider_config`) so that running it against the un-fixed + handler raises and running it against the fixed handler succeeds. + """ + mock_completion_cost.return_value = 0.000050 + mock_get_standard_logging.return_value = {"test": "logging_payload"} + + # A real-shaped Azure / OpenAI Responses API payload — NO `choices`, + # uses `output` and `usage.input_tokens` / `usage.output_tokens`. + responses_api_body = { + "id": "resp_abc123", + "object": "response", + "created_at": 1677652288, + "model": "gpt-4o-2024-08-06", + "status": "completed", + "output": [ + { + "type": "message", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "Hello!", + } + ], + } + ], + "usage": { + "input_tokens": 20, + "output_tokens": 15, + "total_tokens": 35, + }, + } + + mock_httpx_response = self._create_mock_httpx_response(responses_api_body) + mock_logging_obj = self._create_mock_logging_obj() + passthrough_payload = self._create_passthrough_logging_payload() + + kwargs = { + "passthrough_logging_payload": passthrough_payload, + "model": "gpt-4o", + "custom_llm_provider": "openai", + } + + result = OpenAIPassthroughLoggingHandler.openai_passthrough_handler( + httpx_response=mock_httpx_response, + response_body=responses_api_body, + logging_obj=mock_logging_obj, + url_route="https://api.openai.com/v1/responses", + result="", + start_time=self.start_time, + end_time=self.end_time, + cache_hit=False, + request_body={"model": "gpt-4o", "input": "Tell me about AI"}, + **kwargs, + ) + + # Pre-fix this assertion fails — the handler swallows the + # KeyError raised by the chat-completions transformer and falls + # back to the passthrough_chat_handler which yields a different + # response_cost value. Post-fix, the Responses transformer + # succeeds and we get the mocked 0.000050. + assert result is not None + assert result["kwargs"]["response_cost"] == 0.000050 + assert result["kwargs"]["model"] == "gpt-4o" + + # `completion_cost` must be called with the responses call type + # and a `ResponsesAPIResponse` (not a `ModelResponse`). + mock_completion_cost.assert_called_once() + call_kwargs = mock_completion_cost.call_args[1] + assert call_kwargs["call_type"] == "responses" + + from litellm.types.llms.openai import ResponsesAPIResponse + + assert isinstance(call_kwargs["completion_response"], ResponsesAPIResponse), ( + "completion_response must be a ResponsesAPIResponse; passing a " + "chat-completions ModelResponse means the Responses transformer " + "isn't being used and we're back in the bug." + ) + class TestOpenAIPassthroughIntegration: """Integration tests for OpenAI passthrough cost tracking""" @@ -872,6 +982,126 @@ class TestOpenAIPassthroughIntegration: ) assert self.handler.is_openai_route("") == False + def test_is_supported_openai_endpoint_includes_responses_api(self): + """Regression test for the outer dispatch gate. + + `_is_supported_openai_endpoint` is the gate that decides whether the + OpenAI handler runs for a given URL. Before this gate accepted the + Responses API, calls to `/v1/responses` would fail the gate and the + handler's `elif is_responses:` branch was unreachable in the live + success-handler pipeline — every Responses-API call landed in + `LiteLLM_SpendLogs` with zero tokens / zero spend even though the + handler had a Responses branch internally. + + This test exercises the dispatch decision directly so future + refactors of `_is_supported_openai_endpoint` can't silently + remove Responses from the OR-chain without a test failure. + """ + # Responses must be supported on api.openai.com and openai.azure.com. + assert ( + self.handler._is_supported_openai_endpoint( + "https://api.openai.com/v1/responses" + ) + is True + ) + assert ( + self.handler._is_supported_openai_endpoint( + "https://openai.azure.com/v1/responses" + ) + is True + ) + # The other supported endpoints stay supported (no regression). + assert ( + self.handler._is_supported_openai_endpoint( + "https://api.openai.com/v1/chat/completions" + ) + is True + ) + assert ( + self.handler._is_supported_openai_endpoint( + "https://api.openai.com/v1/images/generations" + ) + is True + ) + assert ( + self.handler._is_supported_openai_endpoint( + "https://api.openai.com/v1/images/edits" + ) + is True + ) + # Unsupported OpenAI endpoints (e.g. /v1/models) still return False. + assert ( + self.handler._is_supported_openai_endpoint( + "https://api.openai.com/v1/models" + ) + is False + ) + + @patch( + "litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler.OpenAIPassthroughLoggingHandler.openai_passthrough_handler" + ) + @pytest.mark.asyncio + async def test_success_handler_dispatches_responses_api_to_openai_handler( + self, mock_openai_handler + ): + """End-to-end dispatch test for the Responses API path. + + Pre-fix: `_is_supported_openai_endpoint` returned False for + `/v1/responses` URLs, so the OpenAI handler was never called. + This test would fail (mock never invoked) on the un-fixed + success_handler — passes only when the dispatch gate accepts + Responses URLs. + """ + mock_openai_handler.return_value = { + "result": {"id": "resp_abc123"}, + "kwargs": { + "response_cost": 0.0001, + "model": "gpt-4o", + "custom_llm_provider": "openai", + }, + } + + mock_httpx_response = MagicMock(spec=httpx.Response) + mock_httpx_response.text = ( + '{"id": "resp_abc123", "object": "response", ' + '"output": [], "usage": {"input_tokens": 5, "output_tokens": 3}}' + ) + + mock_logging_obj = AsyncMock() + mock_logging_obj.model_call_details = {} + mock_logging_obj.async_success_handler = AsyncMock() + + passthrough_payload = PassthroughStandardLoggingPayload( + url="https://api.openai.com/v1/responses", + request_body={"model": "gpt-4o", "input": "Hello"}, + request_method="POST", + ) + + await self.handler.pass_through_async_success_handler( + httpx_response=mock_httpx_response, + response_body={ + "id": "resp_abc123", + "object": "response", + "output": [], + "usage": {"input_tokens": 5, "output_tokens": 3}, + }, + logging_obj=mock_logging_obj, + url_route="https://api.openai.com/v1/responses", + result="", + start_time=datetime.now(), + end_time=datetime.now(), + cache_hit=False, + request_body={"model": "gpt-4o", "input": "Hello"}, + passthrough_logging_payload=passthrough_payload, + ) + + # The OpenAI handler MUST have been invoked. Pre-fix the dispatch + # gate filtered Responses URLs out and the mock was never called. + mock_openai_handler.assert_called_once() + # And we can verify it was dispatched with the Responses URL. + call_kwargs = mock_openai_handler.call_args.kwargs + assert call_kwargs["url_route"] == "https://api.openai.com/v1/responses" + @patch( "litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler.OpenAIPassthroughLoggingHandler.openai_passthrough_handler" ) diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index b75fc27e21d..4eab1a4bf61 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -24,11 +24,13 @@ from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( get_vertex_base_url, llm_passthrough_factory_proxy_route, milvus_proxy_route, + mistral_proxy_route, openai_proxy_route, vertex_discovery_proxy_route, vertex_proxy_route, vllm_proxy_route, ) +from litellm.proxy._types import UserAPIKeyAuth from litellm.types.passthrough_endpoints.vertex_ai import VertexPassThroughCredentials @@ -1092,9 +1094,9 @@ class TestVertexAIPassThroughHandler: assert result is not None assert result["result"] is not None - assert result["kwargs"].get("custom_llm_provider") == "gemini", ( - "Google AI Studio embedContent URLs must set custom_llm_provider=gemini, not vertex_ai" - ) + assert ( + result["kwargs"].get("custom_llm_provider") == "gemini" + ), "Google AI Studio embedContent URLs must set custom_llm_provider=gemini, not vertex_ai" assert result["kwargs"].get("model") == "gemini-embedding-2-preview" mock_completion_cost.assert_called_once() @@ -1261,6 +1263,78 @@ async def test_is_streaming_request_fn(): assert await is_streaming_request_fn(mock_request) is True +@pytest.mark.asyncio +async def test_mistral_passthrough_accepts_multipart_without_json_parsing(): + boundary = "----litellm-test-boundary" + body = ( + f"--{boundary}\r\n" + 'Content-Disposition: form-data; name="purpose"\r\n\r\n' + "ocr\r\n" + f"--{boundary}\r\n" + 'Content-Disposition: form-data; name="file"; filename="document.pdf"\r\n' + "Content-Type: application/pdf\r\n\r\n" + "%PDF-1.4 test\r\n" + f"--{boundary}--\r\n" + ).encode("utf-8") + + async def receive(): + return { + "type": "http.request", + "body": body, + "more_body": False, + } + + request = Request( + { + "type": "http", + "method": "POST", + "path": "/mistral/v1/files", + "headers": [ + ( + b"content-type", + f"multipart/form-data; boundary={boundary}".encode("utf-8"), + ) + ], + "query_string": b"", + }, + receive=receive, + ) + + captured_kwargs = {} + + async def fake_endpoint(request, fastapi_response, user_api_key_dict): + return {"ok": True} + + def fake_create_pass_through_route(**kwargs): + captured_kwargs.update(kwargs) + return fake_endpoint + + user_api_key_dict = UserAPIKeyAuth(token="test-key") + + with ( + patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router.get_credentials", + return_value="mistral-test-key", + ), + patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route", + side_effect=fake_create_pass_through_route, + ), + ): + response = await mistral_proxy_route( + endpoint="v1/files", + request=request, + fastapi_response=Response(), + user_api_key_dict=user_api_key_dict, + ) + + assert response == {"ok": True} + assert captured_kwargs["is_streaming_request"] is False + assert captured_kwargs["custom_headers"] == { + "Authorization": "Bearer mistral-test-key" + } + + class TestBedrockLLMProxyRoute: @pytest.mark.asyncio async def test_bedrock_llm_proxy_route_application_inference_profile(self): diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index aef91ed3c77..2632d8af4f1 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -1314,7 +1314,8 @@ async def test_ui_view_session_spend_logs_pagination(client, monkeypatch): assert session_id == "session-123" assert page_size == 1 assert skip == 1 # page=2, page_size=1 - return [mock_spend_logs[1]] + assert 'ORDER BY "startTime" DESC' in sql_query + return [mock_spend_logs[0]] class MockPrismaClient: def __init__(self): @@ -1337,7 +1338,7 @@ async def test_ui_view_session_spend_logs_pagination(client, monkeypatch): assert data["page_size"] == 1 assert data["total_pages"] == 2 assert len(data["data"]) == 1 - assert data["data"][0]["request_id"] == "req2" + assert data["data"][0]["request_id"] == "req1" @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 0f5a0cbe4b6..b45b31cc67c 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -2269,6 +2269,36 @@ class TestHandleLLMApiExceptionDictDetail: assert proxy_exc.message == "Content blocked by guardrail" assert proxy_exc.provider_specific_fields is None + async def test_not_found_error_preserves_404(self): + """NotFoundError with status_code=404 should map to ProxyException code=404.""" + from litellm.exceptions import NotFoundError + + exc = NotFoundError( + message="Model gemini-3.1-flash-lite-preview not found", + model="gemini-3.1-flash-lite-preview", + llm_provider="gemini", + ) + proxy_exc = await self._invoke(exc) + assert proxy_exc.code == "404" + assert "NotFoundError" in proxy_exc.message + + async def test_exception_with_status_code_propagates(self): + """Exception with a statically-set status_code should propagate it.""" + from litellm.llms.vertex_ai.common_utils import VertexAIError + + exc = VertexAIError( + status_code=429, + message="Rate limit exceeded", + ) + proxy_exc = await self._invoke(exc) + assert proxy_exc.code == "429" + + async def test_exception_without_status_code_defaults_to_500(self): + """Exception with no status_code attribute defaults to 500.""" + exc = ValueError("Something broke") + proxy_exc = await self._invoke(exc) + assert proxy_exc.code == "500" + class TestAsyncStreamingDataGeneratorFastPath: """Fast/slow path branching in async_streaming_data_generator.""" diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 8aa839cdfcb..b2e36fd64cf 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -2319,6 +2319,36 @@ async def test_custom_ui_sso_sign_in_handler_config_loading(): os.unlink(config_file_path) +@pytest.mark.asyncio +async def test_load_config_max_budget_env_var_coerced_to_float(tmp_path, monkeypatch): + """ + max_budget configured as os.environ/MAX_BUDGET resolves to a string; + load_config must coerce it to float so the startup check + `litellm.max_budget > 0` doesn't raise TypeError. + """ + from litellm.proxy.proxy_server import ProxyConfig + + monkeypatch.setenv("MAX_BUDGET", "10") + test_config = { + "model_list": [], + "litellm_settings": {"max_budget": "os.environ/MAX_BUDGET"}, + } + config_file = tmp_path / "config.yaml" + config_file.write_text(yaml.dump(test_config)) + + original_max_budget = litellm.max_budget + try: + proxy_config = ProxyConfig() + await proxy_config.load_config( + router=MagicMock(), config_file_path=str(config_file) + ) + assert isinstance(litellm.max_budget, float) + assert litellm.max_budget == 10.0 + assert litellm.max_budget > 0 + finally: + litellm.max_budget = original_max_budget + + @pytest.mark.asyncio async def test_load_environment_variables_direct_and_os_environ(): """ diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py index 07d894d0400..510dcf77afd 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py @@ -1471,3 +1471,191 @@ async def test_affinity_does_not_raise_when_boundary_peer_available(): assert result == [peer] assert request_kwargs.get("_encrypted_content_affinity_pinned") is True + + +@pytest.mark.asyncio +async def test_model_group_affinity_config_enables_encrypted_content_affinity(): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + model_group = "openai.gpt-5.1-codex" + target_deployment = { + "model_name": model_group, + "litellm_params": {"model": "openai/gpt-5.1-codex"}, + "model_info": {"id": "deployment-b"}, + } + healthy_deployments = [ + { + "model_name": model_group, + "litellm_params": {"model": "openai/gpt-5.1-codex"}, + "model_info": {"id": "deployment-a"}, + }, + target_deployment, + ] + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( + "deployment-b", "rs_test" + ) + request_kwargs = { + "input": [{"type": "reasoning", "id": encoded_id}], + "litellm_metadata": {}, + } + check = EncryptedContentAffinityCheck( + enable_global_affinity=False, + model_group_affinity_config={ + model_group: ["encrypted_content_affinity"], + }, + ) + + filtered = await check.async_filter_deployments( + model=model_group, + healthy_deployments=healthy_deployments, + messages=None, + request_kwargs=request_kwargs, + ) + + assert filtered == [target_deployment] + assert request_kwargs["litellm_metadata"]["encrypted_content_affinity_enabled"] + assert request_kwargs.get("_encrypted_content_affinity_pinned") is True + + +@pytest.mark.asyncio +async def test_model_group_affinity_config_does_not_disable_global_encrypted_content_affinity(): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + model_group = "openai.gpt-5.1-codex" + target_deployment = { + "model_name": model_group, + "litellm_params": {"model": "openai/gpt-5.1-codex"}, + "model_info": {"id": "deployment-b"}, + } + healthy_deployments = [ + { + "model_name": model_group, + "litellm_params": {"model": "openai/gpt-5.1-codex"}, + "model_info": {"id": "deployment-a"}, + }, + target_deployment, + ] + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( + "deployment-b", "rs_test" + ) + request_kwargs = { + "input": [{"type": "reasoning", "id": encoded_id}], + "litellm_metadata": {}, + } + check = EncryptedContentAffinityCheck( + enable_global_affinity=True, + model_group_affinity_config={ + model_group: ["deployment_affinity"], + }, + ) + + filtered = await check.async_filter_deployments( + model=model_group, + healthy_deployments=healthy_deployments, + messages=None, + request_kwargs=request_kwargs, + ) + + assert filtered == [target_deployment] + assert request_kwargs["litellm_metadata"]["encrypted_content_affinity_enabled"] + assert request_kwargs.get("_encrypted_content_affinity_pinned") is True + + +@pytest.mark.asyncio +async def test_model_group_encrypted_content_affinity_overrides_global_deployment_affinity(): + from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( + DeploymentAffinityCheck, + ) + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + model_group = "openai.gpt-5.1-codex" + user_api_key_hash = "test-user-key" + deployment_a = { + "model_name": model_group, + "litellm_params": { + "model": "openai/gpt-5.1-codex", + "api_key": "mock-api-key-a", + }, + "model_info": {"id": "deployment-a"}, + } + deployment_b = { + "model_name": model_group, + "litellm_params": { + "model": "openai/gpt-5.1-codex", + "api_key": "mock-api-key-b", + }, + "model_info": {"id": "deployment-b"}, + } + router = litellm.Router( + model_list=[deployment_a, deployment_b], + optional_pre_call_checks=["deployment_affinity"], + model_group_affinity_config={ + model_group: ["encrypted_content_affinity"], + }, + num_retries=0, + ) + + try: + callbacks = router.optional_callbacks or [] + deployment_callback = next( + cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck) + ) + encrypted_content_callback = next( + cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck) + ) + assert callbacks.index(encrypted_content_callback) < callbacks.index( + deployment_callback + ) + assert encrypted_content_callback.enable_global_affinity is False + + cache_key = DeploymentAffinityCheck.get_affinity_cache_key( + model_group=model_group, + user_key=user_api_key_hash, + ) + await deployment_callback.cache.async_set_cache( + key=cache_key, + value={"model_id": "deployment-a"}, + ttl=60, + ) + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( + "deployment-b", "rs_test" + ) + request_kwargs = { + "input": [ + { + "type": "reasoning", + "id": encoded_id, + "encrypted_content": "gAAAAABpnW_yEYmSNEyOG...", + } + ], + "metadata": {"user_api_key_hash": user_api_key_hash}, + "litellm_metadata": {}, + } + + after_deployment_affinity = await deployment_callback.async_filter_deployments( + model=model_group, + healthy_deployments=[deployment_a, deployment_b], + messages=None, + request_kwargs=request_kwargs, + ) + assert after_deployment_affinity == [deployment_a, deployment_b] + + after_encrypted_content_affinity = ( + await encrypted_content_callback.async_filter_deployments( + model=model_group, + healthy_deployments=after_deployment_affinity, + messages=None, + request_kwargs=request_kwargs, + ) + ) + + assert after_encrypted_content_affinity == [deployment_b] + assert request_kwargs.get("_encrypted_content_affinity_pinned") is True + finally: + router.discard() diff --git a/tests/test_litellm/test_anthropic_beta_headers_filtering.py b/tests/test_litellm/test_anthropic_beta_headers_filtering.py index 84867a6e905..9cd27e88c33 100644 --- a/tests/test_litellm/test_anthropic_beta_headers_filtering.py +++ b/tests/test_litellm/test_anthropic_beta_headers_filtering.py @@ -402,6 +402,16 @@ class TestAnthropicBetaHeadersFiltering: test_case["expected"] in filtered ), f"Header '{test_case['input']}' should be mapped to '{test_case['expected']}' for {test_case['provider']}, but got: {filtered}" + def test_filter_and_transform_beta_headers_vertex_ai_keeps_compact(self): + """Vertex AI supports compact context edits, so the compact beta header + must be forwarded instead of stripped (it was previously mapped to null, + which broke compact_20260112 context edits over /v1/messages).""" + filtered = filter_and_transform_beta_headers( + beta_headers=["compact-2026-01-12"], provider="vertex_ai" + ) + + assert filtered == ["compact-2026-01-12"] + def test_null_value_headers_filtered(self): """Test that headers with null values are always filtered out.""" for provider in [ diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index cd235d8de67..e681247959f 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -80,6 +80,256 @@ def test_router_with_model_info_and_model_group(): ) +def test_router_model_group_encrypted_content_affinity_callback_registration(): + from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( + DeploymentAffinityCheck, + ) + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + model_group = "openai.gpt-5.1-codex" + model_group_affinity_config = { + model_group: ["encrypted_content_affinity"], + } + original_callbacks = list(litellm.callbacks) + litellm.callbacks = [] + router = None + + try: + router = litellm.Router( + model_list=[ + { + "model_name": model_group, + "litellm_params": { + "model": "openai/gpt-5.1-codex", + "api_key": "mock-api-key", + }, + } + ], + model_group_affinity_config=model_group_affinity_config, + num_retries=0, + ) + callbacks = router.optional_callbacks or [] + encrypted_content_callbacks = [ + cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck) + ] + deployment_callback = next( + cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck) + ) + assert len(encrypted_content_callbacks) == 1 + assert encrypted_content_callbacks[0].enable_global_affinity is False + assert ( + encrypted_content_callbacks[0].model_group_affinity_config + == model_group_affinity_config + ) + assert callbacks.index(encrypted_content_callbacks[0]) < callbacks.index( + deployment_callback + ) + assert litellm.callbacks.index(encrypted_content_callbacks[0]) < ( + litellm.callbacks.index(deployment_callback) + ) + + router._add_encrypted_content_affinity_check(enable_global_affinity=True) + + callbacks = router.optional_callbacks or [] + encrypted_content_callbacks = [ + cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck) + ] + assert len(encrypted_content_callbacks) == 1 + assert encrypted_content_callbacks[0].enable_global_affinity is True + assert encrypted_content_callbacks[0].router is router + finally: + if router is not None: + router.discard() + litellm.callbacks = original_callbacks + + +@pytest.mark.asyncio +async def test_encrypted_content_affinity_model_group_config_is_additive(): + from litellm.responses.utils import ResponsesAPIRequestUtils + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + model_group = "openai.gpt-5.1-codex" + target_deployment = { + "model_name": model_group, + "litellm_params": {"model": "openai/gpt-5.1-codex"}, + "model_info": {"id": "deployment-b"}, + } + healthy_deployments = [ + { + "model_name": model_group, + "litellm_params": {"model": "openai/gpt-5.1-codex"}, + "model_info": {"id": "deployment-a"}, + }, + target_deployment, + ] + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( + "deployment-b", "rs_test" + ) + + assert EncryptedContentAffinityCheck.has_model_group_affinity_enabled( + {model_group: ["encrypted_content_affinity"]} + ) + assert not EncryptedContentAffinityCheck.has_model_group_affinity_enabled(None) + + per_group_check = EncryptedContentAffinityCheck( + enable_global_affinity=False, + model_group_affinity_config={ + model_group: ["encrypted_content_affinity"], + }, + ) + request_kwargs = { + "input": [{"type": "reasoning", "id": encoded_id}], + "litellm_metadata": {}, + } + filtered = await per_group_check.async_filter_deployments( + model=model_group, + healthy_deployments=healthy_deployments, + messages=None, + request_kwargs=request_kwargs, + ) + + assert filtered == [target_deployment] + assert request_kwargs["litellm_metadata"]["encrypted_content_affinity_enabled"] + + disabled_check = EncryptedContentAffinityCheck( + enable_global_affinity=False, + model_group_affinity_config={ + "other-model-group": ["encrypted_content_affinity"], + }, + ) + disabled_request_kwargs = { + "input": [{"type": "reasoning", "id": encoded_id}], + "litellm_metadata": {}, + } + unfiltered = await disabled_check.async_filter_deployments( + model=model_group, + healthy_deployments=healthy_deployments, + messages=None, + request_kwargs=disabled_request_kwargs, + ) + + assert unfiltered == healthy_deployments + assert "encrypted_content_affinity_enabled" not in disabled_request_kwargs[ + "litellm_metadata" + ] + + global_check = EncryptedContentAffinityCheck( + enable_global_affinity=True, + model_group_affinity_config={ + model_group: ["deployment_affinity"], + }, + ) + global_request_kwargs = { + "input": [{"type": "reasoning", "id": encoded_id}], + "litellm_metadata": {}, + } + globally_filtered = await global_check.async_filter_deployments( + model=model_group, + healthy_deployments=healthy_deployments, + messages=None, + request_kwargs=global_request_kwargs, + ) + + assert globally_filtered == [target_deployment] + assert global_request_kwargs["litellm_metadata"][ + "encrypted_content_affinity_enabled" + ] + + +@pytest.mark.asyncio +async def test_encrypted_content_affinity_takes_priority_over_user_key_affinity(): + from litellm.responses.utils import ResponsesAPIRequestUtils + from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( + DeploymentAffinityCheck, + ) + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + model_group = "openai.gpt-5.1-codex" + user_api_key_hash = "test-user-key" + deployment_a = { + "model_name": model_group, + "litellm_params": { + "model": "openai/gpt-5.1-codex", + "api_key": "mock-api-key-a", + }, + "model_info": {"id": "deployment-a"}, + } + deployment_b = { + "model_name": model_group, + "litellm_params": { + "model": "openai/gpt-5.1-codex", + "api_key": "mock-api-key-b", + }, + "model_info": {"id": "deployment-b"}, + } + original_callbacks = list(litellm.callbacks) + litellm.callbacks = [] + router = None + + try: + router = litellm.Router( + model_list=[deployment_a, deployment_b], + model_group_affinity_config={ + model_group: [ + "deployment_affinity", + "encrypted_content_affinity", + ], + }, + num_retries=0, + ) + callbacks = router.optional_callbacks or [] + deployment_callback = next( + cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck) + ) + encrypted_content_callback = next( + cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck) + ) + assert callbacks.index(encrypted_content_callback) < callbacks.index( + deployment_callback + ) + assert litellm.callbacks.index(encrypted_content_callback) < ( + litellm.callbacks.index(deployment_callback) + ) + + cache_key = DeploymentAffinityCheck.get_affinity_cache_key( + model_group=model_group, + user_key=user_api_key_hash, + ) + await deployment_callback.cache.async_set_cache( + key=cache_key, + value={"model_id": "deployment-a"}, + ttl=60, + ) + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( + "deployment-b", "rs_test" + ) + request_kwargs = { + "input": [{"type": "reasoning", "id": encoded_id}], + "litellm_metadata": {"user_api_key_hash": user_api_key_hash}, + } + + filtered = await router.async_callback_filter_deployments( + model=model_group, + healthy_deployments=[deployment_a, deployment_b], + messages=None, + parent_otel_span=None, + request_kwargs=request_kwargs, + ) + + assert filtered == [deployment_b] + assert request_kwargs.get("_encrypted_content_affinity_pinned") is True + finally: + if router is not None: + router.discard() + litellm.callbacks = original_callbacks + + @pytest.mark.asyncio async def test_arouter_with_tags_and_fallbacks(): """ @@ -4311,6 +4561,48 @@ def test_get_available_deployment_for_pass_through_raises_when_dict_blocked(): ) +def test_initialize_deployment_for_pass_through_keeps_bedrock_iam_deployment(): + """ + Bedrock deployments using IAM/OIDC auth have no api_key; pass-through + init must not raise and drop them from routing (#27728). + """ + import litellm + + router = litellm.Router( + model_list=[ + { + "model_name": "bedrock-claude", + "litellm_params": { + "model": "bedrock/anthropic.claude-3-5-sonnet-20241022-v2:0", + "aws_role_name": "arn:aws:iam::123456789012:role/my-role", + "aws_session_name": "my-session", + "use_in_pass_through": True, + }, + "model_info": {"id": "bedrock-iam-pt"}, + } + ] + ) + assert [m["model_info"]["id"] for m in router.get_model_list()] == [ + "bedrock-iam-pt" + ] + + +def test_initialize_deployment_for_pass_through_sets_credentials_with_api_key(): + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + passthrough_endpoint_router, + ) + + passthrough_endpoint_router.credentials.clear() + router = _router_with_two_pass_through_deployments([False, False]) + assert len(router.get_model_list()) == 2 + assert ( + passthrough_endpoint_router.get_credentials( + custom_llm_provider="openai", region_name=None + ) + == "sk-fake-for-tests" + ) + + def test_get_deployment_credentials_returns_none_for_blocked_deployment(): router = _router_with_two_deployments([True, False]) assert router.get_deployment_credentials(model_id="dep-0") is None diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 4c4d9e1133b..62cb8154b6d 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -700,6 +700,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "cache_read_input_token_cost": {"type": "number"}, "cache_read_input_token_cost_above_200k_tokens": {"type": "number"}, "cache_read_input_token_cost_above_272k_tokens": {"type": "number"}, + "cache_read_input_token_cost_above_512k_tokens": {"type": "number"}, "cache_read_input_token_cost_batches": {"type": "number"}, "cache_creation_input_token_cost_above_1hr_above_200k_tokens": { "type": "number" @@ -721,6 +722,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "input_cost_per_token_above_200k_tokens": {"type": "number"}, "input_cost_per_token_above_256k_tokens": {"type": "number"}, "input_cost_per_token_above_272k_tokens": {"type": "number"}, + "input_cost_per_token_above_512k_tokens": {"type": "number"}, "cache_read_input_token_cost_flex": {"type": "number"}, "cache_read_input_token_cost_priority": {"type": "number"}, "cache_read_input_token_cost_above_200k_tokens_priority": { @@ -811,6 +813,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "output_cost_per_token_above_200k_tokens": {"type": "number"}, "output_cost_per_token_above_256k_tokens": {"type": "number"}, "output_cost_per_token_above_272k_tokens": {"type": "number"}, + "output_cost_per_token_above_512k_tokens": {"type": "number"}, "output_cost_per_image_above_1024_and_1024_pixels": {"type": "number"}, "output_cost_per_image_above_1024_and_1024_pixels_and_premium_image": { "type": "number" @@ -932,6 +935,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "supports_native_streaming": {"type": "boolean"}, "supports_image_size": {"type": "boolean"}, "supports_native_structured_output": {"type": "boolean"}, + "use_openai_responses_path": {"type": "boolean"}, "tiered_pricing": { "type": "array", "items": { diff --git a/ui/litellm-dashboard/src/components/networking.test.ts b/ui/litellm-dashboard/src/components/networking.test.ts index b33c2b741a1..43f85cc8674 100644 --- a/ui/litellm-dashboard/src/components/networking.test.ts +++ b/ui/litellm-dashboard/src/components/networking.test.ts @@ -467,3 +467,50 @@ describe("teamInfoCall", () => { expect(parsed.searchParams.has("team_id")).toBe(false); }); }); + +describe("sessionSpendLogsCall", () => { + const originalFetch = global.fetch; + + beforeEach(() => { + vi.clearAllMocks(); + }); + + afterEach(() => { + global.fetch = originalFetch; + }); + + it("should request the first page with defaults so the caller can page through the session", async () => { + const mockFetch = vi.fn().mockResolvedValue({ + ok: true, + json: vi.fn().mockResolvedValue({ data: [], total: 0, page: 1, page_size: 100, total_pages: 1 }), + } as any); + global.fetch = mockFetch as any; + + await Networking.sessionSpendLogsCall("token", "session-123"); + + expect(mockFetch).toHaveBeenCalledOnce(); + const [url] = mockFetch.mock.calls[0]; + const urlStr = typeof url === "string" ? url : (url as Request).url; + const parsed = typeof url === "string" ? new URL(url, "http://example.com") : new URL((url as Request).url); + + expect(urlStr).toContain("/spend/logs/session/ui"); + expect(parsed.searchParams.get("session_id")).toBe("session-123"); + expect(parsed.searchParams.get("page")).toBe("1"); + expect(parsed.searchParams.get("page_size")).toBe("100"); + }); + + it("should pass explicit page and page_size query params for later pages", async () => { + const mockFetch = vi.fn().mockResolvedValue({ + ok: true, + json: vi.fn().mockResolvedValue({ data: [], total: 250, page: 3, page_size: 100, total_pages: 3 }), + } as any); + global.fetch = mockFetch as any; + + await Networking.sessionSpendLogsCall("token", "session-123", 3, 100); + + const [url] = mockFetch.mock.calls[0]; + const parsed = typeof url === "string" ? new URL(url, "http://example.com") : new URL((url as Request).url); + expect(parsed.searchParams.get("page")).toBe("3"); + expect(parsed.searchParams.get("page_size")).toBe("100"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 9039a387050..b41ff073cb7 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -5616,13 +5616,27 @@ export const teamPermissionsUpdateCall = async (accessToken: string, teamId: str }; /** - * Get all spend logs for a particular session + * Get a page of spend logs for a particular session. + * + * The backend paginates this endpoint (page / page_size, returning + * { data, total, page, page_size, total_pages }). Callers that need the whole + * session should page through total_pages and accumulate the results. */ -export const sessionSpendLogsCall = async (accessToken: string, session_id: string) => { +export const sessionSpendLogsCall = async ( + accessToken: string, + session_id: string, + page: number = 1, + page_size: number = 100, +) => { try { + const params = new URLSearchParams({ + session_id, + page: String(page), + page_size: String(page_size), + }); let url = proxyBaseUrl - ? `${proxyBaseUrl}/spend/logs/session/ui?session_id=${encodeURIComponent(session_id)}` - : `/spend/logs/session/ui?session_id=${encodeURIComponent(session_id)}`; + ? `${proxyBaseUrl}/spend/logs/session/ui?${params.toString()}` + : `/spend/logs/session/ui?${params.toString()}`; const response = await fetch(url, { method: "GET", diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.tsx index 029a9814b81..f1aff8cbce3 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.tsx @@ -28,6 +28,14 @@ export interface LogDetailsDrawerProps { const SIDEBAR_WIDTH_PX = 224; +// Session logs are fetched page-by-page from the paginated backend and +// accumulated so the drawer can show the whole session. page_size is the +// backend maximum (le=100); the page cap bounds the fetch and the +// (un-virtualized) sidebar list for pathological sessions, keeping the most +// recent logs since the endpoint returns newest-first. +const SESSION_PAGE_SIZE = 100; +const MAX_SESSION_PAGES = 50; + /* ------------------------------------------------------------------ */ /* TraceEventRow — compact event row used in both session & non- */ /* session sidebar lists. Extracted to avoid JSX duplication. */ @@ -112,13 +120,39 @@ export function LogDetailsDrawer({ const [isSidebarCollapsed, setIsSidebarCollapsed] = useState(false); const [copiedLeftPanelId, setCopiedLeftPanelId] = useState(false); - const { data: sessionLogs = [] } = useQuery({ + const { data: sessionData } = useQuery({ queryKey: ["sessionLogs", sessionId], queryFn: async () => { - if (!sessionId || !accessToken) return []; - const response = await sessionSpendLogsCall(accessToken, sessionId); - const allSessionLogs: LogEntry[] = response.data || response || []; - return allSessionLogs + if (!sessionId || !accessToken) return { logs: [] as LogEntry[], total: 0 }; + + // Fetch the first page, then page through the rest so sessions with more + // than one page of logs are shown in full (capped for safety). + const firstPage = await sessionSpendLogsCall(accessToken, sessionId, 1, SESSION_PAGE_SIZE); + let rows: LogEntry[] = firstPage.data || firstPage || []; + const pagesToFetch = Math.min(firstPage.total_pages ?? 1, MAX_SESSION_PAGES); + + if (pagesToFetch > 1) { + const BATCH = 5; + const remaining: Awaited>[] = []; + for (let start = 2; start <= pagesToFetch; start += BATCH) { + const end = Math.min(start + BATCH - 1, pagesToFetch); + const batch = await Promise.all( + Array.from({ length: end - start + 1 }, (_, i) => + sessionSpendLogsCall(accessToken, sessionId, start + i, SESSION_PAGE_SIZE), + ), + ); + remaining.push(...batch); + } + for (const page of remaining) { + rows = rows.concat(page.data || []); + } + } + + // Fall back to the accumulated row count (not just the first page) when the + // backend omits total, so the truncation note reflects what was fetched. + const total: number = firstPage.total ?? rows.length; + + const logs = rows .map((row) => ({ ...row, request_duration_ms: row.request_duration_ms ?? Date.parse(row.endTime) - Date.parse(row.startTime), @@ -127,24 +161,49 @@ export function LogDetailsDrawer({ const aIsMcp = MCP_CALL_TYPES.includes(a.call_type) ? 1 : 0; const bIsMcp = MCP_CALL_TYPES.includes(b.call_type) ? 1 : 0; if (aIsMcp !== bIsMcp) return aIsMcp - bIsMcp; - return new Date(a.startTime).getTime() - new Date(b.startTime).getTime(); + // Newest first, matching the all-sessions logs overview. MCP calls + // stay grouped last (above), newest-first within that group too. + return new Date(b.startTime).getTime() - new Date(a.startTime).getTime(); }); + + return { logs, total }; }, enabled: Boolean(open && isSessionMode && sessionId && accessToken), }); + const sessionLogs: LogEntry[] = sessionData?.logs ?? []; + // total reported by the backend; when the page cap truncates the fetch this + // exceeds sessionLogs.length, which drives the "showing most recent" note. + const sessionTotalCount = sessionData?.total ?? sessionLogs.length; + const sessionTruncated = sessionTotalCount > sessionLogs.length; + + // Default selection for a freshly opened session: the most recent log (latest + // startTime). The list is sorted newest-first, but MCP calls are grouped last, + // so the latest log by time is not necessarily sessionLogs[0]; compute it + // explicitly. A clicked/remembered log still wins over this default. + const mostRecentLog = useMemo( + () => + sessionLogs.reduce( + (latest, row) => + !latest || new Date(row.startTime).getTime() > new Date(latest.startTime).getTime() ? row : latest, + null, + ), + [sessionLogs], + ); + const currentLog = useMemo(() => { if (!isSessionMode) return logEntry; if (!sessionLogs.length) return null; + const fallbackLog = mostRecentLog ?? sessionLogs[0]; if (selectedSessionRequestId) { - return sessionLogs.find((row) => row.request_id === selectedSessionRequestId) || sessionLogs[0]; + return sessionLogs.find((row) => row.request_id === selectedSessionRequestId) || fallbackLog; } if (logEntry?.request_id) { const clickedLog = sessionLogs.find((row) => row.request_id === logEntry.request_id); - return clickedLog || sessionLogs[0]; + return clickedLog || fallbackLog; } - return sessionLogs[0]; - }, [isSessionMode, logEntry, selectedSessionRequestId, sessionLogs]); + return fallbackLog; + }, [isSessionMode, logEntry, selectedSessionRequestId, sessionLogs, mostRecentLog]); useEffect(() => { if (!isSessionMode || !sessionLogs.length) return; @@ -152,10 +211,10 @@ export function LogDetailsDrawer({ const fallbackRequestId = logEntry?.request_id && sessionLogs.some((row) => row.request_id === logEntry.request_id) ? logEntry.request_id - : sessionLogs[0].request_id; + : (mostRecentLog ?? sessionLogs[0]).request_id; setSelectedSessionRequestId(fallbackRequestId); } - }, [isSessionMode, logEntry, selectedSessionRequestId, sessionLogs]); + }, [isSessionMode, logEntry, selectedSessionRequestId, sessionLogs, mostRecentLog]); // Reset transient UI state when the drawer opens or closes. useEffect(() => { @@ -327,6 +386,11 @@ export function LogDetailsDrawer({ )}
+ {isSessionMode && sessionTruncated && ( +
+ Showing most recent {logsForList.length} of {sessionTotalCount} +
+ )}
diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 8379d0536a6..203a56f615b 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -25081,6 +25081,12 @@ export interface components { * @default false */ use_litellm_proxy: boolean | null; + /** + * Use Xai Oauth + * @description Use stored xAI OAuth credentials when no xAI API key is configured. + * @default false + */ + use_xai_oauth: boolean | null; /** Vector Store Id */ vector_store_id?: string | null; /** Vertex Credentials */ @@ -32679,6 +32685,12 @@ export interface components { * @default false */ use_litellm_proxy: boolean | null; + /** + * Use Xai Oauth + * @description Use stored xAI OAuth credentials when no xAI API key is configured. + * @default false + */ + use_xai_oauth: boolean | null; /** Vector Store Id */ vector_store_id?: string | null; /** Vertex Credentials */ From f9293d40c4a9a2d3ff2b7fffa618bd9d183c6eef Mon Sep 17 00:00:00 2001 From: michelligabriele Date: Wed, 10 Jun 2026 20:16:58 +0200 Subject: [PATCH 046/185] fix(proxy): self-heal startup/reload prisma reads on engine disconnect (#28803) --- litellm/proxy/db/tool_registry_writer.py | 13 ++- .../cache_settings_endpoints.py | 9 +- litellm/proxy/proxy_server.py | 23 +++- .../search_endpoints/search_tool_registry.py | 11 +- .../proxy/db/test_tool_registry_writer.py | 74 ++++++++++++ .../test_search_tool_management.py | 39 +++++++ .../test_cache_settings_endpoints.py | 35 ++++++ tests/test_litellm/proxy/test_proxy_server.py | 105 ++++++++++++++++++ 8 files changed, 295 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/db/tool_registry_writer.py b/litellm/proxy/db/tool_registry_writer.py index bbcc7396d67..08bc8944b92 100644 --- a/litellm/proxy/db/tool_registry_writer.py +++ b/litellm/proxy/db/tool_registry_writer.py @@ -11,6 +11,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ToolDiscoveryQueueItem +from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry from litellm.repositories.object_permission_repository import ObjectPermissionRepository from litellm.repositories.table_repositories import ToolRepository from litellm.types.tool_management import ( @@ -309,7 +310,11 @@ class ToolPolicyRegistry: async def sync_tool_policy_from_db(self, prisma_client: "PrismaClient") -> None: """Load all tool policies and object-permission blocked_tools from DB.""" try: - tools = await ToolRepository(prisma_client).table.find_many() + tools = await call_with_db_reconnect_retry( + prisma_client, + lambda: ToolRepository(prisma_client).table.find_many(), + reason="sync_tool_policy_from_db_tools_lookup_failure", + ) self._tool_input_policies = { row.tool_name: getattr(row, "input_policy", "untrusted") or "untrusted" for row in tools @@ -319,7 +324,11 @@ class ToolPolicyRegistry: for row in tools } - perms = await ObjectPermissionRepository(prisma_client).table.find_many() + perms = await call_with_db_reconnect_retry( + prisma_client, + lambda: ObjectPermissionRepository(prisma_client).table.find_many(), + reason="sync_tool_policy_from_db_perms_lookup_failure", + ) self._blocked_tools_by_op_id = {} for row in perms: op_id = getattr(row, "object_permission_id", None) diff --git a/litellm/proxy/management_endpoints/cache_settings_endpoints.py b/litellm/proxy/management_endpoints/cache_settings_endpoints.py index d8eb5dfee92..b6ddf2d8e07 100644 --- a/litellm/proxy/management_endpoints/cache_settings_endpoints.py +++ b/litellm/proxy/management_endpoints/cache_settings_endpoints.py @@ -27,6 +27,7 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry from litellm.repositories.table_repositories import CacheConfigRepository from litellm.types.management_endpoints import ( CACHE_SETTINGS_FIELDS, @@ -160,8 +161,12 @@ class CacheSettingsManager: import json try: - cache_config = await CacheConfigRepository(prisma_client).table.find_unique( - where={"id": "cache_config"} + cache_config = await call_with_db_reconnect_retry( + prisma_client, + lambda: CacheConfigRepository(prisma_client).table.find_unique( + where={"id": "cache_config"} + ), + reason="init_cache_settings_in_db_lookup_failure", ) if cache_config is not None and cache_config.cache_settings: # Parse cache settings JSON diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 37a0285b196..709b5dcf7ea 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -311,7 +311,10 @@ from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.container_endpoints.endpoints import router as container_router from litellm.proxy.credential_endpoints.endpoints import router as credential_router from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import SpendLogCleanup -from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler +from litellm.proxy.db.exception_handler import ( + PrismaDBExceptionHandler, + call_with_db_reconnect_retry, +) from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed from litellm.proxy.discovery_endpoints import ui_discovery_endpoints_router from litellm.proxy.fine_tuning_endpoints.endpoints import router as fine_tuning_router @@ -5984,8 +5987,12 @@ class ProxyConfig: """ try: - sso_settings = await SSOConfigRepository(prisma_client).table.find_unique( - where={"id": "sso_config"} + sso_settings = await call_with_db_reconnect_retry( + prisma_client, + lambda: SSOConfigRepository(prisma_client).table.find_unique( + where={"id": "sso_config"} + ), + reason="init_sso_settings_in_db_lookup_failure", ) if sso_settings is not None: sso_settings.sso_settings.pop("role_mappings", None) @@ -6020,9 +6027,13 @@ class ProxyConfig: ) try: - db_record = await ConfigOverridesRepository( - prisma_client - ).table.find_unique(where={"config_type": "hashicorp_vault"}) + db_record = await call_with_db_reconnect_retry( + prisma_client, + lambda: ConfigOverridesRepository(prisma_client).table.find_unique( + where={"config_type": "hashicorp_vault"} + ), + reason="init_hashicorp_vault_config_override_lookup_failure", + ) if db_record is None or db_record.config_value is None: if self._last_hashicorp_vault_config is not None: diff --git a/litellm/proxy/search_endpoints/search_tool_registry.py b/litellm/proxy/search_endpoints/search_tool_registry.py index 588d71b77f9..2ec2533211b 100644 --- a/litellm/proxy/search_endpoints/search_tool_registry.py +++ b/litellm/proxy/search_endpoints/search_tool_registry.py @@ -7,6 +7,7 @@ from typing import List, Optional from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry from litellm.proxy.utils import PrismaClient from litellm.repositories.table_repositories import SearchToolsRepository from litellm.types.search import SearchTool @@ -180,10 +181,12 @@ class SearchToolRegistry: List of search tool configurations """ try: - search_tools_from_db = await SearchToolsRepository( - prisma_client - ).table.find_many( - order={"created_at": "desc"}, + search_tools_from_db = await call_with_db_reconnect_retry( + prisma_client, + lambda: SearchToolsRepository(prisma_client).table.find_many( + order={"created_at": "desc"}, + ), + reason="get_all_search_tools_from_db_lookup_failure", ) search_tools: List[SearchTool] = [] diff --git a/tests/test_litellm/proxy/db/test_tool_registry_writer.py b/tests/test_litellm/proxy/db/test_tool_registry_writer.py index 8074871c3dd..7bf1ffda4fe 100644 --- a/tests/test_litellm/proxy/db/test_tool_registry_writer.py +++ b/tests/test_litellm/proxy/db/test_tool_registry_writer.py @@ -291,3 +291,77 @@ async def test_tool_policy_registry_not_initialized_returns_untrusted(): assert not registry.is_initialized() result = registry.get_effective_policies(["unknown_tool"]) assert result == {"unknown_tool": "untrusted"} + + +@pytest.mark.asyncio +async def test_sync_tool_policy_from_db_retries_on_transport_error_first_read(): + """`ToolPolicyRegistry.sync_tool_policy_from_db` self-heals across one + ClientNotConnectedError on the tools read — the perms read still fires + after the recovery and the registry initializes cleanly.""" + import prisma as prisma_pkg + + registry = ToolPolicyRegistry() + invocations: list = [] + + async def _flaky_find_many(): + invocations.append(None) + if len(invocations) == 1: + raise prisma_pkg.errors.ClientNotConnectedError() + return [] + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_tooltable.find_many = AsyncMock( + side_effect=_flaky_find_many + ) + mock_prisma_client.db.litellm_objectpermissiontable.find_many = AsyncMock( + return_value=[] + ) + mock_prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) + mock_prisma_client._db_auth_reconnect_timeout_seconds = 2.0 + mock_prisma_client._db_auth_reconnect_lock_timeout_seconds = 0.1 + + await registry.sync_tool_policy_from_db(mock_prisma_client) + + assert len(invocations) == 2 + mock_prisma_client.attempt_db_reconnect.assert_awaited_once() + reconnect_kwargs = mock_prisma_client.attempt_db_reconnect.await_args.kwargs + assert ( + reconnect_kwargs["reason"] + == "sync_tool_policy_from_db_tools_lookup_failure" + ) + assert registry.is_initialized() + + +@pytest.mark.asyncio +async def test_sync_tool_policy_from_db_retries_on_transport_error_second_read(): + """Same as above but the blip happens on the perms read — distinct reason + tag in telemetry confirms the second wrap is wired separately.""" + import prisma as prisma_pkg + + registry = ToolPolicyRegistry() + perms_invocations: list = [] + + async def _flaky_perms_find_many(): + perms_invocations.append(None) + if len(perms_invocations) == 1: + raise prisma_pkg.errors.ClientNotConnectedError() + return [] + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_tooltable.find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.litellm_objectpermissiontable.find_many = AsyncMock( + side_effect=_flaky_perms_find_many + ) + mock_prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) + mock_prisma_client._db_auth_reconnect_timeout_seconds = 2.0 + mock_prisma_client._db_auth_reconnect_lock_timeout_seconds = 0.1 + + await registry.sync_tool_policy_from_db(mock_prisma_client) + + assert len(perms_invocations) == 2 + mock_prisma_client.attempt_db_reconnect.assert_awaited_once() + reconnect_kwargs = mock_prisma_client.attempt_db_reconnect.await_args.kwargs + assert ( + reconnect_kwargs["reason"] + == "sync_tool_policy_from_db_perms_lookup_failure" + ) diff --git a/tests/test_litellm/proxy/management_endpoints/search_endpoints/test_search_tool_management.py b/tests/test_litellm/proxy/management_endpoints/search_endpoints/test_search_tool_management.py index ea7e5591f18..f2ccfcd0155 100644 --- a/tests/test_litellm/proxy/management_endpoints/search_endpoints/test_search_tool_management.py +++ b/tests/test_litellm/proxy/management_endpoints/search_endpoints/test_search_tool_management.py @@ -611,6 +611,45 @@ async def test_list_search_tools_db_masking_sensitive_values(monkeypatch): app.dependency_overrides.pop(user_api_key_auth, None) +@pytest.mark.asyncio +async def test_get_all_search_tools_from_db_retries_on_transport_error(): + """`SearchToolRegistry.get_all_search_tools_from_db` self-heals across one + ClientNotConnectedError via call_with_db_reconnect_retry.""" + import prisma + from litellm.proxy.search_endpoints.search_tool_registry import ( + SearchToolRegistry, + ) + + invocations: list = [] + + async def _flaky_find_many(**kwargs): + invocations.append(None) + if len(invocations) == 1: + raise prisma.errors.ClientNotConnectedError() + return [] + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_searchtoolstable.find_many = AsyncMock( + side_effect=_flaky_find_many + ) + mock_prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) + mock_prisma_client._db_auth_reconnect_timeout_seconds = 2.0 + mock_prisma_client._db_auth_reconnect_lock_timeout_seconds = 0.1 + + result = await SearchToolRegistry.get_all_search_tools_from_db( + prisma_client=mock_prisma_client + ) + + assert result == [] + assert len(invocations) == 2 + mock_prisma_client.attempt_db_reconnect.assert_awaited_once() + reconnect_kwargs = mock_prisma_client.attempt_db_reconnect.await_args.kwargs + assert ( + reconnect_kwargs["reason"] + == "get_all_search_tools_from_db_lookup_failure" + ) + + @contextlib.contextmanager def _mock_search_tool_backend(db_tools): """Patch the DB registry, prisma client, and config so /search_tools/list diff --git a/tests/test_litellm/proxy/management_endpoints/test_cache_settings_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_cache_settings_endpoints.py index b892c4e556d..4bdef2e8f96 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_cache_settings_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_cache_settings_endpoints.py @@ -259,6 +259,41 @@ class TestCacheSettingsManager: mock_proxy_config._init_cache.assert_not_called() mock_proxy_config.switch_on_llm_response_caching.assert_not_called() + @pytest.mark.asyncio + async def test_init_cache_settings_in_db_retries_on_transport_error(self): + """`CacheSettingsManager.init_cache_settings_in_db` self-heals across one + ClientNotConnectedError via call_with_db_reconnect_retry.""" + import prisma + + invocations: list = [] + + async def _flaky_find_unique(**kwargs): + invocations.append(None) + if len(invocations) == 1: + raise prisma.errors.ClientNotConnectedError() + return None # No config → function returns early after retry. + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_cacheconfig.find_unique = AsyncMock( + side_effect=_flaky_find_unique + ) + mock_prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) + mock_prisma_client._db_auth_reconnect_timeout_seconds = 2.0 + mock_prisma_client._db_auth_reconnect_lock_timeout_seconds = 0.1 + mock_proxy_config = MagicMock() + + await CacheSettingsManager.init_cache_settings_in_db( + prisma_client=mock_prisma_client, proxy_config=mock_proxy_config + ) + + assert len(invocations) == 2 + mock_prisma_client.attempt_db_reconnect.assert_awaited_once() + reconnect_kwargs = mock_prisma_client.attempt_db_reconnect.await_args.kwargs + assert ( + reconnect_kwargs["reason"] + == "init_cache_settings_in_db_lookup_failure" + ) + # ── Audit-log emission for /cache/settings ──────────────────────────────────── diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index b2e36fd64cf..b1dee205eeb 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -4324,6 +4324,111 @@ async def test_init_sso_settings_in_db_empty_settings(): assert uppercased_settings == {} +@pytest.mark.asyncio +async def test_init_sso_settings_in_db_retries_on_transport_error(): + """`_init_sso_settings_in_db` self-heals across one ClientNotConnectedError + via call_with_db_reconnect_retry — mirrors the auth-path behavior so + startup/reload bursts don't spam the log.""" + import prisma + + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + mock_sso_config = MagicMock() + mock_sso_config.sso_settings = {"GOOGLE_CLIENT_ID": "xxx"} + + invocations: list = [] + + async def _flaky_find_unique(**kwargs): + invocations.append(None) + if len(invocations) == 1: + raise prisma.errors.ClientNotConnectedError() + return mock_sso_config + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_ssoconfig.find_unique = AsyncMock( + side_effect=_flaky_find_unique + ) + mock_prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) + mock_prisma_client._db_auth_reconnect_timeout_seconds = 2.0 + mock_prisma_client._db_auth_reconnect_lock_timeout_seconds = 0.1 + + with patch.object( + proxy_config, "_decrypt_and_set_db_env_variables" + ) as mock_decrypt: + await proxy_config._init_sso_settings_in_db(prisma_client=mock_prisma_client) + + assert len(invocations) == 2 + mock_prisma_client.attempt_db_reconnect.assert_awaited_once() + reconnect_kwargs = mock_prisma_client.attempt_db_reconnect.await_args.kwargs + assert reconnect_kwargs["reason"] == "init_sso_settings_in_db_lookup_failure" + mock_decrypt.assert_called_once() + + +@pytest.mark.asyncio +async def test_init_sso_settings_in_db_propagates_when_reconnect_fails(): + """When reconnect returns False (cooldown / lock contention), the original + ClientNotConnectedError is caught by the function's `except Exception` and + logged — no retry storm, no crash.""" + import prisma + + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_ssoconfig.find_unique = AsyncMock( + side_effect=prisma.errors.ClientNotConnectedError() + ) + mock_prisma_client.attempt_db_reconnect = AsyncMock(return_value=False) + mock_prisma_client._db_auth_reconnect_timeout_seconds = 2.0 + mock_prisma_client._db_auth_reconnect_lock_timeout_seconds = 0.1 + + # Should NOT raise — the function's own try/except swallows the propagated error. + await proxy_config._init_sso_settings_in_db(prisma_client=mock_prisma_client) + + mock_prisma_client.attempt_db_reconnect.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_init_hashicorp_vault_config_override_retries_on_transport_error(): + """`_init_hashicorp_vault_config_override` self-heals across one + ClientNotConnectedError via call_with_db_reconnect_retry.""" + import prisma + + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + proxy_config._last_hashicorp_vault_config = None + + invocations: list = [] + + async def _flaky_find_unique(**kwargs): + invocations.append(None) + if len(invocations) == 1: + raise prisma.errors.ClientNotConnectedError() + return None # No config in DB → function returns early after retry. + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_configoverrides.find_unique = AsyncMock( + side_effect=_flaky_find_unique + ) + mock_prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) + mock_prisma_client._db_auth_reconnect_timeout_seconds = 2.0 + mock_prisma_client._db_auth_reconnect_lock_timeout_seconds = 0.1 + + await proxy_config._init_hashicorp_vault_config_override( + prisma_client=mock_prisma_client + ) + + assert len(invocations) == 2 + mock_prisma_client.attempt_db_reconnect.assert_awaited_once() + reconnect_kwargs = mock_prisma_client.attempt_db_reconnect.await_args.kwargs + assert ( + reconnect_kwargs["reason"] + == "init_hashicorp_vault_config_override_lookup_failure" + ) + + def test_update_config_fields_uppercases_env_vars(monkeypatch): """ Ensure environment variables pulled from DB are uppercased when applied so From a75ed0079cfc1db222e1cd0a3f27402832aa801f Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 10 Jun 2026 11:44:24 -0700 Subject: [PATCH 047/185] chore(ui): make knip recognize .mjs scripts and openapi-typescript (#30052) The knip entry/project globs only matched scripts/**/*.ts, so the two .mjs scripts went unanalyzed and produced "no matches" config hints. openapi-typescript was also reported as unused because gen-api-types.mjs invokes its binary through a dynamic execFileSync path that knip cannot trace statically; ignoreDependencies records that it is genuinely used. --- ui/litellm-dashboard/knip.json | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/ui/litellm-dashboard/knip.json b/ui/litellm-dashboard/knip.json index e95c0acef3f..6f129398981 100644 --- a/ui/litellm-dashboard/knip.json +++ b/ui/litellm-dashboard/knip.json @@ -1,8 +1,9 @@ { "$schema": "https://unpkg.com/knip@5/schema.json", - "entry": ["scripts/**/*.ts"], - "project": ["src/**/*.{ts,tsx}", "tests/**/*.{ts,tsx}", "scripts/**/*.ts", "e2e_tests/**/*.ts"], + "entry": ["scripts/**/*.{ts,mjs}"], + "project": ["src/**/*.{ts,tsx}", "tests/**/*.{ts,tsx}", "scripts/**/*.{ts,mjs}", "e2e_tests/**/*.ts"], "ignore": ["src/lib/http/schema.d.ts"], + "ignoreDependencies": ["openapi-typescript"], "playwright": { "config": "e2e_tests/playwright.config.ts", "entry": ["e2e_tests/**/*.spec.ts", "e2e_tests/**/*.setup.ts", "e2e_tests/globalSetup.ts"] From 410b892f77678f8cbc611d1c699523dc6ae4acd8 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Wed, 10 Jun 2026 12:11:03 -0700 Subject: [PATCH 048/185] fix(register_model): preserve built-in cache pricing when registering custom overrides under unmapped keys (#30044) * fix(spend-tracking): fall back to direct spend-counter increment when reservation reconcile fails When the reservation-reconcile path in `_reconcile_budget_reservation_for_counter_update` hits a Redis error, it now correctly returns an empty set so that `increment_spend_counters` re-runs the direct increment for the affected counters. Previously, the function logged the failure, invalidated the reserved counters, and still returned the reserved counter keys, which caused the caller to skip the direct increment. With the increment skipped and the counter deleted, the next request reseeded the counter from `LiteLLM_VerificationToken.spend`, a column the batched flusher only updates every few seconds, so the enforced cross-pod spend value collapsed to a stale snapshot and budget gating stopped firing for affected keys. Adds a regression test that exercises the failure path with a flaky redis backend and asserts the actual response cost lands in the shared counter. * fix(register_model): preserve built-in cache pricing when registering custom overrides under unmapped keys When a custom-priced model is registered under a key shape that get_model_info cannot resolve (e.g. litellm_params.model set to bedrock/bedrock/us.anthropic.claude-sonnet-4-6 or another non-canonical alias), register_model previously fell back to an empty existing_model. The merged entry then carried only the fields the user set explicitly (input/output cost, provider) and dropped cache pricing. Downstream the cost calculator defaulted cache_creation_input_token_cost and cache_read_input_token_cost to 0, silently dropping the bulk of the bill for cache-heavy Anthropic traffic. register_model now attempts to resolve a canonical built-in entry by stripping provider prefixes, region prefixes, and provider-specific suffixes before giving up. When a variant resolves, its defaults (notably cache pricing) are inherited while the user's explicit overrides still win. When nothing resolves and the user supplied no cache pricing, it logs a warning instead of silently under-billing. * fix(router): inherit built-in cache pricing on deployments with partial custom pricing A deployment configured with only input_cost_per_token and output_cost_per_token under model_info was being registered under its model_info.id with no cache cost fields. The cost calculator then defaulted cache_creation_input_token_cost and cache_read_input_token_cost to 0, silently billing cache_read and cache_creation tokens at zero. For cache-heavy Anthropic traffic this drops the bulk of the bill. When the deployment's litellm_params.model resolves to a built-in cost-map entry, pull the cache pricing fields from there before registering. User-specified cache fields still win on merge; only missing fields are inherited. Pairs with the register_model fallback added earlier in this branch: that handles unmapped key shapes like bedrock/bedrock/x, this handles deploy-id keys whose backend model is mapped. * fix(register_model): inherit only cache pricing on unmapped-key fallback, not provider The unmapped-key fallback in register_model copied the entire resolved built-in entry, so registering openai/command-r-plus inherited the cohere built-in's litellm_provider and get_model_info(custom_llm_provider=openai) could no longer resolve it. Restrict the fallback to the cache-pricing fields, matching the router-side _inherit_builtin_cache_pricing, so the cache-cost dropout stays fixed without clobbering the registered provider. Add a direct unit test for Router._inherit_builtin_cache_pricing so the router coverage check sees it, and pin the fixed spend-counter contract: when reservation reconcile fails the counter must hold the directly incremented cost rather than being left at None. --- litellm/proxy/proxy_server.py | 3 +- litellm/router.py | 47 ++++++ litellm/utils.py | 75 ++++++++++ .../proxy/proxy_server/test_spend_counters.py | 7 +- .../test_budget_reservation_redis_failure.py | 87 +++++++++++ tests/test_litellm/proxy/test_proxy_server.py | 9 +- .../test_register_model_custom_pricing.py | 126 +++++++++++++++- .../test_router_model_cost_isolation.py | 135 ++++++++++++++++++ 8 files changed, 480 insertions(+), 9 deletions(-) create mode 100644 tests/test_litellm/proxy/spend_tracking/test_budget_reservation_redis_failure.py diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 709b5dcf7ea..ba23175c10f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2266,7 +2266,7 @@ async def _reconcile_budget_reservation_for_counter_update( ) except Exception: verbose_proxy_logger.warning( - "Failed to reconcile budget reservation after persisted spend; invalidating reserved counters and continuing", + "Failed to reconcile budget reservation after persisted spend; invalidating reserved counters and falling back to direct increment", exc_info=True, ) try: @@ -2277,6 +2277,7 @@ async def _reconcile_budget_reservation_for_counter_update( verbose_proxy_logger.exception( "Failed to invalidate reserved counters after reservation reconciliation failed" ) + return set() return reserved_counter_keys diff --git a/litellm/router.py b/litellm/router.py index 8966a2fc191..d1c8e227bea 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -7788,6 +7788,39 @@ class Router: return hash_object.hexdigest() + @staticmethod + def _inherit_builtin_cache_pricing( + model_info: dict, backend_model: str, custom_llm_provider: Optional[str] + ) -> None: + """Fill missing cache pricing on a custom-priced deployment entry from + the backend model's built-in cost map entry, so a deployment that + only spells out ``input_cost_per_token``/``output_cost_per_token`` + does not silently bill cache_read/cache_creation at 0. + + User-specified cache fields always win; only ``None``/missing entries + are inherited. No-op when the backend model has no canonical entry. + """ + cache_fields = ( + "cache_creation_input_token_cost", + "cache_creation_input_token_cost_above_1hr", + "cache_creation_input_token_cost_above_200k_tokens", + "cache_read_input_token_cost", + "cache_read_input_token_cost_above_200k_tokens", + ) + if all(model_info.get(f) is not None for f in cache_fields): + return + try: + backend_info = litellm.get_model_info( + model=backend_model, custom_llm_provider=custom_llm_provider + ) + except Exception: + return + for field in cache_fields: + if model_info.get(field) is None: + backend_value = backend_info.get(field) + if backend_value is not None: + model_info[field] = backend_value + def _create_deployment( self, deployment_info: dict, @@ -7816,6 +7849,13 @@ class Router: if deployment.litellm_params.get(field) is not None: _model_info[field] = deployment.litellm_params[field] + if _model_info.get("input_cost_per_token") is not None: + Router._inherit_builtin_cache_pricing( + model_info=_model_info, + backend_model=deployment.litellm_params.model, + custom_llm_provider=deployment.litellm_params.custom_llm_provider, + ) + ## REGISTER MODEL INFO IN LITELLM MODEL COST MAP model_id = deployment.model_info.id if model_id is not None: @@ -8562,6 +8602,13 @@ class Router: if field_value is not None: _model_info_dict[field] = field_value + if _model_info_dict.get("input_cost_per_token") is not None: + Router._inherit_builtin_cache_pricing( + model_info=_model_info_dict, + backend_model=deployment.litellm_params.model, + custom_llm_provider=deployment.litellm_params.custom_llm_provider, + ) + # Register custom pricing in litellm.model_cost. # Mirrors _create_deployment() logic to ensure dynamically-added deployments # (e.g., loaded from DB) also have their custom pricing registered. diff --git a/litellm/utils.py b/litellm/utils.py index 4f4e8d8cb9e..03c628b195f 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2887,6 +2887,61 @@ def _convert_stringified_numbers(value): return value +_BEDROCK_REGION_PREFIXES = ( + "us.", + "eu.", + "apac.", + "jp.", + "au.", + "us-gov.", + "global.", + "ap-northeast-1.", +) + +_CACHE_PRICING_FIELDS = ( + "cache_creation_input_token_cost", + "cache_creation_input_token_cost_above_1hr", + "cache_creation_input_token_cost_above_200k_tokens", + "cache_read_input_token_cost", + "cache_read_input_token_cost_above_200k_tokens", +) + + +def _resolve_builtin_model_cost_entry( + key: str, provider: str +) -> Optional[Dict[str, Any]]: + """Best-effort lookup of a built-in ``model_cost`` entry for a custom key + whose shape ``get_model_info`` cannot resolve (double provider prefixes + like ``bedrock/bedrock/us.anthropic.claude-sonnet-4-6`` or region aliases). + + Returns a copy of the matching entry so the caller can inherit its defaults + (most importantly cache pricing) without mutating the shared built-in. + Returns ``None`` when no safe match exists. + """ + candidates: List[str] = [] + segments = key.split("/") + idx = 0 + while idx < len(segments) - 1 and segments[idx] in LlmProvidersSet: + idx += 1 + candidates.append("/".join(segments[idx:])) + + base = candidates[-1] if candidates else key + for region_prefix in _BEDROCK_REGION_PREFIXES: + if base.startswith(region_prefix): + candidates.append(base[len(region_prefix) :]) + + if provider: + stripped = _strip_model_name(model=base, custom_llm_provider=provider) + if stripped != base: + candidates.append(stripped) + + for candidate in candidates: + entry = litellm.model_cost.get(candidate) + if entry is not None and entry.get("litellm_provider") is not None: + return dict(entry) + return None + + def register_model(model_cost: Union[str, dict]): # noqa: PLR0915 """ Register new / Override existing models (and their pricing) to specific providers. @@ -2933,6 +2988,26 @@ def register_model(model_cost: Union[str, dict]): # noqa: PLR0915 except Exception: existing_model = {} model_cost_key = key + builtin_entry = _resolve_builtin_model_cost_entry( + key=_key_str, provider=provider + ) + if builtin_entry is not None: + for field in _CACHE_PRICING_FIELDS: + if ( + value.get(field) is None + and builtin_entry.get(field) is not None + ): + existing_model[field] = builtin_entry[field] + elif ( + value.get("cache_creation_input_token_cost") is None + and value.get("cache_read_input_token_cost") is None + ): + verbose_logger.warning( + f"register_model: model={key} not in built-in cost map and no " + "prefix/region variant matched; cache cost fields will default " + "to 0. To track cache cost, add cache_creation_input_token_cost " + "and cache_read_input_token_cost to model_info" + ) # ``get_model_info`` returns ``litellm_provider: None`` when the # provider is unknown (e.g. custom deployments registered via # ``Router.add_deployment``). Persisting that None into diff --git a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py index ec8b06d9c97..4e5f13fdf88 100644 --- a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py +++ b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py @@ -192,8 +192,9 @@ async def test_reconcile_budget_reservation_for_counter_update_returns_empty_set async def test_reconcile_budget_reservation_for_counter_update_failure_invalidates( monkeypatch, ): - """Reservation reconcile raising must invalidate reserved counters but - not propagate the exception.""" + """Reservation reconcile raising must invalidate reserved counters, swallow + the exception, and return an empty set so the caller falls back to the + direct spend-counter increment instead of skipping it.""" import litellm.proxy.spend_tracking.budget_reservation as br monkeypatch.setattr( @@ -213,7 +214,7 @@ async def test_reconcile_budget_reservation_for_counter_update_failure_invalidat budget_reservation={"foo": "bar"}, response_cost=1.0 ) - assert result == {"spend:key:abc"} + assert result == set() assert fake_invalidate.called is True diff --git a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation_redis_failure.py b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation_redis_failure.py new file mode 100644 index 00000000000..c123eeeed36 --- /dev/null +++ b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation_redis_failure.py @@ -0,0 +1,87 @@ +""" +Regression test for enforced-spend underreporting when Redis fails during the +budget-reservation reconcile step of ``increment_spend_counters``. + +Production failure mode: a managed Redis returns an intermittent timeout on the +reconcile increment. Reconcile deletes (invalidates) the shared counter and +gives up, but ``increment_spend_counters`` still treats the counter as +"already reconciled" and skips the direct increment. The actual call cost never +lands in the enforced counter, so budgets stop gating until the next cold +reseed pulls a lagging value from the DB. + +The fix makes the reconcile path fall back to the direct increment when it +fails, so the actual cost is always written to the shared counter. +""" + +import pytest + +from litellm.caching import DualCache +from litellm.proxy import proxy_server + + +class _FlakyRedisCache: + def __init__(self) -> None: + self._store: dict = {} + self._increment_calls = 0 + + async def async_increment(self, key, value, **kwargs): + self._increment_calls += 1 + if self._increment_calls == 1: + raise Exception("Redis timeout") + self._store[key] = float(self._store.get(key, 0.0)) + float(value) + return self._store[key] + + async def async_get_cache(self, key, *args, **kwargs): + return self._store.get(key) + + async def async_delete_cache(self, key, *args, **kwargs): + self._store.pop(key, None) + + async def async_set_cache(self, key, value, *args, **kwargs): + self._store[key] = float(value) + return True + + +@pytest.mark.asyncio +async def test_direct_increment_runs_when_reservation_reconcile_hits_redis_failure( + monkeypatch, +): + hashed_token = "hashed_test_token" + counter_key = f"spend:key:{hashed_token}" + reserved_cost = 0.5 + response_cost = 1.0 + + flaky_redis = _FlakyRedisCache() + flaky_redis._store[counter_key] = reserved_cost + + monkeypatch.setattr(proxy_server, "prisma_client", None) + monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache()) + monkeypatch.setattr(proxy_server.spend_counter_cache, "redis_cache", flaky_redis) + proxy_server.spend_counter_cache.in_memory_cache.set_cache( + key=counter_key, value=reserved_cost + ) + + budget_reservation = { + "reserved_cost": reserved_cost, + "finalized": False, + "entries": [ + { + "counter_key": counter_key, + "entity_type": "Key", + "entity_id": hashed_token, + "reserved_cost": reserved_cost, + "applied_adjustment": 0.0, + } + ], + } + + await proxy_server.increment_spend_counters( + token=hashed_token, + team_id=None, + user_id=None, + response_cost=response_cost, + budget_reservation=budget_reservation, + ) + + enforced_spend = await flaky_redis.async_get_cache(key=counter_key) + assert enforced_spend == response_cost diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index b1dee205eeb..9f2c5ffd615 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -6609,7 +6609,12 @@ async def test_increment_spend_counters_finalizes_none_cost_reservation(): @pytest.mark.asyncio -async def test_increment_spend_counters_invalidates_bad_reserved_counter_without_failing(): +async def test_increment_spend_counters_falls_back_to_direct_increment_on_bad_reserved_counter(): + """When the reservation reconcile fails, the reserved counters are + invalidated and the actual response cost must still be written via the + direct increment fallback. Leaving the counter at ``None`` lets the next + request reseed a stale value from the DB and silently stops budget gating, + which is the bug this fix addresses.""" from litellm.caching.dual_cache import DualCache from litellm.proxy.proxy_server import increment_spend_counters @@ -6650,7 +6655,7 @@ async def test_increment_spend_counters_invalidates_bad_reserved_counter_without counter_cache.in_memory_cache.get_cache( key="spend:key:key-bad-reserved-counter" ) - is None + == 0.25 ) finally: ps.spend_counter_cache = orig_counter diff --git a/tests/test_litellm/test_register_model_custom_pricing.py b/tests/test_litellm/test_register_model_custom_pricing.py index 719cb8eecd2..e384d3e1161 100644 --- a/tests/test_litellm/test_register_model_custom_pricing.py +++ b/tests/test_litellm/test_register_model_custom_pricing.py @@ -301,6 +301,126 @@ def test_register_model_strips_none_litellm_provider_from_get_model_info(monkeyp litellm.model_cost.pop(model_key, None) +def test_register_model_inherits_builtin_cache_pricing_for_unmapped_key(): + """Registering a custom override under a key shape that + ``get_model_info`` cannot resolve (e.g. a double provider prefix like + ``bedrock/bedrock/us.anthropic.claude-sonnet-4-6``) must still inherit + the built-in cache pricing for the underlying model. + + Before the fix ``register_model`` fell back to an empty ``existing_model`` + so the merged entry only carried the fields the user set explicitly + (input/output cost). ``cache_creation_input_token_cost`` and + ``cache_read_input_token_cost`` were absent, and the cost calculator + silently charged 0 for every cache token, dropping the bulk of the bill + for cache-heavy Anthropic traffic. + + Regression for the cache-pricing dropout under partial overrides. + """ + from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token + from litellm.types.utils import PromptTokensDetailsWrapper, Usage + + original_model_cost = litellm.model_cost + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + builtin_key = "us.anthropic.claude-sonnet-4-6" + registered_key = f"bedrock/bedrock/{builtin_key}" + builtin = litellm.model_cost[builtin_key] + + assert builtin["cache_creation_input_token_cost"] > 0 + assert builtin["cache_read_input_token_cost"] > 0 + + try: + litellm.register_model( + { + registered_key: { + "input_cost_per_token": builtin["input_cost_per_token"], + "output_cost_per_token": builtin["output_cost_per_token"], + "litellm_provider": "bedrock", + } + } + ) + + registered = litellm.model_cost[registered_key] + assert ( + registered.get("cache_creation_input_token_cost") + == builtin["cache_creation_input_token_cost"] + ) + assert ( + registered.get("cache_read_input_token_cost") + == builtin["cache_read_input_token_cost"] + ) + assert registered["litellm_provider"] == "bedrock" + + usage = Usage( + prompt_tokens=1100, + completion_tokens=100, + total_tokens=1200, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=800, + text_tokens=100, + ), + cache_creation_input_tokens=200, + ) + + input_cost, output_cost = generic_cost_per_token( + model=registered_key, + usage=usage, + custom_llm_provider="bedrock", + ) + + text_only_cost = builtin["input_cost_per_token"] * 100 + expected_input_cost = ( + text_only_cost + + builtin["cache_read_input_token_cost"] * 800 + + builtin["cache_creation_input_token_cost"] * 200 + ) + assert abs(input_cost - expected_input_cost) < 1e-12 + assert abs(output_cost - builtin["output_cost_per_token"] * 100) < 1e-12 + assert input_cost > text_only_cost + 1e-12 + finally: + litellm.model_cost.pop(registered_key, None) + litellm.model_cost = original_model_cost + os.environ.pop("LITELLM_LOCAL_MODEL_COST_MAP", None) + from litellm.utils import _invalidate_model_cost_lowercase_map + + _invalidate_model_cost_lowercase_map() + + +def test_register_model_warns_when_no_builtin_match_for_cache_pricing(caplog): + """When a custom override is registered under a key that neither + ``get_model_info`` nor any prefix/region variant can resolve to a + built-in entry, ``register_model`` must warn that cache cost fields will + default to 0 instead of silently producing an under-billed entry. + """ + import logging + + from litellm._logging import verbose_logger + + registered_key = "bedrock/totally-made-up-model-alias-xyz" + litellm.model_cost.pop(registered_key, None) + + try: + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + litellm.register_model( + { + registered_key: { + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + "litellm_provider": "bedrock", + } + } + ) + + assert any( + registered_key in record.message + and "cache_creation_input_token_cost" in record.message + for record in caplog.records + ), "expected a warning naming the unmapped key and the cache cost fields" + finally: + litellm.model_cost.pop(registered_key, None) + + def test_register_model_router_add_deployment_custom_pricing_applies(): """End-to-end regression for https://github.com/BerriAI/litellm/issues/28336. @@ -344,9 +464,9 @@ def test_register_model_router_add_deployment_custom_pricing_applies(): f"{model_key} / {deployment_model}" ) for k in registered_keys: - assert _check_provider_match(litellm.model_cost[k], "openai") is True, ( - f"custom pricing for {k} was dropped by _check_provider_match" - ) + assert ( + _check_provider_match(litellm.model_cost[k], "openai") is True + ), f"custom pricing for {k} was dropped by _check_provider_match" finally: litellm.model_cost.pop(model_key, None) litellm.model_cost.pop(deployment_model, None) diff --git a/tests/test_litellm/test_router_model_cost_isolation.py b/tests/test_litellm/test_router_model_cost_isolation.py index 9454e03e918..ee64f44d32c 100644 --- a/tests/test_litellm/test_router_model_cost_isolation.py +++ b/tests/test_litellm/test_router_model_cost_isolation.py @@ -402,3 +402,138 @@ def test_should_not_downgrade_chatgpt_shared_key_mode_with_alias_override(): assert bridge_model_info["mode"] == "responses" finally: _restore_model_cost_entries(model_keys) + + +def test_partial_custom_pricing_inherits_builtin_cache_pricing(): + """A deployment that overrides only input/output cost on a cache-supporting + model must still bill cache_read and cache_creation tokens. Before the + fix the deploy-id entry was registered with the user's two fields and + nothing else, so the cost calculator silently billed cache tokens at 0. + Regression for the prompt-caching cost dropout reported by the customer. + """ + backend_model = "anthropic/claude-sonnet-4-5-20250929" + deploy_id = "claude-deploy-partial-pricing" + + builtin_info = litellm.get_model_info(model=backend_model) + builtin_cache_create = builtin_info["cache_creation_input_token_cost"] + builtin_cache_read = builtin_info["cache_read_input_token_cost"] + assert builtin_cache_create is not None and builtin_cache_create > 0 + assert builtin_cache_read is not None and builtin_cache_read > 0 + + model_keys = { + deploy_id: litellm.model_cost.get(deploy_id), + backend_model: copy.deepcopy(litellm.model_cost.get(backend_model)), + } + try: + Router( + model_list=[ + { + "model_name": "claude-custom", + "litellm_params": { + "model": backend_model, + "api_key": "fake-key", + }, + "model_info": { + "id": deploy_id, + "input_cost_per_token": 0.000003, + "output_cost_per_token": 0.000015, + }, + } + ], + ) + + entry = litellm.model_cost[deploy_id] + assert entry["input_cost_per_token"] == 0.000003 + assert entry["output_cost_per_token"] == 0.000015 + assert entry.get("cache_creation_input_token_cost") == builtin_cache_create + assert entry.get("cache_read_input_token_cost") == builtin_cache_read + finally: + _restore_model_cost_entries(model_keys) + + +def test_partial_pricing_does_not_overwrite_explicit_cache_fields(): + """When the user explicitly sets cache_*_input_token_cost on a deployment, + those values must not be replaced by the built-in fallback. + """ + backend_model = "anthropic/claude-sonnet-4-5-20250929" + deploy_id = "claude-deploy-explicit-cache" + + explicit_cache_create = 0.00001 + explicit_cache_read = 0.0000005 + builtin_info = litellm.get_model_info(model=backend_model) + assert builtin_info["cache_creation_input_token_cost"] != explicit_cache_create + assert builtin_info["cache_read_input_token_cost"] != explicit_cache_read + + model_keys = { + deploy_id: litellm.model_cost.get(deploy_id), + backend_model: copy.deepcopy(litellm.model_cost.get(backend_model)), + } + try: + Router( + model_list=[ + { + "model_name": "claude-custom-explicit", + "litellm_params": { + "model": backend_model, + "api_key": "fake-key", + }, + "model_info": { + "id": deploy_id, + "input_cost_per_token": 0.000003, + "output_cost_per_token": 0.000015, + "cache_creation_input_token_cost": explicit_cache_create, + "cache_read_input_token_cost": explicit_cache_read, + }, + } + ], + ) + + entry = litellm.model_cost[deploy_id] + assert entry.get("cache_creation_input_token_cost") == explicit_cache_create + assert entry.get("cache_read_input_token_cost") == explicit_cache_read + finally: + _restore_model_cost_entries(model_keys) + + +def test_inherit_builtin_cache_pricing_fills_only_missing_fields(): + """Direct unit test of the helper: missing cache fields are filled from the + backend model's built-in entry, while an explicitly set cache field and the + user's input/output pricing are left untouched. + """ + backend_model = "anthropic/claude-sonnet-4-5-20250929" + builtin_info = litellm.get_model_info(model=backend_model) + builtin_cache_create = builtin_info["cache_creation_input_token_cost"] + builtin_cache_read = builtin_info["cache_read_input_token_cost"] + assert builtin_cache_create is not None and builtin_cache_create > 0 + assert builtin_cache_read is not None and builtin_cache_read > 0 + + explicit_cache_read = builtin_cache_read + 1 + model_info = { + "input_cost_per_token": 0.000003, + "cache_read_input_token_cost": explicit_cache_read, + } + + Router._inherit_builtin_cache_pricing( + model_info=model_info, + backend_model=backend_model, + custom_llm_provider="anthropic", + ) + + assert model_info["input_cost_per_token"] == 0.000003 + assert model_info["cache_read_input_token_cost"] == explicit_cache_read + assert model_info["cache_creation_input_token_cost"] == builtin_cache_create + + +def test_inherit_builtin_cache_pricing_noop_for_unknown_backend(): + """No canonical entry for the backend model means the helper leaves the + passed-in dict unchanged rather than raising. + """ + model_info = {"input_cost_per_token": 0.000003} + + Router._inherit_builtin_cache_pricing( + model_info=model_info, + backend_model="this-backend-model-does-not-exist-x9y8z7", + custom_llm_provider=None, + ) + + assert model_info == {"input_cost_per_token": 0.000003} From a4a3348801cbb6ea5296b04ff1c6aeb3c10cdd6c Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 10 Jun 2026 12:31:00 -0700 Subject: [PATCH 049/185] [internal copy of #28007] Fix/gcp model garden streaming (#28363) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(vertex): stream Model Garden Gemma/Qwen responses correctly through /v1/messages * test(vertex): cover _CombinedChunkSplitter defensive branches * test(databricks): rename test file to avoid duplicate basename collision * fix(databricks,anthropic): defensive token defaults; document single-mode splitter Address greptile P2 concerns: - databricks: default usage token fields to 0 when constructing ChatCompletionUsageBlock from a partially populated usage block — matches the defensive pattern used in ollama/vertex_ai/cohere/bedrock. - _CombinedChunkSplitter: clarify in the docstring that an instance is single-mode (sync or async, not both), since the two iteration paths hold independent upstream iterator references. Co-authored-by: Claude --------- Co-authored-by: Steven Kessler <9701252+stvnksslr@users.noreply.github.com> Co-authored-by: Claude --- .../adapters/streaming_iterator.py | 98 +++++++++++- litellm/llms/base_llm/base_model_iterator.py | 5 + litellm/llms/databricks/streaming_utils.py | 22 +++ .../test_streaming_iterator_combined_chunk.py | 150 ++++++++++++++++++ .../test_databricks_streaming_utils.py | 64 ++++++++ 5 files changed, 338 insertions(+), 1 deletion(-) create mode 100644 tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_combined_chunk.py create mode 100644 tests/test_litellm/llms/databricks/test_databricks_streaming_utils.py diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py index bacb9f8ddf6..8c20f4c430e 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py @@ -1,5 +1,6 @@ # What is this? ## Translates OpenAI call to Anthropic `/v1/messages` format +import copy import json import traceback from collections import deque @@ -29,6 +30,98 @@ if TYPE_CHECKING: from litellm.types.utils import ModelResponseStream +class _CombinedChunkSplitter: + """ + Splits a streaming chunk that carries BOTH response content and a + ``finish_reason`` into two chunks: a content-only chunk followed by a + finish-only chunk. + + ``AnthropicStreamWrapper`` (via ``translate_streaming_openai_response_to_anthropic``) + assumes content and ``finish_reason`` never arrive in the same chunk — true for + real provider streams, but false for fake-streamed providers (e.g. Vertex AI + Gemma ``:predict``) where ``MockResponseIterator`` collapses the entire response + into a single chunk. Without this split the assumption causes all content to be + silently dropped (only the ``message_delta`` stop event is emitted). + + Supports both sync and async iteration, since ``AnthropicStreamWrapper`` exposes + both ``__next__`` and ``__anext__``. An instance is single-mode: callers must + iterate it either synchronously or asynchronously, never both — the two modes + hold independent iterator references on the upstream stream and mixing them + would advance them out of sync. + """ + + def __init__(self, completion_stream: Any): + self._stream = completion_stream + self._sync_iter: Optional[Iterator[Any]] = None + self._async_iter: Optional[AsyncIterator[Any]] = None + self._buffer: deque = deque() + + @staticmethod + def _is_combined(chunk: Any) -> bool: + """True if ``chunk`` carries response content AND a finish_reason.""" + choices = getattr(chunk, "choices", None) + if not choices: + return False + choice = choices[0] + if getattr(choice, "finish_reason", None) is None: + return False + delta = getattr(choice, "delta", None) + if delta is None: + return False + return bool( + getattr(delta, "content", None) + or getattr(delta, "tool_calls", None) + or getattr(delta, "reasoning_content", None) + or getattr(delta, "thinking_blocks", None) + ) + + @staticmethod + def _split(chunk: Any) -> List[Any]: + """Return ``[chunk]``, or ``[content_chunk, finish_chunk]`` if combined.""" + if not _CombinedChunkSplitter._is_combined(chunk): + return [chunk] + + # Content chunk: keep the delta payload, clear the finish_reason. + content_chunk = copy.deepcopy(chunk) + content_chunk.choices[0].finish_reason = None + + # Finish chunk: keep finish_reason (and usage), clear the delta payload. + finish_chunk = copy.deepcopy(chunk) + finish_delta = finish_chunk.choices[0].delta + finish_delta.content = None + if hasattr(finish_delta, "tool_calls"): + finish_delta.tool_calls = None + if hasattr(finish_delta, "reasoning_content"): + finish_delta.reasoning_content = None + if hasattr(finish_delta, "thinking_blocks"): + finish_delta.thinking_blocks = None + return [content_chunk, finish_chunk] + + def __iter__(self) -> "Iterator[Any]": + return self + + def __next__(self) -> Any: + if self._buffer: + return self._buffer.popleft() + if self._sync_iter is None: + self._sync_iter = iter(self._stream) + chunk = next(self._sync_iter) # propagates StopIteration when exhausted + self._buffer.extend(self._split(chunk)) + return self._buffer.popleft() + + def __aiter__(self) -> "AsyncIterator[Any]": + return self + + async def __anext__(self) -> Any: + if self._buffer: + return self._buffer.popleft() + if self._async_iter is None: + self._async_iter = self._stream.__aiter__() + chunk = await self._async_iter.__anext__() # propagates StopAsyncIteration + self._buffer.extend(self._split(chunk)) + return self._buffer.popleft() + + class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): """ - first chunk return 'message_start' @@ -62,7 +155,10 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): compaction_block: Optional[CompactionBlock] = None, iterations_usage: Optional[List[UsageIteration]] = None, ): - super().__init__(completion_stream) + # Wrap the upstream stream so chunks that carry both content and a + # finish_reason (fake-streamed providers) are split into two — see + # _CombinedChunkSplitter. + super().__init__(_CombinedChunkSplitter(completion_stream)) self.model = model # Mapping of truncated tool names to original names (for OpenAI's 64-char limit) self.tool_name_mapping = tool_name_mapping or {} diff --git a/litellm/llms/base_llm/base_model_iterator.py b/litellm/llms/base_llm/base_model_iterator.py index cf1fd6f786e..bf1bfd06537 100644 --- a/litellm/llms/base_llm/base_model_iterator.py +++ b/litellm/llms/base_llm/base_model_iterator.py @@ -50,6 +50,11 @@ def convert_model_response_to_streaming( model=model_response.model, choices=streaming_choices, ) + # Carry usage onto the streaming chunk so fake-streamed responses + # (e.g. Vertex AI Gemma :predict) still report token counts. + usage = getattr(model_response, "usage", None) + if usage is not None: + setattr(processed_chunk, "usage", usage) return processed_chunk except Exception as e: raise ValueError( diff --git a/litellm/llms/databricks/streaming_utils.py b/litellm/llms/databricks/streaming_utils.py index eebe3182881..7a7330227d6 100644 --- a/litellm/llms/databricks/streaming_utils.py +++ b/litellm/llms/databricks/streaming_utils.py @@ -25,6 +25,28 @@ class ModelResponseIterator: finish_reason = "" usage: Optional[ChatCompletionUsageBlock] = None + # Usage-only final chunk (OpenAI ``stream_options.include_usage``) + # arrives with an empty ``choices`` list — return usage without + # indexing ``choices[0]``. + if len(processed_chunk.choices) == 0: + final_usage = getattr(processed_chunk, "usage", None) + return GenericStreamingChunk( + text="", + tool_use=None, + is_finished=False, + finish_reason="", + usage=( + ChatCompletionUsageBlock( + prompt_tokens=final_usage.prompt_tokens or 0, + completion_tokens=final_usage.completion_tokens or 0, + total_tokens=final_usage.total_tokens or 0, + ) + if final_usage is not None + else None + ), + index=0, + ) + if processed_chunk.choices[0].delta.content is not None: # type: ignore text = processed_chunk.choices[0].delta.content # type: ignore diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_combined_chunk.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_combined_chunk.py new file mode 100644 index 00000000000..f74c5b61300 --- /dev/null +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_combined_chunk.py @@ -0,0 +1,150 @@ +""" +Regression tests for fake-streamed providers routed through `/v1/messages`. + +A fake-streaming provider (e.g. Vertex AI Gemma `:predict`) collapses its whole +response into a single `MockResponseIterator` chunk that carries content text AND a +`finish_reason` together. `AnthropicStreamWrapper` previously dropped all content in +this case — `translate_streaming_openai_response_to_anthropic` sees the finish_reason +and emits only a `message_delta`. `_CombinedChunkSplitter` splits such chunks so the +content survives. +""" + +import asyncio +import json +from types import SimpleNamespace + +from litellm.llms.anthropic.experimental_pass_through.adapters.streaming_iterator import ( + AnthropicStreamWrapper, + _CombinedChunkSplitter, +) +from litellm.llms.base_llm.base_model_iterator import MockResponseIterator +from litellm.types.utils import ( + Choices, + Delta, + Message, + ModelResponse, + ModelResponseStream, + StreamingChoices, + Usage, +) + + +def _build_fake_stream( + content: str, finish_reason: str = "stop" +) -> MockResponseIterator: + """Mimic a Vertex Gemma `:predict` fake stream: one collapsed chunk.""" + model_response = ModelResponse() + model_response.choices = [ + Choices( + index=0, + message=Message(role="assistant", content=content), + finish_reason=finish_reason, + ) + ] + model_response.usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15) + model_response.model = "gemma4" + return MockResponseIterator(model_response=model_response) + + +def _collect_async(wrapper: AnthropicStreamWrapper) -> str: + async def _run() -> str: + out = [] + async for raw in wrapper.async_anthropic_sse_wrapper(): + out.append(raw.decode() if isinstance(raw, bytes) else raw) + return "".join(out) + + return asyncio.run(_run()) + + +def test_fake_stream_content_reaches_anthropic_sse(): + """Content from a collapsed fake-stream chunk must be emitted as a delta.""" + wrapper = AnthropicStreamWrapper( + completion_stream=_build_fake_stream("Hello, the answer is 2."), + model="gemma4", + ) + sse = _collect_async(wrapper) + + assert "content_block_delta" in sse + assert "Hello, the answer is 2." in sse + assert "message_delta" in sse + assert "message_stop" in sse + + +def test_fake_stream_usage_preserved(): + """The finish chunk keeps usage so output_tokens is non-zero.""" + wrapper = AnthropicStreamWrapper( + completion_stream=_build_fake_stream("Two."), + model="gemma4", + ) + sse = _collect_async(wrapper) + + message_delta = next( + json.loads(line[len("data: ") :]) + for block in sse.split("\n\n") + for line in block.splitlines() + if line.startswith("data: ") and '"message_delta"' in line + ) + assert message_delta["usage"]["output_tokens"] == 5 + assert message_delta["usage"]["input_tokens"] == 10 + + +def test_splitter_passes_through_non_combined_chunks(): + """A chunk with content but no finish_reason is not split.""" + chunk = ModelResponseStream( + choices=[ + StreamingChoices( + index=0, delta=Delta(content="partial"), finish_reason=None + ) + ] + ) + chunks = list(_CombinedChunkSplitter(iter([chunk]))) + assert len(chunks) == 1 + assert chunks[0].choices[0].delta.content == "partial" + + +def test_splitter_splits_combined_chunk_into_content_then_finish(): + """A chunk with both content and finish_reason becomes two chunks.""" + chunk = ModelResponseStream( + choices=[ + StreamingChoices(index=0, delta=Delta(content="done"), finish_reason="stop") + ] + ) + content_chunk, finish_chunk = list(_CombinedChunkSplitter(iter([chunk]))) + + assert content_chunk.choices[0].delta.content == "done" + assert content_chunk.choices[0].finish_reason is None + + assert finish_chunk.choices[0].finish_reason == "stop" + assert finish_chunk.choices[0].delta.content is None + + +def test_is_combined_false_when_choices_empty(): + """A metadata-only chunk with no choices is never treated as combined.""" + assert _CombinedChunkSplitter._is_combined(SimpleNamespace(choices=[])) is False + + +def test_is_combined_false_when_delta_missing(): + """A finish chunk whose choice has no delta is not combined.""" + chunk = SimpleNamespace(choices=[SimpleNamespace(finish_reason="stop", delta=None)]) + assert _CombinedChunkSplitter._is_combined(chunk) is False + + +def test_split_clears_reasoning_and_thinking_on_finish_chunk(): + """When the combined delta carries reasoning/thinking, only the content + chunk keeps them — the finish chunk is cleared.""" + delta = SimpleNamespace( + content="hi", + tool_calls=None, + reasoning_content="some reasoning", + thinking_blocks=[{"type": "thinking"}], + ) + chunk = SimpleNamespace( + choices=[SimpleNamespace(finish_reason="stop", delta=delta)] + ) + + content_chunk, finish_chunk = _CombinedChunkSplitter._split(chunk) + + assert content_chunk.choices[0].delta.reasoning_content == "some reasoning" + assert content_chunk.choices[0].delta.thinking_blocks == [{"type": "thinking"}] + assert finish_chunk.choices[0].delta.reasoning_content is None + assert finish_chunk.choices[0].delta.thinking_blocks is None diff --git a/tests/test_litellm/llms/databricks/test_databricks_streaming_utils.py b/tests/test_litellm/llms/databricks/test_databricks_streaming_utils.py new file mode 100644 index 00000000000..5612864a841 --- /dev/null +++ b/tests/test_litellm/llms/databricks/test_databricks_streaming_utils.py @@ -0,0 +1,64 @@ +""" +Regression test for the databricks streaming chunk parser. + +OpenAI-compatible servers (e.g. Vertex AI Model Garden vLLM endpoints) send a final +usage-only chunk with an empty `choices` list when `stream_options.include_usage` is +set. `chunk_parser` previously did `choices[0]` unconditionally, raising +`IndexError` -> `MidStreamFallbackError` and crashing the stream. +""" + +from litellm.llms.databricks.streaming_utils import ModelResponseIterator + + +def test_chunk_parser_handles_empty_choices_usage_chunk(): + """A usage-only final chunk (empty choices) must not raise IndexError.""" + iterator = ModelResponseIterator(streaming_response=None, sync_stream=True) + usage_only_chunk = { + "id": "chatcmpl-x", + "object": "chat.completion.chunk", + "created": 1, + "model": "m", + "choices": [], + "usage": {"prompt_tokens": 20, "completion_tokens": 8, "total_tokens": 28}, + } + + result = iterator.chunk_parser(chunk=usage_only_chunk) + + assert result["text"] == "" + assert result["is_finished"] is False + assert result["usage"] is not None + assert result["usage"]["prompt_tokens"] == 20 + assert result["usage"]["completion_tokens"] == 8 + + +def test_chunk_parser_empty_choices_without_usage(): + """An empty-choices chunk with no usage block returns usage=None, no error.""" + iterator = ModelResponseIterator(streaming_response=None, sync_stream=True) + chunk = { + "id": "chatcmpl-x", + "object": "chat.completion.chunk", + "created": 1, + "model": "m", + "choices": [], + } + + result = iterator.chunk_parser(chunk=chunk) + + assert result["text"] == "" + assert result["usage"] is None + + +def test_chunk_parser_normal_content_chunk_still_works(): + """A regular content chunk is unaffected by the empty-choices guard.""" + iterator = ModelResponseIterator(streaming_response=None, sync_stream=True) + chunk = { + "id": "chatcmpl-x", + "object": "chat.completion.chunk", + "created": 1, + "model": "m", + "choices": [{"index": 0, "delta": {"content": "hi"}, "finish_reason": None}], + } + + result = iterator.chunk_parser(chunk=chunk) + + assert result["text"] == "hi" From 20e453f698dc0758a15a491818411372da041415 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 10 Jun 2026 13:52:26 -0700 Subject: [PATCH 050/185] feat(cli): per-agent `lite claude` / `codex` / `opencode` commands that wrap coding agents through the proxy (#29850) * feat(cli): add `litellm-proxy run -- ` to wrap coding agents through the proxy Wraps Claude Code, Codex, OpenCode, and any other coding agent so all of its LLM traffic routes through a LiteLLM proxy, with the agent-vault style of "just works" DX: one `run -- ` command, auto SSO login when interactive, env-key "agent mode" for containers/CI, and a fail-fast key check against the proxy so bad credentials error immediately instead of deep inside the agent. The wrapped binary is detected by name to pick the right variables. Claude Code gets ANTHROPIC_BASE_URL (the bare proxy root, so it appends /v1/messages) and ANTHROPIC_AUTH_TOKEN, with any stray ANTHROPIC_API_KEY cleared so the proxy token wins. Codex and OpenCode get OPENAI_BASE_URL (proxy + /v1) and OPENAI_API_KEY. Unrecognized commands get both sets so they work either way. `litellm-proxy claude-code` remains as a shortcut for `run -- claude`. The core logic is split into dependency-injected helpers (agent_profile, build_agent_env, verify_proxy_key, run_agent) so env wiring, the preflight, and the launch handoff are unit-tested without monkeypatching, alongside CliRunner tests for auth resolution, agent mode, and auto-login. Mutation-tested the env profiles, preflight, and agent-mode branch to confirm the tests fail when the behavior is broken. https://claude.ai/code/session_0154VpLXW7mMvk5wfbgPRJa6 * Make each coding agent its own litellm-proxy command Replace the `run -- ` interface and the `claude-code` shortcut with top-level commands generated per known agent, so launching is just `litellm-proxy claude`, `litellm-proxy codex`, or `litellm-proxy opencode`, with everything after the agent name forwarded straight to it. This drops the ceremony of `run --` and cuts typing. The `--model`/`--small-fast-model` wrapper flags are gone; pass the agent's own model flag instead, or export the model env vars (the wrapper preserves what you already have set), which keeps the surface minimal and avoids intercepting flags the agent owns. Rename the module to agents.py to match. * fix(cli): route `litellm-proxy codex` through the proxy via a custom provider Codex ignores OPENAI_BASE_URL (it always dials api.openai.com over the Responses WebSocket transport), so the OpenAI env profile alone left `litellm-proxy codex` talking to OpenAI directly instead of the proxy. Point Codex at the proxy with a custom provider passed as `-c` config overrides, and force the HTTP/SSE Responses transport with supports_websockets=false since the proxy does not speak the Responses WebSocket protocol. The provider reads its key from OPENAI_API_KEY, which the agent env already exports. The overrides are injected ahead of the user's args so they precede Codex's subcommand. Claude Code and OpenCode are unaffected; they honor the exported env vars. Adds regression tests for the per-agent launch args and the injection ordering. Co-authored-by: Mateo Wang * Rename litellm-proxy CLI command to lite The proxy management CLI was invoked as litellm-proxy, which is a lot to type for an everyday command. Rename the console script entry point to lite and update the in-CLI usage examples, help text, error messages and docs to match. * fix(sso): stop CLI auth success page from hanging on "Closing..." The CLI opens the SSO success page with webbrowser.open, so the tab is not script-opened and the browser refuses window.close(). The countdown would end on "Closing..." and the tab would sit there forever. Drop the countdown and just show "You can now close this window and return to your terminal." from the start, while still attempting window.close() once so the tab auto-closes in the rare case the browser allows it. Add a regression test asserting the manual-close instruction is always present and the misleading countdown/"Closing..." text is gone. * fix(cli): reattach controlling terminal after SSO login, keep litellm-proxy alias When the first `lite claude` has to log in via browser SSO, completing the login could leave stdin detached from the terminal, so a TUI agent like Claude Code would start in non-interactive mode and exit with "Input must be provided". The wrapper now reopens the controlling terminal onto stdin just before handoff when the session started interactively; piped or redirected input is detected up front and left alone, so agent-mode and non-interactive use are unchanged. Also keep the `litellm-proxy` console script as an alias for `lite` so existing scripts and CI that invoke `litellm-proxy` keep working; both names map to the same CLI. * feat(install): make the curl installer need only curl, not a pre-existing Python The installer now lets uv provision a managed Python 3.13 when no suitable interpreter is found, instead of aborting. The minimum is also bumped from 3.9 to 3.10 to match the package's requires-python (>=3.10), so a system Python 3.9 is no longer selected only for uv tool install to reject it. * feat(cli): add thin litellm[cli] install path (install-cli.sh + brew) for the lite CLI On a developer laptop the `lite` CLI only needs `lite login` and running coding agents through a proxy, but the sole install path was `litellm[proxy]`, which drags in the whole server tree (fastapi, uvicorn, boto3, polars, cryptography, litellm-enterprise). The CLI's heavy imports are all guarded, so it runs on the base SDK plus just rich, pyyaml and requests. Add a `cli` extra carrying exactly those three, a `scripts/install-cli.sh` curl one-liner that installs `litellm[cli]`, and a `BerriAI/homebrew-litellm` tap formula with a release runbook under `packaging/homebrew/`. The installer passes no `--python`, so uv honours litellm's requires-python and provisions a managed interpreter, skipping a too-old (3.9) or too-new (3.14+) system Python instead of failing to resolve. A pyproject thin-contract test asserts the `cli` extra keeps the deps the CLI imports and never leaks a server-only dependency from `proxy`, so the laptop install cannot silently re-bloat * fix(install): let uv pick the Python via --python-preference system Both installers detected a system Python with a floor-only check and forced it with `uv tool install --python `. On a host whose only Python is outside litellm's requires-python (a too-old 3.9 or, increasingly, a too-new 3.14) that forced an incompatible interpreter and the resolve failed. Drop the detection and pass `--python-preference system`: uv reuses a compatible system Python when present and downloads a managed one otherwise, always honouring requires-python * test(router): filter aiohttp unclosed-session gc noise in test_async_fallbacks test_async_fallbacks asserts the last three captured log records are the router's fallback messages. Under the litellm_router_testing job (pytest -k router -n 4) many router tests share the module-level in_memory_llm_clients_cache (max 200, ttl 3600s). Older cached OpenAI/Azure clients get evicted while their aiohttp ClientSession is still open, and when the gc reclaims them aiohttp emits "Unclosed client session"/"Unclosed connector" through the asyncio logger. Those records land in caplog mid-test and push the expected router logs out of the last-three window, so the assertion flips to failing non-deterministically. These warnings are async cleanup noise, not router debug logs, so filter them out exactly like the existing leaked-task warnings before asserting order. The assertion on the three router fallback messages is unchanged. --------- Co-authored-by: Cursor Agent Co-authored-by: Mateo Wang Co-authored-by: Claude --- litellm/litellm_core_utils/cli_token_utils.py | 2 +- litellm/proxy/client/README.md | 14 +- litellm/proxy/client/cli/README.md | 135 +++-- litellm/proxy/client/cli/commands/agents.py | 303 +++++++++++ litellm/proxy/client/cli/commands/auth.py | 2 +- litellm/proxy/client/cli/commands/chat.py | 6 +- litellm/proxy/client/cli/interface.py | 7 +- litellm/proxy/client/cli/main.py | 4 + .../html_forms/cli_sso_success.py | 20 +- packaging/homebrew/README.md | 27 + packaging/homebrew/lite.rb | 33 ++ pyproject.toml | 9 + scripts/install-cli.sh | 128 +++++ scripts/install.sh | 36 +- .../test_basic_python_version.py | 50 ++ tests/local_testing/test_router_debug_logs.py | 6 +- .../proxy/client/cli/test_agents.py | 475 ++++++++++++++++++ .../proxy/client/cli/test_auth_commands.py | 2 +- .../proxy/management_endpoints/test_ui_sso.py | 17 + uv.lock | 10 +- 20 files changed, 1173 insertions(+), 113 deletions(-) create mode 100644 litellm/proxy/client/cli/commands/agents.py create mode 100644 packaging/homebrew/README.md create mode 100644 packaging/homebrew/lite.rb create mode 100755 scripts/install-cli.sh create mode 100644 tests/test_litellm/proxy/client/cli/test_agents.py diff --git a/litellm/litellm_core_utils/cli_token_utils.py b/litellm/litellm_core_utils/cli_token_utils.py index 3776d276912..eb01359cdc0 100644 --- a/litellm/litellm_core_utils/cli_token_utils.py +++ b/litellm/litellm_core_utils/cli_token_utils.py @@ -37,7 +37,7 @@ def get_litellm_gateway_api_key( """ Get the stored CLI API key for use with LiteLLM SDK. - This function reads the token file created by `litellm-proxy login` + This function reads the token file created by `lite login` and returns the API key for use in Python scripts. Args: diff --git a/litellm/proxy/client/README.md b/litellm/proxy/client/README.md index 9fbc6f2197d..c2ce28884c7 100644 --- a/litellm/proxy/client/README.md +++ b/litellm/proxy/client/README.md @@ -338,9 +338,9 @@ sequenceDiagram The CLI provides three authentication commands: -- **`litellm-proxy login`** - Start SSO authentication flow -- **`litellm-proxy logout`** - Clear stored authentication token -- **`litellm-proxy whoami`** - Show current authentication status +- **`lite login`** - Start SSO authentication flow +- **`lite logout`** - Clear stored authentication token +- **`lite whoami`** - Show current authentication status ### Authentication Flow Steps @@ -382,14 +382,14 @@ Once authenticated, the CLI will automatically use the stored token for all requ ```bash # Login -litellm-proxy login +lite login # Use CLI without specifying API key -litellm-proxy models list +lite models list # Check authentication status -litellm-proxy whoami +lite whoami # Logout -litellm-proxy logout +lite logout ``` diff --git a/litellm/proxy/client/cli/README.md b/litellm/proxy/client/cli/README.md index 6ef837cb521..333e2029e46 100644 --- a/litellm/proxy/client/cli/README.md +++ b/litellm/proxy/client/cli/README.md @@ -22,11 +22,11 @@ The CLI can be configured using environment variables or command-line options: Example: ```bash -litellm-proxy version +lite version # or -litellm-proxy --version +lite --version # or -litellm-proxy -v +lite -v ``` ## Commands @@ -40,7 +40,7 @@ The CLI provides several commands for managing models on your LiteLLM proxy serv View all available models: ```bash -litellm-proxy models list [--format table|json] +lite models list [--format table|json] ``` Options: @@ -52,7 +52,7 @@ Options: Get detailed information about all models: ```bash -litellm-proxy models info [options] +lite models info [options] ``` Options: @@ -75,7 +75,7 @@ Default columns: `public_model`, `upstream_model`, `updated_at` Add a new model to the proxy: ```bash -litellm-proxy models add [options] +lite models add [options] ``` Options: @@ -86,7 +86,7 @@ Options: Example: ```bash -litellm-proxy models add gpt-4 -p api_key=sk-123 -p api_base=https://api.openai.com -i description="GPT-4 model" +lite models add gpt-4 -p api_key=sk-123 -p api_base=https://api.openai.com -i description="GPT-4 model" ``` #### Get Model Info @@ -94,7 +94,7 @@ litellm-proxy models add gpt-4 -p api_key=sk-123 -p api_base=https://api.openai. Get information about a specific model: ```bash -litellm-proxy models get [--id MODEL_ID] [--name MODEL_NAME] +lite models get [--id MODEL_ID] [--name MODEL_NAME] ``` Options: @@ -107,7 +107,7 @@ Options: Delete a model from the proxy: ```bash -litellm-proxy models delete +lite models delete ``` #### Update Model @@ -115,7 +115,7 @@ litellm-proxy models delete Update an existing model's configuration: ```bash -litellm-proxy models update [options] +lite models update [options] ``` Options: @@ -128,7 +128,7 @@ Options: Import models from a YAML file: ```bash -litellm-proxy models import models.yaml +lite models import models.yaml ``` Options: @@ -142,31 +142,31 @@ Examples: 1. Import all models from a YAML file: ```bash -litellm-proxy models import models.yaml +lite models import models.yaml ``` 2. Dry run (show what would be imported): ```bash -litellm-proxy models import models.yaml --dry-run +lite models import models.yaml --dry-run ``` 3. Only import models where the model name contains 'gpt': ```bash -litellm-proxy models import models.yaml --only-models-matching-regex gpt +lite models import models.yaml --only-models-matching-regex gpt ``` 4. Only import models with access group containing 'beta': ```bash -litellm-proxy models import models.yaml --only-access-groups-matching-regex beta +lite models import models.yaml --only-access-groups-matching-regex beta ``` 5. Combine both filters: ```bash -litellm-proxy models import models.yaml --only-models-matching-regex gpt --only-access-groups-matching-regex beta +lite models import models.yaml --only-models-matching-regex gpt --only-access-groups-matching-regex beta ``` ### Credentials Management @@ -178,7 +178,7 @@ The CLI provides commands for managing credentials on your LiteLLM proxy server: View all available credentials: ```bash -litellm-proxy credentials list [--format table|json] +lite credentials list [--format table|json] ``` Options: @@ -194,7 +194,7 @@ The table format displays: Create a new credential: ```bash -litellm-proxy credentials create --info --values +lite credentials create --info --values ``` Options: @@ -205,7 +205,7 @@ Options: Example: ```bash -litellm-proxy credentials create azure-cred \ +lite credentials create azure-cred \ --info '{"custom_llm_provider": "azure"}' \ --values '{"api_key": "sk-123", "api_base": "https://example.azure.openai.com"}' ``` @@ -215,7 +215,7 @@ litellm-proxy credentials create azure-cred \ Get information about a specific credential: ```bash -litellm-proxy credentials get +lite credentials get ``` #### Delete Credential @@ -223,7 +223,7 @@ litellm-proxy credentials get Delete a credential: ```bash -litellm-proxy credentials delete +lite credentials delete ``` ### Keys Management @@ -235,7 +235,7 @@ The CLI provides commands for managing API keys on your LiteLLM proxy server: View all API keys: ```bash -litellm-proxy keys list [--format table|json] [options] +lite keys list [--format table|json] [options] ``` Options: @@ -256,7 +256,7 @@ Options: Generate a new API key: ```bash -litellm-proxy keys generate [options] +lite keys generate [options] ``` Options: @@ -274,7 +274,7 @@ Options: Example: ```bash -litellm-proxy keys generate --models gpt-4,gpt-3.5-turbo --spend 100 --duration 24h --key-alias my-key --team-id team123 +lite keys generate --models gpt-4,gpt-3.5-turbo --spend 100 --duration 24h --key-alias my-key --team-id team123 ``` #### Delete Keys @@ -282,7 +282,7 @@ litellm-proxy keys generate --models gpt-4,gpt-3.5-turbo --spend 100 --duration Delete API keys by key or alias: ```bash -litellm-proxy keys delete [--keys ] [--key-aliases ] +lite keys delete [--keys ] [--key-aliases ] ``` Options: @@ -293,7 +293,7 @@ Options: Example: ```bash -litellm-proxy keys delete --keys sk-key1,sk-key2 --key-aliases alias1,alias2 +lite keys delete --keys sk-key1,sk-key2 --key-aliases alias1,alias2 ``` #### Get Key Info @@ -301,7 +301,7 @@ litellm-proxy keys delete --keys sk-key1,sk-key2 --key-aliases alias1,alias2 Get information about a specific API key: ```bash -litellm-proxy keys info --key +lite keys info --key ``` Options: @@ -311,7 +311,7 @@ Options: Example: ```bash -litellm-proxy keys info --key sk-key1 +lite keys info --key sk-key1 ``` ### User Management @@ -323,7 +323,7 @@ The CLI provides commands for managing users on your LiteLLM proxy server: View all users: ```bash -litellm-proxy users list +lite users list ``` #### Get User Info @@ -331,7 +331,7 @@ litellm-proxy users list Get information about a specific user: ```bash -litellm-proxy users get --id +lite users get --id ``` #### Create User @@ -339,7 +339,7 @@ litellm-proxy users get --id Create a new user: ```bash -litellm-proxy users create --email user@example.com --role internal_user --alias "Alice" --team team1 --max-budget 100.0 +lite users create --email user@example.com --role internal_user --alias "Alice" --team team1 --max-budget 100.0 ``` #### Delete User @@ -347,7 +347,7 @@ litellm-proxy users create --email user@example.com --role internal_user --alias Delete one or more users by user_id: ```bash -litellm-proxy users delete +lite users delete ``` ### Chat Commands @@ -359,7 +359,7 @@ The CLI provides commands for interacting with chat models through your LiteLLM Create a chat completion: ```bash -litellm-proxy chat completions [options] +lite chat completions [options] ``` Arguments: @@ -379,12 +379,12 @@ Examples: 1. Simple completion: ```bash -litellm-proxy chat completions gpt-4 -m "user:Hello, how are you?" +lite chat completions gpt-4 -m "user:Hello, how are you?" ``` 2. Multi-message conversation: ```bash -litellm-proxy chat completions gpt-4 \ +lite chat completions gpt-4 \ -m "system:You are a helpful assistant" \ -m "user:What's the capital of France?" \ -m "assistant:The capital of France is Paris." \ @@ -393,7 +393,7 @@ litellm-proxy chat completions gpt-4 \ 3. With generation parameters: ```bash -litellm-proxy chat completions gpt-4 \ +lite chat completions gpt-4 \ -m "user:Write a story" \ --temperature 0.7 \ --max-tokens 500 \ @@ -409,7 +409,7 @@ The CLI provides commands for making direct HTTP requests to your LiteLLM proxy Make an HTTP request to any endpoint: ```bash -litellm-proxy http request [options] +lite http request [options] ``` Arguments: @@ -425,19 +425,46 @@ Examples: 1. List models: ```bash -litellm-proxy http request GET /models +lite http request GET /models ``` 2. Create a chat completion: ```bash -litellm-proxy http request POST /chat/completions -j '{"model": "gpt-4", "messages": [{"role": "user", "content": "Hello"}]}' +lite http request POST /chat/completions -j '{"model": "gpt-4", "messages": [{"role": "user", "content": "Hello"}]}' ``` 3. Test connection with custom headers: ```bash -litellm-proxy http request GET /health/test_connection -H "X-Custom-Header:value" +lite http request GET /health/test_connection -H "X-Custom-Header:value" ``` +### Run a Coding Agent + +Launch a coding agent with all of its LLM traffic routed through your LiteLLM proxy. Each supported agent is its own command, so there is nothing to remember beyond the agent's name: + +```bash +lite claude +lite codex +lite opencode +``` + +Anything you type after the agent name is forwarded to it untouched, so the usual flags keep working: + +```bash +lite claude --resume +lite codex exec "summarize the repo" +``` + +Each command resolves your LiteLLM key (logging in via SSO when none is stored and you are at a terminal; otherwise it expects `LITELLM_PROXY_API_KEY` or `--api-key`), checks the key against the proxy so bad credentials fail immediately instead of deep inside the agent, exports the environment variables the agent reads, then replaces itself with the agent process. + +The right variables are picked per agent. Claude Code gets `ANTHROPIC_BASE_URL` (the proxy root, so it appends `/v1/messages`) and `ANTHROPIC_AUTH_TOKEN`, with any stray `ANTHROPIC_API_KEY` cleared so the proxy token wins. Codex and OpenCode get `OPENAI_BASE_URL` (the proxy plus `/v1`) and `OPENAI_API_KEY`. Codex ignores `OPENAI_BASE_URL`, so it is additionally pointed at the proxy through a custom provider passed as `-c` config overrides (HTTP/SSE Responses transport, since the proxy does not speak the Responses WebSocket protocol). + +Options (these belong to the wrapper, so put them before the agent's own flags): + +- `--skip-verify`: Skip the pre-launch key check (useful offline or with non-standard auth). + +To pin the model, pass the agent's own model flag (for example `lite claude --model my-proxy-model` or `lite codex -m my-proxy-model`), or export the variable the agent reads (`ANTHROPIC_MODEL` / `ANTHROPIC_SMALL_FAST_MODEL` for Claude Code); the wrapper preserves anything you already have set. Whatever model the agent ends up requesting must exist on the proxy, since requests land on the proxy's `/v1/messages` (Anthropic) or `/v1/chat/completions` and `/v1/responses` (OpenAI) endpoints. + ## Environment Variables The CLI respects the following environment variables: @@ -450,37 +477,37 @@ The CLI respects the following environment variables: 1. List all models in table format: ```bash -litellm-proxy models list +lite models list ``` 2. Add a new model with parameters: ```bash -litellm-proxy models add gpt-4 -p api_key=sk-123 -p max_tokens=2048 +lite models add gpt-4 -p api_key=sk-123 -p max_tokens=2048 ``` 3. Get model information in JSON format: ```bash -litellm-proxy models info --format json +lite models info --format json ``` 4. Update model parameters: ```bash -litellm-proxy models update model-123 -p temperature=0.7 -i description="Updated model" +lite models update model-123 -p temperature=0.7 -i description="Updated model" ``` 5. List all credentials in table format: ```bash -litellm-proxy credentials list +lite credentials list ``` 6. Create a new credential for Azure: ```bash -litellm-proxy credentials create azure-prod \ +lite credentials create azure-prod \ --info '{"custom_llm_provider": "azure"}' \ --values '{"api_key": "sk-123", "api_base": "https://prod.azure.openai.com"}' ``` @@ -488,7 +515,7 @@ litellm-proxy credentials create azure-prod \ 7. Make a custom HTTP request: ```bash -litellm-proxy http request POST /chat/completions \ +lite http request POST /chat/completions \ -j '{"model": "gpt-4", "messages": [{"role": "user", "content": "Hello"}]}' \ -H "X-Custom-Header:value" ``` @@ -497,29 +524,29 @@ litellm-proxy http request POST /chat/completions \ ```bash # List users -litellm-proxy users list +lite users list # Get user info -litellm-proxy users get --id u1 +lite users get --id u1 # Create a user -litellm-proxy users create --email a@b.com --role internal_user --alias "Alice" --team team1 --max-budget 100.0 +lite users create --email a@b.com --role internal_user --alias "Alice" --team team1 --max-budget 100.0 # Delete users -litellm-proxy users delete u1 u2 +lite users delete u1 u2 ``` 9. Import models from a YAML file (with filters): ```bash # Only import models where the model name contains 'gpt' -litellm-proxy models import models.yaml --only-models-matching-regex gpt +lite models import models.yaml --only-models-matching-regex gpt # Only import models with access group containing 'beta' -litellm-proxy models import models.yaml --only-access-groups-matching-regex beta +lite models import models.yaml --only-access-groups-matching-regex beta # Combine both filters -litellm-proxy models import models.yaml --only-models-matching-regex gpt --only-access-groups-matching-regex beta +lite models import models.yaml --only-models-matching-regex gpt --only-access-groups-matching-regex beta ``` ## Error Handling diff --git a/litellm/proxy/client/cli/commands/agents.py b/litellm/proxy/client/cli/commands/agents.py new file mode 100644 index 00000000000..f39ffb3e864 --- /dev/null +++ b/litellm/proxy/client/cli/commands/agents.py @@ -0,0 +1,303 @@ +import os +import shutil +import sys +from typing import Callable, Dict, FrozenSet, List, Mapping, Optional, Sequence, Tuple + +import click +import requests + +from .auth import get_stored_api_key, login + +ANTHROPIC_BASE_URL_ENV = "ANTHROPIC_BASE_URL" +ANTHROPIC_AUTH_TOKEN_ENV = "ANTHROPIC_AUTH_TOKEN" +ANTHROPIC_API_KEY_ENV = "ANTHROPIC_API_KEY" +OPENAI_BASE_URL_ENV = "OPENAI_BASE_URL" +OPENAI_API_KEY_ENV = "OPENAI_API_KEY" + +PROFILE_ANTHROPIC = "anthropic" +PROFILE_OPENAI = "openai" + +_KNOWN_AGENTS: Dict[str, Tuple[str, FrozenSet[str]]] = { + "claude": ("Claude Code", frozenset({PROFILE_ANTHROPIC})), + "codex": ("Codex", frozenset({PROFILE_OPENAI})), + "opencode": ("OpenCode", frozenset({PROFILE_OPENAI})), +} + +_INSTALL_DOCS: Dict[str, str] = { + "claude": "https://docs.claude.com/en/docs/claude-code/setup", + "codex": "https://developers.openai.com/codex/cli", + "opencode": "https://opencode.ai/docs", +} + +CODEX_PROXY_PROVIDER = "litellm" + + +class AgentRunError(Exception): + """Raised for any user-actionable failure while preparing to run an agent.""" + + +def agent_profile(command: str) -> Tuple[str, FrozenSet[str]]: + """Return the (display name, env profiles) for a wrapped command. + + Known agents map to the API family they speak. Anything else gets both + families so it works regardless of which env vars the tool reads. + """ + base = os.path.basename(command) + if base in _KNOWN_AGENTS: + return _KNOWN_AGENTS[base] + return base, frozenset({PROFILE_ANTHROPIC, PROFILE_OPENAI}) + + +def build_agent_env( + base_env: Mapping[str, str], + base_url: str, + api_key: str, + profiles: FrozenSet[str], +) -> Dict[str, str]: + """Return a copy of base_env wired to route the agent through the proxy. + + Anthropic clients (Claude Code) append /v1/messages to ANTHROPIC_BASE_URL, + so it stays the bare proxy root; OpenAI clients (Codex, OpenCode) expect the + /v1 suffix on OPENAI_BASE_URL. ANTHROPIC_API_KEY is dropped so a stray + Anthropic key cannot win over the bearer token we set. + """ + env = dict(base_env) + root = base_url.rstrip("/") + if PROFILE_ANTHROPIC in profiles: + env[ANTHROPIC_BASE_URL_ENV] = root + env[ANTHROPIC_AUTH_TOKEN_ENV] = api_key + env.pop(ANTHROPIC_API_KEY_ENV, None) + if PROFILE_OPENAI in profiles: + env[OPENAI_BASE_URL_ENV] = root + "/v1" + env[OPENAI_API_KEY_ENV] = api_key + return env + + +def _codex_proxy_args(base_url: str) -> List[str]: + """Codex `-c` overrides that point it at the proxy. + + Codex ignores OPENAI_BASE_URL (it always dials api.openai.com), so the env + profile alone cannot route it. It does honor a custom provider, so define one + inline; supports_websockets=false forces the HTTP/SSE Responses transport + because the proxy does not speak the Responses WebSocket protocol. The key is + read from OPENAI_API_KEY, which build_agent_env already exports. + """ + root = base_url.rstrip("/") + "/v1" + provider = f"model_providers.{CODEX_PROXY_PROVIDER}" + return [ + "-c", + f'model_provider="{CODEX_PROXY_PROVIDER}"', + "-c", + f'{provider}.name="LiteLLM proxy"', + "-c", + f'{provider}.base_url="{root}"', + "-c", + f'{provider}.env_key="{OPENAI_API_KEY_ENV}"', + "-c", + f'{provider}.wire_api="responses"', + "-c", + f"{provider}.supports_websockets=false", + ] + + +_PROXY_ARGS: Dict[str, Callable[[str], List[str]]] = { + "codex": _codex_proxy_args, +} + + +def agent_launch_args(command: str, base_url: str) -> List[str]: + """Extra CLI args an agent needs to actually honor the proxy. + + Claude Code and OpenCode respect the exported env vars, so they get nothing + here; Codex needs its provider pointed via config overrides. + """ + builder = _PROXY_ARGS.get(os.path.basename(command)) + return builder(base_url) if builder else [] + + +def verify_proxy_key( + base_url: str, + api_key: str, + *, + get: Callable[..., requests.Response] = requests.get, +) -> None: + """Probe the proxy with the key so bad creds fail here, not inside the agent. + + Raises AgentRunError when the proxy is unreachable or rejects the key. Other + non-2xx responses are tolerated; the agent's own call is the real test. + """ + url = base_url.rstrip("/") + "/v1/models" + try: + resp = get(url, headers={"Authorization": f"Bearer {api_key}"}, timeout=10) + except requests.RequestException as e: + raise AgentRunError( + f"Could not reach the LiteLLM proxy at {base_url.rstrip('/')}: {e}. " + "Is it running, and is --base-url (or LITELLM_PROXY_URL) correct?" + ) + if resp.status_code in (401, 403): + raise AgentRunError( + f"LiteLLM rejected your key (HTTP {resp.status_code}). " + "Run `lite login` to refresh it, or pass a valid --api-key." + ) + + +def _exec(path: str, args: Sequence[str], env: Mapping[str, str]) -> None: + os.execvpe(path, list(args), dict(env)) + + +def _restore_controlling_terminal() -> None: + """Reattach the controlling terminal to stdin before handing off to the agent. + + Completing the browser SSO login can leave stdin detached from the terminal, + which makes a TUI agent like Claude Code start in non-interactive mode and + exit immediately. Reopening /dev/tty onto fd 0 gives the agent a live + terminal; when stdin is still a tty (no login happened) this is a no-op. + """ + if sys.stdin.isatty(): + return + try: + fd = os.open("/dev/tty", os.O_RDONLY) + except OSError: + return + try: + os.dup2(fd, 0) + finally: + os.close(fd) + + +def run_agent( + base_url: str, + api_key: str, + command: Sequence[str], + *, + skip_verify: bool = False, + base_env: Optional[Mapping[str, str]] = None, + which: Callable[[str], Optional[str]] = shutil.which, + verify: Callable[[str, str], None] = verify_proxy_key, + launcher: Callable[[str, Sequence[str], Mapping[str, str]], None] = _exec, + reattach_terminal: Optional[Callable[[], None]] = None, +) -> None: + """Validate, wire the environment, and hand off to the agent. + + On success this replaces the current process and never returns. Raises + AgentRunError for missing binaries, an unreachable proxy, or a rejected key. + reattach_terminal, when given, runs just before handoff to restore stdin. + """ + if not command: + raise AgentRunError("Nothing to run.") + + _, profiles = agent_profile(command[0]) + binary = which(command[0]) + if binary is None: + docs = _INSTALL_DOCS.get(os.path.basename(command[0])) + hint = f" Install it first: {docs}" if docs else "" + raise AgentRunError(f"Could not find `{command[0]}` on your PATH.{hint}") + + if not skip_verify: + verify(base_url, api_key) + + env = build_agent_env( + base_env if base_env is not None else os.environ, + base_url, + api_key, + profiles, + ) + extra_args = agent_launch_args(command[0], base_url) + if reattach_terminal is not None: + reattach_terminal() + launcher(binary, [command[0], *extra_args, *command[1:]], env) + + +def _is_interactive() -> bool: + return sys.stdin.isatty() + + +def _resolve_api_key(ctx: click.Context) -> str: + base_url = ctx.obj["base_url"] + api_key = ctx.obj.get("api_key") + if api_key: + return api_key + + if not _is_interactive(): + raise click.ClickException( + "No LiteLLM key found. Set LITELLM_PROXY_API_KEY (or pass --api-key) for " + "non-interactive use, or run `lite login` from a terminal." + ) + + click.echo("No LiteLLM credentials found; starting login...") + ctx.invoke(login) + api_key = get_stored_api_key(expected_base_url=base_url) + if not api_key: + raise click.ClickException( + "Login did not produce an API key; cannot start the agent." + ) + return api_key + + +_SKIP_VERIFY_HELP = "Skip the pre-launch key check against the proxy." + + +def _launch( + ctx: click.Context, binary: str, args: Sequence[str], *, skip_verify: bool +) -> None: + base_url = ctx.obj["base_url"] + started_interactive = _is_interactive() + api_key = _resolve_api_key(ctx) + + display_name, _ = agent_profile(binary) + click.echo( + f"litellm: routing {display_name} through proxy at {base_url.rstrip('/')}" + ) + + try: + run_agent( + base_url, + api_key, + [binary, *args], + skip_verify=skip_verify, + reattach_terminal=( + _restore_controlling_terminal if started_interactive else None + ), + ) + except AgentRunError as e: + raise click.ClickException(str(e)) + + +def _make_agent_command(binary: str, display_name: str) -> click.Command: + @click.command( + name=binary, + context_settings={"ignore_unknown_options": True}, + short_help=f"Run {display_name} through your LiteLLM proxy", + ) + @click.option("--skip-verify", is_flag=True, default=False, help=_SKIP_VERIFY_HELP) + @click.argument("args", nargs=-1, type=click.UNPROCESSED) + @click.pass_context + def _command(ctx: click.Context, skip_verify: bool, args: Sequence[str]) -> None: + _launch(ctx, binary, list(args), skip_verify=skip_verify) + + _command.help = ( + f"Run {display_name} routed through your LiteLLM proxy.\n\n" + f"Logs in with LiteLLM if needed, verifies your key against the proxy, " + f"exports the env vars {binary} reads, then hands off. Any arguments are " + f"forwarded to `{binary}`." + ) + return _command + + +def agent_commands() -> List[click.Command]: + """Build one top-level command per known agent, e.g. `lite claude`.""" + return [ + _make_agent_command(binary, name) + for binary, (name, _profiles) in _KNOWN_AGENTS.items() + ] + + +__all__ = [ + "agent_commands", + "run_agent", + "build_agent_env", + "agent_launch_args", + "verify_proxy_key", + "agent_profile", + "AgentRunError", +] diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index 447837c35e7..b06d86d5965 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -624,7 +624,7 @@ def whoami(): token_data = load_token() if not token_data: - click.echo("❌ Not authenticated. Run 'litellm-proxy login' to authenticate.") + click.echo("❌ Not authenticated. Run 'lite login' to authenticate.") return click.echo("✅ Authenticated") diff --git a/litellm/proxy/client/cli/commands/chat.py b/litellm/proxy/client/cli/commands/chat.py index a078b766107..696e34c3ecd 100644 --- a/litellm/proxy/client/cli/commands/chat.py +++ b/litellm/proxy/client/cli/commands/chat.py @@ -122,13 +122,13 @@ def chat( Examples: # Chat with a specific model - litellm-proxy chat gpt-4 + lite chat gpt-4 # Chat without specifying model (will show model selection) - litellm-proxy chat + lite chat # Chat with custom settings - litellm-proxy chat gpt-4 --temperature 0.9 --system "You are a helpful coding assistant" + lite chat gpt-4 --temperature 0.9 --system "You are a helpful coding assistant" """ console = Console() diff --git a/litellm/proxy/client/cli/interface.py b/litellm/proxy/client/cli/interface.py index eba693dc18e..a32d60aadd9 100644 --- a/litellm/proxy/client/cli/interface.py +++ b/litellm/proxy/client/cli/interface.py @@ -80,6 +80,8 @@ def styled_prompt(): def show_commands(): """Display available commands.""" + from .commands.agents import agent_commands + commands = [ ("login", "Authenticate with the LiteLLM proxy server"), ("logout", "Clear stored authentication"), @@ -91,6 +93,9 @@ def show_commands(): ("keys", "Manage API keys"), ("teams", "Manage teams and team assignments"), ("users", "Manage users"), + ] + commands += [(c.name, c.get_short_help_str()) for c in agent_commands()] + commands += [ ("version", "Show version information"), ("help", "Show this help message"), ("quit", "Exit the interactive session"), @@ -156,7 +161,7 @@ def execute_command(user_input: str, ctx: click.Context): # Execute the command try: # Create a new argument list for click to parse - sys.argv = ["litellm-proxy"] + [command] + args + sys.argv = ["lite"] + [command] + args # Get the command object and invoke it cmd = cli.commands[command] diff --git a/litellm/proxy/client/cli/main.py b/litellm/proxy/client/cli/main.py index be55f79c066..b8c483f4b08 100644 --- a/litellm/proxy/client/cli/main.py +++ b/litellm/proxy/client/cli/main.py @@ -7,6 +7,7 @@ import click from litellm._version import version as litellm_version from litellm.proxy.client.health import HealthManagementClient +from .commands.agents import agent_commands from .commands.auth import get_stored_api_key, login, logout, whoami from .commands.chat import chat from .commands.credentials import credentials @@ -112,6 +113,9 @@ cli.add_command(keys) cli.add_command(teams) # Add the users command group cli.add_command(users) +# Add a top-level command per coding agent (claude, codex, opencode, ...) +for agent_command in agent_commands(): + cli.add_command(agent_command) if __name__ == "__main__": diff --git a/litellm/proxy/common_utils/html_forms/cli_sso_success.py b/litellm/proxy/common_utils/html_forms/cli_sso_success.py index 51f0775d90b..345f3ca5b42 100644 --- a/litellm/proxy/common_utils/html_forms/cli_sso_success.py +++ b/litellm/proxy/common_utils/html_forms/cli_sso_success.py @@ -135,7 +135,7 @@ def render_cli_sso_success_page() -> str: font-size: 14px; }} - .countdown {{ + .status {{ color: #64748b; font-size: 14px; font-weight: 500; @@ -183,23 +183,11 @@ def render_cli_sso_success_page() -> str:

You can now use LiteLLM CLI commands with your authenticated session.

-
This window will close in 3 seconds...
+
You can now close this window and return to your terminal.
- + diff --git a/packaging/homebrew/README.md b/packaging/homebrew/README.md new file mode 100644 index 00000000000..ef441ded304 --- /dev/null +++ b/packaging/homebrew/README.md @@ -0,0 +1,27 @@ +# Homebrew formula for the `lite` CLI + +[`lite.rb`](./lite.rb) is the canonical source for the Homebrew formula that installs the thin LiteLLM CLI (`litellm[cli]`). It lives here so it is versioned with the code, but Homebrew serves formulae from a tap, so it has to be published to the `BerriAI/homebrew-litellm` tap to be installable. + +Once published, end users install with + +```shell +brew install BerriAI/litellm/lite +``` + +which gives them the `lite` command (`lite login`, `lite claude`, `lite models list`, ...) without the proxy server runtime. For the full proxy server, they keep using pip/uv with `litellm[proxy]` or the Docker image. + +## Why a tap and not homebrew-core + +The formula builds the published `litellm` sdist with the `cli` extra and resolves that extra's dependencies from PyPI at build time. homebrew-core forbids network access during `install` and would require every transitive dependency declared as a pinned `resource`, regenerated on each release. For a fast-moving CLI that tradeoff is not worth it, so this stays a tap formula. + +## Release runbook + +The formula can only point at a published artifact, so it activates with the first `litellm` release that ships the `cli` extra (added in [pyproject.toml](../../pyproject.toml)). + +1. Cut a `litellm` release whose `pyproject.toml` includes the `cli` extra and confirm it is on PyPI. +2. Fetch the sdist URL and checksum for that version: `curl -fsSL https://pypi.org/pypi/litellm//json | jq -r '.urls[] | select(.packagetype=="sdist") | "\(.url)\n\(.digests.sha256)"'` +3. Set `url` and `sha256` in `lite.rb` to those values; `version` is parsed from `url`. +4. Copy `lite.rb` into the tap repo under `Formula/lite.rb`, then run `brew install --build-from-source ./Formula/lite.rb` and `brew test lite` to verify a clean build and that `lite --help` works. +5. Commit and push to `BerriAI/homebrew-litellm`. + +Keep `lite.rb` here in sync with the tap copy so the in-repo formula stays the source of truth. diff --git a/packaging/homebrew/lite.rb b/packaging/homebrew/lite.rb new file mode 100644 index 00000000000..d0d61bb5b43 --- /dev/null +++ b/packaging/homebrew/lite.rb @@ -0,0 +1,33 @@ +# Homebrew formula for the thin LiteLLM `lite` CLI (litellm[cli]). +# +# Ships in the BerriAI/homebrew-litellm tap, not homebrew-core: it builds the +# published litellm sdist with the `cli` extra into a dedicated virtualenv and +# pulls the extra's deps from PyPI. That is the low-maintenance path for a +# fast-moving Python CLI; the resource-stanza alternative would need every +# transitive dep re-pinned with a fresh sha256 on each release. +# +# RELEASE STEP (see README.md in this directory): point `url` + `sha256` at the +# PyPI sdist of the first litellm version that ships the `cli` extra. `version` +# is parsed from `url`, and the build installs exactly that version, so the three +# stay in lockstep automatically. +class Lite < Formula + include Language::Python::Virtualenv + + desc "Thin client for the LiteLLM proxy: lite login, lite claude/codex/opencode" + homepage "https://docs.litellm.ai/docs/proxy/management_cli" + url "https://files.pythonhosted.org/packages/source/l/litellm/litellm-REPLACE_AT_RELEASE.tar.gz" + sha256 "REPLACE_AT_RELEASE" + license "MIT" + + depends_on "python@3.13" + + def install + virtualenv_create(libexec, "python3.13") + system libexec/"bin/pip", "install", "#{buildpath}[cli]" + bin.install_symlink libexec/"bin/lite" + end + + test do + assert_match "login", shell_output("#{bin}/lite --help") + end +end diff --git a/pyproject.toml b/pyproject.toml index 28e6f48dc4c..b9d76379faf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -71,6 +71,14 @@ proxy = [ "pyroscope-io>=0.8.16,<1.0; sys_platform != 'win32'", "pydantic-settings>=2.14.1,<3.0", ] +# Thin client install for the `lite` CLI on developer laptops. The CLI's heavy +# imports (fastapi, cryptography, ...) are all guarded, so it runs on the base +# SDK plus just these three; none of the server runtime in `proxy` is pulled in. +cli = [ + "rich>=13.9.4,<14.0", + "pyyaml>=6.0.3,<7.0", + "requests>=2.32.0,<3.0", +] extra_proxy = [ "prisma>=0.11.0,<1.0", "azure-identity>=1.25.2,<2.0", @@ -132,6 +140,7 @@ proxy-runtime = [ [project.scripts] litellm = "litellm:run_server" +lite = "litellm.proxy.client.cli:cli" litellm-proxy = "litellm.proxy.client.cli:cli" [dependency-groups] diff --git a/scripts/install-cli.sh b/scripts/install-cli.sh new file mode 100755 index 00000000000..d147286fcac --- /dev/null +++ b/scripts/install-cli.sh @@ -0,0 +1,128 @@ +#!/usr/bin/env bash +# LiteLLM CLI Installer (the thin `lite` client) +# Usage: curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm/main/scripts/install-cli.sh | sh +# +# Installs only litellm[cli]: the `lite` command for authenticating to a LiteLLM +# proxy and running coding agents (lite claude / codex / opencode) through it. +# None of the proxy server runtime is pulled in. To run a proxy server instead, +# use scripts/install.sh, which installs litellm[proxy]. +# +# Needs only curl: uv is bootstrapped if missing, and uv provisions a compatible +# Python itself (honouring litellm's requires-python), downloading a managed one +# when the host has no suitable interpreter. +# +# NOTE: set -e without pipefail for POSIX sh compatibility (dash on Ubuntu/Debian +# ignores the shebang when invoked as `sh` and does not support `pipefail`). +set -eu + +# NOTE: before merging, this must stay as "litellm[cli]" to install from PyPI. +LITELLM_PACKAGE="litellm[cli]" +UV_VERSION="0.10.9" + +# ── colours ──────────────────────────────────────────────────────────────── +if [ -t 1 ]; then + BOLD='\033[1m' + GREEN='\033[38;2;78;186;101m' + GREY='\033[38;2;153;153;153m' + RESET='\033[0m' +else + BOLD='' GREEN='' GREY='' RESET='' +fi + +info() { printf "${GREY} %s${RESET}\n" "$*"; } +success() { printf "${GREEN} ✔ %s${RESET}\n" "$*"; } +header() { printf "${BOLD} %s${RESET}\n" "$*"; } +die() { printf "\n Error: %s\n\n" "$*" >&2; exit 1; } + +# ── banner ───────────────────────────────────────────────────────────────── +echo "" +cat << 'EOF' + ██╗ ██╗████████╗███████╗ + ██║ ██║╚══██╔══╝██╔════╝ + ██║ ██║ ██║ █████╗ + ██║ ██║ ██║ ██╔══╝ + ███████╗██║ ██║ ███████╗ + ╚══════╝╚═╝ ╚═╝ ╚══════╝ +EOF +printf " ${BOLD}LiteLLM CLI Installer${RESET} ${GREY}the thin 'lite' client for your proxy${RESET}\n\n" + +# ── OS detection ─────────────────────────────────────────────────────────── +OS="$(uname -s)" +ARCH="$(uname -m)" + +case "$OS" in + Darwin) PLATFORM="macOS ($ARCH)" ;; + Linux) PLATFORM="Linux ($ARCH)" ;; + *) die "Unsupported OS: $OS. LiteLLM supports macOS and Linux." ;; +esac + +info "Platform: $PLATFORM" + +# ── uv detection / install ──────────────────────────────────────────────── +UV_BIN="" +CURRENT_UV_VERSION="" +for candidate in uv "$HOME/.local/bin/uv"; do + if command -v "$candidate" >/dev/null 2>&1; then + UV_BIN="$(command -v "$candidate")" + break + elif [ -x "$candidate" ]; then + UV_BIN="$candidate" + break + fi +done + +if [ -n "$UV_BIN" ]; then + CURRENT_UV_VERSION="$("$UV_BIN" --version 2>/dev/null | awk '{print $2}' | head -1 || true)" +fi + +if [ -z "$UV_BIN" ] || [ "${CURRENT_UV_VERSION:-}" != "$UV_VERSION" ]; then + header "Installing uv…" + if [ -n "${CURRENT_UV_VERSION:-}" ]; then + info "Upgrading uv from ${CURRENT_UV_VERSION} to ${UV_VERSION}" + fi + curl -LsSf "https://astral.sh/uv/${UV_VERSION}/install.sh" | env UV_NO_MODIFY_PATH=1 sh \ + || die "uv installation failed. Try manually: curl -LsSf https://astral.sh/uv/${UV_VERSION}/install.sh | sh" + UV_BIN="$HOME/.local/bin/uv" +fi + +# ── install ──────────────────────────────────────────────────────────────── +# --python-preference system: reuse a compatible system Python when present, +# otherwise download a managed one. Either way uv honours litellm's requires-python, +# so a too-old (3.9) or too-new (3.14+) system Python is skipped, not forced. +echo "" +header "Installing litellm[cli]…" +echo "" + +"$UV_BIN" tool install --python-preference system --force "${LITELLM_PACKAGE}" \ + || die "uv tool install failed. Try manually: $UV_BIN tool install '${LITELLM_PACKAGE}'" + +# ── find the lite binary installed by uv tool ────────────────────────────── +SCRIPTS_DIR="$("$UV_BIN" tool dir --bin)" +LITE_BIN="${SCRIPTS_DIR}/lite" + +if [ ! -x "$LITE_BIN" ]; then + die "lite binary not found after install. Try: $UV_BIN tool install '${LITELLM_PACKAGE}'" +fi + +# ── success banner ───────────────────────────────────────────────────────── +echo "" +success "LiteLLM CLI installed" + +installed_ver="$("$LITE_BIN" --version 2>&1 | grep -oE '[0-9]+\.[0-9]+\.[0-9]+' | head -1 || true)" +[ -n "$installed_ver" ] && info "Version: $installed_ver" + +# ── PATH hint ────────────────────────────────────────────────────────────── +if ! command -v lite >/dev/null 2>&1; then + info "Note: add lite to your PATH: export PATH=\"\$PATH:${SCRIPTS_DIR}\"" +fi + +# ── next steps ───────────────────────────────────────────────────────────── +echo "" +header "Next steps:" +echo "" +info " export LITELLM_PROXY_URL=https://your-proxy # point at your gateway" +info " lite login # authenticate via SSO" +info " lite claude # run Claude Code through the proxy" +echo "" +info "Docs: https://docs.litellm.ai/docs/proxy/management_cli" +echo "" diff --git a/scripts/install.sh b/scripts/install.sh index c28d7da872f..06e6249c9ba 100755 --- a/scripts/install.sh +++ b/scripts/install.sh @@ -2,13 +2,13 @@ # LiteLLM Installer # Usage: curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm/main/scripts/install.sh | sh # +# Needs only curl: uv is bootstrapped if missing, and uv provisions a compatible +# Python itself (reusing a suitable system one, else downloading a managed build). +# # NOTE: set -e without pipefail for POSIX sh compatibility (dash on Ubuntu/Debian # ignores the shebang when invoked as `sh` and does not support `pipefail`). set -eu -MIN_PYTHON_MAJOR=3 -MIN_PYTHON_MINOR=9 - # NOTE: before merging, this must stay as "litellm[proxy]" to install from PyPI. LITELLM_PACKAGE="litellm[proxy]" UV_VERSION="0.10.9" @@ -52,27 +52,6 @@ esac info "Platform: $PLATFORM" -# ── Python detection ─────────────────────────────────────────────────────── -PYTHON_BIN="" -for candidate in python3 python; do - if command -v "$candidate" >/dev/null 2>&1; then - major="$("$candidate" -c 'import sys; print(sys.version_info.major)' 2>/dev/null || true)" - minor="$("$candidate" -c 'import sys; print(sys.version_info.minor)' 2>/dev/null || true)" - if [ "${major:-0}" -ge "$MIN_PYTHON_MAJOR" ] && [ "${minor:-0}" -ge "$MIN_PYTHON_MINOR" ]; then - PYTHON_BIN="$(command -v "$candidate")" - info "Python: $("$candidate" --version 2>&1)" - break - fi - fi -done - -if [ -z "$PYTHON_BIN" ]; then - die "Python ${MIN_PYTHON_MAJOR}.${MIN_PYTHON_MINOR}+ is required but not found. - Install it from https://python.org/downloads or via your package manager: - macOS: brew install python@3 - Ubuntu: sudo apt install python3" -fi - # ── uv detection / install ──────────────────────────────────────────────── UV_BIN="" CURRENT_UV_VERSION="" @@ -105,15 +84,18 @@ echo "" header "Installing litellm[proxy]…" echo "" -"$UV_BIN" tool install --python "$PYTHON_BIN" --force "${LITELLM_PACKAGE}" \ - || die "uv tool install failed. Try manually: $UV_BIN tool install --python '$PYTHON_BIN' '${LITELLM_PACKAGE}'" +# --python-preference system: reuse a compatible system Python when present, +# otherwise download a managed one. Either way uv honours litellm's requires-python, +# so a too-old (3.9) or too-new (3.14+) system Python is skipped, not forced. +"$UV_BIN" tool install --python-preference system --force "${LITELLM_PACKAGE}" \ + || die "uv tool install failed. Try manually: $UV_BIN tool install '${LITELLM_PACKAGE}'" # ── find the litellm binary installed by uv tool ─────────────────────────── SCRIPTS_DIR="$("$UV_BIN" tool dir --bin)" LITELLM_BIN="${SCRIPTS_DIR}/litellm" if [ ! -x "$LITELLM_BIN" ]; then - die "litellm binary not found after install. Try: $UV_BIN tool install --python '$PYTHON_BIN' '${LITELLM_PACKAGE}'" + die "litellm binary not found after install. Try: $UV_BIN tool install '${LITELLM_PACKAGE}'" fi # ── success banner ───────────────────────────────────────────────────────── diff --git a/tests/local_testing/test_basic_python_version.py b/tests/local_testing/test_basic_python_version.py index 8308e0d6033..e31c3953714 100644 --- a/tests/local_testing/test_basic_python_version.py +++ b/tests/local_testing/test_basic_python_version.py @@ -92,6 +92,56 @@ def test_package_dependencies(): ) +def test_cli_extra_is_a_thin_client_install(): + """The `cli` extra must install a working `lite` client without dragging in the + proxy server runtime. It therefore has to declare the CLI's real third-party + deps (rich, pyyaml, requests) and must never contain a server-only dependency + from the `proxy` extra; a leak there silently re-bloats the laptop install. + """ + import pathlib + + import litellm + from packaging.requirements import Requirement + + try: + import tomllib as tomli + except ImportError: + try: + import tomli + except ImportError: + pytest.skip("tomli/tomllib not available - skipping dependency check") + + pyproject_path = pathlib.Path(litellm.__file__).parent.parent / "pyproject.toml" + with open(pyproject_path, "rb") as f: + optional_deps = tomli.load(f)["project"]["optional-dependencies"] + + assert "cli" in optional_deps, "Expected a `cli` extra for the thin lite install" + + cli_names = {Requirement(req).name.lower() for req in optional_deps["cli"]} + + missing = {"rich", "pyyaml", "requests"} - cli_names + assert not missing, f"`cli` extra is missing deps the lite CLI imports: {missing}" + + server_only = { + "fastapi", + "uvicorn", + "gunicorn", + "granian", + "starlette", + "boto3", + "polars", + "soundfile", + "mcp", + "cryptography", + "apscheduler", + "rq", + "litellm-enterprise", + "litellm-proxy-extras", + } + leaked = cli_names & server_only + assert not leaked, f"`cli` extra leaks proxy-server deps onto laptops: {leaked}" + + import os import subprocess import time diff --git a/tests/local_testing/test_router_debug_logs.py b/tests/local_testing/test_router_debug_logs.py index f1c7e9d722e..ad807539bf2 100644 --- a/tests/local_testing/test_router_debug_logs.py +++ b/tests/local_testing/test_router_debug_logs.py @@ -82,7 +82,9 @@ def test_async_fallbacks(caplog): asyncio.run(_make_request()) captured_logs = [rec.message for rec in caplog.records] - # on circle ci the captured logs get some async task exception logs - filter them out "Task exception was never retrieved" + # on circle ci the captured logs get async cleanup noise from the gc (leaked + # task warnings, plus aiohttp "Unclosed client session"/"Unclosed connector" + # warnings from cached clients other router tests evicted) - filter it out captured_logs = [ log for log in captured_logs @@ -90,6 +92,8 @@ def test_async_fallbacks(caplog): and "Task was destroyed but it is pending" not in log and "get_available_deployment" not in log and "in the Langfuse queue" not in log + and "Unclosed client session" not in log + and "Unclosed connector" not in log ] print("\n Captured caplog records - ", captured_logs) diff --git a/tests/test_litellm/proxy/client/cli/test_agents.py b/tests/test_litellm/proxy/client/cli/test_agents.py new file mode 100644 index 00000000000..afd1696a89f --- /dev/null +++ b/tests/test_litellm/proxy/client/cli/test_agents.py @@ -0,0 +1,475 @@ +import os +import sys +from unittest.mock import patch + +import click +import pytest +import requests +from click.testing import CliRunner + +sys.path.insert( + 0, os.path.abspath("../../..") +) # Adds the parent directory to the system path + + +from litellm.proxy.client.cli.commands.agents import ( + AgentRunError, + agent_commands, + agent_launch_args, + agent_profile, + build_agent_env, + run_agent, + verify_proxy_key, +) + +AGENTS_MODULE = "litellm.proxy.client.cli.commands.agents" + + +def _agent_command(name): + return next(c for c in agent_commands() if c.name == name) + + +class _FakeResponse: + def __init__(self, status_code): + self.status_code = status_code + + +class TestAgentProfile: + def test_claude_is_anthropic(self): + name, profiles = agent_profile("claude") + assert name == "Claude Code" + assert profiles == frozenset({"anthropic"}) + + def test_claude_full_path_uses_basename(self): + name, profiles = agent_profile("/usr/local/bin/claude") + assert name == "Claude Code" + assert profiles == frozenset({"anthropic"}) + + def test_codex_and_opencode_are_openai(self): + assert agent_profile("codex") == ("Codex", frozenset({"openai"})) + assert agent_profile("opencode") == ("OpenCode", frozenset({"openai"})) + + def test_unknown_command_gets_both_profiles(self): + name, profiles = agent_profile("mytool") + assert name == "mytool" + assert profiles == frozenset({"anthropic", "openai"}) + + +class TestBuildAgentEnv: + def test_anthropic_profile_uses_bare_root_and_bearer(self): + env = build_agent_env( + {}, "http://localhost:4000/", "sk-key", frozenset({"anthropic"}) + ) + assert env["ANTHROPIC_BASE_URL"] == "http://localhost:4000" + assert env["ANTHROPIC_AUTH_TOKEN"] == "sk-key" + assert "OPENAI_BASE_URL" not in env + assert "OPENAI_API_KEY" not in env + + def test_anthropic_profile_drops_existing_api_key(self): + env = build_agent_env( + {"ANTHROPIC_API_KEY": "real-key"}, + "http://localhost:4000", + "sk-key", + frozenset({"anthropic"}), + ) + assert "ANTHROPIC_API_KEY" not in env + + def test_openai_profile_appends_v1(self): + env = build_agent_env( + {}, "http://localhost:4000/", "sk-key", frozenset({"openai"}) + ) + assert env["OPENAI_BASE_URL"] == "http://localhost:4000/v1" + assert env["OPENAI_API_KEY"] == "sk-key" + assert "ANTHROPIC_BASE_URL" not in env + + def test_both_profiles_set_everything(self): + env = build_agent_env( + {}, "http://localhost:4000", "sk-key", frozenset({"anthropic", "openai"}) + ) + assert env["ANTHROPIC_BASE_URL"] == "http://localhost:4000" + assert env["OPENAI_BASE_URL"] == "http://localhost:4000/v1" + assert env["ANTHROPIC_AUTH_TOKEN"] == "sk-key" + assert env["OPENAI_API_KEY"] == "sk-key" + + def test_preserves_unrelated_env_and_does_not_mutate_input(self): + base = {"PATH": "/usr/bin", "ANTHROPIC_API_KEY": "real-key"} + env = build_agent_env( + base, "http://localhost:4000", "sk-key", frozenset({"anthropic"}) + ) + assert env["PATH"] == "/usr/bin" + assert base == {"PATH": "/usr/bin", "ANTHROPIC_API_KEY": "real-key"} + + +class TestAgentLaunchArgs: + def test_claude_and_opencode_get_no_extra_args(self): + assert agent_launch_args("claude", "http://localhost:4000") == [] + assert agent_launch_args("opencode", "http://localhost:4000") == [] + + def test_unknown_agent_gets_no_extra_args(self): + assert agent_launch_args("mytool", "http://localhost:4000") == [] + + def test_codex_points_provider_at_proxy_over_http(self): + args = agent_launch_args("codex", "http://localhost:4000/") + joined = " ".join(args) + assert 'model_provider="litellm"' in args + assert 'model_providers.litellm.base_url="http://localhost:4000/v1"' in args + assert 'model_providers.litellm.env_key="OPENAI_API_KEY"' in args + assert 'model_providers.litellm.wire_api="responses"' in args + assert "model_providers.litellm.supports_websockets=false" in args + assert joined.count("-c") == 6 + + def test_codex_uses_basename(self): + assert agent_launch_args("/usr/local/bin/codex", "http://localhost:4000") == ( + agent_launch_args("codex", "http://localhost:4000") + ) + + +class TestVerifyProxyKey: + def test_ok_status_passes_and_uses_models_endpoint(self): + captured = {} + + def fake_get(url, headers, timeout): + captured["url"] = url + captured["headers"] = headers + return _FakeResponse(200) + + verify_proxy_key("http://localhost:4000/", "sk-key", get=fake_get) + + assert captured["url"] == "http://localhost:4000/v1/models" + assert captured["headers"] == {"Authorization": "Bearer sk-key"} + + @pytest.mark.parametrize("status", [401, 403]) + def test_rejected_key_raises(self, status): + with pytest.raises(AgentRunError, match="rejected your key"): + verify_proxy_key( + "http://localhost:4000", + "sk-key", + get=lambda *a, **k: _FakeResponse(status), + ) + + def test_unreachable_proxy_raises(self): + def boom(*a, **k): + raise requests.ConnectionError("refused") + + with pytest.raises(AgentRunError, match="Could not reach"): + verify_proxy_key("http://localhost:4000", "sk-key", get=boom) + + def test_other_non_2xx_is_tolerated(self): + verify_proxy_key( + "http://localhost:4000", + "sk-key", + get=lambda *a, **k: _FakeResponse(500), + ) + + +class TestRunAgent: + def test_wires_env_and_launches_resolved_binary(self): + calls = {} + + def fake_launcher(path, args, env): + calls["path"] = path + calls["args"] = tuple(args) + calls["env"] = dict(env) + + run_agent( + "http://localhost:4000", + "sk-key", + ["claude", "--resume"], + base_env={"PATH": "/usr/bin", "ANTHROPIC_API_KEY": "leaked"}, + which=lambda name: "/usr/local/bin/claude", + verify=lambda *a: None, + launcher=fake_launcher, + ) + + assert calls["path"] == "/usr/local/bin/claude" + assert calls["args"] == ("claude", "--resume") + env = calls["env"] + assert env["ANTHROPIC_BASE_URL"] == "http://localhost:4000" + assert env["ANTHROPIC_AUTH_TOKEN"] == "sk-key" + assert "ANTHROPIC_API_KEY" not in env + assert "OPENAI_BASE_URL" not in env + + def test_codex_gets_openai_env(self): + calls = {} + run_agent( + "http://localhost:4000", + "sk-key", + ["codex"], + base_env={}, + which=lambda name: "/usr/local/bin/codex", + verify=lambda *a: None, + launcher=lambda p, a, e: calls.update(env=dict(e)), + ) + assert calls["env"]["OPENAI_BASE_URL"] == "http://localhost:4000/v1" + assert calls["env"]["OPENAI_API_KEY"] == "sk-key" + assert "ANTHROPIC_BASE_URL" not in calls["env"] + + def test_codex_injects_proxy_provider_args_before_user_args(self): + calls = {} + run_agent( + "http://localhost:4000", + "sk-key", + ["codex", "exec", "do a thing"], + base_env={}, + which=lambda name: "/usr/local/bin/codex", + verify=lambda *a: None, + launcher=lambda p, a, e: calls.update(args=tuple(a)), + ) + args = calls["args"] + assert args[0] == "codex" + assert args[-2:] == ("exec", "do a thing") + assert 'model_provider="litellm"' in args + assert 'model_providers.litellm.base_url="http://localhost:4000/v1"' in args + # overrides must precede the codex subcommand so codex parses them + assert args.index('model_provider="litellm"') < args.index("exec") + + def test_claude_launches_without_injected_args(self): + calls = {} + run_agent( + "http://localhost:4000", + "sk-key", + ["claude", "--resume"], + base_env={}, + which=lambda name: "/usr/local/bin/claude", + verify=lambda *a: None, + launcher=lambda p, a, e: calls.update(args=tuple(a)), + ) + assert calls["args"] == ("claude", "--resume") + + def test_missing_binary_raises_with_install_hint(self): + with pytest.raises(AgentRunError, match="claude.*Install it first"): + run_agent( + "http://localhost:4000", + "sk-key", + ["claude"], + base_env={}, + which=lambda name: None, + verify=lambda *a: None, + launcher=lambda *a: None, + ) + + def test_skip_verify_does_not_call_verify(self): + verified = [] + launched = [] + run_agent( + "http://localhost:4000", + "sk-key", + ["claude"], + skip_verify=True, + base_env={}, + which=lambda name: "/usr/local/bin/claude", + verify=lambda *a: verified.append(a), + launcher=lambda *a: launched.append(a), + ) + assert verified == [] + assert len(launched) == 1 + + def test_verify_failure_aborts_before_launch(self): + launched = [] + + def boom(*a): + raise AgentRunError("rejected") + + with pytest.raises(AgentRunError): + run_agent( + "http://localhost:4000", + "sk-key", + ["claude"], + base_env={}, + which=lambda name: "/usr/local/bin/claude", + verify=boom, + launcher=lambda *a: launched.append(a), + ) + assert launched == [] + + def test_empty_command_raises(self): + with pytest.raises(AgentRunError): + run_agent("http://localhost:4000", "sk-key", []) + + def test_reattach_terminal_runs_just_before_launch(self): + order = [] + run_agent( + "http://localhost:4000", + "sk-key", + ["claude"], + skip_verify=True, + base_env={}, + which=lambda name: "/usr/local/bin/claude", + launcher=lambda *a: order.append("launch"), + reattach_terminal=lambda: order.append("reattach"), + ) + assert order == ["reattach", "launch"] + + def test_no_reattach_terminal_by_default(self): + order = [] + run_agent( + "http://localhost:4000", + "sk-key", + ["claude"], + skip_verify=True, + base_env={}, + which=lambda name: "/usr/local/bin/claude", + launcher=lambda *a: order.append("launch"), + ) + assert order == ["launch"] + + +class TestAgentCommands: + def setup_method(self): + self.runner = CliRunner() + + def test_one_command_per_known_agent(self): + assert {c.name for c in agent_commands()} == {"claude", "codex", "opencode"} + + def test_claude_launches_with_stored_key_and_forwards_args(self): + captured = {} + + def fake_run_agent(base_url, api_key, command, **kwargs): + captured["base_url"] = base_url + captured["api_key"] = api_key + captured["command"] = list(command) + captured["skip_verify"] = kwargs.get("skip_verify") + + with patch(f"{AGENTS_MODULE}.run_agent", side_effect=fake_run_agent): + result = self.runner.invoke( + _agent_command("claude"), + ["--resume", "-p", "hi"], + obj={"base_url": "http://localhost:4000", "api_key": "sk-key"}, + ) + + assert result.exit_code == 0, result.output + assert captured["api_key"] == "sk-key" + assert captured["command"] == ["claude", "--resume", "-p", "hi"] + assert captured["skip_verify"] is False + assert ( + "routing Claude Code through proxy at http://localhost:4000" + in result.output + ) + + def test_codex_shows_friendly_name(self): + captured = {} + with patch( + f"{AGENTS_MODULE}.run_agent", + side_effect=lambda b, k, c, **kw: captured.update(command=list(c)), + ): + result = self.runner.invoke( + _agent_command("codex"), + ["exec", "do a thing"], + obj={"base_url": "http://localhost:4000", "api_key": "sk-key"}, + ) + assert result.exit_code == 0, result.output + assert captured["command"] == ["codex", "exec", "do a thing"] + assert "routing Codex through proxy" in result.output + + def test_skip_verify_is_consumed_not_forwarded(self): + captured = {} + + def fake_run_agent(base_url, api_key, command, **kwargs): + captured["command"] = list(command) + captured["skip_verify"] = kwargs.get("skip_verify") + + with patch(f"{AGENTS_MODULE}.run_agent", side_effect=fake_run_agent): + result = self.runner.invoke( + _agent_command("claude"), + ["--skip-verify", "--resume"], + obj={"base_url": "http://localhost:4000", "api_key": "sk-key"}, + ) + + assert result.exit_code == 0, result.output + assert captured["skip_verify"] is True + assert captured["command"] == ["claude", "--resume"] + + def test_non_interactive_without_key_errors_clearly(self): + with ( + patch(f"{AGENTS_MODULE}._is_interactive", return_value=False), + patch(f"{AGENTS_MODULE}.run_agent") as mock_run, + ): + result = self.runner.invoke( + _agent_command("claude"), + [], + obj={"base_url": "http://localhost:4000", "api_key": None}, + ) + assert result.exit_code != 0 + assert "LITELLM_PROXY_API_KEY" in result.output + mock_run.assert_not_called() + + def test_interactive_without_key_logs_in_then_launches(self): + captured = {} + + @click.command() + def fake_login(): + pass + + with ( + patch(f"{AGENTS_MODULE}._is_interactive", return_value=True), + patch(f"{AGENTS_MODULE}.login", fake_login), + patch( + f"{AGENTS_MODULE}.get_stored_api_key", return_value="sk-after-login" + ) as mock_get, + patch( + f"{AGENTS_MODULE}.run_agent", + side_effect=lambda base_url, api_key, command, **k: captured.update( + api_key=api_key + ), + ), + ): + result = self.runner.invoke( + _agent_command("claude"), + [], + obj={"base_url": "http://localhost:4000", "api_key": None}, + ) + + assert result.exit_code == 0, result.output + assert captured["api_key"] == "sk-after-login" + mock_get.assert_called_once_with(expected_base_url="http://localhost:4000") + + def test_agent_run_error_becomes_click_error(self): + with patch( + f"{AGENTS_MODULE}.run_agent", + side_effect=AgentRunError("could not reach proxy"), + ): + result = self.runner.invoke( + _agent_command("claude"), + [], + obj={"base_url": "http://localhost:4000", "api_key": "sk-key"}, + ) + assert result.exit_code != 0 + assert "could not reach proxy" in result.output + + def test_interactive_session_reattaches_terminal_before_handoff(self): + from litellm.proxy.client.cli.commands.agents import ( + _restore_controlling_terminal, + ) + + captured = {} + with ( + patch(f"{AGENTS_MODULE}._is_interactive", return_value=True), + patch( + f"{AGENTS_MODULE}.run_agent", + side_effect=lambda b, k, c, **kw: captured.update(kw), + ), + ): + result = self.runner.invoke( + _agent_command("claude"), + [], + obj={"base_url": "http://localhost:4000", "api_key": "sk-key"}, + ) + assert result.exit_code == 0, result.output + assert captured["reattach_terminal"] is _restore_controlling_terminal + + def test_non_interactive_agent_mode_leaves_stdin_alone(self): + captured = {} + with ( + patch(f"{AGENTS_MODULE}._is_interactive", return_value=False), + patch( + f"{AGENTS_MODULE}.run_agent", + side_effect=lambda b, k, c, **kw: captured.update(kw), + ), + ): + result = self.runner.invoke( + _agent_command("claude"), + [], + obj={"base_url": "http://localhost:4000", "api_key": "sk-key"}, + ) + assert result.exit_code == 0, result.output + assert captured["reattach_terminal"] is None diff --git a/tests/test_litellm/proxy/client/cli/test_auth_commands.py b/tests/test_litellm/proxy/client/cli/test_auth_commands.py index 2e738ff900d..4ee8b502aa2 100644 --- a/tests/test_litellm/proxy/client/cli/test_auth_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_auth_commands.py @@ -517,7 +517,7 @@ class TestWhoamiCommand: assert result.exit_code == 0 assert "❌ Not authenticated" in result.output - assert "Run 'litellm-proxy login'" in result.output + assert "Run 'lite login'" in result.output def test_whoami_old_token(self): """Test whoami with old token showing warning""" diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index c763e9c0e98..2efec3e0b34 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -1777,6 +1777,23 @@ class TestHTMLIntegration: assert isinstance(html, str) assert len(html) > 0 + def test_success_page_instructs_manual_close_without_false_countdown(self): + """Browsers refuse window.close() on tabs they did not open via window.open() + (the CLI opens the page with webbrowser.open), so a 'closing in 3...' countdown + is a promise the browser usually can't keep and the page gets stuck on + 'Closing...'. The page must instead always show the manual-close instruction + and never advertise an auto-close that won't happen. + """ + from litellm.proxy.common_utils.html_forms.cli_sso_success import ( + render_cli_sso_success_page, + ) + + html = render_cli_sso_success_page() + + assert "You can now close this window and return to your terminal." in html + assert "Closing..." not in html + assert "This window will close in" not in html + class TestCustomUISSO: """Test the custom UI SSO sign-in handler functionality""" diff --git a/uv.lock b/uv.lock index 2403a7fbf03..1100db783d3 100644 --- a/uv.lock +++ b/uv.lock @@ -3297,6 +3297,11 @@ dependencies = [ caching = [ { name = "diskcache" }, ] +cli = [ + { name = "pyyaml" }, + { name = "requests" }, + { name = "rich" }, +] extra-proxy = [ { name = "a2a-sdk" }, { name = "azure-identity" }, @@ -3528,10 +3533,13 @@ requires-dist = [ { name = "pyroscope-io", marker = "sys_platform != 'win32' and extra == 'proxy'", specifier = ">=0.8.16,<1.0" }, { name = "python-dotenv", specifier = ">=1.0.0,<2.0" }, { name = "python-multipart", marker = "extra == 'proxy'", specifier = ">=0.0.27,<1.0" }, + { name = "pyyaml", marker = "extra == 'cli'", specifier = ">=6.0.3,<7.0" }, { name = "pyyaml", marker = "extra == 'proxy'", specifier = ">=6.0.3,<7.0" }, { name = "redisvl", marker = "python_full_version < '3.14' and extra == 'extra-proxy'", specifier = ">=0.4.1,<1.0" }, + { name = "requests", marker = "extra == 'cli'", specifier = ">=2.32.0,<3.0" }, { name = "resend", marker = "extra == 'extra-proxy'", specifier = ">=2.23.0,<3.0" }, { name = "restrictedpython", marker = "extra == 'proxy'", specifier = ">=8.1,<9.0" }, + { name = "rich", marker = "extra == 'cli'", specifier = ">=13.9.4,<14.0" }, { name = "rich", marker = "extra == 'proxy'", specifier = ">=13.9.4,<14.0" }, { name = "rq", marker = "extra == 'proxy'", specifier = ">=2.7.0,<3.0" }, { name = "semantic-router", marker = "python_full_version < '3.14' and extra == 'semantic-router'", specifier = ">=0.1.15,<1.0" }, @@ -3545,7 +3553,7 @@ requires-dist = [ { name = "uvloop", marker = "sys_platform != 'win32' and extra == 'proxy'", specifier = ">=0.21.0,<1.0" }, { name = "websockets", marker = "extra == 'proxy'", specifier = ">=15.0.1,<16.0" }, ] -provides-extras = ["proxy", "extra-proxy", "utils", "caching", "semantic-router", "mlflow", "grpc", "stt-nvidia-riva", "google", "proxy-runtime"] +provides-extras = ["proxy", "cli", "extra-proxy", "utils", "caching", "semantic-router", "mlflow", "grpc", "stt-nvidia-riva", "google", "proxy-runtime"] [package.metadata.requires-dev] ci = [ From 7899463c6a826e7172427f3e43cce95f8ca547a1 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Wed, 10 Jun 2026 15:22:00 -0700 Subject: [PATCH 051/185] fix(callbacks): forward callback_settings to callback initializers and guard consumers against non-dict values (#30161) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(datadog): pass callback_specific_params so DatadogCostManagementLogger receives cost_tag_keys (#29590) * fix(datadog): pass callback_specific_params so DatadogCostManagementLogger receives cost_tag_keys * test(proxy): regression test that load_config forwards callback_specific_params * fix(proxy): guard lakera_prompt_injection callback_specific_params against non-dict Addresses review feedback: forwarding callback_settings as callback_specific_params (so DatadogCostManagementLogger receives cost_tag_keys) exposed the lakera_prompt_injection branch, which did lakeraAI_Moderation(**callback_specific_params ["lakera_prompt_injection"]) with no type guard. A config like `callback_settings: {lakera_prompt_injection: "any-string"}` then hit `**"any-string"` -> TypeError: argument after ** must be a mapping, not str. Guard the lakera branch with isinstance(dict), matching the existing presidio and datadog_cost_management branches (non-dict values fall back to {}). Add a regression test asserting initialize_callbacks_on_proxy ignores a non-dict value instead of crashing. Co-Authored-By: Claude Opus 4.8 (1M context) * test: inject fake lakera_ai module to avoid importing the real one CI fix for the lakera regression test: it stubbed litellm.proxy.proxy_server with a SimpleNamespace and then monkeypatch.setattr'd the real lakera_ai module, which forces importing it — and lakera_ai does `from litellm.proxy.proxy_server import LiteLLM_TeamTable`, absent on the stub -> ImportError under proxy-infra tests. Inject a fake lakera_ai module into sys.modules instead, so the callbacks branch's `from ...lakera_ai import lakeraAI_Moderation` resolves to the stub without loading the real module. The guard under test (isinstance(dict) in the lakera branch) is unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) --------- Co-authored-by: Claude Opus 4.8 (1M context) * fix(callbacks): guard compression/websearch interceptors against non-dict callback_settings (#30153) #29590 forwards the full callback_settings dict into initialize_callbacks_on_proxy, which activates the compression_interception and websearch_interception consumers. Their initialize_from_proxy_config read the callback_settings subkey without an isinstance(dict) guard, so a non-dict value such as `compression_interception: true` reached from_config_yaml(...).get(...) and aborted proxy startup with AttributeError. #29590 added that guard for lakera_prompt_injection but not for these two Mirror the isinstance(dict) guard already used by the lakera, presidio, and datadog branches so a non-dict value is ignored and the callback initializes with defaults. A parametrized test feeds every callback_settings consumer a non-dict value through initialize_callbacks_on_proxy to catch a future consumer that forgets the guard * fix(callbacks): normalize non-dict callback_specific_params to empty dict A blank callback_settings: key in YAML loads as None, and config.get('callback_settings', {}) returns None because dict.get only falls back to the default when the key is absent. Forwarding that value verbatim to initialize_callbacks_on_proxy made the first '' in callback_specific_params membership test raise TypeError: argument of type 'NoneType' is not iterable, aborting proxy startup. Same failure for any non-dict root such as callback_settings: true. Normalize the value at the function boundary so both callsites (and any future ones) initialize callbacks with their defaults instead of crashing. --------- Co-authored-by: Hedi Daoud <150018939+hdaoud23@users.noreply.github.com> Co-authored-by: Claude Opus 4.8 (1M context) --- .../compression_interception/handler.py | 9 +- .../websearch_interception/handler.py | 9 +- litellm/proxy/common_utils/callback_utils.py | 11 ++- litellm/proxy/proxy_server.py | 1 + .../test_compression_interception_handler.py | 34 +++++++ .../test_websearch_interception_handler.py | 29 ++++++ .../proxy/common_utils/test_callback_utils.py | 99 ++++++++++++++++++- .../proxy/proxy_server/test_proxy_config.py | 92 +++++++++++++++++ 8 files changed, 277 insertions(+), 7 deletions(-) diff --git a/litellm/integrations/compression_interception/handler.py b/litellm/integrations/compression_interception/handler.py index c6ae7d9e82b..8899089500d 100644 --- a/litellm/integrations/compression_interception/handler.py +++ b/litellm/integrations/compression_interception/handler.py @@ -72,8 +72,13 @@ class CompressionInterceptionLogger(CustomLogger): compression_params: CompressionInterceptionConfig = {} if "compression_interception_params" in litellm_settings: compression_params = litellm_settings["compression_interception_params"] - elif "compression_interception" in callback_specific_params: - compression_params = callback_specific_params["compression_interception"] + elif "compression_interception" in callback_specific_params and isinstance( + callback_specific_params["compression_interception"], dict + ): + compression_params = cast( + CompressionInterceptionConfig, + callback_specific_params["compression_interception"], + ) return CompressionInterceptionLogger.from_config_yaml(compression_params) async def async_pre_call_deployment_hook( diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 37528e7dcd5..79f9b16bba0 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -1339,8 +1339,13 @@ class WebSearchInterceptionLogger(CustomLogger): websearch_params: WebSearchInterceptionConfig = {} if "websearch_interception_params" in litellm_settings: websearch_params = litellm_settings["websearch_interception_params"] - elif "websearch_interception" in callback_specific_params: - websearch_params = callback_specific_params["websearch_interception"] + elif "websearch_interception" in callback_specific_params and isinstance( + callback_specific_params["websearch_interception"], dict + ): + websearch_params = cast( + WebSearchInterceptionConfig, + callback_specific_params["websearch_interception"], + ) # Use classmethod to initialize from config return WebSearchInterceptionLogger.from_config_yaml(websearch_params) diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index a65e737f248..c630294c1ec 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -40,8 +40,10 @@ def initialize_callbacks_on_proxy( # noqa: PLR0915 premium_user: bool, config_file_path: str, litellm_settings: dict, - callback_specific_params: dict = {}, + callback_specific_params: Optional[dict] = None, ): + if not isinstance(callback_specific_params, dict): + callback_specific_params = {} from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.logging_callback_manager import ( LoggingCallbackManager, @@ -166,7 +168,12 @@ def initialize_callbacks_on_proxy( # noqa: PLR0915 ) init_params = {} - if "lakera_prompt_injection" in callback_specific_params: + if ( + "lakera_prompt_injection" in callback_specific_params + and isinstance( + callback_specific_params["lakera_prompt_injection"], dict + ) + ): init_params = callback_specific_params["lakera_prompt_injection"] lakera_moderations_object = lakeraAI_Moderation(**init_params) imported_list.append(lakera_moderations_object) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index ba23175c10f..96c9cd1e8fb 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -4079,6 +4079,7 @@ class ProxyConfig: premium_user=premium_user, config_file_path=config_file_path, litellm_settings=litellm_settings, + callback_specific_params=callback_settings, ) elif key == "model_group_settings": diff --git a/tests/test_litellm/integrations/compression_interception/test_compression_interception_handler.py b/tests/test_litellm/integrations/compression_interception/test_compression_interception_handler.py index 56e5a94cd49..ffa81abf86c 100644 --- a/tests/test_litellm/integrations/compression_interception/test_compression_interception_handler.py +++ b/tests/test_litellm/integrations/compression_interception/test_compression_interception_handler.py @@ -32,6 +32,40 @@ def test_initialize_from_proxy_config(): assert logger.compression_target == 789 +def test_initialize_from_proxy_config_ignores_non_dict_callback_specific_params(): + """Regression (#29590): a non-dict value under + callback_settings.compression_interception must not crash initialization. + + Forwarding callback_settings as callback_specific_params activates this + branch; without the isinstance(dict) guard a non-dict value reached + from_config_yaml(...).get(...) and raised AttributeError at proxy startup. + The value is ignored and the logger falls back to defaults. + """ + logger = CompressionInterceptionLogger.initialize_from_proxy_config( + litellm_settings={}, + callback_specific_params={"compression_interception": True}, + ) + + assert logger.enabled is True + assert logger.compression_trigger == 200_000 + + +def test_initialize_from_proxy_config_honors_dict_callback_specific_params(): + """A valid dict under callback_settings.compression_interception is applied.""" + logger = CompressionInterceptionLogger.initialize_from_proxy_config( + litellm_settings={}, + callback_specific_params={ + "compression_interception": { + "enabled": False, + "compression_trigger": 12345, + } + }, + ) + + assert logger.enabled is False + assert logger.compression_trigger == 12345 + + @pytest.mark.asyncio async def test_pre_call_hook_compresses_messages_and_injects_tool(monkeypatch): """Test pre-call hook compresses and stores per-call cache.""" diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py b/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py index 10951265115..c2a502b34eb 100644 --- a/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py +++ b/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py @@ -34,6 +34,35 @@ def test_initialize_from_proxy_config(): assert logger.search_tool_name == "my-search" +def test_initialize_from_proxy_config_ignores_non_dict_callback_specific_params(): + """Regression (#29590): a non-dict value under + callback_settings.websearch_interception must not crash initialization. + + Forwarding callback_settings as callback_specific_params activates this + branch; without the isinstance(dict) guard a non-dict value reached + from_config_yaml(...).get(...) and raised AttributeError at proxy startup. + The value is ignored and the logger falls back to defaults. + """ + logger = WebSearchInterceptionLogger.initialize_from_proxy_config( + litellm_settings={}, + callback_specific_params={"websearch_interception": True}, + ) + + assert logger.search_tool_name is None + + +def test_initialize_from_proxy_config_honors_dict_callback_specific_params(): + """A valid dict under callback_settings.websearch_interception is applied.""" + logger = WebSearchInterceptionLogger.initialize_from_proxy_config( + litellm_settings={}, + callback_specific_params={ + "websearch_interception": {"search_tool_name": "ws-tool"} + }, + ) + + assert logger.search_tool_name == "ws-tool" + + @pytest.mark.asyncio async def test_async_should_run_agentic_loop(): """Test that agentic loop is NOT triggered for wrong provider or missing WebSearch tool""" diff --git a/tests/test_litellm/proxy/common_utils/test_callback_utils.py b/tests/test_litellm/proxy/common_utils/test_callback_utils.py index d328d68dcd4..36ff3f3c399 100644 --- a/tests/test_litellm/proxy/common_utils/test_callback_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_callback_utils.py @@ -1,7 +1,9 @@ import copy import sys import os -from types import SimpleNamespace +from types import ModuleType, SimpleNamespace + +import pytest sys.path.insert( 0, os.path.abspath("../../..") @@ -309,3 +311,98 @@ def test_encrypt_callback_vars_only_encrypts_credential_fields(monkeypatch): assert cv["langfuse_host"] == "https://cloud.langfuse.com" assert cv["langsmith_project"] == "my-proj" assert cv["langsmith_base_url"] == "https://smith.example" + + +def test_initialize_callbacks_on_proxy_lakera_ignores_non_dict_callback_settings( + monkeypatch, +): + """Regression: a non-dict value under callback_settings.lakera_prompt_injection + must not crash initialize_callbacks_on_proxy. + + Forwarding callback_settings as callback_specific_params (so callbacks like + DatadogCostManagementLogger receive their init params) exposes the lakera + branch, which previously did lakeraAI_Moderation(**callback_specific_params[ + "lakera_prompt_injection"]) with no isinstance(dict) guard. For a config like + {"lakera_prompt_injection": "x"} that is `**"x"` -> TypeError: argument after + ** must be a mapping, not str. The branch now guards on isinstance(dict), + matching the presidio / datadog_cost_management branches. + """ + captured = {} + + class _DummyLakera: + def __init__(self, **kwargs): + captured["kwargs"] = kwargs + + # Inject a fake lakera_ai module so the branch's + # `from ...lakera_ai import lakeraAI_Moderation` resolves to our stub without + # importing the real module (which imports proxy_server symbols not present + # under the stubbed proxy_server below). + fake_lakera = ModuleType("litellm.proxy.guardrails.guardrail_hooks.lakera_ai") + fake_lakera.lakeraAI_Moderation = _DummyLakera + monkeypatch.setitem( + sys.modules, + "litellm.proxy.guardrails.guardrail_hooks.lakera_ai", + fake_lakera, + ) + monkeypatch.setitem( + sys.modules, + "litellm.proxy.proxy_server", + SimpleNamespace(prisma_client=None), + ) + + original_callbacks = ( + list(litellm.callbacks) if isinstance(litellm.callbacks, list) else [] + ) + litellm.callbacks = [] + try: + # A non-dict value must be ignored (init_params stays {}), not **-unpacked. + initialize_callbacks_on_proxy( + value=["lakera_prompt_injection"], + premium_user=False, + config_file_path=".", + litellm_settings={}, + callback_specific_params={"lakera_prompt_injection": "any-string"}, + ) + assert captured["kwargs"] == {} + assert any(isinstance(c, _DummyLakera) for c in litellm.callbacks) + finally: + litellm.callbacks = original_callbacks + + +@pytest.mark.parametrize("bad_root", [None, True]) +def test_initialize_callbacks_on_proxy_non_dict_callback_specific_params_root( + monkeypatch, bad_root +): + """Regression: a blank `callback_settings:` key in YAML loads as None (and + `callback_settings: true` as a bool); load_config forwards that value + verbatim as callback_specific_params. Membership tests like + `"compression_interception" in callback_specific_params` then raise + TypeError and abort proxy startup. A non-dict root must be normalized to {} + so the callback initializes with its defaults. + """ + monkeypatch.setitem( + sys.modules, + "litellm.proxy.proxy_server", + SimpleNamespace(prisma_client=None), + ) + from litellm.integrations.compression_interception.handler import ( + CompressionInterceptionLogger, + ) + + original_callbacks = ( + list(litellm.callbacks) if isinstance(litellm.callbacks, list) else [] + ) + litellm.callbacks = [] + try: + initialize_callbacks_on_proxy( + value=["compression_interception"], + premium_user=False, + config_file_path=".", + litellm_settings={}, + callback_specific_params=bad_root, + ) + assert any( + isinstance(c, CompressionInterceptionLogger) for c in litellm.callbacks + ) + finally: + litellm.callbacks = original_callbacks diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index 164538a2757..677d358428d 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -601,6 +601,98 @@ async def test_ProxyConfig_load_config_missing_file_raises(monkeypatch): await pc.load_config(router=None, config_file_path="/no/file.yaml") +@pytest.mark.asyncio +async def test_ProxyConfig_load_config_forwards_callback_specific_params( + tmp_path, monkeypatch +): + """Regression: callback_settings from config must be forwarded to + initialize_callbacks_on_proxy as callback_specific_params. + + Callbacks like DatadogCostManagementLogger read their init params (e.g. + cost_tag_keys) from callback_specific_params[]. If the + argument is dropped at the call site, they silently initialize with empty + params and the configured allowlist never takes effect. + """ + f = tmp_path / "c.yaml" + f.write_text( + "model_list: []\n" + "general_settings: {}\n" + "callback_settings:\n" + " datadog_cost_management:\n" + " cost_tag_keys:\n" + " - capability\n" + " - platform\n" + " - ai_product\n" + "litellm_settings:\n" + ' callbacks: ["datadog_cost_management"]\n' + ) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False) + monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False) + + captured = {} + + def _fake_initialize_callbacks_on_proxy(**kwargs): + captured.update(kwargs) + + monkeypatch.setattr( + "litellm.proxy.proxy_server.initialize_callbacks_on_proxy", + _fake_initialize_callbacks_on_proxy, + ) + + pc = ProxyConfig() + await pc.load_config(router=None, config_file_path=str(f)) + + # The callbacks branch must forward the loaded callback_settings. + assert captured.get("callback_specific_params") == { + "datadog_cost_management": { + "cost_tag_keys": ["capability", "platform", "ai_product"] + } + } + + +@pytest.mark.asyncio +async def test_ProxyConfig_load_config_blank_callback_settings_does_not_crash( + tmp_path, monkeypatch +): + """Regression: `callback_settings:` with no body loads as None because + dict.get() only falls back to the default when the key is absent. The None + was forwarded verbatim to initialize_callbacks_on_proxy, where the first + `"" in callback_specific_params` membership test raised + TypeError: argument of type 'NoneType' is not iterable, aborting startup. + Startup must succeed and the callback must initialize with its defaults. + """ + f = tmp_path / "c.yaml" + f.write_text( + "model_list: []\n" + "general_settings: {}\n" + "callback_settings:\n" + "litellm_settings:\n" + ' callbacks: ["compression_interception"]\n' + ) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False) + monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False) + + from litellm.integrations.compression_interception.handler import ( + CompressionInterceptionLogger, + ) + + original_callbacks = ( + list(litellm.callbacks) if isinstance(litellm.callbacks, list) else [] + ) + litellm.callbacks = [] + try: + pc = ProxyConfig() + await pc.load_config(router=None, config_file_path=str(f)) + + assert any( + isinstance(c, CompressionInterceptionLogger) for c in litellm.callbacks + ) + finally: + litellm.callbacks = original_callbacks + + # --------------------------------------------------------------------------- # ProxyConfig._init_non_llm_configs # --------------------------------------------------------------------------- From 1436ee90928668dd371f2ba1922c78e45af8dcd6 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Wed, 10 Jun 2026 15:56:58 -0700 Subject: [PATCH 052/185] fix(mcp): drop orphaned per-user credential rows when an MCP server is deleted (#30141) --- litellm/proxy/_experimental/mcp_server/db.py | 35 ++++++----- .../mcp_server/test_mcp_env_vars.py | 59 +++++++++++++++++++ 2 files changed, 79 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index c52752940c3..8edb831a9df 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -568,10 +568,12 @@ async def delete_mcp_server( """ Delete the mcp server from the db by server_id - The server-row delete is the commit point. Per-user env var rows have no FK - cascade, so they are cleaned up afterwards on a best-effort basis: a transient - failure there leaves only orphaned rows pointing at a now-missing server and - must not turn a successful delete into a caller-visible error. + The server-row delete is the commit point. Per-user credential and env var + rows have no FK cascade, so they are cleaned up afterwards on a best-effort + basis: a transient failure there leaves only orphaned rows pointing at a + now-missing server and must not turn a successful delete into a + caller-visible error. Each table is cleaned independently so a failure on one + still attempts the other. Returns the deleted mcp server record if it exists, otherwise None """ @@ -581,17 +583,20 @@ async def delete_mcp_server( }, ) if deleted_server is not None: - try: - await prisma_client.db.litellm_mcpuserenvvars.delete_many( - where={"server_id": server_id} - ) - except Exception as e: - verbose_proxy_logger.warning( - "MCP server %s deleted but per-user env var cleanup failed; " - "orphaned rows can be removed on a later delete: %s", - server_id, - e, - ) + for model, label in ( + (prisma_client.db.litellm_mcpusercredentials, "credential"), + (prisma_client.db.litellm_mcpuserenvvars, "env var"), + ): + try: + await model.delete_many(where={"server_id": server_id}) + except Exception as e: + verbose_proxy_logger.warning( + "MCP server %s deleted but per-user %s cleanup failed; " + "orphaned rows can be removed on a later delete: %s", + server_id, + label, + e, + ) return deleted_server diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_env_vars.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_env_vars.py index 19065ff816b..a846ca24739 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_env_vars.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_env_vars.py @@ -872,6 +872,7 @@ def _mock_env_vars_prisma(row=None): prisma.db.litellm_mcpuserenvvars.find_many = AsyncMock(return_value=[]) prisma.db.litellm_mcpuserenvvars.upsert = AsyncMock() prisma.db.litellm_mcpuserenvvars.delete_many = AsyncMock() + prisma.db.litellm_mcpusercredentials.delete_many = AsyncMock() return prisma @@ -1252,6 +1253,64 @@ async def test_delete_mcp_server_succeeds_when_orphan_cleanup_fails(): prisma.db.litellm_mcpuserenvvars.delete_many.assert_awaited_once() +@pytest.mark.asyncio +async def test_delete_mcp_server_removes_orphaned_user_credentials(): + """Deleting a server must also drop every user's stored BYOK/OAuth credential + rows for it; there is no FK cascade, so skipping this leaves encrypted secrets + pointing at a now-missing server.""" + from unittest.mock import AsyncMock + + from litellm.proxy._experimental.mcp_server.db import delete_mcp_server + + prisma = _mock_env_vars_prisma() + prisma.db.litellm_mcpservertable.delete = AsyncMock(return_value=object()) + + await delete_mcp_server(prisma, "srv-1") + + prisma.db.litellm_mcpusercredentials.delete_many.assert_awaited_once() + call = prisma.db.litellm_mcpusercredentials.delete_many.call_args + assert call.kwargs["where"] == {"server_id": "srv-1"} + + +@pytest.mark.asyncio +async def test_delete_mcp_server_skips_credential_cleanup_when_server_missing(): + """A no-op delete (server not found) must not touch the credential table.""" + from unittest.mock import AsyncMock + + from litellm.proxy._experimental.mcp_server.db import delete_mcp_server + + prisma = _mock_env_vars_prisma() + prisma.db.litellm_mcpservertable.delete = AsyncMock(return_value=None) + + result = await delete_mcp_server(prisma, "srv-1") + + assert result is None + prisma.db.litellm_mcpusercredentials.delete_many.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_delete_mcp_server_credential_cleanup_failure_still_cleans_env_vars(): + """Each per-user table is cleaned independently: a failure dropping credential + rows must not skip the env var cleanup (or vice versa), and the delete must + still succeed for the caller.""" + from unittest.mock import AsyncMock + + from litellm.proxy._experimental.mcp_server.db import delete_mcp_server + + deleted = object() + prisma = _mock_env_vars_prisma() + prisma.db.litellm_mcpservertable.delete = AsyncMock(return_value=deleted) + prisma.db.litellm_mcpusercredentials.delete_many = AsyncMock( + side_effect=Exception("connection pool exhausted") + ) + + result = await delete_mcp_server(prisma, "srv-1") + + assert result is deleted + prisma.db.litellm_mcpusercredentials.delete_many.assert_awaited_once() + prisma.db.litellm_mcpuserenvvars.delete_many.assert_awaited_once() + + # ── DB helpers: global env vars encrypted at rest ───────────────────────── From 3bd3951e37a0b3201eac9eb1f858d2253aff56e5 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Wed, 10 Jun 2026 16:06:01 -0700 Subject: [PATCH 053/185] fix(proxy): recover from cached-plan errors by reconnecting the Prisma client (#29983) --- litellm/proxy/utils.py | 59 +++++++----- .../test_prisma_client_get_data.py | 92 +++++++++++++++++-- 2 files changed, 119 insertions(+), 32 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 2bba8bbd604..ebd5b5d90cd 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -3253,40 +3253,49 @@ class PrismaClient: self, sql_query: str, *args ) -> Optional[dict]: """ - Execute a query with automatic fallback for PostgreSQL cached plan errors. + Execute a query, recovering once from PostgreSQL's "cached plan must not + change result type" error. - This handles the "cached plan must not change result type" error that occurs - during rolling deployments when schema changes are applied while old pods - still have cached query plans expecting the old schema. + That error surfaces during rolling deployments when a schema change + invalidates the prepared-statement plans that pooled connections still + hold. Clearing only the server-side plans with DEALLOCATE ALL makes + things worse: Prisma's query engine keeps a per-connection client-side + cache of prepared-statement names, so once the server drops a plan the + engine re-sends a name PostgreSQL no longer recognizes and the + connection breaks with `prepared statement "sN" does not exist`. With a + small pool that connection stays poisoned and every auth lookup fails. - Args: - sql_query: SQL query string to execute + Recreating the Prisma client kills the engine subprocess and drops the + server-side plans and the engine's client-side name cache together, so + the retried query is prepared fresh. We reconnect through + `attempt_db_reconnect`, which is singleflight: when a schema change + poisons every pooled connection at once, the first cached-plan error + recreates the client and the concurrent waiters reuse that single + recreate instead of racing to kill each other's fresh engine. We then + retry the identical query exactly once. - Returns: - Query result or None + The retry reuses the original query byte-for-byte. Mutating the SQL + (e.g. injecting a unique comment) would defeat PostgreSQL's plan cache, + forcing a fresh plan on every request and pegging the database CPU. - Raises: - Original exception if not a cached plan error + If the reconnect is skipped because a recent reconnect is still within + its cooldown, the retry runs against the same connection and may fail + again; the get_data backoff decorator re-runs the lookup and a later + attempt reconnects once the cooldown elapses. """ try: return await self.db.query_first(sql_query, *args) except Exception as e: - error_str = str(e) - if "cached plan must not change result type" in error_str: - # Force PostgreSQL to re-plan by invalidating the cache - # Add a unique comment to make the query different - sql_query_retry = sql_query.replace( - "SELECT", - f"SELECT /* cache_invalidated_{int(time.time() * 1000)} */", - ) - verbose_proxy_logger.warning( - "PostgreSQL cached plan error detected for token lookup, " - "retrying with fresh plan. This may occur during rolling deployments " - "when schema changes are applied." - ) - return await self.db.query_first(sql_query_retry, *args) - else: + if "cached plan must not change result type" not in str(e): raise + verbose_proxy_logger.warning( + "PostgreSQL cached plan error detected for token lookup; " + "recreating the database connection and retrying with the same " + "query. This may occur during rolling deployments when schema " + "changes are applied." + ) + await self.attempt_db_reconnect(reason="postgres_cached_plan_error") + return await self.db.query_first(sql_query, *args) @backoff.on_exception( backoff.expo, diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py index 7e7e98d1360..437984d9273 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py @@ -193,6 +193,7 @@ async def test_query_first_with_cached_plan_fallback_happy_returns_row( ) -> None: expected = {"token": "abc", "team_spend": 1.0, "team_max_budget": 5.0} prisma_client.db.query_first = AsyncMock(return_value=expected) + prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) result = await prisma_client._query_first_with_cached_plan_fallback( "SELECT * FROM x WHERE token = $1", "abc" ) @@ -208,35 +209,110 @@ async def test_query_first_with_cached_plan_fallback_happy_returns_row( "args": ("SELECT * FROM x WHERE token = $1", "abc"), "matches": True, } + prisma_client.attempt_db_reconnect.assert_not_awaited() @pytest.mark.asyncio -async def test_query_first_with_cached_plan_fallback_retries_on_cached_plan_error( +async def test_query_first_with_cached_plan_fallback_reconnects_then_retries_identical_query( prisma_client: PrismaClient, ) -> None: + original_query = 'SELECT * FROM "LiteLLM_VerificationToken" WHERE v.token = $1' expected = {"token": "abc", "team_spend": 1.0, "team_max_budget": 5.0} + manager = MagicMock() + query_first = AsyncMock( + side_effect=[ + RuntimeError("cached plan must not change result type"), + expected, + ] + ) + reconnect = AsyncMock(return_value=True) + manager.attach_mock(query_first, "query_first") + manager.attach_mock(reconnect, "attempt_db_reconnect") + prisma_client.db.query_first = query_first + prisma_client.attempt_db_reconnect = reconnect + + result = await prisma_client._query_first_with_cached_plan_fallback( + original_query, "abc" + ) + + assert result == expected + assert query_first.await_count == 2 + first_call, retry_call = query_first.await_args_list + assert retry_call.args == first_call.args == (original_query, "abc") + reconnect.assert_awaited_once() + assert reconnect.await_args.kwargs.get("force", False) is False + assert [name for name, *_ in manager.mock_calls] == [ + "query_first", + "attempt_db_reconnect", + "query_first", + ] + + +@pytest.mark.asyncio +async def test_query_first_with_cached_plan_fallback_never_deallocates( + prisma_client: PrismaClient, +) -> None: + expected = {"token": "abc"} prisma_client.db.query_first = AsyncMock( side_effect=[ RuntimeError("cached plan must not change result type"), expected, ] ) - result = await prisma_client._query_first_with_cached_plan_fallback( - "SELECT * FROM x WHERE token = $1", "abc" + prisma_client.db.execute_raw = AsyncMock(return_value=0) + prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) + + await prisma_client._query_first_with_cached_plan_fallback("SELECT 1") + + prisma_client.db.execute_raw.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_query_first_with_cached_plan_fallback_propagates_when_retry_also_fails( + prisma_client: PrismaClient, +) -> None: + plan_error = RuntimeError("cached plan must not change result type") + prisma_client.db.query_first = AsyncMock(side_effect=[plan_error, plan_error]) + prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) + + with pytest.raises(RuntimeError, match="cached plan must not change result type"): + await prisma_client._query_first_with_cached_plan_fallback("SELECT 1") + + assert prisma_client.db.query_first.await_count == 2 + prisma_client.attempt_db_reconnect.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_query_first_with_cached_plan_fallback_retries_when_reconnect_returns_false( + prisma_client: PrismaClient, +) -> None: + expected = {"token": "abc"} + prisma_client.db.query_first = AsyncMock( + side_effect=[ + RuntimeError("cached plan must not change result type"), + expected, + ] ) + prisma_client.attempt_db_reconnect = AsyncMock(return_value=False) + + result = await prisma_client._query_first_with_cached_plan_fallback("SELECT 1") + assert result == expected assert prisma_client.db.query_first.await_count == 2 - second_call_sql = prisma_client.db.query_first.await_args_list[1].args[0] - assert "cache_invalidated_" in second_call_sql @pytest.mark.asyncio async def test_query_first_with_cached_plan_fallback_reraises_non_plan_errors( prisma_client: PrismaClient, ) -> None: - prisma_client.db.query_first = AsyncMock(side_effect=RuntimeError("totally unrelated")) + prisma_client.db.query_first = AsyncMock( + side_effect=RuntimeError("totally unrelated") + ) + prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) with pytest.raises(RuntimeError, match="totally unrelated"): await prisma_client._query_first_with_cached_plan_fallback("SELECT 1") + assert prisma_client.db.query_first.await_count == 1 + prisma_client.attempt_db_reconnect.assert_not_awaited() @pytest.mark.asyncio @@ -351,7 +427,9 @@ async def test_get_data_token_find_unique_returns_record( async def test_get_data_token_find_unique_missing_token_raises_401( prisma_client: PrismaClient, ) -> None: - prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None) + prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=None + ) with pytest.raises(HTTPException) as excinfo: await prisma_client.get_data(token="sk-missing", table_name="key") err = excinfo.value From dff25fef449bc3e2051ee2638e323d09d05a850a Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Wed, 10 Jun 2026 16:06:32 -0700 Subject: [PATCH 054/185] feat(proxy): add option to disable server-side prepared statements for DB lookups (#29984) --- litellm/proxy/_types.py | 11 ++ litellm/proxy/proxy_cli.py | 25 +++- tests/test_litellm/proxy/test_proxy_cli.py | 121 ++++++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 + 4 files changed, 159 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 33a1e4179fa..1b594e20d32 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2177,6 +2177,17 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): "`statement_cache_size`). Keys here override any default LiteLLM sets." ), ) + database_disable_prepared_statements: Optional[bool] = Field( + None, + description=( + "Disable server-side prepared statements by setting Prisma's " + "`pgbouncer=true` URL param. Use this for pgbouncer transaction-pooling " + "deployments, or to prevent the 'cached plan must not change result " + "type' error that pooled connections hit during rolling schema " + "migrations. An explicit `pgbouncer` in `database_extra_connection_params` " + "takes precedence." + ), + ) database_type: Optional[Literal["dynamo_db"]] = Field( None, description="to use dynamodb instead of postgres db" ) diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index ae831ef1b53..8c3fa952903 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -44,15 +44,19 @@ def _build_db_connection_url_params( pool_timeout: Optional[Union[int, float]], connect_timeout: Optional[Union[int, float]] = None, socket_timeout: Optional[Union[int, float]] = None, + disable_prepared_statements: bool = False, extra_params: Optional[dict] = None, ) -> dict: """Build the Prisma DATABASE_URL query params controlling connection pool behavior. `connect_timeout` / `socket_timeout` map to the Prisma URL params of the same name (https://www.prisma.io/docs/orm/overview/databases/postgresql) and are - omitted when None so Prisma's defaults apply. `extra_params` is an - untyped passthrough — keys it provides win over the named arguments above, - so it can be used to override any default we set here. + omitted when None so Prisma's defaults apply. `disable_prepared_statements` + sets `pgbouncer=true`, which makes Prisma stop using server-side prepared + statements (pgbouncer transaction-pool compatible; also sidesteps the + "cached plan must not change result type" error during rolling migrations). + `extra_params` is an untyped passthrough — keys it provides win over the + named arguments above, so it can be used to override any default we set here. """ params: dict = { "connection_limit": connection_limit, @@ -63,6 +67,8 @@ def _build_db_connection_url_params( params["connect_timeout"] = connect_timeout if socket_timeout is not None: params["socket_timeout"] = socket_timeout + if disable_prepared_statements: + params["pgbouncer"] = "true" if extra_params: params.update(extra_params) return params @@ -963,6 +969,7 @@ def run_server( # noqa: PLR0915 db_connection_timeout: Optional[Union[int, float]] = 60 db_connect_timeout: Optional[Union[int, float]] = None db_socket_timeout: Optional[Union[int, float]] = None + db_disable_prepared_statements: bool = False db_extra_connection_params: Optional[dict] = None general_settings = {} ### GET DB TOKEN FOR IAM AUTH ### @@ -1083,6 +1090,17 @@ def run_server( # noqa: PLR0915 ) db_connect_timeout = general_settings.get("database_connect_timeout") db_socket_timeout = general_settings.get("database_socket_timeout") + _disable_prepared_statements = general_settings.get( + "database_disable_prepared_statements", False + ) + if isinstance(_disable_prepared_statements, str): + from litellm.secret_managers.main import str_to_bool + + db_disable_prepared_statements = ( + str_to_bool(_disable_prepared_statements) is True + ) + else: + db_disable_prepared_statements = bool(_disable_prepared_statements) db_extra_connection_params = general_settings.get( "database_extra_connection_params" ) @@ -1130,6 +1148,7 @@ def run_server( # noqa: PLR0915 pool_timeout=db_connection_timeout, connect_timeout=db_connect_timeout, socket_timeout=db_socket_timeout, + disable_prepared_statements=db_disable_prepared_statements, extra_params=db_extra_connection_params, ) if os.getenv("DATABASE_URL", None) is not None: diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index 4fb725b7ef3..34c88e2fd33 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -795,6 +795,127 @@ class TestProxyInitializationHelpers: assert appended_params["pgbouncer"] == "true" assert appended_params["statement_cache_size"] == 0 + def test_build_db_connection_url_params_disable_prepared_statements(self): + from litellm.proxy.proxy_cli import _build_db_connection_url_params + + params = _build_db_connection_url_params( + connection_limit=10, + pool_timeout=60, + disable_prepared_statements=True, + ) + assert params["pgbouncer"] == "true" + + def test_build_db_connection_url_params_no_pgbouncer_by_default(self): + from litellm.proxy.proxy_cli import _build_db_connection_url_params + + params = _build_db_connection_url_params( + connection_limit=10, + pool_timeout=60, + ) + assert "pgbouncer" not in params + + def test_build_db_connection_url_params_extra_pgbouncer_overrides_flag(self): + from litellm.proxy.proxy_cli import _build_db_connection_url_params + + params = _build_db_connection_url_params( + connection_limit=10, + pool_timeout=60, + disable_prepared_statements=True, + extra_params={"pgbouncer": "false"}, + ) + assert params["pgbouncer"] == "false" + + @pytest.mark.parametrize( + "config_value, expect_pgbouncer", + [ + (True, True), + (False, False), + ("true", True), + ("false", False), + ("not-a-bool", False), + ], + ) + @patch("subprocess.run") + @patch("atexit.register") + @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") + @patch( + "litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False + ) + def test_disable_prepared_statements_forwarded_to_url( + self, + mock_should_update, + mock_setup_db, + mock_atexit_register, + mock_subprocess_run, + config_value, + expect_pgbouncer, + ): + from click.testing import CliRunner + + from litellm.proxy.proxy_cli import run_server + + runner = CliRunner() + mock_subprocess_run.return_value = MagicMock(returncode=0) + + mock_proxy_module = MagicMock( + app=MagicMock(), + ProxyConfig=MagicMock(), + KeyManagementSettings=MagicMock(), + save_worker_config=MagicMock(), + ) + mock_proxy_module.ProxyConfig.return_value.get_config = AsyncMock( + return_value={ + "general_settings": { + "database_url": "postgresql://test:test@localhost:5432/test", + "database_disable_prepared_statements": config_value, + } + } + ) + + clean_env = { + k: v + for k, v in os.environ.items() + if k not in ("DATABASE_URL", "DIRECT_URL") + } + + with ( + patch.dict(os.environ, clean_env, clear=True), + patch.dict( + "sys.modules", + { + "proxy_server": mock_proxy_module, + "litellm.proxy.proxy_server": mock_proxy_module, + }, + ), + patch( + "litellm.proxy.proxy_cli.ProxyInitializationHelpers._get_default_unvicorn_init_args" + ) as mock_get_args, + patch( + "litellm.proxy.proxy_cli.append_query_params", + side_effect=lambda url, params: str(url), + ) as mock_append_query_params, + ): + mock_get_args.return_value = { + "app": "litellm.proxy.proxy_server:app", + "host": "localhost", + "port": 8000, + } + + result = runner.invoke( + run_server, + ["--local", "--config", "test-config.yaml", "--skip_server_startup"], + ) + + assert ( + result.exit_code == 0 + ), f"exit_code={result.exit_code}, output={result.output}" + mock_append_query_params.assert_called() + appended_params = mock_append_query_params.call_args.args[1] + if expect_pgbouncer: + assert appended_params["pgbouncer"] == "true" + else: + assert "pgbouncer" not in appended_params + @patch("uvicorn.run") @patch("atexit.register") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 203a56f615b..8e470819557 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -22012,6 +22012,11 @@ export interface components { * @default 60 */ database_connection_timeout: number | null; + /** + * Database Disable Prepared Statements + * @description Disable server-side prepared statements by setting Prisma's `pgbouncer=true` URL param. Use this for pgbouncer transaction-pooling deployments, or to prevent the 'cached plan must not change result type' error that pooled connections hit during rolling schema migrations. An explicit `pgbouncer` in `database_extra_connection_params` takes precedence. + */ + database_disable_prepared_statements?: boolean | null; /** * Database Extra Connection Params * @description Escape hatch: extra key/value pairs appended verbatim to the Prisma DATABASE_URL / DIRECT_URL query string (e.g. `sslmode`, `pgbouncer`, `statement_cache_size`). Keys here override any default LiteLLM sets. From b301d306c29d442cd2cb47a809a9620a492b038a Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Wed, 10 Jun 2026 16:33:48 -0700 Subject: [PATCH 055/185] fix(release): stop backport releases from overwriting the latest badge (#30005) create-release published every release with GitHub's default make_latest, which is true, so any newly published stable release claimed the repo "Latest" badge regardless of version. That let a backport like 1.84.6 overwrite a newer line like 1.88.1 as latest. Compute make_latest explicitly: a stable release only claims latest when its version is >= the current latest (via getLatestRelease), backports to an older line publish with make_latest false, and prereleases never claim latest. Version comparison accounts for the maintenance suffix (.postN and legacy -stable.patch.N) so within-line ordering stays correct --- .github/workflows/create-release.yml | 33 ++++++++++++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/.github/workflows/create-release.yml b/.github/workflows/create-release.yml index a726a921a2b..4834775e329 100644 --- a/.github/workflows/create-release.yml +++ b/.github/workflows/create-release.yml @@ -52,6 +52,22 @@ jobs: // are stable maintenance releases, not pre-releases. const isPrerelease = /(?:rc|nightly|alpha|beta|[-.]dev)/i.test(tag); + // A stable release should only claim the repo "latest" badge when its + // version is >= the current latest. Otherwise a backport (e.g. 1.84.6) + // would steal "latest" from a newer line (e.g. 1.88.1). + const versionKey = (rawTag) => { + const m = String(rawTag).match(/^v?(\d+)\.(\d+)\.(\d+)/); + if (!m) return null; + const maintenance = String(rawTag).match(/(?:\.post|\.patch\.)(\d+)/i); + return [Number(m[1]), Number(m[2]), Number(m[3]), maintenance ? Number(maintenance[1]) : 0]; + }; + const isAtLeast = (a, b) => { + for (let i = 0; i < a.length; i++) { + if (a[i] !== b[i]) return a[i] > b[i]; + } + return true; + }; + const cosignSection = [ `## Verify Docker Image Signature`, ``, @@ -90,6 +106,22 @@ jobs: ].join('\n'); try { + let makeLatest = "false"; + const newVersion = versionKey(tag); + if (!isPrerelease && newVersion) { + let latestVersion = null; + try { + const latest = await github.rest.repos.getLatestRelease({ + owner: context.repo.owner, + repo: context.repo.repo, + }); + latestVersion = versionKey(latest.data.tag_name); + } catch (error) { + if (error.status !== 404) throw error; + } + makeLatest = (!latestVersion || isAtLeast(newVersion, latestVersion)) ? "true" : "false"; + } + const response = await github.rest.repos.createRelease({ draft: true, generate_release_notes: true, @@ -108,6 +140,7 @@ jobs: release_id: response.data.id, body: updatedBody, draft: false, + make_latest: makeLatest, }); } catch (error) { From ba72ccf52c2483ab1084d6762c67182c8f1913a5 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 10 Jun 2026 16:34:08 -0700 Subject: [PATCH 056/185] feat: add conventional commits and coding guidelines (#30159) * feat: add guideline for conventional commits * feat: add functional programming coding conventions --- CLAUDE.md | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/CLAUDE.md b/CLAUDE.md index 02a9630b486..758eac7e266 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -52,6 +52,19 @@ Do not put names of customers or customer company names in code, PRs, and issues CI supply-chain safety: Never pipe a remote script into a shell (`curl ... | bash`, `wget ... | sh`); download the artifact to a file, verify its SHA-256 checksum, then install. Pin every external tool to a specific version with a full URL (not `latest` or `stable`). Verify checksums for all downloaded binaries, using the provider's official `.sha256` / `.sha256sum` sidecar when available. These rules apply to every download in CI +Follow these coding conventions for new/updated code (a three-line fix in a legacy file shouldn't trigger huge drive-by refactors): + +- Composition over inheritance +- Never-nester: early returns over deep nesting +- Don't throw; model failures as values (One function (e.g., raise_public) maps error union to existing public exception contracts via exhaustive match + assert_never) +- No mutation; instead of mutable lists and dicts, prefer tuples, NamedTuples, frozen dataclasses, etc. +- Use dependency injection +- Fully typed; no `Any` or coarse types like dict[str, Any]. Every function parameter must be strongly typed +- Use tagged unions + match +- No monster files or god objects + +Follow conventional commits for commit names and PR titles + ## Think Before Coding **Don't assume. Don't hide confusion. Surface tradeoffs** From da9d64b4de4b6927d3496f89fa402490a98bfb10 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Wed, 10 Jun 2026 16:48:11 -0700 Subject: [PATCH 057/185] fix(proxy): return 5xx on DB infra errors during auth; reserve 401 for genuine auth failures (#29986) --- litellm/proxy/auth/auth_exception_handler.py | 10 + litellm/proxy/db/exception_handler.py | 86 ++++++++ .../proxy/auth/test_auth_exception_handler.py | 160 ++++++++++++++ .../proxy/auth/test_user_api_key_auth.py | 119 ++++++++++- .../proxy/db/test_exception_handler.py | 195 ++++++++++++++++++ 5 files changed, 569 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/auth/auth_exception_handler.py b/litellm/proxy/auth/auth_exception_handler.py index f76949f4d11..83f18173182 100644 --- a/litellm/proxy/auth/auth_exception_handler.py +++ b/litellm/proxy/auth/auth_exception_handler.py @@ -168,6 +168,16 @@ class UserAPIKeyAuthExceptionHandler: ) elif isinstance(e, ProxyException): raise e + if PrismaDBExceptionHandler.is_database_service_unavailable_error(e): + raise ProxyException( + message=( + "Service Unavailable, the authentication database is " + "temporarily unreachable. Please retry shortly." + ), + type=ProxyErrorTypes.no_db_connection, + param="None", + code=status.HTTP_503_SERVICE_UNAVAILABLE, + ) raise ProxyException( message="Authentication Error, " + str(e), type=ProxyErrorTypes.auth_error, diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index ab9d341aa51..c500e727595 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -109,6 +109,92 @@ class PrismaDBExceptionHandler: return True return False + @staticmethod + def is_prisma_engine_internal_error(e: Exception) -> bool: + """True iff ``e`` is a non-``PrismaError`` exception raised from inside + prisma-client-py's query-engine layer. + + During the instant a DB connection is torn down, the query engine can + return a malformed error payload (``user_facing_error.meta`` is + ``null``). prisma-client-py's ``handle_response_errors`` then crashes + with ``AttributeError: 'NoneType' object has no attribute 'get'`` + before it can raise the proper P1001 "can't reach database server" + error. That AttributeError carries no connection keyword, so it can't + be matched by message; identify it by its ``prisma.engine`` origin + instead. + + Recognized ``PrismaError`` subclasses are excluded: connectivity ones + are already classified by type/keyword above, and data-layer ones + (the DB IS reachable) must stay 401. + """ + import prisma + + if isinstance(e, prisma.errors.PrismaError): + return False + tb = getattr(e, "__traceback__", None) + while tb is not None: + if tb.tb_frame.f_globals.get("__name__", "").startswith("prisma.engine"): + return True + tb = tb.tb_next + return False + + @staticmethod + def is_database_service_unavailable_error(e: Exception) -> bool: + """True iff the exception means the database could not answer at the + infrastructure level (connection refused, socket/interface failure, + timeout) rather than a genuine auth failure (key not found) or a + data-layer error (the DB IS reachable and rejected the data). + + Auth must answer 401 only for a key the DB confirms is invalid. When + the DB itself is unreachable, the request has to surface as 503 so + callers retry instead of treating valid keys as invalid during an + outage. + + Note: prisma-client-py mislabels the P1001 "can't reach database + server" connectivity failure as a ``DataError`` (a data-layer type), + so a type-only check misses real outages. ``is_database_transport_error`` + keyword-matches the connection message and catches that masquerade, + while genuine data errors (no connection keyword) correctly stay 401. + + The Postgres "cached plan must not change result type" error is matched + here, not in ``is_database_transport_error``: it is a transient stale-DB- + state condition (not an invalid key), but the connection is healthy so it + must not trigger a reconnect. + + A non-``PrismaError`` raised from inside the prisma query engine (e.g. + the ``AttributeError`` from ``handle_response_errors`` when the engine + returns a malformed error payload mid-tear-down) is also treated as + unavailable; see ``is_prisma_engine_internal_error``. + """ + import asyncio + + if PrismaDBExceptionHandler.is_database_connection_error(e): + return True + if PrismaDBExceptionHandler.is_database_transport_error(e): + return True + if PrismaDBExceptionHandler.is_prisma_engine_internal_error(e): + return True + if "cached plan must not change result type" in str(e).lower(): + return True + + # OSError already covers ConnectionError and (Py3.3+) TimeoutError. + # asyncio.TimeoutError is a distinct class before Py3.11. + if isinstance(e, (OSError, asyncio.TimeoutError)): + return True + + try: + import asyncpg + except ImportError: + return False + + return isinstance( + e, + ( + asyncpg.exceptions.PostgresConnectionError, + asyncpg.exceptions.InterfaceError, + ), + ) + @staticmethod def handle_db_exception(e: Exception): """ diff --git a/tests/test_litellm/proxy/auth/test_auth_exception_handler.py b/tests/test_litellm/proxy/auth/test_auth_exception_handler.py index 27f6015e6f4..11e6f483e35 100644 --- a/tests/test_litellm/proxy/auth/test_auth_exception_handler.py +++ b/tests/test_litellm/proxy/auth/test_auth_exception_handler.py @@ -112,6 +112,166 @@ async def test_handle_authentication_error_data_layer_errors_do_not_fall_back( ) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "db_error", + [ + ConnectionError("connection refused"), + TimeoutError("timed out"), + asyncio.TimeoutError(), + OSError("network is unreachable"), + HTTPClientClosedError(), + PrismaError("can't reach database server"), + RawQueryError( + data={ + "user_facing_error": { + "message": "cached plan must not change result type", + "meta": {"table": "t"}, + } + } + ), + ], +) +async def test_handle_authentication_error_db_infra_error_returns_503(db_error): + """Regression for the outage where valid keys got 401 for 4 hours: an + infrastructure-level DB failure during auth must surface as 503 (the DB + could not confirm the key), never as 401 ("Invalid API key").""" + handler = UserAPIKeyAuthExceptionHandler() + + with ( + patch( + "litellm.proxy.proxy_server.proxy_logging_obj.post_call_failure_hook", + new_callable=AsyncMock, + return_value=None, + ), + patch( + "litellm.proxy.auth.auth_exception_handler.seed_request_identity", + ), + patch( + "litellm.proxy.proxy_server.general_settings", + {"allow_requests_on_db_unavailable": False}, + ), + ): + with pytest.raises(ProxyException) as exc_info: + await handler._handle_authentication_error( + db_error, + MagicMock(), + {}, + "/v1/chat/completions", + None, + "sk-valid-but-db-down", + ) + + assert int(exc_info.value.code) == status.HTTP_503_SERVICE_UNAVAILABLE + assert exc_info.value.type == ProxyErrorTypes.no_db_connection + assert "Invalid API key" not in str(exc_info.value.message) + + +@pytest.mark.asyncio +async def test_handle_authentication_error_prisma_engine_teardown_returns_503(): + """Regression for the first-request-of-an-outage edge case: at the instant + the DB socket drops, the prisma query engine returns a malformed error + payload and prisma-client-py crashes with a bare + ``AttributeError: 'NoneType' object has no attribute 'get'`` before it can + raise P1001. That AttributeError reached auth and fell through to 401. It + must surface as 503 like every other infra failure during the outage.""" + from prisma.engine import utils as prisma_engine_utils + + malformed_payload = [ + { + "error": "Can't reach database server", + "user_facing_error": { + "error_code": "P1001", + "message": "Can't reach database server at `localhost`:`5503`", + "meta": None, + }, + } + ] + try: + prisma_engine_utils.handle_response_errors(None, malformed_payload) + raise AssertionError("expected prisma to raise AttributeError") + except AttributeError as e: + teardown_error = e + + handler = UserAPIKeyAuthExceptionHandler() + + with ( + patch( + "litellm.proxy.proxy_server.proxy_logging_obj.post_call_failure_hook", + new_callable=AsyncMock, + return_value=None, + ), + patch( + "litellm.proxy.auth.auth_exception_handler.seed_request_identity", + ), + patch( + "litellm.proxy.proxy_server.general_settings", + {"allow_requests_on_db_unavailable": False}, + ), + ): + with pytest.raises(ProxyException) as exc_info: + await handler._handle_authentication_error( + teardown_error, + MagicMock(), + {}, + "/v1/chat/completions", + None, + "sk-valid-but-db-down", + ) + + assert int(exc_info.value.code) == status.HTTP_503_SERVICE_UNAVAILABLE + assert exc_info.value.type == ProxyErrorTypes.no_db_connection + assert "Invalid API key" not in str(exc_info.value.message) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "auth_error", + [ + # DB returned no row -> get_key_object raises this exact 401. + ProxyException( + message="Authentication Error, Invalid proxy server token passed.", + type=ProxyErrorTypes.token_not_found_in_db, + param="key", + code=status.HTTP_401_UNAUTHORIZED, + ), + # A bare auth failure raised as a plain Exception (e.g. master-key-only + # route) must keep returning 401, not get reclassified as 503. + Exception("Invalid proxy server token passed"), + ], +) +async def test_handle_authentication_error_genuine_auth_failure_stays_401(auth_error): + """Guard against the 503 conversion being too broad: a genuine auth + failure (missing key / wrong key) must still be 401.""" + handler = UserAPIKeyAuthExceptionHandler() + + with ( + patch( + "litellm.proxy.proxy_server.proxy_logging_obj.post_call_failure_hook", + new_callable=AsyncMock, + return_value=None, + ), + patch( + "litellm.proxy.auth.auth_exception_handler.seed_request_identity", + ), + patch( + "litellm.proxy.proxy_server.general_settings", + {"allow_requests_on_db_unavailable": False}, + ), + ): + with pytest.raises(ProxyException) as exc_info: + await handler._handle_authentication_error( + auth_error, + MagicMock(), + {}, + "/v1/chat/completions", + None, + "sk-bad-key", + ) + + assert int(exc_info.value.code) == status.HTTP_401_UNAUTHORIZED + + @pytest.mark.asyncio async def test_handle_authentication_error_budget_exceeded(): handler = UserAPIKeyAuthExceptionHandler() diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 0236646c796..80f12d4459f 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -1696,7 +1696,9 @@ class TestJWTOAuth2Coexistence: assert mock_auto_register.call_args.kwargs["team_id"] == "validated-team" assert mock_auto_register.call_args.kwargs["user_id"] == "validated-user" assert mock_auto_register.call_args.kwargs["org_id"] == "validated-org" - assert mock_auto_register.call_args.kwargs["end_user_id"] == "validated-end-user" + assert ( + mock_auto_register.call_args.kwargs["end_user_id"] == "validated-end-user" + ) assert result.org_id == "validated-org" @pytest.mark.asyncio @@ -3608,3 +3610,118 @@ async def test_user_api_key_auth_does_not_overwrite_end_user_id_set_by_builder() finally: for k, v in originals.items(): setattr(_proxy_server_mod, k, v) + + +def _proxy_attrs_for_db_lookup(): + """Minimal proxy_server attributes for driving the real + ``_user_api_key_auth_builder`` down to the DB key lookup.""" + proxy_logging_obj = MagicMock() + proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + return { + "prisma_client": MagicMock(), + "user_api_key_cache": DualCache(), + "proxy_logging_obj": proxy_logging_obj, + "master_key": "sk-test-master", + "general_settings": {"allow_requests_on_db_unavailable": False}, + "llm_model_list": [], + "llm_router": None, + "open_telemetry_logger": None, + "model_max_budget_limiter": MagicMock(), + "user_custom_auth": None, + "jwt_handler": None, + "litellm_proxy_admin_name": "admin", + } + + +async def _run_builder_with_key_lookup(get_key_object_mock): + """Drive the real auth builder with ``get_key_object`` replaced by the + given mock. Returns the builder result. Patches ``seed_request_identity`` + so the failure path doesn't touch OTEL.""" + from fastapi import Request + from starlette.datastructures import URL + + import litellm.proxy.proxy_server as _proxy_server_mod + from litellm.proxy.auth.user_api_key_auth import _user_api_key_auth_builder + + attrs = _proxy_attrs_for_db_lookup() + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + with ( + patch( + "litellm.proxy.auth.user_api_key_auth.get_key_object", + get_key_object_mock, + ), + patch( + "litellm.proxy.auth.auth_exception_handler.seed_request_identity", + ), + ): + return await _user_api_key_auth_builder( + request=request, + api_key="Bearer sk-db-lookup-test", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={}, + ) + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + +@pytest.mark.asyncio +async def test_builder_returns_503_when_db_lookup_raises_infra_error(): + """End-to-end: a DB infrastructure failure during the key lookup must + propagate past the ``except ProxyException`` guard and surface as 503, + not the 401 that masked the 4-hour outage. Killing the new 503 branch + flips this to 401 and fails the test.""" + get_key_object = AsyncMock(side_effect=ConnectionError("connection refused")) + + with pytest.raises(ProxyException) as exc_info: + await _run_builder_with_key_lookup(get_key_object) + + assert int(exc_info.value.code) == status.HTTP_503_SERVICE_UNAVAILABLE + assert exc_info.value.type == ProxyErrorTypes.no_db_connection + assert "Invalid API key" not in str(exc_info.value.message) + + +@pytest.mark.asyncio +async def test_builder_returns_401_when_db_lookup_reports_missing_key(): + """Regression guard: a genuinely missing key (DB returned no row, which + ``get_key_object`` raises as a 401 ProxyException) must still be 401.""" + missing_key_error = ProxyException( + message="Authentication Error, Invalid proxy server token passed. key=..., not found in db.", + type=ProxyErrorTypes.token_not_found_in_db, + param="key", + code=status.HTTP_401_UNAUTHORIZED, + ) + get_key_object = AsyncMock(side_effect=missing_key_error) + + with pytest.raises(ProxyException) as exc_info: + await _run_builder_with_key_lookup(get_key_object) + + assert int(exc_info.value.code) == status.HTTP_401_UNAUTHORIZED + + +@pytest.mark.asyncio +async def test_builder_succeeds_when_db_lookup_returns_valid_token(): + """Regression guard: a valid key still authenticates. Proves the 503 + conversion only fires on the failure path and never intercepts success.""" + valid_token = UserAPIKeyAuth(api_key="sk-db-lookup-test", token="hashed-valid") + get_key_object = AsyncMock(return_value=valid_token) + + with patch( + "litellm.proxy.auth.user_api_key_auth._return_user_api_key_auth_obj", + new_callable=AsyncMock, + return_value=valid_token, + ) as mock_return: + result = await _run_builder_with_key_lookup(get_key_object) + + assert isinstance(result, UserAPIKeyAuth) + # Reaching the success-assembly return (never the exception handler) + # proves a valid key is unaffected by the 503 conversion. + mock_return.assert_awaited_once() diff --git a/tests/test_litellm/proxy/db/test_exception_handler.py b/tests/test_litellm/proxy/db/test_exception_handler.py index 9dcf5df4aeb..6021c221426 100644 --- a/tests/test_litellm/proxy/db/test_exception_handler.py +++ b/tests/test_litellm/proxy/db/test_exception_handler.py @@ -107,6 +107,201 @@ def test_is_database_connection_generic_errors(): ) +@pytest.mark.parametrize( + "error", + [ + ConnectionError("connection refused"), + TimeoutError("timed out"), + OSError("network is unreachable"), + asyncio.TimeoutError(), + HTTPClientClosedError(), + ClientNotConnectedError(), + PrismaError("can't reach database server"), + PrismaError(), + ], +) +def test_is_database_service_unavailable_error_infra_failures(error): + """Infrastructure-level failures (socket/connection/timeout, prisma + transport, unknown PrismaError) mean the DB could not answer, so auth + must surface 503 instead of treating a valid key as invalid.""" + assert PrismaDBExceptionHandler.is_database_service_unavailable_error(error) is True + + +def test_is_database_service_unavailable_error_prisma_p1001_masquerades_as_dataerror(): + """Real-world regression: prisma-client-py raises the P1001 "can't reach + database server" connectivity failure as a DataError (a data-layer type). + A type-only check would miss it and return 401 during a genuine outage; + the message keyword must still classify it as service-unavailable -> 503.""" + p1001_as_dataerror = DataError( + data={ + "user_facing_error": { + "message": "Can't reach database server at `127.0.0.1`:`5499`", + "meta": {"table": "t"}, + } + } + ) + assert ( + PrismaDBExceptionHandler.is_database_service_unavailable_error( + p1001_as_dataerror + ) + is True + ) + + +def test_is_database_service_unavailable_error_cached_plan_escapes_as_503(): + """Composes with the cached-plan retry: when that recovery fails and the + Postgres "cached plan must not change result type" error escapes (raised by + prisma as a data-layer RawQueryError), it is a transient stale-DB-state + condition, not an invalid key, so it must classify as service-unavailable + -> 503 rather than fall through to 401.""" + cached_plan_error = RawQueryError( + data={ + "user_facing_error": { + "message": "cached plan must not change result type", + "meta": {"table": "t"}, + } + } + ) + assert ( + PrismaDBExceptionHandler.is_database_service_unavailable_error( + cached_plan_error + ) + is True + ) + + +def test_is_database_service_unavailable_error_prisma_engine_malformed_payload(): + """Real-world regression: at the instant the DB socket drops, the prisma + query engine returns a malformed error payload (``user_facing_error.meta`` + is ``null``). prisma-client-py's ``handle_response_errors`` then crashes + with ``AttributeError: 'NoneType' object has no attribute 'get'`` before it + can raise the proper P1001 error. That bare AttributeError has no + connection keyword, so without the prisma-engine-origin check it falls + through to 401 on the first request of an outage. Reproduce the exact + prisma crash and assert it classifies as service-unavailable -> 503.""" + from prisma.engine import utils as prisma_engine_utils + + malformed_payload = [ + { + "error": "Can't reach database server", + "user_facing_error": { + "error_code": "P1001", + "message": "Can't reach database server at `localhost`:`5503`", + "meta": None, + }, + } + ] + with pytest.raises(AttributeError) as exc_info: + prisma_engine_utils.handle_response_errors(None, malformed_payload) + + assert "no attribute 'get'" in str(exc_info.value) + assert ( + PrismaDBExceptionHandler.is_database_service_unavailable_error(exc_info.value) + is True + ) + + +def test_is_prisma_engine_internal_error_excludes_application_attributeerror(): + """The prisma-engine-origin check must stay narrow: a genuine AttributeError + raised by application code (a real bug) must NOT be classified as + service-unavailable, otherwise real bugs would silently become 503s.""" + + def application_bug(): + none_value = None + return none_value.get("oops") + + with pytest.raises(AttributeError) as exc_info: + application_bug() + + assert ( + PrismaDBExceptionHandler.is_prisma_engine_internal_error(exc_info.value) + is False + ) + assert ( + PrismaDBExceptionHandler.is_database_service_unavailable_error(exc_info.value) + is False + ) + + +def test_is_prisma_engine_internal_error_excludes_data_layer_prisma_error(): + """A data-layer ``PrismaError`` (the DB IS reachable and rejected the data) + must stay 401. These are always raised from prisma internals, so the check + excludes any ``PrismaError`` by type before inspecting the traceback.""" + data_layer_error = UniqueViolationError( + data={"user_facing_error": {"meta": {"table": "t"}}} + ) + try: + raise data_layer_error + except UniqueViolationError as e: + assert PrismaDBExceptionHandler.is_prisma_engine_internal_error(e) is False + + +@pytest.mark.parametrize( + "error", + [ + DataError(data={"user_facing_error": {"meta": {"table": "t"}}}), + UniqueViolationError(data={"user_facing_error": {"meta": {"table": "t"}}}), + RecordNotFoundError(data={"user_facing_error": {"meta": {"table": "t"}}}), + Exception("some unrelated error"), + ValueError("bad value"), + ], +) +def test_is_database_service_unavailable_error_excludes_non_infra(error): + """Data-layer errors (the DB IS reachable and answered) and generic + non-DB errors must NOT be classified as service-unavailable, otherwise a + genuine 401 would be masked as a transient 503.""" + assert ( + PrismaDBExceptionHandler.is_database_service_unavailable_error(error) is False + ) + + +def test_is_database_service_unavailable_error_asyncpg(monkeypatch): + """asyncpg connection/interface errors map to service-unavailable. asyncpg + is not a hard dependency, so inject a stand-in module to exercise the + branch deterministically regardless of the install environment.""" + import sys + import types + + fake_asyncpg = types.ModuleType("asyncpg") + fake_exceptions = types.ModuleType("asyncpg.exceptions") + + class PostgresConnectionError(Exception): + pass + + class InterfaceError(Exception): + pass + + class UniqueViolationError(Exception): # data-layer, must stay False + pass + + fake_exceptions.PostgresConnectionError = PostgresConnectionError + fake_exceptions.InterfaceError = InterfaceError + fake_exceptions.UniqueViolationError = UniqueViolationError + fake_asyncpg.exceptions = fake_exceptions + + monkeypatch.setitem(sys.modules, "asyncpg", fake_asyncpg) + monkeypatch.setitem(sys.modules, "asyncpg.exceptions", fake_exceptions) + + assert ( + PrismaDBExceptionHandler.is_database_service_unavailable_error( + PostgresConnectionError("connection reset") + ) + is True + ) + assert ( + PrismaDBExceptionHandler.is_database_service_unavailable_error( + InterfaceError("connection was closed") + ) + is True + ) + assert ( + PrismaDBExceptionHandler.is_database_service_unavailable_error( + UniqueViolationError("duplicate key") + ) + is False + ) + + # Test should_allow_request_on_db_unavailable method @patch( "litellm.proxy.proxy_server.general_settings", From 496f5b98598e35873112da91df6dc2e2990c0fce Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 10 Jun 2026 17:16:36 -0700 Subject: [PATCH 058/185] fix(ui): dev server 404s on migrated-page links because uiBase hardcodes /ui (#30169) * fix(ui): serve migrated-page links unprefixed on the dev server migratedHref and legacyPageHref always prepended /ui, which is where the proxy mounts the static export but not where next dev serves the app (basePath is empty; the app lives at the root on localhost:3000). Every sidebar link to a migrated page and every ?page= bookmark redirect therefore 404'd in dev, and would do so for each page cut over in the App Router migration. uiBase now returns the bare root under NODE_ENV=development. The check is inlined at build time, so production output is unchanged for both the default /ui mount and server_root_path deployments. * test(ui): pin NODE_ENV in production-mode migratedPages tests The production-mode describes relied on vitest defaulting NODE_ENV to test; a developer with NODE_ENV=development exported in their shell would see them fail. Stub it explicitly so the suite is deterministic regardless of ambient environment. --- .../src/utils/migratedPages.test.ts | 46 ++++++++++++++++++- .../src/utils/migratedPages.ts | 5 ++ 2 files changed, 50 insertions(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/src/utils/migratedPages.test.ts b/ui/litellm-dashboard/src/utils/migratedPages.test.ts index e9aceea8148..bd4ad7af5b8 100644 --- a/ui/litellm-dashboard/src/utils/migratedPages.test.ts +++ b/ui/litellm-dashboard/src/utils/migratedPages.test.ts @@ -1,8 +1,13 @@ -import { describe, it, expect, vi, beforeEach } from "vitest"; +import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; describe("migratedHref / legacyPageHref", () => { beforeEach(() => { vi.resetModules(); + vi.stubEnv("NODE_ENV", "test"); + }); + + afterEach(() => { + vi.unstubAllEnvs(); }); it("builds a /ui-rooted path when serverRootPath is /", async () => { @@ -37,9 +42,48 @@ describe("migratedHref / legacyPageHref", () => { }); }); +describe("dev server (NODE_ENV=development)", () => { + beforeEach(() => { + vi.resetModules(); + vi.stubEnv("NODE_ENV", "development"); + }); + + afterEach(() => { + vi.unstubAllEnvs(); + }); + + it("builds root-relative hrefs because next dev serves the app at /, not /ui", async () => { + vi.doMock("@/components/networking", () => ({ serverRootPath: "/" })); + const { migratedHref, legacyPageHref } = await import("./migratedPages"); + + expect(migratedHref("api-reference")).toBe("/api-reference"); + expect(legacyPageHref("models")).toBe("/?page=models"); + }); + + it("ignores serverRootPath, which only applies to proxy-mounted deployments", async () => { + vi.doMock("@/components/networking", () => ({ serverRootPath: "/team-x/" })); + const { migratedHref } = await import("./migratedPages"); + + expect(migratedHref("api-reference")).toBe("/api-reference"); + }); + + it("maps a bare migrated path back to its legacy sidebar key", async () => { + vi.doMock("@/components/networking", () => ({ serverRootPath: "/" })); + const { legacyKeyForPathname } = await import("./migratedPages"); + + expect(legacyKeyForPathname("/api-reference/")).toBe("api_ref"); + expect(legacyKeyForPathname("/")).toBeNull(); + }); +}); + describe("legacyKeyForPathname", () => { beforeEach(() => { vi.resetModules(); + vi.stubEnv("NODE_ENV", "test"); + }); + + afterEach(() => { + vi.unstubAllEnvs(); }); it("maps a migrated path back to its legacy sidebar key (including trailing slash)", async () => { diff --git a/ui/litellm-dashboard/src/utils/migratedPages.ts b/ui/litellm-dashboard/src/utils/migratedPages.ts index 2c27e4fee64..d6f1e7d6f1d 100644 --- a/ui/litellm-dashboard/src/utils/migratedPages.ts +++ b/ui/litellm-dashboard/src/utils/migratedPages.ts @@ -15,6 +15,11 @@ export const MIGRATED_PAGES: Record = { }; function uiBase(): string { + // next dev serves the app at the root; only the proxy mounts the static export under /ui + // (and optionally under server_root_path). Inlined at build time, so production is unaffected. + if (process.env.NODE_ENV === "development") { + return ""; + } const root = serverRootPath && serverRootPath !== "/" ? `/${serverRootPath.replace(/^\/+|\/+$/g, "")}` : ""; return `${root}/ui`; } From 4def6916da7aab5d41b05408bb434ec97fdc6b9d Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 10 Jun 2026 18:37:44 -0700 Subject: [PATCH 059/185] refactor(ui): consolidate dashboard to one shell in the (dashboard) layout (#30166) * refactor(ui): consolidate dashboard to one shell in the (dashboard) layout Moves the legacy ?page= switch page into the (dashboard) route group and hoists Navbar, sidebar, ThemeProvider, and DebugWarningBanner into the shared layout with real props, deleting the degraded duplicate shell that wrapped migrated routes. The active page key now derives from the URL at render time, so navigating between legacy and migrated pages no longer remounts the shell. useProxySettings becomes a React Query hook taking accessToken, shared by the navbar, the AdminPanel arm, and migrated pages; this replaces the lifted proxySettings state and the Navbar setProxySettings prop drilling. The invitation onboarding flow (?invitation_id=) keeps rendering without chrome via a layout escape hatch. Dead dark mode state and the no-op antd ConfigProvider are removed. * fix(ui): include accessToken in useProxySettings query key The queryFn closes over accessToken, so the key must include it for the cache to be honest about its inputs. Settings are instance-global today, which made the omission harmless, but a token change while mounted would have served the cached entry without refetching. * test(ui): point CreateKeyPage test at the moved page The page moved into the (dashboard) route group and no longer renders the navbar (the layout owns chrome now), so the valid-token test asserts the default page content (UserDashboard stub) instead. --- ui/litellm-dashboard/eslint-suppressions.json | 49 -- .../app/(dashboard)/api-reference/page.tsx | 4 +- .../hooks/proxySettings/useProxySettings.ts | 39 +- .../src/app/(dashboard)/layout.tsx | 67 +-- .../src/app/{ => (dashboard)}/page.tsx | 430 ++++++++---------- .../src/components/navbar.test.tsx | 11 +- .../src/components/navbar.tsx | 30 +- .../src/components/public_model_hub.tsx | 10 +- .../src/components/routing_groups/index.tsx | 4 +- .../tests/CreateKeyPage.expiredToken.test.tsx | 8 +- 10 files changed, 256 insertions(+), 396 deletions(-) rename ui/litellm-dashboard/src/app/{ => (dashboard)}/page.tsx (55%) diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 233741652a9..d3169395b4e 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -164,11 +164,6 @@ "count": 2 } }, - "src/app/(dashboard)/layout.tsx": { - "react-hooks/set-state-in-effect": { - "count": 1 - } - }, "src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx": { "no-restricted-imports": { "count": 1 @@ -228,11 +223,6 @@ "count": 1 } }, - "src/app/page.tsx": { - "unused-imports/no-unused-imports": { - "count": 2 - } - }, "src/components/AIHub/AgentHubTableColumns.test.tsx": { "unused-imports/no-unused-imports": { "count": 1 @@ -243,14 +233,6 @@ "count": 1 } }, - "src/components/AIHub/ClaudeCodeMarketplaceTab.tsx": { - "no-restricted-imports": { - "count": 1 - }, - "react-hooks/immutability": { - "count": 1 - } - }, "src/components/AIHub/ModelHubTable.test.tsx": { "max-params": { "count": 1 @@ -303,11 +285,6 @@ "count": 1 } }, - "src/components/AIHub/marketplace_table_columns.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/components/AdminPanel.tsx": { "no-restricted-imports": { "count": 1 @@ -816,11 +793,6 @@ "count": 1 } }, - "src/components/agents/agent_table.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, "src/components/alerting/dynamic_form.tsx": { "no-restricted-imports": { "count": 1 @@ -956,14 +928,6 @@ "count": 1 } }, - "src/components/claude_code_plugins/plugin_info.tsx": { - "no-restricted-imports": { - "count": 1 - }, - "react-hooks/immutability": { - "count": 1 - } - }, "src/components/claude_code_plugins/plugin_table.tsx": { "no-restricted-imports": { "count": 1 @@ -1333,14 +1297,6 @@ "count": 2 } }, - "src/components/mcp_tools/mcp_server_columns.tsx": { - "max-params": { - "count": 1 - }, - "no-restricted-imports": { - "count": 1 - } - }, "src/components/mcp_tools/mcp_server_cost_config.tsx": { "no-restricted-imports": { "count": 1 @@ -1489,11 +1445,6 @@ "count": 1 } }, - "src/components/navbar.tsx": { - "react-hooks/set-state-in-effect": { - "count": 1 - } - }, "src/components/networking.tsx": { "max-params": { "count": 23 diff --git a/ui/litellm-dashboard/src/app/(dashboard)/api-reference/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/api-reference/page.tsx index 02bed1adbe5..a4a4d3d0f43 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/api-reference/page.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/api-reference/page.tsx @@ -1,10 +1,12 @@ "use client"; import APIReferenceView from "@/app/(dashboard)/api-reference/APIReferenceView"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import useProxySettings from "@/app/(dashboard)/hooks/proxySettings/useProxySettings"; const APIReferencePage = () => { - const proxySettings = useProxySettings(); + const { accessToken } = useAuthorized(); + const proxySettings = useProxySettings(accessToken); return ; }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxySettings/useProxySettings.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxySettings/useProxySettings.ts index d4fb3073856..82cefd800f4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxySettings/useProxySettings.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxySettings/useProxySettings.ts @@ -1,21 +1,26 @@ -import { useState, useEffect } from "react"; import { fetchProxySettings } from "@/utils/proxyUtils"; -import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { useQuery } from "@tanstack/react-query"; +import { createQueryKeys } from "../common/queryKeysFactory"; -export default function useProxySettings() { - const { accessToken } = useAuthorized(); - const [proxySettings, setProxySettings] = useState({ - PROXY_BASE_URL: "", - PROXY_LOGOUT_URL: "", - LITELLM_UI_API_DOC_BASE_URL: null as string | null, - }); +export const proxySettingsKeys = createQueryKeys("proxySettings"); - useEffect(() => { - if (!accessToken) return; - fetchProxySettings(accessToken).then((settings) => { - if (settings) setProxySettings(settings); - }); - }, [accessToken]); - - return proxySettings; +export interface ProxySettings { + PROXY_BASE_URL: string; + PROXY_LOGOUT_URL: string; + LITELLM_UI_API_DOC_BASE_URL?: string | null; +} + +const EMPTY_PROXY_SETTINGS: ProxySettings = { + PROXY_BASE_URL: "", + PROXY_LOGOUT_URL: "", + LITELLM_UI_API_DOC_BASE_URL: null, +}; + +export default function useProxySettings(accessToken: string | null): ProxySettings { + const { data } = useQuery({ + queryKey: [...proxySettingsKeys.all, accessToken], + queryFn: () => fetchProxySettings(accessToken), + enabled: Boolean(accessToken), + }); + return data ?? EMPTY_PROXY_SETTINGS; } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx index 5f5c240d025..df5b2ab4511 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx @@ -1,62 +1,63 @@ "use client"; -import React, { Suspense, useEffect, useState } from "react"; +import React, { Suspense, useState } from "react"; import Navbar from "@/components/navbar"; +import LoadingScreen from "@/components/common_components/LoadingScreen"; import { ThemeProvider } from "@/contexts/ThemeContext"; +import { useAuth } from "@/contexts/AuthContext"; import SidebarProvider from "@/app/(dashboard)/components/SidebarProvider"; -import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { useRouter, useSearchParams, usePathname } from "next/navigation"; import { DebugWarningBanner } from "@/components/DebugWarningBanner"; import { MIGRATED_PAGES, migratedHref, legacyPageHref, legacyKeyForPathname } from "@/utils/migratedPages"; -function LayoutContent({ children }: { children: React.ReactNode }) { +function DashboardShell({ children }: { children: React.ReactNode }) { const router = useRouter(); const searchParams = useSearchParams(); const pathname = usePathname(); - const { accessToken } = useAuthorized(); - const [sidebarCollapsed, setSidebarCollapsed] = React.useState(false); - const [page, setPage] = useState(() => { - return legacyKeyForPathname(pathname) || searchParams.get("page") || "api-keys"; - }); + const { accessToken } = useAuth(); + const [sidebarCollapsed, setSidebarCollapsed] = useState(false); - const handleSetPage = (newPage: string) => { + const page = legacyKeyForPathname(pathname) || searchParams.get("page") || "api-keys"; + + const navigateToPage = (newPage: string) => { const migratedRoute = MIGRATED_PAGES[newPage]; router.push(migratedRoute ? migratedHref(migratedRoute) : legacyPageHref(newPage)); - setPage(newPage); }; - useEffect(() => { - setPage(legacyKeyForPathname(pathname) || searchParams.get("page") || "api-keys"); - }, [pathname, searchParams]); + return ( +
+ setSidebarCollapsed((v) => !v)} + /> + +
+
+ +
+
{children}
+
+
+ ); +} - const toggleSidebar = () => setSidebarCollapsed((v) => !v); +function LayoutContent({ children }: { children: React.ReactNode }) { + const searchParams = useSearchParams(); + const { accessToken } = useAuth(); + const isInvitationFlow = Boolean(searchParams.get("invitation_id")); return ( - -
- {}} - accessToken={accessToken} - /> - -
-
- -
-
{children}
-
-
+ + {isInvitationFlow ? children : {children}} ); } export default function Layout({ children }: { children: React.ReactNode }) { return ( - Loading...
}> + }> {children} ); diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/page.tsx similarity index 55% rename from ui/litellm-dashboard/src/app/page.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/page.tsx index 81cce930998..0854b085fae 100644 --- a/ui/litellm-dashboard/src/app/page.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/page.tsx @@ -1,7 +1,6 @@ "use client"; -import SidebarProvider from "@/app/(dashboard)/components/SidebarProvider"; -import OldModelDashboard from "@/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView"; +import ModelsAndEndpointsView from "@/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView"; import PlaygroundPage from "@/app/(dashboard)/playground/page"; import AdminPanel from "@/components/AdminPanel"; import AgentsPanel from "@/components/agents"; @@ -10,6 +9,7 @@ import CacheDashboard from "@/components/cache_dashboard"; import ClaudeCodePluginsPanel from "@/components/claude_code_plugins"; import { teamListCall as v2TeamListCall } from "@/app/(dashboard)/hooks/teams/useTeams"; import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings"; +import useProxySettings from "@/app/(dashboard)/hooks/proxySettings/useProxySettings"; import LoadingScreen from "@/components/common_components/LoadingScreen"; import { CostTrackingSettings } from "@/components/CostTrackingSettings"; import GeneralSettings from "@/components/general_settings"; @@ -19,7 +19,6 @@ import PoliciesPanel from "@/components/policies"; import { Team } from "@/components/key_team_helpers/key_list"; import { MCPServers } from "@/components/mcp_tools"; import ModelHubTable from "@/components/AIHub/ModelHubTable"; -import Navbar from "@/components/navbar"; import { Organization, proxyBaseUrl, getInProductNudgesCall } from "@/components/networking"; import NewUsagePage from "@/components/UsagePage/components/UsagePageView"; import OldTeams from "@/components/OldTeams"; @@ -44,7 +43,6 @@ import { MemoryView } from "@/components/MemoryView"; import WorkflowRuns from "@/components/workflow_runs"; import SpendLogsTable from "@/components/view_logs"; import ViewUserDashboard from "@/components/view_users"; -import { ThemeProvider } from "@/contexts/ThemeContext"; import { useAuth } from "@/contexts/AuthContext"; import { buildLoginUrlWithReturn, @@ -55,16 +53,8 @@ import { } from "@/utils/returnUrlUtils"; import { isAdminRole } from "@/utils/roles"; import { MIGRATED_PAGES, migratedHref } from "@/utils/migratedPages"; -import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; import { useRouter, useSearchParams } from "next/navigation"; import { Suspense, useEffect, useMemo, useRef, useState } from "react"; -import { ConfigProvider, theme } from "antd"; - -interface ProxySettings { - PROXY_BASE_URL: string; - PROXY_LOGOUT_URL: string; - LITELLM_UI_API_DOC_BASE_URL?: string | null; -} function CreateKeyPageContent() { const { authLoading, token, userID, userRole, userEmail, accessToken, premiumUser, setUserRole, setUserEmail } = @@ -74,10 +64,7 @@ function CreateKeyPageContent() { const [keys, setKeys] = useState([]); const [organizations, setOrganizations] = useState([]); const [userModels, setUserModels] = useState([]); - const [proxySettings, setProxySettings] = useState({ - PROXY_BASE_URL: "", - PROXY_LOGOUT_URL: "", - }); + const proxySettings = useProxySettings(accessToken); const router = useRouter(); const searchParams = useSearchParams()!; @@ -96,12 +83,6 @@ function CreateKeyPageContent() { const [showClaudeCodePrompt, setShowClaudeCodePrompt] = useState(false); const [showClaudeCodeModal, setShowClaudeCodeModal] = useState(false); - // Dark mode state - const [isDarkMode, setIsDarkMode] = useState(false); - const toggleDarkMode = () => { - setIsDarkMode(!isDarkMode); - }; - const invitation_id = searchParams.get("invitation_id"); // Parse URL query parameters for pre-filling the create key form @@ -154,33 +135,11 @@ function CreateKeyPageContent() { }; }, [searchParams, autoOpenCreate]); - // Get page from URL, default to 'api-keys' if not present - const [page, setPage] = useState(() => { - return searchParams.get("page") || "api-keys"; - }); - - const updatePage = (newPage: string) => { - const migratedRoute = MIGRATED_PAGES[newPage]; - if (migratedRoute) { - router.push(migratedHref(migratedRoute)); - setPage(newPage); - return; - } - const newSearchParams = new URLSearchParams(searchParams); - newSearchParams.set("page", newPage); - window.history.pushState(null, "", `?${newSearchParams.toString()}`); - setPage(newPage); - }; - - const [sidebarCollapsed, setSidebarCollapsed] = useState(false); + const page = searchParams.get("page") || "api-keys"; // Track if we've already attempted a return URL redirect to prevent race conditions const hasAttemptedReturnRedirectRef = useRef(false); - const toggleSidebar = () => { - setSidebarCollapsed(!sidebarCollapsed); - }; - const addKey = (data: any) => { setKeys((prevData) => (prevData ? [...prevData, data] : [data])); setCreateClicked(() => !createClicked); @@ -349,14 +308,26 @@ function CreateKeyPageContent() { } return ( - }> - - - {invitation_id ? ( + <> + {invitation_id ? ( + + ) : ( + <> + {page == "api-keys" ? ( - ) : ( -
- + ) : page == "llm-playground" ? ( + + ) : page == "users" ? ( + + ) : page == "teams" ? ( + + ) : page == "organizations" ? ( + + ) : page == "admin-panel" ? ( + + ) : page == "logging-and-alerts" ? ( + + ) : page == "budgets" ? ( + + ) : page == "guardrails" ? ( + + ) : page == "policies" ? ( + + ) : page == "agents" ? ( + + ) : page == "prompts" ? ( + + ) : page == "transform-request" ? ( + + ) : page == "router-settings" ? ( + + ) : page == "ui-theme" ? ( + + ) : page == "cost-tracking" ? ( + + ) : page == "model-hub-table" ? ( + isAdminRole(userRole) ? ( + -
-
- -
- {page == "api-keys" ? ( - - ) : page == "models" ? ( - - ) : page == "llm-playground" ? ( - - ) : page == "users" ? ( - - ) : page == "teams" ? ( - - ) : page == "organizations" ? ( - - ) : page == "admin-panel" ? ( - - ) : page == "logging-and-alerts" ? ( - - ) : page == "budgets" ? ( - - ) : page == "guardrails" ? ( - - ) : page == "policies" ? ( - - ) : page == "agents" ? ( - - ) : page == "prompts" ? ( - - ) : page == "transform-request" ? ( - - ) : page == "router-settings" ? ( - - ) : page == "ui-theme" ? ( - - ) : page == "cost-tracking" ? ( - - ) : page == "model-hub-table" ? ( - isAdminRole(userRole) ? ( - - ) : ( - - ) - ) : page == "caching" ? ( - - ) : page == "pass-through-settings" ? ( - - ) : page == "logs" ? ( - - ) : page == "mcp-servers" ? ( - - ) : page == "search-tools" ? ( - - ) : page == "tag-management" ? ( - - ) : page == "skills" || page == "claude-code-plugins" ? ( - - ) : page == "access-groups" ? ( - - ) : page == "projects" ? ( - - ) : page == "vector-stores" ? ( - - ) : page == "tool-policies" ? ( - - ) : page == "workflows" ? ( - - ) : page == "memory" ? ( - - ) : page == "guardrails-monitor" ? ( - - ) : page == "new_usage" ? ( - - ) : ( - - )} -
- - {/* Survey Components */} - - - - {/* Claude Code Components */} - - -
+ ) : ( + + ) + ) : page == "caching" ? ( + + ) : page == "pass-through-settings" ? ( + + ) : page == "logs" ? ( + + ) : page == "mcp-servers" ? ( + + ) : page == "search-tools" ? ( + + ) : page == "tag-management" ? ( + + ) : page == "skills" || page == "claude-code-plugins" ? ( + + ) : page == "access-groups" ? ( + + ) : page == "projects" ? ( + + ) : page == "vector-stores" ? ( + + ) : page == "tool-policies" ? ( + + ) : page == "workflows" ? ( + + ) : page == "memory" ? ( + + ) : page == "guardrails-monitor" ? ( + + ) : page == "new_usage" ? ( + + ) : ( + )} - - - + + {/* Survey Components */} + + + + {/* Claude Code Components */} + + + + )} + ); } diff --git a/ui/litellm-dashboard/src/components/navbar.test.tsx b/ui/litellm-dashboard/src/components/navbar.test.tsx index 274e81db527..bdd0681dfa8 100644 --- a/ui/litellm-dashboard/src/components/navbar.test.tsx +++ b/ui/litellm-dashboard/src/components/navbar.test.tsx @@ -72,7 +72,10 @@ vi.mock("./Navbar/UserDropdown/UserDropdown", async (importOriginal) => { }); vi.mock("@/utils/proxyUtils", () => ({ - fetchProxySettings: vi.fn(), + fetchProxySettings: vi.fn().mockResolvedValue({ + PROXY_BASE_URL: "", + PROXY_LOGOUT_URL: "https://example.com/logout", + }), })); // Mock CommunityEngagementButtons component @@ -137,8 +140,6 @@ Object.defineProperty(window, "location", { describe("Navbar", () => { const defaultProps = { - proxySettings: {}, - setProxySettings: vi.fn(), accessToken: "test-token", isPublicPage: false, }; @@ -298,7 +299,9 @@ describe("Navbar", () => { const cookieUtils = vi.mocked(await import("@/utils/cookieUtils")); expect(cookieUtils.clearTokenCookies).toHaveBeenCalled(); - expect(window.location.href).toBe(""); + await waitFor(() => { + expect(window.location.href).toBe("https://example.com/logout"); + }); }); it("should not render dark mode toggle slider", () => { diff --git a/ui/litellm-dashboard/src/components/navbar.tsx b/ui/litellm-dashboard/src/components/navbar.tsx index e5a1490788c..d9fa0c59e6e 100644 --- a/ui/litellm-dashboard/src/components/navbar.tsx +++ b/ui/litellm-dashboard/src/components/navbar.tsx @@ -6,11 +6,11 @@ import { getProxyBaseUrl } from "@/components/networking"; import { useTheme } from "@/contexts/ThemeContext"; import { clearTokenCookies } from "@/utils/cookieUtils"; import { clearStoredReturnUrl } from "@/utils/returnUrlUtils"; -import { fetchProxySettings } from "@/utils/proxyUtils"; +import useProxySettings from "@/app/(dashboard)/hooks/proxySettings/useProxySettings"; import { DownOutlined, MenuFoldOutlined, MenuUnfoldOutlined } from "@ant-design/icons"; import { Tag } from "antd"; import Link from "next/link"; -import React, { useEffect, useState } from "react"; +import React from "react"; import { BlogDropdown } from "./Navbar/BlogDropdown/BlogDropdown"; import { CommunityEngagementButtons } from "./Navbar/CommunityEngagementButtons/CommunityEngagementButtons"; import { NAV_PRODUCT_LINK_CLASS } from "./Navbar/navProductLinkClass"; @@ -19,8 +19,6 @@ import UserDropdown from "./Navbar/UserDropdown/UserDropdown"; import WorkerDropdown from "./Navbar/WorkerDropdown/WorkerDropdown"; interface NavbarProps { - proxySettings: any; - setProxySettings: React.Dispatch>; accessToken: string | null; isPublicPage: boolean; sidebarCollapsed?: boolean; @@ -28,15 +26,13 @@ interface NavbarProps { } const Navbar: React.FC = ({ - proxySettings, - setProxySettings, accessToken, isPublicPage = false, sidebarCollapsed = false, onToggleSidebar, }) => { const baseUrl = getProxyBaseUrl(); - const [logoutUrl, setLogoutUrl] = useState(""); + const proxySettings = useProxySettings(accessToken); const { logoUrl } = useTheme(); const { data: healthData } = useHealthReadinessDetails(accessToken); const version = healthData?.litellm_version; @@ -47,29 +43,11 @@ const Navbar: React.FC = ({ const imageUrl = logoUrl || `${baseUrl}/get_image`; - useEffect(() => { - const initializeProxySettings = async () => { - if (accessToken) { - const settings = await fetchProxySettings(accessToken); - console.log("response from fetchProxySettings", settings); - if (settings) { - setProxySettings(settings); - } - } - }; - - initializeProxySettings(); - }, [accessToken]); - - useEffect(() => { - setLogoutUrl(proxySettings?.PROXY_LOGOUT_URL || ""); - }, [proxySettings]); - const handleLogout = () => { clearTokenCookies(); localStorage.removeItem("litellm_selected_worker_id"); localStorage.removeItem("litellm_worker_url"); - window.location.href = logoutUrl; + window.location.href = proxySettings.PROXY_LOGOUT_URL || ""; }; const handleWorkerSwitch = (workerId: string) => { diff --git a/ui/litellm-dashboard/src/components/public_model_hub.tsx b/ui/litellm-dashboard/src/components/public_model_hub.tsx index 0421e62d50d..6ad8db19d8d 100644 --- a/ui/litellm-dashboard/src/components/public_model_hub.tsx +++ b/ui/litellm-dashboard/src/components/public_model_hub.tsx @@ -121,7 +121,6 @@ const PublicModelHub: React.FC = ({ accessToken, isEmbedded const [selectedModel, setSelectedModel] = useState(null); const [selectedAgent, setSelectedAgent] = useState(null); const [selectedMcpServer, setSelectedMcpServer] = useState(null); - const [proxySettings, setProxySettings] = useState({}); const [activeTab, setActiveTab] = useState("models"); const [skillHubData, setSkillHubData] = useState([]); const [skillLoading, setSkillLoading] = useState(false); @@ -981,14 +980,7 @@ const PublicModelHub: React.FC = ({ accessToken, isEmbedded
{/* Navigation - only show when not embedded */} - {!isEmbedded && ( - - )} + {!isEmbedded && }
{/* Embedded Explainer - only shown when embedded in dashboard */} diff --git a/ui/litellm-dashboard/src/components/routing_groups/index.tsx b/ui/litellm-dashboard/src/components/routing_groups/index.tsx index cdd2b30c2ed..1ee281bd92a 100644 --- a/ui/litellm-dashboard/src/components/routing_groups/index.tsx +++ b/ui/litellm-dashboard/src/components/routing_groups/index.tsx @@ -6,6 +6,7 @@ import { PlusOutlined, ReloadOutlined, SearchOutlined } from "@ant-design/icons" import { useRoutingGroups, useSaveRoutingGroups } from "@/app/(dashboard)/hooks/routingGroups/useRoutingGroups"; import { useRouterFields } from "@/app/(dashboard)/hooks/router/useRouterFields"; import { useModelHub } from "@/app/(dashboard)/hooks/models/useModels"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import useProxySettings from "@/app/(dashboard)/hooks/proxySettings/useProxySettings"; import RoutingGroupsTable from "./RoutingGroupsTable"; import RoutingGroupModal from "./RoutingGroupModal"; @@ -18,7 +19,8 @@ const RoutingGroups: React.FC = () => { const { data, isLoading, refetch, isFetching } = useRoutingGroups(); const { data: routerFields } = useRouterFields(); const { data: modelHub } = useModelHub(); - const proxySettings = useProxySettings(); + const { accessToken } = useAuthorized(); + const proxySettings = useProxySettings(accessToken); const saveMutation = useSaveRoutingGroups(); const [searchQuery, setSearchQuery] = useState(""); diff --git a/ui/litellm-dashboard/tests/CreateKeyPage.expiredToken.test.tsx b/ui/litellm-dashboard/tests/CreateKeyPage.expiredToken.test.tsx index 2ebde4f295d..60475380b14 100644 --- a/ui/litellm-dashboard/tests/CreateKeyPage.expiredToken.test.tsx +++ b/ui/litellm-dashboard/tests/CreateKeyPage.expiredToken.test.tsx @@ -149,7 +149,7 @@ vi.mock("@/lib/cva.config", () => ({ })); import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; -import CreateKeyPage from "@/app/page"; +import CreateKeyPage from "@/app/(dashboard)/page"; import { AuthProvider } from "@/contexts/AuthContext"; // The page consumes auth state via useAuth(). Wrap it so the hook resolves @@ -242,7 +242,7 @@ describe("CreateKeyPage auth behavior", () => { expect(wroteDeletion).toBe(true); }); - it("does NOT redirect when token is valid and renders the app chrome", async () => { + it("does NOT redirect when token is valid and renders the page content", async () => { // Arrange: valid token in cookie setCookie("token=validtoken"); @@ -269,9 +269,9 @@ describe("CreateKeyPage auth behavior", () => { expect(window.location.replace).not.toHaveBeenCalled(); }); - // And some top-level UI appears (Navbar stub) + // And the default page content appears (UserDashboard stub; chrome now lives in the layout) await waitFor(() => { - expect(screen.getByTestId("navbar")).toBeInTheDocument(); + expect(screen.getByTestId("user-dashboard")).toBeInTheDocument(); }); }); From 6068bb7781b66ea51930f68ed8738ac46f0bdf7d Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Thu, 11 Jun 2026 08:08:21 +0530 Subject: [PATCH 060/185] fix(proxy): align /v1/model/info with router deployments (#30025) * fix(proxy): align /v1/model/info with router deployments Return router model_list entries (including team-scoped models) with team access metadata instead of wildcard-expanded names from get_complete_model_list. Co-authored-by: Cursor * fix(proxy): gate v1 team filter and honor key allowlists Only apply get_all_team_and_direct_access_models for admin or user-bound keys, then intersect with key/team model restrictions to avoid empty lists for service tokens and metadata leaks for restricted keys. Co-authored-by: Cursor * fix(proxy): skip v1 team filter when user row is missing Require a DB-backed user before applying team-access filtering on /v1/model/info, and skip the trailing filter in get_all_team_and_direct_access_models when user context cannot be resolved. Co-authored-by: Cursor * Revert "fix(proxy): skip v1 team filter when user row is missing" This reverts commit 74e1fbd77a981103cd9a4ed1cbdd662f5cbcf209. * fix(proxy): restore legacy v1 model access filtering Keep /v1/model/info on key/team allowlists instead of DB team-membership filtering, while still listing router deployments for team-scoped models. Co-authored-by: Cursor * fix(proxy): drop A2A agent entries from public /v1/model/info list * fix(proxy): scope team BYOK rows on /v1/model/info to caller's teams Listing the full router model_list let any authenticated key without explicit model restrictions enumerate other teams' BYOK deployments (public name, team_id, api_base) via /v1/model/info. Reuse the existing _get_caller_byok_team_scope check so non-admin callers only see global deployments plus their own team's BYOK rows; admins keep the full view. --------- Co-authored-by: Cursor Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 161 ++++++++++----- .../test_team_model_name_translation.py | 185 +++++++++++++++++- .../proxy/test_model_info_default_limits.py | 3 +- tests/test_litellm/proxy/test_proxy_server.py | 3 +- 4 files changed, 294 insertions(+), 58 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 96c9cd1e8fb..267e388112d 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11140,6 +11140,22 @@ async def _get_caller_byok_team_scope( return set(user_row.teams or []) +def _byok_row_outside_caller_teams( + model_info_dict: Dict[str, Any], allowed_team_ids: Optional[Set[str]] +) -> bool: + """Whether a team BYOK row belongs to a team the caller is not a member of. + + `team_id` is only set on team BYOK rows; non-team rows fall through + unaffected. `allowed_team_ids is None` means no scoping (e.g. admins). + """ + if allowed_team_ids is None: + return False + team_id = model_info_dict.get("team_id") + if team_id is None: + return False + return team_id not in allowed_team_ids + + # Hard cap on rows the DB-side BYOK search may pull when results need to be # sorted across the full match set. Without this, an authenticated caller # can hit `/v2/model/info?search=&sortBy=` and force the @@ -11261,15 +11277,7 @@ async def _apply_search_filter_to_models( ) def _is_byok_outside_caller_teams(model_info_dict: Dict[str, Any]) -> bool: - # `team_id` is only set on team BYOK rows. Non-team rows fall - # through unaffected — they are gated by other paths (router - # membership, direct_access, include_team_models). - if allowed_team_ids is None: - return False - team_id = model_info_dict.get("team_id") - if team_id is None: - return False - return team_id not in allowed_team_ids + return _byok_row_outside_caller_teams(model_info_dict, allowed_team_ids) def _model_matches_search(m: Dict[str, Any]) -> bool: # Team BYOK models persist an internal `model_name` @@ -12409,6 +12417,72 @@ async def model_metrics_exceptions( return {"data": response, "exception_types": list(exception_types)} +def _deployment_matches_allowed_model_names( + model: Dict[str, Any], allowed_model_names: Set[str] +) -> bool: + """Match a router deployment against allowed public model names. + + Team-scoped rows store an internal routing key in ``model_name``; callers + with key/team restrictions still refer to the public name in + ``model_info.team_public_model_name``. + """ + if model.get("model_name") in allowed_model_names: + return True + model_info = model.get("model_info") + if not isinstance(model_info, dict): + return False + team_public_model_name = model_info.get("team_public_model_name") + return ( + isinstance(team_public_model_name, str) + and team_public_model_name in allowed_model_names + ) + + +def _get_v1_model_info_allowed_model_names( + user_api_key_dict: UserAPIKeyAuth, + llm_router: Router, +) -> Optional[Set[str]]: + """Return key/team allowlisted public model names, or None if unrestricted.""" + model_access_groups = llm_router.get_model_access_groups() + proxy_model_list = llm_router.get_model_names() + key_models = get_key_models( + user_api_key_dict=user_api_key_dict, + proxy_model_list=proxy_model_list, + model_access_groups=model_access_groups, + ) + team_models = get_team_models( + team_models=user_api_key_dict.team_models, + proxy_model_list=proxy_model_list, + model_access_groups=model_access_groups, + ) + if not key_models and not team_models: + return None + return set( + get_complete_model_list( + key_models=key_models, + team_models=team_models, + proxy_model_list=proxy_model_list, + user_model=user_model, + infer_model_from_keys=general_settings.get("infer_model_from_keys", False), + llm_router=llm_router, + return_wildcard_routes=False, + ) + ) + + +def _filter_v1_model_info_deployments( + all_models: List[dict], + allowed_model_names: Optional[Set[str]], +) -> List[dict]: + if allowed_model_names is None: + return all_models + return [ + model + for model in all_models + if _deployment_matches_allowed_model_names(model, allowed_model_names) + ] + + def _translate_model_name_for_response(model: dict) -> dict: """For team-scoped DB rows, replace `model_name` with the public name in `model_info.team_public_model_name` before returning. The DB column @@ -12578,49 +12652,42 @@ async def model_info_v1( # noqa: PLR0915 ) return {"data": [_deployment_info_dict]} - all_models: List[dict] = [] - model_access_groups: Dict[str, List[str]] = defaultdict(list) - ## CHECK IF MODEL RESTRICTIONS ARE SET AT KEY/TEAM LEVEL ## - if llm_router is None: - proxy_model_list = [] - else: - proxy_model_list = llm_router.get_model_names() - model_access_groups = llm_router.get_model_access_groups() - key_models = get_key_models( + # Return router deployments (same source as /v2/model/info), not wildcard- + # expanded model names from get_complete_model_list(). Team-scoped rows + # use internal routing keys (model_name_{team_id}_{uuid}) and were omitted + # when v1 resolved models only via public model_name strings. + all_models: List[dict] = copy.deepcopy(llm_router.model_list) + allowed_model_names = _get_v1_model_info_allowed_model_names( user_api_key_dict=user_api_key_dict, - proxy_model_list=proxy_model_list, - model_access_groups=model_access_groups, - ) - team_models = get_team_models( - team_models=user_api_key_dict.team_models, - proxy_model_list=proxy_model_list, - model_access_groups=model_access_groups, - ) - all_models_str = get_complete_model_list( - key_models=key_models, - team_models=team_models, - proxy_model_list=proxy_model_list, - user_model=user_model, - infer_model_from_keys=general_settings.get("infer_model_from_keys", False), llm_router=llm_router, ) - if len(all_models_str) > 0: - _relevant_models = [] - for model in all_models_str: - router_models = llm_router.get_model_list(model_name=model) - if router_models is not None: - _relevant_models.extend(router_models) - if llm_model_list is not None: - all_models = copy.deepcopy(_relevant_models) # type: ignore - else: - all_models = [] + all_models = _filter_v1_model_info_deployments( + all_models=all_models, + allowed_model_names=allowed_model_names, + ) - # Reassign each entry: _get_proxy_model_info returns a (possibly new) - # dict via _translate_model_name_for_response, which does NOT mutate in - # place. Binding only the loop variable would drop the public-name swap - # for team-scoped rows and leak the internal routing key (#28382). - all_models = [_get_proxy_model_info(model=model) for model in all_models] + # Team BYOK deployments carry an internal routing key and other teams' + # public name/team_id/api_base; drop the ones the caller cannot access so + # listing the full router model_list does not leak cross-team metadata. + allowed_team_ids = await _get_caller_byok_team_scope( + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + ) + all_models = [ + model + for model in all_models + if not _byok_row_outside_caller_teams( + model.get("model_info") or {}, allowed_team_ids + ) + ] + + all_models = [ + _translate_model_name_for_response( + _enrich_model_info_with_litellm_data(model=model, llm_router=llm_router) + ) + for model in all_models + ] verbose_proxy_logger.debug("all_models: %s", all_models) return {"data": all_models} diff --git a/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py b/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py index 97e5c494916..9757999c85e 100644 --- a/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py +++ b/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py @@ -151,21 +151,24 @@ async def test_model_info_v2_translates_team_model_name(monkeypatch): @pytest.mark.asyncio async def test_model_info_v1_list_path_translates_team_model_name(monkeypatch): - """/v1/model/info list path (no litellm_model_id) must surface the public - name. Covers the list comprehension that assigns _get_proxy_model_info's - return back into all_models (#28382 review).""" + """/v1/model/info list path (no litellm_model_id) must include team-scoped + deployments from the router model list and surface the public name (#28382).""" + team_row = _team_row() + global_row = { + "model_name": "gpt-4o", + "litellm_params": {"model": "gpt-4o"}, + "model_info": {"id": "normal-id-1", "db_model": False}, + } router = MagicMock() - router.get_model_names.return_value = ["team-claude-sonnet"] + router.model_list = [team_row, global_row] + router.get_model_names.return_value = ["gpt-4o"] router.get_model_access_groups.return_value = {} - router.get_model_list.return_value = [_team_row()] monkeypatch.setattr(ps, "user_model", None) - monkeypatch.setattr(ps, "llm_model_list", [_team_row()]) + monkeypatch.setattr(ps, "llm_model_list", router.model_list) monkeypatch.setattr(ps, "llm_router", router) - monkeypatch.setattr(ps, "get_key_models", lambda **kw: []) - monkeypatch.setattr(ps, "get_team_models", lambda **kw: []) monkeypatch.setattr( - ps, "get_complete_model_list", lambda **kw: ["team-claude-sonnet"] + ps, "_enrich_model_info_with_litellm_data", lambda model, **kw: model ) admin = UserAPIKeyAuth( @@ -176,3 +179,167 @@ async def test_model_info_v1_list_path_translates_team_model_name(monkeypatch): names = [m["model_name"] for m in resp["data"]] assert "team-claude-sonnet" in names assert "model_name_team-abc-123_4a6b8" not in names + + +@pytest.mark.asyncio +async def test_model_info_v1_unrestricted_key_returns_all_deployments(monkeypatch): + """Unrestricted keys must see all router deployments (legacy v1 access logic).""" + deployment = { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "global-id-1", "db_model": False}, + } + router = MagicMock() + router.model_list = [deployment] + router.get_model_names.return_value = ["gpt-4"] + router.get_model_access_groups.return_value = {} + + monkeypatch.setattr(ps, "user_model", None) + monkeypatch.setattr(ps, "llm_model_list", router.model_list) + monkeypatch.setattr(ps, "llm_router", router) + monkeypatch.setattr( + ps, "_enrich_model_info_with_litellm_data", lambda model, **kw: model + ) + + caller = UserAPIKeyAuth( + user_id="user-1", + user_role=LitellmUserRoles.INTERNAL_USER, + models=[], + team_models=[], + ) + resp = await ps.model_info_v1(user_api_key_dict=caller, litellm_model_id=None) + + assert [m["model_name"] for m in resp["data"]] == ["gpt-4"] + + +@pytest.mark.asyncio +async def test_model_info_v1_restricted_key_filters_deployments(monkeypatch): + """Key-level model allowlists must filter router deployments.""" + team_row = _team_row() + global_row = { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "global-id-1", "db_model": False}, + } + router = MagicMock() + router.model_list = [team_row, global_row] + router.get_model_names.return_value = ["gpt-4", "team-claude-sonnet"] + router.get_model_access_groups.return_value = {} + + monkeypatch.setattr(ps, "user_model", None) + monkeypatch.setattr(ps, "llm_model_list", router.model_list) + monkeypatch.setattr(ps, "llm_router", router) + monkeypatch.setattr( + ps, "_enrich_model_info_with_litellm_data", lambda model, **kw: model + ) + + caller = UserAPIKeyAuth( + user_id="user-1", + user_role=LitellmUserRoles.INTERNAL_USER, + models=["gpt-4"], + team_models=[], + ) + resp = await ps.model_info_v1(user_api_key_dict=caller, litellm_model_id=None) + + assert [m["model_name"] for m in resp["data"]] == ["gpt-4"] + + +def _other_team_row() -> dict: + return { + "model_name": "model_name_team-other_9f2c1", + "litellm_params": { + "model": "azure/gpt-5.2-low-rpm-testing", + "api_base": "https://team-other-private.example.com", + }, + "model_info": { + "id": "byok-id-other", + "team_id": "team-other", + "team_public_model_name": "team-claude-sonnet", + "db_model": True, + }, + } + + +@pytest.mark.asyncio +async def test_model_info_v1_unrestricted_key_hides_other_team_byok(monkeypatch): + """Unrestricted non-admin keys must not enumerate other teams' BYOK + deployments, but must still see global models and their own team's.""" + team_row = _team_row() + other_team_row = _other_team_row() + global_row = { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "global-id-1", "db_model": False}, + } + router = MagicMock() + router.model_list = [team_row, other_team_row, global_row] + router.get_model_names.return_value = ["gpt-4"] + router.get_model_access_groups.return_value = {} + + prisma_client = MagicMock() + caller_user_row = MagicMock() + caller_user_row.teams = ["team-abc-123"] + prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=caller_user_row + ) + + monkeypatch.setattr(ps, "user_model", None) + monkeypatch.setattr(ps, "llm_model_list", router.model_list) + monkeypatch.setattr(ps, "llm_router", router) + monkeypatch.setattr(ps, "prisma_client", prisma_client) + monkeypatch.setattr( + ps, "_enrich_model_info_with_litellm_data", lambda model, **kw: model + ) + + caller = UserAPIKeyAuth( + user_id="user-1", + user_role=LitellmUserRoles.INTERNAL_USER, + models=[], + team_models=[], + ) + resp = await ps.model_info_v1(user_api_key_dict=caller, litellm_model_id=None) + + returned_ids = {m["model_info"]["id"] for m in resp["data"]} + assert returned_ids == {"global-id-1", "byok-id-1"} + assert "byok-id-other" not in returned_ids + names = [m["model_name"] for m in resp["data"]] + assert "team-claude-sonnet" in names + assert "gpt-4" in names + + +@pytest.mark.asyncio +async def test_model_info_v1_service_key_hides_all_team_byok(monkeypatch): + """A key without a resolvable user (e.g. CI/service token) sees only + global deployments, never any team-scoped BYOK rows.""" + team_row = _team_row() + other_team_row = _other_team_row() + global_row = { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "global-id-1", "db_model": False}, + } + router = MagicMock() + router.model_list = [team_row, other_team_row, global_row] + router.get_model_names.return_value = ["gpt-4"] + router.get_model_access_groups.return_value = {} + + prisma_client = MagicMock() + + monkeypatch.setattr(ps, "user_model", None) + monkeypatch.setattr(ps, "llm_model_list", router.model_list) + monkeypatch.setattr(ps, "llm_router", router) + monkeypatch.setattr(ps, "prisma_client", prisma_client) + monkeypatch.setattr( + ps, "_enrich_model_info_with_litellm_data", lambda model, **kw: model + ) + + caller = UserAPIKeyAuth( + user_id=None, + user_role=LitellmUserRoles.INTERNAL_USER, + team_id="team-abc-123", + models=[], + team_models=[], + ) + resp = await ps.model_info_v1(user_api_key_dict=caller, litellm_model_id=None) + + assert [m["model_info"]["id"] for m in resp["data"]] == ["global-id-1"] diff --git a/tests/test_litellm/proxy/test_model_info_default_limits.py b/tests/test_litellm/proxy/test_model_info_default_limits.py index 641199c96f0..8111a7af006 100644 --- a/tests/test_litellm/proxy/test_model_info_default_limits.py +++ b/tests/test_litellm/proxy/test_model_info_default_limits.py @@ -146,9 +146,9 @@ class TestModelInfoEndpointWithRouter: deployment_dict = deployment.model_dump(exclude_none=True) mock_router = MagicMock() + mock_router.model_list = [deployment_dict] mock_router.get_model_names.return_value = ["model1"] mock_router.get_model_access_groups.return_value = {} - mock_router.get_model_list.return_value = [deployment_dict] user_api_key_dict = UserAPIKeyAuth(api_key="sk-test") @@ -156,6 +156,7 @@ class TestModelInfoEndpointWithRouter: patch("litellm.proxy.proxy_server.llm_router", mock_router), patch("litellm.proxy.proxy_server.llm_model_list", [deployment_dict]), patch("litellm.proxy.proxy_server.user_model", None), + patch("litellm.proxy.proxy_server.prisma_client", None), patch("litellm.proxy.proxy_server.get_key_models", return_value=["model1"]), patch( "litellm.proxy.proxy_server.get_team_models", return_value=["model1"] diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 9f2c5ffd615..9eaccdfcbcd 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -3840,14 +3840,15 @@ async def test_model_info_v1_oci_secrets_not_leaked(): # Mock the llm_router to return our test data mock_router = MagicMock() + mock_router.model_list = [mock_model_data] mock_router.get_model_names.return_value = ["oci-grok-test"] mock_router.get_model_access_groups.return_value = {} - mock_router.get_model_list.return_value = [mock_model_data] # Mock global variables with ( patch("litellm.proxy.proxy_server.llm_router", mock_router), patch("litellm.proxy.proxy_server.llm_model_list", [mock_model_data]), + patch("litellm.proxy.proxy_server.prisma_client", None), patch( "litellm.proxy.proxy_server.general_settings", {"infer_model_from_keys": False}, From 4a3860df1f148486d76093cf95b631e39f888510 Mon Sep 17 00:00:00 2001 From: ishaan-berri <155045088+ishaan-berri@users.noreply.github.com> Date: Wed, 10 Jun 2026 21:20:11 -0700 Subject: [PATCH 061/185] fix: completion_cost AttributeError on streaming Anthropic web_search responses (#26153) (#27346) * fix: coerce server_tool_use dict to ServerToolUse in Usage.__init__ (#26153) * fix: coerce server_tool_use to ServerToolUse in stream_chunk_builder (#26153) * fix: dict/pydantic-tolerant access in tool_call_cost_tracking (#26153) * fix: dict/pydantic-tolerant access in anthropic cost_calculation (#26153) * test: assert ServerToolUse type in existing stream_chunk_builder anthropic web search test * test: regression test for #26153 (stream_chunk_builder server_tool_use type) * test: dict/pydantic safety for tool_call_cost_tracking helper * test: dict/pydantic safety for anthropic web_search cost * refactor: consolidate _get_web_search_requests into shared cost-calc utils * test(realtime): use gpt-realtime; openai retired gpt-4o-realtime-preview OpenAI shut down the gpt-4o-realtime-preview family (incl. the undated alias) on 2026-05-07, causing the live realtime test to fail with a 4000 invalid_request_error.invalid_model close. gpt-realtime is the GA successor; switch the live-call tests to it, matching the base branch. * refactor(types): drop redundant server_tool_use coercion in Usage.__init__ --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../llm_cost_calc/tool_call_cost_tracking.py | 7 +- .../litellm_core_utils/llm_cost_calc/utils.py | 22 ++- .../streaming_chunk_builder_utils.py | 13 +- litellm/llms/anthropic/cost_calculation.py | 14 +- ...est_tool_call_cost_tracking_dict_safety.py | 88 ++++++++++++ ...streaming_chunk_builder_server_tool_use.py | 130 ++++++++++++++++++ .../test_streaming_chunk_builder_utils.py | 5 +- .../test_cost_calculation_dict_safety.py | 94 +++++++++++++ 8 files changed, 360 insertions(+), 13 deletions(-) create mode 100644 tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking_dict_safety.py create mode 100644 tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_server_tool_use.py create mode 100644 tests/test_litellm/llms/anthropic/test_cost_calculation_dict_safety.py diff --git a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py index 8da66d4600d..413ddb71bf8 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py +++ b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py @@ -6,6 +6,7 @@ from typing import Any, Dict, List, Literal, Optional, Tuple import litellm from litellm.constants import OPENAI_FILE_SEARCH_COST_PER_1K_CALLS +from litellm.litellm_core_utils.llm_cost_calc.utils import _get_web_search_requests from litellm.types.llms.openai import ( FileSearchTool, ResponsesAPIResponse, @@ -339,8 +340,7 @@ class StandardBuiltInToolCostTracking: # and _handle_web_search_cost() is never called. if ( hasattr(usage, "server_tool_use") - and usage.server_tool_use is not None - and usage.server_tool_use.web_search_requests is not None + and _get_web_search_requests(usage.server_tool_use) is not None ): return True return False @@ -352,8 +352,7 @@ class StandardBuiltInToolCostTracking: elif usage is not None: if ( hasattr(usage, "server_tool_use") - and usage.server_tool_use is not None - and usage.server_tool_use.web_search_requests is not None + and _get_web_search_requests(usage.server_tool_use) is not None ): return True elif ( diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index f39c942f90f..93049adf75a 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -1,7 +1,7 @@ # What is this? ## Helper utilities for cost_per_token() -from typing import Literal, Optional, Tuple, TypedDict, cast +from typing import Any, Literal, Optional, Tuple, TypedDict, cast import litellm from litellm._logging import verbose_logger @@ -42,6 +42,26 @@ def _get_token_detail_value(details: object, key: str) -> Optional[int]: return value if isinstance(value, int) else None +def _get_web_search_requests(server_tool_use: Any) -> Optional[int]: + """ + Tolerantly read ``web_search_requests`` from a ``server_tool_use`` value + that may be ``None``, a ``dict``, a ``ServerToolUse`` pydantic instance, + or any other object supporting attribute access. + + Returns ``None`` when the value cannot be resolved — callers can + distinguish "absent" from "zero" using ``is None``. + + See https://github.com/BerriAI/litellm/issues/26153 — ``stream_chunk_builder`` + historically left this as a plain ``dict``, which broke direct attribute + access in cost calculation. + """ + if server_tool_use is None: + return None + if isinstance(server_tool_use, dict): + return server_tool_use.get("web_search_requests") + return getattr(server_tool_use, "web_search_requests", None) + + def _is_above_128k(tokens: float) -> bool: if tokens > 128000: return True diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 6257cce9aec..b495b183ec0 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -637,7 +637,18 @@ class ChunkProcessor: hasattr(usage_chunk, "server_tool_use") and usage_chunk.server_tool_use is not None ): - server_tool_use = usage_chunk.server_tool_use + # Coerce dict to ServerToolUse so downstream cost-calc code + # (which accesses .web_search_requests as an attribute) + # doesn't raise AttributeError. Some providers / streaming + # paths leave server_tool_use as a plain dict on the chunk. + if isinstance(usage_chunk.server_tool_use, dict): + server_tool_use = ServerToolUse(**usage_chunk.server_tool_use) + elif isinstance(usage_chunk.server_tool_use, ServerToolUse): + server_tool_use = usage_chunk.server_tool_use + else: + server_tool_use = ServerToolUse.model_validate( + usage_chunk.server_tool_use + ) if ( usage_chunk_dict["prompt_tokens_details"] is not None and getattr( diff --git a/litellm/llms/anthropic/cost_calculation.py b/litellm/llms/anthropic/cost_calculation.py index 3882d8f978c..6a031498dae 100644 --- a/litellm/llms/anthropic/cost_calculation.py +++ b/litellm/llms/anthropic/cost_calculation.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Optional, Tuple from litellm.litellm_core_utils.llm_cost_calc.utils import ( _get_token_base_cost, + _get_web_search_requests, _parse_prompt_tokens_details, calculate_cache_writing_cost, generic_cost_per_token, @@ -110,11 +111,12 @@ def get_cost_for_anthropic_web_search( if model_info is None: return 0.0 - if ( - usage is None - or usage.server_tool_use is None - or usage.server_tool_use.web_search_requests is None - ): + if usage is None: + return 0.0 + web_search_requests = _get_web_search_requests( + getattr(usage, "server_tool_use", None) + ) + if web_search_requests is None: return 0.0 ## Get the cost per web search request @@ -128,5 +130,5 @@ def get_cost_for_anthropic_web_search( return 0.0 ## Calculate the total cost - total_cost = cost_per_web_search_request * usage.server_tool_use.web_search_requests + total_cost = cost_per_web_search_request * web_search_requests return total_cost diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking_dict_safety.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking_dict_safety.py new file mode 100644 index 00000000000..4eee6b59d34 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking_dict_safety.py @@ -0,0 +1,88 @@ +""" +Tests that the cost-tracking call sites tolerate ``server_tool_use`` being +either a ``dict`` or a ``ServerToolUse`` pydantic instance. + +See https://github.com/BerriAI/litellm/issues/26153. +""" + +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( + StandardBuiltInToolCostTracking, + _get_web_search_requests, +) +from litellm.types.utils import ModelResponse, ServerToolUse, Usage + + +class _UsageWithDictServerToolUse: + """ + Tiny stand-in that mimics the broken streaming-rebuild shape: + ``server_tool_use`` is a plain dict. + """ + + def __init__(self, server_tool_use): + self.server_tool_use = server_tool_use + self.prompt_tokens_details = None + + +def test_get_web_search_requests_handles_none(): + assert _get_web_search_requests(None) is None + + +def test_get_web_search_requests_handles_dict(): + assert _get_web_search_requests({"web_search_requests": 5}) == 5 + + +def test_get_web_search_requests_handles_dict_missing_key(): + assert _get_web_search_requests({}) is None + + +def test_get_web_search_requests_handles_pydantic(): + stu = ServerToolUse(web_search_requests=7) + assert _get_web_search_requests(stu) == 7 + + +def test_get_web_search_requests_handles_pydantic_with_none_value(): + stu = ServerToolUse() + assert _get_web_search_requests(stu) is None + + +def test_response_object_includes_web_search_call_with_dict_server_tool_use(): + """ + The exact bug: ``usage.server_tool_use`` is a dict and the check in + ``response_object_includes_web_search_call`` used to crash with + ``AttributeError``. + """ + response = ModelResponse() + usage = _UsageWithDictServerToolUse({"web_search_requests": 2}) + + # Must not raise — and must correctly detect the web search call. + result = StandardBuiltInToolCostTracking.response_object_includes_web_search_call( + response_object=response, usage=usage # type: ignore[arg-type] + ) + assert result is True + + +def test_response_object_includes_web_search_call_with_pydantic_server_tool_use(): + response = ModelResponse() + usage = _UsageWithDictServerToolUse(ServerToolUse(web_search_requests=2)) + + result = StandardBuiltInToolCostTracking.response_object_includes_web_search_call( + response_object=response, usage=usage # type: ignore[arg-type] + ) + assert result is True + + +def test_response_object_includes_web_search_call_with_none_server_tool_use(): + response = ModelResponse() + usage = _UsageWithDictServerToolUse(None) + + result = StandardBuiltInToolCostTracking.response_object_includes_web_search_call( + response_object=response, usage=usage # type: ignore[arg-type] + ) + assert result is False diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_server_tool_use.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_server_tool_use.py new file mode 100644 index 00000000000..4e28d5ba7d2 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_server_tool_use.py @@ -0,0 +1,130 @@ +""" +Regression tests for https://github.com/BerriAI/litellm/issues/26153 + +``stream_chunk_builder`` used to leave ``usage.server_tool_use`` as a plain +``dict`` when reconstructing a streaming response. Downstream cost-calculation +code (``StandardBuiltInToolCostTracking.response_object_includes_web_search_call`` +and ``get_cost_for_anthropic_web_search``) accesses +``usage.server_tool_use.web_search_requests`` as an attribute, which raised +``AttributeError: 'dict' object has no attribute 'web_search_requests'``. + +These tests reconstruct streaming chunks for an Anthropic-style web_search +response and assert: + +1. ``stream_chunk_builder`` returns ``ServerToolUse`` (not ``dict``) for + ``usage.server_tool_use``. +2. ``completion_cost`` runs end-to-end on the rebuilt response without + raising ``AttributeError``. +""" + +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm import completion_cost, stream_chunk_builder +from litellm.types.utils import ( + Delta, + ModelResponseStream, + ServerToolUse, + StreamingChoices, + Usage, +) + + +def _make_text_chunk(text: str) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-test-26153", + created=1700000000, + model="claude-3-haiku-20240307", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta(role="assistant", content=text), + ) + ], + ) + + +def _make_finish_chunk_with_usage_dict_server_tool_use() -> ModelResponseStream: + """Final chunk where server_tool_use is a *dict* — reproduces the bug shape.""" + return ModelResponseStream( + id="chatcmpl-test-26153", + created=1700000000, + model="claude-3-haiku-20240307", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(), + ) + ], + usage=Usage( + prompt_tokens=42, + completion_tokens=11, + total_tokens=53, + # NOTE: passed as a dict on purpose — this is the shape that + # historically slipped through stream_chunk_builder unchanged. + server_tool_use={"web_search_requests": 3}, + ), + ) + + +def test_stream_chunk_builder_coerces_server_tool_use_to_pydantic(): + """ + Regression: stream_chunk_builder must produce ServerToolUse, not dict. + """ + chunks = [ + _make_text_chunk("Otters "), + _make_text_chunk("are great."), + _make_finish_chunk_with_usage_dict_server_tool_use(), + ] + + rebuilt = stream_chunk_builder(chunks) + + assert rebuilt is not None + assert rebuilt.usage is not None # type: ignore[attr-defined] + server_tool_use = rebuilt.usage.server_tool_use # type: ignore[attr-defined] + + assert ( + server_tool_use is not None + ), "server_tool_use should be carried through from the final chunk" + assert isinstance(server_tool_use, ServerToolUse), ( + f"expected ServerToolUse, got {type(server_tool_use).__name__}: " + f"{server_tool_use!r}" + ) + # Attribute access must not raise (this is exactly what was broken). + assert server_tool_use.web_search_requests == 3 + + +def test_completion_cost_does_not_raise_on_streaming_web_search_response(): + """ + Regression: completion_cost(...) must not raise AttributeError when the + response was reconstructed by stream_chunk_builder from a streaming + Anthropic web_search call. + """ + chunks = [ + _make_text_chunk("hello"), + _make_finish_chunk_with_usage_dict_server_tool_use(), + ] + + rebuilt = stream_chunk_builder(chunks) + assert rebuilt is not None + + # The exact dollar amount depends on the model-pricing table; what matters + # for this regression is that it does NOT raise AttributeError on + # `dict has no attribute 'web_search_requests'`. + try: + cost = completion_cost(completion_response=rebuilt) + except AttributeError as e: # pragma: no cover - regression guard + pytest.fail( + "completion_cost raised AttributeError after stream_chunk_builder " + f"(issue #26153 regression): {e}" + ) + + assert isinstance(cost, (int, float)) diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py index 77765340c61..c5794194528 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -520,7 +520,10 @@ def test_stream_chunk_builder_anthropic_web_search(): assert usage.prompt_tokens == 50 assert usage.completion_tokens == 27 assert usage.total_tokens == 77 - assert usage.server_tool_use["web_search_requests"] == 2 + # server_tool_use must be a ServerToolUse pydantic so downstream cost-calc + # (which uses attribute access) works. See issue #26153. + assert isinstance(usage.server_tool_use, ServerToolUse) + assert usage.server_tool_use.web_search_requests == 2 def test_sort_chunks_handles_dict_hidden_params_created_at(): diff --git a/tests/test_litellm/llms/anthropic/test_cost_calculation_dict_safety.py b/tests/test_litellm/llms/anthropic/test_cost_calculation_dict_safety.py new file mode 100644 index 00000000000..70fef0162e6 --- /dev/null +++ b/tests/test_litellm/llms/anthropic/test_cost_calculation_dict_safety.py @@ -0,0 +1,94 @@ +""" +Tests that ``get_cost_for_anthropic_web_search`` tolerates ``server_tool_use`` +being either a ``dict`` or a ``ServerToolUse`` pydantic instance. + +See https://github.com/BerriAI/litellm/issues/26153. +""" + +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.llms.anthropic.cost_calculation import ( + _get_web_search_requests, + get_cost_for_anthropic_web_search, +) +from litellm.types.utils import ModelInfo, ServerToolUse + + +class _UsageWithServerToolUse: + def __init__(self, server_tool_use): + self.server_tool_use = server_tool_use + + +def _make_model_info(cost_per_query: float = 0.01) -> ModelInfo: + info: ModelInfo = { # type: ignore[typeddict-item] + "search_context_cost_per_query": { + "search_context_size_low": cost_per_query, + "search_context_size_medium": cost_per_query, + "search_context_size_high": cost_per_query, + } + } + return info + + +def test_get_web_search_requests_handles_none(): + assert _get_web_search_requests(None) is None + + +def test_get_web_search_requests_handles_dict(): + assert _get_web_search_requests({"web_search_requests": 4}) == 4 + + +def test_get_web_search_requests_handles_dict_missing_key(): + assert _get_web_search_requests({}) is None + + +def test_get_web_search_requests_handles_pydantic(): + assert _get_web_search_requests(ServerToolUse(web_search_requests=2)) == 2 + + +def test_get_cost_for_anthropic_web_search_with_dict_server_tool_use(): + """ + Regression: ``server_tool_use`` was a dict from ``stream_chunk_builder`` and + direct attribute access on it raised ``AttributeError``. + """ + usage = _UsageWithServerToolUse({"web_search_requests": 3}) + info = _make_model_info(cost_per_query=0.01) + + cost = get_cost_for_anthropic_web_search( + model_info=info, usage=usage # type: ignore[arg-type] + ) + + assert cost == pytest.approx(0.03) + + +def test_get_cost_for_anthropic_web_search_with_pydantic_server_tool_use(): + usage = _UsageWithServerToolUse(ServerToolUse(web_search_requests=3)) + info = _make_model_info(cost_per_query=0.01) + + cost = get_cost_for_anthropic_web_search( + model_info=info, usage=usage # type: ignore[arg-type] + ) + + assert cost == pytest.approx(0.03) + + +def test_get_cost_for_anthropic_web_search_with_none_server_tool_use(): + usage = _UsageWithServerToolUse(None) + info = _make_model_info(cost_per_query=0.01) + + cost = get_cost_for_anthropic_web_search( + model_info=info, usage=usage # type: ignore[arg-type] + ) + + assert cost == 0.0 + + +def test_get_cost_for_anthropic_web_search_with_no_usage(): + info = _make_model_info(cost_per_query=0.01) + cost = get_cost_for_anthropic_web_search(model_info=info, usage=None) + assert cost == 0.0 From 7a96b3490d8ac241865fe6930f658aee1575976f Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 10 Jun 2026 21:26:35 -0700 Subject: [PATCH 062/185] [internal copy of #30137] perf(realtime): eliminate redundant per-frame JSON work on OpenAI realtime relay (#30142) * perf(realtime): eliminate redundant per-frame JSON work on OpenAI realtime relay The GA realtime support added in #27110 made backend_to_client_send_messages parse every backend frame up to three times for beta clients (OpenAI-Beta: realtime=v1), build a discarded Pydantic object per frame for logging, and re-serialize even frames that need no translation. For high-frequency response.output_audio.delta frames carrying multi-KB base64 payloads, that serialized CPU work on the hottest relay path drove the latency regression between v1.83.14 and v1.88.1 for gpt-realtime-1.5 and gpt-realtime-2. This parses each frame once via _parse_backend_event and threads the dict into _handle_raw_backend_message, store_message, and _translate_event_to_beta; short-circuits store_message before the Pydantic build for events not in the logged set; returns the original event unchanged from _translate_event_to_beta when no rename applies so the raw frame is forwarded without re-serialization; and only json.dumps when the type is actually renamed. * fix(realtime): widen store_message type hint to accept plain dict The parse-once refactor passes the dict produced by _parse_backend_event into store_message, but the parameter was typed as str | bytes | OpenAIRealtimeEvents (a union of TypedDicts), which mypy does not consider compatible with a plain dict. Add dict to the accepted union; the body already handles it. --------- Co-authored-by: Miguel Armenta --- .../litellm_core_utils/realtime_streaming.py | 180 +++++++------- .../test_realtime_streaming.py | 220 +++++++++++++++++- 2 files changed, 297 insertions(+), 103 deletions(-) diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index 772f058d9bb..4b7f0e22198 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -144,7 +144,7 @@ class RealTimeStreaming: return True return False - def store_message(self, message: Union[str, bytes, OpenAIRealtimeEvents]): + def store_message(self, message: Union[str, bytes, dict, OpenAIRealtimeEvents]): """Store message in list""" if isinstance(message, bytes): message = message.decode("utf-8") @@ -154,22 +154,20 @@ class RealTimeStreaming: else: message_obj = cast(Dict[str, Any], json.loads(cast(str, message))) self._collect_tool_calls_from_response_done(cast(dict, message_obj)) + if not self._should_store_message(message_obj): + return try: event_type = message_obj.get("type", "") if event_type in self._SESSION_EVENT_TYPES: - typed_obj = OpenAIRealtimeStreamSessionEvents(**message_obj) # type: ignore + typed_obj: OpenAIRealtimeEvents = OpenAIRealtimeStreamSessionEvents(**message_obj) # type: ignore else: - # Use the base object as a safe catch-all for all other event types - # (both beta and GA), so unknown/new event names never raise here. + # Catch-all base object so unknown/new event names never raise. typed_obj = OpenAIRealtimeStreamResponseBaseObject(**message_obj) # type: ignore except Exception as e: verbose_logger.debug(f"Error parsing message for logging: {e}") - # Don't re-raise — a parse failure must not drop or delay the message - if self._should_store_message(message_obj): - self.messages.append(message_obj) # type: ignore[arg-type] + self.messages.append(message_obj) # type: ignore[arg-type] return - if self._should_store_message(typed_obj): - self.messages.append(typed_obj) + self.messages.append(typed_obj) def _collect_user_input_from_client_event(self, message: Union[str, dict]) -> None: """Extract user text content from client WebSocket events for spend logging.""" @@ -358,8 +356,7 @@ class RealTimeStreaming: for msg in self._pending_messages_until_setup ) verbose_logger.debug( - "Failed to flush buffered client message after setup: %s " - "(%d buffered message(s) retained)", + "Failed to flush buffered client message after setup: %s (%d buffered message(s) retained)", e, len(unsent), ) @@ -376,8 +373,7 @@ class RealTimeStreaming: return True except Exception as e: verbose_logger.warning( - "Failed to translate %s to beta protocol, forwarding " - "untranslated event to client: %s", + "Failed to translate %s to beta protocol, forwarding untranslated event to client: %s", event.get("type"), e, ) @@ -705,48 +701,48 @@ class RealTimeStreaming: self.store_message(event_str) await self._send_event_to_client(event, event_str) - async def _handle_raw_backend_message(self, raw_response) -> bool: + @staticmethod + def _parse_backend_event(raw_response: str) -> Optional[dict]: + """Parse a backend frame once. Returns None for non-JSON or non-object frames.""" + try: + event = json.loads(raw_response) + except (json.JSONDecodeError, TypeError): + return None + return event if isinstance(event, dict) else None + + async def _handle_raw_backend_message( + self, event_obj: dict, raw_response: str + ) -> bool: """Process a backend message without provider_config (raw path). Returns True if the caller should skip the default store+forward (i.e. continue the loop). """ - try: - event_obj = json.loads(raw_response) + event_type = event_obj.get("type") - # For audio/VAD guardrail path: once the session is ready, tell the backend - # not to auto-respond after VAD detects end-of-speech. We send the - # session.created to the client FIRST so the client is always in sync, then - # inject the session.update so a potential error from the backend doesn't - # arrive before the client sees session.created. - if ( - event_obj.get("type") == "session.created" - and self._has_audio_transcription_guardrails() - ): - self.store_message(raw_response) - await self.websocket.send_text(raw_response) - await self._send_to_backend(self._make_disable_auto_response_message()) - return True + # Send session.created to the client FIRST so it stays in sync, then inject + # the disable-auto-response session.update; otherwise a backend error could + # reach the client before it sees session.created. + if ( + event_type == "session.created" + and self._has_audio_transcription_guardrails() + ): + self.store_message(event_obj) + await self.websocket.send_text(raw_response) + await self._send_to_backend(self._make_disable_auto_response_message()) + return True - if ( - event_obj.get("type") - == "conversation.item.input_audio_transcription.completed" - ): - transcript = event_obj.get("transcript", "") - self._collect_user_input_from_backend_event(event_obj) - ## LOGGING — must happen before continue below - self.store_message(raw_response) - # Forward transcript to client so user sees what they said - await self.websocket.send_text(raw_response) - blocked = await self.run_realtime_guardrails( - transcript, - item_id=event_obj.get("item_id"), - ) - if not blocked: - # Clean — trigger LLM response - await self._send_to_backend(json.dumps({"type": "response.create"})) - return True - except (json.JSONDecodeError, AttributeError): - pass + if event_type == "conversation.item.input_audio_transcription.completed": + transcript = event_obj.get("transcript", "") + self._collect_user_input_from_backend_event(event_obj) + self.store_message(event_obj) + await self.websocket.send_text(raw_response) + blocked = await self.run_realtime_guardrails( + transcript, + item_id=event_obj.get("item_id"), + ) + if not blocked: + await self._send_to_backend(json.dumps({"type": "response.create"})) + return True return False async def backend_to_client_send_messages(self): @@ -779,25 +775,25 @@ class RealTimeStreaming: ) continue else: - handled = await self._handle_raw_backend_message(raw_response) - if handled: - continue - ## LOGGING - self.store_message(raw_response) - - # If the client opted into beta protocol, translate GA event - # names/shapes back to the beta equivalents before forwarding. - if self._client_wants_beta: - try: - event_dict = json.loads(raw_response) - translated = self._translate_event_to_beta(event_dict) - if translated is None: - continue # drop GA-only events (e.g. conversation.item.done) - await self.websocket.send_text(json.dumps(translated)) - except Exception: - await self.websocket.send_text(raw_response) - else: + event = self._parse_backend_event(raw_response) + if event is None: await self.websocket.send_text(raw_response) + continue + + if await self._handle_raw_backend_message(event, raw_response): + continue + self.store_message(event) + + if not self._client_wants_beta: + await self.websocket.send_text(raw_response) + continue + + translated = self._translate_event_to_beta(event) + if translated is None: + continue + await self.websocket.send_text( + raw_response if translated is event else json.dumps(translated) + ) except websockets.exceptions.ConnectionClosed as e: # type: ignore verbose_logger.exception( @@ -927,41 +923,43 @@ class RealTimeStreaming: def _translate_event_to_beta(event: dict) -> Optional[dict]: """Translate a single GA event dict to its beta equivalent. - Returns None if the event should be dropped entirely (e.g. the GA-only - conversation.item.done has no beta counterpart). - Returns the (possibly mutated copy of the) event otherwise. + Returns None when the event must be dropped (the GA-only + conversation.item.done has no beta counterpart). Returns the original + event object unchanged when no translation applies, so the caller can + forward the raw frame without re-serializing; otherwise returns a + translated copy. """ event_type = event.get("type", "") - # conversation.item.done has no beta equivalent — the client already - # received conversation.item.created (translated from .added). if event_type == "conversation.item.done": return None - # Shallow-copy so we don't mutate the stored message + renamed_type = RealTimeStreaming._GA_TO_BETA_EVENT_TYPES.get(event_type) + has_item = isinstance(event.get("item"), dict) + response = event.get("response") + has_response_output = isinstance(response, dict) and isinstance( + response.get("output"), list + ) + if renamed_type is None and not has_item and not has_response_output: + return event + translated = dict(event) - - # Rename the type field - if event_type in RealTimeStreaming._GA_TO_BETA_EVENT_TYPES: - translated["type"] = RealTimeStreaming._GA_TO_BETA_EVENT_TYPES[event_type] - - # Fix content block types inside items (response.done output list, - # conversation.item.created item content, etc.) - if "item" in translated and isinstance(translated["item"], dict): + if renamed_type is not None: + translated["type"] = renamed_type + if has_item: translated["item"] = RealTimeStreaming._translate_item_content_types( dict(translated["item"]) ) - if "response" in translated and isinstance(translated["response"], dict): + if has_response_output: resp = dict(translated["response"]) - if "output" in resp and isinstance(resp["output"], list): - resp["output"] = [ - ( - RealTimeStreaming._translate_item_content_types(dict(o)) - if isinstance(o, dict) - else o - ) - for o in resp["output"] - ] + resp["output"] = [ + ( + RealTimeStreaming._translate_item_content_types(dict(o)) + if isinstance(o, dict) + else o + ) + for o in resp["output"] + ] translated["response"] = resp return translated diff --git a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py index 3424bfd801c..0f8d5cfd85a 100644 --- a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py +++ b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py @@ -773,17 +773,15 @@ async def test_realtime_guardrail_blocks_prompt_injection(): guardrail_items = [ e for e in sent_to_backend if e.get("type") == "conversation.item.create" ] - assert len(guardrail_items) == 1, ( - f"Guardrail should inject a conversation.item.create with violation message, " - f"got: {guardrail_items}" - ) + assert ( + len(guardrail_items) == 1 + ), f"Guardrail should inject a conversation.item.create with violation message, got: {guardrail_items}" response_creates = [ e for e in sent_to_backend if e.get("type") == "response.create" ] - assert len(response_creates) == 1, ( - f"Guardrail should send exactly one response.create to voice the violation, " - f"got: {response_creates}" - ) + assert ( + len(response_creates) == 1 + ), f"Guardrail should send exactly one response.create to voice the violation, got: {response_creates}" # ASSERT 2: error event was sent directly to the client WebSocket sent_to_client = [ @@ -1050,10 +1048,9 @@ async def test_realtime_function_call_output_guardrail_blocks_and_returns_error( # every toolCall with a toolResponse (Gemini/Vertex Live) exit their # pending-tool-call state instead of stalling. The placeholder must NOT # contain any of the blocked content. - assert len(forwarded_tool_outputs) == 1, ( - f"Sanitized function_call_output should be forwarded, got: " - f"{forwarded_tool_outputs}" - ) + assert ( + len(forwarded_tool_outputs) == 1 + ), f"Sanitized function_call_output should be forwarded, got: {forwarded_tool_outputs}" sanitized_item = forwarded_tool_outputs[0]["item"] assert sanitized_item["call_id"] == "call_123" assert "test@example.com" not in sanitized_item["output"] @@ -2110,3 +2107,202 @@ async def test_deferred_setup_caps_non_audio_buffered_bytes(monkeypatch): assert ( streaming._pending_messages_byte_total <= RealTimeStreaming._MAX_BUFFERED_BYTES ) + + +def _beta_client_ws(): + ws = MagicMock() + ws.scope = {"headers": [(b"openai-beta", b"realtime=v1")]} + ws.send_text = AsyncMock() + return ws + + +def _ga_client_ws(): + ws = MagicMock() + ws.scope = {"headers": []} + ws.send_text = AsyncMock() + return ws + + +def _streaming_with(client_ws): + backend_ws = MagicMock() + logging_obj = MagicMock() + logging_obj.async_success_handler = AsyncMock() + logging_obj.success_handler = MagicMock() + return RealTimeStreaming(client_ws, backend_ws, logging_obj) + + +def test_parse_backend_event_returns_none_for_non_json(): + assert RealTimeStreaming._parse_backend_event("not json") is None + + +def test_parse_backend_event_returns_none_for_non_dict_json(): + assert RealTimeStreaming._parse_backend_event("[1, 2, 3]") is None + assert RealTimeStreaming._parse_backend_event('"a string"') is None + + +def test_parse_backend_event_returns_dict(): + parsed = RealTimeStreaming._parse_backend_event('{"type": "x", "v": 1}') + assert parsed == {"type": "x", "v": 1} + + +def test_translate_event_to_beta_returns_identity_when_no_translation(): + """An event with no renamed type and no item/response is returned unchanged + (same object), so the caller can forward the raw frame without re-serializing.""" + ev = {"type": "error", "error": {"message": "boom"}} + out = RealTimeStreaming._translate_event_to_beta(ev) + assert out is ev + + +def test_translate_event_to_beta_preserves_audio_delta_payload(): + payload = "QUJDREVG" * 200 + out = RealTimeStreaming._translate_event_to_beta( + {"type": "response.output_audio.delta", "delta": payload, "event_id": "e1"} + ) + assert out is not None + assert out["type"] == "response.audio.delta" + assert out["delta"] == payload + + +def test_translate_event_to_beta_remaps_response_done_output_content_types(): + out = RealTimeStreaming._translate_event_to_beta( + { + "type": "response.done", + "response": { + "output": [ + { + "type": "message", + "content": [{"type": "output_audio", "transcript": "hi"}], + } + ] + }, + } + ) + assert out is not None + assert out["response"]["output"][0]["content"][0]["type"] == "audio" + + +@pytest.mark.asyncio +async def test_beta_client_receives_translated_audio_delta(): + client_ws = _beta_client_ws() + frame = json.dumps( + {"type": "response.output_audio.delta", "delta": "QUJD", "event_id": "e1"} + ) + backend_ws = MagicMock() + backend_ws.recv = AsyncMock( + side_effect=[frame.encode(), ConnectionClosed(None, None)] + ) + logging_obj = MagicMock() + logging_obj.async_success_handler = AsyncMock() + logging_obj.success_handler = MagicMock() + streaming = RealTimeStreaming(client_ws, backend_ws, logging_obj) + + await streaming.backend_to_client_send_messages() + + assert client_ws.send_text.await_count == 1 + sent = json.loads(client_ws.send_text.await_args.args[0]) + assert sent["type"] == "response.audio.delta" + assert sent["delta"] == "QUJD" + + +@pytest.mark.asyncio +async def test_ga_client_receives_raw_passthrough(): + client_ws = _ga_client_ws() + frame = json.dumps( + {"type": "response.output_audio.delta", "delta": "QUJD", "event_id": "e1"} + ) + backend_ws = MagicMock() + backend_ws.recv = AsyncMock( + side_effect=[frame.encode(), ConnectionClosed(None, None)] + ) + logging_obj = MagicMock() + logging_obj.async_success_handler = AsyncMock() + logging_obj.success_handler = MagicMock() + streaming = RealTimeStreaming(client_ws, backend_ws, logging_obj) + + await streaming.backend_to_client_send_messages() + + assert client_ws.send_text.await_count == 1 + # GA client gets the byte-identical frame, no re-serialization. + assert client_ws.send_text.await_args.args[0] == frame + + +@pytest.mark.asyncio +async def test_beta_client_non_translated_event_forwarded_raw(): + """For a beta client, an event needing no translation is forwarded as the + original raw frame (identity return path), not a re-serialized copy.""" + client_ws = _beta_client_ws() + frame = json.dumps({"type": "error", "error": {"message": "boom"}}) + backend_ws = MagicMock() + backend_ws.recv = AsyncMock( + side_effect=[frame.encode(), ConnectionClosed(None, None)] + ) + logging_obj = MagicMock() + logging_obj.async_success_handler = AsyncMock() + logging_obj.success_handler = MagicMock() + streaming = RealTimeStreaming(client_ws, backend_ws, logging_obj) + + await streaming.backend_to_client_send_messages() + + assert client_ws.send_text.await_count == 1 + assert client_ws.send_text.await_args.args[0] == frame + + +@pytest.mark.asyncio +async def test_beta_client_drops_conversation_item_done(): + client_ws = _beta_client_ws() + frame = json.dumps({"type": "conversation.item.done", "item": {"id": "i1"}}) + backend_ws = MagicMock() + backend_ws.recv = AsyncMock( + side_effect=[frame.encode(), ConnectionClosed(None, None)] + ) + logging_obj = MagicMock() + logging_obj.async_success_handler = AsyncMock() + logging_obj.success_handler = MagicMock() + streaming = RealTimeStreaming(client_ws, backend_ws, logging_obj) + + await streaming.backend_to_client_send_messages() + + assert client_ws.send_text.await_count == 0 + + +def test_store_message_skips_pydantic_for_unlogged_audio_delta(): + """Audio deltas are not in DefaultLoggedRealTimeEventTypes; store_message must + skip the Pydantic build entirely (no append, no validation).""" + streaming = _streaming_with(_ga_client_ws()) + with patch( + "litellm.litellm_core_utils.realtime_streaming.OpenAIRealtimeStreamResponseBaseObject" + ) as base_obj: + streaming.store_message({"type": "response.output_audio.delta", "delta": "x"}) + base_obj.assert_not_called() + assert streaming.messages == [] + + +@pytest.mark.asyncio +async def test_audio_delta_frame_parsed_at_most_once(): + client_ws = _beta_client_ws() + frame = json.dumps( + {"type": "response.output_audio.delta", "delta": "QUJD", "event_id": "e1"} + ) + backend_ws = MagicMock() + backend_ws.recv = AsyncMock( + side_effect=[frame.encode(), ConnectionClosed(None, None)] + ) + logging_obj = MagicMock() + logging_obj.async_success_handler = AsyncMock() + logging_obj.success_handler = MagicMock() + streaming = RealTimeStreaming(client_ws, backend_ws, logging_obj) + + real_loads = json.loads + calls = {"n": 0} + + def counting_loads(*args, **kwargs): + calls["n"] += 1 + return real_loads(*args, **kwargs) + + with patch( + "litellm.litellm_core_utils.realtime_streaming.json.loads", + side_effect=counting_loads, + ): + await streaming.backend_to_client_send_messages() + + assert calls["n"] == 1 From 49ca04d8c3ddea336237ce6f3082dbc26d19e944 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 10 Jun 2026 21:31:08 -0700 Subject: [PATCH 063/185] feat(bedrock): aws_bedrock_project_id for bedrock-mantle project / workspace association (#30163) * feat(bedrock): support aws_bedrock_project_id for bedrock-mantle project association Adds a litellm_params field to associate bedrock-mantle requests with an Amazon Bedrock project, sent as the OpenAI-Project header on the OpenAI-compatible chat and responses paths and as the anthropic-workspace header on the Anthropic messages paths. This lets a single model entry opt into a project-scoped data retention mode (e.g. provider_data_share for Claude Fable 5) while the account-wide setting stays on default. The param is carried via litellm_params only and is explicitly excluded from optional_params so it can never leak into a request body. Fixes #30070 * chore(ui): regenerate schema.d.ts for aws_bedrock_project_id Generated with npm run gen:api after adding the field to LiteLLM_Params * fix(proxy): ban client-supplied aws_bedrock_project_id in request bodies The deployment pins aws_bedrock_project_id so the project's data retention policy applies to its requests. Without this guard an authenticated caller could supply the field in the request body and, since client kwargs win the router merge, run requests under any project reachable with the deployment's shared AWS credentials. Adds the field to _BANNED_REQUEST_BODY_PARAMS so it is rejected at the auth boundary by default while remaining available through the existing admin opt-ins (allow_client_side_credentials proxy-wide or configurable_clientside_auth_params per deployment). --- .../litellm_core_utils/get_litellm_params.py | 1 + .../bedrock/chat/mantle/transformation.py | 24 +++ .../bedrock/messages/mantle_transformation.py | 26 ++- .../bedrock_mantle/chat/transformation.py | 27 +++- .../responses/transformation.py | 2 + litellm/main.py | 1 + litellm/proxy/auth/auth_utils.py | 5 + litellm/types/router.py | 2 + litellm/utils.py | 4 + .../test_litellm/llms/bedrock/test_mantle.py | 148 +++++++++++++++++- ...bedrock_mantle_responses_transformation.py | 20 +++ .../test_bedrock_mantle_transformation.py | 76 +++++++++ .../proxy/auth/test_auth_utils.py | 34 ++++ tests/test_litellm/test_utils.py | 18 +++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 + 15 files changed, 388 insertions(+), 4 deletions(-) diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index f80cb41dc3f..6e655b03fed 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -32,6 +32,7 @@ _OPTIONAL_KWARGS_KEYS = frozenset( "aws_sts_endpoint", "aws_external_id", "aws_bedrock_runtime_endpoint", + "aws_bedrock_project_id", "tpm", "rpm", "use_xai_oauth", diff --git a/litellm/llms/bedrock/chat/mantle/transformation.py b/litellm/llms/bedrock/chat/mantle/transformation.py index ef0199031af..cbed2232be5 100644 --- a/litellm/llms/bedrock/chat/mantle/transformation.py +++ b/litellm/llms/bedrock/chat/mantle/transformation.py @@ -48,6 +48,30 @@ class AmazonMantleConfig(AmazonAnthropicClaudeConfig): region = self._get_aws_region_name(optional_params=optional_params, model=model) return MANTLE_ENDPOINT_TEMPLATE.format(region=region) + def validate_environment( + self, + headers: dict, + model: str, + messages: List[AllMessageValues], + optional_params: dict, + litellm_params: dict, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + ) -> dict: + headers = super().validate_environment( + headers=headers, + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + api_key=api_key, + api_base=api_base, + ) + project_id = litellm_params.get("aws_bedrock_project_id") + if project_id: + headers["anthropic-workspace"] = project_id + return headers + def transform_request( self, model: str, diff --git a/litellm/llms/bedrock/messages/mantle_transformation.py b/litellm/llms/bedrock/messages/mantle_transformation.py index a78f696a057..900d9aa97d8 100644 --- a/litellm/llms/bedrock/messages/mantle_transformation.py +++ b/litellm/llms/bedrock/messages/mantle_transformation.py @@ -6,7 +6,7 @@ AmazonAnthropicClaudeMessagesConfig. Overrides only the URL and model-prefix stripping that are specific to the bedrock-mantle endpoint. """ -from typing import TYPE_CHECKING, Any, Dict, List, Optional +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( AmazonAnthropicClaudeMessagesConfig, @@ -45,6 +45,30 @@ class AmazonMantleMessagesConfig(AmazonAnthropicClaudeMessagesConfig): region = self._get_aws_region_name(optional_params=optional_params, model=model) return MANTLE_ENDPOINT_TEMPLATE.format(region=region) + def validate_anthropic_messages_environment( + self, + headers: dict, + model: str, + messages: List[Any], + optional_params: dict, + litellm_params: dict, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + ) -> Tuple[dict, Optional[str]]: + headers, api_base = super().validate_anthropic_messages_environment( + headers=headers, + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + api_key=api_key, + api_base=api_base, + ) + project_id = litellm_params.get("aws_bedrock_project_id") + if project_id: + headers["anthropic-workspace"] = project_id + return headers, api_base + def transform_anthropic_messages_request( self, model: str, diff --git a/litellm/llms/bedrock_mantle/chat/transformation.py b/litellm/llms/bedrock_mantle/chat/transformation.py index 81a56030a5c..ad37a1990d3 100644 --- a/litellm/llms/bedrock_mantle/chat/transformation.py +++ b/litellm/llms/bedrock_mantle/chat/transformation.py @@ -8,11 +8,12 @@ Auth: AWS Bedrock API key as Bearer token (set via BEDROCK_MANTLE_API_KEY env va or region-aware key via BEDROCK_MANTLE_{REGION}_API_KEY. """ -from typing import Iterator, AsyncIterator, Any, Optional, Tuple, Union +from typing import Iterator, AsyncIterator, Any, List, Optional, Tuple, Union import litellm from litellm._logging import verbose_logger from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.openai import AllMessageValues from ...openai_like.chat.transformation import OpenAILikeChatConfig @@ -48,6 +49,30 @@ class BedrockMantleChatConfig(OpenAILikeChatConfig): dynamic_api_key = api_key or get_secret_str("BEDROCK_MANTLE_API_KEY") return api_base, dynamic_api_key + def validate_environment( + self, + headers: dict, + model: str, + messages: List[AllMessageValues], + optional_params: dict, + litellm_params: dict, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + ) -> dict: + headers = super().validate_environment( + headers=headers, + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + api_key=api_key, + api_base=api_base, + ) + project_id = litellm_params.get("aws_bedrock_project_id") + if project_id: + headers["OpenAI-Project"] = project_id + return headers + def get_supported_openai_params(self, model: str) -> list: base_params = super().get_supported_openai_params(model) try: diff --git a/litellm/llms/bedrock_mantle/responses/transformation.py b/litellm/llms/bedrock_mantle/responses/transformation.py index dfa108833ac..29248e1ca50 100644 --- a/litellm/llms/bedrock_mantle/responses/transformation.py +++ b/litellm/llms/bedrock_mantle/responses/transformation.py @@ -115,6 +115,8 @@ class BedrockMantleResponsesAPIConfig(OpenAIResponsesAPIConfig): ) if api_key: headers["Authorization"] = f"Bearer {api_key}" + if litellm_params.aws_bedrock_project_id: + headers["OpenAI-Project"] = litellm_params.aws_bedrock_project_id return headers def supports_native_file_search(self) -> bool: diff --git a/litellm/main.py b/litellm/main.py index 2c416a595c4..02609217ddb 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -1639,6 +1639,7 @@ def completion( # type: ignore # noqa: PLR0915 tpm=kwargs.get("tpm"), rpm=kwargs.get("rpm"), use_xai_oauth=kwargs.get("use_xai_oauth", False), + aws_bedrock_project_id=kwargs.get("aws_bedrock_project_id"), ) cast(LiteLLMLoggingObj, logging).update_environment_variables( model=model, diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 71cf5197dec..c868d3d22b2 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -271,6 +271,11 @@ _BANNED_REQUEST_BODY_PARAMS: Tuple[str, ...] = ( # tokens) to the attacker's host, or coerces the proxy into # authenticating against the attacker's host with admin secrets. "aws_bedrock_runtime_endpoint", + # Bedrock project/workspace association. Deployments pin this to + # enforce a data-retention policy, so a caller-supplied value would + # re-route the request's retention and accounting to any project + # reachable with the deployment's shared AWS credentials. + "aws_bedrock_project_id", # Provider-specific endpoint overrides that flow into the outbound # request via ``optional_params``. Same threat as ``api_base``: # ``s3_endpoint_url`` redirects Bedrock file uploads to attacker diff --git a/litellm/types/router.py b/litellm/types/router.py index ed858557a61..5047cee424b 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -178,6 +178,7 @@ class CredentialLiteLLMParams(BaseModel): aws_secret_access_key: Optional[str] = None aws_region_name: Optional[str] = None aws_bedrock_runtime_endpoint: Optional[str] = None + aws_bedrock_project_id: Optional[str] = None ## IBM WATSONX ## watsonx_region_name: Optional[str] = None @@ -364,6 +365,7 @@ class LiteLLMParamsTypedDict(TypedDict, total=False): aws_access_key_id: Optional[str] aws_secret_access_key: Optional[str] aws_region_name: Optional[str] + aws_bedrock_project_id: Optional[str] ## AWS S3 VECTORS ## vector_bucket_name: Optional[str] index_name: Optional[str] diff --git a/litellm/utils.py b/litellm/utils.py index 03c628b195f..a0b66234a70 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -3837,6 +3837,10 @@ class PreProcessNonDefaultParams: additional_endpoint_specific_params: List[str], ) -> dict: for k, v in special_params.items(): + if k == "aws_bedrock_project_id": + # sent as a request header (read from litellm_params by the + # bedrock-mantle configs), never as a request body field + continue if k.startswith("aws_") and ( custom_llm_provider != "bedrock" and not custom_llm_provider.startswith("sagemaker") diff --git a/tests/test_litellm/llms/bedrock/test_mantle.py b/tests/test_litellm/llms/bedrock/test_mantle.py index a00057eaa6b..bbefdd621f0 100644 --- a/tests/test_litellm/llms/bedrock/test_mantle.py +++ b/tests/test_litellm/llms/bedrock/test_mantle.py @@ -1,10 +1,17 @@ """ Unit tests for the Bedrock Mantle (Claude Mythos Preview) integration. -Tests cover route detection, URL construction, and config dispatch for both -the /chat/completions and /messages endpoints. +Tests cover route detection, URL construction, config dispatch for both +the /chat/completions and /messages endpoints, and project (workspace) +association via `aws_bedrock_project_id`. """ +import json +from unittest.mock import patch + +import httpx +import pytest + from litellm.llms.bedrock.common_utils import BedrockModelInfo, get_bedrock_chat_config from litellm.llms.bedrock.chat.mantle.transformation import AmazonMantleConfig from litellm.llms.bedrock.messages.mantle_transformation import ( @@ -12,6 +19,32 @@ from litellm.llms.bedrock.messages.mantle_transformation import ( ) +def _anthropic_response(url: str) -> httpx.Response: + return httpx.Response( + status_code=200, + json={ + "id": "msg_test", + "type": "message", + "role": "assistant", + "model": "anthropic.claude-mythos-preview", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + request=httpx.Request("POST", url), + ) + + +def _capture_request(url: str, headers: dict, data) -> dict: + raw_body = data.decode("utf-8") if isinstance(data, bytes) else data or "{}" + return { + "path": httpx.URL(url).path, + "headers": headers, + "body": json.loads(raw_body), + } + + def test_get_bedrock_route_mantle(): assert ( BedrockModelInfo.get_bedrock_route("mantle/anthropic.claude-mythos-preview") @@ -103,3 +136,114 @@ def test_mantle_transform_request_strips_prefix_and_adds_model(): ) assert request["model"] == "anthropic.claude-mythos-preview" assert "mantle/" not in request["model"] + + +def test_mantle_validate_environment_sets_workspace_header(): + config = AmazonMantleConfig() + headers = config.validate_environment( + headers={}, + model="mantle/anthropic.claude-mythos-preview", + messages=[{"role": "user", "content": "Hello"}], + optional_params={}, + litellm_params={"aws_bedrock_project_id": "proj_abc123def456"}, + ) + assert headers["anthropic-workspace"] == "proj_abc123def456" + + +def test_mantle_validate_environment_without_project_id(): + config = AmazonMantleConfig() + headers = config.validate_environment( + headers={}, + model="mantle/anthropic.claude-mythos-preview", + messages=[{"role": "user", "content": "Hello"}], + optional_params={}, + litellm_params={"aws_bedrock_project_id": None}, + ) + assert "anthropic-workspace" not in headers + + +def test_mantle_messages_validate_environment_sets_workspace_header(): + config = AmazonMantleMessagesConfig() + headers, api_base = config.validate_anthropic_messages_environment( + headers={}, + model="mantle/anthropic.claude-mythos-preview", + messages=[{"role": "user", "content": "Hello"}], + optional_params={}, + litellm_params={"aws_bedrock_project_id": "proj_abc123def456"}, + api_base="https://bedrock-mantle.us-east-1.api.aws/anthropic/v1/messages", + ) + assert headers["anthropic-workspace"] == "proj_abc123def456" + assert api_base == "https://bedrock-mantle.us-east-1.api.aws/anthropic/v1/messages" + + +def test_mantle_messages_validate_environment_without_project_id(): + config = AmazonMantleMessagesConfig() + headers, _ = config.validate_anthropic_messages_environment( + headers={}, + model="mantle/anthropic.claude-mythos-preview", + messages=[{"role": "user", "content": "Hello"}], + optional_params={}, + litellm_params={}, + ) + assert "anthropic-workspace" not in headers + + +def test_mantle_completion_sends_workspace_header_and_clean_body(): + import litellm + + requests = [] + + def mock_post(self, url, data=None, headers=None, **kwargs): + requests.append(_capture_request(url=url, headers=headers or {}, data=data)) + return _anthropic_response(url) + + with patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post", mock_post): + response = litellm.completion( + model="bedrock/mantle/anthropic.claude-mythos-preview", + messages=[{"role": "user", "content": "hello"}], + max_tokens=10, + aws_bedrock_project_id="proj_abc123def456", + aws_access_key_id="fake-key", + aws_secret_access_key="fake-secret", + aws_region_name="us-east-1", + ) + + assert response.choices[0].message.content == "ok" + assert len(requests) == 1 + assert requests[0]["path"] == "/anthropic/v1/messages" + assert requests[0]["headers"]["anthropic-workspace"] == "proj_abc123def456" + assert "aws_bedrock_project_id" not in requests[0]["body"] + + +@pytest.mark.asyncio +async def test_mantle_anthropic_messages_sends_workspace_header_and_clean_body(): + import litellm + + requests = [] + + async def mock_post(self, url, data=None, headers=None, **kwargs): + requests.append(_capture_request(url=url, headers=headers or {}, data=data)) + return _anthropic_response(url) + + try: + with patch( + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new=mock_post, + ): + response = await litellm.anthropic_messages( + model="bedrock/mantle/anthropic.claude-mythos-preview", + messages=[{"role": "user", "content": "hello"}], + max_tokens=10, + aws_bedrock_project_id="proj_abc123def456", + aws_access_key_id="fake-key", + aws_secret_access_key="fake-secret", + aws_region_name="us-east-1", + ) + finally: + await litellm.close_litellm_async_clients() + + assert response["content"][0]["text"] == "ok" + assert len(requests) == 1 + assert requests[0]["path"] == "/anthropic/v1/messages" + assert requests[0]["headers"]["anthropic-workspace"] == "proj_abc123def456" + assert "aws_bedrock_project_id" not in requests[0]["body"] diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py index e83992b6bde..c3de29bd9d5 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py @@ -167,6 +167,26 @@ class TestBedrockMantleResponsesAuth: ) assert "Authorization" not in headers + def test_project_id_sets_openai_project_header(self): + cfg = BedrockMantleResponsesAPIConfig() + headers = cfg.validate_environment( + headers={}, + model="openai.gpt-5.5", + litellm_params=GenericLiteLLMParams( + api_key="fake-key", aws_bedrock_project_id="proj_abc123def456" + ), + ) + assert headers["OpenAI-Project"] == "proj_abc123def456" + + def test_no_project_id_no_openai_project_header(self): + cfg = BedrockMantleResponsesAPIConfig() + headers = cfg.validate_environment( + headers={}, + model="openai.gpt-5.5", + litellm_params=GenericLiteLLMParams(api_key="fake-key"), + ) + assert "OpenAI-Project" not in headers + def test_custom_llm_provider(self): cfg = BedrockMantleResponsesAPIConfig() assert cfg.custom_llm_provider == LlmProviders.BEDROCK_MANTLE diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py index 1725aa85d10..deaa0537930 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py @@ -5,11 +5,14 @@ Bedrock Mantle is Amazon Bedrock's OpenAI-compatible inference engine (Project M API docs: https://docs.aws.amazon.com/bedrock/latest/userguide/bedrock-mantle.html """ +import json import os import sys +from unittest.mock import patch sys.path.insert(0, os.path.abspath("../../../../..")) +import httpx import pytest import litellm @@ -96,6 +99,79 @@ class TestBedrockMantleConfig: assert "max_tokens" in params +class TestBedrockMantleProjectHeader: + def test_validate_environment_sets_openai_project_header(self): + cfg = BedrockMantleChatConfig() + headers = cfg.validate_environment( + headers={}, + model="openai.gpt-oss-120b", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={"aws_bedrock_project_id": "proj_abc123def456"}, + api_key="fake-key", + ) + assert headers["OpenAI-Project"] == "proj_abc123def456" + assert headers["Authorization"] == "Bearer fake-key" + + def test_validate_environment_without_project_id(self): + cfg = BedrockMantleChatConfig() + headers = cfg.validate_environment( + headers={}, + model="openai.gpt-oss-120b", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + api_key="fake-key", + ) + assert "OpenAI-Project" not in headers + + def test_completion_sends_openai_project_header_and_clean_body(self): + requests = [] + + def mock_post(self, url, data=None, headers=None, **kwargs): + raw_body = data.decode("utf-8") if isinstance(data, bytes) else data + requests.append( + {"headers": headers or {}, "body": json.loads(raw_body or "{}")} + ) + return httpx.Response( + status_code=200, + json={ + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 1733529600, + "model": "openai.gpt-oss-120b", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "ok"}, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 1, + "completion_tokens": 1, + "total_tokens": 2, + }, + }, + request=httpx.Request("POST", url), + ) + + with patch( + "litellm.llms.custom_httpx.http_handler.HTTPHandler.post", mock_post + ): + response = litellm.completion( + model="bedrock_mantle/openai.gpt-oss-120b", + messages=[{"role": "user", "content": "hello"}], + api_key="fake-key", + aws_bedrock_project_id="proj_abc123def456", + ) + + assert response.choices[0].message.content == "ok" + assert len(requests) == 1 + assert requests[0]["headers"]["OpenAI-Project"] == "proj_abc123def456" + assert "aws_bedrock_project_id" not in requests[0]["body"] + + class TestBedrockMantleProviderResolution: def test_get_llm_provider_resolves_correctly(self): model, provider, _, _ = litellm.get_llm_provider( diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index d4ca55ca16b..32b597376b4 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -1551,6 +1551,40 @@ class TestIsRequestBodySafeBlocksEndpointTargetingFields: ) +class TestIsRequestBodySafeBlocksBedrockProjectOverride: + """``aws_bedrock_project_id`` pins a deployment to a Bedrock project so + that project's data-retention policy applies to its requests. A + caller-supplied value would run the request under any project reachable + with the deployment's shared AWS credentials, bypassing the configured + retention/accounting association.""" + + def test_project_id_in_request_body_is_rejected(self): + with pytest.raises(ValueError, match="aws_bedrock_project_id"): + is_request_body_safe( + request_body={ + "model": "gpt-4", + "aws_bedrock_project_id": "proj_attacker000000", + }, + general_settings={}, + llm_router=None, + model="gpt-4", + ) + + def test_admin_opt_in_proxy_wide_allows_project_id(self): + assert ( + is_request_body_safe( + request_body={ + "model": "gpt-4", + "aws_bedrock_project_id": "proj_byok000000", + }, + general_settings={"allow_client_side_credentials": True}, + llm_router=None, + model="gpt-4", + ) + is True + ) + + # ── is_request_body_safe nested-config recursion (VERIA-6) ──────────────────── diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 62cb8154b6d..b6c9e9c865d 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -4198,3 +4198,21 @@ class TestBedrockBaseModelLabelKeepsTools: ) assert "tools" not in result + + +def test_aws_bedrock_project_id_excluded_from_bedrock_optional_params(): + """`aws_bedrock_project_id` is sent as a bedrock-mantle request header, so it + must never reach optional_params (and from there the request body), while + other aws_* params keep flowing for boto3 auth.""" + from litellm.utils import get_optional_params + + result = get_optional_params( + model="mantle/anthropic.claude-mythos-preview", + custom_llm_provider="bedrock", + max_tokens=10, + aws_bedrock_project_id="proj_abc123def456", + aws_region_name="us-east-1", + ) + + assert "aws_bedrock_project_id" not in result + assert result["aws_region_name"] == "us-east-1" diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 8e470819557..15123bcdbf8 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -24891,6 +24891,8 @@ export interface components { auto_router_embedding_model?: string | null; /** Aws Access Key Id */ aws_access_key_id?: string | null; + /** Aws Bedrock Project Id */ + aws_bedrock_project_id?: string | null; /** Aws Bedrock Runtime Endpoint */ aws_bedrock_runtime_endpoint?: string | null; /** Aws Region Name */ @@ -32495,6 +32497,8 @@ export interface components { auto_router_embedding_model?: string | null; /** Aws Access Key Id */ aws_access_key_id?: string | null; + /** Aws Bedrock Project Id */ + aws_bedrock_project_id?: string | null; /** Aws Bedrock Runtime Endpoint */ aws_bedrock_runtime_endpoint?: string | null; /** Aws Region Name */ From 0d120de785cef131c44fa977e0861af13f9ebfe3 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Thu, 11 Jun 2026 10:00:23 -0700 Subject: [PATCH 064/185] chore(hooks): enforce Conventional Commits and Conventional Branches (#30174) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * chore(hooks): enforce Conventional Commits and Conventional Branches Adds opt-in local git hooks plus a CI PR-title check: - .githooks/commit-msg validates commit subjects against Conventional Commits 1.0.0 (feat|fix|docs|style|refactor|perf|test|build|ci| chore|revert)(scope)!: subject. Merge/revert/fixup!/squash!/amend! messages pass through; --no-verify still works. - .githooks/pre-push validates branch names against Conventional Branches (feature|bugfix|hotfix|release|chore)/desc. Bypasses main, litellm_internal_staging, dependabot/*, gh-readonly-queue/*. Tag pushes and deletions are skipped. - scripts/install_git_hooks.sh sets core.hooksPath=.githooks and is wired up as 'make install-hooks'. Opt-in — not chained into install-dev. - .github/workflows/conventional-commits.yml validates PR titles via amannn/action-semantic-pull-request pinned to v6.1.1's SHA. This is the actual gate since squash-merge uses the PR title as the commit subject. - tests/test_litellm/test_git_hooks.py exercises both hooks via subprocess for accept / reject / bypass / git-generated-message cases. - CONTRIBUTING.md documents the conventions, the install step, the bypass list, and the --no-verify escape hatch. Resolves LIT-3306 Co-Authored-By: Claude Opus 4.7 (1M context) * fix(hooks): address Greptile review on PR #28703 Resolves two findings from the automated code review: 1. CONTRIBUTING.md: shrink the new Conventional Commits / Branches section to a 2-line pointer at docs.litellm.ai. Per the team convention, the full documentation lives in the litellm-docs repo — see BerriAI/litellm-docs#208 for the companion change that adds the section to docs/extras/contributing_code.md. 2. .githooks/commit-msg: tighten the subject regex to also reject an uppercase first letter in the description. CI's subjectPattern is ^(?![A-Z]).+$ so the previous local hook would accept 'feat: Add thing' which would then fail the PR-title check. The local hook is now the strictly tighter of the two gates. Test cases extended to cover both the new rejection and the digit/symbol-start cases that remain allowed. Resolves LIT-3306 Co-Authored-By: Claude Opus 4.7 (1M context) * chore: trigger ci after branch rename * fix(ci): rerun pr title check when bypass label changes amannn/action-semantic-pull-request only honors ignoreLabels if the workflow retriggers on labeled/unlabeled events; without them a red check stays red after a maintainer applies the bypass label. Also point the CONTRIBUTING.md workflow comments at the conventions section, which now sits above the Development Workflow section. --------- Co-authored-by: Yassin Kortam Co-authored-by: Claude Opus 4.7 (1M context) --- .githooks/commit-msg | 75 ++++++ .githooks/pre-push | 92 +++++++ .github/workflows/conventional-commits.yml | 46 ++++ CONTRIBUTING.md | 19 +- Makefile | 8 +- scripts/install_git_hooks.sh | 38 +++ tests/test_litellm/test_git_hooks.py | 286 +++++++++++++++++++++ 7 files changed, 557 insertions(+), 7 deletions(-) create mode 100755 .githooks/commit-msg create mode 100755 .githooks/pre-push create mode 100644 .github/workflows/conventional-commits.yml create mode 100755 scripts/install_git_hooks.sh create mode 100644 tests/test_litellm/test_git_hooks.py diff --git a/.githooks/commit-msg b/.githooks/commit-msg new file mode 100755 index 00000000000..b64e38a2286 --- /dev/null +++ b/.githooks/commit-msg @@ -0,0 +1,75 @@ +#!/usr/bin/env bash +# +# commit-msg — enforce Conventional Commits 1.0.0 +# https://www.conventionalcommits.org/en/v1.0.0/ +# +# Subject format: ()!: +# - must be one of the angular types (feat, fix, ...) +# - () is optional +# - ! is optional and marks a breaking change +# - is mandatory and must be non-empty +# +# Bypass: commit with --no-verify. +# Merge, revert, fixup!, squash!, and amend! messages are passed through. + +set -eu + +COMMIT_MSG_FILE="${1:-}" +if [ -z "$COMMIT_MSG_FILE" ] || [ ! -f "$COMMIT_MSG_FILE" ]; then + echo "commit-msg: missing commit message file" >&2 + exit 1 +fi + +# First non-comment, non-empty line is the subject. +subject="" +while IFS= read -r line || [ -n "$line" ]; do + case "$line" in + ''|'#'*) continue ;; + esac + subject="$line" + break +done < "$COMMIT_MSG_FILE" + +if [ -z "$subject" ]; then + echo "commit-msg: empty commit message" >&2 + exit 1 +fi + +# Pass-through commits generated by git itself. +case "$subject" in + "Merge "*|"Revert \""*|"fixup! "*|"squash! "*|"amend! "*) + exit 0 + ;; +esac + +ALLOWED_TYPES="feat|fix|docs|style|refactor|perf|test|build|ci|chore|revert" +# Description must not start with an uppercase letter — kept in sync with the +# subjectPattern in .github/workflows/conventional-commits.yml so the local +# hook is the strictly tighter of the two gates. (Without this guard, a commit +# like "feat: Add thing" passes locally but fails the PR-title CI check.) +PATTERN="^(${ALLOWED_TYPES})(\([^)]+\))?!?: [^A-Z].*" + +if printf '%s' "$subject" | grep -Eq "$PATTERN"; then + exit 0 +fi + +cat >&2 <()!: + (description must start with a lowercase letter) + + Allowed types: feat, fix, docs, style, refactor, perf, test, build, ci, chore, revert + Examples: + feat(router): add weighted round-robin strategy + fix(bedrock): decouple STS region from aws_region_name + chore(deps): bump black to 26.3.1 + refactor!: drop Python 3.8 support + +See https://www.conventionalcommits.org/en/v1.0.0/ + +To bypass (use sparingly): git commit --no-verify +EOF +exit 1 diff --git a/.githooks/pre-push b/.githooks/pre-push new file mode 100755 index 00000000000..c2267c8501c --- /dev/null +++ b/.githooks/pre-push @@ -0,0 +1,92 @@ +#!/usr/bin/env bash +# +# pre-push — enforce Conventional Branches +# https://conventional-branch.github.io/ +# +# Branch format: / +# must be one of: feature, bugfix, hotfix, release, chore +# +# Protected branches (always allowed): +# - main +# - litellm_internal_staging +# - dependabot/* +# - gh-readonly-queue/* +# +# Tag pushes and branch deletions are skipped. +# Bypass: git push --no-verify. + +set -eu + +ZERO_OID="0000000000000000000000000000000000000000" +ZERO_OID_SHA256="0000000000000000000000000000000000000000000000000000000000000000" +ALLOWED_TYPES="feature|bugfix|hotfix|release|chore" +BRANCH_PATTERN="^(${ALLOWED_TYPES})/.+" + +PROTECTED_NAMES="main litellm_internal_staging" +PROTECTED_PREFIXES="dependabot/ gh-readonly-queue/" + +is_protected() { + branch="$1" + for name in $PROTECTED_NAMES; do + if [ "$branch" = "$name" ]; then + return 0 + fi + done + for prefix in $PROTECTED_PREFIXES; do + case "$branch" in "$prefix"*) return 0 ;; esac + done + return 1 +} + +invalid="" + +while read -r local_ref local_oid remote_ref remote_oid; do + # Branch deletion (no local commit being pushed). + if [ "$local_oid" = "$ZERO_OID" ] || [ "$local_oid" = "$ZERO_OID_SHA256" ]; then + continue + fi + + # Only validate branch pushes; ignore tags and other ref namespaces. + case "$remote_ref" in + refs/heads/*) ;; + *) continue ;; + esac + + branch="${remote_ref#refs/heads/}" + + if is_protected "$branch"; then + continue + fi + + if ! printf '%s' "$branch" | grep -Eq "$BRANCH_PATTERN"; then + invalid="$invalid $branch" + fi +done + +if [ -n "$invalid" ]; then + cat >&2 </ + + Allowed types: feature, bugfix, hotfix, release, chore + Examples: + feature/weighted-round-robin + bugfix/streaming-empty-chunks + chore/bump-deps + hotfix/auth-bypass + + Protected (always allowed): main, litellm_internal_staging, + dependabot/*, gh-readonly-queue/*. + +See https://conventional-branch.github.io/ + +Rename with: git branch -m +To bypass (use sparingly): git push --no-verify +EOF + exit 1 +fi + +exit 0 diff --git a/.github/workflows/conventional-commits.yml b/.github/workflows/conventional-commits.yml new file mode 100644 index 00000000000..69ade24d028 --- /dev/null +++ b/.github/workflows/conventional-commits.yml @@ -0,0 +1,46 @@ +name: Conventional PR Title + +# Squash-merge replaces the merge commit subject with the PR title, so +# enforcing Conventional Commits at the PR-title level is what actually gates +# the commits that land on the default branch. The local commit-msg hook +# (.githooks/commit-msg) is a best-effort assist; this workflow is the gate. +# +# See https://www.conventionalcommits.org/en/v1.0.0/ + +on: + pull_request: + types: [opened, edited, reopened, synchronize, labeled, unlabeled] + +permissions: + pull-requests: read + +jobs: + lint-pr-title: + name: Validate PR title + runs-on: ubuntu-latest + steps: + - name: Check title against Conventional Commits + uses: amannn/action-semantic-pull-request@48f256284bd46cdaab1048c3721360e808335d50 # v6.1.1 + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + with: + # Must mirror the type list in .githooks/commit-msg. + types: | + feat + fix + docs + style + refactor + perf + test + build + ci + chore + revert + requireScope: false + subjectPattern: ^(?![A-Z]).+$ + subjectPatternError: | + The subject "{subject}" must start with a lowercase character. + # Allow merges/reverts that GitHub generates automatically. + ignoreLabels: | + ignore-semantic-pull-request diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 8ac83341f64..2177c764806 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -38,18 +38,25 @@ Before contributing code to LiteLLM, you must sign our [Contributor License Agre git clone https://github.com/YOUR_USERNAME/litellm.git cd litellm -# Create a new branch for your feature -git checkout -b your-feature-branch +# Create a new branch for your feature (see "Commit and Branch Conventions" below) +git checkout -b feature/your-feature # Install development dependencies make install-dev +# Install git hooks that enforce commit + branch conventions (one-time, opt-in) +make install-hooks + # Verify your setup works make help ``` That's it! Your local development environment is ready. +## Commit and Branch Conventions + +Commits follow [Conventional Commits](https://www.conventionalcommits.org/en/v1.0.0/) and branches follow [Conventional Branches](https://conventional-branch.github.io/). Run `make install-hooks` once per clone to enable the local git hooks that enforce these — see the [contributor docs](https://docs.litellm.ai/docs/extras/contributing_code#commit-and-branch-conventions) for the full type list, examples, the protected-branch bypass list, and how to opt out. + ### 2. Development Workflow Here's the recommended workflow for making changes: @@ -67,12 +74,12 @@ make lint # Run unit tests to ensure nothing is broken make test-unit -# Commit your changes +# Commit your changes (must follow Conventional Commits — see above) git add . -git commit -m "Your descriptive commit message" +git commit -m "feat(scope): your descriptive commit message" -# Push and create a PR -git push origin your-feature-branch +# Push and create a PR (branch must follow Conventional Branches — see above) +git push origin feature/your-feature ``` ## Adding Testing diff --git a/Makefile b/Makefile index a00a90da601..3d7b51bc745 100644 --- a/Makefile +++ b/Makefile @@ -5,7 +5,7 @@ test-unit-integrations test-unit-core-utils test-unit-other test-unit-root \ test-proxy-unit-a test-proxy-unit-b test-integration test-unit-helm \ info lint lint-dev format \ - install-dev install-proxy-dev install-test-deps \ + install-dev install-proxy-dev install-test-deps install-hooks \ install-helm-unittest check-circular-imports check-import-safety # Default target @@ -17,6 +17,7 @@ help: @echo " make install-proxy-dev-ci - Install proxy dev dependencies (CI-compatible)" @echo " make install-test-deps - Install the full local test environment" @echo " make install-helm-unittest - Install helm unittest plugin" + @echo " make install-hooks - Install git hooks (Conventional Commits + Branches)" @echo " make format - Apply Black code formatting" @echo " make format-check - Check Black code formatting (matches CI)" @echo " make lint - Run all linting (Ruff, MyPy, Black check, circular imports, import safety)" @@ -68,6 +69,11 @@ install-test-deps: install-proxy-dev install-helm-unittest: helm plugin install https://github.com/helm-unittest/helm-unittest --version v0.4.4 || echo "ignore error if plugin exists" +# Install git hooks that enforce Conventional Commits and Conventional Branches. +# Opt-in: not chained into install-dev. +install-hooks: + ./scripts/install_git_hooks.sh + # Formatting format: install-dev cd litellm && $(UV_RUN) black . && cd .. diff --git a/scripts/install_git_hooks.sh b/scripts/install_git_hooks.sh new file mode 100755 index 00000000000..1e4e3c6de19 --- /dev/null +++ b/scripts/install_git_hooks.sh @@ -0,0 +1,38 @@ +#!/usr/bin/env bash +# +# Install the repo's git hooks by pointing core.hooksPath at .githooks. +# +# Idempotent: re-running just reaffirms the config and refreshes chmod bits. +# Run from anywhere inside the repo. + +set -euo pipefail + +if ! git rev-parse --is-inside-work-tree >/dev/null 2>&1; then + echo "install_git_hooks: not inside a git working tree" >&2 + exit 1 +fi + +repo_root=$(git rev-parse --show-toplevel) +hooks_dir="$repo_root/.githooks" + +if [ ! -d "$hooks_dir" ]; then + echo "install_git_hooks: $hooks_dir does not exist" >&2 + exit 1 +fi + +# Ensure the hook scripts are executable. New clones on case-preserving +# filesystems sometimes lose the exec bit; this normalizes it. +chmod +x "$hooks_dir"/* 2>/dev/null || true + +git config core.hooksPath .githooks + +cat < subprocess.CompletedProcess: + msg_file = tmp_path / "COMMIT_EDITMSG" + msg_file.write_text(subject + "\n", encoding="utf-8") + return subprocess.run( + ["bash", str(_COMMIT_MSG_HOOK), str(msg_file)], + capture_output=True, + text=True, + check=False, + ) + + +def _run_pre_push(stdin: str) -> subprocess.CompletedProcess: + return subprocess.run( + ["bash", str(_PRE_PUSH_HOOK)], + input=stdin, + capture_output=True, + text=True, + check=False, + ) + + +def _ref_line(branch: str, local_oid: str = _NONZERO_OID, remote_oid: str = _ZERO_OID) -> str: + ref = f"refs/heads/{branch}" + return f"{ref} {local_oid} {ref} {remote_oid}\n" + + +# ----- commit-msg ----------------------------------------------------------- + + +@pytest.mark.parametrize( + "subject", + [ + "feat(router): add weighted round-robin strategy", + "fix(bedrock): decouple STS region from aws_region_name", + "chore(deps): bump black to 26.3.1", + "docs: rewrite contributing guide", + "refactor!: drop Python 3.8 support", + "feat(api,proxy)!: rename endpoint", + "test: cover hook bypass list", + "perf(streaming): avoid extra json parse", + "revert: feat(router): add weighted round-robin", + ], +) +def test_commit_msg_accepts_conventional_subjects(tmp_path, subject): + result = _run_commit_msg(subject, tmp_path) + assert result.returncode == 0, ( + f"hook rejected a valid subject:\n subject: {subject!r}\n" + f" stderr: {result.stderr}" + ) + + +@pytest.mark.parametrize( + "subject", + [ + "add stuff", # no type + "feat add router strategy", # missing colon + "feat:add router strategy", # missing space after colon + "feat():", # empty description + "ux: thing", # unknown type + "Feat(router): capital type", # types are lowercase + "feat(router):", # empty description + # Description must start with a lowercase letter — kept in sync with + # the CI workflow's subjectPattern so the local hook never accepts a + # subject that CI will later reject. + "feat: Add thing", + "fix(router): Decouple something", + "chore: BUMP deps", + "feat: A", + ], +) +def test_commit_msg_rejects_invalid_subjects(tmp_path, subject): + result = _run_commit_msg(subject, tmp_path) + assert result.returncode == 1, ( + f"hook accepted an invalid subject:\n subject: {subject!r}\n" + f" stderr: {result.stderr}" + ) + assert "Conventional Commits" in result.stderr + + +@pytest.mark.parametrize( + "subject", + [ + # Lowercase letter — the common case. + "feat: lowercase start is fine", + # The CI's `^(?![A-Z]).+$` rejects only uppercase A-Z, so digits and + # symbols are still allowed; mirror that behavior here. + "feat: 1-based indexing now works", + "fix(deps): @types/node bump", + ], +) +def test_commit_msg_accepts_non_uppercase_starts(tmp_path, subject): + result = _run_commit_msg(subject, tmp_path) + assert result.returncode == 0, ( + f"hook rejected a valid non-uppercase-start subject:\n" + f" subject: {subject!r}\n stderr: {result.stderr}" + ) + + +@pytest.mark.parametrize( + "subject", + [ + "Merge branch 'main' into feature/foo", + 'Revert "feat(router): add weighted round-robin strategy"', + "fixup! feat(router): add weighted round-robin strategy", + "squash! feat(router): add weighted round-robin strategy", + "amend! feat(router): add weighted round-robin strategy", + ], +) +def test_commit_msg_passes_git_generated_messages(tmp_path, subject): + result = _run_commit_msg(subject, tmp_path) + assert result.returncode == 0, ( + f"hook should pass git-generated subject:\n subject: {subject!r}\n" + f" stderr: {result.stderr}" + ) + + +def test_commit_msg_rejects_empty_message(tmp_path): + result = _run_commit_msg("", tmp_path) + assert result.returncode == 1 + assert "empty commit message" in result.stderr + + +def test_commit_msg_skips_comment_only_lines(tmp_path): + # An all-comments file has no subject — should be rejected. + msg_file = tmp_path / "COMMIT_EDITMSG" + msg_file.write_text("# please enter a commit message\n# above this line\n", encoding="utf-8") + result = subprocess.run( + ["bash", str(_COMMIT_MSG_HOOK), str(msg_file)], + capture_output=True, + text=True, + check=False, + ) + assert result.returncode == 1 + assert "empty commit message" in result.stderr + + +def test_commit_msg_uses_first_non_comment_line(tmp_path): + # Real git-generated COMMIT_EDITMSG has a status block prefixed with '#' + # below the subject. Make sure leading comment lines are skipped too. + msg_file = tmp_path / "COMMIT_EDITMSG" + msg_file.write_text( + "# On branch feature/foo\n" + "\n" + "feat(router): add weighted round-robin\n" + "\n" + "# Please enter the commit message...\n", + encoding="utf-8", + ) + result = subprocess.run( + ["bash", str(_COMMIT_MSG_HOOK), str(msg_file)], + capture_output=True, + text=True, + check=False, + ) + assert result.returncode == 0, result.stderr + + +# ----- pre-push ------------------------------------------------------------- + + +@pytest.mark.parametrize( + "branch", + [ + "feature/weighted-round-robin", + "bugfix/streaming-empty-chunks", + "hotfix/auth-bypass", + "release/v1.45.0", + "chore/bump-deps", + "feature/nested/path/ok", # nested slashes after type are fine + ], +) +def test_pre_push_accepts_conventional_branches(branch): + result = _run_pre_push(_ref_line(branch)) + assert result.returncode == 0, ( + f"hook rejected a valid branch:\n branch: {branch!r}\n" + f" stderr: {result.stderr}" + ) + + +@pytest.mark.parametrize( + "branch", + [ + "random-branch-name", + "litellm_fix/optimize-streaming", # legacy pattern is now rejected + "ui/navbar-notifications", # not in the allow list + "feature/", # empty description + "Feature/foo", # type is case-sensitive + "feat/foo", # angular commit type, not branch type + ], +) +def test_pre_push_rejects_non_conventional_branches(branch): + result = _run_pre_push(_ref_line(branch)) + assert result.returncode == 1, ( + f"hook accepted an invalid branch:\n branch: {branch!r}\n" + f" stderr: {result.stderr}" + ) + assert "Conventional Branches" in result.stderr + + +@pytest.mark.parametrize( + "branch", + [ + "main", + "litellm_internal_staging", + "dependabot/github_actions/foo", + "gh-readonly-queue/main/abc123", + ], +) +def test_pre_push_bypasses_protected_branches(branch): + result = _run_pre_push(_ref_line(branch)) + assert result.returncode == 0, ( + f"protected branch was rejected:\n branch: {branch!r}\n" + f" stderr: {result.stderr}" + ) + + +def test_pre_push_skips_tag_pushes(): + line = f"refs/tags/v1 {_NONZERO_OID} refs/tags/v1 {_ZERO_OID}\n" + result = _run_pre_push(line) + assert result.returncode == 0, result.stderr + + +def test_pre_push_skips_branch_deletions(): + # local oid all zeros = deletion + line = f"refs/heads/whatever {_ZERO_OID} refs/heads/whatever {_NONZERO_OID}\n" + result = _run_pre_push(line) + assert result.returncode == 0, result.stderr + + +def test_pre_push_fails_if_any_ref_is_invalid(): + # Mixed batch: one valid, one invalid — entire push should fail. + stdin = _ref_line("feature/ok") + _ref_line("random-bad") + result = _run_pre_push(stdin) + assert result.returncode == 1 + assert "random-bad" in result.stderr + + +def test_pre_push_no_refs_passes(): + # Empty stdin (no refs being pushed) should pass. + result = _run_pre_push("") + assert result.returncode == 0, result.stderr From 012d9f6c0a3f6bbe8d284d2f74fe9b57bdd26835 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Thu, 11 Jun 2026 10:34:26 -0700 Subject: [PATCH 065/185] feat(rate-limiter): allow opting out of v3 TPM reservation and Redis circuit breaker (#30211) --- litellm/caching/redis_cache.py | 16 ++- litellm/constants.py | 3 + .../hooks/parallel_request_limiter_v3.py | 28 +++-- scripts/health_check/health_check_client.py | 4 +- tests/test_litellm/caching/test_dual_cache.py | 61 ++++++++++ .../hooks/test_parallel_request_limiter_v3.py | 108 ++++++++++++++++++ 6 files changed, 208 insertions(+), 12 deletions(-) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index cb9ce475d30..7239bea7853 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -22,6 +22,7 @@ import litellm from litellm._logging import print_verbose, verbose_logger from litellm.constants import ( DEFAULT_REDIS_MAJOR_VERSION, + REDIS_CIRCUIT_BREAKER_ENABLED, REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD, REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT, ) @@ -114,15 +115,23 @@ class RedisCircuitBreaker: OPEN = "open" HALF_OPEN = "half_open" - def __init__(self, failure_threshold: int, recovery_timeout: int) -> None: + def __init__( + self, + failure_threshold: int, + recovery_timeout: int, + enabled: bool = True, + ) -> None: self.failure_threshold = failure_threshold self.recovery_timeout = recovery_timeout + self.enabled = enabled self._failure_count = 0 self._opened_at: Optional[float] = None self._state = self.CLOSED def is_open(self) -> bool: """Returns True if Redis calls should be skipped.""" + if not self.enabled: + return False if self._state == self.HALF_OPEN: # Probe already in flight — fast-fail all concurrent requests. # Only the one call that caused the OPEN→HALF_OPEN transition @@ -136,6 +145,8 @@ class RedisCircuitBreaker: return False def record_failure(self) -> None: + if not self.enabled: + return self._failure_count += 1 self._opened_at = time.time() if self._failure_count >= self.failure_threshold: @@ -149,6 +160,8 @@ class RedisCircuitBreaker: self._state = self.OPEN def record_success(self) -> None: + if not self.enabled: + return if self._state == self.HALF_OPEN: verbose_logger.info("Redis circuit breaker CLOSED — Redis recovered") self._failure_count = 0 @@ -243,6 +256,7 @@ class RedisCache(BaseCache): self._circuit_breaker = RedisCircuitBreaker( failure_threshold=REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD, recovery_timeout=REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT, + enabled=REDIS_CIRCUIT_BREAKER_ENABLED, ) self._setup_health_pings() diff --git a/litellm/constants.py b/litellm/constants.py index 57f55e6c177..ab8e57d735f 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -398,6 +398,9 @@ REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD = int( REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT = int( os.getenv("REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT", 60) ) +REDIS_CIRCUIT_BREAKER_ENABLED = ( + os.getenv("REDIS_CIRCUIT_BREAKER_ENABLED", "true").lower() == "true" +) # Default Redis major version to assume when version cannot be determined # Using 7 as it's the modern version that supports LPOP with count parameter DEFAULT_REDIS_MAJOR_VERSION = int(os.getenv("DEFAULT_REDIS_MAJOR_VERSION", 7)) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 6b70cea65a3..f45c63d1380 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -313,6 +313,14 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): self.window_size = int(os.getenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", 60)) + # When disabled, TPM is enforced post-call from actual usage (pre-v1.82 + # behavior) instead of reserving an estimated budget upfront, shedding + # the extra per-request Redis Lua round-trip and the global-lock + # in-memory fallback that the reservation path incurs. + self.tpm_reservation_enabled = ( + os.getenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", "true").lower() == "true" + ) + # Batch rate limiter (lazy loaded) self._batch_rate_limiter: Optional[Any] = None @@ -2113,17 +2121,19 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # Only check rate limits if we have descriptors with actual limits if descriptors: # First pass: RPM and max_parallel_requests sliding-window check. - # `skip_tpm_check=True` tells should_rate_limit to ignore each - # descriptor's tokens_per_unit so its +1-per-key Lua / in-memory - # increment never touches the :tokens counters — those are owned - # exclusively by the atomic reserve_tpm_tokens path below. Without - # this, every concurrent in-flight request would pre-inflate the - # :tokens counter by 1, shrinking the effective TPM budget by N - # and causing false-positive 429s under bursts. + # When reservation is enabled, `skip_tpm_check=True` tells + # should_rate_limit to ignore each descriptor's tokens_per_unit so + # its +1-per-key Lua / in-memory increment never touches the + # :tokens counters — those are owned exclusively by the atomic + # reserve_tpm_tokens path below. Without this, every concurrent + # in-flight request would pre-inflate the :tokens counter by 1, + # shrinking the effective TPM budget by N and causing + # false-positive 429s under bursts. When reservation is disabled, + # this pass enforces TPM directly from the post-call counters. response = await self.should_rate_limit( descriptors=descriptors, parent_otel_span=user_api_key_dict.parent_otel_span, - skip_tpm_check=True, + skip_tpm_check=self.tpm_reservation_enabled, ) if response["overall_code"] == "OVER_LIMIT": @@ -2153,7 +2163,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ] has_tpm_limits = bool(configured_tpm_limits) - if has_tpm_limits: + if has_tpm_limits and self.tpm_reservation_enabled: min_configured_tpm_limit = min(configured_tpm_limits) # When the configured TPM cap is small enough to constrain the diff --git a/scripts/health_check/health_check_client.py b/scripts/health_check/health_check_client.py index 497fd6271b4..9ef8b934961 100644 --- a/scripts/health_check/health_check_client.py +++ b/scripts/health_check/health_check_client.py @@ -54,7 +54,7 @@ class LiteLLMHealthCheckClient: timeout: Request timeout in seconds (default: 120, matching Go implementation) completion_prompt: Test prompt for chat/completion models embedding_text: Test text for embedding models - custom_auth_header: Optional custom header name for authentication (e.g., "x-ifood-requester-service"). + custom_auth_header: Optional custom header name for authentication (e.g., "x-requester-service"). If provided, uses this header instead of standard "Authorization" header. """ self.base_url = base_url.rstrip("/") @@ -404,7 +404,7 @@ async def main(): yaml_path = os.environ.get("LITELLM_MODELS_YAML") custom_auth_header = os.environ.get( "LITELLM_CUSTOM_AUTH_HEADER" - ) # e.g., "x-ifood-requester-service" + ) # e.g., "x-requester-service" # Debug: Print custom auth header value if set if custom_auth_header: diff --git a/tests/test_litellm/caching/test_dual_cache.py b/tests/test_litellm/caching/test_dual_cache.py index 64774726201..f4f88def78d 100644 --- a/tests/test_litellm/caching/test_dual_cache.py +++ b/tests/test_litellm/caching/test_dual_cache.py @@ -243,6 +243,67 @@ def test_circuit_breaker_half_open_concurrent_calls_are_fast_failed(): ), "concurrent callers should be fast-failed in HALF_OPEN" +def test_circuit_breaker_disabled_never_opens(): + """When disabled, failures never open the circuit and is_open() stays False.""" + from litellm.caching.redis_cache import RedisCircuitBreaker + + cb = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60, enabled=False) + + for _ in range(100): + cb.record_failure() + + assert cb._state == "closed" + assert cb.is_open() is False + + +def test_circuit_breaker_disabled_record_success_leaves_state_untouched(): + """ + A disabled breaker must not mutate state in any state-machine method. Force + a non-default (OPEN) state and assert record_success() returns without + resetting it — the same enabled-guard contract as is_open/record_failure. + """ + from litellm.caching.redis_cache import RedisCircuitBreaker + + cb = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60, enabled=False) + cb._state = "open" + cb._failure_count = 3 + + cb.record_success() + + assert cb._state == "open" + assert cb._failure_count == 3 + + +@pytest.mark.asyncio +async def test_circuit_breaker_disabled_guard_always_calls_method(): + """A disabled breaker lets every guarded call through, even after failures.""" + from litellm.caching.redis_cache import ( + RedisCircuitBreaker, + _redis_circuit_breaker_guard, + ) + + class FakeRedis: + def __init__(self): + self._circuit_breaker = RedisCircuitBreaker( + failure_threshold=1, recovery_timeout=60, enabled=False + ) + self.call_count = 0 + + @_redis_circuit_breaker_guard + async def boom(self): + self.call_count += 1 + raise RuntimeError("redis down") + + fr = FakeRedis() + for _ in range(5): + with pytest.raises(RuntimeError, match="redis down"): + await fr.boom() + + # Every call reached the method body; the breaker never short-circuited. + assert fr.call_count == 5 + assert fr._circuit_breaker.is_open() is False + + @pytest.mark.asyncio async def test_async_increment_cache_returns_none_when_no_in_memory_cache_and_redis_fails(): """ diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index d10311b9f41..ae699ff8e12 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -3187,3 +3187,111 @@ def test_get_key_mcp_rpm_limit_precedence(): none_set = UserAPIKeyAuth(api_key=hash_token("sk-mcp-key")) assert get_key_mcp_rpm_limit(none_set) is None assert get_team_mcp_rpm_limit(none_set) is None + + +def test_tpm_reservation_enabled_by_default(monkeypatch): + """Upfront TPM reservation is on unless explicitly disabled via env.""" + monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()) + ) + assert handler.tpm_reservation_enabled is True + + +@pytest.mark.parametrize("value", ["false", "False", "FALSE"]) +def test_tpm_reservation_disabled_via_env(monkeypatch, value): + monkeypatch.setenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", value) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()) + ) + assert handler.tpm_reservation_enabled is False + + +@pytest.mark.asyncio +async def test_pre_call_hook_reserves_tpm_when_enabled(monkeypatch): + """ + With reservation enabled, the pre-call hook reserves the estimated token + budget upfront and tells should_rate_limit to skip the :tokens counter so + only the reservation path owns it. + """ + monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()) + ) + + user_api_key_dict = UserAPIKeyAuth(api_key=hash_token("sk-tpm"), tpm_limit=10_000) + + should_rate_limit_calls: List[Dict[str, Any]] = [] + original_should_rate_limit = handler.should_rate_limit + + async def spy_should_rate_limit(*args, **kwargs): + should_rate_limit_calls.append(kwargs) + return await original_should_rate_limit(*args, **kwargs) + + reserve_calls: List[int] = [] + original_reserve = handler.reserve_tpm_tokens + + async def spy_reserve(*args, **kwargs): + reserve_calls.append(kwargs.get("estimated_tokens")) + return await original_reserve(*args, **kwargs) + + monkeypatch.setattr(handler, "should_rate_limit", spy_should_rate_limit) + monkeypatch.setattr(handler, "reserve_tpm_tokens", spy_reserve) + + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=handler.internal_usage_cache.dual_cache, + data={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]}, + call_type="completion", + ) + + assert len(reserve_calls) == 1, "reservation must run when enabled" + assert should_rate_limit_calls[0]["skip_tpm_check"] is True + + +@pytest.mark.asyncio +async def test_pre_call_hook_skips_reservation_when_disabled(monkeypatch): + """ + With reservation disabled, the pre-call hook never calls reserve_tpm_tokens + and enforces TPM directly in should_rate_limit (skip_tpm_check=False), the + pre-v1.82 post-call accounting behavior. + """ + monkeypatch.setenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", "false") + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()) + ) + + user_api_key_dict = UserAPIKeyAuth(api_key=hash_token("sk-tpm"), tpm_limit=10_000) + + should_rate_limit_calls: List[Dict[str, Any]] = [] + original_should_rate_limit = handler.should_rate_limit + + async def spy_should_rate_limit(*args, **kwargs): + should_rate_limit_calls.append(kwargs) + return await original_should_rate_limit(*args, **kwargs) + + reserve_calls: List[Any] = [] + + async def spy_reserve(*args, **kwargs): + reserve_calls.append(kwargs) + raise AssertionError("reserve_tpm_tokens must not run when disabled") + + monkeypatch.setattr(handler, "should_rate_limit", spy_should_rate_limit) + monkeypatch.setattr(handler, "reserve_tpm_tokens", spy_reserve) + + data = {"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]} + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=handler.internal_usage_cache.dual_cache, + data=data, + call_type="completion", + ) + + assert reserve_calls == [], "reservation must be skipped when disabled" + assert should_rate_limit_calls[0]["skip_tpm_check"] is False + # No reservation stash leaks into the request metadata. + from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + TPM_RESERVED_TOKENS_KEY, + ) + + assert TPM_RESERVED_TOKENS_KEY not in (data.get("metadata") or {}) From a992ed18df4e746c3230c49cfcabcc962988f470 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Thu, 11 Jun 2026 11:02:42 -0700 Subject: [PATCH 066/185] feat(spend_logs): opt-in native Postgres partitioning for SpendLogs retention (#29466) High-volume deployments see LiteLLM_SpendLogs grow unbounded because retention via DELETE leaves dead tuples that autovacuum cannot reclaim fast enough. With a range-partitioned table, retention drops whole partitions instead: an instant metadata operation that returns disk to the OS immediately. The feature is gated behind general_settings.use_spend_logs_partitioning (default false). With the flag off, the cleanup job never queries the catalog and behaves exactly as today. With it on, the job verifies the table is partitioned, pre-creates upcoming partitions, and drops expired ones; expired rows the drops cannot reach (DEFAULT partition, partitions spanning the cutoff) are still deleted row-wise so retention is never bypassed. If the table is not partitioned it falls back to batched DELETE only. Converting an existing table is a manual, documented operation in db_scripts/partition_spend_logs.sql; db_scripts/unpartition_spend_logs.sql rolls it back. Both scripts rename the old table's indexes aside before recreating them, since a table rename keeps the schema-unique index names and would otherwise silently skip the CREATE INDEX IF NOT EXISTS block. Granularity and pre-create lookahead are tunable via SPEND_LOG_PARTITION_INTERVAL (day/week/month, invalid values fall back to day) and SPEND_LOG_PARTITION_PRECREATE_AHEAD. --- db_scripts/partition_spend_logs.sql | 99 ++++++++ db_scripts/unpartition_spend_logs.sql | 69 ++++++ litellm/constants.py | 4 + litellm/proxy/_types.py | 4 + .../db_transaction_queue/spend_log_cleanup.py | 48 +++- .../spend_logs_partition_manager.py | 208 ++++++++++++++++ .../test_spend_logs_partition_manager.py | 233 ++++++++++++++++++ .../proxy/test_spend_log_cleanup.py | 120 ++++++++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 + 9 files changed, 778 insertions(+), 12 deletions(-) create mode 100644 db_scripts/partition_spend_logs.sql create mode 100644 db_scripts/unpartition_spend_logs.sql create mode 100644 litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py create mode 100644 tests/test_litellm/proxy/db/db_transaction_queue/test_spend_logs_partition_manager.py diff --git a/db_scripts/partition_spend_logs.sql b/db_scripts/partition_spend_logs.sql new file mode 100644 index 00000000000..08fcbddb6f8 --- /dev/null +++ b/db_scripts/partition_spend_logs.sql @@ -0,0 +1,99 @@ +-- Converts an existing LiteLLM_SpendLogs table into a native Postgres +-- range-partitioned table keyed on "startTime". +-- +-- Why: at high request volume, retention via DELETE leaves dead tuples that +-- autovacuum cannot reclaim quickly enough, so the table keeps growing on disk +-- (seen at 450GB+ after ~1 month). With partitioning, retention drops whole +-- partitions, which is instant and returns disk to the OS immediately. +-- +-- This is an opt-in, manual operation. The default LiteLLM schema is NOT +-- partitioned, so existing installs are unaffected until you run this. +-- +-- IMPORTANT +-- * Test on a staging copy first and take a backup. +-- * Postgres cannot convert a populated table to partitioned in place, so this +-- renames the old table aside and creates a fresh partitioned table. +-- * The partition key ("startTime") must be part of the primary key, so the +-- PK becomes composite ("request_id", "startTime"). LiteLLM's write path uses +-- INSERT ... ON CONFLICT DO NOTHING, which is compatible with this. +-- * Choose a partition granularity ("day" is the recommended default for +-- high-volume tables) and keep it consistent with SPEND_LOG_PARTITION_INTERVAL. +-- +-- After running this, enable the feature and set a retention period in +-- proxy_config.yaml: +-- general_settings: +-- use_spend_logs_partitioning: true +-- maximum_spend_logs_retention_period: "30d" +-- The spend-log cleanup job then verifies the table is partitioned and reclaims +-- disk by dropping expired partitions instead of deleting rows. It also +-- pre-creates upcoming partitions on each run. To roll back, see +-- db_scripts/unpartition_spend_logs.sql. + +BEGIN; + +ALTER TABLE "LiteLLM_SpendLogs" RENAME TO "LiteLLM_SpendLogs_legacy"; + +-- Renaming a table does NOT rename its indexes, and index names are unique per +-- schema. Move the legacy table's indexes aside so the CREATE INDEX statements +-- below actually create indexes on the new partitioned table instead of being +-- silently skipped by IF NOT EXISTS, and so the new PK keeps the canonical +-- name instead of getting a "_pkey1" suffix. +ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_pkey" + RENAME TO "LiteLLM_SpendLogs_legacy_pkey"; +ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_startTime_idx" + RENAME TO "LiteLLM_SpendLogs_legacy_startTime_idx"; +ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_startTime_request_id_idx" + RENAME TO "LiteLLM_SpendLogs_legacy_startTime_request_id_idx"; +ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_end_user_idx" + RENAME TO "LiteLLM_SpendLogs_legacy_end_user_idx"; +ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_session_id_idx" + RENAME TO "LiteLLM_SpendLogs_legacy_session_id_idx"; + +CREATE TABLE "LiteLLM_SpendLogs" ( + LIKE "LiteLLM_SpendLogs_legacy" INCLUDING DEFAULTS INCLUDING GENERATED +) PARTITION BY RANGE ("startTime"); + +ALTER TABLE "LiteLLM_SpendLogs" + ADD PRIMARY KEY ("request_id", "startTime"); + +-- Recreate every index Prisma defines on the table. LIKE ... INCLUDING DEFAULTS +-- INCLUDING GENERATED copies columns and defaults but NOT indexes, so without +-- these the admin-UI cost-reporting queries that filter by end_user/session_id +-- fall back to sequential scans. On a partitioned parent these propagate to +-- every current and future partition automatically. +CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_startTime_idx" + ON "LiteLLM_SpendLogs" ("startTime"); + +CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_startTime_request_id_idx" + ON "LiteLLM_SpendLogs" ("startTime", "request_id"); + +CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_end_user_idx" + ON "LiteLLM_SpendLogs" ("end_user"); + +CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_session_id_idx" + ON "LiteLLM_SpendLogs" ("session_id"); + +-- Safety net: any row whose startTime has no explicit partition lands here so +-- writes never fail. The cleanup job never drops the DEFAULT partition. +CREATE TABLE IF NOT EXISTS "LiteLLM_SpendLogs_pdefault" + PARTITION OF "LiteLLM_SpendLogs" DEFAULT; + +COMMIT; + +-- Backfill (optional). Rows route to the correct partition automatically. +-- For large legacy tables, copy in time-bounded batches during a low-traffic +-- window instead of one statement, or simply keep "LiteLLM_SpendLogs_legacy" +-- read-only until its data ages past your retention, then DROP it. +-- +-- Backfilled rows land in the DEFAULT partition until explicit partitions +-- cover their dates. Postgres refuses to create a partition whose range +-- overlaps rows already in DEFAULT, so the cleanup job may log a warning when +-- pre-creating today's partition right after a backfill; it recovers on its +-- own once those dates age out, and future partitions are unaffected because +-- they are always created ahead of writes. +-- +-- INSERT INTO "LiteLLM_SpendLogs" +-- SELECT * FROM "LiteLLM_SpendLogs_legacy" +-- WHERE "startTime" >= now() - interval '30 days'; +-- +-- DROP TABLE "LiteLLM_SpendLogs_legacy"; diff --git a/db_scripts/unpartition_spend_logs.sql b/db_scripts/unpartition_spend_logs.sql new file mode 100644 index 00000000000..0bd82513e4a --- /dev/null +++ b/db_scripts/unpartition_spend_logs.sql @@ -0,0 +1,69 @@ +-- Rolls back db_scripts/partition_spend_logs.sql: converts the native +-- range-partitioned "LiteLLM_SpendLogs" table back into a plain, +-- non-partitioned table matching the default LiteLLM schema. +-- +-- When/why: run this if you want to stop using partition-based retention and +-- return to DELETE-based cleanup, or to restore the original single-column +-- primary key ("request_id") that the partitioned layout had to widen to a +-- composite ("request_id", "startTime"). +-- +-- IMPORTANT +-- * Test on a staging copy first and take a backup. +-- * Postgres cannot convert a partitioned table back in place, so this +-- renames the partitioned table aside and creates a fresh plain table. +-- * The composite PK could in principle hold the same "request_id" in more +-- than one partition, so rows are copied with ON CONFLICT DO NOTHING to +-- restore the single-column PK without failing on such duplicates. +-- * For large tables the INSERT ... SELECT copies every surviving row and may +-- run long; do it during a low-traffic window. +-- * Also remove use_spend_logs_partitioning from proxy_config.yaml (or set it +-- to false) so the cleanup job returns to DELETE-based retention. + +BEGIN; + +ALTER TABLE "LiteLLM_SpendLogs" RENAME TO "LiteLLM_SpendLogs_partitioned"; + +-- Renaming a table does NOT rename its indexes, and index names are unique per +-- schema. Move the partitioned table's indexes aside so the CREATE INDEX +-- statements below actually create indexes on the new plain table instead of +-- being silently skipped by IF NOT EXISTS, and so the new PK keeps the +-- canonical name. +ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_pkey" + RENAME TO "LiteLLM_SpendLogs_partitioned_pkey"; +ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_pkey1" + RENAME TO "LiteLLM_SpendLogs_partitioned_pkey1"; +ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_startTime_idx" + RENAME TO "LiteLLM_SpendLogs_partitioned_startTime_idx"; +ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_startTime_request_id_idx" + RENAME TO "LiteLLM_SpendLogs_partitioned_startTime_request_id_idx"; +ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_end_user_idx" + RENAME TO "LiteLLM_SpendLogs_partitioned_end_user_idx"; +ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_session_id_idx" + RENAME TO "LiteLLM_SpendLogs_partitioned_session_id_idx"; + +CREATE TABLE "LiteLLM_SpendLogs" ( + LIKE "LiteLLM_SpendLogs_partitioned" INCLUDING DEFAULTS INCLUDING GENERATED +); + +ALTER TABLE "LiteLLM_SpendLogs" + ADD PRIMARY KEY ("request_id"); + +CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_startTime_idx" + ON "LiteLLM_SpendLogs" ("startTime"); + +CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_startTime_request_id_idx" + ON "LiteLLM_SpendLogs" ("startTime", "request_id"); + +CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_end_user_idx" + ON "LiteLLM_SpendLogs" ("end_user"); + +CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_session_id_idx" + ON "LiteLLM_SpendLogs" ("session_id"); + +INSERT INTO "LiteLLM_SpendLogs" +SELECT * FROM "LiteLLM_SpendLogs_partitioned" +ON CONFLICT ("request_id") DO NOTHING; + +DROP TABLE "LiteLLM_SpendLogs_partitioned"; + +COMMIT; diff --git a/litellm/constants.py b/litellm/constants.py index ab8e57d735f..a5e3926aa8b 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1497,6 +1497,10 @@ SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES = int( SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS = float( os.getenv("SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.5) ) +SPEND_LOG_PARTITION_INTERVAL = os.getenv("SPEND_LOG_PARTITION_INTERVAL", "day") +SPEND_LOG_PARTITION_PRECREATE_AHEAD = int( + os.getenv("SPEND_LOG_PARTITION_PRECREATE_AHEAD", 7) +) SPEND_LOG_QUEUE_SIZE_THRESHOLD = int(os.getenv("SPEND_LOG_QUEUE_SIZE_THRESHOLD", 100)) SPEND_LOG_QUEUE_POLL_INTERVAL = float(os.getenv("SPEND_LOG_QUEUE_POLL_INTERVAL", 2.0)) SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE = int( diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 1b594e20d32..35e9e0cd74b 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2300,6 +2300,10 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="Maximum retention period for spend logs (e.g., '7d' for 7 days). Logs older than this will be deleted.", ) + use_spend_logs_partitioning: Optional[bool] = Field( + None, + description="If True and LiteLLM_SpendLogs has been converted to a range-partitioned table (db_scripts/partition_spend_logs.sql), retention cleanup drops expired partitions instead of deleting rows, and pre-creates upcoming partitions. Default is False.", + ) mcp_internal_ip_ranges: Optional[List[str]] = Field( None, description="Custom CIDR ranges that define internal/private networks for MCP access control. When set, only these ranges are treated as internal. Defaults to RFC 1918 private ranges (10.0.0.0/8, 172.16.0.0/12, 192.168.0.0/16, 127.0.0.0/8).", diff --git a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py index 9475779cfdf..a4c23937b98 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py +++ b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py @@ -12,19 +12,31 @@ from litellm.constants import ( SPEND_LOG_RUN_LOOPS, ) from litellm.litellm_core_utils.duration_parser import duration_in_seconds +from litellm.proxy.db.db_transaction_queue.spend_logs_partition_manager import ( + SpendLogsPartitionManager, +) from litellm.proxy.utils import PrismaClient class SpendLogCleanup: """ Handles cleaning up old spend logs based on maximum retention period. - Deletes logs in batches to prevent timeouts. + + When LiteLLM_SpendLogs is range-partitioned, expired data is reclaimed by + dropping whole partitions (instant, frees disk immediately). Otherwise it + falls back to deleting logs in batches. Uses PodLockManager to ensure only one pod runs cleanup in multi-pod deployments. """ - def __init__(self, general_settings=None, redis_cache: Optional[RedisCache] = None): + def __init__( + self, + general_settings=None, + redis_cache: Optional[RedisCache] = None, + partition_manager: Optional[SpendLogsPartitionManager] = None, + ): self.batch_size = SPEND_LOG_CLEANUP_BATCH_SIZE self.retention_seconds: Optional[int] = None + self.partition_manager = partition_manager or SpendLogsPartitionManager() from litellm.proxy.proxy_server import general_settings as default_settings self.general_settings = general_settings or default_settings @@ -89,8 +101,8 @@ class SpendLogCleanup: deleted_result = await prisma_client.db.execute_raw( """ DELETE FROM "LiteLLM_SpendLogs" - WHERE "request_id" IN ( - SELECT "request_id" FROM "LiteLLM_SpendLogs" + WHERE ("request_id", "startTime") IN ( + SELECT "request_id", "startTime" FROM "LiteLLM_SpendLogs" WHERE "startTime" < $1::timestamptz LIMIT $2 ) @@ -195,12 +207,32 @@ class SpendLogCleanup: seconds=float(self.retention_seconds) ) verbose_proxy_logger.info( - f"Deleting logs older than {cutoff_date.isoformat()}" + f"Removing logs older than {cutoff_date.isoformat()}" ) - # Perform the actual deletion - total_deleted = await self._delete_old_logs(prisma_client, cutoff_date) - verbose_proxy_logger.info(f"Deleted {total_deleted} logs") + if self.general_settings.get( + "use_spend_logs_partitioning", False + ) and await self.partition_manager.is_partitioned(prisma_client): + await self.partition_manager.ensure_partitions(prisma_client) + dropped = await self.partition_manager.drop_partitions_older_than( + prisma_client, cutoff_date + ) + verbose_proxy_logger.info( + "Dropped %d expired spend-log partitions: %s", + len(dropped), + dropped, + ) + # DROP only reclaims whole expired partitions. Expired rows can + # still sit in the DEFAULT partition (backfill, coverage gaps) + # or in a partition that spans the cutoff, so retention must + # also delete those stragglers row-wise. + total_deleted = await self._delete_old_logs(prisma_client, cutoff_date) + verbose_proxy_logger.info( + f"Deleted {total_deleted} expired logs not covered by dropped partitions" + ) + else: + total_deleted = await self._delete_old_logs(prisma_client, cutoff_date) + verbose_proxy_logger.info(f"Deleted {total_deleted} logs") except Exception as e: # .exception() captures the traceback; str(e) alone on a Prisma/DB diff --git a/litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py b/litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py new file mode 100644 index 00000000000..eee0f862b4e --- /dev/null +++ b/litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py @@ -0,0 +1,208 @@ +""" +Manages native Postgres range partitions for the LiteLLM_SpendLogs table. + +At high request volume, retention via batched DELETE leaves dead tuples that +autovacuum cannot reclaim fast enough, so the table keeps growing on disk. When +the table is range-partitioned on startTime, dropping old data becomes a +DROP TABLE on a whole partition: an instant metadata operation that returns disk +to the OS immediately, with no tombstones and no vacuum. + +This manager only acts when use_spend_logs_partitioning is enabled in +general_settings AND the table is already partitioned (set up via the +db_scripts/partition_spend_logs.sql runbook). Without both, the cleanup job +keeps the batched-DELETE path, so existing deployments are untouched. +""" + +import re +from datetime import date, datetime, timedelta, timezone +from typing import List, Optional, Tuple + +from litellm._logging import verbose_proxy_logger +from litellm.constants import ( + SPEND_LOG_PARTITION_INTERVAL, + SPEND_LOG_PARTITION_PRECREATE_AHEAD, +) + +SPEND_LOGS_TABLE = "LiteLLM_SpendLogs" + +PartitionInterval = str # "day" | "week" | "month" + +VALID_PARTITION_INTERVALS = {"day", "week", "month"} + +_BOUND_UPPER_RE = re.compile(r"TO \('([^']+)'\)") + + +def period_start(day: date, interval: PartitionInterval) -> date: + """First day of the partition period that `day` falls into (UTC).""" + if interval == "day": + return day + if interval == "week": + return day - timedelta(days=day.weekday()) + if interval == "month": + return day.replace(day=1) + raise ValueError(f"Unsupported partition interval: {interval}") + + +def next_period_start(start: date, interval: PartitionInterval) -> date: + if interval == "day": + return start + timedelta(days=1) + if interval == "week": + return start + timedelta(days=7) + if interval == "month": + if start.month == 12: + return start.replace(year=start.year + 1, month=1) + return start.replace(month=start.month + 1) + raise ValueError(f"Unsupported partition interval: {interval}") + + +def partition_name(start: date) -> str: + return f"{SPEND_LOGS_TABLE}_p{start.strftime('%Y%m%d')}" + + +def upcoming_partitions( + today: date, interval: PartitionInterval, ahead: int +) -> List[Tuple[str, date, date]]: + """ + Specs (name, lower_inclusive, upper_exclusive) for the current period plus + the next `ahead` periods, so writes always have a partition to land in. + """ + specs: List[Tuple[str, date, date]] = [] + start = period_start(today, interval) + for _ in range(ahead + 1): + upper = next_period_start(start, interval) + specs.append((partition_name(start), start, upper)) + start = upper + return specs + + +def parse_partition_upper_bound(bound_expr: str) -> Optional[datetime]: + """ + Upper bound of a Postgres partition from its `pg_get_expr(relpartbound)` + string, e.g. "FOR VALUES FROM ('2026-06-01 00:00:00') TO ('2026-06-02 00:00:00')". + Returns None for the DEFAULT partition or anything we cannot parse, so such + partitions are never selected for dropping. + """ + if "DEFAULT" in bound_expr.upper(): + return None + match = _BOUND_UPPER_RE.search(bound_expr) + if match is None: + return None + try: + return datetime.fromisoformat(match.group(1)) + except ValueError: + return None + + +def select_partitions_to_drop( + partitions: List[Tuple[str, Optional[datetime]]], cutoff: datetime +) -> List[str]: + """ + Names of partitions whose entire range is older than `cutoff` (upper bound + <= cutoff). `cutoff` and the bounds are UTC-naive. Partitions without a + parseable upper bound (e.g. DEFAULT) are kept. + """ + return [name for name, upper in partitions if upper is not None and upper <= cutoff] + + +class SpendLogsPartitionManager: + def __init__( + self, + interval: PartitionInterval = SPEND_LOG_PARTITION_INTERVAL, + precreate_ahead: int = SPEND_LOG_PARTITION_PRECREATE_AHEAD, + ): + if interval not in VALID_PARTITION_INTERVALS: + verbose_proxy_logger.warning( + "Invalid SPEND_LOG_PARTITION_INTERVAL %r, falling back to 'day'. " + "Supported values: %s", + interval, + sorted(VALID_PARTITION_INTERVALS), + ) + interval = "day" + self.interval = interval + self.precreate_ahead = precreate_ahead + + async def is_partitioned(self, prisma_client) -> bool: + try: + rows = await prisma_client.db.query_raw( + """ + SELECT EXISTS ( + SELECT 1 + FROM pg_partitioned_table pt + JOIN pg_class c ON c.oid = pt.partrelid + JOIN pg_namespace n ON n.oid = c.relnamespace + WHERE c.relname = $1 + AND n.nspname = current_schema() + ) AS partitioned + """, + SPEND_LOGS_TABLE, + ) + except Exception as e: + verbose_proxy_logger.warning( + "Could not determine if %s is partitioned, assuming it is not: %s", + SPEND_LOGS_TABLE, + e, + ) + return False + return bool(rows and rows[0].get("partitioned")) + + async def ensure_partitions(self, prisma_client) -> List[str]: + """ + Ensure the current and upcoming partitions exist, returning the names + now present. CREATE TABLE IF NOT EXISTS is a no-op for partitions that + already exist, so this list is "ensured present", not "newly created". + """ + ensured: List[str] = [] + for name, lower, upper in upcoming_partitions( + datetime.now(timezone.utc).date(), self.interval, self.precreate_ahead + ): + try: + await prisma_client.db.execute_raw( + f'CREATE TABLE IF NOT EXISTS "{name}" ' + f'PARTITION OF "{SPEND_LOGS_TABLE}" ' + f"FOR VALUES FROM ('{lower.isoformat()}') TO ('{upper.isoformat()}')" + ) + ensured.append(name) + except Exception as e: + verbose_proxy_logger.warning( + "Failed to ensure spend-log partition %s: %s", name, e + ) + return ensured + + async def _list_partitions( + self, prisma_client + ) -> List[Tuple[str, Optional[datetime]]]: + rows = await prisma_client.db.query_raw( + """ + SELECT c.relname AS name, + pg_get_expr(c.relpartbound, c.oid) AS bound + FROM pg_inherits i + JOIN pg_class c ON c.oid = i.inhrelid + JOIN pg_class p ON p.oid = i.inhparent + JOIN pg_namespace n ON n.oid = p.relnamespace + WHERE p.relname = $1 + AND n.nspname = current_schema() + """, + SPEND_LOGS_TABLE, + ) + return [ + (row["name"], parse_partition_upper_bound(row.get("bound") or "")) + for row in rows + ] + + async def drop_partitions_older_than( + self, prisma_client, cutoff: datetime + ) -> List[str]: + """DROP every partition whose whole range is older than `cutoff`.""" + cutoff_naive = cutoff.astimezone(timezone.utc).replace(tzinfo=None) + partitions = await self._list_partitions(prisma_client) + to_drop = select_partitions_to_drop(partitions, cutoff_naive) + dropped: List[str] = [] + for name in to_drop: + try: + await prisma_client.db.execute_raw(f'DROP TABLE IF EXISTS "{name}"') + dropped.append(name) + except Exception as e: + verbose_proxy_logger.warning( + "Failed to drop spend-log partition %s: %s", name, e + ) + return dropped diff --git a/tests/test_litellm/proxy/db/db_transaction_queue/test_spend_logs_partition_manager.py b/tests/test_litellm/proxy/db/db_transaction_queue/test_spend_logs_partition_manager.py new file mode 100644 index 00000000000..289de707387 --- /dev/null +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_spend_logs_partition_manager.py @@ -0,0 +1,233 @@ +""" +Tests for SpendLogsPartitionManager: partition naming/bounds math, retention +selection, the non-partitioned no-op safety path, and the drop/ensure SQL flow. +""" + +from datetime import date, datetime, timezone +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm.proxy.db.db_transaction_queue.spend_logs_partition_manager import ( + SpendLogsPartitionManager, + next_period_start, + parse_partition_upper_bound, + partition_name, + period_start, + select_partitions_to_drop, + upcoming_partitions, +) + + +def test_period_start_per_interval(): + d = date(2026, 6, 3) # a Wednesday + assert period_start(d, "day") == date(2026, 6, 3) + assert period_start(d, "week") == date(2026, 6, 1) # Monday + assert period_start(d, "month") == date(2026, 6, 1) + + +def test_next_period_start_crosses_year_and_month_boundaries(): + assert next_period_start(date(2026, 6, 3), "day") == date(2026, 6, 4) + assert next_period_start(date(2026, 6, 1), "week") == date(2026, 6, 8) + assert next_period_start(date(2026, 12, 1), "month") == date(2027, 1, 1) + + +def test_partition_name_uses_period_start_date(): + assert partition_name(date(2026, 6, 1)) == "LiteLLM_SpendLogs_p20260601" + + +def test_upcoming_partitions_count_and_contiguous_ranges(): + specs = upcoming_partitions(date(2026, 6, 1), "day", ahead=3) + assert len(specs) == 4 # current + 3 ahead + names = [s[0] for s in specs] + assert names == [ + "LiteLLM_SpendLogs_p20260601", + "LiteLLM_SpendLogs_p20260602", + "LiteLLM_SpendLogs_p20260603", + "LiteLLM_SpendLogs_p20260604", + ] + # ranges must be contiguous and half-open: each upper is the next lower + for (_, _, upper), (_, next_lower, _) in zip(specs, specs[1:]): + assert upper == next_lower + + +def test_parse_partition_upper_bound_extracts_to_value(): + bound = "FOR VALUES FROM ('2026-06-01 00:00:00') TO ('2026-06-02 00:00:00')" + assert parse_partition_upper_bound(bound) == datetime(2026, 6, 2, 0, 0, 0) + + +def test_parse_partition_upper_bound_default_is_none(): + assert parse_partition_upper_bound("DEFAULT") is None + assert parse_partition_upper_bound("garbage") is None + + +def test_select_partitions_to_drop_only_fully_expired(): + cutoff = datetime(2026, 6, 10, 0, 0, 0) + partitions = [ + ("p_old", datetime(2026, 6, 9, 0, 0, 0)), # upper < cutoff -> drop + ("p_boundary", datetime(2026, 6, 10, 0, 0, 0)), # upper == cutoff -> drop + ("p_partial", datetime(2026, 6, 11, 0, 0, 0)), # straddles cutoff -> keep + ("p_default", None), # DEFAULT -> keep + ] + assert select_partitions_to_drop(partitions, cutoff) == ["p_old", "p_boundary"] + + +@pytest.mark.asyncio +async def test_is_partitioned_true_and_false(): + mgr = SpendLogsPartitionManager() + + client_true = MagicMock() + client_true.db.query_raw = AsyncMock(return_value=[{"partitioned": True}]) + assert await mgr.is_partitioned(client_true) is True + + client_false = MagicMock() + client_false.db.query_raw = AsyncMock(return_value=[{"partitioned": False}]) + assert await mgr.is_partitioned(client_false) is False + + +@pytest.mark.asyncio +async def test_catalog_queries_are_scoped_to_current_schema(): + """ + Both catalog lookups must filter by current_schema(); otherwise a same-named + table in another schema can flip is_partitioned or return foreign partitions. + """ + mgr = SpendLogsPartitionManager() + client = MagicMock() + client.db.query_raw = AsyncMock(return_value=[]) + + await mgr.is_partitioned(client) + is_partitioned_sql = client.db.query_raw.call_args.args[0] + assert "pg_namespace" in is_partitioned_sql + assert "current_schema()" in is_partitioned_sql + + await mgr._list_partitions(client) + list_sql = client.db.query_raw.call_args.args[0] + assert "pg_namespace" in list_sql + assert "current_schema()" in list_sql + + +@pytest.mark.asyncio +async def test_is_partitioned_swallows_errors_and_returns_false(): + """A catalog query failure must not crash cleanup; fall back to non-partitioned.""" + mgr = SpendLogsPartitionManager() + client = MagicMock() + client.db.query_raw = AsyncMock(side_effect=Exception("db down")) + assert await mgr.is_partitioned(client) is False + + +@pytest.mark.asyncio +async def test_drop_partitions_older_than_drops_expired_only(): + mgr = SpendLogsPartitionManager() + client = MagicMock() + client.db.query_raw = AsyncMock( + return_value=[ + { + "name": "LiteLLM_SpendLogs_p20260601", + "bound": "FOR VALUES FROM ('2026-06-01 00:00:00') TO ('2026-06-02 00:00:00')", + }, + { + "name": "LiteLLM_SpendLogs_p20260609", + "bound": "FOR VALUES FROM ('2026-06-09 00:00:00') TO ('2026-06-10 00:00:00')", + }, + {"name": "LiteLLM_SpendLogs_pdefault", "bound": "DEFAULT"}, + ] + ) + client.db.execute_raw = AsyncMock(return_value=0) + + cutoff = datetime(2026, 6, 5, 0, 0, 0, tzinfo=timezone.utc) + dropped = await mgr.drop_partitions_older_than(client, cutoff) + + assert dropped == ["LiteLLM_SpendLogs_p20260601"] + executed = " ".join(call.args[0] for call in client.db.execute_raw.call_args_list) + assert 'DROP TABLE IF EXISTS "LiteLLM_SpendLogs_p20260601"' in executed + assert "p20260609" not in executed + assert "pdefault" not in executed + + +@pytest.mark.asyncio +async def test_ensure_partitions_issues_create_for_each_period(): + mgr = SpendLogsPartitionManager(interval="day", precreate_ahead=2) + client = MagicMock() + client.db.execute_raw = AsyncMock(return_value=0) + + created = await mgr.ensure_partitions(client) + + assert len(created) == 3 # current + 2 ahead + assert client.db.execute_raw.await_count == 3 + first_sql = client.db.execute_raw.call_args_list[0].args[0] + assert 'PARTITION OF "LiteLLM_SpendLogs"' in first_sql + assert "CREATE TABLE IF NOT EXISTS" in first_sql + + +def test_unsupported_interval_raises(): + with pytest.raises(ValueError): + period_start(date(2026, 6, 1), "year") + with pytest.raises(ValueError): + next_period_start(date(2026, 6, 1), "year") + + +def test_parse_partition_upper_bound_unparseable_to_value_is_none(): + """A TO(...) value that is not a valid timestamp must not raise; return None.""" + assert ( + parse_partition_upper_bound("FOR VALUES FROM ('x') TO ('not-a-date')") is None + ) + + +@pytest.mark.asyncio +async def test_ensure_partitions_continues_when_one_create_fails(): + mgr = SpendLogsPartitionManager(interval="day", precreate_ahead=2) + client = MagicMock() + client.db.execute_raw = AsyncMock(side_effect=[0, Exception("overlap"), 0]) + + created = await mgr.ensure_partitions(client) + + # the failed partition is skipped, the others still created + assert len(created) == 2 + assert client.db.execute_raw.await_count == 3 + + +def test_invalid_interval_falls_back_to_day(): + """ + An invalid interval must not be stored as-is. Otherwise ensure_partitions + raises (via period_start) and aborts the cleanup run before retention drops + old partitions, silently skipping retention. + """ + mgr = SpendLogsPartitionManager(interval="year") + assert mgr.interval == "day" + + +@pytest.mark.asyncio +async def test_invalid_interval_does_not_abort_ensure_partitions(): + """With the fallback, ensure_partitions completes instead of raising ValueError.""" + mgr = SpendLogsPartitionManager(interval="fortnight", precreate_ahead=1) + client = MagicMock() + client.db.execute_raw = AsyncMock(return_value=0) + + created = await mgr.ensure_partitions(client) + + assert len(created) == 2 # current + 1 ahead, day-based fallback + + +@pytest.mark.asyncio +async def test_drop_partitions_continues_when_one_drop_fails(): + mgr = SpendLogsPartitionManager() + client = MagicMock() + client.db.query_raw = AsyncMock( + return_value=[ + { + "name": "LiteLLM_SpendLogs_p20260601", + "bound": "FOR VALUES FROM ('2026-06-01 00:00:00') TO ('2026-06-02 00:00:00')", + }, + { + "name": "LiteLLM_SpendLogs_p20260602", + "bound": "FOR VALUES FROM ('2026-06-02 00:00:00') TO ('2026-06-03 00:00:00')", + }, + ] + ) + client.db.execute_raw = AsyncMock(side_effect=[Exception("locked"), 0]) + + cutoff = datetime(2026, 6, 10, 0, 0, 0, tzinfo=timezone.utc) + dropped = await mgr.drop_partitions_older_than(client, cutoff) + + # both were eligible; the first drop failed so only the second is reported + assert dropped == ["LiteLLM_SpendLogs_p20260602"] diff --git a/tests/test_litellm/proxy/test_spend_log_cleanup.py b/tests/test_litellm/proxy/test_spend_log_cleanup.py index 42bb919295f..a309dd64011 100644 --- a/tests/test_litellm/proxy/test_spend_log_cleanup.py +++ b/tests/test_litellm/proxy/test_spend_log_cleanup.py @@ -183,7 +183,10 @@ async def test_cleanup_old_spend_logs_batch_deletion(): # Check the first call argument call_args_sql = mock_db.execute_raw.call_args_list[0][0][0] assert 'DELETE FROM "LiteLLM_SpendLogs"' in call_args_sql - assert 'WHERE "request_id" IN' in call_args_sql + # must match on the full composite identity: on a partitioned table + # request_id alone is not unique, and deleting by it would let a client + # reusing x-litellm-call-id take out a fresh row alongside the expired one + assert 'WHERE ("request_id", "startTime") IN' in call_args_sql @pytest.mark.asyncio @@ -219,6 +222,109 @@ async def test_cleanup_old_spend_logs_retention_period_cutoff(): ) # Allow 1 second difference for test execution time +@pytest.mark.asyncio +async def test_cleanup_drops_partitions_when_enabled_and_partitioned(): + """ + With use_spend_logs_partitioning enabled and a partitioned table, cleanup + must reclaim disk by dropping partitions AND still delete expired rows the + drops cannot reach (DEFAULT partition, cutoff-spanning partitions), so + retention is never bypassed. + """ + from unittest.mock import AsyncMock, MagicMock + + mock_prisma_client = MagicMock() + mock_prisma_client.db.execute_raw = AsyncMock(return_value=0) + + partition_manager = MagicMock() + partition_manager.is_partitioned = AsyncMock(return_value=True) + partition_manager.ensure_partitions = AsyncMock(return_value=["p1"]) + partition_manager.drop_partitions_older_than = AsyncMock( + return_value=["LiteLLM_SpendLogs_p20260601"] + ) + + cleaner = SpendLogCleanup( + general_settings={ + "maximum_spend_logs_retention_period": "7d", + "use_spend_logs_partitioning": True, + }, + partition_manager=partition_manager, + ) + cleaner.pod_lock_manager = MagicMock() + cleaner.pod_lock_manager.redis_cache = None + + await cleaner.cleanup_old_spend_logs(mock_prisma_client) + + partition_manager.ensure_partitions.assert_awaited_once() + partition_manager.drop_partitions_older_than.assert_awaited_once() + delete_sql = mock_prisma_client.db.execute_raw.call_args_list[0][0][0] + assert 'DELETE FROM "LiteLLM_SpendLogs"' in delete_sql + + +@pytest.mark.asyncio +async def test_cleanup_uses_delete_when_partitioning_not_enabled(): + """ + Even against a partitioned table, the partition path must stay off until + use_spend_logs_partitioning is explicitly enabled, so existing deployments + see zero behavior change. The catalog must not even be queried. + """ + from unittest.mock import AsyncMock, MagicMock + + mock_prisma_client = MagicMock() + mock_prisma_client.db.execute_raw = AsyncMock(side_effect=[10, 0]) + + partition_manager = MagicMock() + partition_manager.is_partitioned = AsyncMock(return_value=True) + partition_manager.ensure_partitions = AsyncMock() + partition_manager.drop_partitions_older_than = AsyncMock() + + cleaner = SpendLogCleanup( + general_settings={"maximum_spend_logs_retention_period": "7d"}, + partition_manager=partition_manager, + ) + cleaner.pod_lock_manager = MagicMock() + cleaner.pod_lock_manager.redis_cache = None + + await cleaner.cleanup_old_spend_logs(mock_prisma_client) + + partition_manager.is_partitioned.assert_not_awaited() + partition_manager.drop_partitions_older_than.assert_not_awaited() + delete_sql = mock_prisma_client.db.execute_raw.call_args_list[0][0][0] + assert 'DELETE FROM "LiteLLM_SpendLogs"' in delete_sql + + +@pytest.mark.asyncio +async def test_cleanup_uses_delete_when_not_partitioned(): + """ + With the feature enabled but the table not actually partitioned (script not + run yet), cleanup must keep using the batched DELETE path. + """ + from unittest.mock import AsyncMock, MagicMock + + mock_prisma_client = MagicMock() + mock_prisma_client.db.execute_raw = AsyncMock(side_effect=[10, 0]) + + partition_manager = MagicMock() + partition_manager.is_partitioned = AsyncMock(return_value=False) + partition_manager.drop_partitions_older_than = AsyncMock() + + cleaner = SpendLogCleanup( + general_settings={ + "maximum_spend_logs_retention_period": "7d", + "use_spend_logs_partitioning": True, + }, + partition_manager=partition_manager, + ) + cleaner.pod_lock_manager = MagicMock() + cleaner.pod_lock_manager.redis_cache = None + + await cleaner.cleanup_old_spend_logs(mock_prisma_client) + + partition_manager.drop_partitions_older_than.assert_not_awaited() + assert mock_prisma_client.db.execute_raw.await_count == 2 + delete_sql = mock_prisma_client.db.execute_raw.call_args_list[0][0][0] + assert 'DELETE FROM "LiteLLM_SpendLogs"' in delete_sql + + @pytest.mark.asyncio async def test_cleanup_old_spend_logs_no_retention_period(): """ @@ -370,7 +476,9 @@ async def test_delete_old_logs_aborts_after_consecutive_failures(monkeypatch): import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module # Lower the threshold so the test is fast and deterministic. - monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 3) + monkeypatch.setattr( + cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 3 + ) monkeypatch.setattr( cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0 ) @@ -400,7 +508,9 @@ async def test_delete_old_logs_resets_consecutive_failures_on_success(monkeypatc intermittent timeouts don't trip the abort threshold.""" import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module - monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 3) + monkeypatch.setattr( + cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 3 + ) monkeypatch.setattr( cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0 ) @@ -471,7 +581,9 @@ async def test_cleanup_releases_lock_after_persistent_batch_failures(monkeypatch must still be released so the next scheduled run isn't permanently blocked.""" import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module - monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 2) + monkeypatch.setattr( + cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 2 + ) monkeypatch.setattr( cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0 ) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 15123bcdbf8..4aeb1cf3e86 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -22175,6 +22175,11 @@ export interface components { * @description decrypt keys with google kms */ use_google_kms?: boolean | null; + /** + * Use Spend Logs Partitioning + * @description If True and LiteLLM_SpendLogs has been converted to a range-partitioned table (db_scripts/partition_spend_logs.sql), retention cleanup drops expired partitions instead of deleting rows, and pre-creates upcoming partitions. Default is False. + */ + use_spend_logs_partitioning?: boolean | null; /** User Header Mappings */ user_header_mappings?: components["schemas"]["UserHeaderMapping"][] | null; /** From 530c0b2326b80fb0fa2cdf583eab03d2f2979f9f Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Thu, 11 Jun 2026 12:07:17 -0700 Subject: [PATCH 067/185] feat(ui): migrate playground to path routing and colocate its files (#30185) * feat(ui): cut playground over to the /ui/playground path route Follows the api-reference recipe: the sidebar and deep links route llm-playground to the path route, ?page=llm-playground redirects, and the legacy switch arm is deleted. The route's page.tsx was already the real implementation, so no view extraction was needed. * refactor(ui): move playground-owned files into its route folder Per the (dashboard) README convention, page-owned code lives in the page's folder: chat_ui/compareUI/complianceUI components, the chat hooks, and the playground-only llm_calls helpers move under (dashboard)/playground/. Modules with non-playground consumers (chat message primitives; fetch_models, chat_completion, responses_api) stay at their lowest common ancestor in src/components/{chat_ui,llm_calls} because legacy pages still import them. eslint-suppressions entries are re-keyed to the new paths so the grandfathered baseline still applies. * test(ui): teach sidebar e2e spec about migrated path routes The sidebar spec asserted ?page= for every item, which the playground cutover correctly broke: the sidebar now links to /ui/playground and the legacy URL redirects there. Drive the expected URL from the migration fixture (now a page-id -> segment map) so future cutovers only add a fixture entry. Also wrap one import line in AgentBuilderView.tsx that the move left unformatted; the changed- files prettier check flagged it. --- .../e2e_tests/fixtures/migratedPages.ts | 21 ++++--- .../tests/navigation/sidebar.spec.ts | 22 +++++-- ui/litellm-dashboard/eslint-suppressions.json | 44 +++++++------- .../src/app/(dashboard)/page.tsx | 3 - .../components}/chat_ui/A2AMetrics.tsx | 0 .../chat_ui/AdditionalModelSettings.test.tsx | 0 .../chat_ui/AdditionalModelSettings.tsx | 0 .../components}/chat_ui/AgentBuilderView.tsx | 18 ++++-- .../chat_ui/AudioRenderer.test.tsx | 2 +- .../components}/chat_ui/AudioRenderer.tsx | 2 +- .../components}/chat_ui/ChatImageRenderer.tsx | 2 +- .../components}/chat_ui/ChatImageUpload.tsx | 0 .../chat_ui/ChatImageUtils.test.tsx | 2 +- .../components}/chat_ui/ChatImageUtils.tsx | 2 +- .../chat_ui/ChatMessageBubble.test.tsx | 10 ++-- .../components}/chat_ui/ChatMessageBubble.tsx | 14 ++--- .../components}/chat_ui/ChatUI.test.tsx | 6 +- .../playground/components}/chat_ui/ChatUI.tsx | 60 +++++++++---------- .../chat_ui/CodeInterpreterOutput.test.tsx | 0 .../chat_ui/CodeInterpreterOutput.tsx | 0 .../chat_ui/CodeInterpreterTool.tsx | 0 .../chat_ui/EndpointSelector.test.tsx | 0 .../components}/chat_ui/EndpointSelector.tsx | 0 .../chat_ui/EndpointUtils.test.tsx | 8 +-- .../components}/chat_ui/EndpointUtils.tsx | 4 +- .../chat_ui/FilePreviewCard.test.tsx | 0 .../components}/chat_ui/FilePreviewCard.tsx | 0 .../chat_ui/RealtimePlayground.tsx | 2 +- .../chat_ui/ResponsesImageRenderer.tsx | 2 +- .../chat_ui/ResponsesImageUpload.tsx | 0 .../chat_ui/ResponsesImageUtils.tsx | 2 +- .../chat_ui/SearchResultsDisplay.tsx | 2 +- .../components}/chat_ui/SessionManagement.tsx | 4 +- .../components}/chat_ui/chatConstants.ts | 2 +- .../components}/compareUI/CompareUI.test.tsx | 6 +- .../components}/compareUI/CompareUI.tsx | 12 ++-- .../components/ComparisonPanel.test.tsx | 6 +- .../compareUI/components/ComparisonPanel.tsx | 6 +- .../components/MessageDisplay.test.tsx | 6 +- .../compareUI/components/MessageDisplay.tsx | 6 +- .../components/MessageInput.test.tsx | 0 .../compareUI/components/MessageInput.tsx | 0 .../components/ModelSelector.test.tsx | 0 .../compareUI/components/ModelSelector.tsx | 0 .../components/UnifiedSelector.test.tsx | 0 .../compareUI/components/UnifiedSelector.tsx | 0 .../compareUI/endpoint_config.test.ts | 2 +- .../components}/compareUI/endpoint_config.ts | 2 +- .../components}/complianceUI/ComplianceUI.tsx | 2 +- .../playground/hooks}/useChatHistory.test.ts | 0 .../playground/hooks}/useChatHistory.ts | 8 +-- .../playground/hooks}/useCodeInterpreter.ts | 4 +- .../playground/llm_calls/a2a_send_message.tsx | 4 +- .../llm_calls/anthropic_messages.tsx | 4 +- .../llm_calls/audio_speech.test.tsx | 0 .../playground/llm_calls/audio_speech.tsx | 2 +- .../llm_calls/audio_transcriptions.test.tsx | 0 .../llm_calls/audio_transcriptions.tsx | 0 .../llm_calls/embeddings_api.test.tsx | 0 .../playground/llm_calls/embeddings_api.tsx | 0 .../playground/llm_calls/fetch_agents.tsx | 2 +- .../playground/llm_calls/image_edits.tsx | 0 .../playground/llm_calls/image_generation.tsx | 0 .../playground/llm_calls/interactions_api.tsx | 0 .../src/app/(dashboard)/playground/page.tsx | 8 +-- .../cost_tracking_settings.test.tsx | 2 +- .../cost_tracking_settings.tsx | 2 +- .../EvaluationSettingsModal.tsx | 2 +- .../MCPSemanticFilterSettings.test.tsx | 2 +- .../MCPSemanticFilterSettings.tsx | 2 +- .../Fallbacks/AddFallbacks.test.tsx | 4 +- .../RouterSettings/Fallbacks/AddFallbacks.tsx | 2 +- .../Fallbacks/Fallbacks.test.tsx | 4 +- .../add_model/ComplexityRouterConfig.tsx | 2 +- .../add_model/RouterConfigBuilder.tsx | 2 +- .../add_model/add_auto_router_tab.tsx | 2 +- .../cache_settings/CacheFieldRenderer.tsx | 2 +- .../src/components/chat/ChatMessages.tsx | 4 +- .../src/components/chat/ChatPage.tsx | 6 +- .../chat_ui/CodeSnippets.test.tsx | 0 .../{playground => }/chat_ui/CodeSnippets.tsx | 2 +- .../chat_ui/MCPEventsDisplay.tsx | 2 +- .../chat_ui/ReasoningContent.tsx | 0 .../chat_ui/ResponseMetrics.tsx | 0 .../chat_ui/mode_endpoint_mapping.tsx | 0 .../{playground => }/chat_ui/types.ts | 0 .../common_components/ModelSelector.tsx | 2 +- .../RouterSettingsAccordion.tsx | 2 +- .../edit_auto_router_modal.tsx | 2 +- .../llm_calls/chat_completion.test.tsx | 0 .../llm_calls/chat_completion.tsx | 2 +- .../llm_calls/code_interpreter_handler.ts | 0 .../llm_calls/fetch_models.tsx | 2 +- .../llm_calls/responses_api.test.tsx | 0 .../llm_calls/responses_api.tsx | 4 +- .../conversation_panel/MessageBubble.tsx | 2 +- .../conversation_panel/types.ts | 2 +- .../conversation_panel/useConversation.ts | 2 +- .../src/components/public_model_hub.tsx | 6 +- .../S3VectorsConfig.test.tsx | 4 +- .../S3VectorsConfig.tsx | 2 +- .../VectorStoreForm.tsx | 2 +- .../src/utils/migratedPages.test.ts | 7 +++ .../src/utils/migratedPages.ts | 1 + 104 files changed, 214 insertions(+), 184 deletions(-) rename ui/litellm-dashboard/src/{components/playground => app/(dashboard)/playground/components}/chat_ui/A2AMetrics.tsx (100%) rename ui/litellm-dashboard/src/{components/playground => app/(dashboard)/playground/components}/chat_ui/AdditionalModelSettings.test.tsx (100%) rename ui/litellm-dashboard/src/{components/playground => app/(dashboard)/playground/components}/chat_ui/AdditionalModelSettings.tsx (100%) rename ui/litellm-dashboard/src/{components/playground => app/(dashboard)/playground/components}/chat_ui/AgentBuilderView.tsx (98%) rename ui/litellm-dashboard/src/{components/playground => app/(dashboard)/playground/components}/chat_ui/AudioRenderer.test.tsx (88%) rename ui/litellm-dashboard/src/{components/playground => app/(dashboard)/playground/components}/chat_ui/AudioRenderer.tsx (90%) rename ui/litellm-dashboard/src/{components/playground => app/(dashboard)/playground/components}/chat_ui/ChatImageRenderer.tsx (95%) rename ui/litellm-dashboard/src/{components/playground => app/(dashboard)/playground/components}/chat_ui/ChatImageUpload.tsx (100%) rename ui/litellm-dashboard/src/{components/playground => app/(dashboard)/playground/components}/chat_ui/ChatImageUtils.test.tsx (99%) rename ui/litellm-dashboard/src/{components/playground => app/(dashboard)/playground/components}/chat_ui/ChatImageUtils.tsx (96%) rename ui/litellm-dashboard/src/{components/playground => app/(dashboard)/playground/components}/chat_ui/ChatMessageBubble.test.tsx (96%) rename ui/litellm-dashboard/src/{components/playground => app/(dashboard)/playground/components}/chat_ui/ChatMessageBubble.tsx (93%) rename ui/litellm-dashboard/src/{components/playground => app/(dashboard)/playground/components}/chat_ui/ChatUI.test.tsx (98%) rename ui/litellm-dashboard/src/{components/playground => app/(dashboard)/playground/components}/chat_ui/ChatUI.tsx (97%) rename ui/litellm-dashboard/src/{components/playground => app/(dashboard)/playground/components}/chat_ui/CodeInterpreterOutput.test.tsx (100%) rename ui/litellm-dashboard/src/{components/playground => app/(dashboard)/playground/components}/chat_ui/CodeInterpreterOutput.tsx (100%) rename ui/litellm-dashboard/src/{components/playground => app/(dashboard)/playground/components}/chat_ui/CodeInterpreterTool.tsx (100%) rename ui/litellm-dashboard/src/{components/playground => app/(dashboard)/playground/components}/chat_ui/EndpointSelector.test.tsx (100%) rename ui/litellm-dashboard/src/{components/playground => app/(dashboard)/playground/components}/chat_ui/EndpointSelector.tsx (100%) rename ui/litellm-dashboard/src/{components/playground => app/(dashboard)/playground/components}/chat_ui/EndpointUtils.test.tsx (95%) rename ui/litellm-dashboard/src/{components/playground => app/(dashboard)/playground/components}/chat_ui/EndpointUtils.tsx (82%) rename ui/litellm-dashboard/src/{components/playground => app/(dashboard)/playground/components}/chat_ui/FilePreviewCard.test.tsx (100%) rename ui/litellm-dashboard/src/{components/playground => app/(dashboard)/playground/components}/chat_ui/FilePreviewCard.tsx (100%) rename ui/litellm-dashboard/src/{components/playground => app/(dashboard)/playground/components}/chat_ui/RealtimePlayground.tsx (99%) rename ui/litellm-dashboard/src/{components/playground => app/(dashboard)/playground/components}/chat_ui/ResponsesImageRenderer.tsx (94%) rename ui/litellm-dashboard/src/{components/playground => app/(dashboard)/playground/components}/chat_ui/ResponsesImageUpload.tsx (100%) rename ui/litellm-dashboard/src/{components/playground => app/(dashboard)/playground/components}/chat_ui/ResponsesImageUtils.tsx (96%) rename ui/litellm-dashboard/src/{components/playground => app/(dashboard)/playground/components}/chat_ui/SearchResultsDisplay.tsx (98%) rename ui/litellm-dashboard/src/{components/playground => app/(dashboard)/playground/components}/chat_ui/SessionManagement.tsx (96%) rename ui/litellm-dashboard/src/{components/playground => app/(dashboard)/playground/components}/chat_ui/chatConstants.ts (95%) rename ui/litellm-dashboard/src/{components/playground => app/(dashboard)/playground/components}/compareUI/CompareUI.test.tsx (96%) rename ui/litellm-dashboard/src/{components/playground => app/(dashboard)/playground/components}/compareUI/CompareUI.tsx (98%) rename ui/litellm-dashboard/src/{components/playground => app/(dashboard)/playground/components}/compareUI/components/ComparisonPanel.test.tsx (93%) rename ui/litellm-dashboard/src/{components/playground => app/(dashboard)/playground/components}/compareUI/components/ComparisonPanel.tsx (97%) rename ui/litellm-dashboard/src/{components/playground => app/(dashboard)/playground/components}/compareUI/components/MessageDisplay.test.tsx (94%) rename ui/litellm-dashboard/src/{components/playground => app/(dashboard)/playground/components}/compareUI/components/MessageDisplay.tsx (96%) rename ui/litellm-dashboard/src/{components/playground => app/(dashboard)/playground/components}/compareUI/components/MessageInput.test.tsx (100%) rename ui/litellm-dashboard/src/{components/playground => app/(dashboard)/playground/components}/compareUI/components/MessageInput.tsx (100%) rename ui/litellm-dashboard/src/{components/playground => app/(dashboard)/playground/components}/compareUI/components/ModelSelector.test.tsx (100%) rename ui/litellm-dashboard/src/{components/playground => app/(dashboard)/playground/components}/compareUI/components/ModelSelector.tsx (100%) rename ui/litellm-dashboard/src/{components/playground => app/(dashboard)/playground/components}/compareUI/components/UnifiedSelector.test.tsx (100%) rename ui/litellm-dashboard/src/{components/playground => app/(dashboard)/playground/components}/compareUI/components/UnifiedSelector.tsx (100%) rename ui/litellm-dashboard/src/{components/playground => app/(dashboard)/playground/components}/compareUI/endpoint_config.test.ts (98%) rename ui/litellm-dashboard/src/{components/playground => app/(dashboard)/playground/components}/compareUI/endpoint_config.ts (98%) rename ui/litellm-dashboard/src/{components/playground => app/(dashboard)/playground/components}/complianceUI/ComplianceUI.tsx (99%) rename ui/litellm-dashboard/src/{components/playground/chat_ui => app/(dashboard)/playground/hooks}/useChatHistory.test.ts (100%) rename ui/litellm-dashboard/src/{components/playground/chat_ui => app/(dashboard)/playground/hooks}/useChatHistory.ts (98%) rename ui/litellm-dashboard/src/{components/playground/chat_ui => app/(dashboard)/playground/hooks}/useCodeInterpreter.ts (88%) rename ui/litellm-dashboard/src/{components => app/(dashboard)}/playground/llm_calls/a2a_send_message.tsx (98%) rename ui/litellm-dashboard/src/{components => app/(dashboard)}/playground/llm_calls/anthropic_messages.tsx (96%) rename ui/litellm-dashboard/src/{components => app/(dashboard)}/playground/llm_calls/audio_speech.test.tsx (100%) rename ui/litellm-dashboard/src/{components => app/(dashboard)}/playground/llm_calls/audio_speech.tsx (96%) rename ui/litellm-dashboard/src/{components => app/(dashboard)}/playground/llm_calls/audio_transcriptions.test.tsx (100%) rename ui/litellm-dashboard/src/{components => app/(dashboard)}/playground/llm_calls/audio_transcriptions.tsx (100%) rename ui/litellm-dashboard/src/{components => app/(dashboard)}/playground/llm_calls/embeddings_api.test.tsx (100%) rename ui/litellm-dashboard/src/{components => app/(dashboard)}/playground/llm_calls/embeddings_api.tsx (100%) rename ui/litellm-dashboard/src/{components => app/(dashboard)}/playground/llm_calls/fetch_agents.tsx (98%) rename ui/litellm-dashboard/src/{components => app/(dashboard)}/playground/llm_calls/image_edits.tsx (100%) rename ui/litellm-dashboard/src/{components => app/(dashboard)}/playground/llm_calls/image_generation.tsx (100%) rename ui/litellm-dashboard/src/{components => app/(dashboard)}/playground/llm_calls/interactions_api.tsx (100%) rename ui/litellm-dashboard/src/components/{playground => }/chat_ui/CodeSnippets.test.tsx (100%) rename ui/litellm-dashboard/src/components/{playground => }/chat_ui/CodeSnippets.tsx (99%) rename ui/litellm-dashboard/src/components/{playground => }/chat_ui/MCPEventsDisplay.tsx (99%) rename ui/litellm-dashboard/src/components/{playground => }/chat_ui/ReasoningContent.tsx (100%) rename ui/litellm-dashboard/src/components/{playground => }/chat_ui/ResponseMetrics.tsx (100%) rename ui/litellm-dashboard/src/components/{playground => }/chat_ui/mode_endpoint_mapping.tsx (100%) rename ui/litellm-dashboard/src/components/{playground => }/chat_ui/types.ts (100%) rename ui/litellm-dashboard/src/components/{playground => }/llm_calls/chat_completion.test.tsx (100%) rename ui/litellm-dashboard/src/components/{playground => }/llm_calls/chat_completion.tsx (99%) rename ui/litellm-dashboard/src/components/{playground => }/llm_calls/code_interpreter_handler.ts (100%) rename ui/litellm-dashboard/src/components/{playground => }/llm_calls/fetch_models.tsx (94%) rename ui/litellm-dashboard/src/components/{playground => }/llm_calls/responses_api.test.tsx (100%) rename ui/litellm-dashboard/src/components/{playground => }/llm_calls/responses_api.tsx (98%) diff --git a/ui/litellm-dashboard/e2e_tests/fixtures/migratedPages.ts b/ui/litellm-dashboard/e2e_tests/fixtures/migratedPages.ts index f2ba66147ea..749c3cde179 100644 --- a/ui/litellm-dashboard/e2e_tests/fixtures/migratedPages.ts +++ b/ui/litellm-dashboard/e2e_tests/fixtures/migratedPages.ts @@ -1,16 +1,23 @@ /** - * Source of truth for the App Router migration smoke (tests/migration/migratedPages.spec.ts). + * Source of truth for the App Router migration E2E suites. * - * Add a route segment here once its migration has MERGED to the branch under test. - * Both suites pick it up automatically: - * - default mount: npm run e2e:migration - * - server-root-path mount: SERVER_ROOT_PATH=/ npm run e2e:migration:root + * Add an entry (legacy sidebar page id -> route segment) once a page's migration + * has MERGED to the branch under test. Consumers pick it up automatically: + * - migration smoke (tests/migration/migratedPages.spec.ts), via MIGRATED_E2E_SEGMENTS: + * default mount: npm run e2e:migration + * server-root-path mount: SERVER_ROOT_PATH=/ npm run e2e:migration:root + * - navigation specs that assert per-page URLs (tests/navigation/sidebar.spec.ts) * * Keep this in lockstep with MIGRATED_PAGES in src/utils/migratedPages.ts. - * Pending (uncomment as each PR lands): playground, and the leaf-pages batch + * Pending (add as each PR lands): the leaf-pages batch * (budgets, caching, cost-tracking, guardrails, guardrails-monitor, logs, * mcp-servers, memory, policies, projects, prompts, search-tools, skills, * tag-management, tool-policies, transform-request, ui-theme, vector-stores, * workflows, access-groups). */ -export const MIGRATED_E2E_SEGMENTS: string[] = ["api-reference"]; +export const MIGRATED_E2E_PAGES: Record = { + api_ref: "api-reference", + "llm-playground": "playground", +}; + +export const MIGRATED_E2E_SEGMENTS: string[] = [...new Set(Object.values(MIGRATED_E2E_PAGES))]; diff --git a/ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts index b8fb95b764d..7ac2e7df39d 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts +++ b/ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts @@ -4,11 +4,23 @@ import { ADMIN_STORAGE_PATH } from "../../constants"; import { Page } from "../../fixtures/pages"; import { menuLabelToPage } from "../../fixtures/menuMappings"; import { navigateToPage } from "../../helpers/navigation"; +import { MIGRATED_E2E_PAGES } from "../../fixtures/migratedPages"; +import type { Page as PlaywrightPage } from "@playwright/test"; const sidebarButtons = { [Role.ProxyAdmin]: ["Virtual Keys", "Playground", "Models", "Usage", "Teams", "Internal Users", "AI Hub"], }; +/** Migrated pages live at a path route; legacy pages keep the ?page= query param. */ +async function expectPageUrl(page: PlaywrightPage, pageKey: string): Promise { + const migratedSegment = MIGRATED_E2E_PAGES[pageKey]; + if (migratedSegment) { + await expect(page).toHaveURL(new RegExp(`/ui/${migratedSegment}/?($|\\?)`)); + } else { + await expect(page).toHaveURL(new RegExp(`[?&]page=${pageKey}(&|$)`)); + } +} + const roles = [{ role: Role.ProxyAdmin, storage: ADMIN_STORAGE_PATH }]; for (const { role, storage } of roles) { @@ -35,8 +47,7 @@ for (const { role, storage } of roles) { await tab.click(); - // Verify URL contains the correct page query parameter - await expect(page).toHaveURL(new RegExp(`[?&]page=${expectedPage}(&|$)`)); + await expectPageUrl(page, expectedPage); } }); @@ -50,13 +61,14 @@ for (const { role, storage } of roles) { // Test direct navigation to verify the helper function works await navigateToPage(page, Page.ApiKeys); - await expect(page).toHaveURL(new RegExp(`[?&]page=${Page.ApiKeys}(&|$)`)); + await expectPageUrl(page, Page.ApiKeys); await navigateToPage(page, Page.Models); - await expect(page).toHaveURL(new RegExp(`[?&]page=${Page.Models}(&|$)`)); + await expectPageUrl(page, Page.Models); + // Migrated page: /ui?page=llm-playground redirects to the path route await navigateToPage(page, Page.LlmPlayground); - await expect(page).toHaveURL(new RegExp(`[?&]page=${Page.LlmPlayground}(&|$)`)); + await expectPageUrl(page, Page.LlmPlayground); }); }); } diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index d3169395b4e..358064c00af 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -1533,7 +1533,7 @@ "count": 1 } }, - "src/components/playground/chat_ui/AdditionalModelSettings.tsx": { + "src/app/(dashboard)/playground/components/chat_ui/AdditionalModelSettings.tsx": { "no-restricted-imports": { "count": 1 }, @@ -1541,17 +1541,17 @@ "count": 2 } }, - "src/components/playground/chat_ui/AgentBuilderView.tsx": { + "src/app/(dashboard)/playground/components/chat_ui/AgentBuilderView.tsx": { "react-hooks/set-state-in-effect": { "count": 5 } }, - "src/components/playground/chat_ui/ChatImageUtils.test.tsx": { + "src/app/(dashboard)/playground/components/chat_ui/ChatImageUtils.test.tsx": { "max-nested-callbacks": { "count": 1 } }, - "src/components/playground/chat_ui/ChatUI.tsx": { + "src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx": { "no-restricted-imports": { "count": 1 }, @@ -1562,17 +1562,17 @@ "count": 13 } }, - "src/components/playground/chat_ui/CodeInterpreterOutput.tsx": { + "src/app/(dashboard)/playground/components/chat_ui/CodeInterpreterOutput.tsx": { "no-restricted-syntax": { "count": 2 } }, - "src/components/playground/chat_ui/CodeInterpreterTool.tsx": { + "src/app/(dashboard)/playground/components/chat_ui/CodeInterpreterTool.tsx": { "no-restricted-imports": { "count": 1 } }, - "src/components/playground/chat_ui/RealtimePlayground.tsx": { + "src/app/(dashboard)/playground/components/chat_ui/RealtimePlayground.tsx": { "react-hooks/immutability": { "count": 2 }, @@ -1580,22 +1580,22 @@ "count": 1 } }, - "src/components/playground/compareUI/CompareUI.tsx": { + "src/app/(dashboard)/playground/components/compareUI/CompareUI.tsx": { "react-hooks/set-state-in-effect": { "count": 1 } }, - "src/components/playground/compareUI/components/ModelSelector.tsx": { + "src/app/(dashboard)/playground/components/compareUI/components/ModelSelector.tsx": { "no-restricted-imports": { "count": 1 } }, - "src/components/playground/complianceUI/ComplianceUI.tsx": { + "src/app/(dashboard)/playground/components/complianceUI/ComplianceUI.tsx": { "react-hooks/preserve-manual-memoization": { "count": 3 } }, - "src/components/playground/llm_calls/a2a_send_message.tsx": { + "src/app/(dashboard)/playground/llm_calls/a2a_send_message.tsx": { "max-params": { "count": 2 }, @@ -1603,27 +1603,27 @@ "count": 2 } }, - "src/components/playground/llm_calls/anthropic_messages.tsx": { + "src/app/(dashboard)/playground/llm_calls/anthropic_messages.tsx": { "max-params": { "count": 1 } }, - "src/components/playground/llm_calls/audio_speech.tsx": { + "src/app/(dashboard)/playground/llm_calls/audio_speech.tsx": { "max-params": { "count": 1 } }, - "src/components/playground/llm_calls/audio_transcriptions.tsx": { + "src/app/(dashboard)/playground/llm_calls/audio_transcriptions.tsx": { "max-params": { "count": 1 } }, - "src/components/playground/llm_calls/chat_completion.tsx": { + "src/components/llm_calls/chat_completion.tsx": { "max-params": { "count": 1 } }, - "src/components/playground/llm_calls/embeddings_api.tsx": { + "src/app/(dashboard)/playground/llm_calls/embeddings_api.tsx": { "max-params": { "count": 1 }, @@ -1631,22 +1631,22 @@ "count": 1 } }, - "src/components/playground/llm_calls/fetch_agents.tsx": { + "src/app/(dashboard)/playground/llm_calls/fetch_agents.tsx": { "no-restricted-syntax": { "count": 1 } }, - "src/components/playground/llm_calls/image_edits.tsx": { + "src/app/(dashboard)/playground/llm_calls/image_edits.tsx": { "max-params": { "count": 1 } }, - "src/components/playground/llm_calls/image_generation.tsx": { + "src/app/(dashboard)/playground/llm_calls/image_generation.tsx": { "max-params": { "count": 1 } }, - "src/components/playground/llm_calls/interactions_api.tsx": { + "src/app/(dashboard)/playground/llm_calls/interactions_api.tsx": { "max-params": { "count": 1 }, @@ -1654,7 +1654,7 @@ "count": 1 } }, - "src/components/playground/llm_calls/responses_api.tsx": { + "src/components/llm_calls/responses_api.tsx": { "max-params": { "count": 1 } @@ -2250,4 +2250,4 @@ "count": 1 } } -} \ No newline at end of file +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/page.tsx index 0854b085fae..f26e5dd5ef2 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/page.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/page.tsx @@ -1,7 +1,6 @@ "use client"; import ModelsAndEndpointsView from "@/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView"; -import PlaygroundPage from "@/app/(dashboard)/playground/page"; import AdminPanel from "@/components/AdminPanel"; import AgentsPanel from "@/components/agents"; import BudgetPanel from "@/components/budgets/budget_panel"; @@ -354,8 +353,6 @@ function CreateKeyPageContent() { premiumUser={premiumUser} teams={teams} /> - ) : page == "llm-playground" ? ( - ) : page == "users" ? ( { it("should render the audio renderer", () => { diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/AudioRenderer.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/AudioRenderer.tsx similarity index 90% rename from ui/litellm-dashboard/src/components/playground/chat_ui/AudioRenderer.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/AudioRenderer.tsx index 48a766283e6..4b5a1596c1d 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/AudioRenderer.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/AudioRenderer.tsx @@ -1,5 +1,5 @@ import React from "react"; -import { MessageType } from "./types"; +import { MessageType } from "@/components/chat_ui/types"; interface AudioRendererProps { message: MessageType; diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatImageRenderer.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatImageRenderer.tsx similarity index 95% rename from ui/litellm-dashboard/src/components/playground/chat_ui/ChatImageRenderer.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatImageRenderer.tsx index 49adaaed8c1..f4449bee803 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatImageRenderer.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatImageRenderer.tsx @@ -1,6 +1,6 @@ import React from "react"; import Image from "next/image"; -import { MessageType } from "./types"; +import { MessageType } from "@/components/chat_ui/types"; import { shouldShowChatAttachedImage } from "./ChatImageUtils"; import { FilePdfOutlined } from "@ant-design/icons"; diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatImageUpload.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatImageUpload.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/playground/chat_ui/ChatImageUpload.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatImageUpload.tsx diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatImageUtils.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatImageUtils.test.tsx similarity index 99% rename from ui/litellm-dashboard/src/components/playground/chat_ui/ChatImageUtils.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatImageUtils.test.tsx index ecc4914b0ab..ab6c542d709 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatImageUtils.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatImageUtils.test.tsx @@ -5,7 +5,7 @@ import { createChatDisplayMessage, shouldShowChatAttachedImage, } from "./ChatImageUtils"; -import { MessageType } from "./types"; +import { MessageType } from "@/components/chat_ui/types"; describe("ChatImageUtils", () => { beforeEach(() => { diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatImageUtils.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatImageUtils.tsx similarity index 96% rename from ui/litellm-dashboard/src/components/playground/chat_ui/ChatImageUtils.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatImageUtils.tsx index 368ba3825ed..7a6340d9fde 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatImageUtils.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatImageUtils.tsx @@ -1,4 +1,4 @@ -import { MessageType } from "./types"; +import { MessageType } from "@/components/chat_ui/types"; export interface ChatMultimodalContent { type: "text" | "image_url"; diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatMessageBubble.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.test.tsx similarity index 96% rename from ui/litellm-dashboard/src/components/playground/chat_ui/ChatMessageBubble.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.test.tsx index 7c4c56d5ade..647258b3d48 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatMessageBubble.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.test.tsx @@ -1,8 +1,8 @@ import { render, screen } from "@testing-library/react"; import { describe, it, expect, vi } from "vitest"; import ChatMessageBubble from "./ChatMessageBubble"; -import { EndpointType } from "./mode_endpoint_mapping"; -import { MessageType } from "./types"; +import { EndpointType } from "@/components/chat_ui/mode_endpoint_mapping"; +import { MessageType } from "@/components/chat_ui/types"; // Mock child components to isolate bubble rendering logic vi.mock("react-markdown", () => ({ @@ -17,13 +17,13 @@ vi.mock("react-syntax-highlighter/dist/esm/styles/prism", () => ({ coy: {}, })); -vi.mock("./ReasoningContent", () => ({ +vi.mock("@/components/chat_ui/ReasoningContent", () => ({ default: ({ reasoningContent }: { reasoningContent: string }) => (
{reasoningContent}
), })); -vi.mock("./MCPEventsDisplay", () => ({ +vi.mock("@/components/chat_ui/MCPEventsDisplay", () => ({ default: ({ events }: { events: unknown[] }) =>
{events.length} events
, })); @@ -33,7 +33,7 @@ vi.mock("./SearchResultsDisplay", () => ({ ), })); -vi.mock("./ResponseMetrics", () => ({ +vi.mock("@/components/chat_ui/ResponseMetrics", () => ({ default: ({ timeToFirstToken }: { timeToFirstToken?: number }) => (
TTFT: {timeToFirstToken}
), diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatMessageBubble.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.tsx similarity index 93% rename from ui/litellm-dashboard/src/components/playground/chat_ui/ChatMessageBubble.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.tsx index b5bff3b70fa..27226757c91 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatMessageBubble.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.tsx @@ -3,19 +3,19 @@ import React from "react"; import ReactMarkdown from "react-markdown"; import { Prism as SyntaxHighlighter } from "react-syntax-highlighter"; import { coy } from "react-syntax-highlighter/dist/esm/styles/prism"; -import { CodeInterpreterResult } from "../llm_calls/code_interpreter_handler"; +import { CodeInterpreterResult } from "@/components/llm_calls/code_interpreter_handler"; import A2AMetrics from "./A2AMetrics"; import AudioRenderer from "./AudioRenderer"; import ChatImageRenderer from "./ChatImageRenderer"; import CodeInterpreterOutput from "./CodeInterpreterOutput"; -import { EndpointType } from "./mode_endpoint_mapping"; -import MCPEventsDisplay from "./MCPEventsDisplay"; -import type { MCPEvent } from "../../mcp_tools/types"; -import ReasoningContent from "./ReasoningContent"; -import ResponseMetrics from "./ResponseMetrics"; +import { EndpointType } from "@/components/chat_ui/mode_endpoint_mapping"; +import MCPEventsDisplay from "@/components/chat_ui/MCPEventsDisplay"; +import type { MCPEvent } from "@/components/mcp_tools/types"; +import ReasoningContent from "@/components/chat_ui/ReasoningContent"; +import ResponseMetrics from "@/components/chat_ui/ResponseMetrics"; import ResponsesImageRenderer from "./ResponsesImageRenderer"; import { SearchResultsDisplay } from "./SearchResultsDisplay"; -import { MessageType } from "./types"; +import { MessageType } from "@/components/chat_ui/types"; interface ChatMessageBubbleProps { message: MessageType; diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx similarity index 98% rename from ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx index 62bda4933ab..9da3e3a4a08 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx @@ -1,15 +1,15 @@ import { act, fireEvent, render, screen, waitFor } from "@testing-library/react"; import { beforeEach, describe, expect, it, vi } from "vitest"; import ChatUI from "./ChatUI"; -import * as fetchModelsModule from "../llm_calls/fetch_models"; +import * as fetchModelsModule from "@/components/llm_calls/fetch_models"; // Mock the fetchAvailableModels function -vi.mock("../llm_calls/fetch_models", () => ({ +vi.mock("@/components/llm_calls/fetch_models", () => ({ fetchAvailableModels: vi.fn(), })); // Mock other networking functions that cause errors -vi.mock("../networking", () => ({ +vi.mock("@/components/networking", () => ({ tagListCall: vi.fn().mockResolvedValue({ data: [] }), vectorStoreListCall: vi.fn().mockResolvedValue({ data: [] }), getGuardrailsList: vi.fn().mockResolvedValue({ data: [] }), diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx similarity index 97% rename from ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx index 5ba2df607cc..db46eb30cb8 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx @@ -28,28 +28,28 @@ import ReactMarkdown from "react-markdown"; import { Prism as SyntaxHighlighter } from "react-syntax-highlighter"; import { coy } from "react-syntax-highlighter/dist/esm/styles/prism"; import { v4 as uuidv4 } from "uuid"; -import GuardrailSelector from "../../guardrails/GuardrailSelector"; -import PolicySelector from "../../policies/PolicySelector"; -import MCPToolArgumentsForm, { MCPToolArgumentsFormRef } from "../../mcp_tools/MCPToolArgumentsForm"; -import { MCPServer } from "../../mcp_tools/types"; -import { ByokCredentialModal } from "../../mcp_tools/ByokCredentialModal"; -import NotificationsManager from "../../molecules/notifications_manager"; -import { callMCPTool, fetchMCPServers, fetchMCPToolsets, listMCPTools } from "../../networking"; -import { MCPToolset } from "../../mcp_tools/types"; -import TagSelector from "../../tag_management/TagSelector"; -import VectorStoreSelector from "../../vector_store_management/VectorStoreSelector"; -import { makeA2ASendMessageRequest } from "../llm_calls/a2a_send_message"; -import { makeAnthropicMessagesRequest } from "../llm_calls/anthropic_messages"; -import { makeOpenAIAudioSpeechRequest } from "../llm_calls/audio_speech"; -import { makeOpenAIAudioTranscriptionRequest } from "../llm_calls/audio_transcriptions"; -import { makeOpenAIChatCompletionRequest } from "../llm_calls/chat_completion"; -import { makeOpenAIEmbeddingsRequest } from "../llm_calls/embeddings_api"; -import { Agent, fetchAvailableAgents } from "../llm_calls/fetch_agents"; -import { fetchAvailableModels, ModelGroup } from "../llm_calls/fetch_models"; -import { makeOpenAIImageEditsRequest } from "../llm_calls/image_edits"; -import { makeOpenAIImageGenerationRequest } from "../llm_calls/image_generation"; -import { makeOpenAIResponsesRequest } from "../llm_calls/responses_api"; -import { makeInteractionsRequest } from "../llm_calls/interactions_api"; +import GuardrailSelector from "@/components/guardrails/GuardrailSelector"; +import PolicySelector from "@/components/policies/PolicySelector"; +import MCPToolArgumentsForm, { MCPToolArgumentsFormRef } from "@/components/mcp_tools/MCPToolArgumentsForm"; +import { MCPServer } from "@/components/mcp_tools/types"; +import { ByokCredentialModal } from "@/components/mcp_tools/ByokCredentialModal"; +import NotificationsManager from "@/components/molecules/notifications_manager"; +import { callMCPTool, fetchMCPServers, fetchMCPToolsets, listMCPTools } from "@/components/networking"; +import { MCPToolset } from "@/components/mcp_tools/types"; +import TagSelector from "@/components/tag_management/TagSelector"; +import VectorStoreSelector from "@/components/vector_store_management/VectorStoreSelector"; +import { makeA2ASendMessageRequest } from "../../llm_calls/a2a_send_message"; +import { makeAnthropicMessagesRequest } from "../../llm_calls/anthropic_messages"; +import { makeOpenAIAudioSpeechRequest } from "../../llm_calls/audio_speech"; +import { makeOpenAIAudioTranscriptionRequest } from "../../llm_calls/audio_transcriptions"; +import { makeOpenAIChatCompletionRequest } from "@/components/llm_calls/chat_completion"; +import { makeOpenAIEmbeddingsRequest } from "../../llm_calls/embeddings_api"; +import { Agent, fetchAvailableAgents } from "../../llm_calls/fetch_agents"; +import { fetchAvailableModels, ModelGroup } from "@/components/llm_calls/fetch_models"; +import { makeOpenAIImageEditsRequest } from "../../llm_calls/image_edits"; +import { makeOpenAIImageGenerationRequest } from "../../llm_calls/image_generation"; +import { makeOpenAIResponsesRequest } from "@/components/llm_calls/responses_api"; +import { makeInteractionsRequest } from "../../llm_calls/interactions_api"; import A2AMetrics from "./A2AMetrics"; import AdditionalModelSettings from "./AdditionalModelSettings"; import AudioRenderer from "./AudioRenderer"; @@ -59,23 +59,23 @@ import ChatImageUpload from "./ChatImageUpload"; import { createChatDisplayMessage, createChatMultimodalMessage } from "./ChatImageUtils"; import CodeInterpreterOutput from "./CodeInterpreterOutput"; import CodeInterpreterTool from "./CodeInterpreterTool"; -import { generateCodeSnippet } from "./CodeSnippets"; +import { generateCodeSnippet } from "@/components/chat_ui/CodeSnippets"; import EndpointSelector from "./EndpointSelector"; import FilePreviewCard from "./FilePreviewCard"; import ChatMessageBubble from "./ChatMessageBubble"; -import MCPEventsDisplay from "./MCPEventsDisplay"; -import { EndpointType, getEndpointType } from "./mode_endpoint_mapping"; -import ReasoningContent from "./ReasoningContent"; -import ResponseMetrics, { TokenUsage } from "./ResponseMetrics"; +import MCPEventsDisplay from "@/components/chat_ui/MCPEventsDisplay"; +import { EndpointType, getEndpointType } from "@/components/chat_ui/mode_endpoint_mapping"; +import ReasoningContent from "@/components/chat_ui/ReasoningContent"; +import ResponseMetrics, { TokenUsage } from "@/components/chat_ui/ResponseMetrics"; import ResponsesImageRenderer from "./ResponsesImageRenderer"; import ResponsesImageUpload from "./ResponsesImageUpload"; import { createDisplayMessage, createMultimodalMessage } from "./ResponsesImageUtils"; import { SearchResultsDisplay } from "./SearchResultsDisplay"; import SessionManagement from "./SessionManagement"; import RealtimePlayground from "./RealtimePlayground"; -import { A2ATaskMetadata, MessageType } from "./types"; -import { useCodeInterpreter } from "./useCodeInterpreter"; -import { useChatHistory } from "./useChatHistory"; +import { A2ATaskMetadata, MessageType } from "@/components/chat_ui/types"; +import { useCodeInterpreter } from "../../hooks/useCodeInterpreter"; +import { useChatHistory } from "../../hooks/useChatHistory"; import { getSecureItem, setSecureItem } from "@/utils/secureStorage"; const { TextArea } = Input; diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/CodeInterpreterOutput.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/CodeInterpreterOutput.test.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/playground/chat_ui/CodeInterpreterOutput.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/CodeInterpreterOutput.test.tsx diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/CodeInterpreterOutput.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/CodeInterpreterOutput.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/playground/chat_ui/CodeInterpreterOutput.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/CodeInterpreterOutput.tsx diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/CodeInterpreterTool.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/CodeInterpreterTool.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/playground/chat_ui/CodeInterpreterTool.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/CodeInterpreterTool.tsx diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/EndpointSelector.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/EndpointSelector.test.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/playground/chat_ui/EndpointSelector.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/EndpointSelector.test.tsx diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/EndpointSelector.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/EndpointSelector.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/playground/chat_ui/EndpointSelector.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/EndpointSelector.tsx diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/EndpointUtils.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/EndpointUtils.test.tsx similarity index 95% rename from ui/litellm-dashboard/src/components/playground/chat_ui/EndpointUtils.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/EndpointUtils.test.tsx index 6eb481381ac..2bb1395428b 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/EndpointUtils.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/EndpointUtils.test.tsx @@ -1,10 +1,10 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; -import type { ModelGroup } from "../llm_calls/fetch_models"; +import type { ModelGroup } from "@/components/llm_calls/fetch_models"; import { determineEndpointType } from "./EndpointUtils"; -import { EndpointType } from "./mode_endpoint_mapping"; +import { EndpointType } from "@/components/chat_ui/mode_endpoint_mapping"; // Mock the getEndpointType function -vi.mock("./mode_endpoint_mapping", () => ({ +vi.mock("@/components/chat_ui/mode_endpoint_mapping", () => ({ EndpointType: { IMAGE: "image", VIDEO: "video", @@ -32,7 +32,7 @@ vi.mock("./mode_endpoint_mapping", () => ({ })); // Import the mocked function -import { getEndpointType } from "./mode_endpoint_mapping"; +import { getEndpointType } from "@/components/chat_ui/mode_endpoint_mapping"; describe("determineEndpointType", () => { beforeEach(() => { diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/EndpointUtils.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/EndpointUtils.tsx similarity index 82% rename from ui/litellm-dashboard/src/components/playground/chat_ui/EndpointUtils.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/EndpointUtils.tsx index de337e5638f..84579610a41 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/EndpointUtils.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/EndpointUtils.tsx @@ -1,5 +1,5 @@ -import { ModelGroup } from "../llm_calls/fetch_models"; -import { EndpointType, getEndpointType } from "./mode_endpoint_mapping"; +import { ModelGroup } from "@/components/llm_calls/fetch_models"; +import { EndpointType, getEndpointType } from "@/components/chat_ui/mode_endpoint_mapping"; /** * Determines the appropriate endpoint type based on the selected model diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/FilePreviewCard.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/FilePreviewCard.test.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/playground/chat_ui/FilePreviewCard.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/FilePreviewCard.test.tsx diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/FilePreviewCard.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/FilePreviewCard.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/playground/chat_ui/FilePreviewCard.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/FilePreviewCard.tsx diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/RealtimePlayground.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/RealtimePlayground.tsx similarity index 99% rename from ui/litellm-dashboard/src/components/playground/chat_ui/RealtimePlayground.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/RealtimePlayground.tsx index 88c6efd87bd..68a150be8c0 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/RealtimePlayground.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/RealtimePlayground.tsx @@ -3,7 +3,7 @@ import { AudioMutedOutlined, AudioOutlined, CloseCircleOutlined, SendOutlined, SoundOutlined } from "@ant-design/icons"; import { Button, Input, Select, Typography } from "antd"; import React, { useCallback, useEffect, useRef, useState } from "react"; -import { getProxyBaseUrl } from "../../networking"; +import { getProxyBaseUrl } from "@/components/networking"; import { OPEN_AI_VOICE_SELECT_OPTIONS } from "./chatConstants"; const { Text } = Typography; diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/ResponsesImageRenderer.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ResponsesImageRenderer.tsx similarity index 94% rename from ui/litellm-dashboard/src/components/playground/chat_ui/ResponsesImageRenderer.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ResponsesImageRenderer.tsx index a8707ebd21b..d459d638c0a 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/ResponsesImageRenderer.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ResponsesImageRenderer.tsx @@ -1,5 +1,5 @@ import React from "react"; -import { MessageType } from "./types"; +import { MessageType } from "@/components/chat_ui/types"; import { shouldShowAttachedImage } from "./ResponsesImageUtils"; import { FilePdfOutlined } from "@ant-design/icons"; diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/ResponsesImageUpload.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ResponsesImageUpload.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/playground/chat_ui/ResponsesImageUpload.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ResponsesImageUpload.tsx diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/ResponsesImageUtils.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ResponsesImageUtils.tsx similarity index 96% rename from ui/litellm-dashboard/src/components/playground/chat_ui/ResponsesImageUtils.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ResponsesImageUtils.tsx index 50dee5c86f6..04dfa39d5c6 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/ResponsesImageUtils.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ResponsesImageUtils.tsx @@ -1,4 +1,4 @@ -import { MessageType, MultimodalContent } from "./types"; +import { MessageType, MultimodalContent } from "@/components/chat_ui/types"; export const convertImageToBase64 = (file: File): Promise => { return new Promise((resolve, reject) => { diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/SearchResultsDisplay.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/SearchResultsDisplay.tsx similarity index 98% rename from ui/litellm-dashboard/src/components/playground/chat_ui/SearchResultsDisplay.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/SearchResultsDisplay.tsx index 966dbb67fc6..ba84fe95f80 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/SearchResultsDisplay.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/SearchResultsDisplay.tsx @@ -1,6 +1,6 @@ import React, { useState } from "react"; import { Button } from "antd"; -import { VectorStoreSearchResponse } from "./types"; +import { VectorStoreSearchResponse } from "@/components/chat_ui/types"; import { DatabaseOutlined, FileTextOutlined, DownOutlined, RightOutlined } from "@ant-design/icons"; interface SearchResultsDisplayProps { diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/SessionManagement.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/SessionManagement.tsx similarity index 96% rename from ui/litellm-dashboard/src/components/playground/chat_ui/SessionManagement.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/SessionManagement.tsx index e782845bfc3..87f9d797643 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/SessionManagement.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/SessionManagement.tsx @@ -1,8 +1,8 @@ import React from "react"; import { Switch, Tooltip } from "antd"; import { InfoCircleOutlined, CopyOutlined } from "@ant-design/icons"; -import { EndpointType } from "./mode_endpoint_mapping"; -import NotificationsManager from "../../molecules/notifications_manager"; +import { EndpointType } from "@/components/chat_ui/mode_endpoint_mapping"; +import NotificationsManager from "@/components/molecules/notifications_manager"; interface SessionManagementProps { endpointType: string; diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/chatConstants.ts b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/chatConstants.ts similarity index 95% rename from ui/litellm-dashboard/src/components/playground/chat_ui/chatConstants.ts rename to ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/chatConstants.ts index 9592a521250..2b59fbad2ee 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/chatConstants.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/chatConstants.ts @@ -1,4 +1,4 @@ -import { EndpointType } from "./mode_endpoint_mapping"; +import { EndpointType } from "@/components/chat_ui/mode_endpoint_mapping"; export const OPEN_AI_VOICES = { ALLOY: "alloy", diff --git a/ui/litellm-dashboard/src/components/playground/compareUI/CompareUI.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/CompareUI.test.tsx similarity index 96% rename from ui/litellm-dashboard/src/components/playground/compareUI/CompareUI.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/CompareUI.test.tsx index 8e6da976cbd..4278cb6a0e4 100644 --- a/ui/litellm-dashboard/src/components/playground/compareUI/CompareUI.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/CompareUI.test.tsx @@ -2,13 +2,13 @@ import { render, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; import CompareUI from "./CompareUI"; -import { makeOpenAIChatCompletionRequest } from "../llm_calls/chat_completion"; +import { makeOpenAIChatCompletionRequest } from "@/components/llm_calls/chat_completion"; -vi.mock("../llm_calls/fetch_models", () => ({ +vi.mock("@/components/llm_calls/fetch_models", () => ({ fetchAvailableModels: vi.fn().mockResolvedValue([{ model_group: "gpt-4" }, { model_group: "gpt-3.5-turbo" }]), })); -vi.mock("../llm_calls/chat_completion", () => ({ +vi.mock("@/components/llm_calls/chat_completion", () => ({ makeOpenAIChatCompletionRequest: vi.fn().mockResolvedValue(undefined), })); diff --git a/ui/litellm-dashboard/src/components/playground/compareUI/CompareUI.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/CompareUI.tsx similarity index 98% rename from ui/litellm-dashboard/src/components/playground/compareUI/CompareUI.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/CompareUI.tsx index ce738e4cd62..e41ee3ba8fa 100644 --- a/ui/litellm-dashboard/src/components/playground/compareUI/CompareUI.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/CompareUI.tsx @@ -7,12 +7,12 @@ import { useEffect, useMemo, useState } from "react"; import { v4 as uuidv4 } from "uuid"; import ChatImageUpload from "../chat_ui/ChatImageUpload"; import { createChatDisplayMessage, createChatMultimodalMessage } from "../chat_ui/ChatImageUtils"; -import type { TokenUsage } from "../chat_ui/ResponseMetrics"; -import type { MessageType, VectorStoreSearchResponse } from "../chat_ui/types"; -import { makeOpenAIChatCompletionRequest } from "../llm_calls/chat_completion"; -import { fetchAvailableModels } from "../llm_calls/fetch_models"; -import { Agent, fetchAvailableAgents } from "../llm_calls/fetch_agents"; -import { makeA2AStreamMessageRequest } from "../llm_calls/a2a_send_message"; +import type { TokenUsage } from "@/components/chat_ui/ResponseMetrics"; +import type { MessageType, VectorStoreSearchResponse } from "@/components/chat_ui/types"; +import { makeOpenAIChatCompletionRequest } from "@/components/llm_calls/chat_completion"; +import { fetchAvailableModels } from "@/components/llm_calls/fetch_models"; +import { Agent, fetchAvailableAgents } from "../../llm_calls/fetch_agents"; +import { makeA2AStreamMessageRequest } from "../../llm_calls/a2a_send_message"; import { ComparisonPanel } from "./components/ComparisonPanel"; import { MessageInput } from "./components/MessageInput"; import { diff --git a/ui/litellm-dashboard/src/components/playground/compareUI/components/ComparisonPanel.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/ComparisonPanel.test.tsx similarity index 93% rename from ui/litellm-dashboard/src/components/playground/compareUI/components/ComparisonPanel.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/ComparisonPanel.test.tsx index 2aef47f71d9..c07ad367606 100644 --- a/ui/litellm-dashboard/src/components/playground/compareUI/components/ComparisonPanel.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/ComparisonPanel.test.tsx @@ -18,15 +18,15 @@ vi.mock("./UnifiedSelector", () => ({ ), })); -vi.mock("../../../tag_management/TagSelector", () => ({ +vi.mock("@/components/tag_management/TagSelector", () => ({ default: () =>
TagSelector
, })); -vi.mock("../../../vector_store_management/VectorStoreSelector", () => ({ +vi.mock("@/components/vector_store_management/VectorStoreSelector", () => ({ default: () =>
VectorStoreSelector
, })); -vi.mock("../../../guardrails/GuardrailSelector", () => ({ +vi.mock("@/components/guardrails/GuardrailSelector", () => ({ default: () =>
GuardrailSelector
, })); diff --git a/ui/litellm-dashboard/src/components/playground/compareUI/components/ComparisonPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/ComparisonPanel.tsx similarity index 97% rename from ui/litellm-dashboard/src/components/playground/compareUI/components/ComparisonPanel.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/ComparisonPanel.tsx index 1074172974f..6ddb5e27947 100644 --- a/ui/litellm-dashboard/src/components/playground/compareUI/components/ComparisonPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/ComparisonPanel.tsx @@ -3,9 +3,9 @@ import { useState } from "react"; import { ComparisonInstance } from "../CompareUI"; import { MessageDisplay } from "./MessageDisplay"; import { UnifiedSelector } from "./UnifiedSelector"; -import TagSelector from "../../../tag_management/TagSelector"; -import VectorStoreSelector from "../../../vector_store_management/VectorStoreSelector"; -import GuardrailSelector from "../../../guardrails/GuardrailSelector"; +import TagSelector from "@/components/tag_management/TagSelector"; +import VectorStoreSelector from "@/components/vector_store_management/VectorStoreSelector"; +import GuardrailSelector from "@/components/guardrails/GuardrailSelector"; import { Checkbox, Divider, Popover, Slider } from "antd"; import { SelectorOption, EndpointConfig, isAgentEndpoint, getComparisonSelection } from "../endpoint_config"; diff --git a/ui/litellm-dashboard/src/components/playground/compareUI/components/MessageDisplay.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/MessageDisplay.test.tsx similarity index 94% rename from ui/litellm-dashboard/src/components/playground/compareUI/components/MessageDisplay.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/MessageDisplay.test.tsx index c635e9555f6..72a1d41f9fe 100644 --- a/ui/litellm-dashboard/src/components/playground/compareUI/components/MessageDisplay.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/MessageDisplay.test.tsx @@ -1,15 +1,15 @@ import { render } from "@testing-library/react"; import { describe, expect, it, vi } from "vitest"; -import type { MessageType } from "../../chat_ui/types"; +import type { MessageType } from "@/components/chat_ui/types"; import { MessageDisplay } from "./MessageDisplay"; -vi.mock("../../chat_ui/ReasoningContent", () => ({ +vi.mock("@/components/chat_ui/ReasoningContent", () => ({ default: ({ reasoningContent }: { reasoningContent: string }) => (
{reasoningContent}
), })); -vi.mock("../../chat_ui/ResponseMetrics", () => ({ +vi.mock("@/components/chat_ui/ResponseMetrics", () => ({ default: () =>
ResponseMetrics
, })); diff --git a/ui/litellm-dashboard/src/components/playground/compareUI/components/MessageDisplay.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/MessageDisplay.tsx similarity index 96% rename from ui/litellm-dashboard/src/components/playground/compareUI/components/MessageDisplay.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/MessageDisplay.tsx index 088e45b91b3..2e4868d3527 100644 --- a/ui/litellm-dashboard/src/components/playground/compareUI/components/MessageDisplay.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/MessageDisplay.tsx @@ -4,10 +4,10 @@ import ReactMarkdown from "react-markdown"; import { Prism as SyntaxHighlighter } from "react-syntax-highlighter"; import { coy } from "react-syntax-highlighter/dist/esm/styles/prism"; import ChatImageRenderer from "../../chat_ui/ChatImageRenderer"; -import ReasoningContent from "../../chat_ui/ReasoningContent"; -import ResponseMetrics from "../../chat_ui/ResponseMetrics"; +import ReasoningContent from "@/components/chat_ui/ReasoningContent"; +import ResponseMetrics from "@/components/chat_ui/ResponseMetrics"; import { SearchResultsDisplay } from "../../chat_ui/SearchResultsDisplay"; -import type { MessageType } from "../../chat_ui/types"; +import type { MessageType } from "@/components/chat_ui/types"; interface MessageDisplayProps { messages: MessageType[]; diff --git a/ui/litellm-dashboard/src/components/playground/compareUI/components/MessageInput.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/MessageInput.test.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/playground/compareUI/components/MessageInput.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/MessageInput.test.tsx diff --git a/ui/litellm-dashboard/src/components/playground/compareUI/components/MessageInput.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/MessageInput.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/playground/compareUI/components/MessageInput.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/MessageInput.tsx diff --git a/ui/litellm-dashboard/src/components/playground/compareUI/components/ModelSelector.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/ModelSelector.test.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/playground/compareUI/components/ModelSelector.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/ModelSelector.test.tsx diff --git a/ui/litellm-dashboard/src/components/playground/compareUI/components/ModelSelector.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/ModelSelector.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/playground/compareUI/components/ModelSelector.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/ModelSelector.tsx diff --git a/ui/litellm-dashboard/src/components/playground/compareUI/components/UnifiedSelector.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/UnifiedSelector.test.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/playground/compareUI/components/UnifiedSelector.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/UnifiedSelector.test.tsx diff --git a/ui/litellm-dashboard/src/components/playground/compareUI/components/UnifiedSelector.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/UnifiedSelector.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/playground/compareUI/components/UnifiedSelector.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/UnifiedSelector.tsx diff --git a/ui/litellm-dashboard/src/components/playground/compareUI/endpoint_config.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/endpoint_config.test.ts similarity index 98% rename from ui/litellm-dashboard/src/components/playground/compareUI/endpoint_config.test.ts rename to ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/endpoint_config.test.ts index 67ecf32fc2a..6463b2f9e5e 100644 --- a/ui/litellm-dashboard/src/components/playground/compareUI/endpoint_config.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/endpoint_config.test.ts @@ -12,7 +12,7 @@ import { getComparisonSelection, hasValidSelection, } from "./endpoint_config"; -import { Agent } from "../llm_calls/fetch_agents"; +import { Agent } from "../../llm_calls/fetch_agents"; describe("endpoint_config", () => { it("should export EndpointId constants", () => { diff --git a/ui/litellm-dashboard/src/components/playground/compareUI/endpoint_config.ts b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/endpoint_config.ts similarity index 98% rename from ui/litellm-dashboard/src/components/playground/compareUI/endpoint_config.ts rename to ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/endpoint_config.ts index c8e39f3ad03..2d1ebdaedd6 100644 --- a/ui/litellm-dashboard/src/components/playground/compareUI/endpoint_config.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/endpoint_config.ts @@ -3,7 +3,7 @@ * Add new endpoints here to extend the comparison functionality. */ -import { Agent } from "../llm_calls/fetch_agents"; +import { Agent } from "../../llm_calls/fetch_agents"; // Endpoint identifiers export const EndpointId = { diff --git a/ui/litellm-dashboard/src/components/playground/complianceUI/ComplianceUI.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/complianceUI/ComplianceUI.tsx similarity index 99% rename from ui/litellm-dashboard/src/components/playground/complianceUI/ComplianceUI.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/components/complianceUI/ComplianceUI.tsx index 448ddb5d32b..cf619057c38 100644 --- a/ui/litellm-dashboard/src/components/playground/complianceUI/ComplianceUI.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/complianceUI/ComplianceUI.tsx @@ -9,7 +9,7 @@ import { import { getGuardrailsList, testPoliciesAndGuardrails } from "@/components/networking"; import PolicySelector, { getPolicyOptionEntries } from "@/components/policies/PolicySelector"; import { Policy } from "@/components/policies/types"; -import { makeOpenAIChatCompletionRequest } from "../llm_calls/chat_completion"; +import { makeOpenAIChatCompletionRequest } from "@/components/llm_calls/chat_completion"; import { AlertTriangle, BarChart3, diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/playground/hooks/useChatHistory.test.ts similarity index 100% rename from ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.test.ts rename to ui/litellm-dashboard/src/app/(dashboard)/playground/hooks/useChatHistory.test.ts diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.ts b/ui/litellm-dashboard/src/app/(dashboard)/playground/hooks/useChatHistory.ts similarity index 98% rename from ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.ts rename to ui/litellm-dashboard/src/app/(dashboard)/playground/hooks/useChatHistory.ts index d0da281bcc7..6e2a263f599 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/useChatHistory.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/hooks/useChatHistory.ts @@ -1,8 +1,8 @@ import React, { useState, useEffect } from "react"; -import { MessageType, A2ATaskMetadata } from "./types"; -import { TokenUsage } from "./ResponseMetrics"; -import { MCPEvent } from "../../mcp_tools/types"; -import { truncateString } from "../../../utils/textUtils"; +import { MessageType, A2ATaskMetadata } from "@/components/chat_ui/types"; +import { TokenUsage } from "@/components/chat_ui/ResponseMetrics"; +import { MCPEvent } from "@/components/mcp_tools/types"; +import { truncateString } from "@/utils/textUtils"; export interface UseChatHistoryReturn { // State diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/useCodeInterpreter.ts b/ui/litellm-dashboard/src/app/(dashboard)/playground/hooks/useCodeInterpreter.ts similarity index 88% rename from ui/litellm-dashboard/src/components/playground/chat_ui/useCodeInterpreter.ts rename to ui/litellm-dashboard/src/app/(dashboard)/playground/hooks/useCodeInterpreter.ts index 430bee33be0..9f32ef072e1 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/useCodeInterpreter.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/hooks/useCodeInterpreter.ts @@ -4,7 +4,7 @@ */ import { useState, useCallback } from "react"; -import { CodeInterpreterResult } from "../llm_calls/code_interpreter_handler"; +import { CodeInterpreterResult } from "@/components/llm_calls/code_interpreter_handler"; export interface UseCodeInterpreterReturn { // State @@ -54,4 +54,4 @@ export function useCodeInterpreter(): UseCodeInterpreterReturn { } // Re-export the type for convenience -export type { CodeInterpreterResult } from "../llm_calls/code_interpreter_handler"; +export type { CodeInterpreterResult } from "@/components/llm_calls/code_interpreter_handler"; diff --git a/ui/litellm-dashboard/src/components/playground/llm_calls/a2a_send_message.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/a2a_send_message.tsx similarity index 98% rename from ui/litellm-dashboard/src/components/playground/llm_calls/a2a_send_message.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/a2a_send_message.tsx index 42c1c4cc81f..4e01dce0d9d 100644 --- a/ui/litellm-dashboard/src/components/playground/llm_calls/a2a_send_message.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/a2a_send_message.tsx @@ -2,8 +2,8 @@ // A2A Protocol (JSON-RPC 2.0) implementation for sending messages to agents import { v4 as uuidv4 } from "uuid"; -import { getProxyBaseUrl, getGlobalLitellmHeaderName } from "../../networking"; -import { A2ATaskMetadata } from "../chat_ui/types"; +import { getProxyBaseUrl, getGlobalLitellmHeaderName } from "@/components/networking"; +import { A2ATaskMetadata } from "@/components/chat_ui/types"; interface A2AMessagePart { kind: "text"; diff --git a/ui/litellm-dashboard/src/components/playground/llm_calls/anthropic_messages.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/anthropic_messages.tsx similarity index 96% rename from ui/litellm-dashboard/src/components/playground/llm_calls/anthropic_messages.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/anthropic_messages.tsx index 5570c7408fa..11e7a5e1601 100644 --- a/ui/litellm-dashboard/src/components/playground/llm_calls/anthropic_messages.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/anthropic_messages.tsx @@ -1,6 +1,6 @@ import Anthropic from "@anthropic-ai/sdk"; -import { MessageType } from "../chat_ui/types"; -import { TokenUsage } from "../chat_ui/ResponseMetrics"; +import { MessageType } from "@/components/chat_ui/types"; +import { TokenUsage } from "@/components/chat_ui/ResponseMetrics"; import { getProxyBaseUrl } from "@/components/networking"; import NotificationManager from "@/components/molecules/notifications_manager"; diff --git a/ui/litellm-dashboard/src/components/playground/llm_calls/audio_speech.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/audio_speech.test.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/playground/llm_calls/audio_speech.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/audio_speech.test.tsx diff --git a/ui/litellm-dashboard/src/components/playground/llm_calls/audio_speech.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/audio_speech.tsx similarity index 96% rename from ui/litellm-dashboard/src/components/playground/llm_calls/audio_speech.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/audio_speech.tsx index c5d4ae4d686..eda5d6ed66a 100644 --- a/ui/litellm-dashboard/src/components/playground/llm_calls/audio_speech.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/audio_speech.tsx @@ -1,7 +1,7 @@ import openai from "openai"; import { getProxyBaseUrl } from "@/components/networking"; import NotificationManager from "@/components/molecules/notifications_manager"; -import type { OpenAIVoice } from "../chat_ui/chatConstants"; +import type { OpenAIVoice } from "../components/chat_ui/chatConstants"; export async function makeOpenAIAudioSpeechRequest( input: string, diff --git a/ui/litellm-dashboard/src/components/playground/llm_calls/audio_transcriptions.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/audio_transcriptions.test.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/playground/llm_calls/audio_transcriptions.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/audio_transcriptions.test.tsx diff --git a/ui/litellm-dashboard/src/components/playground/llm_calls/audio_transcriptions.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/audio_transcriptions.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/playground/llm_calls/audio_transcriptions.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/audio_transcriptions.tsx diff --git a/ui/litellm-dashboard/src/components/playground/llm_calls/embeddings_api.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/embeddings_api.test.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/playground/llm_calls/embeddings_api.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/embeddings_api.test.tsx diff --git a/ui/litellm-dashboard/src/components/playground/llm_calls/embeddings_api.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/embeddings_api.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/playground/llm_calls/embeddings_api.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/embeddings_api.tsx diff --git a/ui/litellm-dashboard/src/components/playground/llm_calls/fetch_agents.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/fetch_agents.tsx similarity index 98% rename from ui/litellm-dashboard/src/components/playground/llm_calls/fetch_agents.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/fetch_agents.tsx index 258adcc93b2..0dc589188f5 100644 --- a/ui/litellm-dashboard/src/components/playground/llm_calls/fetch_agents.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/fetch_agents.tsx @@ -1,6 +1,6 @@ // fetch_agents.tsx -import { getProxyBaseUrl, getGlobalLitellmHeaderName, modelInfoCall } from "../../networking"; +import { getProxyBaseUrl, getGlobalLitellmHeaderName, modelInfoCall } from "@/components/networking"; export interface Agent { agent_id: string; diff --git a/ui/litellm-dashboard/src/components/playground/llm_calls/image_edits.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/image_edits.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/playground/llm_calls/image_edits.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/image_edits.tsx diff --git a/ui/litellm-dashboard/src/components/playground/llm_calls/image_generation.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/image_generation.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/playground/llm_calls/image_generation.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/image_generation.tsx diff --git a/ui/litellm-dashboard/src/components/playground/llm_calls/interactions_api.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/interactions_api.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/playground/llm_calls/interactions_api.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/interactions_api.tsx diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/page.tsx index abcbe80a382..c9bca86dc23 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/page.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/page.tsx @@ -1,10 +1,10 @@ "use client"; import { useState, useEffect } from "react"; -import AgentBuilderView from "@/components/playground/chat_ui/AgentBuilderView"; -import ChatUI from "@/components/playground/chat_ui/ChatUI"; -import CompareUI from "@/components/playground/compareUI/CompareUI"; -import ComplianceUI from "@/components/playground/complianceUI/ComplianceUI"; +import AgentBuilderView from "@/app/(dashboard)/playground/components/chat_ui/AgentBuilderView"; +import ChatUI from "@/app/(dashboard)/playground/components/chat_ui/ChatUI"; +import CompareUI from "@/app/(dashboard)/playground/components/compareUI/CompareUI"; +import ComplianceUI from "@/app/(dashboard)/playground/components/complianceUI/ComplianceUI"; import { TabGroup, TabList, Tab, TabPanels, TabPanel } from "@tremor/react"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { fetchProxySettings } from "@/utils/proxyUtils"; diff --git a/ui/litellm-dashboard/src/components/CostTrackingSettings/cost_tracking_settings.test.tsx b/ui/litellm-dashboard/src/components/CostTrackingSettings/cost_tracking_settings.test.tsx index b0cf4f8262a..89711a098fe 100644 --- a/ui/litellm-dashboard/src/components/CostTrackingSettings/cost_tracking_settings.test.tsx +++ b/ui/litellm-dashboard/src/components/CostTrackingSettings/cost_tracking_settings.test.tsx @@ -33,7 +33,7 @@ vi.mock("./pricing_calculator/index", () => ({ default: () =>
Pricing Calculator
, })); -vi.mock("../playground/llm_calls/fetch_models", () => ({ +vi.mock("@/components/llm_calls/fetch_models", () => ({ fetchAvailableModels: vi.fn().mockResolvedValue([]), })); diff --git a/ui/litellm-dashboard/src/components/CostTrackingSettings/cost_tracking_settings.tsx b/ui/litellm-dashboard/src/components/CostTrackingSettings/cost_tracking_settings.tsx index 9dade715cbc..d9cca4d3c23 100644 --- a/ui/litellm-dashboard/src/components/CostTrackingSettings/cost_tracking_settings.tsx +++ b/ui/litellm-dashboard/src/components/CostTrackingSettings/cost_tracking_settings.tsx @@ -24,7 +24,7 @@ import { DocsMenu } from "../HelpLink"; import HowItWorks from "./how_it_works"; import { useDiscountConfig } from "./use_discount_config"; import { useMarginConfig } from "./use_margin_config"; -import { fetchAvailableModels, ModelGroup } from "../playground/llm_calls/fetch_models"; +import { fetchAvailableModels, ModelGroup } from "@/components/llm_calls/fetch_models"; const DOCS_LINKS = [ { label: "Custom pricing for models", href: "https://docs.litellm.ai/docs/proxy/custom_pricing" }, diff --git a/ui/litellm-dashboard/src/components/GuardrailsMonitor/EvaluationSettingsModal.tsx b/ui/litellm-dashboard/src/components/GuardrailsMonitor/EvaluationSettingsModal.tsx index 90a2c0d597d..0edfa65dfe8 100644 --- a/ui/litellm-dashboard/src/components/GuardrailsMonitor/EvaluationSettingsModal.tsx +++ b/ui/litellm-dashboard/src/components/GuardrailsMonitor/EvaluationSettingsModal.tsx @@ -1,7 +1,7 @@ import { CloseOutlined, PlayCircleOutlined } from "@ant-design/icons"; import { Button, Modal, Select, Input } from "antd"; import React, { useEffect, useState } from "react"; -import { fetchAvailableModels, type ModelGroup } from "@/components/playground/llm_calls/fetch_models"; +import { fetchAvailableModels, type ModelGroup } from "@/components/llm_calls/fetch_models"; const DEFAULT_PROMPT = `Evaluate whether this guardrail's decision was correct. Analyze the user input, the guardrail action taken, and determine if it was appropriate. diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPSemanticFilterSettings/MCPSemanticFilterSettings.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPSemanticFilterSettings/MCPSemanticFilterSettings.test.tsx index a2288f40039..39542945c45 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPSemanticFilterSettings/MCPSemanticFilterSettings.test.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPSemanticFilterSettings/MCPSemanticFilterSettings.test.tsx @@ -14,7 +14,7 @@ vi.mock("@/app/(dashboard)/hooks/mcpSemanticFilterSettings/useUpdateMCPSemanticF useUpdateMCPSemanticFilterSettings: vi.fn(), })); -vi.mock("@/components/playground/llm_calls/fetch_models", () => ({ +vi.mock("@/components/llm_calls/fetch_models", () => ({ fetchAvailableModels: vi.fn().mockResolvedValue([]), })); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPSemanticFilterSettings/MCPSemanticFilterSettings.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPSemanticFilterSettings/MCPSemanticFilterSettings.tsx index c62e752565c..38b1420f97e 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPSemanticFilterSettings/MCPSemanticFilterSettings.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/MCPSemanticFilterSettings/MCPSemanticFilterSettings.tsx @@ -21,7 +21,7 @@ import { } from "antd"; import { QuestionCircleOutlined, CheckCircleOutlined, SaveOutlined } from "@ant-design/icons"; import { useEffect, useState } from "react"; -import { fetchAvailableModels, ModelGroup } from "@/components/playground/llm_calls/fetch_models"; +import { fetchAvailableModels, ModelGroup } from "@/components/llm_calls/fetch_models"; import MCPSemanticFilterTestPanel from "./MCPSemanticFilterTestPanel"; import { getCurlCommand, runSemanticFilterTest, TestResult } from "./semanticFilterTestUtils"; diff --git a/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/AddFallbacks.test.tsx b/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/AddFallbacks.test.tsx index 0e05c141570..1253532f269 100644 --- a/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/AddFallbacks.test.tsx +++ b/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/AddFallbacks.test.tsx @@ -2,9 +2,9 @@ import { render, screen, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; import AddFallbacks, { Fallbacks } from "./AddFallbacks"; -import * as fetchModelsModule from "../../../playground/llm_calls/fetch_models"; +import * as fetchModelsModule from "@/components/llm_calls/fetch_models"; -vi.mock("../../../playground/llm_calls/fetch_models", () => ({ +vi.mock("@/components/llm_calls/fetch_models", () => ({ fetchAvailableModels: vi.fn(), })); diff --git a/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/AddFallbacks.tsx b/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/AddFallbacks.tsx index 8b86ff862f9..8147189ddcc 100644 --- a/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/AddFallbacks.tsx +++ b/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/AddFallbacks.tsx @@ -9,7 +9,7 @@ import { Button } from "antd"; import React, { useEffect, useState } from "react"; import MessageManager from "@/components/molecules/message_manager"; import NotificationManager from "../../../molecules/notifications_manager"; -import { fetchAvailableModels, ModelGroup } from "../../../playground/llm_calls/fetch_models"; +import { fetchAvailableModels, ModelGroup } from "@/components/llm_calls/fetch_models"; import { AddFallbacksModal } from "./AddFallbacksModal"; import { FallbackGroup } from "./FallbackGroupConfig"; import { FallbackSelectionForm } from "./FallbackSelectionForm"; diff --git a/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/Fallbacks.test.tsx b/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/Fallbacks.test.tsx index 323b71bd5c4..513abb78590 100644 --- a/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/Fallbacks.test.tsx +++ b/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/Fallbacks.test.tsx @@ -3,14 +3,14 @@ import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; import Fallbacks from "./Fallbacks"; import * as networkingModule from "../../../networking"; -import * as fetchModelsModule from "../../../playground/llm_calls/fetch_models"; +import * as fetchModelsModule from "@/components/llm_calls/fetch_models"; vi.mock("../../../networking", () => ({ getCallbacksCall: vi.fn(), setCallbacksCall: vi.fn(), })); -vi.mock("../../../playground/llm_calls/fetch_models", () => ({ +vi.mock("@/components/llm_calls/fetch_models", () => ({ fetchAvailableModels: vi.fn(), })); diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx index a20cb969e33..d826f3df32c 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx @@ -1,7 +1,7 @@ import { InfoCircleOutlined } from "@ant-design/icons"; import { Select as AntdSelect, Card, Divider, Space, Tooltip, Typography } from "antd"; import React from "react"; -import { ModelGroup } from "../playground/llm_calls/fetch_models"; +import { ModelGroup } from "@/components/llm_calls/fetch_models"; const { Text } = Typography; diff --git a/ui/litellm-dashboard/src/components/add_model/RouterConfigBuilder.tsx b/ui/litellm-dashboard/src/components/add_model/RouterConfigBuilder.tsx index 17c962c7e25..08acf993e2c 100644 --- a/ui/litellm-dashboard/src/components/add_model/RouterConfigBuilder.tsx +++ b/ui/litellm-dashboard/src/components/add_model/RouterConfigBuilder.tsx @@ -14,7 +14,7 @@ import { Typography, } from "antd"; import React, { useEffect, useState } from "react"; -import { ModelGroup } from "../playground/llm_calls/fetch_models"; +import { ModelGroup } from "@/components/llm_calls/fetch_models"; const { Text } = Typography; diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx index 6d01e9bf74e..9e9a097e3fe 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx @@ -6,7 +6,7 @@ import { modelAvailableCall } from "../networking"; import ConnectionErrorDisplay from "./model_connection_test"; import { all_admin_roles } from "@/utils/roles"; import { handleAddAutoRouterSubmit } from "./handle_add_auto_router_submit"; -import { fetchAvailableModels, ModelGroup } from "../playground/llm_calls/fetch_models"; +import { fetchAvailableModels, ModelGroup } from "@/components/llm_calls/fetch_models"; import RouterConfigBuilder from "./RouterConfigBuilder"; import ComplexityRouterConfig from "./ComplexityRouterConfig"; import NotificationManager from "../molecules/notifications_manager"; diff --git a/ui/litellm-dashboard/src/components/cache_settings/CacheFieldRenderer.tsx b/ui/litellm-dashboard/src/components/cache_settings/CacheFieldRenderer.tsx index eeabda23f9f..6608b09d261 100644 --- a/ui/litellm-dashboard/src/components/cache_settings/CacheFieldRenderer.tsx +++ b/ui/litellm-dashboard/src/components/cache_settings/CacheFieldRenderer.tsx @@ -4,7 +4,7 @@ import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { NumberInput, TextInput } from "@tremor/react"; import { Select } from "antd"; import React, { useEffect, useState } from "react"; -import { fetchAvailableModels, ModelGroup } from "../playground/llm_calls/fetch_models"; +import { fetchAvailableModels, ModelGroup } from "@/components/llm_calls/fetch_models"; import NumericalInput from "../shared/numerical_input"; interface CacheFieldRendererProps { diff --git a/ui/litellm-dashboard/src/components/chat/ChatMessages.tsx b/ui/litellm-dashboard/src/components/chat/ChatMessages.tsx index 67b55697481..53877be1737 100644 --- a/ui/litellm-dashboard/src/components/chat/ChatMessages.tsx +++ b/ui/litellm-dashboard/src/components/chat/ChatMessages.tsx @@ -7,8 +7,8 @@ import ReactMarkdown from "react-markdown"; import remarkGfm from "remark-gfm"; import { Prism as SyntaxHighlighter } from "react-syntax-highlighter"; import { coy } from "react-syntax-highlighter/dist/esm/styles/prism"; -import ReasoningContent from "../playground/chat_ui/ReasoningContent"; -import MCPEventsDisplay from "../playground/chat_ui/MCPEventsDisplay"; +import ReasoningContent from "@/components/chat_ui/ReasoningContent"; +import MCPEventsDisplay from "@/components/chat_ui/MCPEventsDisplay"; import { ChatMessage } from "./types"; const { Panel } = Collapse; diff --git a/ui/litellm-dashboard/src/components/chat/ChatPage.tsx b/ui/litellm-dashboard/src/components/chat/ChatPage.tsx index 0e729c8ae9d..2a7fdb75eb8 100644 --- a/ui/litellm-dashboard/src/components/chat/ChatPage.tsx +++ b/ui/litellm-dashboard/src/components/chat/ChatPage.tsx @@ -27,9 +27,9 @@ import ChatMessages from "./ChatMessages"; import MCPConnectPicker from "./MCPConnectPicker"; import MCPAppsPanel from "./MCPAppsPanel"; import MCPCredentialsTab from "./MCPCredentialsTab"; -import { fetchAvailableModels } from "../playground/llm_calls/fetch_models"; -import { makeOpenAIChatCompletionRequest } from "../playground/llm_calls/chat_completion"; -import { makeOpenAIResponsesRequest } from "../playground/llm_calls/responses_api"; +import { fetchAvailableModels } from "@/components/llm_calls/fetch_models"; +import { makeOpenAIChatCompletionRequest } from "@/components/llm_calls/chat_completion"; +import { makeOpenAIResponsesRequest } from "@/components/llm_calls/responses_api"; import type { MCPEvent } from "./types"; import { getProxyBaseUrl } from "@/components/networking"; import { useUIConfig } from "@/app/(dashboard)/hooks/uiConfig/useUIConfig"; diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/CodeSnippets.test.tsx b/ui/litellm-dashboard/src/components/chat_ui/CodeSnippets.test.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/playground/chat_ui/CodeSnippets.test.tsx rename to ui/litellm-dashboard/src/components/chat_ui/CodeSnippets.test.tsx diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/CodeSnippets.tsx b/ui/litellm-dashboard/src/components/chat_ui/CodeSnippets.tsx similarity index 99% rename from ui/litellm-dashboard/src/components/playground/chat_ui/CodeSnippets.tsx rename to ui/litellm-dashboard/src/components/chat_ui/CodeSnippets.tsx index d38d4dff7a7..576b094b5d3 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/CodeSnippets.tsx +++ b/ui/litellm-dashboard/src/components/chat_ui/CodeSnippets.tsx @@ -1,6 +1,6 @@ import { MessageType } from "./types"; import { EndpointType } from "./mode_endpoint_mapping"; -import { MCPServer } from "../../mcp_tools/types"; +import { MCPServer } from "@/components/mcp_tools/types"; interface CodeGenMetadata { tags?: string[]; diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/MCPEventsDisplay.tsx b/ui/litellm-dashboard/src/components/chat_ui/MCPEventsDisplay.tsx similarity index 99% rename from ui/litellm-dashboard/src/components/playground/chat_ui/MCPEventsDisplay.tsx rename to ui/litellm-dashboard/src/components/chat_ui/MCPEventsDisplay.tsx index cd7ecbf266f..e9319169d20 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/MCPEventsDisplay.tsx +++ b/ui/litellm-dashboard/src/components/chat_ui/MCPEventsDisplay.tsx @@ -1,6 +1,6 @@ import React from "react"; import { Typography, Collapse } from "antd"; -import type { MCPEvent } from "../../mcp_tools/types"; +import type { MCPEvent } from "@/components/mcp_tools/types"; const { Text } = Typography; const { Panel } = Collapse; diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/ReasoningContent.tsx b/ui/litellm-dashboard/src/components/chat_ui/ReasoningContent.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/playground/chat_ui/ReasoningContent.tsx rename to ui/litellm-dashboard/src/components/chat_ui/ReasoningContent.tsx diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/ResponseMetrics.tsx b/ui/litellm-dashboard/src/components/chat_ui/ResponseMetrics.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/playground/chat_ui/ResponseMetrics.tsx rename to ui/litellm-dashboard/src/components/chat_ui/ResponseMetrics.tsx diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/mode_endpoint_mapping.tsx b/ui/litellm-dashboard/src/components/chat_ui/mode_endpoint_mapping.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/playground/chat_ui/mode_endpoint_mapping.tsx rename to ui/litellm-dashboard/src/components/chat_ui/mode_endpoint_mapping.tsx diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/types.ts b/ui/litellm-dashboard/src/components/chat_ui/types.ts similarity index 100% rename from ui/litellm-dashboard/src/components/playground/chat_ui/types.ts rename to ui/litellm-dashboard/src/components/chat_ui/types.ts diff --git a/ui/litellm-dashboard/src/components/common_components/ModelSelector.tsx b/ui/litellm-dashboard/src/components/common_components/ModelSelector.tsx index 24d6696dec6..060a8f0f111 100644 --- a/ui/litellm-dashboard/src/components/common_components/ModelSelector.tsx +++ b/ui/litellm-dashboard/src/components/common_components/ModelSelector.tsx @@ -2,7 +2,7 @@ import React, { useState, useEffect, useRef } from "react"; import { TextInput, Text } from "@tremor/react"; import { Select } from "antd"; import { RobotOutlined } from "@ant-design/icons"; -import { fetchAvailableModels, ModelGroup } from "../playground/llm_calls/fetch_models"; +import { fetchAvailableModels, ModelGroup } from "@/components/llm_calls/fetch_models"; interface ModelSelectorProps { accessToken: string; diff --git a/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.tsx b/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.tsx index 74b968b5b03..0aa274b5749 100644 --- a/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.tsx +++ b/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.tsx @@ -5,7 +5,7 @@ import RouterSettingsForm, { RouterSettingsFormValue } from "../router_settings/ import { Fallbacks } from "../Settings/RouterSettings/Fallbacks/AddFallbacks"; import { FallbackSelectionForm } from "../Settings/RouterSettings/Fallbacks/FallbackSelectionForm"; import { FallbackGroup } from "../Settings/RouterSettings/Fallbacks/FallbackGroupConfig"; -import { fetchAvailableModels, ModelGroup } from "../playground/llm_calls/fetch_models"; +import { fetchAvailableModels, ModelGroup } from "@/components/llm_calls/fetch_models"; export interface RouterSettingsAccordionValue { router_settings: { diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx index e4e6ebd1da9..caa061147d8 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx @@ -2,7 +2,7 @@ import React, { useEffect, useState } from "react"; import { Modal, Form, Button, Select as AntdSelect } from "antd"; import { Text, TextInput } from "@tremor/react"; import { modelAvailableCall, modelPatchUpdateCall } from "../networking"; -import { fetchAvailableModels, ModelGroup } from "../playground/llm_calls/fetch_models"; +import { fetchAvailableModels, ModelGroup } from "@/components/llm_calls/fetch_models"; import RouterConfigBuilder from "../add_model/RouterConfigBuilder"; import NotificationsManager from "../molecules/notifications_manager"; diff --git a/ui/litellm-dashboard/src/components/playground/llm_calls/chat_completion.test.tsx b/ui/litellm-dashboard/src/components/llm_calls/chat_completion.test.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/playground/llm_calls/chat_completion.test.tsx rename to ui/litellm-dashboard/src/components/llm_calls/chat_completion.test.tsx diff --git a/ui/litellm-dashboard/src/components/playground/llm_calls/chat_completion.tsx b/ui/litellm-dashboard/src/components/llm_calls/chat_completion.tsx similarity index 99% rename from ui/litellm-dashboard/src/components/playground/llm_calls/chat_completion.tsx rename to ui/litellm-dashboard/src/components/llm_calls/chat_completion.tsx index a45c3036f59..54d9273c463 100644 --- a/ui/litellm-dashboard/src/components/playground/llm_calls/chat_completion.tsx +++ b/ui/litellm-dashboard/src/components/llm_calls/chat_completion.tsx @@ -3,7 +3,7 @@ import { ChatCompletionMessageParam } from "openai/resources/chat/completions"; import { TokenUsage } from "../chat_ui/ResponseMetrics"; import { VectorStoreSearchResponse } from "../chat_ui/types"; import { getProxyBaseUrl } from "@/components/networking"; -import { MCPServer, MCPToolset, type MCPEvent } from "../../mcp_tools/types"; +import { MCPServer, MCPToolset, type MCPEvent } from "@/components/mcp_tools/types"; export async function makeOpenAIChatCompletionRequest( chatHistory: { role: string; content: string | any[] }[], diff --git a/ui/litellm-dashboard/src/components/playground/llm_calls/code_interpreter_handler.ts b/ui/litellm-dashboard/src/components/llm_calls/code_interpreter_handler.ts similarity index 100% rename from ui/litellm-dashboard/src/components/playground/llm_calls/code_interpreter_handler.ts rename to ui/litellm-dashboard/src/components/llm_calls/code_interpreter_handler.ts diff --git a/ui/litellm-dashboard/src/components/playground/llm_calls/fetch_models.tsx b/ui/litellm-dashboard/src/components/llm_calls/fetch_models.tsx similarity index 94% rename from ui/litellm-dashboard/src/components/playground/llm_calls/fetch_models.tsx rename to ui/litellm-dashboard/src/components/llm_calls/fetch_models.tsx index 14ba8c2381b..6be6b65502a 100644 --- a/ui/litellm-dashboard/src/components/playground/llm_calls/fetch_models.tsx +++ b/ui/litellm-dashboard/src/components/llm_calls/fetch_models.tsx @@ -1,6 +1,6 @@ // fetch_models.ts -import { modelHubCall } from "../../networking"; +import { modelHubCall } from "@/components/networking"; export interface ModelGroup { model_group: string; diff --git a/ui/litellm-dashboard/src/components/playground/llm_calls/responses_api.test.tsx b/ui/litellm-dashboard/src/components/llm_calls/responses_api.test.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/playground/llm_calls/responses_api.test.tsx rename to ui/litellm-dashboard/src/components/llm_calls/responses_api.test.tsx diff --git a/ui/litellm-dashboard/src/components/playground/llm_calls/responses_api.tsx b/ui/litellm-dashboard/src/components/llm_calls/responses_api.tsx similarity index 98% rename from ui/litellm-dashboard/src/components/playground/llm_calls/responses_api.tsx rename to ui/litellm-dashboard/src/components/llm_calls/responses_api.tsx index 4e88a356cf3..d3dd866b36b 100644 --- a/ui/litellm-dashboard/src/components/playground/llm_calls/responses_api.tsx +++ b/ui/litellm-dashboard/src/components/llm_calls/responses_api.tsx @@ -3,8 +3,8 @@ import { MessageType } from "../chat_ui/types"; import { TokenUsage } from "../chat_ui/ResponseMetrics"; import { getProxyBaseUrl } from "@/components/networking"; import NotificationManager from "@/components/molecules/notifications_manager"; -import type { MCPEvent } from "../../mcp_tools/types"; -import { MCPServer, MCPToolset } from "../../mcp_tools/types"; +import type { MCPEvent } from "@/components/mcp_tools/types"; +import { MCPServer, MCPToolset } from "@/components/mcp_tools/types"; import { CodeInterpreterResult, CodeInterpreterState, diff --git a/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/conversation_panel/MessageBubble.tsx b/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/conversation_panel/MessageBubble.tsx index 12189f2b2c3..7d58465751e 100644 --- a/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/conversation_panel/MessageBubble.tsx +++ b/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/conversation_panel/MessageBubble.tsx @@ -3,7 +3,7 @@ import { RobotOutlined, UserOutlined } from "@ant-design/icons"; import ReactMarkdown from "react-markdown"; import { Prism as SyntaxHighlighter } from "react-syntax-highlighter"; import { coy } from "react-syntax-highlighter/dist/esm/styles/prism"; -import ResponseMetrics from "../../../playground/chat_ui/ResponseMetrics"; +import ResponseMetrics from "@/components/chat_ui/ResponseMetrics"; import { Message } from "./types"; interface MessageBubbleProps { diff --git a/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/conversation_panel/types.ts b/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/conversation_panel/types.ts index b56f7b2f582..33d6f2dc7cc 100644 --- a/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/conversation_panel/types.ts +++ b/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/conversation_panel/types.ts @@ -1,4 +1,4 @@ -import { TokenUsage } from "../../../playground/chat_ui/ResponseMetrics"; +import { TokenUsage } from "@/components/chat_ui/ResponseMetrics"; export interface Message { role: string; diff --git a/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/conversation_panel/useConversation.ts b/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/conversation_panel/useConversation.ts index 7b595c777c1..e55d8bdeadf 100644 --- a/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/conversation_panel/useConversation.ts +++ b/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/conversation_panel/useConversation.ts @@ -1,6 +1,6 @@ import { useState, useRef, useEffect } from "react"; import NotificationsManager from "../../../molecules/notifications_manager"; -import { TokenUsage } from "../../../playground/chat_ui/ResponseMetrics"; +import { TokenUsage } from "@/components/chat_ui/ResponseMetrics"; import { Message } from "./types"; import { convertToDotPrompt, extractVariables } from "../utils"; import { getProxyBaseUrl, getGlobalLitellmHeaderName } from "../../../networking"; diff --git a/ui/litellm-dashboard/src/components/public_model_hub.tsx b/ui/litellm-dashboard/src/components/public_model_hub.tsx index 6ad8db19d8d..5299ec0fce8 100644 --- a/ui/litellm-dashboard/src/components/public_model_hub.tsx +++ b/ui/litellm-dashboard/src/components/public_model_hub.tsx @@ -18,9 +18,9 @@ import { } from "./networking"; import { Plugin } from "./claude_code_plugins/types"; import SkillHubDashboard from "./AIHub/SkillHubDashboard"; -import { generateCodeSnippet } from "./playground/chat_ui/CodeSnippets"; -import { getEndpointType } from "./playground/chat_ui/mode_endpoint_mapping"; -import { MessageType } from "./playground/chat_ui/types"; +import { generateCodeSnippet } from "@/components/chat_ui/CodeSnippets"; +import { getEndpointType } from "@/components/chat_ui/mode_endpoint_mapping"; +import { MessageType } from "@/components/chat_ui/types"; import { getProviderLogoAndName } from "./provider_info_helpers"; const { TabPane } = Tabs; diff --git a/ui/litellm-dashboard/src/components/vector_store_management/S3VectorsConfig.test.tsx b/ui/litellm-dashboard/src/components/vector_store_management/S3VectorsConfig.test.tsx index 459f7ed5dd5..b43c74d71a0 100644 --- a/ui/litellm-dashboard/src/components/vector_store_management/S3VectorsConfig.test.tsx +++ b/ui/litellm-dashboard/src/components/vector_store_management/S3VectorsConfig.test.tsx @@ -1,10 +1,10 @@ import { render, screen, fireEvent, waitFor, act } from "@testing-library/react"; import { describe, it, expect, vi, beforeEach } from "vitest"; import S3VectorsConfig from "./S3VectorsConfig"; -import * as fetchModels from "../playground/llm_calls/fetch_models"; +import * as fetchModels from "@/components/llm_calls/fetch_models"; // Mock fetchAvailableModels -vi.mock("../playground/llm_calls/fetch_models", () => ({ +vi.mock("@/components/llm_calls/fetch_models", () => ({ fetchAvailableModels: vi.fn(), })); diff --git a/ui/litellm-dashboard/src/components/vector_store_management/S3VectorsConfig.tsx b/ui/litellm-dashboard/src/components/vector_store_management/S3VectorsConfig.tsx index 0568f982bab..9f2c8f5e45e 100644 --- a/ui/litellm-dashboard/src/components/vector_store_management/S3VectorsConfig.tsx +++ b/ui/litellm-dashboard/src/components/vector_store_management/S3VectorsConfig.tsx @@ -1,7 +1,7 @@ import React, { useState, useEffect } from "react"; import { Alert, Form, Input, Select, Tooltip } from "antd"; import { InfoCircleOutlined } from "@ant-design/icons"; -import { fetchAvailableModels, ModelGroup } from "../playground/llm_calls/fetch_models"; +import { fetchAvailableModels, ModelGroup } from "@/components/llm_calls/fetch_models"; interface S3VectorsConfigProps { accessToken: string | null; diff --git a/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreForm.tsx b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreForm.tsx index 538292a2437..4f0433025b7 100644 --- a/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreForm.tsx +++ b/ui/litellm-dashboard/src/components/vector_store_management/VectorStoreForm.tsx @@ -10,7 +10,7 @@ import { getProviderSpecificFields, VectorStoreFieldConfig, } from "../vector_store_providers"; -import { fetchAvailableModels, ModelGroup } from "../playground/llm_calls/fetch_models"; +import { fetchAvailableModels, ModelGroup } from "@/components/llm_calls/fetch_models"; import NotificationsManager from "../molecules/notifications_manager"; interface VectorStoreFormProps { diff --git a/ui/litellm-dashboard/src/utils/migratedPages.test.ts b/ui/litellm-dashboard/src/utils/migratedPages.test.ts index bd4ad7af5b8..58958f7b649 100644 --- a/ui/litellm-dashboard/src/utils/migratedPages.test.ts +++ b/ui/litellm-dashboard/src/utils/migratedPages.test.ts @@ -40,6 +40,13 @@ describe("migratedHref / legacyPageHref", () => { expect(MIGRATED_PAGES.api_ref).toBe("api-reference"); expect(MIGRATED_PAGES["api-reference"]).toBe("api-reference"); }); + + it("maps the llm-playground sidebar id to the playground route", async () => { + vi.doMock("@/components/networking", () => ({ serverRootPath: "/" })); + const { MIGRATED_PAGES } = await import("./migratedPages"); + + expect(MIGRATED_PAGES["llm-playground"]).toBe("playground"); + }); }); describe("dev server (NODE_ENV=development)", () => { diff --git a/ui/litellm-dashboard/src/utils/migratedPages.ts b/ui/litellm-dashboard/src/utils/migratedPages.ts index d6f1e7d6f1d..d911bf0566a 100644 --- a/ui/litellm-dashboard/src/utils/migratedPages.ts +++ b/ui/litellm-dashboard/src/utils/migratedPages.ts @@ -12,6 +12,7 @@ export const MIGRATED_PAGES: Record = { api_ref: "api-reference", // Legacy alias: older bookmarks used the hyphenated ?page=api-reference form. "api-reference": "api-reference", + "llm-playground": "playground", }; function uiBase(): string { From a2c916fb45e238020f48e665beb43a8a360f2641 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Thu, 11 Jun 2026 13:20:21 -0700 Subject: [PATCH 068/185] feat(ui): migrate projects and access-groups to path routes (#30226) * feat(ui): cut projects and access-groups over to path routes Same recipe as playground (#30185): MIGRATED_PAGES entries route the sidebar and redirect the legacy ?page= URLs, the switch arms are deleted, and the e2e fixture grows two entries. Both components were already zero-prop and self-fetching via React Query hooks, so the route wrappers are trivial. * refactor(ui): move Projects and AccessGroups components into their route folders Both folders were imported only by the legacy switch, so they colocate wholesale under (dashboard)/{projects,access-groups}/components. Their React Query hooks stay in the shared (dashboard)/hooks layer. eslint suppressions are re-keyed to the new paths. * test(ui): enable enable_projects_ui in e2e global setup The projects migration smoke clicks the Projects sidebar link, which only renders when the enterprise-gated enable_projects_ui setting is on; the seeded e2e database starts with it off, so the locator timed out in both e2e_ui_testing jobs. CI already launches the proxy with LITELLM_LICENSE for premium UI coverage, so flip the setting in globalSetup via the same /update/ui_settings call the admin UI toggle makes, failing loudly if the PATCH is rejected. * test(ui): use Playwright request context instead of raw fetch in global setup The frontend lint bans raw fetch() outside src/lib/http/; the e2e convention for proxy API calls is Playwright's APIRequestContext, as in routerSettings.spec.ts. --- .../e2e_tests/fixtures/migratedPages.ts | 6 ++++-- ui/litellm-dashboard/e2e_tests/globalSetup.ts | 17 ++++++++++++++++- ui/litellm-dashboard/eslint-suppressions.json | 8 ++++---- .../AccessGroupsDetailsPage.test.tsx | 2 +- .../components}/AccessGroupsDetailsPage.tsx | 2 +- .../AccessGroupsModal/AccessGroupBaseForm.tsx | 0 .../AccessGroupCreateModal.tsx | 0 .../AccessGroupsModal/AccessGroupEditModal.tsx | 0 .../components}/AccessGroupsPage.test.tsx | 2 +- .../components}/AccessGroupsPage.tsx | 6 +++--- .../access-groups/components}/types.ts | 0 .../src/app/(dashboard)/access-groups/page.tsx | 9 +++++++++ .../src/app/(dashboard)/page.tsx | 6 ------ .../components}/ProjectDetailsPage.test.tsx | 2 +- .../projects/components}/ProjectDetailsPage.tsx | 2 +- .../components}/ProjectKeysSection.test.tsx | 2 +- .../projects/components}/ProjectKeysSection.tsx | 0 .../components}/ProjectKeysTable.test.tsx | 2 +- .../projects/components}/ProjectKeysTable.tsx | 2 +- .../ProjectModals/CreateProjectModal.test.tsx | 2 +- .../ProjectModals/CreateProjectModal.tsx | 0 .../ProjectModals/EditProjectModal.test.tsx | 2 +- .../ProjectModals/EditProjectModal.tsx | 0 .../ProjectModals/ProjectBaseForm.test.tsx | 2 +- .../ProjectModals/ProjectBaseForm.tsx | 6 +++--- .../ProjectModals/projectFormUtils.test.ts | 0 .../ProjectModals/projectFormUtils.ts | 0 .../projects/components}/ProjectsPage.test.tsx | 2 +- .../projects/components}/ProjectsPage.tsx | 0 .../src/app/(dashboard)/projects/page.tsx | 9 +++++++++ .../src/utils/migratedPages.test.ts | 8 ++++++++ ui/litellm-dashboard/src/utils/migratedPages.ts | 2 ++ ui/litellm-dashboard/tsconfig.tsbuildinfo | 2 +- 33 files changed, 71 insertions(+), 32 deletions(-) rename ui/litellm-dashboard/src/{components/AccessGroups => app/(dashboard)/access-groups/components}/AccessGroupsDetailsPage.test.tsx (99%) rename ui/litellm-dashboard/src/{components/AccessGroups => app/(dashboard)/access-groups/components}/AccessGroupsDetailsPage.tsx (99%) rename ui/litellm-dashboard/src/{components/AccessGroups => app/(dashboard)/access-groups/components}/AccessGroupsModal/AccessGroupBaseForm.tsx (100%) rename ui/litellm-dashboard/src/{components/AccessGroups => app/(dashboard)/access-groups/components}/AccessGroupsModal/AccessGroupCreateModal.tsx (100%) rename ui/litellm-dashboard/src/{components/AccessGroups => app/(dashboard)/access-groups/components}/AccessGroupsModal/AccessGroupEditModal.tsx (100%) rename ui/litellm-dashboard/src/{components/AccessGroups => app/(dashboard)/access-groups/components}/AccessGroupsPage.test.tsx (99%) rename ui/litellm-dashboard/src/{components/AccessGroups => app/(dashboard)/access-groups/components}/AccessGroupsPage.tsx (97%) rename ui/litellm-dashboard/src/{components/AccessGroups => app/(dashboard)/access-groups/components}/types.ts (100%) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/access-groups/page.tsx rename ui/litellm-dashboard/src/{components/Projects => app/(dashboard)/projects/components}/ProjectDetailsPage.test.tsx (98%) rename ui/litellm-dashboard/src/{components/Projects => app/(dashboard)/projects/components}/ProjectDetailsPage.tsx (99%) rename ui/litellm-dashboard/src/{components/Projects => app/(dashboard)/projects/components}/ProjectKeysSection.test.tsx (97%) rename ui/litellm-dashboard/src/{components/Projects => app/(dashboard)/projects/components}/ProjectKeysSection.tsx (100%) rename ui/litellm-dashboard/src/{components/Projects => app/(dashboard)/projects/components}/ProjectKeysTable.test.tsx (98%) rename ui/litellm-dashboard/src/{components/Projects => app/(dashboard)/projects/components}/ProjectKeysTable.tsx (94%) rename ui/litellm-dashboard/src/{components/Projects => app/(dashboard)/projects/components}/ProjectModals/CreateProjectModal.test.tsx (96%) rename ui/litellm-dashboard/src/{components/Projects => app/(dashboard)/projects/components}/ProjectModals/CreateProjectModal.tsx (100%) rename ui/litellm-dashboard/src/{components/Projects => app/(dashboard)/projects/components}/ProjectModals/EditProjectModal.test.tsx (97%) rename ui/litellm-dashboard/src/{components/Projects => app/(dashboard)/projects/components}/ProjectModals/EditProjectModal.tsx (100%) rename ui/litellm-dashboard/src/{components/Projects => app/(dashboard)/projects/components}/ProjectModals/ProjectBaseForm.test.tsx (99%) rename ui/litellm-dashboard/src/{components/Projects => app/(dashboard)/projects/components}/ProjectModals/ProjectBaseForm.tsx (98%) rename ui/litellm-dashboard/src/{components/Projects => app/(dashboard)/projects/components}/ProjectModals/projectFormUtils.test.ts (100%) rename ui/litellm-dashboard/src/{components/Projects => app/(dashboard)/projects/components}/ProjectModals/projectFormUtils.ts (100%) rename ui/litellm-dashboard/src/{components/Projects => app/(dashboard)/projects/components}/ProjectsPage.test.tsx (98%) rename ui/litellm-dashboard/src/{components/Projects => app/(dashboard)/projects/components}/ProjectsPage.tsx (100%) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/projects/page.tsx diff --git a/ui/litellm-dashboard/e2e_tests/fixtures/migratedPages.ts b/ui/litellm-dashboard/e2e_tests/fixtures/migratedPages.ts index 749c3cde179..38b9a875828 100644 --- a/ui/litellm-dashboard/e2e_tests/fixtures/migratedPages.ts +++ b/ui/litellm-dashboard/e2e_tests/fixtures/migratedPages.ts @@ -11,13 +11,15 @@ * Keep this in lockstep with MIGRATED_PAGES in src/utils/migratedPages.ts. * Pending (add as each PR lands): the leaf-pages batch * (budgets, caching, cost-tracking, guardrails, guardrails-monitor, logs, - * mcp-servers, memory, policies, projects, prompts, search-tools, skills, + * mcp-servers, memory, policies, prompts, search-tools, skills, * tag-management, tool-policies, transform-request, ui-theme, vector-stores, - * workflows, access-groups). + * workflows). */ export const MIGRATED_E2E_PAGES: Record = { api_ref: "api-reference", "llm-playground": "playground", + projects: "projects", + "access-groups": "access-groups", }; export const MIGRATED_E2E_SEGMENTS: string[] = [...new Set(Object.values(MIGRATED_E2E_PAGES))]; diff --git a/ui/litellm-dashboard/e2e_tests/globalSetup.ts b/ui/litellm-dashboard/e2e_tests/globalSetup.ts index 0b3fa7e8807..661155b761f 100644 --- a/ui/litellm-dashboard/e2e_tests/globalSetup.ts +++ b/ui/litellm-dashboard/e2e_tests/globalSetup.ts @@ -1,4 +1,4 @@ -import { chromium, expect } from "@playwright/test"; +import { chromium, expect, request } from "@playwright/test"; import { users, Role, STORAGE_PATHS } from "./fixtures/users"; import * as fs from "fs"; @@ -6,6 +6,21 @@ async function globalSetup() { const browser = await chromium.launch(); const rootPath = process.env.SERVER_ROOT_PATH ?? ""; + // The Projects sidebar item is hidden unless the enterprise-gated + // enable_projects_ui setting is on, and the seeded DB starts with it off. + // The proxy runs with LITELLM_LICENSE in CI, so enable it the same way + // the admin UI toggle does; the projects migration smoke needs the link. + const masterKey = process.env.LITELLM_MASTER_KEY || "sk-1234"; + const api = await request.newContext(); + const settingsRes = await api.patch(`http://localhost:4000${rootPath}/update/ui_settings`, { + headers: { Authorization: `Bearer ${masterKey}` }, + data: { enable_projects_ui: true }, + }); + if (!settingsRes.ok()) { + throw new Error(`Enabling enable_projects_ui failed (${settingsRes.status()}): ${await settingsRes.text()}`); + } + await api.dispose(); + for (const role of Object.values(Role)) { const { email, password } = users[role]; const storagePath = STORAGE_PATHS[role]; diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 358064c00af..8f865915e41 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -472,22 +472,22 @@ "count": 4 } }, - "src/components/Projects/ProjectDetailsPage.tsx": { + "src/app/(dashboard)/projects/components/ProjectDetailsPage.tsx": { "no-restricted-imports": { "count": 1 } }, - "src/components/Projects/ProjectKeysSection.tsx": { + "src/app/(dashboard)/projects/components/ProjectKeysSection.tsx": { "react-hooks/set-state-in-effect": { "count": 1 } }, - "src/components/Projects/ProjectModals/ProjectBaseForm.tsx": { + "src/app/(dashboard)/projects/components/ProjectModals/ProjectBaseForm.tsx": { "react-hooks/set-state-in-effect": { "count": 2 } }, - "src/components/Projects/ProjectsPage.tsx": { + "src/app/(dashboard)/projects/components/ProjectsPage.tsx": { "react-hooks/set-state-in-effect": { "count": 1 } diff --git a/ui/litellm-dashboard/src/components/AccessGroups/AccessGroupsDetailsPage.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/components/AccessGroupsDetailsPage.test.tsx similarity index 99% rename from ui/litellm-dashboard/src/components/AccessGroups/AccessGroupsDetailsPage.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/access-groups/components/AccessGroupsDetailsPage.test.tsx index ee8a8d0ffc5..cf41f623fd6 100644 --- a/ui/litellm-dashboard/src/components/AccessGroups/AccessGroupsDetailsPage.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/components/AccessGroupsDetailsPage.test.tsx @@ -3,7 +3,7 @@ import { AccessGroupResponse } from "@/app/(dashboard)/hooks/accessGroups/useAcc import { screen } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; -import { renderWithProviders } from "../../../tests/test-utils"; +import { renderWithProviders } from "../../../../../tests/test-utils"; import { AccessGroupDetail } from "./AccessGroupsDetailsPage"; vi.mock("@/app/(dashboard)/hooks/accessGroups/useAccessGroupDetails"); diff --git a/ui/litellm-dashboard/src/components/AccessGroups/AccessGroupsDetailsPage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/components/AccessGroupsDetailsPage.tsx similarity index 99% rename from ui/litellm-dashboard/src/components/AccessGroups/AccessGroupsDetailsPage.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/access-groups/components/AccessGroupsDetailsPage.tsx index ae0cd8cd61b..72a89093bdb 100644 --- a/ui/litellm-dashboard/src/components/AccessGroups/AccessGroupsDetailsPage.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/components/AccessGroupsDetailsPage.tsx @@ -17,7 +17,7 @@ import { } from "antd"; import { ArrowLeftIcon, BotIcon, EditIcon, KeyIcon, LayersIcon, ServerIcon, UsersIcon } from "lucide-react"; import { useState } from "react"; -import DefaultProxyAdminTag from "../common_components/DefaultProxyAdminTag"; +import DefaultProxyAdminTag from "@/components/common_components/DefaultProxyAdminTag"; import { AccessGroupEditModal } from "./AccessGroupsModal/AccessGroupEditModal"; const { Title, Text } = Typography; diff --git a/ui/litellm-dashboard/src/components/AccessGroups/AccessGroupsModal/AccessGroupBaseForm.tsx b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/components/AccessGroupsModal/AccessGroupBaseForm.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/AccessGroups/AccessGroupsModal/AccessGroupBaseForm.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/access-groups/components/AccessGroupsModal/AccessGroupBaseForm.tsx diff --git a/ui/litellm-dashboard/src/components/AccessGroups/AccessGroupsModal/AccessGroupCreateModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/components/AccessGroupsModal/AccessGroupCreateModal.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/AccessGroups/AccessGroupsModal/AccessGroupCreateModal.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/access-groups/components/AccessGroupsModal/AccessGroupCreateModal.tsx diff --git a/ui/litellm-dashboard/src/components/AccessGroups/AccessGroupsModal/AccessGroupEditModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/components/AccessGroupsModal/AccessGroupEditModal.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/AccessGroups/AccessGroupsModal/AccessGroupEditModal.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/access-groups/components/AccessGroupsModal/AccessGroupEditModal.tsx diff --git a/ui/litellm-dashboard/src/components/AccessGroups/AccessGroupsPage.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/components/AccessGroupsPage.test.tsx similarity index 99% rename from ui/litellm-dashboard/src/components/AccessGroups/AccessGroupsPage.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/access-groups/components/AccessGroupsPage.test.tsx index d50811949f5..7c8aaa2b785 100644 --- a/ui/litellm-dashboard/src/components/AccessGroups/AccessGroupsPage.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/components/AccessGroupsPage.test.tsx @@ -65,7 +65,7 @@ vi.mock("./AccessGroupsModal/AccessGroupCreateModal", () => ({ ) : null, })); -vi.mock("../common_components/IconActionButton/TableIconActionButtons/TableIconActionButton", () => ({ +vi.mock("@/components/common_components/IconActionButton/TableIconActionButtons/TableIconActionButton", () => ({ default: ({ variant, tooltipText, onClick }: { variant: string; tooltipText: string; onClick: () => void }) => (
)} + {endpointData.timeout !== undefined && endpointData.timeout !== null && ( +
+ Request Timeout +
{endpointData.timeout}s
+
+ )}
Authentication Required {endpointData.auth ? "Yes" : "No"} diff --git a/ui/litellm-dashboard/src/components/pass_through_settings.tsx b/ui/litellm-dashboard/src/components/pass_through_settings.tsx index 3d53de70727..66b41f092c5 100644 --- a/ui/litellm-dashboard/src/components/pass_through_settings.tsx +++ b/ui/litellm-dashboard/src/components/pass_through_settings.tsx @@ -38,6 +38,7 @@ export interface passThroughItem { headers: object; include_subpath?: boolean; cost_per_request?: number; + timeout?: number; auth?: boolean; methods?: string[]; guardrails?: Record; diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 3786b6e728b..4797e41a62b 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -22203,6 +22203,11 @@ export interface components { * @description Set-up pass-through endpoints for provider-specific endpoints. Docs - https://docs.litellm.ai/docs/proxy/pass_through */ pass_through_endpoints?: components["schemas"]["PassThroughGenericEndpoint"][] | null; + /** + * Pass Through Request Timeout + * @description Default upstream request timeout in seconds for native and custom pass-through endpoints that use pass_through_request. Defaults to 600 when unset. + */ + pass_through_request_timeout?: number | null; /** * Reject Clientside Metadata Tags * @description When set to True, rejects requests that contain client-side 'metadata.tags' to prevent users from influencing budgets by sending different tags. Tags can only be inherited from the API key metadata. @@ -27993,6 +27998,11 @@ export interface components { * @description The URL to which requests for this path should be forwarded. */ target: string; + /** + * Timeout + * @description Upstream request timeout in seconds for this pass-through endpoint. If unset, uses general_settings.pass_through_request_timeout (default 600). + */ + timeout?: number | null; }; /** * PassThroughGuardrailSettings From 729b005e4e00ef43f6187baf7fcd8f2737c764b8 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Fri, 12 Jun 2026 20:19:30 +0530 Subject: [PATCH 085/185] fix(google_genai): preserve complete SSE events in Vertex/Gemini image streaming (#30270) * fix(google_genai): preserve complete SSE events in image streaming Use iter_lines/aiter_lines instead of byte chunking so large inlineData base64 payloads from Vertex/Gemini streamGenerateContent are not split across events, which caused truncated JSON and SDK parse failures. Co-authored-by: Cursor * fix(google_genai): buffer SSE lines until event delimiter Assemble multi-field SSE events on blank-line boundaries instead of terminating each field line individually. Co-authored-by: Cursor * fix(tests): update google_ai_studio mocks from aiter_bytes to aiter_lines Streaming iterator was changed to use iter_lines/aiter_lines instead of iter_bytes/aiter_bytes. Update the two mocked streaming responses in test_google_ai_studio.py to match. Co-authored-by: Cursor --------- Co-authored-by: Cursor --- litellm/google_genai/streaming_iterator.py | 54 ++++++-- .../test_google_genai_streaming_iterator.py | 128 ++++++++++++++++++ .../test_google_ai_studio.py | 33 ++--- 3 files changed, 186 insertions(+), 29 deletions(-) create mode 100644 tests/test_litellm/google_genai/test_google_genai_streaming_iterator.py diff --git a/litellm/google_genai/streaming_iterator.py b/litellm/google_genai/streaming_iterator.py index 3e97b480779..a8d0e5976f0 100644 --- a/litellm/google_genai/streaming_iterator.py +++ b/litellm/google_genai/streaming_iterator.py @@ -18,6 +18,42 @@ else: GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ = PassThroughEndpointLogging() +def _encode_google_genai_sse_event(event_lines: List[str]) -> bytes: + return ("\n".join(event_lines) + "\n\n").encode("utf-8") + + +def _next_google_genai_sse_chunk(line_iter) -> bytes: + event_lines: List[str] = [] + while True: + try: + line = next(line_iter) + except StopIteration: + if event_lines: + return _encode_google_genai_sse_event(event_lines) + raise + if line == "": + if event_lines: + return _encode_google_genai_sse_event(event_lines) + continue + event_lines.append(line) + + +async def _anext_google_genai_sse_chunk(line_iter) -> bytes: + event_lines: List[str] = [] + while True: + try: + line = await line_iter.__anext__() + except StopAsyncIteration: + if event_lines: + return _encode_google_genai_sse_event(event_lines) + raise + if line == "": + if event_lines: + return _encode_google_genai_sse_event(event_lines) + continue + event_lines.append(line) + + class BaseGoogleGenAIGenerateContentStreamingIterator: """ Base class for Google GenAI Generate Content streaming iterators that provides common logic @@ -91,18 +127,17 @@ class GoogleGenAIGenerateContentStreamingIterator( self.generate_content_provider_config = generate_content_provider_config self.litellm_metadata = litellm_metadata self.custom_llm_provider = custom_llm_provider - # Store the iterator once to avoid multiple stream consumption - self.stream_iterator = response.iter_bytes() + # Gemini streamGenerateContent uses SSE line framing; iter_lines keeps + # large inlineData payloads (e.g. image/jpeg) intact within one event. + self.stream_iterator = response.iter_lines() def __iter__(self): return self def __next__(self): try: - # Get the next chunk from the stored iterator - chunk = next(self.stream_iterator) + chunk = _next_google_genai_sse_chunk(self.stream_iterator) self.collected_chunks.append(chunk) - # Just yield raw bytes return chunk except StopIteration: raise StopIteration @@ -147,18 +182,17 @@ class AsyncGoogleGenAIGenerateContentStreamingIterator( self.generate_content_provider_config = generate_content_provider_config self.litellm_metadata = litellm_metadata self.custom_llm_provider = custom_llm_provider - # Store the async iterator once to avoid multiple stream consumption - self.stream_iterator = response.aiter_bytes() + # Gemini streamGenerateContent uses SSE line framing; aiter_lines keeps + # large inlineData payloads (e.g. image/jpeg) intact within one event. + self.stream_iterator = response.aiter_lines() def __aiter__(self): return self async def __anext__(self): try: - # Get the next chunk from the stored async iterator - chunk = await self.stream_iterator.__anext__() + chunk = await _anext_google_genai_sse_chunk(self.stream_iterator) self.collected_chunks.append(chunk) - # Just yield raw bytes return chunk except StopAsyncIteration: await self._handle_async_streaming_logging() diff --git a/tests/test_litellm/google_genai/test_google_genai_streaming_iterator.py b/tests/test_litellm/google_genai/test_google_genai_streaming_iterator.py new file mode 100644 index 00000000000..d74a05ec59c --- /dev/null +++ b/tests/test_litellm/google_genai/test_google_genai_streaming_iterator.py @@ -0,0 +1,128 @@ +import json +from unittest.mock import MagicMock + +import pytest + +from litellm.google_genai.streaming_iterator import ( + AsyncGoogleGenAIGenerateContentStreamingIterator, + GoogleGenAIGenerateContentStreamingIterator, +) +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + + +def _large_inline_data_event() -> str: + payload = { + "candidates": [ + { + "content": { + "parts": [ + { + "inlineData": { + "mimeType": "image/jpeg", + "data": "A" * 20000, + } + } + ] + } + } + ] + } + return f"data: {json.dumps(payload)}" + + +@pytest.mark.asyncio +async def test_async_streaming_iterator_yields_complete_sse_events(): + """Large inlineData must not be split across byte-chunk boundaries.""" + mock_response = MagicMock() + + async def _aiter_lines(): + yield _large_inline_data_event() + + mock_response.aiter_lines = _aiter_lines + + iterator = AsyncGoogleGenAIGenerateContentStreamingIterator( + response=mock_response, + model="gemini-3.1-flash-image-preview", + logging_obj=MagicMock(spec=LiteLLMLoggingObj), + generate_content_provider_config=MagicMock(), + litellm_metadata={}, + custom_llm_provider="gemini", + ) + + chunk = await iterator.__anext__() + assert chunk.startswith(b"data: ") + assert chunk.endswith(b"\n\n") + assert ( + json.loads(chunk[len(b"data: ") : -2])["candidates"][0]["content"]["parts"][0][ + "inlineData" + ]["mimeType"] + == "image/jpeg" + ) + + +def test_sync_streaming_iterator_yields_complete_sse_events(): + mock_response = MagicMock() + mock_response.iter_lines.return_value = iter([_large_inline_data_event()]) + + iterator = GoogleGenAIGenerateContentStreamingIterator( + response=mock_response, + model="gemini-3.1-flash-image-preview", + logging_obj=MagicMock(spec=LiteLLMLoggingObj), + generate_content_provider_config=MagicMock(), + litellm_metadata={}, + custom_llm_provider="gemini", + ) + + chunk = next(iterator) + assert chunk.startswith(b"data: ") + assert chunk.endswith(b"\n\n") + assert json.loads(chunk[len(b"data: ") : -2])["candidates"][0]["content"]["parts"][ + 0 + ]["inlineData"]["data"].startswith("A") + + +@pytest.mark.asyncio +async def test_async_streaming_iterator_preserves_multi_field_sse_event(): + mock_response = MagicMock() + + async def _aiter_lines(): + yield "event: message" + yield 'data: {"text":"hi"}' + yield "" + + mock_response.aiter_lines = _aiter_lines + + iterator = AsyncGoogleGenAIGenerateContentStreamingIterator( + response=mock_response, + model="gemini-test", + logging_obj=MagicMock(spec=LiteLLMLoggingObj), + generate_content_provider_config=MagicMock(), + litellm_metadata={}, + custom_llm_provider="gemini", + ) + + chunk = await iterator.__anext__() + assert chunk == b'event: message\ndata: {"text":"hi"}\n\n' + + +@pytest.mark.asyncio +async def test_async_streaming_iterator_forwards_sse_comment_events(): + mock_response = MagicMock() + + async def _aiter_lines(): + yield ": keepalive" + yield "" + + mock_response.aiter_lines = _aiter_lines + + iterator = AsyncGoogleGenAIGenerateContentStreamingIterator( + response=mock_response, + model="gemini-test", + logging_obj=MagicMock(spec=LiteLLMLoggingObj), + generate_content_provider_config=MagicMock(), + litellm_metadata={}, + custom_llm_provider="gemini", + ) + + chunk = await iterator.__anext__() + assert chunk == b": keepalive\n\n" diff --git a/tests/unified_google_tests/test_google_ai_studio.py b/tests/unified_google_tests/test_google_ai_studio.py index afe237a4e5b..3e40fa41089 100644 --- a/tests/unified_google_tests/test_google_ai_studio.py +++ b/tests/unified_google_tests/test_google_ai_studio.py @@ -74,12 +74,6 @@ async def test_mock_stream_generate_content_with_tools(): }, } - # Convert to bytes as expected by the streaming iterator - raw_chunks = [ - f"data: {json.dumps(mock_response_chunk)}\n\n".encode(), - b"data: [DONE]\n\n", - ] - # Mock the HTTP handler with unittest.mock.patch( "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", @@ -90,12 +84,15 @@ async def test_mock_stream_generate_content_with_tools(): mock_response.status_code = 200 mock_response.headers = {"content-type": "application/json"} - # Mock the aiter_bytes method to return our chunks as bytes - async def mock_aiter_bytes(): - for chunk in raw_chunks: - yield chunk + # Mock aiter_lines: yield one line at a time (no trailing newlines), + # with a blank line between events, matching httpx aiter_lines behaviour. + async def mock_aiter_lines(): + yield f"data: {json.dumps(mock_response_chunk)}" + yield "" + yield "data: [DONE]" + yield "" - mock_response.aiter_bytes = mock_aiter_bytes + mock_response.aiter_lines = mock_aiter_lines mock_post.return_value = mock_response print( @@ -328,9 +325,6 @@ async def test_validate_post_request_parameters(): } ] - # Mock response for the HTTP request - raw_chunks = [b"data: [DONE]\n\n"] - # Mock the HTTP handler to capture the request with unittest.mock.patch( "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", @@ -341,12 +335,13 @@ async def test_validate_post_request_parameters(): mock_response.status_code = 200 mock_response.headers = {"content-type": "application/json"} - # Mock the aiter_bytes method - async def mock_aiter_bytes(): - for chunk in raw_chunks: - yield chunk + # Mock aiter_lines: yield one line at a time (no trailing newlines), + # with a blank line between events, matching httpx aiter_lines behaviour. + async def mock_aiter_lines(): + yield "data: [DONE]" + yield "" - mock_response.aiter_bytes = mock_aiter_bytes + mock_response.aiter_lines = mock_aiter_lines mock_post.return_value = mock_response print("\n--- Testing POST request parameters validation ---") From 7d1f68e72a9bc0aa0f9b69d8f2a8a9647b49f3be Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Fri, 12 Jun 2026 22:19:09 +0530 Subject: [PATCH 086/185] fix(proxy): populate access_via_team_ids on /v1/model/info (#30274) * fix(proxy): populate access_via_team_ids on /v1/model/info Team metadata enrichment previously only ran on /v2/model/info with include_team_models=true, leaving /v1/model/info without access_via_team_ids for project model-picker flows. Co-authored-by: Cursor * docs(dashboard): sync OpenAPI schema for /v1/model/info query params Add include_team_models and teamId to the generated schema for /model/info and /v1/model/info after the proxy endpoint gained team-access filtering. Co-authored-by: Cursor * fix(proxy): always return direct_access on /v1/model/info Set direct_access to true or false on every enriched model so clients can filter without treating a missing field as ambiguous. Co-authored-by: Cursor * perf(proxy): fail fast when teamId is set without a connected DB on /v1/model/info Raise the db_not_connected error before building, enriching, and translating the model list instead of after, so a teamId query against a proxy with no database no longer wastes the full enrichment pipeline. * fix(proxy): fail fast when include_team_models is set without a database include_team_models=True relies on _populate_team_access_on_models to set direct_access/access_via_team_ids, which only runs when a database is connected. Without one, _filter_models_to_user_accessible discarded every model and the endpoint returned an empty list with HTTP 200. Mirror the teamId guard so the request fails fast with a clear db_not_connected error before any model-list work. * fix(proxy): populate direct_access on single-model /model/info lookup The /v1/model/info list path populates model_info.direct_access (and access_via_team_ids) when a database is connected, but the litellm_model_id single-model lookup returned early without it. This made the two endpoints disagree, breaking the parity assertion in test_get_specific_model. Run the same population on the single-model path so both responses match. * fix(proxy): apply no-DB fast-fail before litellm_model_id branch The teamId/include_team_models no-DB guard sat after the litellm_model_id early return, so ?litellm_model_id=X&teamId=Y with no DB returned 200 with unpopulated access fields instead of the 500 raised on every other path. Move the guard ahead of the branch so the fast-fail is uniform. * fix(proxy): apply teamId/include_team_models filters on single-model lookup The litellm_model_id early-return branch in model_info_v1 populated the team access fields but returned before the teamId and include_team_models filters ran, so a single-model lookup surfaced the deployment regardless of team access when the DB was connected. Run both filters on the single-model list before returning so the documented query params behave the same with and without litellm_model_id. --------- Co-authored-by: Cursor Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 116 ++++++-- .../test_team_model_name_translation.py | 250 ++++++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 18 ++ 3 files changed, 368 insertions(+), 16 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index ea2aa8fb01e..1d8cbb6fe0a 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11031,16 +11031,26 @@ def get_direct_access_models( return direct_access_models -async def get_all_team_and_direct_access_models( +def _filter_models_to_user_accessible(all_models: List[Dict]) -> List[Dict]: + """Keep only deployments the caller can use via direct access or team membership.""" + return [ + _model + for _model in all_models + if _model.get("model_info", {}).get("direct_access", False) + or _model.get("model_info", {}).get("access_via_team_ids", []) + ] + + +async def _populate_team_access_on_models( user_api_key_dict: UserAPIKeyAuth, prisma_client: PrismaClient, llm_router: Router, all_models: List[Dict], ) -> List[Dict]: """ - Get all models across all teams user is in. + Populate `model_info.access_via_team_ids` and `model_info.direct_access` + without filtering the model list. """ - user_teams: Optional[Union[List[str], Literal["*"]]] = None direct_access_models: List[str] = [] if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: @@ -11059,7 +11069,6 @@ async def get_all_team_and_direct_access_models( user_db_object=user_object, llm_router=llm_router, ) - ## ADD ACCESS_VIA_TEAM_IDS TO ALL MODELS if user_teams is not None: team_models = await get_all_team_models( user_teams=user_teams, @@ -11082,23 +11091,33 @@ async def get_all_team_and_direct_access_models( model_id, [] ) - ## ADD DIRECT_ACCESS TO RELEVANT MODELS - + direct_access_model_ids = set(direct_access_models) for _model in all_models: model_id = _model.get("model_info", {}).get("id", None) - if model_id is not None and model_id in direct_access_models: - _model["model_info"]["direct_access"] = True + if model_id is not None: + _model["model_info"]["direct_access"] = model_id in direct_access_model_ids - ## FILTER OUT MODELS THAT ARE NOT IN DIRECT_ACCESS_MODELS OR ACCESS_VIA_TEAM_IDS - only show user models they can call - all_models = [ - _model - for _model in all_models - if _model.get("model_info", {}).get("direct_access", False) - or _model.get("model_info", {}).get("access_via_team_ids", []) - ] return all_models +async def get_all_team_and_direct_access_models( + user_api_key_dict: UserAPIKeyAuth, + prisma_client: PrismaClient, + llm_router: Router, + all_models: List[Dict], +) -> List[Dict]: + """ + Get all models across all teams user is in. + """ + all_models = await _populate_team_access_on_models( + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + llm_router=llm_router, + all_models=all_models, + ) + return _filter_models_to_user_accessible(all_models) + + def _enrich_model_info_with_litellm_data( model: Dict[str, Any], debug: bool = False, llm_router: Optional[Router] = None ) -> Dict[str, Any]: @@ -12633,6 +12652,14 @@ def _get_proxy_model_info(model: dict) -> dict: async def model_info_v1( # noqa: PLR0915 user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), litellm_model_id: Optional[str] = None, + include_team_models: Optional[bool] = fastapi.Query( + False, + description="When true, filter to deployments the caller can use via direct access or team membership.", + ), + teamId: Optional[str] = fastapi.Query( + None, + description="Filter models by team ID. Returns models with direct_access=True or teamId in access_via_team_ids", + ), ): """ Provides more info about each model in /models, including config.yaml descriptions (except api key and api base) @@ -12642,6 +12669,11 @@ async def model_info_v1( # noqa: PLR0915 - When litellm_model_id is passed, it will return the info for that specific model - When litellm_model_id is not passed, it will return the info for all models + - include_team_models: When true, filter to deployments the caller can use (same as /v2/model/info). + - teamId: Filter to models accessible by the given team. + + Each model in the list response includes `model_info.access_via_team_ids` and + `model_info.direct_access` when the proxy database is connected. Returns: Returns a dictionary containing information about each model. @@ -12668,6 +12700,12 @@ async def model_info_v1( # noqa: PLR0915 """ global llm_model_list, general_settings, user_config_file_path, proxy_config, llm_router, user_model + # Unit tests call this handler directly; FastAPI normally resolves Query defaults. + if not isinstance(include_team_models, bool): + include_team_models = False + if not isinstance(teamId, str): + teamId = None + if user_model is not None: # user is trying to get specific model from litellm router try: @@ -12704,6 +12742,14 @@ async def model_info_v1( # noqa: PLR0915 }, ) + if prisma_client is None and ( + include_team_models or (teamId is not None and teamId.strip()) + ): + raise HTTPException( + status_code=500, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) + if litellm_model_id is not None: # user is trying to get specific model from litellm router deployment_info = llm_router.get_deployment(model_id=litellm_model_id) @@ -12717,7 +12763,25 @@ async def model_info_v1( # noqa: PLR0915 _deployment_info_dict = _get_proxy_model_info( model=deployment_info.model_dump(exclude_none=True) ) - return {"data": [_deployment_info_dict]} + single_model_list: List[dict] = [_deployment_info_dict] + if prisma_client is not None: + single_model_list = await _populate_team_access_on_models( + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + llm_router=llm_router, + all_models=single_model_list, + ) + if include_team_models: + single_model_list = _filter_models_to_user_accessible(single_model_list) + if teamId is not None and teamId.strip(): + single_model_list = await _filter_models_by_team_id( + all_models=single_model_list, + team_id=teamId.strip(), + prisma_client=prisma_client, + llm_router=llm_router, + user_api_key_dict=user_api_key_dict, + ) + return {"data": single_model_list} # Return router deployments (same source as /v2/model/info), not wildcard- # expanded model names from get_complete_model_list(). Team-scoped rows @@ -12749,6 +12813,17 @@ async def model_info_v1( # noqa: PLR0915 ) ] + if prisma_client is not None: + all_models = await _populate_team_access_on_models( + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + llm_router=llm_router, + all_models=all_models, + ) + + if include_team_models: + all_models = _filter_models_to_user_accessible(all_models) + all_models = [ _translate_model_name_for_response( _enrich_model_info_with_litellm_data(model=model, llm_router=llm_router) @@ -12756,6 +12831,15 @@ async def model_info_v1( # noqa: PLR0915 for model in all_models ] + if teamId is not None and teamId.strip(): + all_models = await _filter_models_by_team_id( + all_models=all_models, + team_id=teamId.strip(), + prisma_client=cast(PrismaClient, prisma_client), + llm_router=llm_router, + user_api_key_dict=user_api_key_dict, + ) + verbose_proxy_logger.debug("all_models: %s", all_models) return {"data": all_models} diff --git a/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py b/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py index 9757999c85e..6a8e0d15d8b 100644 --- a/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py +++ b/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py @@ -279,6 +279,11 @@ async def test_model_info_v1_unrestricted_key_hides_other_team_byok(monkeypatch) prisma_client = MagicMock() caller_user_row = MagicMock() caller_user_row.teams = ["team-abc-123"] + caller_user_row.model_dump.return_value = { + "user_id": "user-1", + "teams": ["team-abc-123"], + "models": [], + } prisma_client.db.litellm_usertable.find_unique = AsyncMock( return_value=caller_user_row ) @@ -287,6 +292,7 @@ async def test_model_info_v1_unrestricted_key_hides_other_team_byok(monkeypatch) monkeypatch.setattr(ps, "llm_model_list", router.model_list) monkeypatch.setattr(ps, "llm_router", router) monkeypatch.setattr(ps, "prisma_client", prisma_client) + monkeypatch.setattr(ps, "get_all_team_models", AsyncMock(return_value={})) monkeypatch.setattr( ps, "_enrich_model_info_with_litellm_data", lambda model, **kw: model ) @@ -343,3 +349,247 @@ async def test_model_info_v1_service_key_hides_all_team_byok(monkeypatch): resp = await ps.model_info_v1(user_api_key_dict=caller, litellm_model_id=None) assert [m["model_info"]["id"] for m in resp["data"]] == ["global-id-1"] + + +@pytest.mark.asyncio +async def test_model_info_v1_populates_access_via_team_ids(monkeypatch): + """`/v1/model/info` must populate access_via_team_ids when the DB is connected.""" + team_id = "team-abc-123" + team_row = _team_row() + global_row = { + "model_name": "gpt-4o", + "litellm_params": {"model": "gpt-4o"}, + "model_info": {"id": "global-id-1", "db_model": False}, + } + router = MagicMock() + router.model_list = [team_row, global_row] + router.get_model_names.return_value = ["gpt-4o", "team-claude-sonnet"] + router.get_model_access_groups.return_value = {} + router.get_model_ids.return_value = ["global-id-1"] + + prisma_client = MagicMock() + + async def _fake_populate(**kwargs): + for model in kwargs["all_models"]: + model_id = model["model_info"]["id"] + if model_id == "byok-id-1": + model["model_info"]["access_via_team_ids"] = [team_id] + model["model_info"]["direct_access"] = False + elif model_id == "global-id-1": + model["model_info"]["direct_access"] = True + return kwargs["all_models"] + + monkeypatch.setattr(ps, "user_model", None) + monkeypatch.setattr(ps, "llm_model_list", router.model_list) + monkeypatch.setattr(ps, "llm_router", router) + monkeypatch.setattr(ps, "prisma_client", prisma_client) + monkeypatch.setattr(ps, "_populate_team_access_on_models", _fake_populate) + monkeypatch.setattr( + ps, "_enrich_model_info_with_litellm_data", lambda model, **kw: model + ) + + admin = UserAPIKeyAuth( + user_id="u", user_role=LitellmUserRoles.PROXY_ADMIN, team_models=[] + ) + resp = await ps.model_info_v1(user_api_key_dict=admin, litellm_model_id=None) + + by_id = {m["model_info"]["id"]: m for m in resp["data"]} + assert by_id["byok-id-1"]["model_info"]["access_via_team_ids"] == [team_id] + assert by_id["byok-id-1"]["model_info"]["direct_access"] is False + assert by_id["global-id-1"]["model_info"]["direct_access"] is True + + +@pytest.mark.asyncio +async def test_populate_team_access_sets_direct_access_false_by_default(monkeypatch): + """Team-accessible models without direct access must return direct_access=false.""" + team_row = _team_row() + global_row = { + "model_name": "gpt-4o", + "litellm_params": {"model": "gpt-4o"}, + "model_info": {"id": "global-id-1", "db_model": False}, + } + router = MagicMock() + router.get_model_ids.return_value = ["global-id-1"] + monkeypatch.setattr( + ps, + "get_all_team_models", + AsyncMock(return_value={"byok-id-1": ["team-abc-123"]}), + ) + + admin = UserAPIKeyAuth( + user_id="u", user_role=LitellmUserRoles.PROXY_ADMIN, team_models=[] + ) + result = await ps._populate_team_access_on_models( + user_api_key_dict=admin, + prisma_client=MagicMock(), + llm_router=router, + all_models=[team_row, global_row], + ) + + by_id = {m["model_info"]["id"]: m for m in result} + assert by_id["byok-id-1"]["model_info"]["direct_access"] is False + assert by_id["global-id-1"]["model_info"]["direct_access"] is True + + +@pytest.mark.asyncio +async def test_model_info_v1_team_id_without_db_fails_fast(monkeypatch): + """`teamId` without a connected DB raises 500 before any enrichment work runs.""" + router = MagicMock() + router.model_list = [_team_row()] + + enrich_spy = MagicMock(side_effect=lambda model, **kw: model) + + monkeypatch.setattr(ps, "user_model", None) + monkeypatch.setattr(ps, "llm_model_list", router.model_list) + monkeypatch.setattr(ps, "llm_router", router) + monkeypatch.setattr(ps, "prisma_client", None) + monkeypatch.setattr(ps, "_enrich_model_info_with_litellm_data", enrich_spy) + + admin = UserAPIKeyAuth( + user_id="u", user_role=LitellmUserRoles.PROXY_ADMIN, team_models=[] + ) + + with pytest.raises(ps.HTTPException) as exc_info: + await ps.model_info_v1( + user_api_key_dict=admin, litellm_model_id=None, teamId="team-abc-123" + ) + + assert exc_info.value.status_code == 500 + assert "DB not connected" in exc_info.value.detail["error"] + enrich_spy.assert_not_called() + + +@pytest.mark.asyncio +async def test_model_info_v1_include_team_models_without_db_fails_fast(monkeypatch): + """`include_team_models` without a connected DB raises 500 instead of silently + returning an empty list (the access fields can only be populated from the DB).""" + router = MagicMock() + router.model_list = [_team_row()] + + enrich_spy = MagicMock(side_effect=lambda model, **kw: model) + + monkeypatch.setattr(ps, "user_model", None) + monkeypatch.setattr(ps, "llm_model_list", router.model_list) + monkeypatch.setattr(ps, "llm_router", router) + monkeypatch.setattr(ps, "prisma_client", None) + monkeypatch.setattr(ps, "_enrich_model_info_with_litellm_data", enrich_spy) + + admin = UserAPIKeyAuth( + user_id="u", user_role=LitellmUserRoles.PROXY_ADMIN, team_models=[] + ) + + with pytest.raises(ps.HTTPException) as exc_info: + await ps.model_info_v1( + user_api_key_dict=admin, litellm_model_id=None, include_team_models=True + ) + + assert exc_info.value.status_code == 500 + assert "DB not connected" in exc_info.value.detail["error"] + enrich_spy.assert_not_called() + + +@pytest.mark.asyncio +async def test_model_info_v1_litellm_model_id_team_id_without_db_fails_fast( + monkeypatch, +): + """`litellm_model_id` + `teamId` without a connected DB must raise 500 too, not + return 200 with a model dict missing direct_access/access_via_team_ids.""" + router = MagicMock() + router.model_list = [_team_row()] + + monkeypatch.setattr(ps, "user_model", None) + monkeypatch.setattr(ps, "llm_model_list", router.model_list) + monkeypatch.setattr(ps, "llm_router", router) + monkeypatch.setattr(ps, "prisma_client", None) + + admin = UserAPIKeyAuth( + user_id="u", user_role=LitellmUserRoles.PROXY_ADMIN, team_models=[] + ) + + with pytest.raises(ps.HTTPException) as exc_info: + await ps.model_info_v1( + user_api_key_dict=admin, + litellm_model_id="byok-id-1", + teamId="team-abc-123", + ) + + assert exc_info.value.status_code == 500 + assert "DB not connected" in exc_info.value.detail["error"] + router.get_deployment.assert_not_called() + + +@pytest.mark.asyncio +async def test_model_info_v1_litellm_model_id_include_team_models_filters_inaccessible( + monkeypatch, +): + """`litellm_model_id` + `include_team_models` must drop a model the caller cannot + use instead of returning it unconditionally from the single-model lookup.""" + team_row = _team_row() + + router = MagicMock() + deployment = MagicMock() + deployment.model_dump.return_value = team_row + router.get_deployment.return_value = deployment + + async def _fake_populate(**kwargs): + for model in kwargs["all_models"]: + model["model_info"]["direct_access"] = False + model["model_info"]["access_via_team_ids"] = [] + return kwargs["all_models"] + + monkeypatch.setattr(ps, "user_model", None) + monkeypatch.setattr(ps, "llm_model_list", [team_row]) + monkeypatch.setattr(ps, "llm_router", router) + monkeypatch.setattr(ps, "prisma_client", MagicMock()) + monkeypatch.setattr(ps, "_get_proxy_model_info", lambda model: team_row) + monkeypatch.setattr(ps, "_populate_team_access_on_models", _fake_populate) + + caller = UserAPIKeyAuth( + user_id="u", user_role=LitellmUserRoles.INTERNAL_USER, team_models=[] + ) + resp = await ps.model_info_v1( + user_api_key_dict=caller, + litellm_model_id="byok-id-1", + include_team_models=True, + ) + + assert resp["data"] == [] + + +@pytest.mark.asyncio +async def test_model_info_v1_litellm_model_id_team_id_applies_team_filter(monkeypatch): + """`litellm_model_id` + `teamId` must run the teamId filter on the single model + rather than returning it regardless of the team's access.""" + team_row = _team_row() + + router = MagicMock() + deployment = MagicMock() + deployment.model_dump.return_value = team_row + router.get_deployment.return_value = deployment + + async def _fake_populate(**kwargs): + return kwargs["all_models"] + + team_filter = AsyncMock(return_value=[]) + + monkeypatch.setattr(ps, "user_model", None) + monkeypatch.setattr(ps, "llm_model_list", [team_row]) + monkeypatch.setattr(ps, "llm_router", router) + monkeypatch.setattr(ps, "prisma_client", MagicMock()) + monkeypatch.setattr(ps, "_get_proxy_model_info", lambda model: team_row) + monkeypatch.setattr(ps, "_populate_team_access_on_models", _fake_populate) + monkeypatch.setattr(ps, "_filter_models_by_team_id", team_filter) + + admin = UserAPIKeyAuth( + user_id="u", user_role=LitellmUserRoles.PROXY_ADMIN, team_models=[] + ) + resp = await ps.model_info_v1( + user_api_key_dict=admin, + litellm_model_id="byok-id-1", + teamId="other-team", + ) + + assert resp["data"] == [] + team_filter.assert_awaited_once() + assert team_filter.await_args.kwargs["team_id"] == "other-team" + assert team_filter.await_args.kwargs["all_models"] == [team_row] diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 4797e41a62b..2e24e83cefa 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -7338,6 +7338,11 @@ export interface paths { * * - When litellm_model_id is passed, it will return the info for that specific model * - When litellm_model_id is not passed, it will return the info for all models + * - include_team_models: When true, filter to deployments the caller can use (same as /v2/model/info). + * - teamId: Filter to models accessible by the given team. + * + * Each model in the list response includes `model_info.access_via_team_ids` and + * `model_info.direct_access` when the proxy database is connected. * * Returns: * Returns a dictionary containing information about each model. @@ -16565,6 +16570,11 @@ export interface paths { * * - When litellm_model_id is passed, it will return the info for that specific model * - When litellm_model_id is not passed, it will return the info for all models + * - include_team_models: When true, filter to deployments the caller can use (same as /v2/model/info). + * - teamId: Filter to models accessible by the given team. + * + * Each model in the list response includes `model_info.access_via_team_ids` and + * `model_info.direct_access` when the proxy database is connected. * * Returns: * Returns a dictionary containing information about each model. @@ -42440,6 +42450,10 @@ export interface operations { parameters: { query?: { litellm_model_id?: string | null; + /** @description When true, filter to deployments the caller can use via direct access or team membership. */ + include_team_models?: boolean | null; + /** @description Filter models by team ID. Returns models with direct_access=True or teamId in access_via_team_ids */ + teamId?: string | null; }; header?: never; path?: never; @@ -53705,6 +53719,10 @@ export interface operations { parameters: { query?: { litellm_model_id?: string | null; + /** @description When true, filter to deployments the caller can use via direct access or team membership. */ + include_team_models?: boolean | null; + /** @description Filter models by team ID. Returns models with direct_access=True or teamId in access_via_team_ids */ + teamId?: string | null; }; header?: never; path?: never; From 079c136742f78442ff660aa49b1e39379a32ae6b Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Fri, 12 Jun 2026 22:19:25 +0530 Subject: [PATCH 087/185] chore(oss): litellm oss staging 120626 (#30292) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(bedrock): add bedrock mantle gemma 4 models (#30264) * feat(bedrock): add bedrock mantle gemma 4 models * test(bedrock): harden mantle local cost fixture * feat(responses): enable the responses API for the Tensormesh provider (#30209) * feat(responses): enable the responses API for the Tensormesh provider * Update litellm/llms/openai_like/providers.json Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --------- Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * fix(langfuse_otel): mark LLM spans as generations (#30250) * fix(bedrock): stop stream_chunk_size leaking into invoke request bodies (#30240) stream_chunk_size is a LiteLLM-internal knob for re-chunking the HTTP response stream. The invoke transformations splat optional_params into the provider request body without dropping it, and Bedrock rejects unknown fields, so any bedrock/invoke request that sets the parameter fails with ValidationException: stream_chunk_size: Extra inputs are not permitted. Drop it in the invoke dispatcher (covers cohere, titan, mistral, meta, ai21) and in the Claude messages-format request builder (the route used for bedrock/invoke Anthropic models) * fix(bedrock): stop buffering streamed tool-call argument deltas (#30231) * fix(bedrock): stop buffering streamed tool-call argument deltas Two issues made Bedrock tool-use streaming arrive as a single end-of-stream burst through LiteLLM while plain text streamed fine. First, the anthropic-beta allowlist mapped fine-grained-tool-streaming-2025-05-14 to null for bedrock and bedrock_converse, so the header was silently stripped. Without that beta, Anthropic models on Bedrock buffer tool input server-side and emit all toolUse.input deltas at once (verified against converse-stream and invoke-with-response-stream directly). Bedrock accepts the beta via additionalModelRequestFields.anthropic_beta, so it is now forwarded. Second, the streaming reads re-chunked the AWS event stream with iter_bytes(chunk_size=1024). httpx's ByteChunker only releases full 1024-byte blocks, so the small early events (messageStart, contentBlockStart, first deltas) sat in the buffer until enough bytes accumulated, pushing time-to-first-byte from ~1.4s to ~8.5s on buffered tool-use streams. The default is now no re-chunking; an explicit stream_chunk_size is still honored. * test(bedrock): cover explicit stream_chunk_size on sync invoke path * test(bedrock): cover stream_chunk_size plumbing through converse completion * test(bedrock): cover stream_chunk_size default in legacy BedrockLLM streaming * test(bedrock): merge converse handler tests into existing mapped test file pytest imports test modules by basename in non-package test dirs, so the new tests/test_litellm/llms/bedrock/chat/test_converse_handler.py collided with the pre-existing tests/test_litellm/llms/chat/test_converse_handler.py and broke collection in CI. Move the new tests into the existing file * feat(otel): emit v2 cost breakdown + stamp tracer scope version (#30156) Read the StandardLoggingPayload cost_breakdown into a typed LLMCost on LLMCallSpanData and emit each component under litellm.cost.* (absent components omitted, so spans stay sparse). Stamp litellm.__version__ as the instrumentation scope version so every v2 span carries a deterministic scope.version. Tests under tests/test_litellm/integrations/otel/. * fix(proxy): cancel in-flight upstream LLM request on client disconnect (opt-in) (#30223) * fix(proxy): cancel in-flight upstream LLM request on client disconnect (opt-in) On the non-streaming path, base_process_llm_request awaited the LLM call with no disconnect monitoring; when the HTTP client went away the upstream request kept running until completion or request_timeout (6000s default), holding a backend slot (e.g. a vLLM GPU slot) for output nobody would read Add an opt-in general_settings.cancel_on_disconnect flag, default off, so the default code path is unchanged. When enabled, a receive-based watcher task observes http.disconnect and cancels the asyncio.gather driving the upstream call. The resulting CancelledError is converted to HTTPException 499 only when the disconnect event is set, so server-initiated cancellations still propagate as-is. The 499 then flows through _handle_llm_api_exception like any other failure, meaning post_call_failure_hook still releases max_parallel_requests slots and fires spend and alerting callbacks; it is logged at info level instead of a full traceback Also removes the dead check_request_disconnection helper in proxy_server.py (zero call sites) along with its behavior-pin tests Builds on the receive-based design from #25776 Addresses #13774. Re-fixes #22805 (regressed after the #14295 revert) Co-authored-by: CreateRandom <18438707+CreateRandom@users.noreply.github.com> * fix(proxy): scope 499 quiet logging to disconnects and harden watcher Address the two P2 findings from the Greptile review on #30223. The info-level logging in _log_llm_api_exception now applies only to the disconnect-specific HTTPException (status 499 plus the shared _CLIENT_DISCONNECT_DETAIL message), so any other 499 raised by hooks or guardrails keeps its full traceback. The disconnect watcher now catches exceptions from request.receive() (e.g. a transport reset) and logs a warning instead of dying silently, making the degradation to no-op visible; a test pins that the LLM call is not cancelled in that case --------- Co-authored-by: kursad Co-authored-by: CreateRandom <18438707+CreateRandom@users.noreply.github.com> * fix(bedrock): grant aws-external-anthropic:* in OIDC session policy for claude_platform (#30200) (#30205) The inline STS session policy passed to assume_role_with_web_identity acts as an IAM PERMISSION CEILING — effective permissions are the intersection of the role's identity policies and this policy. Any action not listed is silently denied even when the IAM role grants it. #27678 added the bedrock/claude_platform/ route but its service-side action namespace is aws-external-anthropic:*, not bedrock:*. Without a matching statement here, every claude_platform request via OIDC (GCP federation, EKS Pod Identity webhook, etc.) 403s with 'no session policy allows the aws-external-anthropic:CreateInference action' — even with a fully permissive identity policy. Add a second ClaudePlatformLiteLLM statement covering CreateInference, CreateBatchInference, CancelBatchInference, DeleteBatchInference, CountTokens, Get*, List*. Keep aws:SecureTransport=true parity with the bedrock statement. Static creds + IRSA flow through different code paths and are not affected. Fixes #30200 * fix(proxy): set Retry-After header on RouterRateLimitError 429 responses (#30098) * Set Retry-After header on RouterRateLimitError responses When all deployments for a model are in cooldown, the proxy returns a 429 whose cooldown timing is only available by parsing the error message string. RouterRateLimitError already carries cooldown_time, so expose it as a standard retry-after header in _handle_llm_api_exception. The value is rounded up so clients never retry before the cooldown window ends. Fixes #27823. * Set Retry-After after response-headers hook so cooldown wins The cooldown-derived retry-after was assigned before the post_call_response_headers_hook merge, so a callback returning a retry-after key (including a stale or empty value) silently clobbered it. Move the RouterRateLimitError block after the callback merge so the cooldown value is authoritative for this error type. * fix(router): route aspeech through async_function_with_fallbacks (#30104) * fix(router): route aspeech through async_function_with_fallbacks Router.aspeech selected a deployment and awaited litellm.aspeech directly, so TTS requests got no retry on failure and no failover to backup deployments; the except block only fired an exception alert and re-raised. Every other router endpoint (acompletion, aembedding, atranscription, arerank) already delegates to async_function_with_fallbacks Mirror the atranscription pattern: move deployment selection and the litellm.aspeech call into a private _aspeech method, then have the public aspeech set kwargs["original_function"] = self._aspeech and await self.async_function_with_fallbacks(**kwargs). _aspeech also picks up the shared _get_async_openai_model_client helper and the same total/success/fail call accounting the sibling endpoints use Fixes #27778. * fix(router): apply deployment kwargs and rpm semaphore in _aspeech Bring _aspeech fully in line with _atranscription: call _update_kwargs_with_deployment so deployment metadata, model_info, timeout, and default litellm params flow into the request, and wrap the litellm.aspeech call with the max_parallel_requests semaphore plus async_routing_strategy_pre_call_checks so TTS respects rpm limits the same way the other router endpoints do Also add a unit test that exercises _aspeech directly and asserts the deployment metadata reaches the underlying call * fix(slack_alerting): stop false-positive hanging request alerts for requests below the alerting threshold (#30106) * fix(slack_alerting): skip hanging request alerts below the threshold The hanging request check alerted on any cached request whose completion status was not yet recorded, with no minimum age check. Since the background loop runs every alerting_threshold / 2 seconds, any request that happened to be in flight at a check fired a "hanging - Ns+ request time" alert even if it was only seconds old, producing a steady stream of false positives. Add a created_at timestamp to HangingRequestData, stamped when the request enters the hanging request cache, and skip requests younger than alerting_threshold without evicting them, so a later check can still alert if they never complete. Extend the cache TTL from threshold + 60s to 1.5x threshold + 60s; with the age check, entries only become alertable after threshold seconds, and the check period is threshold / 2, so the old TTL could evict a genuinely hanging request before any check saw it cross the threshold. Fixes #27855. * fix(slack_alerting): alert once per hanging request The min-age gate stops false positives for young in-flight requests, but a genuinely hanging request still re-alerted on every checker tick within the cache TTL. With the wider TTL (1.5x threshold + 60s) that is 1-2 extra Slack notifications per stuck request at the default 600s threshold. Flag a HangingRequestData entry as alerted once its alert fires and skip flagged entries on later ticks, so each hang produces exactly one alert. The cache reference is mutated in place, so the TTL is untouched and still handles cleanup. Adds a regression test asserting one alert across multiple ticks. Fixes #27855. * fix(health): treat all-proxy-models keys as unrestricted in /health (#30087) * fix(health): treat all-proxy-models keys as unrestricted in /health A key granted all model permissions stores the literal "all-proxy-models" marker in its models list. The /health access filter compared that marker against real model_names, so the model list filtered down to nothing and the WebUI health check returned healthy_count=0, unhealthy_count=0 with HTTP 503. Skip the filter (both the live path and the background-cache model_id scoping) when the marker is present, matching how auth_checks treats SpecialModelNames.all_proxy_models. Fixes #29744. * fix(health): resolve all-team-models sentinel to the team allowlist Same failure shape as the all-proxy-models case: a key carrying the literal "all-team-models" entry matches no real model_name, so the /health access filter would zero out the model list. Resolve the sentinel to the key's team models when team_id is set, matching get_key_models in model_checks.py. Without a team_id the sentinel stays unresolved and matches nothing, denying rather than widening access, mirroring _resolve_key_models_for_auth_check. * feat(proxy): auto-enable drop_params for Claude Code requests (#30218) * feat(proxy): auto-enable drop_params for Claude Code requests Claude Code identifies itself with a claude-cli/ user agent and sends Anthropic-specific params (top_k, thinking, etc.) on every request. When the proxy routes those requests to a non-Anthropic provider, the unsupported params fail the call unless drop_params is configured. Detect the Claude Code user agent in add_litellm_data_to_request and default drop_params to true for those requests, without overriding an explicit drop_params value sent by the caller. * feat(proxy): respect operator litellm_settings drop_params over Claude Code default An explicit drop_params in the operator's litellm_settings (true or false) now suppresses the Claude Code user agent default, so an operator who deliberately configured drop_params: false keeps strict param validation for Claude Code clients too. The auto-default only fills the gap when neither the request body nor the config sets a value. * fix(snowflake): migrate to native endpoints with auto-routing for Claude models (#29964) * fix(snowflake): migrate to native Cortex REST API endpoints Replaces the legacy /api/v2/cortex/inference:complete endpoint with the native OpenAI-compatible /api/v2/cortex/v1/chat/completions endpoint, fixing error 390142 (Incoming request does not contain a valid payload) when using model: snowflake/ in LiteLLM proxy. Changes: - litellm/llms/snowflake/chat/transformation.py: route to native /cortex/v1/chat/completions, remove Snowflake-specific tool_spec payload transformation, remove content_list response handling, add stream to supported params - litellm/llms/snowflake/anthropic/transformation.py (new): SnowflakeCortexAnthropicConfig routes Claude models to /cortex/v1/messages with anthropic-version header and Anthropic->OpenAI response transform - tests: 29 unit tests covering URL routing, auth headers, payload format, and response parsing * fix(snowflake): map max_tokens to max_completion_tokens for native endpoint * fix: handle multi-turn tool conversations and OpenAI→Anthropic tool format conversion - _extract_system_and_messages now preserves tool_calls from assistant messages and converts them to Anthropic tool_use content blocks - tool role messages are converted to user role with tool_result content blocks (as required by Anthropic Messages API) - Added _transform_tools_to_anthropic() to convert OpenAI tool format (type/function/parameters) to Anthropic format (name/input_schema) - Added comprehensive tests for multi-turn tool conversations Addresses review feedback on PR #29964 * test: add coverage for malformed JSON and non-string tool arguments * fix(tests): update chat transformation tests for native OpenAI-compatible endpoint * style: apply black formatting * fix: resolve mypy type errors in anthropic transformation * fix: correct mypy type: ignore error codes (attr-defined) * fix: use max_tokens instead of max_completion_tokens for Snowflake endpoint compatibility * refactor: merge Anthropic config into unified SnowflakeConfig with auto-routing - Remove separate SnowflakeCortexAnthropicConfig and anthropic/ directory - SnowflakeConfig now auto-routes based on model name: - Claude models → /messages endpoint (Anthropic format) - All others → /chat/completions endpoint (OpenAI format) - No new provider needed (stays as SNOWFLAKE = 'snowflake') - Tool message transformation for Claude: tool_calls → tool_use blocks, tool role → user with tool_result - OpenAI → Anthropic tool format conversion (parameters → input_schema) - Addresses Greptile feedback about unwired SnowflakeCortexAnthropicConfig * fix: use max_completion_tokens for /chat/completions (Snowflake deprecated max_tokens on this endpoint) * fix(tests): update assertions for Claude auto-routing to /messages endpoint * fix(snowflake): add tool_choice conversion and preserve max_completion_tokens in Anthropic path * fix(snowflake): use ChatCompletionMessageToolCall objects and strip model prefix on OpenAI path * fix(snowflake): collect multiple system messages to prevent guardrail override * chore: remove committed .pyc files and add __pycache__ to .gitignore * fix: remove unused Union import * fix: restore original .gitignore (accidentally replaced in earlier commit) * feat(snowflake): add streaming response handler for both Anthropic and OpenAI SSE formats * fix: remove unused AsyncIterator and Iterator imports * fix: add missing total_tokens to ChatCompletionUsageBlock * fix(snowflake): coalesce consecutive tool results into single user message for Anthropic * fix(snowflake): handle message_start event for streaming input_tokens tracking * fix: evict last deleted model in multi-instance deployments (#28608) * fix: evict last deleted model in multi-instance deployments _delete_deployment had an early return when db_models was empty, preventing eviction of the last deleted model during reconciliation. - Remove len(db_models)==0 early return from _delete_deployment - Return None (not []) from _get_models_from_db on DB failure so callers can distinguish a transient failure from a genuinely empty DB - Guard _update_llm_router against None to skip updates on DB failure Fixes #28443 * test: remove dead MagicMock assignment in type_mismatch test * fix: update test to pass [] not None to _update_llm_router test_ProxyConfig__update_llm_router_bad_proxy_logging_raises was passing None as new_models to get through to the proxy_logging_obj check, but the None guard we added now returns early before reaching that path. Pass [] instead so the test exercises the intended AttributeError case. Signed-off-by: Rudra Dudhat * chore: regenerate API types to sync schema.d.ts with proxy OpenAPI spec Signed-off-by: Rudra Dudhat --------- Signed-off-by: Rudra Dudhat * fix: invalidate Redis spend counter on /key/reset_spend (#29694) * fix: set Redis spend counter to reset_to value on /key/reset_spend Previously, the Redis spend counter was always set to 0.0 after a reset, even when reset_to was a non-zero value (partial reset). This caused the budget to be under-enforced for up to 60 seconds until the counter expired and fell through to the DB. Now the counter is set to the actual reset_to value, so partial resets are reflected correctly and budget enforcement is consistent. * test: update reset_key_spend test to match direct cache set The implementation now sets spend_counter_cache directly instead of calling _invalidate_spend_counter. Update the test to verify the in_memory_cache.set_cache call with the correct key, value, and ttl. --------- Co-authored-by: michaelxer * fix: add scaleway models pricing (#27659) * fix: Add embeddings support for Scaleway provider * fix: resolve merge conflicts * fix(main): clarify backend route handling for Swagger static assets (#30196) * fix(main): clarify backend route handling for Swagger static assets * fix(allowlist): add BACKEND_MOUNT_PATHS for Swagger static assets * fix(voyage): route multimodal embeddings to correct endpoint (#30193) * fix(voyage): route multimodal embeddings to correct endpoint * test(voyage): cover multimodal embedding edge cases * test(voyage): cover api key fallback * fix(voyage): raise early on missing api key and malformed image url * test(voyage): cover utils routing and helper * fix(voyage): route supported openai params for multimodal models * style: apply black formatting * fix(ui): infer Azure API version from API base (#30204) * fix(ui): infer Azure API version from API base * fix(ui): address Azure API version feedback * Update litellm/llms/snowflake/chat/transformation.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * feat(datadog): add team-scoped Datadog callback support (#29947) Enable teams to configure their own Datadog credentials via POST /team/{team_id}/callback, following the same pattern as Langfuse. * Merge pull request #29528 from aanchal22/litellm_byok-alias-merge fix(proxy): atomic merge for team model aliases and team.models on BYOK create * feat: add EmpirioLabs as an OpenAI-compatible provider (#30278) Co-authored-by: Adam Dalloul * fix: resolve failing tests and lint in snowflake/team endpoints - Black-format snowflake/chat/transformation.py to fix lint failure - Update Anthropic config test to expect default max_tokens of 4096 (matches implementation) - Add AsyncMock + execute_raw mock to team_model_add cache-refresh pin test - Add model_dump mock and patch cache/logging in test_uses_atomic_array_append_with_dedup Co-Authored-By: Claude Sonnet 4.6 * fix(test): update test_db_error_new_model_check for new _delete_deployment logic _delete_deployment no longer short-circuits on empty db_models — it now treats [] as a valid empty-DB state and proceeds to check config models. Mock get_config to return the two router deployments so they appear in combined_id_list and are protected, which matches the real-world scenario where a DB error occurs but the models are config-backed. Co-Authored-By: Claude Sonnet 4.6 * feat(proxy): register cancel_on_disconnect in ConfigGeneralSettings and config list (#30295) * feat(proxy): register cancel_on_disconnect in ConfigGeneralSettings and config list Follow-up to #30223 per maintainer review: documents the flag in ConfigGeneralSettings with a short description and adds it to allowed_args in get_config_list so the UI and /config/list expose it. A test pins that /config/list returns the field with type Boolean, which requires both registrations to be present * chore(ui): regenerate schema.d.ts for cancel_on_disconnect --------- Co-authored-by: kursad * fix(datadog): never fall back to env DD_API_KEY for caller-supplied destinations Team/key-scoped Datadog loggers could be pointed at an arbitrary dd_agent_host or dd_site while omitting dd_api_key, causing the proxy's global DD_API_KEY to be sent as the DD-API-KEY header to that destination. Gate the env-var fallback behind an allow_env_credentials flag, set to False when the destination is caller-supplied, mirroring the existing langfuse/langsmith pattern. --------- Signed-off-by: Rudra Dudhat Co-authored-by: Emerson Gomes Co-authored-by: daitran-tensormesh Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Co-authored-by: Muspi Merol Co-authored-by: fangkang Co-authored-by: Chris Hoogeboom Co-authored-by: kursadlacin Co-authored-by: kursad Co-authored-by: CreateRandom <18438707+CreateRandom@users.noreply.github.com> Co-authored-by: hcl Co-authored-by: Filippo Menghi <113345637+Cyberfilo@users.noreply.github.com> Co-authored-by: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Co-authored-by: sfc-gh-nashukla Co-authored-by: Rudra Dudhat Co-authored-by: Michael <52305679+michaelxer@users.noreply.github.com> Co-authored-by: michaelxer Co-authored-by: Quentin Champenois <26109239+Quentinchampenois@users.noreply.github.com> Co-authored-by: mauriceberentsen Co-authored-by: lost9999 <56498264+lost9999@users.noreply.github.com> Co-authored-by: GaetanVDB07 <86427581+GaetanVDB07@users.noreply.github.com> Co-authored-by: Aanchal Khandelwal Co-authored-by: Adam Dalloul Co-authored-by: Adam Dalloul Co-authored-by: Claude Sonnet 4.6 --- backend/main.py | 11 +- backend/routes/allowlist.py | 6 + litellm/__init__.py | 3 + litellm/_lazy_imports_registry.py | 5 + litellm/anthropic_beta_headers_config.json | 4 +- .../SlackAlerting/hanging_request_check.py | 30 +- litellm/integrations/datadog/datadog.py | 81 +- .../datadog/datadog_team_handler.py | 124 +++ .../integrations/langfuse/langfuse_otel.py | 1 + litellm/integrations/otel/mappers/genai.py | 14 + litellm/integrations/otel/model/payloads.py | 48 + .../integrations/otel/plumbing/providers.py | 6 +- .../get_supported_openai_params.py | 9 + .../initialize_dynamic_callback_params.py | 8 + litellm/litellm_core_utils/litellm_logging.py | 40 +- litellm/llms/bedrock/base_aws_llm.py | 46 +- litellm/llms/bedrock/chat/converse_handler.py | 6 +- litellm/llms/bedrock/chat/invoke_handler.py | 10 +- .../anthropic_claude3_transformation.py | 1 + .../base_invoke_transformation.py | 1 + litellm/llms/openai_like/providers.json | 12 +- litellm/llms/snowflake/chat/transformation.py | 827 +++++++++++++----- .../embedding/transformation_multimodal.py | 183 ++++ ...odel_prices_and_context_window_backup.json | 54 +- litellm/proxy/_types.py | 4 + litellm/proxy/common_request_processing.py | 77 +- .../health_endpoints/_health_endpoints.py | 25 +- litellm/proxy/litellm_pre_call_utils.py | 29 + .../key_management_endpoints.py | 25 +- .../management_endpoints/team_endpoints.py | 33 +- litellm/proxy/proxy_server.py | 60 +- litellm/router.py | 110 ++- litellm/types/integrations/slack_alerting.py | 3 + litellm/types/utils.py | 6 + litellm/utils.py | 14 + model_prices_and_context_window.json | 214 +++++ provider_endpoints_support.json | 21 +- proxy_server_config.yaml | 1 + .../openai_like/test_empiriolabs_provider.py | 63 ++ tests/local_testing/test_config.py | 18 +- .../test_router_endpoints.py | 90 ++ .../test_hanging_request_check.py | 88 +- .../datadog/test_datadog_team_handler.py | 263 ++++++ .../otel/test_otel_v2_components.py | 92 ++ .../integrations/otel/test_otel_v2_emitter.py | 36 + .../integrations/test_langfuse_otel.py | 3 + .../test_base_invoke_transformation.py | 41 + .../llms/bedrock/chat/test_invoke_handler.py | 128 ++- .../test_web_identity_session_policy.py | 176 ++++ .../test_bedrock_mantle_transformation.py | 66 ++ .../llms/chat/test_converse_handler.py | 80 ++ .../openai_like/test_tensormesh_provider.py | 14 + .../test_snowflake_chat_transformation.py | 159 ++-- .../test_snowflake_native_endpoints.py | 718 +++++++++++++++ .../test_voyage_multimodal_embedding.py | 306 +++++++ .../health_endpoints/test_health_endpoints.py | 132 +++ .../test_key_management_endpoints.py | 176 ++-- .../test_team_endpoints.py | 5 +- .../test_team_model_alias_merge.py | 83 ++ .../proxy/proxy_server/test_lifecycle.py | 60 +- .../proxy/proxy_server/test_proxy_config.py | 2 +- .../proxy/test_common_request_processing.py | 267 +++++- .../proxy/test_component_allowlists.py | 47 +- .../proxy/test_litellm_pre_call_utils.py | 62 ++ tests/test_litellm/proxy/test_proxy_server.py | 173 +++- .../test_anthropic_beta_headers_filtering.py | 14 + .../provider_specific_fields.test.tsx | 132 ++- .../add_model/provider_specific_fields.tsx | 37 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 + 69 files changed, 5055 insertions(+), 633 deletions(-) create mode 100644 litellm/integrations/datadog/datadog_team_handler.py create mode 100644 litellm/llms/voyage/embedding/transformation_multimodal.py create mode 100644 tests/litellm/llms/openai_like/test_empiriolabs_provider.py create mode 100644 tests/test_litellm/integrations/datadog/test_datadog_team_handler.py create mode 100644 tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py create mode 100644 tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py create mode 100644 tests/test_litellm/llms/snowflake/test_snowflake_native_endpoints.py create mode 100644 tests/test_litellm/llms/voyage/test_voyage_multimodal_embedding.py create mode 100644 tests/test_litellm/proxy/management_endpoints/test_team_model_alias_merge.py diff --git a/backend/main.py b/backend/main.py index 4092cd63f69..292ece48e7d 100644 --- a/backend/main.py +++ b/backend/main.py @@ -20,7 +20,11 @@ DatabaseURLSettings.from_env().apply_to_env() from litellm.proxy.proxy_server import app -from backend.routes.allowlist import BACKEND_EXACT_PATHS, BACKEND_PATH_PREFIXES +from backend.routes.allowlist import ( + BACKEND_EXACT_PATHS, + BACKEND_MOUNT_PATHS, + BACKEND_PATH_PREFIXES, +) def _is_backend_route(route) -> bool: @@ -29,8 +33,9 @@ def _is_backend_route(route) -> bool: if path is None: return False if isinstance(route, Mount): - # Static UI mounts are served by the dedicated UI container, not here. - return False + # The dashboard UI static mounts are served by the dedicated UI container. + # Only Mounts in the backend allowlist (e.g. swagger docs) remain on backend. + return path in BACKEND_MOUNT_PATHS if path in BACKEND_EXACT_PATHS: return True return any(path.startswith(prefix) for prefix in BACKEND_PATH_PREFIXES) diff --git a/backend/routes/allowlist.py b/backend/routes/allowlist.py index 610ba3dbd69..d1a576aeb33 100644 --- a/backend/routes/allowlist.py +++ b/backend/routes/allowlist.py @@ -133,3 +133,9 @@ BACKEND_EXACT_PATHS: frozenset[str] = frozenset( "/fallback/login", } ) + +BACKEND_MOUNT_PATHS: frozenset[str] = frozenset( + { + "/swagger", # API documentation static assets belong to the backend + } +) diff --git a/litellm/__init__.py b/litellm/__init__.py index e5bc785ed3b..d5fbb41c462 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1731,6 +1731,9 @@ if TYPE_CHECKING: from .llms.voyage.embedding.transformation_contextual import ( VoyageContextualEmbeddingConfig as VoyageContextualEmbeddingConfig, ) + from .llms.voyage.embedding.transformation_multimodal import ( + VoyageMultimodalEmbeddingConfig as VoyageMultimodalEmbeddingConfig, + ) from .llms.infinity.embedding.transformation import ( InfinityEmbeddingConfig as InfinityEmbeddingConfig, ) diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index bace54ffad1..6073b6b2833 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -223,6 +223,7 @@ LLM_CONFIG_NAMES = ( "GenAIHubOrchestrationConfig", "VoyageEmbeddingConfig", "VoyageContextualEmbeddingConfig", + "VoyageMultimodalEmbeddingConfig", "InfinityEmbeddingConfig", "PerplexityEmbeddingConfig", "AzureAIStudioConfig", @@ -903,6 +904,10 @@ _LLM_CONFIGS_IMPORT_MAP = { ".llms.voyage.embedding.transformation_contextual", "VoyageContextualEmbeddingConfig", ), + "VoyageMultimodalEmbeddingConfig": ( + ".llms.voyage.embedding.transformation_multimodal", + "VoyageMultimodalEmbeddingConfig", + ), "InfinityEmbeddingConfig": ( ".llms.infinity.embedding.transformation", "InfinityEmbeddingConfig", diff --git a/litellm/anthropic_beta_headers_config.json b/litellm/anthropic_beta_headers_config.json index a0d63f5043c..11fdb26e42d 100644 --- a/litellm/anthropic_beta_headers_config.json +++ b/litellm/anthropic_beta_headers_config.json @@ -75,7 +75,7 @@ "effort-2025-11-24": "effort-2025-11-24", "fast-mode-2026-02-01": null, "files-api-2025-04-14": null, - "fine-grained-tool-streaming-2025-05-14": null, + "fine-grained-tool-streaming-2025-05-14": "fine-grained-tool-streaming-2025-05-14", "interleaved-thinking-2025-05-14": null, "mcp-client-2025-11-20": null, "mcp-client-2025-04-04": null, @@ -106,7 +106,7 @@ "effort-2025-11-24": "effort-2025-11-24", "fast-mode-2026-02-01": null, "files-api-2025-04-14": null, - "fine-grained-tool-streaming-2025-05-14": null, + "fine-grained-tool-streaming-2025-05-14": "fine-grained-tool-streaming-2025-05-14", "interleaved-thinking-2025-05-14": null, "mcp-client-2025-11-20": null, "mcp-client-2025-04-04": null, diff --git a/litellm/integrations/SlackAlerting/hanging_request_check.py b/litellm/integrations/SlackAlerting/hanging_request_check.py index d2f70c9caf1..98f1eb2d551 100644 --- a/litellm/integrations/SlackAlerting/hanging_request_check.py +++ b/litellm/integrations/SlackAlerting/hanging_request_check.py @@ -8,6 +8,7 @@ Notes: """ import asyncio +import time from typing import TYPE_CHECKING, Any, Optional import litellm @@ -36,11 +37,15 @@ class AlertingHangingRequestCheck: slack_alerting_object: SlackAlerting, ): self.slack_alerting_object = slack_alerting_object + # checks run every alerting_threshold / 2 seconds, so entries must + # stay cached for at least 1.5x the threshold to guarantee a check + # happens after they cross it + self.hanging_request_cache_ttl = int( + self.slack_alerting_object.alerting_threshold * 1.5 + + HANGING_ALERT_BUFFER_TIME_SECONDS + ) self.hanging_request_cache = InMemoryCache( - default_ttl=int( - self.slack_alerting_object.alerting_threshold - + HANGING_ALERT_BUFFER_TIME_SECONDS - ), + default_ttl=self.hanging_request_cache_ttl, ) async def add_request_to_hanging_request_check( @@ -76,10 +81,7 @@ class AlertingHangingRequestCheck: await self.hanging_request_cache.async_set_cache( key=hanging_request_data.request_id, value=hanging_request_data, - ttl=int( - self.slack_alerting_object.alerting_threshold - + HANGING_ALERT_BUFFER_TIME_SECONDS - ), + ttl=self.hanging_request_cache_ttl, ) return @@ -111,6 +113,9 @@ class AlertingHangingRequestCheck: if hanging_request_data is None: continue + if hanging_request_data.alerted: + continue + request_status = ( await proxy_logging_obj.internal_usage_cache.async_get_cache( key="request_status:{}".format(hanging_request_data.request_id), @@ -127,12 +132,21 @@ class AlertingHangingRequestCheck: ) continue + request_age_seconds = time.time() - hanging_request_data.created_at + if request_age_seconds < self.slack_alerting_object.alerting_threshold: + # in-flight but below the alerting threshold; keep it cached + # so a later check can alert if it never completes + continue + ################ # Send the Alert on Slack ################ await self.send_hanging_request_alert( hanging_request_data=hanging_request_data ) + # flag so the entry is skipped on later ticks; one alert per hang, + # with the existing TTL still handling cleanup + hanging_request_data.alerted = True return diff --git a/litellm/integrations/datadog/datadog.py b/litellm/integrations/datadog/datadog.py index 79a9219a39c..b0cd0eb1172 100644 --- a/litellm/integrations/datadog/datadog.py +++ b/litellm/integrations/datadog/datadog.py @@ -92,12 +92,26 @@ class DataDogLogger( # Class variables or attributes def __init__( self, + dd_api_key: Optional[str] = None, + dd_site: Optional[str] = None, + dd_agent_host: Optional[str] = None, + dd_agent_port: Optional[str] = None, + allow_env_credentials: bool = True, **kwargs, ): """ Initializes the datadog logger, checks if the correct env variables are set - Required environment variables (Direct API): + Args: + dd_api_key: Datadog API key. Falls back to DD_API_KEY env var when allow_env_credentials is True. + dd_site: Datadog site (e.g. "us5.datadoghq.com"). Falls back to DD_SITE env var. + dd_agent_host: Hostname or IP of DataDog agent. Falls back to LITELLM_DD_AGENT_HOST env var. + dd_agent_port: Port of DataDog agent (default: 10518). Falls back to LITELLM_DD_AGENT_PORT env var. + allow_env_credentials: When False, the API key is never read from DD_API_KEY env var. Set to + False for team/key-scoped loggers whose destination (dd_agent_host/dd_site) is caller-supplied, + so the proxy's global DD_API_KEY is never sent to an untrusted host. + + Required environment variables (Direct API) when kwargs not provided: `DD_API_KEY` - your datadog api key `DD_SITE` - your datadog site, example = `"us5.datadoghq.com"` @@ -130,12 +144,21 @@ class DataDogLogger( ) # Configure DataDog endpoint (Agent or Direct API) - # Use LITELLM_DD_AGENT_HOST to avoid conflicts with ddtrace's DD_AGENT_HOST - dd_agent_host = os.getenv("LITELLM_DD_AGENT_HOST") - if dd_agent_host: - self._configure_dd_agent(dd_agent_host=dd_agent_host) + # Prefer explicit kwargs, then fall back to env vars + resolved_agent_host = dd_agent_host or os.getenv("LITELLM_DD_AGENT_HOST") + if resolved_agent_host: + self._configure_dd_agent( + dd_agent_host=resolved_agent_host, + dd_agent_port=dd_agent_port, + dd_api_key=dd_api_key, + allow_env_credentials=allow_env_credentials, + ) else: - self._configure_dd_direct_api() + self._configure_dd_direct_api( + dd_api_key=dd_api_key, + dd_site=dd_site, + allow_env_credentials=allow_env_credentials, + ) # Optional override for testing dd_base_url = get_datadog_base_url_from_env() @@ -172,34 +195,60 @@ class DataDogLogger( ).model_dump() return dict_datadog_params - def _configure_dd_agent(self, dd_agent_host: str) -> None: + def _configure_dd_agent( + self, + dd_agent_host: str, + dd_agent_port: Optional[str] = None, + dd_api_key: Optional[str] = None, + allow_env_credentials: bool = True, + ) -> None: """ Configure DataDog Agent for log forwarding Args: dd_agent_host: Hostname or IP of DataDog agent + dd_agent_port: Port of DataDog agent. Falls back to LITELLM_DD_AGENT_PORT env var (default: 10518). + dd_api_key: Datadog API key. Falls back to DD_API_KEY env var when allow_env_credentials is True. Optional when using agent. + allow_env_credentials: When False, never read the API key from DD_API_KEY env var. """ - dd_agent_port = os.getenv( + resolved_port = dd_agent_port or os.getenv( "LITELLM_DD_AGENT_PORT", "10518" ) # default port for logs - self.intake_url = f"http://{dd_agent_host}:{dd_agent_port}/api/v2/logs" - self.DD_API_KEY = os.getenv("DD_API_KEY") # Optional when using agent + self.intake_url = f"http://{dd_agent_host}:{resolved_port}/api/v2/logs" + self.DD_API_KEY = dd_api_key or ( + os.getenv("DD_API_KEY") if allow_env_credentials else None + ) # Optional when using agent verbose_logger.debug(f"Datadog: Using DD Agent at {self.intake_url}") - def _configure_dd_direct_api(self) -> None: + def _configure_dd_direct_api( + self, + dd_api_key: Optional[str] = None, + dd_site: Optional[str] = None, + allow_env_credentials: bool = True, + ) -> None: """ Configure direct DataDog API connection + Args: + dd_api_key: Datadog API key. Falls back to DD_API_KEY env var when allow_env_credentials is True. + dd_site: Datadog site. Falls back to DD_SITE env var. + allow_env_credentials: When False, never read the API key from DD_API_KEY env var. + Raises: - Exception: If required environment variables are not set + Exception: If required credentials are not provided via args or env vars """ - if os.getenv("DD_API_KEY", None) is None: + resolved_api_key = dd_api_key or ( + os.getenv("DD_API_KEY") if allow_env_credentials else None + ) + resolved_site = dd_site or os.getenv("DD_SITE") + + if resolved_api_key is None: raise Exception("DD_API_KEY is not set, set 'DD_API_KEY=<>") - if os.getenv("DD_SITE", None) is None: + if resolved_site is None: raise Exception("DD_SITE is not set in .env, set 'DD_SITE=<>") - self.DD_API_KEY = os.getenv("DD_API_KEY") - self.intake_url = f"https://http-intake.logs.{os.getenv('DD_SITE')}/api/v2/logs" + self.DD_API_KEY = resolved_api_key + self.intake_url = f"https://http-intake.logs.{resolved_site}/api/v2/logs" async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): """ diff --git a/litellm/integrations/datadog/datadog_team_handler.py b/litellm/integrations/datadog/datadog_team_handler.py new file mode 100644 index 00000000000..3a5b73fc005 --- /dev/null +++ b/litellm/integrations/datadog/datadog_team_handler.py @@ -0,0 +1,124 @@ +""" +DataDog Team Handler + +Used to get the DataDogLogger for a given request. +Handles Key/Team Based Datadog Logging, following the same pattern as LangFuseHandler. +""" + +from typing import TYPE_CHECKING, Any, Dict, Optional, TypedDict + +from litellm._logging import verbose_logger +from litellm.litellm_core_utils.litellm_logging import StandardCallbackDynamicParams + +from .datadog import DataDogLogger + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import DynamicLoggingCache +else: + DynamicLoggingCache = Any + + +class DatadogLoggingConfig(TypedDict): + dd_api_key: Optional[str] + dd_site: Optional[str] + dd_agent_host: Optional[str] + dd_agent_port: Optional[str] + + +class DataDogHandler: + @staticmethod + def get_datadog_logger_for_request( + standard_callback_dynamic_params: StandardCallbackDynamicParams, + in_memory_dynamic_logger_cache: DynamicLoggingCache, + ) -> DataDogLogger: + """ + Get a team-scoped DataDogLogger for a given request. + + Resolves and caches per-team DataDogLogger instances using DynamicLoggingCache, + keyed by the team's DD credentials. Each unique set of credentials gets its own + logger instance with its own batch/flush loop. + + Note: This handler is only called when team-scoped DD credentials are present. + The global (env-var based) DataDogLogger is managed separately by + _init_custom_logger_compatible_class via _in_memory_loggers. + """ + _credentials = DataDogHandler.get_dynamic_datadog_logging_config( + standard_callback_dynamic_params=standard_callback_dynamic_params, + ) + credentials_dict = dict(_credentials) + + # check if datadog logger is already cached + temp_datadog_logger = in_memory_dynamic_logger_cache.get_cache( + credentials=credentials_dict, service_name="datadog" + ) + + # if not cached, create a new datadog logger and cache it + if temp_datadog_logger is None: + temp_datadog_logger = ( + DataDogHandler._create_datadog_logger_from_credentials( + credentials=credentials_dict, + in_memory_dynamic_logger_cache=in_memory_dynamic_logger_cache, + ) + ) + + return temp_datadog_logger + + @staticmethod + def _create_datadog_logger_from_credentials( + credentials: Dict, + in_memory_dynamic_logger_cache: DynamicLoggingCache, + ) -> DataDogLogger: + """ + Create a DataDogLogger from the credentials and cache it. + """ + # When the destination is caller-supplied (dd_agent_host/dd_site), never fall back to the + # proxy's DD_API_KEY env var, otherwise it would be sent to a team-controlled host. + allow_env_credentials = ( + credentials.get("dd_agent_host") is None + and credentials.get("dd_site") is None + ) + datadog_logger = DataDogLogger( + dd_api_key=credentials.get("dd_api_key"), + dd_site=credentials.get("dd_site"), + dd_agent_host=credentials.get("dd_agent_host"), + dd_agent_port=credentials.get("dd_agent_port"), + allow_env_credentials=allow_env_credentials, + ) + in_memory_dynamic_logger_cache.set_cache( + credentials=credentials, + service_name="datadog", + logging_obj=datadog_logger, + ) + verbose_logger.debug( + "Datadog: Created and cached new DataDogLogger for team-scoped credentials" + ) + return datadog_logger + + @staticmethod + def get_dynamic_datadog_logging_config( + standard_callback_dynamic_params: StandardCallbackDynamicParams, + ) -> DatadogLoggingConfig: + """ + Get the Datadog logging config for a given request from dynamic params. + """ + return DatadogLoggingConfig( + dd_api_key=standard_callback_dynamic_params.get("dd_api_key"), + dd_site=standard_callback_dynamic_params.get("dd_site"), + dd_agent_host=standard_callback_dynamic_params.get("dd_agent_host"), + dd_agent_port=standard_callback_dynamic_params.get("dd_agent_port"), + ) + + @staticmethod + def _dynamic_datadog_credentials_are_passed( + standard_callback_dynamic_params: StandardCallbackDynamicParams, + ) -> bool: + """ + Check if dynamic Datadog credentials are passed in standard_callback_dynamic_params. + """ + if ( + standard_callback_dynamic_params.get("dd_api_key") is not None + or standard_callback_dynamic_params.get("dd_site") is not None + or standard_callback_dynamic_params.get("dd_agent_host") is not None + ): + return True + return False diff --git a/litellm/integrations/langfuse/langfuse_otel.py b/litellm/integrations/langfuse/langfuse_otel.py index b96ec72b04e..7370bcdf934 100644 --- a/litellm/integrations/langfuse/langfuse_otel.py +++ b/litellm/integrations/langfuse/langfuse_otel.py @@ -43,6 +43,7 @@ class LangfuseOtelLogger(OpenTelemetry): """ _utils.set_attributes(span, kwargs, response_obj, LangfuseLLMObsOTELAttributes) + span.set_attribute("langfuse.observation.type", "generation") ######################################################### # Set Langfuse specific attributes diff --git a/litellm/integrations/otel/mappers/genai.py b/litellm/integrations/otel/mappers/genai.py index 6c61feced4d..d4f14e97a7a 100644 --- a/litellm/integrations/otel/mappers/genai.py +++ b/litellm/integrations/otel/mappers/genai.py @@ -63,6 +63,20 @@ class GenAIMapper: # routing) onto the boundary-born LLM span — stamp it directly here. LiteLLM.PROVIDER_MODEL: lambda d: d.identity.provider_model or None, f"{LiteLLM.COST_PREFIX}total": lambda d: d.response_cost, + # Per-component cost breakdown (from the StandardLoggingPayload + # ``cost_breakdown``). Each component is omitted when the source didn't + # report it, so spans stay sparse rather than carrying zeros. + f"{LiteLLM.COST_PREFIX}input": lambda d: d.cost.input, + f"{LiteLLM.COST_PREFIX}output": lambda d: d.cost.output, + f"{LiteLLM.COST_PREFIX}cache_read": lambda d: d.cost.cache_read, + f"{LiteLLM.COST_PREFIX}cache_creation": lambda d: d.cost.cache_creation, + f"{LiteLLM.COST_PREFIX}tool_usage": lambda d: d.cost.tool_usage, + f"{LiteLLM.COST_PREFIX}original": lambda d: d.cost.original, + f"{LiteLLM.COST_PREFIX}discount_amount": lambda d: d.cost.discount_amount, + f"{LiteLLM.COST_PREFIX}discount_percent": lambda d: d.cost.discount_percent, + f"{LiteLLM.COST_PREFIX}margin_fixed_amount": lambda d: d.cost.margin_fixed_amount, + f"{LiteLLM.COST_PREFIX}margin_percent": lambda d: d.cost.margin_percent, + f"{LiteLLM.COST_PREFIX}margin_total_amount": lambda d: d.cost.margin_total_amount, LiteLLM.REQUEST_STREAMING: lambda d: d.is_streaming, } diff --git a/litellm/integrations/otel/model/payloads.py b/litellm/integrations/otel/model/payloads.py index bbef40ba374..82b7df5922c 100644 --- a/litellm/integrations/otel/model/payloads.py +++ b/litellm/integrations/otel/model/payloads.py @@ -34,6 +34,7 @@ __all__ = [ "RequestIdentity", "GuardrailSpanData", "LLMCallSpanData", + "LLMCost", "LLMRequestParams", "LLMUsage", "MCPToolCallSpanData", @@ -91,6 +92,49 @@ class LLMUsage: total_tokens: int | None = None +@dataclass(frozen=True) +class LLMCost: + """Per-component cost breakdown, from the StandardLoggingPayload + ``cost_breakdown`` (``litellm.types.utils.CostBreakdown``). + + Each field is the USD cost of one component, or ``None`` when the source did + not report it — so the mapper omits absent components instead of emitting 0. + The final (post-discount/post-margin) total is carried separately on + ``LLMCallSpanData.response_cost``. Free-form ``additional_costs`` are not + surfaced here: span attributes are scalar and there is no agreed key shape + for them yet. + """ + + input: float | None = None + output: float | None = None + cache_read: float | None = None + cache_creation: float | None = None + tool_usage: float | None = None + original: float | None = None + discount_amount: float | None = None + discount_percent: float | None = None + margin_fixed_amount: float | None = None + margin_percent: float | None = None + margin_total_amount: float | None = None + + @classmethod + def from_breakdown(cls, breakdown: Mapping[str, object] | None) -> "LLMCost": + b = breakdown or {} + return cls( + input=as_float(b.get("input_cost")), + output=as_float(b.get("output_cost")), + cache_read=as_float(b.get("cache_read_cost")), + cache_creation=as_float(b.get("cache_creation_cost")), + tool_usage=as_float(b.get("tool_usage_cost")), + original=as_float(b.get("original_cost")), + discount_amount=as_float(b.get("discount_amount")), + discount_percent=as_float(b.get("discount_percent")), + margin_fixed_amount=as_float(b.get("margin_fixed_amount")), + margin_percent=as_float(b.get("margin_percent")), + margin_total_amount=as_float(b.get("margin_total_amount")), + ) + + @dataclass(frozen=True) class SpanError: error_type: str | None = None @@ -255,6 +299,7 @@ class LLMCallSpanData: server: ServerInfo | None identity: RequestIdentity is_streaming: bool | None = None + cost: LLMCost = field(default_factory=LLMCost) tools: tuple[ToolDefinition, ...] = () # Raw messages and response, needed by vendor mappers (OpenInference, # Langfuse, Weave) that stamp message-level attributes. ``messages_in`` is @@ -302,6 +347,9 @@ class LLMCallSpanData: finish_reasons=finish_reasons, error=_parse_error(payload), response_cost=as_float(payload.get("response_cost")), + cost=LLMCost.from_breakdown( + cast("Mapping[str, object] | None", payload.get("cost_breakdown")) + ), server=ServerInfo.from_api_base(context.api_base), identity=context.identity, is_streaming=as_bool(payload.get("stream")), diff --git a/litellm/integrations/otel/plumbing/providers.py b/litellm/integrations/otel/plumbing/providers.py index 40a0e41b905..4c98802479a 100644 --- a/litellm/integrations/otel/plumbing/providers.py +++ b/litellm/integrations/otel/plumbing/providers.py @@ -17,6 +17,7 @@ from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( ) from opentelemetry.trace import Span, SpanKind, Tracer +from litellm._version import version as litellm_version from litellm.integrations.otel.model.config import ExporterSpec, OpenTelemetryV2Config from litellm.integrations.otel.model.semconv import LiteLLM from litellm.integrations.otel.model.spans import LiteLLMSpanKind @@ -207,7 +208,10 @@ def build_tracer_provider( def get_tracer(provider: TracerProvider, name: str = "litellm") -> Tracer: - return provider.get_tracer(name) + # Stamp the instrumentation scope with the LiteLLM package version so every + # emitted span carries a deterministic ``scope.version`` (the standard OTel + # location for the emitting library's version) for downstream consumers. + return provider.get_tracer(name, litellm_version) def in_memory_provider( diff --git a/litellm/litellm_core_utils/get_supported_openai_params.py b/litellm/litellm_core_utils/get_supported_openai_params.py index 23b51faafc7..65c238344e9 100644 --- a/litellm/litellm_core_utils/get_supported_openai_params.py +++ b/litellm/litellm_core_utils/get_supported_openai_params.py @@ -295,6 +295,15 @@ def get_supported_openai_params( # noqa: PLR0915 elif custom_llm_provider == "predibase": return litellm.PredibaseConfig().get_supported_openai_params(model=model) elif custom_llm_provider == "voyage": + if ( + request_type == "embeddings" + and litellm.VoyageMultimodalEmbeddingConfig.is_multimodal_embeddings(model) + ): + return ( + litellm.VoyageMultimodalEmbeddingConfig().get_supported_openai_params( + model=model + ) + ) return litellm.VoyageEmbeddingConfig().get_supported_openai_params(model=model) elif custom_llm_provider == "infinity": return litellm.InfinityEmbeddingConfig().get_supported_openai_params( diff --git a/litellm/litellm_core_utils/initialize_dynamic_callback_params.py b/litellm/litellm_core_utils/initialize_dynamic_callback_params.py index a89dae52316..949076aabf3 100644 --- a/litellm/litellm_core_utils/initialize_dynamic_callback_params.py +++ b/litellm/litellm_core_utils/initialize_dynamic_callback_params.py @@ -53,11 +53,19 @@ _supported_callback_params = [ "braintrust_host", "slack_webhook_url", "lunary_public_key", + "dd_api_key", + "dd_site", + "dd_agent_host", + "dd_agent_port", ] _request_blocked_callback_params = { "gcs_bucket_name", "gcs_path_service_account", + "dd_api_key", + "dd_site", + "dd_agent_host", + "dd_agent_port", } diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index b2db334d5ff..2cc8e794d40 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -381,13 +381,14 @@ class Logging(LiteLLMLoggingBaseClass): List[Union[str, Callable, CustomLogger]] ] = dynamic_async_failure_callbacks - # Process dynamic callbacks - self.process_dynamic_callbacks() - ## DYNAMIC LANGFUSE / GCS / logging callback KEYS ## self.standard_callback_dynamic_params: StandardCallbackDynamicParams = ( self.initialize_standard_callback_dynamic_params(kwargs) ) + + # Process dynamic callbacks (after standard_callback_dynamic_params is initialized, + # so team-scoped credentials are available for callback initialization) + self.process_dynamic_callbacks() self.standard_built_in_tools_params: StandardBuiltInToolsParams = ( self.initialize_standard_built_in_tools_params(kwargs) ) @@ -482,8 +483,21 @@ class Logging(LiteLLMLoggingBaseClass): isinstance(callback, str) and callback in litellm._known_custom_logger_compatible_callbacks ): + # For callbacks that support team-scoped credentials (e.g. datadog), + # pass only the relevant dynamic params as custom_logger_init_args. + _custom_logger_init_args: Optional[dict] = None + if callback == "datadog": + _custom_logger_init_args = { + k: v + for k, v in self.standard_callback_dynamic_params.items() + if k.startswith("dd_") + } + callback_class = _init_custom_logger_compatible_class( - callback, internal_usage_cache=None, llm_router=None # type: ignore + callback, # type: ignore[arg-type] + internal_usage_cache=None, + llm_router=None, # type: ignore + custom_logger_init_args=_custom_logger_init_args, ) if callback_class is not None: processed_list.append(callback_class) @@ -3941,6 +3955,24 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 _in_memory_loggers.append(_prometheus_logger) return _prometheus_logger # type: ignore elif logging_integration == "datadog": + # Check if team-scoped credentials are provided + _dd_api_key = custom_logger_init_args.get("dd_api_key") + _dd_site = custom_logger_init_args.get("dd_site") + _dd_agent_host = custom_logger_init_args.get("dd_agent_host") + _dd_agent_port = custom_logger_init_args.get("dd_agent_port") + + if _dd_api_key or _dd_site or _dd_agent_host: + # Team-scoped credentials: use DynamicLoggingCache for per-credential isolation + from litellm.integrations.datadog.datadog_team_handler import ( + DataDogHandler, + ) + + return DataDogHandler.get_datadog_logger_for_request( + standard_callback_dynamic_params=custom_logger_init_args, # type: ignore + in_memory_dynamic_logger_cache=in_memory_dynamic_logger_cache, + ) + + # Global (env-var based): reuse cached instance for callback in _in_memory_loggers: if isinstance(callback, DataDogLogger): return callback # type: ignore diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index b1b06829387..2c9ea187912 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -861,14 +861,58 @@ class BaseAWSLLM: with tracer.trace("boto3.client(sts)"): sts_client = boto3.client("sts", **sts_client_kwargs) + # The session policy is an IAM PERMISSION CEILING — effective + # permissions are the intersection of the role's identity policies + # and this policy. Any action not listed here is silently denied + # even when the IAM role grants it. So every Bedrock route we + # support needs a matching action statement, or it 403s on OIDC + # auth only (static creds + IRSA take other code paths). # https://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRoleWithWebIdentity.html # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sts/client/assume_role_with_web_identity.html + bedrock_session_policy = { + "Version": "2012-10-17", + "Statement": [ + { + "Sid": "BedrockLiteLLM", + "Effect": "Allow", + "Action": [ + "bedrock:InvokeModel", + "bedrock:InvokeModelWithResponseStream", + "bedrock:ApplyGuardrail", + "bedrock:GetGuardrail", + "bedrock:ListGuardrails", + ], + "Resource": "*", + "Condition": {"Bool": {"aws:SecureTransport": "true"}}, + }, + # Claude Platform on AWS (added by #27678 for the + # ``bedrock/claude_platform/`` route) lives under + # a separate IAM action namespace; without these entries + # the OIDC path 403s on every claude_platform request + # even with a fully permissive identity policy (#30200). + { + "Sid": "ClaudePlatformLiteLLM", + "Effect": "Allow", + "Action": [ + "aws-external-anthropic:CreateInference", + "aws-external-anthropic:CreateBatchInference", + "aws-external-anthropic:CancelBatchInference", + "aws-external-anthropic:DeleteBatchInference", + "aws-external-anthropic:CountTokens", + "aws-external-anthropic:Get*", + "aws-external-anthropic:List*", + ], + "Resource": "*", + "Condition": {"Bool": {"aws:SecureTransport": "true"}}, + }, + ], + } assume_role_params = { "RoleArn": aws_role_name, "RoleSessionName": aws_session_name, "WebIdentityToken": oidc_token, "DurationSeconds": 3600, - "Policy": '{"Version":"2012-10-17","Statement":[{"Sid":"BedrockLiteLLM","Effect":"Allow","Action":["bedrock:InvokeModel","bedrock:InvokeModelWithResponseStream","bedrock:ApplyGuardrail","bedrock:GetGuardrail","bedrock:ListGuardrails"],"Resource":"*","Condition":{"Bool":{"aws:SecureTransport":"true"}}}]}', + "Policy": json.dumps(bedrock_session_policy, separators=(",", ":")), } # Add ExternalId parameter if provided diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index 388947a4e9b..7e1020000f4 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -32,7 +32,7 @@ def make_sync_call( logging_obj: LiteLLMLoggingObject, json_mode: Optional[bool] = False, fake_stream: bool = False, - stream_chunk_size: int = 1024, + stream_chunk_size: Optional[int] = None, ): if client is None: client = _get_httpx_client() # Create a new client if none provided @@ -108,7 +108,7 @@ class BedrockConverseLLM(BaseAWSLLM): fake_stream: bool = False, json_mode: Optional[bool] = False, api_key: Optional[str] = None, - stream_chunk_size: int = 1024, + stream_chunk_size: Optional[int] = None, ) -> CustomStreamWrapper: request_data = await litellm.AmazonConverseConfig()._async_transform_request( model=model, @@ -268,7 +268,7 @@ class BedrockConverseLLM(BaseAWSLLM): ): ## SETUP ## stream = optional_params.pop("stream", None) - stream_chunk_size = optional_params.pop("stream_chunk_size", 1024) + stream_chunk_size = optional_params.pop("stream_chunk_size", None) unencoded_model_id = optional_params.pop("model_id", None) fake_stream = optional_params.pop("fake_stream", False) json_mode = optional_params.get("json_mode", False) diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index 7a9916f1f31..0a1322a751e 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -197,7 +197,7 @@ async def make_call( fake_stream: bool = False, json_mode: Optional[bool] = False, bedrock_invoke_provider: Optional[litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL] = None, - stream_chunk_size: int = 1024, + stream_chunk_size: Optional[int] = None, ): try: if client is None: @@ -294,7 +294,7 @@ def make_sync_call( fake_stream: bool = False, json_mode: Optional[bool] = False, bedrock_invoke_provider: Optional[litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL] = None, - stream_chunk_size: int = 1024, + stream_chunk_size: Optional[int] = None, ): try: if client is None: @@ -790,7 +790,7 @@ class BedrockLLM(BaseAWSLLM): ## SETUP ## stream = optional_params.pop("stream", None) - stream_chunk_size = optional_params.pop("stream_chunk_size", 1024) + stream_chunk_size = optional_params.pop("stream_chunk_size", None) provider = self.get_bedrock_invoke_provider(model) modelId = self.get_bedrock_model_id( @@ -1203,7 +1203,7 @@ class BedrockLLM(BaseAWSLLM): extra_headers: Optional[dict] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[AsyncHTTPHandler] = None, - stream_chunk_size: int = 1024, + stream_chunk_size: Optional[int] = None, ) -> Union[ModelResponse, CustomStreamWrapper]: transformed_request = ( await litellm.AmazonAnthropicClaudeConfig().async_transform_request( @@ -1350,7 +1350,7 @@ class BedrockLLM(BaseAWSLLM): logger_fn=None, headers={}, client: Optional[AsyncHTTPHandler] = None, - stream_chunk_size: int = 1024, + stream_chunk_size: Optional[int] = None, ) -> CustomStreamWrapper: # The call is not made here; instead, we prepare the necessary objects for the stream. diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py index 4887cbd23be..79153c3ceff 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py @@ -215,6 +215,7 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): anthropic_request.pop("model", None) anthropic_request.pop("stream", None) + anthropic_request.pop("stream_chunk_size", None) output_format = anthropic_request.pop("output_format", None) output_config_format = pop_bedrock_invoke_output_config_format( anthropic_request diff --git a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py index 43850440072..6bb2da1ad44 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py @@ -150,6 +150,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): ) -> dict: ## SETUP ## stream = optional_params.pop("stream", None) + optional_params.pop("stream_chunk_size", None) custom_prompt_dict: dict = litellm_params.pop("custom_prompt_dict", None) or {} hf_model_name = litellm_params.get("hf_model_name", None) diff --git a/litellm/llms/openai_like/providers.json b/litellm/llms/openai_like/providers.json index 13d22488838..303e9ba8f9e 100644 --- a/litellm/llms/openai_like/providers.json +++ b/litellm/llms/openai_like/providers.json @@ -131,7 +131,8 @@ "base_class": "openai_gpt", "param_mappings": { "max_completion_tokens": "max_tokens" - } + }, + "supported_endpoints": ["/v1/chat/completions", "/v1/responses"] }, "parasail": { "base_url": "https://api.parasail.io/v1", @@ -141,5 +142,14 @@ "special_handling": { "force_store_false": true } + }, + "empiriolabs": { + "base_url": "https://api.empiriolabs.ai/v1", + "api_key_env": "EMPIRIOLABS_API_KEY", + "api_base_env": "EMPIRIOLABS_API_BASE", + "param_mappings": { + "max_completion_tokens": "max_tokens" + }, + "supported_endpoints": ["/v1/chat/completions", "/v1/responses"] } } diff --git a/litellm/llms/snowflake/chat/transformation.py b/litellm/llms/snowflake/chat/transformation.py index 23bb6f44757..ed30522876a 100644 --- a/litellm/llms/snowflake/chat/transformation.py +++ b/litellm/llms/snowflake/chat/transformation.py @@ -1,17 +1,32 @@ """ -Support for Snowflake REST API +Snowflake Cortex REST API — Chat Transformation + +Routes to native Cortex REST API endpoints based on model: + - Claude models → POST /api/v2/cortex/v1/messages (Anthropic format) + - All other models → POST /api/v2/cortex/v1/chat/completions (OpenAI format) + +Ref: https://docs.snowflake.com/en/user-guide/snowflake-cortex/cortex-rest-api """ import json -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union +from typing import TYPE_CHECKING, Any, Dict, List, Optional import httpx -from litellm.types.llms.openai import AllMessageValues -from litellm.types.utils import ChatCompletionMessageToolCall, Function, ModelResponse +from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk +from litellm.types.utils import ( + ChatCompletionMessageToolCall, + ChatCompletionUsageBlock, + Choices, + Function, + GenericStreamingChunk, + Message, + ModelResponse, + Usage, +) +from ...base_llm.base_model_iterator import BaseModelResponseIterator from ...openai_like.chat.transformation import OpenAIGPTConfig - from ..utils import SnowflakeBaseConfig if TYPE_CHECKING: @@ -21,69 +36,343 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +ANTHROPIC_VERSION = "2023-06-01" + +_CLAUDE_MODEL_PREFIXES = ( + "claude-", + "claude_", +) + + +def _is_claude_model(model: str) -> bool: + """Return True if model name (after stripping snowflake/ prefix) is a Claude model.""" + name = model.lower().removeprefix("snowflake/") + return any(name.startswith(p) for p in _CLAUDE_MODEL_PREFIXES) + class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig): """ - Reference: https://docs.snowflake.com/en/user-guide/snowflake-cortex/cortex-llm-rest-api + Snowflake Cortex REST API — unified provider. - Snowflake Cortex LLM REST API supports function calling with specific models (e.g., Claude 3.5 Sonnet). - This config handles transformation between OpenAI format and Snowflake's tool_spec format. + Auto-routes based on model name: + - Claude models → /api/v2/cortex/v1/messages (Anthropic Messages format) + - All others → /api/v2/cortex/v1/chat/completions (OpenAI format) + + Auth: + PAT: api_key="pat/" → X-Snowflake-Authorization-Token-Type: PROGRAMMATIC_ACCESS_TOKEN + JWT: api_key="" → X-Snowflake-Authorization-Token-Type: KEYPAIR_JWT """ @classmethod def get_config(cls): return super().get_config() - def _transform_tool_calls_from_snowflake_to_openai( - self, content_list: List[Dict[str, Any]] - ) -> Tuple[str, Optional[List[ChatCompletionMessageToolCall]]]: + def get_supported_openai_params(self, model: str) -> List[str]: + params = [ + "temperature", + "max_tokens", + "max_completion_tokens", + "top_p", + "stream", + "tools", + "tool_choice", + ] + if _is_claude_model(model): + params.append("thinking") + return params + + def get_complete_url( + self, + api_base: Optional[str], + api_key: Optional[str], + model: str, + optional_params: dict, + litellm_params: dict, + stream: Optional[bool] = None, + ) -> str: + api_base = self._get_api_base(api_base, optional_params) + if _is_claude_model(model): + return f"{api_base}/cortex/v1/messages" + return f"{api_base}/cortex/v1/chat/completions" + + def validate_environment( + self, + headers: dict, + model: str, + messages: List[AllMessageValues], + optional_params: dict, + litellm_params: dict, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + ) -> dict: + headers = super().validate_environment( + headers=headers, + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + api_key=api_key, + api_base=api_base, + ) + if _is_claude_model(model): + headers["anthropic-version"] = ANTHROPIC_VERSION + return headers + + def _transform_tools_to_anthropic(self, tools: List[Dict]) -> List[Dict]: """ - Transform Snowflake tool calls to OpenAI format. + Convert tools from OpenAI format to Anthropic format. - Args: - content_list: Snowflake's content_list array containing text and tool_use items + OpenAI: {"type": "function", "function": {"name": ..., "parameters": {...}}} + Anthropic: {"name": ..., "description": ..., "input_schema": {...}} + """ + anthropic_tools = [] + for tool in tools: + if tool.get("type") == "function" and "function" in tool: + func = tool["function"] + anthropic_tool: Dict[str, Any] = { + "name": func.get("name", ""), + } + if "description" in func: + anthropic_tool["description"] = func["description"] + if "parameters" in func: + anthropic_tool["input_schema"] = func["parameters"] + else: + anthropic_tool["input_schema"] = { + "type": "object", + "properties": {}, + } + anthropic_tools.append(anthropic_tool) + else: + anthropic_tools.append(tool) + return anthropic_tools - Returns: - Tuple of (text_content, tool_calls) + def _extract_system_and_messages( + self, messages: List[AllMessageValues] + ) -> tuple[Optional[str], List[Dict]]: + """ + Split messages into system prompt and conversation turns for Anthropic format. - Snowflake format in content_list: - { - "type": "tool_use", - "tool_use": { - "tool_use_id": "tooluse_...", - "name": "get_weather", - "input": {"location": "Paris"} - } + - system messages → collected and joined (preserves guardrail prompts) + - assistant messages with tool_calls → tool_use content blocks + - tool role messages → user role with tool_result content blocks + """ + system_parts: List[str] = [] + conversation: List[Dict] = [] + + for msg in messages: + if isinstance(msg, dict): + role = msg.get("role", "") + content: Any = msg.get("content", "") + else: + role = getattr(msg, "role", "") + content = getattr(msg, "content", "") + + if role == "system": + if isinstance(content, str) and content: + system_parts.append(content) + elif isinstance(content, list): + system_parts.append( + "\n".join( + b.get("text", "") + for b in content + if b.get("type") == "text" + ) + ) + elif role == "assistant": + tool_calls = ( + msg.get("tool_calls") + if isinstance(msg, dict) + else getattr(msg, "tool_calls", None) + ) + if tool_calls: # type: ignore[truthy-bool] + content_blocks: List[Dict[str, Any]] = [] + if content: + content_blocks.append({"type": "text", "text": content}) + for tc in tool_calls: # type: ignore[attr-defined] + func = ( + tc.get("function", {}) + if isinstance(tc, dict) + else getattr(tc, "function", {}) + ) + tc_id = ( + tc.get("id", "") + if isinstance(tc, dict) + else getattr(tc, "id", "") + ) + func_name = ( + func.get("name", "") + if isinstance(func, dict) + else getattr(func, "name", "") + ) + func_args = ( + func.get("arguments", "{}") + if isinstance(func, dict) + else getattr(func, "arguments", "{}") + ) + try: + input_data = ( + json.loads(func_args) + if isinstance(func_args, str) + else func_args + ) + except (json.JSONDecodeError, TypeError): + input_data = {} + content_blocks.append( + { + "type": "tool_use", + "id": tc_id, + "name": func_name, + "input": input_data, + } + ) + conversation.append( + {"role": "assistant", "content": content_blocks} + ) + else: + conversation.append({"role": "assistant", "content": content}) + elif role == "tool": + tool_call_id = ( + msg.get("tool_call_id", "") + if isinstance(msg, dict) + else getattr(msg, "tool_call_id", "") + ) + tool_content = ( + content if isinstance(content, str) else json.dumps(content) + ) + tool_result_block = { + "type": "tool_result", + "tool_use_id": tool_call_id, + "content": tool_content, + } + if ( + conversation + and conversation[-1]["role"] == "user" + and isinstance(conversation[-1]["content"], list) + and conversation[-1]["content"] + and conversation[-1]["content"][0].get("type") == "tool_result" + ): + conversation[-1]["content"].append(tool_result_block) + else: + conversation.append( + {"role": "user", "content": [tool_result_block]} + ) + else: + conversation.append({"role": role, "content": content}) + + system: Optional[str] = "\n\n".join(system_parts) if system_parts else None + return system, conversation + + def transform_request( + self, + model: str, + messages: List[AllMessageValues], + optional_params: dict, + litellm_params: dict, + headers: dict, + ) -> dict: + stream: bool = optional_params.pop("stream", False) or False + extra_body = optional_params.pop("extra_body", {}) + + if _is_claude_model(model): + return self._transform_request_anthropic( + model, messages, optional_params, stream, extra_body + ) + return self._transform_request_openai( + model, messages, optional_params, stream, extra_body + ) + + def _transform_request_openai( + self, + model: str, + messages: List[AllMessageValues], + optional_params: dict, + stream: bool, + extra_body: dict, + ) -> dict: + """OpenAI format for /chat/completions endpoint.""" + max_tokens = optional_params.pop("max_tokens", None) + max_completion_tokens = optional_params.pop("max_completion_tokens", None) + resolved_max = max_completion_tokens or max_tokens + + body: dict = { + "model": model.removeprefix("snowflake/"), + "messages": messages, + "stream": stream, + **optional_params, + **extra_body, } - OpenAI format (returned tool_calls): - ChatCompletionMessageToolCall( - id="tooluse_...", - type="function", - function=Function(name="get_weather", arguments='{"location": "Paris"}') - ) + if resolved_max is not None: + body["max_completion_tokens"] = resolved_max + + return body + + def _transform_tool_choice_to_anthropic(self, tool_choice: Any) -> Dict[str, Any]: """ - text_content = "" - tool_calls: List[ChatCompletionMessageToolCall] = [] + Convert tool_choice from OpenAI format to Anthropic format. - for idx, content_item in enumerate(content_list): - if content_item.get("type") == "text": - text_content += content_item.get("text", "") + OpenAI string values: "auto", "required", "none" + OpenAI dict: {"type": "function", "function": {"name": "..."}} + Anthropic: {"type": "auto"}, {"type": "any"}, {"type": "tool", "name": "..."} + """ + if isinstance(tool_choice, str): + mapping = { + "auto": {"type": "auto"}, + "required": {"type": "any"}, + "none": {"type": "none"}, + } + return mapping.get(tool_choice, {"type": "auto"}) + elif isinstance(tool_choice, dict): + if tool_choice.get("type") == "function": + func = tool_choice.get("function", {}) + return {"type": "tool", "name": func.get("name", "")} + return tool_choice + return {"type": "auto"} - ## TOOL CALLING - elif content_item.get("type") == "tool_use": - tool_use_data = content_item.get("tool_use", {}) - tool_call = ChatCompletionMessageToolCall( - id=tool_use_data.get("tool_use_id", ""), - type="function", - function=Function( - name=tool_use_data.get("name", ""), - arguments=json.dumps(tool_use_data.get("input", {})), - ), - ) - tool_calls.append(tool_call) + def _transform_request_anthropic( + self, + model: str, + messages: List[AllMessageValues], + optional_params: dict, + stream: bool, + extra_body: dict, + ) -> dict: + """Anthropic Messages format for /messages endpoint.""" + system, conversation = self._extract_system_and_messages(messages) - return text_content, tool_calls if tool_calls else None + if "tools" in optional_params: + optional_params["tools"] = self._transform_tools_to_anthropic( + optional_params["tools"] + ) + + if "tool_choice" in optional_params: + optional_params["tool_choice"] = self._transform_tool_choice_to_anthropic( + optional_params["tool_choice"] + ) + + max_completion_tokens = optional_params.pop("max_completion_tokens", None) + if max_completion_tokens and "max_tokens" not in optional_params: + optional_params["max_tokens"] = max_completion_tokens + + model_name = model.removeprefix("snowflake/") + + body: Dict[str, Any] = { + "model": model_name, + "messages": conversation, + "stream": stream, + **optional_params, + **extra_body, + } + + if system is not None: + body["system"] = system + + if "max_tokens" not in body: + body["max_tokens"] = ( + 4096 # reasonable default; Anthropic API max varies by model + ) + + return body def transform_response( self, @@ -99,6 +388,24 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig): api_key: Optional[str] = None, json_mode: Optional[bool] = None, ) -> ModelResponse: + if _is_claude_model(model): + return self._transform_response_anthropic( + model, raw_response, model_response, logging_obj, request_data, messages + ) + return self._transform_response_openai( + model, raw_response, model_response, logging_obj, request_data, messages + ) + + def _transform_response_openai( + self, + model: str, + raw_response: httpx.Response, + model_response: ModelResponse, + logging_obj: LiteLLMLoggingObj, + request_data: dict, + messages: List[AllMessageValues], + ) -> ModelResponse: + """Parse standard OpenAI chat completions response.""" response_json = raw_response.json() logging_obj.post_call( @@ -108,180 +415,278 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig): additional_args={"complete_input_dict": request_data}, ) - ## RESPONSE TRANSFORMATION - # Snowflake returns content_list (not content) with tool_use objects - # We need to transform this to OpenAI's format with content + tool_calls - if "choices" in response_json and len(response_json["choices"]) > 0: - choice = response_json["choices"][0] - if "message" in choice and "content_list" in choice["message"]: - content_list = choice["message"]["content_list"] - ( - text_content, - tool_calls, - ) = self._transform_tool_calls_from_snowflake_to_openai(content_list) - - # Update the choice message with OpenAI format - choice["message"]["content"] = text_content - if tool_calls: - choice["message"]["tool_calls"] = tool_calls - - # Remove Snowflake-specific content_list - del choice["message"]["content_list"] - returned_response = ModelResponse(**response_json) - returned_response.model = "snowflake/" + (returned_response.model or "") if model is not None: returned_response._hidden_params["model"] = model + return returned_response - def get_complete_url( - self, - api_base: Optional[str], - api_key: Optional[str], - model: str, - optional_params: dict, - litellm_params: dict, - stream: Optional[bool] = None, - ) -> str: - """ - If api_base is not provided, use the default DeepSeek /chat/completions endpoint. - """ - - api_base = self._get_api_base(api_base, optional_params) - - return f"{api_base}/cortex/inference:complete" - - def _transform_tools(self, tools: List[Dict[str, Any]]) -> List[Dict[str, Any]]: - """ - Transform OpenAI tool format to Snowflake tool format. - - Args: - tools: List of tools in OpenAI format - - Returns: - List of tools in Snowflake format - - OpenAI format: - { - "type": "function", - "function": { - "name": "get_weather", - "description": "...", - "parameters": {...} - } - } - - Snowflake format: - { - "tool_spec": { - "type": "generic", - "name": "get_weather", - "description": "...", - "input_schema": {...} - } - } - """ - snowflake_tools: List[Dict[str, Any]] = [] - for tool in tools: - if tool.get("type") == "function": - function = tool.get("function", {}) - snowflake_tool: Dict[str, Any] = { - "tool_spec": { - "type": "generic", - "name": function.get("name"), - "input_schema": function.get( - "parameters", - {"type": "object", "properties": {}}, - ), - } - } - # Add description if present - if "description" in function: - snowflake_tool["tool_spec"]["description"] = function["description"] - - snowflake_tools.append(snowflake_tool) - - return snowflake_tools - - def _transform_tool_choice( - self, tool_choice: Union[str, Dict[str, Any]] - ) -> Dict[str, Any]: - """ - Transform OpenAI tool_choice format to Snowflake format. - - Snowflake requires tool_choice to be an object, not a string. - Ref: https://docs.snowflake.com/en/developer-guide/snowflake-rest-api/reference/cortex-inference#post--api-v2-cortex-inference-complete-req-body-schema - - Args: - tool_choice: Tool choice in OpenAI format (str or dict) - - Returns: - Tool choice in Snowflake format (always an object, never a string) - - OpenAI format (string): - "auto", "required", "none" - - OpenAI format (dict): - {"type": "function", "function": {"name": "get_weather"}} - - Snowflake format: - {"type": "auto"} / {"type": "any"} / {"type": "none"} - {"type": "tool", "name": ["get_weather"]} - - Snowflake's API (like Anthropic) requires tool_choice as an object - with a "type" field, not as a bare string. - """ - if isinstance(tool_choice, str): - # Snowflake requires object format, not string. - # Map OpenAI string values to Snowflake object format. - # "required" maps to "any" (Snowflake/Anthropic convention). - _type_map = { - "auto": "auto", - "required": "any", - "none": "none", - } - mapped_type = _type_map.get(tool_choice, tool_choice) - return {"type": mapped_type} - - if isinstance(tool_choice, dict): - if tool_choice.get("type") == "function": - function_name = tool_choice.get("function", {}).get("name") - if function_name: - return { - "type": "tool", - "name": [function_name], # Snowflake expects array - } - - return tool_choice - - def transform_request( + def _transform_response_anthropic( self, model: str, + raw_response: httpx.Response, + model_response: ModelResponse, + logging_obj: LiteLLMLoggingObj, + request_data: dict, messages: List[AllMessageValues], - optional_params: dict, - litellm_params: dict, - headers: dict, - ) -> dict: - stream: bool = optional_params.pop("stream", None) or False - extra_body = optional_params.pop("extra_body", {}) + ) -> ModelResponse: + """Parse Anthropic Messages response into OpenAI format.""" + response_json = raw_response.json() - ## TOOL CALLING - # Transform tools from OpenAI format to Snowflake's tool_spec format - tools = optional_params.pop("tools", None) - if tools: - optional_params["tools"] = self._transform_tools(tools) + logging_obj.post_call( + input=messages, + api_key="", + original_response=response_json, + additional_args={"complete_input_dict": request_data}, + ) - # Transform tool_choice from OpenAI format to Snowflake's tool name array format - tool_choice = optional_params.pop("tool_choice", None) - if tool_choice: - optional_params["tool_choice"] = self._transform_tool_choice(tool_choice) + text_content = "" + tool_calls = [] - return { - "model": model, - "messages": messages, - "stream": stream, - **optional_params, - **extra_body, + for block in response_json.get("content", []): + if block.get("type") == "text": + text_content += block.get("text", "") + elif block.get("type") == "tool_use": + tool_calls.append( + ChatCompletionMessageToolCall( + id=block.get("id", ""), + type="function", + function=Function( + name=block.get("name", ""), + arguments=json.dumps(block.get("input", {})), + ), + ) + ) + + _stop_reason_map = { + "end_turn": "stop", + "max_tokens": "length", + "tool_use": "tool_calls", + "stop_sequence": "stop", } + finish_reason = _stop_reason_map.get( + response_json.get("stop_reason", "end_turn"), "stop" + ) + + message = Message(content=text_content or None, role="assistant") + if tool_calls: + message.tool_calls = tool_calls + + choice = Choices( + finish_reason=finish_reason, + index=0, + message=message, + ) + + usage_data = response_json.get("usage", {}) + usage = Usage( + prompt_tokens=usage_data.get("input_tokens", 0), + completion_tokens=usage_data.get("output_tokens", 0), + total_tokens=usage_data.get("input_tokens", 0) + + usage_data.get("output_tokens", 0), + ) + + model_response.choices = [choice] + model_response.usage = usage # type: ignore[attr-defined] + model_response.model = "snowflake/" + response_json.get("model", model) + model_response.id = response_json.get("id", "") + + if model is not None: + model_response._hidden_params["model"] = model + + return model_response + + def get_model_response_iterator( + self, + streaming_response: Any, + sync_stream: bool, + json_mode: Optional[bool] = False, + ) -> Any: + return SnowflakeStreamingHandler( + streaming_response=streaming_response, + sync_stream=sync_stream, + json_mode=json_mode, + ) + + +class SnowflakeStreamingHandler(BaseModelResponseIterator): + """ + Parse streaming events from both Snowflake endpoints. + + - /chat/completions: OpenAI SSE format (has "choices" key) + - /messages: Anthropic SSE format (has "type" key like content_block_delta) + """ + + def __init__( + self, + streaming_response: Any, + sync_stream: bool, + json_mode: Optional[bool] = False, + ): + super().__init__(streaming_response=streaming_response, sync_stream=sync_stream) + self._tool_index = 0 + self._tool_id = "" + self._tool_name = "" + self._input_tokens = 0 + + def chunk_parser(self, chunk: dict) -> GenericStreamingChunk: + if "choices" in chunk: + return self._parse_openai_chunk(chunk) + return self._parse_anthropic_chunk(chunk) + + def _parse_openai_chunk(self, chunk: dict) -> GenericStreamingChunk: + choices = chunk.get("choices", []) + if not choices: + return GenericStreamingChunk( + text="", + is_finished=False, + finish_reason="", + usage=None, + index=0, + tool_use=None, + ) + + choice = choices[0] + delta = choice.get("delta", {}) + finish_reason = choice.get("finish_reason") or "" + text = delta.get("content") or "" + + tool_use = None + tool_calls = delta.get("tool_calls") + if tool_calls: + tc = tool_calls[0] + func = tc.get("function", {}) + tool_use = ChatCompletionToolCallChunk( + id=tc.get("id", ""), + type="function", + function={ + "name": func.get("name", ""), + "arguments": func.get("arguments", ""), + }, + index=tc.get("index", 0), + ) + + return GenericStreamingChunk( + text=text, + is_finished=finish_reason != "", + finish_reason=finish_reason, + usage=None, + index=choice.get("index", 0), + tool_use=tool_use, + ) + + def _parse_anthropic_chunk(self, chunk: dict) -> GenericStreamingChunk: + event_type = chunk.get("type", "") + + if event_type == "message_start": + message = chunk.get("message", {}) + usage_data = message.get("usage", {}) + self._input_tokens = usage_data.get("input_tokens", 0) + return GenericStreamingChunk( + text="", + is_finished=False, + finish_reason="", + usage=None, + index=0, + tool_use=None, + ) + + elif event_type == "content_block_delta": + delta = chunk.get("delta", {}) + delta_type = delta.get("type", "") + + if delta_type == "text_delta": + return GenericStreamingChunk( + text=delta.get("text", ""), + is_finished=False, + finish_reason="", + usage=None, + index=chunk.get("index", 0), + tool_use=None, + ) + elif delta_type == "input_json_delta": + return GenericStreamingChunk( + text="", + is_finished=False, + finish_reason="", + usage=None, + index=chunk.get("index", 0), + tool_use=ChatCompletionToolCallChunk( + id=self._tool_id, + type="function", + function={ + "name": self._tool_name, + "arguments": delta.get("partial_json", ""), + }, + index=self._tool_index, + ), + ) + + elif event_type == "content_block_start": + content_block = chunk.get("content_block", {}) + if content_block.get("type") == "tool_use": + self._tool_id = content_block.get("id", "") + self._tool_name = content_block.get("name", "") + self._tool_index = chunk.get("index", 0) + return GenericStreamingChunk( + text="", + is_finished=False, + finish_reason="", + usage=None, + index=chunk.get("index", 0), + tool_use=ChatCompletionToolCallChunk( + id=self._tool_id, + type="function", + function={"name": self._tool_name, "arguments": ""}, + index=self._tool_index, + ), + ) + + elif event_type == "message_delta": + delta = chunk.get("delta", {}) + stop_reason = delta.get("stop_reason", "") + usage_data = chunk.get("usage", {}) + _stop_map = { + "end_turn": "stop", + "max_tokens": "length", + "tool_use": "tool_calls", + "stop_sequence": "stop", + } + usage = None + if usage_data or self._input_tokens: + output_t = usage_data.get("output_tokens", 0) + input_t = self._input_tokens or usage_data.get("input_tokens", 0) + usage = ChatCompletionUsageBlock( + prompt_tokens=input_t, + completion_tokens=output_t, + total_tokens=input_t + output_t, + ) + return GenericStreamingChunk( + text="", + is_finished=True, + finish_reason=_stop_map.get(stop_reason, "stop"), + usage=usage, + index=0, + tool_use=None, + ) + + elif event_type == "message_stop": + return GenericStreamingChunk( + text="", + is_finished=True, + finish_reason="stop", + usage=None, + index=0, + tool_use=None, + ) + + return GenericStreamingChunk( + text="", + is_finished=False, + finish_reason="", + usage=None, + index=0, + tool_use=None, + ) diff --git a/litellm/llms/voyage/embedding/transformation_multimodal.py b/litellm/llms/voyage/embedding/transformation_multimodal.py new file mode 100644 index 00000000000..55e221b065b --- /dev/null +++ b/litellm/llms/voyage/embedding/transformation_multimodal.py @@ -0,0 +1,183 @@ +""" +Transform request/response for Voyage multimodal embeddings. + +Voyage multimodal models use /v1/multimodalembeddings and accept `inputs` +containing content blocks, unlike standard Voyage embeddings which use +/v1/embeddings and a string/list `input` field. +""" + +from typing import Any, Dict, List, Optional, Union + +import httpx + +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig +from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.openai import AllEmbeddingInputValues, AllMessageValues +from litellm.types.utils import EmbeddingResponse, Usage + + +class VoyageMultimodalEmbeddingError(BaseLLMException): + def __init__( + self, + status_code: int, + message: str, + headers: Union[dict, httpx.Headers] = {}, + ): + self.status_code = status_code + self.message = message + self.request = httpx.Request( + method="POST", url="https://api.voyageai.com/v1/multimodalembeddings" + ) + self.response = httpx.Response(status_code=status_code, request=self.request) + super().__init__( + status_code=status_code, + message=message, + headers=headers, + ) + + +class VoyageMultimodalEmbeddingConfig(BaseEmbeddingConfig): + """ + Reference: https://docs.voyageai.com/reference/multimodal-embeddings-api + """ + + @staticmethod + def is_multimodal_embeddings(model: str) -> bool: + return "multimodal" in model.lower() + + def get_complete_url( + self, + api_base: Optional[str], + api_key: Optional[str], + model: str, + optional_params: dict, + litellm_params: dict, + stream: Optional[bool] = None, + ) -> str: + if api_base: + if not api_base.endswith("/multimodalembeddings"): + api_base = f"{api_base}/multimodalembeddings" + return api_base + return "https://api.voyageai.com/v1/multimodalembeddings" + + def get_supported_openai_params(self, model: str) -> list: + return ["dimensions"] + + def map_openai_params( + self, + non_default_params: dict, + optional_params: dict, + model: str, + drop_params: bool, + ) -> dict: + if "dimensions" in non_default_params: + optional_params["output_dimension"] = non_default_params["dimensions"] + return optional_params + + def validate_environment( + self, + headers: dict, + model: str, + messages: List[AllMessageValues], + optional_params: dict, + litellm_params: dict, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + ) -> dict: + if api_key is None: + api_key = ( + get_secret_str("VOYAGE_API_KEY") + or get_secret_str("VOYAGE_AI_API_KEY") + or get_secret_str("VOYAGE_AI_TOKEN") + ) + if not api_key: + raise ValueError( + "Voyage API key is required for multimodal embeddings. " + "Set VOYAGE_API_KEY / VOYAGE_AI_API_KEY / VOYAGE_AI_TOKEN " + "or pass `api_key` explicitly." + ) + return {"Authorization": f"Bearer {api_key}"} + + def _normalize_content_item(self, item: Dict[str, Any]) -> Dict[str, Any]: + item_type = item.get("type") + if item_type == "image_url": + image_url = item.get("image_url") + if isinstance(image_url, dict): + image_url = image_url.get("url") + if image_url is None: + raise ValueError( + "Voyage multimodal embeddings require a non-empty `image_url`. " + "Got an image content block without a `url`." + ) + if isinstance(image_url, str) and image_url.startswith("data:image/"): + _, _, encoded = image_url.partition(",") + return {"type": "image_base64", "image_base64": encoded} + return {"type": "image_url", "image_url": image_url} + return item + + def _normalize_input_item(self, item: Any) -> Dict[str, Any]: + if isinstance(item, str): + return {"content": [{"type": "text", "text": item}]} + if isinstance(item, dict) and "content" in item: + content = item.get("content") or [] + return { + **item, + "content": [ + self._normalize_content_item(content_item) + for content_item in content + ], + } + return item + + def transform_embedding_request( + self, + model: str, + input: AllEmbeddingInputValues, + optional_params: dict, + headers: dict, + ) -> dict: + inputs = input if isinstance(input, list) else [input] + return { + "inputs": [self._normalize_input_item(item) for item in inputs], + "model": model, + **optional_params, + } + + def transform_embedding_response( + self, + model: str, + raw_response: httpx.Response, + model_response: EmbeddingResponse, + logging_obj: LiteLLMLoggingObj, + api_key: Optional[str] = None, + request_data: dict = {}, + optional_params: dict = {}, + litellm_params: dict = {}, + ) -> EmbeddingResponse: + try: + raw_response_json = raw_response.json() + except Exception: + raise VoyageMultimodalEmbeddingError( + message=raw_response.text, status_code=raw_response.status_code + ) + + model_response.model = raw_response_json.get("model") + model_response.data = raw_response_json.get("data") + model_response.object = raw_response_json.get("object") + + usage_payload = raw_response_json.get("usage", {}) + total_tokens = usage_payload.get("total_tokens", 0) + model_response.usage = Usage( + prompt_tokens=total_tokens, + total_tokens=total_tokens, + ) + return model_response + + def get_error_class( + self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] + ) -> BaseLLMException: + return VoyageMultimodalEmbeddingError( + message=error_message, status_code=status_code, headers=headers + ) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 01a01ea7a76..76a7c0640af 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -35852,7 +35852,17 @@ "max_input_tokens": 32000, "max_tokens": 32000, "mode": "embedding", - "output_cost_per_token": 0.0 + "output_cost_per_token": 0.0, + "supports_vision": true + }, + "voyage/voyage-multimodal-3.5": { + "input_cost_per_token": 1.2e-07, + "litellm_provider": "voyage", + "max_input_tokens": 32000, + "max_tokens": 32000, + "mode": "embedding", + "output_cost_per_token": 0.0, + "supports_vision": true }, "wandb/openai/gpt-oss-120b": { "max_tokens": 131072, @@ -41753,6 +41763,48 @@ "supports_tool_choice": true, "supports_vision": true }, + "bedrock_mantle/google.gemma-4-31b": { + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 4e-07, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 256000, + "max_output_tokens": 256000, + "max_tokens": 256000, + "mode": "chat", + "supports_function_calling": true, + "supports_parallel_function_calling": false, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "bedrock_mantle/google.gemma-4-26b-a4b": { + "input_cost_per_token": 1.3e-07, + "output_cost_per_token": 4e-07, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 256000, + "max_output_tokens": 256000, + "max_tokens": 256000, + "mode": "chat", + "supports_function_calling": true, + "supports_parallel_function_calling": false, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "bedrock_mantle/google.gemma-4-e2b": { + "input_cost_per_token": 4e-08, + "output_cost_per_token": 8e-08, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supports_function_calling": true, + "supports_parallel_function_calling": false, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, "volcengine/doubao-seed-2-0-pro-260215": { "litellm_provider": "volcengine", "max_input_tokens": 256000, diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index e6ffe71a971..493e09e3af1 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2226,6 +2226,10 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="max response size in MB, if a response is larger than this size it will be rejected", ) + cancel_on_disconnect: Optional[bool] = Field( + None, + description="cancel the in-flight upstream LLM request (non-streaming) when the client disconnects, freeing backend capacity (e.g. a vLLM GPU slot); the request is logged as a 499 failure", + ) infer_model_from_keys: Optional[bool] = Field( None, description="for `/models` endpoint, infers available model based on environment keys (e.g. OPENAI_API_KEY)", diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index d81668a7804..90ad0f28808 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1,6 +1,7 @@ import asyncio import json import logging +import math import time import traceback from datetime import datetime @@ -49,6 +50,7 @@ from litellm.proxy.route_llm_request import route_request from litellm.proxy.utils import ProxyLogging from litellm.router import Router from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.router import RouterRateLimitError from litellm.types.utils import ServerToolUse # Type alias for streaming chunk serializer (chunk after hooks + cost injection -> wire format) @@ -556,6 +558,64 @@ def _has_attribute_error_in_chain(exc: Exception) -> bool: return False +_CLIENT_DISCONNECT_DETAIL = "Client disconnected the request" + + +def _log_llm_api_exception(e: Exception) -> None: + if ( + getattr(e, "status_code", None) == 499 + and getattr(e, "detail", None) == _CLIENT_DISCONNECT_DETAIL + ): + verbose_proxy_logger.info( + "litellm.proxy.proxy_server._handle_llm_api_exception(): client disconnected, upstream LLM request cancelled" + ) + return + verbose_proxy_logger.exception( + f"litellm.proxy.proxy_server._handle_llm_api_exception(): Exception occured - {str(e)}" + ) + + +async def _cancel_llm_call_on_client_disconnect( + request: Request, + llm_api_call: "asyncio.Future[Any]", + disconnect_event: asyncio.Event, +) -> None: + try: + while True: + message = await request.receive() + if message["type"] == "http.disconnect": + disconnect_event.set() + llm_api_call.cancel() + return + except Exception as exc: + verbose_proxy_logger.warning( + "cancel_on_disconnect: request.receive() raised %s; " + "upstream LLM call will not be cancelled on disconnect", + exc, + ) + + +async def _await_llm_call_cancelling_on_disconnect( + request: Request, + llm_api_call: "asyncio.Future[Any]", +) -> Any: + disconnect_event = asyncio.Event() + monitor = asyncio.create_task( + _cancel_llm_call_on_client_disconnect(request, llm_api_call, disconnect_event) + ) + try: + return await llm_api_call + except asyncio.CancelledError: + if disconnect_event.is_set(): + raise HTTPException( + status_code=499, + detail=_CLIENT_DISCONNECT_DETAIL, + ) + raise + finally: + monitor.cancel() + + class ProxyBaseLLMRequestProcessing: def __init__(self, data: dict): self.data = data @@ -1224,7 +1284,12 @@ class ProxyBaseLLMRequestProcessing: *tasks ) # run the moderation check in parallel to the actual llm api call - responses = await llm_responses + if general_settings.get("cancel_on_disconnect", False): + responses = await _await_llm_call_cancelling_on_disconnect( + request, llm_responses + ) + else: + responses = await llm_responses response = responses[1] @@ -2067,6 +2132,10 @@ class ProxyBaseLLMRequestProcessing: e, ) + def _apply_router_cooldown_retry_after(self, headers: dict, e: Exception) -> None: + if isinstance(e, RouterRateLimitError) and e.cooldown_time > 0: + headers["retry-after"] = str(math.ceil(e.cooldown_time)) + async def _handle_llm_api_exception( self, e: Exception, @@ -2075,9 +2144,7 @@ class ProxyBaseLLMRequestProcessing: version: Optional[str] = None, ): """Raises ProxyException (OpenAI API compatible) if an exception is raised""" - verbose_proxy_logger.exception( - f"litellm.proxy.proxy_server._handle_llm_api_exception(): Exception occured - {str(e)}" - ) + _log_llm_api_exception(e) # Allow callbacks to transform the error response transformed_exception = await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, @@ -2148,6 +2215,8 @@ class ProxyBaseLLMRequestProcessing: except Exception: pass + self._apply_router_cooldown_retry_after(headers, e) + if isinstance(e, HTTPException): raw_detail = getattr(e, "detail", str(e)) message, structured_fields = _serialize_http_exception_detail(raw_detail) diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 507e8e4d4da..e0d018d4344 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -24,6 +24,7 @@ from litellm.proxy._types import ( LitellmUserRoles, ProxyErrorTypes, ProxyException, + SpecialModelNames, UserAPIKeyAuth, WebhookEvent, ) @@ -1074,8 +1075,26 @@ async def health_endpoint( # response but NOT in the background-cache /health response. This is # surfaced via the "warnings" field below so operators can fix the # missing model_info.id rather than guess at the discrepancy. - if len(user_api_key_dict.models) > 0: - allowed_models = set(user_api_key_dict.models) + # Keys granted SpecialModelNames.all_proxy_models carry the literal + # "all-proxy-models" entry, which matches no real model_name; treat + # them as unrestricted instead of filtering the list down to nothing. + # Keys granted SpecialModelNames.all_team_models inherit the parent + # team's allowlist (same semantics as get_key_models in + # model_checks.py). Without a team_id the sentinel cannot resolve and + # stays in the list, matching nothing; denied rather than + # unrestricted, mirroring _resolve_key_models_for_auth_check. + accessible_models = list(user_api_key_dict.models) + if ( + SpecialModelNames.all_team_models.value in accessible_models + and user_api_key_dict.team_id is not None + ): + accessible_models = list(user_api_key_dict.team_models) + restrict_to_allowed_models = ( + len(accessible_models) > 0 + and SpecialModelNames.all_proxy_models.value not in accessible_models + ) + if restrict_to_allowed_models: + allowed_models = set(accessible_models) _llm_model_list = [ m for m in _llm_model_list if m.get("model_name") in allowed_models ] @@ -1087,7 +1106,7 @@ async def health_endpoint( # other healthy model would still report healthy_count > 0 and # the targeted-503 path would never fire. targeted_ids = _resolve_targeted_model_ids(_llm_model_list, model, model_id) - if len(user_api_key_dict.models) > 0: + if restrict_to_allowed_models: allowed_model_ids = { (m.get("model_info") or {}).get("id") for m in _llm_model_list diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 7666b23f2af..fca395f889c 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -397,6 +397,32 @@ def get_chain_id_from_headers(headers: Optional[Dict[str, str]]) -> Optional[str ) +def is_claude_code_user_agent(user_agent: str) -> bool: + """Claude Code identifies itself as ``claude-cli/ ...``; the IDE + extensions and the Agent SDK run through the same CLI and share that prefix.""" + return user_agent.startswith("claude-cli/") + + +def should_auto_drop_params_for_claude_code( + user_agent: str, data: dict, proxy_config: ProxyConfig +) -> bool: + """drop_params defaults to on for Claude Code so its Anthropic-specific + params (e.g. thinking) don't fail requests routed to non-Anthropic + providers. An explicit drop_params from the caller or in the operator's + ``litellm_settings`` always wins over this default.""" + if not is_claude_code_user_agent(user_agent): + return False + if "drop_params" in data: + return False + config = getattr(proxy_config, "config", None) + litellm_settings = ( + config.get("litellm_settings") if isinstance(config, dict) else None + ) + return not ( + isinstance(litellm_settings, dict) and "drop_params" in litellm_settings + ) + + def safe_add_api_version_from_query_params(data: dict, request: Request): try: if hasattr(request, "query_params"): @@ -1742,6 +1768,9 @@ async def add_litellm_data_to_request( # noqa: PLR0915 user_agent = request.headers["user-agent"] data[_metadata_variable_name]["user_agent"] = user_agent + if should_auto_drop_params_for_claude_code(user_agent, data, proxy_config): + data["drop_params"] = True + # Merge caller-supplied tags (x-litellm-tags header, data["tags"] root-level) # into request metadata for tag-based routing and spend attribution. tags = LiteLLMProxyRequestSetup.add_request_tag_to_metadata( diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 2f239c8da84..c980f6f5260 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -4794,12 +4794,27 @@ async def reset_key_spend_fn( proxy_logging_obj=proxy_logging_obj, ) - try: - from litellm.proxy.proxy_server import _invalidate_spend_counter + # Set Redis spend counter to the new value so get_current_spend() + # returns the correct amount immediately instead of the stale pre-reset value. + # We use reset_to (not 0.0) so partial resets are reflected correctly. + from litellm.proxy.proxy_server import spend_counter_cache - await _invalidate_spend_counter(counter_key=f"spend:key:{hashed_api_key}") - except Exception: - pass + _counter_key = f"spend:key:{hashed_api_key}" + spend_counter_cache.in_memory_cache.set_cache( + key=_counter_key, value=reset_to, ttl=60 + ) + if spend_counter_cache.redis_cache is not None: + try: + await spend_counter_cache.redis_cache.async_set_cache( + key=_counter_key, value=reset_to, ttl=60 + ) + except Exception as redis_err: + verbose_proxy_logger.warning( + "Failed to update spend counter %s in Redis: %s. " + "Budget checks may use stale value until counter expires.", + _counter_key, + redis_err, + ) max_budget = updated_key.max_budget budget_reset_at = updated_key.budget_reset_at diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index c894813ada4..1a0a57c71fd 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -4850,15 +4850,34 @@ async def team_model_add( detail={"error": "Only proxy admin or team admin can modify team models"}, ) - updated_models = add_new_models_to_team(team_obj=team_obj, new_models=data.models) - # Update team. `include` mirrors the relations the auth path consumes - # off the cached team object so that `_refresh_cached_team` doesn't - # null them out — see object_permission_utils.validate_key_search_tools_against_team - # and the MCP/agent authz paths, which treat a missing object_permission - # as "no team-level restriction". + # Atomic array append with dedup at the database level so concurrent + # BYOK model creates don't overwrite each other's team.models entries. + # When the team currently has models=[] (unrestricted access), the + # CASE expression inserts the 'all-proxy-models' sentinel first. + models_to_add = list(data.models) + await prisma_client.db.execute_raw( + 'UPDATE "LiteLLM_TeamTable" ' + "SET models = (" + " SELECT ARRAY(SELECT DISTINCT unnest(" + " CASE WHEN cardinality(COALESCE(models, ARRAY[]::text[])) = 0 " + " THEN ARRAY['all-proxy-models']::text[] " + " ELSE models " + " END || $1::text[]" + " ))" + ") " + "WHERE team_id = $2", + models_to_add, + data.team_id, + ) + # Re-fetch via update (write-routed) instead of find_unique (read-routed) + # to avoid returning stale data from a read replica. The models column + # was already set by execute_raw above; this just retrieves the row from + # the writer and lets Prisma bump updated_at. + # `include` mirrors the relations the auth path consumes off the cached + # team object so that `_refresh_cached_team` doesn't null them out. updated_team = await TeamRepository(prisma_client).table.update( where={"team_id": data.team_id}, - data={"models": updated_models}, + data={"updated_at": datetime.now(timezone.utc)}, include={"object_permission": True}, # type: ignore ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 1d8cbb6fe0a..0d6374fec69 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1942,34 +1942,6 @@ db_writer_client: Optional[AsyncHTTPHandler] = None ### logger ### -async def check_request_disconnection(request: Request, llm_api_call_task): - """ - Asynchronously checks if the request is disconnected at regular intervals. - If the request is disconnected - - cancel the litellm.router task - - raises an HTTPException with status code 499 and detail "Client disconnected the request". - - Parameters: - - request: Request: The request object to check for disconnection. - Returns: - - None - """ - - # only run this function for 10 mins -> if these don't get cancelled -> we don't want the server to have many while loops - start_time = time.time() - while time.time() - start_time < 600: - await asyncio.sleep(1) - if await request.is_disconnected(): - # cancel the LLM API Call task if any passed - this is passed from individual providers - # Example OpenAI, Azure, VertexAI etc - llm_api_call_task.cancel() - - raise HTTPException( - status_code=499, - detail="Client disconnected the request", - ) - - def _resolve_typed_dict_type(typ): """Resolve the actual TypedDict class from a potentially wrapped type.""" from typing_extensions import _TypedDictMeta # type: ignore @@ -4920,9 +4892,12 @@ class ProxyConfig: combined_id_list = [] ## BASE CASES ## - # if llm_router is None or db_models is empty, return 0 - if llm_router is None or len(db_models) == 0: + if llm_router is None: return 0 + # NOTE: db_models may be legitimately empty when all DB models have been deleted. + # Do NOT short-circuit on len(db_models) == 0 — we must still evict any + # DB-sourced deployments that are no longer in the DB. The caller + # (_update_llm_router) already guards against None (transient fetch failure). ## DB MODELS ## for m in db_models: @@ -5072,6 +5047,15 @@ class ProxyConfig: ) try: + # new_models is None when _get_models_from_db failed (transient DB error). + # Skip the update entirely so we don't evict valid deployments. + if new_models is None: + verbose_proxy_logger.warning( + "_update_llm_router: DB model fetch returned None (transient failure). " + "Skipping router update to preserve existing deployments." + ) + return + models_list: list = new_models if isinstance(new_models, list) else [] if llm_router is None and master_key is not None: verbose_proxy_logger.debug(f"len new_models: {len(models_list)}") @@ -5774,18 +5758,25 @@ class ProxyConfig: # Check if the object type is in the list (supports both str and enum values) return any(str(obj) == object_type_str for obj in supported_db_objects) - async def _get_models_from_db(self, prisma_client: PrismaClient) -> list: + async def _get_models_from_db(self, prisma_client: PrismaClient) -> Optional[list]: + """ + Fetch all model deployments from the DB. + + Returns: + - list: the rows (may be empty if no models exist) + - None: signals a DB fetch *failure* — callers must not treat this + as "all models deleted" and must not evict existing router deployments. + """ try: new_models = await ModelRepository(prisma_client).table.find_many() + return new_models except Exception as e: verbose_proxy_logger.exception( "litellm.proxy_server.py::add_deployment() - Error getting new models from DB - {}".format( str(e) ) ) - new_models = [] - - return new_models + return None async def add_deployment( self, @@ -14775,6 +14766,7 @@ async def get_config_list( "always_include_stream_usage": {"type": "Boolean"}, "forward_client_headers_to_llm_api": {"type": "Boolean"}, "mcp_required_fields": {"type": "List"}, + "cancel_on_disconnect": {"type": "Boolean"}, } return_val = [] diff --git a/litellm/router.py b/litellm/router.py index 34f67c11873..80584858311 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -4112,47 +4112,13 @@ class Router: ``` """ try: + kwargs["model"] = model kwargs["input"] = input kwargs["voice"] = voice - - deployment = await self.async_get_available_deployment( - model=model, - messages=[{"role": "user", "content": "prompt"}], - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) + kwargs["original_function"] = self._aspeech self._update_kwargs_before_fallbacks(model=model, kwargs=kwargs) - data = deployment["litellm_params"].copy() - data["model"] - for k, v in self.default_litellm_params.items(): - if ( - k not in kwargs - ): # prioritize model-specific params > default router params - kwargs[k] = v - elif k == "metadata": - kwargs[k].update(v) + response = await self.async_function_with_fallbacks(**kwargs) - potential_model_client = self._get_client( - deployment=deployment, kwargs=kwargs, client_type="async" - ) - # check if provided keys == client keys # - dynamic_api_key = kwargs.get("api_key", None) - if ( - dynamic_api_key is not None - and potential_model_client is not None - and dynamic_api_key != potential_model_client.api_key - ): - model_client = None - else: - model_client = potential_model_client - - response = await litellm.aspeech( - **{ - **data, - "client": model_client, - **kwargs, - } - ) return response except Exception as e: asyncio.create_task( @@ -4165,6 +4131,76 @@ class Router: ) raise e + async def _aspeech(self, model: str, input: str, voice: str, **kwargs): + model_name = model + try: + verbose_router_logger.debug( + f"Inside _aspeech()- model: {model}; kwargs: {kwargs}" + ) + parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs) + deployment = await self.async_get_available_deployment( + model=model, + messages=[{"role": "user", "content": "prompt"}], + specific_deployment=kwargs.pop("specific_deployment", None), + request_kwargs=kwargs, + ) + + self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) + data = deployment["litellm_params"].copy() + model_client = self._get_async_openai_model_client( + deployment=deployment, + kwargs=kwargs, + ) + + self.total_calls[model_name] += 1 + response = litellm.aspeech( + **{ + **data, + "input": input, + "voice": voice, + "client": model_client, + **kwargs, + } + ) + + ### CONCURRENCY-SAFE RPM CHECKS ### + rpm_semaphore = self._get_client( + deployment=deployment, + kwargs=kwargs, + client_type="max_parallel_requests", + ) + + if rpm_semaphore is not None and isinstance( + rpm_semaphore, asyncio.Semaphore + ): + async with rpm_semaphore: + """ + - Check rpm limits before making the call + - If allowed, increment the rpm limit (allows global value to be updated, concurrency-safe) + """ + await self.async_routing_strategy_pre_call_checks( + deployment=deployment, parent_otel_span=parent_otel_span + ) + response = await response + else: + await self.async_routing_strategy_pre_call_checks( + deployment=deployment, parent_otel_span=parent_otel_span + ) + response = await response + + self.success_calls[model_name] += 1 + verbose_router_logger.info( + f"litellm.aspeech(model={model_name})\033[32m 200 OK\033[0m" + ) + return response + except Exception as e: + verbose_router_logger.info( + f"litellm.aspeech(model={model_name})\033[31m Exception {str(e)}\033[0m" + ) + if model_name is not None: + self.fail_calls[model_name] += 1 + raise e + async def arerank(self, model: str, **kwargs): try: kwargs["model"] = model diff --git a/litellm/types/integrations/slack_alerting.py b/litellm/types/integrations/slack_alerting.py index 078e7953ad8..4786dbab101 100644 --- a/litellm/types/integrations/slack_alerting.py +++ b/litellm/types/integrations/slack_alerting.py @@ -1,4 +1,5 @@ import os +import time from datetime import datetime as dt from enum import Enum from typing import Any, Dict, List, Literal, Optional, Set, Union @@ -201,6 +202,8 @@ class HangingRequestData(BaseModel): key_alias: Optional[str] = None team_alias: Optional[str] = None alerting_metadata: Optional[dict] = None + created_at: float = Field(default_factory=time.time) + alerted: bool = False class AlertTypeConfig(LiteLLMPydanticObjectBase): diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 644ad2cb905..d3dc7eadb94 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3060,6 +3060,12 @@ class StandardCallbackDynamicParams(TypedDict, total=False): wandb_api_key: Optional[str] weave_project_id: Optional[str] + # Datadog dynamic params + dd_api_key: Optional[str] + dd_site: Optional[str] + dd_agent_host: Optional[str] + dd_agent_port: Optional[str] + # Logging settings turn_off_message_logging: Optional[bool] # when true will not log messages litellm_disabled_callbacks: Optional[List[str]] diff --git a/litellm/utils.py b/litellm/utils.py index 46d48279198..4c67abdf937 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -3583,6 +3583,15 @@ def get_optional_params_embeddings( # noqa: PLR0915 drop_params=drop_params if drop_params is not None else False, ) ) + elif litellm.VoyageMultimodalEmbeddingConfig.is_multimodal_embeddings(model): + optional_params = ( + litellm.VoyageMultimodalEmbeddingConfig().map_openai_params( + non_default_params=non_default_params, + optional_params={}, + model=model, + drop_params=drop_params if drop_params is not None else False, + ) + ) else: optional_params = litellm.VoyageEmbeddingConfig().map_openai_params( non_default_params=non_default_params, @@ -8666,6 +8675,11 @@ class ProviderConfigManager: ) ): return litellm.VoyageContextualEmbeddingConfig() + elif ( + litellm.LlmProviders.VOYAGE == provider + and litellm.VoyageMultimodalEmbeddingConfig.is_multimodal_embeddings(model) + ): + return litellm.VoyageMultimodalEmbeddingConfig() elif litellm.LlmProviders.VOYAGE == provider: return litellm.VoyageEmbeddingConfig() elif litellm.LlmProviders.TRITON == provider: diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 8cdde5ac82a..b181df94131 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -39511,6 +39511,178 @@ "litellm_provider": "fireworks_ai", "mode": "chat" }, + "scaleway/qwen/qwen3.5-397b-a17b": { + "input_cost_per_token": 6e-07, + "litellm_provider": "scaleway", + "max_input_tokens": 256000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 3.6e-06, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true + }, + "scaleway/qwen/qwen3.6-35b-a3b": { + "input_cost_per_token": 2.5e-07, + "litellm_provider": "scaleway", + "max_input_tokens": 256000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "supports_function_calling": true, + "supports_vision": true, + "supports_reasoning": true + }, + "scaleway/qwen/qwen3-235b-a22b-instruct-2507": { + "input_cost_per_token": 7.5e-07, + "litellm_provider": "scaleway", + "max_input_tokens": 256000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 2.25e-06, + "supports_function_calling": true + }, + "scaleway/qwen/qwen3-embedding-8b": { + "input_cost_per_token": 1e-07, + "litellm_provider": "scaleway", + "mode": "embedding", + "output_cost_per_token": 0.0 + }, + "scaleway/qwen/qwen3-coder-30b-a3b-instruct": { + "input_cost_per_token": 2e-07, + "litellm_provider": "scaleway", + "max_input_tokens": 128000, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 8e-07, + "supports_function_calling": true + }, + "scaleway/openai/gpt-oss-120b": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "scaleway", + "max_input_tokens": 128000, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 6e-07, + "supports_function_calling": true + }, + "scaleway/openai/whisper-large-v3": { + "input_cost_per_audio_token": 0.0, + "litellm_provider": "scaleway", + "mode": "audio_transcription", + "output_cost_per_token": 0.0 + }, + "scaleway/google/gemma-4-26b-a4b-it": { + "input_cost_per_token": 2.5e-07, + "litellm_provider": "scaleway", + "max_input_tokens": 256000, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 5e-07, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_vision": true + }, + "scaleway/google/gemma-3-27b-it": { + "input_cost_per_token": 2.5e-07, + "litellm_provider": "scaleway", + "max_input_tokens": 40000, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 5e-07, + "supports_function_calling": true, + "supports_vision": true + }, + "scaleway/hcompany/holo2-30b-a3b": { + "input_cost_per_token": 3e-07, + "litellm_provider": "scaleway", + "max_input_tokens": 22000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 7e-07, + "supports_reasoning": true, + "supports_vision": true + }, + "scaleway/mistralai/mistral-medium-3.5-128b": { + "input_cost_per_token": 1.5e-06, + "litellm_provider": "scaleway", + "max_input_tokens": 256000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 7.5e-06, + "supports_reasoning": true, + "supports_function_calling": true, + "supports_vision": true, + "supports_tool_choice": true + }, + "scaleway/mistralai/devstral-2-123b-instruct-2512": { + "input_cost_per_token": 4e-07, + "litellm_provider": "scaleway", + "max_input_tokens": 200000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 2e-06, + "supports_function_calling": true + }, + "scaleway/mistralai/voxtral-small-24b-2507": { + "input_cost_per_audio_token": 1.5e-07, + "input_cost_per_token": 1.5e-07, + "litellm_provider": "scaleway", + "max_input_tokens": 32000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 3.5e-07, + "supports_audio_input": true + }, + "scaleway/mistralai/mistral-small-3.2-24b-instruct-2506": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "scaleway", + "max_input_tokens": 128000, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 3.5e-07, + "supports_function_calling": true, + "supports_vision": true + }, + "scaleway/mistralai/pixtral-12b-2409": { + "input_cost_per_token": 2e-07, + "litellm_provider": "scaleway", + "max_input_tokens": 128000, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "chat", + "output_cost_per_token": 2e-07, + "supports_vision": true, + "supports_function_calling": true + }, + "scaleway/BAAI/bge-multilingual-gemma2": { + "input_cost_per_token": 1e-07, + "litellm_provider": "scaleway", + "mode": "embedding", + "output_cost_per_token": 0.0 + }, + "scaleway/meta/llama-3.3-70b-instruct": { + "input_cost_per_token": 9e-07, + "litellm_provider": "scaleway", + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 9e-07, + "supports_function_calling": true + }, "novita/deepseek/deepseek-v3.2": { "litellm_provider": "novita", "mode": "chat", @@ -41793,6 +41965,48 @@ "supports_tool_choice": true, "supports_vision": true }, + "bedrock_mantle/google.gemma-4-31b": { + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 4e-07, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 256000, + "max_output_tokens": 256000, + "max_tokens": 256000, + "mode": "chat", + "supports_function_calling": true, + "supports_parallel_function_calling": false, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "bedrock_mantle/google.gemma-4-26b-a4b": { + "input_cost_per_token": 1.3e-07, + "output_cost_per_token": 4e-07, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 256000, + "max_output_tokens": 256000, + "max_tokens": 256000, + "mode": "chat", + "supports_function_calling": true, + "supports_parallel_function_calling": false, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "bedrock_mantle/google.gemma-4-e2b": { + "input_cost_per_token": 4e-08, + "output_cost_per_token": 8e-08, + "litellm_provider": "bedrock_mantle", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "supports_function_calling": true, + "supports_parallel_function_calling": false, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, "volcengine/doubao-seed-2-0-pro-260215": { "litellm_provider": "volcengine", "max_input_tokens": 256000, diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 6caab585ac9..2ad2b3ec982 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -2086,7 +2086,7 @@ "chat_completions": true, "messages": true, "responses": true, - "embeddings": false, + "embeddings": true, "image_generations": false, "audio_transcriptions": true, "audio_speech": false, @@ -2153,7 +2153,7 @@ "endpoints": { "chat_completions": true, "messages": true, - "responses": false, + "responses": true, "embeddings": false, "image_generations": false, "audio_transcriptions": false, @@ -2752,6 +2752,23 @@ "batches": false, "rerank": false } + }, + "empiriolabs": { + "display_name": "EmpirioLabs (`empiriolabs`)", + "url": "https://docs.litellm.ai/docs/providers/empiriolabs", + "endpoints": { + "chat_completions": true, + "messages": false, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": false + } } }, "endpoints": { diff --git a/proxy_server_config.yaml b/proxy_server_config.yaml index d0730094ce1..f5f4e1956d4 100644 --- a/proxy_server_config.yaml +++ b/proxy_server_config.yaml @@ -230,6 +230,7 @@ general_settings: # background_health_checks: true # use_shared_health_check: true # health_check_interval: 30 + # cancel_on_disconnect: true # cancel the in-flight upstream LLM request (non-streaming) when the client disconnects, freeing backend capacity (e.g. a vLLM GPU slot) # database_url: "postgresql://:@:/" # [OPTIONAL] use for token-based auth to proxy pass_through_endpoints: diff --git a/tests/litellm/llms/openai_like/test_empiriolabs_provider.py b/tests/litellm/llms/openai_like/test_empiriolabs_provider.py new file mode 100644 index 00000000000..58f5e47d09e --- /dev/null +++ b/tests/litellm/llms/openai_like/test_empiriolabs_provider.py @@ -0,0 +1,63 @@ +""" +Unit tests for the EmpirioLabs OpenAI-like provider. +""" + +import os +import sys + +sys.path.insert( + 0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../..")) +) + +from litellm.llms.openai_like.dynamic_config import create_config_class +from litellm.llms.openai_like.json_loader import JSONProviderRegistry + +EMPIRIOLABS_BASE_URL = "https://api.empiriolabs.ai/v1" + + +def _get_config(): + provider = JSONProviderRegistry.get("empiriolabs") + assert provider is not None + config_class = create_config_class(provider) + return config_class() + + +def test_empiriolabs_provider_registered(): + provider = JSONProviderRegistry.get("empiriolabs") + assert provider is not None + assert provider.base_url == EMPIRIOLABS_BASE_URL + assert provider.api_key_env == "EMPIRIOLABS_API_KEY" + assert provider.api_base_env == "EMPIRIOLABS_API_BASE" + + +def test_empiriolabs_resolves_env_api_key(monkeypatch): + config = _get_config() + monkeypatch.setenv("EMPIRIOLABS_API_KEY", "test-key") + api_base, api_key = config._get_openai_compatible_provider_info(None, None) + assert api_base == EMPIRIOLABS_BASE_URL + assert api_key == "test-key" + + +def test_empiriolabs_maps_max_completion_tokens(): + config = _get_config() + params = config.map_openai_params( + non_default_params={"max_completion_tokens": 256}, + optional_params={}, + model="empiriolabs/qwen3-7-plus", + drop_params=False, + ) + assert params.get("max_tokens") == 256 + assert "max_completion_tokens" not in params + + +def test_empiriolabs_complete_url_appends_endpoint(): + config = _get_config() + url = config.get_complete_url( + api_base=EMPIRIOLABS_BASE_URL, + api_key="test-key", + model="empiriolabs/qwen3-7-plus", + optional_params={}, + litellm_params={}, + stream=False, + ) + assert url == f"{EMPIRIOLABS_BASE_URL}/chat/completions" diff --git a/tests/local_testing/test_config.py b/tests/local_testing/test_config.py index 2c5d04d3815..e4d0ffb4408 100644 --- a/tests/local_testing/test_config.py +++ b/tests/local_testing/test_config.py @@ -224,8 +224,22 @@ async def test_db_error_new_model_check(): model_info={"id": deployment.model_info.id}, ) - db_models = [] - deleted_deployments = await pc._delete_deployment(db_models=db_models) + # Mock get_config to return the two deployments as config-backed models so + # they appear in combined_id_list and are not evicted when db_models is empty + # (simulates the real-world case: DB error returns [], but models live in config). + config_model_list = [ + deployment.to_json(exclude_none=True), + deployment_2.to_json(exclude_none=True), + ] + from unittest.mock import AsyncMock, patch + + with patch.object( + pc, + "get_config", + new=AsyncMock(return_value={"model_list": config_model_list}), + ): + db_models = [] + deleted_deployments = await pc._delete_deployment(db_models=db_models) assert deleted_deployments == 0 assert init_len_list == len(llm_router.model_list) diff --git a/tests/router_unit_tests/test_router_endpoints.py b/tests/router_unit_tests/test_router_endpoints.py index c170972d984..658ad4f3b5c 100644 --- a/tests/router_unit_tests/test_router_endpoints.py +++ b/tests/router_unit_tests/test_router_endpoints.py @@ -198,6 +198,96 @@ async def test_audio_speech_router(mode): assert test_logger.standard_logging_object["model_group"] == "tts" +@pytest.mark.asyncio +async def test_aspeech_fallbacks_on_deployment_failure(): + router = Router( + model_list=[ + { + "model_name": "tts-main", + "litellm_params": {"model": "openai/tts-1", "api_key": "fake-key"}, + }, + { + "model_name": "tts-backup", + "litellm_params": {"model": "openai/tts-1-hd", "api_key": "fake-key"}, + }, + ], + fallbacks=[{"tts-main": ["tts-backup"]}], + num_retries=0, + ) + + called_models = [] + + async def mock_aspeech(*args, **kwargs): + called_models.append(kwargs["model"]) + if kwargs["model"] == "openai/tts-1": + raise litellm.InternalServerError( + message="deployment down", + llm_provider="openai", + model="tts-1", + ) + return MagicMock() + + with patch("litellm.aspeech", side_effect=mock_aspeech): + response = await router.aspeech( + model="tts-main", + input="the quick brown fox jumped over the lazy dogs", + voice="alloy", + ) + + assert response is not None + assert called_models == ["openai/tts-1", "openai/tts-1-hd"] + + +@pytest.mark.asyncio +async def test_aspeech_success_returns_response(): + router = Router( + model_list=[ + { + "model_name": "tts", + "litellm_params": {"model": "openai/tts-1", "api_key": "fake-key"}, + }, + ] + ) + + mock_response = MagicMock() + with patch("litellm.aspeech", return_value=mock_response) as mock_aspeech: + response = await router.aspeech( + model="tts", + input="the quick brown fox jumped over the lazy dogs", + voice="alloy", + ) + + assert response is mock_response + mock_aspeech.assert_called_once() + assert mock_aspeech.call_args.kwargs["model"] == "openai/tts-1" + + +@pytest.mark.asyncio +async def test_aspeech_sets_deployment_metadata(): + router = Router( + model_list=[ + { + "model_name": "tts", + "litellm_params": {"model": "openai/tts-1", "api_key": "fake-key"}, + }, + ] + ) + + mock_response = MagicMock() + with patch("litellm.aspeech", return_value=mock_response) as mock_aspeech: + response = await router._aspeech( + model="tts", + input="the quick brown fox jumped over the lazy dogs", + voice="alloy", + ) + + assert response is mock_response + metadata = mock_aspeech.call_args.kwargs["metadata"] + assert metadata["deployment"] == "openai/tts-1" + assert metadata["deployment_model_name"] == "tts" + assert metadata["model_info"]["id"] is not None + + @pytest.mark.asyncio() async def test_rerank_endpoint(model_list): from litellm.types.utils import RerankResponse diff --git a/tests/test_litellm/integrations/SlackAlerting/test_hanging_request_check.py b/tests/test_litellm/integrations/SlackAlerting/test_hanging_request_check.py index 0bece97b6f0..063aabd309b 100644 --- a/tests/test_litellm/integrations/SlackAlerting/test_hanging_request_check.py +++ b/tests/test_litellm/integrations/SlackAlerting/test_hanging_request_check.py @@ -1,6 +1,7 @@ import json import os import sys +import time from typing import Optional from unittest.mock import AsyncMock, MagicMock, patch @@ -35,13 +36,13 @@ class TestAlertingHangingRequestCheck: async def test_init_creates_cache_with_correct_ttl(self, mock_slack_alerting): """ Test that initialization creates a hanging request cache with correct TTL. - The TTL should be alerting_threshold + buffer time. + The TTL should be 1.5x alerting_threshold + buffer time, so entries + survive long enough to be checked after crossing the threshold. """ checker = AlertingHangingRequestCheck(slack_alerting_object=mock_slack_alerting) - # The cache should be created with TTL = alerting_threshold + buffer time - expected_ttl = ( - mock_slack_alerting.alerting_threshold + 60 + expected_ttl = int( + mock_slack_alerting.alerting_threshold * 1.5 + 60 ) # HANGING_ALERT_BUFFER_TIME_SECONDS assert checker.hanging_request_cache.default_ttl == expected_ttl @@ -208,13 +209,14 @@ class TestAlertingHangingRequestCheck: Test send_alerts_for_hanging_requests when request is actually hanging. Should send alert for requests that haven't completed within threshold. """ - # Add a hanging request to the cache + # Add a hanging request that is older than the alerting threshold hanging_data = HangingRequestData( request_id="hanging_request_999", model="gpt-4", api_base="https://api.openai.com/v1", key_alias="test_key", team_alias="test_team", + created_at=time.time() - 301, ) await hanging_request_checker.hanging_request_cache.async_set_cache( key="hanging_request_999", value=hanging_data, ttl=300 @@ -236,6 +238,82 @@ class TestAlertingHangingRequestCheck: # Verify alert was sent for hanging request hanging_request_checker.slack_alerting_object.send_alert.assert_called_once() + @pytest.mark.asyncio + async def test_send_alerts_for_hanging_requests_alerts_once_per_hang( + self, hanging_request_checker + ): + """ + A single hanging request must alert exactly once even though the + checker tick revisits it on every run within the cache TTL. + """ + hanging_data = HangingRequestData( + request_id="hanging_once_555", + model="gpt-4", + api_base="https://api.openai.com/v1", + created_at=time.time() - 301, + ) + await hanging_request_checker.hanging_request_cache.async_set_cache( + key="hanging_once_555", value=hanging_data, ttl=300 + ) + + with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy: + mock_internal_cache = AsyncMock() + mock_internal_cache.async_get_cache.return_value = None + mock_proxy.internal_usage_cache = mock_internal_cache + + hanging_request_checker.hanging_request_cache.async_get_oldest_n_keys = ( + AsyncMock(return_value=["hanging_once_555"]) + ) + + for _ in range(3): + await hanging_request_checker.send_alerts_for_hanging_requests() + + assert hanging_request_checker.slack_alerting_object.send_alert.call_count == 1 + cached = await hanging_request_checker.hanging_request_cache.async_get_cache( + key="hanging_once_555" + ) + assert cached is not None + assert cached.alerted is True + + @pytest.mark.asyncio + async def test_send_alerts_for_hanging_requests_skips_request_younger_than_threshold( + self, hanging_request_checker + ): + """ + Test that an in-flight request younger than the alerting threshold + does not trigger an alert and stays in the cache for later checks. + """ + hanging_data = HangingRequestData( + request_id="young_request_123", + model="gpt-4", + api_base="https://api.openai.com/v1", + ) + await hanging_request_checker.hanging_request_cache.async_set_cache( + key="young_request_123", value=hanging_data, ttl=300 + ) + + with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy: + # Mock internal usage cache to return None (request still in flight) + mock_internal_cache = AsyncMock() + mock_internal_cache.async_get_cache.return_value = None + mock_proxy.internal_usage_cache = mock_internal_cache + + hanging_request_checker.hanging_request_cache.async_get_oldest_n_keys = ( + AsyncMock(return_value=["young_request_123"]) + ) + + await hanging_request_checker.send_alerts_for_hanging_requests() + + # No alert for a request below the threshold, and it must remain + # cached so a later check can alert if it never completes + hanging_request_checker.slack_alerting_object.send_alert.assert_not_called() + assert ( + await hanging_request_checker.hanging_request_cache.async_get_cache( + key="young_request_123" + ) + is not None + ) + @pytest.mark.asyncio async def test_send_alerts_for_hanging_requests_with_missing_hanging_data( self, hanging_request_checker diff --git a/tests/test_litellm/integrations/datadog/test_datadog_team_handler.py b/tests/test_litellm/integrations/datadog/test_datadog_team_handler.py new file mode 100644 index 00000000000..772e993c132 --- /dev/null +++ b/tests/test_litellm/integrations/datadog/test_datadog_team_handler.py @@ -0,0 +1,263 @@ +""" +Tests for team-scoped Datadog callback support. + +Verifies that DataDogLogger can be instantiated with per-team credentials +(dd_api_key, dd_site) instead of relying solely on environment variables, +and that the DataDogHandler correctly resolves and caches per-team loggers. +""" + +from unittest.mock import patch + +import pytest + +from litellm.integrations.datadog.datadog import DataDogLogger +from litellm.integrations.datadog.datadog_team_handler import ( + DataDogHandler, + DatadogLoggingConfig, +) +from litellm.litellm_core_utils.specialty_caches.dynamic_logging_cache import ( + DynamicLoggingCache, +) +from litellm.types.utils import StandardCallbackDynamicParams + + +@pytest.fixture +def datadog_env(monkeypatch): + """Set global DD env vars for the default/global logger.""" + monkeypatch.setenv("DD_API_KEY", "global_api_key") + monkeypatch.setenv("DD_SITE", "us1.datadoghq.com") + + +class TestDataDogLoggerCredentialKwargs: + """Test that DataDogLogger accepts credentials as kwargs.""" + + def test_init_with_explicit_credentials(self): + """Logger should use explicit kwargs instead of env vars.""" + with patch("asyncio.create_task"): + logger = DataDogLogger( + dd_api_key="team_api_key", + dd_site="eu1.datadoghq.com", + ) + + assert logger.DD_API_KEY == "team_api_key" + assert "eu1.datadoghq.com" in logger.intake_url + + def test_init_falls_back_to_env_vars(self, datadog_env): + """Logger should fall back to env vars when no kwargs provided.""" + with patch("asyncio.create_task"): + logger = DataDogLogger() + + assert logger.DD_API_KEY == "global_api_key" + assert "us1.datadoghq.com" in logger.intake_url + + def test_init_kwargs_override_env_vars(self, datadog_env): + """Explicit kwargs should take precedence over env vars.""" + with patch("asyncio.create_task"): + logger = DataDogLogger( + dd_api_key="override_key", + dd_site="ap1.datadoghq.com", + ) + + assert logger.DD_API_KEY == "override_key" + assert "ap1.datadoghq.com" in logger.intake_url + + def test_init_with_agent_credentials(self): + """Logger should use agent mode when dd_agent_host is provided.""" + with patch("asyncio.create_task"): + logger = DataDogLogger( + dd_agent_host="dd-agent.local", + dd_agent_port="8125", + dd_api_key="agent_api_key", + ) + + assert "dd-agent.local:8125" in logger.intake_url + assert logger.DD_API_KEY == "agent_api_key" + + def test_init_raises_without_credentials(self, monkeypatch): + """Logger should raise if no credentials are available.""" + monkeypatch.delenv("DD_API_KEY", raising=False) + monkeypatch.delenv("DD_SITE", raising=False) + monkeypatch.delenv("LITELLM_DD_AGENT_HOST", raising=False) + + with pytest.raises(Exception, match="DD_API_KEY"): + with patch("asyncio.create_task"): + DataDogLogger() + + def test_agent_mode_does_not_leak_env_api_key_when_disallowed(self, datadog_env): + """With allow_env_credentials=False, the agent logger must not pick up DD_API_KEY env var.""" + with patch("asyncio.create_task"): + logger = DataDogLogger( + dd_agent_host="attacker.example.com", + allow_env_credentials=False, + ) + + assert logger.DD_API_KEY is None + assert "attacker.example.com" in logger.intake_url + + def test_direct_api_mode_does_not_leak_env_api_key_when_disallowed( + self, datadog_env + ): + """With allow_env_credentials=False and no explicit key, init must fail rather than reuse env key.""" + with pytest.raises(Exception, match="DD_API_KEY"): + with patch("asyncio.create_task"): + DataDogLogger( + dd_site="attacker.example.com", + allow_env_credentials=False, + ) + + +class TestDataDogHandler: + """Test that DataDogHandler resolves the correct logger per team.""" + + def test_creates_team_logger_with_dynamic_credentials(self, datadog_env): + """Should create a new logger when team credentials are provided.""" + cache = DynamicLoggingCache() + params = StandardCallbackDynamicParams( + dd_api_key="team_a_key", + dd_site="eu1.datadoghq.com", + ) + + with patch("asyncio.create_task"): + result = DataDogHandler.get_datadog_logger_for_request( + standard_callback_dynamic_params=params, + in_memory_dynamic_logger_cache=cache, + ) + + assert result.DD_API_KEY == "team_a_key" + assert "eu1.datadoghq.com" in result.intake_url + + def test_caches_team_logger(self, datadog_env): + """Same team credentials should return the same cached logger instance.""" + cache = DynamicLoggingCache() + params = StandardCallbackDynamicParams( + dd_api_key="team_b_key", + dd_site="us5.datadoghq.com", + ) + + with patch("asyncio.create_task"): + result1 = DataDogHandler.get_datadog_logger_for_request( + standard_callback_dynamic_params=params, + in_memory_dynamic_logger_cache=cache, + ) + result2 = DataDogHandler.get_datadog_logger_for_request( + standard_callback_dynamic_params=params, + in_memory_dynamic_logger_cache=cache, + ) + + assert result1 is result2 + + def test_different_teams_get_different_loggers(self, datadog_env): + """Different team credentials should create separate logger instances.""" + cache = DynamicLoggingCache() + + params_a = StandardCallbackDynamicParams( + dd_api_key="team_a_key", + dd_site="us1.datadoghq.com", + ) + params_b = StandardCallbackDynamicParams( + dd_api_key="team_b_key", + dd_site="eu1.datadoghq.com", + ) + + with patch("asyncio.create_task"): + result_a = DataDogHandler.get_datadog_logger_for_request( + standard_callback_dynamic_params=params_a, + in_memory_dynamic_logger_cache=cache, + ) + result_b = DataDogHandler.get_datadog_logger_for_request( + standard_callback_dynamic_params=params_b, + in_memory_dynamic_logger_cache=cache, + ) + + assert result_a is not result_b + assert result_a.DD_API_KEY == "team_a_key" + assert result_b.DD_API_KEY == "team_b_key" + + def test_partial_agent_config_does_not_leak_env_api_key(self, datadog_env): + """A team-supplied dd_agent_host without dd_api_key must not exfiltrate the proxy DD_API_KEY.""" + cache = DynamicLoggingCache() + params = StandardCallbackDynamicParams( + dd_agent_host="attacker.example.com", + ) + + with patch("asyncio.create_task"): + result = DataDogHandler.get_datadog_logger_for_request( + standard_callback_dynamic_params=params, + in_memory_dynamic_logger_cache=cache, + ) + + assert result.DD_API_KEY is None + assert "attacker.example.com" in result.intake_url + + def test_partial_site_config_does_not_leak_env_api_key(self, datadog_env): + """A team-supplied dd_site without dd_api_key must not exfiltrate the proxy DD_API_KEY.""" + cache = DynamicLoggingCache() + params = StandardCallbackDynamicParams( + dd_site="attacker.example.com", + ) + + with pytest.raises(Exception, match="DD_API_KEY"): + with patch("asyncio.create_task"): + DataDogHandler.get_datadog_logger_for_request( + standard_callback_dynamic_params=params, + in_memory_dynamic_logger_cache=cache, + ) + + def test_full_team_config_still_uses_supplied_key(self, datadog_env): + """When a team supplies its own key alongside a custom site, that key (not the env key) is used.""" + cache = DynamicLoggingCache() + params = StandardCallbackDynamicParams( + dd_api_key="team_key", + dd_site="eu1.datadoghq.com", + ) + + with patch("asyncio.create_task"): + result = DataDogHandler.get_datadog_logger_for_request( + standard_callback_dynamic_params=params, + in_memory_dynamic_logger_cache=cache, + ) + + assert result.DD_API_KEY == "team_key" + assert "eu1.datadoghq.com" in result.intake_url + + def test_request_blocked_callback_params_includes_dd(self): + """DD params should be blocked from request-level metadata (security).""" + from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( + _request_blocked_callback_params, + ) + + assert "dd_api_key" in _request_blocked_callback_params + assert "dd_site" in _request_blocked_callback_params + assert "dd_agent_host" in _request_blocked_callback_params + assert "dd_agent_port" in _request_blocked_callback_params + + +class TestDynamicCredentialDetection: + """Test that _dynamic_datadog_credentials_are_passed works correctly.""" + + def test_no_credentials(self): + params = StandardCallbackDynamicParams() + assert DataDogHandler._dynamic_datadog_credentials_are_passed(params) is False + + def test_dd_api_key_only(self): + params = StandardCallbackDynamicParams(dd_api_key="key") + assert DataDogHandler._dynamic_datadog_credentials_are_passed(params) is True + + def test_dd_site_only(self): + params = StandardCallbackDynamicParams(dd_site="site") + assert DataDogHandler._dynamic_datadog_credentials_are_passed(params) is True + + def test_dd_agent_host_only(self): + params = StandardCallbackDynamicParams(dd_agent_host="host") + assert DataDogHandler._dynamic_datadog_credentials_are_passed(params) is True + + +class TestStandardCallbackDynamicParamsIncludesDatadog: + """Verify that Datadog params are in the allow-list.""" + + def test_dd_params_in_annotations(self): + annotations = StandardCallbackDynamicParams.__annotations__ + assert "dd_api_key" in annotations + assert "dd_site" in annotations + assert "dd_agent_host" in annotations + assert "dd_agent_port" in annotations diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_components.py b/tests/test_litellm/integrations/otel/test_otel_v2_components.py index 86d84bd8100..9d81c193af1 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_components.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_components.py @@ -29,6 +29,7 @@ from litellm.integrations.otel.plumbing.metrics import ( from litellm.integrations.otel.model.payloads import ( # noqa: E402 GuardrailSpanData, LLMCallSpanData, + LLMCost, LLMRequestParams, LLMUsage, ProxyRequestSpanData, @@ -224,6 +225,97 @@ def test_genai_mapper_all_request_params(): assert attrs["server.port"] == 443 +def test_genai_mapper_cost_breakdown(): + from litellm.integrations.otel.model.semconv import LiteLLM + + data = LLMCallSpanData( + operation=GenAIOperation.CHAT, + provider="anthropic", + request_model="claude-sonnet-4-6", + response_model=None, + response_id=None, + request_params=LLMRequestParams(), + usage=LLMUsage(), + finish_reasons=(), + error=None, + response_cost=0.012, + server=None, + identity=RequestIdentity(call_id=None), + cost=LLMCost( + input=0.004, + output=0.006, + cache_read=0.001, + cache_creation=0.0, + tool_usage=0.0005, + original=0.013, + discount_amount=0.001, + discount_percent=0.077, + margin_total_amount=0.0, + # margin_fixed_amount / margin_percent left unset on purpose + ), + ) + attrs = GenAIMapper().map(data) + assert attrs[f"{LiteLLM.COST_PREFIX}total"] == 0.012 + assert attrs[f"{LiteLLM.COST_PREFIX}input"] == 0.004 + assert attrs[f"{LiteLLM.COST_PREFIX}output"] == 0.006 + assert attrs[f"{LiteLLM.COST_PREFIX}cache_read"] == 0.001 + assert attrs[f"{LiteLLM.COST_PREFIX}cache_creation"] == 0.0 + assert attrs[f"{LiteLLM.COST_PREFIX}tool_usage"] == 0.0005 + assert attrs[f"{LiteLLM.COST_PREFIX}original"] == 0.013 + assert attrs[f"{LiteLLM.COST_PREFIX}discount_amount"] == 0.001 + assert attrs[f"{LiteLLM.COST_PREFIX}discount_percent"] == 0.077 + assert attrs[f"{LiteLLM.COST_PREFIX}margin_total_amount"] == 0.0 + # Components the source did not report are omitted, not zero-filled. + assert f"{LiteLLM.COST_PREFIX}margin_fixed_amount" not in attrs + assert f"{LiteLLM.COST_PREFIX}margin_percent" not in attrs + + +def test_genai_mapper_cost_breakdown_absent(): + # No cost_breakdown → only the rolled-up total (from response_cost) emits. + from litellm.integrations.otel.model.semconv import LiteLLM + + attrs = GenAIMapper().map(_full_llm_call()) + assert attrs[f"{LiteLLM.COST_PREFIX}total"] == 0.002 + assert not any( + k.startswith(LiteLLM.COST_PREFIX) and k != f"{LiteLLM.COST_PREFIX}total" + for k in attrs + ) + + +def test_llm_cost_from_breakdown_maps_costbreakdown_keys(): + cost = LLMCost.from_breakdown( + { + "input_cost": 0.004, + "output_cost": 0.006, + "cache_read_cost": 0.001, + "cache_creation_cost": 0.002, + "tool_usage_cost": 0.0005, + "original_cost": 0.013, + "discount_amount": 0.001, + "discount_percent": 0.077, + "margin_fixed_amount": 0.0, + "margin_percent": 0.1, + "margin_total_amount": 0.0011, + "total_cost": 0.012, # carried on response_cost, not LLMCost + } + ) + assert cost.input == 0.004 + assert cost.output == 0.006 + assert cost.cache_read == 0.001 + assert cost.cache_creation == 0.002 + assert cost.tool_usage == 0.0005 + assert cost.original == 0.013 + assert cost.discount_amount == 0.001 + assert cost.discount_percent == 0.077 + assert cost.margin_fixed_amount == 0.0 + assert cost.margin_percent == 0.1 + assert cost.margin_total_amount == 0.0011 + + +def test_llm_cost_from_breakdown_none_is_empty(): + assert LLMCost.from_breakdown(None) == LLMCost() + + def test_genai_mapper_guardrail_and_service(): from litellm.integrations.otel.model.semconv import LiteLLM diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py b/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py index 2dbedda1ab6..48190a798da 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_emitter.py @@ -57,6 +57,42 @@ def _engine(legacy_compat=True): return SpanEmitter(tracer, cfg), exporter +def test_llm_call_span_cost_breakdown(): + engine, exporter = _engine() + data = LLMCallSpanData.from_standard_logging_payload( + _payload( + cost_breakdown={ + "input_cost": 0.004, + "output_cost": 0.006, + "cache_read_cost": 0.001, + "total_cost": 0.011, + } + ) + ) + engine.emit(SpanRole.LLM_CALL, data) + (span,) = exporter.get_finished_spans() + a = span.attributes + # The rolled-up total stays sourced from response_cost. + assert a[f"{LiteLLM.COST_PREFIX}total"] == 0.002 + # Per-component breakdown now rides the span. + assert a[f"{LiteLLM.COST_PREFIX}input"] == 0.004 + assert a[f"{LiteLLM.COST_PREFIX}output"] == 0.006 + assert a[f"{LiteLLM.COST_PREFIX}cache_read"] == 0.001 + # Unreported components are omitted, not zero-filled. + assert f"{LiteLLM.COST_PREFIX}margin_total_amount" not in a + + +def test_tracer_scope_carries_litellm_version(): + from litellm._version import version as litellm_version + + cfg = OpenTelemetryV2Config(exporter="in_memory") + provider, exporter = providers.in_memory_provider(cfg) + tracer = providers.get_tracer(provider, "litellm-test") + tracer.start_span("probe").end() + (span,) = exporter.get_finished_spans() + assert span.instrumentation_scope.version == litellm_version + + def test_llm_call_span_golden(): engine, exporter = _engine() data = LLMCallSpanData.from_standard_logging_payload(_payload()) diff --git a/tests/test_litellm/integrations/test_langfuse_otel.py b/tests/test_litellm/integrations/test_langfuse_otel.py index 44853d9dce5..3aade7514e4 100644 --- a/tests/test_litellm/integrations/test_langfuse_otel.py +++ b/tests/test_litellm/integrations/test_langfuse_otel.py @@ -114,6 +114,9 @@ class TestLangfuseOtelIntegration: mock_set_attributes.assert_called_once_with( mock_span, mock_kwargs, mock_response, LangfuseLLMObsOTELAttributes ) + mock_span.set_attribute.assert_any_call( + "langfuse.observation.type", "generation" + ) def test_set_langfuse_environment_attribute(self): """Test that Langfuse environment is set correctly when environment variable is present.""" diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py new file mode 100644 index 00000000000..aff89f02ff2 --- /dev/null +++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py @@ -0,0 +1,41 @@ +import json +import os +import sys + +import pytest + +sys.path.insert( + 0, os.path.abspath("../../../../../..") +) # Adds the parent directory to the system path + +from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import ( + AmazonAnthropicClaudeConfig, +) +from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import ( + AmazonInvokeConfig, +) + + +@pytest.mark.parametrize( + "config,model", + [ + (AmazonInvokeConfig, "anthropic.claude-3-sonnet-20240229-v1:0"), + (AmazonInvokeConfig, "amazon.titan-text-express-v1"), + (AmazonInvokeConfig, "mistral.mistral-7b-instruct-v0:2"), + (AmazonAnthropicClaudeConfig, "anthropic.claude-sonnet-4-6"), + ], +) +def test_transform_request_drops_stream_chunk_size(config, model): + """stream_chunk_size is a LiteLLM-internal knob for re-chunking the HTTP + response stream. Leaking it into the provider request body makes Bedrock + reject the whole request: ValidationException 'stream_chunk_size: Extra + inputs are not permitted'.""" + request_body = config().transform_request( + model=model, + messages=[{"role": "user", "content": "hi"}], + optional_params={"stream": True, "stream_chunk_size": 2048, "max_tokens": 10}, + litellm_params={}, + headers={}, + ) + + assert "stream_chunk_size" not in json.dumps(request_body) diff --git a/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py b/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py index a415d550215..61987d25d9c 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py +++ b/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py @@ -1,12 +1,21 @@ import os import sys +from unittest.mock import AsyncMock, MagicMock +import pytest sys.path.insert( 0, os.path.abspath("../../../../..") ) # Adds the parent directory to the system path -from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder +import litellm +from litellm.llms.bedrock.chat.invoke_handler import ( + AWSEventStreamDecoder, + BedrockLLM, + make_call, + make_sync_call, +) +from litellm.llms.custom_httpx.http_handler import HTTPHandler def test_transform_thinking_blocks_with_redacted_content(): @@ -200,3 +209,120 @@ def test_bedrock_converse_streaming_consistent_id(): assert ( response.id == expected_id ), "All chunk IDs must match the one captured from the messageStart event" + + +@pytest.mark.asyncio +async def test_make_call_does_not_rechunk_stream_by_default(): + """Re-chunking the event stream into fixed 1024-byte blocks holds small + early events (messageStart, contentBlockStart) in httpx's ByteChunker until + 1024 bytes accumulate, delaying time-to-first-chunk by the whole generation + when Bedrock trickles bytes (e.g. buffered tool-use streams).""" + response = MagicMock() + response.status_code = 200 + client = MagicMock() + client.post = AsyncMock(return_value=response) + + await make_call( + client=client, + api_base="https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-sonnet-4-6/converse-stream", + headers={}, + data="{}", + model="anthropic.claude-sonnet-4-6", + messages=[], + logging_obj=MagicMock(), + ) + + response.aiter_bytes.assert_called_once_with(chunk_size=None) + + +@pytest.mark.asyncio +async def test_make_call_honors_explicit_stream_chunk_size(): + response = MagicMock() + response.status_code = 200 + client = MagicMock() + client.post = AsyncMock(return_value=response) + + await make_call( + client=client, + api_base="https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-sonnet-4-6/converse-stream", + headers={}, + data="{}", + model="anthropic.claude-sonnet-4-6", + messages=[], + logging_obj=MagicMock(), + stream_chunk_size=2048, + ) + + response.aiter_bytes.assert_called_once_with(chunk_size=2048) + + +def test_make_sync_call_does_not_rechunk_stream_by_default(): + response = MagicMock() + response.status_code = 200 + client = MagicMock() + client.post = MagicMock(return_value=response) + + make_sync_call( + client=client, + api_base="https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-sonnet-4-6/converse-stream", + headers={}, + data="{}", + signed_json_body=None, + model="anthropic.claude-sonnet-4-6", + messages=[], + logging_obj=MagicMock(), + ) + + response.iter_bytes.assert_called_once_with(chunk_size=None) + + +def test_make_sync_call_honors_explicit_stream_chunk_size(): + response = MagicMock() + response.status_code = 200 + client = MagicMock() + client.post = MagicMock(return_value=response) + + make_sync_call( + client=client, + api_base="https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-sonnet-4-6/converse-stream", + headers={}, + data="{}", + signed_json_body=None, + model="anthropic.claude-sonnet-4-6", + messages=[], + logging_obj=MagicMock(), + stream_chunk_size=2048, + ) + + response.iter_bytes.assert_called_once_with(chunk_size=2048) + + +def test_legacy_bedrock_llm_streaming_does_not_rechunk_by_default(): + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.iter_bytes = MagicMock(return_value=iter([])) + client = HTTPHandler() + client.post = MagicMock(return_value=mock_response) + + BedrockLLM().completion( + model="cohere.command-text-v14", + messages=[{"role": "user", "content": "hi"}], + api_base=None, + custom_prompt_dict={}, + model_response=litellm.ModelResponse(), + print_verbose=lambda *args, **kwargs: None, + encoding=litellm.encoding, + logging_obj=MagicMock(), + optional_params={ + "stream": True, + "aws_access_key_id": "fake", + "aws_secret_access_key": "fake", + "aws_region_name": "us-east-1", + }, + acompletion=False, + timeout=None, + litellm_params={}, + client=client, + ) + + mock_response.iter_bytes.assert_called_once_with(chunk_size=None) diff --git a/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py b/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py new file mode 100644 index 00000000000..2cf1fa16e91 --- /dev/null +++ b/tests/test_litellm/llms/bedrock/test_web_identity_session_policy.py @@ -0,0 +1,176 @@ +""" +Regression for #30200. + +``_auth_with_web_identity_token`` passes an inline ``Policy`` to +``sts.assume_role_with_web_identity``. In AWS IAM an STS session policy +acts as a PERMISSION CEILING — effective permissions are the +intersection of the role's identity policies and this policy, so any +action not listed here 403s on OIDC-auth requests only (static creds +and IRSA flow through different paths). + +The original policy only granted ``bedrock:*`` actions. When +``#27678`` added the ``bedrock/claude_platform/`` route, the +service-side action namespace was ``aws-external-anthropic:*``, not +``bedrock:*``, so every claude_platform call via OIDC silently denied +with:: + + User: arn:aws:sts::ACCOUNT:assumed-role/... + is not authorized to perform: aws-external-anthropic:CreateInference + on resource: arn:aws:aws-external-anthropic:... + because no session policy allows the + aws-external-anthropic:CreateInference action + +— even with a fully permissive identity policy. + +Tests below intercept the kwargs handed to +``assume_role_with_web_identity``, parse the embedded ``Policy`` JSON, +and assert that both the original bedrock statement and the new +claude_platform statement are present and cover every documented +action. +""" + +import json +from datetime import datetime, timedelta, timezone +from unittest.mock import MagicMock, patch + +import pytest + +# Actions the Claude Platform on AWS service is documented to call. +# Source: AWS IAM action reference + the #27678 surface area. +_CLAUDE_PLATFORM_ACTIONS = { + "aws-external-anthropic:CreateInference", + "aws-external-anthropic:CreateBatchInference", + "aws-external-anthropic:CancelBatchInference", + "aws-external-anthropic:DeleteBatchInference", + "aws-external-anthropic:CountTokens", + "aws-external-anthropic:Get*", + "aws-external-anthropic:List*", +} + + +def _captured_policy() -> dict: + """Run _auth_with_web_identity_token under mocks + return the parsed + Policy dict that was actually sent to STS.""" + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + base = BaseAWSLLM() + + mock_sts = MagicMock() + mock_sts.assume_role_with_web_identity.return_value = { + "Credentials": { + "AccessKeyId": "k", + "SecretAccessKey": "s", + "SessionToken": "t", + "Expiration": datetime.now(timezone.utc) + timedelta(hours=1), + }, + "PackedPolicySize": 0, + } + + with ( + patch("boto3.client", return_value=mock_sts), + patch( + "litellm.llms.bedrock.base_aws_llm.get_secret", + return_value="oidc-jwt-token", + ), + ): + base._auth_with_web_identity_token( + aws_web_identity_token="/path/to/token", + aws_role_name="arn:aws:iam::123456789012:role/litellm-bedrock-role", + aws_session_name="test-session", + aws_region_name="us-east-1", + aws_sts_endpoint=None, + ) + + mock_sts.assume_role_with_web_identity.assert_called_once() + kwargs = mock_sts.assume_role_with_web_identity.call_args.kwargs + policy_str = kwargs["Policy"] + return json.loads(policy_str) + + +def _statement_by_sid(policy: dict, sid: str) -> dict: + for stmt in policy["Statement"]: + if stmt.get("Sid") == sid: + return stmt + raise AssertionError( + f"Sid={sid!r} not found in session policy; " + f"saw {[s.get('Sid') for s in policy['Statement']]}" + ) + + +class TestWebIdentitySessionPolicyShape: + def test_policy_parses_as_valid_iam_document(self): + policy = _captured_policy() + assert policy["Version"] == "2012-10-17" + assert isinstance(policy["Statement"], list) + assert len(policy["Statement"]) >= 2 + + def test_bedrock_statement_actions_preserved(self): + """The original bedrock action set must still be granted — + regression for the pre-existing bedrock/* routes.""" + policy = _captured_policy() + bedrock_stmt = _statement_by_sid(policy, "BedrockLiteLLM") + actions = set(bedrock_stmt["Action"]) + for required in ( + "bedrock:InvokeModel", + "bedrock:InvokeModelWithResponseStream", + ): + assert required in actions, f"{required} missing from BedrockLiteLLM" + + +class TestClaudePlatformActionsCovered: + """The #30200 bug: every action in the claude_platform service + namespace must appear in the session policy or OIDC requests 403.""" + + @pytest.mark.parametrize("action", sorted(_CLAUDE_PLATFORM_ACTIONS)) + def test_claude_platform_action_present(self, action: str): + policy = _captured_policy() + # Action may live in any Statement — search across all. + all_actions: set = set() + for stmt in policy["Statement"]: + stmt_actions = stmt.get("Action") + if isinstance(stmt_actions, str): + all_actions.add(stmt_actions) + elif isinstance(stmt_actions, list): + all_actions.update(stmt_actions) + assert action in all_actions, ( + f"{action} missing from session policy — " + f"bedrock/claude_platform/* requests will 403 on OIDC auth" + ) + + def test_claude_platform_statement_allows(self): + policy = _captured_policy() + stmt = _statement_by_sid(policy, "ClaudePlatformLiteLLM") + assert stmt["Effect"] == "Allow" + assert stmt["Resource"] == "*" + + def test_no_aws_external_anthropic_statement_collision(self): + """Don't accidentally grant a `*` action that would broaden the + ceiling beyond what the documented actions require.""" + policy = _captured_policy() + stmt = _statement_by_sid(policy, "ClaudePlatformLiteLLM") + actions = stmt["Action"] + if isinstance(actions, str): + actions = [actions] + assert "aws-external-anthropic:*" not in actions, ( + "session policy must not grant aws-external-anthropic:* — " + "the ceiling should match the documented action set" + ) + + +class TestPolicyTransportConditions: + def test_bedrock_statement_keeps_secure_transport_condition(self): + policy = _captured_policy() + bedrock_stmt = _statement_by_sid(policy, "BedrockLiteLLM") + cond = bedrock_stmt.get("Condition") or {} + assert cond.get("Bool", {}).get("aws:SecureTransport") == "true" + + def test_claude_platform_statement_carries_secure_transport_condition(self): + """The new statement should match the existing one's hardening + posture — TLS-only, same as bedrock.""" + policy = _captured_policy() + stmt = _statement_by_sid(policy, "ClaudePlatformLiteLLM") + cond = stmt.get("Condition") or {} + assert cond.get("Bool", {}).get("aws:SecureTransport") == "true", ( + "ClaudePlatformLiteLLM must require aws:SecureTransport=true " + "to keep parity with the bedrock statement" + ) diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py index 061d378f757..6fb02113a45 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py @@ -20,6 +20,23 @@ from litellm.llms.bedrock_mantle.chat.transformation import BedrockMantleChatCon from litellm.types.utils import LlmProviders +@pytest.fixture +def local_cost_map(monkeypatch): + original_model_cost = litellm.model_cost + original_bedrock_mantle_models = set(litellm.bedrock_mantle_models) + try: + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "true") + litellm.model_cost = litellm.get_model_cost_map(url="") + litellm.get_model_info.cache_clear() + litellm.add_known_models() + yield + finally: + litellm.model_cost = original_model_cost + litellm.bedrock_mantle_models.clear() + litellm.bedrock_mantle_models.update(original_bedrock_mantle_models) + litellm.get_model_info.cache_clear() + + class TestBedrockMantleProviderRegistration: def test_provider_enum_exists(self): assert LlmProviders.BEDROCK_MANTLE == "bedrock_mantle" @@ -310,3 +327,52 @@ class TestBedrockMantlePricing: litellm.add_known_models() info = litellm.get_model_info("bedrock_mantle/openai.gpt-oss-120b") assert info["max_input_tokens"] == 131072 + + +@pytest.mark.parametrize( + "model_id,input_cost,output_cost,max_tokens", + [ + ("google.gemma-4-31b", 1.4e-07, 4e-07, 256000), + ("google.gemma-4-26b-a4b", 1.3e-07, 4e-07, 256000), + ("google.gemma-4-e2b", 4e-08, 8e-08, 128000), + ], +) +def test_gemma_4_bedrock_mantle_model_metadata( + local_cost_map, model_id, input_cost, output_cost, max_tokens +): + full_model_name = f"bedrock_mantle/{model_id}" + info = litellm.get_model_info(full_model_name) + + assert info["mode"] == "chat" + assert info["input_cost_per_token"] == pytest.approx(input_cost) + assert info["output_cost_per_token"] == pytest.approx(output_cost) + assert info["max_input_tokens"] == max_tokens + assert info["max_output_tokens"] == max_tokens + assert info["supports_function_calling"] is True + assert info["supports_reasoning"] is True + assert info["supports_tool_choice"] is True + assert info["supports_vision"] is True + assert ( + litellm.supports_parallel_function_calling( + model=full_model_name, custom_llm_provider="bedrock_mantle" + ) + is False + ) + + +@pytest.mark.parametrize( + "model_id", + [ + "google.gemma-4-31b", + "google.gemma-4-26b-a4b", + "google.gemma-4-e2b", + ], +) +def test_gemma_4_models_register_under_bedrock_mantle(local_cost_map, model_id): + full_model_name = f"bedrock_mantle/{model_id}" + + assert full_model_name in litellm.bedrock_mantle_models + + resolved_model, provider, _, _ = litellm.get_llm_provider(full_model_name) + assert provider == "bedrock_mantle" + assert resolved_model == model_id diff --git a/tests/test_litellm/llms/chat/test_converse_handler.py b/tests/test_litellm/llms/chat/test_converse_handler.py index b636ea468ca..2a3db5982ef 100644 --- a/tests/test_litellm/llms/chat/test_converse_handler.py +++ b/tests/test_litellm/llms/chat/test_converse_handler.py @@ -1,10 +1,14 @@ import os import sys +from unittest.mock import MagicMock import pytest +import litellm from litellm.llms.bedrock.chat import BedrockConverseLLM +from litellm.llms.bedrock.chat.converse_handler import make_sync_call from litellm.llms.bedrock.common_utils import _get_all_bedrock_regions +from litellm.llms.custom_httpx.http_handler import HTTPHandler sys.path.insert( 0, os.path.abspath("../../../../..") @@ -133,3 +137,79 @@ class TestBedrockRegionInModelPath: assert model_id == "moonshotai.kimi-k2.5" # explicitly set region is preserved assert optional_params["aws_region_name"] == "eu-west-1" + + +def _stream_completion_with_spied_iter_bytes(model: str, **kwargs) -> MagicMock: + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.iter_bytes = MagicMock(return_value=iter([])) + client = HTTPHandler() + client.post = MagicMock(return_value=mock_response) + + litellm.completion( + model=model, + messages=[{"role": "user", "content": "hi"}], + stream=True, + client=client, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + **kwargs, + ) + return mock_response.iter_bytes + + +def test_make_sync_call_does_not_rechunk_stream_by_default(): + """Re-chunking the event stream into fixed 1024-byte blocks holds small + early events in httpx's ByteChunker until 1024 bytes accumulate, delaying + time-to-first-chunk by the whole generation when Bedrock trickles bytes + (e.g. buffered tool-use streams).""" + response = MagicMock() + response.status_code = 200 + client = MagicMock() + client.post = MagicMock(return_value=response) + + make_sync_call( + client=client, + api_base="https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-sonnet-4-6/converse-stream", + headers={}, + data="{}", + model="anthropic.claude-sonnet-4-6", + messages=[], + logging_obj=MagicMock(), + ) + + response.iter_bytes.assert_called_once_with(chunk_size=None) + + +def test_make_sync_call_honors_explicit_stream_chunk_size(): + response = MagicMock() + response.status_code = 200 + client = MagicMock() + client.post = MagicMock(return_value=response) + + make_sync_call( + client=client, + api_base="https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-sonnet-4-6/converse-stream", + headers={}, + data="{}", + model="anthropic.claude-sonnet-4-6", + messages=[], + logging_obj=MagicMock(), + stream_chunk_size=2048, + ) + + response.iter_bytes.assert_called_once_with(chunk_size=2048) + + +def test_completion_plumbs_stream_chunk_size_through_converse(): + iter_bytes_spy = _stream_completion_with_spied_iter_bytes( + model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0" + ) + iter_bytes_spy.assert_called_once_with(chunk_size=None) + + iter_bytes_spy = _stream_completion_with_spied_iter_bytes( + model="bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", + stream_chunk_size=2048, + ) + iter_bytes_spy.assert_called_once_with(chunk_size=2048) diff --git a/tests/test_litellm/llms/openai_like/test_tensormesh_provider.py b/tests/test_litellm/llms/openai_like/test_tensormesh_provider.py index 09248a779c5..c94b2cbfa80 100644 --- a/tests/test_litellm/llms/openai_like/test_tensormesh_provider.py +++ b/tests/test_litellm/llms/openai_like/test_tensormesh_provider.py @@ -79,6 +79,20 @@ class TestTensormeshProviderConfig: matching the text_completion flag in provider_endpoints_support.json.""" assert "tensormesh" in litellm.openai_text_completion_compatible_providers + def test_tensormesh_responses_api_enabled(self): + """Tensormesh declares /v1/responses in supported_endpoints, so litellm + resolves a responses config for it.""" + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + from litellm.utils import ProviderConfigManager + + assert JSONProviderRegistry.supports_responses_api("tensormesh") is True + config = ProviderConfigManager.get_provider_responses_api_config( + provider="tensormesh", + model="tensormesh/openai/gpt-oss-120b", + ) + assert config is not None + assert config.custom_llm_provider == "tensormesh" + def test_tensormesh_router_config(self): """Test that tensormesh can be used in Router configuration""" from litellm import Router diff --git a/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py b/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py index 31e1c61d6ac..a182656e4a8 100644 --- a/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py +++ b/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py @@ -26,11 +26,13 @@ class TestSnowflakeToolTransformation: def test_transform_request_with_tools(self): """ - Test that OpenAI tool format is correctly transformed to Snowflake's tool_spec format. + Test that OpenAI tool format is passed through as-is to the native endpoint. + + The native /chat/completions endpoint accepts standard OpenAI tool format + directly — no Snowflake-specific tool_spec transformation needed. """ config = SnowflakeConfig() - # OpenAI format tools tools = [ { "type": "function", @@ -58,113 +60,94 @@ class TestSnowflakeToolTransformation: optional_params = {"tools": tools} transformed_request = config.transform_request( - model="claude-3-5-sonnet", + model="llama3.1-70b", messages=[{"role": "user", "content": "What's the weather?"}], optional_params=optional_params, litellm_params={}, headers={}, ) - # Verify tools were transformed to Snowflake format assert "tools" in transformed_request assert len(transformed_request["tools"]) == 1 - - snowflake_tool = transformed_request["tools"][0] - assert "tool_spec" in snowflake_tool - assert snowflake_tool["tool_spec"]["type"] == "generic" - assert snowflake_tool["tool_spec"]["name"] == "get_weather" - assert ( - snowflake_tool["tool_spec"]["description"] - == "Get the current weather in a given location" - ) - assert "input_schema" in snowflake_tool["tool_spec"] - assert snowflake_tool["tool_spec"]["input_schema"]["type"] == "object" - assert "location" in snowflake_tool["tool_spec"]["input_schema"]["properties"] + assert transformed_request["tools"] == tools + assert "tool_spec" not in json.dumps(transformed_request) def test_transform_request_with_tool_choice(self): """ - Test that OpenAI tool_choice format is correctly transformed to Snowflake format. + Test that OpenAI tool_choice format is passed through as-is to the native endpoint. """ config = SnowflakeConfig() - # OpenAI format tool_choice tool_choice = {"type": "function", "function": {"name": "get_weather"}} optional_params = {"tool_choice": tool_choice} transformed_request = config.transform_request( - model="claude-3-5-sonnet", + model="llama3.1-70b", messages=[{"role": "user", "content": "What's the weather?"}], optional_params=optional_params, litellm_params={}, headers={}, ) - # Verify tool_choice was transformed to Snowflake format assert "tool_choice" in transformed_request - assert transformed_request["tool_choice"]["type"] == "tool" - assert transformed_request["tool_choice"]["name"] == [ - "get_weather" - ] # Array format + assert transformed_request["tool_choice"] == tool_choice def test_transform_request_with_string_tool_choice(self): """ - Test that string tool_choice values are transformed to Snowflake object format. + Test that string tool_choice values are passed through as-is to the native endpoint. - Snowflake's API (like Anthropic) requires tool_choice as an object - with a "type" field, not as a bare string. OpenAI's "required" maps - to Snowflake's "any". + The native /chat/completions endpoint accepts OpenAI-style string + tool_choice values directly ("auto", "required", "none"). """ config = SnowflakeConfig() - expected_mappings = { - "auto": {"type": "auto"}, - "required": {"type": "any"}, - "none": {"type": "none"}, - } - - for value, expected in expected_mappings.items(): + for value in ["auto", "required", "none"]: optional_params = {"tool_choice": value} transformed_request = config.transform_request( - model="claude-3-5-sonnet", + model="llama3.1-70b", messages=[{"role": "user", "content": "Test"}], optional_params=optional_params, litellm_params={}, headers={}, ) - assert transformed_request["tool_choice"] == expected, ( - f"tool_choice='{value}' should be transformed to {expected}, " + assert transformed_request["tool_choice"] == value, ( + f"tool_choice='{value}' should pass through unchanged, " f"got {transformed_request['tool_choice']}" ) def test_transform_response_with_tool_calls(self): """ - Test that Snowflake's content_list with tool_use is transformed to OpenAI format. + Test that standard OpenAI tool_calls response format is parsed correctly. + + The native /chat/completions endpoint returns standard OpenAI format. """ config = SnowflakeConfig() - # Mock Snowflake response with tool call - mock_snowflake_response = { + mock_response = { + "id": "chatcmpl-123", + "object": "chat.completion", + "model": "llama3.1-70b", "choices": [ { + "index": 0, "message": { - "content_list": [ - {"type": "text", "text": ""}, + "role": "assistant", + "content": None, + "tool_calls": [ { - "type": "tool_use", - "tool_use": { - "tool_use_id": "tooluse_abc123", + "id": "call_abc123", + "type": "function", + "function": { "name": "get_weather", - "input": { - "location": "Paris, France", - "unit": "celsius", - }, + "arguments": json.dumps({"location": "Paris, France", "unit": "celsius"}), }, - }, - ] - } + } + ], + }, + "finish_reason": "tool_calls", } ], "usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, @@ -172,7 +155,7 @@ class TestSnowflakeToolTransformation: response = httpx.Response( status_code=200, - json=mock_snowflake_response, + json=mock_response, headers={"Content-Type": "application/json"}, ) @@ -183,7 +166,7 @@ class TestSnowflakeToolTransformation: logging_obj = MagicMock() result = config.transform_response( - model="claude-3-5-sonnet", + model="llama3.1-70b", raw_response=response, model_response=model_response, logging_obj=logging_obj, @@ -194,61 +177,50 @@ class TestSnowflakeToolTransformation: encoding={}, ) - # General assertions assert isinstance(result, ModelResponse) assert len(result.choices) == 1 - choice = result.choices[0] - assert isinstance(choice, litellm.Choices) - - # Message and tool_calls assertions - message = choice.message - assert isinstance(message, litellm.Message) - assert hasattr(message, "tool_calls") - assert isinstance(message.tool_calls, list) + message = result.choices[0].message + assert message.tool_calls is not None assert len(message.tool_calls) == 1 - # Specific tool_call assertions tool_call = message.tool_calls[0] - assert isinstance(tool_call, litellm.utils.ChatCompletionMessageToolCall) - assert tool_call.id == "tooluse_abc123" + assert tool_call.id == "call_abc123" assert tool_call.type == "function" assert tool_call.function.name == "get_weather" - # Verify arguments are properly JSON serialized arguments = json.loads(tool_call.function.arguments) assert arguments["location"] == "Paris, France" assert arguments["unit"] == "celsius" - # Verify content_list was removed and content was set - assert message.content == "" - def test_transform_response_with_mixed_content(self): """ - Test that responses with both text and tool calls are handled correctly. + Test that responses with both text content and tool calls are parsed correctly. """ config = SnowflakeConfig() - # Mock Snowflake response with text and tool call - mock_snowflake_response = { + mock_response = { + "id": "chatcmpl-456", + "object": "chat.completion", + "model": "llama3.1-70b", "choices": [ { + "index": 0, "message": { - "content_list": [ + "role": "assistant", + "content": "Let me check the weather for you.", + "tool_calls": [ { - "type": "text", - "text": "Let me check the weather for you. ", - }, - { - "type": "tool_use", - "tool_use": { - "tool_use_id": "tooluse_xyz789", + "id": "call_xyz789", + "type": "function", + "function": { "name": "get_weather", - "input": {"location": "Tokyo, Japan"}, + "arguments": json.dumps({"location": "Tokyo, Japan"}), }, - }, - ] - } + } + ], + }, + "finish_reason": "tool_calls", } ], "usage": {"prompt_tokens": 15, "completion_tokens": 25, "total_tokens": 40}, @@ -256,7 +228,7 @@ class TestSnowflakeToolTransformation: response = httpx.Response( status_code=200, - json=mock_snowflake_response, + json=mock_response, headers={"Content-Type": "application/json"}, ) @@ -267,7 +239,7 @@ class TestSnowflakeToolTransformation: logging_obj = MagicMock() result = config.transform_response( - model="claude-3-5-sonnet", + model="llama3.1-70b", raw_response=response, model_response=model_response, logging_obj=logging_obj, @@ -278,11 +250,8 @@ class TestSnowflakeToolTransformation: encoding={}, ) - # Verify text content was extracted message = result.choices[0].message - assert message.content == "Let me check the weather for you. " - - # Verify tool call was also extracted + assert message.content == "Let me check the weather for you." assert len(message.tool_calls) == 1 assert message.tool_calls[0].function.name == "get_weather" @@ -341,7 +310,7 @@ class TestSnowflakeToolTransformation: Test that tools and tool_choice are in supported params. """ config = SnowflakeConfig() - supported_params = config.get_supported_openai_params("claude-3-5-sonnet") + supported_params = config.get_supported_openai_params("llama3.1-70b") assert "tools" in supported_params assert "tool_choice" in supported_params @@ -392,8 +361,8 @@ class TestSnowFlakeCompletion: assert "00000" in post_kwargs["headers"]["Authorization"] # account id was used assert "AAAA-BBBB" in post_kwargs["url"] - # is completion - assert post_kwargs["url"].endswith("cortex/inference:complete") + # uses native endpoint + assert post_kwargs["url"].endswith("cortex/v1/chat/completions") @patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post") def test_snowflake_pat_key_account_id(self, mock_post): diff --git a/tests/test_litellm/llms/snowflake/test_snowflake_native_endpoints.py b/tests/test_litellm/llms/snowflake/test_snowflake_native_endpoints.py new file mode 100644 index 00000000000..fb21e2e6f6b --- /dev/null +++ b/tests/test_litellm/llms/snowflake/test_snowflake_native_endpoints.py @@ -0,0 +1,718 @@ +""" +Tests for Snowflake Cortex native endpoint migration. + +Covers: + - SnowflakeConfig with auto-routing: + - Non-Claude models → /chat/completions (OpenAI format) + - Claude models → /messages (Anthropic format) + +Run: + pytest tests/test_litellm/llms/snowflake/test_snowflake_native_endpoints.py -v +""" + +import json +from unittest.mock import MagicMock, patch + +import httpx +import pytest + +from litellm.llms.snowflake.chat.transformation import ( + SnowflakeConfig, + _is_claude_model, +) +from litellm.types.utils import ModelResponse + + +# ─── Fixtures ────────────────────────────────────────────────────────────── + +ACCOUNT_ID = "myaccount" +API_BASE = f"https://{ACCOUNT_ID}.snowflakecomputing.com" +PAT_TOKEN = "pat/my-secret-pat-token" +JWT_TOKEN = "eyJhbGciOiJSUzI1NiJ9.test" + + +def _mock_logging(): + m = MagicMock() + m.post_call = MagicMock() + return m + + +def _make_openai_response(content: str = "Hello!") -> httpx.Response: + body = { + "id": "chatcmpl-abc123", + "object": "chat.completion", + "model": "llama3.1-70b", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": content}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + } + return httpx.Response(200, json=body) + + +def _make_anthropic_response(content: str = "Hello!") -> httpx.Response: + body = { + "id": "msg_abc123", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5", + "content": [{"type": "text", "text": content}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 10, "output_tokens": 5}, + } + return httpx.Response(200, json=body) + + +# ─── SnowflakeConfig (OpenAI-compatible) ─────────────────────────────────── + +class TestSnowflakeConfigURL: + def setup_method(self): + self.cfg = SnowflakeConfig() + + def test_url_with_account_id_in_optional_params(self): + optional_params = {"account_id": ACCOUNT_ID} + url = self.cfg.get_complete_url( + api_base=None, + api_key=JWT_TOKEN, + model="snowflake/llama3.1-70b", + optional_params=optional_params, + litellm_params={}, + ) + assert url == f"https://{ACCOUNT_ID}.snowflakecomputing.com/api/v2/cortex/v1/chat/completions" + + def test_url_with_explicit_api_base(self): + url = self.cfg.get_complete_url( + api_base=API_BASE, + api_key=JWT_TOKEN, + model="snowflake/llama3.1-70b", + optional_params={}, + litellm_params={}, + ) + assert url.endswith("/api/v2/cortex/v1/chat/completions") + assert "cortex/inference:complete" not in url + + def test_url_never_uses_legacy_endpoint(self): + url = self.cfg.get_complete_url( + api_base=API_BASE, + api_key=JWT_TOKEN, + model="snowflake/llama3.1-70b", + optional_params={}, + litellm_params={}, + ) + assert "inference:complete" not in url + assert "/v1/chat/completions" in url + + def test_url_works_for_claude_models(self): + url = self.cfg.get_complete_url( + api_base=API_BASE, + api_key=JWT_TOKEN, + model="snowflake/claude-sonnet-4-5", + optional_params={}, + litellm_params={}, + ) + assert "/cortex/v1/messages" in url + + def test_url_works_for_llama_models(self): + url = self.cfg.get_complete_url( + api_base=API_BASE, + api_key=JWT_TOKEN, + model="snowflake/llama3.1-70b", + optional_params={}, + litellm_params={}, + ) + assert "/cortex/v1/chat/completions" in url + + +class TestSnowflakeConfigAuth: + def setup_method(self): + self.cfg = SnowflakeConfig() + + def test_pat_auth_strips_prefix_and_sets_header(self): + headers = self.cfg.validate_environment( + headers={}, + model="snowflake/llama3.1-70b", + messages=[], + optional_params={}, + litellm_params={}, + api_key=PAT_TOKEN, + ) + assert headers["X-Snowflake-Authorization-Token-Type"] == "PROGRAMMATIC_ACCESS_TOKEN" + assert headers["Authorization"] == "Bearer my-secret-pat-token" + + def test_jwt_auth_sets_keypair_header(self): + headers = self.cfg.validate_environment( + headers={}, + model="snowflake/llama3.1-70b", + messages=[], + optional_params={}, + litellm_params={}, + api_key=JWT_TOKEN, + ) + assert headers["X-Snowflake-Authorization-Token-Type"] == "KEYPAIR_JWT" + assert headers["Authorization"] == f"Bearer {JWT_TOKEN}" + + def test_missing_api_key_raises(self): + with pytest.raises(ValueError, match="Missing Snowflake JWT key"): + self.cfg.validate_environment( + headers={}, + model="snowflake/llama3.1-70b", + messages=[], + optional_params={}, + litellm_params={}, + api_key=None, + ) + + +class TestSnowflakeConfigRequest: + def setup_method(self): + self.cfg = SnowflakeConfig() + self.messages = [{"role": "user", "content": "hello"}] + + def test_request_uses_openai_tool_format(self): + tools = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get weather", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}, + }, + } + ] + body = self.cfg.transform_request( + model="snowflake/llama3.1-70b", + messages=self.messages, + optional_params={"tools": tools}, + litellm_params={}, + headers={}, + ) + assert body["tools"] == tools + assert "tool_spec" not in json.dumps(body) + + def test_stream_defaults_to_false(self): + body = self.cfg.transform_request( + model="snowflake/llama3.1-70b", + messages=self.messages, + optional_params={}, + litellm_params={}, + headers={}, + ) + assert body["stream"] is False + + def test_stream_true_passes_through(self): + body = self.cfg.transform_request( + model="snowflake/llama3.1-70b", + messages=self.messages, + optional_params={"stream": True}, + litellm_params={}, + headers={}, + ) + assert body["stream"] is True + + def test_supported_params_includes_stream(self): + params = self.cfg.get_supported_openai_params("snowflake/llama3.1-70b") + assert "stream" in params + + def test_no_content_list_in_request(self): + body = self.cfg.transform_request( + model="snowflake/llama3.1-70b", + messages=self.messages, + optional_params={}, + litellm_params={}, + headers={}, + ) + assert "content_list" not in body + + +class TestSnowflakeConfigResponse: + def setup_method(self): + self.cfg = SnowflakeConfig() + + def test_standard_response_parsed(self): + raw = _make_openai_response("Hello from Snowflake!") + result = self.cfg.transform_response( + model="snowflake/llama3.1-70b", + raw_response=raw, + model_response=ModelResponse(), + logging_obj=_mock_logging(), + request_data={}, + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + encoding=None, + ) + assert result.choices[0].message.content == "Hello from Snowflake!" + assert result.model.startswith("snowflake/") + + def test_model_prefixed_with_snowflake(self): + raw = _make_openai_response() + result = self.cfg.transform_response( + model="snowflake/llama3.1-70b", + raw_response=raw, + model_response=ModelResponse(), + logging_obj=_mock_logging(), + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=None, + ) + assert result.model.startswith("snowflake/") + + +# ─── SnowflakeConfig ──────────────────────────────────────── + +class TestAnthropicConfigURL: + def setup_method(self): + self.cfg = SnowflakeConfig() + + def test_url_routes_to_messages_endpoint(self): + url = self.cfg.get_complete_url( + api_base=API_BASE, + api_key=PAT_TOKEN, + model="snowflake/claude-sonnet-4-5", + optional_params={}, + litellm_params={}, + ) + assert url.endswith("/api/v2/cortex/v1/messages") + assert "chat/completions" not in url + assert "inference:complete" not in url + + def test_url_with_account_id(self): + url = self.cfg.get_complete_url( + api_base=None, + api_key=PAT_TOKEN, + model="snowflake/claude-sonnet-4-5", + optional_params={"account_id": ACCOUNT_ID}, + litellm_params={}, + ) + assert f"https://{ACCOUNT_ID}.snowflakecomputing.com/api/v2/cortex/v1/messages" == url + + +class TestAnthropicConfigAuth: + def setup_method(self): + self.cfg = SnowflakeConfig() + + def test_anthropic_version_header_set(self): + headers = self.cfg.validate_environment( + headers={}, + model="snowflake/claude-sonnet-4-5", + messages=[], + optional_params={}, + litellm_params={}, + api_key=PAT_TOKEN, + ) + assert headers["anthropic-version"] == "2023-06-01" + + def test_pat_auth_and_anthropic_version_combined(self): + headers = self.cfg.validate_environment( + headers={}, + model="snowflake/claude-sonnet-4-5", + messages=[], + optional_params={}, + litellm_params={}, + api_key=PAT_TOKEN, + ) + assert headers["X-Snowflake-Authorization-Token-Type"] == "PROGRAMMATIC_ACCESS_TOKEN" + assert headers["anthropic-version"] == "2023-06-01" + assert "Bearer" in headers["Authorization"] + + +class TestAnthropicConfigRequest: + def setup_method(self): + self.cfg = SnowflakeConfig() + + def test_system_message_extracted_to_top_level(self): + messages = [ + {"role": "system", "content": "You are helpful."}, + {"role": "user", "content": "Hello"}, + ] + body = self.cfg.transform_request( + model="snowflake/claude-sonnet-4-5", + messages=messages, + optional_params={}, + litellm_params={}, + headers={}, + ) + assert body["system"] == "You are helpful." + assert all(m["role"] != "system" for m in body["messages"]) + assert body["messages"][0] == {"role": "user", "content": "Hello"} + + def test_model_prefix_stripped(self): + body = self.cfg.transform_request( + model="snowflake/claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + headers={}, + ) + assert body["model"] == "claude-sonnet-4-5" + assert "snowflake/" not in body["model"] + + def test_max_tokens_defaulted_when_missing(self): + body = self.cfg.transform_request( + model="snowflake/claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + headers={}, + ) + assert "max_tokens" in body + assert body["max_tokens"] == 4096 + + def test_max_tokens_not_overridden_when_provided(self): + body = self.cfg.transform_request( + model="snowflake/claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + optional_params={"max_tokens": 500}, + litellm_params={}, + headers={}, + ) + assert body["max_tokens"] == 500 + + def test_no_system_key_when_no_system_message(self): + body = self.cfg.transform_request( + model="snowflake/claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + headers={}, + ) + assert "system" not in body + + +class TestAnthropicConfigResponse: + def setup_method(self): + self.cfg = SnowflakeConfig() + + def test_anthropic_response_to_openai_format(self): + raw = _make_anthropic_response("Hi there!") + result = self.cfg.transform_response( + model="snowflake/claude-sonnet-4-5", + raw_response=raw, + model_response=ModelResponse(), + logging_obj=_mock_logging(), + request_data={}, + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + encoding=None, + ) + assert result.choices[0].message.content == "Hi there!" + assert result.choices[0].finish_reason == "stop" + + def test_usage_tokens_mapped(self): + raw = _make_anthropic_response() + result = self.cfg.transform_response( + model="snowflake/claude-sonnet-4-5", + raw_response=raw, + model_response=ModelResponse(), + logging_obj=_mock_logging(), + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=None, + ) + assert result.usage.prompt_tokens == 10 + assert result.usage.completion_tokens == 5 + assert result.usage.total_tokens == 15 + + def test_stop_reason_end_turn_maps_to_stop(self): + raw = _make_anthropic_response() + result = self.cfg.transform_response( + model="snowflake/claude-sonnet-4-5", + raw_response=raw, + model_response=ModelResponse(), + logging_obj=_mock_logging(), + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=None, + ) + assert result.choices[0].finish_reason == "stop" + + def test_tool_use_block_mapped_to_tool_calls(self): + body = { + "id": "msg_tool", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5", + "content": [ + { + "type": "tool_use", + "id": "toolu_01", + "name": "get_weather", + "input": {"city": "Paris"}, + } + ], + "stop_reason": "tool_use", + "usage": {"input_tokens": 20, "output_tokens": 10}, + } + raw = httpx.Response(200, json=body) + result = self.cfg.transform_response( + model="snowflake/claude-sonnet-4-5", + raw_response=raw, + model_response=ModelResponse(), + logging_obj=_mock_logging(), + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=None, + ) + assert result.choices[0].finish_reason == "tool_calls" + tool_calls = result.choices[0].message.tool_calls + assert len(tool_calls) == 1 + assert tool_calls[0].function.name == "get_weather" + assert json.loads(tool_calls[0].function.arguments) == {"city": "Paris"} + + +# ─── Model detection helper ──────────────────────────────────────────────── + +class TestIsClaudeModel: + def test_claude_model_detected(self): + assert _is_claude_model("snowflake/claude-sonnet-4-5") is True + assert _is_claude_model("claude-3-haiku") is True + assert _is_claude_model("snowflake/claude-opus-4") is True + + def test_non_claude_not_detected(self): + assert _is_claude_model("snowflake/llama3.1-70b") is False + assert _is_claude_model("snowflake/mistral-large") is False + assert _is_claude_model("snowflake/deepseek-r1") is False + assert _is_claude_model("snowflake/snowflake-arctic") is False + + +# ─── Anthropic Tool Transformation Tests ────────────────────────────────── + +class TestAnthropicToolTransformation: + def setup_method(self): + self.cfg = SnowflakeConfig() + + def test_openai_tools_converted_to_anthropic_format(self): + messages = [{"role": "user", "content": "What's the weather?"}] + tools = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get current weather", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, + } + ] + body = self.cfg.transform_request( + model="snowflake/claude-sonnet-4-5", + messages=messages, + optional_params={"tools": tools}, + litellm_params={}, + headers={}, + ) + assert len(body["tools"]) == 1 + tool = body["tools"][0] + assert tool["name"] == "get_weather" + assert tool["description"] == "Get current weather" + assert "input_schema" in tool + assert tool["input_schema"]["properties"]["city"]["type"] == "string" + assert "function" not in tool + assert "type" not in tool + + def test_tools_already_in_anthropic_format_pass_through(self): + messages = [{"role": "user", "content": "hi"}] + tools = [{"name": "my_tool", "input_schema": {"type": "object", "properties": {}}}] + body = self.cfg.transform_request( + model="snowflake/claude-sonnet-4-5", + messages=messages, + optional_params={"tools": tools}, + litellm_params={}, + headers={}, + ) + assert body["tools"] == tools + + +class TestAnthropicMultiTurnToolMessages: + def setup_method(self): + self.cfg = SnowflakeConfig() + + def test_assistant_tool_calls_converted_to_tool_use_blocks(self): + messages = [ + {"role": "user", "content": "What's the weather in Paris?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_123", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city": "Paris"}', + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_123", + "content": "Sunny, 22°C", + }, + {"role": "user", "content": "Thanks!"}, + ] + body = self.cfg.transform_request( + model="snowflake/claude-sonnet-4-5", + messages=messages, + optional_params={}, + litellm_params={}, + headers={}, + ) + msgs = body["messages"] + assert msgs[0] == {"role": "user", "content": "What's the weather in Paris?"} + + assistant_msg = msgs[1] + assert assistant_msg["role"] == "assistant" + assert isinstance(assistant_msg["content"], list) + assert assistant_msg["content"][0]["type"] == "tool_use" + assert assistant_msg["content"][0]["id"] == "call_123" + assert assistant_msg["content"][0]["name"] == "get_weather" + assert assistant_msg["content"][0]["input"] == {"city": "Paris"} + + tool_result_msg = msgs[2] + assert tool_result_msg["role"] == "user" + assert tool_result_msg["content"][0]["type"] == "tool_result" + assert tool_result_msg["content"][0]["tool_use_id"] == "call_123" + assert tool_result_msg["content"][0]["content"] == "Sunny, 22°C" + + assert msgs[3] == {"role": "user", "content": "Thanks!"} + + def test_assistant_with_text_and_tool_calls(self): + messages = [ + {"role": "user", "content": "Check weather"}, + { + "role": "assistant", + "content": "Let me check that for you.", + "tool_calls": [ + { + "id": "call_456", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city": "London"}', + }, + } + ], + }, + ] + body = self.cfg.transform_request( + model="snowflake/claude-sonnet-4-5", + messages=messages, + optional_params={}, + litellm_params={}, + headers={}, + ) + assistant_msg = body["messages"][1] + assert assistant_msg["content"][0] == {"type": "text", "text": "Let me check that for you."} + assert assistant_msg["content"][1]["type"] == "tool_use" + assert assistant_msg["content"][1]["name"] == "get_weather" + + def test_tool_role_never_in_output(self): + messages = [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": None, + "tool_calls": [{"id": "c1", "type": "function", "function": {"name": "f", "arguments": "{}"}}], + }, + {"role": "tool", "tool_call_id": "c1", "content": "result"}, + ] + body = self.cfg.transform_request( + model="snowflake/claude-sonnet-4-5", + messages=messages, + optional_params={}, + litellm_params={}, + headers={}, + ) + for msg in body["messages"]: + assert msg["role"] != "tool" + + def test_malformed_json_in_tool_arguments_handled_gracefully(self): + messages = [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_bad", + "type": "function", + "function": {"name": "broken_tool", "arguments": "not valid json{{{"}, + } + ], + }, + ] + body = self.cfg.transform_request( + model="snowflake/claude-sonnet-4-5", + messages=messages, + optional_params={}, + litellm_params={}, + headers={}, + ) + assistant_msg = body["messages"][1] + tool_use_block = assistant_msg["content"][0] + assert tool_use_block["type"] == "tool_use" + assert tool_use_block["name"] == "broken_tool" + assert tool_use_block["input"] == {} + + def test_non_string_tool_arguments_pass_through(self): + messages = [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_dict", + "type": "function", + "function": {"name": "dict_tool", "arguments": {"already": "parsed"}}, + } + ], + }, + ] + body = self.cfg.transform_request( + model="snowflake/claude-sonnet-4-5", + messages=messages, + optional_params={}, + litellm_params={}, + headers={}, + ) + tool_use_block = body["messages"][1]["content"][0] + assert tool_use_block["input"] == {"already": "parsed"} + + def test_tool_result_with_non_string_content(self): + messages = [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": None, + "tool_calls": [{"id": "c1", "type": "function", "function": {"name": "f", "arguments": "{}"}}], + }, + {"role": "tool", "tool_call_id": "c1", "content": {"result_key": "result_value"}}, + ] + body = self.cfg.transform_request( + model="snowflake/claude-sonnet-4-5", + messages=messages, + optional_params={}, + litellm_params={}, + headers={}, + ) + tool_result = body["messages"][2]["content"][0] + assert tool_result["type"] == "tool_result" + assert json.loads(tool_result["content"]) == {"result_key": "result_value"} diff --git a/tests/test_litellm/llms/voyage/test_voyage_multimodal_embedding.py b/tests/test_litellm/llms/voyage/test_voyage_multimodal_embedding.py new file mode 100644 index 00000000000..f283e7fe0df --- /dev/null +++ b/tests/test_litellm/llms/voyage/test_voyage_multimodal_embedding.py @@ -0,0 +1,306 @@ +import json +from unittest.mock import MagicMock + +import pytest + + +class TestVoyageMultimodalEmbeddings: + def test_multimodal_model_detection(self): + from litellm.llms.voyage.embedding.transformation_multimodal import ( + VoyageMultimodalEmbeddingConfig, + ) + + assert VoyageMultimodalEmbeddingConfig.is_multimodal_embeddings( + "voyage-multimodal-3.5" + ) + assert VoyageMultimodalEmbeddingConfig.is_multimodal_embeddings( + "voyage-multimodal-3" + ) + assert not VoyageMultimodalEmbeddingConfig.is_multimodal_embeddings("voyage-4") + + def test_multimodal_embedding_url_generation(self): + from litellm.llms.voyage.embedding.transformation_multimodal import ( + VoyageMultimodalEmbeddingConfig, + ) + + config = VoyageMultimodalEmbeddingConfig() + assert ( + config.get_complete_url(None, None, "voyage-multimodal-3.5", {}, {}) + == "https://api.voyageai.com/v1/multimodalembeddings" + ) + assert ( + config.get_complete_url( + "https://custom.api.com", None, "voyage-multimodal-3.5", {}, {} + ) + == "https://custom.api.com/multimodalembeddings" + ) + assert ( + config.get_complete_url( + "https://custom.api.com/multimodalembeddings", + None, + "voyage-multimodal-3.5", + {}, + {}, + ) + == "https://custom.api.com/multimodalembeddings" + ) + + def test_multimodal_embedding_request_transformation(self): + from litellm.llms.voyage.embedding.transformation_multimodal import ( + VoyageMultimodalEmbeddingConfig, + ) + + config = VoyageMultimodalEmbeddingConfig() + data_uri = "data:image/png;base64,AAAA" + request = config.transform_embedding_request( + "voyage-multimodal-3.5", + [ + { + "content": [ + {"type": "text", "text": "Describe this"}, + {"type": "image_url", "image_url": {"url": data_uri}}, + {"type": "image_url", "image_url": "https://example.com/a.png"}, + ] + } + ], + {"input_type": "document", "output_dimension": 512}, + {}, + ) + + assert request["model"] == "voyage-multimodal-3.5" + assert "inputs" in request + assert "input" not in request + assert request["input_type"] == "document" + assert request["output_dimension"] == 512 + assert request["inputs"][0]["content"][1] == { + "type": "image_base64", + "image_base64": "AAAA", + } + assert request["inputs"][0]["content"][2] == { + "type": "image_url", + "image_url": "https://example.com/a.png", + } + + def test_multimodal_embedding_string_input_transformation(self): + from litellm.llms.voyage.embedding.transformation_multimodal import ( + VoyageMultimodalEmbeddingConfig, + ) + + config = VoyageMultimodalEmbeddingConfig() + request = config.transform_embedding_request( + "voyage-multimodal-3.5", "hello", {}, {} + ) + assert request["inputs"] == [ + {"content": [{"type": "text", "text": "hello"}]} + ] + + def test_multimodal_embedding_response_transformation(self): + from litellm.llms.voyage.embedding.transformation_multimodal import ( + VoyageMultimodalEmbeddingConfig, + ) + from litellm.types.utils import EmbeddingResponse + + config = VoyageMultimodalEmbeddingConfig() + response_payload = { + "object": "list", + "data": [ + {"object": "embedding", "embedding": [0.1, 0.2], "index": 0} + ], + "model": "voyage-multimodal-3.5", + "usage": { + "text_tokens": 2, + "image_pixels": 0, + "video_pixels": 0, + "total_tokens": 2, + }, + } + raw_response = MagicMock() + raw_response.json.return_value = response_payload + raw_response.status_code = 200 + raw_response.text = json.dumps(response_payload) + + model_response = EmbeddingResponse() + transformed = config.transform_embedding_response( + "voyage-multimodal-3.5", raw_response, model_response, MagicMock() + ) + + assert transformed.model == "voyage-multimodal-3.5" + assert transformed.object == "list" + assert transformed.data == response_payload["data"] + assert transformed.usage.prompt_tokens == 2 + assert transformed.usage.total_tokens == 2 + + def test_provider_config_manager_routes_multimodal_models(self): + import litellm + from litellm.llms.voyage.embedding.transformation_multimodal import ( + VoyageMultimodalEmbeddingConfig, + ) + from litellm.utils import ProviderConfigManager + + config = ProviderConfigManager.get_provider_embedding_config( + model="voyage-multimodal-3.5", provider=litellm.LlmProviders.VOYAGE + ) + + assert isinstance(config, VoyageMultimodalEmbeddingConfig) + + def test_map_openai_params_dimensions(self): + from litellm.llms.voyage.embedding.transformation_multimodal import ( + VoyageMultimodalEmbeddingConfig, + ) + + config = VoyageMultimodalEmbeddingConfig() + assert config.get_supported_openai_params("voyage-multimodal-3.5") == [ + "dimensions" + ] + optional_params = config.map_openai_params( + {"dimensions": 512}, {}, "voyage-multimodal-3.5", False + ) + assert optional_params == {"output_dimension": 512} + assert ( + config.map_openai_params({}, {}, "voyage-multimodal-3.5", False) == {} + ) + + def test_validate_environment_uses_api_key(self): + from litellm.llms.voyage.embedding.transformation_multimodal import ( + VoyageMultimodalEmbeddingConfig, + ) + + config = VoyageMultimodalEmbeddingConfig() + headers = config.validate_environment( + {}, "voyage-multimodal-3.5", [], {}, {}, api_key="test-key" + ) + assert headers == {"Authorization": "Bearer test-key"} + + def test_validate_environment_uses_secret_fallback(self, monkeypatch): + import litellm.llms.voyage.embedding.transformation_multimodal as module + from litellm.llms.voyage.embedding.transformation_multimodal import ( + VoyageMultimodalEmbeddingConfig, + ) + + def fake_get_secret(name): + return "secret-key" if name == "VOYAGE_AI_API_KEY" else None + + monkeypatch.setattr(module, "get_secret_str", fake_get_secret) + config = VoyageMultimodalEmbeddingConfig() + headers = config.validate_environment( + {}, "voyage-multimodal-3.5", [], {}, {}, api_key=None + ) + assert headers == {"Authorization": "Bearer secret-key"} + + def test_validate_environment_raises_without_api_key(self, monkeypatch): + import litellm.llms.voyage.embedding.transformation_multimodal as module + from litellm.llms.voyage.embedding.transformation_multimodal import ( + VoyageMultimodalEmbeddingConfig, + ) + + monkeypatch.setattr(module, "get_secret_str", lambda name: None) + config = VoyageMultimodalEmbeddingConfig() + with pytest.raises(ValueError) as exc_info: + config.validate_environment( + {}, "voyage-multimodal-3.5", [], {}, {}, api_key=None + ) + assert "VOYAGE_API_KEY" in str(exc_info.value) + + def test_normalize_image_url_dict_missing_url_raises(self): + from litellm.llms.voyage.embedding.transformation_multimodal import ( + VoyageMultimodalEmbeddingConfig, + ) + + config = VoyageMultimodalEmbeddingConfig() + with pytest.raises(ValueError) as exc_info: + config._normalize_content_item({"type": "image_url", "image_url": {}}) + assert "image_url" in str(exc_info.value) + + def test_is_multimodal_embeddings_helper(self): + from litellm.llms.voyage.embedding.transformation_multimodal import ( + VoyageMultimodalEmbeddingConfig, + ) + + assert VoyageMultimodalEmbeddingConfig.is_multimodal_embeddings( + "voyage-multimodal-3" + ) + assert VoyageMultimodalEmbeddingConfig.is_multimodal_embeddings( + "VOYAGE-MULTIMODAL-3.5" + ) + assert not VoyageMultimodalEmbeddingConfig.is_multimodal_embeddings( + "voyage-3.5" + ) + + def test_utils_routing_via_provider_config_and_dimensions(self): + import litellm + from litellm.llms.voyage.embedding.transformation_multimodal import ( + VoyageMultimodalEmbeddingConfig, + ) + from litellm.utils import ( + ProviderConfigManager, + get_optional_params_embeddings, + ) + + config = ProviderConfigManager.get_provider_embedding_config( + model="voyage-multimodal-3.5", provider=litellm.LlmProviders.VOYAGE + ) + assert isinstance(config, VoyageMultimodalEmbeddingConfig) + + optional_params = get_optional_params_embeddings( + model="voyage-multimodal-3.5", + dimensions=1024, + custom_llm_provider="voyage", + drop_params=True, + ) + assert optional_params.get("output_dimension") == 1024 + + def test_get_supported_openai_params_voyage_routes_multimodal(self): + from litellm.litellm_core_utils.get_supported_openai_params import ( + get_supported_openai_params, + ) + + multimodal_params = get_supported_openai_params( + model="voyage-multimodal-3.5", + custom_llm_provider="voyage", + request_type="embeddings", + ) + assert multimodal_params == ["dimensions"] + + standard_params = get_supported_openai_params( + model="voyage-3.5", + custom_llm_provider="voyage", + request_type="embeddings", + ) + assert "dimensions" in standard_params + assert "encoding_format" in standard_params + + def test_passthrough_non_content_input(self): + from litellm.llms.voyage.embedding.transformation_multimodal import ( + VoyageMultimodalEmbeddingConfig, + ) + + config = VoyageMultimodalEmbeddingConfig() + request = config.transform_embedding_request( + "voyage-multimodal-3.5", [{"foo": "bar"}], {}, {} + ) + assert request["inputs"] == [{"foo": "bar"}] + + def test_error_response_transformation_and_error_class(self): + from litellm.llms.voyage.embedding.transformation_multimodal import ( + VoyageMultimodalEmbeddingConfig, + VoyageMultimodalEmbeddingError, + ) + from litellm.types.utils import EmbeddingResponse + + config = VoyageMultimodalEmbeddingConfig() + raw_response = MagicMock() + raw_response.json.side_effect = ValueError("not json") + raw_response.status_code = 400 + raw_response.text = "bad request" + + with pytest.raises(VoyageMultimodalEmbeddingError) as exc_info: + config.transform_embedding_response( + "voyage-multimodal-3.5", raw_response, EmbeddingResponse(), MagicMock() + ) + assert exc_info.value.status_code == 400 + assert exc_info.value.message == "bad request" + + error = config.get_error_class("rate limited", 429, {"x-test": "1"}) + assert isinstance(error, VoyageMultimodalEmbeddingError) + assert error.status_code == 429 + assert error.message == "rate limited" diff --git a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py index d31cfdc39bd..a04ad5598df 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -1221,6 +1221,138 @@ async def test_health_endpoint_filters_model_list_by_user_access(): }, f"health_endpoint did not scope model_list to caller access: {returned_names}" +@pytest.mark.asyncio +async def test_health_endpoint_keeps_full_model_list_for_all_proxy_models(): + """ + A key granted all model permissions carries the literal + "all-proxy-models" entry in user_api_key_dict.models. It matches no real + model_name, so the access filter must be skipped entirely; otherwise the + model list filters down to nothing and /health reports 0/0 counts. + """ + from litellm.proxy._types import SpecialModelNames, UserAPIKeyAuth + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + full_model_list = [ + { + "model_name": "model-a", + "litellm_params": {"model": "openai/gpt-4o"}, + "model_info": {"id": "id-a"}, + }, + { + "model_name": "model-b", + "litellm_params": {"model": "openai/gpt-4o"}, + "model_info": {"id": "id-b"}, + }, + ] + + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-test-key", + models=[SpecialModelNames.all_proxy_models.value], + ) + + captured: dict = {} + + async def fake_perform(**kwargs): + captured["model_list"] = kwargs["model_list"] + return { + "healthy_endpoints": [], + "unhealthy_endpoints": [], + "healthy_count": 0, + "unhealthy_count": 0, + } + + with ( + patch("litellm.proxy.proxy_server.llm_model_list", full_model_list), + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.proxy.proxy_server.use_background_health_checks", False), + patch("litellm.proxy.proxy_server.user_model", None), + patch("litellm.proxy.proxy_server.health_check_results", {}), + patch("litellm.proxy.proxy_server.health_check_details", True), + patch("litellm.proxy.proxy_server.health_check_concurrency", 1), + patch( + "litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save", + side_effect=fake_perform, + ), + ): + from fastapi import Response + + await health_endpoint(response=Response(), user_api_key_dict=user_api_key_dict) + + returned_names = {m["model_name"] for m in captured["model_list"]} + assert returned_names == { + "model-a", + "model-b", + }, f"all-proxy-models key should health-check every model: {returned_names}" + + +@pytest.mark.asyncio +async def test_health_endpoint_resolves_all_team_models_to_team_allowlist(): + """ + A key granted "all-team-models" carries the literal sentinel in + user_api_key_dict.models, which matches no real model_name. With a + team_id the sentinel must resolve to the team's allowlist (same + semantics as get_key_models); otherwise the filter would zero out the + model list just like the all-proxy-models case. + """ + from litellm.proxy._types import SpecialModelNames, UserAPIKeyAuth + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + full_model_list = [ + { + "model_name": "model-a", + "litellm_params": {"model": "openai/gpt-4o"}, + "model_info": {"id": "id-a"}, + }, + { + "model_name": "model-b", + "litellm_params": {"model": "openai/gpt-4o"}, + "model_info": {"id": "id-b"}, + }, + ] + + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-test-key", + models=[SpecialModelNames.all_team_models.value], + team_id="team-1", + team_models=["model-b"], + ) + + captured: dict = {} + + async def fake_perform(**kwargs): + captured["model_list"] = kwargs["model_list"] + return { + "healthy_endpoints": [], + "unhealthy_endpoints": [], + "healthy_count": 0, + "unhealthy_count": 0, + } + + with ( + patch("litellm.proxy.proxy_server.llm_model_list", full_model_list), + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.proxy.proxy_server.use_background_health_checks", False), + patch("litellm.proxy.proxy_server.user_model", None), + patch("litellm.proxy.proxy_server.health_check_results", {}), + patch("litellm.proxy.proxy_server.health_check_details", True), + patch("litellm.proxy.proxy_server.health_check_concurrency", 1), + patch( + "litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save", + side_effect=fake_perform, + ), + ): + from fastapi import Response + + await health_endpoint(response=Response(), user_api_key_dict=user_api_key_dict) + + returned_names = {m["model_name"] for m in captured["model_list"]} + assert returned_names == { + "model-b" + }, f"all-team-models key should health-check the team's models: {returned_names}" + + @pytest.mark.asyncio async def test_health_endpoint_filters_background_cache_by_user_access(): """ diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 046971d033b..ed04b9e30dd 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -6555,14 +6555,20 @@ async def test_reset_key_spend_success(monkeypatch): patch( "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object" ) as mock_delete_cache, - patch( - "litellm.proxy.proxy_server._invalidate_spend_counter" - ) as mock_invalidate, ): mock_hash_token.return_value = hashed_key mock_check_admin.return_value = None mock_delete_cache.return_value = None + # Mock spend_counter_cache to verify direct cache set instead of + # _invalidate_spend_counter (removed in favour of atomic cache write). + mock_spend_counter_cache = MagicMock() + mock_spend_counter_cache.redis_cache = None + monkeypatch.setattr( + "litellm.proxy.proxy_server.spend_counter_cache", + mock_spend_counter_cache, + ) + user_api_key_dict = UserAPIKeyAuth( user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", @@ -6582,7 +6588,9 @@ async def test_reset_key_spend_success(monkeypatch): assert response["max_budget"] == 200.0 mock_prisma_client.db.litellm_verificationtoken.update.assert_called_once() mock_delete_cache.assert_awaited_once() - mock_invalidate.assert_awaited_once_with(counter_key=f"spend:key:{hashed_key}") + mock_spend_counter_cache.in_memory_cache.set_cache.assert_called_once_with( + key=f"spend:key:{hashed_key}", value=50.0, ttl=60 + ) @pytest.mark.asyncio @@ -11853,83 +11861,83 @@ async def test_ghsa_q775_default_team_id_does_not_grant_session_token_exemption( assert str(code) == "400" assert "cannot exceed" in msg.lower() - - -@pytest.mark.asyncio -async def test_prepare_key_update_data_budget_duration_null_clears_fields(): - """ - When budget_duration is explicitly set to null, prepare_key_update_data - should produce budget_duration=None and budget_reset_at=None so Prisma - clears them in the DB. - """ - existing_key = LiteLLM_VerificationToken( - token="test-token", - key_alias="test-key", - models=[], - user_id="test-user", - team_id=None, - metadata={}, - ) - - update_request = UpdateKeyRequest(key="test-token", budget_duration=None) - - result = await prepare_key_update_data( - data=update_request, existing_key_row=existing_key - ) - - assert "budget_duration" in result - assert result["budget_duration"] is None - assert "budget_reset_at" in result - assert result["budget_reset_at"] is None - - -@pytest.mark.asyncio -async def test_prepare_key_update_data_budget_duration_not_sent_excluded(): - """ - When budget_duration is NOT sent in the request (unset), it should not - appear in the result dict at all — the existing DB value stays unchanged. - """ - existing_key = LiteLLM_VerificationToken( - token="test-token", - key_alias="test-key", - models=[], - user_id="test-user", - team_id=None, - metadata={}, - ) - - update_request = UpdateKeyRequest(key="test-token", models=["gpt-4"]) - - result = await prepare_key_update_data( - data=update_request, existing_key_row=existing_key - ) - - assert "budget_duration" not in result - assert "budget_reset_at" not in result - - -@pytest.mark.asyncio -async def test_prepare_key_update_data_budget_duration_valid_sets_reset(): - """ - When budget_duration is set to a valid duration string, both - budget_duration and budget_reset_at should be populated. - """ - existing_key = LiteLLM_VerificationToken( - token="test-token", - key_alias="test-key", - models=[], - user_id="test-user", - team_id=None, - metadata={}, - ) - - update_request = UpdateKeyRequest(key="test-token", budget_duration="30d") - - result = await prepare_key_update_data( - data=update_request, existing_key_row=existing_key - ) - - assert result["budget_duration"] == "30d" - assert result["budget_reset_at"] is not None - - + + +@pytest.mark.asyncio +async def test_prepare_key_update_data_budget_duration_null_clears_fields(): + """ + When budget_duration is explicitly set to null, prepare_key_update_data + should produce budget_duration=None and budget_reset_at=None so Prisma + clears them in the DB. + """ + existing_key = LiteLLM_VerificationToken( + token="test-token", + key_alias="test-key", + models=[], + user_id="test-user", + team_id=None, + metadata={}, + ) + + update_request = UpdateKeyRequest(key="test-token", budget_duration=None) + + result = await prepare_key_update_data( + data=update_request, existing_key_row=existing_key + ) + + assert "budget_duration" in result + assert result["budget_duration"] is None + assert "budget_reset_at" in result + assert result["budget_reset_at"] is None + + +@pytest.mark.asyncio +async def test_prepare_key_update_data_budget_duration_not_sent_excluded(): + """ + When budget_duration is NOT sent in the request (unset), it should not + appear in the result dict at all — the existing DB value stays unchanged. + """ + existing_key = LiteLLM_VerificationToken( + token="test-token", + key_alias="test-key", + models=[], + user_id="test-user", + team_id=None, + metadata={}, + ) + + update_request = UpdateKeyRequest(key="test-token", models=["gpt-4"]) + + result = await prepare_key_update_data( + data=update_request, existing_key_row=existing_key + ) + + assert "budget_duration" not in result + assert "budget_reset_at" not in result + + +@pytest.mark.asyncio +async def test_prepare_key_update_data_budget_duration_valid_sets_reset(): + """ + When budget_duration is set to a valid duration string, both + budget_duration and budget_reset_at should be populated. + """ + existing_key = LiteLLM_VerificationToken( + token="test-token", + key_alias="test-key", + models=[], + user_id="test-user", + team_id=None, + metadata={}, + ) + + update_request = UpdateKeyRequest(key="test-token", budget_duration="30d") + + result = await prepare_key_update_data( + data=update_request, existing_key_row=existing_key + ) + + assert result["budget_duration"] == "30d" + assert result["budget_reset_at"] is not None + + diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index d4bc3841668..f0198320f22 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -1609,7 +1609,8 @@ async def test_team_model_add_delete_refresh_team_cache(endpoint_name): patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_logging, patch( - "litellm.proxy.management_endpoints.team_endpoints._cache_team_object" + "litellm.proxy.management_endpoints.team_endpoints._cache_team_object", + new_callable=AsyncMock, ) as mock_cache_team, ): mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( @@ -1618,7 +1619,7 @@ async def test_team_model_add_delete_refresh_team_cache(endpoint_name): mock_prisma_client.db.litellm_teamtable.update = AsyncMock( return_value=updated_team ) - mock_cache_team.return_value = None + mock_prisma_client.db.execute_raw = AsyncMock(return_value=None) if endpoint_name == "team_model_add": await team_model_add( diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_model_alias_merge.py b/tests/test_litellm/proxy/management_endpoints/test_team_model_alias_merge.py new file mode 100644 index 00000000000..45405ba78d6 --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/test_team_model_alias_merge.py @@ -0,0 +1,83 @@ +""" +Tests for atomic team model operations during BYOK model creation. + +Regression tests for https://github.com/BerriAI/litellm/issues/22594 +Concurrent BYOK model creates must not overwrite each other's entries +in team.models. +""" + +import os +import sys +from unittest.mock import AsyncMock, MagicMock + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.proxy._types import ( + LitellmUserRoles, + TeamModelAddRequest, + UserAPIKeyAuth, +) + + +class TestTeamModelAddAtomicAppend: + """Verify team_model_add uses atomic SQL for the models array append.""" + + @pytest.mark.asyncio + async def test_uses_atomic_array_append_with_dedup(self): + """team_model_add must call execute_raw with DISTINCT unnest SQL.""" + from unittest.mock import patch + + from litellm.proxy.management_endpoints.team_endpoints import team_model_add + + mock_request = MagicMock() + mock_user = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test_user" + ) + + existing_team = MagicMock() + existing_team.model_dump.return_value = { + "team_id": "team-1", + "models": ["existing-model"], + } + + updated_team = MagicMock() + updated_team.team_id = "team-1" + updated_team.model_dump.return_value = { + "team_id": "team-1", + "models": ["existing-model", "new-model"], + } + + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, + patch( + "litellm.proxy.management_endpoints.team_endpoints._cache_team_object", + new_callable=AsyncMock, + ), + patch("litellm.proxy.proxy_server.user_api_key_cache"), + patch("litellm.proxy.proxy_server.proxy_logging_obj"), + ): + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock( + return_value=existing_team + ) + mock_prisma.db.execute_raw = AsyncMock(return_value=None) + mock_prisma.db.litellm_teamtable.update = AsyncMock( + return_value=updated_team + ) + + await team_model_add( + data=TeamModelAddRequest(team_id="team-1", models=["new-model"]), + http_request=mock_request, + user_api_key_dict=mock_user, + ) + + mock_prisma.db.execute_raw.assert_called_once() + sql = mock_prisma.db.execute_raw.call_args[0][0] + assert "DISTINCT unnest" in sql + assert "all-proxy-models" in sql + assert mock_prisma.db.execute_raw.call_args[0][1] == ["new-model"] + assert mock_prisma.db.execute_raw.call_args[0][2] == "team-1" + + # Should use write-routed update to re-fetch, not find_unique + mock_prisma.db.litellm_teamtable.update.assert_called_once() diff --git a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py index 0b733401b59..1bc761df5c5 100644 --- a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py +++ b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py @@ -9,7 +9,6 @@ Pins covered: - ``initialize`` - ``load_from_azure_key_vault`` - ``cost_tracking`` -- ``check_request_disconnection`` - ``_resolve_typed_dict_type`` - ``_resolve_pydantic_type`` - ``get_litellm_model_info`` @@ -26,7 +25,7 @@ from typing import List, Optional, Union from unittest.mock import AsyncMock, MagicMock, patch import pytest -from fastapi import FastAPI, HTTPException +from fastapi import FastAPI from pydantic import BaseModel from typing_extensions import TypedDict @@ -35,7 +34,6 @@ from litellm.proxy.proxy_server import ( _initialize_shared_aiohttp_session, _resolve_pydantic_type, _resolve_typed_dict_type, - check_request_disconnection, cleanup_router_config_variables, cost_tracking, get_litellm_model_info, @@ -324,62 +322,6 @@ def test_cost_tracking_no_op_when_prisma_missing(monkeypatch): assert litellm._async_success_callback == [] -# --------------------------------------------------------------------------- -# check_request_disconnection -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_check_request_disconnection_cancels_task_and_raises_499(monkeypatch): - monkeypatch.setattr(ps.asyncio, "sleep", AsyncMock(return_value=None)) - - request = MagicMock() - request.is_disconnected = AsyncMock(return_value=True) - task = MagicMock() - - raised_status = None - try: - await check_request_disconnection(request=request, llm_api_call_task=task) - except HTTPException as exc: - raised_status = exc.status_code - - observed = { - "raised_status": raised_status, - "cancel_called": task.cancel.called, - "is_async": inspect.iscoroutinefunction(check_request_disconnection), - } - assert normalize(observed) == { - "raised_status": 499, - "cancel_called": True, - "is_async": True, - } - - -@pytest.mark.asyncio -async def test_check_request_disconnection_invalid_when_connected_times_out(monkeypatch): - """With a connected request the function loops for up to 10 minutes — - wrap in wait_for and assert it times out. Patch ``asyncio.sleep`` so the - loop spins without real wall-clock waits.""" - import litellm.proxy.proxy_server as ps - - request = MagicMock() - request.is_disconnected = AsyncMock(return_value=False) - task = MagicMock() - - _real_sleep = asyncio.sleep - - async def _instant_sleep(_seconds): - await _real_sleep(0) - - monkeypatch.setattr(ps.asyncio, "sleep", _instant_sleep) - - with pytest.raises(asyncio.TimeoutError): - await asyncio.wait_for( - check_request_disconnection(request=request, llm_api_call_task=task), - timeout=0.05, - ) - - # --------------------------------------------------------------------------- # _resolve_typed_dict_type # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index 677d358428d..592232f45f5 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -980,7 +980,7 @@ async def test_ProxyConfig__update_llm_router_bad_proxy_logging_raises(monkeypat # Passing None for proxy_logging_obj triggers AttributeError in _add_general_settings_from_db_config # when it calls proxy_logging_obj.update_values. with pytest.raises(AttributeError): - await pc._update_llm_router(new_models=None, proxy_logging_obj=None) # type: ignore[arg-type] + await pc._update_llm_router(new_models=[], proxy_logging_obj=None) # type: ignore[arg-type] # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 0b3a31d2de4..ec186ffa795 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -1,6 +1,7 @@ +import asyncio import copy import datetime -from typing import AsyncGenerator +from typing import AsyncGenerator, Optional from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -15,6 +16,8 @@ from litellm.integrations.opentelemetry import UserAPIKeyAuth from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, ProxyConfig, + _await_llm_call_cancelling_on_disconnect, + _cancel_llm_call_on_client_disconnect, _extract_error_from_sse_chunk, _get_cost_breakdown_from_logging_obj, _has_attribute_error_in_chain, @@ -2412,6 +2415,77 @@ class TestHandleLLMApiExceptionDictDetail: assert proxy_exc.code == "500" +class TestHandleLLMApiExceptionRetryAfter: + """RouterRateLimitError cooldown_time must surface as a retry-after header.""" + + async def _invoke(self, exc: Exception, callback_headers: Optional[dict] = None): + from litellm.proxy._types import ProxyException, UserAPIKeyAuth + + processor = ProxyBaseLLMRequestProcessing(data={}) + user_api_key_dict = UserAPIKeyAuth(api_key="sk-test") + proxy_logging_obj = MagicMock() + proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock( + return_value=callback_headers or {} + ) + + try: + await processor._handle_llm_api_exception( + e=exc, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + ) + except ProxyException as raised: + return raised + raise AssertionError("ProxyException was not raised") + + async def test_handle_llm_api_exception_sets_retry_after_from_cooldown_time(self): + from litellm.types.router import RouterRateLimitError + + exc = RouterRateLimitError( + model="gpt-4", + cooldown_time=42.3, + enable_pre_call_checks=False, + cooldown_list=[], + ) + proxy_exc = await self._invoke(exc) + assert proxy_exc.headers["retry-after"] == "43" + assert proxy_exc.code == "429" + + async def test_handle_llm_api_exception_skips_retry_after_when_cooldown_is_zero( + self, + ): + from litellm.types.router import RouterRateLimitError + + exc = RouterRateLimitError( + model="gpt-4", + cooldown_time=0, + enable_pre_call_checks=False, + cooldown_list=[], + ) + proxy_exc = await self._invoke(exc) + assert "retry-after" not in proxy_exc.headers + + async def test_handle_llm_api_exception_no_retry_after_for_plain_exception(self): + proxy_exc = await self._invoke(ValueError("some other failure")) + assert "retry-after" not in proxy_exc.headers + + async def test_handle_llm_api_exception_retry_after_survives_callback_headers(self): + from litellm.types.router import RouterRateLimitError + + exc = RouterRateLimitError( + model="gpt-4", + cooldown_time=42.3, + enable_pre_call_checks=False, + cooldown_list=[], + ) + proxy_exc = await self._invoke( + exc, callback_headers={"retry-after": "", "x-custom": "1"} + ) + assert proxy_exc.headers["retry-after"] == "43" + assert proxy_exc.headers["x-custom"] == "1" + + class TestAsyncStreamingDataGeneratorFastPath: """Fast/slow path branching in async_streaming_data_generator.""" @@ -2482,6 +2556,197 @@ class TestAsyncStreamingDataGeneratorFastPath: ProxyLogging._callback_capabilities_cache.clear() +class TestCancelOnDisconnect: + """ + Coverage for the opt-in `general_settings.cancel_on_disconnect` flag: + cancelling the in-flight upstream LLM call when the HTTP client disconnects + (issue #13774), without changing the default code path and without skipping + failure accounting (post_call_failure_hook) on the resulting 499. + """ + + def _request(self, messages: list) -> Request: + async def receive(): + if messages: + return messages.pop(0) + await asyncio.Event().wait() + + return Request(scope={"type": "http", "headers": []}, receive=receive) + + async def test_monitor_cancels_llm_call_and_sets_event_on_disconnect(self): + request = self._request( + [ + {"type": "http.request", "body": b"", "more_body": False}, + {"type": "http.disconnect"}, + ] + ) + llm_call = asyncio.get_running_loop().create_future() + disconnect_event = asyncio.Event() + + await _cancel_llm_call_on_client_disconnect( + request, llm_call, disconnect_event + ) + + assert llm_call.cancelled() + assert disconnect_event.is_set() + + async def test_monitor_is_noop_while_client_stays_connected(self): + request = self._request( + [{"type": "http.request", "body": b"", "more_body": False}] + ) + llm_call = asyncio.get_running_loop().create_future() + disconnect_event = asyncio.Event() + + monitor = asyncio.create_task( + _cancel_llm_call_on_client_disconnect(request, llm_call, disconnect_event) + ) + await asyncio.sleep(0.01) + + assert not monitor.done() + assert not llm_call.cancelled() + assert not disconnect_event.is_set() + monitor.cancel() + + async def test_monitor_survives_receive_failure_without_cancelling(self): + """If request.receive() fails (e.g. transport reset) the watcher must + degrade to a no-op instead of crashing or cancelling the LLM call.""" + + async def receive(): + raise RuntimeError("transport reset") + + request = Request(scope={"type": "http", "headers": []}, receive=receive) + llm_call = asyncio.get_running_loop().create_future() + disconnect_event = asyncio.Event() + + await _cancel_llm_call_on_client_disconnect( + request, llm_call, disconnect_event + ) + + assert not llm_call.cancelled() + assert not disconnect_event.is_set() + + async def test_cancellation_without_disconnect_reraises_cancelled_error(self): + """A CancelledError that is NOT client-initiated (e.g. server shutdown) + must propagate as-is instead of being masked as a 499.""" + request = self._request([]) + llm_call = asyncio.get_running_loop().create_future() + llm_call.cancel() + + with pytest.raises(asyncio.CancelledError): + await _await_llm_call_cancelling_on_disconnect(request, llm_call) + + async def _drive_base_process_llm_request( + self, monkeypatch, general_settings: dict, llm_call, request: Request + ): + from litellm.proxy._types import UserAPIKeyAuth + + logging_obj = MagicMock() + logging_obj.litellm_call_id = "test-cancel-on-disconnect" + logging_obj._defer_async_logging = False + logging_obj._on_deferred_stream_complete = None + logging_obj.cost_breakdown = None + + processor = ProxyBaseLLMRequestProcessing( + data={"model": "fake-model", "litellm_logging_obj": logging_obj} + ) + + proxy_logging_obj = MagicMock(spec=ProxyLogging) + proxy_logging_obj.during_call_hook = AsyncMock(return_value=None) + proxy_logging_obj.update_request_status = AsyncMock(return_value=None) + proxy_logging_obj.post_call_success_hook = AsyncMock( + side_effect=lambda data, user_api_key_dict, response: response + ) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock( + return_value=None + ) + + async def fake_route_request(**kwargs): + return llm_call() + + monkeypatch.setattr( + litellm.proxy.common_request_processing, + "route_request", + fake_route_request, + ) + + return await processor.base_process_llm_request( + request=request, + fastapi_response=Response(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + route_type="acompletion", + proxy_logging_obj=proxy_logging_obj, + general_settings=general_settings, + proxy_config=MagicMock(spec=ProxyConfig), + skip_pre_call_logic=True, + ) + + async def test_disconnect_ignored_when_flag_disabled(self, monkeypatch): + upstream_cancelled = asyncio.Event() + model_response = litellm.ModelResponse() + + async def llm_call(): + try: + await asyncio.sleep(0.05) + return model_response + except asyncio.CancelledError: + upstream_cancelled.set() + raise + + result = await self._drive_base_process_llm_request( + monkeypatch, + general_settings={}, + llm_call=llm_call, + request=self._request([{"type": "http.disconnect"}]), + ) + + assert result is model_response + assert not upstream_cancelled.is_set() + + async def test_disconnect_cancels_upstream_when_flag_enabled(self, monkeypatch): + upstream_cancelled = asyncio.Event() + + async def llm_call(): + try: + await asyncio.sleep(5) + return litellm.ModelResponse() + except asyncio.CancelledError: + upstream_cancelled.set() + raise + + with pytest.raises(HTTPException) as exc_info: + await self._drive_base_process_llm_request( + monkeypatch, + general_settings={"cancel_on_disconnect": True}, + llm_call=llm_call, + request=self._request([{"type": "http.disconnect"}]), + ) + + assert exc_info.value.status_code == 499 + assert upstream_cancelled.is_set() + + async def test_499_still_fires_post_call_failure_hook(self): + """Regression guard: the 499 path must NOT bypass post_call_failure_hook, + which releases max_parallel_requests slots and fires spend/alerting + callbacks (cf. #14457; P1 review finding on #25776/#27146).""" + from litellm.proxy._types import ProxyException, UserAPIKeyAuth + + processor = ProxyBaseLLMRequestProcessing(data={}) + proxy_logging_obj = MagicMock() + proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={}) + + with pytest.raises(ProxyException) as exc_info: + await processor._handle_llm_api_exception( + e=HTTPException( + status_code=499, detail="Client disconnected the request" + ), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + proxy_logging_obj=proxy_logging_obj, + ) + + assert exc_info.value.code == "499" + proxy_logging_obj.post_call_failure_hook.assert_awaited_once() + + class TestAllmPassthroughRoutePostCallGuardrails: """ Regression: non-streaming allm_passthrough_route responses are httpx.Response objects. diff --git a/tests/test_litellm/proxy/test_component_allowlists.py b/tests/test_litellm/proxy/test_component_allowlists.py index ad25856b972..926ce3bee66 100644 --- a/tests/test_litellm/proxy/test_component_allowlists.py +++ b/tests/test_litellm/proxy/test_component_allowlists.py @@ -42,7 +42,11 @@ _REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", if _REPO_ROOT not in sys.path: sys.path.insert(0, _REPO_ROOT) -from backend.routes.allowlist import BACKEND_EXACT_PATHS, BACKEND_PATH_PREFIXES +from backend.routes.allowlist import ( + BACKEND_EXACT_PATHS, + BACKEND_MOUNT_PATHS, + BACKEND_PATH_PREFIXES, +) from gateway.routes.allowlist import GATEWAY_EXACT_PATHS, GATEWAY_PATH_PREFIXES from litellm.proxy.proxy_server import app @@ -88,3 +92,44 @@ def test_gateway_plus_backend_covers_full_app(): f"Update gateway/routes/allowlist.py or backend/routes/allowlist.py to cover:\n " + "\n ".join(sorted(uncovered)) ) + + +def test_backend_mount_paths_defined(): + """BACKEND_MOUNT_PATHS constant must exist and be a frozenset.""" + assert isinstance(BACKEND_MOUNT_PATHS, frozenset), \ + f"BACKEND_MOUNT_PATHS must be a frozenset, got {type(BACKEND_MOUNT_PATHS)}" + assert len(BACKEND_MOUNT_PATHS) > 0, \ + "BACKEND_MOUNT_PATHS must contain at least one Mount path" + + +def test_swagger_mount_in_backend_allowlist(): + """The /swagger Mount must be in BACKEND_MOUNT_PATHS.""" + assert "/swagger" in BACKEND_MOUNT_PATHS, \ + "/swagger Mount path must be in BACKEND_MOUNT_PATHS" + + +def test_backend_keeps_swagger_mount(): + """Verify that Mounts in BACKEND_MOUNT_PATHS are kept on the backend.""" + backend_mounts = { + getattr(r, "path") + for r in app.router.routes + if isinstance(r, Mount) and getattr(r, "path", None) in BACKEND_MOUNT_PATHS + } + assert "/swagger" in backend_mounts, \ + "/swagger Mount is expected on the proxy app and should be in BACKEND_MOUNT_PATHS" + + +def test_backend_drops_non_allowlisted_mounts(): + """Verify that Mounts NOT in BACKEND_MOUNT_PATHS would be dropped from backend.""" + all_mounts = { + getattr(r, "path") + for r in app.router.routes + if isinstance(r, Mount) and getattr(r, "path", None) is not None + } + non_backend_mounts = all_mounts - BACKEND_MOUNT_PATHS + + assert len(non_backend_mounts) > 0, \ + "Expected at least one non-backend Mount (e.g., /ui, /_next) to verify filtering logic" + for mount_path in non_backend_mounts: + assert mount_path not in BACKEND_MOUNT_PATHS, \ + f"Mount {mount_path} should not be in BACKEND_MOUNT_PATHS" diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index f336c632546..09cc7a51caf 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -4603,3 +4603,65 @@ def test_apply_overrides_provider_prefix_in_model_skips_router_lookup( assert data["api_base"] == "https://hotel-eastus.openai.azure.com/" assert data["api_key"] == "key-hotel-eastus" router.get_deployment_by_model_group_name.assert_not_called() + + +def _make_request_mock(path: str, headers: dict) -> MagicMock: + request_mock = MagicMock(spec=Request) + request_mock.url = MagicMock() + request_mock.url.path = path + request_mock.url.__str__.return_value = f"http://localhost{path}" + request_mock.method = "POST" + request_mock.query_params = {} + request_mock.headers = headers + request_mock.client = MagicMock() + request_mock.client.host = "127.0.0.1" + return request_mock + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "user_agent, request_drop_params, operator_drop_params, expected_drop_params", + [ + ("claude-cli/2.0.69 (external, cli)", None, None, True), + ("claude-cli/1.0.44 (external, sdk-py)", None, None, True), + ("claude-cli/2.0.69 (external, cli)", False, None, False), + ("claude-cli/2.0.69 (external, cli)", None, False, None), + ("claude-cli/2.0.69 (external, cli)", None, True, None), + ("PostmanRuntime/7.53.0", None, None, None), + (None, None, None, None), + ], +) +async def test_add_litellm_data_to_request_claude_code_drop_params( + user_agent, request_drop_params, operator_drop_params, expected_drop_params +): + """Claude Code sends Anthropic-specific params that fail on non-Anthropic + providers, so its user agent must turn on drop_params automatically, + without overriding an explicit caller value, an explicit operator-level + litellm_settings value, or affecting other clients. + """ + headers = {"Content-Type": "application/json"} + if user_agent is not None: + headers["user-agent"] = user_agent + request_mock = _make_request_mock("/v1/messages", headers) + + data = {"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}]} + if request_drop_params is not None: + data["drop_params"] = request_drop_params + + proxy_config = MagicMock() + proxy_config.config = ( + {"litellm_settings": {"drop_params": operator_drop_params}} + if operator_drop_params is not None + else {"litellm_settings": {}} + ) + + updated = await add_litellm_data_to_request( + data=data, + request=request_mock, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=proxy_config, + general_settings={}, + version="test-version", + ) + + assert updated.get("drop_params") == expected_drop_params diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 9eaccdfcbcd..baf1f145612 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -1928,23 +1928,6 @@ async def test_delete_deployment_type_mismatch(): # Create mock ProxyConfig instance pc = ProxyConfig() - pc.get_config = MagicMock( - return_value={ - "model_list": [ - { - "model_name": "openai-gpt-4o", - "litellm_params": {"model": "gpt-4o"}, - "model_info": {"id": 12345678}, - }, - { - "model_name": "openai-gpt-4o", - "litellm_params": {"model": "gpt-4o"}, - "model_info": {"id": 12345679}, - }, - ] - } - ) - # Mock llm_router with string IDs (this is the source of the type mismatch) mock_llm_router = MagicMock() mock_llm_router.get_model_ids.return_value = [ @@ -1963,11 +1946,23 @@ async def test_delete_deployment_type_mismatch(): mock_llm_router.delete_deployment = MagicMock(side_effect=mock_delete_deployment) - # Mock get_config to return empty config (no config models) async def mock_get_config(config_file_path): - return {} + return { + "model_list": [ + { + "model_name": "openai-gpt-4o", + "litellm_params": {"model": "gpt-4o"}, + "model_info": {"id": 12345678}, + }, + { + "model_name": "openai-gpt-4o", + "litellm_params": {"model": "gpt-4o"}, + "model_info": {"id": 12345679}, + }, + ] + } - pc.get_config = MagicMock(side_effect=mock_get_config) + pc.get_config = AsyncMock(side_effect=mock_get_config) # Patch the global llm_router with ( @@ -1977,20 +1972,29 @@ async def test_delete_deployment_type_mismatch(): # Call the function under test deleted_count = await pc._delete_deployment(db_models=[]) - # Assertions: Models 12345678 and 12345679 should NOT be deleted - # because they exist in combined_id_list (as integers) even though - # router has them as strings + # The two SHA-hash models have no corresponding entry in combined_id_list + # and must be evicted. + assert ( + deleted_count == 2 + ), f"Expected 2 deletions (SHA-hash models), got {deleted_count}" + assert ( + "a96e12e76b36a57cfae57a41288eb41567629cac89b4828c6f7074afc3534695" + in deleted_ids + ) + assert ( + "a40186dd0fdb9b7282380277d7f57044d29de95bfbfcd7f4322b3493702d5cd3" + in deleted_ids + ) - # The function should delete the other 2 models that are not in combined_id_list - assert deleted_count == 0, f"Expected 0 deletions, got {deleted_count}" - - # Verify that 12345678 and 12345679 were NOT deleted - assert ( - "12345678" not in deleted_ids - ), f"Model 12345678 should NOT be deleted. Deleted IDs: {deleted_ids}" - assert ( - "12345679" not in deleted_ids - ), f"Model 12345679 should NOT be deleted. Deleted IDs: {deleted_ids}" + # Models 12345678 and 12345679 exist in the config (as integers); str() + # conversion in _delete_deployment makes them match the router's string IDs, + # so they must NOT be evicted. + assert ( + "12345678" not in deleted_ids + ), f"Model 12345678 should NOT be deleted. Deleted IDs: {deleted_ids}" + assert ( + "12345679" not in deleted_ids + ), f"Model 12345679 should NOT be deleted. Deleted IDs: {deleted_ids}" @pytest.mark.asyncio @@ -7937,3 +7941,106 @@ class TestSortModelsByDisplayName: all_models=models, sort_by="model_name", sort_order="asc" ) assert [m["model_name"] for m in sorted_models] == ["alpha", "beta"] + + +class TestDeleteDeploymentSync: + @pytest.mark.asyncio + async def test_delete_deployment_evicts_model_when_all_db_models_deleted(self): + """ + Regression test for #28443. + When all DB models are deleted, _delete_deployment must evict them from + the router. The old code returned 0 early when db_models was empty. + """ + from unittest.mock import AsyncMock, MagicMock, patch + + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + mock_router = MagicMock() + mock_router.get_model_ids.return_value = ["model-id-to-evict"] + mock_router.delete_deployment.return_value = MagicMock() + + with patch("litellm.proxy.proxy_server.llm_router", mock_router): + with patch.object( + proxy_config, "get_config", AsyncMock(return_value={"model_list": []}) + ): + count = await proxy_config._delete_deployment(db_models=[]) + + mock_router.delete_deployment.assert_called_once_with(id="model-id-to-evict") + assert count == 1 + + @pytest.mark.asyncio + async def test_update_llm_router_skips_update_on_db_fetch_failure(self): + """ + When _get_models_from_db returns None (transient DB failure), _update_llm_router + must return early without touching the router. + """ + from unittest.mock import AsyncMock, MagicMock, patch + + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + mock_router = MagicMock() + + with patch("litellm.proxy.proxy_server.llm_router", mock_router): + with patch.object(proxy_config, "get_config", AsyncMock(return_value={})): + await proxy_config._update_llm_router( + new_models=None, proxy_logging_obj=MagicMock() + ) + + mock_router.delete_deployment.assert_not_called() + mock_router.upsert_deployment.assert_not_called() + + @pytest.mark.asyncio + async def test_get_models_from_db_returns_none_on_exception(self): + """ + _get_models_from_db must return None (not []) when the DB raises an exception, + so callers can distinguish a transient failure from a genuinely empty DB. + """ + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + mock_prisma = MagicMock() + mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock( + side_effect=Exception("DB connection lost") + ) + + result = await proxy_config._get_models_from_db(prisma_client=mock_prisma) + + assert ( + result is None + ), f"Expected None on DB failure to signal fetch error, got {result!r}" + + +def test_get_config_list_includes_cancel_on_disconnect(monkeypatch): + """Follow-up to #30223: the flag must be discoverable via /config/list, + which requires both the ConfigGeneralSettings field and the allowed_args + entry in get_config_list; missing either silently hides it from the UI.""" + import types + from unittest.mock import AsyncMock, MagicMock + + from fastapi.testclient import TestClient + + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.proxy_server import app + + mock_prisma = MagicMock() + mock_config_table = MagicMock() + mock_config_table.find_first = AsyncMock(return_value=None) + mock_prisma.db = types.SimpleNamespace(litellm_config=mock_config_table) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN + ) + try: + client = TestClient(app) + resp = client.get("/config/list", params={"config_type": "general_settings"}) + assert resp.status_code == 200, resp.text + fields = {item["field_name"]: item for item in resp.json()} + assert "cancel_on_disconnect" in fields + assert fields["cancel_on_disconnect"]["field_type"] == "Boolean" + finally: + app.dependency_overrides.clear() diff --git a/tests/test_litellm/test_anthropic_beta_headers_filtering.py b/tests/test_litellm/test_anthropic_beta_headers_filtering.py index 9cd27e88c33..59bab22de74 100644 --- a/tests/test_litellm/test_anthropic_beta_headers_filtering.py +++ b/tests/test_litellm/test_anthropic_beta_headers_filtering.py @@ -412,6 +412,20 @@ class TestAnthropicBetaHeadersFiltering: assert filtered == ["compact-2026-01-12"] + @pytest.mark.parametrize("provider", ["bedrock_converse", "bedrock"]) + def test_fine_grained_tool_streaming_forwarded_for_bedrock(self, provider): + """Bedrock honors fine-grained-tool-streaming-2025-05-14 via + additionalModelRequestFields.anthropic_beta. Stripping it (previously + mapped to null) silently re-enables Anthropic's server-side buffering of + tool-call argument deltas, so streamed tool args arrive in a single + end-of-stream burst instead of incrementally.""" + filtered = filter_and_transform_beta_headers( + beta_headers=["fine-grained-tool-streaming-2025-05-14"], + provider=provider, + ) + + assert filtered == ["fine-grained-tool-streaming-2025-05-14"] + def test_null_value_headers_filtered(self): """Test that headers with null values are always filtered out.""" for provider in [ diff --git a/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.test.tsx b/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.test.tsx index 4590121acf2..c2b730bf46a 100644 --- a/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.test.tsx @@ -1,5 +1,5 @@ import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; -import { render, screen, waitFor } from "@testing-library/react"; +import { fireEvent, render, screen, waitFor } from "@testing-library/react"; import { Form } from "antd"; import { beforeAll, describe, expect, it, vi } from "vitest"; import { Providers } from "../provider_info_helpers"; @@ -215,4 +215,134 @@ describe("ProviderSpecificFields", () => { expect(baseModelInput).toBeInTheDocument(); }); }); + + it("sets Azure API version from the API base query parameter", async () => { + const queryClient = createQueryClient(); + render( + +
+ + +
, + ); + + const apiBaseInput = await screen.findByPlaceholderText("https://..."); + const apiVersionInput = await screen.findByPlaceholderText("2023-07-01-preview"); + + fireEvent.change(apiBaseInput, { + target: { + value: + "https://test-resource.openai.azure.com/openai/deployments/gpt-4/chat/completions?api_version=2024-10-21", + }, + }); + + await waitFor(() => { + expect(apiVersionInput).toHaveValue("2024-10-21"); + }); + }); + + it("sets Azure API version from the hyphenated API base query parameter", async () => { + const queryClient = createQueryClient(); + render( + +
+ + +
, + ); + + const apiBaseInput = await screen.findByPlaceholderText("https://..."); + const apiVersionInput = await screen.findByPlaceholderText("2023-07-01-preview"); + + fireEvent.change(apiBaseInput, { + target: { + value: + "https://test-resource.openai.azure.com/openai/deployments/gpt-4/chat/completions?api-version=2024-10-21", + }, + }); + + await waitFor(() => { + expect(apiVersionInput).toHaveValue("2024-10-21"); + }); + }); + + it("clears an inferred Azure API version when the API base has no version parameter", async () => { + const queryClient = createQueryClient(); + render( + +
+ + +
, + ); + + const apiBaseInput = await screen.findByPlaceholderText("https://..."); + const apiVersionInput = await screen.findByPlaceholderText("2023-07-01-preview"); + + fireEvent.change(apiBaseInput, { + target: { + value: + "https://test-resource.openai.azure.com/openai/deployments/gpt-4/chat/completions?api-version=2024-10-21", + }, + }); + + await waitFor(() => { + expect(apiVersionInput).toHaveValue("2024-10-21"); + }); + + fireEvent.change(apiBaseInput, { + target: { + value: "https://test-resource.openai.azure.com/openai/deployments/gpt-4/chat/completions", + }, + }); + + await waitFor(() => { + expect(apiVersionInput).toHaveValue(""); + }); + }); + + it("preserves a manually edited Azure API version when the API base has no version parameter", async () => { + const queryClient = createQueryClient(); + render( + +
+ + +
, + ); + + const apiBaseInput = await screen.findByPlaceholderText("https://..."); + const apiVersionInput = await screen.findByPlaceholderText("2023-07-01-preview"); + + fireEvent.change(apiBaseInput, { + target: { + value: + "https://test-resource.openai.azure.com/openai/deployments/gpt-4/chat/completions?api-version=2024-10-21", + }, + }); + + await waitFor(() => { + expect(apiVersionInput).toHaveValue("2024-10-21"); + }); + + fireEvent.change(apiVersionInput, { + target: { + value: "2025-01-01-preview", + }, + }); + + await waitFor(() => { + expect(apiVersionInput).toHaveValue("2025-01-01-preview"); + }); + + fireEvent.change(apiBaseInput, { + target: { + value: "https://test-resource.openai.azure.com/openai/deployments/gpt-4/chat/completions", + }, + }); + + await waitFor(() => { + expect(apiVersionInput).toHaveValue("2025-01-01-preview"); + }); + }); }); diff --git a/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx b/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx index 24df0ac21ef..045a9b0c1b6 100644 --- a/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx +++ b/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx @@ -28,6 +28,18 @@ export interface CredentialValues { value: string; } +const getApiVersionFromApiBase = (apiBase: string): string | null => { + const queryStartIndex = apiBase.indexOf("?"); + if (queryStartIndex === -1) { + return null; + } + + const queryString = apiBase.slice(queryStartIndex + 1).split("#")[0]; + const searchParams = new URLSearchParams(queryString); + + return searchParams.get("api_version") || searchParams.get("api-version"); +}; + const mapFieldMetadataToUiField = (field: ProviderCredentialFieldMetadata): ProviderCredentialField => { const type: ProviderCredentialField["type"] = field.field_type === "password" @@ -167,6 +179,30 @@ const ProviderSpecificFields: React.FC = ({ selecte return mapped; }, [selectedProviderEnum, selectedProvider, providerMetadata]); + const hasApiVersionField = React.useMemo(() => allFields.some((field) => field.key === "api_version"), [allFields]); + const lastInferredApiVersionRef = React.useRef(null); + + const handleApiBaseChange = React.useCallback( + (event: React.ChangeEvent) => { + if (!hasApiVersionField) { + return; + } + + const apiVersion = getApiVersionFromApiBase(event.target.value); + if (apiVersion) { + lastInferredApiVersionRef.current = apiVersion; + form.setFieldsValue({ api_version: apiVersion }); + return; + } + + if (form.getFieldValue("api_version") === lastInferredApiVersionRef.current) { + form.setFieldsValue({ api_version: "" }); + } + lastInferredApiVersionRef.current = null; + }, + [form, hasApiVersionField], + ); + const handleUpload = { name: "file", accept: ".json", @@ -261,6 +297,7 @@ const ProviderSpecificFields: React.FC = ({ selecte placeholder={field.placeholder} type={field.type === "password" ? "password" : "text"} defaultValue={field.defaultValue} + onChange={field.key === "api_base" ? handleApiBaseChange : undefined} /> )} diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 2e24e83cefa..d95a918a3b0 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -22062,6 +22062,11 @@ export interface components { * @description run health checks in background */ background_health_checks?: boolean | null; + /** + * Cancel On Disconnect + * @description cancel the in-flight upstream LLM request (non-streaming) when the client disconnects, freeing backend capacity (e.g. a vLLM GPU slot); the request is logged as a 499 failure + */ + cancel_on_disconnect?: boolean | null; /** * Completion Model * @description proxy level default model for all chat completion calls From 2893f9b67b741922a84702040763dcefb8ba2820 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 12 Jun 2026 13:11:54 -0700 Subject: [PATCH 088/185] feat(ui): migrate policies, guardrails, prompts, tool-policies, and skills to path routes (#30263) * feat(ui): cut policies, guardrails, prompts, tool-policies, and skills over to path routes Continues the page-by-page App Router migration. All five legacy switch arms passed only accessToken/userRole, so each route wrapper is a thin useAuthorized() + render. skills keeps a claude-code-plugins alias in MIGRATED_PAGES because the old switch matched both page ids, mirroring the api_ref/api-reference precedent. * refactor(ui): colocate the prompts panel under its route The new route wrapper was its only importer, so the 32-file folder moves wholesale into (dashboard)/prompts/components; tree-escaping relative imports (networking, molecules, common_components) become @/components aliases and the suppressions baseline is re-keyed. policies, guardrails, claude_code_plugins, and ToolPoliciesView stay at src/components: each has consumers on other pages (playground selectors, AI Hub, public model hub), so their shared/page splits go in the colocation follow-up. * fix(ui): move the PromptsPanel file along with its folder @/components/prompts resolved to the prompts.tsx FILE next to the prompts/ folder, not the folder itself; the colocation moved only the folder, so the wrapper's ./components import and the panel's ./prompts/* imports both broke and next build failed. Move the panel in as components/index.tsx and fix its now-escaping relative imports. Caught by next build; tsc --noEmit missed it because incremental mode reused a stale tsbuildinfo. * test(ui): lock skills alias resolution in legacyKeyForPathname Both skills and claude-code-plugins map to the skills segment, and sidebar highlighting depends on first-match-wins returning the sidebar key; assert it so a future reorder of MIGRATED_PAGES cannot silently break highlighting. Mirrors the api_ref/api-reference assertion. Flagged by Greptile. --- .../e2e_tests/fixtures/migratedPages.ts | 8 ++- ui/litellm-dashboard/eslint-suppressions.json | 70 +++++++++---------- .../src/app/(dashboard)/guardrails/page.tsx | 9 +++ .../src/app/(dashboard)/page.tsx | 15 ---- .../src/app/(dashboard)/policies/page.tsx | 9 +++ .../(dashboard)/prompts/components}/README.md | 0 .../prompts/components}/add_prompt_form.tsx | 4 +- .../(dashboard)/prompts/components/index.tsx} | 12 ++-- .../components}/prompt_editor_view.tsx | 0 .../DeveloperMessageCard.tsx | 0 .../prompt_editor_view/DotpromptViewTab.tsx | 0 .../prompt_editor_view/ModelConfigCard.tsx | 2 +- .../prompt_editor_view/PromptCodeSnippets.tsx | 2 +- .../prompt_editor_view/PromptEditorHeader.tsx | 0 .../prompt_editor_view/PromptMessagesCard.tsx | 0 .../prompt_editor_view/PublishModal.tsx | 0 .../prompt_editor_view/ToolsCard.test.tsx | 0 .../prompt_editor_view/ToolsCard.tsx | 0 .../VersionHistorySidePanel.test.tsx | 6 +- .../VersionHistorySidePanel.tsx | 2 +- .../conversation_panel/EmptyState.tsx | 0 .../conversation_panel/MessageBubble.tsx | 0 .../conversation_panel/MessageInput.tsx | 0 .../conversation_panel/MessageList.tsx | 0 .../conversation_panel/VariableInput.tsx | 0 .../conversation_panel/VariableWarning.tsx | 0 .../conversation_panel/index.tsx | 0 .../conversation_panel/types.ts | 0 .../conversation_panel/useConversation.ts | 4 +- .../components}/prompt_editor_view/index.tsx | 4 +- .../components}/prompt_editor_view/types.ts | 0 .../prompt_editor_view/utils.test.ts | 0 .../components}/prompt_editor_view/utils.ts | 0 .../prompts/components}/prompt_info.tsx | 2 +- .../prompts/components}/prompt_table.tsx | 0 .../prompts/components}/prompt_utils.tsx | 0 .../prompts/components}/tool_modal.tsx | 0 .../prompts/components}/variable_textarea.tsx | 0 .../src/app/(dashboard)/prompts/page.tsx | 9 +++ .../src/app/(dashboard)/skills/page.tsx | 9 +++ .../app/(dashboard)/tool-policies/page.tsx | 9 +++ .../src/utils/migratedPages.test.ts | 16 +++++ .../src/utils/migratedPages.ts | 7 ++ 43 files changed, 128 insertions(+), 71 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/guardrails/page.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/policies/page.tsx rename ui/litellm-dashboard/src/{components/prompts => app/(dashboard)/prompts/components}/README.md (100%) rename ui/litellm-dashboard/src/{components/prompts => app/(dashboard)/prompts/components}/add_prompt_form.tsx (96%) rename ui/litellm-dashboard/src/{components/prompts.tsx => app/(dashboard)/prompts/components/index.tsx} (94%) rename ui/litellm-dashboard/src/{components/prompts => app/(dashboard)/prompts/components}/prompt_editor_view.tsx (100%) rename ui/litellm-dashboard/src/{components/prompts => app/(dashboard)/prompts/components}/prompt_editor_view/DeveloperMessageCard.tsx (100%) rename ui/litellm-dashboard/src/{components/prompts => app/(dashboard)/prompts/components}/prompt_editor_view/DotpromptViewTab.tsx (100%) rename ui/litellm-dashboard/src/{components/prompts => app/(dashboard)/prompts/components}/prompt_editor_view/ModelConfigCard.tsx (97%) rename ui/litellm-dashboard/src/{components/prompts => app/(dashboard)/prompts/components}/prompt_editor_view/PromptCodeSnippets.tsx (99%) rename ui/litellm-dashboard/src/{components/prompts => app/(dashboard)/prompts/components}/prompt_editor_view/PromptEditorHeader.tsx (100%) rename ui/litellm-dashboard/src/{components/prompts => app/(dashboard)/prompts/components}/prompt_editor_view/PromptMessagesCard.tsx (100%) rename ui/litellm-dashboard/src/{components/prompts => app/(dashboard)/prompts/components}/prompt_editor_view/PublishModal.tsx (100%) rename ui/litellm-dashboard/src/{components/prompts => app/(dashboard)/prompts/components}/prompt_editor_view/ToolsCard.test.tsx (100%) rename ui/litellm-dashboard/src/{components/prompts => app/(dashboard)/prompts/components}/prompt_editor_view/ToolsCard.tsx (100%) rename ui/litellm-dashboard/src/{components/prompts => app/(dashboard)/prompts/components}/prompt_editor_view/VersionHistorySidePanel.test.tsx (98%) rename ui/litellm-dashboard/src/{components/prompts => app/(dashboard)/prompts/components}/prompt_editor_view/VersionHistorySidePanel.tsx (98%) rename ui/litellm-dashboard/src/{components/prompts => app/(dashboard)/prompts/components}/prompt_editor_view/conversation_panel/EmptyState.tsx (100%) rename ui/litellm-dashboard/src/{components/prompts => app/(dashboard)/prompts/components}/prompt_editor_view/conversation_panel/MessageBubble.tsx (100%) rename ui/litellm-dashboard/src/{components/prompts => app/(dashboard)/prompts/components}/prompt_editor_view/conversation_panel/MessageInput.tsx (100%) rename ui/litellm-dashboard/src/{components/prompts => app/(dashboard)/prompts/components}/prompt_editor_view/conversation_panel/MessageList.tsx (100%) rename ui/litellm-dashboard/src/{components/prompts => app/(dashboard)/prompts/components}/prompt_editor_view/conversation_panel/VariableInput.tsx (100%) rename ui/litellm-dashboard/src/{components/prompts => app/(dashboard)/prompts/components}/prompt_editor_view/conversation_panel/VariableWarning.tsx (100%) rename ui/litellm-dashboard/src/{components/prompts => app/(dashboard)/prompts/components}/prompt_editor_view/conversation_panel/index.tsx (100%) rename ui/litellm-dashboard/src/{components/prompts => app/(dashboard)/prompts/components}/prompt_editor_view/conversation_panel/types.ts (100%) rename ui/litellm-dashboard/src/{components/prompts => app/(dashboard)/prompts/components}/prompt_editor_view/conversation_panel/useConversation.ts (97%) rename ui/litellm-dashboard/src/{components/prompts => app/(dashboard)/prompts/components}/prompt_editor_view/index.tsx (99%) rename ui/litellm-dashboard/src/{components/prompts => app/(dashboard)/prompts/components}/prompt_editor_view/types.ts (100%) rename ui/litellm-dashboard/src/{components/prompts => app/(dashboard)/prompts/components}/prompt_editor_view/utils.test.ts (100%) rename ui/litellm-dashboard/src/{components/prompts => app/(dashboard)/prompts/components}/prompt_editor_view/utils.ts (100%) rename ui/litellm-dashboard/src/{components/prompts => app/(dashboard)/prompts/components}/prompt_info.tsx (99%) rename ui/litellm-dashboard/src/{components/prompts => app/(dashboard)/prompts/components}/prompt_table.tsx (100%) rename ui/litellm-dashboard/src/{components/prompts => app/(dashboard)/prompts/components}/prompt_utils.tsx (100%) rename ui/litellm-dashboard/src/{components/prompts => app/(dashboard)/prompts/components}/tool_modal.tsx (100%) rename ui/litellm-dashboard/src/{components/prompts => app/(dashboard)/prompts/components}/variable_textarea.tsx (100%) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/prompts/page.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/skills/page.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/tool-policies/page.tsx diff --git a/ui/litellm-dashboard/e2e_tests/fixtures/migratedPages.ts b/ui/litellm-dashboard/e2e_tests/fixtures/migratedPages.ts index a56ce79d8f1..17a27d451df 100644 --- a/ui/litellm-dashboard/e2e_tests/fixtures/migratedPages.ts +++ b/ui/litellm-dashboard/e2e_tests/fixtures/migratedPages.ts @@ -10,8 +10,7 @@ * * Keep this in lockstep with MIGRATED_PAGES in src/utils/migratedPages.ts. * Pending (add as each PR lands): the leaf-pages batch - * (caching, cost-tracking, guardrails, logs, policies, prompts, skills, - * tool-policies, transform-request, ui-theme). + * (caching, cost-tracking, logs, transform-request, ui-theme). */ export const MIGRATED_E2E_PAGES: Record = { api_ref: "api-reference", @@ -26,6 +25,11 @@ export const MIGRATED_E2E_PAGES: Record = { "tag-management": "tag-management", "vector-stores": "vector-stores", memory: "memory", + policies: "policies", + guardrails: "guardrails", + prompts: "prompts", + "tool-policies": "tool-policies", + skills: "skills", }; export const MIGRATED_E2E_SEGMENTS: string[] = [...new Set(Object.values(MIGRATED_E2E_PAGES))]; diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index b838736ba26..53dde4d28a4 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -1774,7 +1774,22 @@ "count": 2 } }, - "src/components/prompts.tsx": { + "src/app/(dashboard)/prompts/components/add_prompt_form.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/app/(dashboard)/prompts/components/prompt_editor_view/DeveloperMessageCard.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/app/(dashboard)/prompts/components/prompt_editor_view/ModelConfigCard.tsx": { + "no-restricted-imports": { + "count": 1 + } + }, + "src/app/(dashboard)/prompts/components/prompt_editor_view/PromptCodeSnippets.tsx": { "no-restricted-imports": { "count": 1 }, @@ -1782,75 +1797,52 @@ "count": 1 } }, - "src/components/prompts/add_prompt_form.tsx": { + "src/app/(dashboard)/prompts/components/prompt_editor_view/PromptEditorHeader.tsx": { "no-restricted-imports": { "count": 1 } }, - "src/components/prompts/prompt_editor_view/DeveloperMessageCard.tsx": { + "src/app/(dashboard)/prompts/components/prompt_editor_view/PromptMessagesCard.tsx": { "no-restricted-imports": { "count": 1 } }, - "src/components/prompts/prompt_editor_view/ModelConfigCard.tsx": { + "src/app/(dashboard)/prompts/components/prompt_editor_view/PublishModal.tsx": { "no-restricted-imports": { "count": 1 } }, - "src/components/prompts/prompt_editor_view/PromptCodeSnippets.tsx": { - "no-restricted-imports": { - "count": 1 - }, - "react-hooks/set-state-in-effect": { - "count": 1 - } - }, - "src/components/prompts/prompt_editor_view/PromptEditorHeader.tsx": { + "src/app/(dashboard)/prompts/components/prompt_editor_view/ToolsCard.tsx": { "no-restricted-imports": { "count": 1 } }, - "src/components/prompts/prompt_editor_view/PromptMessagesCard.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, - "src/components/prompts/prompt_editor_view/PublishModal.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, - "src/components/prompts/prompt_editor_view/ToolsCard.tsx": { - "no-restricted-imports": { - "count": 1 - } - }, - "src/components/prompts/prompt_editor_view/VersionHistorySidePanel.test.tsx": { + "src/app/(dashboard)/prompts/components/prompt_editor_view/VersionHistorySidePanel.test.tsx": { "max-nested-callbacks": { "count": 1 } }, - "src/components/prompts/prompt_editor_view/VersionHistorySidePanel.tsx": { + "src/app/(dashboard)/prompts/components/prompt_editor_view/VersionHistorySidePanel.tsx": { "react-hooks/immutability": { "count": 1 } }, - "src/components/prompts/prompt_editor_view/conversation_panel/MessageInput.tsx": { + "src/app/(dashboard)/prompts/components/prompt_editor_view/conversation_panel/MessageInput.tsx": { "no-restricted-imports": { "count": 1 } }, - "src/components/prompts/prompt_editor_view/conversation_panel/index.tsx": { + "src/app/(dashboard)/prompts/components/prompt_editor_view/conversation_panel/index.tsx": { "no-restricted-imports": { "count": 1 } }, - "src/components/prompts/prompt_editor_view/conversation_panel/useConversation.ts": { + "src/app/(dashboard)/prompts/components/prompt_editor_view/conversation_panel/useConversation.ts": { "no-restricted-syntax": { "count": 1 } }, - "src/components/prompts/prompt_info.tsx": { + "src/app/(dashboard)/prompts/components/prompt_info.tsx": { "no-restricted-imports": { "count": 1 }, @@ -1858,7 +1850,7 @@ "count": 2 } }, - "src/components/prompts/prompt_table.tsx": { + "src/app/(dashboard)/prompts/components/prompt_table.tsx": { "no-restricted-imports": { "count": 1 } @@ -2249,5 +2241,13 @@ "react/display-name": { "count": 1 } + }, + "src/app/(dashboard)/prompts/components/index.tsx": { + "no-restricted-imports": { + "count": 1 + }, + "react-hooks/set-state-in-effect": { + "count": 1 + } } } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/page.tsx new file mode 100644 index 00000000000..4e7fa88f70f --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/page.tsx @@ -0,0 +1,9 @@ +"use client"; + +import GuardrailsPanel from "@/components/guardrails"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; + +export default function Guardrails() { + const { accessToken, userRole } = useAuthorized(); + return ; +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/page.tsx index 9758786331f..9318bc1b332 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/page.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/page.tsx @@ -4,15 +4,12 @@ import ModelsAndEndpointsView from "@/app/(dashboard)/models-and-endpoints/Model import AdminPanel from "@/components/AdminPanel"; import AgentsPanel from "@/components/agents"; import CacheDashboard from "@/components/cache_dashboard"; -import ClaudeCodePluginsPanel from "@/components/claude_code_plugins"; import { teamListCall as v2TeamListCall } from "@/app/(dashboard)/hooks/teams/useTeams"; import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings"; import useProxySettings from "@/app/(dashboard)/hooks/proxySettings/useProxySettings"; import LoadingScreen from "@/components/common_components/LoadingScreen"; import { CostTrackingSettings } from "@/components/CostTrackingSettings"; import GeneralSettings from "@/components/general_settings"; -import GuardrailsPanel from "@/components/guardrails"; -import PoliciesPanel from "@/components/policies"; import { Team } from "@/components/key_team_helpers/key_list"; import ModelHubTable from "@/components/AIHub/ModelHubTable"; import { Organization, proxyBaseUrl, getInProductNudgesCall } from "@/components/networking"; @@ -21,7 +18,6 @@ import OldTeams from "@/components/OldTeams"; import { fetchUserModels, CreateKeyPrefillData } from "@/components/organisms/create_key_button"; import Organizations, { fetchOrganizations } from "@/components/organizations"; import PassThroughSettings from "@/components/pass_through_settings"; -import PromptsPanel from "@/components/prompts"; import PublicModelHub from "@/components/public_model_hub"; import Settings from "@/components/settings"; import { SurveyPrompt, SurveyModal, ClaudeCodePrompt, ClaudeCodeModal } from "@/components/survey"; @@ -29,7 +25,6 @@ import TransformRequestPanel from "@/components/transform_request"; import UIThemeSettings from "@/components/ui_theme_settings"; import Usage from "@/components/usage"; import UserDashboard from "@/components/user_dashboard"; -import ToolPoliciesView from "@/components/ToolPoliciesView"; import SpendLogsTable from "@/components/view_logs"; import ViewUserDashboard from "@/components/view_users"; import { useAuth } from "@/contexts/AuthContext"; @@ -377,14 +372,8 @@ function CreateKeyPageContent() { ) : page == "logging-and-alerts" ? ( - ) : page == "guardrails" ? ( - - ) : page == "policies" ? ( - ) : page == "agents" ? ( - ) : page == "prompts" ? ( - ) : page == "transform-request" ? ( ) : page == "router-settings" ? ( @@ -428,10 +417,6 @@ function CreateKeyPageContent() { accessToken={accessToken} premiumUser={premiumUser} /> - ) : page == "skills" || page == "claude-code-plugins" ? ( - - ) : page == "tool-policies" ? ( - ) : page == "new_usage" ? ( ) : ( diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/page.tsx new file mode 100644 index 00000000000..eb7840d8795 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/page.tsx @@ -0,0 +1,9 @@ +"use client"; + +import PoliciesPanel from "@/components/policies"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; + +export default function Policies() { + const { accessToken, userRole } = useAuthorized(); + return ; +} diff --git a/ui/litellm-dashboard/src/components/prompts/README.md b/ui/litellm-dashboard/src/app/(dashboard)/prompts/components/README.md similarity index 100% rename from ui/litellm-dashboard/src/components/prompts/README.md rename to ui/litellm-dashboard/src/app/(dashboard)/prompts/components/README.md diff --git a/ui/litellm-dashboard/src/components/prompts/add_prompt_form.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/components/add_prompt_form.tsx similarity index 96% rename from ui/litellm-dashboard/src/components/prompts/add_prompt_form.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/prompts/components/add_prompt_form.tsx index cdb77bb66fc..48623bbda60 100644 --- a/ui/litellm-dashboard/src/components/prompts/add_prompt_form.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/components/add_prompt_form.tsx @@ -3,8 +3,8 @@ import { Modal, Form, Select, Upload, Button, Divider } from "antd"; import { TextInput } from "@tremor/react"; import { UploadOutlined } from "@ant-design/icons"; import type { UploadFile, UploadProps } from "antd"; -import { convertPromptFileToJson, createPromptCall } from "../networking"; -import NotificationsManager from "../molecules/notifications_manager"; +import { convertPromptFileToJson, createPromptCall } from "@/components/networking"; +import NotificationsManager from "@/components/molecules/notifications_manager"; const { Option } = Select; diff --git a/ui/litellm-dashboard/src/components/prompts.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/components/index.tsx similarity index 94% rename from ui/litellm-dashboard/src/components/prompts.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/prompts/components/index.tsx index 1e0155a7738..3430d9808d1 100644 --- a/ui/litellm-dashboard/src/components/prompts.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/components/index.tsx @@ -2,12 +2,12 @@ import React, { useState, useEffect } from "react"; import { Button } from "@tremor/react"; import { Modal, Select } from "antd"; -import { getPromptsList, PromptSpec, ListPromptsResponse, deletePromptCall } from "./networking"; -import PromptTable from "./prompts/prompt_table"; -import PromptInfoView from "./prompts/prompt_info"; -import AddPromptForm from "./prompts/add_prompt_form"; -import PromptEditorView from "./prompts/prompt_editor_view"; -import NotificationsManager from "./molecules/notifications_manager"; +import { getPromptsList, PromptSpec, ListPromptsResponse, deletePromptCall } from "@/components/networking"; +import PromptTable from "./prompt_table"; +import PromptInfoView from "./prompt_info"; +import AddPromptForm from "./add_prompt_form"; +import PromptEditorView from "./prompt_editor_view"; +import NotificationsManager from "@/components/molecules/notifications_manager"; import { isAdminRole, isProxyAdminRole } from "@/utils/roles"; interface PromptsProps { diff --git a/ui/litellm-dashboard/src/components/prompts/prompt_editor_view.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/prompts/prompt_editor_view.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view.tsx diff --git a/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/DeveloperMessageCard.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/DeveloperMessageCard.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/prompts/prompt_editor_view/DeveloperMessageCard.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/DeveloperMessageCard.tsx diff --git a/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/DotpromptViewTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/DotpromptViewTab.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/prompts/prompt_editor_view/DotpromptViewTab.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/DotpromptViewTab.tsx diff --git a/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/ModelConfigCard.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/ModelConfigCard.tsx similarity index 97% rename from ui/litellm-dashboard/src/components/prompts/prompt_editor_view/ModelConfigCard.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/ModelConfigCard.tsx index aa564160ddf..66ddb90bea3 100644 --- a/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/ModelConfigCard.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/ModelConfigCard.tsx @@ -2,7 +2,7 @@ import React, { useState } from "react"; import { Text } from "@tremor/react"; import { Input } from "antd"; import { SettingsIcon } from "lucide-react"; -import ModelSelector from "../../common_components/ModelSelector"; +import ModelSelector from "@/components/common_components/ModelSelector"; interface ModelConfigCardProps { model: string; diff --git a/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/PromptCodeSnippets.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/PromptCodeSnippets.tsx similarity index 99% rename from ui/litellm-dashboard/src/components/prompts/prompt_editor_view/PromptCodeSnippets.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/PromptCodeSnippets.tsx index 89b74b88bc4..3d52b3c03e5 100644 --- a/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/PromptCodeSnippets.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/PromptCodeSnippets.tsx @@ -4,7 +4,7 @@ import { CodeOutlined } from "@ant-design/icons"; import { Button as TremorButton, Text } from "@tremor/react"; import { Prism as SyntaxHighlighter } from "react-syntax-highlighter"; import { coy } from "react-syntax-highlighter/dist/esm/styles/prism"; -import NotificationsManager from "../../molecules/notifications_manager"; +import NotificationsManager from "@/components/molecules/notifications_manager"; interface PromptCodeSnippetsProps { promptId: string; diff --git a/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/PromptEditorHeader.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/PromptEditorHeader.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/prompts/prompt_editor_view/PromptEditorHeader.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/PromptEditorHeader.tsx diff --git a/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/PromptMessagesCard.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/PromptMessagesCard.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/prompts/prompt_editor_view/PromptMessagesCard.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/PromptMessagesCard.tsx diff --git a/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/PublishModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/PublishModal.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/prompts/prompt_editor_view/PublishModal.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/PublishModal.tsx diff --git a/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/ToolsCard.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/ToolsCard.test.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/prompts/prompt_editor_view/ToolsCard.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/ToolsCard.test.tsx diff --git a/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/ToolsCard.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/ToolsCard.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/prompts/prompt_editor_view/ToolsCard.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/ToolsCard.tsx diff --git a/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/VersionHistorySidePanel.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/VersionHistorySidePanel.test.tsx similarity index 98% rename from ui/litellm-dashboard/src/components/prompts/prompt_editor_view/VersionHistorySidePanel.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/VersionHistorySidePanel.test.tsx index b0346e03a2d..c76c64b89a6 100644 --- a/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/VersionHistorySidePanel.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/VersionHistorySidePanel.test.tsx @@ -1,11 +1,11 @@ import { describe, it, expect, vi, beforeEach, afterEach, type Mock } from "vitest"; import { render, screen, fireEvent, act, waitFor } from "@testing-library/react"; import VersionHistorySidePanel from "./VersionHistorySidePanel"; -import { getPromptVersions } from "../../networking"; -import type { PromptSpec } from "../../networking"; +import { getPromptVersions } from "@/components/networking"; +import type { PromptSpec } from "@/components/networking"; // Mock the networking function -vi.mock("../../networking", () => ({ +vi.mock("@/components/networking", () => ({ getPromptVersions: vi.fn(), })); diff --git a/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/VersionHistorySidePanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/VersionHistorySidePanel.tsx similarity index 98% rename from ui/litellm-dashboard/src/components/prompts/prompt_editor_view/VersionHistorySidePanel.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/VersionHistorySidePanel.tsx index 851bdb78ad2..97fe70ba3e1 100644 --- a/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/VersionHistorySidePanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/VersionHistorySidePanel.tsx @@ -1,6 +1,6 @@ import { Drawer, List, Skeleton, Tag, Typography } from "antd"; import React, { useEffect, useState } from "react"; -import { getPromptVersions, PromptSpec } from "../../networking"; +import { getPromptVersions, PromptSpec } from "@/components/networking"; const { Text } = Typography; diff --git a/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/conversation_panel/EmptyState.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/conversation_panel/EmptyState.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/prompts/prompt_editor_view/conversation_panel/EmptyState.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/conversation_panel/EmptyState.tsx diff --git a/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/conversation_panel/MessageBubble.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/conversation_panel/MessageBubble.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/prompts/prompt_editor_view/conversation_panel/MessageBubble.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/conversation_panel/MessageBubble.tsx diff --git a/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/conversation_panel/MessageInput.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/conversation_panel/MessageInput.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/prompts/prompt_editor_view/conversation_panel/MessageInput.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/conversation_panel/MessageInput.tsx diff --git a/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/conversation_panel/MessageList.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/conversation_panel/MessageList.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/prompts/prompt_editor_view/conversation_panel/MessageList.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/conversation_panel/MessageList.tsx diff --git a/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/conversation_panel/VariableInput.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/conversation_panel/VariableInput.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/prompts/prompt_editor_view/conversation_panel/VariableInput.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/conversation_panel/VariableInput.tsx diff --git a/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/conversation_panel/VariableWarning.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/conversation_panel/VariableWarning.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/prompts/prompt_editor_view/conversation_panel/VariableWarning.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/conversation_panel/VariableWarning.tsx diff --git a/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/conversation_panel/index.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/conversation_panel/index.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/prompts/prompt_editor_view/conversation_panel/index.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/conversation_panel/index.tsx diff --git a/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/conversation_panel/types.ts b/ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/conversation_panel/types.ts similarity index 100% rename from ui/litellm-dashboard/src/components/prompts/prompt_editor_view/conversation_panel/types.ts rename to ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/conversation_panel/types.ts diff --git a/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/conversation_panel/useConversation.ts b/ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/conversation_panel/useConversation.ts similarity index 97% rename from ui/litellm-dashboard/src/components/prompts/prompt_editor_view/conversation_panel/useConversation.ts rename to ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/conversation_panel/useConversation.ts index e55d8bdeadf..d8632170c69 100644 --- a/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/conversation_panel/useConversation.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/conversation_panel/useConversation.ts @@ -1,9 +1,9 @@ import { useState, useRef, useEffect } from "react"; -import NotificationsManager from "../../../molecules/notifications_manager"; +import NotificationsManager from "@/components/molecules/notifications_manager"; import { TokenUsage } from "@/components/chat_ui/ResponseMetrics"; import { Message } from "./types"; import { convertToDotPrompt, extractVariables } from "../utils"; -import { getProxyBaseUrl, getGlobalLitellmHeaderName } from "../../../networking"; +import { getProxyBaseUrl, getGlobalLitellmHeaderName } from "@/components/networking"; export const useConversation = (prompt: any, accessToken: string | null) => { const [isLoading, setIsLoading] = useState(false); diff --git a/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/index.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/index.tsx similarity index 99% rename from ui/litellm-dashboard/src/components/prompts/prompt_editor_view/index.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/index.tsx index c8c572468f8..046805c15b8 100644 --- a/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/index.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/index.tsx @@ -1,7 +1,7 @@ import React, { useState } from "react"; import ToolModal from "../tool_modal"; -import NotificationsManager from "../../molecules/notifications_manager"; -import { createPromptCall, updatePromptCall, getPromptInfo } from "../../networking"; +import NotificationsManager from "@/components/molecules/notifications_manager"; +import { createPromptCall, updatePromptCall, getPromptInfo } from "@/components/networking"; import { PromptType, PromptEditorViewProps, Tool } from "./types"; import { convertToDotPrompt, parseExistingPrompt } from "./utils"; import PromptEditorHeader from "./PromptEditorHeader"; diff --git a/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/types.ts b/ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/types.ts similarity index 100% rename from ui/litellm-dashboard/src/components/prompts/prompt_editor_view/types.ts rename to ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/types.ts diff --git a/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/utils.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/utils.test.ts similarity index 100% rename from ui/litellm-dashboard/src/components/prompts/prompt_editor_view/utils.test.ts rename to ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/utils.test.ts diff --git a/ui/litellm-dashboard/src/components/prompts/prompt_editor_view/utils.ts b/ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/utils.ts similarity index 100% rename from ui/litellm-dashboard/src/components/prompts/prompt_editor_view/utils.ts rename to ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_editor_view/utils.ts diff --git a/ui/litellm-dashboard/src/components/prompts/prompt_info.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_info.tsx similarity index 99% rename from ui/litellm-dashboard/src/components/prompts/prompt_info.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_info.tsx index a5e76542ebd..f96445c1a20 100644 --- a/ui/litellm-dashboard/src/components/prompts/prompt_info.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_info.tsx @@ -29,7 +29,7 @@ import { } from "@/components/networking"; import { copyToClipboard as utilCopyToClipboard } from "@/utils/dataUtils"; import { CheckIcon, CopyIcon } from "lucide-react"; -import NotificationsManager from "../molecules/notifications_manager"; +import NotificationsManager from "@/components/molecules/notifications_manager"; import PromptCodeSnippets from "./prompt_editor_view/PromptCodeSnippets"; import { extractModel, extractTemplateVariables, getBasePromptId, getCurrentVersion } from "./prompt_utils"; diff --git a/ui/litellm-dashboard/src/components/prompts/prompt_table.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_table.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/prompts/prompt_table.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_table.tsx diff --git a/ui/litellm-dashboard/src/components/prompts/prompt_utils.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_utils.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/prompts/prompt_utils.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/prompts/components/prompt_utils.tsx diff --git a/ui/litellm-dashboard/src/components/prompts/tool_modal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/components/tool_modal.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/prompts/tool_modal.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/prompts/components/tool_modal.tsx diff --git a/ui/litellm-dashboard/src/components/prompts/variable_textarea.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/components/variable_textarea.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/prompts/variable_textarea.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/prompts/components/variable_textarea.tsx diff --git a/ui/litellm-dashboard/src/app/(dashboard)/prompts/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/prompts/page.tsx new file mode 100644 index 00000000000..59c194b0855 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/prompts/page.tsx @@ -0,0 +1,9 @@ +"use client"; + +import PromptsPanel from "./components"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; + +export default function Prompts() { + const { accessToken, userRole } = useAuthorized(); + return ; +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/skills/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/skills/page.tsx new file mode 100644 index 00000000000..bd2b12c73b0 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/skills/page.tsx @@ -0,0 +1,9 @@ +"use client"; + +import ClaudeCodePluginsPanel from "@/components/claude_code_plugins"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; + +export default function Skills() { + const { accessToken, userRole } = useAuthorized(); + return ; +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/tool-policies/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/tool-policies/page.tsx new file mode 100644 index 00000000000..6aaebaab959 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/tool-policies/page.tsx @@ -0,0 +1,9 @@ +"use client"; + +import ToolPoliciesView from "@/components/ToolPoliciesView"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; + +export default function ToolPolicies() { + const { accessToken, userRole } = useAuthorized(); + return ; +} diff --git a/ui/litellm-dashboard/src/utils/migratedPages.test.ts b/ui/litellm-dashboard/src/utils/migratedPages.test.ts index 656b91a6c32..8471c8a2567 100644 --- a/ui/litellm-dashboard/src/utils/migratedPages.test.ts +++ b/ui/litellm-dashboard/src/utils/migratedPages.test.ts @@ -75,6 +75,19 @@ describe("migratedHref / legacyPageHref", () => { expect(MIGRATED_PAGES["vector-stores"]).toBe("vector-stores"); expect(MIGRATED_PAGES.memory).toBe("memory"); }); + + it("maps the policies, guardrails, prompts, tool-policies, and skills ids to their routes", async () => { + vi.doMock("@/components/networking", () => ({ serverRootPath: "/" })); + const { MIGRATED_PAGES } = await import("./migratedPages"); + + expect(MIGRATED_PAGES.policies).toBe("policies"); + expect(MIGRATED_PAGES.guardrails).toBe("guardrails"); + expect(MIGRATED_PAGES.prompts).toBe("prompts"); + expect(MIGRATED_PAGES["tool-policies"]).toBe("tool-policies"); + expect(MIGRATED_PAGES.skills).toBe("skills"); + // Old bookmarks used ?page=claude-code-plugins for the same panel. + expect(MIGRATED_PAGES["claude-code-plugins"]).toBe("skills"); + }); }); describe("dev server (NODE_ENV=development)", () => { @@ -128,6 +141,9 @@ describe("legacyKeyForPathname", () => { // Resolves to the sidebar key api_ref, not the hyphenated alias, so highlighting works. expect(legacyKeyForPathname("/ui/api-reference")).toBe("api_ref"); expect(legacyKeyForPathname("/ui/api-reference/")).toBe("api_ref"); + // Same for skills: the claude-code-plugins alias maps to the same segment, + // and first-match-wins iteration must keep returning the sidebar key. + expect(legacyKeyForPathname("/ui/skills")).toBe("skills"); }); it("returns null for a not-yet-migrated path", async () => { diff --git a/ui/litellm-dashboard/src/utils/migratedPages.ts b/ui/litellm-dashboard/src/utils/migratedPages.ts index 8f8a21d97c6..f51a7af73c2 100644 --- a/ui/litellm-dashboard/src/utils/migratedPages.ts +++ b/ui/litellm-dashboard/src/utils/migratedPages.ts @@ -23,6 +23,13 @@ export const MIGRATED_PAGES: Record = { "tag-management": "tag-management", "vector-stores": "vector-stores", memory: "memory", + policies: "policies", + guardrails: "guardrails", + prompts: "prompts", + "tool-policies": "tool-policies", + skills: "skills", + // Legacy alias: the old switch matched ?page=claude-code-plugins for the same panel. + "claude-code-plugins": "skills", }; function uiBase(): string { From 40301820e7d5df289bf3112929d1d6dacac84f46 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 12 Jun 2026 15:35:15 -0700 Subject: [PATCH 089/185] feat(ui): migrate caching, cost-tracking, transform-request, ui-theme, and logs to path routes (#30267) * feat(ui): cut caching, cost-tracking, transform-request, ui-theme, and logs over to path routes Completes the simple-leaf portion of the page-by-page App Router migration. All five legacy switch arms passed only identity props (accessToken/userRole/userID, plus token/premiumUser for caching and logs), all of which useAuthorized() provides, so each route wrapper is a thin useAuthorized() + render. MIGRATED_PAGES routes the sidebar and redirects the legacy ?page= URLs; the e2e fixture picks all five up in the migration smoke and sidebar specs automatically. * refactor(ui): colocate caching, cost-tracking, transform-request, and ui-theme components Each had the legacy switch as its only importer. caching takes its whole closure (cache_dashboard, cache_health, cache_settings, response_time_indicator); CostTrackingSettings moves as the cost-tracking components folder; the transform-request and ui-theme single-file panels move under their routes. view_logs stays at src/components: six other pages (guardrails monitor, tool policies, pass-through, MCP toolsets, usage) import it. Suppressions re-keyed. * chore: retrigger ci e2e_ui_testing failed on three specs unrelated to this PR's pages (team-info tabs, MCP create form) and local_testing_part1 on test_batch_completions; all pass on the pre-merge commit and none touch files in this diff. --- .../e2e_tests/fixtures/migratedPages.ts | 9 +++- ui/litellm-dashboard/eslint-suppressions.json | 44 +++++++++---------- .../caching}/components/cache_dashboard.tsx | 6 +-- .../caching}/components/cache_health.tsx | 0 .../cache_settings/CacheFieldGroup.test.tsx | 0 .../cache_settings/CacheFieldGroup.tsx | 0 .../CacheFieldRenderer.test.tsx | 0 .../cache_settings/CacheFieldRenderer.tsx | 2 +- .../cache_settings/RedisTypeSelector.test.tsx | 0 .../cache_settings/RedisTypeSelector.tsx | 0 .../cache_settings/cacheSettingsUtils.ts | 0 .../components/cache_settings/index.tsx | 4 +- .../components/response_time_indicator.tsx | 0 .../src/app/(dashboard)/caching/page.tsx | 17 +++++++ .../components}/add_margin_form.test.tsx | 4 +- .../components}/add_margin_form.tsx | 2 +- .../components}/add_provider_form.test.tsx | 4 +- .../components}/add_provider_form.tsx | 2 +- .../cost_tracking_settings.test.tsx | 6 +-- .../components}/cost_tracking_settings.tsx | 2 +- .../components}/how_it_works.test.tsx | 2 +- .../components}/how_it_works.tsx | 0 .../cost-tracking/components}/index.ts | 0 .../pricing_calculator/index.test.tsx | 2 +- .../components}/pricing_calculator/index.tsx | 0 .../multi_cost_results.test.tsx | 2 +- .../pricing_calculator/multi_cost_results.tsx | 0 .../multi_export_dropdown.test.tsx | 2 +- .../multi_export_dropdown.tsx | 0 .../multi_export_utils.test.ts | 0 .../pricing_calculator/multi_export_utils.ts | 0 .../components}/pricing_calculator/types.ts | 0 .../use_multi_cost_estimate.test.ts | 0 .../use_multi_cost_estimate.ts | 0 .../provider_discount_table.test.tsx | 2 +- .../components}/provider_discount_table.tsx | 2 +- .../provider_display_helpers.test.ts | 2 +- .../components}/provider_display_helpers.ts | 2 +- .../provider_margin_table.test.tsx | 2 +- .../components}/provider_margin_table.tsx | 2 +- .../cost-tracking/components}/types.ts | 0 .../components}/use_discount_config.test.ts | 2 +- .../components}/use_discount_config.ts | 4 +- .../components}/use_margin_config.test.ts | 2 +- .../components}/use_margin_config.ts | 4 +- .../app/(dashboard)/cost-tracking/page.tsx | 9 ++++ .../src/app/(dashboard)/logs/page.tsx | 17 +++++++ .../src/app/(dashboard)/page.tsx | 27 ------------ .../TransformRequestPanel.tsx} | 4 +- .../(dashboard)/transform-request/page.tsx | 9 ++++ .../(dashboard)/ui-theme/UIThemeSettings.tsx} | 2 +- .../src/app/(dashboard)/ui-theme/page.tsx | 9 ++++ .../src/utils/migratedPages.test.ts | 11 +++++ .../src/utils/migratedPages.ts | 5 +++ 54 files changed, 141 insertions(+), 86 deletions(-) rename ui/litellm-dashboard/src/{ => app/(dashboard)/caching}/components/cache_dashboard.tsx (98%) rename ui/litellm-dashboard/src/{ => app/(dashboard)/caching}/components/cache_health.tsx (100%) rename ui/litellm-dashboard/src/{ => app/(dashboard)/caching}/components/cache_settings/CacheFieldGroup.test.tsx (100%) rename ui/litellm-dashboard/src/{ => app/(dashboard)/caching}/components/cache_settings/CacheFieldGroup.tsx (100%) rename ui/litellm-dashboard/src/{ => app/(dashboard)/caching}/components/cache_settings/CacheFieldRenderer.test.tsx (100%) rename ui/litellm-dashboard/src/{ => app/(dashboard)/caching}/components/cache_settings/CacheFieldRenderer.tsx (98%) rename ui/litellm-dashboard/src/{ => app/(dashboard)/caching}/components/cache_settings/RedisTypeSelector.test.tsx (100%) rename ui/litellm-dashboard/src/{ => app/(dashboard)/caching}/components/cache_settings/RedisTypeSelector.tsx (100%) rename ui/litellm-dashboard/src/{ => app/(dashboard)/caching}/components/cache_settings/cacheSettingsUtils.ts (100%) rename ui/litellm-dashboard/src/{ => app/(dashboard)/caching}/components/cache_settings/index.tsx (98%) rename ui/litellm-dashboard/src/{ => app/(dashboard)/caching}/components/response_time_indicator.tsx (100%) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/caching/page.tsx rename ui/litellm-dashboard/src/{components/CostTrackingSettings => app/(dashboard)/cost-tracking/components}/add_margin_form.test.tsx (97%) rename ui/litellm-dashboard/src/{components/CostTrackingSettings => app/(dashboard)/cost-tracking/components}/add_margin_form.tsx (98%) rename ui/litellm-dashboard/src/{components/CostTrackingSettings => app/(dashboard)/cost-tracking/components}/add_provider_form.test.tsx (96%) rename ui/litellm-dashboard/src/{components/CostTrackingSettings => app/(dashboard)/cost-tracking/components}/add_provider_form.tsx (97%) rename ui/litellm-dashboard/src/{components/CostTrackingSettings => app/(dashboard)/cost-tracking/components}/cost_tracking_settings.test.tsx (97%) rename ui/litellm-dashboard/src/{components/CostTrackingSettings => app/(dashboard)/cost-tracking/components}/cost_tracking_settings.tsx (99%) rename ui/litellm-dashboard/src/{components/CostTrackingSettings => app/(dashboard)/cost-tracking/components}/how_it_works.test.tsx (98%) rename ui/litellm-dashboard/src/{components/CostTrackingSettings => app/(dashboard)/cost-tracking/components}/how_it_works.tsx (100%) rename ui/litellm-dashboard/src/{components/CostTrackingSettings => app/(dashboard)/cost-tracking/components}/index.ts (100%) rename ui/litellm-dashboard/src/{components/CostTrackingSettings => app/(dashboard)/cost-tracking/components}/pricing_calculator/index.test.tsx (98%) rename ui/litellm-dashboard/src/{components/CostTrackingSettings => app/(dashboard)/cost-tracking/components}/pricing_calculator/index.tsx (100%) rename ui/litellm-dashboard/src/{components/CostTrackingSettings => app/(dashboard)/cost-tracking/components}/pricing_calculator/multi_cost_results.test.tsx (99%) rename ui/litellm-dashboard/src/{components/CostTrackingSettings => app/(dashboard)/cost-tracking/components}/pricing_calculator/multi_cost_results.tsx (100%) rename ui/litellm-dashboard/src/{components/CostTrackingSettings => app/(dashboard)/cost-tracking/components}/pricing_calculator/multi_export_dropdown.test.tsx (98%) rename ui/litellm-dashboard/src/{components/CostTrackingSettings => app/(dashboard)/cost-tracking/components}/pricing_calculator/multi_export_dropdown.tsx (100%) rename ui/litellm-dashboard/src/{components/CostTrackingSettings => app/(dashboard)/cost-tracking/components}/pricing_calculator/multi_export_utils.test.ts (100%) rename ui/litellm-dashboard/src/{components/CostTrackingSettings => app/(dashboard)/cost-tracking/components}/pricing_calculator/multi_export_utils.ts (100%) rename ui/litellm-dashboard/src/{components/CostTrackingSettings => app/(dashboard)/cost-tracking/components}/pricing_calculator/types.ts (100%) rename ui/litellm-dashboard/src/{components/CostTrackingSettings => app/(dashboard)/cost-tracking/components}/pricing_calculator/use_multi_cost_estimate.test.ts (100%) rename ui/litellm-dashboard/src/{components/CostTrackingSettings => app/(dashboard)/cost-tracking/components}/pricing_calculator/use_multi_cost_estimate.ts (100%) rename ui/litellm-dashboard/src/{components/CostTrackingSettings => app/(dashboard)/cost-tracking/components}/provider_discount_table.test.tsx (99%) rename ui/litellm-dashboard/src/{components/CostTrackingSettings => app/(dashboard)/cost-tracking/components}/provider_discount_table.tsx (98%) rename ui/litellm-dashboard/src/{components/CostTrackingSettings => app/(dashboard)/cost-tracking/components}/provider_display_helpers.test.ts (98%) rename ui/litellm-dashboard/src/{components/CostTrackingSettings => app/(dashboard)/cost-tracking/components}/provider_display_helpers.ts (93%) rename ui/litellm-dashboard/src/{components/CostTrackingSettings => app/(dashboard)/cost-tracking/components}/provider_margin_table.test.tsx (99%) rename ui/litellm-dashboard/src/{components/CostTrackingSettings => app/(dashboard)/cost-tracking/components}/provider_margin_table.tsx (99%) rename ui/litellm-dashboard/src/{components/CostTrackingSettings => app/(dashboard)/cost-tracking/components}/types.ts (100%) rename ui/litellm-dashboard/src/{components/CostTrackingSettings => app/(dashboard)/cost-tracking/components}/use_discount_config.test.ts (99%) rename ui/litellm-dashboard/src/{components/CostTrackingSettings => app/(dashboard)/cost-tracking/components}/use_discount_config.ts (97%) rename ui/litellm-dashboard/src/{components/CostTrackingSettings => app/(dashboard)/cost-tracking/components}/use_margin_config.test.ts (99%) rename ui/litellm-dashboard/src/{components/CostTrackingSettings => app/(dashboard)/cost-tracking/components}/use_margin_config.ts (97%) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/page.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/logs/page.tsx rename ui/litellm-dashboard/src/{components/transform_request.tsx => app/(dashboard)/transform-request/TransformRequestPanel.tsx} (98%) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/transform-request/page.tsx rename ui/litellm-dashboard/src/{components/ui_theme_settings.tsx => app/(dashboard)/ui-theme/UIThemeSettings.tsx} (98%) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/ui-theme/page.tsx diff --git a/ui/litellm-dashboard/e2e_tests/fixtures/migratedPages.ts b/ui/litellm-dashboard/e2e_tests/fixtures/migratedPages.ts index 17a27d451df..19b90848e60 100644 --- a/ui/litellm-dashboard/e2e_tests/fixtures/migratedPages.ts +++ b/ui/litellm-dashboard/e2e_tests/fixtures/migratedPages.ts @@ -9,8 +9,8 @@ * - navigation specs that assert per-page URLs (tests/navigation/sidebar.spec.ts) * * Keep this in lockstep with MIGRATED_PAGES in src/utils/migratedPages.ts. - * Pending (add as each PR lands): the leaf-pages batch - * (caching, cost-tracking, logs, transform-request, ui-theme). + * Pending (add as each PR lands): admin-panel, logging-and-alerts, + * model-hub-table, and usage (#30268). */ export const MIGRATED_E2E_PAGES: Record = { api_ref: "api-reference", @@ -30,6 +30,11 @@ export const MIGRATED_E2E_PAGES: Record = { prompts: "prompts", "tool-policies": "tool-policies", skills: "skills", + caching: "caching", + "cost-tracking": "cost-tracking", + "transform-request": "transform-request", + "ui-theme": "ui-theme", + logs: "logs", }; export const MIGRATED_E2E_SEGMENTS: string[] = [...new Set(Object.values(MIGRATED_E2E_PAGES))]; diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 53dde4d28a4..6dc9be0fc90 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -293,77 +293,77 @@ "count": 1 } }, - "src/components/CostTrackingSettings/add_margin_form.tsx": { + "src/app/(dashboard)/cost-tracking/components/add_margin_form.tsx": { "no-restricted-imports": { "count": 1 } }, - "src/components/CostTrackingSettings/add_provider_form.tsx": { + "src/app/(dashboard)/cost-tracking/components/add_provider_form.tsx": { "no-restricted-imports": { "count": 1 } }, - "src/components/CostTrackingSettings/cost_tracking_settings.tsx": { + "src/app/(dashboard)/cost-tracking/components/cost_tracking_settings.tsx": { "no-restricted-imports": { "count": 1 } }, - "src/components/CostTrackingSettings/how_it_works.tsx": { + "src/app/(dashboard)/cost-tracking/components/how_it_works.tsx": { "no-restricted-imports": { "count": 1 } }, - "src/components/CostTrackingSettings/pricing_calculator/multi_cost_results.test.tsx": { + "src/app/(dashboard)/cost-tracking/components/pricing_calculator/multi_cost_results.test.tsx": { "unused-imports/no-unused-imports": { "count": 1 } }, - "src/components/CostTrackingSettings/pricing_calculator/multi_cost_results.tsx": { + "src/app/(dashboard)/cost-tracking/components/pricing_calculator/multi_cost_results.tsx": { "no-restricted-imports": { "count": 1 } }, - "src/components/CostTrackingSettings/pricing_calculator/multi_export_dropdown.test.tsx": { + "src/app/(dashboard)/cost-tracking/components/pricing_calculator/multi_export_dropdown.test.tsx": { "unused-imports/no-unused-imports": { "count": 1 } }, - "src/components/CostTrackingSettings/pricing_calculator/multi_export_dropdown.tsx": { + "src/app/(dashboard)/cost-tracking/components/pricing_calculator/multi_export_dropdown.tsx": { "no-restricted-imports": { "count": 1 } }, - "src/components/CostTrackingSettings/pricing_calculator/use_multi_cost_estimate.ts": { + "src/app/(dashboard)/cost-tracking/components/pricing_calculator/use_multi_cost_estimate.ts": { "no-restricted-syntax": { "count": 1 } }, - "src/components/CostTrackingSettings/provider_discount_table.test.tsx": { + "src/app/(dashboard)/cost-tracking/components/provider_discount_table.test.tsx": { "unused-imports/no-unused-imports": { "count": 1 } }, - "src/components/CostTrackingSettings/provider_discount_table.tsx": { + "src/app/(dashboard)/cost-tracking/components/provider_discount_table.tsx": { "no-restricted-imports": { "count": 1 } }, - "src/components/CostTrackingSettings/provider_display_helpers.test.ts": { + "src/app/(dashboard)/cost-tracking/components/provider_display_helpers.test.ts": { "unused-imports/no-unused-imports": { "count": 1 } }, - "src/components/CostTrackingSettings/provider_margin_table.tsx": { + "src/app/(dashboard)/cost-tracking/components/provider_margin_table.tsx": { "no-restricted-imports": { "count": 1 } }, - "src/components/CostTrackingSettings/use_discount_config.ts": { + "src/app/(dashboard)/cost-tracking/components/use_discount_config.ts": { "no-restricted-syntax": { "count": 2 } }, - "src/components/CostTrackingSettings/use_margin_config.ts": { + "src/app/(dashboard)/cost-tracking/components/use_margin_config.ts": { "no-restricted-syntax": { "count": 2 } @@ -826,7 +826,7 @@ "count": 1 } }, - "src/components/cache_dashboard.tsx": { + "src/app/(dashboard)/caching/components/cache_dashboard.tsx": { "no-restricted-imports": { "count": 1 }, @@ -837,22 +837,22 @@ "count": 2 } }, - "src/components/cache_health.tsx": { + "src/app/(dashboard)/caching/components/cache_health.tsx": { "no-restricted-imports": { "count": 1 } }, - "src/components/cache_settings/CacheFieldRenderer.tsx": { + "src/app/(dashboard)/caching/components/cache_settings/CacheFieldRenderer.tsx": { "no-restricted-imports": { "count": 1 } }, - "src/components/cache_settings/RedisTypeSelector.tsx": { + "src/app/(dashboard)/caching/components/cache_settings/RedisTypeSelector.tsx": { "no-restricted-imports": { "count": 1 } }, - "src/components/cache_settings/index.tsx": { + "src/app/(dashboard)/caching/components/cache_settings/index.tsx": { "no-restricted-imports": { "count": 1 }, @@ -1993,12 +1993,12 @@ "count": 1 } }, - "src/components/transform_request.tsx": { + "src/app/(dashboard)/transform-request/TransformRequestPanel.tsx": { "no-restricted-imports": { "count": 1 } }, - "src/components/ui_theme_settings.tsx": { + "src/app/(dashboard)/ui-theme/UIThemeSettings.tsx": { "no-restricted-imports": { "count": 1 }, diff --git a/ui/litellm-dashboard/src/components/cache_dashboard.tsx b/ui/litellm-dashboard/src/app/(dashboard)/caching/components/cache_dashboard.tsx similarity index 98% rename from ui/litellm-dashboard/src/components/cache_dashboard.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/caching/components/cache_dashboard.tsx index 874cb43276e..99656f0db4a 100644 --- a/ui/litellm-dashboard/src/components/cache_dashboard.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/caching/components/cache_dashboard.tsx @@ -16,11 +16,11 @@ import { Text, } from "@tremor/react"; import React, { useEffect, useState } from "react"; -import NotificationsManager from "./molecules/notifications_manager"; -import UsageDatePicker from "./shared/usage_date_picker"; +import NotificationsManager from "@/components/molecules/notifications_manager"; +import UsageDatePicker from "@/components/shared/usage_date_picker"; import { RefreshIcon } from "@heroicons/react/outline"; -import { adminGlobalCacheActivity, cachingHealthCheckCall } from "./networking"; +import { adminGlobalCacheActivity, cachingHealthCheckCall } from "@/components/networking"; // Import the new component import { CacheHealthTab } from "./cache_health"; diff --git a/ui/litellm-dashboard/src/components/cache_health.tsx b/ui/litellm-dashboard/src/app/(dashboard)/caching/components/cache_health.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/cache_health.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/caching/components/cache_health.tsx diff --git a/ui/litellm-dashboard/src/components/cache_settings/CacheFieldGroup.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/caching/components/cache_settings/CacheFieldGroup.test.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/cache_settings/CacheFieldGroup.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/caching/components/cache_settings/CacheFieldGroup.test.tsx diff --git a/ui/litellm-dashboard/src/components/cache_settings/CacheFieldGroup.tsx b/ui/litellm-dashboard/src/app/(dashboard)/caching/components/cache_settings/CacheFieldGroup.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/cache_settings/CacheFieldGroup.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/caching/components/cache_settings/CacheFieldGroup.tsx diff --git a/ui/litellm-dashboard/src/components/cache_settings/CacheFieldRenderer.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/caching/components/cache_settings/CacheFieldRenderer.test.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/cache_settings/CacheFieldRenderer.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/caching/components/cache_settings/CacheFieldRenderer.test.tsx diff --git a/ui/litellm-dashboard/src/components/cache_settings/CacheFieldRenderer.tsx b/ui/litellm-dashboard/src/app/(dashboard)/caching/components/cache_settings/CacheFieldRenderer.tsx similarity index 98% rename from ui/litellm-dashboard/src/components/cache_settings/CacheFieldRenderer.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/caching/components/cache_settings/CacheFieldRenderer.tsx index 6608b09d261..27d9fc57200 100644 --- a/ui/litellm-dashboard/src/components/cache_settings/CacheFieldRenderer.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/caching/components/cache_settings/CacheFieldRenderer.tsx @@ -5,7 +5,7 @@ import { NumberInput, TextInput } from "@tremor/react"; import { Select } from "antd"; import React, { useEffect, useState } from "react"; import { fetchAvailableModels, ModelGroup } from "@/components/llm_calls/fetch_models"; -import NumericalInput from "../shared/numerical_input"; +import NumericalInput from "@/components/shared/numerical_input"; interface CacheFieldRendererProps { field: any; diff --git a/ui/litellm-dashboard/src/components/cache_settings/RedisTypeSelector.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/caching/components/cache_settings/RedisTypeSelector.test.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/cache_settings/RedisTypeSelector.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/caching/components/cache_settings/RedisTypeSelector.test.tsx diff --git a/ui/litellm-dashboard/src/components/cache_settings/RedisTypeSelector.tsx b/ui/litellm-dashboard/src/app/(dashboard)/caching/components/cache_settings/RedisTypeSelector.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/cache_settings/RedisTypeSelector.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/caching/components/cache_settings/RedisTypeSelector.tsx diff --git a/ui/litellm-dashboard/src/components/cache_settings/cacheSettingsUtils.ts b/ui/litellm-dashboard/src/app/(dashboard)/caching/components/cache_settings/cacheSettingsUtils.ts similarity index 100% rename from ui/litellm-dashboard/src/components/cache_settings/cacheSettingsUtils.ts rename to ui/litellm-dashboard/src/app/(dashboard)/caching/components/cache_settings/cacheSettingsUtils.ts diff --git a/ui/litellm-dashboard/src/components/cache_settings/index.tsx b/ui/litellm-dashboard/src/app/(dashboard)/caching/components/cache_settings/index.tsx similarity index 98% rename from ui/litellm-dashboard/src/components/cache_settings/index.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/caching/components/cache_settings/index.tsx index c7d8c579af3..7de49e08ace 100644 --- a/ui/litellm-dashboard/src/components/cache_settings/index.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/caching/components/cache_settings/index.tsx @@ -1,7 +1,7 @@ import React, { useState, useEffect, useCallback } from "react"; import { Button, Accordion, AccordionHeader, AccordionBody } from "@tremor/react"; -import { getCacheSettingsCall, testCacheConnectionCall, updateCacheSettingsCall } from "../networking"; -import NotificationsManager from "../molecules/notifications_manager"; +import { getCacheSettingsCall, testCacheConnectionCall, updateCacheSettingsCall } from "@/components/networking"; +import NotificationsManager from "@/components/molecules/notifications_manager"; import RedisTypeSelector from "./RedisTypeSelector"; import CacheFieldRenderer from "./CacheFieldRenderer"; import { gatherFormValues, groupFieldsByCategory } from "./cacheSettingsUtils"; diff --git a/ui/litellm-dashboard/src/components/response_time_indicator.tsx b/ui/litellm-dashboard/src/app/(dashboard)/caching/components/response_time_indicator.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/response_time_indicator.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/caching/components/response_time_indicator.tsx diff --git a/ui/litellm-dashboard/src/app/(dashboard)/caching/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/caching/page.tsx new file mode 100644 index 00000000000..0ef88ec9eb5 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/caching/page.tsx @@ -0,0 +1,17 @@ +"use client"; + +import CacheDashboard from "./components/cache_dashboard"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; + +export default function Caching() { + const { accessToken, userRole, userId, token, premiumUser } = useAuthorized(); + return ( + + ); +} diff --git a/ui/litellm-dashboard/src/components/CostTrackingSettings/add_margin_form.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/add_margin_form.test.tsx similarity index 97% rename from ui/litellm-dashboard/src/components/CostTrackingSettings/add_margin_form.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/add_margin_form.test.tsx index 9d261e1b686..21ee41936c1 100644 --- a/ui/litellm-dashboard/src/components/CostTrackingSettings/add_margin_form.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/add_margin_form.test.tsx @@ -2,11 +2,11 @@ import React from "react"; import { describe, it, expect, vi, beforeEach } from "vitest"; import { screen } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; -import { renderWithProviders } from "../../../tests/test-utils"; +import { renderWithProviders } from "../../../../../tests/test-utils"; import AddMarginForm from "./add_margin_form"; import { MarginConfig } from "./types"; -vi.mock("../provider_info_helpers", () => ({ +vi.mock("@/components/provider_info_helpers", () => ({ Providers: { OpenAI: "OpenAI", Anthropic: "Anthropic", diff --git a/ui/litellm-dashboard/src/components/CostTrackingSettings/add_margin_form.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/add_margin_form.tsx similarity index 98% rename from ui/litellm-dashboard/src/components/CostTrackingSettings/add_margin_form.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/add_margin_form.tsx index a3900eab257..56b34d6a68b 100644 --- a/ui/litellm-dashboard/src/components/CostTrackingSettings/add_margin_form.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/add_margin_form.tsx @@ -2,7 +2,7 @@ import React from "react"; import { TextInput, Button } from "@tremor/react"; import { Select as AntdSelect, Form, Tooltip, Radio } from "antd"; import { InfoCircleOutlined } from "@ant-design/icons"; -import { Providers, provider_map, providerLogoMap } from "../provider_info_helpers"; +import { Providers, provider_map, providerLogoMap } from "@/components/provider_info_helpers"; import { MarginConfig } from "./types"; import { handleImageError } from "./provider_display_helpers"; diff --git a/ui/litellm-dashboard/src/components/CostTrackingSettings/add_provider_form.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/add_provider_form.test.tsx similarity index 96% rename from ui/litellm-dashboard/src/components/CostTrackingSettings/add_provider_form.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/add_provider_form.test.tsx index e0e5600126b..48d23d4645d 100644 --- a/ui/litellm-dashboard/src/components/CostTrackingSettings/add_provider_form.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/add_provider_form.test.tsx @@ -2,11 +2,11 @@ import React from "react"; import { describe, it, expect, vi, beforeEach } from "vitest"; import { screen } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; -import { renderWithProviders } from "../../../tests/test-utils"; +import { renderWithProviders } from "../../../../../tests/test-utils"; import AddProviderForm from "./add_provider_form"; import { DiscountConfig } from "./types"; -vi.mock("../provider_info_helpers", () => ({ +vi.mock("@/components/provider_info_helpers", () => ({ Providers: { OpenAI: "OpenAI", Anthropic: "Anthropic", diff --git a/ui/litellm-dashboard/src/components/CostTrackingSettings/add_provider_form.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/add_provider_form.tsx similarity index 97% rename from ui/litellm-dashboard/src/components/CostTrackingSettings/add_provider_form.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/add_provider_form.tsx index bb11acb83aa..61ba3194607 100644 --- a/ui/litellm-dashboard/src/components/CostTrackingSettings/add_provider_form.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/add_provider_form.tsx @@ -2,7 +2,7 @@ import React from "react"; import { TextInput, Button } from "@tremor/react"; import { Select as AntdSelect, Form, Tooltip } from "antd"; import { InfoCircleOutlined } from "@ant-design/icons"; -import { Providers, provider_map, providerLogoMap } from "../provider_info_helpers"; +import { Providers, provider_map, providerLogoMap } from "@/components/provider_info_helpers"; import { DiscountConfig } from "./types"; import { handleImageError } from "./provider_display_helpers"; diff --git a/ui/litellm-dashboard/src/components/CostTrackingSettings/cost_tracking_settings.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/cost_tracking_settings.test.tsx similarity index 97% rename from ui/litellm-dashboard/src/components/CostTrackingSettings/cost_tracking_settings.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/cost_tracking_settings.test.tsx index 89711a098fe..0e1c7da92ba 100644 --- a/ui/litellm-dashboard/src/components/CostTrackingSettings/cost_tracking_settings.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/cost_tracking_settings.test.tsx @@ -2,7 +2,7 @@ import React from "react"; import { describe, it, expect, vi, beforeEach } from "vitest"; import { screen } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; -import { renderWithProviders } from "../../../tests/test-utils"; +import { renderWithProviders } from "../../../../../tests/test-utils"; import CostTrackingSettings from "./cost_tracking_settings"; // Mock sub-hooks so we can control their state without network calls @@ -37,7 +37,7 @@ vi.mock("@/components/llm_calls/fetch_models", () => ({ fetchAvailableModels: vi.fn().mockResolvedValue([]), })); -vi.mock("../HelpLink", () => ({ +vi.mock("@/components/HelpLink", () => ({ DocsMenu: () => null, })); @@ -45,7 +45,7 @@ vi.mock("./how_it_works", () => ({ default: () =>
How It Works
, })); -vi.mock("../provider_info_helpers", () => ({ +vi.mock("@/components/provider_info_helpers", () => ({ Providers: { OpenAI: "OpenAI" }, provider_map: { OpenAI: "openai" }, providerLogoMap: {}, diff --git a/ui/litellm-dashboard/src/components/CostTrackingSettings/cost_tracking_settings.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/cost_tracking_settings.tsx similarity index 99% rename from ui/litellm-dashboard/src/components/CostTrackingSettings/cost_tracking_settings.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/cost_tracking_settings.tsx index d9cca4d3c23..22ea8d8d517 100644 --- a/ui/litellm-dashboard/src/components/CostTrackingSettings/cost_tracking_settings.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/cost_tracking_settings.tsx @@ -20,7 +20,7 @@ import ProviderMarginTable from "./provider_margin_table"; import AddMarginForm from "./add_margin_form"; import PricingCalculator from "./pricing_calculator/index"; import { ExclamationCircleOutlined } from "@ant-design/icons"; -import { DocsMenu } from "../HelpLink"; +import { DocsMenu } from "@/components/HelpLink"; import HowItWorks from "./how_it_works"; import { useDiscountConfig } from "./use_discount_config"; import { useMarginConfig } from "./use_margin_config"; diff --git a/ui/litellm-dashboard/src/components/CostTrackingSettings/how_it_works.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/how_it_works.test.tsx similarity index 98% rename from ui/litellm-dashboard/src/components/CostTrackingSettings/how_it_works.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/how_it_works.test.tsx index fa608f555ce..711a8795f15 100644 --- a/ui/litellm-dashboard/src/components/CostTrackingSettings/how_it_works.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/how_it_works.test.tsx @@ -2,7 +2,7 @@ import React from "react"; import { describe, it, expect, vi, beforeEach } from "vitest"; import { screen } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; -import { renderWithProviders } from "../../../tests/test-utils"; +import { renderWithProviders } from "../../../../../tests/test-utils"; import HowItWorks from "./how_it_works"; vi.mock("@/app/(dashboard)/api-reference/components/CodeBlock", () => ({ diff --git a/ui/litellm-dashboard/src/components/CostTrackingSettings/how_it_works.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/how_it_works.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/CostTrackingSettings/how_it_works.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/how_it_works.tsx diff --git a/ui/litellm-dashboard/src/components/CostTrackingSettings/index.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/index.ts similarity index 100% rename from ui/litellm-dashboard/src/components/CostTrackingSettings/index.ts rename to ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/index.ts diff --git a/ui/litellm-dashboard/src/components/CostTrackingSettings/pricing_calculator/index.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/pricing_calculator/index.test.tsx similarity index 98% rename from ui/litellm-dashboard/src/components/CostTrackingSettings/pricing_calculator/index.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/pricing_calculator/index.test.tsx index 3e39e87a4b1..e7a858196c0 100644 --- a/ui/litellm-dashboard/src/components/CostTrackingSettings/pricing_calculator/index.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/pricing_calculator/index.test.tsx @@ -2,7 +2,7 @@ import React from "react"; import { describe, it, expect, vi, beforeEach } from "vitest"; import { screen, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; -import { renderWithProviders } from "../../../../tests/test-utils"; +import { renderWithProviders } from "../../../../../../tests/test-utils"; import PricingCalculator from "./index"; import type { ModelEntry } from "./types"; import type { MultiModelResult } from "./types"; diff --git a/ui/litellm-dashboard/src/components/CostTrackingSettings/pricing_calculator/index.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/pricing_calculator/index.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/CostTrackingSettings/pricing_calculator/index.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/pricing_calculator/index.tsx diff --git a/ui/litellm-dashboard/src/components/CostTrackingSettings/pricing_calculator/multi_cost_results.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/pricing_calculator/multi_cost_results.test.tsx similarity index 99% rename from ui/litellm-dashboard/src/components/CostTrackingSettings/pricing_calculator/multi_cost_results.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/pricing_calculator/multi_cost_results.test.tsx index a4ca0b01e79..6dc9309b5e3 100644 --- a/ui/litellm-dashboard/src/components/CostTrackingSettings/pricing_calculator/multi_cost_results.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/pricing_calculator/multi_cost_results.test.tsx @@ -2,7 +2,7 @@ import React from "react"; import { describe, it, expect, vi, beforeEach } from "vitest"; import { screen, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; -import { renderWithProviders } from "../../../../tests/test-utils"; +import { renderWithProviders } from "../../../../../../tests/test-utils"; import MultiCostResults from "./multi_cost_results"; import type { MultiModelResult } from "./types"; import type { CostEstimateResponse } from "../types"; diff --git a/ui/litellm-dashboard/src/components/CostTrackingSettings/pricing_calculator/multi_cost_results.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/pricing_calculator/multi_cost_results.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/CostTrackingSettings/pricing_calculator/multi_cost_results.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/pricing_calculator/multi_cost_results.tsx diff --git a/ui/litellm-dashboard/src/components/CostTrackingSettings/pricing_calculator/multi_export_dropdown.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/pricing_calculator/multi_export_dropdown.test.tsx similarity index 98% rename from ui/litellm-dashboard/src/components/CostTrackingSettings/pricing_calculator/multi_export_dropdown.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/pricing_calculator/multi_export_dropdown.test.tsx index 20495c44311..02940dd1325 100644 --- a/ui/litellm-dashboard/src/components/CostTrackingSettings/pricing_calculator/multi_export_dropdown.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/pricing_calculator/multi_export_dropdown.test.tsx @@ -2,7 +2,7 @@ import React from "react"; import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; import { screen, fireEvent } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; -import { renderWithProviders } from "../../../../tests/test-utils"; +import { renderWithProviders } from "../../../../../../tests/test-utils"; import MultiExportDropdown from "./multi_export_dropdown"; import type { MultiModelResult } from "./types"; diff --git a/ui/litellm-dashboard/src/components/CostTrackingSettings/pricing_calculator/multi_export_dropdown.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/pricing_calculator/multi_export_dropdown.tsx similarity index 100% rename from ui/litellm-dashboard/src/components/CostTrackingSettings/pricing_calculator/multi_export_dropdown.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/pricing_calculator/multi_export_dropdown.tsx diff --git a/ui/litellm-dashboard/src/components/CostTrackingSettings/pricing_calculator/multi_export_utils.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/pricing_calculator/multi_export_utils.test.ts similarity index 100% rename from ui/litellm-dashboard/src/components/CostTrackingSettings/pricing_calculator/multi_export_utils.test.ts rename to ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/pricing_calculator/multi_export_utils.test.ts diff --git a/ui/litellm-dashboard/src/components/CostTrackingSettings/pricing_calculator/multi_export_utils.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/pricing_calculator/multi_export_utils.ts similarity index 100% rename from ui/litellm-dashboard/src/components/CostTrackingSettings/pricing_calculator/multi_export_utils.ts rename to ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/pricing_calculator/multi_export_utils.ts diff --git a/ui/litellm-dashboard/src/components/CostTrackingSettings/pricing_calculator/types.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/pricing_calculator/types.ts similarity index 100% rename from ui/litellm-dashboard/src/components/CostTrackingSettings/pricing_calculator/types.ts rename to ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/pricing_calculator/types.ts diff --git a/ui/litellm-dashboard/src/components/CostTrackingSettings/pricing_calculator/use_multi_cost_estimate.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/pricing_calculator/use_multi_cost_estimate.test.ts similarity index 100% rename from ui/litellm-dashboard/src/components/CostTrackingSettings/pricing_calculator/use_multi_cost_estimate.test.ts rename to ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/pricing_calculator/use_multi_cost_estimate.test.ts diff --git a/ui/litellm-dashboard/src/components/CostTrackingSettings/pricing_calculator/use_multi_cost_estimate.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/pricing_calculator/use_multi_cost_estimate.ts similarity index 100% rename from ui/litellm-dashboard/src/components/CostTrackingSettings/pricing_calculator/use_multi_cost_estimate.ts rename to ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/pricing_calculator/use_multi_cost_estimate.ts diff --git a/ui/litellm-dashboard/src/components/CostTrackingSettings/provider_discount_table.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/provider_discount_table.test.tsx similarity index 99% rename from ui/litellm-dashboard/src/components/CostTrackingSettings/provider_discount_table.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/provider_discount_table.test.tsx index 130b7adffe4..c1c43ebdb4f 100644 --- a/ui/litellm-dashboard/src/components/CostTrackingSettings/provider_discount_table.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/provider_discount_table.test.tsx @@ -2,7 +2,7 @@ import React from "react"; import { describe, it, expect, vi, beforeEach } from "vitest"; import { screen, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; -import { renderWithProviders } from "../../../tests/test-utils"; +import { renderWithProviders } from "../../../../../tests/test-utils"; import ProviderDiscountTable from "./provider_discount_table"; vi.mock("@heroicons/react/outline", () => ({ diff --git a/ui/litellm-dashboard/src/components/CostTrackingSettings/provider_discount_table.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/provider_discount_table.tsx similarity index 98% rename from ui/litellm-dashboard/src/components/CostTrackingSettings/provider_discount_table.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/provider_discount_table.tsx index 43c052b9e5c..d802f6d83dd 100644 --- a/ui/litellm-dashboard/src/components/CostTrackingSettings/provider_discount_table.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/provider_discount_table.tsx @@ -1,7 +1,7 @@ import React, { useState } from "react"; import { TextInput, Icon, Text } from "@tremor/react"; import { TrashIcon, PencilAltIcon, CheckIcon, XIcon } from "@heroicons/react/outline"; -import { SimpleTable } from "../common_components/simple_table"; +import { SimpleTable } from "@/components/common_components/simple_table"; import { DiscountConfig } from "./types"; import { getProviderDisplayInfo, handleImageError } from "./provider_display_helpers"; diff --git a/ui/litellm-dashboard/src/components/CostTrackingSettings/provider_display_helpers.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/provider_display_helpers.test.ts similarity index 98% rename from ui/litellm-dashboard/src/components/CostTrackingSettings/provider_display_helpers.test.ts rename to ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/provider_display_helpers.test.ts index 9668f07c2c5..c7b93c6f825 100644 --- a/ui/litellm-dashboard/src/components/CostTrackingSettings/provider_display_helpers.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/provider_display_helpers.test.ts @@ -1,7 +1,7 @@ import { describe, it, expect, vi, beforeEach } from "vitest"; import { getProviderDisplayInfo, getProviderBackendValue, handleImageError } from "./provider_display_helpers"; -vi.mock("../provider_info_helpers", () => ({ +vi.mock("@/components/provider_info_helpers", () => ({ Providers: { OpenAI: "OpenAI", Anthropic: "Anthropic", diff --git a/ui/litellm-dashboard/src/components/CostTrackingSettings/provider_display_helpers.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/provider_display_helpers.ts similarity index 93% rename from ui/litellm-dashboard/src/components/CostTrackingSettings/provider_display_helpers.ts rename to ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/provider_display_helpers.ts index dc61a9d6218..cd088da09da 100644 --- a/ui/litellm-dashboard/src/components/CostTrackingSettings/provider_display_helpers.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/provider_display_helpers.ts @@ -1,4 +1,4 @@ -import { Providers, provider_map, providerLogoMap } from "../provider_info_helpers"; +import { Providers, provider_map, providerLogoMap } from "@/components/provider_info_helpers"; export interface ProviderDisplayInfo { displayName: string; diff --git a/ui/litellm-dashboard/src/components/CostTrackingSettings/provider_margin_table.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/provider_margin_table.test.tsx similarity index 99% rename from ui/litellm-dashboard/src/components/CostTrackingSettings/provider_margin_table.test.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/provider_margin_table.test.tsx index 3f0ab4ae16b..e1b17dea23d 100644 --- a/ui/litellm-dashboard/src/components/CostTrackingSettings/provider_margin_table.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/provider_margin_table.test.tsx @@ -2,7 +2,7 @@ import React from "react"; import { describe, it, expect, vi, beforeEach } from "vitest"; import { screen } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; -import { renderWithProviders } from "../../../tests/test-utils"; +import { renderWithProviders } from "../../../../../tests/test-utils"; import ProviderMarginTable from "./provider_margin_table"; vi.mock("@heroicons/react/outline", () => ({ diff --git a/ui/litellm-dashboard/src/components/CostTrackingSettings/provider_margin_table.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/provider_margin_table.tsx similarity index 99% rename from ui/litellm-dashboard/src/components/CostTrackingSettings/provider_margin_table.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/provider_margin_table.tsx index bee7a1219d2..b2baccc510f 100644 --- a/ui/litellm-dashboard/src/components/CostTrackingSettings/provider_margin_table.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/provider_margin_table.tsx @@ -1,7 +1,7 @@ import React, { useState } from "react"; import { TextInput, Icon, Text } from "@tremor/react"; import { TrashIcon, PencilAltIcon, CheckIcon, XIcon } from "@heroicons/react/outline"; -import { SimpleTable } from "../common_components/simple_table"; +import { SimpleTable } from "@/components/common_components/simple_table"; import { MarginConfig } from "./types"; import { getProviderDisplayInfo, handleImageError } from "./provider_display_helpers"; diff --git a/ui/litellm-dashboard/src/components/CostTrackingSettings/types.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/types.ts similarity index 100% rename from ui/litellm-dashboard/src/components/CostTrackingSettings/types.ts rename to ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/types.ts diff --git a/ui/litellm-dashboard/src/components/CostTrackingSettings/use_discount_config.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/use_discount_config.test.ts similarity index 99% rename from ui/litellm-dashboard/src/components/CostTrackingSettings/use_discount_config.test.ts rename to ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/use_discount_config.test.ts index 967be542a81..d0ebb8ee7c7 100644 --- a/ui/litellm-dashboard/src/components/CostTrackingSettings/use_discount_config.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/use_discount_config.test.ts @@ -18,7 +18,7 @@ vi.mock("./provider_display_helpers", () => ({ }), })); -vi.mock("../provider_info_helpers", () => ({ +vi.mock("@/components/provider_info_helpers", () => ({ Providers: { OpenAI: "OpenAI", Anthropic: "Anthropic", diff --git a/ui/litellm-dashboard/src/components/CostTrackingSettings/use_discount_config.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/use_discount_config.ts similarity index 97% rename from ui/litellm-dashboard/src/components/CostTrackingSettings/use_discount_config.ts rename to ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/use_discount_config.ts index 5ed00ce1cbc..c9b4f47a7b8 100644 --- a/ui/litellm-dashboard/src/components/CostTrackingSettings/use_discount_config.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/use_discount_config.ts @@ -1,9 +1,9 @@ import { useState, useCallback } from "react"; import { getProxyBaseUrl, getGlobalLitellmHeaderName } from "@/components/networking"; -import NotificationsManager from "../molecules/notifications_manager"; +import NotificationsManager from "@/components/molecules/notifications_manager"; import { DiscountConfig } from "./types"; import { getProviderBackendValue } from "./provider_display_helpers"; -import { Providers } from "../provider_info_helpers"; +import { Providers } from "@/components/provider_info_helpers"; export interface UseDiscountConfigProps { accessToken: string | null; diff --git a/ui/litellm-dashboard/src/components/CostTrackingSettings/use_margin_config.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/use_margin_config.test.ts similarity index 99% rename from ui/litellm-dashboard/src/components/CostTrackingSettings/use_margin_config.test.ts rename to ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/use_margin_config.test.ts index 8f9085de539..88a865e4fa2 100644 --- a/ui/litellm-dashboard/src/components/CostTrackingSettings/use_margin_config.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/use_margin_config.test.ts @@ -18,7 +18,7 @@ vi.mock("./provider_display_helpers", () => ({ }), })); -vi.mock("../provider_info_helpers", () => ({ +vi.mock("@/components/provider_info_helpers", () => ({ Providers: { OpenAI: "OpenAI", Anthropic: "Anthropic", diff --git a/ui/litellm-dashboard/src/components/CostTrackingSettings/use_margin_config.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/use_margin_config.ts similarity index 97% rename from ui/litellm-dashboard/src/components/CostTrackingSettings/use_margin_config.ts rename to ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/use_margin_config.ts index 0af70b070d5..4994e9e6678 100644 --- a/ui/litellm-dashboard/src/components/CostTrackingSettings/use_margin_config.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/components/use_margin_config.ts @@ -1,9 +1,9 @@ import { useState, useCallback } from "react"; import { getProxyBaseUrl, getGlobalLitellmHeaderName } from "@/components/networking"; -import NotificationsManager from "../molecules/notifications_manager"; +import NotificationsManager from "@/components/molecules/notifications_manager"; import { MarginConfig } from "./types"; import { getProviderBackendValue } from "./provider_display_helpers"; -import { Providers } from "../provider_info_helpers"; +import { Providers } from "@/components/provider_info_helpers"; export interface UseMarginConfigProps { accessToken: string | null; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/page.tsx new file mode 100644 index 00000000000..c72fed4c594 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/page.tsx @@ -0,0 +1,9 @@ +"use client"; + +import { CostTrackingSettings } from "./components"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; + +export default function CostTracking() { + const { accessToken, userRole, userId } = useAuthorized(); + return ; +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/logs/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/logs/page.tsx new file mode 100644 index 00000000000..88909e3b87f --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/logs/page.tsx @@ -0,0 +1,17 @@ +"use client"; + +import SpendLogsTable from "@/components/view_logs"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; + +export default function Logs() { + const { accessToken, userRole, userId, token, premiumUser } = useAuthorized(); + return ( + + ); +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/page.tsx index 9318bc1b332..864007b4fe8 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/page.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/page.tsx @@ -3,12 +3,10 @@ import ModelsAndEndpointsView from "@/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView"; import AdminPanel from "@/components/AdminPanel"; import AgentsPanel from "@/components/agents"; -import CacheDashboard from "@/components/cache_dashboard"; import { teamListCall as v2TeamListCall } from "@/app/(dashboard)/hooks/teams/useTeams"; import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings"; import useProxySettings from "@/app/(dashboard)/hooks/proxySettings/useProxySettings"; import LoadingScreen from "@/components/common_components/LoadingScreen"; -import { CostTrackingSettings } from "@/components/CostTrackingSettings"; import GeneralSettings from "@/components/general_settings"; import { Team } from "@/components/key_team_helpers/key_list"; import ModelHubTable from "@/components/AIHub/ModelHubTable"; @@ -21,11 +19,8 @@ import PassThroughSettings from "@/components/pass_through_settings"; import PublicModelHub from "@/components/public_model_hub"; import Settings from "@/components/settings"; import { SurveyPrompt, SurveyModal, ClaudeCodePrompt, ClaudeCodeModal } from "@/components/survey"; -import TransformRequestPanel from "@/components/transform_request"; -import UIThemeSettings from "@/components/ui_theme_settings"; import Usage from "@/components/usage"; import UserDashboard from "@/components/user_dashboard"; -import SpendLogsTable from "@/components/view_logs"; import ViewUserDashboard from "@/components/view_users"; import { useAuth } from "@/contexts/AuthContext"; import { @@ -374,14 +369,8 @@ function CreateKeyPageContent() { ) : page == "agents" ? ( - ) : page == "transform-request" ? ( - ) : page == "router-settings" ? ( - ) : page == "ui-theme" ? ( - - ) : page == "cost-tracking" ? ( - ) : page == "model-hub-table" ? ( isAdminRole(userRole) ? ( ) - ) : page == "caching" ? ( - ) : page == "pass-through-settings" ? ( - ) : page == "logs" ? ( - ) : page == "new_usage" ? ( ) : ( diff --git a/ui/litellm-dashboard/src/components/transform_request.tsx b/ui/litellm-dashboard/src/app/(dashboard)/transform-request/TransformRequestPanel.tsx similarity index 98% rename from ui/litellm-dashboard/src/components/transform_request.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/transform-request/TransformRequestPanel.tsx index cc68972d009..04d1701de3f 100644 --- a/ui/litellm-dashboard/src/components/transform_request.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/transform-request/TransformRequestPanel.tsx @@ -2,8 +2,8 @@ import React, { useState } from "react"; import { Button } from "antd"; import { CopyOutlined } from "@ant-design/icons"; import { Title } from "@tremor/react"; -import { transformRequestCall } from "./networking"; -import NotificationsManager from "./molecules/notifications_manager"; +import { transformRequestCall } from "@/components/networking"; +import NotificationsManager from "@/components/molecules/notifications_manager"; interface TransformRequestPanelProps { accessToken: string | null; } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/transform-request/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/transform-request/page.tsx new file mode 100644 index 00000000000..55289af3e43 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/transform-request/page.tsx @@ -0,0 +1,9 @@ +"use client"; + +import TransformRequestPanel from "./TransformRequestPanel"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; + +export default function TransformRequest() { + const { accessToken } = useAuthorized(); + return ; +} diff --git a/ui/litellm-dashboard/src/components/ui_theme_settings.tsx b/ui/litellm-dashboard/src/app/(dashboard)/ui-theme/UIThemeSettings.tsx similarity index 98% rename from ui/litellm-dashboard/src/components/ui_theme_settings.tsx rename to ui/litellm-dashboard/src/app/(dashboard)/ui-theme/UIThemeSettings.tsx index b68b0aeb1a6..2b70a0e8c96 100644 --- a/ui/litellm-dashboard/src/components/ui_theme_settings.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/ui-theme/UIThemeSettings.tsx @@ -2,7 +2,7 @@ import React, { useState, useEffect } from "react"; import { Card, Title, Text, TextInput, Button } from "@tremor/react"; import { useTheme } from "@/contexts/ThemeContext"; import { getProxyBaseUrl, getGlobalLitellmHeaderName } from "@/components/networking"; -import NotificationsManager from "./molecules/notifications_manager"; +import NotificationsManager from "@/components/molecules/notifications_manager"; interface UIThemeSettingsProps { userID: string | null; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/ui-theme/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/ui-theme/page.tsx new file mode 100644 index 00000000000..e80caa22c74 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/ui-theme/page.tsx @@ -0,0 +1,9 @@ +"use client"; + +import UIThemeSettings from "./UIThemeSettings"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; + +export default function UITheme() { + const { accessToken, userRole, userId } = useAuthorized(); + return ; +} diff --git a/ui/litellm-dashboard/src/utils/migratedPages.test.ts b/ui/litellm-dashboard/src/utils/migratedPages.test.ts index 8471c8a2567..1183c81d05f 100644 --- a/ui/litellm-dashboard/src/utils/migratedPages.test.ts +++ b/ui/litellm-dashboard/src/utils/migratedPages.test.ts @@ -88,6 +88,17 @@ describe("migratedHref / legacyPageHref", () => { // Old bookmarks used ?page=claude-code-plugins for the same panel. expect(MIGRATED_PAGES["claude-code-plugins"]).toBe("skills"); }); + + it("maps the caching, cost-tracking, transform-request, ui-theme, and logs ids to their routes", async () => { + vi.doMock("@/components/networking", () => ({ serverRootPath: "/" })); + const { MIGRATED_PAGES } = await import("./migratedPages"); + + expect(MIGRATED_PAGES.caching).toBe("caching"); + expect(MIGRATED_PAGES["cost-tracking"]).toBe("cost-tracking"); + expect(MIGRATED_PAGES["transform-request"]).toBe("transform-request"); + expect(MIGRATED_PAGES["ui-theme"]).toBe("ui-theme"); + expect(MIGRATED_PAGES.logs).toBe("logs"); + }); }); describe("dev server (NODE_ENV=development)", () => { diff --git a/ui/litellm-dashboard/src/utils/migratedPages.ts b/ui/litellm-dashboard/src/utils/migratedPages.ts index f51a7af73c2..c54b0473c02 100644 --- a/ui/litellm-dashboard/src/utils/migratedPages.ts +++ b/ui/litellm-dashboard/src/utils/migratedPages.ts @@ -30,6 +30,11 @@ export const MIGRATED_PAGES: Record = { skills: "skills", // Legacy alias: the old switch matched ?page=claude-code-plugins for the same panel. "claude-code-plugins": "skills", + caching: "caching", + "cost-tracking": "cost-tracking", + "transform-request": "transform-request", + "ui-theme": "ui-theme", + logs: "logs", }; function uiBase(): string { From 76b4c4b1118b2b4e7abca529ea35907548e647c8 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 12 Jun 2026 15:35:48 -0700 Subject: [PATCH 090/185] fix(ui): gate dashboard layout on ui config load so deep links work under SERVER_ROOT_PATH (#30312) * fix(ui): gate dashboard layout on ui config load so deep links work under SERVER_ROOT_PATH * test(ui): create ui config deferred per test so the pending state stays repeatable --- .../src/app/(dashboard)/layout.test.tsx | 78 +++++++++++++++++++ .../src/app/(dashboard)/layout.tsx | 6 +- 2 files changed, 83 insertions(+), 1 deletion(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx diff --git a/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx new file mode 100644 index 00000000000..af68d9f87e9 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx @@ -0,0 +1,78 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { render, screen, waitFor } from "@testing-library/react"; +import { AuthProvider } from "@/contexts/AuthContext"; +import Layout from "./layout"; + +vi.mock("next/navigation", () => ({ + useRouter: vi.fn(() => ({ push: vi.fn(), replace: vi.fn() })), + useSearchParams: vi.fn(() => new URLSearchParams()), + usePathname: vi.fn(() => "/ui/guardrails"), +})); + +vi.mock("@/components/navbar", () => ({ + default: () =>
, +})); + +vi.mock("@/app/(dashboard)/components/SidebarProvider", () => ({ + default: () =>
, +})); + +vi.mock("@/components/DebugWarningBanner", () => ({ + DebugWarningBanner: () => null, +})); + +vi.mock("@/contexts/ThemeContext", () => ({ + ThemeProvider: ({ children }: { children: React.ReactNode }) => <>{children}, +})); + +vi.mock("@/components/common_components/LoadingScreen", () => ({ + default: () =>
, +})); + +type Deferred = { promise: Promise; resolve: () => void }; + +const createDeferred = (): Deferred => { + let resolve!: () => void; + const promise = new Promise((r) => { + resolve = r; + }); + return { promise, resolve }; +}; + +let pendingUiConfig: Deferred; + +vi.mock("@/components/networking", async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + getUiConfig: vi.fn(() => pendingUiConfig.promise), + setGlobalLitellmHeaderName: vi.fn(), + }; +}); + +describe("(dashboard) Layout", () => { + beforeEach(() => { + vi.clearAllMocks(); + pendingUiConfig = createDeferred(); + }); + + it("does not mount route content until getUiConfig has resolved", async () => { + render( + + +
+ + , + ); + + await waitFor(() => expect(screen.getByTestId("loading-screen")).toBeTruthy()); + expect(screen.queryByTestId("page-content")).toBeNull(); + expect(screen.queryByTestId("navbar")).toBeNull(); + + pendingUiConfig.resolve(); + + await waitFor(() => expect(screen.getByTestId("page-content")).toBeTruthy()); + expect(screen.getByTestId("navbar")).toBeTruthy(); + expect(screen.queryByTestId("loading-screen")).toBeNull(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx index df5b2ab4511..b32bed44a87 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/layout.tsx @@ -45,9 +45,13 @@ function DashboardShell({ children }: { children: React.ReactNode }) { function LayoutContent({ children }: { children: React.ReactNode }) { const searchParams = useSearchParams(); - const { accessToken } = useAuth(); + const { accessToken, authLoading } = useAuth(); const isInvitationFlow = Boolean(searchParams.get("invitation_id")); + if (authLoading) { + return ; + } + return ( {isInvitationFlow ? children : {children}} From d258e022d18d702216140dc8e4d9aab434ce9f31 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 12 Jun 2026 16:16:27 -0700 Subject: [PATCH 091/185] feat(ui): cut admin-panel, logging-and-alerts, model-hub-table, and usage over to path routes (#30268) admin-panel pulls proxySettings from the shared useProxySettings query hook (dropping the last reader of the legacy page's copy), the model hub wrapper keeps the admin-vs-public branch as an early return, and the usage wrapper feeds NewUsagePage from the useTeams and useOrganizations query hooks instead of the lifted switch state. new_usage maps to the /usage segment while the old ?page=usage report keeps its legacy arm, asserted in the unit test so the two cannot be confused. --- .../e2e_tests/fixtures/migratedPages.ts | 6 +++-- .../src/app/(dashboard)/admin-panel/page.tsx | 11 ++++++++ .../(dashboard)/logging-and-alerts/page.tsx | 9 +++++++ .../app/(dashboard)/model-hub-table/page.tsx | 14 +++++++++++ .../src/app/(dashboard)/page.tsx | 25 ------------------- .../src/app/(dashboard)/usage/page.tsx | 13 ++++++++++ .../src/utils/migratedPages.test.ts | 12 +++++++++ .../src/utils/migratedPages.ts | 5 ++++ 8 files changed, 68 insertions(+), 27 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/admin-panel/page.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/logging-and-alerts/page.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/model-hub-table/page.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/usage/page.tsx diff --git a/ui/litellm-dashboard/e2e_tests/fixtures/migratedPages.ts b/ui/litellm-dashboard/e2e_tests/fixtures/migratedPages.ts index 19b90848e60..d0dde5a8155 100644 --- a/ui/litellm-dashboard/e2e_tests/fixtures/migratedPages.ts +++ b/ui/litellm-dashboard/e2e_tests/fixtures/migratedPages.ts @@ -9,8 +9,6 @@ * - navigation specs that assert per-page URLs (tests/navigation/sidebar.spec.ts) * * Keep this in lockstep with MIGRATED_PAGES in src/utils/migratedPages.ts. - * Pending (add as each PR lands): admin-panel, logging-and-alerts, - * model-hub-table, and usage (#30268). */ export const MIGRATED_E2E_PAGES: Record = { api_ref: "api-reference", @@ -35,6 +33,10 @@ export const MIGRATED_E2E_PAGES: Record = { "transform-request": "transform-request", "ui-theme": "ui-theme", logs: "logs", + "admin-panel": "admin-panel", + "logging-and-alerts": "logging-and-alerts", + "model-hub-table": "model-hub-table", + new_usage: "usage", }; export const MIGRATED_E2E_SEGMENTS: string[] = [...new Set(Object.values(MIGRATED_E2E_PAGES))]; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/admin-panel/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/admin-panel/page.tsx new file mode 100644 index 00000000000..aac835b02fc --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/admin-panel/page.tsx @@ -0,0 +1,11 @@ +"use client"; + +import AdminPanel from "@/components/AdminPanel"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import useProxySettings from "@/app/(dashboard)/hooks/proxySettings/useProxySettings"; + +export default function AdminPanelPage() { + const { accessToken } = useAuthorized(); + const proxySettings = useProxySettings(accessToken); + return ; +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/logging-and-alerts/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/logging-and-alerts/page.tsx new file mode 100644 index 00000000000..8232e391259 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/logging-and-alerts/page.tsx @@ -0,0 +1,9 @@ +"use client"; + +import Settings from "@/components/settings"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; + +export default function LoggingAndAlerts() { + const { accessToken, userRole, userId, premiumUser } = useAuthorized(); + return ; +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/model-hub-table/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/model-hub-table/page.tsx new file mode 100644 index 00000000000..7327d332fbd --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/model-hub-table/page.tsx @@ -0,0 +1,14 @@ +"use client"; + +import ModelHubTable from "@/components/AIHub/ModelHubTable"; +import PublicModelHub from "@/components/public_model_hub"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { isAdminRole } from "@/utils/roles"; + +export default function ModelHubTablePage() { + const { accessToken, userRole, premiumUser } = useAuthorized(); + if (!isAdminRole(userRole)) { + return ; + } + return ; +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/page.tsx index 864007b4fe8..45ec0f62357 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/page.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/page.tsx @@ -1,23 +1,17 @@ "use client"; import ModelsAndEndpointsView from "@/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView"; -import AdminPanel from "@/components/AdminPanel"; import AgentsPanel from "@/components/agents"; import { teamListCall as v2TeamListCall } from "@/app/(dashboard)/hooks/teams/useTeams"; import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings"; -import useProxySettings from "@/app/(dashboard)/hooks/proxySettings/useProxySettings"; import LoadingScreen from "@/components/common_components/LoadingScreen"; import GeneralSettings from "@/components/general_settings"; import { Team } from "@/components/key_team_helpers/key_list"; -import ModelHubTable from "@/components/AIHub/ModelHubTable"; import { Organization, proxyBaseUrl, getInProductNudgesCall } from "@/components/networking"; -import NewUsagePage from "@/components/UsagePage/components/UsagePageView"; import OldTeams from "@/components/OldTeams"; import { fetchUserModels, CreateKeyPrefillData } from "@/components/organisms/create_key_button"; import Organizations, { fetchOrganizations } from "@/components/organizations"; import PassThroughSettings from "@/components/pass_through_settings"; -import PublicModelHub from "@/components/public_model_hub"; -import Settings from "@/components/settings"; import { SurveyPrompt, SurveyModal, ClaudeCodePrompt, ClaudeCodeModal } from "@/components/survey"; import Usage from "@/components/usage"; import UserDashboard from "@/components/user_dashboard"; @@ -30,7 +24,6 @@ import { normalizeUrlForCompare, storeReturnUrl, } from "@/utils/returnUrlUtils"; -import { isAdminRole } from "@/utils/roles"; import { MIGRATED_PAGES, migratedHref } from "@/utils/migratedPages"; import { useRouter, useSearchParams } from "next/navigation"; import { Suspense, useEffect, useMemo, useRef, useState } from "react"; @@ -43,7 +36,6 @@ function CreateKeyPageContent() { const [keys, setKeys] = useState([]); const [organizations, setOrganizations] = useState([]); const [userModels, setUserModels] = useState([]); - const proxySettings = useProxySettings(accessToken); const router = useRouter(); const searchParams = useSearchParams()!; @@ -363,25 +355,10 @@ function CreateKeyPageContent() { userRole={userRole} premiumUser={premiumUser} /> - ) : page == "admin-panel" ? ( - - ) : page == "logging-and-alerts" ? ( - ) : page == "agents" ? ( ) : page == "router-settings" ? ( - ) : page == "model-hub-table" ? ( - isAdminRole(userRole) ? ( - - ) : ( - - ) ) : page == "pass-through-settings" ? ( - ) : page == "new_usage" ? ( - ) : ( ; +} diff --git a/ui/litellm-dashboard/src/utils/migratedPages.test.ts b/ui/litellm-dashboard/src/utils/migratedPages.test.ts index 1183c81d05f..7e74fa1eec4 100644 --- a/ui/litellm-dashboard/src/utils/migratedPages.test.ts +++ b/ui/litellm-dashboard/src/utils/migratedPages.test.ts @@ -99,6 +99,18 @@ describe("migratedHref / legacyPageHref", () => { expect(MIGRATED_PAGES["ui-theme"]).toBe("ui-theme"); expect(MIGRATED_PAGES.logs).toBe("logs"); }); + + it("maps the admin-panel, logging-and-alerts, model-hub-table, and new_usage ids to their routes", async () => { + vi.doMock("@/components/networking", () => ({ serverRootPath: "/" })); + const { MIGRATED_PAGES } = await import("./migratedPages"); + + expect(MIGRATED_PAGES["admin-panel"]).toBe("admin-panel"); + expect(MIGRATED_PAGES["logging-and-alerts"]).toBe("logging-and-alerts"); + expect(MIGRATED_PAGES["model-hub-table"]).toBe("model-hub-table"); + // new_usage routes to /usage; the legacy ?page=usage report keeps its switch arm. + expect(MIGRATED_PAGES.new_usage).toBe("usage"); + expect(MIGRATED_PAGES.usage).toBeUndefined(); + }); }); describe("dev server (NODE_ENV=development)", () => { diff --git a/ui/litellm-dashboard/src/utils/migratedPages.ts b/ui/litellm-dashboard/src/utils/migratedPages.ts index c54b0473c02..46cdd0d7476 100644 --- a/ui/litellm-dashboard/src/utils/migratedPages.ts +++ b/ui/litellm-dashboard/src/utils/migratedPages.ts @@ -35,6 +35,11 @@ export const MIGRATED_PAGES: Record = { "transform-request": "transform-request", "ui-theme": "ui-theme", logs: "logs", + "admin-panel": "admin-panel", + "logging-and-alerts": "logging-and-alerts", + "model-hub-table": "model-hub-table", + // The modern usage dashboard; the old ?page=usage report stays on the legacy switch. + new_usage: "usage", }; function uiBase(): string { From f49707bc66ff1ec3e9c8c72a0f15dc3d4a10bfa5 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Fri, 12 Jun 2026 17:29:46 -0700 Subject: [PATCH 092/185] fix(otel): cap metric attribute cardinality with include/exclude lists (#30257) * fix(otel): cap metric attribute cardinality with include/exclude lists OTEL metrics stamped every per-request hidden_params and metadata.* field onto each gen_ai.client.* sample, so near-unique values created one metric time series per request and backends like Splunk Observability Cloud throttled and dropped the data. Add an attributes block under callback_settings.otel with mutually-exclusive include_list (allowlist) and exclude_list (denylist), validated against the known attribute names at startup and applied once to the metric attributes in _record_metrics. Spans are untouched, and with no config every attribute is still emitted so existing setups are unaffected. Resolves LIT-3600 * fix(otel): resolve metric attribute filter from callback_settings The proxy usually constructs the OpenTelemetry logger without forwarding the attributes kwarg, while the filter lives under litellm.callback_settings["otel"]["attributes"]. __init__ only read the kwarg, so the recording instance kept config.attributes=None and shipped metrics at full cardinality even when the filter was configured; a live proxy run exposed this. Fall back to the global at init for the base otel logger, and add a regression test that drives the real success hook through the callback_settings path (the unit tests passed before because they injected the config directly). * fix(otel): reject gen_ai.token.type from metric attribute filter lists gen_ai.token.type was a member of VALID_METRIC_ATTRIBUTE_NAMES, so an operator could list it in include_list or exclude_list and pass startup validation. The attribute is injected into the input/output token series after _filter_metric_attributes runs, so the filter never sees it and the request silently has no effect. Reject it loudly from either list instead, matching the contract that a non-actionable attribute name fails fast rather than falling through to a no-op. It stays a structural discriminator on the token-usage histogram. * fix(otel): resolve metric attribute filter lazily at record time The proxy constructs the OpenTelemetry logger before it populates litellm.callback_settings["otel"]["attributes"], so resolving the filter at __init__ left config.attributes None and shipped metrics at full cardinality. A live proxy run confirmed the leak. Resolve the filter on the first metric record instead, when callback_settings is populated, while still validating an explicit config eagerly so a bad SDK config fails at startup. The regression test now constructs the logger before populating callback_settings to mirror that ordering, so it fails if the filter is resolved too early. * fix(otel): don't cache invalid filter on lazy callback_settings path On the lazy callback_settings resolution path, _ensure_metric_attribute_filter wrote self.config.attributes before validating it. When validation then failed, _metric_attr_filter_resolved stayed False while config.attributes held the bad filter, so the next record skipped the callback_settings re-read and re-raised the stale error indefinitely; fixing the misconfiguration required a restart. Drop the premature write and resolve from the local value. A subsequent record now re-reads callback_settings, so a corrected config takes effect without a restart. The write was dead on the success path anyway, since the resolved frozensets are what the filter reads. --- litellm/integrations/opentelemetry.py | 168 +++++++++-- .../integrations/test_opentelemetry.py | 282 ++++++++++++++++++ 2 files changed, 430 insertions(+), 20 deletions(-) diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index 24780eb4bfc..fc37b6a34d8 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -1,7 +1,18 @@ import os from dataclasses import dataclass, field from datetime import datetime -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Union, cast +from typing import ( + TYPE_CHECKING, + Any, + Dict, + FrozenSet, + List, + Optional, + Set, + Tuple, + Union, + cast, +) import litellm from litellm._logging import verbose_logger @@ -82,6 +93,88 @@ _VALID_CAPTURE_MODES = { CAPTURE_MODE_SPAN_AND_EVENT, } +METRIC_METADATA_KEYS: Tuple[str, ...] = ( + "user_api_key_hash", + "user_api_key_alias", + "user_api_key_team_id", + "user_api_key_org_id", + "user_api_key_user_id", + "user_api_key_team_alias", + "user_api_key_user_email", + "spend_logs_metadata", + "requester_ip_address", + "requester_metadata", + "user_api_key_end_user_id", + "prompt_management_metadata", + "applied_guardrails", + "mcp_tool_call_metadata", + "vector_store_request_metadata", +) + +TOKEN_TYPE_ATTRIBUTE: str = "gen_ai.token.type" + +VALID_METRIC_ATTRIBUTE_NAMES: FrozenSet[str] = frozenset( + ( + "gen_ai.operation.name", + "gen_ai.system", + "gen_ai.request.model", + "gen_ai.framework", + "hidden_params", + ) + + tuple(f"metadata.{key}" for key in METRIC_METADATA_KEYS) +) + + +@dataclass(frozen=True) +class OTELMetricAttributeFilter: + include_list: Optional[List[str]] = None + exclude_list: Optional[List[str]] = None + + +def _build_metric_attribute_filter(value: Any) -> OTELMetricAttributeFilter: + if isinstance(value, OTELMetricAttributeFilter): + return value + if not isinstance(value, dict): + raise ValueError( + "otel.attributes must be a mapping with optional 'include_list' / " + f"'exclude_list', got {type(value).__name__}" + ) + return OTELMetricAttributeFilter( + include_list=value.get("include_list"), + exclude_list=value.get("exclude_list"), + ) + + +def _resolve_metric_attribute_filter( + attributes: Optional[OTELMetricAttributeFilter], +) -> Tuple[Optional[FrozenSet[str]], Optional[FrozenSet[str]]]: + if attributes is None: + return None, None + include = attributes.include_list or None + exclude = attributes.exclude_list or None + if include and exclude: + raise ValueError( + "otel.attributes: include_list and exclude_list are mutually exclusive" + ) + requested = include or exclude or [] + if TOKEN_TYPE_ATTRIBUTE in requested: + raise ValueError( + f"otel.attributes: {TOKEN_TYPE_ATTRIBUTE} is a structural token-usage " + "discriminator and cannot be filtered" + ) + unknown = sorted( + name for name in requested if name not in VALID_METRIC_ATTRIBUTE_NAMES + ) + if unknown: + raise ValueError( + f"otel.attributes: unknown attribute name(s) {unknown}. " + f"Valid names: {sorted(VALID_METRIC_ATTRIBUTE_NAMES)}" + ) + return ( + frozenset(include) if include else None, + frozenset(exclude) if exclude else None, + ) + def _normalize_team_metadata_keys(value: Any) -> List[str]: """Coerce a team-metadata allowlist from a list or comma-separated string. @@ -117,6 +210,9 @@ class OpenTelemetryConfig: # under ``litellm.team.metadata``. Empty by default so none of a team's # metadata leaves the process until explicitly allowlisted. baggage_team_metadata_keys: List[str] = field(default_factory=list) + # Prometheus-style include/exclude control over which attributes are stamped + # on emitted metrics, to cap metric cardinality. + attributes: Optional[OTELMetricAttributeFilter] = None def __post_init__(self) -> None: # If endpoint is specified but exporter is still the default "console", @@ -211,15 +307,29 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): **kwargs, ): team_metadata_keys_override = kwargs.pop("baggage_team_metadata_keys", None) + metric_attributes_override = kwargs.pop("attributes", None) if config is None: config = OpenTelemetryConfig.from_env() if team_metadata_keys_override is not None: config.baggage_team_metadata_keys = _normalize_team_metadata_keys( team_metadata_keys_override ) + if metric_attributes_override is not None: + config.attributes = _build_metric_attribute_filter( + metric_attributes_override + ) self.config = config self.callback_name = callback_name + # Resolved on first metric record, not here: the proxy populates + # callback_settings.otel.attributes after this logger is constructed, so + # reading it now would miss it. An explicit config is validated eagerly so + # a bad config still fails at startup. + self._metric_attr_include: Optional[FrozenSet[str]] = None + self._metric_attr_exclude: Optional[FrozenSet[str]] = None + self._metric_attr_filter_resolved = False + if config.attributes is not None: + self._ensure_metric_attribute_filter() self.OTEL_EXPORTER = self.config.exporter self.OTEL_ENDPOINT = self.config.endpoint self.OTEL_HEADERS = self.config.headers @@ -1318,6 +1428,38 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): return None return safe_dumps(filtered) + def _ensure_metric_attribute_filter(self) -> None: + """Resolve the include/exclude filter once, falling back to the proxy's + callback_settings.otel.attributes when no explicit config was passed.""" + if self._metric_attr_filter_resolved: + return + attributes = self.config.attributes + if attributes is None and self.callback_name in (None, "otel"): + otel_settings = (litellm.callback_settings or {}).get("otel") or {} + raw = ( + otel_settings.get("attributes") + if isinstance(otel_settings, dict) + else None + ) + if raw is not None: + attributes = _build_metric_attribute_filter(raw) + ( + self._metric_attr_include, + self._metric_attr_exclude, + ) = _resolve_metric_attribute_filter(attributes) + self._metric_attr_filter_resolved = True + + def _filter_metric_attributes(self, attrs: Dict[str, Any]) -> Dict[str, Any]: + if not self._metric_attr_filter_resolved: + self._ensure_metric_attribute_filter() + if self._metric_attr_include is not None: + return {k: v for k, v in attrs.items() if k in self._metric_attr_include} + if self._metric_attr_exclude is not None: + return { + k: v for k, v in attrs.items() if k not in self._metric_attr_exclude + } + return attrs + def _record_metrics(self, kwargs, response_obj, start_time, end_time): duration_s = (end_time - start_time).total_seconds() params = kwargs.get("litellm_params") or {} @@ -1336,23 +1478,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): std_log = kwargs.get("standard_logging_object") md = getattr(std_log, "metadata", None) or (std_log or {}).get("metadata", {}) - for key in [ - "user_api_key_hash", - "user_api_key_alias", - "user_api_key_team_id", - "user_api_key_org_id", - "user_api_key_user_id", - "user_api_key_team_alias", - "user_api_key_user_email", - "spend_logs_metadata", - "requester_ip_address", - "requester_metadata", - "user_api_key_end_user_id", - "prompt_management_metadata", - "applied_guardrails", - "mcp_tool_call_metadata", - "vector_store_request_metadata", - ]: + for key in METRIC_METADATA_KEYS: value = md.get(key) if value is None: continue @@ -1368,6 +1494,8 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): if hidden_params: common_attrs["hidden_params"] = safe_dumps(hidden_params) + common_attrs = self._filter_metric_attributes(common_attrs) + if self._operation_duration_histogram: self._operation_duration_histogram.record( duration_s, attributes=common_attrs @@ -1377,8 +1505,8 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): and (usage := response_obj.get("usage")) and self._token_usage_histogram ): - in_attrs = {**common_attrs, "gen_ai.token.type": "input"} - out_attrs = {**common_attrs, "gen_ai.token.type": "output"} + in_attrs = {**common_attrs, TOKEN_TYPE_ATTRIBUTE: "input"} + out_attrs = {**common_attrs, TOKEN_TYPE_ATTRIBUTE: "output"} self._token_usage_histogram.record( usage.get("prompt_tokens", 0), attributes=in_attrs ) diff --git a/tests/test_litellm/integrations/test_opentelemetry.py b/tests/test_litellm/integrations/test_opentelemetry.py index 0601f9c0eef..e47e437a131 100644 --- a/tests/test_litellm/integrations/test_opentelemetry.py +++ b/tests/test_litellm/integrations/test_opentelemetry.py @@ -19,9 +19,11 @@ from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +import litellm from litellm.integrations.opentelemetry import ( OpenTelemetry, OpenTelemetryConfig, + OTELMetricAttributeFilter, OTELSemconvCategory, _normalize_team_metadata_keys, ) @@ -5301,6 +5303,8 @@ class TestEndProxySpanLitellmMetadataFallback(unittest.TestCase): otel._end_proxy_span_from_kwargs(kwargs, end_time=datetime.now()) mock_span.end.assert_called_once() + + class TestOpenTelemetryInferenceIdentityAttributes(unittest.TestCase): """team_metadata, http.route, and both model names (the user-facing model_group alias and the dispatched provider model) must land on the @@ -5467,3 +5471,281 @@ class TestOpenTelemetryTeamMetadataKeysConfig(unittest.TestCase): ): cfg = OpenTelemetryConfig(baggage_team_metadata_keys=["from_arg"]) assert cfg.baggage_team_metadata_keys == ["from_arg"] + + +class TestOpenTelemetryMetricAttributeFiltering(unittest.TestCase): + """LIT-3600: include/exclude control over which attributes are stamped on + emitted metrics, to cap metric cardinality. These drive the real + _handle_success -> _record_metrics path through an in-memory reader and + read attributes straight off the recorded data points, so they fail if the + filtering feature is reverted and pass only when it works end to end.""" + + HERE = os.path.dirname(__file__) + POLL_INTERVAL = 0.05 + POLL_TIMEOUT = 2.0 + DURATION_METRIC = "gen_ai.client.operation.duration" + TOKEN_METRIC = "gen_ai.client.token.usage" + + # High-cardinality attributes the captured fixture emits by default. Each is + # a member of VALID_METRIC_ATTRIBUTE_NAMES and is present on the recorded + # metric when no filter is configured (verified by the backward-compat test). + HIGH_CARDINALITY_KEYS = ( + "hidden_params", + "metadata.user_api_key_hash", + "metadata.requester_ip_address", + "metadata.requester_metadata", + "metadata.applied_guardrails", + ) + RETAINED_LOW_CARDINALITY_KEY = "gen_ai.request.model" + + def _load_fixtures(self): + with open( + os.path.join(self.HERE, "open_telemetry", "data", "captured_kwargs.json") + ) as f: + kwargs = json.load(f) + with open( + os.path.join(self.HERE, "open_telemetry", "data", "captured_response.json") + ) as f: + response_obj = json.load(f) + return kwargs, response_obj + + def _record(self, attributes): + """Run a real success hook with metrics enabled and return the reader.""" + metric_reader = InMemoryMetricReader() + meter_provider = MeterProvider(metric_readers=[metric_reader]) + tracer_provider = TracerProvider() + tracer_provider.add_span_processor(SimpleSpanProcessor(InMemorySpanExporter())) + otel = OpenTelemetry( + config=OpenTelemetryConfig( + exporter="console", enable_metrics=True, attributes=attributes + ), + tracer_provider=tracer_provider, + meter_provider=meter_provider, + ) + otel.tracer = tracer_provider.get_tracer(__name__) + + kwargs, response_obj = self._load_fixtures() + start = datetime.utcnow() + end = start + timedelta(seconds=1) + otel._handle_success(kwargs, response_obj, start, end) + return metric_reader + + def _keysets(self, reader, metric_name): + """Attribute-key sets, one per recorded data point of `metric_name`.""" + deadline = time.time() + self.POLL_TIMEOUT + while time.time() < deadline: + data = reader.get_metrics_data() + if data and hasattr(data, "resource_metrics"): + for rm in data.resource_metrics: + for sm in rm.scope_metrics: + for m in sm.metrics: + if m.name == metric_name: + return [ + set(dp.attributes.keys()) + for dp in m.data.data_points + ] + time.sleep(self.POLL_INTERVAL) + return None + + def test_exclude_list_strips_high_cardinality_keys_across_metrics(self): + """The bug: high-cardinality metadata/hidden_params explode metric + cardinality. With exclude_list set, none of them reach any data point, + while the retained low-cardinality model attribute survives. Asserted + on both the duration and token-usage histograms.""" + reader = self._record( + OTELMetricAttributeFilter(exclude_list=list(self.HIGH_CARDINALITY_KEYS)) + ) + excluded = set(self.HIGH_CARDINALITY_KEYS) + + for metric_name in (self.DURATION_METRIC, self.TOKEN_METRIC): + keysets = self._keysets(reader, metric_name) + self.assertTrue(keysets, f"{metric_name} was not recorded") + for keys in keysets: + self.assertTrue( + excluded.isdisjoint(keys), + f"{metric_name} leaked excluded keys: {excluded & keys}", + ) + self.assertIn(self.RETAINED_LOW_CARDINALITY_KEY, keys) + + def test_include_list_allows_only_listed_attributes(self): + """An allowlist caps emitted attributes to exactly the listed set. + gen_ai.token.type is a structural discriminator added to the token + histogram after filtering, so it is the only key permitted beyond the + allowlist, and only on that metric.""" + include = ["gen_ai.request.model", "gen_ai.system"] + reader = self._record(OTELMetricAttributeFilter(include_list=include)) + allowed = set(include) + + duration_keysets = self._keysets(reader, self.DURATION_METRIC) + self.assertTrue(duration_keysets, "duration metric was not recorded") + for keys in duration_keysets: + self.assertEqual(keys, allowed) + + token_keysets = self._keysets(reader, self.TOKEN_METRIC) + self.assertTrue(token_keysets, "token-usage metric was not recorded") + for keys in token_keysets: + self.assertEqual(keys - {"gen_ai.token.type"}, allowed) + + def test_no_filter_preserves_high_cardinality_keys(self): + """Backward compatibility: with no attributes config, every + high-cardinality key the fixture carries is still stamped on the + metric, so existing customers who rely on them are unaffected.""" + reader = self._record(None) + expected = set(self.HIGH_CARDINALITY_KEYS) + + for metric_name in (self.DURATION_METRIC, self.TOKEN_METRIC): + keysets = self._keysets(reader, metric_name) + self.assertTrue(keysets, f"{metric_name} was not recorded") + for keys in keysets: + self.assertTrue( + expected.issubset(keys), + f"{metric_name} dropped {expected - keys} by default", + ) + self.assertIn(self.RETAINED_LOW_CARDINALITY_KEY, keys) + + def test_proxy_callback_settings_attributes_applied_without_kwarg(self): + """Regression for the proxy path: the OpenTelemetry logger is constructed + before the proxy populates litellm.callback_settings['otel']['attributes'], + and without the attributes kwarg, so the filter must be resolved at record + time rather than at __init__. Otherwise metrics ship at full cardinality + (the bug the live proxy surfaced; constructing with the kwarg, or with + callback_settings already set, hid it).""" + previous = litellm.callback_settings + litellm.callback_settings = {} # not yet populated when the logger is built + try: + metric_reader = InMemoryMetricReader() + meter_provider = MeterProvider(metric_readers=[metric_reader]) + tracer_provider = TracerProvider() + tracer_provider.add_span_processor( + SimpleSpanProcessor(InMemorySpanExporter()) + ) + otel = OpenTelemetry( + config=OpenTelemetryConfig(exporter="console", enable_metrics=True), + tracer_provider=tracer_provider, + meter_provider=meter_provider, + ) + otel.tracer = tracer_provider.get_tracer(__name__) + # The proxy sets this only after the logger already exists. + litellm.callback_settings = { + "otel": { + "attributes": {"exclude_list": list(self.HIGH_CARDINALITY_KEYS)} + } + } + kwargs, response_obj = self._load_fixtures() + start = datetime.utcnow() + otel._handle_success( + kwargs, response_obj, start, start + timedelta(seconds=1) + ) + finally: + litellm.callback_settings = previous + + excluded = set(self.HIGH_CARDINALITY_KEYS) + for metric_name in (self.DURATION_METRIC, self.TOKEN_METRIC): + keysets = self._keysets(metric_reader, metric_name) + self.assertTrue(keysets, f"{metric_name} was not recorded") + for keys in keysets: + self.assertTrue( + excluded.isdisjoint(keys), + f"{metric_name} leaked {excluded & keys} via callback_settings", + ) + self.assertIn(self.RETAINED_LOW_CARDINALITY_KEY, keys) + + def test_callback_settings_validation_failure_is_not_sticky(self): + """On the lazy callback_settings path a validation failure must not cache + the bad config. Once the operator corrects + callback_settings['otel']['attributes'], the next record resolves the + fixed filter instead of re-raising the stale error until a restart.""" + previous = litellm.callback_settings + litellm.callback_settings = { + "otel": { + "attributes": { + "include_list": ["gen_ai.system"], + "exclude_list": ["hidden_params"], + } + } + } + try: + otel = OpenTelemetry(config=OpenTelemetryConfig(exporter="console")) + attrs = {"gen_ai.system": "openai", "hidden_params": "{}"} + + with self.assertRaises(ValueError): + otel._filter_metric_attributes(attrs) + + litellm.callback_settings = { + "otel": {"attributes": {"exclude_list": ["hidden_params"]}} + } + filtered = otel._filter_metric_attributes(attrs) + finally: + litellm.callback_settings = previous + + self.assertEqual(filtered, {"gen_ai.system": "openai"}) + + def test_include_and_exclude_together_raise_value_error(self): + with self.assertRaises(ValueError): + OpenTelemetry( + config=OpenTelemetryConfig( + exporter="console", + attributes=OTELMetricAttributeFilter( + include_list=["gen_ai.system"], + exclude_list=["hidden_params"], + ), + ) + ) + + def test_unknown_include_name_raises_value_error(self): + with self.assertRaises(ValueError): + OpenTelemetry( + config=OpenTelemetryConfig( + exporter="console", + attributes=OTELMetricAttributeFilter( + include_list=["not.a.real.attribute"] + ), + ) + ) + + def test_unknown_exclude_name_raises_value_error(self): + with self.assertRaises(ValueError): + OpenTelemetry( + config=OpenTelemetryConfig( + exporter="console", + attributes=OTELMetricAttributeFilter( + exclude_list=["metadata.does_not_exist"] + ), + ) + ) + + def test_dict_attributes_kwarg_path_validates(self): + """The YAML/kwargs entry point (a plain dict) flows through + _build_metric_attribute_filter and hits the same validation.""" + with self.assertRaises(ValueError): + OpenTelemetry( + attributes={ + "include_list": ["gen_ai.system"], + "exclude_list": ["hidden_params"], + } + ) + + def test_no_filter_returns_attrs_object_unchanged(self): + """The no-config path is a hot-path no-op: it returns the same dict + object, so default emission pays zero copy cost. Locking identity makes + a future refactor that always copies/filters trip here.""" + otel = OpenTelemetry(config=OpenTelemetryConfig(exporter="console")) + attrs = {"gen_ai.request.model": "m", "hidden_params": "{}"} + self.assertIs(otel._filter_metric_attributes(attrs), attrs) + + def test_token_type_discriminator_rejected_from_either_list(self): + """gen_ai.token.type is a structural discriminator stamped onto the + input/output token series after filtering; it cannot be filtered without + collapsing the two series into one. Listing it in include_list or + exclude_list is rejected loudly at startup rather than silently ignored, + so an operator gets an error instead of a no-op.""" + for attributes in ( + OTELMetricAttributeFilter(exclude_list=["gen_ai.token.type"]), + OTELMetricAttributeFilter(include_list=["gen_ai.token.type"]), + ): + with self.assertRaises(ValueError): + OpenTelemetry( + config=OpenTelemetryConfig( + exporter="console", attributes=attributes + ) + ) From 5047eaf7f0e7151891d7edbf19f92eb0004ff274 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 12 Jun 2026 17:44:04 -0700 Subject: [PATCH 093/185] fix(proxy): return deprecated-key lookup result directly in get_data combined view (#30327) The grace-period branch assigned the recursive get_data result (a finished LiteLLM_VerificationTokenView) back into the variable that the combined-view dict normalization then subscripts, raising TypeError on every request made with a rotated key inside its grace window; auth surfaced that as a 401. Return the recursive result directly instead. Regression test drives the full get_data flow: old hash misses the view, deprecated table resolves to the active token, and the call must return the view object --- litellm/proxy/utils.py | 8 +++- .../test_prisma_client_get_data.py | 38 +++++++++++++++++++ 2 files changed, 44 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 4aa555164b0..98d57229a52 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -3692,7 +3692,10 @@ class PrismaClient: db=self.db, hashed_token=hashed_token ) if active_token_id: - response = await self.get_data( + # The recursive call returns a finished + # LiteLLM_VerificationTokenView; the dict + # normalization below would crash subscripting it. + deprecated_response = await self.get_data( token=active_token_id, table_name="combined_view", query_type="find_unique", @@ -3700,10 +3703,11 @@ class PrismaClient: proxy_logging_obj=proxy_logging_obj, check_deprecated=False, ) - if response is not None: + if deprecated_response is not None: verbose_proxy_logger.debug( "Deprecated key used during grace period" ) + return deprecated_response if response is not None: if response["team_models"] is None: diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py index 437984d9273..08d1ef619a7 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py @@ -15,6 +15,7 @@ from __future__ import annotations import hashlib import json +from datetime import datetime, timedelta, timezone from types import SimpleNamespace from typing import Any from unittest.mock import AsyncMock, MagicMock @@ -22,6 +23,7 @@ from unittest.mock import AsyncMock, MagicMock import pytest from fastapi import HTTPException +from litellm.proxy._types import LiteLLM_VerificationTokenView from litellm.proxy.utils import PrismaClient @@ -476,3 +478,39 @@ async def test_get_data_logs_and_raises_on_db_error( ) with pytest.raises(RuntimeError, match="network split"): await prisma_client.get_data(token="sk-broken", table_name="key") + + +@pytest.mark.asyncio +async def test_get_data_combined_view_returns_view_for_deprecated_key( + prisma_client: PrismaClient, +) -> None: + """Grace-period rotation, full get_data flow: the old hash misses the + combined view, the deprecated-key table resolves it to the active token, + and get_data must return the recursive lookup's finished view instead of + re-running dict normalization on it (which raised TypeError and turned + every grace-period request into a 401).""" + old_hash = "hashed-old-token-grace-e2e" + active_hash = "hashed-active-token-grace-e2e" + active_row = { + "token": active_hash, + "team_models": None, + "team_blocked": None, + "team_members_with_roles": None, + "user_id": None, + "expires": None, + } + prisma_client.db.query_first = AsyncMock(side_effect=[None, active_row]) + prisma_client.db.litellm_deprecatedverificationtoken = MagicMock() + prisma_client.db.litellm_deprecatedverificationtoken.find_first = AsyncMock( + return_value=SimpleNamespace( + active_token_id=active_hash, + revoke_at=datetime.now(timezone.utc) + timedelta(hours=1), + ) + ) + + response = await prisma_client.get_data( + token=old_hash, table_name="combined_view", query_type="find_unique" + ) + + assert isinstance(response, LiteLLM_VerificationTokenView) + assert response.token == active_hash From d96ab467f1dba5e4dbe02de3d5af62ec710c44fd Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 12 Jun 2026 17:48:00 -0700 Subject: [PATCH 094/185] chore(deps): bump vitest, brace-expansion, pypdf and tornado (#30220) * chore(deps): bump aiohttp to 3.14.1 and vitest to 3.2.6 Lockfile-only bump for aiohttp (3.13.5 -> 3.14.1, within the existing pyproject constraint) and dashboard devDependency bumps for vitest, @vitest/coverage-v8, @vitest/ui (3.2.4 -> 3.2.6) plus transitive brace-expansion (5.0.5 -> 5.0.6). Clears the currently published advisories flagged by osv.dev against uv.lock and the dashboard lockfile. Verified: 154 custom_httpx unit tests and all 3943 dashboard vitest tests pass; live proxy completion and streaming calls succeed on the bumped venv * chore(deps): raise aiohttp floor to 3.14.0 The lockfile bump alone only protects environments built from uv.lock. Raising the pyproject floor extends the same minimum to package consumers installing litellm from PyPI, and prevents a future lockfile regeneration from resolving below 3.14.0 * Revert "chore(deps): raise aiohttp floor to 3.14.0" This reverts commit d6c1c9dc0c8664c015a5dabbde2469539bd247fd. * revert(deps): roll back aiohttp to 3.13.5 vcrpy is incompatible with aiohttp >= 3.14 (the aiohttp_stubs module imports a symbol removed in 3.14) and the upstream fix is merged but unreleased, so every cassette-based test suite fails on 3.14. Hold aiohttp at 3.13.5 until a vcrpy release ships; the vitest and brace-expansion bumps stay * chore(deps): bump pypdf to 6.13.1 and tornado to 6.5.7 Lockfile-only bumps clearing the advisories published for both since this branch was opened * chore(deps): add regression guards for the bumped versions Raise the pypdf floor to 6.12.0 (direct dependency, applies to package consumers too) and add uv constraint-dependencies for the transitive pins: tornado >= 6.5.6, and aiohttp held in [3.13.5, 3.14) so a lockfile regeneration can neither fall back below the current version nor move onto 3.14 while vcrpy is incompatible. Constraints live in [tool.uv] and only affect this repo's resolution, not published metadata. Verified: uv lock -P with each out-of-range version fails to resolve; in-range resolutions unchanged (pypdf 6.13.1, tornado 6.5.7, aiohttp 3.13.5) --- pyproject.toml | 6 +- ui/litellm-dashboard/package-lock.json | 336 ++++++++++++------------- ui/litellm-dashboard/package.json | 6 +- uv.lock | 36 +-- 4 files changed, 196 insertions(+), 188 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index b9d76379faf..6429b810969 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -133,7 +133,7 @@ proxy-runtime = [ "mangum>=0.17.0,<1.0", "azure-ai-contentsafety>=1.0.0,<2.0", "azure-storage-file-datalake>=12.20.0,<13.0", - "pypdf>=6.10.2,<7.0; python_version < '3.14'", + "pypdf>=6.12.0,<7.0; python_version < '3.14'", "llm-sandbox>=0.3.39,<1.0", "detect-secrets>=1.5.0,<2.0", ] @@ -240,6 +240,10 @@ requires = ["uv_build==0.11.8"] build-backend = "uv_build" [tool.uv] +constraint-dependencies = [ + "tornado>=6.5.6", + "aiohttp>=3.13.5,<3.14", +] default-groups = ["dev"] required-version = ">=0.10.9" exclude-newer = "3 days" diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json index 568f6b288d5..dff6c25a9a4 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -49,8 +49,8 @@ "@types/react-copy-to-clipboard": "5.0.7", "@types/react-dom": "18.3.7", "@types/react-syntax-highlighter": "15.5.13", - "@vitest/coverage-v8": "3.2.4", - "@vitest/ui": "3.2.4", + "@vitest/coverage-v8": "3.2.6", + "@vitest/ui": "3.2.6", "autoprefixer": "10.4.24", "eslint": "9.39.2", "eslint-config-next": "16.2.6", @@ -64,7 +64,7 @@ "tailwindcss": "3.4.19", "typescript": "5.9.3", "typescript-eslint": "8.60.1", - "vitest": "3.2.4" + "vitest": "3.2.6" }, "engines": { "node": ">=20.9.0", @@ -2843,9 +2843,9 @@ } }, "node_modules/@rollup/rollup-android-arm-eabi": { - "version": "4.60.3", - "resolved": "https://registry.npmjs.org/@rollup/rollup-android-arm-eabi/-/rollup-android-arm-eabi-4.60.3.tgz", - "integrity": "sha512-x35CNW/ANXG3hE/EZpRU8MXX1JDN86hBb2wMGAtltkz7pc6cxgjpy1OMMfDosOQ+2hWqIkag/fGok1Yady9nGw==", + "version": "4.61.1", + "resolved": "https://registry.npmjs.org/@rollup/rollup-android-arm-eabi/-/rollup-android-arm-eabi-4.61.1.tgz", + "integrity": "sha512-JnBB8MdXj45cajvTuO5FmPlvFVJRQgvrz1uSEl3NwqFnReAPGwb8EanbGi4z2nRaqLzjJSv5/JmycoTKlRZxHA==", "cpu": [ "arm" ], @@ -2857,9 +2857,9 @@ ] }, "node_modules/@rollup/rollup-android-arm64": { - "version": "4.60.3", - "resolved": "https://registry.npmjs.org/@rollup/rollup-android-arm64/-/rollup-android-arm64-4.60.3.tgz", - "integrity": "sha512-xw3xtkDApIOGayehp2+Rz4zimfkaX65r4t47iy+ymQB2G4iJCBBfj0ogVg5jpvjpn8UWn/+q9tprxleYeNp3Hw==", + "version": "4.61.1", + "resolved": "https://registry.npmjs.org/@rollup/rollup-android-arm64/-/rollup-android-arm64-4.61.1.tgz", + "integrity": "sha512-Jx2g7iSjw4AOT0HDPHM9RV3GNjRXwybWtSFZiZAYUTjUwjVrYIwq3kBf+LnhqJlzXFAqTAh2F7IGI+O568exPw==", "cpu": [ "arm64" ], @@ -2871,9 +2871,9 @@ ] }, "node_modules/@rollup/rollup-darwin-arm64": { - "version": "4.60.3", - "resolved": "https://registry.npmjs.org/@rollup/rollup-darwin-arm64/-/rollup-darwin-arm64-4.60.3.tgz", - "integrity": "sha512-vo6Y5Qfpx7/5EaamIwi0WqW2+zfiusVihKatLvtN1VFVy3D13uERk/6gZLU1UiHRL6fDXqj/ELIeVRGnvcTE1g==", + "version": "4.61.1", + "resolved": "https://registry.npmjs.org/@rollup/rollup-darwin-arm64/-/rollup-darwin-arm64-4.61.1.tgz", + "integrity": "sha512-0F1L/Z3Eqv8mT2n3dCpeO8GcTvHvVqkP5/t6DMsn0KzhYVcg+s7Ncl5DS8qjKYEeio6Az0Gt6nyBORay5qIlCA==", "cpu": [ "arm64" ], @@ -2885,9 +2885,9 @@ ] }, "node_modules/@rollup/rollup-darwin-x64": { - "version": "4.60.3", - "resolved": "https://registry.npmjs.org/@rollup/rollup-darwin-x64/-/rollup-darwin-x64-4.60.3.tgz", - "integrity": "sha512-D+0QGcZhBzTN82weOnsSlY7V7+RMmPuF1CkbxyMAGE8+ZHeUjyb76ZiWmBlCu//AQQONvxcqRbwZTajZKqjuOw==", + "version": "4.61.1", + "resolved": "https://registry.npmjs.org/@rollup/rollup-darwin-x64/-/rollup-darwin-x64-4.61.1.tgz", + "integrity": "sha512-qLttcH871ujY4YcVfUSShhOw+CsoTatYz8gRbHO7Bb92QH059/P0y5do1KMs41fY0BpD2x4AJH/gID0zFiqVKQ==", "cpu": [ "x64" ], @@ -2899,9 +2899,9 @@ ] }, "node_modules/@rollup/rollup-freebsd-arm64": { - "version": "4.60.3", - "resolved": "https://registry.npmjs.org/@rollup/rollup-freebsd-arm64/-/rollup-freebsd-arm64-4.60.3.tgz", - "integrity": "sha512-6HnvHCT7fDyj6R0Ph7A6x8dQS/S38MClRWeDLqc0MdfWkxjiu1HSDYrdPhqSILzjTIC/pnXbbJbo+ft+gy/9hQ==", + "version": "4.61.1", + "resolved": "https://registry.npmjs.org/@rollup/rollup-freebsd-arm64/-/rollup-freebsd-arm64-4.61.1.tgz", + "integrity": "sha512-fUI4RapGE0Oh3mb8mgfvC1O2nU1RpDZUKnDQm3xB1Ipg7C2wTs5Kstz7G2uWK99a8S2yTMq8/P4uycwNa0nJyw==", "cpu": [ "arm64" ], @@ -2913,9 +2913,9 @@ ] }, "node_modules/@rollup/rollup-freebsd-x64": { - "version": "4.60.3", - "resolved": "https://registry.npmjs.org/@rollup/rollup-freebsd-x64/-/rollup-freebsd-x64-4.60.3.tgz", - "integrity": "sha512-KHLgC3WKlUYW3ShFKnnosZDOJ0xjg9zp7au3sIm2bs/tGBeC2ipmvRh/N7JKi0t9Ue20C0dpEshi8WUubg+cnA==", + "version": "4.61.1", + "resolved": "https://registry.npmjs.org/@rollup/rollup-freebsd-x64/-/rollup-freebsd-x64-4.61.1.tgz", + "integrity": "sha512-H5YrdvJaDtI/U9/emrD4b++xkvp3y/JvOe4rizHbxvkyMfRS/CiRYdji+Pl8D0brEaNFWUh1drQxgAGIl6Xudw==", "cpu": [ "x64" ], @@ -2927,9 +2927,9 @@ ] }, "node_modules/@rollup/rollup-linux-arm-gnueabihf": { - "version": "4.60.3", - "resolved": "https://registry.npmjs.org/@rollup/rollup-linux-arm-gnueabihf/-/rollup-linux-arm-gnueabihf-4.60.3.tgz", - "integrity": "sha512-DV6fJoxEYWJOvaZIsok7KrYl0tPvga5OZ2yvKHNNYyk/2roMLqQAbGhr78EQ5YhHpnhLKJD3S1WFusAkmUuV5g==", + "version": "4.61.1", + "resolved": "https://registry.npmjs.org/@rollup/rollup-linux-arm-gnueabihf/-/rollup-linux-arm-gnueabihf-4.61.1.tgz", + "integrity": "sha512-Q8CBCCQtDFrYtXoeUXSrnFXKOnyUhx6bz+SkL6A0E7V8kAiCJ5pamq1WtbfpVGhR5TSpXY6ak3avmDc5fHTyJA==", "cpu": [ "arm" ], @@ -2941,9 +2941,9 @@ ] }, "node_modules/@rollup/rollup-linux-arm-musleabihf": { - "version": "4.60.3", - "resolved": "https://registry.npmjs.org/@rollup/rollup-linux-arm-musleabihf/-/rollup-linux-arm-musleabihf-4.60.3.tgz", - "integrity": "sha512-mQKoJAzvuOs6F+TZybQO4GOTSMUu7v0WdxEk24krQ/uUxXoPTtHjuaUuPmFhtBcM4K0ons8nrE3JyhTuCFtT/w==", + "version": "4.61.1", + "resolved": "https://registry.npmjs.org/@rollup/rollup-linux-arm-musleabihf/-/rollup-linux-arm-musleabihf-4.61.1.tgz", + "integrity": "sha512-nwnhk1581l0FBVellGcVCAT0Oi06onEA3WB53sf01VO3I0UPBkMH9sXONYME2K0ovXcNayJfNtHfm6mpJElatQ==", "cpu": [ "arm" ], @@ -2955,9 +2955,9 @@ ] }, "node_modules/@rollup/rollup-linux-arm64-gnu": { - "version": "4.60.3", - "resolved": "https://registry.npmjs.org/@rollup/rollup-linux-arm64-gnu/-/rollup-linux-arm64-gnu-4.60.3.tgz", - "integrity": "sha512-Whjj2qoiJ6+OOJMGptTYazaJvjOJm+iKHpXQM1P3LzGjt7Ff++Tp7nH4N8J/BUA7R9IHfDyx4DJIflifwnbmIA==", + "version": "4.61.1", + "resolved": "https://registry.npmjs.org/@rollup/rollup-linux-arm64-gnu/-/rollup-linux-arm64-gnu-4.61.1.tgz", + "integrity": "sha512-x5Xr49hwt3hdW75UOZm3395YwwzPyauktslv29KpWL/T+vVAzoT3azLcTWv0eMciBNrx+DYjH4paehHoLpPvpg==", "cpu": [ "arm64" ], @@ -2969,9 +2969,9 @@ ] }, "node_modules/@rollup/rollup-linux-arm64-musl": { - "version": "4.60.3", - "resolved": "https://registry.npmjs.org/@rollup/rollup-linux-arm64-musl/-/rollup-linux-arm64-musl-4.60.3.tgz", - "integrity": "sha512-4YTNHKqGng5+yiZt3mg77nmyuCfmNfX4fPmyUapBcIk+BdwSwmCWGXOUxhXbBEkFHtoN5boLj/5NON+u5QC9tg==", + "version": "4.61.1", + "resolved": "https://registry.npmjs.org/@rollup/rollup-linux-arm64-musl/-/rollup-linux-arm64-musl-4.61.1.tgz", + "integrity": "sha512-unMS3H73DpaoPyyEVPjGKleM/s0mkmsauTENpw4INQY8y4+IuLNjkueQ5QCtC0D3N38Y38yhAU8OoZ20S2Tm6w==", "cpu": [ "arm64" ], @@ -2983,9 +2983,9 @@ ] }, "node_modules/@rollup/rollup-linux-loong64-gnu": { - "version": "4.60.3", - "resolved": "https://registry.npmjs.org/@rollup/rollup-linux-loong64-gnu/-/rollup-linux-loong64-gnu-4.60.3.tgz", - "integrity": "sha512-SU3kNlhkpI4UqlUc2VXPGK9o886ZsSeGfMAX2ba2b8DKmMXq4AL7KUrkSWVbb7koVqx41Yczx6dx5PNargIrEA==", + "version": "4.61.1", + "resolved": "https://registry.npmjs.org/@rollup/rollup-linux-loong64-gnu/-/rollup-linux-loong64-gnu-4.61.1.tgz", + "integrity": "sha512-zNZzGRnAhwjFEYmvphJRV5XaQGjs62cCmeYYHUT//NbvEnHauw+I85nGG+SiVg5ld4GX8D1IbKIX+ozITQnhMQ==", "cpu": [ "loong64" ], @@ -2997,9 +2997,9 @@ ] }, "node_modules/@rollup/rollup-linux-loong64-musl": { - "version": "4.60.3", - "resolved": "https://registry.npmjs.org/@rollup/rollup-linux-loong64-musl/-/rollup-linux-loong64-musl-4.60.3.tgz", - "integrity": "sha512-6lDLl5h4TXpB1mTf2rQWnAk/LcXrx9vBfu/DT5TIPhvMhRWaZ5MxkIc8u4lJAmBo6klTe1ywXIUHFjylW505sg==", + "version": "4.61.1", + "resolved": "https://registry.npmjs.org/@rollup/rollup-linux-loong64-musl/-/rollup-linux-loong64-musl-4.61.1.tgz", + "integrity": "sha512-LdpWGL8X209B2SIvWjqlc8VZgM6PKfontSerGepuldQmHYrAOtnMCXeJkxXGbC+PPZVOuu5czJo7fNV6aeW8rQ==", "cpu": [ "loong64" ], @@ -3011,9 +3011,9 @@ ] }, "node_modules/@rollup/rollup-linux-ppc64-gnu": { - "version": "4.60.3", - "resolved": "https://registry.npmjs.org/@rollup/rollup-linux-ppc64-gnu/-/rollup-linux-ppc64-gnu-4.60.3.tgz", - "integrity": "sha512-BMo8bOw8evlup/8G+cj5xWtPyp93xPdyoSN16Zy90Q2QZ0ZYRhCt6ZJSwbrRzG9HApFabjwj2p25TUPDWrhzqQ==", + "version": "4.61.1", + "resolved": "https://registry.npmjs.org/@rollup/rollup-linux-ppc64-gnu/-/rollup-linux-ppc64-gnu-4.61.1.tgz", + "integrity": "sha512-EC5kTtNaNGOmbMGqar8dvJy6y/hg99GAwjfBz++pxZhQATXGcRjd6c5en5wcbru0vkRmiMGsQKdMJOOf6sza4g==", "cpu": [ "ppc64" ], @@ -3025,9 +3025,9 @@ ] }, "node_modules/@rollup/rollup-linux-ppc64-musl": { - "version": "4.60.3", - "resolved": "https://registry.npmjs.org/@rollup/rollup-linux-ppc64-musl/-/rollup-linux-ppc64-musl-4.60.3.tgz", - "integrity": "sha512-E0L8X1dZN1/Rph+5VPF6Xj2G7JJvMACVXtamTJIDrVI44Y3K+G8gQaMEAavbqCGTa16InptiVrX6eM6pmJ+7qA==", + "version": "4.61.1", + "resolved": "https://registry.npmjs.org/@rollup/rollup-linux-ppc64-musl/-/rollup-linux-ppc64-musl-4.61.1.tgz", + "integrity": "sha512-8hiwp6D4acEcNK78I4rP0/XtS1sknWIAMJBPdR4l6zUtyTm5KiTDr5bXmWt4foY7nAN7AThDHgkLIEZOWKbzWw==", "cpu": [ "ppc64" ], @@ -3039,9 +3039,9 @@ ] }, "node_modules/@rollup/rollup-linux-riscv64-gnu": { - "version": "4.60.3", - "resolved": "https://registry.npmjs.org/@rollup/rollup-linux-riscv64-gnu/-/rollup-linux-riscv64-gnu-4.60.3.tgz", - "integrity": "sha512-oZJ/WHaVfHUiRAtmTAeo3DcevNsVvH8mbvodjZy7D5QKvCefO371SiKRpxoDcCxB3PTRTLayWBkvmDQKTcX/sw==", + "version": "4.61.1", + "resolved": "https://registry.npmjs.org/@rollup/rollup-linux-riscv64-gnu/-/rollup-linux-riscv64-gnu-4.61.1.tgz", + "integrity": "sha512-10dh/h/BqA7DuMPWSxkR8uks18FRwnwOEqr5zOTEl+NOwP/OMzKX8OFR/Of9xxDA7D5qef1Nzar5WDD2kCCr1g==", "cpu": [ "riscv64" ], @@ -3053,9 +3053,9 @@ ] }, "node_modules/@rollup/rollup-linux-riscv64-musl": { - "version": "4.60.3", - "resolved": "https://registry.npmjs.org/@rollup/rollup-linux-riscv64-musl/-/rollup-linux-riscv64-musl-4.60.3.tgz", - "integrity": "sha512-Dhbyh7j9FybM3YaTgaHmVALwA8AkUwTPccyCQ79TG9AJUsMQqgN1DDEZNr4+QUfwiWvLDumW5vdwzoeUF+TNxQ==", + "version": "4.61.1", + "resolved": "https://registry.npmjs.org/@rollup/rollup-linux-riscv64-musl/-/rollup-linux-riscv64-musl-4.61.1.tgz", + "integrity": "sha512-YKJ5lg35DP17gcAOggnihe+APw9HLyj1Xn7gsmGumBJAUDa6NGXNixJzmkWLhcK9TOuuyQjdamzvJefkO7qHZQ==", "cpu": [ "riscv64" ], @@ -3067,9 +3067,9 @@ ] }, "node_modules/@rollup/rollup-linux-s390x-gnu": { - "version": "4.60.3", - "resolved": "https://registry.npmjs.org/@rollup/rollup-linux-s390x-gnu/-/rollup-linux-s390x-gnu-4.60.3.tgz", - "integrity": "sha512-cJd1X5XhHHlltkaypz1UcWLA8AcoIi1aWhsvaWDskD1oz2eKCypnqvTQ8ykMNI0RSmm7NkTdSqSSD7zM0xa6Ig==", + "version": "4.61.1", + "resolved": "https://registry.npmjs.org/@rollup/rollup-linux-s390x-gnu/-/rollup-linux-s390x-gnu-4.61.1.tgz", + "integrity": "sha512-Mlil5G2Jj6a7B3LWGctg+XPL9vdXYuzCtNXfxOQ0nPjc2m6ueUktocPGH9bnAM0bNRKb/bAWTujUU7IJQdQA+g==", "cpu": [ "s390x" ], @@ -3081,9 +3081,9 @@ ] }, "node_modules/@rollup/rollup-linux-x64-gnu": { - "version": "4.60.3", - "resolved": "https://registry.npmjs.org/@rollup/rollup-linux-x64-gnu/-/rollup-linux-x64-gnu-4.60.3.tgz", - "integrity": "sha512-DAZDBHQfG2oQuhY7mc6I3/qB4LU2fQCjRvxbDwd/Jdvb9fypP4IJ4qmtu6lNjes6B531AI8cg1aKC2di97bUxA==", + "version": "4.61.1", + "resolved": "https://registry.npmjs.org/@rollup/rollup-linux-x64-gnu/-/rollup-linux-x64-gnu-4.61.1.tgz", + "integrity": "sha512-bVWIOIk6pV01p4CdUbPP7CJ/434z+OooYjDuFcR+44N35YvKUC66G8MGnvcWx5mWKW3g61J+t74l3Kj15Kwn2Q==", "cpu": [ "x64" ], @@ -3095,9 +3095,9 @@ ] }, "node_modules/@rollup/rollup-linux-x64-musl": { - "version": "4.60.3", - "resolved": "https://registry.npmjs.org/@rollup/rollup-linux-x64-musl/-/rollup-linux-x64-musl-4.60.3.tgz", - "integrity": "sha512-cRxsE8c13mZOh3vP+wLDxpQBRrOHDIGOWyDL93Sy0Ga8y515fBcC2pjUfFwUe5T7tqvTvWbCpg1URM/AXdWIXA==", + "version": "4.61.1", + "resolved": "https://registry.npmjs.org/@rollup/rollup-linux-x64-musl/-/rollup-linux-x64-musl-4.61.1.tgz", + "integrity": "sha512-qy5pBvZbqNFheBz61R1rzsezjm0J7O2oNGoWtGoY89SZYLUfxAJTBAqDChqAIdB4rCiIbi9nF7yZ83GnNiLwSw==", "cpu": [ "x64" ], @@ -3109,9 +3109,9 @@ ] }, "node_modules/@rollup/rollup-openbsd-x64": { - "version": "4.60.3", - "resolved": "https://registry.npmjs.org/@rollup/rollup-openbsd-x64/-/rollup-openbsd-x64-4.60.3.tgz", - "integrity": "sha512-QaWcIgRxqEdQdhJqW4DJctsH6HCmo5vHxY0krHSX4jMtOqfzC+dqDGuHM87bu4H8JBeibWx7jFz+h6/4C8wA5Q==", + "version": "4.61.1", + "resolved": "https://registry.npmjs.org/@rollup/rollup-openbsd-x64/-/rollup-openbsd-x64-4.61.1.tgz", + "integrity": "sha512-E83TXjI4zm0+5f2qO+UOudaCYIhYwpJ5jq6YCZNIZ+6CbfhKrkAGezeiASBL9ElxAxFsRS9ZhESv8mfnj6TKeg==", "cpu": [ "x64" ], @@ -3123,9 +3123,9 @@ ] }, "node_modules/@rollup/rollup-openharmony-arm64": { - "version": "4.60.3", - "resolved": "https://registry.npmjs.org/@rollup/rollup-openharmony-arm64/-/rollup-openharmony-arm64-4.60.3.tgz", - "integrity": "sha512-AaXwSvUi3QIPtroAUw1t5yHGIyqKEXwH54WUocFolZhpGDruJcs8c+xPNDRn4XiQsS7MEwnYsHW2l0MBLDMkWg==", + "version": "4.61.1", + "resolved": "https://registry.npmjs.org/@rollup/rollup-openharmony-arm64/-/rollup-openharmony-arm64-4.61.1.tgz", + "integrity": "sha512-fbWnKqVkjrJN38vNe3ahkbk6iejS/3b0Nt7EEtPpE6RBacZcGXNKbzfHN3GUUlXOPghUg0j6XUGrtjX9z1sIvA==", "cpu": [ "arm64" ], @@ -3137,9 +3137,9 @@ ] }, "node_modules/@rollup/rollup-win32-arm64-msvc": { - "version": "4.60.3", - "resolved": "https://registry.npmjs.org/@rollup/rollup-win32-arm64-msvc/-/rollup-win32-arm64-msvc-4.60.3.tgz", - "integrity": "sha512-65LAKM/bAWDqKNEelHlcHvm2V+Vfb8C6INFxQXRHCvaVN1rJfwr4NvdP4FyzUaLqWfaCGaadf6UbTm8xJeYfEg==", + "version": "4.61.1", + "resolved": "https://registry.npmjs.org/@rollup/rollup-win32-arm64-msvc/-/rollup-win32-arm64-msvc-4.61.1.tgz", + "integrity": "sha512-ArMl38iVAbk0New1ogihQNY6iphLi4ZaRsa037gUzv5yeKPY8TD3Dmy4x2RNC1VztU/uqm+G+/RwFrSka3Oy2g==", "cpu": [ "arm64" ], @@ -3151,9 +3151,9 @@ ] }, "node_modules/@rollup/rollup-win32-ia32-msvc": { - "version": "4.60.3", - "resolved": "https://registry.npmjs.org/@rollup/rollup-win32-ia32-msvc/-/rollup-win32-ia32-msvc-4.60.3.tgz", - "integrity": "sha512-EEM2gyhBF5MFnI6vMKdX1LAosE627RGBzIoGMdLloPZkXrUN0Ckqgr2Qi8+J3zip/8NVVro3/FjB+tjhZUgUHA==", + "version": "4.61.1", + "resolved": "https://registry.npmjs.org/@rollup/rollup-win32-ia32-msvc/-/rollup-win32-ia32-msvc-4.61.1.tgz", + "integrity": "sha512-0mYtjHS9ucAbcATycCNK9IGBk/cCe/ma7EmSLGZdsxnOA8cjRIyU04wDpVAD9NiOfLUR9KTxdiO53uOkherqjQ==", "cpu": [ "ia32" ], @@ -3165,9 +3165,9 @@ ] }, "node_modules/@rollup/rollup-win32-x64-gnu": { - "version": "4.60.3", - "resolved": "https://registry.npmjs.org/@rollup/rollup-win32-x64-gnu/-/rollup-win32-x64-gnu-4.60.3.tgz", - "integrity": "sha512-E5Eb5H/DpxaoXH++Qkv28RcUJboMopmdDUALBczvHMf7hNIxaDZqwY5lK12UK1BHacSmvupoEWGu+n993Z0y1A==", + "version": "4.61.1", + "resolved": "https://registry.npmjs.org/@rollup/rollup-win32-x64-gnu/-/rollup-win32-x64-gnu-4.61.1.tgz", + "integrity": "sha512-gK1iCEPfpoSG9wfBihXxvBMi8ZfcWffYkEsC/Eih+iFENTaewvNcrEQ69lIOWYO5pePHKLHHO7nq5AILGO/HQQ==", "cpu": [ "x64" ], @@ -3179,9 +3179,9 @@ ] }, "node_modules/@rollup/rollup-win32-x64-msvc": { - "version": "4.60.3", - "resolved": "https://registry.npmjs.org/@rollup/rollup-win32-x64-msvc/-/rollup-win32-x64-msvc-4.60.3.tgz", - "integrity": "sha512-hPt/bgL5cE+Qp+/TPHBqptcAgPzgj46mPcg/16zNUmbQk0j+mOEQV/+Lqu8QRtDV3Ek95Q6FeFITpuhl6OTsAA==", + "version": "4.61.1", + "resolved": "https://registry.npmjs.org/@rollup/rollup-win32-x64-msvc/-/rollup-win32-x64-msvc-4.61.1.tgz", + "integrity": "sha512-X+zaP2x+j4RXGfbp/seSoRHWnPxzApilDszisZxbYH5C/jTxFhCtDNdPGZb9lJyYPs24wGxruPF7Y+sIXt9Gzw==", "cpu": [ "x64" ], @@ -3567,9 +3567,9 @@ "license": "MIT" }, "node_modules/@types/estree": { - "version": "1.0.8", - "resolved": "https://registry.npmjs.org/@types/estree/-/estree-1.0.8.tgz", - "integrity": "sha512-dWHzHa2WqEXI/O1E9OjrocMTKJl2mSrEolh1Iomrv6U+JuNwaHXsXx9bLu5gG7BUWFIN0skIQJQ/L1rIex4X6w==", + "version": "1.0.9", + "resolved": "https://registry.npmjs.org/@types/estree/-/estree-1.0.9.tgz", + "integrity": "sha512-GhdPgy1el4/ImP05X05Uw4cw2/M93BCUmnEvWZNStlCzEKME4Fkk+YpoA5OiHNQmoS7Cafb8Xa3Pya8m1Qrzeg==", "license": "MIT" }, "node_modules/@types/estree-jsx": { @@ -4245,9 +4245,9 @@ ] }, "node_modules/@vitest/coverage-v8": { - "version": "3.2.4", - "resolved": "https://registry.npmjs.org/@vitest/coverage-v8/-/coverage-v8-3.2.4.tgz", - "integrity": "sha512-EyF9SXU6kS5Ku/U82E259WSnvg6c8KTjppUncuNdm5QHpe17mwREHnjDzozC8x9MZ0xfBUFSaLkRv4TMA75ALQ==", + "version": "3.2.6", + "resolved": "https://registry.npmjs.org/@vitest/coverage-v8/-/coverage-v8-3.2.6.tgz", + "integrity": "sha512-LsAdmUapA0qSN306d8+zOyawM0hFm2m2Hg9IwVNIKBm+qJV8cijiq2c+gxKZcB1HCfIWAy+0qEZDCUQA58A1cw==", "dev": true, "license": "MIT", "dependencies": { @@ -4269,8 +4269,8 @@ "url": "https://opencollective.com/vitest" }, "peerDependencies": { - "@vitest/browser": "3.2.4", - "vitest": "3.2.4" + "@vitest/browser": "3.2.6", + "vitest": "3.2.6" }, "peerDependenciesMeta": { "@vitest/browser": { @@ -4279,15 +4279,15 @@ } }, "node_modules/@vitest/expect": { - "version": "3.2.4", - "resolved": "https://registry.npmjs.org/@vitest/expect/-/expect-3.2.4.tgz", - "integrity": "sha512-Io0yyORnB6sikFlt8QW5K7slY4OjqNX9jmJQ02QDda8lyM6B5oNgVWoSoKPac8/kgnCUzuHQKrSLtu/uOqqrig==", + "version": "3.2.6", + "resolved": "https://registry.npmjs.org/@vitest/expect/-/expect-3.2.6.tgz", + "integrity": "sha512-1+7q9BtaKzEmO+fmNT3kYvoNn5Y71XWAx2Q5HRim4tTVRQVRv4uJFAQ5FbK0OPUeNP/WmVCpxYxoJdvuHVjzBQ==", "dev": true, "license": "MIT", "dependencies": { "@types/chai": "^5.2.2", - "@vitest/spy": "3.2.4", - "@vitest/utils": "3.2.4", + "@vitest/spy": "3.2.6", + "@vitest/utils": "3.2.6", "chai": "^5.2.0", "tinyrainbow": "^2.0.0" }, @@ -4296,13 +4296,13 @@ } }, "node_modules/@vitest/mocker": { - "version": "3.2.4", - "resolved": "https://registry.npmjs.org/@vitest/mocker/-/mocker-3.2.4.tgz", - "integrity": "sha512-46ryTE9RZO/rfDd7pEqFl7etuyzekzEhUbTW3BvmeO/BcCMEgq59BKhek3dXDWgAj4oMK6OZi+vRr1wPW6qjEQ==", + "version": "3.2.6", + "resolved": "https://registry.npmjs.org/@vitest/mocker/-/mocker-3.2.6.tgz", + "integrity": "sha512-EZOrpDbkKotFAP7wPAQV1UIyoGOk4oX7ynWhBhLB7v+meMHbQhU16oPpIYGTTe4oFlhpryGpgpcZP/sin3hYuw==", "dev": true, "license": "MIT", "dependencies": { - "@vitest/spy": "3.2.4", + "@vitest/spy": "3.2.6", "estree-walker": "^3.0.3", "magic-string": "^0.30.17" }, @@ -4323,9 +4323,9 @@ } }, "node_modules/@vitest/pretty-format": { - "version": "3.2.4", - "resolved": "https://registry.npmjs.org/@vitest/pretty-format/-/pretty-format-3.2.4.tgz", - "integrity": "sha512-IVNZik8IVRJRTr9fxlitMKeJeXFFFN0JaB9PHPGQ8NKQbGpfjlTx9zO4RefN8gp7eqjNy8nyK3NZmBzOPeIxtA==", + "version": "3.2.6", + "resolved": "https://registry.npmjs.org/@vitest/pretty-format/-/pretty-format-3.2.6.tgz", + "integrity": "sha512-lb7XXXzmm2h2ASzFnRvQpDo6onT1NmMJA3tkGTWiBFtRJ9lxGY3d3mm/Apt36gej2bkkOVLL/yTOtufDaFa/jA==", "dev": true, "license": "MIT", "dependencies": { @@ -4336,13 +4336,13 @@ } }, "node_modules/@vitest/runner": { - "version": "3.2.4", - "resolved": "https://registry.npmjs.org/@vitest/runner/-/runner-3.2.4.tgz", - "integrity": "sha512-oukfKT9Mk41LreEW09vt45f8wx7DordoWUZMYdY/cyAk7w5TWkTRCNZYF7sX7n2wB7jyGAl74OxgwhPgKaqDMQ==", + "version": "3.2.6", + "resolved": "https://registry.npmjs.org/@vitest/runner/-/runner-3.2.6.tgz", + "integrity": "sha512-HYcoSj1w5tcgUnzoF0HcyaAQjpA1gj9ftUJ7iSJSuipc02jW9gKkigwZbjFldAfYHA1fa8UZVRftdMY5msWM9Q==", "dev": true, "license": "MIT", "dependencies": { - "@vitest/utils": "3.2.4", + "@vitest/utils": "3.2.6", "pathe": "^2.0.3", "strip-literal": "^3.0.0" }, @@ -4351,13 +4351,13 @@ } }, "node_modules/@vitest/snapshot": { - "version": "3.2.4", - "resolved": "https://registry.npmjs.org/@vitest/snapshot/-/snapshot-3.2.4.tgz", - "integrity": "sha512-dEYtS7qQP2CjU27QBC5oUOxLE/v5eLkGqPE0ZKEIDGMs4vKWe7IjgLOeauHsR0D5YuuycGRO5oSRXnwnmA78fQ==", + "version": "3.2.6", + "resolved": "https://registry.npmjs.org/@vitest/snapshot/-/snapshot-3.2.6.tgz", + "integrity": "sha512-H+ZjNTWGpObenh0YnlBctAPnJSI20P81PL8BPzWpx54YXLLTm8hEsWawtcYLMrwvpK48hGxLLbCS+1KRXhsKhw==", "dev": true, "license": "MIT", "dependencies": { - "@vitest/pretty-format": "3.2.4", + "@vitest/pretty-format": "3.2.6", "magic-string": "^0.30.17", "pathe": "^2.0.3" }, @@ -4366,9 +4366,9 @@ } }, "node_modules/@vitest/spy": { - "version": "3.2.4", - "resolved": "https://registry.npmjs.org/@vitest/spy/-/spy-3.2.4.tgz", - "integrity": "sha512-vAfasCOe6AIK70iP5UD11Ac4siNUNJ9i/9PZ3NKx07sG6sUxeag1LWdNrMWeKKYBLlzuK+Gn65Yd5nyL6ds+nw==", + "version": "3.2.6", + "resolved": "https://registry.npmjs.org/@vitest/spy/-/spy-3.2.6.tgz", + "integrity": "sha512-oq6BbH68WzcWmwtBrU9nqLeaXTR4XwJF7FSLkKEZo4i6eoXcrxjcwSuTvWBIRUTC6VC72nXYunzqgZA+IKdtxg==", "dev": true, "license": "MIT", "dependencies": { @@ -4379,13 +4379,13 @@ } }, "node_modules/@vitest/ui": { - "version": "3.2.4", - "resolved": "https://registry.npmjs.org/@vitest/ui/-/ui-3.2.4.tgz", - "integrity": "sha512-hGISOaP18plkzbWEcP/QvtRW1xDXF2+96HbEX6byqQhAUbiS5oH6/9JwW+QsQCIYON2bI6QZBF+2PvOmrRZ9wA==", + "version": "3.2.6", + "resolved": "https://registry.npmjs.org/@vitest/ui/-/ui-3.2.6.tgz", + "integrity": "sha512-mATfG3zVdhobE9U1rIpvtYD3DGuSSxqZ3Aj/8ityGqKXy8YDJ9BoAjZmAz6dZ1IZ1xI5V+MerkCczvVa+3QK9Q==", "dev": true, "license": "MIT", "dependencies": { - "@vitest/utils": "3.2.4", + "@vitest/utils": "3.2.6", "fflate": "^0.8.2", "flatted": "^3.3.3", "pathe": "^2.0.3", @@ -4397,17 +4397,17 @@ "url": "https://opencollective.com/vitest" }, "peerDependencies": { - "vitest": "3.2.4" + "vitest": "3.2.6" } }, "node_modules/@vitest/utils": { - "version": "3.2.4", - "resolved": "https://registry.npmjs.org/@vitest/utils/-/utils-3.2.4.tgz", - "integrity": "sha512-fB2V0JFrQSMsCo9HiSq3Ezpdv4iYaXRG1Sx8edX3MwxfyNn83mKiGzOcH+Fkxt4MHxr3y42fQi1oeAInqgX2QA==", + "version": "3.2.6", + "resolved": "https://registry.npmjs.org/@vitest/utils/-/utils-3.2.6.tgz", + "integrity": "sha512-lI23nIs4bnT3T8NIoh+vFaz5s2/DdP0Jgt2jxwgWljvwn82cLJtyi/If+fjFyoLMGIOz0U/fKvWE0d4jsNQEfg==", "dev": true, "license": "MIT", "dependencies": { - "@vitest/pretty-format": "3.2.4", + "@vitest/pretty-format": "3.2.6", "loupe": "^3.1.4", "tinyrainbow": "^2.0.0" }, @@ -4996,9 +4996,9 @@ } }, "node_modules/brace-expansion": { - "version": "5.0.5", - "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-5.0.5.tgz", - "integrity": "sha512-VZznLgtwhn+Mact9tfiwx64fA9erHH/MCXEUfB/0bX/6Fz6ny5EGTXYltMocqg4xFAQZtnO3DHWWXi8RiuN7cQ==", + "version": "5.0.6", + "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-5.0.6.tgz", + "integrity": "sha512-kLpxurY4Z4r9sgMsyG0Z9uzsBlgiU/EFKhj/h91/8yHu0edo7XuixOIH3VcJ8kkxs6/jPzoI6U9Vj3WqbMQ94g==", "dev": true, "license": "MIT", "dependencies": { @@ -6796,9 +6796,9 @@ } }, "node_modules/fflate": { - "version": "0.8.2", - "resolved": "https://registry.npmjs.org/fflate/-/fflate-0.8.2.tgz", - "integrity": "sha512-cPJU47OaAoCbg0pBvzsgpTPhmhqI5eJjh/JIu8tPj5q+T7iLvW/JAYUqmE7KOB4R1ZyEhzBaIQpQpardBF5z8A==", + "version": "0.8.3", + "resolved": "https://registry.npmjs.org/fflate/-/fflate-0.8.3.tgz", + "integrity": "sha512-tbZNuJrLwGUp3zshBtdy4W+ORxZuIh8a5ilyIEQDC5rY1f3U20JMry0Ll3WBzU58EZKsEuJFXhb5gwv8CsPvgA==", "dev": true, "license": "MIT" }, @@ -11828,13 +11828,13 @@ } }, "node_modules/rollup": { - "version": "4.60.3", - "resolved": "https://registry.npmjs.org/rollup/-/rollup-4.60.3.tgz", - "integrity": "sha512-pAQK9HalE84QSm4Po3EmWIZPd3FnjkShVkiMlz1iligWYkWQ7wHYd1PF/T7QZ5TVSD6uSTon5gBVMSM4JfBV+A==", + "version": "4.61.1", + "resolved": "https://registry.npmjs.org/rollup/-/rollup-4.61.1.tgz", + "integrity": "sha512-I4KW6iuRpuu2uHBLraZ1wNZe0DP7lnRha+VJ9tNaYVaVgKhW0aI3h4RYnoRPeql0flHm/Co55b7snEDcOfOJrA==", "dev": true, "license": "MIT", "dependencies": { - "@types/estree": "1.0.8" + "@types/estree": "1.0.9" }, "bin": { "rollup": "dist/bin/rollup" @@ -11844,31 +11844,31 @@ "npm": ">=8.0.0" }, "optionalDependencies": { - "@rollup/rollup-android-arm-eabi": "4.60.3", - "@rollup/rollup-android-arm64": "4.60.3", - "@rollup/rollup-darwin-arm64": "4.60.3", - "@rollup/rollup-darwin-x64": "4.60.3", - "@rollup/rollup-freebsd-arm64": "4.60.3", - "@rollup/rollup-freebsd-x64": "4.60.3", - "@rollup/rollup-linux-arm-gnueabihf": "4.60.3", - "@rollup/rollup-linux-arm-musleabihf": "4.60.3", - "@rollup/rollup-linux-arm64-gnu": "4.60.3", - "@rollup/rollup-linux-arm64-musl": "4.60.3", - "@rollup/rollup-linux-loong64-gnu": "4.60.3", - "@rollup/rollup-linux-loong64-musl": "4.60.3", - "@rollup/rollup-linux-ppc64-gnu": "4.60.3", - "@rollup/rollup-linux-ppc64-musl": "4.60.3", - "@rollup/rollup-linux-riscv64-gnu": "4.60.3", - "@rollup/rollup-linux-riscv64-musl": "4.60.3", - "@rollup/rollup-linux-s390x-gnu": "4.60.3", - "@rollup/rollup-linux-x64-gnu": "4.60.3", - "@rollup/rollup-linux-x64-musl": "4.60.3", - "@rollup/rollup-openbsd-x64": "4.60.3", - "@rollup/rollup-openharmony-arm64": "4.60.3", - "@rollup/rollup-win32-arm64-msvc": "4.60.3", - "@rollup/rollup-win32-ia32-msvc": "4.60.3", - "@rollup/rollup-win32-x64-gnu": "4.60.3", - "@rollup/rollup-win32-x64-msvc": "4.60.3", + "@rollup/rollup-android-arm-eabi": "4.61.1", + "@rollup/rollup-android-arm64": "4.61.1", + "@rollup/rollup-darwin-arm64": "4.61.1", + "@rollup/rollup-darwin-x64": "4.61.1", + "@rollup/rollup-freebsd-arm64": "4.61.1", + "@rollup/rollup-freebsd-x64": "4.61.1", + "@rollup/rollup-linux-arm-gnueabihf": "4.61.1", + "@rollup/rollup-linux-arm-musleabihf": "4.61.1", + "@rollup/rollup-linux-arm64-gnu": "4.61.1", + "@rollup/rollup-linux-arm64-musl": "4.61.1", + "@rollup/rollup-linux-loong64-gnu": "4.61.1", + "@rollup/rollup-linux-loong64-musl": "4.61.1", + "@rollup/rollup-linux-ppc64-gnu": "4.61.1", + "@rollup/rollup-linux-ppc64-musl": "4.61.1", + "@rollup/rollup-linux-riscv64-gnu": "4.61.1", + "@rollup/rollup-linux-riscv64-musl": "4.61.1", + "@rollup/rollup-linux-s390x-gnu": "4.61.1", + "@rollup/rollup-linux-x64-gnu": "4.61.1", + "@rollup/rollup-linux-x64-musl": "4.61.1", + "@rollup/rollup-openbsd-x64": "4.61.1", + "@rollup/rollup-openharmony-arm64": "4.61.1", + "@rollup/rollup-win32-arm64-msvc": "4.61.1", + "@rollup/rollup-win32-ia32-msvc": "4.61.1", + "@rollup/rollup-win32-x64-gnu": "4.61.1", + "@rollup/rollup-win32-x64-msvc": "4.61.1", "fsevents": "~2.3.2" } }, @@ -13342,9 +13342,9 @@ } }, "node_modules/vite": { - "version": "7.3.2", - "resolved": "https://registry.npmjs.org/vite/-/vite-7.3.2.tgz", - "integrity": "sha512-Bby3NOsna2jsjfLVOHKes8sGwgl4TT0E6vvpYgnAYDIF/tie7MRaFthmKuHx1NSXjiTueXH3do80FMQgvEktRg==", + "version": "7.3.5", + "resolved": "https://registry.npmjs.org/vite/-/vite-7.3.5.tgz", + "integrity": "sha512-KuOaNhcnGFN2zIPGA7wRmzF+lJA1sea7rHq17aiJ++9lzY1WWG6Jpwqwe1KNbRVPIqHmr8GLYx7jbrQcN/7/ww==", "dev": true, "license": "MIT", "dependencies": { @@ -13455,20 +13455,20 @@ } }, "node_modules/vitest": { - "version": "3.2.4", - "resolved": "https://registry.npmjs.org/vitest/-/vitest-3.2.4.tgz", - "integrity": "sha512-LUCP5ev3GURDysTWiP47wRRUpLKMOfPh+yKTx3kVIEiu5KOMeqzpnYNsKyOoVrULivR8tLcks4+lga33Whn90A==", + "version": "3.2.6", + "resolved": "https://registry.npmjs.org/vitest/-/vitest-3.2.6.tgz", + "integrity": "sha512-xejya+bT/j/+R/AGa1XOfRxLmNUlLtlwjRsFUILF+xHfzElmGcmFydy2gqqIrd62ptIEfwVMofd19uNWD9L7Nw==", "dev": true, "license": "MIT", "dependencies": { "@types/chai": "^5.2.2", - "@vitest/expect": "3.2.4", - "@vitest/mocker": "3.2.4", - "@vitest/pretty-format": "^3.2.4", - "@vitest/runner": "3.2.4", - "@vitest/snapshot": "3.2.4", - "@vitest/spy": "3.2.4", - "@vitest/utils": "3.2.4", + "@vitest/expect": "3.2.6", + "@vitest/mocker": "3.2.6", + "@vitest/pretty-format": "^3.2.6", + "@vitest/runner": "3.2.6", + "@vitest/snapshot": "3.2.6", + "@vitest/spy": "3.2.6", + "@vitest/utils": "3.2.6", "chai": "^5.2.0", "debug": "^4.4.1", "expect-type": "^1.2.1", @@ -13498,8 +13498,8 @@ "@edge-runtime/vm": "*", "@types/debug": "^4.1.12", "@types/node": "^18.0.0 || ^20.0.0 || >=22.0.0", - "@vitest/browser": "3.2.4", - "@vitest/ui": "3.2.4", + "@vitest/browser": "3.2.6", + "@vitest/ui": "3.2.6", "happy-dom": "*", "jsdom": "*" }, diff --git a/ui/litellm-dashboard/package.json b/ui/litellm-dashboard/package.json index eb6211a91d1..7187ec6da4b 100644 --- a/ui/litellm-dashboard/package.json +++ b/ui/litellm-dashboard/package.json @@ -64,8 +64,8 @@ "@types/react-copy-to-clipboard": "5.0.7", "@types/react-dom": "18.3.7", "@types/react-syntax-highlighter": "15.5.13", - "@vitest/coverage-v8": "3.2.4", - "@vitest/ui": "3.2.4", + "@vitest/coverage-v8": "3.2.6", + "@vitest/ui": "3.2.6", "autoprefixer": "10.4.24", "eslint": "9.39.2", "eslint-config-next": "16.2.6", @@ -79,7 +79,7 @@ "tailwindcss": "3.4.19", "typescript": "5.9.3", "typescript-eslint": "8.60.1", - "vitest": "3.2.4" + "vitest": "3.2.6" }, "overrides": { "prismjs": "1.30.0", diff --git a/uv.lock b/uv.lock index 1100db783d3..0efaa74cddb 100644 --- a/uv.lock +++ b/uv.lock @@ -9,7 +9,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-06-05T23:18:37.734017Z" +exclude-newer = "2026-06-10T00:35:00.40525Z" exclude-newer-span = "P3D" [manifest] @@ -18,6 +18,10 @@ members = [ "litellm-enterprise", "litellm-proxy-extras", ] +constraints = [ + { name = "aiohttp", specifier = ">=3.13.5,<3.14" }, + { name = "tornado", specifier = ">=6.5.6" }, +] [[package]] name = "a2a-sdk" @@ -3529,7 +3533,7 @@ requires-dist = [ { name = "pydantic-settings", marker = "extra == 'proxy'", specifier = ">=2.14.1,<3.0" }, { name = "pyjwt", marker = "extra == 'proxy'", specifier = ">=2.13.0,<3.0" }, { name = "pynacl", marker = "extra == 'proxy'", specifier = ">=1.6.2,<2.0" }, - { name = "pypdf", marker = "python_full_version < '3.14' and extra == 'proxy-runtime'", specifier = ">=6.10.2,<7.0" }, + { name = "pypdf", marker = "python_full_version < '3.14' and extra == 'proxy-runtime'", specifier = ">=6.12.0,<7.0" }, { name = "pyroscope-io", marker = "sys_platform != 'win32' and extra == 'proxy'", specifier = ">=0.8.16,<1.0" }, { name = "python-dotenv", specifier = ">=1.0.0,<2.0" }, { name = "python-multipart", marker = "extra == 'proxy'", specifier = ">=0.0.27,<1.0" }, @@ -6059,14 +6063,14 @@ wheels = [ [[package]] name = "pypdf" -version = "6.10.2" +version = "6.13.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "typing-extensions", marker = "python_full_version < '3.11'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/7b/3f/9f2167401c2e94833ca3b69535bad89e533b5de75fefe4197a2c224baec2/pypdf-6.10.2.tar.gz", hash = "sha256:7d09ce108eff6bf67465d461b6ef352dcb8d84f7a91befc02f904455c6eea11d", size = 5315679, upload-time = "2026-04-15T16:37:36.978Z" } +sdist = { url = "https://files.pythonhosted.org/packages/15/d9/9d12fa0d9660d03320725ff686c961b645a4218940a82296e1272d9e1ff0/pypdf-6.13.1.tar.gz", hash = "sha256:4841d8a4c1589e5833915dc0c7ddfacff80a2e0bcbeb5d1e681fecaa1674b03a", size = 6477811, upload-time = "2026-06-08T11:01:49.344Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/0c/d6/1d5c60cc17bbdf37c1552d9c03862fc6d32c5836732a0415b2d637edc2d0/pypdf-6.10.2-py3-none-any.whl", hash = "sha256:aa53be9826655b51c96741e5d7983ca224d898ac0a77896e64636810517624aa", size = 336308, upload-time = "2026-04-15T16:37:34.851Z" }, + { url = "https://files.pythonhosted.org/packages/fe/dd/8f03e0a5788a5d1feb4550617c3e6db5e9099eaee248a3e482ddaeacbbb0/pypdf-6.13.1-py3-none-any.whl", hash = "sha256:e555e4ce3f561ef069307622f1374136ba964ca6ca24f24158701decaf83ed9b", size = 346259, upload-time = "2026-06-08T11:01:47.741Z" }, ] [[package]] @@ -7582,19 +7586,19 @@ wheels = [ [[package]] name = "tornado" -version = "6.5.5" +version = "6.5.7" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/f8/f1/3173dfa4a18db4a9b03e5d55325559dab51ee653763bb8745a75af491286/tornado-6.5.5.tar.gz", hash = "sha256:192b8f3ea91bd7f1f50c06955416ed76c6b72f96779b962f07f911b91e8d30e9", size = 516006, upload-time = "2026-03-10T21:31:02.067Z" } +sdist = { url = "https://files.pythonhosted.org/packages/64/24/95ec527ad67b76d59299e5465b3935d05e4294b7e0290a3924b7487df30b/tornado-6.5.7.tar.gz", hash = "sha256:66c513a76cda70d53907bc27cf1447557699c2e95aa48ba27a442ff61c3ddfc2", size = 519252, upload-time = "2026-06-08T17:34:51.232Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/59/8c/77f5097695f4dd8255ecbd08b2a1ed8ba8b953d337804dd7080f199e12bf/tornado-6.5.5-cp39-abi3-macosx_10_9_universal2.whl", hash = "sha256:487dc9cc380e29f58c7ab88f9e27cdeef04b2140862e5076a66fb6bb68bb1bfa", size = 445983, upload-time = "2026-03-10T21:30:44.28Z" }, - { url = "https://files.pythonhosted.org/packages/ab/5e/7625b76cd10f98f1516c36ce0346de62061156352353ef2da44e5c21523c/tornado-6.5.5-cp39-abi3-macosx_10_9_x86_64.whl", hash = "sha256:65a7f1d46d4bb41df1ac99f5fcb685fb25c7e61613742d5108b010975a9a6521", size = 444246, upload-time = "2026-03-10T21:30:46.571Z" }, - { url = "https://files.pythonhosted.org/packages/b2/04/7b5705d5b3c0fab088f434f9c83edac1573830ca49ccf29fb83bf7178eec/tornado-6.5.5-cp39-abi3-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:e74c92e8e65086b338fd56333fb9a68b9f6f2fe7ad532645a290a464bcf46be5", size = 447229, upload-time = "2026-03-10T21:30:48.273Z" }, - { url = "https://files.pythonhosted.org/packages/34/01/74e034a30ef59afb4097ef8659515e96a39d910b712a89af76f5e4e1f93c/tornado-6.5.5-cp39-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:435319e9e340276428bbdb4e7fa732c2d399386d1de5686cb331ec8eee754f07", size = 448192, upload-time = "2026-03-10T21:30:51.22Z" }, - { url = "https://files.pythonhosted.org/packages/be/00/fe9e02c5a96429fce1a1d15a517f5d8444f9c412e0bb9eadfbe3b0fc55bf/tornado-6.5.5-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:3f54aa540bdbfee7b9eb268ead60e7d199de5021facd276819c193c0fb28ea4e", size = 448039, upload-time = "2026-03-10T21:30:53.52Z" }, - { url = "https://files.pythonhosted.org/packages/82/9e/656ee4cec0398b1d18d0f1eb6372c41c6b889722641d84948351ae19556d/tornado-6.5.5-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:36abed1754faeb80fbd6e64db2758091e1320f6bba74a4cf8c09cd18ccce8aca", size = 447445, upload-time = "2026-03-10T21:30:55.541Z" }, - { url = "https://files.pythonhosted.org/packages/5a/76/4921c00511f88af86a33de770d64141170f1cfd9c00311aea689949e274e/tornado-6.5.5-cp39-abi3-win32.whl", hash = "sha256:dd3eafaaeec1c7f2f8fdcd5f964e8907ad788fe8a5a32c4426fbbdda621223b7", size = 448582, upload-time = "2026-03-10T21:30:57.142Z" }, - { url = "https://files.pythonhosted.org/packages/2c/23/f6c6112a04d28eed765e374435fb1a9198f73e1ec4b4024184f21faeb1ad/tornado-6.5.5-cp39-abi3-win_amd64.whl", hash = "sha256:6443a794ba961a9f619b1ae926a2e900ac20c34483eea67be4ed8f1e58d3ef7b", size = 448990, upload-time = "2026-03-10T21:30:58.857Z" }, - { url = "https://files.pythonhosted.org/packages/b7/c8/876602cbc96469911f0939f703453c1157b0c826ecb05bdd32e023397d4e/tornado-6.5.5-cp39-abi3-win_arm64.whl", hash = "sha256:2c9a876e094109333f888539ddb2de4361743e5d21eece20688e3e351e4990a6", size = 448016, upload-time = "2026-03-10T21:31:00.43Z" }, + { url = "https://files.pythonhosted.org/packages/02/dc/c7043cab6fed8ae159fc1923ce829ada35c4dbd797d408a43858ffaf9639/tornado-6.5.7-cp39-abi3-macosx_10_9_universal2.whl", hash = "sha256:148b2eb15c2c765a50796172c1e499649b35f30d2e3c3d3e15913cfa56bfb163", size = 448543, upload-time = "2026-06-08T17:34:38.052Z" }, + { url = "https://files.pythonhosted.org/packages/92/4f/090b1431e5a43df696feceffc268c5383cc079ecb5f08ce58f917109aafe/tornado-6.5.7-cp39-abi3-macosx_10_9_x86_64.whl", hash = "sha256:9da38de27f1da3b78a966f0dae12b5a1ea9afe72ca805d84ff06508272ddf100", size = 446707, upload-time = "2026-06-08T17:34:39.594Z" }, + { url = "https://files.pythonhosted.org/packages/37/d8/ef374952fd5da67d4463122c2b8e5a96536ec10b4b339254c6dcde81d01c/tornado-6.5.7-cp39-abi3-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:8d759e71906ee783f8867b93bf26a265743da4c1e2f4a018464c1ba019862972", size = 449774, upload-time = "2026-06-08T17:34:41.204Z" }, + { url = "https://files.pythonhosted.org/packages/35/37/d434c73f4c6e014b745b9b37085f34f40c022f007efff3d7fe65991899f3/tornado-6.5.7-cp39-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8a46347a18f23fb92b396beebe0fb78f61dda0cc302445202c16203d8a18848b", size = 450745, upload-time = "2026-06-08T17:34:42.531Z" }, + { url = "https://files.pythonhosted.org/packages/b6/2b/56b9aff361d7f1ab728a805ec7d7ea835f8807afa9f5cc690ea0e630efb9/tornado-6.5.7-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:7778b30bef919231265e91c69963ce0f49a1e9c07ac900bbe75b19ce2575ba92", size = 450578, upload-time = "2026-06-08T17:34:43.787Z" }, + { url = "https://files.pythonhosted.org/packages/02/30/a7444fb23aa76860a14198fab96ac79f1866b0a6e19e26c4381b0938e50f/tornado-6.5.7-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:e726f0c75da7726eec023aa62751ff8878bd2737e34fbdd33b1ae5897d2200f5", size = 449985, upload-time = "2026-06-08T17:34:45.326Z" }, + { url = "https://files.pythonhosted.org/packages/5c/42/5f0e56c01e8d9d36f4e23f367b85ae6cae0c1ecddd5e6977d8388ad27488/tornado-6.5.7-cp39-abi3-win32.whl", hash = "sha256:f8de3bf12d3efdd0cbe7c8887868198f8a91415e3f29fcf258d9b8eb7b1d9ae4", size = 451047, upload-time = "2026-06-08T17:34:46.784Z" }, + { url = "https://files.pythonhosted.org/packages/c9/a4/b393076ffb21b469eec5b328a0534cf03a3b90bfc6b1f09507cdd075d938/tornado-6.5.7-cp39-abi3-win_amd64.whl", hash = "sha256:de942f843533a039ef9fa3d9c88c7cd8a7c94553fb5ad0154270989b3d99a2c4", size = 451485, upload-time = "2026-06-08T17:34:48.248Z" }, + { url = "https://files.pythonhosted.org/packages/71/2e/7b1c769803121b809112cf9a00681c472eae1d80e32d7ec0e0bd61d0d0e1/tornado-6.5.7-cp39-abi3-win_arm64.whl", hash = "sha256:ff934fce95643af5f11efdae618eaa73d469dc588641e5c8d19295a0c65c4796", size = 450506, upload-time = "2026-06-08T17:34:49.702Z" }, ] [[package]] From e5a3083c2e21bf789c43400968b87e43033845cb Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 12 Jun 2026 17:56:33 -0700 Subject: [PATCH 095/185] refactor(ui): remove unreachable /chat page (#30178) The /ui/chat route is not linked from anywhere: no sidebar entry, no redirect, and no backend reference. It is only reachable by typing the URL by hand. Delete the route (src/app/chat) and its components (src/components/chat), which nothing else imports, and drop the deleted files' entries from the eslint suppressions baseline. --- ui/litellm-dashboard/eslint-suppressions.json | 42 - ui/litellm-dashboard/src/app/chat/page.tsx | 27 - .../src/components/chat/ChatMessages.tsx | 590 ------- .../src/components/chat/ChatPage.tsx | 1512 ----------------- .../src/components/chat/ConversationList.tsx | 450 ----- .../src/components/chat/MCPAppsPanel.tsx | 726 -------- .../src/components/chat/MCPConnectPicker.tsx | 171 -- .../src/components/chat/MCPCredentialsTab.tsx | 166 -- .../src/components/chat/types.ts | 25 - .../src/components/chat/useChatHistory.ts | 213 --- 10 files changed, 3922 deletions(-) delete mode 100644 ui/litellm-dashboard/src/app/chat/page.tsx delete mode 100644 ui/litellm-dashboard/src/components/chat/ChatMessages.tsx delete mode 100644 ui/litellm-dashboard/src/components/chat/ChatPage.tsx delete mode 100644 ui/litellm-dashboard/src/components/chat/ConversationList.tsx delete mode 100644 ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx delete mode 100644 ui/litellm-dashboard/src/components/chat/MCPConnectPicker.tsx delete mode 100644 ui/litellm-dashboard/src/components/chat/MCPCredentialsTab.tsx delete mode 100644 ui/litellm-dashboard/src/components/chat/types.ts delete mode 100644 ui/litellm-dashboard/src/components/chat/useChatHistory.ts diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 6dc9be0fc90..369efb7e338 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -860,48 +860,6 @@ "count": 1 } }, - "src/components/chat/ChatMessages.tsx": { - "react-hooks/refs": { - "count": 1 - } - }, - "src/components/chat/ChatPage.tsx": { - "max-params": { - "count": 2 - }, - "react-hooks/set-state-in-effect": { - "count": 2 - }, - "unused-imports/no-unused-imports": { - "count": 1 - } - }, - "src/components/chat/ConversationList.tsx": { - "react-hooks/set-state-in-effect": { - "count": 1 - } - }, - "src/components/chat/MCPAppsPanel.tsx": { - "max-nested-callbacks": { - "count": 2 - }, - "react-hooks/set-state-in-effect": { - "count": 2 - } - }, - "src/components/chat/MCPCredentialsTab.tsx": { - "no-restricted-imports": { - "count": 1 - }, - "react-hooks/set-state-in-effect": { - "count": 1 - } - }, - "src/components/chat/useChatHistory.ts": { - "react-hooks/set-state-in-effect": { - "count": 3 - } - }, "src/components/claude_code_plugins.tsx": { "no-restricted-imports": { "count": 1 diff --git a/ui/litellm-dashboard/src/app/chat/page.tsx b/ui/litellm-dashboard/src/app/chat/page.tsx deleted file mode 100644 index 5046f162877..00000000000 --- a/ui/litellm-dashboard/src/app/chat/page.tsx +++ /dev/null @@ -1,27 +0,0 @@ -"use client"; - -import { Suspense } from "react"; -import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; -import ChatPage from "@/components/chat/ChatPage"; - -// ChatPage uses useSearchParams() which requires a Suspense boundary for static export. -const ChatPageContent = () => { - const { accessToken, userRole, userId, userEmail } = useAuthorized(); - - return ( - - ); -}; - -const ChatPageRoute = () => ( - - - -); - -export default ChatPageRoute; diff --git a/ui/litellm-dashboard/src/components/chat/ChatMessages.tsx b/ui/litellm-dashboard/src/components/chat/ChatMessages.tsx deleted file mode 100644 index 53877be1737..00000000000 --- a/ui/litellm-dashboard/src/components/chat/ChatMessages.tsx +++ /dev/null @@ -1,590 +0,0 @@ -"use client"; - -import { ToolOutlined, CopyOutlined, CheckOutlined, EditOutlined } from "@ant-design/icons"; -import { Collapse, Tooltip } from "antd"; -import React, { useEffect, useRef, useState } from "react"; -import ReactMarkdown from "react-markdown"; -import remarkGfm from "remark-gfm"; -import { Prism as SyntaxHighlighter } from "react-syntax-highlighter"; -import { coy } from "react-syntax-highlighter/dist/esm/styles/prism"; -import ReasoningContent from "@/components/chat_ui/ReasoningContent"; -import MCPEventsDisplay from "@/components/chat_ui/MCPEventsDisplay"; -import { ChatMessage } from "./types"; - -const { Panel } = Collapse; - -// Keys whose values must be redacted in tool args display -const REDACTED_KEY_PATTERNS = /token|key|secret|password|auth/i; - -function redactSensitiveValues(obj: Record): Record { - const result: Record = {}; - for (const [k, v] of Object.entries(obj)) { - if (REDACTED_KEY_PATTERNS.test(k)) { - result[k] = "[redacted]"; - } else if (Array.isArray(v)) { - result[k] = v.map((item) => - item !== null && typeof item === "object" && !Array.isArray(item) - ? redactSensitiveValues(item as Record) - : item, - ); - } else if (v !== null && typeof v === "object") { - result[k] = redactSensitiveValues(v as Record); - } else { - result[k] = v; - } - } - return result; -} - -function formatTimestamp(ts: number): string { - const d = new Date(ts); - const hh = String(d.getHours()).padStart(2, "0"); - const mm = String(d.getMinutes()).padStart(2, "0"); - return `${hh}:${mm}`; -} - -// Shared markdown code renderer matching ReasoningContent style. -// react-markdown v9 removed the `inline` prop; detect fenced blocks via language className. -function MarkdownCodeRenderer({ - node, - className, - children, - ...props -}: React.ComponentPropsWithoutRef<"code"> & { node?: unknown }) { - const match = /language-(\w+)/.exec(className || ""); - return match ? ( - } - language={match[1]} - PreTag="div" - className="rounded-md my-2" - {...(props as Record)} - > - {String(children).replace(/\n$/, "")} - - ) : ( - - {children} - - ); -} - -// ------- Sub-components ------- - -interface UserBubbleProps { - message: ChatMessage; - onEdit?: (messageId: string, newContent: string) => void; - isStreaming?: boolean; -} - -function UserBubble({ message, onEdit, isStreaming }: UserBubbleProps) { - const [hovered, setHovered] = useState(false); - const [editing, setEditing] = useState(false); - const [editValue, setEditValue] = useState(message.content); - const textareaRef = useRef(null); - - useEffect(() => { - if (editing && textareaRef.current) { - textareaRef.current.focus(); - textareaRef.current.selectionStart = textareaRef.current.value.length; - } - }, [editing]); - - // Auto-resize textarea - useEffect(() => { - const ta = textareaRef.current; - if (!ta) return; - ta.style.height = "auto"; - ta.style.height = `${ta.scrollHeight}px`; - }, [editValue, editing]); - - const handleSave = () => { - const trimmed = editValue.trim(); - if (trimmed && trimmed !== message.content && onEdit) { - onEdit(message.id, trimmed); - } - setEditing(false); - }; - - const handleKeyDown = (e: React.KeyboardEvent) => { - if (e.key === "Enter" && !e.shiftKey) { - e.preventDefault(); - handleSave(); - } - if (e.key === "Escape") { - setEditValue(message.content); - setEditing(false); - } - }; - - if (editing) { - return ( -
-
-